[FP8][MoE] Honor UE8M0 activation scales in Triton MoE

[kaixih]

Motivation

DeepSeek-V4 checkpoint metadata declares scale_fmt: ue8m0 for its FP8 path. The official DSV4 reference applies this policy when quantizing activations before every expert Linear, so both FC1 and FC2 require power-of-two activation scales. The Triton MoE path previously dropped that policy and used continuous amax / 448 FP32 scales instead.

Test

Made Triton MoE recognize use_scale_ue8m0 and apply power-of-two activation scales to both FC1 and FC2, then ran the following model-level A/B evaluation.

• Machine: 8×NVIDIA B200
• Model: DeepSeek-V4-Flash
• Runtime: lmsysorg/sglang:latest (sha256:7b6a35df9839fd593a94a1eaee82d7777f472225d9f3ad1f8a2e0cb2bd1785d0)
• Parallelism: TP8 / EP1
• Expert path: the checkpoint's packed FP4 expert weights were converted once to FP8 at load time with SGLANG_DSV4_FP4_DEQUANT=1, allowing the Triton FP8 MoE path to be evaluated
• MoE: --moe-a2a-backend none; runtime logs confirmed the Triton MoE implementation

Benchmark / protocol	Baseline (c039e1a7)	PR (47f3c98c)	Delta
GSM8K, temperature 0, 1 run	1266 / 1319 (95.9818%)	1264 / 1319 (95.8302%)	-0.1516 pp
GPQA Diamond, temperature 1, top-p 1, max thinking, 10 paired samples/question	1676 / 1980 (84.6465%)	1672 / 1980 (84.4444%)	-0.2020 pp

---

CI States

Latest PR Test (Base): ✅ Run #31798419785
Latest PR Test (Extra): ❌ Run #31798423103

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[kaixih]

More context:

On Blackwell, native block-scaled MMA backends can consume E8M0 scale factors directly. The Triton MoE path here performs FP8 dot products and applies FP32 scales outside the dot; it does not use the block-scaled MMA interface that consumes packed E8M0 scale operands. However, when the checkpoint declares scale_fmt: ue8m0, those FP32 scale values should still be restricted to powers of two. This is not expected to materially affect model accuracy, but it removes a numerical-contract mismatch and helps align Triton with native block-scaled backends and FP8 trainers such as Miles/Megatron during RL (example).

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[kaixih]

@nvpohanh @maocheng23 can you help review or find right people?

[b8zhong]

Question: why don't we use trtllm-gen MoE?

[kaixih]

For pure SGLang inference with SGLANG_DSV4_FP4_DEQUANT=1, this PR changes only the dynamic MoE activation-scale policy:

State	Weight scales	Activation scales
Before this PR	Power-of-two	Continuous
After this PR	Power-of-two	Power-of-two

The load-time FP4-to-FP8 conversion already derives power-of-two FP8 weight scales from the checkpoint's serialized E8M0 scales.

[kaixih]

Question: why don't we use trtllm-gen MoE?

This PR does not decide the MoE backend; it follows SGLang’s existing runner selection. Our use case sets SGLANG_DSV4_FP4_DEQUANT=1, which converts the expert weights to block FP8. With moe_runner_backend=auto and moe_a2a_backend=none, the existing FP8 runner selection falls through to Triton (code). Therefore, this PR fixes the UE8M0 activation-scale behavior of the backend actually selected for this configuration.

[b8zhong]

I see. Though unless there are issues using that backend, I'd also recommend directly selecting it with --moe-runner-backend=flashinfer_trtllm (or the routed version)

[b8zhong]

@kaixih Please resolve conflicts

[b8zhong]

/rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp8_h200.py

[b8zhong]

/rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

[b8zhong]

/rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py

[github-actions]

Results for /rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp8_h200.py:

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

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp8_h200.py

[github-actions]

Results for /rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py:

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

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

[github-actions]

Results for /rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py:

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

cd test/ && python3 registered/models_e2e/test_dsa_glm52_tp_mtp.py