[Bugfix][KVPool] Cover physical DCP shards in scheduler lookup - #15591
[Bugfix][KVPool] Cover physical DCP shards in scheduler lookup#15591zmc1997 wants to merge 1 commit 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 the KVPool scheduler lookup logic where incomplete DCP shard coverage could lead to premature cache hit reporting. By expanding the lookup keys to cover every DCP rank alongside existing PP and TP ranks, the system now ensures stricter validation for external prefix-cache hits. Note that this change updates the serialized key format, which may impact existing KVPool entries. 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:
[kv_transfer][Misc] Update metadata key format to include colons and expand lookup keys by DCP rankSuggested PR Summary:
### What this PR does / why we need it?
This PR updates the string representation of metadata keys in the Ascend KV pool to use colons for `@pcp:` and `@dcp:` fields, ensuring consistent formatting. Additionally, it updates `_expand_lookup_keys_by_rank` in `pool_worker.py` to support expansion by `dcp_rank` alongside `pp_rank` and `head_or_tp_rank`.
Feedback:
An optimization is suggested for `_expand_lookup_keys_by_rank` to avoid redundant string replacements inside the nested loops by hoisting independent field replacements to their respective outer loops.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Tested with updated unit tests in `test_metadata.py`.| for pp_rank in range(self.pp_size): | ||
| for tp_rank in range(self.get_group_tp_size(group_id)): | ||
| for key in keys: | ||
| tp_key = self._replace_key_field(key, "head_or_tp_rank", tp_rank) | ||
| expanded.append(self._replace_key_field(tp_key, "pp_rank", pp_rank)) | ||
| for dcp_rank in range(self.dcp_size): | ||
| for head_or_tp_rank in range(num_head_or_tp_ranks): | ||
| for key in keys: | ||
| rank_key = self._replace_key_field(key, "dcp", dcp_rank) | ||
| rank_key = self._replace_key_field(rank_key, "head_or_tp_rank", head_or_tp_rank) | ||
| expanded.append(self._replace_key_field(rank_key, "pp_rank", pp_rank)) |
There was a problem hiding this comment.
Performance Optimization: Avoid Redundant String Replacements in Nested Loops
In the current implementation, _replace_key_field is called 3 times for every single key in the innermost loop. Since string manipulation (using .find() and slicing) is relatively expensive in Python, and this lookup function is executed on every scheduler step, we can significantly optimize this by hoisting the independent field replacements to their respective outer loops.
By replacing @pp_rank: in the pp_rank loop, @dcp: in the dcp_rank loop, and @head_or_tp_rank: in the head_or_tp_rank loop, we reduce the total number of string replacement operations from
| for pp_rank in range(self.pp_size): | |
| for tp_rank in range(self.get_group_tp_size(group_id)): | |
| for key in keys: | |
| tp_key = self._replace_key_field(key, "head_or_tp_rank", tp_rank) | |
| expanded.append(self._replace_key_field(tp_key, "pp_rank", pp_rank)) | |
| for dcp_rank in range(self.dcp_size): | |
| for head_or_tp_rank in range(num_head_or_tp_ranks): | |
| for key in keys: | |
| rank_key = self._replace_key_field(key, "dcp", dcp_rank) | |
| rank_key = self._replace_key_field(rank_key, "head_or_tp_rank", head_or_tp_rank) | |
| expanded.append(self._replace_key_field(rank_key, "pp_rank", pp_rank)) | |
| for pp_rank in range(self.pp_size): | |
| pp_keys = [self._replace_key_field(key, "pp_rank", pp_rank) for key in keys] | |
| for dcp_rank in range(self.dcp_size): | |
| dcp_keys = [self._replace_key_field(key, "dcp", dcp_rank) for key in pp_keys] | |
| for head_or_tp_rank in range(num_head_or_tp_ranks): | |
| for key in dcp_keys: | |
| expanded.append(self._replace_key_field(key, "head_or_tp_rank", head_or_tp_rank)) |
Signed-off-by: zmc1997 <40617288+zmc1997@users.noreply.github.com>
7244ec3 to
0774245
Compare
|
The code-side pre-commit failure has been fixed, and pre-commit/select-tests now pass. Selected tests require the repository ready label, but the PR author does not have permission to add it. Maintainers, please add ready to run the selected tests. |
What this PR does / why we need it?
Main-branch counterpart of #15506. Addresses #15505.
KVPoolWorker.lookup_scheduler()must check every rank-specific key beforereporting an external prefix-cache hit. The previous expansion covered PP and
head_or_tp_rank, but not every DCP shard, so a DCP0 hit could be reportedeven when another DCP key required by the subsequent load was missing.
This PR:
field:valueformat for PCP and DCP key fields as the otherrank fields;
[rank_shard][block]hit reduction.The all-rank lookup still assumes
PCP=1; this change does not add unvalidatedPCP expansion behavior.
Does this PR introduce any user-facing change?
Yes. DCP deployments now report an external cache hit only when every required
DCP shard exists.
The serialized key format changes from
@pcp0@dcp0to@pcp:0@dcp:0. Workers should be rolled together, and existing KVPool entrieswritten with the old format will not be reused.
How was this patch tested?
The same implementation was validated in a multi-rank DCP deployment under
KVPool eviction pressure. No dedicated DCP lookup test file is added.