[ROCm][MoE] Fix expert_map vs AITER expert_mask for non-AITER experts under EP (#49758)

Signed-off-by: Rohan Potdar <rohan.potdar@amd.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/kernels/test_minimax_m3_amd_ops.py

@@ -425,14 +425,14 @@ def test_mxfp8_linear_emulation_bf16_at_load(
 
 
 # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ────────────────
-# Regression for the EP + aiter-master-switch interaction: under expert
-# parallelism ``RoutedExperts.expert_map`` hands the experts either the 0/1
-# ``expert_mask`` (aiter master ON, ``rocm_aiter_fmoe_enabled``) or vLLM's -1
-# index map (master OFF). ``AiterMxfp8Experts.apply`` must forward the right 0/1
-# mask to aiter in BOTH cases. The old code always rebuilt the mask via
-# ``(expert_map >= 0)``; on the already-0/1 mask that collapses to all-ones (no
-# experts masked out) and EP output becomes garbage (no accuracy).
-def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_experts):
+# The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER
+# experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and
+# everyone else the canonical -1 index map, keyed on the resolved experts kernel
+# rather than the global VLLM_ROCM_USE_AITER switch. ``AiterMxfp8Experts.apply``
+# forwards that mask to aiter unchanged (it no longer rebuilds it via
+# ``(expert_map >= 0)``, which collapsed an already-0/1 mask to all-ones and
+# produced EP garbage).
+def _capture_expert_mask(expert_mask, *, global_num_experts):
     """Drive the real ``AiterMxfp8Experts.apply`` mask branch and capture the
     ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``."""
     from types import SimpleNamespace
@@ -444,9 +444,6 @@ def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_expe
     )
 
     experts = object.__new__(AiterMxfp8Experts)  # bypass heavy __init__
-    experts.moe_config = SimpleNamespace(
-        rocm_aiter_fmoe_enabled=rocm_aiter_fmoe_enabled
-    )
     experts.quant_config = SimpleNamespace(gemm1_clamp_limit=None)
     experts.w1_scale_val = None
     experts.w2_scale_val = None
@@ -474,7 +471,7 @@ def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_expe
             topk_ids=ti,
             activation=None,
             global_num_experts=global_num_experts,
-            expert_map=expert_map,
+            expert_map=expert_mask,
             a1q_scale=None,
             a2_scale=None,
             workspace13=None,
@@ -486,33 +483,45 @@ def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_expe
 
 
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
-def test_aiter_mxfp8_ep_expert_mask_both_master_modes():
-    """Both aiter-master forms must yield the SAME correct 0/1 aiter mask;
-    guards the EP+master regression (mask must not collapse to all-ones)."""
-    from vllm.model_executor.layers.fused_moe.expert_map_manager import (
-        determine_expert_map,
+def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged():
+    """AiterMxfp8Experts.apply forwards the precomputed 0/1 mask to aiter as-is
+    (no re-derivation that would collapse an already-0/1 mask to all-ones)."""
+    # 0/1 local-expert mask over global ids + trailing sentinel (rank owns 0..3).
+    ep_mask = torch.tensor(
+        [1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32, device=DEVICE
     )
+    got = _capture_expert_mask(ep_mask, global_num_experts=8)
+    assert torch.equal(got.cpu().to(torch.int32), ep_mask.cpu().to(torch.int32))
+    assert got.sum().item() == 4  # the 4 local experts, not all 9
 
-    E, ep_size, ep_rank = 8, 2, 0  # rank owns global experts 0..3
-    # master OFF: vLLM's -1 index map
-    _, idx_map, _ = determine_expert_map(ep_size, ep_rank, E, return_expert_mask=False)
-    # master ON: 0/1 mask (+ trailing sentinel) that RoutedExperts forwards
-    _, _, ep_mask = determine_expert_map(ep_size, ep_rank, E, return_expert_mask=True)
 
-    idx_map = idx_map.to(DEVICE)
-    ep_mask = ep_mask.to(DEVICE)
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
+def test_routed_experts_expert_map_delegates_to_kernel():
+    """RoutedExperts.expert_map returns the 0/1 mask only for kernels that set
+    ``consumes_expert_mask`` (AITER), and the canonical -1 map otherwise -- keyed
+    on the resolved kernel, not the global aiter switch."""
+    from types import SimpleNamespace
 
-    # Expected aiter expert_mask: 0/1 over global ids + trailing sentinel slot.
-    expected = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32)
+    from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 
-    got_off = _capture_expert_mask(
-        idx_map, rocm_aiter_fmoe_enabled=False, global_num_experts=E
-    )
-    got_on = _capture_expert_mask(
-        ep_mask, rocm_aiter_fmoe_enabled=True, global_num_experts=E
-    )
+    canonical = torch.tensor([0, 1, 2, 3, -1, -1, -1, -1], dtype=torch.int32)
+    mask = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32)
 
-    assert torch.equal(got_off.cpu().to(torch.int32), expected)
-    # master ON forwards the prebuilt mask unchanged (NOT collapsed to all-ones)
-    assert torch.equal(got_on.cpu().to(torch.int32), ep_mask.cpu().to(torch.int32))
-    assert got_on.sum().item() == 4  # exactly the 4 local experts, not all 9
+    def resolve(consumes_mask, *, has_moe_kernel=True):
+        moe_kernel = (
+            SimpleNamespace(
+                fused_experts=SimpleNamespace(consumes_expert_mask=consumes_mask)
+            )
+            if has_moe_kernel
+            else None
+        )
+        layer = SimpleNamespace(
+            _expert_map=canonical,
+            expert_mask=mask,
+            quant_method=SimpleNamespace(moe_kernel=moe_kernel),
+        )
+        return RoutedExperts.expert_map.fget(layer)
+
+    assert torch.equal(resolve(True), mask)  # AITER kernel -> 0/1 mask
+    assert torch.equal(resolve(False), canonical)  # non-AITER -> canonical map
+    assert torch.equal(resolve(False, has_moe_kernel=False), canonical)  # non-modular

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

