[DSv4 Perf] Adaptive topk width, 1.0% E2E throughput improvement

[yewentao256]

Purpose

Let's assume token num = 2

Originally: Always use maximum length

c128a_max_compressed
→ 1,048,576 / 128
→ 8,192

So buffer

global_decode_buffer[:2]
→ shape = [2, 8192]
→ stride = 8192

kernel loop for 8192 times

Now based on current context

32,768 / 128
→ 256
→ active_topk_width = 256

buffer.view(-1)[: 2 × 256]
→ view(2, 256)
→ shape = [2, 256]
→ stride = 256

Kernel only loops for 256 times

Test

cd /home/yewentao256/vllm-source
source /home/yewentao256/wt_env.sh
export MODEL=deepseek-ai/DeepSeek-V4-Pro

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
/home/yewentao256/.venv/bin/python -m vllm.entrypoints.cli.main serve "$MODEL" \
  --served-model-name dsv4-pro \
  --tensor-parallel-size 8 \
  --enable-expert-parallel \
  --max-model-len 1048576 \
  --max-num-batched-tokens 8192 \
  --max-num-seqs 512 \
  --gpu-memory-utilization 0.95 \
  --kv-cache-dtype fp8_ds_mla \
  --no-enable-prefix-caching \
  --trust-remote-code \
  --port 8000 \
  2>&1 | tee /home/yewentao256/dsv4-c128a-prefill-results/diff-server.log

Acc

lm_eval run \
  --model local-completions \
  --model_args model="$MODEL",base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=100,tokenized_requests=False \
  --tasks gsm8k \
  --num_fewshot 5 \
  --output_path /home/yewentao256/dsv4-c128a-lm-eval-results/results \
  --log_samples \
  --seed 0 \
  2>&1 | tee /home/yewentao256/dsv4-c128a-lm-eval-results/lm-eval.log

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9492|±  |0.0060|
|     |       |strict-match    |     5|exact_match|↑  |0.9484|±  |0.0061|

Perf

vllm bench serve \
  --backend openai \
  --base-url http://127.0.0.1:8000 \
  --endpoint /v1/completions \
  --model "$MODEL" \
  --served-model-name dsv4-pro \
  --dataset-name random \
  --random-input-len 8192 \
  --random-output-len 1 \
  --random-range-ratio 0 \
  --num-prompts 32 \
  --num-warmups 4 \
  --request-rate inf \
  --max-concurrency 1 \
  --ignore-eos \

# now
============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  11.84     
Total input tokens:                      262144    
Total generated tokens:                  32        
Request throughput (req/s):              2.70      
Output token throughput (tok/s):         2.70      
Peak output token throughput (tok/s):    3.00      
Peak concurrent requests:                4.00      
Total token throughput (tok/s):          22151.38  
---------------Time to First Token----------------
Mean TTFT (ms):                          369.74    
Median TTFT (ms):                        369.99    
P99 TTFT (ms):                           374.28    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          369.74    
Median E2EL (ms):                        369.99    
P99 E2EL (ms):                           374.28    
==================================================

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  11.82     
Total input tokens:                      262144    
Total generated tokens:                  32        
Request throughput (req/s):              2.71      
Output token throughput (tok/s):         2.71      
Peak output token throughput (tok/s):    3.00      
Peak concurrent requests:                4.00      
Total token throughput (tok/s):          22186.10  
---------------Time to First Token----------------
Mean TTFT (ms):                          369.17    
Median TTFT (ms):                        369.42    
P99 TTFT (ms):                           372.82    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          369.17    
Median E2EL (ms):                        369.42    
P99 E2EL (ms):                           372.82    
==================================================

# main
============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  11.94     
Total input tokens:                      262144    
Total generated tokens:                  32        
Request throughput (req/s):              2.68      
Output token throughput (tok/s):         2.68      
Peak output token throughput (tok/s):    3.00      
Peak concurrent requests:                4.00      
Total token throughput (tok/s):          21959.34  
---------------Time to First Token----------------
Mean TTFT (ms):                          372.97    
Median TTFT (ms):                        373.45    
P99 TTFT (ms):                           377.40    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          372.97    
Median E2EL (ms):                        373.45    
P99 E2EL (ms):                           377.40    
==================================================
============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  11.95     
Total input tokens:                      262144    
Total generated tokens:                  32        
Request throughput (req/s):              2.68      
Output token throughput (tok/s):         2.68      
Peak output token throughput (tok/s):    3.00      
Peak concurrent requests:                4.00      
Total token throughput (tok/s):          21945.26  
---------------Time to First Token----------------
Mean TTFT (ms):                          373.21    
Median TTFT (ms):                        373.45    
P99 TTFT (ms):                           377.10    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          373.21    
Median E2EL (ms):                        373.45    
P99 E2EL (ms):                           377.10    
==================================================

[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.}