[DSv4 Perf] Skip topk and router when not needed, 3.4% E2E TTFT improvement for Decode case

[yewentao256]

Purpose

Skip topk and router when not needed

Originally:

Input
-> K compressor (and write K cache)
-> Query proj wq_b
-> Query RoPE + quantization
->  calculation for Query-K logits
->  topk logits
-> return index

Now

-> K compressor (and write K cache)
-> directly choose all candidates if candidates num <= topk num
-> return index

Part of #45861

Test

vllm serve deepseek-ai/DeepSeek-V4-Flash --tensor-parallel-size 4 --enable-expert-parallel --attention-backend FLASHMLA_SPARSE_DSV4 --attention-config '{"use_fp4_indexer_cache":true}' --kv-cache-dtype fp8 --tokenizer-mode deepseek_v4 --all2all-backend allgather_reducescatter --port 8003

Acc

lm_eval --model  local-completions --model_args "base_url=http://127.0.0.1:8003/v1/completions,model=deepseek-ai/DeepSeek-V4-Flash,num_concurrent=1024" --tasks gsm8k
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9507|±  |0.0060|
|     |       |strict-match    |     5|exact_match|↑  |0.9515|±  |0.0059|

Perf

Max concurrency 1

vllm bench serve   --model deepseek-ai/DeepSeek-V4-Flash   --dataset-name random   --host 127.0.0.1   --port 8003   --random-input-len 2   --random-output-len 2048   --request-rate inf   --max-concurrency 1   --num-prompts 4   --num-warmups 2   --ignore-eos   --temperature 0   --seed 0

# now
============ Serving Benchmark Result ============
Successful requests:                     4         
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  45.81     
Total input tokens:                      8         
Total generated tokens:                  8192      
Request throughput (req/s):              0.09      
Output token throughput (tok/s):         178.84    
Peak output token throughput (tok/s):    182.00    
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          179.01    
---------------Time to First Token----------------
Mean TTFT (ms):                          40.96     
Median TTFT (ms):                        41.30     
P99 TTFT (ms):                           43.14     
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          5.57      
Median TPOT (ms):                        5.57      
P99 TPOT (ms):                           5.58      
---------------Inter-token Latency----------------
Mean ITL (ms):                           5.57      
Median ITL (ms):                         5.55      
P99 ITL (ms):                            5.79      
==================================================

# main
============ Serving Benchmark Result ============
Successful requests:                     4         
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  45.61     
Total input tokens:                      8         
Total generated tokens:                  8192      
Request throughput (req/s):              0.09      
Output token throughput (tok/s):         179.61    
Peak output token throughput (tok/s):    182.00    
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          179.78    
---------------Time to First Token----------------
Mean TTFT (ms):                          42.38     
Median TTFT (ms):                        42.18     
P99 TTFT (ms):                           43.64     
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          5.55      
Median TPOT (ms):                        5.55      
P99 TPOT (ms):                           5.56      
---------------Inter-token Latency----------------
Mean ITL (ms):                           5.55      
Median ITL (ms):                         5.52      
P99 ITL (ms):                            5.77      
==================================================

heavy decode

# now
============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  20.27     
Total input tokens:                      256       
Total generated tokens:                  131072    
Request throughput (req/s):              6.32      
Output token throughput (tok/s):         6467.01   
Peak output token throughput (tok/s):    6656.00   
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          6479.64   
---------------Time to First Token----------------
Mean TTFT (ms):                          222.31    
Median TTFT (ms):                        231.81    
P99 TTFT (ms):                           285.95    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          19.57     
Median TPOT (ms):                        19.56     
P99 TPOT (ms):                           19.63     
---------------Inter-token Latency----------------
Mean ITL (ms):                           19.57     
Median ITL (ms):                         19.44     
P99 ITL (ms):                            21.70     
==================================================

# main
============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  20.27     
Total input tokens:                      256       
Total generated tokens:                  131072    
Request throughput (req/s):              6.32      
Output token throughput (tok/s):         6466.61   
Peak output token throughput (tok/s):    6656.00   
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          6479.24   
---------------Time to First Token----------------
Mean TTFT (ms):                          229.54    
Median TTFT (ms):                        234.82    
P99 TTFT (ms):                           292.12    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          19.57     
Median TPOT (ms):                        19.56     
P99 TPOT (ms):                           19.63     
---------------Inter-token Latency----------------
Mean ITL (ms):                           19.57     
Median ITL (ms):                         19.53     
P99 ITL (ms):                            21.46     
==================================================

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[mgoin]

TTFT is better but the benchmark takes slightly longer/decode is slower. Would be best to isolate this concern

[yewentao256]

@mgoin @MatthewBonanni
Added benchmark for more prompts (decode heavy case) in PR description.
The small regression is just from fluctuations

[MatthewBonanni]

Updated benchmark results look good to me