[Perf][DSv4.1] Keep activations MXFP8 between DeepGEMM kernels

[gcanlin]

Purpose

DeepGEMM's DeepSeek-V4.1 kernels can hand FP8 activations to each other, but vLLM re-quantized their BF16 outputs before every dense projection: 127 standalone MXFP8 quantize launches per decode step at C1 (TP4), plus the quantize half of the Mega-MoE input staging. This PR wires the FP8 outputs DeepGEMM already provides to the consumers, and runs those consumers on DeepGEMM.

One decoder layer, decode step. [q] = standalone MXFP8 quantize launch, sf = the UE8M0 group scales.

Before — every DeepGEMM kernel returns BF16, and a FlashInfer quantize kernel re-creates MXFP8 before each consumer:

attention sublayer
  mega_mhc (post+pre+RMSNorm) ──bf16──┬──▶ compressor score / indexer weights GEMMs (bf16, unchanged)
                                      └──▶ [q] ──mxfp8+sf──▶ FlashInfer GEMM  fused_wqa_wkv
  attention ──▶ fp8_einsum wo_a ──bf16──▶ [q] ──mxfp8+sf──▶ FlashInfer GEMM  wo_b ──▶ reduce-scatter

FFN sublayer
  mega_mhc (post+pre+RMSNorm) ──bf16──┬──▶ router gate + top-k (bf16, unchanged)
                                      └──▶ prepare_megamoe = [q] into symm buffer + pack top-k ──▶ fp8_fp4_mega_moe

After — the producers emit MXFP8 with packed scales directly, the consumers run on DeepGEMM fp8_gemm_nt:

attention sublayer
  mega_mhc ──bf16──────────────────────▶ compressor score / indexer weights GEMMs (unchanged)
           ──mxfp8+sf (y_fp8, y_gemm_sf)─▶ DeepGEMM GEMM  fused_wqa_wkv        (not under SP: input is all-gathered in bf16)
  attention ──▶ fp8_einsum wo_a ──mxfp8+sf ((z, sf) output)──▶ DeepGEMM GEMM  wo_b ──▶ reduce-scatter

FFN sublayer
  mega_mhc ──bf16─────────────────────────────────────▶ router gate + top-k (unchanged)
           ──mxfp8+sf written straight into the symm buffer (y_fp8, y_routed_sf, y_shared_sf)
                                                       ──▶ prepare_megamoe = pack top-k only ──▶ fp8_fp4_mega_moe

per layer	before	after
standalone quantize launches	3 (+ the quantize half of prepare_megamoe)	1 under SP (fused_wqa_wkv), 0 otherwise
dense GEMM kernel	FlashInfer CuTe-DSL ×2	DeepGEMM fp8_gemm_nt (1,1,32) ×2
bytes between mHC / einsum and their consumers	bf16 out + fp8 in	fp8 out only

Pieces:

• DeepGemmMxfp8LinearKernel (2D fp8_gemm_nt, recipe (1,1,32)); its activations carry a new QuantKey kMxfp8DynamicDeepGemm (int32 packed UE8M0 scales) so a FlashInfer-layout producer can never be routed to a DeepGEMM consumer or vice versa. The model marks wo_b / fused_wqa_wkv with deep_gemm_activations, and ModelOpt's post-load selects the kernel for them the same way it does for is_bmm layers.
• mega_mhc wrapper: optional y_fp8/y_gemm_sf and y_fp8/y_routed_sf/y_shared_sf outputs (DeepGEMM already implements them), reported as a typed MegaMhcFp8Outputs; the MoE leg is passed explicitly (DeepseekV4MoE.forward(staged=...) → experts.forward_staged) and stage_megamoe_routing repacks only the top-k.
• deep_gemm_fp8_o_proj(fp8_z=True): fp8_einsum emits (z, sf); also in the mega-attention backend's _o_proj.
• Under sequence parallel the attention input is all-gathered in BF16 after the norm, so fused_wqa_wkv stays on FlashInfer there (FP8 all-gather is follow-up work).
• The new fp8_gemm_nt variants are pre-compiled by the existing DeepGEMM startup warmup (layers found through the deep_gemm_warmup_provider protocol).
• Binds wherever SM100 + DeepGEMM are available (checked at construction; other platforms keep their kernels with a logged reason). No new flag.

