diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 60fbf53d4a..182672597a 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index c1fbb11a81..93a4088fc4 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index c737109f47..3b80b5b1e3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 7dba2d271a..d2d84878e1 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 1eb96d5594..87fc3b1c36 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/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 diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index a7b4d28480..7ef64e97a9 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/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, )