[BugFix][MoE] Fix MegaMoe prefill buffer sizing - #14439
weijinqian0 merged 7 commits into
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 bug in CANN MegaMoe buffer sizing where prefill batches could exceed allocated capacity when 'enable_prefill_mc2' was disabled. By centralizing the eligibility logic and ensuring the buffer is sized based on actual batched token limits, the changes ensure stability for eager prefill scenarios. The update also includes defensive programming to catch and report capacity violations explicitly. 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
|
|
👋 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! |
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][Feature] Refactor CANN MegaMoe activation and add token limit validationSuggested PR Summary:
### What this PR does / why we need it?
This PR refactors the CANN MegaMoe activation logic by introducing a centralized `use_cann_megamoe` helper function across various modules, replacing direct checks on `_MEGA_MOE_SUPPORTED`. It also adds a safety check in `_apply_cann_mega_moe` to raise a `ValueError` if the number of tokens per rank exceeds the allocated symmetric buffer capacity.
### Does this PR introduce _any_ user-facing change?
No, this is an internal refactoring and validation improvement.
### How was this patch tested?
CI/CD tests.Feedback Summary:
One review comment was kept which suggests using getattr to safely access max_num_tokens_per_rank on self.token_dispatcher to prevent potential AttributeErrors in mock or testing environments.
| num_tokens = fused_experts_input.hidden_states.shape[0] | ||
| num_max_tokens = self.token_dispatcher.max_num_tokens_per_rank | ||
| if num_tokens > num_max_tokens: | ||
| raise ValueError( | ||
| f"MegaMoe received {num_tokens} tokens per rank, but its symmetric buffer " | ||
| f"was allocated for at most {num_max_tokens}. Increase max_num_batched_tokens " | ||
| "or disable fused MC2." | ||
| ) |
There was a problem hiding this comment.
Using direct attribute access self.token_dispatcher.max_num_tokens_per_rank can raise an AttributeError if the attribute is not set or initialized on the dispatcher (e.g., in certain testing or mock environments). To be robust and consistent with how this attribute is accessed elsewhere in this file (such as using getattr with a default value in _init_mega_moe_symm_buffer), it is safer to use getattr here as well.
| num_tokens = fused_experts_input.hidden_states.shape[0] | |
| num_max_tokens = self.token_dispatcher.max_num_tokens_per_rank | |
| if num_tokens > num_max_tokens: | |
| raise ValueError( | |
| f"MegaMoe received {num_tokens} tokens per rank, but its symmetric buffer " | |
| f"was allocated for at most {num_max_tokens}. Increase max_num_batched_tokens " | |
| "or disable fused MC2." | |
| ) | |
| num_tokens = fused_experts_input.hidden_states.shape[0] | |
| 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: | |
| raise ValueError( | |
| f"MegaMoe received {num_tokens} tokens per rank, but its symmetric buffer " | |
| f"was allocated for at most {num_max_tokens}. Increase max_num_batched_tokens " | |
| "or disable fused MC2." | |
| ) |
d81fd73 to
42279c5
Compare
f77f1ea to
588a676
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
588a676 to
2b2e768
Compare
|
/rerun [Bot]: rerun completed. Rerun (failed jobs only):
|
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
b641690 to
54b0f4d
Compare
|
/rerun [Bot]: rerun completed. Rerun (failed jobs only):
|
|
/rerun Rerun (failed jobs only):
|
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
54b0f4d to
a68edb1
Compare
| @patch("vllm_ascend.quantization.methods.w4a8.get_current_vllm_config", new=lambda: None) | ||
| @patch("vllm_ascend.quantization.methods.w4a8.use_cann_megamoe", new=lambda _: False) | ||
| @patch("vllm_ascend.quantization.methods.w4a8.get_ascend_config") | ||
| @patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz") |
There was a problem hiding this comment.
it should be w4a8.w4a8 rather than w4a8
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
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 BF16/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>
Mock the centralized MegaMoe selector together with current config evaluation in weight-loading tests. Extend forward-context stubs with use_mega_moe so the existing non-MegaMoe assertions remain isolated from the new shared state. Signed-off-by: Zheng Shoujian <zheng.shoujian@outlook.com>
800f92a to
e72d2c3
Compare
…are profiles (#15478) ## What - add the final MoE/compilation/sampling hardware capabilities and a `MoECommPolicy` - migrate MoE communication selection, graph fusion gates, routing, shared experts, token dispatch, SwiGLU-OAI MX quant, and top-k/top-p sampling away from direct device identity checks - preserve the latest upstream Situ, hierarchical communication, SFA, DSV4, sampler, and CANN MegaMoe paths - model the new A3-only CANN MegaMoe consumer with `HardwareCapability.CANN_MEGAMOE` and update the exact capability matrix ## Why Shared business logic should consume hardware capabilities and implementation policies instead of branching on Ascend device identities. This keeps hardware detection/profile registration as the single identity boundary and makes the final consumers explicit. This is a refactor only; there is no user-visible behavior change. ## Series Device HAL / Hardware Profile refactor 7/N. Depends on merged #15407 (`d4ebe8a0e0fdf56135507259f90f80a45f736cab`) and is rebased onto current `main` (`30f54b5c341f015bcc124022e62156d3e5ef5919`). This base includes the GLM5.3 documentation-link fix from #15508 and the MegaMoe prefill update from #14439. ## Testing - `ruff check` on all 14 changed Python files - `ruff format --check` on all 14 changed Python files - `python3 -m compileall -q` on all 14 changed Python files - `git diff --check upstream/main..HEAD` - production-code scan for newly introduced direct device-identity branches - E2E run https://github.com/vllm-project/vllm-ascend/actions/runs/33579196584: 23 successful jobs, 4 skipped, zero failures; CPU, all selected NPU jobs, 310P, upstream, pre-commit, DCO, Read the Docs, and `ci-gate` succeeded The local static checks above passed. The development server does not provide the complete pytest/mypy/vLLM environment, so targeted unit tests and Python 3.10/3.11/3.12 mypy comparison are covered by the successful CI run and are not claimed as local results. --------- Signed-off-by: frost_mourne <2906339855@qq.com>
### What this PR does / why we need it? CANN MegaMoe allocates its 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 eager prefill could still select MegaMoe through fused MC2. A larger prefill batch could therefore exceed the allocated capacity and produce invalid outputs. This PR: - Sizes the MegaMoe buffer from `max_num_batched_tokens` when MegaMoe is actually eligible. - Centralizes MegaMoe eligibility in `use_cann_megamoe` so communication selection, buffer sizing, DP token handling, and BF16/W4A8/W8A8 weight preparation use the same decision, while preserving the `dispatch_ffn_combine` fallback. - Keeps DP token counts uniform for MegaMoe, without disabling the uneven-token optimization for `dispatch_ffn_combine`. - Raises an explicit error if a MegaMoe input exceeds the allocated per-rank capacity. Fixes vllm-project#14273. Related release-branch PR: vllm-project#14358. ### Does this PR introduce _any_ user-facing change? Yes. MegaMoe now allocates enough symmetric-buffer capacity for eager prefill batches when prefill MC2 is disabled. Non-MegaMoe configurations continue to use their existing communication path, and an invalid capacity is reported explicitly instead of producing incorrect output. ### How was this patch tested? - `git diff --check` - `python3 -m py_compile vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/ops/fused_moe/routed_experts.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.py` - `ruff check vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/ops/fused_moe/routed_experts.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.py` - The underlying buffer-sizing fix was validated end-to-end with DeepSeek V4 W8A8 inference on Ascend A3 using eager prefill, fused MC2 enabled, and prefill MC2 disabled. A prefill batch larger than the decode graph capture capacity completed without non-finite output and produced the expected response. This final main-branch adaptation was verified with the static checks above and was not rerun end-to-end. No unit test was added because the failure depends on CANN MegaMoe symmetric-buffer allocation and collective execution across an NPU communication group. - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Zheng Shoujian <zheng.shoujian@outlook.com>
…are profiles (vllm-project#15478) ## What - add the final MoE/compilation/sampling hardware capabilities and a `MoECommPolicy` - migrate MoE communication selection, graph fusion gates, routing, shared experts, token dispatch, SwiGLU-OAI MX quant, and top-k/top-p sampling away from direct device identity checks - preserve the latest upstream Situ, hierarchical communication, SFA, DSV4, sampler, and CANN MegaMoe paths - model the new A3-only CANN MegaMoe consumer with `HardwareCapability.CANN_MEGAMOE` and update the exact capability matrix ## Why Shared business logic should consume hardware capabilities and implementation policies instead of branching on Ascend device identities. This keeps hardware detection/profile registration as the single identity boundary and makes the final consumers explicit. This is a refactor only; there is no user-visible behavior change. ## Series Device HAL / Hardware Profile refactor 7/N. Depends on merged vllm-project#15407 (`d4ebe8a0e0fdf56135507259f90f80a45f736cab`) and is rebased onto current `main` (`30f54b5c341f015bcc124022e62156d3e5ef5919`). This base includes the GLM5.3 documentation-link fix from vllm-project#15508 and the MegaMoe prefill update from vllm-project#14439. ## Testing - `ruff check` on all 14 changed Python files - `ruff format --check` on all 14 changed Python files - `python3 -m compileall -q` on all 14 changed Python files - `git diff --check upstream/main..HEAD` - production-code scan for newly introduced direct device-identity branches - E2E run https://github.com/vllm-project/vllm-ascend/actions/runs/33579196584: 23 successful jobs, 4 skipped, zero failures; CPU, all selected NPU jobs, 310P, upstream, pre-commit, DCO, Read the Docs, and `ci-gate` succeeded The local static checks above passed. The development server does not provide the complete pytest/mypy/vLLM environment, so targeted unit tests and Python 3.10/3.11/3.12 mypy comparison are covered by the successful CI run and are not claimed as local results. --------- Signed-off-by: frost_mourne <2906339855@qq.com>
### What this PR does / why we need it? CANN MegaMoe allocates its 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 eager prefill could still select MegaMoe through fused MC2. A larger prefill batch could therefore exceed the allocated capacity and produce invalid outputs. This PR: - Sizes the MegaMoe buffer from `max_num_batched_tokens` when MegaMoe is actually eligible. - Centralizes MegaMoe eligibility in `use_cann_megamoe` so communication selection, buffer sizing, DP token handling, and BF16/W4A8/W8A8 weight preparation use the same decision, while preserving the `dispatch_ffn_combine` fallback. - Keeps DP token counts uniform for MegaMoe, without disabling the uneven-token optimization for `dispatch_ffn_combine`. - Raises an explicit error if a MegaMoe input exceeds the allocated per-rank capacity. Fixes vllm-project#14273. Related release-branch PR: vllm-project#14358. ### Does this PR introduce _any_ user-facing change? Yes. MegaMoe now allocates enough symmetric-buffer capacity for eager prefill batches when prefill MC2 is disabled. Non-MegaMoe configurations continue to use their existing communication path, and an invalid capacity is reported explicitly instead of producing incorrect output. ### How was this patch tested? - `git diff --check` - `python3 -m py_compile vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/ops/fused_moe/routed_experts.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.py` - `ruff check vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/ops/fused_moe/routed_experts.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.py` - The underlying buffer-sizing fix was validated end-to-end with DeepSeek V4 W8A8 inference on Ascend A3 using eager prefill, fused MC2 enabled, and prefill MC2 disabled. A prefill batch larger than the decode graph capture capacity completed without non-finite output and produced the expected response. This final main-branch adaptation was verified with the static checks above and was not rerun end-to-end. No unit test was added because the failure depends on CANN MegaMoe symmetric-buffer allocation and collective execution across an NPU communication group. - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Zheng Shoujian <zheng.shoujian@outlook.com>
…are profiles (vllm-project#15478) ## What - add the final MoE/compilation/sampling hardware capabilities and a `MoECommPolicy` - migrate MoE communication selection, graph fusion gates, routing, shared experts, token dispatch, SwiGLU-OAI MX quant, and top-k/top-p sampling away from direct device identity checks - preserve the latest upstream Situ, hierarchical communication, SFA, DSV4, sampler, and CANN MegaMoe paths - model the new A3-only CANN MegaMoe consumer with `HardwareCapability.CANN_MEGAMOE` and update the exact capability matrix ## Why Shared business logic should consume hardware capabilities and implementation policies instead of branching on Ascend device identities. This keeps hardware detection/profile registration as the single identity boundary and makes the final consumers explicit. This is a refactor only; there is no user-visible behavior change. ## Series Device HAL / Hardware Profile refactor 7/N. Depends on merged vllm-project#15407 (`d4ebe8a0e0fdf56135507259f90f80a45f736cab`) and is rebased onto current `main` (`30f54b5c341f015bcc124022e62156d3e5ef5919`). This base includes the GLM5.3 documentation-link fix from vllm-project#15508 and the MegaMoe prefill update from vllm-project#14439. ## Testing - `ruff check` on all 14 changed Python files - `ruff format --check` on all 14 changed Python files - `python3 -m compileall -q` on all 14 changed Python files - `git diff --check upstream/main..HEAD` - production-code scan for newly introduced direct device-identity branches - E2E run https://github.com/vllm-project/vllm-ascend/actions/runs/33579196584: 23 successful jobs, 4 skipped, zero failures; CPU, all selected NPU jobs, 310P, upstream, pre-commit, DCO, Read the Docs, and `ci-gate` succeeded The local static checks above passed. The development server does not provide the complete pytest/mypy/vLLM environment, so targeted unit tests and Python 3.10/3.11/3.12 mypy comparison are covered by the successful CI run and are not claimed as local results. --------- Signed-off-by: frost_mourne <2906339855@qq.com>
What this PR does / why we need it?
CANN MegaMoe allocates its symmetric communication buffer once from the MC2 token capacity. When
enable_prefill_mc2is disabled, that capacity was derived from the decode graph capture size even though eager prefill could still select MegaMoe through fused MC2. A larger prefill batch could therefore exceed the allocated capacity and produce invalid outputs.This PR:
max_num_batched_tokenswhen MegaMoe is actually eligible.use_cann_megamoeso communication selection, buffer sizing, DP token handling, and BF16/W4A8/W8A8 weight preparation use the same decision, while preserving thedispatch_ffn_combinefallback.dispatch_ffn_combine.Fixes #14273.
Related release-branch PR: #14358.
Does this PR introduce any user-facing change?
Yes. MegaMoe now allocates enough symmetric-buffer capacity for eager prefill batches when prefill MC2 is disabled. Non-MegaMoe configurations continue to use their existing communication path, and an invalid capacity is reported explicitly instead of producing incorrect output.
How was this patch tested?
git diff --checkpython3 -m py_compile vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/ops/fused_moe/routed_experts.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.pyruff check vllm_ascend/ascend_forward_context.py vllm_ascend/ops/fused_moe/moe_comm_method.py vllm_ascend/ops/fused_moe/routed_experts.py vllm_ascend/quantization/methods/w4a8.py vllm_ascend/quantization/methods/w8a8_dynamic.py vllm_ascend/utils.pyNo unit test was added because the failure depends on CANN MegaMoe symmetric-buffer allocation and collective execution across an NPU communication group.