Duplicate check: searched open PRs for mega_mhc y_fp8, y_gemm_sf, fp8_einsum, DeepGemmMxfp8LinearKernel, prepare_megamoe_inputs, wo_b quantize, and by touched file. #57679 (restores the wq_b query-quant fusion) is complementary and untouched here; #56751 (skinny GEMM, M ≤ 4) quantizes in-kernel for the same projections at tiny M, a separate mechanism; #55505 (input-scale contract) does not conflict.

Test Plan

pytest -v -s tests/kernels/test_dsv41_deepgemm_fp8_chain.py   # SM100

vllm serve deepseek-ai/DeepSeek-V4.1-Flash --tensor-parallel-size 4 --enable-expert-parallel \
  --kernel-config '{"moe_backend":"deep_gemm_mega_moe"}' --attention-backend FLASHINFER_MLA_SPARSE_DSV41 \
  --kv-cache-dtype fp8 --tokenizer-mode deepseek_v41 --max-model-len 32768 --max-num-batched-tokens 8192 --max-num-seqs 64

Baseline: the same command on main. vllm bench serve random 4096/256, ignore-eos, one server at a time; GSM8K via lm-eval 0.4.12 (5-shot chat, thinking off, greedy, 1319 questions).

Test Result

4×B200, TP4 + EP4, MegaMoE + SP, FlashInfer sparse MLA, fp8 KV.

• Unit tests: 21 passed (DeepGEMM kernel vs FlashInfer CuTe-DSL at M ∈ {1,4,7,64,300}; pre-quantized input bit-identical to in-kernel quant; mega_mhc FP8 outputs bit-exact MXFP8 of its BF16 output, MoE layout bit-identical to prepare_megamoe_inputs; einsum FP8 output vs quantized BF16 output; routing-only staging).
• Greedy outputs identical on 8/8 prompts × 64 tokens.
• C1 decode step (torch profiler, 61 steps): GPU busy 7.64 → 6.85 ms (−10 %), quantize launches 130 → 90.

lane	baseline	PR	change
C1 ×3 TPOT	6.58 / 6.58 / 6.58 ms	6.50 / 6.50 / 6.50 ms	−1.2 %
C4 ×2 tok/s (TPOT)	517 / 520 (6.97 / 6.94 ms)	531 / 530 (6.83 / 6.86 ms)	+2.4 % (−1.6 %)
C4 ×2 TTFT	201 / 199 ms	184 / 183 ms	−8 %
C16 ×3 tok/s (TPOT)	1347 / 1346 / 1340 (9.92 / 9.88 / 10.03 ms)	1365 / 1357 / 1368 (9.85 / 9.86 / 9.73 ms)	+1.5 % (−1.3 %)
C64 ×1 tok/s (TTFT)	1848 (2360 ms)	1880 (2124 ms)	+1.7 % (−10 %)
GSM8K flexible / strict	0.9621–0.9659 / 0.7817–0.7832 (3 baselines)	0.9644 / 0.7877	unchanged

After the plumbing refactor (kernel selection through ModelOpt post-load, typed mHC outputs, explicit staged MoE input): 21 unit tests pass, C1 TPOT 6.50 / 6.50 / 6.50 ms, GSM8K 0.9636 / 0.7915, greedy outputs coherent on 12 prompts.

At C1 the host loop (~6.5 ms/step) bounds TPOT, so most of the GPU saving shows from C4 up. The FP8-output mega_mhc variants JIT-compile on first use per prefill shape, like the BF16 variants today (cache persists under DG_JIT_CACHE_DIR).

Developed with AI assistance (Claude Code); the human submitter reviewed every changed line and ran the measurements.

🤖 Generated with Claude Code

[chilkotiKartik]

Architecture Review: MXFP8 Activation Retention between DeepGEMM Kernels

Review Highlights:

1. Quantization Round-Trip Elimination:
   • Retaining activations in Microscaling FP8 (MXFP8 / OCP FP8) format between contiguous DeepGEMM projections eliminates unnecessary FP8 -> BF16 -> FP8 dequant/requant overhead.
2. Scale Factor Alignment:
   • Ensure the block scale factors (e.g. 1 scale per 32 elements) are correctly forwarded and aligned along the inner reduction dimension $K$ of the subsequent layer.
3. Numerical Precision:
   • Retaining MXFP8 representations across linear layers matches DeepSeek-V3/DSv4.1 numerical parity with minimal loss in perplexity while maximizing Tensor Core throughput.

Great performance enhancement for DeepSeek model serving on Hopper / Blackwell architectures!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gcanlin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

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

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gcanlin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork