[feat] FlashInfer CuteDSL MegaMoE integration (#54049)

Signed-off-by: jdebache <jdebache@nvidia.com>

docs/design/moe_kernel_features.md

@@ -37,6 +37,7 @@ th {
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
 | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
 | flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
+| passthrough | standard | nvfp4,mxfp4 | G | N | Y | [`PassThroughPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough.PassThroughPrepareAndFinalize] |
 
 !!! info "Table key"
     1. All types: mxfp4, nvfp4, int4, int8, fp8

docs/serving/expert_parallel_deployment.md

@@ -27,6 +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) |
 | `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
 
 ## Single Node Deployment

tests/distributed/test_engram_dp_shard.py

@@ -511,7 +511,7 @@ def _check_dummy_hash_model_forward(
 
     model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model)
     torch.nn.Module.__init__(model)
-    model.use_mega_moe = False
+    model.use_native_mega_moe = False
     model.use_sequence_parallel = False
     model.fuse_mhc_all_reduce = False
     model.engram_hash = state

tests/kernels/mhc/test_mhc_kernels.py

@@ -167,7 +167,7 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):
 
     monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams)
     draft = SimpleNamespace(
-        use_mega_moe=False,
+        use_native_mega_moe=False,
         use_sequence_parallel=False,
         hc_mult=hc_mult,
         layers=[make_layer(mix) for mix in mixes],

tests/kernels/moe/test_flashinfer_moe_ep.py

@@ -0,0 +1,307 @@
+# 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.
+
+Each test spawns one process per GPU, bootstraps vLLM's expert-parallel
+groups, builds a ``FlashInferMoeEp`` adapter from bf16 expert weights at the
+MoE geometry of a real model, warms it up and runs a forward whose inputs are
+chosen so that every quantization step in the megakernel is exact, for the
+CuTeDSL kernel (NVFP4 activations and weights, E4M3 scales per 16) as well as
+the DeepGEMM kernel (FP8 activations and FP4 weights, power-of-two scales per
+32). The expected output is then known in closed form and compared bit for
+bit.
+"""
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from tests.distributed.eplb_utils import distributed_run, set_env_vars_and_device
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM
+from vllm.distributed.parallel_state import ensure_model_parallel_initialized
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    RoutingMethodType,
+)
+from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import (
+    FlashInferMoeEpExperts,
+)
+from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import (
+    FlashInferMoeEp,
+    FlashInferMoeEpWeights,
+    apply_topk_in_fc1,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Static
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_deep_gemm
+
+
+@dataclass(frozen=True)
+class Geometry:
+    name: str
+    hidden_size: int
+    intermediate_size: int
+    num_experts: int
+    top_k: int
+
+
+# Routed-expert geometry of the models the branch was evaluated on. DeepGEMM
+# keeps one UE8M0 scale byte per 32 elements and requires the per-token scale
+# rows to be TMA aligned, so hidden and intermediate sizes must be multiples
+# of 512 (FlashInfer's own validation only checks 128); all four qualify.
+GEOMETRIES = (
+    Geometry("dsv4-pro", 7168, 3072, 384, 6),
+    Geometry("dsv4-flash", 4096, 2048, 256, 6),
+    Geometry("mistral-small-4", 4096, 2048, 128, 4),
+    Geometry("mistral-large-3", 7168, 4096, 128, 4),
+)
+BACKENDS = (FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM)
+
+NUM_TOKENS_PER_RANK = 4
+MAX_TOKENS_PER_RANK = 128
+
+# Every block scale below is either 1 or a power of two, so the FP4 codes
+# and the FP8 values survive quantization exactly: NVFP4 scales a 16-block by
+# absmax / 6 in E4M3, DeepGEMM scales a 32-block by absmax / 6 (FP4) or
+# absmax / 448 (FP8) rounded up to a power of two. A 6.0 in a block pins the
+# FP4 scale to 1; 2.0, 6.0 and 144 are exact FP8 values at any such scale.
+#
+# Token t carries X_SIGNAL at hidden column 16 * t and a 6.0 sentinel next to
+# it; every weight column the sentinel meets is zero.
+X_SIGNAL = 2.0
+X_SENTINEL = 6.0
+# Gate and up rows are identical, so h = silu(12) * 12 = 144 = 6 * 24
+# whichever half the kernel treats as the gate, and 144 requantizes exactly.
+W13_DIAG = 6.0
+H_SIGNAL = 144.0
+# Expert e writes E2M1_CODES[e % 7] * H_SIGNAL into output channel e. The 6.0
+# sentinel sits next to a zero intermediate channel and pins the block scale.
+W2_SENTINEL = 6.0
+E2M1_CODES = (0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+
+
+def _expert_out(expert: int) -> float:
+    return E2M1_CODES[expert % len(E2M1_CODES)]
+
+
+def _topk_weights(top_k: int) -> tuple[float, ...]:
+    # The DeepSeek-V4 routing policy applies the router weight before fc2, so
+    # weight * 144 must requantize exactly too: powers of two do.
+    return tuple(2.0 ** -(k + 1) for k in range(top_k))
+
+
+def _routing_method(backend: str) -> RoutingMethodType:
+    if backend == FLASHINFER_MOE_EP_DEEP_GEMM:
+        return RoutingMethodType.DeepseekV4
+    return RoutingMethodType.Default
+
+
+def _routing(
+    geometry: Geometry, rank: int, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """top_k distinct experts per token, strided across ranks so dispatch goes
+    remote."""
+    base = (
+        torch.arange(NUM_TOKENS_PER_RANK, device=device) + rank * NUM_TOKENS_PER_RANK
+    ) % geometry.num_experts
+    stride = geometry.num_experts // geometry.top_k
+    offsets = torch.arange(geometry.top_k, device=device) * stride
+    topk_ids = ((base[:, None] + offsets[None, :]) % geometry.num_experts).to(
+        torch.int32
+    )
+    topk_weights = torch.tensor(
+        [_topk_weights(geometry.top_k)] * NUM_TOKENS_PER_RANK,
+        dtype=torch.float32,
+        device=device,
+    )
+    return topk_ids, topk_weights
+
+
+def _hidden_states(geometry: Geometry, device: torch.device) -> torch.Tensor:
+    x = torch.zeros(
+        NUM_TOKENS_PER_RANK, geometry.hidden_size, dtype=torch.bfloat16, device=device
+    )
+    for t in range(NUM_TOKENS_PER_RANK):
+        x[t, 16 * t] = X_SIGNAL
+        x[t, 16 * t + 1] = X_SENTINEL
+    return x
+
+
+def _local_expert_weights(
+    geometry: Geometry, rank: int, num_local_experts: int, device: torch.device
+) -> FlashInferMoeEpWeights:
+    w13 = torch.zeros(
+        num_local_experts,
+        2 * geometry.intermediate_size,
+        geometry.hidden_size,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    w2 = torch.zeros(
+        num_local_experts,
+        geometry.hidden_size,
+        geometry.intermediate_size,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    token_columns = [16 * t for t in range(NUM_TOKENS_PER_RANK)]
+    for local in range(num_local_experts):
+        expert = rank * num_local_experts + local
+        w13[local, 0, token_columns] = W13_DIAG
+        w13[local, geometry.intermediate_size, token_columns] = W13_DIAG
+        w2[local, expert, 0] = _expert_out(expert)
+        w2[local, expert, 2] = W2_SENTINEL
+    return FlashInferMoeEpWeights(w13=w13, w2=w2)
+
+
+def _expected_output(
+    geometry: Geometry, topk_ids: torch.Tensor, topk_weights: torch.Tensor
+) -> torch.Tensor:
+    expected = torch.zeros(
+        NUM_TOKENS_PER_RANK,
+        geometry.hidden_size,
+        dtype=torch.float32,
+        device=topk_ids.device,
+    )
+    for t in range(NUM_TOKENS_PER_RANK):
+        for k in range(geometry.top_k):
+            expert = int(topk_ids[t, k])
+            weight = float(topk_weights[t, k])
+            expected[t, expert] += weight * _expert_out(expert) * H_SIGNAL
+    return expected.to(torch.bfloat16)
+
+
+def _moe_config(
+    vllm_config: VllmConfig,
+    geometry: Geometry,
+    world_size: int,
+    device: torch.device,
+    backend: str,
+) -> FusedMoEConfig:
+    parallel = FusedMoEParallelConfig.make(
+        tp_size_=1,
+        pcp_size_=1,
+        dp_size_=world_size,
+        sp_size_=1,
+        vllm_parallel_config=vllm_config.parallel_config,
+    )
+    return FusedMoEConfig(
+        num_experts=geometry.num_experts,
+        experts_per_token=geometry.top_k,
+        hidden_dim=geometry.hidden_size,
+        intermediate_size=geometry.intermediate_size,
+        num_local_experts=geometry.num_experts // world_size,
+        num_logical_experts=geometry.num_experts,
+        activation=MoEActivation.SILU,
+        device=device,
+        routing_method=_routing_method(backend),
+        moe_parallel_config=parallel,
+        in_dtype=torch.bfloat16,
+        moe_backend=backend,
+        max_num_tokens=MAX_TOKENS_PER_RANK,
+    )
+
+
+def _check_expert_views_alias_kernel_memory(
+    adapter: FlashInferMoeEp,
+    geometry: Geometry,
+    rank: int,
+    num_local_experts: int,
+    device: torch.device,
+) -> None:
+    """EPLB moves experts through ``expert_weight_views``; the kernel must see it.
+
+    Swap local experts 0 and 1 in place on every rank, then route through the
+    swapped physical slots: the output must be the one the unswapped logical
+    experts would have produced.
+    """
+    for view in adapter.expert_weight_views():
+        first = view[0].clone()
+        view[0].copy_(view[1])
+        view[1].copy_(first)
+    torch.distributed.barrier()
+
+    topk_ids, topk_weights = _routing(geometry, rank, device)
+    slot = topk_ids % num_local_experts
+    physical_ids = torch.where(slot == 0, topk_ids + 1, topk_ids)
+    physical_ids = torch.where(slot == 1, topk_ids - 1, physical_ids)
+    output = adapter(_hidden_states(geometry, device), physical_ids, topk_weights)
+    torch.accelerator.synchronize()
+    assert torch.equal(output, _expected_output(geometry, topk_ids, topk_weights))
+
+
+def _init_warmup_forward(
+    env: dict[str, str], world_size: int, geometry: Geometry, backend: str
+) -> None:
+    set_env_vars_and_device(env)
+    vllm_config = VllmConfig()
+    vllm_config.parallel_config.data_parallel_size = world_size
+    vllm_config.parallel_config.enable_expert_parallel = True
+    vllm_config.kernel_config.moe_backend = backend
+
+    with set_current_vllm_config(vllm_config):
+        ensure_model_parallel_initialized(
+            tensor_model_parallel_size=1, pipeline_model_parallel_size=1
+        )
+        rank = torch.distributed.get_rank()
+        device = torch.device(f"cuda:{rank}")
+        moe = _moe_config(vllm_config, geometry, world_size, device, backend)
+        supported, reason = FlashInferMoeEpExperts.is_supported_config(
+            FlashInferMoeEpExperts,
+            moe,
+            kMxfp4Static,
+            None,
+            mk.FusedMoEActivationFormat.Standard,
+        )
+        assert supported, reason
+
+        adapter = FlashInferMoeEp(
+            moe,
+            _local_expert_weights(geometry, rank, moe.num_local_experts, device),
+            apply_topk_in_fc1=apply_topk_in_fc1(moe),
+        )
+        try:
+            adapter.warmup()
+
+            topk_ids, topk_weights = _routing(geometry, rank, device)
+            output = adapter(_hidden_states(geometry, device), topk_ids, topk_weights)
+            torch.accelerator.synchronize()
+
+            expected = _expected_output(geometry, topk_ids, topk_weights)
+            assert output.dtype == expected.dtype and output.shape == expected.shape
+            assert torch.equal(output, expected), (
+                f"rank {rank}: max abs diff "
+                f"{(output.float() - expected.float()).abs().max().item()}, "
+                f"nonzero output channels per token "
+                f"{[torch.nonzero(row).flatten().tolist() for row in output]}"
+            )
+
+            if backend == FLASHINFER_MOE_EP_CUTEDSL:
+                _check_expert_views_alias_kernel_memory(
+                    adapter, geometry, rank, moe.num_local_experts, device
+                )
+        finally:
+            adapter.destroy()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="The FlashInfer MoE-EP megakernels require SM100",
+)
+@pytest.mark.parametrize("geometry", GEOMETRIES, ids=[g.name for g in GEOMETRIES])
+@pytest.mark.parametrize("backend", BACKENDS, ids=("cutedsl", "deep_gemm"))
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_init_warmup_forward_is_exact_on_grid_inputs(
+    world_size: int, backend: str, geometry: Geometry
+):
+    pytest.importorskip("flashinfer.moe_ep")
+    if backend == FLASHINFER_MOE_EP_DEEP_GEMM and not has_deep_gemm():
+        pytest.skip("DeepGEMM is not available")
+    if torch.accelerator.device_count() < world_size:
+        pytest.skip(f"Need at least {world_size} GPUs")
+    distributed_run(_init_warmup_forward, world_size, geometry, backend)

tests/models/test_deepseek_v4_fi_moe_ep.py

@@ -1,116 +1,90 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the flashinfer moe_ep backend plumbing.
 
-Everything here runs without a GPU or a flashinfer install: the flashinfer
-modules the helpers import lazily are replaced with capture fakes.
-"""
-
-import sys
-from dataclasses import dataclass, field
-from types import ModuleType, SimpleNamespace
-from typing import Any
+from dataclasses import replace
+from types import SimpleNamespace
 
 import pytest
 import torch
 
+from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm.config.kernel import (
     FLASHINFER_MOE_EP_BACKENDS,
+    FLASHINFER_MOE_EP_CUTEDSL,
+    FLASHINFER_MOE_EP_DEEP_GEMM,
     MEGA_MOE_BACKENDS,
+    NATIVE_MEGA_MOE_BACKENDS,
+    PASSTHROUGH_ALL2ALL_BACKEND,
+    bind_passthrough_all2all_backend,
     validate_flashinfer_moe_ep_model,
 )
-from vllm.utils.flashinfer_moe_ep import (
-    _E2M1_LUT,
-    FI_MOE_EP_BACKEND_SPECS,
-    _dequant_fp4_ue8m0_gran32,
-    build_fi_mega_config,
-    fi_moe_ep_backend_spec,
-    make_fi_moe_ep_bootstrap,
-    megakernel_runtime_requirements,
+from vllm.engine.arg_utils import EngineArgs
+from vllm.model_executor.layers.fused_moe import flashinfer_moe_ep as fi_ep
+from vllm.model_executor.layers.fused_moe import modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
+    maybe_make_prepare_finalize,
 )
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEQuantConfig,
+    RoutingMethodType,
+    nvfp4_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import (
+    FlashInferMoeEpExperts,
+    epilogue_from_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+    FLASHINFER_MOE_EP_MXFP4_BACKENDS,
+    Mxfp4MoeBackend,
+    make_mxfp4_moe_quant_config,
+    map_mxfp4_backend,
+    mxfp4_round_up_hidden_size_and_intermediate_size,
+)
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+    backend_to_kernel_cls as mxfp4_backend_to_kernel_cls,
+)
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+    NvFp4MoeBackend,
+    convert_to_nvfp4_moe_kernel_format,
+    map_nvfp4_backend,
+)
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+    backend_to_kernel_cls as nvfp4_backend_to_kernel_cls,
+)
+from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import (
+    PassThroughPrepareAndFinalize,
+)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kMxfp4Static,
+    kNvfp4Dynamic,
+    kNvfp4Static,
+)
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
-@dataclass
-class _FakeBootstrapConfig:
-    world_size: int
-    rank: int
-    process_group: Any = None
-    auto_bootstrap: bool = True
-    device: int | None = field(default=None, kw_only=True)
+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)
 
 
-@dataclass
-class _FakeDeepGemmMegaMoeConfig:
-    intermediate_size: int
-    top_k: int
-    activation_clamp: float | None
-    fast_math: bool
-
-
-@dataclass
-class _FakeNvfp4CutedslMegaMoeConfig:
-    intermediate_size: int
-    top_k: int
-    activation_clamp: float | None
-    fast_math: bool
-
-
-@dataclass
-class _FakeMegaConfig:
-    megakernel: Any
-    preprocess_weights: bool
-    quantize_input: bool
-
-
-@pytest.fixture
-def fake_flashinfer(monkeypatch):
-    """Install a minimal fake flashinfer.moe_ep for the lazy imports."""
-    moe_ep = ModuleType("flashinfer.moe_ep")
-    core = ModuleType("flashinfer.moe_ep.core")
-    runtime = ModuleType("flashinfer.moe_ep.core.runtime")
-    flashinfer = ModuleType("flashinfer")
-    fake_attrs: dict[ModuleType, dict[str, Any]] = {
-        moe_ep: {
-            "BootstrapConfig": _FakeBootstrapConfig,
-            "DeepGemmMegaMoeConfig": _FakeDeepGemmMegaMoeConfig,
-            "Nvfp4CutedslMegaMoeConfig": _FakeNvfp4CutedslMegaMoeConfig,
-            "MegaConfig": _FakeMegaConfig,
-            "core": core,
-        },
-        runtime: {"TORCH_DIST": "torch_dist", "NVSHMEM": "nvshmem"},
-        flashinfer: {"moe_ep": moe_ep},
-        core: {"runtime": runtime},
-    }
-    for mod, attrs in fake_attrs.items():
-        for attr, value in attrs.items():
-            setattr(mod, attr, value)
-
-    for name, mod in {
-        "flashinfer": flashinfer,
-        "flashinfer.moe_ep": moe_ep,
-        "flashinfer.moe_ep.core": core,
-        "flashinfer.moe_ep.core.runtime": runtime,
-    }.items():
-        monkeypatch.setitem(sys.modules, name, mod)
-    return moe_ep
-
-
-def test_fi_backend_strings_are_registered_mega_moe_backends():
-    assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
-    assert FLASHINFER_MOE_EP_BACKENDS < MEGA_MOE_BACKENDS
-
-
-@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS))
-def test_fi_moe_ep_backend_rejected_for_non_dsv4(moe_backend):
-    """An FI moe_ep backend with a non-DSv4 model must fail at config time
-    instead of silently falling through to the generic FusedMoE path."""
+def test_only_deep_gemm_backend_is_dsv4_specific():
+    validate_flashinfer_moe_ep_model(
+        FLASHINFER_MOE_EP_CUTEDSL,
+        ["MixtralForCausalLM"],
+    )
     with pytest.raises(ValueError, match="only supported for DeepSeek-V4"):
