[Perf] dsv41: Mega-mHC and the wo_a einsum write MXFP8 for wqkv_a and wo_b by kpham-sgl · Pull Request #45632 · sgl-project/sglang

kpham-sgl · 2026-10-21T16:39:35Z

Motivation

In DeepSeek-V4.1 Mega attention decode, wqkv_a and wo_b quantize their bf16 input in a separate kernel before the MXFP8 GEMM. The producers of those inputs (Mega-mHC and the wo_a fp8_einsum) can write the MXFP8 input directly.

Modifications

LinearBase.deepgemm_mxfp8_input: MQALayer sets it on wqkv_a and wo_b before load; Fp8LinearMethod.process_weights_after_loading packs a DeepGEMM weight scale beside the FlashInfer one and keeps the flag only when it did. A DeepGEMM-packed activation pair is a plain (fp8, int32 scale) tuple, which apply routes to the existing deepgemm_w8a8_mxfp8_linear_with_fallback.
mhc_mega_boundary takes y_fp8 / y_gemm_sf out-buffers; the Mega-mHC post fills them for the next layer's wqkv_a.
The Mega wo_a fp8_einsum writes its FP8 output pair for wo_b.
On whenever the layer can be prepared (Blackwell, DeepGEMM UE8M0, FlashInfer dense backend, aligned shapes); decode/verify only (no extend, no SP), at most 2048 rows. Every other linear is unchanged.

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ❌ Run #39703628265
Latest PR Test (Extra): ❌ Run #39703628339
Latest PR Test (AMD ROCm 10): ❌ Run #39703628142

… wo_b With SGLANG_OPT_DSV41_FUSED_MXFP8_PRODUCERS=1, Mega-mHC also writes wqkv_a's E4M3 input with DeepGEMM-packed UE8M0 scales, and the Mega wo_a fp8_einsum writes wo_b's input through its FP8 epilogue; both linears then run DeepGEMM 1d1d on the pre-quantized input. Decode and verify only, up to 2048 rows. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

kpham-sgl · 2026-10-21T18:19:17Z

(codex)

P1 — Weight reload: fp8_utils.py:1022 replaces the packed-scale tensor, leaving captured CUDA graphs using old scales. Use copy_or_rebind_param to refresh the existing storage.
P2 — Native MXFP8: fp8.py:791 returns without preparing the DeepGEMM scales, so the fusion never activates for native MXFP8. Prepare them on this path too.
Drop the manual script: it bypasses production loading/dispatch and only prints reference errors. It can even report PASS for NaN outputs: max(0.0, nan, nan) == 0.0.

Source review only; GPU numerics/performance were not run.

…ative MXFP8 copy_or_rebind_param keeps weight_scale_inv_deepgemm_input's storage stable across weight reloads so captured CUDA graphs see refreshed scales. The native MXFP8 checkpoint path now prepares the packed scale too, so the fused producers activate there as well. Drop the manual numeric script.

kpham-sgl · 2026-10-21T23:07:25Z

/rerun-test test_deepseek_v41_flash_gb300_tp4ep4.py test_deepseek_v41_flash_gb300_dp4ep4.py test_deepseek_v41_flash_gb300_encoder_bcg.py test_fp8_utils.py test_fp8_blockwise_linear_backends.py test_deepgemm_ue8m0_requant.py test_mhc_kernels.py test_mega_attention.py

github-actions · 2026-10-21T23:07:47Z

🚀 4-gpu-gb300 (3 tests): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_tp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_encoder_bcg.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_utils.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/attention/dsv4/test_mega_attention.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

kpham-sgl · 2026-10-21T23:29:29Z

/rerun-test test_deepseek_v41_flash_gb300_tp4ep4.py test_deepseek_v41_flash_gb300_dp4ep4.py test_deepseek_v41_flash_gb300_encoder_bcg.py test_fp8_utils.py test_fp8_blockwise_linear_backends.py test_deepgemm_ue8m0_requant.py test_mhc_kernels.py test_mega_attention.py

Fp8LinearMethod.prepare_deepgemm_mxfp8_input(layer) packs the scale and says whether the layer takes an Mxfp8DeepGemmInput; MQALayer.setup_mega_attention calls it for wqkv_a and wo_b and stores the two booleans. That replaces the opt-in attribute on the layer, the hasattr probes on the derived scale and the duck-typed helper in the mHC module.

github-actions · 2026-10-21T23:29:51Z

🚀 4-gpu-gb300 (3 tests): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_tp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_encoder_bcg.py

