Commit 1586d41f authored by yangshj1's avatar yangshj1
Browse files

fix connector metadata

parent c6b7a44b
...@@ -155,3 +155,4 @@ class SchedulerOutput: ...@@ -155,3 +155,4 @@ class SchedulerOutput:
# KV Cache Connector metadata. # KV Cache Connector metadata.
kv_connector_metadata: Optional[KVConnectorMetadata] = None kv_connector_metadata: Optional[KVConnectorMetadata] = None
kv_lmcache_connector_metadata: Optional[KVConnectorMetadata] = None
...@@ -89,6 +89,7 @@ class Scheduler(SchedulerInterface): ...@@ -89,6 +89,7 @@ class Scheduler(SchedulerInterface):
self.connector = KVConnectorFactory.create_connector_v1( self.connector = KVConnectorFactory.create_connector_v1(
config=self.vllm_config, role=KVConnectorRole.SCHEDULER) config=self.vllm_config, role=KVConnectorRole.SCHEDULER)
self.lmcache_connector = None
if envs.VLLM_LMCACHE_ENABLE: if envs.VLLM_LMCACHE_ENABLE:
lmcache_config = copy.deepcopy(self.vllm_config) lmcache_config = copy.deepcopy(self.vllm_config)
lmcache_config.kv_transfer_config = KVTransferConfig(kv_connector="LMCacheConnectorV1", kv_role="kv_both") lmcache_config.kv_transfer_config = KVTransferConfig(kv_connector="LMCacheConnectorV1", kv_role="kv_both")
...@@ -602,7 +603,7 @@ class Scheduler(SchedulerInterface): ...@@ -602,7 +603,7 @@ class Scheduler(SchedulerInterface):
if self.lmcache_connector is not None: if self.lmcache_connector is not None:
meta = self.lmcache_connector.build_connector_meta(scheduler_output) meta = self.lmcache_connector.build_connector_meta(scheduler_output)
scheduler_output.kv_connector_metadata = meta scheduler_output.kv_lmcache_connector_metadata = meta
events = self.kv_cache_manager.take_events() events = self.kv_cache_manager.take_events()
if events: if events:
...@@ -1034,7 +1035,7 @@ class Scheduler(SchedulerInterface): ...@@ -1034,7 +1035,7 @@ class Scheduler(SchedulerInterface):
if self.lmcache_connector is not None: if self.lmcache_connector is not None:
meta = self.lmcache_connector.build_connector_meta(scheduler_output) meta = self.lmcache_connector.build_connector_meta(scheduler_output)
scheduler_output.kv_connector_metadata = meta scheduler_output.kv_lmcache_connector_metadata = meta
events = self.kv_cache_manager.take_events() events = self.kv_cache_manager.take_events()
if events: if events:
......
...@@ -1747,7 +1747,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): ...@@ -1747,7 +1747,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
lmcache_connector = get_lmcache_connector() lmcache_connector = get_lmcache_connector()
if lmcache_connector is not None: if lmcache_connector is not None:
lmcache_connector.bind_connector_metadata( lmcache_connector.bind_connector_metadata(
scheduler_output.kv_connector_metadata) scheduler_output.kv_lmcache_connector_metadata)
lmcache_connector.start_load_kv(get_forward_context()) lmcache_connector.start_load_kv(get_forward_context())
@staticmethod @staticmethod
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment