change position of block_size

Author: bukeao

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

@@ -1009,8 +1009,8 @@ def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
         assert self.copy_blocks is not None
 
         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", self.block_size)
+        self.copy_blocks(self.block_size, self.host_xfer_buffers, self.device_kv_caches,
+                         local_block_ids, local_block_ids, "h2d")
         if logger.isEnabledFor(logging.DEBUG):
             logger.debug(
                 "synced recved kv of request[%s] to device kv buffer,"
@@ -1029,8 +1029,8 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata):
                     "local_block_ids: %s. ", req_id,
                     ",".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", self.block_size)
+            self.copy_blocks(self.block_size, self.device_kv_caches, self.host_xfer_buffers,
+                             meta.local_block_ids, meta.local_block_ids, "d2h")
 
     def get_finished(self) -> tuple[set[str], set[str]]:
         """

vllm/v1/worker/hpu_model_runner.py

@@ -2515,12 +2515,12 @@ def _swap_out_hpu_blocks(
     cpu_cache[cpu_block_indices] = hpu_cache[hpu_block_indices].cpu()
 
 def copy_kv_blocks(
+    block_size: int,
     src_kv_caches: dict[str, torch.Tensor],
     dst_kv_caches: dict[str, torch.Tensor],
     src_block_ids: list[int],
     dst_block_ids: list[int],
     direction: Literal["h2d", "d2h"],
-    block_size: int
 ) -> None:
     """Copy kv blocks between different buffers."""
     if not src_kv_caches or not dst_kv_caches or \