Use dsv3 optimized routing fused_topk_deepseek instead of moe_fused_gate

[leejnau]

Motivation

flashinfer has an optimized routing kernel for DeepSeek V3: flashinfer-ai/flashinfer#2099
The API was renamed to fused_topk_deepseek here: flashinfer-ai/flashinfer#2181

Modifications

Replace the call to moe_fused_gate with fused_topk_deepseek.

Accuracy Tests

Server Command:

python3 -m sglang.launch_server --model-path nvidia/DeepSeek-R1-0528-FP4-V2 --tensor-parallel-size=8 --cuda-graph-max-bs 256 --max-running-requests 256 --mem-fraction-static 0.85 --ep-size 8 --enable-symm-mem --moe-runner-backend flashinfer_cutlass --quantization modelopt_fp4

Client Benchmark Command:

python3 benchmark/gsm8k/bench_sglang.py \
  --num-shots 8 \
  --num-questions 1316 \
  --parallel 1316

Before:

Accuracy: 0.954
Invalid: 0.000
Latency: 143.606 s
Output throughput: 1009.301 token/s

After:

Accuracy: 0.960
Invalid: 0.000
Latency: 70.387 s
Output throughput: 2060.780 token/s

Benchmarking and Profiling

python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 2048 --random-input 1024 --random-output 1024 --random-range-ratio 1 --max-concurrency 2048

Before:

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 2048      
Successful requests:                     2048      
Benchmark duration (s):                  319.04    
Total input tokens:                      2097152   
Total input text tokens:                 2097152   
Total input vision tokens:               0         
Total generated tokens:                  2097152   
Total generated tokens (retokenized):    2091930   
Request throughput (req/s):              6.42      
Input token throughput (tok/s):          6573.25   
Output token throughput (tok/s):         6573.25   
Peak output token throughput (tok/s):    9728.00   
Peak concurrent requests:                2048      
Total token throughput (tok/s):          13146.50  
Concurrency:                             1154.13   
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   179793.11 
Median E2E Latency (ms):                 179816.14 
---------------Time to First Token----------------
Mean TTFT (ms):                          142579.01 
Median TTFT (ms):                        142477.90 
P99 TTFT (ms):                           282920.44 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          36.38     
Median TPOT (ms):                        36.24     
P99 TPOT (ms):                           38.56     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           36.38     
Median ITL (ms):                         34.66     
P95 ITL (ms):                            36.35     
P99 ITL (ms):                            39.03     
Max ITL (ms):                            3984.88   
==================================================

After:

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 2048      
Successful requests:                     2048      
Benchmark duration (s):                  312.63    
Total input tokens:                      2097152   
Total input text tokens:                 2097152   
Total input vision tokens:               0         
Total generated tokens:                  2097152   
Total generated tokens (retokenized):    2090916   
Request throughput (req/s):              6.55      
Input token throughput (tok/s):          6708.17   
Output token throughput (tok/s):         6708.17   
Peak output token throughput (tok/s):    8449.00   
Peak concurrent requests:                2048      
Total token throughput (tok/s):          13416.34  
Concurrency:                             1153.51   
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   176083.48 
Median E2E Latency (ms):                 176110.58 
---------------Time to First Token----------------
Mean TTFT (ms):                          139678.72 
Median TTFT (ms):                        139542.14 
P99 TTFT (ms):                           277228.84 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          35.59     
Median TPOT (ms):                        35.50     
P99 TPOT (ms):                           37.81     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           35.59     
Median ITL (ms):                         33.88     
P95 ITL (ms):                            36.34     
P99 ITL (ms):                            38.77     
Max ITL (ms):                            3960.43   
==================================================

The kernel itself is approximately 4x faster.

Before:

After:

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[trevor-m]

Thanks! I left some comments.

Can you also add a unit test which instantiates the TopK class such that this new kernel is called and verify the output?

[trevor-m]

Thanks, looks good!

Can you run pre-commit to fix the linting? https://docs.sglang.io/developer_guide/contribution_guide.html#format-code-with-pre-commit

[leejnau]

@Fridge003 Is it ok to merge?

[Fridge003]

@leejnau Is this PR depending on Flashinfer with version newer than 0.5.3?

[leejnau]

@leejnau Is this PR depending on Flashinfer with version newer than 0.5.3?

For the optimized path (fused_topk_deepseek) it relies upon flashinfer Release v0.6.0rc1
https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.6.0rc1

• Rename noauxtc to fused_topk_deepseek by @nv-yunzheq in Rename noauxtc to fused_topk_deepseek flashinfer-ai/flashinfer#2181

However, the check for this API availability means that we should be able to merge this without waiting for flashinfer Release v0.6.0rc1. If you'd prefer to wait for flashinfer Release v0.6.0rc1 that's fine too.

[Fridge003]

@leejnau We need to wait for 0.6.0 update. By then this PR can be thoroughly tested

[Fridge003]

@leejnau Can you please test this PR on DeepSeek R1/V3.2 with the GPQA dataset with this command

python3 -m sglang.test.run_eval --port 30000 --eval-name gpqa --num-examples 198 --max-tokens 128000 --repeat 8 --thinking-mode deepseek-v3

I just tested gpqa on dpsk fp4 with gb200 and the accuracy dropped like ~1%. Not sure whether that's fluctuation

[leejnau]

@leejnau Can you please test this PR on DeepSeek R1/V3.2 with the GPQA dataset with this command

python3 -m sglang.test.run_eval --port 30000 --eval-name gpqa --num-examples 198 --max-tokens 128000 --repeat 8 --thinking-mode deepseek-v3

I just tested gpqa on dpsk fp4 with gb200 and the accuracy dropped like ~1%. Not sure whether that's fluctuation

I ran with server command:

python3 -m sglang.launch_server --model-path nvidia/DeepSeek-R1-0528-FP4-V2 --tensor-parallel-size=4 --cuda-graph-max-bs 256 --max-running-requests 256 --mem-fraction-static 0.85 --ep-size 4 --enable-symm-mem --moe-runner-backend flashinfer_cutlass --quantization modelopt_fp4

and the client command you requested above. I can reproduce the accuracy drop with fused_topk_deepseek:

commit 21e4828f0eec47eabb6b966bc4050771e6316c37 (HEAD -> fused_topk_routing, origin/fused_topk_routing)
"mean_score": 0.7651515151515151

commit 6037267f5bc14105965d00ecd712e8aa5e7aaff4 (HEAD -> main, origin/main, origin/HEAD)
"mean_score": 0.7916666666666667

[yzh119]

Might be relevant to flashinfer-ai/flashinfer#2325 which is fixed in flashinfer v0.6.1.

[Fridge003]

Just tested again with v0.6.1, and the accuracy restored to expected number
@leejnau @yzh119

[Fridge003]

Waiting for upgrade of flashinfer

[Fridge003]

@leejnau The accept length in this test drops, please take a look
https://github.com/sgl-project/sglang/actions/runs/21086349532/job/60703667670?pr=15347

We can relax the accept length to 2.8, since the prior runs hit like 2.92. 2.9 might be too tight for the CI
https://github.com/sgl-project/sglang/blob/main/test/srt/test_deepseek_v3_mtp.py#L86

[Fridge003]

/rerun-stage unit-test-backend-8-gpu-h200

[github-actions]

✅ Triggered unit-test-backend-8-gpu-h200 to run independently (skipping dependencies).

[github-actions]

🔗 View workflow run