feat: extended EPLB support for Mistral Large 3 and additional MoE backends (#48355)

Signed-off-by: jdebache <[email protected]>

tests/distributed/test_eplb_quant_scale_consistency.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""EPLB rearrangement consistency tests for quant-method derived state.
+
+EPLB rearranges every registered Parameter of a RoutedExperts in place,
+sliced along dim 0 (see RoutedExperts.get_expert_weights). Quant methods
+must therefore register all derived per-expert tensors as Parameters and
+alias the same storage in their FusedMoEQuantConfig, so the kernels observe
+rearranged values with no extra bookkeeping. These tests verify that
+contract: simulating a rearrangement on the registered Parameters must be
+indistinguishable from loading a checkpoint with permuted experts.
+"""
+
+import pytest
+import torch
+
+from vllm.config import VllmConfig, set_current_vllm_config
+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.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import (  # noqa: E501
+    CompressedTensorsW4A4Nvfp4MoEMethod,
+)
+from vllm.platforms import current_platform
+
+EPLB_NVFP4_BACKENDS = ["flashinfer_cutedsl", "flashinfer_trtllm"]
+
+NUM_EXPERTS = 8
+HIDDEN_SIZE = 128
+INTERMEDIATE_SIZE = 256
+EXPERT_PERMUTATION = [3, 0, 5, 1, 7, 2, 6, 4]
+
+QUANT_CONFIG_TENSORS = (
+    "w1_scale",
+    "w2_scale",
+    "g1_alphas",
+    "g2_alphas",
+    "a1_gscale",
+    "a2_gscale",
+)
+
+
+class _RoutedExpertsStub(torch.nn.Module):
+    """Minimal ``RoutedExperts`` stand-in for NVFP4 weight processing."""
+
+    def __init__(self, moe_config: FusedMoEConfig) -> None:
+        super().__init__()
+        self.moe_config = moe_config
+        self.activation = moe_config.activation
+
+    def _expert_routing_tables(
+        self,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+        return None
+
+
+def _make_moe_config(backend: str) -> FusedMoEConfig:
+    parallel_config = FusedMoEParallelConfig(
+        tp_size=1,
+        pcp_size=1,
+        dp_size=1,
+        ep_size=1,
+        tp_rank=0,
+        pcp_rank=0,
+        dp_rank=0,
+        ep_rank=0,
+        sp_size=1,
+        use_ep=True,
+        all2all_backend="allgather_reducescatter",
+        enable_eplb=True,
+    )
+    return FusedMoEConfig(
+        num_experts=NUM_EXPERTS,
+        experts_per_token=2,
+        hidden_dim=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+        num_local_experts=NUM_EXPERTS,
+        num_logical_experts=NUM_EXPERTS,
+        activation=MoEActivation.SILU,
+        device="cuda",
+        routing_method=RoutingMethodType.TopK,
+        moe_parallel_config=parallel_config,
+        in_dtype=torch.bfloat16,
+        moe_backend=backend,
+    )
+
+
+def _make_raw_weights(
+    device: torch.device, generator: torch.Generator
+) -> dict[str, torch.Tensor]:
+    """Random raw NVFP4 checkpoint tensors, indexable per-expert on dim 0."""
+    e, h, i = NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE
+
+    def packed(*shape: int) -> torch.Tensor:
+        return torch.randint(
+            0, 256, shape, dtype=torch.uint8, device=device, generator=generator
+        )
+
+    def block_scale(*shape: int) -> torch.Tensor:
+        return torch.randint(
+            0, 128, shape, dtype=torch.uint8, device=device, generator=generator
+        ).view(torch.float8_e4m3fn)
+
+    def global_scale(*shape: int) -> torch.Tensor:
+        return torch.rand(shape, device=device, generator=generator) + 0.5
+
+    return {
+        "w13_weight_packed": packed(e, 2 * i, h // 2),
+        "w2_weight_packed": packed(e, h, i // 2),
+        "w13_weight_scale": block_scale(e, 2 * i, h // 16),
+        "w2_weight_scale": block_scale(e, h, i // 16),
+        "w13_weight_global_scale": global_scale(e, 2),
+        "w2_weight_global_scale": global_scale(e),
+        "w13_input_global_scale": global_scale(e, 2),
+        "w2_input_global_scale": global_scale(e),
+    }
+
+
+def _build_processed_layer(
+    backend: str, raw: dict[str, torch.Tensor], device: torch.device
+) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]:
+    """Create the method + layer, load ``raw`` and run the real
+    process_weights_after_loading (which builds the flashinfer kernel)."""
+    moe_config = _make_moe_config(backend)
+    method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False)
+    layer = _RoutedExpertsStub(moe_config).to(device)
+    method.create_weights(
+        layer,
+        num_experts=NUM_EXPERTS,
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size_per_partition=INTERMEDIATE_SIZE,
+        params_dtype=torch.bfloat16,
+        weight_loader=lambda *args, **kwargs: None,
+    )
+    layer = layer.to(device)
+    with torch.no_grad():
+        for name, value in raw.items():
+            getattr(layer, name).copy_(value)
+    method.process_weights_after_loading(layer)
+    return method, layer
+
+
+def _assert_tensors_equal(
+    actual: torch.Tensor, expected: torch.Tensor, name: str
+) -> None:
+    if actual.dtype == torch.float8_e4m3fn:
+        assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), name
+    elif actual.dtype == torch.uint8:
+        assert torch.equal(actual, expected), name
+    else:
+        torch.testing.assert_close(actual, expected, msg=lambda m: f"{name}: {m}")
+
+
+def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None:
+    """Permute experts of every registered Parameter in place, the way
+    EPLB's in-place rearrangement moves expert slices along dim 0."""
+    with torch.no_grad():
+        for _, param in layer.named_parameters():
+            param.copy_(param[perm])
+
+
+def _ensure_world1_distributed() -> None:
+    """The enable_eplb config makes weight processing all-reduce activation
+    scale amaxes over the EP group, so a (single-rank) EP group must exist."""
+    from vllm.distributed.parallel_state import (
+        ensure_model_parallel_initialized,
+        init_distributed_environment,
+    )
+
+    if not torch.distributed.is_initialized():
+        init_distributed_environment(
+            world_size=1,
+            rank=0,
+            distributed_init_method="tcp://127.0.0.1:0",
+            local_rank=0,
+        )
+    ensure_model_parallel_initialized(1, 1)
+
+
+@pytest.mark.parametrize("backend", EPLB_NVFP4_BACKENDS)
+def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None:
+    if not (
+        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+    ):
+        pytest.skip("NVFP4 CuteDSL/TRTLLM MoE backends require Blackwell (SM100).")
+
+    device = torch.device("cuda:0")
+    torch.accelerator.set_device_index(device)
+    perm = torch.tensor(EXPERT_PERMUTATION, device=device)
+
+    with set_current_vllm_config(VllmConfig()):
+        _ensure_world1_distributed()
+        generator = torch.Generator(device=device).manual_seed(1234)
+        raw = _make_raw_weights(device, generator)
+        raw_permuted = {name: value[perm].contiguous() for name, value in raw.items()}
+
+        try:
+            method, layer = _build_processed_layer(backend, raw, device)
+        except ValueError as exc:
+            pytest.skip(f"{backend} NVFP4 MoE backend unavailable: {exc}")
+        ref_method, ref_layer = _build_processed_layer(backend, raw_permuted, device)
+
+        # The EPLB contract: every registered Parameter must be an
+        # expert-major contiguous tensor so get_expert_weights can view it
+        # as (E, -1) and rearrange expert slices in place.
+        for name, param in layer.named_parameters():
+            assert param.is_contiguous(), f"{name} is not contiguous"
+            assert param.shape[0] == NUM_EXPERTS, (
+                f"{name} is not expert-major: {tuple(param.shape)}"
+            )
+
+        # Derived per-expert scales must live in registered Parameters that
+        # the quant config aliases; anything else goes stale on rearrangement.
+        quant_config = method.moe_quant_config
+        assert quant_config is not None
+        params = dict(layer.named_parameters())
+        assert quant_config.g1_alphas.data_ptr() == (
+            params["w13_weight_scale_2"].data_ptr()
+        )
+        assert quant_config.g2_alphas.data_ptr() == (
+            params["w2_weight_scale_2"].data_ptr()
+        )
+        assert quant_config.w1_scale.data_ptr() == (
+            params["w13_weight_scale"].data_ptr()
+        )
+        assert quant_config.w2_scale.data_ptr() == params["w2_weight_scale"].data_ptr()
+
+        _simulate_eplb_rearrangement(layer, perm)
+
+        # After rearrangement, all registered state must equal what a fresh
+        # load of the permuted experts produces.
+        ref_params = dict(ref_layer.named_parameters())
+        assert params.keys() == ref_params.keys()
+        for name in params:
+            _assert_tensors_equal(params[name], ref_params[name], name)
+
+        # And so must the kernel-visible quant config tensors (the CuteDSL
+        # MMA scale views alias the registered Parameters' storage).
+        ref_quant_config = ref_method.moe_quant_config
+        assert ref_quant_config is not None
+        for name in QUANT_CONFIG_TENSORS:
+            actual = getattr(quant_config, name)
+            expected = getattr(ref_quant_config, name)
+            assert (actual is None) == (expected is None), name
+            if actual is not None:
+                _assert_tensors_equal(actual, expected, f"quant_config.{name}")
+
+
+def test_fp8_per_tensor_alphas_registered_and_aliased() -> None:
+    """CPU-only: the fp8 per-tensor oracle must register the fused
+    (w_scale * a_scale) products as layer Parameters aliased by the quant
+    config, so EPLB rearrangement keeps kernels consistent."""
+    from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+        Fp8MoeBackend,
+        make_fp8_moe_quant_config,
+    )
+
+    e = NUM_EXPERTS
+    generator = torch.Generator().manual_seed(1234)
+
+    def scales() -> tuple[torch.Tensor, torch.Tensor]:
+        w_scale = torch.rand((e,), generator=generator) + 0.5
+        a_scale = torch.rand((), generator=generator) + 0.5
+        return w_scale, a_scale.expand(e).contiguous()
+
+    w1_scale, a1_scale = scales()
+    w2_scale, a2_scale = scales()
+
+    layer = torch.nn.Module()
+    quant_config = make_fp8_moe_quant_config(
+        fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        layer=layer,
+    )
+
+    params = dict(layer.named_parameters())
+    assert quant_config.g1_alphas.data_ptr() == params["g1_alphas"].data_ptr()
+    assert quant_config.g2_alphas.data_ptr() == params["g2_alphas"].data_ptr()
+
+    perm = torch.tensor(EXPERT_PERMUTATION)
+    ref_config = make_fp8_moe_quant_config(
+        fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS,
+        w1_scale=w1_scale[perm].contiguous(),
+        w2_scale=w2_scale[perm].contiguous(),
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        layer=None,
+    )
+    _simulate_eplb_rearrangement(layer, perm)
+
+    torch.testing.assert_close(quant_config.g1_alphas, ref_config.g1_alphas)
+    torch.testing.assert_close(quant_config.g2_alphas, ref_config.g2_alphas)

tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py

@@ -193,7 +193,22 @@ def test_flashinfer_cutedsl_fp4_moe(
             hidden_states, score, topk, renormalize=False
         )
 
-        fake_layer = SimpleNamespace(activation=activation)
+        activation = MoEActivation.RELU2_NO_MUL
+        moe_config = FusedMoEConfig(
+            num_experts=e,
+            experts_per_token=topk,
+            hidden_dim=k,
+            intermediate_size=n,
+            num_local_experts=e,
+            num_logical_experts=e,
+            activation=activation,
+            device="cuda",
+            moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+            in_dtype=dtype,
+            routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
+        )
+        fake_layer = SimpleNamespace(activation=activation, moe_config=moe_config)
         a1_scale = torch.ones(1, device="cuda", dtype=torch.float32)
         a2_scale = torch.ones(1, device="cuda", dtype=torch.float32)
         (
@@ -230,20 +245,6 @@ def test_flashinfer_cutedsl_fp4_moe(
             gemm1_beta=beta,
             gemm1_clamp_limit=limit,
         )
-        moe_config = FusedMoEConfig(
-            num_experts=e,
-            experts_per_token=topk,
-            hidden_dim=k,
-            intermediate_size=n,
-            num_local_experts=e,
-            num_logical_experts=e,
-            activation=activation,
-            device="cuda",
-            moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
-            in_dtype=dtype,
-            routing_method=RoutingMethodType.TopK,
-            max_num_tokens=next_power_of_2(m),
-        )
 
         cutedsl_experts = mk.FusedMoEKernel(
             maybe_make_prepare_finalize(

tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py

@@ -19,7 +19,12 @@ def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
     monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
 
     num_experts = 3
-    layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False))
+    layer = SimpleNamespace(
+        activation=SimpleNamespace(is_gated=False),
+        moe_config=SimpleNamespace(
+            moe_parallel_config=SimpleNamespace(enable_eplb=False)
+        ),
+    )
     w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
     w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
     w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -108,8 +108,10 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
     )
     monkeypatch.setattr(
         eplb,
-        "is_mixture_of_experts",
-        lambda loaded_model: getattr(loaded_model, "is_moe", False),
+        "get_mixture_of_experts_model",
+        lambda loaded_model: (
+            loaded_model if getattr(loaded_model, "is_moe", False) else None
+        ),
     )
 
     runner = _make_runner(is_last_pp_rank=False)
