Commit ff5888f99 for llama.cpp
commit ff5888f999f24e2f9eaf6537b4eb6b85fd559188
Author: gianni-cor <gianfrancocordella@gmail.com>
Date: Thu Oct 8 12:14:35 2026 +0200
vulkan : fix TOP_K for +inf/NaN inputs and k = 1 on negative values (#30107)
The bucket search in topk_nary_search.comp started from the range
[0, 0xFF800000), which ends just below the ordered-uint mapping of +inf,
so +inf and NaN were never counted. A workgroup block with fewer than k
countable values left the ballot empty and the shader read uninitialized
shared state (hang/device lost on NVIDIA, wrong indices on AMD), and a few
+inf in a block were selected without being counted, dropping real top
values.
Map NaN to -inf on input, start from [0, 0xFFFFFFFF) so every value is
counted, and clamp the top bucket's end (2^32) instead of wrapping to 0.
The k = 1 path compared float bits as signed integers, which orders
negative values backwards; compare floats instead.
Add test_top_k_inf to test-backend-ops: negative values, fewer than k
+inf and many -inf, for k = 1, 10, 40.
Assisted-by: Claude Opus 5.5
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp b/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp
index 0b757f38e..3f6296bda 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp
@@ -60,7 +60,13 @@ void topk(const uint row) {
if (gl_GlobalInvocationID.x < p.ncols_input) {
if (p.first_pass != 0) {
const uint row_offset = row * p.ncols_input;
- dst_row[tid] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(data_a[row_offset + gl_GlobalInvocationID.x]));
+ // NaN ranks lowest, like -inf, so that every value has a place in
+ // the ordering the search below counts
+ float a = float(data_a[row_offset + gl_GlobalInvocationID.x]);
+ if (isnan(a)) {
+ a = uintBitsToFloat(0xFF800000);
+ }
+ dst_row[tid] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(a));
} else {
const uint row_offset = row * p.ncols_input;
dst_row[tid] = data_s[row_offset + gl_GlobalInvocationID.x];
@@ -76,8 +82,10 @@ void topk(const uint row) {
if (tid < s) {
ivec2 a = dst_row[tid];
ivec2 b = dst_row[tid + s];
+ // compare as floats: the bit patterns of negative values
+ // order the other way as integers
if (a.x >= p.orig_ncols ||
- b.x < p.orig_ncols && b.y > a.y) {
+ b.x < p.orig_ncols && intBitsToFloat(b.y) > intBitsToFloat(a.y)) {
dst_row[tid] = b;
}
}
@@ -95,9 +103,11 @@ void topk(const uint row) {
int shift = 32 - SUBGROUP_SIZE_LOG2;
uint mask = ((1 << SUBGROUP_SIZE_LOG2) - 1) << shift;
- // The current range.
+ // The current range, [range_min, range_max). It starts as every value
+ // (+inf maps to 0xFF800000 and NaN was replaced by -inf), so the
+ // buckets always hold at least limit values.
uint range_min = 0;
- uint range_max = 0xFF800000;
+ uint range_max = 0xFFFFFFFF;
// How many are above the current range, and how many we need to find.
uint total = 0;
uint limit = min(p.k, p.ncols_input - gl_WorkGroupID.x * BLOCK_SIZE);
@@ -138,8 +148,12 @@ void topk(const uint row) {
total = sh_total;
// Update the range, and break if we've found the K-th largest.
- range_max = range_min + ((min_idx + 1) << shift);
- range_min = range_min + (min_idx << shift);
+ // The end of the top bucket wraps past 2^32, clamp it instead.
+ range_min = range_min + (uint(min_idx) << shift);
+ range_max = range_min + (1u << shift);
+ if (range_max < range_min) {
+ range_max = 0xFFFFFFFF;
+ }
if (total == p.k) {
break;
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 709997680..c5b1cd935 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -7042,6 +7042,46 @@ struct test_top_k : public test_case {
}
};
+// top_k over rows like log-probabilities: distinct negative values, fewer
+// than k +inf (none for k = 1, so the expected indices are unique) and many
+// -inf (masked tokens)
+struct test_top_k_inf : public test_top_k {
+ test_top_k_inf(std::array<int64_t, 4> ne, int k)
+ : test_top_k(GGML_TYPE_F32, ne, k, false) {}
+
+ std::string vars() override {
+ return test_top_k::vars() + ",inf=1";
+ }
+
+ // compare only the output: the input holds infinities, which err() would
+ // read as indices
+ bool run_whole_graph() override { return true; }
+
+ void initialize_tensors(ggml_context * ctx) override {
+ std::random_device rd;
+ std::default_random_engine rng(rd());
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+ 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] = -1.0f - i;
+ }
+ std::shuffle(data.begin(), data.end(), rng);
+ const int n_pinf = k / 2;
+ for (int i = 0; i < t->ne[0]; i++) {
+ if (i < n_pinf) {
+ data[i] = INFINITY;
+ } else if (i % 3 == 0) {
+ data[i] = -INFINITY;
+ }
+ }
+ std::shuffle(data.begin(), data.end(), rng);
+ ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float));
+ }
+ }
+ }
+};
+
// qwen4exp QSA indexer top-k fusion: expand per-block scores to cells, add the f16 mask, top-k.
struct test_topk_qsa : public test_case {
const int64_t n_blocks;
@@ -11179,6 +11219,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {n, 2, 1, 3}, k, true));
}
}
+ for (int k : {1, 10, 40}) {
+ test_cases.emplace_back(new test_top_k_inf({4096, 2, 1, 1}, k));
+ test_cases.emplace_back(new test_top_k_inf({248320, 1, 1, 1}, k));
+ }
+
for (int i = 0; i < 20; ++i) {
for (int k : {1, 2, 3, 7, 15, 100, 500, 1023, 9999}) {
if (k <= 1<<i) {