[MoE] Add FlashInfer SM120 W4A8 MegaMoE backend for DeepSeek-V4-Flash

[Euynaheh]

Summary

This PR adds an opt-in flashinfer_megamoe runner backend for DeepSeek-V4-Flash with MXFP4 expert weights and MXFP8 activations on SM120 GPUs.

The backend calls FlashInfer's public flashinfer.moe_ep interface and replaces the routed MoE sequence

dispatch -> FC1 -> FC2 -> combine

with a FlashInfer MoEEpLayer.

The initial supported configuration is:

--tp-size 4 \
--dp-size 4 \
--enable-dp-attention \
--ep-size 4 \
--moe-runner-backend flashinfer_megamoe \
--flashinfer-megamoe-max-num-tokens 8192

With DP attention enabled, each rank owns a different DP token stream when entering the MoE layers. The adapter communicates the maximum active token count across ranks to FlashInfer and supports ranks with zero local tokens.

Existing MoE backends are unchanged unless flashinfer_megamoe is selected explicitly.

Implementation

This PR:

• Adds flashinfer_megamoe to the MoE runner backend choices.
• Adds --flashinfer-megamoe-max-num-tokens for the symmetric workspace capacity.
• Constructs FlashInfer PrequantizedMoEWeights, BootstrapConfig, FleetParams, and MoEEpLayer objects from SGLang-loaded checkpoint weights.
• Preserves the canonical [gate; up] W13 layout expected by the MegaMoE preprocessor.
• Uses zero-copy byte dtype views for MXFP4 weights and UE8M0 scales.
• Registers CUDA Graph-compatible staging and compute custom operations.
• Reuses FlashInfer's stable output buffer.
• Drops loader-side expert tensors after FlashInfer has created its preprocessed weight pack, avoiding retention of a second expert-weight copy.
• Handles different DP token counts and empty ranks.
• Computes the shared expert separately and disables shared-expert fusion for this backend.
• Applies the routed scaling factor consistently with SGLang's existing MoE path.
• Validates the per-rank token capacity before execution.

Validation and unsupported combinations

The argument resolver rejects configurations that are not currently supported:

• --enable-two-batch-overlap
• External MoE dispatch backends such as DeepEP or MORI
• EPLB
• Redundant or replicated experts
• Attention TP or MoE TP greater than 1 after parallelism groups are derived
• Uneven expert placement

The token capacity must be large enough for the maximum number of active tokens on any EP rank. The backend reports an explicit error when the configured capacity is exceeded.

All benchmark requests completed without crashes, OOMs, watchdog failures, CUDA/NCCL errors, or HTTP 5xx responses.

Tests

python3 -m pytest -q \
  test/registered/unit/layers/moe/test_flashinfer_megamoe_adapter.py
pre-commit run ruff --files <changed-files>

pre-commit run ruff-format --files <changed-files>

Result:

6 passed

ruff passed

ruff-format passed

The unit tests cover:

• Argument-resolution defaults
• External dispatch rejection
• EPLB rejection
• TBO rejection
• Zero-copy MXFP4 dtype views

FlashInfer dependency

This PR requires a FlashInfer version that provides the SM120 MXFP4-weight × MXFP8-activation flashinfer.moe_ep backend.

The prerequisite FlashInfer infrastructure PR has been merged:

• feat(moe_ep): SM120 MXFP8 swap-AB CuTeDSL MegaMoE kernel flashinfer-ai/flashinfer#4387

We're working on the CI of MXFP4 x MXFP8 MegaMoE in FlashInfer. Will be merged soon:

• [SM120] Add W4A8 CuTeDSL MegaMoE backend flashinfer-ai/flashinfer#4632

Current limitations

• SM120 only
• MXFP4 expert weights with MXFP8 activations
• Even, non-replicated expert placement
• No TBO
• No EPLB
• No external MoE A2A backend
• Shared-expert fusion is disabled
• The configured per-rank token capacity is fixed when the FlashInfer layer is constructed

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.
Latest PR Test (AMD ROCm 7.2): ➖ No AMD PR run found for this commit.