perf: reduce Xavier GPU transfer padding with persistent small slabs

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

@@ -68,6 +68,19 @@ def __init__(self, actor, caches: Dict[str, torch.Tensor], budget: int):
         )
         self.recv_buffer = torch.zeros_like(self.send_buffer)
         self.recv_ref = xo.buffer_ref(actor.address, self.recv_buffer)
+        # Keep stable views for small warm-cache transfers. Reuse their buffer
+        # identities, rather than creating and registering a slice per request.
+        small_bytes = min(self.slab_bytes, 256 * 1024)
+        self.send_buffers = {self.slab_bytes: self.send_buffer}
+        self.recv_buffers = {self.slab_bytes: self.recv_buffer}
+        self.recv_refs = {self.slab_bytes: self.recv_ref}
+        if small_bytes < self.slab_bytes:
+            self.send_buffers[small_bytes] = self.send_buffer[:small_bytes]
+            self.recv_buffers[small_bytes] = self.recv_buffer[:small_bytes]
+            self.recv_refs[small_bytes] = xo.buffer_ref(
+                actor.address, self.recv_buffers[small_bytes]
+            )
+        self._small_slab_streak: Dict[int, int] = {}
         self.send_lock, self.recv_lock = asyncio.Lock(), asyncio.Lock()
         self.tasks: set[asyncio.Task] = set()
         self.closing = False
@@ -104,6 +117,9 @@ async def _close(self):
                 dict(self.metrics, cache=self.store.stats()),
             )
         finally:
+            self.recv_refs.clear()
+            self.recv_buffers.clear()
+            self.send_buffers.clear()
             self.recv_ref = None
             self.caches.clear()
 
@@ -128,7 +144,7 @@ async def stage(self, entries):
                         continue
                     cache = self.caches[layer]
                     blocks = cache.index_select(
-                        0, torch.tensor(ids, device=cache.device)
+                        0, torch.tensor(ids, dtype=torch.long, device=cache.device)
                     )
                     self.store.stage(layer, block_keys, blocks)
                     copied = True
@@ -174,13 +190,13 @@ def locations(self, reads):
 
     async def send(self, reads, remote_ref, slab_bytes):
         async with self.send_lock:
-            if slab_bytes != self.slab_bytes:
+            if slab_bytes not in self.send_buffers:
                 raise ValueError("Xavier peer transfer slab sizes differ")
             locations = self.locations(reads)
             if any(tier != "gpu" for tier in locations.values()):
                 raise ValueError("GPU transfer requested for a CPU snapshot")
             size = sum(read.nbytes for read in reads)
-            if size > self.slab_bytes:
+            if size > slab_bytes:
                 raise ValueError("Xavier GPU batch exceeds transfer slab")
             offset = 0
             for read in reads:
@@ -196,10 +212,22 @@ async def send(self, reads, remote_ref, slab_bytes):
                 )
                 offset = end
             await asyncio.to_thread(torch.cuda.synchronize, self.device)
-            await xo.copy_to([self.send_buffer], [remote_ref])
-            self.metrics["wire_bytes"] += self.slab_bytes
+            await xo.copy_to([self.send_buffers[slab_bytes]], [remote_ref])
+            self.metrics["wire_bytes"] += slab_bytes
             self.metrics["useful_bytes"] += size
 
+    def _select_slab_bytes(self, rank: int, size: int) -> int:
+        candidate = min(n for n in self.recv_refs if n >= size)
+        if candidate == self.slab_bytes:
+            self._small_slab_streak.pop(rank, None)
+            return candidate
+        # xoscar caches only the latest registration on each peer channel.
+        # Require consecutive small batches so full/small tails do not churn
+        # registrations. Sustained small workloads still reuse the small view.
+        streak = min(2, self._small_slab_streak.get(rank, 0) + 1)
+        self._small_slab_streak[rank] = streak
+        return candidate if streak == 2 else self.slab_bytes
+
     async def load(self, ranks):
         from .transfer import TransferActor
 
@@ -242,8 +270,9 @@ async def load(self, ranks):
                     for batch in batch_reads(selected, max_bytes=self.slab_bytes):
                         size = sum(read.nbytes for read in batch)
                         if tier == "gpu":
+                            slab_bytes = self._select_slab_bytes(rank, size)
                             await sender.send_gpu_request_v1(
-                                batch, self.recv_ref, self.slab_bytes
+                                batch, self.recv_refs[slab_bytes], slab_bytes
                             )
                             payload = self.recv_buffer[:size]
                         else:
