@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0f0cb918b7

[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>

Rohan Potdar committed Aug 11, 2026 at 15:33 UTC 0f0cb918b7c6d08aabb86ec4eff4bae05f6736d9

6 files changed +70 -59

tests/kernels/test_minimax_m3_amd_ops.py

+45 -36

@@ -425,14 +425,14 @@ def test_mxfp8_linear_emulation_bf16_at_load(

425

426

427 # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ────────────────

428 -# Regression for the EP + aiter-master-switch interaction: under expert

429 -# parallelism ``RoutedExperts.expert_map`` hands the experts either the 0/1

430 -# ``expert_mask`` (aiter master ON, ``rocm_aiter_fmoe_enabled``) or vLLM's -1

431 -# index map (master OFF). ``AiterMxfp8Experts.apply`` must forward the right 0/1

432 -# mask to aiter in BOTH cases. The old code always rebuilt the mask via

433 -# ``(expert_map >= 0)``; on the already-0/1 mask that collapses to all-ones (no

434 -# experts masked out) and EP output becomes garbage (no accuracy).

435 -def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_experts):

428 +# The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER

429 +# experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and

430 +# everyone else the canonical -1 index map, keyed on the resolved experts kernel

431 +# rather than the global VLLM_ROCM_USE_AITER switch. ``AiterMxfp8Experts.apply``

432 +# forwards that mask to aiter unchanged (it no longer rebuilds it via

433 +# ``(expert_map >= 0)``, which collapsed an already-0/1 mask to all-ones and

434 +# produced EP garbage).

435 +def _capture_expert_mask(expert_mask, *, global_num_experts):

436 """Drive the real ``AiterMxfp8Experts.apply`` mask branch and capture the

437 ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``."""

438 from types import SimpleNamespace

@@ -444,9 +444,6 @@ def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_expe

444 )

445

446 experts = object.__new__(AiterMxfp8Experts) # bypass heavy __init__

447 - experts.moe_config = SimpleNamespace(

448 - rocm_aiter_fmoe_enabled=rocm_aiter_fmoe_enabled

449 - )

447 experts.quant_config = SimpleNamespace(gemm1_clamp_limit=None)

448 experts.w1_scale_val = None

449 experts.w2_scale_val = None

@@ -474,7 +471,7 @@ def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_expe

471 topk_ids=ti,

472 activation=None,

473 global_num_experts=global_num_experts,

477 - expert_map=expert_map,

474 + expert_map=expert_mask,

475 a1q_scale=None,

476 a2_scale=None,

477 workspace13=None,

@@ -486,33 +483,45 @@ def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_expe

483

484

485 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")

489 -def test_aiter_mxfp8_ep_expert_mask_both_master_modes():

490 - """Both aiter-master forms must yield the SAME correct 0/1 aiter mask;

491 - guards the EP+master regression (mask must not collapse to all-ones)."""

492 - from vllm.model_executor.layers.fused_moe.expert_map_manager import (

493 - determine_expert_map,

486 +def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged():

487 + """AiterMxfp8Experts.apply forwards the precomputed 0/1 mask to aiter as-is

488 + (no re-derivation that would collapse an already-0/1 mask to all-ones)."""

489 + # 0/1 local-expert mask over global ids + trailing sentinel (rank owns 0..3).

490 + ep_mask = torch.tensor(

491 + [1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32, device=DEVICE

492 )

493 + got = _capture_expert_mask(ep_mask, global_num_experts=8)

494 + assert torch.equal(got.cpu().to(torch.int32), ep_mask.cpu().to(torch.int32))

495 + assert got.sum().item() == 4 # the 4 local experts, not all 9

496

496 - E, ep_size, ep_rank = 8, 2, 0 # rank owns global experts 0..3

497 - # master OFF: vLLM's -1 index map

498 - _, idx_map, _ = determine_expert_map(ep_size, ep_rank, E, return_expert_mask=False)

499 - # master ON: 0/1 mask (+ trailing sentinel) that RoutedExperts forwards

500 - _, _, ep_mask = determine_expert_map(ep_size, ep_rank, E, return_expert_mask=True)

497

502 - idx_map = idx_map.to(DEVICE)

503 - ep_mask = ep_mask.to(DEVICE)

498 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")

499 +def test_routed_experts_expert_map_delegates_to_kernel():

500 + """RoutedExperts.expert_map returns the 0/1 mask only for kernels that set

501 + ``consumes_expert_mask`` (AITER), and the canonical -1 map otherwise -- keyed

502 + on the resolved kernel, not the global aiter switch."""

503 + from types import SimpleNamespace

504

505 - # Expected aiter expert_mask: 0/1 over global ids + trailing sentinel slot.

506 - expected = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32)

505 + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts

506

508 - got_off = _capture_expert_mask(

509 - idx_map, rocm_aiter_fmoe_enabled=False, global_num_experts=E

510 - )

511 - got_on = _capture_expert_mask(

512 - ep_mask, rocm_aiter_fmoe_enabled=True, global_num_experts=E

513 - )

507 + canonical = torch.tensor([0, 1, 2, 3, -1, -1, -1, -1], dtype=torch.int32)

508 + mask = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32)

509 +

510 + def resolve(consumes_mask, *, has_moe_kernel=True):

511 + moe_kernel = (

512 + SimpleNamespace(

513 + fused_experts=SimpleNamespace(consumes_expert_mask=consumes_mask)

514 + )

515 + if has_moe_kernel

516 + else None

517 + )

518 + layer = SimpleNamespace(

519 + _expert_map=canonical,

520 + expert_mask=mask,

521 + quant_method=SimpleNamespace(moe_kernel=moe_kernel),

522 + )

523 + return RoutedExperts.expert_map.fget(layer)

524

515 - assert torch.equal(got_off.cpu().to(torch.int32), expected)

516 - # master ON forwards the prebuilt mask unchanged (NOT collapsed to all-ones)

517 - assert torch.equal(got_on.cpu().to(torch.int32), ep_mask.cpu().to(torch.int32))

518 - assert got_on.sum().item() == 4 # exactly the 4 local experts, not all 9

525 + assert torch.equal(resolve(True), mask) # AITER kernel -> 0/1 mask

526 + assert torch.equal(resolve(False), canonical) # non-AITER -> canonical map

527 + assert torch.equal(resolve(False, has_moe_kernel=False), canonical) # non-modular

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

+5 -15

@@ -62,6 +62,8 @@ def is_aiter_mxfp8_moe_available() -> bool:

62 class AiterMxfp8Experts(Mxfp8TritonExpertsBase):

63 """MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950)."""

65 + consumes_expert_mask = True

66 +

67 @property

68 def quant_dtype(self) -> torch.dtype | str | None:

69 return self.quant_config.quant_dtype

@@ -148,21 +150,9 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):

150 limit = self.quant_config.gemm1_clamp_limit

151 swiglu_limit = 0.0 if limit is None else float(limit)

152

151 - # Under EP, aiter expects ``expert_mask``: a 0/1 *local-expert* mask over

152 - # global ids with a trailing fake-expert sentinel slot (shape

153 - # ``[global_num_experts + 1]``), from which it derives the global->local

154 - # compaction. What ``RoutedExperts.expert_map`` hands us depends on the

155 - # aiter master switch (``rocm_aiter_fmoe_enabled``).

156 - # Branching on the (static) master flag — not the tensor contents —

157 - # keeps this HIP-graph/torch.compile safe (no data-dependent sync).

158 - # ``None`` under pure TP.

159 - if expert_map is None:

160 - expert_mask = None

161 - elif self.moe_config.rocm_aiter_fmoe_enabled:

162 - expert_mask = expert_map

163 - else:

164 - local_mask = (expert_map >= 0).to(torch.int32)

