PR #37960: [Performance] Fused rmsnorm fp8 quant for DeepSeek

• Repository: vllm-project/vllm
• Author: tianrengao
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/37960

Description (problem / solution / changelog)

fixes: https://github.com/vllm-project/vllm/issues/37832

• Fused Triton kernel: Combines RMSNorm + FP8 per-block quantization into a single kernel for q_b_proj, eliminating the intermediate DRAM round-trip between norm and quant.
• Native aten decomposition: Replaces the opaque CUDA CustomOp with native aten ops for norm-adjacent linears, enabling Inductor to automatically fuse normalize + quantize (requires companion PyTorch PR #178237 to fix Reduction+Pointwise fusion).

Summary: Introduces a custom Triton kernel that fuses RMSNorm and per-block FP8 quantization into a single kernel for norm-adjacent linear layers, along with selective native quantization that enables Inductor to fuse normalize + quantize operations. Together these reduce E2E latency significantly on DeepSeek-R1 671B FP8 with TP=8. see:

Companion PyTorch PR: [inductor] Enable Reduction+Pointwise fusion when reduction output has opaque consumers (fixes Inductor scheduler to allow fusing Reduction epilogues when the output feeds an opaque consumer like cutlass_scaled_mm).

Problem

Under torch.compile with Inductor, DeepSeek-R1 FP8 inference has a performance gap between RMSNorm and per-block FP8 quantization:

1. The QuantFP8 CustomOp is opaque to Inductor, preventing fusion with the preceding RMSNorm. This forces an unnecessary DRAM round-trip (write normalized output, read it back for quantization).

2. Even when quantization uses native aten ops (Inductor-visible), the variance reduction (rnumel=hidden_dim) and quant reduction (rnumel=128) have different iteration spaces. Inductor scheduler requires rnumel1 == rnumel2 (simd.py:1337) and MixOrderReduction requires no producer-consumer relationship (scheduler.py:296). Neither can fuse them.

Solution (3 layers)

1. Selective native quantization (fp8.py + fp8_utils.py)

For norm-adjacent linear layers (q_a_proj, q_b_proj, kv_a_proj, kv_b_proj), replace the opaque QuantFP8 CustomOp with pure aten ops (_quantize_input_native). Non-norm-adjacent layers (o_proj, MoE) keep the optimized CUDA CustomOp.

Key technique: transposed-amax produces column-major scales without .contiguous() copy — compute amax on permute(1,0,2) view so output [G,B] contiguous = [B,G] col-major via .t() view.

2. Fused RMSNorm+FP8Quant Triton kernel (fp8_utils.py)

_fused_rmsnorm_fp8_quant_kernel — a two-pass persistent reduction:

• Pass 1: Read x, compute variance = mean(x^2) per row
• Pass 2: Read x again (L1 cache hit), normalize * weight, per-group amax -> scale -> FP8 cast, write column-major scales directly

Registered as vllm::fused_rmsnorm_fp8_quant custom_op. Replaces 2 Inductor kernels (triton_red + triton_per) with 1 custom kernel.

Applied to q_b_proj (preceded by q_a_layernorm). kv_b_proj cannot be fused because kv_c_normed is stored in the KV cache.

3. Model integration (deepseek_v2.py + mla.py)

_set_fused_norm() stores RMSNorm weight/eps on the linear ops quant op. The MLA forward skips standalone q_a_layernorm when the downstream q_b_proj has fused norm (norm is done inside the fused kernel).

E2E Codegen Verification

DeepSeek-R1 671B FP8, TP=8, compiled graph (rank0):

fused_rmsnorm_fp8_quant calls: 6 (q_b_proj across layers) per_token_group_fp8_quant calls: 2 (non-norm-adjacent, CustomOp) triton_per (native quant): 9 (kv_b, q_a, kv_a) Scale layout: (s72, 12) strides (1, s72) — column-major confirmed

Benchmark

DeepSeek-R1 671B FP8, TP=8, BS=1, input=128, output=8 (vllm bench latency): rate = 15, num requests=600 Baseline (all CUDA CustomOp): 98.15 ms This PR (fused Triton kernel): 95.91 ms (-2.24 ms, -2.3%)

Numerical Correctness

BS=1: q_diff=0.0 vs float32 reference (exact match) BS=8: q_diff=0.0 (Triton reduction ordering, within FP8 precision) BS=64: q_diff=0.0 (within FP8 precision bounds)

The fused kernel operates in float32 throughout (no intermediate bf16 truncation), giving equal or better precision than separate norm->quant.

Test Plan:

Codegen verification (quick, ~5 min compile)

VLLM_SKIP_NATIVE_QUANT=0 TORCH_LOGS=output_code \
vllm bench latency --model deepseek-ai/DeepSeek-R1 \
  --tensor-parallel-size 8 \
  --compilation-config '"'"'{"custom_ops": ["-rms_norm"]}'"'"' \
  --input-len 128 --output-len 2 --batch-size 1 \
  --num-iters 1 --num-iters-warmup 0 2>&1 | \
  grep "fused_rmsnorm_fp8_quant"
# Expected: 6 fused_rmsnorm_fp8_quant calls in compiled graph

A/B benchmark

for SKIP in 1 0; do
  echo "=== SKIP=$SKIP ==="
  VLLM_SKIP_NATIVE_QUANT=$SKIP \
  vllm bench latency --model deepseek-ai/DeepSeek-R1 \
    --tensor-parallel-size 8 \
    --compilation-config '"'"'{"custom_ops": ["-rms_norm"]}'"'"' \
    --input-len 128 --output-len 8 --batch-size 1 \
    --num-iters 3 --num-iters-warmup 1
done
# SKIP=1 (baseline) vs SKIP=0 (fused): expect ~2% improvement

Standalone numerical correctness test

python test_fused_rmsnorm_quant.py
# Tests fused kernel vs float32 reference across BS=1/8/64
# and hidden_dim=512/1536/7168

NOTE: Requires companion PyTorch PR for Inductor scheduler fix (choices.py + scheduler.py). Without it, the native quant path for kv_b/q_a/kv_a will produce extra triton_poi kernels. ' 2>&1

• 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.

Changed files

• vllm/config/compilation.py (modified, +18/-0)
• vllm/model_executor/layers/mla.py (modified, +11/-1)
• vllm/model_executor/layers/quantization/fp8.py (modified, +14/-0)
• vllm/model_executor/layers/quantization/utils/fp8_utils.py (modified, +207/-4)
• vllm/model_executor/models/deepseek_v2.py (modified, +37/-0)