[chaunceyjiang]

Overview

Enable existing native 8/16/32-head kernels in SM90 BF16 sparse FlashMLA for sufficiently large query batches, avoiding unnecessary computation from padding heads to 64.

#61615

Claims

• Reduce head-padding overhead in the BF16 sparse attention path.
• Preserve the padded path for small query batches, where native small-head kernels can be slower.
• Improve output throughput by 10.0% in the recorded GLM-5.3-Flash serving benchmark.

Validation

Serving benchmark: GLM-5.3-Flash, TP4 with EP, BF16 KV cache, 16,384 input tokens, 512 output tokens, 128 requests, concurrency 32.

Metric	main	This PR	Change
Output throughput	258.99 tok/s	284.97 tok/s	+10.0%
Mean TTFT	11,161.42 ms	9,981.46 ms	−10.6%
Mean TPOT	101.74 ms	92.77 ms	−8.8%
Mean end-to-end latency	63,150.57 ms	57,385.93 ms	−9.1%

vllm serve \                                                   
  zai-org/GLM-5.3-Flash \                                                                            
  --host 127.0.0.1 \                                                                                         
  --port 8123 \                                                       
  --seed 0 \                                               
  --tensor-parallel-size 4 \
  --max-model-len 32768 \
  --max-num-seqs 128 \
  --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.90 \
  --kv-cache-dtype bfloat16 \
  --no-enable-prefix-caching \
  --max-cudagraph-capture-size 128 \
  --limit-mm-per-prompt '{"image":0,"video":0}' -ep --attention-backend FLASHMLA_SPARSE

vllm bench serve --backend vllm --model zai-org/GLM-5.3-Flash \
--host 127.0.0.1 --port 8123 --dataset-name random --random-input-len 16384 --random-output-len 512 \
--random-range-ratio 0 --num-prompts 128 --max-concurrency 32 --request-rate inf --num-warmups 32 --seed 42 \
--temperature 0 --ignore-eos --percentile-metrics ttft,tpot,itl,e2el \
--metric-percentiles 50,90,99 --save-result --save-detailed

main

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Maximum request concurrency:             32        
Benchmark duration (s):                  253.05    
Total input tokens:                      2097152   
Total generated tokens:                  65536     
Request throughput (req/s):              0.51      
Output token throughput (tok/s):         258.99    
Peak output token throughput (tok/s):    1696.00   
Peak concurrent requests:                35.00     
Total token throughput (tok/s):          8546.64   
---------------Time to First Token----------------
Mean TTFT (ms):                          11161.42  
Median TTFT (ms):                        5847.39   
P50 TTFT (ms):                           5847.39   
P90 TTFT (ms):                           33109.19  
P99 TTFT (ms):                           52259.67  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          101.74    
Median TPOT (ms):                        112.62    
P50 TPOT (ms):                           112.62    
P90 TPOT (ms):                           113.17    
P99 TPOT (ms):                           118.99    
---------------Inter-token Latency----------------
Mean ITL (ms):                           101.74    
Median ITL (ms):                         20.60     
P50 ITL (ms):                            20.60     
P90 ITL (ms):                            253.46    
P99 ITL (ms):                            848.57    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          63150.57  
Median E2EL (ms):                        62729.48  
P50 E2EL (ms):                           62729.48  
P90 E2EL (ms):                           90909.01  
P99 E2EL (ms):                           110017.84 
==================================================

this pr

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Maximum request concurrency:             32        
Benchmark duration (s):                  229.98    
Total input tokens:                      2097152   
Total generated tokens:                  65536     
Request throughput (req/s):              0.56      
Output token throughput (tok/s):         284.97    
Peak output token throughput (tok/s):    1696.00   
Peak concurrent requests:                35.00     
Total token throughput (tok/s):          9403.97   
---------------Time to First Token----------------
Mean TTFT (ms):                          9981.46   
Median TTFT (ms):                        5236.98   
P50 TTFT (ms):                           5236.98   
P90 TTFT (ms):                           29585.98  
P99 TTFT (ms):                           46692.16  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          92.77     
Median TPOT (ms):                        102.48    
P50 TPOT (ms):                           102.48    
P90 TPOT (ms):                           102.81    
P99 TPOT (ms):                           107.98    
---------------Inter-token Latency----------------
Mean ITL (ms):                           92.77     
Median ITL (ms):                         20.48     
P50 ITL (ms):                            20.48     
P90 ITL (ms):                            234.20    
P99 ITL (ms):                            756.93    
----------------End-to-end Latency----------------
Mean E2EL (ms):                          57385.93  
Median E2EL (ms):                        57243.95  
P50 E2EL (ms):                           57243.95  
P90 E2EL (ms):                           82115.96  
P99 E2EL (ms):                           99186.91  
==================================================

lm_eval run \ 
  --model local-completions \
  --model_args model=zai-org/GLM-5.3-Flash,base_url=http://127.0.0.1:8123/v1/completions,num_concurrent=64,tokenized_requests=False \
  --tasks gsm8k \
  --num_fewshot 10 \
  --batch_size 1 \
  --gen_kwargs temperature=0 \
  --output_path /home/chauncey/vllm/downloads/glm53-flash_profile/gsm8k_lm_eval \
  --log_samples \
  --seed 12345


|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    10|exact_match|↑  |0.9393|±  |0.0066|
|     |       |strict-match    |    10|exact_match|↑  |0.9378|±  |0.0067|

Details

Native heads are used on SM90 when the actual head count is 8/16/32 and the query-token count exceeds:

• SM count for 8 or 16 heads.
• Twice the SM count for 32 heads.

These are empirical performance thresholds. Each query launches one CTA, and small query batches can favor the padded kernel despite its extra head computation.

The implementation uses the actual head count, including heads gathered by DCP, and slices both queries and optional output buffers before calling FlashMLA. Existing padded workspaces remain available for fallback.

Related work and duplicate check: Reviewed open PRs matching FlashMLA small heads and SM90 sparse. #59418 (#59418) adds a Triton decoder that directly reads packed FP8 KV caches, while #57673 (#57673) adds an opt-in Triton backend. This PR integrates existing native BF16 kernels into the current FlashMLA sparse backend.

AI assistance was used for implementation, validation, and preparation of this description.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[chaunceyjiang]

/ci run

[github-actions]

✅ Triggered Buildkite CI #100101 for commit 729f28a937d9.

[chaunceyjiang]

/ci run

[github-actions]

✅ Triggered Buildkite CI #100167 for commit 0b39be0ffd3d.