Repository navigation
Conversation
Align persistent global and gathered PCP slot buffers and the padding scalar with the normal block-table dtype during initialization. Add gather and dummy buffer regression coverage. Validation: Python syntax and git diff --check passed. Runtime tests and format.sh ci remain unverified because PyTorch, pytest and pre-commit are unavailable. AI assistance: Codex. Signed-off-by: recky-c <252898884+recky-c@users.noreply.github.com>
AI assistance: Codex. Comment-only change; git diff --check passed. format.sh ci is blocked by missing pre-commit. Signed-off-by: recky-c <252898884+recky-c@users.noreply.github.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 addresses a critical dtype mismatch between MRV2 PCP output buffers and Ascend block-table slots. By ensuring the PCP buffers inherit the dtype from the block tables during initialization, the change prevents incorrect cache-write metadata generation on PD producers. The implementation preserves existing persistent buffer behavior required for graph capture while ensuring consistency across the operator pipeline. 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! |
![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][BugFix] Align PCP slot buffer dtypes with block tablesSuggested PR Summary:
### What this PR does / why we need it?
This PR aligns the dtypes of PCP-local output buffers (`_global_batch_slot_mappings`, `_gathered_kv_slot_mappings`, and `_pad_slot_id`) with the Ascend block-table slot mappings dtype to prevent potential issues with cache-write operators. Additionally, feedback was provided to defensively check if `_pad_slot_id` is not None and avoid redundant allocations if the dtypes already match.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
A unit test `test_pcp_slot_buffers_match_block_tables` was added in `tests/ut/worker/test_pcp_manager_v2.py` to verify that the buffers are correctly initialized with matching dtypes.| if block_tables is not None: | ||
| slot_dtype = block_tables.slot_mappings.dtype | ||
| if self._global_batch_slot_mappings is not None: | ||
| self._global_batch_slot_mappings = torch.empty_like( | ||
| self._global_batch_slot_mappings, dtype=slot_dtype | ||
| ) | ||
| if self._gathered_kv_slot_mappings is not None: | ||
| self._gathered_kv_slot_mappings = torch.empty_like( | ||
| self._gathered_kv_slot_mappings, dtype=slot_dtype | ||
| ) | ||
| self._pad_slot_id = self._pad_slot_id.to(slot_dtype) |
There was a problem hiding this comment.
To adhere to defensive programming practices and optimize memory usage, we should:
- Ensure
self._pad_slot_idis notNonebefore calling.to(slot_dtype)to prevent potential runtime errors. - Avoid redundant tensor allocations by checking if the existing dtypes of
self._global_batch_slot_mappingsandself._gathered_kv_slot_mappingsalready matchslot_dtypebefore callingtorch.empty_like.
| if block_tables is not None: | |
| slot_dtype = block_tables.slot_mappings.dtype | |
| if self._global_batch_slot_mappings is not None: | |
| self._global_batch_slot_mappings = torch.empty_like( | |
| self._global_batch_slot_mappings, dtype=slot_dtype | |
| ) | |
| if self._gathered_kv_slot_mappings is not None: | |
| self._gathered_kv_slot_mappings = torch.empty_like( | |
| self._gathered_kv_slot_mappings, dtype=slot_dtype | |
| ) | |
| self._pad_slot_id = self._pad_slot_id.to(slot_dtype) | |
| if block_tables is not None: | |
| slot_dtype = block_tables.slot_mappings.dtype | |
| if self._global_batch_slot_mappings is not None and self._global_batch_slot_mappings.dtype != slot_dtype: | |
| self._global_batch_slot_mappings = torch.empty_like( | |
| self._global_batch_slot_mappings, dtype=slot_dtype | |
| ) | |
| if self._gathered_kv_slot_mappings is not None and self._gathered_kv_slot_mappings.dtype != slot_dtype: | |
| self._gathered_kv_slot_mappings = torch.empty_like( | |
| self._gathered_kv_slot_mappings, dtype=slot_dtype | |
| ) | |
| if self._pad_slot_id is not None: | |
| self._pad_slot_id = self._pad_slot_id.to(slot_dtype) |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
1 participant
What this PR does / why we need it?
MRV2 PCP supplies its own output buffers to compute_slot_mappings, bypassing the normal Ascend int32 slot buffer. The upstream PCP buffers are int64, while the C8 store_kv_block_metadata kernel reads slots as int32, which can generate incorrect cache-write metadata on PD producers.
Initialize the Ascend PCP global and gathered slot buffers with block_tables.slot_mappings.dtype, and align the padding scalar. Allocation happens once before graph capture; normal computation, dummy inputs and graph padding continue to reuse the persistent buffers. No upstream vLLM, attention call sites or native operators are changed.
Related: #16664 already fixes the dtype issue at the operator call site. This draft proposes an alternative fix at the PCP buffer source for review, rather than an independent issue. It does not include that PR's separate PCP+DCP SFA slot-length fix. Maintainers should coordinate the overlapping dtype changes.
Does this PR introduce any user-facing change?
Intended to fix incorrect C8 cache writes caused by PCP slot dtype mismatches. No configuration changes. This follows the existing Ascend block-table dtype and capacity assumptions; it does not add range checks or claim all slot paths in the repository use int32.
How was this patch tested?
Implementation and PR preparation used AI assistance (Codex).