diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py new file mode 100644 index 000000000000..859d0fecd581 --- /dev/null +++ b/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 diff --git a/vllm/v1/simple_kv_offload/copy_backend.py b/vllm/v1/simple_kv_offload/copy_backend.py index 114f26973767..58de7a7e9eff 100644 --- a/vllm/v1/simple_kv_offload/copy_backend.py +++ b/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,20 @@ 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) + # Stores read the live KV cache -> STREAM (paired with the compute-done + # wait in get_finished); loads read stable pinned host memory -> ANY. + 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 +75,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 +112,19 @@ 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 + 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)) diff --git a/vllm/v1/simple_kv_offload/cuda_mem_ops.py b/vllm/v1/simple_kv_offload/cuda_mem_ops.py index b4c68aff3ca9..69b1677e0ac9 100644 --- a/vllm/v1/simple_kv_offload/cuda_mem_ops.py +++ b/vllm/v1/simple_kv_offload/cuda_mem_ops.py @@ -13,6 +13,12 @@ logger = init_logger(__name__) +# CUmemcpySrcAccessOrder values (CUDA driver API). STREAM(1): source read in +# stream order, safe when the source may still be written. ANY(3): source may +# be read early, only safe for a stable source (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. @@ -106,8 +112,8 @@ class BatchMemcpyParams(NamedTuple): dst_bases: np.ndarray # [num_layers] uint64 bpb: np.ndarray # [num_layers] uint64 — bytes per block num_layers: int - # CUDA only: one attributes entry with srcAccessOrder=ANY. Unused on - # ROCm (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0. + # CUDA only: one attributes entry carrying srcAccessOrder. Unused on ROCm + # (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0. attrs: _CUmemcpyAttributes attrs_idx: ctypes.c_size_t # NOTE: cuMemcpyBatchAsync_v2() removed fail_idx field, but we use @@ -120,6 +126,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 +144,7 @@ def build_params( dst_bases.append(d.data_ptr()) bpb.append(s_bpb) - # ``srcAccessOrder=3`` == CU_MEMCPY_SRC_ACCESS_ORDER_ANY / - # hipMemcpySrcAccessOrderAny. 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), diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index d33e5f762049..9cb9c02ed7c5 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -57,6 +57,10 @@ def __init__( # Metadata for the current step self._connector_metadata: SimpleCPUOffloadMetadata | None = None + # Compute-done event recorded before each store; reused across steps + # (get_finished runs once per step, copy queue is FIFO). + self._store_compute_done: torch.Event | None = None + # Pending event index sets, populated in bind_connector_metadata self._pending_load_event_indices: set[int] = set() self._pending_store_event_indices: set[int] = set() @@ -206,9 +210,11 @@ def get_finished( ) -> tuple[set[str] | None, set[str] | None]: """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. + Stores (GPU->CPU) read the live KV cache, which the compute stream may + still be writing under v1 overlapped execution, so they are ordered + after a compute-done event recorded on the current stream. Loads + (CPU->GPU) read stable pinned host memory and launch immediately. See + #45704 for the bug and #39306 for the srcAccessOrder rationale. Returns: tuple of (finished_sending, finished_recving). @@ -218,7 +224,6 @@ def get_finished( # (1) Submit transfers metadata = self._connector_metadata if metadata is not None: - # Launch loads (CPU->GPU). if metadata.load_cpu_blocks: self._backend.launch_copy( metadata.load_cpu_blocks, @@ -227,14 +232,17 @@ def get_finished( event_idx=metadata.load_event, events_list=self._load_events, ) - # Launch stores (GPU->CPU). if metadata.store_gpu_blocks: + if self._store_compute_done is None: + self._store_compute_done = torch.Event() + self._store_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=self._store_compute_done, ) # (2) Track completed transfer events