add block_size

Author: bukeao

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

@@ -1010,7 +1010,7 @@ def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
 
         local_block_ids = meta.local_block_ids
         self.copy_blocks(self.host_xfer_buffers, self.device_kv_caches,
-                         local_block_ids, local_block_ids, "h2d")
+                         local_block_ids, local_block_ids, "h2d", self.block_size)
         if logger.isEnabledFor(logging.DEBUG):
             logger.debug(
                 "synced recved kv of request[%s] to device kv buffer,"
@@ -1030,7 +1030,7 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata):
                     ",".join(map(str, meta.local_block_ids)))
             # blocking
             self.copy_blocks(self.device_kv_caches, self.host_xfer_buffers,
-                             meta.local_block_ids, meta.local_block_ids, "d2h")
+                             meta.local_block_ids, meta.local_block_ids, "d2h", self.block_size)
 
     def get_finished(self) -> tuple[set[str], set[str]]:
         """