@@ -253,7 +282,11 @@ async def load(self, ranks):
                         for read, blocks in unpack_reads(payload, batch):
                             cache = self.caches[read.layer]
                             cache[
-                                torch.tensor(read.destinations, device=cache.device)
+                                torch.tensor(
+                                    read.destinations,
+                                    dtype=torch.long,
+                                    device=cache.device,
+                                )
                             ] = blocks.to(cache.device, non_blocking=True)
                         # Both producer's IPC ownership and receiver slab reuse
                         # require the cache writes to finish before acknowledging.

xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py

@@ -21,10 +21,14 @@ def runtime(monkeypatch, gpu_slots=1):
     r.device = torch.device("cpu")
     r.caches = {"K": torch.zeros(8, 2, dtype=torch.bfloat16)}
     r.store = TieredKVSnapshotStore(2, gpu_slots * 4, 4, r.device)
+    r._small_slab_streak = {}
     r.slab_bytes = 16
     r.send_buffer = torch.zeros(16, dtype=torch.uint8)
     r.recv_buffer = torch.zeros_like(r.send_buffer)
     r.recv_ref = r.recv_buffer
+    r.send_buffers = {r.slab_bytes: r.send_buffer}
+    r.recv_buffers = {r.slab_bytes: r.recv_buffer}
+    r.recv_refs = {r.slab_bytes: r.recv_ref}
     r.send_lock, r.recv_lock = asyncio.Lock(), asyncio.Lock()
     r.tasks, r.closing = set(), False
     r.metrics = dict(gpu_batches=0, cpu_batches=0, wire_bytes=0, useful_bytes=0)
@@ -38,8 +42,13 @@ def stage(r, key):
 
 
 @pytest.mark.asyncio
-async def test_mixed_tier_load_preserves_bits_and_destinations(monkeypatch):
+@pytest.mark.parametrize("small_slab", [False, True])
+async def test_mixed_tier_load_preserves_bits_and_destinations(monkeypatch, small_slab):
     source, dest = runtime(monkeypatch), runtime(monkeypatch)
+    if small_slab:
+        source.send_buffers[4] = source.send_buffer[:4]
+        dest.recv_buffers[4] = dest.recv_buffer[:4]
+        dest.recv_refs[4] = dest.recv_buffers[4]
     stage(source, 1)
     stage(source, 2)
     assert source.store.tiers == {1: "cpu", 2: "gpu"}
@@ -64,6 +73,55 @@ async def copy(buffers, refs):
     assert source.metrics["wire_bytes"] == 16
 
 
+@pytest.mark.asyncio
+async def test_load_switches_after_repeated_small_batches_and_reuses_views(monkeypatch):
+    source, dest = runtime(monkeypatch, gpu_slots=2), runtime(monkeypatch)
+    stage(source, 1)
+    stage(source, 2)
+    source.send_buffers[4] = source.send_buffer[:4]
+    dest.recv_buffers[4] = dest.recv_buffer[:4]
+    dest.recv_refs[4] = dest.recv_buffers[4]
+    buffers_seen = []
+
+    async def copy(buffers, refs):
+        buffers_seen.append(buffers[0])
+        assert buffers[0].numel() == refs[0].numel()
+        refs[0].copy_(buffers[0])
+
+    monkeypatch.setattr(xo, "copy_to", copy)
+    peer = SimpleNamespace(
+        gpu_snapshot_locations_v1=AsyncMock(side_effect=source.locations),
+        send_gpu_request_v1=AsyncMock(side_effect=source.send),
+    )
+    monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer))
+    await dest.run(dest.load, {0: {"K": {1: 0}}})
+    await dest.run(dest.load, {0: {"K": {2: 1}}})
+    await dest.run(dest.load, {0: {"K": {2: 1}}})
+    await dest.run(dest.load, {0: {"K": {1: 2, 2: 3}}})
+    assert [b.numel() for b in buffers_seen] == [16, 4, 4, 16]
+    assert buffers_seen[1] is buffers_seen[2] is source.send_buffers[4]
+    assert buffers_seen[0] is buffers_seen[3] is source.send_buffer
+    assert dest.caches["K"][:4].tolist() == [[1, -1], [2, -2], [1, -1], [2, -2]]
+    assert source.metrics["wire_bytes"] == 40
+    assert source.metrics["useful_bytes"] == 20
+
+
+@pytest.mark.asyncio
+async def test_send_rejects_payload_exceeding_selected_slab(monkeypatch):
+    r = runtime(monkeypatch, gpu_slots=2)
+    stage(r, 1)
+    stage(r, 2)
+    r.send_buffers[4] = r.send_buffer[:4]
+    copy = AsyncMock()
+    monkeypatch.setattr(xo, "copy_to", copy)
+    reads = [LayerRead("K", [1, 2], [0, 1], (2,), torch.bfloat16)]
+    with pytest.raises(ValueError, match="exceeds transfer slab"):
+        await r.send(reads, r.recv_ref, 4)
+    with pytest.raises(ValueError, match="slab sizes differ"):
+        await r.send(reads, r.recv_ref, 8)
+    copy.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_cancel_waits_for_transfer_before_buffer_reuse(monkeypatch):
     r = runtime(monkeypatch)
@@ -105,6 +163,7 @@ async def test_close_drains_operations_and_rejects_new_work(monkeypatch):
     release.set()
     await asyncio.gather(task, close)
     assert not r.caches and r.recv_ref is None
+    assert not r.recv_refs and not r.recv_buffers and not r.send_buffers
 
 
 @pytest.mark.asyncio
@@ -256,6 +315,7 @@ async def test_cancelled_close_finishes_ipc_cleanup(monkeypatch):
     with pytest.raises(asyncio.CancelledError):
         await closing
     assert not r.caches and r.recv_ref is None
+    assert not r.recv_refs and not r.recv_buffers and not r.send_buffers
     await r.close()
 
 
@@ -272,6 +332,7 @@ async def test_gpu_packing_preserves_all_bf16_bits(monkeypatch, device):
     r.store.publish(list(range(8)), {"K"})
     r.slab_bytes = 131072
     r.send_buffer = torch.empty(r.slab_bytes, dtype=torch.uint8, device=device)
+    r.send_buffers = {r.slab_bytes: r.send_buffer}
     received = torch.empty_like(r.send_buffer)
 
     async def copy(buffers, refs):
@@ -538,3 +599,70 @@ def fail(value, dim, ids):
     assert r.store.ready == {1}
     assert r.store.read("K", [1]).tolist() == [[1, -1]]
     assert "V" not in r.store.blocks[1]
+
+
+@pytest.mark.parametrize("block_bytes", [1024, 4096, 8192])
+def test_constructor_creates_persistent_slab_views(monkeypatch, block_bytes):
+    from .. import gpu_transfer
+
+    # Keep real torch allocations/views on CPU while exercising the CUDA-only
+    # constructor; only the cache device guard and NIXL reference are substituted.
+    class Cache:
+        is_cuda = True
+        device = torch.device("cpu")
+
+        def __getitem__(self, index):
+            return torch.empty(block_bytes, dtype=torch.uint8)
+
+        def element_size(self):
+            return 1
+
+    monkeypatch.setattr(gpu_transfer, "version", lambda name: "0.11.1")
+    refs = []
+
+    def buffer_ref(address, buffer):
+        ref = SimpleNamespace(address=address, buffer=buffer)
+        refs.append(ref)
+        return ref
+
+    monkeypatch.setattr(xo, "buffer_ref", buffer_ref)
+    actor = SimpleNamespace(
+        address="nixl://127.0.0.1:1234", _snapshot_store=SimpleNamespace(capacity=8)
+    )
+    transfer = GPUTransfer(actor, {"K": Cache()}, block_bytes * 2)
+    expected_slab = block_bytes * 64
+    assert transfer.slab_bytes == expected_slab
+    expected_keys = {expected_slab, min(expected_slab, 262144)}
+    assert set(transfer.send_buffers) == expected_keys
+    assert set(transfer.recv_buffers) == expected_keys
+    assert set(transfer.recv_refs) == expected_keys
+    assert len(refs) == len(expected_keys)
+    assert transfer.send_buffers[expected_slab] is transfer.send_buffer
+    assert transfer.recv_buffers[expected_slab] is transfer.recv_buffer
+    assert transfer.recv_refs[expected_slab] is transfer.recv_ref
+    for size in expected_keys:
+        for buffers, full in (
+            (transfer.send_buffers, transfer.send_buffer),
+            (transfer.recv_buffers, transfer.recv_buffer),
+        ):
+            view = buffers[size]
+            assert view.numel() == size
+            assert view.data_ptr() == full.data_ptr()
+            assert (
+                view.untyped_storage().data_ptr() == full.untyped_storage().data_ptr()
+            )
+            assert view.storage_offset() == 0
+        ref = transfer.recv_refs[size]
+        assert ref.address == actor.address
+        assert ref.buffer is transfer.recv_buffers[size]
+
+
+def test_slab_selection_avoids_alternating_churn_per_peer(monkeypatch):
+    r = runtime(monkeypatch)
+    r.recv_refs[4] = r.recv_buffer[:4]
+    assert [r._select_slab_bytes(0, size) for size in [8, 4] * 5] == [16] * 10
+    assert r._select_slab_bytes(1, 4) == 16
+    assert r._select_slab_bytes(0, 4) == 4
+    assert r._select_slab_bytes(0, 8) == 16
+    assert r._select_slab_bytes(1, 4) == 4
+    assert r._select_slab_bytes(0, 4) == 16

xinference/model/llm/vllm/xavier/test/test_review_regressions.py

