[kpham-sgl]

Motivation

With MegaGate on (SGLANG_OPT_DEEPGEMM_MEGA_GATE=1) the gate writes the routing into MegaMoE's buffer, but every MoE layer still runs mega_moe_pre_dispatch to quantize its input: one pass over the bf16 rows per layer that writes the FP8 rows, the per-row routed scales and the paged shared-expert scales. On SM100 those rows come out of the Mega-mHC kernel, which can already emit them as FP8 in both scale layouts (deep_gemm.mega_mhc takes y_routed_sf, y_shared_sf and shared_sf_block_m). Follow-up to #43262 and #46071.

Modifications

• mega_moe.py: mega_moe_input_target resolves the MoE's buffer views for a batch that will take the MegaGate path (fp8xfp4 with fused shared experts); MegaMoeInput.mark tags the normed rows; run_mega_routed_experts skips the pre-dispatch when the routing is in the buffer and the rows carry the tag for that buffer and row count.
• deepseek_v4_mhc.py: HcNextBoundary.mega_moe_input; _post_fusion asks the next MoE for its target and hands the views to mhc_mega_boundary.
• deepseek_v4_boundary.py: the MoE-side read boundary carries the resolver with its forward batch.
• mhc_mega.py: moe_input through to deep_gemm.mega_mhc.

No change with MegaGate off. Prefill, non-fused shared experts, routing that stays out of the buffer and narrowed rows all keep the pre-dispatch.

Accuracy Tests

GSM8K 5-shot, 1,314 examples, native acceptance, DP attention + EP 4 on GB300, MegaGate + v41_new + --enable-fp32-lm-head: 0.9711 with this PR and 0.9711 on the base commit. The GB300 dp2ep2 and tp2ep2 e2e test files pass (GSM8K 0.975 both, MMMU-Pro 0.815 / 0.79).

Speed Tests and Profiling

Performance neutral: 0.02 to 0.05 ms per decode step (0.1 to 0.4%), at the edge of what one pair of runs resolves.

DeepSeek-V4.1-Flash decode, DP attention + EP 8 on two GB300 nodes, MegaGate + v41_new + --enable-fp32-lm-head, forced balanced routing, one cached 126,976-token prefix plus 4,096 own tokens per request, 3 drafts at simulated acceptance 3.25, a fixed batch of N requests per rank; step time per rank in ms from the scheduler's own per-step record over a 120 s window, same commit with and without this PR:

N per rank	base ea226cf	this PR	diff
8	10.761	10.745	-0.017
16	13.029	12.976	-0.053
32	13.654	13.622	-0.033
64	15.483	15.433	-0.050
128	20.141	20.123	-0.018
256	28.867	28.844	-0.023

GPU traces (20 steps, rank 0) show why: mega_moe_pre_dispatch_kernel drops from 42 to 3 launches per cycle (the 3 are the DSpark draft MoE layers), 0.103 -> 0.007 ms at 64 per rank and 0.207 -> 0.012 at 256; but the FP8-emitting sm100_mega_mhc_impl template is slower than the bf16-only one at 256 rows, 1.192 -> 1.313 ms per cycle over 39 layers (about 3 us per launch), and equal at 64 rows (0.600 -> 0.585). Net kernel time -0.11 ms per cycle at 64 per rank and -0.07 at 256. The gain would be about 0.2 ms per step at 256 per rank if the FP8-emitting template cost the same as the bf16 one.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #40028315844
Latest PR Test (Extra): ❌ Run #40028315489
Latest PR Test (AMD ROCm 10): ❌ Run #40028315589

[kpham-sgl]

/rerun-test test/registered/kernels/ops/moe/test_mega_gate.py test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py test/registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py

[github-actions]

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

cd test/ && python3 registered/kernels/ops/moe/test_mega_gate.py

🚀 ubuntu-latest (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
cd test/ && python3 registered/unit/layers/test_dsv41_late_layer_tail_readers.py

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

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py

[kpham-sgl]

Correctness, head aca8787, DeepSeek-V4.1-Flash on GB300 (the 2-gpu-gb300 rerun was cancelled for queue time and run on our cluster from the same test files; the 4-gpu-b200 and ubuntu-latest reruns passed above):

Test	Setting	Result
test/registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py	as in CI	OK, GSM8K 0.975, MMMU-Pro 0.815
test/registered/e2e/models/test_deepseek_v41_flash_gb300_tp2ep2.py	as in CI	OK, GSM8K 0.975, MMMU-Pro 0.79
GSM8K 5-shot, 1,314 examples, native acceptance, DP attention + EP 4, MegaGate + v41_new + --enable-fp32-lm-head	this PR	0.9711
same	base ea226cf	0.9711