diff --git a/lmcache/v1/multiprocess/transfer_context/base.py b/lmcache/v1/multiprocess/transfer_context/base.py index bd359cb1e3a..0a91a1d0d58 100644 --- a/lmcache/v1/multiprocess/transfer_context/base.py +++ b/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()` diff --git a/lmcache/v1/multiprocess/transfer_context/worker_transfer.py b/lmcache/v1/multiprocess/transfer_context/worker_transfer.py index 0f0830c467e..abf804006d4 100644 --- a/lmcache/v1/multiprocess/transfer_context/worker_transfer.py +++ b/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,14 +1240,17 @@ 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 ) @@ -1254,9 +1258,10 @@ def submit_store( 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 diff --git a/tests/v1/multiprocess/test_engine_driven_async_copy_lifetime.py b/tests/v1/multiprocess/test_engine_driven_async_copy_lifetime.py new file mode 100644 index 00000000000..46e8154cd82 --- /dev/null +++ b/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}" + )