]> git.djapps.eu Git - pkg/ggml/sources/ggml/commit
vulkan: fix top_k bug when there are ties in the input (llama/17659)
authorJeff Bolz <redacted>
Fri, 5 Dec 2025 21:03:19 +0000 (15:03 -0600)
committerGeorgi Gerganov <redacted>
Thu, 11 Dec 2025 13:32:56 +0000 (15:32 +0200)
commite95bb22ddd83eb0f65391fdbd6dead8927f31211
tree3c40ab6b9b9425f27c555ff3812d5c480ecfbc4e
parenta61f64a0175e8b83116e6c8f2678e0c6d5f290ab
vulkan: fix top_k bug when there are ties in the input (llama/17659)

* vulkan: Reduce temporary memory usage for TOP_K

- Compute row size for the temp buffer based on the output of the first pass.
- Update shader addressing math to use the output row size
- Pass the output row size as "ncols_output", what used to be "ncols_output" is now "k"

For the common case of K=40 and src0=(200000,1,1,1), this reduces the temporary buffer
from about 3.2MB to 500KB.

* vulkan: fix top_k bug when there are ties in the input

I noticed by inspection a bug in the vulkan top_k shader where if the least
value in the top_k appears multiple times we could end up writing those extra
copies out rather than some larger values (if the larger values are on higher
numbered threads).

I rewrote the test verification to handle this case, where the final index set
is not necessarily the same.

* Update tests/test-backend-ops.cpp

Co-authored-by: Georgi Gerganov <redacted>
---------

Co-authored-by: Georgi Gerganov <redacted>
src/ggml-vulkan/ggml-vulkan.cpp
src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp
tests/test-backend-ops.cpp