vulkan: large top-k support via argsort

[antoinezambelli]

Overview

Extend Vulkan TOP_K support to larger k values by falling back to the existing descending ARGSORT implementation when no optimized Top-K pipeline is available.

The existing optimized path remains unchanged and is still used whenever it supports the requested k. For larger values, each input row is sorted in descending order and the first k indices are copied into the output tensor.

Qwen3.8-Flash-Next configures its attention indexer with top_k=2048. The current qwen4exp graph also includes the compression-block tail, producing k=2051, which exceeds the existing Vulkan implementation’s supported range.

This enables that workload without adding a model-specific path, changing the model graph, or introducing another shader.

Additional information

Implementation

• Factor the existing argsort configuration, capability check, and dispatch logic into reusable helpers.

• Continue using the specialized Top-K pipelines whenever available.

• Otherwise:

  • sort each input row in descending order with Vulkan argsort;
  • copy the first k indices from each sorted row into the output tensor.

• Preserve the existing contiguity and Vulkan argsort capability requirements.

• Add coverage for large k and multi-row inputs.

Validation

Focused Vulkan cases

The parent revision reported all four cases as unsupported:

TOP_K ne=[1024,1,1,1], k=1024: not supported
TOP_K ne=[2048,2,1,1], k=1024: not supported
TOP_K ne=[4096,1,1,1], k=2048: not supported
TOP_K ne=[8192,2,1,1], k=2051: not supported

With this change:

4/4 tests passed

Full backend suite

17278/17278 tests passed
Backend Vulkan0: OK

Vulkan validation layers

The focused large-Top-K cases completed without validation errors.

The only emitted validation warning was the existing SPIR-V WorkgroupSize best-practice warning, which was also reproduced on the unchanged optimized small-Top-K path.

Existing-path regression checks

The existing optimized case remained unchanged within measurement noise:

TOP_K ne=[8192,1,1,1], k=16
parent:    5.00–5.02 us/run
candidate: 5.00–5.02 us/run

Qwen3.6-27B prompt processing and generation throughput were also unchanged within normal run-to-run variation. A matched Wikitext-2 check produced identical results:

parent:    PPL = 6.0616 +/- 0.71603
candidate: PPL = 6.0616 +/- 0.71603

Application-level check

Qwen3.8-Flash-Next has a separate large PLE tensor loading issue on Vulkan iGPU systems. For this end-to-end check, both the parent and candidate builds were layered on the same selective lazy-loading implementation proposed in #27837.

That loader work was needed only to fit this model on the test system. It is independent of Vulkan Top-K and is not included in this PR. The binary difference between the two validation builds was verified to be identical to this PR's Top-K diff.

Qwen3.8-Flash-Next UD-Q4_K_XL was then exercised over two Vulkan RPC workers at 262,144 context with Q8 K/V caches. A five-run stateful tool-calling smoke test completed successfully with 5/5 validated results.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Codex-assisted development

[ggml-gh-bot]

Hi @antoinezambelli, 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.

---

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

[jeffbolznv]

Qwen3.8-Flash-Next exercises this boundary with larger Top-K operations.

Can you share more details on the real-world workload?

[antoinezambelli]

Qwen3.8-Flash-Next exercises this boundary with larger Top-K operations.

Can you share more details on the real-world workload?

At its core, I ran a forge scenario - a pretty basic 3 step agentic workflow. It has since progressed to more complex scenarios. Happy to pull more data if you're looking for something specific!

• Qwen3.8-Flash-Next UD-Q4_K_XL, served across two RADV gfx1151 Strix Halo systems using Vulkan RPC.
• Runtime used a 262,144-token context and Q8 K/V caches.

[0cc4m]

I'm not sure what you're trying to say there. The question was what happens in the model that is not yet supported, and the answer is it has TOP_K operations with k 2048, we currently only support up to 1024.

[antoinezambelli]

I'm not sure what you're trying to say there. The question was what happens in the model that is not yet supported, and the answer is it has TOP_K operations with k 2048, we currently only support up to 1024.

Oops, I misunderstood workload to mean end-to-end application. Sorry about that - PR overview updated as well.

You got it right though. The model-side gap is Qwen3.8-Flash-Next’s attention indexer requesting TOP_K with a nominal k=2048, beyond the existing Vulkan path.

The current qwen4exp graph also includes the compression-block tail, producing k=2051. The Forge/RPC details were only the end-to-end correctness check.

[jeffbolznv]

What are the input dimensions?

[jeffbolznv]

With this change I'm getting a devicelost in test-backend-ops.exe -o TOP_K(type=f32,ne=[32779,1,2,1],k=9999,ties=0)

[antoinezambelli]

Closing in favor of the purpose-built large Top-K implementations in #28032 and #28036 . Thanks to both of you for picking this up, and for the validation at larger input sizes - I missed that in my original test coverage.