Skip to content

Conversation

@saikat-royc
Copy link
Collaborator

Description

Start with a short description of what the PR does and how this is a change from
the past.

This commit addresses a race condition in the CPU KV cache where a cached block could be prematurely evicted if multiple requests were relying on it and one of the requests finished early.

The rest of the description includes relevant details and context, examples:

  • why is this change being made,
  • the problem being solved and any relevant context,
  • why this is a good solution,
  • some information about the specific implementation,

The LocalCPUBackend now uses a reference counting mechanism for pinned keys:

  • Replaced the pinned_keys set with a pin_counts dictionary to track active references for each cached block.

  • The contains method increments the reference count when a key is pinned.

  • The maybe_unpin_keys method decrements the count and only truly unpins a key when its reference count reaches zero.

  • The eviction logic in the add() method now respects these reference counts preventing eviction of actively used blocks.

  • Additionally, the unpin_keys method was renamed to maybe_unpin_keys to better reflect its conditional nature.

  • shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@saikat-royc saikat-royc changed the title fix pin/unpin cpu backend keys logic [TPU host offoad] fix pin/unpin cpu backend keys logic Nov 6, 2025
This commit addresses a race condition in the CPU KV cache where a cached block could be
prematurely evicted if multiple requests were relying on it and one of the requests finished
early.

The LocalCPUBackend now uses a reference counting mechanism for pinned keys:
- Replaced the pinned_keys set with a `pin_counts` dictionary to track active
  references for each cached block.
- The contains method increments the reference count when a key is pinned.
- The maybe_unpin_keys method decrements the count and only truly unpins a
  key when its reference count reaches zero.
- The eviction logic in the add() method now respects these reference counts
  preventing eviction of actively used blocks.
- Additionally, the unpin_keys method was renamed to maybe_unpin_keys
  to better reflect its conditional nature.

Signed-off-by: Saikat Roychowdhury <saikat.royc85@gmail.com>
@saikat-royc saikat-royc merged commit 6565a5d into vllm-project:cpu-offloading/dev Nov 7, 2025
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants