[BugFix] Fix SP/PCP token layout in MoE AllGather - #16477
Conversation
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 critical shape mismatch issue occurring in MoE AllGather operations when sequence parallelism and PCP are simultaneously enabled. By refining how token metadata is processed and expanding it to align with the actual DP/PCP/TP group structure, the fix ensures consistent tensor shapes. Additionally, the changes clean up the communication path by removing unnecessary PCP collectives, streamlining the execution flow for the SP path. 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
|
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
|
👋 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! |
Signed-off-by: wzx0726 <278573478+wzx0726@users.noreply.github.com>
60bd7bc to
11a8eec
Compare
|
/rerun Rerun (failed jobs only):
|
|
@weiguihua2 PTAL |
| if len(local_sizes) != ep_group.world_size: | ||
| pcp_size = get_pcp_group().world_size | ||
| dp_size = get_dp_group().world_size | ||
| if len(local_sizes) * pcp_size != ep_group.world_size or len(local_sizes) % dp_size: |
There was a problem hiding this comment.
Should this side report an error or return None? Reporting an error is inconsistent with previous behavior.
There was a problem hiding this comment.
The path has been changed to the original None path.
Signed-off-by: wzx0726 <278573478+wzx0726@users.noreply.github.com>
|
/rerun Rerun (failed jobs only):
|
5 participants
What this PR does / why we need it?
Fix sequence-parallel MoE ALLGATHER with PCP enabled. SP length metadata describes DP x SP, while the EP group spans DP x PCP x TP. Discarding the shorter metadata selects a DP-only reshape that cannot represent the gathered tensor.
Extend the existing
_get_ep_local_sizeshelper to expand lengths in DP/PCP/TP order when the topology matches. When the layout cannot be expanded, retain the existingNonereturn and caller fallback instead of introducing a new exception. Existing gather, reduce-scatter and fake implementations share the mapping. Remove the redundant PCP collectives from the EP sequence-parallel prepare/finalize path. The non-SP path is retained. No new persistent state, flags or public interfaces are introduced.Does this PR introduce any user-facing change?
Fixes a shape error in the SP+PCP MoE ALLGATHER path and removes duplicate communication. No configuration changes.
How was this patch tested?
Added regression tests for multiple DP/PCP/SP combinations, uneven token counts, token ordering, gather/reduce-scatter round trips, fake shapes, fallback for unmatched topology and absence of redundant PCP collectives.
The latest compatibility update changes the unmatched-layout regression test to assert
None.git diff --checkand Python AST syntax checks passed for the update, and both changed-file diffs were reviewed.bash format.sh cicould not run its checks because pre-commit is not installed in the editing environment. Runtime tests have not been rerun on this exact PR head; targeted unit tests and NPU validation remain pending.