perf: reduce Xavier cold-prefill snapshot overhead

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:

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)