[Revert] Revert "[Feature][MRV1][MRV2] Refactor Host-Side Parameter Updates for ACL Graph Replay." (#16425) - #16726
Conversation
…r ACL Graph Replay. (vllm-project#16425)" This reverts commit 484a5f9.
|
👋 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 [Feature] 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 reverts the changes introduced in PR #16425, which implemented an UpdatableGraph design for host-side parameter updates during ACL graph replay. The revert restores the previous graph update logic across the attention and compilation modules, ensuring the codebase returns to its prior stable state while removing the newly introduced abstraction. 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:\n\nmarkdown\n[Attention][Misc] Refactor graph capture and replay to use standard NPUGraph\n\n\nSuggested PR Summary:\n\nmarkdown\n### What this PR does / why we need it?\nThis pull request refactors the graph capture and replay mechanism by removing the custom `UpdatableGraph` and replacing it with standard `torch.npu.NPUGraph` and manual graph parameter updates. This simplifies the compilation stack and aligns it with standard PyTorch NPU graph APIs.\n\nFeedback:\n- In `vllm_ascend/utils.py`, `weak_ref_tensors` raises a `ValueError` if the input is `None`. Since `block_table` can be `None` in standard non-paged FIA attention, calling `weak_ref_tensors(block_table)` will cause a runtime crash. It should safely return `None` when `tensors` is `None`.\n- In `vllend_ascend/attention/utils.py`, `update_paged_attention_graph_param` should accept an optional `workspace` parameter and reuse it to avoid redundant per-layer workspace allocations.\n\n### Does this PR introduce _any_ user-facing change?\nNo.\n\n### How was this patch tested?\nCI passed with updated unit tests.\n
| return tensors | ||
| ret = IntermediateTensors({key: weak_ref_tensor(val) for key, val in tensors.tensors.items()}) | ||
| return ret | ||
| raise ValueError("Invalid type for tensors") |
There was a problem hiding this comment.
The function weak_ref_tensors raises a ValueError if the input tensors is None. However, in vllm_ascend/attention/attention_v1.py, weak_ref_tensors(block_table) is called where block_table can be None (e.g., in standard non-paged FIA attention). This will cause a runtime crash during graph capture. We should safely return None when tensors is None to prevent this crash.
| raise ValueError("Invalid type for tensors") | |
| if tensors is None: | |
| return None | |
| raise ValueError("Invalid type for tensors") |
| def update_paged_attention_graph_param( | ||
| update_stream, | ||
| handle, | ||
| event, | ||
| param: PagedAttentionGraphParam, | ||
| block_table: torch.Tensor, | ||
| seq_lens: torch.Tensor, | ||
| ) -> None: |
There was a problem hiding this comment.
The update_paged_attention_graph_param function currently calls torch_npu._npu_paged_attention_get_workspace on every invocation. When called in a loop over multiple layers during graph updates, this leads to redundant per-layer workspace allocations (each backed by a ~36-54MB buffer), causing significant memory churn and performance overhead. Adding an optional workspace parameter allows reusing a single cached workspace across layers, matching the optimization used in the main paged attention update path.
def update_paged_attention_graph_param(
update_stream,
handle,
event,
param: PagedAttentionGraphParam,
block_table: torch.Tensor,
seq_lens: torch.Tensor,
workspace: torch.Tensor | None = None,
) -> None:| workspace = torch_npu._npu_paged_attention_get_workspace( | ||
| query=query, | ||
| key_cache=key_cache, | ||
| value_cache=value_cache, | ||
| num_kv_heads=num_kv_heads, | ||
| num_heads=num_heads, | ||
| scale_value=scale, | ||
| block_table=block_table, | ||
| context_lens=seq_lens, | ||
| out=output, | ||
| ) |
There was a problem hiding this comment.
Use the passed-in workspace if provided, instead of always calling _npu_paged_attention_get_workspace to allocate a new one.
| workspace = torch_npu._npu_paged_attention_get_workspace( | |
| query=query, | |
| key_cache=key_cache, | |
| value_cache=value_cache, | |
| num_kv_heads=num_kv_heads, | |
| num_heads=num_heads, | |
| scale_value=scale, | |
| block_table=block_table, | |
| context_lens=seq_lens, | |
| out=output, | |
| ) | |
| if workspace is None: | |
| workspace = torch_npu._npu_paged_attention_get_workspace( | |
| query=query, | |
| key_cache=key_cache, | |
| value_cache=value_cache, | |
| num_kv_heads=num_kv_heads, | |
| num_heads=num_heads, | |
| scale_value=scale, | |
| block_table=block_table, | |
| context_lens=seq_lens, | |
| out=output, | |
| ) |
Qwen3_5ForConditionalGeneration hybrid VL hits MRv2 encoder graph capture (CUDA stream assert) and hybrid KV copy. Keep it on V1 by default. Pin dflash2 PIECEWISE acceptance to V1 after vllm-project#16726 broke dummy propose without set_forward_context; V2 eager remains covered. Signed-off-by: yjyang62 <yjyang62@users.noreply.github.com> Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
### What this PR does / why we need it? On vLLM 0.28.0, Ascend MRV2 PCP+DP is rejected during parallel configuration validation. Allowing configuration alone is insufficient: dispatch still synchronizes the global token count before PCP partitioning, so a 44-token prefill with PCP2 reaches `DPMetadata.make` with 44 recorded tokens but a 22-token local batch. - Scope the configuration workaround to vLLM 0.28.0, Ascend MRV2 explicitly enabled, DP > 1, PCP > 1, and DCP = 1. Preserve other validation and restore the real PCP size on every exit. - Reuse the existing rank-segment rules to compute the PCP execution count before dispatch, including replicated decode tokens and uneven prefill partitions. - Pass only the computed count through a call-scoped context to the original dispatch function. Preserve global request state, original DP synchronization, and its consistency checks. Based on upstream main `628fac6d8`. The linear commit preserves #16832's revert of default MRV2 whitelist selection; MRV2 remains explicitly enabled through the environment. The graph replay revert is already upstream in #16726 and is not included. Existing upstream MRV2 selection and KV-cache/preemption changes are retained. ### Does this PR introduce _any_ user-facing change? Enables the scoped PCP+DP configuration and fixes its dispatch token-count mismatch on vLLM 0.28.0. No new user-facing flags. ### How was this patch tested? - On revision `fb0554b7f`: the PR was linearized and successfully rebased onto the refreshed upstream main using the same operation as CI. The PR diff is byte-identical to the resolved merge diff. Ruff lint/format and AST syntax checks passed for all seven PR Python files; `git diff upstream/main --check` passed. The local `format.sh ci` attempt stopped because `pre-commit` is unavailable. Full checks are being rerun by GitHub CI. The previous run passed pre-commit but stopped during CI rebase, before unit tests started. - Before this merge: Ruff lint, Ruff format check, and `git diff --check` passed. The six transplanted adaptation functions and the added unit test/fixture definitions were also compared structurally against the tested local implementation. - On the prior local base (`bddbab4c3` plus the adaptation): 121 unit tests passed, 2 skipped in the Ascend development container, covering PCP partition counts, zero counts, dummy/no-PCP paths, version gating, context isolation/restoration, and configuration validation. The corresponding unit regression tests are included in this PR; the earlier E2E test expansion is excluded. - Earlier service smoke on that local line: MiniMax + Eagle3, TP2/DP2/PCP2, seq16, FlashComm1 and FULL_DECODE_ONLY completed the original 44-token request and four different-length concurrent requests. This preceded the count-only context simplification. The exact main-based PR revision has not yet been rerun in the remote runtime. Import/override compatibility, single-rank and multi-rank model smoke remain pending on this revision. Full accuracy, MTP/MLA coverage and performance regression validation are not claimed; keep this PR in draft. - vLLM main: vllm-project/vllm@84030bb Signed-off-by: wzx0726 <278573478+wzx0726@users.noreply.github.com> Co-authored-by: wzx0726 <278573478+wzx0726@users.noreply.github.com>
2 participants
Revert of PR #16425 (merged onto
main).Original PR: #16425
Original author: @zhiyu-wa
Merge commit:
484a5f9aa7de9e1a51843ab577bc3c3ee401353fWhat this PR does / why we need it?
Implements the
UpdatableGraphdesign discussed in #13058.Refactors the
FIAandSpeculative Decodingto work with the newUpdatableGraphdesign.Manually validated on A3 with the following scenarios:
Manually validated on A5 with the following scenarios:
Does this PR introduce any user-facing change?
No.
How was this patch tested?
Temporarily validated on some model.