ggml : add ggml_top_k

[ggerganov]

Add a dedicated top-k op so that it can be more efficiently optimized by backend implementations. The old implementation is renamed to ggml_argsort_top_k.

TODO:

• Allow unsorted output (ggml : add ggml_top_k #17365 (comment)) (see: c63ecde)
• Do not rely on op_params
• CUDA (will be added in sampling : add support for backend sampling #17004)
• Metal

Next PRs:

• Vulkan
• etc.

[am17an]

Does this operator expect the top-K elements to be sorted?

[ORippler]

Does this operator expect the top-K elements to be sorted?

I feel it should not be sorted, as algorithmically we are performing a selection, and depending on the algorithm the outcome of this selection is unordered:

https://leimao.github.io/blog/CPU-TopK-Algorithm/
https://nvidia.github.io/cccl/cub/api/structcub_1_1DeviceTopK.html#overview

Should one wish to sort, one could easily do GGML_OP_TOP_K -> GGML_OP_ARGSORT

[ggerganov]

Does this operator expect the top-K elements to be sorted?

In principle it does not have to expect the elements to be sorted. However the current implementation sorts them in descending order in order to be able to verify correctness with test-backend-ops. If we allow arbitrary order I am not sure how we would verify correctness.

[ORippler]

If we allow arbitrary order I am not sure how we would verify correctness.

By treating them as sets rather than lists? We could use std::unordered_set for this

[am17an]

If we allow arbitrary order I am not sure how we would verify correctness.

By treating them as sets rather than lists? We could use std::unordered_set for this

Currently, test-backend-ops relies on NMSE of outputs rather than cardinality checks, but I guess that can be changed.

[slaren]

It would be ok to add an overrideable error function to test_case. Leave NMSE as the default and override it in test_top_k to compare as a set.

[jeffbolznv]

What are common tensor shapes and values of k we should optimize for?

Does this operation support non-contiguous rows?

[ggerganov]

@jeffbolznv This will be used in #17004 to do top-k sampling efficiently on the GPU. The typical shapes are:

• large src[0]->ne[0] (i.e. up to vocab size)
• small k (usually 1, 10, 40)

Support for non-contiguous rows is not necessary for now - will add asserts for that.

[jeffbolznv]

OK, understood. When you get a chance, please rebase, I'll implement something based on #17313.

[CISC]

I guess these are temporary until all backends support are in place? Add a TODO?

[ggerganov]

Not 100% sure yet - keeping the expert order deterministic might be necessary. And using ggml_top_k here would likely not make a big difference performance wise since the arrays are very small.

[CISC]

Completely unnecessary for the expert group selection at least.

[am17an]

Please don't change it anywhere else as it will also break fusion in all backends for topk-moe

[ggerganov]

since we store ascending int numbers in our array,

The values are also shuffled:

llama.cpp/tests/test-backend-ops.cpp

Lines 5034 to 5042 in c63ecde

	// initialize with unique values to avoid ties
	for (int64_t r = 0; r < ggml_nrows(t); r++) {
	std::vector<float> data(t->ne[0]);
	for (int i = 0; i < t->ne[0]; i++) {
	data[i] = i;
	}
	std::shuffle(data.begin(), data.end(), rng);
	ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float));
	}

So top 1 could be any number.

[jeffbolznv]

Vulkan support is ready in #17418.