Repository navigation
Conversation
Contributor
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a bug in the routing replay capturing mechanism where TP-local expert IDs were incorrectly handled when DP=1. Previously, the system treated the local shard as the complete batch, leading to stale expert IDs in the device buffer for other TP ranks. The changes introduce a gathering step to reconstruct the full batch before writing to the buffer, ensuring accurate expert routing replay. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
wuruiyun
changed the title
Contributor
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [BugFix] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][BugFix] Gather TP-sharded topk_ids in single DP routing captureSuggested PR Summary:
### What this PR does / why we need it?
This PR addresses an issue in single DP (DP == 1) scenarios with TP > 1 where sharded communication types (All2All, MC2, Fused MC2) split the token dimension across the TP group. This caused `topk_ids` to only hold the current rank's shard, leaving other rows in `device_buffer` stale and leading to incorrect expert IDs during routing replay. This change gathers the shards back across the TP group using `dist.all_gather` when these conditions are met.
Feedback has been provided to call `.contiguous()` on `topk_ids` before gathering to prevent potential data corruption or crashes on Ascend NPUs.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Tested via new unit tests `test_single_dp_tp_sharded` and `test_single_dp_tp_allgather` in `tests/ut/patch/worker/test_patch_routed_experts_capture.py`.
Ronald1995
reviewed
Oct 13, 2026
6 tasks done
wuruiyun
force-pushed
the
fix/issue-15451
branch
from
October 13, 2026 10:59
8c2cd2d to
e5bb84e
Compare
Ronald1995
reviewed
Oct 14, 2026
wuruiyun
force-pushed
the
fix/issue-15451
branch
from
October 15, 2026 04:29
e5bb84e to
2d62647
Compare
Ronald1995
reviewed
Oct 15, 2026
…n DP=1 Fixes vllm-project#15451. The routed-experts (R3 / routing replay) capture patch assumed that `topk_ids` covers the whole batch whenever `ctx.dp_metadata is None`: if ctx.dp_metadata is None: # single dp token_num_per_dp = topk_ids.shape[0] That is wrong. `ctx.dp_metadata` describes data parallelism only, but the All2All and MC2 strategies also split the token dimension inside the TP group: `PrepareAndFinalizeWithAll2All.prepare` pads to the TP size, and `PrepareAndFinalizeWithMC2.prepare` pads to `padded_num_tokens = ceil(max_tokens_across_dp / tp_size) * tp_size`; both then `tensor_split` the result across the TP ranks. The capture hook runs after `prepare()` (`AscendRoutedExperts.forward_impl` -> `prepare()` -> `BaseRouter._select_experts` -> `capture_fn(topk_ids)`), so with TP > 1 it only ever sees this rank's shard. Capturing that shard as the complete batch wrote just `device_buffer[:batch / tp_size]`, and every remaining row kept its previous contents. On a freshly started worker those rows stay zero; after warmup they hold another batch's expert ids. R3 then replays them. Measured on Qwen3-30B-A3B, TP=2 + EP, DP=1, 6-token prefill (shard length 3), before the fix: routed_experts shape=(13, 48, 8) rows(all-zero) = [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0] Gather the per-rank shards back before writing the buffer: - a new `_TP_SHARDED_COMM_TYPES` set names the strategies that split (`ALLTOALL`, `MC2`, `FUSED_MC2`). `ALLGATHER` is deliberately excluded: it gathers the batch across the EP/DP group before routing, so its `topk_ids` already spans the whole batch. - when `tp_size > 1` and the active comm type is in that set, the DP=1 branch all-gathers the shards into the padded batch first. The number of rows to rebuild is read from the padding target `prepare()` used, not inferred from this rank's shard: the two only agree for an even split, and All2All splits a batch with `num_tokens >= tp_size` unevenly whenever it is not a multiple of the TP size. `_padded_batch_rows` therefore returns - `max(num_tokens, tp_size)` for All2All, which pads to the TP size and only when the batch is smaller than it, and - `max(num_tokens, padded_num_tokens)` for MC2/FUSED_MC2, where `num_tokens` is the pre-pad batch size `prepare()` recorded and `padded_num_tokens` is `_EXTRA_CTX.padded_num_tokens`. Under sequence parallelism the wrapper has already sharded the sequence to a multiple of the TP size, so `padded_num_tokens` is the whole answer there. The helper takes no rank-local input at all, so its result is rank-invariant by construction, and it raises instead of falling back to `shard_rows * tp_size` when the metadata is missing -- guessing there would silently bring the rank-dependent length back. The gather decision itself is made only from rank-invariant values (`tp_size` and the comm type) for the same reason: the previous `total_rows > shard_rows` guard could in principle let some ranks gather while others skipped. Measured for All2All with `num_tokens=241, tp=2` (uneven split): both ranks derive 241, where a `shard_rows * tp_size` derivation gave 242 on rank 0 and 240 on rank 1. Even splits (MC2 with any batch, All2All with a multiple of the TP size, sequence parallel) are unchanged, and the two derivations produce byte-identical captures on the All2All 240- and 241-token prefills. `_gather_tp_sharded_topk_ids` also passes `topk_ids.contiguous()` to the collective. Two alternatives from the issue are deliberately not used: - `n != full_num_tokens` is false for a single-token batch, where MC2 still pads to `tp_size` and splits, so the gather would be skipped and the second rank's row would stay stale. - `prepare_finalize.num_tokens` is the local sequence-parallel shard when `replace_allreduce` is set, not the full batch. Writing the whole gathered tensor, instead of trimming it to a per-rank token count, avoids a second token-count estimate: the padding `prepare()` appends is always a suffix, and the consumer only reads `device_buffer[:num_tokens]` for the step. Note (pre-existing, not changed here): `All2All.prepare` computes `pad_size = tp_size - num_tokens`, which pads only when the batch is smaller than the TP group. A batch with `num_tokens >= tp_size` and `num_tokens % tp_size != 0` is therefore split into unequal chunks. Deriving the total above keeps this capture consistent with `All2All.finalize`, which uses the same pre-split shape, but if the intent of that line was "pad to a multiple of the TP size" it should read `(-num_tokens) % tp_size`. Worth a separate look. Tests: `tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py` gains `test_moe_routing_replay_covers_every_token`, which drives both TP-sharded strategies from one engine and asserts the row count plus that no row is empty. The strategy is picked by `num_tokens <= mc2_tokens_capacity` and, in eager mode without CUDA-graph sizes, that capacity is `max_num_reqs * decode_query_len` (measured: 8 with `max_num_seqs=8`). The test therefore pins `max_num_seqs=8` and uses two prompts: a 6-token one that routes through MC2 and a 49-token one that routes through ALLTOALL (both confirmed with `MoE comm method selected: ... method=MoECommType.鈥). Eager mode is deliberate: graph warm-up would pre-fill the capture buffer, so an unwritten row would hold stale expert ids instead of zeros and the assertion could no longer tell the two apart. Measured on Qwen3-30B-A3B: on the unpatched code both prompts fail, on `rows [3, 4, 5] of 13` for the MC2 prompt and on `rows [25..48] of 56` for the ALLTOALL one; with this patch both pass. The new case is restricted to Qwen3-30B-A3B: what it pins is model independent, and the case above already loads Qwen3.5-35B-A3B, so loading it a second time would only double the file's runtime. No unit tests are added; the existing unit suite is unchanged and still passes. Signed-off-by: wuruiyun <1251540982@qq.com>
wuruiyun
force-pushed
the
fix/issue-15451
branch
from
October 15, 2026 07:38
2d62647 to
44e25d5
Compare
linfeng-yuan
previously approved these changes
Oct 15, 2026
…ategies Signed-off-by: wuruiyun <1251540982@qq.com>
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.
What this PR does / why we need it?
Fixes #15451. The routed-experts (R3 / routing replay) capture patch normalises
topk_idsfor dataparallelism, but the strategies that split the token dimension inside the TP group were never handled
for
DP == 1.The patch only reconstructed the batch when
ctx.dp_metadata is not None, so withDP == 1it treated theshard as the complete batch:
Only
device_buffer[:batch / tp_size]was written. The remaining rows kept whatever they held before:zeros on a freshly started worker, another batch's expert ids after warmup. R3 then replays them, and
nothing about it is visible in the response text.
Measured on Qwen3-30B-A3B, TP=2 + EP, DP=1, 6-token prefill (shard length 3):
The fix gathers the per-rank shards back before writing the buffer:
_TP_SHARDED_COMM_TYPESnames the strategies whoseprepare()splits inside the TP group(
ALLTOALL,MC2,FUSED_MC2).ALLGATHERis deliberately absent: it gathers the batch across theEP/DP group before routing, so its
topk_idsalready spans the whole batch and gathering again wouldbe wrong.
dp_metadata is Nonebranch, whentp_size > 1and the active comm type is in that set, theshards are all-gathered into the padded batch before the buffer write.
Does this PR introduce any user-facing change?
How was this patch tested?
============================= test session starts ==============================
=== environment ===
vllm-ascend : 44e25d5 [Bugfix][Worker] Fix routing replay capturing TP-local expert ids when DP=1
vllm : c4df37dfcc [ROCm][BugFix][The Rock] Fix mori build for the rock (#59372)
model : Qwen/Qwen3-30B-A3B, TP=2 + EP, DP=1, cudagraph_capture_sizes=[1,2,4,8]
test : the pre-existing test_moe_routing_replay, extended (two prompts, row checks)
=== WITH the fix ===
=========== 1 passed, 1 deselected in 161.69s (0:02:41) ===========
=== WITHOUT the fix (upstream capture patch) ===
f"rows {empty_rows} of {expected_rows} carry no routing data for a "
E AssertionError: rows [4, 5] of 10 carry no routing data for a 6-token prompt
=========== 1 failed, 1 deselected in 153.45s (0:02:33) ===========