vulkan: support TOP_K with large K

[jeffbolznv]

Overview

Similar to the N-ary search histogram, but multipass through global memory, and then a final pass writes the top K values.

Perf on RTX 5090:

  TOP_K(type=f32,ne=[2048,1,1,1],k=2048,ties=0):               73728 runs -    14.57 us/run -       16 kB/run -    1.05 GB/s
  TOP_K(type=f32,ne=[65000,1,1,1],k=2048,ties=0):              65536 runs -    16.58 us/run -      261 kB/run -   15.07 GB/s
  TOP_K(type=f32,ne=[200000,1,1,1],k=2048,ties=0):             65536 runs -    17.33 us/run -      789 kB/run -   43.43 GB/s
  TOP_K(type=f32,ne=[2048,16,1,1],k=2048,ties=0):              73728 runs -    14.67 us/run -      256 kB/run -   16.64 GB/s
  TOP_K(type=f32,ne=[65000,16,1,1],k=2048,ties=0):             32032 runs -    33.07 us/run -     4190 kB/run -  120.84 GB/s
  TOP_K(type=f32,ne=[200000,16,1,1],k=2048,ties=0):                    15948 runs -    72.31 us/run -    12628 kB/run -  166.54 GB/s

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, algorithm was my idea, but codex implemented it.

[jeffbolznv]

Linking to #28032 and #28005.

[0cc4m]

Interestingly, seems our implementations have different strengths. These are test-backend-ops perf results of the ops extracted from Qwen 3.8 Flash Next:

Results (us/run, lower is better)

nv-spark (GB10 Blackwell)

impl	decode [2051,1]	prefill [2051,512]
master (CPU)	267.48	13804.27
argsort (Vulkan)	49.92	11177.47
0cc4m (Vulkan)	34.09	423.43
jeffbolznv (Vulkan)	22.81	1202.86
jeffbolznv-tuned (Vulkan)	21.54	1033.66

amd-395 (Radeon 8060S RDNA3.5)

impl	decode [2051,1]	prefill [2051,512]
master (CPU)	516.75	13645.21
argsort (Vulkan)	65.42	14094.80
0cc4m (Vulkan)	36.20	374.25
jeffbolznv (Vulkan)	15.52	700.44
jeffbolznv-tuned (Vulkan)	15.71	526.51

llama-bench t/s, Qwen3.8 Flash Next IQ4_XS, -ngl 99 -fa 1

nv-spark (GB10)

impl	pp512 d0	tg128 d0	pp512 d16k	tg128 d16k	pp512 d32k	tg128 d32k
master	408.88	23.69	321.44	13.73	293.67	12.49
argsort	688.87	26.15	421.89	21.76	278.82	18.94
0cc4m	690.03	26.16	587.10	22.44	528.86	19.88
0cc4m no-fuse	694.60	26.14	574.67	22.06	507.57	18.93
jeffbolznv	697.93	26.21	566.70	21.99	494.25	19.25
jeffbolznv-tuned	693.53	26.16	568.81	22.07	496.02	19.35

amd-395 (Radeon 8060S)

impl	pp512 d0	tg128 d0	pp512 d16k	tg128 d16k	pp512 d32k	tg128 d32k
master	198.93	22.18	134.07	15.35	119.66	13.09
argsort	309.32	24.01	179.99	17.96	128.45	14.80
0cc4m	302.01	23.60	217.24	18.22	186.66	15.16
0cc4m no-fuse	304.10	24.11	225.11	18.03	177.35	14.73
jeffbolznv	300.94	23.91	213.17	18.07	168.74	14.83
jeffbolznv-tuned	304.12	24.11	225.95	18.25	171.23	14.93

[jeffbolznv]

In your tests, is 2051 the value of K or the value of ne0?

[0cc4m]

That's K, yes.

decode: TOP_K(ne=[2051,1,1,1], sources=f32[16384,1,1,1])
prefill: TOP_K(ne=[2051,512,1,1], sources=f32[16384,512,1,1])

[jeffbolznv]

I added some tuning to do fewer workgroups. This recovers the perf on my system but not sure it'll be as fast with fewer SMs.

[0cc4m]

I updated my comment, added a tuned test and actual model tests. My branch is still faster in prefill, and also slightly faster in decode, mostly due to the fusion though.

[jeffbolznv]

OK, closing in favor of yours.