[ROCm] fixed aiter master flag and expert parallelism compatibility on minimax-m3-mxfp8 (#47158)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
Hongxia Yang committed
Jul 8, 2026 at 03:26 UTC
2c64b4c1cc18bc25ff98f1835993d7545a8b394c
2 files changed
+107
-12
tests/kernels/test_minimax_m3_amd_ops.py
+94
@@ -422,3 +422,97 @@ def test_mxfp8_linear_emulation_bf16_at_load(
422
out = kernel.apply_weights(layer, x)
423
assert out.dtype == act_dtype # dtype-match preserved (no tl.dot/F.linear crash)
424
assert _relerr(out.float(), out_ref.float()) < 2e-2
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):
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
439
+ from unittest import mock
440
+
441
+ from vllm._aiter_ops import rocm_aiter_ops
442
+ from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import (
443
+ AiterMxfp8Experts,
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
+ )
450
+ experts.quant_config = SimpleNamespace(gemm1_clamp_limit=None)
451
+ experts.w1_scale_val = None
452
+ experts.w2_scale_val = None
453
+
454
+ captured = {}
455
+
456
+ def _fake_fused_moe(hidden_states, w1, w2, tw, ti, *, expert_mask, **kw):
457
+ captured["expert_mask"] = expert_mask
458
+ return torch.zeros_like(hidden_states)
459
+
460
+ w1 = torch.zeros(1, device=DEVICE)
461
+ w2 = torch.zeros(1, device=DEVICE)
462
+ out = torch.zeros(4, 8, device=DEVICE, dtype=torch.bfloat16)
463
+ hidden = torch.zeros(4, 8, device=DEVICE, dtype=torch.bfloat16)
464
+ tw = torch.ones(4, 2, device=DEVICE)
465
+ ti = torch.zeros(4, 2, dtype=torch.int32, device=DEVICE)
466
+
467
+ with mock.patch.object(rocm_aiter_ops, "fused_moe", side_effect=_fake_fused_moe):
468
+ experts.apply(
469
+ output=out,
470
+ hidden_states=hidden,
471
+ w1=w1,
472
+ w2=w2,
473
+ topk_weights=tw,
474
+ topk_ids=ti,
475
+ activation=None,
476
+ global_num_experts=global_num_experts,
477
+ expert_map=expert_map,
478
+ a1q_scale=None,
479
+ a2_scale=None,
480
+ workspace13=None,
481
+ workspace2=None,
482
+ expert_tokens_meta=None,
483
+ apply_router_weight_on_input=False,
484
+ )
485
+ return captured["expert_mask"]
486
+
487
+
488
+@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,
494
+ )
495
+
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)
501
+
502
+ idx_map = idx_map.to(DEVICE)
503
+ ep_mask = ep_mask.to(DEVICE)
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)
507
+
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
+ )
514
+
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
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
+13
-12
@@ -78,9 +78,6 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
78
79
@staticmethod
80
def _supports_parallel_config(moe_parallel_config) -> bool:
81
- # Both TP (expert_map=None) and EP are supported: apply() forwards the
82
- # expert_map as aiter's ``expert_mask`` (the per-rank local-expert
83
- # selection), mirroring the native rocm_aiter_moe path.
81
return True
82
83
@staticmethod
@@ -129,17 +126,21 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
126
limit = self.quant_config.gemm1_clamp_limit
127
swiglu_limit = 0.0 if limit is None else float(limit)
128
132
- # Under EP, aiter expects ``expert_mask`` as a 0/1 *local-expert* mask
133
- # over global ids with a trailing fake-expert sentinel slot
134
- # (shape ``[global_num_experts + 1]``), NOT vLLM's expert_map (a
135
- # global->local index map with -1 for non-local). Convert it; aiter
136
- # derives the global->local compaction from the mask itself. ``None``
137
- # under pure TP.
138
- if expert_map is not None:
129
+ # Under EP, aiter expects ``expert_mask``: a 0/1 *local-expert* mask over
130
+ # global ids with a trailing fake-expert sentinel slot (shape
131
+ # ``[global_num_experts + 1]``), from which it derives the global->local
132
+ # compaction. What ``RoutedExperts.expert_map`` hands us depends on the
133
+ # aiter master switch (``rocm_aiter_fmoe_enabled``).
134
+ # Branching on the (static) master flag — not the tensor contents —
135
+ # keeps this HIP-graph/torch.compile safe (no data-dependent sync).
136
+ # ``None`` under pure TP.
137
+ if expert_map is None:
138
+ expert_mask = None
139
+ elif self.moe_config.rocm_aiter_fmoe_enabled:
140
+ expert_mask = expert_map
141
+ else:
142
local_mask = (expert_map >= 0).to(torch.int32)
143
expert_mask = torch.cat([local_mask, local_mask.new_zeros(1)])
141
- else:
142
- expert_mask = None
144
145
# Route through the graph-safe ``rocm_aiter_fused_moe`` custom op so the
146
# call is captured under HIP graphs / torch.compile (a direct