[Experiment][MoE] H200 NVFP4 MegaMoE via flashinfer.megamoe

.buildkite/test_areas/distributed.yaml

@@ -428,12 +428,16 @@ steps:
   - vllm/v1/executor/uniproc_executor.py
   - vllm/v1/worker/gpu_worker.py
   - tests/distributed/test_mnnvl_alltoall.py
+  - tests/kernels/moe/test_flashinfer_moe_ep.py
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
 
   commands:
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
   # test with torchrun tp=2 and dp=4 with ep
   - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
+  - pytest -v -s kernels/moe/test_flashinfer_moe_ep.py -k sm90_nvfp4
   mirror:
     amd:
       label: ":amd: (MI355) Distributed DP + EP"

.buildkite/test_areas/kernels.yaml

@@ -906,12 +906,14 @@ steps:
   - csrc/quantization/w8a8/cutlass/moe/
   - csrc/moe/
   - tests/kernels/moe
+  - tests/models/test_deepseek_v4_fi_moe_ep.py
   - vllm/model_executor/layers/fused_moe/
   - vllm/model_executor/layers/quantization/
   - vllm/distributed/device_communicators/
   - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
+    - pytest -v -s models/test_deepseek_v4_fi_moe_ep.py -k sm90_nvfp4
   mirror:
     amd:
       label: ":amd: (MI355) FusedMoE Layer Kernels"

docs/serving/expert_parallel_deployment.md

@@ -27,7 +27,7 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
 | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
 | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
 | `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
-| `passthrough` | Any | Internal, not accepted on the CLI: bound automatically for MoE backends that dispatch and combine themselves, so no all-to-all runs in the framework | `--moe-backend flashinfer_moe_ep_*` (FlashInfer MoE-EP megakernels) |
+| `passthrough` | Any | Internal, not accepted on the CLI: bound automatically for MoE backends that dispatch and combine themselves; excluded unquantized layers retain framework collectives | `--moe-backend flashinfer_moe_ep_*` (FlashInfer MoE-EP megakernels) |
 | `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
 
 ## Single Node Deployment
@@ -80,6 +80,25 @@ vllm serve deepseek-ai/DeepSeek-V3-0324 \
     --enable-expert-parallel         # Enable expert parallelism
 ```
 
+### Hopper NVFP4 MegaMoE
+
+On H20/H200 with four or eight NVLink-connected GPUs, ModelOpt NVFP4 checkpoints
+with a SwiGLU clamp (`swiglu_limit`) can run FlashInfer's SM90 MegaMoE (NVFP4
+weights, FP8 activations) with `--moe-backend flashinfer_moe_ep_cutedsl`. It needs
+a FlashInfer build that provides [`flashinfer.megamoe`](https://github.com/flashinfer-ai/flashinfer/tree/main/flashinfer/megamoe) and does not support EPLB,
+microbatching/dual batch overlap or runtime weight transfer. Unquantized MoE layers, such as
+BF16 MTP experts, keep the default backend.
+Gate/up must share each expert's global weight scale; local scales must remain
+exactly representable during weight preparation.
+This backend is opt-in: `auto` keeps its existing selection, and unsupported
+explicit configurations fail at startup.
+
+```bash
+vllm serve nvidia/GLM-5.3-Flash-NVFP4 \
+    --data-parallel-size 8 --enable-expert-parallel \
+    --moe-backend flashinfer_moe_ep_cutedsl
+```
+
 ## Multi-Node Deployment
 
 For multi-node deployment, use the DeepEP communication kernel with one of two modes (see [Backend Selection Guide](#backend-selection-guide) above).

tests/kernels/moe/test_flashinfer_moe_ep.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end tests for the FlashInfer MoE-EP adapter on real Blackwell GPUs.
+"""End-to-end tests for the FlashInfer MoE-EP adapters on Hopper and Blackwell.
 
 Each test spawns one process per GPU, bootstraps vLLM's expert-parallel
 groups, builds a ``FlashInferMoeEp`` adapter from bf16 expert weights at the
@@ -367,3 +367,150 @@ def destroy(self) -> None:
         else moe_ep.DeepGemmMegaMoeConfig
     )
     assert isinstance(mega_config.megakernel, expected)
+
+
+@torch.inference_mode()
+def _sm90_nvfp4_mixed_moe(env: dict[str, str], world: int, tp_size: int) -> None:
+    from vllm.forward_context import get_forward_context, set_forward_context
+    from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config
+    from vllm.v1.worker.workspace import init_workspace_manager
+
+    set_env_vars_and_device(env)
+    rank = torch.distributed.get_rank()
+    device = torch.device("cuda", rank)
+    dp_size = world // tp_size
+    dp_rank = rank // tp_size
+    cfg = VllmConfig()
+    cfg.parallel_config.tensor_parallel_size = tp_size
+    cfg.parallel_config.data_parallel_size = dp_size
+    cfg.parallel_config.data_parallel_rank = dp_rank
+    cfg.parallel_config.enable_expert_parallel = True
+    cfg.parallel_config.all2all_backend = "passthrough"
+    cfg.kernel_config.moe_backend = "flashinfer_moe_ep_cutedsl"
+    cfg.scheduler_config.max_num_batched_tokens = 512 * tp_size
+    init_workspace_manager(device)
+    hidden, intermediate, experts, topk = 4096, 2048, 288, 8
+    with set_current_vllm_config(cfg):
+        ensure_model_parallel_initialized(tp_size, 1)
+        for quantized in (True, False):
+            with torch.device(device):
+                layer = FusedMoEFactory(
+                    num_experts=experts,
+                    top_k=topk,
+                    hidden_size=hidden,
+                    intermediate_size=intermediate,
+                    params_dtype=torch.bfloat16,
+                    quant_config=(
+                        ModelOptNvFp4Config(is_checkpoint_nvfp4_serialized=True)
+                        if quantized
+                        else None
+                    ),
+                    swiglu_limit=10.0,
+                    is_sequence_parallel=tp_size > 1,
+                    prefix=f"operator_{quantized}",
+                )
+            routed = layer.routed_experts
+            local = experts // world
+            routed.w13_weight.zero_()
+            routed.w2_weight.zero_()
+            for i in range(local):
+                e = rank * local + i
+                w1 = 2.0 / (1 + e % 2)
+                w2 = 1.75 * (1 + e % 3)
+                routed.w13_weight[i, 0, 0] = 4 if quantized else w1
+                routed.w13_weight[i, intermediate, 0] = 4 if quantized else w1
+                routed.w2_weight[i, e, 0] = 4 if quantized else w2
+            if quantized:
+                routed.w13_weight_scale.fill_(128)
+                routed.w2_weight_scale.fill_(224)
+                for i in range(local):
+                    e = rank * local + i
+                    routed.w13_weight_scale_2[i].fill_(1 / (128 * (1 + e % 2)))
+                    routed.w2_weight_scale_2[i].fill_((1 + e % 3) / 256)
+                routed.w13_input_scale.fill_(1)
+                routed.w2_input_scale.fill_(1)
+            routed.quant_method.process_weights_after_loading(routed)
+
+            def case(routed, lengths, graph=False, quantized=quantized):
+                n = (lengths[dp_rank] + tp_size - 1) // tp_size
+                x = torch.zeros(n, hidden, dtype=torch.bfloat16, device=device)
+                x[:, 0] = torch.tensor([2.0, 4.0, 8.0], device=device)[
+                    torch.arange(n, device=device) % 3
+                ]
+                x[:, 1] = 7.0
+                ids = (
+                    (
+                        (torch.arange(n, device=device) + rank * 7)[:, None]
+                        + torch.arange(topk, device=device)[None, :] * (experts // topk)
+                    )
+                    .remainder(experts)
+                    .to(torch.int32)
+                )
+                weights = (
+                    torch.tensor([2.0 ** -(i + 1) for i in range(topk)], device=device)
+                    .expand(n, -1)
+                    .contiguous()
+                )
+                expected = torch.zeros(n, hidden, device=device)
+                for j in range(topk):
+                    e = ids[:, j].long()
+                    g = (x[:, 0].float() * 2 / (1 + e % 2)).clamp(max=10)
+                    value = torch.nn.functional.silu(g) * g
+                    value *= weights[:, j]
+                    if quantized:
+                        # FI's intermediate E4M3 quantizer uses UE8M0 scales.
+                        scale = torch.exp2(torch.ceil(torch.log2(value.abs() / 448)))
+                        value = (value / scale).to(torch.float8_e4m3fn).float() * scale
+                    value *= 1.75 * (1 + e % 3)
+                    expected.scatter_add_(1, e[:, None], value[:, None])
+                with set_forward_context(
+                    None,
+                    cfg,
+                    num_tokens=lengths[dp_rank],
+                    num_tokens_across_dp=torch.tensor(lengths, dtype=torch.int32),
+                ):
+                    metadata = get_forward_context().dp_metadata
+                    assert metadata is not None
+                    with metadata.sp_local_sizes(tp_size):
+                        y = routed.forward_modular(x, weights, ids)
+                        torch.accelerator.synchronize()
+                        if graph:
+                            torch.distributed.barrier()
+                            captured = torch.cuda.CUDAGraph()
+                            with torch.cuda.graph(captured):
+                                y = routed.forward_modular(x, weights, ids)
+                            captured.replay()
+                            captured.replay()
+                            torch.accelerator.synchronize()
+                assert torch.isfinite(y).all()
+                torch.testing.assert_close(
+                    y, expected.to(torch.bfloat16), rtol=0.008, atol=0.001
+                )
+                torch.distributed.barrier()
+                if rank == 0:
+                    print(
+                        f"PASS NVFP4={quantized} lengths={lengths} graph={graph}",
+                        flush=True,
+                    )
+
+            for n in (1, 16, 64, 257):
+                case(routed, [n * tp_size] * dp_size)
+            case(routed, [0] + [3 * i * tp_size for i in range(1, dp_size)])
+            case(routed, [16 * tp_size] * dp_size, graph=True)
+            case(routed, [257 * tp_size] * dp_size, graph=True)
+            del routed, layer
+            torch.accelerator.memory.empty_cache()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability(90), reason="Requires SM90"
+)
+@pytest.mark.parametrize("world_size,tp_size", [(4, 1), (8, 1), (4, 2)])
+def test_sm90_nvfp4_and_bf16_mtp_experts(world_size: int, tp_size: int):
+    if torch.cuda.get_device_properties().multi_processor_count not in (78, 132):
+        pytest.skip("FlashInfer SM90 NVFP4 requires 78 or 132 SMs")
+    if torch.accelerator.device_count() < world_size:
+        pytest.skip(f"Requires {world_size} NVLink-connected Hopper GPUs")
+    pytest.importorskip("flashinfer.megamoe")
+    distributed_run(_sm90_nvfp4_mixed_moe, world_size, tp_size)

tests/models/test_deepseek_v4_fi_moe_ep.py

@@ -66,6 +66,91 @@
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a Hopper GPU")
+def test_sm90_nvfp4_supports_the_glm_geometry():
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_device_capability(90):
+        pytest.skip("Requires a Hopper GPU")
+    moe = make_dummy_moe_config(
+        num_experts=288,
+        num_local_experts=36,
+        experts_per_token=8,
+        hidden_dim=4096,
+        intermediate_size=2048,
+    )
+    moe.moe_backend = FLASHINFER_MOE_EP_CUTEDSL
+    moe.swiglu_limit = 10.0
+    moe.moe_parallel_config = replace(
+        moe.moe_parallel_config,
+        use_ep=True,
+        ep_size=8,
+        dp_size=8,
+        all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND,
+    )
+    cfg = VllmConfig()
+    with set_current_vllm_config(cfg):
+        if torch.cuda.get_device_properties().multi_processor_count not in (78, 132):
+            assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0]
+            return
+        assert _is_supported(moe, kNvfp4Static, kNvfp4Dynamic) == (True, "")
+        cfg.parallel_config.ubatch_size = 2
+        assert not cfg.parallel_config.enable_dbo
+        assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0]
+        cfg.parallel_config.ubatch_size = 1
+        assert not _is_supported(moe, kMxfp4Static, None)[0]
+        moe.moe_parallel_config.ep_size = 2
+        assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0]
+        moe.moe_parallel_config.ep_size = 8
+        moe.hidden_dim = 4352
+        assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0]
+
+
+def test_sm90_nvfp4_marlin_byte_order():
+    from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep_sm90_nvfp4 import (
+        nvfp4_to_marlin,
+    )
+
+    packed = torch.tensor(
+        [0x10, 0x32, 0x54, 0x76, 0x98, 0xBA, 0xDC, 0xFE], dtype=torch.uint8
+    ).view(1, 1, 8)
+    expected = torch.tensor(
+        [0x04, 0x15, 0x26, 0x37, 0x8C, 0x9D, 0xAE, 0xBF], dtype=torch.uint8
+    ).view(1, 1, 8)
+    assert torch.equal(nvfp4_to_marlin(packed), expected)
+
+
+@pytest.mark.parametrize("as_bytes", [False, True])
+def test_sm90_nvfp4_scale_rebalance_preserves_dequantization(as_bytes):
+    from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep_sm90_nvfp4 import (
+        rebalance_nvfp4_scales,
+    )
+
+    scales = torch.tensor([0.5, 32, 224, 448], dtype=torch.float8_e4m3fn)
+    scales = scales.view(1, 2, 2).repeat(2, 1, 1)
+    global_scales = torch.tensor([0.5, 2.0])
+    encoded, adjusted = rebalance_nvfp4_scales(
+        scales.view(torch.uint8) if as_bytes else scales, global_scales
+    )
+    assert encoded.dtype == torch.uint8
+    assert encoded.view(torch.float8_e4m3fn).float().amax() <= 448 / 6
+    assert torch.equal(
+        encoded.view(torch.float8_e4m3fn).float() * adjusted[:, None, None],
+        scales.float() * global_scales[:, None, None],
+    )
+
+
+def test_sm90_nvfp4_rejects_lossy_scale_rebalance():
+    from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep_sm90_nvfp4 import (
+        rebalance_nvfp4_scales,
+    )
+
+    scales = torch.tensor([2**-9, 448], dtype=torch.float8_e4m3fn).view(1, 1, 2)
+    with pytest.raises(ValueError, match="exactly"):
+        rebalance_nvfp4_scales(scales, torch.ones(1))
+
+
 def test_flashinfer_backends_are_megakernels_outside_the_native_model_path():
     assert FLASHINFER_MOE_EP_BACKENDS <= MEGA_MOE_BACKENDS
     assert FLASHINFER_MOE_EP_BACKENDS.isdisjoint(NATIVE_MEGA_MOE_BACKENDS)
@@ -169,15 +254,17 @@ def _is_supported(moe, weight_key, activation_key) -> tuple[bool, str]:
 
 
 @pytest.fixture
-def megakernel_host(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
-    """A vLLM config without DBO/EPLB/weight transfer, on a supported device."""
-    config = SimpleNamespace(
-        weight_transfer_config=None,
-        parallel_config=SimpleNamespace(enable_dbo=False, enable_eplb=False),
-    )
+def megakernel_host(monkeypatch: pytest.MonkeyPatch):
+    """A Blackwell host without DBO/EPLB/weight transfer."""
+    from vllm.config import VllmConfig
+    from vllm.platforms.interface import DeviceCapability
+
+    config = VllmConfig()
     monkeypatch.setattr(fi_ep, "get_current_vllm_config", lambda: config)
     monkeypatch.setattr(
-        FlashInferMoeEpExperts, "_supports_current_device", staticmethod(lambda: True)
+        fi_ep.current_platform,
+        "get_device_capability",
+        lambda *args, **kwargs: DeviceCapability(10, 0),
     )
     return config
 
@@ -347,6 +434,39 @@ def test_passthrough_all2all_backend_selects_the_pass_through_stages():
     assert isinstance(prepare_finalize, PassThroughPrepareAndFinalize)
 
 
+def test_passthrough_keeps_dispatch_for_unquantized_mtp_experts():
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import (
+        MoEPrepareAndFinalizeNaiveDPEPModular,
+    )
+
+    moe = make_dummy_moe_config()
+    moe.moe_parallel_config = replace(
+        moe.moe_parallel_config,
+        use_ep=True,
+        ep_size=8,
+        dp_size=8,
+        all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND,
+    )
+    stages = maybe_make_prepare_finalize(moe, FusedMoEQuantConfig.make())
+    assert isinstance(stages, MoEPrepareAndFinalizeNaiveDPEPModular)
+
+
+def test_passthrough_tp_only_bf16_does_not_dispatch_across_dp():
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import (
+        MoEPrepareAndFinalizeNoDPEPModular,
+    )
+
+    moe = make_dummy_moe_config()
+    moe.moe_parallel_config = replace(
+        moe.moe_parallel_config,
+        use_ep=True,
+        ep_size=4,
+        all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND,
+    )
+    stages = maybe_make_prepare_finalize(moe, FusedMoEQuantConfig.make())
+    assert isinstance(stages, MoEPrepareAndFinalizeNoDPEPModular)
+
+
 def test_passthrough_all2all_backend_is_not_a_cli_choice():
     """Users cannot pick the passthrough all2all backend explicitly."""
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())

vllm/config/kernel.py

@@ -312,9 +312,10 @@ class KernelConfig:
     - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
     - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
     - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
-    - "flashinfer_moe_ep_cutedsl": Use FlashInfer's CuTeDSL MoE-EP
+    - "flashinfer_moe_ep_cutedsl": Use FlashInfer's MoE-EP
       mega-kernel with NVFP4 weights (MXFP4 checkpoints are requantized at
-      load); requires Blackwell, expert parallelism, and NVSHMEM
+      load on Blackwell). Requires expert parallelism; Blackwell uses NVSHMEM,
+      Hopper NVFP4 uses Torch symmetric memory on four or eight NVLink GPUs
     - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x
       (RTX Pro 6000 / DGX Spark)
     - "b12x": Use b12x FP4 MoE kernels on SM12x