[ouershua2299]

Co-authors: @AichenF @jychen21 @YijiaZhao
Submitted by: xuechengw

Dependency scope: opt-in SM90 support that requires a FlashInfer build containing flashinfer-ai/flashinfer#6343 and flashinfer-ai/flashinfer#6517, both merged into FlashInfer main (validated at da992ed9919604c5eb6dc255c63f5e89a466452c). The default FlashInfer 0.7.2 pin and Docker/CI dependencies remain unchanged. An older build gives an explicit load-time error only when this path is selected; every other SM90 path is unaffected. The cookbook includes a fixed-revision source-install recipe. Out-of-box support follows a coordinated FlashInfer version bump, which needs no change to this code. The branch is rebased onto main 6cac4dc3fe.

Motivation

On Hopper, --moe-a2a-backend megamoe covers FP8 checkpoints only. FlashInfer flashinfer-ai/flashinfer#6343 and flashinfer-ai/flashinfer#6517 add flashinfer.megamoe, an SM90 MegaMoE kernel for NVFP4 weights with FP8 activations (dispatch, both expert GEMMs, SwiGLU and combine in one persistent kernel, with a split variant for large batches). This PR wires it into the existing megamoe path for ModelOpt NVFP4 checkpoints, starting with nvidia/DeepSeek-V4-Flash-NVFP4.

On 8x H200 and 8x H20, TP8 / attention-DP8 / EP8, it is faster than DeepEP + DeepGEMM FP8 (W8A8) in prefill and decode (H200: 1.34x prefill TTFT, 1.13x decode ITL; H20: 1.27x / 1.42x) and at parity with the SM90 FP8 MegaMoE path while using NVFP4 weights.

Modifications

1. SM90 NVFP4 MegaMoE path (layers/moe/mega_moe_sm90_fi_nvfp4.py, layers/moe/mega_moe.py, layers/moe/fused_moe_triton/layer.py):
   • Converts the ModelOpt expert weights at load into the format transform_nvfp4_weights_for_mega_moe expects (Marlin nibble order, one global scale per expert, local scales rebalanced into the kernel's range without changing the dequantized values).
   • Forward: DeepGEMM's SM90 pre-dispatch writes FP8 activations into the FlashInfer buffer, then commit_staged() and nvfp4_mega_moe. Under DP attention every rank launches with the same padded token count, as FlashInfer requires.
   • Prewarms the FlashInfer JIT modules before CUDA-graph capture; rejects EPLB and FlashInfer builds without [SM90] MegaMoE: fix workspace cleanup; enable push dispatch on H200; support other models flashinfer-ai/flashinfer#6517.
2. ModelOpt NVFP4 MoE method (layers/quantization/modelopt_quant.py): on SM90 with --moe-a2a-backend megamoe, builds the experts through the module above.
3. Environment variables (environ.py, environment_variables.mdx): SGLANG_FLASHINFER_MEGAMOE_SM90_PREWARM, SGLANG_FLASHINFER_MEGAMOE_SM90_KERNEL_FAMILY.
4. Tests: weekly 8x H200 e2e test (cookbook launch + sgl-eval GSM8K), skipped when FlashInfer lacks [SM90] MegaMoE: fix workspace cleanup; enable push dispatch on H200; support other models flashinfer-ai/flashinfer#6517.
5. Docs: DeepSeek-V4 cookbook note with the launch command and a fixed-revision FlashInfer install.
6. Prefill CUDA graphs (layers/dp_attention.py): when the FI path is captured into a prefill graph it latches the shared-bucket flag, so every DP rank replays one MAX_LEN bucket (FlashInfer needs the same token count on every rank) and batches with an idle rank fall back to eager.

Local CPU validation on the branch in the CI image without GPUs: the layers/moe, layers/quantization and server_args unit suites show no new failures against main 6cac4dc3fe (the same 17 GPU-dependent cases fail on both); pre-commit passed for all changed files.

Test environment: one host, 8x NVIDIA H200 141 GB (NVLink); image lmsysorg/sglang:nightly-dev-cu13-20261021-c1c287a4 with this branch on top of 6cac4dc3fe; FlashInfer built from main da992ed as a wheel (pip wheel "flashinfer-python @ git+https://github.com/flashinfer-ai/flashinfer@da992ed...") and installed into the image, with no PYTHONPATH overlay and no disabled version checks; model nvidia/DeepSeek-V4-Flash-NVFP4; sgl-eval 0.1.3. Unless noted: --tp 8 --dp 8 --enable-dp-attention --moe-a2a-backend megamoe.

Review update (rebased onto main ca4bedfffa). The prefill-graph change was checked on 8x H200 (image nightly-dev-cu13-20261025-65633129, FlashInfer da992ed) with uneven DP bursts: one long request with 7 idle ranks, mixes crossing the 256-token fused/split boundary, alternating 120/600 and 150/450-token prefills one per rank, and 7 short + 1 long. Default config (eager prefill) and --cuda-graph-backend-prefill breakable, before and after the change: no hang, GSM8K answers unchanged within noise. --cuda-graph-backend-prefill full does not start on DeepSeek-V4 (its attention backend does not capture EXTEND), with or without this PR.

Dependency compatibility. With the image's released FlashInfer 0.7.2 wheel (no flashinfer.megamoe), selecting this path raises the expected load-time error naming flashinfer-ai/flashinfer#6343 / flashinfer-ai/flashinfer#6517, the new e2e test skips (8 skipped), and the unit suites above behave as on main.

Accuracy Tests

Registered tests on the environment above:

python -m pytest test/registered/e2e/models/test_deepseek_v4_flash_nvfp4_megamoe_h200.py

8 passed in 545.37s           # e2e (8xH200), GSM8K gate 0.975

GSM8K with the registered GSM8KMixin settings (sgl-eval run gsm8k --num-examples 200 --num-threads 128 --no-thinking, sgl-eval 0.1.3), the cookbook command for this PR and the launch commands below for the baselines, 8x H200:

Backend	Checkpoint	GSM8K
FlashInfer SM90 NVFP4 MegaMoE (this PR)	nvidia/DeepSeek-V4-Flash-NVFP4	0.965
DeepEP + DeepGEMM FP8 (W8A8)	deepseek-ai/DeepSeek-V4-Flash	0.975
SM90 FP8 MegaMoE	sgl-project/DeepSeek-V4-Flash-FP8	0.980

Across the runs of this PR (cookbook launch and the e2e gate, before and after rebasing) GSM8K ranged 0.965-0.980.

Speed Tests and Profiling

DeepSeek-V4-Flash on 8x H200 and 8x H20-3e, TP8 / attention-DP8 / EP8, decode CUDA graphs on, sglang.bench_serving random dataset with fixed lengths (--random-range-ratio 1), seed 20260811, warmup requests equal to the concurrency, cache flushed between cases, median of 3 runs per point. Ratios are FlashInfer MegaMoE over the baseline, geometric mean over input <= 8k and concurrency <= 256; > 1 means FlashInfer is better.

Metric	H200 vs DeepEP+DeepGEMM FP8	H200 vs DeepEP+Humming W4A8	H200 vs SM90 FP8 MegaMoE	H20 vs DeepEP+DeepGEMM FP8	H20 vs DeepEP+Humming W4A8	H20 vs SM90 FP8 MegaMoE
Prefill-only input throughput	1.29	1.37 †	0.97	1.20	1.22	1.01
Prefill-only TTFT	1.34	1.41 †	0.98	1.27	1.30	1.01
Decode-only ITL	1.13	0.97 †	1.06	1.42	1.07	1.05
Unified ITL	1.11	n/a †	1.06	1.36	1.04	1.04
Unified TTFT	1.30	n/a †	0.99	1.25	1.28	1.01
Unified output throughput	1.17	n/a †	1.03	1.30	1.11	1.03

Workloads: decode-only input 256 / output 512, concurrency 8-256 (1-32 tokens per rank); prefill-only output 1, input 1k-8k; unified 1k/256, 2k/512, 4k/128 and 8k/128.

† Humming versions: H200 uses humming-kernels 0.1.18 (latest), H20 uses 0.1.12 (SGLang's current pin). On H200, 0.1.12 faults in prefill (its tile heuristic picks block_k=256, and the indexed FP8 group-128 GEMM hits a misaligned address) and its group-128 FP8 GEMMs run at about 17 TFLOPS; 0.1.18 fixes both, so the H200 Humming baseline was rerun with it. Each H200 Humming ratio is against this PR measured on the same node in the same session (prefill TTFT within 0.5%, decode ITL within 1.2% of the "This PR" column). The H200 unified workloads were not rerun with 0.1.18. On H20, 0.1.12 uses a tuned tile table that avoids the fault.

Per-point results behind the prefill-only and decode-only rows; values in parentheses are the speedup of this PR over that baseline (> 1 is better).

Prefill-only (output 1), median TTFT (ms):

GPU	Input (concurrency)	This PR	DeepEP + DeepGEMM FP8	DeepEP + Humming W4A8	SM90 FP8 MegaMoE ‡
H200	1k (64)	880	1,157 (1.32x)	1,220 (1.39x) †	907 (1.00x)
H200	2k (128)	2,445	3,324 (1.36x)	3,318 (1.39x) †	2,365 (0.96x)
H200	4k (32)	1,536	2,049 (1.33x)	2,156 (1.41x) †	1,545 (0.97x)
H200	8k (16)	1,715	2,289 (1.33x)	2,403 (1.45x) †	1,755 (1.00x)
H20	1k (64)	1,203	1,574 (1.31x)	1,618 (1.35x)	1,215 (1.00x)
H20	2k (128)	3,695	4,458 (1.21x)	4,666 (1.26x)	3,718 (1.00x)
H20	4k (32)	2,353	3,005 (1.28x)	3,070 (1.30x)	2,376 (1.01x)
H20	8k (16)	2,672	3,398 (1.27x)	3,479 (1.30x)	2,650 (1.01x)

‡ Measured in a separate session; its ratio is against this PR rerun in that session (TTFT within 4% of the "This PR" column).

Decode-only (input 256 / output 512), median ITL (ms):

GPU	Concurrency (tokens per rank)	This PR	DeepEP + DeepGEMM FP8	DeepEP + Humming W4A8	SM90 FP8 MegaMoE
H200	8 (1)	9.21	10.45 (1.13x)	9.72 (1.06x) †	9.71 (1.05x)
H200	16 (2)	10.32	11.71 (1.13x)	10.45 (1.02x) †	10.90 (1.06x)
H200	32 (4)	11.61	13.29 (1.14x)	11.25 (0.97x) †	12.50 (1.08x)
H200	64 (8)	13.02	15.32 (1.18x)	12.39 (0.96x) †	14.24 (1.09x)
H200	128 (16)	15.09	17.07 (1.13x)	13.79 (0.92x) †	15.98 (1.06x)
H200	256 (32)	17.26	18.61 (1.08x)	15.51 (0.91x) †	17.62 (1.02x)
H20	8 (1)	11.09	14.98 (1.35x)	12.61 (1.14x)	11.58 (1.04x)
H20	16 (2)	12.51	17.25 (1.38x)	13.70 (1.10x)	13.06 (1.04x)
H20	32 (4)	14.27	20.42 (1.43x)	15.16 (1.06x)	14.95 (1.05x)
H20	64 (8)	16.52	24.32 (1.47x)	17.12 (1.04x)	17.43 (1.06x)
H20	128 (16)	19.52	27.75 (1.42x)	19.87 (1.02x)	19.93 (1.02x)
H20	256 (32)	21.33	30.76 (1.44x)	22.98 (1.08x)	23.03 (1.08x)

Launch commands (all four backends)

Common flags for every backend:

COMMON="--trust-remote-code --tp 8 --dp 8 --enable-dp-attention --load-balance-method round_robin \
  --chunked-prefill-size 8192 --disable-shared-experts-fusion --disable-prefill-cuda-graph \
  --cuda-graph-max-bs-decode 256 --mem-fraction-static 0.70"
# unified workloads: --max-running-requests 256 --max-total-tokens 131072 (DeepEP dispatch cap 128)
# prefill-only / decode-only workloads: --max-running-requests 2048 --max-total-tokens 262144 (DeepEP dispatch cap 256)

FlashInfer SM90 NVFP4 MegaMoE (this PR):

SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=4096 \
python -m sglang.launch_server --model-path nvidia/DeepSeek-V4-Flash-NVFP4 $COMMON \
  --moe-a2a-backend megamoe --max-running-requests 256 --max-total-tokens 131072

DeepEP + DeepGEMM FP8 (W8A8; FP4 experts dequantized to FP8 at load):

SGLANG_DSV4_FP4_DEQUANT=1 SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128 \
python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V4-Flash $COMMON \
  --moe-a2a-backend deepep --deepep-mode auto --max-running-requests 256 --max-total-tokens 131072

DeepEP + Humming W4A8 (FP8 dispatch, MXFP4 weights):

# H200: pip install --no-deps "humming-kernels[cu13]==0.1.18" (H20: the image's 0.1.12)
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128 \
python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V4-Flash $COMMON \
  --moe-a2a-backend deepep --moe-runner-backend humming --deepep-mode auto \
  --max-running-requests 256 --max-total-tokens 131072

SM90 FP8 MegaMoE (cookbook FP8 checkpoint):

SGLANG_DSV4_FP4_EXPERTS=0 SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=4096 \
python -m sglang.launch_server --model-path sgl-project/DeepSeek-V4-Flash-FP8 $COMMON \
  --moe-a2a-backend megamoe --max-running-requests 256 --max-total-tokens 131072

Benchmark client (one workload; each runs 3 times):

python -m sglang.bench_serving --backend sglang --dataset-name random \
  --random-input-len 256 --random-output-len 512 --random-range-ratio 1 \
  --num-prompts 32 --max-concurrency 32 --warmup-requests 32 --flush-cache --seed 20260811

Workloads: decode-only input 256 / output 512, concurrency 8–256 (1–32 tokens per rank); prefill-only output 1,
input 1k–8k; unified 1k/256, 2k/512, 4k/128 and 8k/128.

The tables above were measured before the final rebase with the FP4 experts of deepseek-ai/DeepSeek-V4-Flash, which give FlashInfer byte-identical weights to nvidia/DeepSeek-V4-Flash-NVFP4. Rerunning five points on this branch with FlashInfer main da992ed gives the same numbers (this branch / tables: decode ITL C32 1.004, C256 0.992; prefill 8k input throughput 1.036; unified 1k/256 and 8k/128 output throughput 1.007 / 1.031).

The adapter adds one pre-dispatch kernel (about 2 µs per layer); no extra pad, copy or cast kernels run in the decode
graph.

Scope and Limitations

• SM90 only, opt-in through --moe-a2a-backend megamoe with ModelOpt NVFP4 experts. The integration is not DeepSeek-V4-specific: it uses the generic MegaMoE hook, and FlashInfer plans any legal geometry (tuned for DeepSeek-V4-Flash, DeepSeek-V4-Pro and MiMo-V2.5-Pro). It needs a model with a SwiGLU limit (the kernel's clamped SwiGLU) and has been validated only on nvidia/DeepSeek-V4-Flash-NVFP4.
• Against DeepEP + Humming W4A8 with humming-kernels 0.1.18 on H200, prefill is faster (1.41x TTFT) but decode above 4 tokens per rank is 3-9% slower.
• EPLB is not supported on this path.
• A coordinated FlashInfer dependency update is required before out-of-box use.

Checklist

• Format your code according to the Format code with pre-commit.
• Add tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

🤖 Generated with Claude Code

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit 232e96a.

Approve preview

[Fridge003]

(From GPT6)
[P1] Uneven DP prefill CUDA graph replay violates FlashInfer’s cross-rank contract
_rank_invariant_tokens() returns the local graph’s token count during capture. However, prefill_graph_tolerates_sum_len() still allows MegaMoE ranks to replay different prefill graph buckets. For example, ranks replaying 128-token and 512-token graphs would launch the fused and split kernel families respectively, while FlashInfer requires all ranks to use the same token count and matching kernels. This mismatch can cause collective hangs or incorrect execution.

Please enforce a shared prefill graph bucket across all EP ranks for this backend, or temporarily reject this configuration. Regression coverage should include uneven token counts, idle ranks, and buckets crossing the fused/split boundary.

[Fridge003]

Move this test to weekly

[ouershua2299]

Done, moved to the weekly stage (8-gpu-h200) in 232e96a.

[Fridge003]

Remove this test

[ouershua2299]

Removed in 232e96a.

[ouershua2299]

@Fridge003 Thanks for the review. We could not reproduce a hang or wrong output from uneven prefill graph replay on DeepSeek-V4-Flash, but the concern is valid against FlashInfer's contract (same token count on every EP rank), so we changed it as you suggested in 232e96a:

• While the FI NVFP4 MegaMoE path is captured into a prefill graph, it now latches the existing shared-bucket flag (the same one a captured DP gather sets, via a new note_shared_bucket_in_prefill_graph() in dp_attention.py). prefill_graph_tolerates_sum_len() then returns False, so every DP rank replays one MAX_LEN bucket, and batches with an idle rank fall back to eager, where this path already pads to the gathered max token count. Decode graphs already replay one shared bucket on all DP ranks, and the default config (prefill graphs disabled) is unchanged.

Test setup: 8x H200, branch rebased onto main ca4bedfffa, image nightly-dev-cu13-20261025-65633129, FlashInfer da992ed, nvidia/DeepSeek-V4-Flash-NVFP4, cookbook flags. The client sends concurrent GSM8K prompts padded to different lengths so DP ranks get uneven token counts in the same step: 1 long request with 7 idle ranks, 3 / 8 / 20 mixed lengths crossing the 256-token fused/split boundary, 120/600 and 150/450 tokens alternating one per rank, 7 short + 1 long, and 1 short request.

Prefill graph	Before (5e5c949)	After (232e96a)
disabled (default)	pass, 2 rounds	pass, 2 rounds, same answers
--cuda-graph-backend-prefill breakable	pass, 2 + 5 rounds, ranks replay per-rank buckets	pass, 2 + 5 rounds, all ranks log the shared MAX_LEN bucket
--cuda-graph-backend-prefill full	fails at startup	fails at startup

"pass" means no request timed out (no cross-rank hang) and greedy GSM8K answers stayed within noise (e.g., 150/450 alternating: 13-16/16 before, 14-16/16 after). With full, DeepSeek-V4's attention backend raises NotImplementedError: unsupported forward_mode=EXTEND during prefill graph capture, before the MoE runs, so breakable is the reachable prefill-graph path for this model. Our guess for why the old behavior did not fail: in the observed steps the ranks replaying graphs mostly landed above 256 tokens, where they select the same split plan even with different M.

[kabishou11]

Thanks for keeping the FlashInfer megamoe wiring explicit (source build with #6343/#6517, not PyPI 0.7.2).

Our offline DeepInfer Hopper fork stays on DeepEP V2 + humming for stock MXFP4 DeepSeek-V4.1-Flash and will not vendor flashinfer.megamoe. flashinfer#6999 has now merged on the separate moe_ep NVFP4×FP8 surface; we treat that as a watch item too, still without a CUDA drop. Once this PR merges we can re-test on H200 against a humming JSONL, only for a ModelOpt NVFP4 checkpoint if we actually have one.

Question: after #6999 landed, should this PR stay on flashinfer.megamoe (#6517), or do you intend to move Hopper NVFP4 MegaMoE over to moe_ep?