perf: batch Xavier KV loads across scheduled requests

xinference/model/llm/vllm/xavier/gpu_transfer.py

@@ -84,7 +84,14 @@ def __init__(self, actor, caches: Dict[str, torch.Tensor], budget: int):
         self.send_lock, self.recv_lock = asyncio.Lock(), asyncio.Lock()
         self.tasks: set[asyncio.Task] = set()
         self.closing = False
-        self.metrics = dict(gpu_batches=0, cpu_batches=0, wire_bytes=0, useful_bytes=0)
+        self.metrics = dict(
+            gpu_batches=0,
+            cpu_batches=0,
+            wire_bytes=0,
+            useful_bytes=0,
+            load_calls=0,
+            load_requests=0,
+        )
 
     async def run(self, function, *args):
         if self.closing:
@@ -122,6 +129,15 @@ async def _close(self):
             self.send_buffers.clear()
             self.recv_ref = None
             self.caches.clear()
+            self.send_buffer = self.recv_buffer = None
+            self.store.blocks.clear()
+            self.store.ready.clear()
+            self.store.tiers.clear()
+            self.store.logical_dtypes.clear()
+            self.store.leases.clear()
+            self.store.evicted.clear()
+            self.store.counts.update(gpu=0, cpu=0)
+            self._small_slab_streak.clear()
 
     async def stage(self, entries):
         if self.closing:
@@ -229,8 +245,39 @@ def _select_slab_bytes(self, rank: int, size: int) -> int:
         return candidate if streak == 2 else self.slab_bytes
 
     async def load(self, ranks):
+        return await self.load_requests([ranks])
+
+    async def load_requests(self, requests):
         from .transfer import TransferActor
 
+        # Index by destination, not hash: one cached block may feed multiple
+        # requests' slots. Reject conflicting writes before any transfer starts.
+        ranks = {}
+        owners = {}
+        request_bytes = {}
+        for request in requests:
+            for rank, layers in request.items():
+                size = sum(
+                    len(mapping)
+                    * self.caches[name][0].numel()
+                    * self.caches[name].element_size()
+                    for name, mapping in layers.items()
+                )
+                request_bytes[rank] = max(request_bytes.get(rank, 0), size)
+                for name, mapping in layers.items():
+                    for key, destination in mapping.items():
+                        target = (name, destination)
+                        if target in owners:
+                            if owners[target] != key:
+                                raise ValueError("Conflicting Xavier KV destinations")
+                            continue
+                        owners[target] = key
+                        ranks.setdefault(rank, {}).setdefault(name, {})[
+                            destination
+                        ] = key
+
+        self.metrics["load_calls"] += 1
+        self.metrics["load_requests"] += len(requests)
         async with self.recv_lock:
             for rank, layers in ranks.items():
                 sender = await xo.actor_ref(
@@ -240,8 +287,8 @@ async def load(self, ranks):
                 reads = [
                     LayerRead(
                         name,
-                        list(mapping),
                         list(mapping.values()),
+                        list(mapping),
                         tuple(self.caches[name].shape[1:]),
                         self.caches[name].dtype,
                     )
@@ -267,7 +314,16 @@ async def load(self, ranks):
                                     read.dtype,
                                 )
                             )
-                    for batch in batch_reads(selected, max_bytes=self.slab_bytes):
+                    # Combining small requests must not turn each transfer back
+                    # into a padded full-slab copy. Large requests retain the
+                    # original batch bound to avoid fragmenting cold prefills.
+                    limit = self.slab_bytes
+                    if tier == "gpu":
+                        limit = min(
+                            (n for n in self.recv_refs if n >= request_bytes[rank]),
+                            default=self.slab_bytes,
+                        )
+                    for batch in batch_reads(selected, max_bytes=limit):
                         size = sum(read.nbytes for read in batch)
                         if tier == "gpu":
                             slab_bytes = self._select_slab_bytes(rank, size)
@@ -318,11 +374,13 @@ async def send_gpu_request_v1(self, reads, remote_ref, slab_bytes):
         runtime = self._gpu_transfer
         return await runtime.run(runtime.send, reads, remote_ref, slab_bytes)
 
-    async def load_gpu_request_v1(self, ranks):
+    async def load_gpu_requests_v1(self, requests):
         runtime = self._gpu_transfer
-        return await runtime.run(runtime.load, ranks)
+        return await runtime.run(runtime.load_requests, requests)
 
     async def close_gpu_caches_v1(self):
         runtime = getattr(self, "_gpu_transfer", None)
         if runtime is not None:
+            # Keep the closed runtime as a tombstone: close() reuses its task,
+            # and map_gpu_caches_v1 must not re-register after shutdown/failure.
             await runtime.close()