[Bugfix][KVConnector] Fix SimpleCPUOffloadConnector GPU->CPU store race

tests/v1/simple_kv_offload/test_worker.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worker-side unit tests for SimpleCPUOffloadConnector.
+
+Covers the GPU->CPU store cross-stream synchronization: the store copy must be
+ordered after the compute stream that writes the KV blocks, otherwise it can
+read partially written / stale blocks and silently corrupt the CPU cache.
+"""
+
+from __future__ import annotations
+
+import time
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_cuda_alike():
+    pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
+
+from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
+from vllm.v1.simple_kv_offload.cuda_mem_ops import (
+    CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
+    CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
+    build_params,
+    pin_tensor,
+)
+from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata
+from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker
+
+NUM_BLOCKS = 64
+BLOCK_BYTES = 4096
+ITERS = 30
+# Keep the compute stream busy so the KV write lands late; this makes the
+# store-vs-compute race deterministic instead of timing-dependent.
+SLEEP_CYCLES = 50_000_000
+
+
+def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:
+    gpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cuda")}
+    cpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu")}
+    pin_tensor(cpu["k"])
+    low_pri, _ = torch.cuda.Stream.priority_range()
+    backend = DmaCopyBackend()
+    backend.init(
+        gpu,
+        cpu,
+        gpu["k"].device,
+        torch.cuda.Stream(priority=low_pri),
+        torch.cuda.Stream(priority=low_pri),
+    )
+    return backend, gpu["k"], cpu["k"]
+
+
+def _drive_store(
+    backend: DmaCopyBackend,
+    gpu: torch.Tensor,
+    cpu: torch.Tensor,
+    *,
+    with_barrier: bool,
+) -> int:
+    """Run ITERS store cycles; return how many landed corrupted in the CPU pool.
+
+    Each cycle writes a unique value on a compute stream (after a deliberate
+    delay) and then issues the GPU->CPU store. The store is issued *after* the
+    write in host program order, mirroring the connector's deferred-store
+    assumption. Only the compute-done event creates a real device-side
+    happens-before edge.
+    """
+    block_ids = list(range(gpu.shape[0]))
+    compute_stream = torch.cuda.Stream()
+    corrupt = 0
+    for it in range(ITERS):
+        val = (it % 126) + 1  # 1..126; distinct from the zero-initialized pool
+        with torch.cuda.stream(compute_stream):
+            torch.cuda._sleep(SLEEP_CYCLES)
+            gpu.fill_(val)
+
+        wait_event = None
+        if with_barrier:
+            wait_event = torch.Event()
+            wait_event.record(compute_stream)
+
+        store_events: list[tuple[int, torch.Event]] = []
+        backend.launch_copy(
+            block_ids,
+            block_ids,
+            is_store=True,
+            event_idx=it,
+            events_list=store_events,
+            wait_event=wait_event,
+        )
+
+        deadline = time.time() + 10.0
+        while not store_events and time.time() < deadline:
+            time.sleep(0.0005)
+        assert store_events, "background copy was never enqueued"
+        store_events[0][1].synchronize()
+
+        if int((cpu[:, 0].to(torch.int32) != val).sum().item()):
+            corrupt += 1
+    return corrupt
+
+
+def test_store_orders_after_compute_write():
+    """The store must wait for the compute event; without it, it races.
+
+    Asserts both directions so the test is self-validating: the no-barrier
+    control must actually corrupt (proving the race window is exercised), and
+    the fixed path with the compute-done event must be clean.
+    """
+    backend, gpu, cpu = _make_backend()
+    try:
+        control = _drive_store(backend, gpu, cpu, with_barrier=False)
+        fixed = _drive_store(backend, gpu, cpu, with_barrier=True)
+    finally:
+        backend.shutdown()
+
+    assert control > 0, (
+        "no-barrier store did not race the compute write; the test no longer "
+        "exercises the hazard it is meant to guard"
+    )
+    assert fixed == 0, f"store raced compute even with the barrier: {fixed} corrupt"
+
+
+class _RecordingBackend:
+    """Captures launch_copy calls without touching the GPU."""
+
+    def __init__(self) -> None:
+        self.calls: list[dict] = []
+
+    def launch_copy(
+        self,
+        src_blocks,
+        dst_blocks,
+        is_store,
+        event_idx,
+        events_list,
+        wait_event=None,
+    ) -> None:
+        self.calls.append({"is_store": is_store, "wait_event": wait_event})
+
+
+def test_get_finished_passes_wait_event_for_store_only():
+    """get_finished gates stores on a compute-done event but not loads."""
+    worker = SimpleCPUOffloadWorker(
+        vllm_config=None, kv_cache_config=None, cpu_capacity_bytes=0
+    )
+    recording = _RecordingBackend()
+    worker._backend = recording
+    worker._connector_metadata = SimpleCPUOffloadMetadata(
+        load_event=0,
+        load_gpu_blocks=[0],
+        load_cpu_blocks=[0],
+        store_event=1,
+        store_gpu_blocks=[1],
+        store_cpu_blocks=[1],
+    )
+
+    worker.get_finished(set())
+
+    store_calls = [c for c in recording.calls if c["is_store"]]
+    load_calls = [c for c in recording.calls if not c["is_store"]]
+    assert len(store_calls) == 1
+    assert len(load_calls) == 1
+    assert isinstance(store_calls[0]["wait_event"], torch.Event)
+    assert load_calls[0]["wait_event"] is None
+
+
+def test_build_params_src_access_order():
+    """build_params defaults to ANY and honors an explicit STREAM override."""
+    gpu = {"k": torch.zeros((4, 64), dtype=torch.int8, device="cuda")}
+    cpu = {"k": torch.zeros((4, 64), dtype=torch.int8, device="cpu")}
+    stream = torch.cuda.Stream()
+
+    default = build_params(gpu, cpu, stream)
+    assert default.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_ANY
+
+    ordered = build_params(
+        gpu, cpu, stream, src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
+    )
+    assert ordered.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_STREAM

vllm/v1/simple_kv_offload/copy_backend.py

@@ -12,6 +12,8 @@
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.v1.simple_kv_offload.cuda_mem_ops import (
+    CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
+    CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
     BatchMemcpyParams,
     build_params,
     copy_blocks,
@@ -43,8 +45,22 @@ def init(
         self._load_stream = load_stream
         self._store_stream = store_stream
 
-        self._store_params = build_params(gpu_caches, cpu_caches, store_stream)
-        self._load_params = build_params(cpu_caches, gpu_caches, load_stream)
+        # Store (GPU->CPU): source is the live GPU KV cache, still written by
+        # the compute stream -> STREAM ordering (paired with a wait on the
+        # compute-done event in launch_copy). Load (CPU->GPU): source is stable
+        # pinned host memory -> ANY is safe and lets the driver pipeline reads.
+        self._store_params = build_params(
+            gpu_caches,
+            cpu_caches,
+            store_stream,
+            src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
+        )
+        self._load_params = build_params(
+            cpu_caches,
+            gpu_caches,
+            load_stream,
+            src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
+        )
 
         self._queue = queue.SimpleQueue()
         self._thread = threading.Thread(
@@ -61,11 +77,20 @@ def launch_copy(
         is_store: bool,
         event_idx: int,
         events_list: list[tuple[int, torch.Event]],
+        wait_event: torch.Event | None = None,
     ) -> None:
         params = self._store_params if is_store else self._load_params
         assert params is not None and self._queue is not None
         self._queue.put(
-            (src_blocks, dst_blocks, params, is_store, event_idx, events_list)
+            (
+                src_blocks,
+                dst_blocks,
+                params,
+                is_store,
+                event_idx,
+                events_list,
+                wait_event,
+            )
         )
 
     def shutdown(self) -> None:
@@ -89,9 +114,23 @@ def _copy_loop(
             item = q.get()
             if item is None:
                 return
-            src_blocks, dst_blocks, params, is_store, event_idx, events_list = item
-            copy_blocks(src_blocks, dst_blocks, params)
+            (
+                src_blocks,
+                dst_blocks,
+                params,
+                is_store,
+                event_idx,
+                events_list,
+                wait_event,
+            ) = item
             stream = store_stream if is_store else load_stream
+            # Gate the copy on the compute-done event so the GPU->CPU store
+            # cannot read KV blocks before the kernels that wrote them have
+            # retired on the compute stream. The wait is enqueued on the
+            # transfer stream before the copy, so the two are ordered on-device.
+            if wait_event is not None:
+                stream.wait_event(wait_event)
+            copy_blocks(src_blocks, dst_blocks, params)
             event = torch.Event()
             event.record(stream)
             events_list.append((event_idx, event))

vllm/v1/simple_kv_offload/cuda_mem_ops.py

@@ -13,6 +13,14 @@
 
 logger = init_logger(__name__)
 
+# CUmemcpySrcAccessOrder values (CUDA driver API).
+#   STREAM (1): source reads are ordered after prior writes on the stream the
+#       copy runs on -> safe when the source may still be written.
+#   ANY (3): driver may read the source before prior writes complete -> only
+#       safe when the source is guaranteed stable (e.g. pinned host memory).
+CU_MEMCPY_SRC_ACCESS_ORDER_STREAM = 1
+CU_MEMCPY_SRC_ACCESS_ORDER_ANY = 3
+
 
 def pin_tensor(tensor: torch.Tensor) -> None:
     """Pin a CPU tensor via cudaHostRegister.
@@ -120,6 +128,7 @@ def build_params(
     src_caches: dict[str, torch.Tensor],
     dst_caches: dict[str, torch.Tensor],
     stream: torch.cuda.Stream,
+    src_access_order: int = CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
 ) -> BatchMemcpyParams:
     global _batch_memcpy_fn
     if _batch_memcpy_fn is None:
@@ -137,10 +146,12 @@ def build_params(
         dst_bases.append(d.data_ptr())
         bpb.append(s_bpb)
 
-    # ``srcAccessOrder=3`` == CU_MEMCPY_SRC_ACCESS_ORDER_ANY /
-    # hipMemcpySrcAccessOrderAny. See
+    # srcAccessOrder controls whether the DMA may read the source before prior
+    # writes to it complete. STREAM(1) for GPU->CPU stores (live KV cache is
+    # still being written by the compute stream); ANY(3) for CPU->GPU loads
+    # (source is stable pinned host memory). See
     # https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f  # noqa: E501
-    attrs = _CUmemcpyAttributes(srcAccessOrder=3)
+    attrs = _CUmemcpyAttributes(srcAccessOrder=src_access_order)
 
     return BatchMemcpyParams(
         src_bases=np.array(src_bases, dtype=np.uint64),

vllm/v1/simple_kv_offload/worker.py

@@ -207,8 +207,13 @@ def get_finished(
         """Submit transfers and report completed events to the scheduler.
 
         Called after model execution. The manager only schedules stores for
-        blocks whose KV data is confirmed computed, so we launch both loads
-        and stores immediately — no deferral or cross-stream sync needed.
+        blocks whose KV data is confirmed computed on the host (scheduler
+        bookkeeping). That is NOT a device-side guarantee: under v1 overlapped
+        execution the host runs ahead of the GPU and the KV-write kernels may
+        still be in flight on the compute stream. Loads read from stable pinned
+        host memory and are safe to launch immediately; GPU->CPU stores read
+        the live KV cache and must be ordered after the compute stream, so we
+        record a compute-done event here and have the store stream wait on it.
 
         Returns:
             tuple of (finished_sending, finished_recving).
@@ -218,7 +223,9 @@ def get_finished(
         # (1) Submit transfers
         metadata = self._connector_metadata
         if metadata is not None:
-            # Launch loads (CPU->GPU).
+            # Launch loads (CPU->GPU). Source is stable pinned host memory, so
+            # no cross-stream barrier is needed; load->use is gated separately
+            # by the load event.
             if metadata.load_cpu_blocks:
                 self._backend.launch_copy(
                     metadata.load_cpu_blocks,
@@ -227,14 +234,21 @@ def get_finished(
                     event_idx=metadata.load_event,
                     events_list=self._load_events,
                 )
-            # Launch stores (GPU->CPU).
+            # Launch stores (GPU->CPU). Record a compute-done event on the
+            # current (compute) stream so the store waits until every kernel
+            # enqueued so far — including the KV writes for the blocks being
+            # stored — has retired. Without this the store can read a partially
+            # written / stale block, silently corrupting the CPU cache.
             if metadata.store_gpu_blocks:
+                compute_done = torch.Event()
+                compute_done.record(torch.cuda.current_stream())
                 self._backend.launch_copy(
                     metadata.store_gpu_blocks,
                     metadata.store_cpu_blocks,
                     is_store=True,
                     event_idx=metadata.store_event,
                     events_list=self._store_events,
+                    wait_event=compute_done,
                 )
 
         # (2) Track completed transfer events