Repository navigation
[ROCm][BugFix] Trim MoE GEMM input to valid token prefix - #53108
Open
amd-weisun wants to merge 12 commits into
Open
amd-weisun wants to merge 12 commits into
amd-weisun wants to merge 12 commits into
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
amd-weisun
marked this pull request as ready for review
August 20, 2026 13:02
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
amd-weisun
marked this pull request as draft
August 20, 2026 21:06
amd-weisun
force-pushed
the
fix/aiter-moe-trim-to-valid-tokens
branch
from
August 20, 2026 23:25
b4c7bd0 to
3970b3d
Compare
amd-weisun
marked this pull request as ready for review
August 20, 2026 23:26
amd-weisun
requested review from
mgoin,
pavanimajety and
zyongye
as code owners
August 20, 2026 23:26
1 task done
|
✅ @amd-weisun, CI is now available for this PR.
|
Contributor
|
@amd-weisun can you try to build a UT that'll drive this corruption scenario so we can be sure it stays gone as we move forward |
Contributor
|
Hi @amd-weisun, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
1 similar comment
Contributor
|
Hi @amd-weisun, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
…g AITER AITER's fused_moe kernel silently corrupts its own output (returns all-zero rows) when the allocated input buffer (hidden_states.shape[0]) is much larger than the real number of valid tokens it is asked to process. This has been confirmed independently of MoRI or any other EP dispatcher, via a standalone microbenchmark that calls AITER's fused_moe directly and compares against a from-scratch reference GEMM. This condition is common with EP dispatchers (e.g. MoRI) that pre-allocate a fixed-size receive buffer sized for the worst case (max_num_batched_tokens * world_size), while the actual per-step valid token count is usually far smaller (a handful of tokens per rank in normal decode-heavy serving). AITER's fused_moe() call receives the full pre-allocated buffer shape even though num_local_tokens correctly identifies only a small valid prefix within it. Fix: in AiterExperts.apply(), when num_local_tokens is available and smaller than the input buffer, slice hidden_states/topk_weights/ topk_ids/a1q_scale down to the valid prefix before calling rocm_aiter_fused_experts(), then scatter the (smaller) result back into a zero-initialized output buffer of the original shape. Padding rows beyond the valid prefix are never read downstream (num_local_tokens already tells consumers where the real data ends), so zero-filling them is semantically a no-op. Applied unconditionally: this is a real correctness defect, not a behavior change some configurations might prefer to opt out of, so there is no feature flag to disable it. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
Signed-off-by: amd-weisun <Wei.Sun@amd.com>
Drives AiterExperts.apply() directly with a synthetic post-dispatch-shaped buffer (most rows unwritten padding, a small valid prefix) at the row count where AITER's fused_moe kernel is known to corrupt its own output. Verified on MI300 hardware that reverting the trim makes the large-buffer case fail with non-finite output, confirming this catches the regression. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
mypy can't see across the _is_uniform_full_graph_batch() helper that batch_descriptor is non-None when it returns True, so assert it explicitly at the call site. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
Padding content past the valid prefix is never read downstream (the fix's own commit message notes zero-filling it is "semantically a no-op"), so asserting it stays exactly zero checks an implementation detail, not a real contract. The whole-buffer isfinite() check already covers the actual correctness/safety concern for that region. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
Adds tests/kernels/moe/test_mori_aiter_ep_trim_regression.py, which drives real MoRI InterNodeV1 dispatch, real AiterExperts.apply() (AITER GEMM), and real MoRI combine directly (bypassing FusedMoEFactory, matching test_moe_layer_multinode.py's approach), repeating dispatch/combine rounds on the same reused mori_op buffers to reproduce the buffer-trim regression. Verified on real 4-node/EP32 InterNodeV1 hardware via torchrun (two independent node allocations): with the trim fix present, the test passes across repeated rounds; with it reverted, real tokens' combined output goes all-zero after buffer reuse and the test fails with an AssertionError, matching the corruption mechanism this branch fixes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
amd-weisun
force-pushed
the
fix/aiter-moe-trim-to-valid-tokens
branch
from
September 21, 2026 15:48
788a75b to
2eb5a39
Compare
Move mechanism/rationale explanations out of code and into the PR description; keep only what's needed to run/understand each test. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
Rename test_mori_aiter_ep_trim_regression.py to test_rocm_mori_moe_ep.py and skip on non-ROCm platforms (previously only gated on mori availability). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
Contributor
|
@amd-weisun can you check on the mori test if we can drive this on a single node with 8 gpus split into two groups? |
Auto-select single-node (world_size == GPU_PER_NODE, split into 2 logical MoRI nodes) vs. genuine multi-node based on world_size, so the test no longer requires a real multi-host cluster. Bump rounds 20 to reliably catch the bug at single-node scale, where it's lower-probability per round than at real multi-node scale. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
This was referenced Sep 30, 2026
Merged
Author
|
/ci run |
|
✅ Triggered Buildkite CI #92166 for commit |
Collaborator
CI selector (shadow): 2 test steps (6 jobs) instead of 65 (88 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 (2)
Would skip (today's rules run them) (63)
Would add (today's rules do not run them) (0)none AMD mirrors: would skip (69)
AMD mirrors: would add (3)
2 changed files · base |
With the dispatch-output trim in MoriPrepareAndFinalize.prepare(), the trim in AiterExperts.apply() has nothing left to do for MoRI and only adds a second device-to-host sync and a second capture policy. Remove it and its helpers, drop the unit test that only exercised it, and soften the MoRI test's assertion message since either trim alone avoids the failure. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Signed-off-by: amd-weisun <Wei.Sun@amd.com>
Author
|
/ci run |
|
✅ Triggered Buildkite CI #92636 for commit |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
AITER's ROCm
fused_moe(AiterExperts) silently corrupts its own output(returns all-zero rows) when the allocated input buffer is much larger than
the real number of valid tokens it's asked to process. This is the normal
case for EP dispatchers like MoRI, which pre-allocate a worst-case buffer
(
max_num_batched_tokens * world_size) while real per-step token counts areusually a handful per rank.
AiterExperts.apply()passed the full bufferstraight to
rocm_aiter_fused_experts()even thoughnum_local_tokensidentifies the real valid prefix.
Confirmed independent of any EP dispatcher via a standalone microbenchmark
calling AITER's
fused_moe()directly against a from-scratch reference GEMM:AITER zeros the valid rows whenever the buffer-to-valid-token ratio is large
enough, regardless of routing pattern.
Fix: in
AiterExperts.apply(), slicehidden_states/topk_weights/topk_ids/a1q_scaledown to the valid prefix before callingrocm_aiter_fused_experts(), then scatter the result into a zero-initializedoutput buffer of the original shape (padding rows are never read downstream).
Applied unconditionally — no feature flag, since this is a correctness fix,
not an opt-in behavior change.
Under CUDA-graph capture, a live
.item()read ofnum_local_tokensisillegal, so the bound is instead a static worst-case value
(
graph_bs * topk * dispatch_group_size) under full-graph capture with auniform batch, or trimming is skipped entirely when no safe bound exists
(e.g. piecewise cudagraphs).
Test Plan
Hy3-preview (192 experts, top-8, BF16) deployment,
--all2all-backend mori_high_throughput, non-disaggregated, under DP-basedwideEP (
-tp 1 --data-parallel-size <N>) and TP-based EP (-tp <N>, no DP).Two checks per config: a 10-prompt known-answer probe (greedy) and full
GSM8K-1319 (5-shot, greedy).
At DP EP32 (4 nodes), ran a full 2x2 ablation — neither fix / each fix alone
/ both together — against a companion MoRI
combine()fix(#52639 ), since both defects trigger simultaneously
there.
Test Result
EP32 ablation (4 nodes) — both fixes needed together:
Neither fix alone reaches parity:
combine()alone makes no measurabledifference on the unfixed AITER defect; this trim alone recovers most of the
way but leaves visible incoherence. Both are independently necessary and
insufficient — only combined reaches the ~93% ceiling.
Full matrix, both fixes :
Unit Tests
tests/kernels/moe/test_rocm_mori_moe_ep.pyDrives the real production data flow directly — real MoRI
InterNodeV1dispatch + real AITER GEMM (AiterExperts.apply()) + real MoRI combine — across real multi-node hardware, reusing dispatch buffers across repeated (big round, small round) cycles to mimic real serving buffer reuse. Proves the actual real-world consequence: without the fix, a real, valid token's final combined output comes back wrong (all-zero), not just that some GEMM buffer contains garbage. Requires ≥2 real nodes with RDMA +moriinstalled — not routine-CI-friendly, intended as a hardware verification test.Run (one command per node,
<0..N-1>= that node's rank,<node0-ip>= rank-0 node's address, works at 2 nodes/EP16 or 4 nodes/EP32):Results:
1 passed1 failed—AssertionError: MoRI dispatch/combine returned all-zero output for real tokens after buffer reuse1 passed1 failed— same assertionReproduced across multiple independent runs/node allocations in both directions. Note: real EP16 e2e serving doesn't show a visible GSM8K accuracy drop, but this test reproduces the underlying defect at EP16 too — the corruption mechanism isn't exclusively an EP32-scale phenomenon, it's just rare/diluted enough in ordinary EP16 traffic not to move the aggregate accuracy metric.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.