🚀 1-gpu-h100 (3 tests): ⏳ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_utils.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py

🚀 4-gpu-b200 (2 tests): ❌ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/attention/dsv4/test_mega_attention.py

🚀 1-gpu-5090 (1 test): ❌ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 ubuntu-latest (1 test): ❌ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

kpham-sgl · 2026-10-21T23:30:37Z

/rerun-test test_deepseek_v41_flash_gb300_tp4ep4.py test_deepseek_v41_flash_gb300_dp4ep4.py test_deepseek_v41_flash_gb300_encoder_bcg.py test_fp8_utils.py test_fp8_blockwise_linear_backends.py test_deepgemm_ue8m0_requant.py test_mhc_kernels.py test_mega_attention.py

github-actions · 2026-10-21T23:30:59Z

🚀 4-gpu-gb300 (3 tests): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_tp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_encoder_bcg.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_utils.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/attention/dsv4/test_mega_attention.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

Drop Mxfp8DeepGemmInput: a DeepGEMM-packed activation pair is a plain (fp8, int32 scale) tuple, as swiglu_oai_mxfp8_deepgemm already produces, and Fp8LinearMethod.apply routes it to the existing deepgemm_w8a8_mxfp8_linear_with_fallback. The prepared layer stores the packed weight scale under the same weight_scale_inv_deepgemm name a DeepGEMM-backend layer uses, through one pack_mxfp8_weight_scale_for_deepgemm helper. mhc_mega_boundary takes the fp8 pair as out-buffers instead of changing its return arity; both producers allocate it through empty_deepgemm_mxfp8_pair and gate on fused_mxfp8_rows. combine hands rows back to any caller without a quantized list, and HcQuantized.pair is named for what it now holds.

kpham-sgl · 2026-10-22T00:05:34Z

/rerun-test test_deepseek_v41_flash_gb300_tp4ep4.py test_deepseek_v41_flash_gb300_dp4ep4.py test_deepseek_v41_flash_gb300_encoder_bcg.py test_fp8_utils.py test_fp8_blockwise_linear_backends.py test_deepgemm_ue8m0_requant.py test_mhc_kernels.py test_mega_attention.py test_minimax_fused_qkv_index_gemm.py

github-actions · 2026-10-22T00:05:55Z

🚀 4-gpu-gb300 (3 tests): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_tp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_encoder_bcg.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_utils.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py

🚀 4-gpu-b200 (3 tests): ❌ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/attention/dsv4/test_mega_attention.py
cd test/ && python3 registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

Drop SGLANG_OPT_DSV41_FUSED_MXFP8_PRODUCERS. prepare_deepgemm_mxfp8_input is the condition, and the wqkv_a half already sits behind the Mega-mHC fused post.

kpham-sgl · 2026-10-22T00:18:18Z

/rerun-test test_deepseek_v41_flash_gb300_tp4ep4.py test_deepseek_v41_flash_gb300_dp4ep4.py test_deepseek_v41_flash_gb300_encoder_bcg.py test_fp8_utils.py test_fp8_blockwise_linear_backends.py test_deepgemm_ue8m0_requant.py test_mhc_kernels.py test_mega_attention.py test_minimax_fused_qkv_index_gemm.py

github-actions · 2026-10-22T00:18:40Z

🚀 4-gpu-gb300 (3 tests): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_tp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_encoder_bcg.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_utils.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py

🚀 4-gpu-b200 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/attention/dsv4/test_mega_attention.py
cd test/ && python3 registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

post_load_weights runs inside load_weights, before the quant methods process their layers, so packing the scale from setup_mega_attention found no block_fp8_mxfp8_ready yet (GB300 job 3927). LinearBase.deepgemm_mxfp8_input is the request: MQALayer sets it in __init__ for wqkv_a and wo_b, Fp8LinearMethod.process_weights_after_loading keeps it True only once the packed scale exists, and the producers read it as the fact.

kpham-sgl · 2026-10-22T01:06:34Z

/rerun-test test_deepseek_v41_flash_gb300_tp4ep4.py test_deepseek_v41_flash_gb300_dp4ep4.py test_deepseek_v41_flash_gb300_encoder_bcg.py test_fp8_utils.py test_fp8_blockwise_linear_backends.py test_deepgemm_ue8m0_requant.py test_mhc_kernels.py test_mega_attention.py test_minimax_fused_qkv_index_gemm.py

github-actions · 2026-10-22T01:06:56Z

🚀 4-gpu-gb300 (3 tests): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_tp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp4ep4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_encoder_bcg.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_utils.py
cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py