-        validate_flashinfer_moe_ep_model(moe_backend, ["MixtralForCausalLM"])
-
-
-@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS))
-def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend):
-    validate_flashinfer_moe_ep_model(moe_backend, ["DeepseekV4ForCausalLM"])
+        validate_flashinfer_moe_ep_model(
+            FLASHINFER_MOE_EP_DEEP_GEMM,
+            ["MixtralForCausalLM"],
+        )
+    validate_flashinfer_moe_ep_model(
+        FLASHINFER_MOE_EP_DEEP_GEMM,
+        ["DeepseekV4ForCausalLM"],
+    )
 
 
 @pytest.mark.parametrize(
@@ -118,8 +92,7 @@ def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend):
     [["KimiK3ForConditionalGeneration"], ["MixtralForCausalLM"]],
 )
 def test_native_deep_gemm_mega_moe_not_arch_gated(architectures):
-    """VLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it);
-    models validate their own constraints at construction time."""
+    """Models validate native deep_gemm mega constraints at construction time."""
     validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures)
 
 
@@ -128,11 +101,7 @@ def test_non_fi_backend_ignores_architectures():
 
 
 @pytest.mark.parametrize("moe_backend", sorted(MEGA_MOE_BACKENDS))
-def test_all_mega_backends_get_sequence_parallel_moe(moe_backend):
-    """Every mega backend must qualify for sequence-parallel MoE at
-    TP>1/EP: the predicate once matched only the native backend string,
-    which silently ran the fi backends full-batch with an all-reduce on
-    every rank — 0.42-0.65x native e2e at TP8."""
+def test_token_sharding_backends_enable_dsv4_sequence_parallel(moe_backend: str):
     from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel
 
     vllm_config = SimpleNamespace(
@@ -147,113 +116,262 @@ def test_all_mega_backends_get_sequence_parallel_moe(moe_backend):
     assert _use_sequence_parallel(vllm_config)
 
 
-def test_fi_moe_ep_backend_spec_kernel_and_nvshmem_contract():
-    dg = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm")
-    assert dg.megakernel == "deep_gemm_mega"
-    assert not dg.needs_nvshmem
-
-    cd = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl")
-    assert cd.megakernel == "nvfp4_cutedsl"
-    assert cd.needs_nvshmem
-
-    with pytest.raises(ValueError, match="not a flashinfer moe_ep backend"):
-        fi_moe_ep_backend_spec("deep_gemm_mega_moe")
+def test_dsv4_requests_pre_fc2_router_weight_placement():
+    dsv4 = SimpleNamespace(routing_method=RoutingMethodType.DeepseekV4)
+    default = SimpleNamespace(routing_method=RoutingMethodType.Default)
+    assert fi_ep.apply_topk_in_fc1(dsv4)
+    assert not fi_ep.apply_topk_in_fc1(default)
 
 
-def test_megakernel_runtime_requirements(fake_flashinfer):
-    dg = megakernel_runtime_requirements(
-        fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm")
+def test_backend_specs_preserve_weight_format_contracts():
+    cutedsl = fi_ep.flashinfer_moe_ep_backend_spec(FLASHINFER_MOE_EP_CUTEDSL)
+    assert cutedsl.kernel == "cutedsl"
+    assert cutedsl.weight_formats == frozenset({"nvfp4", "mxfp4"})
+
+    deep_gemm = fi_ep.flashinfer_moe_ep_backend_spec(FLASHINFER_MOE_EP_DEEP_GEMM)
+    assert deep_gemm.kernel == "deep_gemm"
+    assert deep_gemm.weight_formats == frozenset({"mxfp4"})
+
+
+def _megakernel_moe(moe_backend: str, **overrides) -> SimpleNamespace:
+    fields = dict(
+        moe_backend=moe_backend,
+        moe_parallel_config=SimpleNamespace(
+            use_ep=True, use_batched_activation_format=False
+        ),
+        is_act_and_mul=True,
+        activation=MoEActivation.SILU,
+        has_hash_routing=False,
+        routing_method=RoutingMethodType.DeepseekV4,
+        router_logits_dtype=torch.float32,
+        hidden_dim=256,
+        is_lora_enabled=False,
+        skip_final_all_reduce=False,
+        in_dtype=torch.bfloat16,
+        has_bias=False,
+        swiglu_alpha=None,
+        swiglu_beta=None,
     )
-    assert dg == frozenset({"torch_dist"})
+    fields.update(overrides)
+    return SimpleNamespace(**fields)
 
-    cd = megakernel_runtime_requirements(
-        fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl")
+
+def _is_supported(moe, weight_key, activation_key) -> tuple[bool, str]:
+    supported, reason = FlashInferMoeEpExperts.is_supported_config(
+        FlashInferMoeEpExperts,
+        moe,
+        weight_key,
+        activation_key,
+        mk.FusedMoEActivationFormat.Standard,
     )
-    assert cd == frozenset({"torch_dist", "nvshmem"})
+    return supported, reason or ""
 
 
-def test_bootstrap_pins_the_device_vllm_bound(fake_flashinfer, monkeypatch):
-    """The runtime must not rederive the device from LOCAL_RANK/rank: under a
-    remapped CUDA_VISIBLE_DEVICES that ordinal points at the wrong GPU
-    (CUDA_ERROR_ILLEGAL_ADDRESS in the weight transforms). vLLM passes the
-    device it already bound via BootstrapConfig.device."""
-    import vllm.utils.flashinfer_moe_ep as mod
-
-    pg = object()
+@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),
+    )
+    monkeypatch.setattr(fi_ep, "get_current_vllm_config", lambda: config)
     monkeypatch.setattr(
-        mod,
-        "get_ep_group",
-        lambda: SimpleNamespace(world_size=4, rank_in_group=2, device_group=pg),
+        FlashInferMoeEpExperts, "_supports_current_device", staticmethod(lambda: True)
     )
-    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3)
-
-    bootstrap = make_fi_moe_ep_bootstrap()
-
-    assert bootstrap.world_size == 4
-    assert bootstrap.rank == 2
-    assert bootstrap.process_group is pg
-    assert bootstrap.auto_bootstrap is False
-    assert bootstrap.device == 3
+    return config
 
 
-def test_build_fi_mega_config_selects_kernel_config(fake_flashinfer):
-    dg = build_fi_mega_config(
-        intermediate_size=2048,
-        top_k=8,
-        activation_clamp=7.0,
-        megakernel="deep_gemm_mega",
+def test_megakernel_requires_expert_parallel(megakernel_host):
+    """The oracle rejects the megakernel through the experts' support check."""
+    moe = _megakernel_moe(FLASHINFER_MOE_EP_CUTEDSL)
+    assert _is_supported(moe, kNvfp4Static, kNvfp4Dynamic) == (True, "")
+
+    moe.moe_parallel_config.use_ep = False
+    supported, reason = _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)
+    assert not supported and "parallel config" in reason
+
+
+def test_megakernel_reports_its_own_constraints(megakernel_host):
+    """Constraints the framework does not know about surface as reasons."""
+    _, reason = _is_supported(
+        _megakernel_moe(FLASHINFER_MOE_EP_CUTEDSL), kNvfp4Static, None
     )
-    assert isinstance(dg.megakernel, _FakeDeepGemmMegaMoeConfig)
-    assert dg.megakernel.intermediate_size == 2048
-    assert dg.megakernel.top_k == 8
-    assert dg.megakernel.activation_clamp == 7.0
-    assert dg.preprocess_weights and dg.quantize_input
+    assert "A16 activations" in reason
 
-    cd = build_fi_mega_config(
-        intermediate_size=2048,
-        top_k=8,
-        activation_clamp=None,
-        megakernel="nvfp4_cutedsl",
+    deep_gemm = _megakernel_moe(
+        FLASHINFER_MOE_EP_DEEP_GEMM, routing_method=RoutingMethodType.Renormalize
     )
-    assert isinstance(cd.megakernel, _FakeNvfp4CutedslMegaMoeConfig)
+    _, reason = _is_supported(deep_gemm, kMxfp4Static, None)
+    assert "routing method" in reason
 
-    with pytest.raises(ValueError, match="Unsupported fi_moe_ep megakernel"):
-        build_fi_mega_config(
-            intermediate_size=2048,
-            top_k=8,
-            activation_clamp=None,
-            megakernel="deep_gemm",
+    megakernel_host.parallel_config.enable_eplb = True
+    _, reason = _is_supported(
+        _megakernel_moe(FLASHINFER_MOE_EP_DEEP_GEMM), kMxfp4Static, None
+    )
+    assert reason.endswith("EPLB")
+
+
+def test_megakernels_are_oracle_backends():
+    """Quant methods reach the megakernel through the NVFP4/MXFP4 oracles."""
+    nvfp4 = map_nvfp4_backend(FLASHINFER_MOE_EP_CUTEDSL)
+    assert nvfp4 is NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL
+    assert nvfp4_backend_to_kernel_cls(nvfp4) == [FlashInferMoeEpExperts]
+
+    assert map_mxfp4_backend(FLASHINFER_MOE_EP_DEEP_GEMM) == [
+        Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM
+    ]
+    for backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
+        assert mxfp4_backend_to_kernel_cls(backend) == [FlashInferMoeEpExperts]
+        # The megakernel takes the model's own dimensions.
+        assert mxfp4_round_up_hidden_size_and_intermediate_size(
+            backend, 2880, 2880
+        ) == (2880, 2880)
+
+    quant_config = make_mxfp4_moe_quant_config(
+        Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, w1_scale=None, w2_scale=None
+    )
+    assert quant_config is not None and quant_config.use_nvfp4_w4a4
+
+
+def test_megakernel_is_a_modular_kernel_with_pass_through_stages():
+    """The megakernel plugs into ``FusedMoEKernel`` like any other experts impl.
+
+    Routing stays outside (top-k ids in), dispatch and combine happen inside,
+    so prepare passes tokens through untouched and finalize has nothing to do.
+    """
+    moe = make_dummy_moe_config(
+        num_experts=8, experts_per_token=2, hidden_dim=256, intermediate_size=128
+    )
+    quant_config = FusedMoEQuantConfig.make("nvfp4", weight_dtype="nvfp4")
+    kernel = mk.FusedMoEKernel(
+        PassThroughPrepareAndFinalize(),
+        FlashInferMoeEpExperts(moe, quant_config),
+    )
+
+    assert not kernel.is_monolithic
+    assert kernel.prepare_finalize.output_is_reduced()
+    # Generic stages: the int32 top-k ids the megakernel wants are its concern.
+    assert kernel.prepare_finalize.topk_indices_dtype() is None
+    assert kernel.fused_experts.expects_unquantized_inputs
+    assert isinstance(
+        kernel.fused_experts.finalize_weight_and_reduce_impl(), TopKWeightAndReduceNoOP
+    )
+    assert kernel.fused_experts.workspace_shapes(
+        16, 256, 256, 2, 8, 8, None, MoEActivation.SILU
+    ) == ((0,), (0,), (16, 256))
+
+    tokens = torch.zeros(4, 256, dtype=torch.bfloat16)
+    prepared = kernel.prepare_finalize.prepare(
+        tokens,
+        torch.ones(4, 2),
+        torch.zeros(4, 2, dtype=torch.int32),
+        8,
+        None,
+        False,
+        quant_config,
+    )
+    assert prepared[0] is tokens and all(item is None for item in prepared[1:])
+
+
+def test_epilogue_takes_the_weight_global_scales_from_the_quant_config():
+    """The canonical ``g1/g2_alphas`` become the per-expert fc1/fc2 alphas and
+    keep aliasing the layer's tensors, so EPLB permutations reach the kernel."""
+    g1_alphas = torch.tensor([1.0, 2.0])
+    g2_alphas = torch.tensor([3.0, 4.0])
+    quant_config = nvfp4_moe_quant_config(
+        g1_alphas=g1_alphas,
+        g2_alphas=g2_alphas,
+        a1_gscale=torch.ones(2),
+        a2_gscale=torch.ones(2),
+        w1_scale=torch.ones(2, 4, 1),
+        w2_scale=torch.ones(2, 2, 1),
+    )
+
+    epilogue = epilogue_from_quant_config(quant_config)
+
+    assert epilogue.fc1_alpha is not None and epilogue.fc2_alpha is not None
+    assert epilogue.fc1_alpha.data_ptr() == g1_alphas.data_ptr()
+    assert epilogue.fc2_alpha.data_ptr() == g2_alphas.data_ptr()
+    # Activations are quantized dynamically: no norm constants.
+    assert epilogue.input_norm_const == 1.0 and epilogue.fc1_norm_const is None
+
+
+def test_epilogue_rejects_unloaded_weight_global_scales():
+    quant_config = nvfp4_moe_quant_config(
+        g1_alphas=torch.tensor([1.0, float("nan")]),
+        g2_alphas=torch.ones(2),
+        a1_gscale=torch.ones(2),
+        a2_gscale=torch.ones(2),
+        w1_scale=torch.ones(2, 4, 1),
+        w2_scale=torch.ones(2, 2, 1),
+    )
+    with pytest.raises(ValueError, match="g1_alphas"):
+        epilogue_from_quant_config(quant_config)
+
+
+def test_mxfp4_checkpoints_keep_the_default_epilogue():
+    """MXFP4 has no per-tensor global scales; the megakernel needs no alphas."""
+    quant_config = FusedMoEQuantConfig.make("nvfp4", weight_dtype="mxfp4")
+    assert epilogue_from_quant_config(quant_config) == fi_ep.FlashInferMoeEpEpilogue()
+
+
+def test_nvfp4_kernel_format_passes_megakernel_weights_through():
+    """The megakernel takes the canonical tensors as they are; the oracle must
+    not reject or reshape them."""
+    tensors = tuple(torch.arange(4.0).reshape(2, 2) + i for i in range(8))
+    converted = convert_to_nvfp4_moe_kernel_format(
+        NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL,
+        SimpleNamespace(),
+        *tensors,
+        is_act_and_mul=True,
+    )
+    assert all(out is src for out, src in zip(converted, tensors))
+
+
+def test_passthrough_all2all_backend_selects_the_pass_through_stages():
+    """``all2all_backend="passthrough"`` is how the oracles get the
+    pass-through prepare/finalize."""
+    moe = make_dummy_moe_config(
+        num_experts=8, experts_per_token=2, hidden_dim=256, intermediate_size=128
+    )
+    moe.moe_parallel_config = replace(
+        moe.moe_parallel_config, all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND
+    )
+    assert not moe.moe_parallel_config.use_all2all_kernels
+
+    prepare_finalize = maybe_make_prepare_finalize(
+        moe, FusedMoEQuantConfig.make("nvfp4", weight_dtype="nvfp4")
+    )
+
+    assert isinstance(prepare_finalize, PassThroughPrepareAndFinalize)
+
+
+def test_passthrough_all2all_backend_is_not_a_cli_choice():
+    """Users cannot pick the passthrough all2all backend explicitly."""
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    with pytest.raises(SystemExit):
+        parser.parse_args(["--all2all-backend", PASSTHROUGH_ALL2ALL_BACKEND])
+    parser.parse_args(["--all2all-backend", "deepep_low_latency"])
+
+
+def test_megakernel_moe_backend_binds_the_all2all_backend():
+    """A megakernel deployment gets its all2all backend by default."""
+    parallel = SimpleNamespace(all2all_backend="allgather_reducescatter")
+    bind_passthrough_all2all_backend(FLASHINFER_MOE_EP_CUTEDSL, parallel)
+    assert parallel.all2all_backend == PASSTHROUGH_ALL2ALL_BACKEND
+    bind_passthrough_all2all_backend(FLASHINFER_MOE_EP_CUTEDSL, parallel)
+
+    with pytest.raises(ValueError, match="cannot be paired"):
+        bind_passthrough_all2all_backend(
+            FLASHINFER_MOE_EP_DEEP_GEMM,
+            SimpleNamespace(all2all_backend="deepep_low_latency"),
+        )
+    with pytest.raises(ValueError, match="only valid with"):
+        bind_passthrough_all2all_backend(
+            "flashinfer_trtllm",
+            SimpleNamespace(all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND),
         )
 
-
-def test_ckpt_uses_nvfp4_experts_reads_moe_quant_algo():
-    from vllm.models.deepseek_v4.nvidia.fi_moe import ckpt_uses_nvfp4_experts
-
-    nvfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo="NVFP4"))
-    assert ckpt_uses_nvfp4_experts(nvfp4)
-
-    mxfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo=None))
-    assert not ckpt_uses_nvfp4_experts(mxfp4)
-
-    no_algo = SimpleNamespace(quant_config=SimpleNamespace())
-    assert not ckpt_uses_nvfp4_experts(no_algo)
-
-
-def test_dequant_fp4_ue8m0_gran32_decodes_lut_and_scales():
-    """One 32-element scale group per row: low nibble is the even element,
-    high nibble the odd one, ue8m0 scale applies to the whole group."""
-    packed = torch.arange(32, dtype=torch.uint8).reshape(2, 16)
-    sf = torch.tensor([[127], [128]], dtype=torch.uint8)  # 2**0, 2**1
-
-    out = _dequant_fp4_ue8m0_gran32(packed, sf)
-
-    assert out.shape == (2, 32)
-    assert out.dtype == torch.bfloat16
-    expected = torch.empty(2, 32)
-    for row in range(2):
-        for col in range(16):
-            byte = int(packed[row, col])
-            expected[row, 2 * col] = _E2M1_LUT[byte & 0x0F]
-            expected[row, 2 * col + 1] = _E2M1_LUT[byte >> 4]
-        expected[row] *= 2.0**row
-    assert torch.equal(out, expected.to(torch.bfloat16))
+    untouched = SimpleNamespace(all2all_backend="deepep_low_latency")
+    bind_passthrough_all2all_backend("flashinfer_trtllm", untouched)
+    assert untouched.all2all_backend == "deepep_low_latency"

tests/models/test_deepseek_v4_mega_moe.py

@@ -657,7 +657,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert(
 
     moe = DeepseekV4MoE.__new__(DeepseekV4MoE)
     torch.nn.Module.__init__(moe)
-    moe.use_mega_moe = True
+    moe.use_native_mega_moe = True
     moe.gate = FakeGate()
     moe.experts = FakeExperts()
     moe.shared_experts = FakeSharedExperts()

tests/models/test_dspark_mla.py

@@ -241,7 +241,7 @@ def test_v41_dspark_loads_linear_scales(
         linear_scale_name=scale_name,
         pad_shared_expert=False,
         model=SimpleNamespace(
-            layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))],
+            layers=[SimpleNamespace(ffn=SimpleNamespace(use_native_mega_moe=False))],
             confidence_head=None,
         ),
         named_parameters=lambda: [(runtime_name, param)],

vllm/config/kernel.py

@@ -3,7 +3,7 @@
 import contextlib
 from collections.abc import Callable, Iterable
 from dataclasses import asdict, fields
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Final, Literal
 
 from pydantic import Field, field_validator
 
@@ -12,6 +12,7 @@ from vllm.logger import init_logger
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
+    from vllm.config.parallel import ParallelConfig
 
 logger = init_logger(__name__)
 
@@ -128,10 +129,10 @@ MoEBackend = Literal[
     "flashinfer_trtllm",
     "flashinfer_cutlass",
     "flashinfer_cutedsl",
+    "flashinfer_moe_ep_cutedsl",
     "flashinfer_b12x",
     "b12x",
     "flashinfer_moe_ep_mega_deep_gemm",
-    "flashinfer_moe_ep_mega_cutedsl",
     "marlin",
     "humming",
     "triton_unfused",
@@ -143,21 +144,28 @@ MoEBackend = Literal[
     "rdna3",
 ]
 
-# Backends that run the mega-MoE model path through the flashinfer moe_ep
-# runtime. Only architectures in FLASHINFER_MOE_EP_ARCHITECTURES wire up
-# these experts.
+FLASHINFER_MOE_EP_CUTEDSL = "flashinfer_moe_ep_cutedsl"
+FLASHINFER_MOE_EP_DEEP_GEMM = "flashinfer_moe_ep_mega_deep_gemm"
+
+# Backends implemented by the shared FlashInfer MoE-EP adapter.
 FLASHINFER_MOE_EP_BACKENDS = frozenset(
     {
-        "flashinfer_moe_ep_mega_deep_gemm",
-        "flashinfer_moe_ep_mega_cutedsl",
+        FLASHINFER_MOE_EP_CUTEDSL,
+        FLASHINFER_MOE_EP_DEEP_GEMM,
     }
 )
+# Those backends dispatch and combine themselves, so the framework must not
+# run a separate all2all around them.
+PASSTHROUGH_ALL2ALL_BACKEND: Final = "passthrough"
 
-# Backends that run a mega-MoE model path (fused expert module plus
-# prepare_megamoe routing): vLLM's native deep_gemm path, which any model
-# with a mega-MoE module may use (DeepSeek-V4, Kimi K3), plus the flashinfer
-# moe_ep variants.
-MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS
+# "Native" here meaning that the backend does not use the fused MoE path.
+# Currently only relevant for the vLLM provided DeepGEMM MegaMoE.,
+# TODO: migrate to the fused MoE path ultimately.
+NATIVE_MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"})
+
+# Every megakernel backend. The FlashInfer variants run through the standard
+# fused-MoE model path; the quantization method binds the kernel.
+MEGA_MOE_BACKENDS = NATIVE_MEGA_MOE_BACKENDS | FLASHINFER_MOE_EP_BACKENDS
 
 SparseIndexerTopkBackend = Literal[
     "auto",
@@ -170,8 +178,8 @@ SparseIndexerTopkBackend = Literal[
     "aiter",
 ]
 
-# Architectures whose model code wires up the flashinfer moe_ep experts. MTP
-# and DSpark draft variants inherit the setting from these target models.
+# DeepGEMM consumes the DeepSeek-V4 MXFP4 checkpoint recipe. CuTeDSL is
+# selected by quantization methods and is not architecture-specific.
 FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
     {
         "DeepseekV4ForCausalLM",
@@ -181,11 +189,48 @@ FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
 )
 
 
+def bind_passthrough_all2all_backend(
+    moe_backend: str, parallel_config: "ParallelConfig"
+) -> None:
+    """Bind ``all2all_backend="passthrough"`` to MoE backends that own their
+    expert-parallel communication (currently the FlashInfer MoE-EP megakernels).
+
+    Such a deployment gets ``passthrough`` when the all2all backend is left at
+    its default and rejects any other explicit choice; the value is meaningless
+    with a MoE backend that relies on the framework to communicate. The CLI does
+    not offer the value; the early return keeps repeated
+    ``set_platform_defaults`` calls idempotent.
+    """
+    a2a = parallel_config.all2all_backend
+    if moe_backend in FLASHINFER_MOE_EP_BACKENDS:
+        if a2a == PASSTHROUGH_ALL2ALL_BACKEND:
+            return
+        if a2a != "allgather_reducescatter":
+            raise ValueError(
+                f"moe_backend={moe_backend!r} dispatches and combines itself and "
+                f"cannot be paired with all2all_backend={a2a!r}; leave "
+                f"--all2all-backend unset."
+            )
+        logger.info_once(
+            "moe_backend=%r: using all2all_backend=%r (dispatch and combine run "
+            "inside the MoE backend).",
+            moe_backend,
+            PASSTHROUGH_ALL2ALL_BACKEND,
+        )
+        parallel_config.all2all_backend = PASSTHROUGH_ALL2ALL_BACKEND
+    elif a2a == PASSTHROUGH_ALL2ALL_BACKEND:
+        raise ValueError(
+            f"all2all_backend={a2a!r} is only valid with a moe_backend that "
+            f"dispatches and combines itself ({sorted(FLASHINFER_MOE_EP_BACKENDS)}),"
+            f" got moe_backend={moe_backend!r}."
+        )
+
+
 def validate_flashinfer_moe_ep_model(
     moe_backend: str, architectures: Iterable[str]
 ) -> None:
-    """Reject flashinfer moe_ep backends for models that lack the FI path."""
-    if moe_backend not in FLASHINFER_MOE_EP_BACKENDS:
+    """Reject model-specific FlashInfer MoE-EP backends."""
+    if moe_backend != FLASHINFER_MOE_EP_DEEP_GEMM:
         return
     if not any(arch in FLASHINFER_MOE_EP_ARCHITECTURES for arch in architectures):
         raise ValueError(
@@ -255,6 +300,9 @@ 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
+      mega-kernel with NVFP4 weights (MXFP4 checkpoints are requantized at
+      load); requires Blackwell, expert parallelism, and NVSHMEM
     - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x
       (RTX Pro 6000 / DGX Spark)
     - "b12x": Use b12x FP4 MoE kernels on SM12x
@@ -262,10 +310,6 @@ class KernelConfig:
       expert-parallel mega-MoE with the DeepGEMM megakernel, which consumes an
       MXFP4 checkpoint verbatim (Blackwell, requires expert parallel;
       DeepSeek-V4 only)
-    - "flashinfer_moe_ep_mega_cutedsl": Same, with the CuteDSL megakernel
-      (additionally requires NVSHMEM). The checkpoint selects the weight path:
-      an NVFP4 checkpoint is consumed prequantized, MXFP4 weights are
-      requantized at load
     - "marlin": Use Marlin kernels (weight-only quantization)
     - "humming": Use Humming Mixed Precision kernels
     - "triton_unfused": Use Triton unfused MoE kernels
@@ -391,6 +435,7 @@ class KernelConfig:
             validate_flashinfer_moe_ep_model(
                 self.moe_backend, vllm_config.model_config.architectures
             )
+        bind_passthrough_all2all_backend(self.moe_backend, vllm_config.parallel_config)
 
         platform_op_priority = current_platform.get_default_ir_op_priority(vllm_config)
         logger.debug(

vllm/config/parallel.py

@@ -56,6 +56,7 @@ All2AllBackend = Literal[
     "flashinfer_all2allv",  # temporary alias for flashinfer_nvlink_two_sided
     "flashinfer_nvlink_two_sided",
     "flashinfer_nvlink_one_sided",
+    "passthrough",
 ]
 
 
@@ -211,7 +212,10 @@ class ParallelConfig:
     - "moonep": MoonEP balanced EP with dynamic redundant experts (NVLink)
     - "nixl_ep": Use nixl-ep kernels
     - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels
-    - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl"""
+    - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl
+    - "passthrough": No all2all at all: the MoE backend dispatches and combines
+      itself (e.g. the FlashInfer MoE-EP megakernels). Bound automatically for
+      such MoE backends; not selectable from the CLI."""
 
     max_parallel_loading_workers: int | None = Field(default=None, ge=1)
     """Maximum number of parallel loading workers when loading model
@@ -742,6 +746,7 @@ class ParallelConfig:
                 "deepep_low_latency",
                 "deepep_v2",
                 "flashinfer_nvlink_one_sided",
+                "passthrough",
                 "mori_high_throughput",
                 "mori_low_latency",
                 "nixl_ep",

vllm/distributed/device_communicators/all2all.py

@@ -49,6 +49,20 @@ if has_flashinfer_nvlink_one_sided():
 logger = init_logger(__name__)
 
 
+class PassThroughAll2AllManager(All2AllManagerBase):
+    """Placeholder for ``all2all_backend="passthrough"``.
+
+    The MoE backend dispatches and combines itself, so there is no all2all to
+    manage.
+    """
+
+    def get_handle(self, kwargs):
+        raise RuntimeError(
+            "passthrough has no all2all handle: dispatch and combine run "
+            "inside the MoE backend."
+        )
+
+
 class AgRsAll2AllManager(All2AllManagerBase):
     """An implementation of all2all communication based on
     all-gather (dispatch) and reduce-scatter (combine).

vllm/distributed/device_communicators/cuda_communicator.py

@@ -255,6 +255,12 @@ class CudaCommunicator(DeviceCommunicatorBase):
                 from .all2all import FlashInferNVLinkOneSidedManager
 
                 self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group)
+            elif self.all2all_backend == "passthrough":
+                from .all2all import PassThroughAll2AllManager
+
+                self.all2all_manager = PassThroughAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
             else:
                 raise ValueError(f"Unknown all2all backend: {self.all2all_backend}")
 

vllm/engine/arg_utils.py

@@ -80,6 +80,7 @@ from vllm.config.cache import (
 )
 from vllm.config.device import Device
 from vllm.config.kernel import (
+    PASSTHROUGH_ALL2ALL_BACKEND,
     IrOpPriorityConfig,
     LinearBackend,
     MoEBackend,
@@ -1104,6 +1105,10 @@ class EngineArgs:
 
         # Parallel arguments
         parallel_kwargs = get_kwargs(ParallelConfig)
+        # Bound from --moe-backend in KernelConfig.set_platform_defaults().
+        parallel_kwargs["all2all_backend"]["choices"].remove(
+            PASSTHROUGH_ALL2ALL_BACKEND
+        )
         parallel_group = parser.add_argument_group(
             title="ParallelConfig",
             description=ParallelConfig.__doc__,

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -31,6 +31,9 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one
 from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import (  # noqa: E501
     FlashInferNVLinkTwoSidedPrepareAndFinalize,
 )
+from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import (
+    PassThroughPrepareAndFinalize,
+)
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import (
     has_deep_ep,
@@ -185,6 +188,10 @@ def maybe_make_prepare_finalize(
     all2all_manager: Any | None = None,
     input_dtype: torch.dtype | None = None,
 ) -> FusedMoEPrepareAndFinalize | None:
+    if moe.moe_parallel_config.use_passthrough_all2all:
+        # The experts dispatch and combine themselves, in every EP topology.
+        return PassThroughPrepareAndFinalize()
+
     if not moe.moe_parallel_config.use_all2all_kernels:
         if not allow_new_interface:
             return None

vllm/model_executor/layers/fused_moe/config.py

@@ -1068,6 +1068,12 @@ class FusedMoEParallelConfig:
             and self.all2all_backend == "allgather_reducescatter"
         )
 
+    @property
+    def use_passthrough_all2all(self):
+        # Not gated on use_all2all_kernels: the experts dispatch and combine
+        # internally in every EP topology, including TP-only.
+        return self.all2all_backend == "passthrough"
+
     @property
     def use_mori_kernels(self):
         return self.use_all2all_kernels and self.all2all_backend in (
@@ -1530,6 +1536,10 @@ class FusedMoEConfig:
     def use_deepep_ll_kernels(self):
         return self.moe_parallel_config.use_deepep_ll_kernels
 
+    @property
+    def use_passthrough_all2all(self):
+        return self.moe_parallel_config.use_passthrough_all2all
+
     @property
     def use_mori_kernels(self):
         return self.moe_parallel_config.use_mori_kernels

vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlashInfer MoE-EP megakernels as a modular-kernel experts implementation.
+
+The megakernel routes on vLLM's top-k output and then dispatches, computes and
+combines in one launch, so it is paired with a pass-through prepare/finalize.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import (
+    FlashInferMoeEp,
+    FlashInferMoeEpEpilogue,
+    FlashInferMoeEpWeights,
+    apply_topk_in_fc1,
+    flashinfer_moe_ep_unsupported_reasons,
+    supports_current_device,
+    validate_flashinfer_moe_ep_layer,
+)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp4Static,
+    kNvfp4Static,
+    require_finite_positive_scales,
+)
+from vllm.model_executor.utils import replace_parameter
+
+if TYPE_CHECKING:
+    pass
+
+# Layer parameters the megakernel reads. After ``process_weights_after_loading``
+# they alias the megakernel's own tensors, so EPLB permutes exactly what it uses.
+KERNEL_WEIGHT_NAMES = ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale")
+
+
+_supported_weight_quant_schemes: frozenset[QuantKey] = frozenset(
+    (kNvfp4Static, kMxfp4Static)
+)
+
+_supported_activations: frozenset[MoEActivation] = frozenset((MoEActivation.SILU,))
+
+
+class FlashInferMoeEpExperts(mk.FusedMoEExpertsModular):
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ) -> None:
+        super().__init__(moe_config, quant_config)
+        self._adapter: FlashInferMoeEp | None = None
+
+    @staticmethod
+    def is_supported_config(
+        cls: type[mk.FusedMoEExperts],
+        moe_config: FusedMoEConfig,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+        activation_format: mk.FusedMoEActivationFormat,
+    ) -> tuple[bool, str | None]:
+        supported, reason = mk.FusedMoEExperts.is_supported_config(
+            cls, moe_config, weight_key, activation_key, activation_format
+        )
+        if not supported:
+            return False, reason
+        reasons = flashinfer_moe_ep_unsupported_reasons(
+            moe_config, weight_key, activation_key
+        )
+        if reasons:
+            return False, f"kernel does not support {', '.join(reasons)}"
+        return True, None
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return supports_current_device()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return False
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return weight_key in _supported_weight_quant_schemes
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation in _supported_activations
+
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        return (
+            moe_parallel_config.use_ep
+            and not moe_parallel_config.use_batched_activation_format
+        )
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceNoOP()
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        return (0,), (0,), (M, K)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Build the megakernel from the layer's canonical weights.
+
+        The oracle has already put the weights in kernel format; the per-expert
+        weight global scales arrive through the quant config and become the
+        epilogue alphas.
+        """
+        validate_flashinfer_moe_ep_layer(layer)  # type: ignore[arg-type]
+        adapter = FlashInferMoeEp(
+            self.moe_config,
+            FlashInferMoeEpWeights(
+                w13=layer.w13_weight,
+                w2=layer.w2_weight,
+                w13_scale=getattr(layer, "w13_weight_scale", None),
+                w2_scale=getattr(layer, "w2_weight_scale", None),
+            ),
+            epilogue_from_quant_config(self.quant_config),
+            apply_topk_in_fc1=apply_topk_in_fc1(
+                self.moe_config,
+                apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            ),
+        )
+        adapter.warmup()
+        for name, tensor in zip(KERNEL_WEIGHT_NAMES, adapter.kernel_weights()):
+            replace_parameter(layer, name, tensor)
+        self._adapter = adapter
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ) -> None:
+        if self._adapter is None:
+            raise RuntimeError("FlashInfer MoE-EP weights have not been processed")
+        output.copy_(
+            self._adapter(hidden_states, topk_ids.to(torch.int32), topk_weights)
+        )
+
+
+def epilogue_from_quant_config(
+    quant_config: FusedMoEQuantConfig,
+) -> FlashInferMoeEpEpilogue:
+    """Per-expert epilogue constants from the canonical quant config.
+
+    The weight global scales (``g1_alphas``/``g2_alphas``) become the fc1/fc2
+    alphas. Activations are quantized dynamically inside the megakernel, so the
+    checkpoint's input scales drop out. MXFP4 checkpoints carry no global scales
+    and keep the default epilogue. The alphas alias the layer's parameters, so
+    EPLB permutations reach the kernel.
+    """
+    fc1_alpha, fc2_alpha = quant_config.g1_alphas, quant_config.g2_alphas
+    if fc1_alpha is None and fc2_alpha is None:
+        return FlashInferMoeEpEpilogue()
+    assert fc1_alpha is not None and fc2_alpha is not None
+    require_finite_positive_scales("g1_alphas", fc1_alpha)
+    require_finite_positive_scales("g2_alphas", fc2_alpha)
+    if fc1_alpha.dim() != 1 or fc2_alpha.dim() != 1:
+        raise ValueError("the megakernel needs one weight global scale per expert")
+    return FlashInferMoeEpEpilogue(
+        fc1_alpha=fc1_alpha.float().contiguous(),
+        fc2_alpha=fc2_alpha.float().contiguous(),
+    )

vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py

@@ -0,0 +1,481 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import contextlib
+import sys
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal, cast
+
+import torch
+
+from vllm.config import get_current_vllm_config
+from vllm.config.kernel import (
+    FLASHINFER_MOE_EP_BACKENDS,
+    FLASHINFER_MOE_EP_CUTEDSL,
+    FLASHINFER_MOE_EP_DEEP_GEMM,
+)
+from vllm.distributed import get_ep_group
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp4Static,
+    kNvfp4Static,
+)
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import _import_deep_gemm
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
+    from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+
+FlashInferMoeEpKernel = Literal["cutedsl", "deep_gemm"]
+FlashInferMoeEpWeightFormat = Literal["nvfp4", "mxfp4"]
+_WEIGHT_FORMATS: dict[QuantKey | None, FlashInferMoeEpWeightFormat] = {
+    kNvfp4Static: "nvfp4",
+    kMxfp4Static: "mxfp4",
+}
+
+_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)})
+
+
+def supports_current_device() -> bool:
+    capability = current_platform.get_device_capability()
+    return capability is not None and (
+        (capability.major, capability.minor) in _SUPPORTED_CAPABILITIES
+    )
+
+
+_E2M1_LUT = (
+    0.0,
+    0.5,
+    1.0,
+    1.5,
+    2.0,
+    3.0,
+    4.0,
+    6.0,
+    -0.0,
+    -0.5,
+    -1.0,
+    -1.5,
+    -2.0,
+    -3.0,
+    -4.0,
+    -6.0,
+)
+
+
+@dataclass(frozen=True)
+class FlashInferMoeEpBackendSpec:
+    kernel: FlashInferMoeEpKernel
+    weight_formats: frozenset[FlashInferMoeEpWeightFormat]
+
+
+@dataclass(frozen=True)
+class FlashInferMoeEpWeights:
+    w13: torch.Tensor
+    w2: torch.Tensor
+    w13_scale: torch.Tensor | None = None
+    w2_scale: torch.Tensor | None = None
+
+
+@dataclass(frozen=True)
+class FlashInferMoeEpEpilogue:
+    input_norm_const: float = 1.0
+    fc1_alpha: torch.Tensor | None = None
+    fc2_alpha: torch.Tensor | None = None
+    fc1_norm_const: torch.Tensor | None = None
+
+
+_BACKEND_SPECS = {
+    FLASHINFER_MOE_EP_CUTEDSL: FlashInferMoeEpBackendSpec(
+        kernel="cutedsl",
+        weight_formats=frozenset({"nvfp4", "mxfp4"}),
+    ),
+    FLASHINFER_MOE_EP_DEEP_GEMM: FlashInferMoeEpBackendSpec(
+        kernel="deep_gemm",
+        weight_formats=frozenset({"mxfp4"}),
+    ),
+}
+
+assert frozenset(_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
+
+
+def is_flashinfer_moe_ep_backend(moe_backend: str) -> bool:
+    return moe_backend in _BACKEND_SPECS
+
+
+def flashinfer_moe_ep_backend_spec(
+    moe_backend: str,
+) -> FlashInferMoeEpBackendSpec:
+    try:
+        return _BACKEND_SPECS[moe_backend]
+    except KeyError:
+        raise ValueError(
+            f"{moe_backend!r} is not a FlashInfer MoE-EP backend; expected "
+            f"one of {sorted(_BACKEND_SPECS)}"
+        ) from None
+
+
+def flashinfer_moe_ep_unsupported_reasons(
+    moe: FusedMoEConfig,
+    weight_key: QuantKey | None,
+    activation_key: QuantKey | None,
+) -> tuple[str, ...]:
+    """Deployment constraints of the megakernel beyond the framework's checks.
+
+    ``activation_key`` is ``None`` for MXFP4 checkpoints (the kernel quantizes
+    bf16 activations itself) but marks a W4A16 recipe for NVFP4 weights.
+    """
+    spec = flashinfer_moe_ep_backend_spec(moe.moe_backend)
+    unsupported: list[str] = []
+    weight_format = _WEIGHT_FORMATS.get(weight_key)
+    if weight_format is None:
+        unsupported.append(f"weight scheme {weight_key}")
+    elif weight_format not in spec.weight_formats:
+        unsupported.append(f"{weight_format.upper()} weights")
+    if weight_key == kNvfp4Static and activation_key is None:
+        unsupported.append("A16 activations with NVFP4 weights")
+    if moe.skip_final_all_reduce:
+        unsupported.append("skip_final_all_reduce")
+    if moe.in_dtype != torch.bfloat16:
+        unsupported.append(f"activation dtype {moe.in_dtype}")
+    if moe.has_bias:
+        unsupported.append("expert bias")
+    if moe.swiglu_alpha is not None or moe.swiglu_beta is not None:
+        unsupported.append("custom SwiGLU alpha or beta")
+    if (
+        spec.kernel == "deep_gemm"
+        and moe.routing_method is not RoutingMethodType.DeepseekV4
+    ):
+        unsupported.append(f"routing method {moe.routing_method.name}")
+
+    vllm_config = get_current_vllm_config()
+    if vllm_config.weight_transfer_config is not None:
+        unsupported.append("runtime weight transfer")
+    if vllm_config.parallel_config.enable_dbo:
+        unsupported.append("dual batch overlap")
+    if spec.kernel == "deep_gemm" and vllm_config.parallel_config.enable_eplb:
+        unsupported.append("EPLB")
+    return tuple(unsupported)
+
+
+def validate_flashinfer_moe_ep_layer(layer: RoutedExperts) -> None:
+    if layer.expert_map_manager.placement_strategy != "linear":
+        raise ValueError(
+            f"{layer.moe_config.moe_backend} requires contiguous linear "
+            "expert placement"
+        )
+
+
+def apply_topk_in_fc1(
+    moe: FusedMoEConfig,
+    *,
+    apply_router_weight_on_input: bool = False,
+) -> bool:
+    return (
+        apply_router_weight_on_input
+        or moe.routing_method is RoutingMethodType.DeepseekV4
+    )
+
+
+def _dequant_mxfp4_ue8m0_gran32(
+    packed: torch.Tensor,
+    scale: torch.Tensor,
+) -> torch.Tensor:
+    raw = packed.view(torch.uint8)
+    lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device)
+    values = torch.empty(
+        raw.shape[0],
+        raw.shape[1] * 2,
+        dtype=torch.float32,
+        device=raw.device,
+    )
+    values[:, ::2] = lut[(raw & 0x0F).to(torch.int64)]
+    values[:, 1::2] = lut[(raw >> 4).to(torch.int64)]
+    encoded_scale = (
+        scale
+        if scale.dtype == torch.float8_e8m0fnu
+        else scale.view(torch.float8_e8m0fnu)
+    )
+    decoded_scale = encoded_scale.float()
+    return (values * decoded_scale.repeat_interleave(32, dim=-1)).to(torch.bfloat16)
+
+
+def _dequant_mxfp4_experts(
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+) -> torch.Tensor:
+    num_experts, output_size, packed_input_size = weight.shape
+    result = torch.empty(
+        num_experts,
+        output_size,
+        packed_input_size * 2,
+        dtype=torch.bfloat16,
+        device=weight.device,
+    )
+    for expert_id in range(num_experts):
+        result[expert_id] = _dequant_mxfp4_ue8m0_gran32(
+            weight[expert_id], scale[expert_id]
+        )
+    return result
+
+
+def mxfp4_moe_ep_weights(
+    moe_backend: str,
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+) -> FlashInferMoeEpWeights:
+    spec = flashinfer_moe_ep_backend_spec(moe_backend)
+    if spec.kernel == "deep_gemm":
+        return FlashInferMoeEpWeights(w13, w2, w13_scale, w2_scale)
+    return FlashInferMoeEpWeights(
+        w13=_dequant_mxfp4_experts(w13, w13_scale),
+        w2=_dequant_mxfp4_experts(w2, w2_scale),
+    )
+
+
+class _FlashInferMoeEpApi:
+    def __init__(self) -> None:
+        try:
+            from flashinfer.moe_ep import (
+                BootstrapConfig,
+                DeepGemmMegaMoeConfig,
+                FleetParams,
+                MegaConfig,
+                MoEEpMegaLayer,
+                MoEEpTensors,
+                MoEWeightPack,
+                Nvfp4CutedslMegaMoeConfig,
+            )
+        except ImportError as err:
+            raise RuntimeError(
+                "FlashInfer MoE-EP requires the public API shipped with "
+                "FlashInfer 0.7.0 or newer"
+            ) from err
+
+        self.BootstrapConfig = BootstrapConfig
+        self.DeepGemmMegaMoeConfig = DeepGemmMegaMoeConfig
+        self.FleetParams = FleetParams
+        self.MegaConfig = MegaConfig
+        self.MoEEpMegaLayer = MoEEpMegaLayer
+        self.MoEEpTensors = MoEEpTensors
+        self.MoEWeightPack = MoEWeightPack
+        self.Nvfp4CutedslMegaMoeConfig = Nvfp4CutedslMegaMoeConfig
+
+
+def _load_flashinfer_moe_ep_api() -> _FlashInferMoeEpApi:
+    return _FlashInferMoeEpApi()
+
+
+def _expose_deep_gemm_to_flashinfer() -> None:
+    deep_gemm = _import_deep_gemm()
+    if deep_gemm is None:
+        raise RuntimeError(
+            "FlashInfer MoE-EP DeepGEMM requires a usable DeepGEMM installation"
+        )
+    sys.modules.setdefault("deep_gemm", deep_gemm)
+
+
+class FlashInferMoeEp:
+    def __init__(
+        self,
+        moe: FusedMoEConfig,
+        weights: FlashInferMoeEpWeights,
+        epilogue: FlashInferMoeEpEpilogue | None = None,
+        *,
+        apply_topk_in_fc1: bool,
+    ) -> None:
+        spec = flashinfer_moe_ep_backend_spec(moe.moe_backend)
+        if spec.kernel == "deep_gemm":
+            _expose_deep_gemm_to_flashinfer()
+        api = _load_flashinfer_moe_ep_api()
+        if epilogue is None:
+            epilogue = FlashInferMoeEpEpilogue()
+        ep_group = get_ep_group()
+        bootstrap = api.BootstrapConfig(
+            world_size=ep_group.world_size,
+            rank=ep_group.rank_in_group,
+            process_group=ep_group.device_group,
+            device=torch.accelerator.current_device_index(),
+        )
+        fleet_params = api.FleetParams(
+            num_experts=moe.num_experts,
+            max_tokens_per_rank=moe.max_num_tokens,
+            token_hidden_size=moe.hidden_dim,
+        )
+        weight_pack = api.MoEWeightPack(
+            w13=weights.w13,
+            w2=weights.w2,
+            w13_scale=weights.w13_scale,
+            w2_scale=weights.w2_scale,
+        )
+        if spec.kernel == "cutedsl":
+            megakernel = api.Nvfp4CutedslMegaMoeConfig(
+                intermediate_size=moe.intermediate_size,
+                top_k=moe.experts_per_token,
+                gate_up_clamp=moe.swiglu_limit,
+                fast_math=True,
+                apply_topk_in_fc1=apply_topk_in_fc1,
+                enable_in_kernel_fc2_reduce=False,
+                combine_dtype="bf16",
+                input_norm_const=epilogue.input_norm_const,
+                fc1_alpha=None,
+                fc2_alpha=None,
+                fc1_norm_const=None,
+                knobs=None,
+            )
+        else:
+            megakernel = api.DeepGemmMegaMoeConfig(
+                intermediate_size=moe.intermediate_size,
+                top_k=moe.experts_per_token,
+                activation_clamp=moe.swiglu_limit,
+                fast_math=True,
+            )
+        backend = api.MegaConfig(
+            megakernel=megakernel,
+            quantize_input=True,
+            preprocess_weights=True,
+        )
+
+        self._moe_ep_tensors_cls = api.MoEEpTensors
+        self._mega_layer: Any | None = api.MoEEpMegaLayer(
+            bootstrap,
+            fleet_params,
+            weight_pack,
+            backend,
+        )
+        self._device = weights.w13.device
+        self._hidden_size = moe.hidden_dim
+        self._max_num_tokens = moe.max_num_tokens
+        self._num_experts = moe.num_experts
+        self._top_k = moe.experts_per_token
+        self._epilogue: FlashInferMoeEpEpilogue | None = epilogue
+        self._return_workspace_view = moe.routing_method is RoutingMethodType.DeepseekV4
+
+    @property
+    def can_overlap_shared_experts(self) -> bool:
+        return False
+
+    @property
+    def output_is_reduced(self) -> bool:
+        return True
+
+    @property
+    def topk_indices_dtype(self) -> torch.dtype:
+        return torch.int32
+
+    @property
+    def is_monolithic(self) -> bool:
+        return False
+
+    def _tensors(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+    ) -> Any:
+        epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue)
+        return self._moe_ep_tensors_cls(
+            hidden_states=hidden_states,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            fc1_alpha=epilogue.fc1_alpha,
+            fc2_alpha=epilogue.fc2_alpha,
+            fc1_norm_const=epilogue.fc1_norm_const,
+        )
+
+    def __call__(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        if self._mega_layer is None:
+            raise RuntimeError("FlashInfer MoE-EP layer was destroyed")
+        if hidden_states.shape[0] > self._max_num_tokens:
+            raise ValueError(
+                f"FlashInfer MoE-EP got {hidden_states.shape[0]} tokens, but "
+                f"the workspace supports at most {self._max_num_tokens}"
+            )
+        tensors = self._tensors(hidden_states, topk_ids, topk_weights)
+        if self._return_workspace_view and getattr(
+            self._mega_layer, "supports_output_view", False
+        ):
+            return self._mega_layer.forward(tensors, return_workspace_view=True)
+        return self._mega_layer(tensors)
+
+    @torch.inference_mode()
+    def warmup(self) -> None:
+        if self._mega_layer is None:
+            return
+        hidden_states = torch.zeros(
+            1,
+            self._hidden_size,
+            dtype=torch.bfloat16,
+            device=self._device,
+        )
+        topk_ids = (
+            torch.arange(self._top_k, dtype=torch.int32, device=self._device)
+            .mul_(self._num_experts // self._top_k)
+            .view(1, self._top_k)
+        )
+        topk_weights = torch.full(
+            (1, self._top_k),
+            1.0 / self._top_k,
+            dtype=torch.float32,
+            device=self._device,
+        )
+        self._mega_layer.warmup(self._tensors(hidden_states, topk_ids, topk_weights))
+
+    def kernel_weights(
+        self,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Kernel-resident fc1 weight, fc1 scales, fc2 weight, fc2 scales.
+
+        Contiguous per-expert byte tensors that alias the memory the megakernel
+        reads, so a caller may permute experts in place (EPLB).
+        """
+        if self._mega_layer is None:
+            raise RuntimeError("FlashInfer MoE-EP layer was destroyed")
+        (fc1_weight, fc1_sf), (fc2_weight, fc2_sf) = self._mega_layer._transformed
+        return (
+            _expert_major_bytes(fc1_weight),
+            _expert_major_bytes(fc1_sf),
+            _expert_major_bytes(fc2_weight),
+            _expert_major_bytes(fc2_sf),
+        )
+
+    def expert_weight_views(self) -> tuple[torch.Tensor, ...]:
+        """Every per-expert tensor the kernel reads: weights plus epilogue."""
+        epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue)
+        vectors = (epilogue.fc1_alpha, epilogue.fc2_alpha, epilogue.fc1_norm_const)
+        return (*self.kernel_weights(), *(v for v in vectors if v is not None))
+
+    def destroy(self) -> None:
+        mega_layer = self._mega_layer
+        if mega_layer is None:
+            return
+        mega_layer.destroy()
+        self._mega_layer = None
+        self._epilogue = None
+
+    def __del__(self) -> None:
+        with contextlib.suppress(Exception):
+            self.destroy()
+
+
+def _expert_major_bytes(tensor: torch.Tensor) -> torch.Tensor:
+    """Contiguous storage of a per-expert kernel tensor, 1-byte dtypes as uint8.
+
+    The CuTeDSL kernel keeps (E, K, N) transpose views over (E, N, K) storage.
+    """
+    if not tensor.is_contiguous():
+        tensor = tensor.transpose(1, 2)
+    if not tensor.is_contiguous():
+        raise ValueError(f"expert tensor is not expert-major: {tensor.stride()}")
+    return tensor.view(torch.uint8) if tensor.element_size() == 1 else tensor

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

@@ -8,7 +8,11 @@ import torch
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm import envs
 from vllm.config import get_current_vllm_config
-from vllm.config.kernel import MoEBackend
+from vllm.config.kernel import (
+    FLASHINFER_MOE_EP_CUTEDSL,
+    FLASHINFER_MOE_EP_DEEP_GEMM,
+    MoEBackend,
+)
 from vllm.config.quantization import QuantizationConfigArgs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
@@ -145,6 +149,9 @@ class Mxfp4MoeBackend(Enum):
     EMULATION = "EMULATION"
     # Humming
     HUMMING = "HUMMING"
+    # FlashInfer MoE-EP megakernels: routing outside, dispatch/GEMMs/combine fused.
+    FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL"
+    FLASHINFER_MOE_EP_DEEP_GEMM = "FLASHINFER_MOE_EP_DEEP_GEMM"
 
 
 # Backends that share the same TRTLLM weight format
@@ -165,6 +172,13 @@ B12X_BACKENDS = (
     Mxfp4MoeBackend.B12X_MXFP4_BF16,
 )
 
+# Megakernels consume the checkpoint's MXFP4 tensors through the shared
+# FlashInfer MoE-EP adapter; vLLM does not convert their weight layout.
+FLASHINFER_MOE_EP_MXFP4_BACKENDS = (
+    Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL,
+    Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM,
+)
+
 
 def backend_to_kernel_cls(
     backend: Mxfp4MoeBackend,
@@ -291,6 +305,13 @@ def backend_to_kernel_cls(
 
         return [OCP_MXQuantizationEmulationTritonExperts]
 
+    elif backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
+        from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import (  # noqa: E501
+            FlashInferMoeEpExperts,
+        )
+
+        return [FlashInferMoeEpExperts]
+
     else:
         raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}")
 
@@ -314,6 +335,10 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
             Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
         ],
         "flashinfer_cutlass_afp8": [Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8],
+        "flashinfer_moe_ep_cutedsl": [Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL],
+        "flashinfer_moe_ep_mega_deep_gemm": [
+            Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM
+        ],
         "triton": [Mxfp4MoeBackend.TRITON],
         "triton_unfused": [Mxfp4MoeBackend.TRITON_UNFUSED],
         "humming": [Mxfp4MoeBackend.HUMMING],
@@ -735,6 +760,9 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
     activation: MoEActivation | None = None,
 ) -> tuple[int, int]:
     """Round up hidden_size and intermediate_size based on backend requirements."""
+    if backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
+        # The megakernel takes the model's dimensions and validates them itself.
+        return hidden_size, intermediate_size
     if backend in B12X_BACKENDS:
         # b12x plans for the exact model dimensions. B12xExperts validates the
         # required MXFP4 block alignment before selecting the backend.
@@ -1827,6 +1855,29 @@ def convert_weight_to_mxfp4_moe_kernel_format(
             w2_bias=w2_bias,
             _cache_permute_indices=_cache_permute_indices,
         )
+    elif mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
+        from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import (
+            mxfp4_moe_ep_weights,
+        )
+
+        # DeepGEMM consumes MXFP4 directly; CuTeDSL gets bf16 and requantizes.
+        weights = mxfp4_moe_ep_weights(
+            FLASHINFER_MOE_EP_CUTEDSL
+            if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL
+            else FLASHINFER_MOE_EP_DEEP_GEMM,
+            w13_weight,
+            w2_weight,
+            w13_weight_scale,
+            w2_weight_scale,
+        )
+        return (
+            weights.w13,
+            weights.w2,
+            weights.w13_scale,
+            weights.w2_scale,
+            w13_bias,
+            w2_bias,
+        )
     elif mxfp4_backend == Mxfp4MoeBackend.CPU:
         from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
             prepare_mxfp4_moe_layer_for_cpu,
@@ -1874,6 +1925,9 @@ def make_mxfp4_moe_quant_config(
     layer: "RoutedExperts | None" = None,
 ) -> FusedMoEQuantConfig | None:
     """Create a FusedMoEQuantConfig for the given MXFP4 backend."""
+    if mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
+        # The megakernel quantizes activations itself and owns its scales.
+        return FusedMoEQuantConfig.make("nvfp4", weight_dtype="mxfp4")
     if mxfp4_backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8:
         return mxfp4_mxfp8_moe_quant_config(
             w1_bias=w1_bias,

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

@@ -50,6 +50,7 @@ class NvFp4MoeBackend(Enum):
     MARLIN = "MARLIN"
     HUMMING = "HUMMING"
     EMULATION = "EMULATION"
+    FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL"
 
 
 FLASHINFER_NVFP4_MOE_BACKENDS = [
@@ -148,6 +149,12 @@ def backend_to_kernel_cls(
         )
 
         return [Nvfp4QuantizationEmulationTritonExperts]
+    elif backend == NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL:
+        from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import (  # noqa: E501
+            FlashInferMoeEpExperts,
+        )
+
+        return [FlashInferMoeEpExperts]
     else:
         raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}")
 
@@ -160,6 +167,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend:
         "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM,
         "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS,
         "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+        "flashinfer_moe_ep_cutedsl": NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL,
         "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X,
         "marlin": NvFp4MoeBackend.MARLIN,
         "humming": NvFp4MoeBackend.HUMMING,
@@ -207,6 +215,7 @@ def select_nvfp4_moe_backend(
         NvFp4MoeBackend.FLASHINFER_TRTLLM,
         NvFp4MoeBackend.FLASHINFER_CUTLASS,
         NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+        NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL,
         NvFp4MoeBackend.VLLM_CUTLASS,
         NvFp4MoeBackend.MARLIN,
         NvFp4MoeBackend.EMULATION,
@@ -281,8 +290,9 @@ def select_nvfp4_moe_backend(
                 f"Model sets swiglu_limit={config.swiglu_limit}, but the "
                 f"explicitly requested moe_backend={runner_backend!r} does "
                 f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', "
-                f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', "
-                f"'b12x', 'marlin', or 'humming' instead."
+                f"'flashinfer_cutlass', 'flashinfer_cutedsl', "
+                f"'flashinfer_moe_ep_cutedsl', 'cutlass', 'b12x', 'marlin', "
+                f"or 'humming' instead."
             )
         return _return_or_raise(
             requested_backend, config, weight_key, activation_key, activation_format
@@ -491,6 +501,9 @@ def convert_to_nvfp4_moe_kernel_format(
         # for other experts - other selection strategies may be used.
         a13_scale = 1.0 / a13_scale.max().to(torch.float32)
         a2_scale = 1.0 / a2_scale.max().to(torch.float32)
+    elif nvfp4_backend == NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL:
+        # The megakernel transforms the checkpoint tensors itself at init.
+        pass
     else:
         raise ValueError(f"Unknown NvFp4 backend for MoE: {nvfp4_backend}")
 

vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Prepare/finalize stages that do no expert-parallel communication.
+
+Selected by ``all2all_backend="passthrough"``: the experts implementation
+dispatches, combines and reduces itself, so nothing is left for these stages.
+"""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import (
+    MoEPrepareAndFinalizeNoDPEPModular,
+)
+
+
+class PassThroughPrepareAndFinalize(MoEPrepareAndFinalizeNoDPEPModular):
+    """Hands activations and routing straight to experts that own their EP
+    communication."""
+
+    def output_is_reduced(self) -> bool:
+        return True
+
+    def supports_deferred_moe_finalize(self) -> bool:
+        # The experts apply the top-k weights and combine internally; there
+        # is no finalize left for a consumer to take over.
+        return False
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.PrepareResultType:
+        return a1, None, None, None, None

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py

@@ -24,10 +24,15 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
     XPUExpertsMxFp4,
 )
+from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import (
+    is_flashinfer_moe_ep_backend,
+)
 from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
     B12X_BACKENDS,
+    FLASHINFER_MOE_EP_MXFP4_BACKENDS,
     Mxfp4MoeBackend,
+    convert_weight_to_mxfp4_moe_kernel_format,
     make_mxfp4_moe_kernel,
     make_mxfp4_moe_quant_config,
     select_mxfp4_moe_backend,
@@ -38,7 +43,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
     prepare_moe_fp4_layer_for_marlin,
 )
-from vllm.model_executor.utils import set_weight_attrs
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
 from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
@@ -49,12 +54,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
         super().__init__(moe)
         self.group_size = 32
         self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
+        self.experts_cls: type[mk.FusedMoEExperts] | None = None
+
         # Backend selection must match the weight preparation below: CUTLASS
-        # swizzles scales, b12x and XPU consume checkpoint packing, and Marlin
-        # repacks weights and scales.
+        # swizzles scales, b12x and XPU consume checkpoint packing, Marlin
+        # repacks weights and scales, and the megakernels take the checkpoint
+        # tensors as they are.
         self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
-        self.experts_cls: type[mk.FusedMoEExperts]
-        if moe.moe_backend == "b12x":
+        if moe.moe_backend == "b12x" or is_flashinfer_moe_ep_backend(moe.moe_backend):
             self.mxfp4_backend, experts_cls = select_mxfp4_moe_backend(moe)
             assert experts_cls is not None
             self.experts_cls = experts_cls
@@ -70,6 +77,10 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
             logger.info_once("Using MarlinExperts for MXFP4 MoE")
             self.experts_cls = MarlinExperts
 
+    @property
+    def supports_eplb(self) -> bool:
+        return self.mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -198,6 +209,21 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
             layer.w2_weight_scale = torch.nn.Parameter(
                 torch.stack(swizzled_w2), requires_grad=False
             )
+        elif self.mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS:
+            w13, w2, w13_scale, w2_scale, _, _ = (
+                convert_weight_to_mxfp4_moe_kernel_format(
+                    mxfp4_backend=self.mxfp4_backend,
+                    layer=layer,
+                    w13_weight=layer.w13_weight,
+                    w2_weight=layer.w2_weight,
+                    w13_weight_scale=layer.w13_weight_scale,
+                    w2_weight_scale=layer.w2_weight_scale,
+                )
+            )
+            replace_parameter(layer, "w13_weight", w13)
+            replace_parameter(layer, "w2_weight", w2)
+            replace_parameter(layer, "w13_weight_scale", w13_scale)
+            replace_parameter(layer, "w2_weight_scale", w2_scale)
         elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu():
             pass
         else:
@@ -211,6 +237,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
 
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
         if self.moe_quant_config is not None:
+            assert self.experts_cls is not None
             self.moe_kernel = make_mxfp4_moe_kernel(
                 moe_quant_config=self.moe_quant_config,
                 moe_config=self.moe,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py

@@ -38,6 +38,7 @@ _EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset(
         NvFp4MoeBackend.FLASHINFER_CUTEDSL,
         NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
         NvFp4MoeBackend.FLASHINFER_TRTLLM,
+        NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL,
     }
 )
 
@@ -61,7 +62,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
             weight_key=kNvfp4Static,
             activation_key=None if use_a16 else kNvfp4Dynamic,
         )
-
         self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
             self.nvfp4_backend
         )
@@ -155,7 +155,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
         set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
 
         w2_weight_scale_2 = torch.nn.Parameter(
-            torch.empty(num_experts, dtype=torch.float32), requires_grad=False
+            torch.empty(num_experts, dtype=torch.float32),
+            requires_grad=False,
         )
         layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2)
         extra_weight_attrs.update(
@@ -175,7 +176,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
         set_weight_attrs(w13_input_scale, extra_weight_attrs)
 
         w2_input_scale = torch.nn.Parameter(
-            torch.empty(num_experts, dtype=torch.float32), requires_grad=False
+            torch.empty(num_experts, dtype=torch.float32),
+            requires_grad=False,
         )
         layer.register_parameter("w2_input_global_scale", w2_input_scale)
         extra_weight_attrs.update(

vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py

@@ -4,6 +4,7 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
     FusedMoeWeightScaleSupported,
@@ -53,6 +54,11 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):
 
     def __init__(self, moe) -> None:
         super().__init__(moe)
+        if moe.moe_backend in FLASHINFER_MOE_EP_BACKENDS:
+            raise ValueError(
+                f"moe_backend={moe.moe_backend!r} is not supported for AutoRound "
+                "(INC) MXFP4 MoE checkpoints"
+            )
         self.group_size = 32
         # Backend selection must stay consistent with the weight preparation in
         # process_weights_after_loading / get_fused_moe_quant_config, which use

vllm/model_executor/layers/quantization/mxfp4.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+
 import torch
 
 from vllm.logger import init_logger
@@ -783,7 +784,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             )
             self.moe_kernel.fused_experts.process_weights_after_loading(layer)
 
-    def process_weights_after_loading(self, layer):
+    def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
             return
 

vllm/model_executor/layers/quantization/utils/quant_utils.py

@@ -25,6 +25,11 @@ INT4_DTYPE = scalar_types.uint4b8
 INT8_DTYPE = scalar_types.uint8b128
 
 
+def require_finite_positive_scales(name: str, value: torch.Tensor) -> None:
+    if not torch.isfinite(value).all() or not (value > 0).all():
+        raise ValueError(f"{name} must contain finite positive scales")
+
+
 def _dtype_abbr(dtype: torch.dtype | ScalarType) -> str:
     """Return a stable short name for torch and ScalarType dtypes."""
     if isinstance(dtype, ScalarType):

vllm/models/deepseek_v4/amd/dspark.py

@@ -11,8 +11,9 @@ ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for
     tilelang / triton / torch) instead of calling the tilelang kernels directly,
     and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER
     or TileLang fused MHC is available; False only on the torch fallback);
-  * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` /
-    ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py).
+  * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping``,
+    ``use_native_mega_moe`` and ``finalize_mega_moe_weights`` do not exist in
+    amd/model.py).
 
 Everything else — the semi-autoregressive drafting hooks, the Markov head, the
 sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is

vllm/models/deepseek_v4/nvidia/dspark.py

@@ -17,7 +17,7 @@ import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.config import VllmConfig, get_current_vllm_config
-from vllm.config.kernel import MEGA_MOE_BACKENDS
+from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -99,7 +99,9 @@ class DSparkDeepseekV4Model(nn.Module):
         self.rms_norm_eps = config.rms_norm_eps
         self.num_hidden_layers = config.num_hidden_layers
         self.target_layer_ids = tuple(config.dspark_target_layer_ids)
-        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+        self.use_native_mega_moe = (
+            vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS
+        )
         self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
 
         self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
@@ -240,7 +242,7 @@ class DSparkDeepseekV4Model(nn.Module):
             inputs_embeds = sp_shard(inputs_embeds)
             input_ids = sp_shard(input_ids)
         mega_gate_metadata = None
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             mega_gate_metadata = prepare_mega_gate_routing_metadata(
                 input_ids,
                 has_hash_routing=False,
@@ -437,8 +439,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
         are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
         """
         first_layer = self.model.layers[0]
-        use_mega_moe = first_layer.ffn.use_mega_moe
-        if use_mega_moe:
+        use_native_mega_moe = first_layer.ffn.use_native_mega_moe
+        if use_native_mega_moe:
             expert_mapping = make_deepseek_v4_expert_params_mapping(
                 self.config.n_routed_experts
             )

vllm/models/deepseek_v4/nvidia/fi_moe.py

@@ -1,397 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""FlashInfer ``moe_ep`` expert module for DeepSeek V4."""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Any
-
-import torch
-import torch.nn as nn
-
-from vllm.model_executor.utils import set_weight_attrs
-from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts
-from vllm.utils.flashinfer_moe_ep import (
-    build_fi_mega_layer,
-    ensure_fi_moe_ep_runtime,
-    fi_moe_ep_backend_spec,
-    make_fi_moe_ep_bootstrap,
-    mega_moe_weight_pack_from_params,
-)
-
-if TYPE_CHECKING:
-    from flashinfer.moe_ep import MoEEpMegaLayer
-
-    from vllm.config import VllmConfig
-    from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MLP
-
-_MOE_SKIP_PADDING: bool | None = None
-
-
-def ckpt_uses_nvfp4_experts(vllm_config: VllmConfig) -> bool:
-    """True when the loaded checkpoint quantizes experts with modelopt NVFP4
-    (e2m1 + fp8-e4m3 per-16 block scales + per-tensor weight_scale_2), i.e.
-    the recipe the nvfp4_cutedsl prequantized-weights path consumes verbatim.
-
-    DeepSeek-V4 NVFP4-expert checkpoints declare ``moe_quant_algo: NVFP4`` in
-    their ``quantization_config``, which DeepseekV4FP8Config surfaces as
-    ``moe_quant_algo``.
-    """
-    return getattr(vllm_config.quant_config, "moe_quant_algo", "") == "NVFP4"
-
-
-def resolve_mega_moe_is_padding(num_tokens: int) -> torch.Tensor | None:
-    from vllm.forward_context import get_forward_context, is_forward_context_available
-
-    global _MOE_SKIP_PADDING
-    if _MOE_SKIP_PADDING is None:
-        import vllm.envs as envs
-
-        _MOE_SKIP_PADDING = bool(envs.VLLM_MOE_SKIP_PADDING)
-    if not _MOE_SKIP_PADDING or not is_forward_context_available():
-        return None
-    is_padding = get_forward_context().is_padding
-    if is_padding is None:
-        return None
-    return is_padding[:num_tokens]
-
-
-def apply_mega_moe_routing_preprocess(
-    topk_ids: torch.Tensor,
-    *,
-    is_padding: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Padding-only routing preprocess (EPLB hooks go here later)."""
-    if is_padding is not None:
-        topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)
-    return topk_ids
-
-
-def nvfp4_prequant_pack_and_alphas(
-    w13_weight: torch.Tensor,
-    w13_weight_scale: torch.Tensor,
-    w13_weight_scale_2: torch.Tensor,  # (E_local, 2) fp32: [:,0]=gate(w1), [:,1]=up(w3)
-    w2_weight: torch.Tensor,
-    w2_weight_scale: torch.Tensor,
-    w2_weight_scale_2: torch.Tensor,  # (E_local,) fp32
-    *,
-    intermediate_size: int,
-):
-    """NVFP4 checkpoint params -> (MoEWeightPack, fc1_alpha, fc2_alpha).
-
-    Activation quantization is fully dynamic, so the checkpoint's static
-    ``input_scale`` drops out and each GEMM's epilogue alpha reduces to the
-    weight's per-tensor ``weight_scale_2``.
-
-    fc1_alpha is one scalar per expert, but gate (w1) and up (w3) carry their
-    own ``weight_scale_2``. When they differ, the ratio is folded into the up
-    half's e4m3 block scales — only if the fold round-trips exactly, since a
-    lossy rescale would silently change the model.
-    """
-    from flashinfer.moe_ep import MoEWeightPack
-
-    inter = intermediate_size
-    gate_s2 = w13_weight_scale_2[:, 0].float()
-    up_s2 = w13_weight_scale_2[:, 1].float()
-    if (gate_s2 <= 0).any() or (up_s2 <= 0).any() or (w2_weight_scale_2 <= 0).any():
-        raise ValueError(
-            "nvfp4 prequant: non-positive weight_scale_2 loaded — checkpoint "
-            "scale tensors missing or loader routed them wrong."
-        )
-
-    w13_scale = w13_weight_scale
-    if not torch.equal(gate_s2, up_s2):
-        ratio = up_s2 / gate_s2  # (E_local,)
-        up_sf = w13_scale[:, inter:, :].float()
-        folded = up_sf * ratio[:, None, None]
-        folded_e4m3 = folded.to(torch.float8_e4m3fn)
-        if not torch.equal(folded_e4m3.float(), folded):
-            raise ValueError(
-                "NVFP4 MoE cannot merge the gate and up weight_scale_2 values "
-                "because their ratio is not exactly representable in "
-                "float8_e4m3fn. Use a checkpoint quantized with a shared "
-                "gate/up weight_scale_2."
-            )
-        w13_scale = w13_scale.clone()
-        w13_scale[:, inter:, :] = folded_e4m3
-
-    fc1_alpha = gate_s2.clone().contiguous()
-    fc2_alpha = w2_weight_scale_2.float().clone().contiguous()
-    pack = MoEWeightPack(
-        w13=w13_weight,
-        w2=w2_weight,
-        w13_scale=w13_scale,
-        w2_scale=w2_weight_scale,
-    )
-    return pack, fc1_alpha, fc2_alpha
-
-
-class DeepseekV4MegaMoEExpertsFI(DeepseekV4MegaMoEExperts):
-    """Same weight layout/loader as the native mega experts, FI compute path."""
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        *,
-        activation_clamp: float | None = None,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(vllm_config, **kwargs)
-        self._vllm_config = vllm_config
-        self._activation_clamp = activation_clamp
-        self._mega_layer: MoEEpMegaLayer | None = None
-        self._fast_ctx: tuple[Any, Any, Any, int, bool] | None = None
-        self._epilogue_alphas: tuple[torch.Tensor, torch.Tensor] | None = None
-        self._nvfp4_prequant = ckpt_uses_nvfp4_experts(vllm_config)
-        if self._nvfp4_prequant:
-            megakernel = fi_moe_ep_backend_spec(
-                vllm_config.kernel_config.moe_backend
-            ).megakernel
-            if megakernel != "nvfp4_cutedsl":
-                raise ValueError(
-                    "NVFP4-quantized expert checkpoint requires "
-                    "moe_backend=flashinfer_moe_ep_mega_cutedsl, got a "
-                    f"backend using megakernel {megakernel!r} "
-                    "(deep_gemm consumes the MXFP4 checkpoint instead)."
-                )
-            self._realloc_nvfp4_params()
-
-    def _realloc_nvfp4_params(self) -> None:
-        """Swap the mx-recipe scale params for the NVFP4 checkpoint's:
-        fp8-e4m3 per-16 block scales plus the per-tensor second-level
-        scales (weight_scale_2, input_scale) modelopt exports."""
-        n_e = self.num_local_experts
-        inter = self.intermediate_size
-        hidden = self.hidden_size
-        attrs = {"weight_loader": self.weight_loader}
-
-        def _param(shape: tuple, dtype: torch.dtype) -> nn.Parameter:
-            p = nn.Parameter(torch.zeros(*shape, dtype=dtype), requires_grad=False)
-            set_weight_attrs(p, attrs)
-            return p
-
-        self.w13_weight_scale = _param(
-            (n_e, 2 * inter, hidden // 16), torch.float8_e4m3fn
-        )
-        self.w13_weight_scale.quant_method = "block"
-        self.w2_weight_scale = _param((n_e, hidden, inter // 16), torch.float8_e4m3fn)
-        self.w2_weight_scale.quant_method = "block"
-        # (E, 2): column 0 = gate (w1), column 1 = up (w3).
-        self.w13_weight_scale_2 = _param((n_e, 2), torch.float32)
-        self.w2_weight_scale_2 = _param((n_e,), torch.float32)
-        # Loaded for completeness; unused — activation quant is dynamic
-        # (see nvfp4_prequant_pack_and_alphas).
-        self.w13_input_scale = _param((n_e, 2), torch.float32)
-        self.w2_input_scale = _param((n_e,), torch.float32)
-
-    def weight_loader(
-        self,
-        param: nn.Parameter,
-        loaded_weight: torch.Tensor,
-        weight_name: str,
-        shard_id: str,
-        expert_id: int,
-        return_success: bool = False,
-    ) -> bool | None:
-        # NVFP4 checkpoint second-level scalars route here; everything
-        # else (packed weights, block scales) matches the base layout.
-        if "weight_scale_2" in weight_name or "input_scale" in weight_name:
-            local_expert_ids = self._map_global_expert_id(expert_id)
-            if not local_expert_ids:
-                return False if return_success else None
-            value = loaded_weight.reshape(()).to(torch.float32)
-            for local_expert_id in local_expert_ids:
-                if shard_id in ("w1", "w3"):
-                    if "w13_" not in weight_name:
-                        return False if return_success else None
-                    param.data[local_expert_id, 0 if shard_id == "w1" else 1] = value
-                elif shard_id == "w2":
-                    if "w2_" not in weight_name:
-                        return False if return_success else None
-                    param.data[local_expert_id] = value
-                else:
-                    raise ValueError(f"Unsupported expert shard id: {shard_id}")
-            return True if return_success else None
-        return super().weight_loader(
-            param,
-            loaded_weight,
-            weight_name,
-            shard_id,
-            expert_id,
-            return_success,
-        )
-
-    def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
-        # The FlashInfer megakernel has no shared-expert fusion; the caller's
-        # serial shared MLP path handles shared_experts.
-        if self._mega_layer is not None:
-            return
-        if self.w13_weight is None:
-            return
-
-        self._check_runtime_supported()
-        ensure_fi_moe_ep_runtime(self._vllm_config)
-
-        if self._nvfp4_prequant:
-            # NVFP4 checkpoint: hand the packed weights + both scale
-            # planes straight to the backend (no dequant->requant);
-            # per-expert globals become fc1/fc2 epilogue alphas staged
-            # at every forward via MoEEpTensors.
-            weights, fc1_alpha, fc2_alpha = nvfp4_prequant_pack_and_alphas(
-                self.w13_weight.data,
-                self.w13_weight_scale.data,
-                self.w13_weight_scale_2.data,
-                self.w2_weight.data,
-                self.w2_weight_scale.data,
-                self.w2_weight_scale_2.data,
-                intermediate_size=self.intermediate_size,
-            )
-            self._epilogue_alphas = (fc1_alpha, fc2_alpha)
-        else:
-            weights = mega_moe_weight_pack_from_params(
-                self.w13_weight,
-                self.w13_weight_scale,
-                self.w2_weight,
-                self.w2_weight_scale,
-                megakernel=fi_moe_ep_backend_spec(
-                    self._vllm_config.kernel_config.moe_backend
-                ).megakernel,
-            )
-        self._mega_layer = build_fi_mega_layer(
-            make_fi_moe_ep_bootstrap(),
-            vllm_config=self._vllm_config,
-            num_experts=self.num_experts,
-            max_tokens_per_rank=self.max_num_tokens,
-            hidden_size=self.hidden_size,
-            intermediate_size=self.intermediate_size,
-            top_k=self.top_k,
-            activation_clamp=self._activation_clamp,
-            weights=weights,
-        )
-        del weights
-        # Allocate (or attach to) the pooled workspace before first
-        # forward so warmup/capture never hits the lazy path.
-        self._mega_layer._ensure_workspace()
-        self.w13_weight = None
-        self.w13_weight_scale = None
-        self.w2_weight = None
-        self.w2_weight_scale = None
-        if self._nvfp4_prequant:
-            self.w13_weight_scale_2 = None
-            self.w2_weight_scale_2 = None
-            self.w13_input_scale = None
-            self.w2_input_scale = None
-
-    # EPLB is rejected for these backends in validate_fi_moe_ep_config, so
-    # nothing constructs an EplbState and none of these run. They are kept
-    # as explicit errors rather than deleted because inheriting the native
-    # implementations would be worse: set_eplb_state would record a map
-    # this forward path never consults, silently splitting routing from
-    # weights, and get_expert_weights reads transformed-weight attributes
-    # the flashinfer path releases after preprocess.
-    _EPLB_UNSUPPORTED = (
-        "EPLB is not supported with the flashinfer moe_ep backends; "
-        "validate_fi_moe_ep_config should have rejected this "
-        "configuration at startup."
-    )
-
-    def set_eplb_state(
-        self,
-        moe_layer_idx: int,
-        expert_load_view: torch.Tensor,
-        logical_to_physical_map: torch.Tensor,
-        logical_replica_count: torch.Tensor,
-    ) -> None:
-        raise NotImplementedError(self._EPLB_UNSUPPORTED)
-
-    def get_expert_weights(self) -> list[torch.Tensor]:
-        raise NotImplementedError(self._EPLB_UNSUPPORTED)
-
-    def update_expert_map(self) -> None:
-        raise NotImplementedError(self._EPLB_UNSUPPORTED)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        *,
-        activation_clamp: float | None,
-        fast_math: bool = True,
-    ) -> torch.Tensor:
-        # fast_math is a native deep_gemm knob; the FI kernels have no
-        # equivalent toggle, so it is accepted for signature parity only.
-        if hidden_states.shape[0] > self.max_num_tokens:
-            raise ValueError(
-                f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
-                f"but the symmetric buffer was sized for {self.max_num_tokens}."
-            )
-
-        from flashinfer.moe_ep import MoEEpTensors
-
-        num_tokens = hidden_states.shape[0]
-        is_padding = resolve_mega_moe_is_padding(num_tokens)
-        topk_ids = apply_mega_moe_routing_preprocess(
-            topk_ids,
-            is_padding=is_padding,
-        )
-
-        # Fast path: after the first successful full forward the layer is
-        # immutable, so skip MoEEpMegaLayer.forward()'s per-call validation
-        # and go straight to the kernel backend's stage_inputs + compute.
-        alphas = self._epilogue_alphas
-        fc1_alpha = alphas[0] if alphas is not None else None
-        fc2_alpha = alphas[1] if alphas is not None else None
-
-        fast = self._fast_ctx
-        if fast is not None:
-            kernel, workspace, transformed, hidden_size, zero_copy = fast
-            t = MoEEpTensors(
-                hidden_states=hidden_states,
-                topk_ids=topk_ids,
-                topk_weights=topk_weights,
-                fc1_alpha=fc1_alpha,
-                fc2_alpha=fc2_alpha,
-            )
-            kernel.stage_inputs(t, workspace, quantize_input=True)
-            if zero_copy:
-                # Zero-copy (cutedsl backends): consume the workspace [:n]
-                # view directly — valid under stream ordering until the next
-                # MoE layer's launch on the shared workspace.
-                return kernel.compute(workspace, transformed, output=None)
-            # deep_gemm_mega's compute() requires a real output tensor.
-            out = torch.empty(
-                num_tokens,
-                hidden_size,
-                dtype=torch.bfloat16,
-                device=hidden_states.device,
-            )
-            return kernel.compute(workspace, transformed, output=out)
-
-        ensure_fi_moe_ep_runtime(self._vllm_config)
-        self.finalize_weights()
-        assert self._mega_layer is not None
-
-        y = self._mega_layer.forward(
-            MoEEpTensors(
-                hidden_states=hidden_states,
-                topk_ids=topk_ids,
-                topk_weights=topk_weights,
-                fc1_alpha=fc1_alpha,
-                fc2_alpha=fc2_alpha,
-            )
-        )
-        layer = self._mega_layer
-        if hidden_states.dtype == torch.bfloat16:
-            self._fast_ctx = (
-                layer._kernel,
-                layer._ensure_workspace(),
-                layer._transformed,
-                layer._fleet_params.token_hidden_size,
-                # zero-copy output views are a cutedsl-backend contract
-                layer._kernel.kernel_name() != "deep_gemm_mega",
-            )
-        return y
-
-
-DeepseekV4MegaMoEExpertsFI.weight_loader.supports_moe_loading = True  # type: ignore[attr-defined]

vllm/models/deepseek_v4/nvidia/model.py

@@ -11,7 +11,7 @@ import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
-from vllm.config.kernel import MEGA_MOE_BACKENDS
+from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_ep_group,
     get_pp_group,
@@ -85,10 +85,6 @@ from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.utils.flashinfer_moe_ep import (
-    is_fi_moe_ep_backend,
-    validate_fi_moe_ep_config,
-)
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
@@ -807,10 +803,11 @@ class DeepseekV4MoE(nn.Module):
         self.use_sequence_parallel = use_sequence_parallel
         self.reduce_results = reduce_results
         moe_backend = vllm_config.kernel_config.moe_backend
-        validate_fi_moe_ep_config(vllm_config)
-        self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS
-        self.use_fi_mega_moe = is_fi_moe_ep_backend(moe_backend)
-        if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
+        self.use_native_mega_moe = moe_backend in NATIVE_MEGA_MOE_BACKENDS
+        if (
+            self.use_native_mega_moe
+            and not vllm_config.parallel_config.enable_expert_parallel
+        ):
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE currently requires expert parallel. "
                 "Enable it with --enable-expert-parallel, or pick a different "
@@ -832,15 +829,15 @@ class DeepseekV4MoE(nn.Module):
         self.swiglu_limit = config.swiglu_limit
         self.renormalize = config.norm_topk_prob
         self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
-        if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
+        if self.use_native_mega_moe and self.scoring_func != "sqrtsoftplus":
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
             )
-        if self.use_mega_moe and not self.renormalize:
+        if self.use_native_mega_moe and not self.renormalize:
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE requires normalized top-k probabilities."
             )
-        if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
+        if self.use_native_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
                 f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
@@ -864,7 +861,9 @@ class DeepseekV4MoE(nn.Module):
             image_sentinel_lo if getattr(config, "vision_n_layers", 0) > 0 else 0
         )
         is_hash_moe = extract_layer_index(prefix) < num_hash_layers
-        self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
+        self.hash_indices_dtype = (
+            torch.int64 if self.use_native_mega_moe else torch.int32
+        )
         if is_hash_moe:
             # hash MoE doesn't use e_score_correction_bias
             # Use randint instead of empty to avoid garbage values causing
@@ -908,12 +907,12 @@ class DeepseekV4MoE(nn.Module):
                 hidden_act=config.hidden_act,
                 swiglu_limit=self.swiglu_limit,
                 quant_config=quant_config,
-                reduce_results=self.use_mega_moe,
+                reduce_results=self.use_native_mega_moe,
                 is_sequence_parallel=use_sequence_parallel,
                 prefix=f"{prefix}.shared_experts",
             )
 
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             self._init_mega_moe_experts(vllm_config, config, prefix)
         else:
             self._init_fused_moe_experts(vllm_config, config, quant_config, prefix)
@@ -950,28 +949,13 @@ class DeepseekV4MoE(nn.Module):
         # Native DeepGEMM fusion requires each EP rank to own the complete
         # shared MLP. Sequence parallel replicates those weights while sharding
         # tokens. TP=1 is also naturally replicated. With PP+TP the shared MLP
-        # remains tensor-sharded, so retain the serial path. The FlashInfer
-        # megakernel has no shared-expert fusion, so it keeps the serial path.
+        # remains tensor-sharded, so retain the serial path.
         fuse_shared_experts = bool(
             self.shared_experts is not None
             and not envs.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION
             and (self.use_sequence_parallel or self.tp_size == 1)
-            and not self.use_fi_mega_moe
         )
 
-        activation_clamp = (
-            float(self.swiglu_limit) if self.swiglu_limit is not None else None
-        )
-        if self.use_fi_mega_moe:
-            # Deferred: fi_moe subclasses DeepseekV4MegaMoEExperts, so a
-            # module-level import here would be circular.
-            from vllm.models.deepseek_v4.nvidia.fi_moe import (
-                DeepseekV4MegaMoEExpertsFI,
-            )
-
-            experts_cls: type[DeepseekV4MegaMoEExperts] = DeepseekV4MegaMoEExpertsFI
-        else:
-            experts_cls = DeepseekV4MegaMoEExperts
         expert_kwargs: dict[str, typing.Any] = dict(
             num_experts=self.n_physical_experts,
             num_local_experts=self.n_local_physical_experts,
@@ -983,9 +967,7 @@ class DeepseekV4MoE(nn.Module):
             num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0),
             prefix=f"{prefix}.experts",
         )
-        if self.use_fi_mega_moe:
-            expert_kwargs["activation_clamp"] = activation_clamp
-        self.experts = experts_cls(vllm_config, **expert_kwargs)
+        self.experts = DeepseekV4MegaMoEExperts(vllm_config, **expert_kwargs)
 
     def _init_fused_moe_experts(
         self,
@@ -1048,7 +1030,7 @@ class DeepseekV4MoE(nn.Module):
         if bias_vl is not None and input_ids is None:
             raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.")
 
-        if not self.use_mega_moe:
+        if not self.use_native_mega_moe:
             return self._forward_fused_moe(hidden_states, input_ids)
 
         org_shape = hidden_states.shape
@@ -1139,7 +1121,7 @@ class DeepseekV4MoE(nn.Module):
         return final_hidden_states.view(org_shape)
 
     def finalize_mega_moe_weights(self) -> None:
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             self.experts.finalize_weights(self.shared_experts)
 
 
@@ -1179,12 +1161,15 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
 
 def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
     parallel_config = vllm_config.parallel_config
-    use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+    moe_backend = vllm_config.kernel_config.moe_backend
+    moe_needs_token_sharded_input = (
+        moe_backend in MEGA_MOE_BACKENDS or parallel_config.data_parallel_size > 1
+    )
     return (
         parallel_config.pipeline_parallel_size == 1
         and parallel_config.enable_expert_parallel
         and parallel_config.tensor_parallel_size > 1
-        and (use_mega_moe or parallel_config.data_parallel_size > 1)
+        and moe_needs_token_sharded_input
     )
 
 
@@ -1425,9 +1410,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
         self.config = config
         self.quant_config = quant_config
         self.parallel_config = vllm_config.parallel_config
-        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+        self.use_native_mega_moe = (
+            vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS
+        )
         self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
-        if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
+        if (
+            self.use_native_mega_moe
+            and not vllm_config.parallel_config.enable_expert_parallel
+        ):
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE currently requires expert parallel. "
                 "Enable it with --enable-expert-parallel, or pick a different "
@@ -1564,7 +1554,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
 
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             input_ids = input_ids.to(torch.int64)
 
         full_num_tokens = positions.shape[0]
@@ -1578,7 +1568,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             input_ids = sp_shard(input_ids)
 
         mega_gate_metadata = None
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             mega_gate_metadata = prepare_mega_gate_routing_metadata(
                 input_ids,
                 has_hash_routing=self.has_local_hash_moe,
@@ -1802,7 +1792,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
-        if first_layer.ffn.use_mega_moe:
+        if first_layer.ffn.use_native_mega_moe:
             return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)

vllm/models/deepseek_v4/nvidia/mtp.py

@@ -20,7 +20,7 @@ import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
-from vllm.config.kernel import MEGA_MOE_BACKENDS
+from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -95,7 +95,9 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         self.config = config
         quant_config = vllm_config.quant_config
         self.rms_norm_eps = config.rms_norm_eps
-        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+        self.use_native_mega_moe = (
+            vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS
+        )
 
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -186,7 +188,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
             inputs_embeds
         ).unsqueeze(-2)
         mega_gate_metadata = None
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             routing_input_ids = input_ids
             if self.mtp_block.use_sequence_parallel:
                 routing_input_ids = sp_shard(routing_input_ids)
@@ -384,7 +386,7 @@ class DeepSeekV4MTP(nn.Module):
 
         # Pre-compute expert mapping ONCE.
         first_layer = next(iter(self.model.layers.values()))
-        if first_layer.mtp_block.ffn.use_mega_moe:
+        if first_layer.mtp_block.ffn.use_native_mega_moe:
             expert_mapping = make_deepseek_v4_expert_params_mapping(
                 self.config.n_routed_experts
             )

vllm/models/deepseek_v41/nvidia/dspark.py

@@ -18,7 +18,7 @@ import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.config import VllmConfig, get_current_vllm_config
-from vllm.config.kernel import MEGA_MOE_BACKENDS
+from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -81,7 +81,9 @@ class DSparkDeepseekV4Model(nn.Module):
         self.rms_norm_eps = config.rms_norm_eps
         self.num_hidden_layers = config.num_hidden_layers
         self.target_layer_ids = tuple(config.dspark_target_layer_ids)
-        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+        self.use_native_mega_moe = (
+            vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS
+        )
         self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
 
         self.num_dspark_layers = (
@@ -214,7 +216,7 @@ class DSparkDeepseekV4Model(nn.Module):
             inputs_embeds = sp_shard(inputs_embeds)
             input_ids = sp_shard(input_ids)
         mega_gate_metadata = None
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             mega_gate_metadata = prepare_mega_gate_routing_metadata(
                 input_ids,
                 has_hash_routing=False,
@@ -371,7 +373,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
         are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
         """
         first_layer = self.model.layers[0]
-        use_mega_moe = first_layer.ffn.use_mega_moe
+        use_native_mega_moe = first_layer.ffn.use_native_mega_moe
         # Draft MoE layers use the dspark_* expert counts, not the
         # backbone's (see DeepseekV4MoE and the reference
         # ModelArgs.get_moe_config).
@@ -379,7 +381,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
             getattr(self.config, "dspark_n_routed_experts", 0)
             or self.config.n_routed_experts
         )
-        if use_mega_moe:
+        if use_native_mega_moe:
             expert_mapping = make_deepseek_v4_expert_params_mapping(n_draft_experts)
         else:
             expert_mapping = fused_moe_make_expert_params_mapping(

vllm/models/deepseek_v41/nvidia/model.py

@@ -11,7 +11,7 @@ import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
-from vllm.config.kernel import MEGA_MOE_BACKENDS
+from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_engram_dp_size,
     get_pp_group,
@@ -174,7 +174,7 @@ class DeepseekV4MoE(DeepseekV4MoEBase):
 
     def defers_finalize(self, num_tokens: int) -> bool:
         return (
-            not self.use_mega_moe
+            not self.use_native_mega_moe
             and self.experts.moe_config.should_defer_moe_finalize(num_tokens)
         )
 
@@ -244,12 +244,14 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
 
 def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
     parallel_config = vllm_config.parallel_config
-    use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+    moe_needs_token_sharded_input = (
+        vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+    )
     return (
         parallel_config.pipeline_parallel_size == 1
         and parallel_config.enable_expert_parallel
         and parallel_config.tensor_parallel_size > 1
-        and (use_mega_moe or parallel_config.data_parallel_size > 1)
+        and (moe_needs_token_sharded_input or parallel_config.data_parallel_size > 1)
     )
 
 
@@ -631,9 +633,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
         self.config = config
         self.quant_config = quant_config
         self.parallel_config = vllm_config.parallel_config
-        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+        self.use_native_mega_moe = (
+            vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS
+        )
         self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
-        if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
+        if (
+            self.use_native_mega_moe
+            and not vllm_config.parallel_config.enable_expert_parallel
+        ):
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE currently requires expert parallel. "
                 "Enable it with --enable-expert-parallel, or pick a different "
@@ -865,7 +872,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
 
-        if self.use_mega_moe:
+        if self.use_native_mega_moe:
             input_ids = input_ids.to(torch.int64)
 
         # Engram n-gram hashes for the whole (flattened) batch, computed once
@@ -1015,7 +1022,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
     def _mega_gate_metadata(
         self, input_ids: torch.Tensor | None
     ) -> MegaGateRoutingMetadata | None:
-        if not self.use_mega_moe:
+        if not self.use_native_mega_moe:
             return None
         assert input_ids is not None
         return prepare_mega_gate_routing_metadata(
@@ -1331,7 +1338,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
-        if first_layer.ffn.use_mega_moe:
+        if first_layer.ffn.use_native_mega_moe:
             return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)

vllm/utils/flashinfer_moe_ep.py

@@ -1,339 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""FlashInfer ``moe_ep`` helpers for DeepSeek V4 vLLM integration."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
-
-import torch
-import torch.nn as nn
-
-from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS
-from vllm.distributed import get_ep_group
-from vllm.platforms import current_platform
-
-if TYPE_CHECKING:
-    from flashinfer.moe_ep import BootstrapConfig, MoEEpMegaLayer
-
-    from vllm.config import VllmConfig
-
-
-# Every mega kernel is Blackwell-only, so the arch is a property of the family
-# rather than of a particular backend: validate it against the live device
-# instead of encoding it in the backend name. An explicit allowlist so new
-# archs (SM110/SM120, Rubin) fail loudly until flashinfer supports them.
-FI_MOE_EP_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)})
-
-
-@dataclass(frozen=True)
-class FiMoeEpBackendSpec:
-    """Static properties of one ``flashinfer_moe_ep_*`` backend string.
-
-    The backend names the kernel *family*; the arch comes from the device and
-    the weight handling from the checkpoint, so this only has to carry which
-    megakernel to build and whether it needs NVSHMEM in the runtime set.
-    """
-
-    megakernel: str
-    needs_nvshmem: bool
-
-
-FI_MOE_EP_BACKEND_SPECS: dict[str, FiMoeEpBackendSpec] = {
-    # Consumes an MXFP4 checkpoint verbatim (e2m1 weights + E8M0 per-32
-    # scales) -- the same recipe the native deep_gemm mega path uses.
-    "flashinfer_moe_ep_mega_deep_gemm": FiMoeEpBackendSpec(
-        megakernel="deep_gemm_mega",
-        needs_nvshmem=False,
-    ),
-    # The checkpoint picks the weight path, not the kernel: an NVFP4
-    # checkpoint is consumed prequantized (no round trip), while MXFP4 weights
-    # are dequantized to bf16 and requantized. See
-    # models/deepseek_v4/nvidia/fi_moe.py:ckpt_uses_nvfp4_experts.
-    "flashinfer_moe_ep_mega_cutedsl": FiMoeEpBackendSpec(
-        megakernel="nvfp4_cutedsl",
-        needs_nvshmem=True,
-    ),
-}
-
-assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
-
-_FI_RUNTIME_HANDLE: Any = None
-
-
-def is_fi_moe_ep_backend(moe_backend: str) -> bool:
-    return moe_backend in FI_MOE_EP_BACKEND_SPECS
-
-
-def fi_moe_ep_backend_spec(moe_backend: str) -> FiMoeEpBackendSpec:
-    try:
-        return FI_MOE_EP_BACKEND_SPECS[moe_backend]
-    except KeyError:
-        raise ValueError(
-            f"{moe_backend!r} is not a flashinfer moe_ep backend; expected "
-            f"one of {sorted(FI_MOE_EP_BACKEND_SPECS)}"
-        ) from None
-
-
-def validate_fi_moe_ep_config(vllm_config: VllmConfig) -> None:
-    """Config-time checks for the mega-MoE backends, native and flashinfer."""
-    moe_backend = vllm_config.kernel_config.moe_backend
-    if not is_fi_moe_ep_backend(moe_backend):
-        return
-
-    # flashinfer validates the arch too, but not until the layer constructor
-    # runs during weight load; check here so the error names the flag the user
-    # actually typed.
-    capability = current_platform.get_device_capability()
-    if capability is not None:
-        cc = (capability.major, capability.minor)
-        if cc not in FI_MOE_EP_SUPPORTED_CAPABILITIES:
-            supported = ", ".join(
-                f"{m}.{n}" for m, n in sorted(FI_MOE_EP_SUPPORTED_CAPABILITIES)
-            )
-            raise ValueError(
-                f"moe_backend={moe_backend!r} is only supported on compute "
-                f"capability {supported} (SM100/SM103), but this device is "
-                f"{cc[0]}.{cc[1]}."
-            )
-
-    if vllm_config.parallel_config.enable_eplb:
-        raise NotImplementedError(
-            f"EPLB is not supported with moe_backend={moe_backend!r}: the "
-            "flashinfer moe_ep experts neither apply the logical-to-physical "
-            "expert map nor report per-expert load, so rebalancing would move "
-            "weights without moving routing. Use "
-            "moe_backend=deep_gemm_mega_moe to run the mega path with EPLB."
-        )
-
-
-def make_fi_moe_ep_bootstrap() -> BootstrapConfig:
-    from flashinfer.moe_ep import BootstrapConfig
-
-    ep = get_ep_group()
-    return BootstrapConfig(
-        world_size=ep.world_size,
-        rank=ep.rank_in_group,
-        process_group=ep.device_group,
-        auto_bootstrap=False,
-        # vLLM already bound this worker's (possibly remapped) device;
-        # without this the runtime would rebind to cuda:LOCAL_RANK|rank and
-        # launch weight transforms against another device's pointers.
-        device=torch.accelerator.current_device_index(),
-    )
-
-
-def megakernel_runtime_requirements(spec: FiMoeEpBackendSpec) -> frozenset[str]:
-    from flashinfer.moe_ep.core.runtime import NVSHMEM, TORCH_DIST
-
-    if spec.needs_nvshmem:
-        return frozenset({TORCH_DIST, NVSHMEM})
-    return frozenset({TORCH_DIST})
-
-
-def ensure_fi_moe_ep_runtime(vllm_config: VllmConfig) -> None:
-    """Acquire the process-wide flashinfer moe_ep runtime once per worker."""
-    global _FI_RUNTIME_HANDLE
-    if _FI_RUNTIME_HANDLE is not None:
-        return
-
-    from flashinfer.moe_ep import bootstrap_moe_ep_runtime
-
-    bootstrap = make_fi_moe_ep_bootstrap()
-    spec = fi_moe_ep_backend_spec(vllm_config.kernel_config.moe_backend)
-    _FI_RUNTIME_HANDLE = bootstrap_moe_ep_runtime(
-        bootstrap,
-        megakernel_runtime_requirements(spec),
-    )
-
-
-def finalize_fi_moe_ep_runtime() -> None:
-    """Release the process-wide flashinfer moe_ep runtime."""
-    global _FI_RUNTIME_HANDLE
-    if _FI_RUNTIME_HANDLE is None:
-        return
-
-    from flashinfer.moe_ep import finalize_moe_ep_runtime
-
-    finalize_moe_ep_runtime(_FI_RUNTIME_HANDLE)
-    _FI_RUNTIME_HANDLE = None
-
-
-_E2M1_LUT = (
-    0.0,
-    0.5,
-    1.0,
-    1.5,
-    2.0,
-    3.0,
-    4.0,
-    6.0,
-    -0.0,
-    -0.5,
-    -1.0,
-    -1.5,
-    -2.0,
-    -3.0,
-    -4.0,
-    -6.0,
-)
-
-
-def _dequant_fp4_ue8m0_gran32(
-    packed: torch.Tensor, sf_ue8m0: torch.Tensor
-) -> torch.Tensor:
-    """[rows, K//2] packed e2m1 + [rows, K//32] ue8m0-uint8 scales -> bf16 [rows, K]."""
-    raw = packed.view(torch.uint8)
-    lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device)
-    vals = torch.empty(
-        raw.shape[0], raw.shape[1] * 2, dtype=torch.float32, device=raw.device
-    )
-    vals[:, ::2] = lut[(raw & 0x0F).to(torch.int64)]
-    vals[:, 1::2] = lut[(raw >> 4).to(torch.int64)]
-    sf = (sf_ue8m0.to(torch.int32) << 23).view(torch.float32)
-    return (vals * sf.repeat_interleave(32, dim=-1)).to(torch.bfloat16)
-
-
-def _dequant_expert_weights_to_bf16(
-    weight: torch.Tensor, scale: torch.Tensor
-) -> torch.Tensor:
-    """[E, N, K//2] fp4 + [E, N, K//32] ue8m0 -> [E, N, K] bf16 (expert loop)."""
-    num_experts, n, k_half = weight.shape
-    out = torch.empty(
-        num_experts, n, k_half * 2, dtype=torch.bfloat16, device=weight.device
-    )
-    for e in range(num_experts):
-        out[e] = _dequant_fp4_ue8m0_gran32(weight[e], scale[e])
-    return out
-
-
-def mega_moe_weight_pack_from_params(
-    w13_weight: nn.Parameter,
-    w13_weight_scale: nn.Parameter,
-    w2_weight: nn.Parameter,
-    w2_weight_scale: nn.Parameter,
-    *,
-    megakernel: str = "deep_gemm_mega",
-):
-    from flashinfer.moe_ep import MoEWeightPack
-
-    if megakernel == "deep_gemm_mega":
-        # Same fp4-e2m1 + ue8m0-per-32 recipe as the native path: pass verbatim,
-        # flashinfer runs the identical deep_gemm transform.
-        return MoEWeightPack(
-            w13=w13_weight.data,
-            w2=w2_weight.data,
-            w13_scale=w13_weight_scale.data,
-            w2_scale=w2_weight_scale.data,
-        )
-    # The cutedsl kernel quantizes with its own recipe (nvfp4
-    # e2m1+e4m3-per-16): dequantize the checkpoint fp4 to bf16 and let the
-    # backend preprocess requantize. Double quantization: outputs are close to
-    # but not bit-identical with the native path.
-    return MoEWeightPack(
-        w13=_dequant_expert_weights_to_bf16(w13_weight.data, w13_weight_scale.data),
-        w2=_dequant_expert_weights_to_bf16(w2_weight.data, w2_weight_scale.data),
-    )
-
-
-def build_fi_mega_config(
-    *,
-    intermediate_size: int,
-    top_k: int,
-    activation_clamp: float | None,
-    megakernel: str,
-):
-    from flashinfer.moe_ep import (
-        DeepGemmMegaMoeConfig,
-        MegaConfig,
-        Nvfp4CutedslMegaMoeConfig,
-    )
-
-    # fast_math selects approximate exp/rcp in DeepGEMM's fused SwiGLU
-    # epilogue; the cutedsl kernels accept it for API parity only.
-    if megakernel == "deep_gemm_mega":
-        mk = DeepGemmMegaMoeConfig(
-            intermediate_size=intermediate_size,
-            top_k=top_k,
-            activation_clamp=activation_clamp,
-            fast_math=True,
-        )
-    elif megakernel == "nvfp4_cutedsl":
-        mk = Nvfp4CutedslMegaMoeConfig(
-            intermediate_size=intermediate_size,
-            top_k=top_k,
-            activation_clamp=activation_clamp,
-            fast_math=True,
-        )
-    else:
-        raise ValueError(f"Unsupported fi_moe_ep megakernel {megakernel!r}")
-
-    return MegaConfig(
-        megakernel=mk,
-        preprocess_weights=True,
-        quantize_input=True,
-    )
-
-
-# All MoE layers share one symmetric workspace, like the native path's
-# class-level DeepseekV4MegaMoEExperts._symm_buffer_cache. Without this the
-# fi path allocates one symm buffer PER LAYER (43x memory + cold working
-# sets); the workspace is stateless across forwards (kernel tail-cleans) and
-# layers execute sequentially on one stream, so sharing is safe.
-def build_fi_mega_layer(
-    bootstrap: BootstrapConfig,
-    *,
-    vllm_config: VllmConfig,
-    num_experts: int,
-    max_tokens_per_rank: int,
-    hidden_size: int,
-    intermediate_size: int,
-    top_k: int,
-    activation_clamp: float | None,
-    weights,
-) -> MoEEpMegaLayer:
-    from flashinfer.moe_ep import FleetParams, MoEEpLayer
-
-    megakernel = fi_moe_ep_backend_spec(
-        vllm_config.kernel_config.moe_backend
-    ).megakernel
-    mega_config = build_fi_mega_config(
-        intermediate_size=intermediate_size,
-        top_k=top_k,
-        activation_clamp=activation_clamp,
-        megakernel=megakernel,
-    )
-    layer = MoEEpLayer(
-        bootstrap=bootstrap,
-        fleet_params=FleetParams(
-            num_experts=num_experts,
-            max_tokens_per_rank=max_tokens_per_rank,
-            token_hidden_size=hidden_size,
-        ),
-        weights=weights,
-        backend=mega_config,
-    )
-    from flashinfer.moe_ep import MoEEpMegaLayer
-
-    if not isinstance(layer, MoEEpMegaLayer):
-        raise TypeError(
-            f"fi_moe_ep expected MoEEpMegaLayer, got {type(layer).__name__}"
-        )
-    return layer
-
-
-__all__ = [
-    "FI_MOE_EP_BACKEND_SPECS",
-    "FiMoeEpBackendSpec",
-    "build_fi_mega_config",
-    "build_fi_mega_layer",
-    "ensure_fi_moe_ep_runtime",
-    "fi_moe_ep_backend_spec",
-    "finalize_fi_moe_ep_runtime",
-    "is_fi_moe_ep_backend",
-    "make_fi_moe_ep_bootstrap",
-    "mega_moe_weight_pack_from_params",
-    "megakernel_runtime_requirements",
-    "validate_fi_moe_ep_config",
-]