165 - expert_mask = torch.cat([local_mask, local_mask.new_zeros(1)])

153 + # RoutedExperts.expert_map hands AITER experts the precomputed 0/1

154 + # expert_mask (with trailing sentinel) instead of the vLLM expert_map.

155 + expert_mask = expert_map

156

157 # Route through the graph-safe ``rocm_aiter_fused_moe`` custom op so the

158 # call is captured under HIP graphs / torch.compile (a direct

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

+4 -4

@@ -240,7 +240,7 @@ def rocm_aiter_fused_experts(

240 moe_config: FusedMoEConfig,

241 activation: MoEActivation = MoEActivation.SILU,

242 apply_router_weight_on_input: bool = False,

243 - expert_map: torch.Tensor | None = None,

243 + expert_mask: torch.Tensor | None = None,

244 quant_config: FusedMoEQuantConfig | None = None,

245 a1q_scale: torch.Tensor | None = None,

246 num_local_tokens: torch.Tensor | None = None,

@@ -271,8 +271,6 @@ def rocm_aiter_fused_experts(

271 topk_weights = topk_weights.to(torch.float32)

272 topk_ids = topk_ids.to(torch.int32)

273

274 - expert_mask = expert_map if expert_map is not None else None

275 -

274 # w8a8 per-channel quantization

275 if (

276 quant_config.per_act_token_quant

@@ -420,6 +418,8 @@ def rocm_aiter_fused_experts(

418

419

420 class AiterExperts(mk.FusedMoEExpertsModular):

421 + consumes_expert_mask = True

422 +

423 @property

424 def expects_unquantized_inputs(self) -> bool:

425 # When paired with MoRI, the prepare/finalize handles FP8

@@ -550,7 +550,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):

550 topk_ids=topk_ids,

551 activation=activation,

552 apply_router_weight_on_input=apply_router_weight_on_input,

553 - expert_map=expert_map,

553 + expert_mask=expert_map,

554 quant_config=self.quant_config,

555 moe_config=self.moe_config,

556 a1q_scale=a1q_scale,

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -470,6 +470,12 @@ class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):

470

471 # TODO: add supported activations method (return string)

472 class FusedMoEExperts(ABC):

473 + # ROCm AITER kernels consume a 0/1 local-expert mask (with a trailing

474 + # sentinel slot); every other backend consumes the canonical -1/local-slot

475 + # expert_map. RoutedExperts.expert_map reads this flag to pick which to hand

476 + # the active experts kernel.

477 + consumes_expert_mask: bool = False

478 +

479 def __init__(

480 self,

481 moe_config: FusedMoEConfig,

vllm/model_executor/layers/fused_moe/routed_experts.py

+9 -3

@@ -226,9 +226,15 @@ class RoutedExperts(PluggableLayer):

226

227 @property

228 def expert_map(self) -> torch.Tensor | None:

229 - return (

230 - self._expert_map if not self.rocm_aiter_fmoe_enabled else self.expert_mask

231 - )

229 + # AITER fused-MoE kernels consume the 0/1 expert_mask; every other

230 + # backend consumes the canonical -1/local-slot map. Ask the active

231 + # experts kernel which it wants (only AITER sets consumes_expert_mask)

232 + # rather than keying on the global VLLM_ROCM_USE_AITER switch, which

233 + # does not reflect the per-layer kernel actually selected.

234 + moe_kernel = getattr(self.quant_method, "moe_kernel", None)

235 + if moe_kernel is not None and moe_kernel.fused_experts.consumes_expert_mask:

236 + return self.expert_mask

237 + return self._expert_map

238

239 def update_expert_map_info(self):

240 # Update local attributes from ExpertMapManager

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

+1 -1

@@ -1026,7 +1026,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

1026 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1027 quant_config=self.moe_quant_config,

1028 moe_config=layer.moe_config,

1029 - expert_map=layer.expert_map,

1029 + expert_mask=layer.expert_mask,

1030 )

1031

1032