diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 37a9458654..51e9a12d3c 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/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() 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 66be05474c..4b1b2e406a 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -31,7 +31,14 @@ def runtime(monkeypatch, gpu_slots=1): 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) + r.metrics = dict( + gpu_batches=0, + cpu_batches=0, + wire_bytes=0, + useful_bytes=0, + load_calls=0, + load_requests=0, + ) r.actor = SimpleNamespace(_world_addresses=["source"]) return r @@ -273,7 +280,7 @@ def test_gpu_connector_failure_releases_lease(connector, connector_module): connector._get_connector_metadata = ( lambda: connector_module.XavierConnectorMetadata(load_requests=[request]) ) - connector._load_gpu_request = AsyncMock(side_effect=RuntimeError("peer failed")) + connector._load_gpu_requests = AsyncMock(side_effect=RuntimeError("peer failed")) connector._release_load_request = AsyncMock() with pytest.raises(RuntimeError, match="peer failed"): connector.start_load_kv(SimpleNamespace()) @@ -666,3 +673,280 @@ def test_slab_selection_avoids_alternating_churn_per_peer(monkeypatch): assert r._select_slab_bytes(0, 8) == 16 assert r._select_slab_bytes(1, 4) == 4 assert r._select_slab_bytes(0, 4) == 16 + + +@pytest.mark.asyncio +async def test_batch_load_preserves_shared_keys_and_bounds_transfers(monkeypatch): + source, dest = runtime(monkeypatch), runtime(monkeypatch) + stage(source, 1) + stage(source, 2) + sent = [] + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + async def send(reads, ref, size): + assert sum(r.nbytes for r in reads) <= size + sent.append(reads) + await source.send(reads, ref, size) + + monkeypatch.setattr(xo, "copy_to", copy) + peer = SimpleNamespace( + gpu_snapshot_locations_v1=AsyncMock(side_effect=source.locations), + send_gpu_request_v1=AsyncMock(side_effect=send), + ) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + dest.actor.read_request_blocks_v1 = AsyncMock( + side_effect=lambda rank, reads: pack_reads(source.store, reads) + ) + requests = [{0: {"K": {2: i}}} for i in range(6)] + requests += [{0: {"K": {1: 6}}}, {0: {"K": {1: 7, 2: 0}}}] + await dest.run(dest.load_requests, requests) + assert dest.caches["K"].tolist() == [[2, -2]] * 6 + [[1, -1]] * 2 + peer.gpu_snapshot_locations_v1.assert_awaited_once() + assert len(sent) == 2 + assert [len(r.keys) for batch in sent for r in batch] == [4, 2] + assert dest.metrics["gpu_batches"] == 2 + assert dest.metrics["cpu_batches"] == 1 + + +@pytest.mark.asyncio +async def test_batch_conflicting_destinations_fail_before_transfer(monkeypatch): + r = runtime(monkeypatch) + lookup = AsyncMock() + monkeypatch.setattr(xo, "actor_ref", lookup) + with pytest.raises(ValueError, match="Conflicting"): + await r.run(r.load_requests, [{0: {"K": {1: 0}}}, {1: {"K": {2: 0}}}]) + lookup.assert_not_awaited() + assert r.caches["K"].count_nonzero() == 0 + + +@pytest.mark.asyncio +async def test_connector_batch_cancellation_holds_all_leases( + connector, connector_module +): + requests = [ + connector_module.XavierLoadRequest(str(i), {}, lease=str(i)) for i in range(2) + ] + entered, release = asyncio.Event(), asyncio.Event() + + async def load(entries): + assert entries == requests + entered.set() + await release.wait() + + connector._load_gpu_requests = AsyncMock(side_effect=load) + connector._release_load_request = AsyncMock() + task = asyncio.create_task(connector._load_gpu_batch(requests)) + await entered.wait() + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + connector._release_load_request.assert_not_awaited() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert connector._release_load_request.await_count == 2 + + +@pytest.mark.asyncio +async def test_connector_batch_attempts_all_lease_releases(connector, connector_module): + requests = [ + connector_module.XavierLoadRequest(str(i), {}, lease=str(i)) for i in range(2) + ] + connector._load_gpu_requests = AsyncMock() + connector._release_load_request = AsyncMock( + side_effect=[RuntimeError("release failed"), None] + ) + with pytest.raises(RuntimeError, match="release failed"): + await connector._load_gpu_batch(requests) + assert connector._release_load_request.await_count == 2 + + +@pytest.mark.asyncio +async def test_connector_batch_builds_one_actor_call( + connector, connector_module, monkeypatch +): + requests = [ + connector_module.XavierLoadRequest( + str(i), {0: {123: i}}, local_transfers_by_group={0: {0: {123: i}}} + ) + for i in range(2) + ] + connector._registered_kv_caches = {"layer": torch.zeros(8, 2)} + transfer = SimpleNamespace(load_gpu_requests_v1=AsyncMock()) + connector._ensure_gpu_cache_mapping = AsyncMock(return_value=transfer) + monkeypatch.setattr(torch.cuda, "synchronize", lambda *args: None) + await connector._load_gpu_requests(requests) + transfer.load_gpu_requests_v1.assert_awaited_once_with( + [{0: {"layer": {123: 0}}}, {0: {"layer": {123: 1}}}] + ) + + +@pytest.mark.asyncio +async def test_batch_small_requests_keep_small_slab_bound(monkeypatch): + source, dest = runtime(monkeypatch, gpu_slots=2), runtime(monkeypatch) + stage(source, 2) + source.send_buffers[8] = source.send_buffer[:8] + dest.recv_buffers[8] = dest.recv_buffer[:8] + dest.recv_refs[8] = dest.recv_buffers[8] + + sizes = [] + + async def copy(buffers, refs): + sizes.append(buffers[0].numel()) + 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_requests, [{0: {"K": {2: i}}} for i in range(6)]) + assert dest.caches["K"][:6].tolist() == [[2, -2]] * 6 + assert peer.send_gpu_request_v1.await_count == 3 + assert sizes == [16, 8, 8] + assert source.metrics["wire_bytes"] == 32 + assert dest.metrics["load_calls"] == 1 + assert dest.metrics["load_requests"] == 6 + + +@pytest.mark.asyncio +async def test_batch_groups_multiple_source_ranks_without_losing_destinations( + monkeypatch, +): + sources = [runtime(monkeypatch), runtime(monkeypatch)] + dest = runtime(monkeypatch) + dest.actor._world_addresses = ["first", "second"] + for key, source in enumerate(sources, 1): + stage(source, key) + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + peers = { + address: SimpleNamespace( + gpu_snapshot_locations_v1=AsyncMock(side_effect=source.locations), + send_gpu_request_v1=AsyncMock(side_effect=source.send), + ) + for address, source in zip(dest.actor._world_addresses, sources) + } + monkeypatch.setattr( + xo, "actor_ref", AsyncMock(side_effect=lambda address, uid: peers[address]) + ) + await dest.run( + dest.load_requests, + [ + {0: {"K": {1: 0}}}, + {1: {"K": {2: 1}}}, + {0: {"K": {1: 2}}}, + {1: {"K": {2: 3}}}, + ], + ) + assert dest.caches["K"][:4].tolist() == [[1, -1], [2, -2], [1, -1], [2, -2]] + for peer in peers.values(): + peer.gpu_snapshot_locations_v1.assert_awaited_once() + peer.send_gpu_request_v1.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_fence", [False, True]) +async def test_close_releases_snapshot_storage_and_is_idempotent( + monkeypatch, fail_fence +): + from ..gpu_transfer import GPUTransferMixin + + r = runtime(monkeypatch) + stage(r, 1) + stage(r, 2) + assert r.store.reserve("held", [1, 2]) + calls = [] + + def fence(*args): + calls.append(args) + if fail_fence: + raise RuntimeError("fence failed") + + monkeypatch.setattr(torch.cuda, "synchronize", fence) + actor = SimpleNamespace(_gpu_transfer=r, _snapshot_store=r.store) + for _ in range(2): + if fail_fence: + with pytest.raises(RuntimeError, match="fence failed"): + await GPUTransferMixin.close_gpu_caches_v1(actor) + else: + await GPUTransferMixin.close_gpu_caches_v1(actor) + assert len(calls) == 1 + assert actor._gpu_transfer is r + assert not r.store.blocks and not r.store.ready and not r.store.tiers + assert not r.store.logical_dtypes and not r.store.leases and not r.store.evicted + assert r.store.counts == {"gpu": 0, "cpu": 0} + assert r.send_buffer is None and r.recv_buffer is None + assert not r.recv_refs and not r.caches + with pytest.raises(RuntimeError, match="already registered"): + GPUTransferMixin.map_gpu_caches_v1(actor, {}, 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("duplicate_rank", [False, True]) +async def test_batch_large_small_and_cross_rank_duplicates(monkeypatch, duplicate_rank): + 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] + dest.actor._world_addresses = ["first", "second"] + sent = [] + + async def copy(buffers, refs): + sent.append(buffers[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), + ) + lookup = AsyncMock(return_value=peer) + monkeypatch.setattr(xo, "actor_ref", lookup) + if duplicate_rank: + requests = [{0: {"K": {1: 3}}}, {1: {"K": {1: 3}}}] + await dest.run(dest.load_requests, requests) + assert dest.caches["K"][3].tolist() == [1, -1] + assert len(sent) == 1 + else: + requests = [{0: {"K": {1: 0, 2: 1}}}] + requests += [{0: {"K": {1: 2}}}, {0: {"K": {2: 3}}}] + await dest.run(dest.load_requests, requests) + assert sent == [16] # largest request needs full slab; no 4-byte splits + assert dest.caches["K"][:4].tolist() == [[1, -1], [2, -2]] * 2 + lookup.assert_awaited_once() + assert lookup.await_args.kwargs["address"] == "first" + peer.gpu_snapshot_locations_v1.assert_awaited_once() + + +def test_alloc_filters_only_null_positions_in_hybrid_groups( + connector, connector_module +): + for request_id, key, destination in [("a", 11, 4), ("b", 22, 5)]: + connector._requests_need_load[request_id] = connector_module.XavierLoadRequest( + request_id, {0: {key: 0, key + 1: 1}}, lease=request_id + ) + groups = ( + [SimpleNamespace(is_null=False), SimpleNamespace(is_null=False)], + [SimpleNamespace(is_null=True), SimpleNamespace(is_null=False)], + ) + blocks = SimpleNamespace( + blocks=groups, + get_block_ids=lambda d=destination: ([d, d + 2], [0, d]), + ) + connector.update_state_after_alloc( + SimpleNamespace(request_id=request_id), blocks, 32 + ) + mapping = connector._requests_need_load[request_id].local_transfers_by_group + assert mapping[0] == {0: {key: destination, key + 1: destination + 2}} + assert mapping[1] == {0: {key + 1: destination}} 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 59cd378d2b..d036d0dc20 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py @@ -389,7 +389,7 @@ async def test_gpu_connector_mapping_staging_and_load_fences( stage_gpu_requests_v1=AsyncMock( side_effect=lambda *args: calls.append("stage") ), - load_gpu_request_v1=AsyncMock(side_effect=lambda *args: calls.append("load")), + load_gpu_requests_v1=AsyncMock(side_effect=lambda *args: calls.append("load")), close_gpu_caches_v1=AsyncMock(), ) connector._get_transfer_ref = AsyncMock(return_value=transfer) @@ -409,9 +409,9 @@ async def test_gpu_connector_mapping_staging_and_load_fences( load = connector_module.XavierLoadRequest( "r", {1: {111: 0}}, local_transfers_by_group={0: {1: {111: 3}}} ) - await connector._load_gpu_request(load) + await connector._load_gpu_requests([load]) assert calls == ["fence", "load", "fence"] - transfer.load_gpu_request_v1.assert_awaited_once_with({1: {"layer": {111: 3}}}) + transfer.load_gpu_requests_v1.assert_awaited_once_with([{1: {"layer": {111: 3}}}]) assert transfer.map_gpu_caches_v1.await_count == 1 connector._get_connector_metadata = ( lambda: connector_module.XavierConnectorMetadata(store_requests=[request]) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 261e64abac..626f25e35f 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -171,11 +171,12 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non if not metadata.load_requests: return + if getattr(self, "_gpu_budget", None) is not None: + self._call(self._load_gpu_batch(metadata.load_requests)) + return + try: - if getattr(self, "_gpu_budget", None) is not None: - for request in metadata.load_requests: - self._call(self._load_gpu_request(request)) - elif self._registered_kv_caches: + if self._registered_kv_caches: for request in metadata.load_requests: self._load_request_blocks(request) else: @@ -335,6 +336,14 @@ def update_state_after_alloc( return block_ids_by_group = self._normalize_block_groups(blocks.get_block_ids()) + # Sliding-window groups use shared null blocks for skipped positions. + # They are placeholders, not writable destinations for remote KV. + null_positions = { + (group_id, index) + for group_id, group in enumerate(getattr(blocks, "blocks", ())) + for index, block in enumerate(group) + if block.is_null + } local_transfers_by_group: Dict[int, Dict[int, Dict[int, int]]] = {} for group_id, group_block_ids in enumerate(block_ids_by_group): updated: Dict[int, Dict[int, int]] = {} @@ -346,6 +355,7 @@ def update_state_after_alloc( updated[from_rank] = { remote_block_id: group_block_ids[placeholder] for remote_block_id, placeholder in remote_to_placeholder.items() + if (group_id, placeholder) not in null_positions } local_transfers_by_group[group_id] = updated @@ -758,21 +768,41 @@ async def _stage_gpu_requests(self, requests): torch.cuda.synchronize() await transfer.stage_gpu_requests_v1(entries) - async def _load_gpu_request(self, request): + async def _load_gpu_batch(self, requests): + from .gpu_transfer import finish_before_cancel + + async def load_and_release(): + try: + await self._load_gpu_requests(requests) + finally: + results = await asyncio.gather( + *(self._release_load_request(r) for r in requests if r.lease), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + raise result + + # A cancelled caller must not release leases while the peer is writing. + await finish_before_cancel(asyncio.create_task(load_and_release())) + + async def _load_gpu_requests(self, requests): transfer = await self._ensure_gpu_cache_mapping() - ranks = {} - for rank in request.transfers: - layers = {} - for name, cache in self._registered_kv_caches.items(): - for layer, _ in self._iter_kv_tensors(name, cache): - mapping = self._get_local_transfer_map(request, layer, rank) - if mapping: - layers[layer] = mapping - if layers: - ranks[rank] = layers - # Drain prior EngineCore kernels before another process writes recycled slots. + entries = [] + for request in requests: + ranks = {} + for rank in request.transfers: + layers = {} + for name, cache in self._registered_kv_caches.items(): + for layer, _ in self._iter_kv_tensors(name, cache): + mapping = self._get_local_transfer_map(request, layer, rank) + if mapping: + layers[layer] = mapping + if layers: + ranks[rank] = layers + entries.append(ranks) torch.cuda.synchronize() - await transfer.load_gpu_request_v1(ranks) + await transfer.load_gpu_requests_v1(entries) torch.cuda.synchronize() @staticmethod