@@ -138,7 +140,7 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
         "init_model_state",
         lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1),
     )
-    monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True)
+    monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model)
 
     runner = _make_runner(is_last_pp_rank=False)
     mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True)
@@ -153,7 +155,7 @@ def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch):
     FakeEplbState.instances.clear()
     FakeEplbState.from_mapping_kwargs = None
     monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
-    monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True)
+    monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model)
 
     runner = _make_runner(model=SimpleNamespace(is_moe=True))
     mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)

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

@@ -628,6 +628,17 @@ def make_fp8_moe_quant_config(
         and block_shape is None
     ):
         assert a1_scale is not None and a2_scale is not None
+        g1_alphas = w1_scale * a1_scale
+        g2_alphas = w2_scale * a2_scale
+        if layer is not None:
+            layer.register_parameter(
+                "g1_alphas", torch.nn.Parameter(g1_alphas, requires_grad=False)
+            )
+            layer.register_parameter(
+                "g2_alphas", torch.nn.Parameter(g2_alphas, requires_grad=False)
+            )
+            g1_alphas = layer.g1_alphas
+            g2_alphas = layer.g2_alphas
         return fp8_w8a8_moe_quant_config(
             w1_scale=w1_scale,
             w2_scale=w2_scale,
@@ -637,8 +648,8 @@ def make_fp8_moe_quant_config(
             a2_scale=a2_scale,
             a1_gscale=(1.0 / a1_scale),
             a2_gscale=(1.0 / a2_scale),
-            g1_alphas=(w1_scale * a1_scale).squeeze(),
-            g2_alphas=(w2_scale * a2_scale).squeeze(),
+            g1_alphas=g1_alphas,
+            g2_alphas=g2_alphas,
             gemm1_clamp_limit=swiglu_limit,
         )
     # MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are

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

@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import (
 )
 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+    nvfp4_swizzled_scale_to_cutedsl_mma_view,
     prepare_nvfp4_moe_layer_for_fi_or_cutlass,
     prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
 )
