Repository navigation
[ROCm][MoE] Split AiterExperts by weight-quantization key - #60322
AndreasKaratzas merged 8 commits into
Conversation
Break up the single catch-all AiterExperts modular-kernel class into
five narrower leaf classes, one per weight-quantization scheme:
- AiterUnquantizedExperts (None, None)
- AiterFp8BlockExperts (kFp8Static128BlockSym, kFp8Dynamic128Sym)
- AiterFp8TensorExperts (kFp8StaticTensorSym, {Static,Dynamic}TensorSym)
- AiterFp8ChannelExperts (kFp8StaticChannelSym, kFp8DynamicTokenSym)
- AiterMxfp4Experts (kMxfp4Static, {None,kMxfp4Dynamic}), gfx950-only
AiterExperts becomes a shared, non-instantiable base class (all other
methods -- apply(), workspace_shapes(), device/activation/parallel-config
gating, weight-conversion dispatch -- stay on the base since they are
already scheme-agnostic at runtime). This mirrors the existing pattern
used by the Cutlass family (CutlassExpertsFp8Base, CutlassExpertsNvfp4,
CutlassExpertsMxfp4), where AiterExperts was the remaining catch-all
outlier.
Update the oracle backend_to_kernel_cls() mappings (fp8.py, mxfp4.py,
unquantized.py) to return the appropriate leaf class(es) instead of the
old monolithic class; existing multi-class-list consumers already
iterate correctly so no further changes were needed there.
Update the modular-kernel test harness (mk_objects.py) to register each
leaf class with its own supported dtypes instead of one class registered
fp8_types-only -- this was silently excluding the unquantized and MXFP4
cases from ever being exercised by the harness. Convert `is AiterExperts`
identity checks in test_modular_kernel_combinations.py to
`issubclass(_, AiterExperts)` (safe: Config.fused_experts_type is a
required, non-Optional field, never None), and update direct-construction
call sites / parametrize tables in test_mori_moe.py, test_rocm_aiter_moe.py,
and test_topk_softplus_sqrt.py to reference the correct leaf class.
ROCm/AITER-only change; no CUDA or other-backend code paths touched.
Test plan:
pytest -x -v tests/kernels/moe/test_modular_kernel_combinations.py -k "aiter"
pytest -x -v tests/kernels/moe/test_rocm_aiter_moe.py
pytest -x -v tests/kernels/moe/test_mori_moe.py
pytest -x -v tests/kernels/moe/test_moe_layer.py -k "mori or aiter"
pytest -x -v tests/kernels/moe/test_topk_softplus_sqrt.py
All green on MI300 (gfx942) in the dv_DistriDPBasic container
(HIP_VISIBLE_DEVICES=0,1,2,3). Full tests/kernels/moe/ suite also run
as a regression sanity check; all failures traced to a pre-existing,
unrelated triton-version environment issue or full-suite resource
contention, not this change.
MXFP4 leaf (AiterMxfp4Experts) is a mechanical, behavior-preserving
relocation of the existing gfx950/gfx1250 gate, verified via code
review and the gfx950-independent gate-logic unit test
(test_aiter_mxfp4_quant_scheme_support_matches_gfx950); full numeric
MXFP4 verification remains blocked on MI350/gfx950 hardware access,
consistent with the existing status of steps 1b/2b in the issue's
test-coverage plan.
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
…eme split Package-level re-export and the design doc's kernel table still referenced the now-abstract AiterExperts base class instead of the new concrete leaf classes: - vllm/model_executor/layers/fused_moe/__init__.py re-exported AiterExperts directly. Since it's no longer instantiable, this is a behavior change for anything constructing it via this import path. Following the existing convention for CutlassExpertsFp8Base (whose abstract base is never re-exported, only its concrete leaves), export the 5 concrete leaf classes instead. - docs/design/moe_kernel_features.md's "Fused Experts Kernels" table listed `AiterExperts` as the rocm aiter moe row's representative modular class; update to list the 5 leaf classes with proper mkdocstrings cross-reference links, matching every other row in the table. Found via a full-repo grep for all AiterExperts references (not just the files already known to reference it) plus a check for any reflection-based class registries (__subclasses__, etc.) that might also need updating -- none found. No functional/test behavior change: nothing in-repo constructs AiterExperts via the top-level package import or relies on class-based reflection. Signed-off-by: Divakar Verma <divakar.verma@amd.com>
|
Documentation preview: https://vllm--60322.org.readthedocs.build/en/60322/ |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The claimed MXFP4 modular-harness coverage remains unreachable, and the updated feature matrix contains incorrect capability information.
Review effort: Balanced
Findings: 1
Open (2)
What changed in this PR
Splits the ROCm AITER MoE implementation into quantization-specific expert classes and updates selection, exports, tests, and documentation accordingly.
Changes:
- Adds five quantization-specific
AiterExpertssubclasses. - Updates backend oracles and tests to select the appropriate subclass.
- Registers the subclasses in the modular-kernel test harness.
| File | Description |
|---|---|
vllm/model_executor/layers/fused_moe/oracle/unquantized.py |
Selects the unquantized AITER class. |
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py |
Selects the MXFP4 AITER class. |
vllm/model_executor/layers/fused_moe/oracle/fp8.py |
Selects among FP8 AITER classes. |
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py |
Introduces quantization-specific subclasses. |
vllm/model_executor/layers/fused_moe/__init__.py |
Re-exports the new classes. |
tests/kernels/moe/test_topk_softplus_sqrt.py |
Updates expected MXFP4 class selection. |
tests/kernels/moe/test_rocm_aiter_moe.py |
Updates MXFP4 support checks. |
tests/kernels/moe/test_mori_moe.py |
Chooses the appropriate MoRI expert class. |
tests/kernels/moe/test_modular_kernel_combinations.py |
Recognizes AITER subclasses in test logic. |
tests/kernels/moe/modular_kernel_tools/mk_objects.py |
Registers per-scheme AITER classes. |
docs/design/moe_kernel_features.md |
Links the new expert classes. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| (AiterFp8BlockExperts, fp8_types), | ||
| (AiterFp8TensorExperts, fp8_types), | ||
| (AiterFp8ChannelExperts, fp8_types), | ||
| (AiterMxfp4Experts, mxfp4_types), |
The AiterExperts leaf classes are FusedMoEExpertsModular subclasses (Modular: N -> Y), AiterUnquantizedExperts supports the unquantized (None, None) scheme, and the shared base's _supports_activation() accepts MoEActivation.SITU. Update the table row to reflect this. Signed-off-by: Divakar Verma <divakar.verma@amd.com>
|
🔥🔥🔥 |
|
/ci run |
|
✅ Triggered Buildkite CI #93301 for commit |
| output.copy_(result) | ||
|
|
||
|
|
||
| class AiterUnquantizedExperts(AiterExperts): |
There was a problem hiding this comment.
Is there any other logic that is weight quant key-specific from the parent class that could be ported in the children? Let me know if it makes sense or not / this could also be left as is or done later.
things like:
vllm/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
Lines 322 to 339 in d9503dc
vllm/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
Lines 391 to 406 in d9503dc
this is simply with the goal of making AiterExperts less obscure and more easily debuggable.
There was a problem hiding this comment.
Also if there is, any restriction on weight quant key / activation quant key / activation / gate_mode combination.
There was a problem hiding this comment.
Regarding, porting it down to the children: Yeah, that can be moved down, mainly the quant_method derivation. It is basically the same split we just did for _supports_quant_scheme(). We can keep the scope narrow in this PR since apply() didn't strictly need to change for the split to be correct, but I'm happy to push this as a follow-up.
The gate_mode logic is a bit more tangled since it checks activation == SITU before weight quant key, so it doesn't split as cleanly per class. I'd lean toward leaving that one in the shared function.
CI selector (shadow): 75 test steps (126 jobs) instead of 69 (92 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (75)
Would skip (today's rules run them) (36)
Would add (today's rules do not run them) (42)
AMD mirrors: would skip (2)
AMD mirrors: would add (76)
13 changed files · base |
|
This pull request has merge conflicts that must be resolved before it can be |
There was a problem hiding this comment.
Tested on MI355 (gfx950), PR head 04f4d2cdeb, VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1, GPUs 0-3.
Test (author's plan + tests/rocm/test_moe_padding.py) |
PR head | main |
|---|---|---|
test_modular_kernel_combinations.py -k aiter |
13 passed, 18 failed (all test_aiter_moe_padding_matrix) |
padding matrix: 12 passed, 6 failed (pre-existing gfx950 CK device_gemm ... does not support this GEMM problem) |
test_rocm_aiter_moe.py (incl. 3 gfx950 MXFP4 tests) |
70 passed | |
test_mori_moe.py |
14 passed | |
test_moe_layer.py -k "mori or aiter" |
3 passed, 20 skipped (generic config skips) | |
test_topk_softplus_sqrt.py |
1672 passed, 4 skipped | |
tests/rocm/test_moe_padding.py |
7 passed |
The AITER cases in the modular-kernel sweep are the same 10 on both sides (same configs, now spread across the right leaves).
-
tests/kernels/moe/test_modular_kernel_combinations.py:952**fused_experts_type=AiterExperts,
test_aiter_moe_padding_matrix(#59333, brought in by your merge of main) still builds aConfigwith the base class. The base is no longer inEXPERT_INFO, soconfig.is_valid()→expert_info(AiterExperts)fails withAssertionError(mk_objects.py:162). All 18 params fail at PR head in about 2 s. Pick the leaf from the quant config, e.g. a small{None: AiterUnquantizedExperts, ...}lookup or a helper that returns the first leaf whose_supports_quant_scheme(*config.fp8_quant_key_pair())is true. This also contradicts the "no overlap with #59332/#59333/#59334" line in the description. -
Merge with current main (conflict + two new
AiterExpertsdependencies)**
git merge-treeagainstorigin/main(2ddde19e8e) has a conflict in thetest_modular_kernel_combinations.pyimports. Since your last merge, main also added:
- #59334
test_aiter_moe_token_padding*:fused_experts_type=AiterExperts(main L1116). This will hit the sameexpert_infoassertion as point 1, so it needs the same leaf lookup. - #60633
modular_kernel_tools/common.py:737:if fe_name == "AiterExperts" and (...). This compares the class__name__, so after the split it never matches. Nothing fails, but the AITER weight pre-shuffle stops and every AITER harness case falls back to JIT-compiling thepreshuffle_offmodule (the exact slowdown #60633 fixed). Switch toissubclass(fe_type, AiterExperts).
(Points 1 and the #59334 case come from a test run and a traced code path. )
- PR description, "Bugs found and fixed along the way": The claim isn't true. Both the merge base (
d9503dc503) and current main already registerAiterExpertswithcommon_float_types, with a comment saying it covers(None, None). Collect-only shows the same 10 AITER sweep cases before and after this PR. The split doesn't change harness coverage, so please drop or reword this section. The commit message makes the same claim and also says MXFP4 is "now exercised", which it isn't (point 4).
Resolves a conflict in tests/kernels/moe/test_modular_kernel_combinations.py between this branch's AiterExperts split and main's new HIP-graph token padding tests. Since AiterExperts is now a non-instantiable base class, updates main's two new tests (test_aiter_moe_padding_matrix and test_aiter_moe_token_padding_garbage_rows) to select the concrete per-scheme leaf class instead of instantiating AiterExperts directly. Also fixes a latent bug this branch's split exposed in modular_kernel_tools/common.py: _maybe_convert_weights_for_experts matched on the literal class name "AiterExperts" to decide whether to apply the gfx950 weight-preshuffle workaround, which no longer matches any instantiated class post-split. Switched to an issubclass(fe_type, AiterExperts) check so the workaround still applies to AiterUnquantizedExperts/AiterFp8BlockExperts as before. Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…mem_cudagraph_pool Main's use_cumem_cudagraph_pool (vllm/config/vllm.py) unconditionally reads model_config.enable_sleep_mode/sleep_mode_backend/sleep_mode_offload_cudagraph. The modular-kernel test harness's Config.make_env_data() fakes model_config with a bare SimpleNamespace that predates these fields, so any multi-GPU modular-kernel test that initializes a process group now raises AttributeError. Add the three fields with their real ModelConfig defaults. Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2386480 to
f4a451d
Compare
…binations test_modular_kernel_combinations_multigpu always needs world_size=2 and self-skips via its own GPU-count check otherwise. The only AMD CI job covering this file runs on the single-GPU mi300_1 pool, so these cases have never actually executed in CI - they show green while silently skipping every run. Add a dedicated mi300_2 job (num_gpus: 2) that runs just the -k "multigpu" cases, and deselect them from the mi300_1 job so the two jobs don't both claim the same test IDs. This surfaced while debugging a real failure caught running these tests locally on MI300 (2 GPUs): a stale model_config test stub (fixed in the AiterExperts-split commit on this branch) broke every multi-GPU modular-kernel test with an AttributeError on sleep_mode_offload_cudagraph. That bug had been silently unreachable in CI for the same reason - no job ever ran these cases with enough GPUs to hit it. Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Divakar Verma <divakar.verma@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #94317 for commit |
|
/amd-ci run nightly --allow-stale |
|
✅ Triggered Buildkite AMD CI #14512 for commit
|
…est_areas Some AMD test groups run the same command in test-amd.yaml and in the test_areas mirror, but with a different version or flags. Use the more recently changed side in both places: - (MI300) Quantization: torchao 0.17.0 -> 0.18.0 in test_areas, matching test-amd.yaml (vllm-project#53009). - (MI355 DPX) Language Models (Extended Generation): use the fix-rocm-7.0-warp-size-constexpr mamba branch in test-amd.yaml, matching test_areas (vllm-project#57877). - (MI300) MoE Kernels: run test_modular_kernel_combinations with AITER enabled and without the multi-GPU cases in test_areas, matching test-amd.yaml (vllm-project#57074, vllm-project#60322). - (MI355 DPX) Basic Correctness Models: run with VLLM_TARGET_TEST_SUITE=MI355 in test_areas, as test-amd.yaml does. Signed-off-by: aarushjain29 <aarushi.jain2@amd.com>
…est_areas Some AMD test groups run the same command in test-amd.yaml and in the test_areas mirror, but with a different version or flags. Use the more recently changed side in both places: - (MI300) Quantization: torchao 0.17.0 -> 0.18.0 in test_areas, matching test-amd.yaml (vllm-project#53009). - (MI355 DPX) Language Models (Extended Generation): use the fix-rocm-7.0-warp-size-constexpr mamba branch in test-amd.yaml, matching test_areas (vllm-project#57877). - (MI300) MoE Kernels: run test_modular_kernel_combinations with AITER enabled and without the multi-GPU cases in test_areas, matching test-amd.yaml (vllm-project#57074, vllm-project#60322). Signed-off-by: aarushjain29 <aarushi.jain2@amd.com>
…est_areas Some AMD test groups run the same command in test-amd.yaml and in the test_areas mirror, but with a different version or flags. Use the more recently changed side in both places: - (MI300) Quantization: torchao 0.17.0 -> 0.18.0 in test_areas, matching test-amd.yaml (vllm-project#53009) and the NVIDIA step. - (MI355 DPX) Language Models (Extended Generation): use the fix-rocm-7.0-warp-size-constexpr mamba branch in test-amd.yaml, matching test_areas (vllm-project#57877). - (MI300) MoE Kernels: run test_modular_kernel_combinations with AITER enabled and without the multi-GPU cases in test_areas, matching test-amd.yaml (vllm-project#57074, vllm-project#60322). - (MI355 DPX) Basic Correctness Models: drop VLLM_TARGET_TEST_SUITE=MI355 from test-amd.yaml, matching test_areas. It is only read by test_models_distributed, which needs 2 GPUs and detects MI355 itself. Signed-off-by: aarushjain29 <aarushi.jain2@amd.com>
…ct#60322) Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <AndreasKaratzas@users.noreply.github.com>


Summary
Split AITER experts by weight quantization key. Splits the single catch-all
AiterExpertsmodular-kernel class into five narrower leaf classes, one per weight-quantization scheme:AiterUnquantizedExperts—(None, None)AiterFp8BlockExperts—(kFp8Static128BlockSym, kFp8Dynamic128Sym)AiterFp8TensorExperts—(kFp8StaticTensorSym, {Static,Dynamic}TensorSym)AiterFp8ChannelExperts—(kFp8StaticChannelSym, kFp8DynamicTokenSym)AiterMxfp4Experts—(kMxfp4Static, {None, kMxfp4Dynamic}), gfx950-onlyAiterExpertsbecomes a shared, non-instantiable base class. All other methods (apply(),workspace_shapes(), device/activation/parallel-config gating, weight-conversion dispatch) stay on the base since they're already scheme-agnostic at runtime — only_supports_quant_scheme()meaningfully diverges per scheme, so that's the only thing that moved into the leaves. This is intentionally scoped narrowly:apply()'s internalquant_methodderivation, theasm_moe_tkw1dispatch, and MXFP4'sgate_mode/bias selection all remain in the shared base, still keyed off runtimequant_configflags rather than owned per-leaf. Movingquant_methoddown is a reasonable follow-up;gate_modeis more tangled (it checksactivation == SITUbefore weight-quant key) and likely stays shared. Issue #54966's "own apply()/scale contract per leaf" isn't done by this step — only the scheme-gating (_supports_quant_scheme) is.This mirrors the pattern already used by the Cutlass family (
CutlassExpertsFp8Base,CutlassExpertsNvfp4,CutlassExpertsMxfp4), whereAiterExpertswas the remaining monolithic outlier. Part of #54966Test plan
MI300 (gfx942),
dv_DistriDPBasiccontainer,HIP_VISIBLE_DEVICES=0,1,2,3:All green. Also ran the full
tests/kernels/moe/suite as a broader regression sanity check (ROCm/AITER-only change; no CUDA or other-backend code paths touched); all observed failures traced to a pre-existing, unrelated triton-version environment issue (is_hip_gfx1250import error) or full-suite resource contention, not this change — confirmed by rerunning the affected tests standalone.MI350 (gfx950), same test files — all green, including the three gfx950-only numeric MXFP4 tests:
test_aiter_mxfp4_quant_scheme_support_matches_gfx950— PASSEDtest_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy— PASSED, 100% match (16384/16384) against the dequantized MXFP4 referencetest_aiter_fused_moe_mi350_mxfp4_w4a16_determinism— PASSEDMerged with main (resolving the conflict in
test_modular_kernel_combinations.py) and re-ran the full test plan on MI355 (gfx950), twice, on two disjoint GPU sets (HIP_VISIBLE_DEVICES=0,1,2,3and4,5,6,7), identical results both times:test_modular_kernel_combinations.py -k "aiter": 48 passed, 6 failed, 6 skipped. The 6 failures (test_aiter_moe_padding_matrix[{none,hidden}-{fp8_channel_token,fp8_tensor_tensor,fp8_tensor_token}]) are a pre-existing AITER CKdevice_gemmGEMM-shape limitation on gfx950 for non-block fp8 schemes at the unpaddedintermediate_size=384shape, unrelated to this split — these are brand-new tests added by main ([ROCm][CI][AiterExperts] Add test coverage for hidden/intermediate padding correctness #59333/[ROCm][CI] Add coverage for AiterExperts token padding: inf/nan garbage rows #59334) hitting real gfx950 hardware behavior for the first time; they reproduce identically on a cleanmaincheckout on this hardware.test_rocm_aiter_moe.py+test_mori_moe.py+test_topk_softplus_sqrt.py: 1756 passed, 4 skipped.test_moe_layer.py -k "mori or aiter": 3 passed, 20 skipped.Known limitations / follow-up
AiterMxfp4Experts's harness registration (mk_objects.py) doesn't currently produce any modular-kernel test case —MK_QUANT_CONFIGShas noquant_dtype="mxfp4"entry, so it's registered but unreachable. Real MXFP4 modular-kernel coverage is tracked separately in [ROCm][MoE] Add AITER MXFP4 W4A16/W4A4 modular-kernel test coverage #60886._supports_quant_scheme()to the leaves (see Summary); other scheme-specific runtime dispatch stays in the shared base for now.Bugs found and fixed along the way
test_aiter_moe_padding_matrix,test_aiter_moe_token_padding_garbage_rows) that build aConfigwithfused_experts_type=AiterExpertsdirectly. Since the split makesAiterExpertsa non-instantiable base no longer registered in the harness, this failsconfig.is_valid()(expert_info(AiterExperts)→AssertionError). Fixed both call sites to pick the correct per-scheme leaf class from the quant config instead.modular_kernel_tools/common.py's_maybe_convert_weights_for_expertsgated the gfx950 AITER weight-preshuffle workaround ([CI][ROCm] Reuse AITER's prebuilt MoE module in the modular-kernel harness #60633) onfe_name == "AiterExperts"— a literal class-name string match that no longer matches any instantiated class after the split, silently turning off the preshuffle and falling back to JIT-compilingpreshuffle_offfor every AITER harness case (the exact slowdown [CI][ROCm] Reuse AITER's prebuilt MoE module in the modular-kernel harness #60633 fixed). Switched toissubclass(fe_type, AiterExperts).model_config.sleep_mode_offload_cudagraphread inVllmConfig.use_cumem_cudagraph_poolthat the modular-kernel harness's fakemodel_config(a bareSimpleNamespace) didn't have, breaking any multi-GPU modular-kernel test. Added the missing fields with their realModelConfigdefaults.model_configbug was independently found and fixed in [Bugfix][ROCm] Fix stale model_config test stub breaking MoE multi-GPU combos; add mi300_2 CI job #61260, which also discovered that it was never actually caught in CI: the only AMD job covering this file runs on the single-GPUmi300_1pool, sotest_modular_kernel_combinations_multigpu(which needsworld_size=2) self-skips via its own GPU-count check on every run. [Bugfix][ROCm] Fix stale model_config test stub breaking MoE multi-GPU combos; add mi300_2 CI job #61260 added a dedicatedmi300_2job for these cases; that job is now included in this PR (latest commit) instead, so [Bugfix][ROCm] Fix stale model_config test stub breaking MoE multi-GPU combos; add mi300_2 CI job #61260 is being closed in favor of landing it here.Why this isn't a duplicate
No PR performs this exact split. #59333/#59334 (merged into main after this PR was opened) do instantiate
AiterExpertsdirectly in two new tests — as noted above, merging main surfaced this as a real incompatibility, now fixed by picking the correct leaf class in both tests.AI assistance disclosure
This change was developed with AI assistance (Claude Code),