[Attention][Perf][Kernel] Improve topKperRow for large context decode path - DeepSeek-V3.2 sparse attention

[LopezCastroRoberto]

Summary

This PR integrates FlashInfer's radix-based top-k kernel as an alternative implementation for the large context top-k operation in the sparse attention indexer, specifically for DeepSeek-V3.2 models.

Kernel adapted from: flashinfer-ai/flashinfer#2215

Microbenchmark study

E2E results

Example on NVIDIA B200:

vllm serve nvidia/DeepSeek-V3.2-NVFP4 -tp 4
vllm bench serve --backend vllm --model nvidia/DeepSeek-V3.2-NVFP4 --input-len 128000 --output-len 4096 --num-prompts 1

MAIN:

============ Serving Benchmark Result ============
Successful requests:                     1         
Failed requests:                         0         
Benchmark duration (s):                  59.15     
Total input tokens:                      128000    
Total generated tokens:                  4096      
Request throughput (req/s):              0.02      
Output token throughput (tok/s):         69.24     
Peak output token throughput (tok/s):    71.00     
Peak concurrent requests:                1.00      
Total token throughput (tok/s):          2233.14   
---------------Time to First Token----------------
Mean TTFT (ms):                          717.80    
Median TTFT (ms):                        717.80    
P99 TTFT (ms):                           717.80    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          14.27     
Median TPOT (ms):                        14.27     
P99 TPOT (ms):                           14.27     
---------------Inter-token Latency----------------
Mean ITL (ms):                           14.27     
Median ITL (ms):                         14.27     
P99 ITL (ms):                            14.57     
==================================================

PR:

============ Serving Benchmark Result ============
Successful requests:                     1         
Failed requests:                         0         
Benchmark duration (s):                  53.62     
Total input tokens:                      128000    
Total generated tokens:                  4096      
Request throughput (req/s):              0.02      
Output token throughput (tok/s):         76.39     
Peak output token throughput (tok/s):    80.00     
Peak concurrent requests:                1.00      
Total token throughput (tok/s):          2463.46   
---------------Time to First Token----------------
Mean TTFT (ms):                          732.15    
Median TTFT (ms):                        732.15    
P99 TTFT (ms):                           732.15    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          12.92     
Median TPOT (ms):                        12.92     
P99 TPOT (ms):                           12.92     
---------------Inter-token Latency----------------
Mean ITL (ms):                           12.92     
Median ITL (ms):                         12.72     
P99 ITL (ms):                            13.81     
==================================================

In this example, the current PR improves MAIN throughput by ~10%

Here is a more general analysis across different sequence lengths on NVIDIA B300:

vllm bench serve --backend vllm --model nvidia/DeepSeek-V3.2-NVFP4 --input-len seq_len --output-len 4096 --num-prompts 1

Accuracy

python tests/evals/gsm8k/gsm8k_eval.py

MAIN:

Results:
Accuracy: 0.926
Invalid responses: 0.000
Total latency: 54.086 s
Questions per second: 24.387
Total output tokens: 121416
Output tokens per second: 2244.889

PR:

Results:
Accuracy: 0.929
Invalid responses: 0.000
Total latency: 52.035 s
Questions per second: 25.348
Total output tokens: 121881
Output tokens per second: 2342.299

[gemini-code-assist]

Code Review

This pull request replaces the custom large_context_topk kernel with flashinfer.top_k_ragged_transform for handling top-k operations in the large context decode path. The changes primarily involve updating sparse_attn_indexer.py to use the FlashInfer function and passing a new offsets_buffer. Corresponding changes are made for API compatibility in the ROCm path. The tests are also updated to validate the new implementation. My review found a critical issue in the test file where a new test function shadows an existing one due to having the same name, and also misuses a pytest parameter. I've provided a suggestion to fix this.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LopezCastroRoberto.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgoin]

Interesting that the gsm8k eval is ~15% faster even though it is only ~2k context length and has especially short prefills (due to high prefix cache hit rate)

[LopezCastroRoberto]

Interesting that the gsm8k eval is ~15% faster even though it is only ~2k context length and has especially short prefills (due to high prefix cache hit rate)

@mgoin Yeah, there isn’t actually a remarkable speedup on GSM8K. My goal wasn’t to demonstrate performance improvements on this test, but simply to verify that accuracy was preserved. It looks like I probably just copied one of the later runs (second/third/fourth) from several consecutive executions I did for this test, rather than the initial run.

e.g.,
python tests/evals/gsm8k/gsm8k_eval.py

First execution:

Results:
Accuracy: 0.929
Invalid responses: 0.000
Total latency: 52.035 s
Questions per second: 25.348
Total output tokens: 121881
Output tokens per second: 2342.299

Second execution:

Accuracy: 0.929
Invalid responses: 0.000
Total latency: 46.138 s
Questions per second: 28.588
Total output tokens: 122155
Output tokens per second: 2647.619

I agree this can be confusing. I will update it in the PR description.

[LopezCastroRoberto]

However, there may be a minor improvement for 2k contexts due to the updated heuristic introduced in this PR. Specifically, large_context_topk (integrated in a previous PR) is selected instead of top_k_per_row_decode. This choice is based on the microbenchmark results described above.

[LopezCastroRoberto]

This could be fairly expensive, any way to avoid? I noticed it isn't included in the benchmarking code. Maybe the kernel itself could zero only the required sections just for RadixRowState?
cc: @mgoin

I’ve fused the topk_workspace.zero_() op directly into the kernel. At the kernel level, there’s no remarkable overhead from this change. I also added a few additional test points to further validate that the heuristic behaves as expected.

[LucasWilkinson]

do we know how much perf the cudagraph specialization brings e2e? it adds quite a bit of complexity, just wondering if its worth it? how hard would it be optimize the topk for the shorter contexts?