@@ -506,6 +507,10 @@ def make_nvfp4_moe_quant_config(
             gemm1_clamp_limit=swiglu_limit,
         )
 
+    if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
+        w13_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w13_scale)
+        w2_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w2_scale)
+
     # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas.
     # The expert's process_weights_after_loading will fuse activation
     # scales in-place. Since the quant config references the same tensor

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

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+    NvFp4MoeBackend,
     convert_to_nvfp4_moe_kernel_format,
     is_global_sf_supported_for_nvfp4_backend,
     make_nvfp4_moe_kernel,
@@ -30,6 +31,15 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
 
+# NVFP4 backends that have been verified to support EPLB for this quantization recipe.
+_EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset(
+    {
+        NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+        NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
+        NvFp4MoeBackend.FLASHINFER_TRTLLM,
+    }
+)
+
 logger = init_logger(__name__)
 
 
@@ -54,6 +64,10 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
             self.nvfp4_backend
         )
 
+    @property
+    def supports_eplb(self) -> bool:
+        return self.nvfp4_backend in _EPLB_SUPPORTED_NVFP4_BACKENDS
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -223,8 +237,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
         replace_parameter(layer, "w13_weight_scale", w13_scale)
         replace_parameter(layer, "w2_weight", w2)
         replace_parameter(layer, "w2_weight_scale", w2_scale)
-        layer.w13_weight_scale_2 = w13_scale_2
-        layer.w2_weight_scale_2 = w2_scale_2
+        replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
+        replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
         layer.w13_input_scale = a13_scale
         layer.w2_input_scale = a2_scale
 

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

@@ -290,7 +290,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
             assert self.input_quant.strategy == QuantizationStrategy.TENSOR
             assert w13_input_scale is not None and w2_input_scale is not None
             w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
-                w13_input_scale, w2_input_scale
+                w13_input_scale,
+                w2_input_scale,
+                layer.moe_config.moe_parallel_config.enable_eplb,
             )
             replace_parameter(layer, "w13_input_scale", w13_input_scale)
             replace_parameter(layer, "w2_input_scale", w2_input_scale)

vllm/model_executor/layers/quantization/fp8.py

@@ -709,7 +709,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             assert not self.block_quant
             assert w13_input_scale is not None and w2_input_scale is not None
             w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
-                w13_input_scale, w2_input_scale
+                w13_input_scale,
+                w2_input_scale,
+                layer.moe_config.moe_parallel_config.enable_eplb,
             )
             replace_parameter(layer, "w13_input_scale", w13_input_scale)
             replace_parameter(layer, "w2_input_scale", w2_input_scale)

vllm/model_executor/layers/quantization/modelopt.py

@@ -899,7 +899,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
 
         # Per tensor kernels require single activation scale. Use the max.
         w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
-            w13_input_scale, w2_input_scale
+            w13_input_scale,
+            w2_input_scale,
+            layer.moe_config.moe_parallel_config.enable_eplb,
         )
         replace_parameter(layer, "w13_input_scale", w13_input_scale)
         replace_parameter(layer, "w2_input_scale", w2_input_scale)

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

@@ -15,6 +15,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
     swizzle_blockscale,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    amax_for_moe_activation_quant,
