vulkan: support TOP_K when k exceeds the workgroup limit

[Anothers]

Problem

On Vulkan, TOP_K falls back to the CPU once k goes past 1024. Qwen3.8-Flash-Next
(qwen4exp) hits that on 12 attention layers for every token it decodes, so past about 1K of
context the model is doing 12 round trips per token. It is a real problem for me — this is
my daily driver and every session runs past that point.

The limit is in supports_op:

uint32_t min_pipeline = (uint32_t)log2f(float(op->ne[0])) + 1;
if (min_pipeline >= num_topk_pipelines || !device->pipeline_topk_f32[min_pipeline]) {
    return false;
}

op is the destination, so op->ne[0] is k. ggml_vk_topk reduces a row one workgroup at
a time and the last pass has to hold all k candidates in a single workgroup, which caps at
1024 invocations. qwen4exp's QSA indexer asks for
min(n_kv, indexer_top_k + compress_ratio - 1) = min(n_kv, 2051), so below ~1K of context
k is the window and fits, and above it k pins at 2051 and the scheduler moves all twelve
nodes to the CPU.

It shows up as unsupported rather than slow:

$ test-backend-ops test -o TOP_K
TOP_K(type=f32,ne=[54822,33,1,1],k=2051,ties=0): not supported [Vulkan0]

Change

The comment already in supports_op suggested the way out:

// We could potentially support larger, using argsort to sort the
// whole thing. Not clear if this is needed.

So: when k is larger than one workgroup can hold, sort the whole row with the existing
argsort_large pipelines and keep the leading k. That is a valid answer because
ggml_top_k does not order its output — the CPU reference swaps the first two entries to
make the point:

std::partial_sort(tmp, tmp + top_k, tmp + ne00, cmp_top_k{src_data});
std::copy(tmp, tmp + top_k, dst_data);
// emphasize that the order is not important
if (top_k > 1) { std::swap(dst_data[0], dst_data[1]); }

The argsort dispatch was only reachable through ggml_vk_argsort, which took its buffers
from the tensors, so it is split into ggml_vk_argsort_rows taking explicit subbuffers.
ggml_vk_argsort becomes a wrapper and is otherwise unchanged. prealloc_x holds the sort
scratch followed by the permutation, and one copyBuffer with a region per row gathers the
leading k into the destination.

Why it is gated on row count

A full sort is more work than the selection it replaces, so it only pays for itself while
the row count is low enough that the round trip dominates. Measured here with
ncols 54822, k 2051:

rows	what submits it	GPU vs CPU fallback
3	an MTP decode step	11% faster
1024	a prefill ubatch	10% slower

The bound is 32 — above any decode batch, well below any prefill ubatch. My first version
had no bound and traded a 10% prefill regression for the decode win; it was only visible
because prefill was measured too. llama-bench -p 0 -n 128 measures generation alone and
hides it.

Numbers

gfx1151 (Radeon 8060S, RADV, Mesa 26.1.7), Qwen3.8-Flash-Next UD-Q4_K_XL,
llama-bench -fa 1 -n 128 -d .... Same build both columns; "before" only has the
supports_op relaxation reverted.

depth	before	after
0	25.72	25.78
1024	23.63	25.21	+6.7%
4096	21.70	24.31	+12.0%
16384	18.72	21.61	+15.4%
32768	16.53	19.04	+15.2%
65536	12.58	14.01	+11.4%

The gain starts where k crosses the limit and not before: at depth 0 the window is still
narrower than k, so the node never takes the new path.

On the server, with the MTP draft head and a 52K prompt, five runs each: decode
21.46 -> 23.39 (+9.0%), prefill 192.0 -> 190.0 (unchanged).

Note this is a slope here, not the cliff #27856 reports on HIP. This is an APU, so the
fallback pays for a synchronisation rather than a transfer. #27466 fixes the CUDA/HIP side
of the same limit; I did not find anything covering Vulkan.

Tests

test-backend-ops -o TOP_K passes, 523/523. The k = 9999 cases already in the suite now
run on Vulkan instead of being skipped as unsupported.

Added cases in the shape qwen4exp produces — k = 2051 over 54822 columns, with and
without ties, at 1, 3, 32 and 33 rows. 33 is over the bound, so the fallback side is covered
too.

AI disclosure

The code was written with AI assistance. I reviewed it, I understand it, and I will maintain
it.

[ggml-gh-bot]

Hi @Anothers, thanks for your contribution!

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

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

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

[0cc4m]

Already covered by #28032. Please check ongoing work before opening PRs. This PR is also basically identical to the already superseded #28005.