@@ -1,5 +1,6 @@
 # Copyright 2022-2026 Xinference Holdings Pte. Ltd
 # Licensed under the Apache License, Version 2.0.
+import asyncio
 from collections import deque
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, Mock
@@ -420,3 +421,56 @@ async def test_gpu_connector_mapping_staging_and_load_fences(
     )
     connector.save_kv_layer("layer", cache, None)
     assert connector._pending_store_requests["r"] is request
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("fail_first", [False, True])
+async def test_gpu_mapping_serializes_callers_and_retries(
+    connector, connector_module, monkeypatch, fail_first
+):
+    from torch.multiprocessing import reductions
+
+    entered = asyncio.Event()
+    release = asyncio.Event()
+    calls = []
+
+    async def register(descriptors, budget):
+        calls.append((descriptors, budget))
+        if len(calls) == 1:
+            entered.set()
+            await release.wait()
+            if fail_first:
+                raise RuntimeError("registration failed")
+
+    transfer = SimpleNamespace(map_gpu_caches_v1=register)
+    connector._get_transfer_ref = AsyncMock(return_value=transfer)
+    connector._registered_kv_caches = {"K": object()}
+    connector._gpu_budget = 256
+    tensor = SimpleNamespace(is_cuda=True)
+    connector._iter_kv_tensors = lambda *args: [("K", tensor)]
+    monkeypatch.setattr(connector_module, "block_major_view", lambda t, n: t)
+    monkeypatch.setattr(reductions, "reduce_tensor", lambda t: (None, ("descriptor",)))
+    monkeypatch.setattr(torch.cuda, "synchronize", Mock())
+
+    first = asyncio.create_task(connector._ensure_gpu_cache_mapping())
+    await asyncio.wait_for(entered.wait(), timeout=2)
+    second = asyncio.create_task(connector._ensure_gpu_cache_mapping())
+    try:
+        await asyncio.sleep(0)
+        assert len(calls) == 1
+        assert not connector._gpu_cache_mapped
+    finally:
+        release.set()
+        results = await asyncio.wait_for(
+            asyncio.gather(first, second, return_exceptions=True), timeout=2
+        )
+    assert results[1] is transfer
+    if fail_first:
+        assert isinstance(results[0], RuntimeError)
+        assert str(results[0]) == "registration failed"
+    else:
+        assert results[0] is transfer
+    assert len(calls) == (2 if fail_first else 1)
+    assert connector._gpu_cache_mapped
+    assert await connector._ensure_gpu_cache_mapping() is transfer
+    assert len(calls) == (2 if fail_first else 1)

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

@@ -125,6 +125,7 @@ def __init__(
             )
         self._gpu_budget = self._xavier_config.get("gpu_cache_bytes")
         self._gpu_cache_mapped = False
+        self._gpu_mapping_lock = asyncio.Lock()
         self._kv_schema: Optional[XavierKVSchema] = None
         self._rank = int(self._xavier_config.get("rank", 0))
         self._is_producer = self._kv_transfer_config.is_kv_producer
@@ -722,19 +723,24 @@ async def _ensure_gpu_cache_mapping(self):
         from torch.multiprocessing.reductions import reduce_tensor
 
         transfer = await self._get_transfer_ref()
-        if not self._gpu_cache_mapped:
-            descriptors = {}
-            for name, cache in self._registered_kv_caches.items():
-                for layer, tensor in self._iter_kv_tensors(name, cache):
-                    tensor = block_major_view(tensor, self._num_cache_blocks)
-                    if not tensor.is_cuda:
-                        raise ValueError("Xavier GPU transfer requires CUDA KV caches")
-                    _, descriptors[layer] = reduce_tensor(tensor)
-            if not descriptors:
-                raise ValueError("Xavier GPU transfer requires registered KV caches")
-            torch.cuda.synchronize()
-            await transfer.map_gpu_caches_v1(descriptors, self._gpu_budget)
-            self._gpu_cache_mapped = True
+        async with self._gpu_mapping_lock:
+            if not self._gpu_cache_mapped:
+                descriptors = {}
+                for name, cache in self._registered_kv_caches.items():
+                    for layer, tensor in self._iter_kv_tensors(name, cache):
+                        tensor = block_major_view(tensor, self._num_cache_blocks)
+                        if not tensor.is_cuda:
+                            raise ValueError(
+                                "Xavier GPU transfer requires CUDA KV caches"
+                            )
+                        _, descriptors[layer] = reduce_tensor(tensor)
+                if not descriptors:
+                    raise ValueError(
+                        "Xavier GPU transfer requires registered KV caches"
+                    )
+                torch.cuda.synchronize()
+                await transfer.map_gpu_caches_v1(descriptors, self._gpu_budget)
+                self._gpu_cache_mapped = True
         return transfer
 
     async def _stage_gpu_requests(self, requests):