[Perf] Improve fp8 quant in mla; replace ReduceSum with ReduceScatterSum

[IwakuraRein]

Purpose

• Choose ReduceScatterSum over ReduceSum when sizes are the same

  

• Avoid calling fp8 quant kernel twice when using fp8 attention by creating a stage buffer for decode_ql_nope and decode_q_pe. This also eliminates the torch.cat in the flashinfer mla backend.

  

Test Plan

Launch command

  --dtype auto --kv-cache-dtype fp8 \
  --tensor-parallel-size 8 \
  --swap-space 16 --max-num-seqs 1024 --trust-remote-code --max-model-len 10240 --gpu-memory-utilization 0.95 \
  --max-num-batched-tokens 16384 --async-scheduling \
  --max-cudagraph-capture-size 1024 --compilation_config.cudagraph_mode FULL_DECODE_ONLY

Evaluation command

lm-eval --model local-completions --tasks gsm8k --model_args model=nvidia/DeepSeek-R1-0528-FP4-v2,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False --num_fewshot 20

Models

• Test with DeepSeek R1 0528 checkpoint and 8 H200.
• Test with DeepSeek R1 0528 FP4 V2 checkpoint and 4 GB200.

Test Result

• Hopper, VLLM_ATTENTION_BACKEND=FLASHMLA

  |Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
  |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
  |gsm8k|      3|flexible-extract|    20|exact_match|↑  |0.9530|±  |0.0058|
  |     |       |strict-match    |    20|exact_match|↑  |0.9522|±  |0.0059|
  

• Blackwell, VLLM_ATTENTION_BACKEND=FLASHINFER_MLA

  |Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
  |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
  |gsm8k|      3|flexible-extract|    20|exact_match|_  |0.9538|_  |0.0058|
  |     |       |strict-match    |    20|exact_match|_  |0.9522|_  |0.0059|
  

Perf gain: 5% projected improvement in decode when using flashinfer mla backend when serving DeepSeek R1 FP4 on 4 GB200 with DEP4

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[heheda12345]

CC @MatthewBonanni @LucasWilkinson

[MatthewBonanni]

Thanks for this contribution! Can you include some benchmark results (or a snippet of a profile) to get a sense of the speedup?

[pavanimajety]

LGTM, minor comments

[pavanimajety]

@IwakuraRein Could you please add in the PR description why ReduceScatterSum is better than ReduceSum when sizes are same? or post some comparison results?