From 544c6d9c11bb284e10e192ab0bf2fa99597a34bc Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sat, 3 Oct 2026 18:43:25 +0800 Subject: [PATCH 1/3] perf: use persistent small slabs for Xavier GPU transfers --- .../model/llm/vllm/xavier/gpu_transfer.py | 26 ++++++-- .../llm/vllm/xavier/test/test_gpu_transfer.py | 63 ++++++++++++++++++- 2 files changed, 82 insertions(+), 7 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 7760698c83..69583047b5 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -68,6 +68,18 @@ 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.send_lock, self.recv_lock = asyncio.Lock(), asyncio.Lock() self.tasks: set[asyncio.Task] = set() self.closing = False @@ -104,6 +116,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() @@ -174,13 +189,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,8 +211,8 @@ 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 async def load(self, ranks): @@ -242,8 +257,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 = min(n for n in self.recv_refs if n >= 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: diff --git a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py index 9861933427..9134b7bbfd 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -25,6 +25,9 @@ def runtime(monkeypatch, gpu_slots=1): 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 +41,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"} @@ -61,7 +69,55 @@ async def copy(buffers, refs): assert dest.caches["K"][:3].count_nonzero() == 0 assert dest.metrics["gpu_batches"] == dest.metrics["cpu_batches"] == 1 assert source.metrics["useful_bytes"] == 4 - assert source.metrics["wire_bytes"] == 16 + assert source.metrics["wire_bytes"] == (4 if small_slab else 16) + + +@pytest.mark.asyncio +async def test_load_selects_smallest_slab_and_reuses_buffer_objects(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": {1: 2, 2: 3}}}) + assert [b.numel() for b in buffers_seen] == [4, 4, 16] + assert buffers_seen[0] is buffers_seen[1] is source.send_buffers[4] + assert buffers_seen[2] is source.send_buffer + assert dest.caches["K"][:4].tolist() == [[1, -1], [2, -2], [1, -1], [2, -2]] + assert source.metrics["wire_bytes"] == 24 + assert source.metrics["useful_bytes"] == 16 + + +@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 @@ -105,6 +161,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 +313,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 +330,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): From aaea732813734e17c69a26fe2a3cdc220f36af78 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 03:57:25 +0800 Subject: [PATCH 2/3] fix: serialize Xavier GPU cache registration --- .../model/llm/vllm/xavier/gpu_transfer.py | 8 ++- .../xavier/test/test_review_regressions.py | 54 +++++++++++++++++++ .../model/llm/vllm/xavier/v1_connector.py | 32 ++++++----- 3 files changed, 79 insertions(+), 15 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 69583047b5..ed7715ccf2 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -143,7 +143,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 @@ -269,7 +269,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. diff --git a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py index 77038bffa9..59cd378d2b 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/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) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 560f6400b3..261e64abac 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/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): From c72606288427e2a0c9a14c218a805fa8a8e803d6 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 04:21:23 +0800 Subject: [PATCH 3/3] fix: avoid Xavier slab registration churn on alternating batches --- .../model/llm/vllm/xavier/gpu_transfer.py | 15 +++- .../llm/vllm/xavier/test/test_gpu_transfer.py | 83 +++++++++++++++++-- 2 files changed, 90 insertions(+), 8 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index ed7715ccf2..37a9458654 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -80,6 +80,7 @@ def __init__(self, actor, caches: Dict[str, torch.Tensor], budget: int): 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 @@ -215,6 +216,18 @@ async def send(self, reads, remote_ref, slab_bytes): 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 @@ -257,7 +270,7 @@ 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 = min(n for n in self.recv_refs if n >= size) + slab_bytes = self._select_slab_bytes(rank, size) await sender.send_gpu_request_v1( batch, self.recv_refs[slab_bytes], slab_bytes ) diff --git a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py index 9134b7bbfd..66be05474c 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -21,6 +21,7 @@ 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) @@ -69,11 +70,11 @@ async def copy(buffers, refs): assert dest.caches["K"][:3].count_nonzero() == 0 assert dest.metrics["gpu_batches"] == dest.metrics["cpu_batches"] == 1 assert source.metrics["useful_bytes"] == 4 - assert source.metrics["wire_bytes"] == (4 if small_slab else 16) + assert source.metrics["wire_bytes"] == 16 @pytest.mark.asyncio -async def test_load_selects_smallest_slab_and_reuses_buffer_objects(monkeypatch): +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) @@ -95,13 +96,14 @@ async def copy(buffers, refs): 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] == [4, 4, 16] - assert buffers_seen[0] is buffers_seen[1] is source.send_buffers[4] - assert buffers_seen[2] is source.send_buffer + 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"] == 24 - assert source.metrics["useful_bytes"] == 16 + assert source.metrics["wire_bytes"] == 40 + assert source.metrics["useful_bytes"] == 20 @pytest.mark.asyncio @@ -597,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