updated

Signed-off-by: [email protected] <[email protected]>

Author: robertgshaw2-redhat

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -159,7 +159,7 @@ def get_external_prefix_cache_blocks(
         Get the external prefix cache blocks from the connector.
 
         This function may change the state of the connector, which will
-        be used by `attach_connector_meta` later.
+        be used by `build_connector_meta` later.
 
         This function will also allocate/free the blocks dynamically when  
         there is remote cache hit.
@@ -178,13 +178,12 @@ def get_external_prefix_cache_blocks(
         pass
 
     @abstractmethod
-    def attach_connector_meta(
-            self, scheduler_output: SchedulerOutput) -> SchedulerOutput:
+    def build_connector_meta(
+            self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
         """
-        Attach the connector metadata to the request object.
+        Build the connector metadata for this step.
 
-        This function should NOT modify other fields in the scheduler_output 
-        except the `connector_metadata` field.
+        This function should NOT modify fields in the scheduler_output.
         Also, calling this function will reset the state of the connector.
 
         Args:

vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py

@@ -221,7 +221,7 @@ def get_external_prefix_cache_blocks(
         """Get the external prefix cache blocks from the connector.
 
         This function may change the state of the connector, which will be 
-        used by `attach_connector_meta` later.
+        used by `build_connector_meta` later.
 
         Args:
             request (Request): the request object.
@@ -283,12 +283,11 @@ def get_external_prefix_cache_blocks(
         else:
             return computed_blocks
 
-    def attach_connector_meta(
-            self, scheduler_output: SchedulerOutput) -> SchedulerOutput:
-        """Attach the connector metadata to the request object.
+    def build_connector_meta(
+            self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
+        """Build the connector metadata for this step.
 
-        This function should NOT modify other fields in the scheduler_output 
-        except the `kv_connector_metadata` field.
+        This function should NOT modify any fields in the scheduler_output.
         Also, calling this function will reset the state of the connector.
 
         Args:
@@ -304,10 +303,9 @@ def attach_connector_meta(
                 # store and load status
                 if not self.found_match_for_request(request):
                     meta.add_request(request, self._block_size, is_store=True)
-        scheduler_output.kv_connector_metadata = meta
 
         self._requests_need_load.clear()
-        return scheduler_output
+        return meta
 
     # ==============================
     # Helper functions

vllm/v1/core/sched/scheduler.py

@@ -453,8 +453,8 @@ def schedule(self) -> SchedulerOutput:
         # 2. Wrap up all the KV cache load / save ops into an opaque object
         # 3. Clear the internal states of the connector
         if self.connector is not None:
-            scheduler_output = self.connector.attach_connector_meta(
-                scheduler_output)
+            meta = self.connector.build_connector_meta(scheduler_output)
+            scheduler_output.kv_connector_metadata = meta
 
         # Advance the number of computed tokens for the request AFTER
         # the request is scheduled.