[Feat] [NVIDIA] moe: SM90 NVFP4 MegaMoE via FlashInfer

docs/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx

@@ -119,7 +119,7 @@ The Playground is where you experiment with **SGLang features beyond the verifie
 The knobs come in two flavors:
 
 - **Built-in SGLang features** — parallelism overrides (TP / CP / DP-Attention — DP-Attention's value is the DP degree, with `off` to disable), MoE backend + EP, reasoning / tool-call parsers, speculative-decoding presets, prefill/decode disaggregation, HiCache tiers, and HiSparse hierarchical sparse attention (decode-role only — the card appears once PD-Disagg mode is set to decode).
-- **DeepSeek-V4 specific features** — MegaMoE W4A8 / W4A4 fused kernel (Blackwell only; Hopper SM90 uses a separate all-FP8 MegaMoE path — see Configuration Tips below).
+- **DeepSeek-V4 specific features** — MegaMoE W4A8 / W4A4 fused kernel (Blackwell only; on Hopper SM90, MegaMoE runs either the all-FP8 path or the FlashInfer NVFP4 path — see Configuration Tips below).
 
 Lines highlighted **green** are added by your overrides; lines with **red strikethrough** were in the verified base but stripped by an override. When no override differs from the base cell, the playground inherits the base's **Verified** badge; any actual change flips it to **Not Verified** until the new configuration is run end-to-end and submitted back.
 
@@ -374,7 +374,7 @@ sglang serve \
   --tp 8
 ```
 
-Requires Blackwell (SM100+). The MTP layer in this checkpoint stays
+Requires Blackwell (SM100+), except the SM90 NVFP4 MegaMoE path in Configuration Tips below. The MTP layer in this checkpoint stays
 MXFP4-packed and is routed through the `Mxfp4FlashinferTrtllmMoEMethod` path
 automatically.
 
@@ -449,7 +449,7 @@ Two variants are exposed:
   Higher throughput with negligible accuracy drop (~89.5 GPQA on Pro).
 
 Notes:
-- The W4A8 / W4A4 variants above are **Blackwell-only** (B200 / B300 / GB200 / GB300). On **Hopper (SM90, H100 / H200)** use the all-FP8 MegaMoE path described below instead.
+- The W4A8 / W4A4 variants above are **Blackwell-only** (B200 / B300 / GB200 / GB300). On **Hopper (SM90, H100 / H200)** use the all-FP8 or the FlashInfer NVFP4 MegaMoE path described below instead.
 - MegaMoE is **only wired into the `high-throughput` recipe** on Blackwell (per [sgl-project/sglang#26451](https://github.com/sgl-project/sglang/pull/26451)). The chip is hidden on `low-latency` and `balanced` — switch to `high-throughput` to expose it.
 - When running MegaMoE, don't set `--moe-runner-backend` manually.
 - Adjust `SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK` based on your workload and memory usage. Setting higher number of tokens for MegaMoE requires more HBM space (recommended: 8320 for high-throughput).
@@ -479,6 +479,36 @@ the MegaMoE path processes per rank (i.e. per GPU); the MegaMoE path is only use
 for batches at or below this cap. The right value depends on your parallelism /
 token-split scheme, and larger values reserve more HBM.
 
+**SM90 (Hopper) NVFP4 MegaMoE via FlashInfer (Experimental)**
+
+With the NVFP4 hybrid checkpoints (see **NVFP4 Hybrid Checkpoints** above), `--moe-a2a-backend megamoe`
+on SM90 runs the FlashInfer SM90 NVFP4 MegaMoE kernel (`flashinfer.megamoe`, W4A8: NVFP4
+weights, FP8 activations); the expert tensors are repacked for the kernel at load time.
+It needs a FlashInfer build that contains
+[flashinfer-ai/flashinfer#6343](https://github.com/flashinfer-ai/flashinfer/pull/6343) and
+[flashinfer-ai/flashinfer#6517](https://github.com/flashinfer-ai/flashinfer/pull/6517); until a FlashInfer
+release includes them, install the validated `main` revision from source (the kernels are JIT-compiled):
+
+```bash Command
+pip install --no-build-isolation --no-deps \
+  "flashinfer-python @ git+https://github.com/flashinfer-ai/flashinfer@da992ed9919604c5eb6dc255c63f5e89a466452c"
+```
+
+Verified end to end with `nvidia/DeepSeek-V4-Flash-NVFP4` on 8x H200; EPLB is not supported on this path.
+
+```bash Command
+SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=4096 \
+sglang serve \
+  --model-path nvidia/DeepSeek-V4-Flash-NVFP4 \
+  --tp 8 --dp 8 --enable-dp-attention \
+  --moe-a2a-backend megamoe \
+  --disable-shared-experts-fusion \
+  --chunked-prefill-size 8192
+```
+
+All fused and split kernel variants are JIT-compiled at startup, sharded across ranks
+(`SGLANG_FLASHINFER_MEGAMOE_SM90_PREWARM=0` defers compilation to first use).
+
 **GB300 PD-Disagg cross-pod MNNVL**
 
 On some GB300 clusters with cross-pod KV transfer over NVLink, mooncake may

docs/docs/references/environment_variables.mdx

@@ -361,6 +361,16 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Permit CuTe DSL MegaMOE in-kernel FC2 reduction for NVFP4 W4A4, NVFP4 W4A16, and MXFP8. FlashInfer's tuning profile may still select external reduction. BF16 atomic accumulation is nondeterministic; leave disabled for reproducibility. Requires BF16 combine.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>"false"</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_FLASHINFER_MEGAMOE_SM90_PREWARM</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>SM90 FlashInfer NVFP4 MegaMoE (<code>--moe-a2a-backend megamoe</code> with NVFP4 checkpoints): JIT-compile every fused / split kernel variant reachable under the per-rank token cap at startup, sharded across ranks. Set to <code>false</code> to compile on first use.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>"true"</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_FLASHINFER_MEGAMOE_SM90_KERNEL_FAMILY</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>SM90 FlashInfer NVFP4 MegaMoE kernel family. <code>auto</code> uses the fused kernel for at most 256 tokens per rank and the split kernel above; <code>fused</code> and <code>split</code> force one family.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>"auto"</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Number of SMs used for DeepEP combine when single batch overlap is enabled</td>

python/sglang/srt/environ.py

@@ -1219,6 +1219,11 @@ class Envs:
     # traffic with a small accuracy tradeoff and require FC2 reduce outside the
     # kernel.
     SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE = EnvStr("bf16")
+    # SM90 NVFP4 MegaMoE (flashinfer.megamoe) under --moe-a2a-backend megamoe.
+    # Builds every fused / split JIT module reachable for the token cap at startup.
+    SGLANG_FLASHINFER_MEGAMOE_SM90_PREWARM = EnvBool(True)
+    # "auto" (fused for M <= 256 tokens per rank, split above), "fused", or "split".
+    SGLANG_FLASHINFER_MEGAMOE_SM90_KERNEL_FAMILY = EnvStr("auto")
     # Enable per-token FP32 activation scaling for serialized ModelOpt FP4 with
     # FlashInfer TRT-LLM or CuTe DSL v2 MoE.
     SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION = EnvBool(False)

python/sglang/srt/layers/dp_attention.py

@@ -978,6 +978,11 @@ def _note_dp_gather_in_prefill_graph() -> None:
         dp.prefill_graph_has_dp_gather = True
 
 
+def note_shared_bucket_in_prefill_graph() -> None:
+    """For collectives that need the same token count on every DP rank."""
+    _note_dp_gather_in_prefill_graph()
+
+
 def note_rank_barrier_in_prefill_graph() -> None:
     """For breaks that capture graphs of their own behind TP-group barriers."""
     dp = get_flags().dp

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -414,6 +414,7 @@ def __init__(
         # Set by the quant method when it repacks experts for MegaMoE.
         self._mega_moe_weights_built = False
         self._mega_moe_nvfp4 = False
+        self._mega_moe_sm90_fi_nvfp4 = False
         # Read at construction: draft_model_build_scope applies the draft's own value.
         self._mega_moe_w4a4 = get_exec().moe.enable_w4a4_mxfp4_megamoe
         self._pending_fp8_shared_weights: dict[tuple[int, str], torch.Tensor] = {}

python/sglang/srt/layers/moe/mega_moe.py

@@ -162,7 +162,13 @@ def is_mega_moe_experts_ready(experts) -> bool:
     if not experts._mega_moe_weights_built:
         return False
     if _device_sm == 90:
-        return is_sm90_fp8_mega_moe_available(experts)
+        from sglang.srt.layers.moe.mega_moe_sm90_fi_nvfp4 import (
+            is_sm90_fi_nvfp4_mega_moe_available,
+        )
+
+        return is_sm90_fp8_mega_moe_available(experts) or (
+            is_sm90_fi_nvfp4_mega_moe_available(experts)
+        )
     # The SM100 mega kernels exist for compute capability 10.x only.
     return _device_sm // 10 == 10
 
@@ -444,6 +450,31 @@ def run_mega_routed_experts(
         f"--cuda-graph-max-bs-decode / --chunked-prefill-size accordingly"
     )
 
+    if _device_sm == 90 and experts._mega_moe_sm90_fi_nvfp4:
+        from sglang.srt.layers.moe.mega_moe_sm90_fi_nvfp4 import (
+            run_sm90_fi_nvfp4_mega_routed,
+        )
+
+        if num_tokens > 0:
+            ids, wts = topk_ids, topk_weights
+        else:
+            ids = hidden_states.new_empty((0, top_k), dtype=torch.int32)
+            wts = hidden_states.new_empty((0, top_k), dtype=torch.float32)
+        return run_sm90_fi_nvfp4_mega_routed(
+            experts=experts,
+            hidden_states=hidden_states,
+            topk_ids=ids,
+            topk_weights=wts,
+            num_tokens=num_tokens,
+            ep_group=ep_group,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            top_k=top_k,
+            activation_clamp=activation_clamp,
+            routed_scaling_factor=routed_scaling_factor,
+        )
+
     mma_type = _mega_moe_mma_type(experts)
     buf = _get_mega_moe_symm_buffer(
         ep_group,