vulkan : fix TOP_K for +inf/NaN inputs and k = 1 on negative values

[gianni-cor]

Overview

Fixes two bugs in the Vulkan TOP_K shader (topk_nary_search.comp).

+inf / NaN inputs. The shader finds the k-th largest value of each workgroup's block by mapping the floats to ordered uints and counting them into buckets, narrowing the range each step. The initial range was [0, 0xFF800000), which ends just below the mapping of +inf, so +inf and NaN were never counted. When a block has fewer than k countable values, no bucket reaches the limit, the ballot is empty, and the shader continues with uninitialized shared state (sh_min_idx, sh_total): on NVIDIA the dispatch hangs until the driver resets the device, on AMD it returns wrong indices. With only a few +inf in a block, they still pass the final >= range_min selection without having been counted, so real top values are dropped.

k = 1 on negative values. The k = 1 fast path compared the float bit patterns as signed integers, which orders negative values backwards (-2.0 beats -1.0), so the result is wrong whenever a block's maximum is negative.

Fix.

• NaN is mapped to -inf when the input is read, so it ranks lowest.
• The initial range is [0, 0xFFFFFFFF), so every value, +inf included, falls in a bucket and the search always finds one.
• The end of the top bucket, range_min + (SUBGROUP_SIZE << shift) = 2^32, is clamped instead of wrapping to 0 (previously only reachable with finite values above ~2^113; +inf now lands in that bucket).
• The k = 1 path compares floats.

Additional information

How it shows up. In a downstream application built on llama.cpp b11018, on an RTX 5090 (NVIDIA Vulkan), loading an MTP model with parallel: 1 and then parallel: 2 in one process lost the device on the second load. The faulting submission was the TOP_K node of the MTP draft context's backend sampler, run while that context was being created, most likely its warmup decode on logits computed from uninitialized (NaN) data. With this fix the crash no longer reproduces.

Standalone repro (ggml_top_k on the Vulkan backend, 248320 columns, k = 10):

Row	RTX 5090, before	Radeon 8060S (RADV), before	After (both)
finite values	correct	correct	correct
100 × +inf	4 of the +inf indices + 6 finite	4 of the +inf indices + 6 finite	the first 10 +inf indices
all NaN	hangs	indices 0-9	indices 0-9
all +inf	hangs	indices 0-9	indices 0-9
248000 × NaN, 320 finite	hangs	NaN indices 0-9	top 10 of the finite values

Tests. New test_top_k_inf case in test-backend-ops: rows of distinct negative values with fewer than k +inf (none for k = 1, so the expected indices
are unique) and many -inf, for k = 1, 10, 40. On master 9881906, oe Radeon 8060S:

• before: 525/531 TOP_K cases pass; the 6 new cases fail (k = 1 = 10/40 on the +inf);
• after: 531/531.

NaN is not in the test because the CPU reference (std::partial_sort with >) has no defined result for it. CUDA is not affected: it has its own implementation and passes the new cases.

this ports tetherto#350 to master

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. Claude (Opus 5.5) debugged the device loss, found the shader bugs, wrote the fix and the test, and drafted this description; I reviewed the change, the results and the text.

[ggml-gh-bot]

Hi @gianni-cor, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• Maintainers cannot push to this PR: Please enable Allow edits by maintainers. If this PR comes from an organization-owned fork, that option is not available on GitHub; please re-open the PR from a fork owned by your personal account.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.