+)
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.fused_moe import RoutedExperts
@@ -129,12 +132,13 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
     and interleaves w13 gate/linear rows for gated activations. Non-gated
     activations use a single w13 projection and keep its row order unchanged.
     """
-    from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
-
     # Global scaling factors (same as other FlashInfer backends).
     num_experts = w13.shape[0]
-    a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
-    a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
+    enable_eplb = layer.moe_config.moe_parallel_config.enable_eplb
+    a13_scale = amax_for_moe_activation_quant(a13_scale, enable_eplb).repeat(
+        num_experts
+    )
+    a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(num_experts)
 
     if layer.activation.is_gated:
         w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(
@@ -145,29 +149,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
         w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
         w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)
 
-    # Convert w13 scale factors: linear → swizzled → MMA layout.
     w13_scale = swizzle_blockscale(w13_scale)
-    E, M_padded, K_sf_padded = w13_scale.shape
-    w13_scale_flat = w13_scale.reshape(E * M_padded, K_sf_padded)
-    w13_scale = convert_sf_to_mma_layout(
-        w13_scale_flat,
-        m=M_padded,
-        k=K_sf_padded * 16,
-        num_groups=E,
-        sf_vec_size=16,
-    )
-
-    # Convert w2 scale factors: linear → swizzled → MMA layout.
     w2_scale = swizzle_blockscale(w2_scale)
-    E, M_padded, K_sf_padded = w2_scale.shape
-    w2_scale_flat = w2_scale.reshape(E * M_padded, K_sf_padded)
-    w2_scale = convert_sf_to_mma_layout(
-        w2_scale_flat,
-        m=M_padded,
-        k=K_sf_padded * 16,
-        num_groups=E,
-        sf_vec_size=16,
-    )
 
     return (
         w13,
@@ -181,6 +164,31 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
     )
 
 
+def nvfp4_swizzled_scale_to_cutedsl_mma_view(scale: torch.Tensor) -> torch.Tensor:
+    """View a swizzled (E, M_padded, K_sf_padded) block-scale tensor in the
+    MMA layout expected by the CuteDSL MoE kernel.
+
+    The returned tensor aliases `scale`'s storage, so in-place updates of the
+    registered Parameter (weight reloads, EPLB rearrangement) are visible to
+    the kernel with no extra bookkeeping.
+    """
+    from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
+
+    num_experts, m_padded, k_sf_padded = scale.shape
+    mma_view = convert_sf_to_mma_layout(
+        scale.reshape(num_experts * m_padded, k_sf_padded),
+        m=m_padded,
+        k=k_sf_padded * 16,
+        num_groups=num_experts,
+        sf_vec_size=16,
+    )
+    assert mma_view.data_ptr() == scale.data_ptr(), (
+        "convert_sf_to_mma_layout no longer returns a view of its input; "
+        "the quant config would go stale after weight updates."
+    )
+    return mma_view
+
+
 def prepare_static_weights_for_trtllm_fp4_moe(
     # args_dequant,
     # args,
@@ -363,8 +371,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
     # For some FI kernels, the input scales are shared by all experts.
     if is_global_sf_supported_for_nvfp4_backend(backend):
         num_experts = w13.shape[0]
-        a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
-        a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
+        enable_eplb = layer.moe_config.moe_parallel_config.enable_eplb
+        a13_scale = amax_for_moe_activation_quant(a13_scale, enable_eplb).repeat(
+            num_experts
+        )
+        a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(
+            num_experts
+        )
     else:
         a13_scale = a13_scale.max(dim=1).values.to(torch.float32)
 

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

@@ -14,6 +14,7 @@ import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    amax_for_moe_activation_quant,
     get_fp8_min_max,
 )
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -1437,6 +1438,7 @@ def process_fp8_weight_tensor_strategy_moe(
 def process_fp8_input_tensor_strategy_moe(
     w13_input_scale: torch.Tensor,
     w2_input_scale: torch.Tensor,
+    enable_eplb: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Process moe input scales for tensor-wise quantization strategy."""
 
@@ -1447,4 +1449,7 @@ def process_fp8_input_tensor_strategy_moe(
             "for each layer."
         )
 
-    return w13_input_scale.max(), w2_input_scale.max()
+    return (
+        amax_for_moe_activation_quant(w13_input_scale, enable_eplb),
+        amax_for_moe_activation_quant(w2_input_scale, enable_eplb),
+    )

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

@@ -62,6 +62,28 @@ def amax_for_moe_weight_quant(amax: torch.Tensor, moe_tp_size: int) -> torch.Ten
     return amax
 
 
+def amax_for_moe_activation_quant(
+    a_scale: torch.Tensor, enable_eplb: bool
+) -> torch.Tensor:
+    """Reduce a per-expert activation scale to one value shared by all experts.
+
+    Note: when EPLB is enabled and since this quantization scales get
+    folded into the per-expert dequantization alphas, we can only
+    ensure that the quant/dequant scales match by having a single
+    quantization scale shared across all ranks.
+    """
+    a_max = a_scale.max().to(torch.float32)
+    if enable_eplb:
+        from vllm.distributed.parallel_state import get_ep_group
+
+        torch.distributed.all_reduce(
+            a_max,
+            op=torch.distributed.ReduceOp.MAX,
+            group=get_ep_group().device_group,
+        )
+    return a_max
+
+
 def get_fp8_min_max() -> tuple[float, float]:
     """Get the min and max values for FP8 quantization."""
     # Using the default value (240.0) from pytorch will cause accuracy

vllm/model_executor/models/interfaces.py

@@ -5,7 +5,6 @@ import asyncio
 from collections.abc import (
     AsyncGenerator,
     Callable,
-    Iterable,
     Mapping,
     MutableSequence,
     Sequence,
@@ -949,7 +948,7 @@ class MixtureOfExperts(Protocol):
     num_redundant_experts: int
     """Number of redundant experts in this model."""
 
-    moe_layers: Iterable["MoERunner"]
+    moe_layers: Sequence["MoERunner"]
     """List of MoE layers in this model."""
 
     def set_eplb_state(
@@ -992,6 +991,37 @@ class MixtureOfExperts(Protocol):
     ) -> None: ...
 
 
+def get_mixture_of_experts_model(model: object) -> MixtureOfExperts | None:
+    """Return the MixtureOfExperts contained within an arbitrary model.
+
+    - If the model itself is a MixtureOfExperts, return the model directly.
+    - If the model is a multi-modal model, and its `language_model` is a
+      MixtureOfExperts, return the `language_model`.
+    - If neither, return None.
+
+    Args:
+        model: Model being served.
+
+    Returns:
+        The MixtureOfExperts instance contained within the model, or None.
+    """
+
+    if is_mixture_of_experts(model):
+        return model
+
+    if isinstance(model, SupportsMultiModal):
+        try:
+            mm_language_model = model.get_language_model()
+            return (
+                mm_language_model if is_mixture_of_experts(mm_language_model) else None
+            )
+        except NotImplementedError:
+            logger.info_once("Cannot fetch language_model from MultiModal model")
+            return None
+
+    return None
+
+
 def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
     return (
         isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0

vllm/v1/worker/gpu/eplb_utils.py

@@ -12,23 +12,12 @@ from vllm.config import ModelConfig
 from vllm.distributed.eplb.eplb_state import EplbState
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import (
-    SupportsMultiModal,
-    is_mixture_of_experts,
+    get_mixture_of_experts_model,
 )
 
 logger = init_logger(__name__)
 
 
-def _unwrap_moe(model: nn.Module) -> nn.Module:
-    # VLM wrappers (e.g. KimiK25ForConditionalGeneration) hold the MoE
-    # language model under `.language_model` but don't implement
-    # MixtureOfExperts themselves. Mirror the V1 path
-    # (see vllm/v1/worker/gpu_model_runner.py, PR #39805).
-    if not is_mixture_of_experts(model) and isinstance(model, SupportsMultiModal):
-        return model.get_language_model()
-    return model
-
-
 def step_eplb_after(*, is_dummy: bool = False) -> Callable:
     """Step EPLB after a model runner method completes successfully."""
 
@@ -78,7 +67,8 @@ class EPLBController:
             return False
 
         draft_model = speculator.model
-        if not is_mixture_of_experts(draft_model):
+        draft_moe_model = get_mixture_of_experts_model(draft_model)
+        if draft_moe_model is None:
             return False
 
         assert not self.parallel_config.enable_elastic_ep, (
@@ -88,7 +78,7 @@ class EPLBController:
         assert speculative_config.draft_model_config is not None
         assert self.state is not None
         self.state.add_model(
-            draft_model,
+            draft_moe_model,
             speculative_config.draft_model_config,
         )
         speculator.set_eplb_state(self.state)
@@ -104,13 +94,15 @@ class EPLBController:
         if not self.parallel_config.enable_eplb or load_dummy_weights:
             return False
 
-        model = _unwrap_moe(model)
-        if not is_mixture_of_experts(model):
+        moe_model = get_mixture_of_experts_model(model)
+        if moe_model is None:
             return False
 
-        logger.info_once("EPLB is enabled for model %s.", model_config.model)
+        logger.info_once(
+            "EPLB is enabled for MoE part of model %s.", model_config.model
+        )
         assert self.state is not None
-        self.state.add_model(model, model_config)
+        self.state.add_model(moe_model, model_config)
         self._has_registered_models = True
         return True
 
@@ -154,11 +146,11 @@ class EPLBController:
         expanded_physical_to_logical: torch.Tensor,
         old_num_physical_experts: int,
     ) -> None:
-        model = _unwrap_moe(model)
-        assert is_mixture_of_experts(model)
+        moe_model = get_mixture_of_experts_model(model)
+        assert moe_model is not None
 
         self.state = EplbState.from_mapping(
-            model=model,
+            model=moe_model,
             model_config=model_config,
             device=self.device,
             parallel_config=self.parallel_config,

vllm/v1/worker/gpu_model_runner.py

@@ -82,7 +82,7 @@ from vllm.model_executor.models.interfaces import (
     SupportsMRoPE,
     SupportsMultiModal,
     SupportsXDRoPE,
-    is_mixture_of_experts,
+    get_mixture_of_experts_model,
     supports_eagle3,
     supports_mrope,
     supports_multimodal_pruning,
@@ -5364,9 +5364,14 @@ class GPUModelRunner(
                     if hasattr(self.drafter, "load_model"):
                         self.drafter.load_model(self.model)
                     if (
-                        hasattr(self.drafter, "model")
-                        and is_mixture_of_experts(self.drafter.model)
-                        and self.parallel_config.enable_eplb
+                        self.parallel_config.enable_eplb
+                        and hasattr(self.drafter, "model")
+                        and (
+                            drafter_moe_model := get_mixture_of_experts_model(
+                                self.drafter.model
+                            )
+                        )
+                        is not None
                     ):
                         assert not self.parallel_config.enable_elastic_ep, (
                             "Elastic EP is not supported with drafter model."
@@ -5375,7 +5380,7 @@ class GPUModelRunner(
                         assert spec_config is not None
                         assert spec_config.draft_model_config is not None
                         logger.info_once(
-                            "EPLB is enabled for drafter model %s.",
+                            "EPLB is enabled for MoE part of drafter model %s.",
                             spec_config.draft_model_config.model,
                         )
                         if self.eplb_state is None:
@@ -5383,7 +5388,7 @@ class GPUModelRunner(
                                 self.parallel_config, self.device
                             )
                         self.eplb_state.add_model(
-                            self.drafter.model,
+                            drafter_moe_model,
                             spec_config.draft_model_config,
                         )
                         assert hasattr(self.drafter, "set_eplb_state")
@@ -5396,21 +5401,15 @@ class GPUModelRunner(
                 # VLM models (e.g. KimiK25ForConditionalGeneration) wrap the
                 # actual MoE language model but don't implement
                 # MixtureOfExperts themselves.
-                moe_candidate = self.model
-                if not is_mixture_of_experts(moe_candidate) and isinstance(
-                    moe_candidate, SupportsMultiModal
-                ):
-                    moe_candidate = moe_candidate.get_language_model()
-                if is_mixture_of_experts(moe_candidate):
-                    self._moe_model = moe_candidate
+                self._moe_model = get_mixture_of_experts_model(self.model)
 
                 if (
-                    self._moe_model is not None
-                    and self.parallel_config.enable_eplb
+                    self.parallel_config.enable_eplb
                     and not load_dummy_weights
+                    and self._moe_model is not None
                 ):
                     logger.info_once(
-                        "EPLB is enabled for model %s.",
+                        "EPLB is enabled for MoE part of model %s.",
                         self.model_config.model,
                     )
                     assert self.eplb_state is not None
@@ -5450,8 +5449,8 @@ class GPUModelRunner(
         )  # Temporary hack for dynamic res video w/o support for bs>1 yet
 
         if (
-            self._moe_model is not None
-            and self.parallel_config.enable_eplb
+            self.parallel_config.enable_eplb
+            and self._moe_model is not None
             and not load_dummy_weights
             and self.eplb_state is not None
             and self.eplb_state.is_async