A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Feat] Support chunked prefill for LMCache connector (#14505)
Signed-off-by: YaoJiayi <120040070@link.cuhk.edu.cn>
Jiayi Yao committed
Mar 8, 2025 at 21:30 UTC
6d7f037748b2e7df64f3318e54101a1c80016f3c
1 file changed
+9
-19
vllm/distributed/kv_transfer/kv_connector/lmcache_connector.py
+9
-19
@@ -38,7 +38,8 @@ class LMCacheConnector(KVConnectorBase):
38
from lmcache.integration.vllm.utils import ENGINE_NAME
39
from lmcache.integration.vllm.vllm_adapter import (
40
RetrieveStatus, StoreStatus, init_lmcache_engine,
41
- lmcache_retrieve_kv, lmcache_should_store, lmcache_store_kv)
41
+ lmcache_retrieve_kv, lmcache_should_retrieve, lmcache_should_store,
42
+ lmcache_store_kv)
43
logger.info("Initializing LMCacheConfig under kv_transfer_config %s",
44
self.transfer_config)
45
@@ -54,6 +55,7 @@ class LMCacheConnector(KVConnectorBase):
55
self.cache_config = config.cache_config
56
self.lmcache_retrieve_kv = lmcache_retrieve_kv
57
self.lmcache_store_kv = lmcache_store_kv
58
+ self.lmcache_should_retrieve = lmcache_should_retrieve
59
self.lmcache_should_store = lmcache_should_store
60
self.store_status = StoreStatus
61
self.retrieve_status = RetrieveStatus
@@ -65,15 +67,11 @@ class LMCacheConnector(KVConnectorBase):
67
) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,
68
"ModelInputForGPUWithSamplingMetadata"]:
69
68
- hidden_or_intermediate_states = None
69
-
70
- # TODO (Jiayi): Need to support chunked prefill
71
- retrieve_status = self.retrieve_status.PREFILL
72
-
73
- model_input, bypass_model_exec = self.lmcache_retrieve_kv(
74
- model_executable, model_input, self.cache_config, kv_caches,
75
- retrieve_status)
76
-
70
+ retrieve_status = self.lmcache_should_retrieve(model_input)
71
+ model_input, bypass_model_exec, hidden_or_intermediate_states =\
72
+ self.lmcache_retrieve_kv(
73
+ model_executable, model_input, self.cache_config, kv_caches,
74
+ retrieve_status)
75
return hidden_or_intermediate_states, bypass_model_exec, model_input
76
77
def send_kv_caches_and_hidden_states(
@@ -84,15 +82,7 @@ class LMCacheConnector(KVConnectorBase):
82
hidden_or_intermediate_states: Union[torch.Tensor,
83
IntermediateTensors],
84
) -> None:
87
- num_reqs = 0
88
- seq_group_list = model_input.sampling_metadata.seq_groups
89
- assert seq_group_list is not None
90
- for seq_group in seq_group_list:
91
- seq_ids = seq_group.seq_ids
92
- for seq_id in seq_ids:
93
- num_reqs += 1
94
-
95
- # TODO (Jiayi): Only normal prefill is supported for now
85
+
86
store_status = self.lmcache_should_store(model_input)
87
self.lmcache_store_kv(
88
self.model_config,