[KingsleyZhang123]

Purpose

Fix the edge case in async kv load, where the request tokens left for compute are less then one block. Current implementation will still use 1 block and cache_block will fail.

Test Plan

Tested with customized implementation with prompt: "The capitol of France is ", which only has 7 tokens, while 1 block size is 16.

Test Result

Will return response successfully

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[liuzijing2014]

cc @robertgshaw2-redhat if can use a quick review

[yeqcharlotte]

thanks! @KingsleyZhang123 - could you add a small minimal-repo unit test to guard this?

we ran into this edge case in internal disagg testing. cc: @robertgshaw2-redhat @njhill could you help to take a look?

[houseroad]

Looks reasonable. Could you add some unittest if possible.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[njhill]

/gemini review

[gemini-code-assist]

Code Review

This pull request correctly addresses an edge case in the asynchronous KV cache loading logic. The fix ensures that num_computed_tokens is accurately capped by the actual number of tokens in the request, which is crucial for the correctness of subsequent scheduling decisions, especially when request.num_tokens is less than the allocated block capacity (e.g., less than block_size). The change is clear and the problem description is well-articulated.

I have one minor suggestion to improve the clarity of the inline comment.

Summary of Findings

• Comment Clarity: The inline comment for the added line could be more general to better describe its function of capping num_computed_tokens to the actual number of tokens in the request, rather than just the specific case of being less than one block. I have provided a code suggestion for this.

Merge Readiness

The core logic of the fix is sound and addresses an important correctness issue. There's one minor suggestion for improving comment clarity. Once that's considered, the PR should be in good shape for merging. As an AI, I am not authorized to directly approve pull requests, and users should have others review and approve this code before merging.

[gemini-code-assist]

The current comment "# Handle the case where num request tokens less then one block." is a bit specific. While true for one scenario, the fix is more general: it ensures num_computed_tokens doesn't exceed request.num_tokens, regardless of whether request.num_tokens is less than one block or spans multiple blocks but is less than the total capacity of those blocks.

Could we make the comment slightly more general to better reflect the line's purpose? For example:

# Cap num_computed_tokens at the actual number of tokens in the request.

This would clarify that num_computed_tokens is being adjusted to the true length of the content being cached, rather than just the allocated block capacity.

Suggested change

	# Handle the case where num request tokens less then one block.
	num_computed_tokens = min(num_computed_tokens, request.num_tokens)
	# Cap num_computed_tokens at the actual number of tokens in the request.

[njhill]

@KingsleyZhang123 sorry for taking some time to review this - I have been looking at it more closely and trying to understand how it could happen. As far as I can see this shouldn't happen with the existing connector implementations - in particular the NixlConnector should only return token counts that are multiples of the block size and smaller than the total number of request tokens. Did you encounter this with a different connector implementation?

So it would be great if you could add a small unit test that exposes the problem without your fix, as @yeqcharlotte suggested.

I'm not saying that the fix is wrong, just that we may want to make further adjustments - possibly by more explicitly enforcing assumptions/constraints on connectors w.r.t. block alignment, and/or clarifying cases where this additional min may be necessary. For example, it could in theory also address other cases where the allocated blocks exceed the number of request tokens, not just in the request.num_tokens < block_size case ... per gemini's last comment above.

[gemini-code-assist]

Thanks @njhill for the detailed review and questions for the author.

Just a quick note regarding the reference to "gemini's last comment above": my previous comment (comment 5) was unfortunately just a notification about reaching a daily quota limit and did not contain any technical feedback on this pull request. Apologies if that caused any confusion!

I agree with the suggestion from both you and @yeqcharlotte that adding a minimal unit test to reproduce the issue without the fix would be very valuable for demonstrating the problem and ensuring the fix is robust.

[KingsleyZhang123]

@njhill Thanks for reviewing, I think the NixlConnector is rounding down the blocks so this won't happen. however in our implementations we are rounding up the blocks we need, so we do have the case where left tokens are smaller than a block. I'll add a follow unit test for this change.

From scheduler's perspective, it should accept both round-up & round-down usecases IMO, let me know your thoughts on this.

[njhill]

Thanks @KingsleyZhang123 for the clarification. It's not just smaller than a block, but any case where the allocated blocks exceed the request tokens right?