🚀 4-gpu-b200 (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/attention/dsv4/test_mega_attention.py
cd test/ && python3 registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

kpham-sgl · 2026-10-22T01:23:30Z

Ran the three registered GB300 accuracy tests on 4x GB300 at d463e21, as CI does (cd test && python3 registered/e2e/models/<file>):

Test	GSM8K	MMMU-Pro
test_deepseek_v41_flash_gb300_tp4ep4	0.980	0.765
test_deepseek_v41_flash_gb300_dp4ep4	0.985	0.770
test_deepseek_v41_flash_gb300_encoder_bcg	0.980	0.785

All OK. The fused wo_b path was active: each log shows the DeepGEMM JIT compile of the dense GEMM at wo_b's shape (N=5120, K=2048 under TP4, K=8192 under DP4), which nothing else triggers with fp8_gemm_runner_backend=flashinfer_cutedsl. The wqkv_a half only runs inside the Mega-mHC fused post (SGLANG_OPT_DSV41_MEGA_MHC_PREFILL), which no registered test sets.

The previous head 57860e4 failed at weight load: setup_mega_attention runs inside load_weights, before the loader's process_weights_after_loading pass, so block_fp8_mxfp8_ready did not exist yet. d463e21 moves the request to LinearBase.deepgemm_mxfp8_input, honoured by Fp8LinearMethod.process_weights_after_loading. Cancelled the queued GB300 CI rerun since this covers it; the unit and kernel reruns on d463e21 passed.

kpham-sgl requested review from Alisehen, AniZpZ, BBuf, DarkSharpness, Edwardf0t1, FlamingoPg, HaiShaw, HydraQYH, OrangeRedeng, b8zhong, celve, ch-wan, mmangkad and yuan-luo as code owners October 21, 2026 16:39

github-actions Bot added deepseek npu jit-kernel labels Oct 21, 2026

dsv41: Run the fused MXFP8 producers whenever the layer can be prepared

57860e4

Drop SGLANG_OPT_DSV41_FUSED_MXFP8_PRODUCERS. prepare_deepgemm_mxfp8_input is the condition, and the wqkv_a half already sits behind the Mega-mHC fused post.

kpham-sgl requested review from Fridge003, Ying1123, ispobock and merrymercy as code owners October 22, 2026 01:05

kpham-sgl merged commit aa10927 into sgl-project:main Oct 22, 2026
109 of 126 checks passed

This was referenced Oct 22, 2026

[Perf] dsv41: Read Engram shards directly over NVLink #45679

Closed

[Superseded] Add opt-in MXFP8 output for DeepSeek WO-A #45677

Closed

bingxche mentioned this pull request Oct 22, 2026

[CI Monitor] Daily Report - 2026-10-22 bingxche/sglang-ci-bot#241

Open

BBuf mentioned this pull request Oct 22, 2026

Allow configurable absmax floor in FP8 epilogues sgl-project/DeepGEMM#111

Closed

bingxche mentioned this pull request Oct 23, 2026

[CI Monitor] Daily Report - 2026-10-23 bingxche/sglang-ci-bot#242

Open

chaohu2000 mentioned this pull request Oct 27, 2026

feat(dsv41): enable DeepInfer features on PPU SGLang (v0.5.22) flytiger-eco/sglang-for-sail#165

Open

5 tasks

xulunfan mentioned this pull request Oct 29, 2026

[Perf] [SM120] fp8: Use native block32 DeepGEMM for DSV4.1 #47210

Open

8 of 11 tasks

Smallfu666 mentioned this pull request Nov 3, 2026

[Bug] RadixArk/Muse-Glimmer-NVFP4 fails at load since #45632 (MXFP8 lm_head), and on Hopper also at DeepGEMM warmup #48127

Closed

5 tasks done

Conversation

kpham-sgl commented Oct 21, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

CI States

Uh oh!

kpham-sgl commented Oct 21, 2026

Uh oh!

kpham-sgl commented Oct 21, 2026

Uh oh!

github-actions Bot commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Oct 21, 2026

Uh oh!

github-actions Bot commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Oct 21, 2026

Uh oh!

github-actions Bot commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Oct 22, 2026

Uh oh!

github-actions Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Oct 22, 2026

Uh oh!

github-actions Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Oct 22, 2026

Uh oh!

github-actions Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Oct 22, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

kpham-sgl commented Oct 21, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Oct 21, 2026 •

edited

Loading

github-actions Bot commented Oct 21, 2026 •

edited

Loading

github-actions Bot commented Oct 21, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading