[KV Pool]get_num_new_matched_tokens return 0 if token length < block_size

[Pz1116]

What this PR does / why we need it?

Currently, we call lookup_client for looking up token hit in KV Pool, however, when token length < block size, the key will be empty and there is no point to lookup in KV Pool backend since there will never be a hit.
Hence, add early return in get_num_new_matched_tokens when token_len < block_size

Does this PR introduce any user-facing change?

How was this patch tested?

• vLLM version: v0.16.0
• vLLM main: vllm-project/vllm@4034c3d

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 introduces a performance optimization to the KV Pool lookup mechanism. By adding a conditional early exit, it prevents redundant calls to the lookup client when the token length is too small to yield any matches, thereby improving efficiency and reducing computational overhead.

Highlights

• KV Pool Optimization: Implemented an early return in the get_num_new_matched_tokens function to prevent unnecessary lookup calls when the token_len is less than or equal to block_size, as no hits are possible under this condition.

🧠 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.

Changelog

• vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_scheduler.py
  • Added an early return condition in get_num_new_matched_tokens to immediately return 0, False if token_len is less than or equal to self._block_size.

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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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]

👋 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.

[gemini-code-assist]

Code Review

This pull request aims to optimize the KV pool lookup by adding an early return in get_num_new_matched_tokens when the token length is smaller than the block size, avoiding a lookup that would never succeed. My review focuses on the correctness of this optimization. I've found a potential issue where the condition token_len <= self._block_size is too strict and may prevent valid cache lookups for prompts that are exactly one block long. I've suggested changing it to token_len < self._block_size to align with the stated intent and avoid missing potential cache hits.

[gemini-code-assist]

The condition token_len <= self._block_size seems too strict. If token_len is exactly self._block_size, it represents one full block of tokens that could potentially be found in the KV pool. By returning early, you prevent this lookup from happening.

Based on the pull request description ('when token length < block size'), the condition should likely be token_len < self._block_size. This would correctly skip lookups for partial blocks while allowing lookups for single full blocks.

Suggested change

	if token_len <= self._block_size:
	return 0, False
	if token_len < self._block_size:
	return 0, False

[LCAIZJ]

LGTM

[LL1024LL]

Hi, I'm experiencing an issue after upgrading to v0.17.0-rc1.

The API response's prompt_tokens_details.cached_tokens field is now returning null for all requests, whereas in the previous version it
showed the number of cached tokens from the KV pool.

After investigation, I found that this PR (#7146) added an early return when token_len < block_size, which prevents the load_spec from
being created for short sequences. Since load_spec is the only mechanism that tracks kvpool_cached_tokens, this causes the metrics to be
lost.

Is this the intended behavior? Should the metrics tracking be decoupled from the KV pool lookup optimization?

For context:

• Most of my requests are shorter than block_size (128 tokens)
• The KV pool is still functioning (I can see transfer activity in logs)
• Only the metrics reporting is missing

Would it be possible to create a load_spec with kvpool_cached_tokens=0 for short sequences to preserve metrics tracking while keeping the
lookup optimization?

[Pz1116]

Hi, I'm experiencing an issue after upgrading to v0.17.0-rc1.

The API response's prompt_tokens_details.cached_tokens field is now returning null for all requests, whereas in the previous version it showed the number of cached tokens from the KV pool.

After investigation, I found that this PR (#7146) added an early return when token_len < block_size, which prevents the load_spec from being created for short sequences. Since load_spec is the only mechanism that tracks kvpool_cached_tokens, this causes the metrics to be lost.

Is this the intended behavior? Should the metrics tracking be decoupled from the KV pool lookup optimization?

For context:

• Most of my requests are shorter than block_size (128 tokens)
• The KV pool is still functioning (I can see transfer activity in logs)
• Only the metrics reporting is missing

Would it be possible to create a load_spec with kvpool_cached_tokens=0 for short sequences to preserve metrics tracking while keeping the lookup optimization?

load_spec with kvpool_cached_tokens=0 is doable, but I'm curious that where does prompt_tokens_details.cached_tokens take value from? Since LoadSpec is a class defined internally by AscendstoreConnector while prompt_tokens_details.cached_tokens seems like a metric from vLLM.