[Feat][KV offloading][WIP] The prototype implementation of a KV offloader used in CPU KV server #22608
Conversation
Signed-off-by: ApostaC <[email protected]>
Signed-off-by: ApostaC <[email protected]>
|
👋 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 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 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a prototype implementation of a KV offloader for a CPU KV server using LMCache. The implementation includes a new abstract class BlockingKVInterface and a concrete class LMCacheBlockingKVMgr. My review identified a few critical issues in LMCacheBlockingKVMgr that need to be addressed. These include a potential crash on CPU-only servers due to a hardcoded .cuda() call, an incorrect boundary check for the worker rank, and a bug in the lookup_internal method that prevents it from correctly iterating over all ranks.
| (num_blocks, 1)) * block_size).flatten() | ||
|
|
||
| # TODO: compatibility with multiple cuda devices | ||
| return slot_mapping[:len(token_ids)].cuda() |
There was a problem hiding this comment.
The .cuda() call here will cause a crash if this code is run on a CPU-only machine, which is the expected environment for a 'CPU KV server'. The LMCacheEngine should handle the necessary data transfers between devices. Please remove the explicit .cuda() call and let the underlying library manage device placement.
| return slot_mapping[:len(token_ids)].cuda() | |
| return slot_mapping[:len(token_ids)] |
| def lookup_internal(self, token_ids: list[int], pin: bool) -> int: | ||
| lengths = [] | ||
| for i in range(self.world_size): | ||
| length = self.lmcache_engines[0].lookup(token_ids, pin=pin) |
There was a problem hiding this comment.
There is a bug in this loop. It iterates with i from 0 to self.world_size - 1, but it always accesses self.lmcache_engines[0]. It should use self.lmcache_engines[i] to query the engine for each rank.
| length = self.lmcache_engines[0].lookup(token_ids, pin=pin) | |
| length = self.lmcache_engines[i].lookup(token_ids, pin=pin) |
| if rank > self.world_size: | ||
| raise ValueError( | ||
| f"Rank {rank} exceeds world size {self.world_size}.") |
There was a problem hiding this comment.
The check for rank validity is incorrect. Ranks are 0-indexed, so a rank equal to self.world_size is also out of bounds. The condition should be rank >= self.world_size.
| if rank > self.world_size: | |
| raise ValueError( | |
| f"Rank {rank} exceeds world size {self.world_size}.") | |
| if rank >= self.world_size: | |
| raise ValueError( | |
| f"Rank {rank} exceeds world size {self.world_size}.") |
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
This PR is part of #22605.
This PR gives a (WIP) prototype implementation for the KV cache offloader running in the CPU KV process.
The functionality in this PR will be replaced with a vLLM-native implementation in the future
Test Plan
Test Result
(Optional) Documentation Update