[v0.25.1rc][BugFix][MoE] Fix MegaMoe prefill buffer sizing - #14519
QwertyJack wants to merge 4 commits into
Conversation
MegaMoe allocates a symmetric communication buffer once from the MC2 token capacity. When enable_prefill_mc2 is disabled, that capacity was derived from the decode graph capture size even though fused MC2 can still select MegaMoe for eager prefill. Larger prefill batches could exceed the configured buffer and produce invalid outputs. Size capacity from max_num_batched_tokens whenever A3 MegaMoe is actually eligible, while retaining existing graph sizing for dispatch_ffn_combine and other MC2 paths. Centralize that eligibility in use_cann_megamoe so communication selection, buffer sizing, DP token uniformity, and W4A8/W8A8 weight preparation cannot disagree. Reject inputs exceeding the allocated MegaMoe buffer while preserving the uneven-token optimization for dispatch_ffn_combine. Signed-off-by: Zheng Shoujian <zheng.shoujian@outlook.com> (cherry picked from commit 8a173ed)
The MegaMoe config helper now validates moe_intermediate_size, but the forward-context test fixture did not provide it. Use a supported value so the existing W4A8 and W8A8 eligibility cases continue to exercise their intended quantization checks. Signed-off-by: Zheng Shoujian <zheng.shoujian@outlook.com> (cherry picked from commit 8b2fbbd)
Signed-off-by: QwertyJack <7554089+QwertyJack@users.noreply.github.com>
github-actions
Bot
added
module:tests
module:ops
module:core
module:quantization
labels
|
👋 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! |
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 bug in the CANN MegaMoe implementation where buffer capacity was incorrectly derived from decode graph capture sizes rather than the actual prefill batch requirements. By centralizing the eligibility logic and correctly sizing buffers based on max_num_batched_tokens, the fix prevents potential output corruption and provides explicit error feedback when capacity limits are exceeded. 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][Feature] Integrate CANN MegaMoe constraints and use_cann_megamoe helperSuggested PR Summary:
### What this PR does / why we need it?
This PR refactors the CANN MegaMoe integration by introducing a centralized `use_cann_megamoe` helper function. It adds validation for `moe_intermediate_size` to comply with CANN 9.1.0 tiling constraints (must be between 1024 and 3072, and a multiple of 512). It also updates the quantization paths (W4A8 and W8A8) and utility functions to use this helper, and adds a safeguard in `_apply_cann_mega_moe` to prevent buffer overflow by raising a `ValueError` if token counts exceed the allocated limit.
Feedback: One review comment suggests using `getattr` when accessing `max_num_tokens_per_rank` on `self.token_dispatcher` to prevent potential `AttributeError` or `TypeError` if the attribute is undefined or `None`.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Tested with the newly added unit test `test_set_mc2_tokens_capacity_megamoe_uses_max_num_batched_tokens` in `tests/ut/test_ascend_forward_context.py`.| num_max_tokens = self.token_dispatcher.max_num_tokens_per_rank | ||
| if num_tokens > num_max_tokens: |
There was a problem hiding this comment.
Accessing self.token_dispatcher.max_num_tokens_per_rank directly can raise an AttributeError if the attribute is not defined on the dispatcher (as defensively handled on line 331). Additionally, if num_max_tokens is None or not set, comparing it directly with num_tokens will raise a TypeError. We should use getattr with a default of None and ensure it is not None before performing the comparison.
| num_max_tokens = self.token_dispatcher.max_num_tokens_per_rank | |
| if num_tokens > num_max_tokens: | |
| num_max_tokens = getattr(self.token_dispatcher, "max_num_tokens_per_rank", None) | |
| if num_max_tokens is not None and num_tokens > num_max_tokens: |
2 participants
What this PR does / why we need it?
CANN 9.1 makes
cann_ops_transformeravailable, so the v0.25 fused-MC2 path can select CANN MegaMoe instead of the legacydispatch_ffn_combineimplementation. Whenenable_prefill_mc2is disabled, MegaMoe was still selected for eager prefill, but its symmetric buffer capacity was derived from the decode graph capture size. A larger prefill batch could therefore exceed the allocated capacity and produce invalid output instead of failing explicitly.This PR backports the fix from #14358 to
releases/v0.25.1rc:max_num_batched_tokenswhenever MegaMoe is eligible.dispatch_ffn_combinefallback.dispatch_ffn_combine.Refs #14273.
Related #14358 and #14439.
Does this PR introduce any user-facing change?
Yes. Eligible MegaMoe configurations now allocate enough symmetric-buffer capacity for eager prefill batches when prefill MC2 is disabled. Non-MegaMoe configurations retain their existing communication path, and an invalid capacity is reported explicitly instead of causing silent corruption.
How was this patch tested?
git diff --check upstream/releases/v0.25.1rc...HEADpython3 -m py_compile vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.py tests/ut/test_ascend_forward_context.pyruff check tests/ut/test_ascend_forward_context.py vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.pypython3 -m pytest tests/ut/test_ascend_forward_context.py -q(26 passed)The underlying capacity fix was also validated on the affected v0.25 stack with a two-node DeepSeek V4 W8A8 PD deployment. The original path reproduced invalid target logits and EngineDead after a 477-token prefill exceeded the 32-token-per-rank MegaMoe buffer. With capacity derived from
max_num_batched_tokens, a short request, the 477-token request, and three concurrency-16 rounds completed successfully (48/48HTTP 200 with valid JSON and no NaN, token-0, all--1,507035, traceback, or EngineDead signature). A separate 32-rank HCCL all-reduce smoke test also returned the expected value on every rank.The final backport and focused regression test were verified with the unit and static checks above; the full two-node run was not repeated after packaging the commits for this PR.