diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 2b7e20b044..15fb169d83 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -178,8 +178,15 @@ async def stage(self, entries): failed = False copied = False last_reused_keys = None + # These complete snapshots were copied on this call's stream. They may + # be reused before our final fence, unlike unpublished work of a peer call. + staged_keys: set[int] = set() try: for layers in entries: + if self._can_stage_blocks(layers, staged_keys): + copied = self._stage_blocks(layers, staged_keys) or copied + last_reused_keys = None + continue for layer, (block_keys, ids) in layers.items(): # Unpublished layers may still be copying in another stage # awaiting its fence; only published data can skip the copy. @@ -222,6 +229,64 @@ async def stage(self, entries): if key in self.store.blocks and key not in self.store.ready: self.store._drop(key) + def _can_stage_blocks(self, layers, staged_keys): + if not layers or set(layers) != set(self.caches): + return False + keys, ids = next(iter(layers.values())) + return ( + bool(keys) + and len(keys) == len(ids) + and len(keys) == len(set(keys)) + and all(value == (keys, ids) for value in layers.values()) + and all( + key not in self.store.blocks + or ( + (key in self.store.ready or key in staged_keys) + and self.caches.keys() <= self.store.blocks[key].keys() + ) + for key in keys + ) + ) + + def _stage_blocks(self, layers, staged_keys): + keys, ids = next(iter(layers.values())) + count = max( + 1, min(MAX_REQUEST_BLOCKS, MAX_REQUEST_BYTES // self.store.block_bytes) + ) + pending_keys, pending_ids = [], [] + copied = False + + def flush(): + nonlocal copied + if not pending_keys: + return + index = torch.tensor(pending_ids, dtype=torch.long, device=self.device) + values = { + name: cache.index_select(0, index) + for name, cache in self.caches.items() + } + self.store.stage_blocks(pending_keys, values) + staged_keys.update(key for key in pending_keys if key in self.store.blocks) + copied = True + pending_keys.clear() + pending_ids.clear() + + for key, source in zip(keys, ids): + if key in self.store.ready or key in staged_keys: + flush() + # Admission of preceding blocks may have evicted this hit. + if key in self.store.blocks and ( + key in self.store.ready or key in staged_keys + ): + self.store.touch(key) + continue + pending_keys.append(key) + pending_ids.append(source) + if len(pending_keys) == count: + flush() + flush() + return copied + def locations(self, reads): result = {} for read in reads: 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 539551c5e9..ab9c04d283 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -343,20 +343,24 @@ async def copy(buffers, refs): @pytest.mark.asyncio @pytest.mark.parametrize("failure", [torch.cuda.OutOfMemoryError, RuntimeError]) +@pytest.mark.parametrize("packed", [False, True]) async def test_staging_failure_drains_and_drops_only_unpublished( - monkeypatch, failure, assert_gpu_lru_consistent + monkeypatch, failure, packed, assert_gpu_lru_consistent ): r = runtime(monkeypatch, gpu_slots=3) stage(r, 1) assert r.store.reserve("2:live", [1]) - original = r.store.stage + method = "stage_blocks" if packed else "stage" + original = getattr(r.store, method) + if not packed: + monkeypatch.setattr(r, "_can_stage_blocks", lambda layers, staged_keys: False) calls = [] - def copy(layer, keys, tensors): - original(layer, keys, tensors) + def copy(*args): + original(*args) raise failure("copy failed after partial staging") - monkeypatch.setattr(r.store, "stage", copy) + monkeypatch.setattr(r.store, method, copy) monkeypatch.setattr(torch.cuda, "synchronize", lambda device: calls.append("sync")) await r.run(r.stage, [{"K": ([1, 2], [0, 1])}]) assert calls == ["sync"] @@ -1310,3 +1314,179 @@ async def fail(): assert "background write failed" in caplog.text assert not r.tasks await r.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "byte_limit,block_limit,expected_batches", + [ + (16, 64, [[11, 12], [13, 14], [15]]), + (1024, 2, [[11, 12], [13, 14], [15]]), + (4, 64, [[11], [12], [13], [14], [15]]), + ], +) +async def test_complete_block_staging_is_bounded_and_preserves_content( + monkeypatch, byte_limit, block_limit, expected_batches +): + from .. import gpu_transfer + + r = runtime(monkeypatch, gpu_slots=8) + r.store.block_bytes = 8 + r.caches["K"] = torch.arange(16, dtype=torch.bfloat16).reshape(8, 2) + r.caches["V"] = -r.caches["K"] + monkeypatch.setattr(gpu_transfer, "MAX_REQUEST_BYTES", byte_limit) + monkeypatch.setattr(gpu_transfer, "MAX_REQUEST_BLOCKS", block_limit) + batches = [] + original = r.store.stage_blocks + + def capture(keys, layers): + batches.append(list(keys)) + return original(keys, layers) + + monkeypatch.setattr(r.store, "stage_blocks", capture) + await r.stage( + [{name: ([11, 12, 13, 14, 15], [7, 5, 3, 1, 0]) for name in r.caches}] + ) + assert batches == expected_batches + assert r.store.publish([11, 12, 13, 14, 15], {"K", "V"}) == [11, 12, 13, 14, 15] + assert r.store.read("K", [11, 12, 13, 14, 15]).tolist() == [ + [14, 15], + [10, 11], + [6, 7], + [2, 3], + [0, 1], + ] + assert torch.equal( + r.store.read("V", [11, 12, 13, 14, 15]), + -r.store.read("K", [11, 12, 13, 14, 15]), + ) + r.caches["K"].zero_() + assert r.store.read("K", [11]).tolist() == [[14, 15]] + + +@pytest.mark.asyncio +async def test_layer_specific_block_indices_use_existing_staging(monkeypatch): + r = runtime(monkeypatch, gpu_slots=2) + r.store.block_bytes = 8 + r.caches["K"][3] = torch.tensor([3, -3]) + r.caches["V"] = torch.ones_like(r.caches["K"]) + + packed_calls = [] + + def unexpected(*args): + packed_calls.append(args) + + monkeypatch.setattr(r.store, "stage_blocks", unexpected) + await r.stage([{"K": ([9], [3]), "V": ([9], [1])}]) + assert not packed_calls + assert r.store.read("K", [9]).tolist() == [[3, -3]] + assert r.store.read("V", [9]).tolist() == [[1, 1]] + + +@pytest.mark.asyncio +async def test_packed_partial_hit_gathers_only_missing_blocks(monkeypatch): + r = runtime(monkeypatch, gpu_slots=4) + stage(r, 1) + stage(r, 3) + r.caches["K"][2] = torch.tensor([2, -2]) + r.caches["K"][4] = torch.tensor([4, -4]) + gathered = [] + original = torch.Tensor.index_select + + def gather(value, dim, index): + gathered.extend(index.tolist()) + return original(value, dim, index) + + monkeypatch.setattr(torch.Tensor, "index_select", gather) + await r.stage([{"K": ([1, 2, 3, 4], [0, 2, 0, 4])}]) + assert gathered == [2, 4] + assert list(r.store.blocks) == [1, 2, 3, 4] + assert r.store.read("K", [1, 2, 3, 4]).tolist() == [ + [1, -1], + [2, -2], + [3, -3], + [4, -4], + ] + + +@pytest.mark.asyncio +async def test_packed_hit_evicted_by_preceding_admission_is_copied(monkeypatch): + r = runtime(monkeypatch, gpu_slots=1) + r.store.cpu_capacity = 1 + stage(r, 1) + stage(r, 2) + r.caches["K"][2] = torch.tensor([2, -2]) + r.caches["K"][3] = torch.tensor([3, -3]) + r.caches["K"][4] = torch.tensor([4, -4]) + await r.stage([{"K": ([3, 4, 2], [3, 4, 2])}]) + assert list(r.store.blocks) == [4, 2] + assert r.store.publish([3, 4, 2], {"K"}) == [4, 2] + assert r.store.read("K", [2]).tolist() == [[2, -2]] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("same_call", [False, True]) +async def test_unpublished_prefix_reuse_is_scoped_to_staging_call( + monkeypatch, same_call +): + r = runtime(monkeypatch, gpu_slots=8) + r.caches["K"] = torch.arange(16, dtype=torch.bfloat16).reshape(8, 2) + batches, gathered = [], [] + original_stage, original_gather = r.store.stage_blocks, torch.Tensor.index_select + + def capture(keys, layers): + batches.append(list(keys)) + original_stage(keys, layers) + + def gather(value, dim, index): + gathered.extend(index.tolist()) + return original_gather(value, dim, index) + + monkeypatch.setattr(r.store, "stage_blocks", capture) + monkeypatch.setattr(torch.Tensor, "index_select", gather) + entries = [{"K": ([1, 2, 3], [1, 2, 3])}, {"K": ([1, 2, 3, 4], [1, 2, 3, 4])}] + if same_call: + await r.stage(entries) + else: + for entry in entries: + await r.stage([entry]) + assert not r.store.ready + assert batches == ([[1, 2, 3], [4]] if same_call else [[1, 2, 3]]) + assert gathered == ([1, 2, 3, 4] if same_call else [1, 2, 3, 1, 2, 3, 4]) + assert torch.equal(r.store.read("K", [1, 2, 3, 4]), r.caches["K"][1:5]) + + +@pytest.mark.asyncio +async def test_call_local_hit_evicted_by_prior_chunk_is_gathered_again(monkeypatch): + r = runtime(monkeypatch, gpu_slots=1) + r.store.cpu_capacity = 1 + r.caches["K"] = torch.arange(16, dtype=torch.bfloat16).reshape(8, 2) + await r.stage([{"K": ([1, 2], [1, 2])}, {"K": ([3, 4, 2], [3, 4, 2])}]) + assert list(r.store.blocks) == [4, 2] + assert r.store.publish([4, 2], {"K"}) == [4, 2] + assert torch.equal(r.store.read("K", [2]), r.caches["K"][2:3]) + + +@pytest.mark.asyncio +async def test_later_chunk_failure_drops_all_unpublished_chunks( + monkeypatch, assert_gpu_lru_consistent +): + from .. import gpu_transfer + + r = runtime(monkeypatch, gpu_slots=8) + stage(r, 7) + monkeypatch.setattr(gpu_transfer, "MAX_REQUEST_BLOCKS", 2) + original, batches = r.store.stage_blocks, [] + + def fail(keys, layers): + batches.append(list(keys)) + original(keys, layers) + if len(batches) == 2: + raise RuntimeError("second chunk failed") + + monkeypatch.setattr(r.store, "stage_blocks", fail) + await r.stage([{"K": ([1, 2, 3, 4], [1, 2, 3, 4])}]) + assert batches == [[1, 2], [3, 4]] + assert set(r.store.blocks) == r.store.ready == {7} + assert r.store._block_sizes == {7: 4} + assert_gpu_lru_consistent(r.store) diff --git a/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py b/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py index 6573b7a58d..8136f804f8 100644 --- a/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py +++ b/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py @@ -361,3 +361,116 @@ def test_incremental_size_survives_demotion_eviction_and_key_reuse(): assert s._block_sizes[key] == sum( value.numel() * value.element_size() for value in layers.values() ) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_packed_blocks_own_storage_and_preserve_mixed_dtypes(device): + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA required") + target = torch.device(device) + s = TieredKVSnapshotStore(2, 64, 32, target) + layers = { + "K": torch.arange(8, device=target, dtype=torch.float32).reshape(2, 2, 2), + "V": torch.arange(4, device=target, dtype=torch.bfloat16).reshape(2, 1, 2), + "index": torch.tensor([[7], [9]], device=target, dtype=torch.int64), + } + expected = {name: value.cpu().clone() for name, value in layers.items()} + s.stage_blocks([1, 2], layers) + assert s.publish([1, 2], set(layers)) == [1, 2] + for name, value in layers.items(): + value.zero_() + assert torch.equal(s.read(name, [1, 2]), expected[name]) + first, second = s.blocks[1][name], s.blocks[2][name] + assert first.untyped_storage().data_ptr() != second.untyped_storage().data_ptr() + assert s.logical_dtypes[1][name] == expected[name].dtype + assert s._block_sizes == {1: 28, 2: 28} + s.stage_blocks([3], {name: value[:1] for name, value in layers.items()}) + assert s.tiers == {1: "cpu", 2: "gpu", 3: "gpu"} + for name in layers: + assert torch.equal(s.read(name, [1]), expected[name][:1]) + + +def test_packed_blocks_respect_leases_and_reject_invalid_batch(): + s = store(gpu=1, cpu=1) + layers = {"K": torch.tensor([[1.0, 2.0]]), "V": torch.tensor([[3.0, 4.0]])} + s.stage_blocks([1], layers) + s.publish([1], set(layers)) + assert s.reserve("2:gpu", [1]) + s.stage_blocks([2], layers) + s.publish([2], set(layers)) + assert s.reserve("2:cpu", [2]) + s.stage_blocks([3], layers) + assert set(s.blocks) == {1, 2} + with pytest.raises(ValueError, match="new snapshot"): + s.stage_blocks([1], layers) + with pytest.raises(ValueError, match="full-block size"): + s.stage_blocks([4], {"K": torch.ones(1, 5)}) + with pytest.raises(ValueError, match="length"): + s.stage_blocks([4, 5], layers) + with pytest.raises(ValueError, match="unique"): + s.stage_blocks( + [4, 4], {name: value.repeat(2, 1) for name, value in layers.items()} + ) + s.stage_blocks([], {}) + assert set(s.blocks) == {1, 2} + assert s.read("K", [1, 2]).tolist() == [[1.0, 2.0], [1.0, 2.0]] + + +def test_packed_copy_failure_drops_incomplete_block(monkeypatch): + s = store() + original = torch.Tensor.to + + def fail(value, *args, **kwargs): + if kwargs.get("copy"): + raise RuntimeError("copy failed") + return original(value, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "to", fail) + with pytest.raises(RuntimeError, match="copy failed"): + s.stage_blocks([1], {"K": torch.ones(1, 2)}) + assert not s.blocks and not s._block_sizes and not s._gpu_lru + assert s.counts == {"cpu": 0, "gpu": 0} + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_packed_bf16_preserves_all_bit_patterns(device): + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA required") + bits = torch.arange(65536, dtype=torch.int32).to(torch.int16).reshape(2, -1) + values = bits.to(device).view(torch.bfloat16) + s = TieredKVSnapshotStore(2, 262144, 131072, torch.device(device)) + s.stage_blocks([1, 2], {"K": values, "V": values}) + for name in ("K", "V"): + assert torch.equal(s.read(name, [1, 2]).view(torch.int16), bits) + # Compatible layers share only their own block allocation. + assert ( + s.blocks[1]["K"].untyped_storage().data_ptr() + == s.blocks[1]["V"].untyped_storage().data_ptr() + ) + assert s.blocks[1]["K"].untyped_storage().nbytes() == 131072 + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_packed_same_dtype_layers_with_different_shapes(device): + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA required") + s = TieredKVSnapshotStore(2, 80, 40, torch.device(device)) + layers = { + "K": torch.arange(12, device=device, dtype=torch.float32).reshape(2, 2, 3), + "V": torch.arange(8, device=device, dtype=torch.float32).reshape(2, 4, 1), + } + s.stage_blocks([1, 2], layers) + for name, value in layers.items(): + assert torch.equal(s.read(name, [1, 2]), value.cpu()) + assert s.blocks[1][name].shape == value.shape[1:] + assert ( + s.blocks[1]["K"].untyped_storage().data_ptr() + == s.blocks[1]["V"].untyped_storage().data_ptr() + ) + assert s.blocks[1]["K"].untyped_storage().nbytes() == 40 + assert ( + s.blocks[1]["K"].untyped_storage().data_ptr() + != s.blocks[2]["K"].untyped_storage().data_ptr() + ) + s.blocks[1]["K"].zero_() + assert torch.equal(s.read("V", [1, 2]), layers["V"].cpu()) diff --git a/xinference/model/llm/vllm/xavier/tiered_snapshot.py b/xinference/model/llm/vllm/xavier/tiered_snapshot.py index 3167454824..457fd2a77f 100644 --- a/xinference/model/llm/vllm/xavier/tiered_snapshot.py +++ b/xinference/model/llm/vllm/xavier/tiered_snapshot.py @@ -166,6 +166,61 @@ def stage( self._drop(key) raise + def stage_blocks(self, keys: List[int], layers: Dict[str, torch.Tensor]): + """Copy complete new blocks once per dtype, with independent ownership.""" + if not keys: + return + if len(keys) != len(set(keys)): + raise ValueError("Packed staging requires unique snapshot keys") + if not layers or any(len(value) != len(keys) for value in layers.values()): + raise ValueError("Snapshot keys and tensors differ in length") + groups: Dict[ + Tuple[torch.device, torch.dtype], List[Tuple[str, torch.Tensor]] + ] = {} + nbytes = sum( + value[0].numel() * value.element_size() for value in layers.values() + ) + if nbytes > self.block_bytes: + raise ValueError("Snapshot exceeds configured full-block size") + if any(key in self.blocks for key in keys): + raise ValueError("Packed staging requires new snapshot keys") + for name, value in layers.items(): + groups.setdefault((value.device, value.dtype), []).append((name, value)) + packed = [] + for group in groups.values(): + values = torch.cat( + [value.reshape(len(keys), -1) for _, value in group], dim=1 + ) + names = [name for name, _ in group] + shapes = [value.shape[1:] for _, value in group] + sizes = [value[0].numel() for _, value in group] + packed.append((names, shapes, sizes, values)) + dtypes = {name: value.dtype for name, value in layers.items()} + pinned = set().union(*self.leases.values()) if self.leases else set() + for index, key in enumerate(keys): + if not self._admit(key, pinned): + continue + try: + target = ( + self.gpu_device if self.tiers[key] == "gpu" else torch.device("cpu") + ) + content: Dict[str, torch.Tensor] = {} + for names, shapes, sizes, values in packed: + # Layers share only this block's allocation. Evicting one + # block cannot retain the rest of the staging batch. + block = values[index].detach().to(target, copy=True) + parts = ( + value.reshape(shape) + for value, shape in zip(block.split(sizes), shapes) + ) + content.update(zip(names, parts)) + self.blocks[key] = content + self.logical_dtypes[key] = dtypes.copy() + self._block_sizes[key] = nbytes + except Exception: + self._drop(key) + raise + def reserve(self, lease: str, keys: List[int]) -> bool: if not super().reserve(lease, keys): return False