Skip to content

[Main][Bugfix][KV Pool]Fix KV transfer Put Logic#7717

Open
Pz1116 wants to merge 4 commits intovllm-project:mainfrom
Pz1116:put_all_keys
Open

[Main][Bugfix][KV Pool]Fix KV transfer Put Logic#7717
Pz1116 wants to merge 4 commits intovllm-project:mainfrom
Pz1116:put_all_keys

Conversation

@Pz1116
Copy link
Collaborator

@Pz1116 Pz1116 commented Mar 27, 2026

What this PR does / why we need it?

Before when we do put for KV Pool, we find the first non-existing key and put all the blocks starting from that index; however, if the prefix cache blocks is from another request, and some of the blocks are evicted due to LRU, we will be putting blocks that still exist in the pool, and causing MooncakeStore printing unnecessary logs in master service.

What this PR does:

  1. Now we lookup all the keys and only put the ones that are missing.
  2. Fix lookup_scheduler in pool_worker so it handles GQA correctly.
  3. Fixes a few existing typos
  4. Add UT, written by codex

Does this PR introduce any user-facing change?

How was this patch tested?

@gemini-code-assist
Copy link
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 optimizes the KV pool transfer mechanism by refining how blocks are identified and sent to the store. By shifting from a sequential prefix-based check to an explicit existence check for all keys, the system avoids redundant data transfers and unnecessary log noise. Additionally, the PR includes necessary refactoring for type consistency and improved test coverage to ensure robust behavior in the KV transfer pipeline.

Highlights

  • KV Transfer Logic: Updated the KV transfer process to check for the existence of all keys individually, ensuring only missing blocks are sent to the store, which prevents redundant logging and unnecessary operations.
  • Code Quality and Fixes: Corrected several typos in class names (e.g., LasyerMultiBlockReqMeta to LayerMultiBlockReqMeta) and improved the lookup mechanism to return a list of booleans for better handling of block states.
  • Testing: Added comprehensive unit tests to verify that the sending threads correctly identify and only process missing keys.

🧠 New Feature in Public Preview: You can now enable Memory 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 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 counter productive. 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.

Copy link
Contributor

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

Choose a reason for hiding this comment

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

Code Review

This pull request refactors the KV cache lookup mechanism in the Ascend store to return a boolean list for key existence, allowing for the filtering of non-contiguous missing blocks. It also fixes a typo in LayerMultiBlockReqMeta and introduces unit tests for the sending threads. Feedback was provided to ensure the keys list is filtered in KVCacheStoreLayerSendingThread to maintain accurate logging.

Suggested PR Title:

[vllm-ascend][Distributed][BugFix] Refactor KV cache lookup to support non-contiguous missing blocks

Suggested PR Summary:

### What this PR does / why we need it?
This PR refactors the KV cache lookup mechanism in the Ascend store. It changes the `lookup` method to return a list of booleans indicating the existence of each key, allowing the sending threads to filter and transmit only the specific blocks missing from the cache. This replaces the previous logic which only supported skipping a prefix of blocks. Additionally, it fixes a typo in `LayerMultiBlockReqMeta` and adds unit tests to verify the new filtering logic.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
New unit tests were added in `tests/ut/distributed/mooncake/test_kv_transfer.py` to verify that both `KVCacheStoreSendingThread` and `KVCacheStoreLayerSendingThread` correctly identify and process only missing keys.

Comment on lines +346 to +348
starts = [starts[index] for index in missing_indices]
ends = [ends[index] for index in missing_indices]
key_list = [key_list[index] for index in missing_indices]
Copy link
Contributor

Choose a reason for hiding this comment

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

high

For consistency with KVCacheStoreSendingThread._handle_request and to ensure logging is accurate, the keys list should also be filtered to only include the missing keys. Currently, len(keys) in the subsequent log message refers to the number of keys before filtering, which is misleading as it doesn't reflect the actual number of blocks being stored.

Suggested change
starts = [starts[index] for index in missing_indices]
ends = [ends[index] for index in missing_indices]
key_list = [key_list[index] for index in missing_indices]
starts = [starts[index] for index in missing_indices]
ends = [ends[index] for index in missing_indices]
key_list = [key_list[index] for index in missing_indices]
keys = [keys[index] for index in missing_indices]

@github-actions
Copy link
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.

Signed-off-by: Pz1116 <zpbzpb123123@gmail.com>
Co-authored-by: DreamerLeader <2270923832@qq.com>
Co-authored-by: fems14 <1804143737@qq.com>
Pz1116 added 3 commits March 27, 2026 13:34
Co-authored-by: DreamerLeader <2270923832@qq.com>
Co-authored-by: fems14 <1804143737@qq.com>

Signed-off-by: Pz1116 <zpbzpb123123@gmail.com>
Co-authored-by: DreamerLeader <2270923832@qq.com>
Co-authored-by: fems14 <1804143737@qq.com>

Signed-off-by: Pz1116 <zpbzpb123123@gmail.com>
Co-authored-by: DreamerLeader <2270923832@qq.com>
Co-authored-by: fems14 <1804143737@qq.com>

Signed-off-by: Pz1116 <zpbzpb123123@gmail.com>
@Pz1116
Copy link
Collaborator Author

Pz1116 commented Mar 27, 2026

please take a look @LCAIZJ @fems14 @baxingpiaochong

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant