[Feature] Support KV cache offloading and disagg prefill with LMCache connector. #12953
Conversation
|
👋 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 🚀 |
YaoJiayi
changed the title
KuntaiDu
requested review from
youkaichao,
comaniac,
robertgshaw2-redhat and
simon-mo
|
Do we still need lmcache_vllm ? or just need lmcache ? |
Signed-off-by: YaoJiayi <120040070@link.cuhk.edu.cn>
Signed-off-by: YaoJiayi <120040070@link.cuhk.edu.cn>
Signed-off-by: YaoJiayi <120040070@link.cuhk.edu.cn>
Signed-off-by: YaoJiayi <120040070@link.cuhk.edu.cn>
Signed-off-by: YaoJiayi <120040070@link.cuhk.edu.cn>
YaoJiayi
force-pushed
the
localdev/lmcache-connector
branch
from
1cff5c6 to
12b713e
Compare
lmcache-vllm repo is not needed if this PR gets merged. |
| model_executable: torch.nn.Module, | ||
| model_input: "ModelInputForGPUWithSamplingMetadata", | ||
| kv_caches: List[torch.Tensor], | ||
| hidden_or_intermediate_states: Union[torch.Tensor, |
There was a problem hiding this comment.
why not send hidden_or_intermediate_states to remote cache
There was a problem hiding this comment.
Currently LMCache assumes that the user only stores KV caches. I guess the API can be extended but it requires some API change. @YaoJiayi does this align with what you are thinking?
There was a problem hiding this comment.
This is correct:) The last token will be re-prefilled for disagg prefill for now.
| model_executable: torch.nn.Module, | ||
| model_input: "ModelInputForGPUWithSamplingMetadata", | ||
| kv_caches: List[torch.Tensor], | ||
| hidden_or_intermediate_states: Union[torch.Tensor, |
There was a problem hiding this comment.
Currently LMCache assumes that the user only stores KV caches. I guess the API can be extended but it requires some API change. @YaoJiayi does this align with what you are thinking?
… connector. (vllm-project#12953) Signed-off-by: Louis Ulmer <ulmerlouis@gmail.com>
LMCache (https://github.com/LMCache/LMCache/tree/dev) uses the
kv_transferinterface to support both KV cache offloading and disagg prefill.The original interfaces
recv_kv_caches_and_hidden_statesandsend_kv_caches_and_hidden_statesinkv_connectorare used as wrappers to calllmcache_retrieve_kv(retrieves the KV from local cpu, local disk, or remote storage to vllm paged memory) andlmcache_store_kv(extracts the KV from vllm paged memory to local cpu, local disk, or remote storage) respectively.