@@ -62,6 +62,8 @@ def is_aiter_mxfp8_moe_available() -> bool:
 class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
     """MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950)."""
 
+    consumes_expert_mask = True
+
     @property
     def quant_dtype(self) -> torch.dtype | str | None:
         return self.quant_config.quant_dtype
@@ -148,21 +150,9 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
         limit = self.quant_config.gemm1_clamp_limit
         swiglu_limit = 0.0 if limit is None else float(limit)
 
-        # Under EP, aiter expects ``expert_mask``: a 0/1 *local-expert* mask over
-        # global ids with a trailing fake-expert sentinel slot (shape
-        # ``[global_num_experts + 1]``), from which it derives the global->local
-        # compaction. What ``RoutedExperts.expert_map`` hands us depends on the
-        # aiter master switch (``rocm_aiter_fmoe_enabled``).
-        # Branching on the (static) master flag — not the tensor contents —
-        # keeps this HIP-graph/torch.compile safe (no data-dependent sync).
-        # ``None`` under pure TP.
-        if expert_map is None:
-            expert_mask = None
-        elif self.moe_config.rocm_aiter_fmoe_enabled:
-            expert_mask = expert_map
-        else:
-            local_mask = (expert_map >= 0).to(torch.int32)
-            expert_mask = torch.cat([local_mask, local_mask.new_zeros(1)])
+        # RoutedExperts.expert_map hands AITER experts the precomputed 0/1
+        # expert_mask (with trailing sentinel) instead of the vLLM expert_map.
+        expert_mask = expert_map
 
         # Route through the graph-safe ``rocm_aiter_fused_moe`` custom op so the
         # call is captured under HIP graphs / torch.compile (a direct

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

@@ -240,7 +240,7 @@ def rocm_aiter_fused_experts(
     moe_config: FusedMoEConfig,
     activation: MoEActivation = MoEActivation.SILU,
     apply_router_weight_on_input: bool = False,
-    expert_map: torch.Tensor | None = None,
+    expert_mask: torch.Tensor | None = None,
     quant_config: FusedMoEQuantConfig | None = None,
     a1q_scale: torch.Tensor | None = None,
     num_local_tokens: torch.Tensor | None = None,
@@ -271,8 +271,6 @@ def rocm_aiter_fused_experts(
     topk_weights = topk_weights.to(torch.float32)
     topk_ids = topk_ids.to(torch.int32)
 
-    expert_mask = expert_map if expert_map is not None else None
-
     # w8a8 per-channel quantization
     if (
         quant_config.per_act_token_quant
@@ -420,6 +418,8 @@ def rocm_aiter_fused_experts(
 
 
 class AiterExperts(mk.FusedMoEExpertsModular):
+    consumes_expert_mask = True
+
     @property
     def expects_unquantized_inputs(self) -> bool:
         # When paired with MoRI, the prepare/finalize handles FP8
@@ -550,7 +550,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):
             topk_ids=topk_ids,
             activation=activation,
             apply_router_weight_on_input=apply_router_weight_on_input,
-            expert_map=expert_map,
+            expert_mask=expert_map,
             quant_config=self.quant_config,
             moe_config=self.moe_config,
             a1q_scale=a1q_scale,

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -470,6 +470,12 @@ class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):
 
 # TODO: add supported activations method (return string)
 class FusedMoEExperts(ABC):
+    # ROCm AITER kernels consume a 0/1 local-expert mask (with a trailing
+    # sentinel slot); every other backend consumes the canonical -1/local-slot
+    # expert_map. RoutedExperts.expert_map reads this flag to pick which to hand
+    # the active experts kernel.
+    consumes_expert_mask: bool = False
+
     def __init__(
         self,
         moe_config: FusedMoEConfig,

vllm/model_executor/layers/fused_moe/routed_experts.py

@@ -226,9 +226,15 @@ class RoutedExperts(PluggableLayer):
 
     @property
     def expert_map(self) -> torch.Tensor | None:
-        return (
-            self._expert_map if not self.rocm_aiter_fmoe_enabled else self.expert_mask
-        )
+        # AITER fused-MoE kernels consume the 0/1 expert_mask; every other
+        # backend consumes the canonical -1/local-slot map. Ask the active
+        # experts kernel which it wants (only AITER sets consumes_expert_mask)
+        # rather than keying on the global VLLM_ROCM_USE_AITER switch, which
+        # does not reflect the per-layer kernel actually selected.
+        moe_kernel = getattr(self.quant_method, "moe_kernel", None)
+        if moe_kernel is not None and moe_kernel.fused_experts.consumes_expert_mask:
+            return self.expert_mask
+        return self._expert_map
 
     def update_expert_map_info(self):
         # Update local attributes from ExpertMapManager

vllm/model_executor/layers/quantization/quark/quark_moe.py

@@ -1026,7 +1026,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
             apply_router_weight_on_input=layer.apply_router_weight_on_input,
             quant_config=self.moe_quant_config,
             moe_config=layer.moe_config,
-            expert_map=layer.expert_map,
+            expert_mask=layer.expert_mask,
         )