Skip to content

[Bugfix][KVPool] Cover physical DCP shards in scheduler lookup - #15591

Open
zmc1997 wants to merge 1 commit into
vllm-project:mainfrom
zmc1997:codex/fix-kvpool-dcp-lookup-main
Open

[Bugfix][KVPool] Cover physical DCP shards in scheduler lookup#15591
zmc1997 wants to merge 1 commit into
vllm-project:mainfrom
zmc1997:codex/fix-kvpool-dcp-lookup-main

Conversation

@zmc1997

@zmc1997 zmc1997 commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

What this PR does / why we need it?

Main-branch counterpart of #15506. Addresses #15505.

KVPoolWorker.lookup_scheduler() must check every rank-specific key before
reporting 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 reported
even when another DCP key required by the subsequent load was missing.

This PR:

  • uses the same field:value format for PCP and DCP key fields as the other
    rank fields;
  • expands lookup keys over every DCP rank for each PP and head/TP rank;
  • keeps each rank shard's block/layer keys contiguous for the existing
    [rank_shard][block] hit reduction.

The all-rank lookup still assumes PCP=1; this change does not add unvalidated
PCP 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@dcp0 to
@pcp:0@dcp:0. Workers should be rolled together, and existing KVPool entries
written with the old format will not be reused.

How was this patch tested?

PYTHONPATH=. python3 tests/ut/distributed/ascend_store/test_metadata.py
Ran 58 tests ... OK

python3 -m py_compile +  vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/metadata.py +  vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_worker.py

git diff --check

The same implementation was validated in a multi-rank DCP deployment under
KVPool eviction pressure. No dedicated DCP lookup test file is added.

@zmc1997
zmc1997 requested a review from LCAIZJ as a code owner September 2, 2026 12:40
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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

  • DCP Shard Lookup Expansion: Updated KVPoolWorker.lookup_scheduler() to include all DCP shards in the lookup process, ensuring external prefix-cache hits are only reported when all required shards are present.
  • Serialized Key Format Update: Standardized the serialized key format for PCP and DCP fields from '@pcp0@dcp0' to '@pcp:0@dcp:0' to maintain consistency with other rank fields.
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 Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@github-actions

github-actions Bot commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

  • A PR should do only one thing, smaller PRs enable faster reviews.
  • Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
  • Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.


Tip

💡 Consider Linking a Related Issue or RFC

Your 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:

  • Fixes #<issue_number>
  • Closes #<issue_number>
  • Resolves #<issue_number>
  • Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

Suggested PR Title:

[kv_transfer][Misc] Update metadata key format to include colons and expand lookup keys by DCP rank

Suggested 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`.

Comment on lines 2219 to +2225
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))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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 $3 \times PP \times DCP \times TP \times K$ to $(1 + DCP + DCP \times TP) \times PP \times K$ (where $K$ is the number of keys). For typical distributed configurations, this results in a ~50-60% reduction in string replacement overhead.

Suggested change
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))

@LCAIZJ LCAIZJ added ready-all run all e2e test for pr and removed module:tests labels Sep 2, 2026
Signed-off-by: zmc1997 <40617288+zmc1997@users.noreply.github.com>
@zmc1997
zmc1997 force-pushed the codex/fix-kvpool-dcp-lookup-main branch from 7244ec3 to 0774245 Compare September 2, 2026 12:47
@zmc1997

zmc1997 commented Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

module:tests ready-all run all e2e test for pr

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants