[ROCm][MoE] Fix expert_map vs AITER expert_mask for non-AITER experts under EP - #49758
Conversation
6cffcaa to
af1f447
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
7009110 to
1211225
Compare
… under EP RoutedExperts.expert_map returns the AITER 0/1 expert_mask whenever AITER fused MoE is enabled (VLLM_ROCM_USE_AITER=1), regardless of which experts kernel runs. Non-AITER kernels (Triton moe_wna16, MXFP4 emulation, Marlin, ...) expect the canonical global->local expert_map (-1 for non-local) and misread the 0/1 mask as slot indices, collapsing every token onto local slot 0/1 -> garbage under EP. Add FusedMoEExperts.expects_aiter_expert_mask (True only for the AITER experts that consume the mask) and RoutedExperts.expert_map_for(experts), which hands each kernel the representation it expects. Wire it into the modular method and the compressed-tensors MXFP4 method; the WNA16 non-modular path uses the raw map directly. Validated: Kimi-K2.6 INT4 (compressed-tensors WNA16), MI350X gfx950, TP4, --enable-expert-parallel, gsm8k 5-shot exact-match 0.00 -> 0.95. Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1211225 to
7361ba5
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
This PR currently introduces a slight 1-2% regression under AITER and AITER FP8 MoE under EP due to the recomputation of the mask. I'll work on a follow-up PR to make expert_mask a functional of |
|
cc @tjtanaa could you PTAL? |
| @property | ||
| def expert_map(self) -> torch.Tensor | None: | ||
| return ( | ||
| self._expert_map if not self.rocm_aiter_fmoe_enabled else self.expert_mask |
There was a problem hiding this comment.
Since we are storing self.moe_config, we can use this to determine if we are using aiter moe_backend or not to pick if we want to return self._expert_map or self._expert_mask.
There was a problem hiding this comment.
with this, I think we don't have to modify vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
and
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
There was a problem hiding this comment.
Well that would work when we explicitly pass in --moe-backend aiter, but not when it's left to the auto default and it's just picking the default based on the oracle and the quant dtype/gfx arch etc.
Also in this instance, for example, vllm serve moonshotai/Kimi-K2.6 --moe-backend aiter actually dispatches to the FlyDSL MXFP8 MoE, which does not take the 0/1 mask, so --moe-backend aiter is actually misleading/insufficient; you'd have to rederive the actual backend from the oracles anyway.
There was a problem hiding this comment.
@tjtanaa could you PTAL? I went ahead with the refactor, but this touches the routed_experts and MoE classes, so may need additional review?
…forward derive) RoutedExperts.expert_map now returns the precomputed 0/1 expert_mask to the AITER CK experts (which set consumes_expert_mask) and the canonical -1 map to everyone else, keyed on the resolved experts kernel rather than the global VLLM_ROCM_USE_AITER switch. The AITER experts forward that mask to aiter unchanged, eliminating the per-forward derive (cat[(expert_map>=0), 0]); quark's non-modular path passes layer.expert_mask directly. Makes expert_mask a live buffer again. Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
✅ @Rohan138, CI is now available for this PR.
|
|
/ci run |
|
✅ Triggered Buildkite CI #83212 for commit |
NOT FOR FILING. The same fix merged upstream as vllm-project#49758 (0f0cb91, 2026-08-11). This commit is on the branch only so its tip equals the tree the rest of the series was built and tested against; drop it when rebasing onto a main that contains 0f0cb91. Under EP, `RoutedExperts.expert_map` handed the experts either AITER's 0/1 `expert_mask` or vLLM's -1 index map, keyed on `rocm_aiter_fmoe_enabled` -- a global environment switch. That is the wrong predicate: an explicit `--moe-backend` makes the oracle return that backend before it reaches its AITER branch, so `VLLM_ROCM_USE_AITER_MOE=1` with `--moe-backend triton` handed Triton the mask and produced silently wrong output. vllm-project#49758 keys on `consumes_expert_mask` on the resolved kernel, which is the better predicate; this version keys on the requested `--moe-backend`, so it still assumes AITER under `auto`. Both close the explicit-backend case. Measured at EP=4 on DeepEP high-throughput: max |delta| 1.57, batch variant in both arms (32/32 with `VLLM_BATCH_INVARIANT` set and 32/32 without), and gone at DP=1. The both-arms result is the point: the symptom is not a determinism problem, so no batch-invariance test will catch it. Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
… under EP (vllm-project#49758) Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
3 participants
RoutedExperts.expert_mapreturned the AITER 0/1expert_maskwhenever AITERfused MoE is enabled (
VLLM_ROCM_USE_AITER=1), regardless of which expertskernel actually runs. Non-AITER kernels (Triton
moe_wna16, MXFP4 emulation,Marlin, …) expect the canonical global→local
expert_map(-1for non-local)and misread the 0/1 mask as slot indices → every token collapses onto local
slot 0/1 → garbage under EP.
Fix:
RoutedExperts.expert_mapasks the resolved experts kernel whichrepresentation it consumes. Only the AITER CK kernels —
AiterExperts(viarocm_aiter_fused_experts) andAiterMxfp8Experts— setconsumes_expert_maskand receive the precomputed 0/1
expert_mask; every other kernel gets thecanonical
_expert_map. This keys on the actual per-layer kernel rather thanthe global
VLLM_ROCM_USE_AITERswitch (which does not reflect the kernel aWNA16 / Marlin / emulation layer ends up on). The AITER experts forward the
precomputed mask to aiter unchanged — no per-forward derive — and quark's
non-modular path passes
layer.expert_maskdirectly.expert_maskstays alive buffer; every quant method keeps passing
layer.expert_mapunchanged.Accuracy (Kimi-K2.6 INT4 / compressed-tensors WNA16, MI350X gfx950, TP4):
plain
--enable-expert-parallelgsm8k 5-shot 0.00 → 0.95; DP2×TP4 EP viamori_high_throughput(modular kernel) 0.925/0.955. The mask handed to theAITER kernels is byte-identical to the previously-derived
expert_mask, soAITER paths (fp8/mxfp4/mxfp8/unquantized/quark) are behavior-unchanged.
Perf (DeepSeek-R1-0528-MXFP4, TP8 +
--enable-expert-parallel, gfx950): anearlier derive-based revision already showed no e2e regression (mean TPOT within
run noise, GPU-kernel time flat); this revision removes the derive entirely — the
AITER experts read the precomputed
expert_mask, so the four tiny per-MoE-layerkernels (fill/compare/cast/cat) are gone. Trace re-confirmation to follow.
Draft — AI assistance was used; the submitter is validating and iterating.