Backport engine-driven async-copy lifetime synchronization

lmcache/v1/multiprocess/transfer_context/base.py

@@ -739,6 +739,12 @@ def scatter_cpu_to_paged_kv(
         ValueError: If ``block_ids`` is shorter than
             ``len(chunks) * blocks_per_chunk``, or group layers disagree on
             physical block stride.
+
+    Notes:
+        Raw-pointer transfers from locally pinned temporary buffers complete
+        before this function returns or propagates a transfer exception.
+        Caller-owned pinned buffers remain asynchronous: the caller must keep
+        them alive and synchronize before consuming the destination KV.
     """
     # First Party
     from lmcache import device_ops
@@ -823,6 +829,9 @@ def scatter_cpu_to_paged_kv(
     if not selected_block_ids:
         return
 
+    # Only the ptr-only branch below pins temporaries; the tensor branch
+    # hands torch the chunks directly and keeps them alive itself.
+    dynamically_pinned = False
     if _LMC_OPS_BLOCK_TRANSFER_ACCEPTS_TENSOR:
         # Python fallback: accepts tensor list directly for all params.
         paged_arg = normalized
@@ -846,8 +855,8 @@ def scatter_cpu_to_paged_kv(
         # Defensive check: Ensure all incoming CPU chunks are pinned memory.
         # Otherwise, the underlying CUDA kernel may throw an Illegal
         # Memory Access error during H2D transfer.
-        pinned_copy = False
-        if not all(chunk.is_pinned() for chunk in chunks):
+        dynamically_pinned = not all(chunk.is_pinned() for chunk in chunks)
+        if dynamically_pinned:
             logger.warning(
                 "Received unpinned CPU tensors in scatter_cpu_to_paged_kv. "
                 "Dynamically pinning memory now, which may incur additional"
@@ -857,7 +866,6 @@ def scatter_cpu_to_paged_kv(
                 chunk.pin_memory() if not chunk.is_pinned() else chunk
                 for chunk in chunks
             ]
-            pinned_copy = True
 
         # Compiled C++/CUDA/XPU: requires int64 pointer tensor and list[int].
         _ptrs_np = np.array(
@@ -875,35 +883,41 @@ def scatter_cpu_to_paged_kv(
         req_blocks_per_obj = bpw
         total_chunks = len(chunks)
 
-        for i in range(0, total_chunks, MAX_OBJECTS):
-            # Slice objects and block IDs for this batch
-            batch_objs_ptrs = objs_arg[i : i + MAX_OBJECTS]
+        h2d_started = False
+        try:
+            for i in range(0, total_chunks, MAX_OBJECTS):
+                # Slice objects and block IDs for this batch
+                batch_objs_ptrs = objs_arg[i : i + MAX_OBJECTS]
 
-            start_block = i * req_blocks_per_obj
-            end_block = min(
-                (i + MAX_OBJECTS) * req_blocks_per_obj, len(selected_block_ids)
-            )
-            batch_blocks = block_ids_arg[start_block:end_block]
-            batch_skip = max(0, skip_prefix_window - start_block)
-            if batch_skip >= len(batch_blocks):
-                continue
+                start_block = i * req_blocks_per_obj
+                end_block = min(
+                    (i + MAX_OBJECTS) * req_blocks_per_obj, len(selected_block_ids)
+                )
+                batch_blocks = block_ids_arg[start_block:end_block]
+                batch_skip = max(0, skip_prefix_window - start_block)
+                if batch_skip >= len(batch_blocks):
+                    continue
 
-            # Execute transfer for this batch
-            device_ops.multi_layer_block_kv_transfer(
-                paged_arg,
-                batch_objs_ptrs,
-                batch_blocks,
-                get_device(normalized),
-                lmcache_native.TransferDirection.H2D,
-                shape_desc,
-                window_tokens,
-                engine_kv_format,
-                batch_skip,
-            )
-        # Dynamically pinned tensors are local temporaries whose data pointers
-        # must remain valid until the asynchronous transfer has completed.
-        if pinned_copy:
-            torch_dev.synchronize()
+                # A native call can enqueue a copy before raising.
+                h2d_started = True
+                device_ops.multi_layer_block_kv_transfer(
+                    paged_arg,
+                    batch_objs_ptrs,
+                    batch_blocks,
+                    get_device(normalized),
+                    lmcache_native.TransferDirection.H2D,
+                    shape_desc,
+                    window_tokens,
+                    engine_kv_format,
+                    batch_skip,
+                )
+        finally:
+            if dynamically_pinned and h2d_started:
+                # Torch cannot track the source lifetime through raw pointers.
+                # Drain copies before releasing local temporaries, including
+                # when a subsequent batch raises. Caller-owned pinned buffers
+                # retain their asynchronous lifetime contract.
+                torch_dev.synchronize()
     # Fast path: The async GPU copy might still be in progress.
     # We intentionally omit synchronization here for performance.
     # WARNING: The caller MUST explicitly call `torch_dev.synchronize()`

lmcache/v1/multiprocess/transfer_context/worker_transfer.py

@@ -1226,6 +1226,7 @@ def submit_store(
         if self._worker_groups:
             torch_dev.synchronize()
             group_out_buffers: list[list[torch.Tensor]] | None = None
+            gather_started = False
             prepared = False
             try:
                 group_result = self._engine_driven_context.prepare_store_grouped(
@@ -1239,24 +1240,28 @@ def submit_store(
                     group_future: MessagingFuture[bool] = MessagingFuture()
                     group_future.set_result(True)
                     return group_future
+                gather_started = True
                 cpu_groups = self._gather_group_payloads(
                     kv_caches,
                     block_ids,
                     out_buffers=group_out_buffers,
                     group_chunk_indices=group_chunk_indices,
                 )
-                if group_out_buffers is not None:
-                    torch_dev.synchronize()
+                # Both SHM and pickle transports consume data produced by an
+                # asynchronous device-to-host gather. Commit only after every
+                # group payload is complete.
+                torch_dev.synchronize()
                 ok = self._engine_driven_context.commit_store_grouped(
                     key, instance_id, cpu_groups
                 )
                 if not ok:
                     self._abort_store_safely(key, instance_id)
             except Exception:
                 logger.exception("Failed to store engine-driven hybrid chunks")
-                if group_out_buffers is not None:
+                if gather_started:
                     # A failed gather can leave asynchronous writes targeting
-                    # SHM views. Drain them before the server frees the slots.
+                    # SHM views or temporary pickle buffers. Drain them before
+                    # either storage is released.
                     torch_dev.synchronize()
                 if prepared:
                     self._abort_store_safely(key, instance_id)
@@ -1267,6 +1272,7 @@ def submit_store(
 
         torch_dev.synchronize()
         out_buffers: list[torch.Tensor] | None = None
+        gather_started = False
         prepared = False
         try:
             legacy_result = self._engine_driven_context.prepare_store(key, instance_id)
@@ -1279,6 +1285,7 @@ def submit_store(
                 future: MessagingFuture[bool] = MessagingFuture()
                 future.set_result(True)
                 return future
+            gather_started = True
             cpu_chunks = gather_paged_kv_to_cpu(
                 kv_caches,
                 _single_group_block_ids(block_ids),
@@ -1288,20 +1295,19 @@ def submit_store(
                 out=out_buffers,
                 chunk_indices=chunk_indices,
             )
-            if out_buffers is not None:
-                # SHM path uses async device->CPU copies; complete them before commit.
-                torch_dev.synchronize()
+            # Both SHM and pickle transports consume data produced by an
+            # asynchronous device-to-host gather.
+            torch_dev.synchronize()
             ok = self._engine_driven_context.commit_store(key, instance_id, cpu_chunks)
             if not ok:
                 self._abort_store_safely(key, instance_id)
         except Exception:
             logger.exception("Failed to store engine-driven chunks")
-            if out_buffers is not None:
+            if gather_started:
                 torch_dev.synchronize()
             if prepared:
                 self._abort_store_safely(key, instance_id)
             ok = False
-
         future = MessagingFuture()
         future.set_result(ok)
         return future

tests/v1/multiprocess/test_engine_driven_async_copy_lifetime.py

@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Async device<->CPU copies must complete before their buffers are reused.
+
+Pickle serialization must wait for gathered data, and raw-pointer scatter must
+retain locally pinned buffers until its copies complete, including on errors.
+Caller-owned pinned buffers keep their asynchronous transfer contract.
+
+These tests assert the ordering contract without a GPU, so they run in the
+CPU-only unit CI. The hardware reproductions live alongside them in
+``test_engine_driven_transfer.py`` and skip without CUDA.
+"""
+
+# Standard
+from contextlib import nullcontext
+from typing import cast
+from unittest.mock import MagicMock, patch
+
+# Third Party
+import pytest
+import torch
+
+
+def test_scatter_syncs_before_releasing_dynamically_pinned_chunks() -> None:
+    """Unpinned input must be synced before scatter returns.
+
+    ``scatter_cpu_to_paged_kv`` pins unpinned chunks into temporaries and
+    launches async H2D reads on them through raw pointers, which torch's
+    stream tracking cannot see. Returning drops the last reference, so the
+    caching host allocator can hand that memory to the next caller while the
+    copies are in flight. The documented caller-side synchronize cannot cover
+    it -- by then the temporaries are gone -- so scatter must sync itself.
+    """
+    # First Party
+    from lmcache.v1.multiprocess.transfer_context import base
+
+    kv = {f"layer_{i}": torch.zeros(2, 4, 4, 2, 8) for i in range(2)}
+
+    # Mock chunks, not real tensors: pin_memory() needs an accelerator, so the
+    # ptr-only branch cannot execute for real on the CPU-only unit CI. What we
+    # are pinning down is the ordering contract, not the copy itself.
+    def _unpinned_chunk() -> MagicMock:
+        c = MagicMock()
+        c.is_pinned.return_value = False
+        c.pin_memory.return_value = c
+        c.data_ptr.return_value = 0
+        return c
+
+    chunks = [_unpinned_chunk()]
+
+    # Pin the ptr-only path: only that branch pins temporaries, and whether the
+    # compiled op takes tensors varies by build. scatter imports device_ops
+    # inside the function, so patch it at source.
+    with (
+        patch.object(base, "_LMC_OPS_BLOCK_TRANSFER_ACCEPTS_TENSOR", False),
+        patch.object(base, "torch_dev") as dev,
+        patch("lmcache.device_ops") as ops,
+    ):
+        # cast: the mocks stand in for tensors on purpose (see above).
+        base.scatter_cpu_to_paged_kv(
+            kv, list(range(4)), cast(list[torch.Tensor], chunks), 4
+        )
+        assert ops.multi_layer_block_kv_transfer.called, (
+            "fixture must reach the async H2D launches"
+        )
+        assert dev.synchronize.called, (
+            "scatter must complete async H2D before releasing the temporaries "
+            "it pinned; otherwise the host allocator reuses them mid-copy"
+        )
+
+
+@pytest.mark.parametrize("pinned", [False, True])
+@pytest.mark.parametrize("fail_batch", [None, 1, 2])
+def test_scatter_drains_owned_temporaries_when_a_transfer_fails(
+    pinned: bool, fail_batch: int | None
+) -> None:
+    """Raw-pointer transfers cannot outlive locally pinned source buffers."""
+    # First Party
+    from lmcache.v1.multiprocess.transfer_context import base
+
+    kv = {f"layer_{i}": torch.zeros(2, 4, 4, 2, 8) for i in range(2)}
+    chunks = [MagicMock() for _ in range(5)]
+    for chunk in chunks:
+        chunk.is_pinned.return_value = pinned
+        chunk.pin_memory.return_value = chunk
+        chunk.data_ptr.return_value = 0
+    order: list[str] = []
+    launches = 0
+
+    def transfer(*_args: object, **_kwargs: object) -> None:
+        nonlocal launches
+        launches += 1
+        order.append("launch")
+        if launches == fail_batch:
+            raise RuntimeError("transfer batch failed")
+
+    with (
+        patch.object(base, "_LMC_OPS_BLOCK_TRANSFER_ACCEPTS_TENSOR", False),
+        patch.object(base, "torch_dev") as dev,
+        patch("lmcache.device_ops") as ops,
+    ):
+        ops.multi_layer_block_kv_transfer.side_effect = transfer
+        dev.synchronize.side_effect = lambda: order.append("sync")
+        expected_error = (
+            pytest.raises(RuntimeError, match="transfer batch failed")
+            if fail_batch is not None
+            else nullcontext()
+        )
+        with expected_error:
+            base.scatter_cpu_to_paged_kv(
+                kv, [0, 1, 2, 3] * 5, cast(list[torch.Tensor], chunks), 4
+            )
+        order.append("caller")
+
+    expected_launches = fail_batch if fail_batch is not None else 2
+    assert order == (
+        ["launch"] * expected_launches + ([] if pinned else ["sync"]) + ["caller"]
+    )
+
+
+def test_scatter_does_not_sync_when_every_block_is_skipped() -> None:
+    """An entirely cached prefix does not launch or synchronize a transfer."""
+    # First Party
+    from lmcache.v1.multiprocess.transfer_context import base
+
+    kv = {f"layer_{i}": torch.zeros(2, 4, 4, 2, 8) for i in range(2)}
+    chunk = MagicMock()
+    chunk.is_pinned.return_value = False
+    chunk.pin_memory.return_value = chunk
+    chunk.data_ptr.return_value = 0
+
+    with (
+        patch.object(base, "_LMC_OPS_BLOCK_TRANSFER_ACCEPTS_TENSOR", False),
+        patch.object(base, "torch_dev") as dev,
+        patch("lmcache.device_ops") as ops,
+    ):
+        base.scatter_cpu_to_paged_kv(
+            kv,
+            list(range(4)),
+            cast(list[torch.Tensor], [chunk]),
+            4,
+            skip_first_n_tokens=16,
+        )
+        ops.multi_layer_block_kv_transfer.assert_not_called()
+        dev.synchronize.assert_not_called()
+
+
+def test_pickle_store_syncs_before_commit_serializes() -> None:
+    """The pickle path must sync before commit_store reads the buffers.
+
+    Gather issues async device->CPU copies into fresh buffers, and the pickle
+    transport serializes them immediately in ``commit_store``. Syncing only
+    when ``out_buffers`` is given (the SHM path) leaves pickle serializing a
+    buffer that is still being written.
+    """
+    # First Party
+    from lmcache.v1.multiprocess.transfer_context import worker_transfer
+
+    order: list[str] = []
+    ctx = worker_transfer.EngineDrivenTransferContext()
+    ctx._engine_driven_context = MagicMock()
+    ctx._engine_driven_context.prepare_store.return_value = None  # pickle mode
+
+    def _commit(*_a: object, **_k: object) -> bool:
+        order.append("commit")
+        return True
+
+    ctx._engine_driven_context.commit_store.side_effect = _commit
+    ctx._layout_hints = None
+    ctx._engine_kv_format = None
+
+    def _gather(*_a: object, **_k: object) -> list[torch.Tensor]:
+        order.append("gather")
+        return [torch.zeros(1)]
+
+    with (
+        patch.object(worker_transfer, "torch_dev") as dev,
+        patch.object(worker_transfer, "gather_paged_kv_to_cpu", side_effect=_gather),
+    ):
+        dev.synchronize.side_effect = lambda *a, **k: order.append("sync")
+        ctx.submit_store(
+            "req",
+            MagicMock(),  # key
+            1,  # instance_id
+            {"layer_0": torch.zeros(2, 4, 4, 2, 8)},
+            [[0, 1, 2, 3]],
+            MagicMock(),  # event (unused on this transport)
+            4,  # blocks_in_chunk
+        )
+
+    # A sync must fall BETWEEN gather and commit. submit_store also syncs
+    # before prepare_store, so merely finding a "sync" proves nothing -- that
+    # earlier one is why guarding this on out_buffers went unnoticed.
+    gathered, committed = order.index("gather"), order.index("commit")
+    assert any(
+        i for i, step in enumerate(order) if step == "sync" and gathered < i < committed
+    ), (
+        "pickle store must synchronize after gather and before commit_store "
+        f"serializes the buffers, got {order}"
+    )