[CI]test: restore ACL graph, MoE reduction, and Gumbel sampling - #16862
U1stRsouland wants to merge 1 commit into
Conversation
Reapply the ACL graph implementation reverted by vllm-project#16726, restore the pre-vllm-project#16550 MoE reduction path, and restore the pre-vllm-project#16269 MRV2 Gumbel sampler to isolate Qwen3-235B A3 three-node PD accuracy regression. Signed-off-by: U1stRsouland <19800362117@163.com>
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 focuses on reverting and restoring specific components related to ACL graph management, MoE reduction, and Gumbel sampling. The primary goal is to provide a stable baseline to investigate and resolve accuracy regressions observed in Qwen3-235B A3 three-node PD configurations. 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. |
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Attention][Feature] Introduce UpdatableGraph for dynamic attention metadata updatesSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces `UpdatableGraph` to replace the legacy manual graph parameter update logic for attention metadata on Ascend NPU. It refactors the attention backend (`AscendAttentionBackendImpl`) to register tasks and resources with `UpdatableGraph`, simplifying the codebase and removing redundant workspace allocations. Additionally, it refactors the Gumbel sampling path to use `gumbel_sample` directly, removes the custom `maybe_all_reduce` operators in favor of standard TP collectives, and updates the `weak_ref_tensors` utility to recursively handle nested containers.
Feedback:
- In `vllm_ascend/compilation/updatable_graph.py`, there is a redundant call to `factory()` when initializing a new capture resource with `use_max_workspace=True`. Using `elif use_max_workspace:` instead of `if use_max_workspace:` avoids this unnecessary allocation.
- In `vllm_ascend/worker/v2/spec_decode/autoregressive/speculator.py`, using `as_strided` to expand `block_table` is unsafe and can lead to out-of-bounds memory access. It is safer to pad the tensor using `torch.nn.functional.pad`.
### Does this PR introduce _any_ user-facing change?
No, these are internal refactorings and performance optimizations for Ascend NPU.
### How was this patch tested?
The changes were tested by updating existing unit tests and end-to-end tests (e.g., `test_gumbel_sampling.py`, `test_acl_graph.py`, `test_fused_moe.py`) to align with the new `UpdatableGraph` and refactored MoE reduction contracts.| if use_max_workspace: | ||
| # Some models mix attention layer shapes under the same graph size. | ||
| # During capture, keep the largest required workspace for that size. | ||
| candidate_workspace = factory() | ||
| if ( | ||
| candidate_workspace.numel() * candidate_workspace.element_size() | ||
| > self.capture_resources[key].numel() * self.capture_resources[key].element_size() | ||
| ): | ||
| self.capture_resources[key] = candidate_workspace |
There was a problem hiding this comment.
When key is not in self.capture_resources, it is initialized by calling factory(). If use_max_workspace is True, the code immediately calls factory() again to get a candidate_workspace and compares it to the newly created resource. This second call is completely redundant when the resource is first initialized, resulting in an unnecessary memory allocation on the NPU.\n\nUsing elif use_max_workspace: instead of if use_max_workspace: avoids this redundant allocation.
elif use_max_workspace:\n # Some models mix attention layer shapes under the same graph size.\n # During capture, keep the largest required workspace for that size.\n candidate_workspace = factory()\n if (\n candidate_workspace.numel() * candidate_workspace.element_size()\n > self.capture_resources[key].numel() * self.capture_resources[key].element_size()\n ):\n self.capture_resources[key] = candidate_workspace| if block_table is not None and block_table.shape[0] < num_reqs_padded: | ||
| block_table = block_table.as_strided((num_reqs_padded, block_table.shape[1]), block_table.stride()) |
There was a problem hiding this comment.
Using as_strided to expand the block_table to num_reqs_padded rows when block_table.shape[0] < num_reqs_padded is unsafe. Since as_strided does not allocate new storage, accessing rows beyond block_table.shape[0] with a non-zero stride will result in out-of-bounds memory access of the underlying storage. This can lead to silent data corruption, NPU driver errors, or segmentation faults.\n\nA safer approach is to use torch.nn.functional.pad to pad the tensor along the first dimension with zeros, which allocates safe, properly bounded memory.
if block_table is not None and block_table.shape[0] < num_reqs_padded:\n pad_size = num_reqs_padded - block_table.shape[0]\n block_table = torch.nn.functional.pad(block_table, (0, 0, 0, pad_size))|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
Reapply the ACL graph implementation reverted by #16726, restore the pre-#16550 MoE reduction path, and restore the pre-#16269 MRV2 Gumbel sampler to isolate Qwen3-235B A3 three-node PD accuracy regression.
What this PR does / why we need it?
Does this PR introduce any user-facing change?
How was this patch tested?