@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 82ede09a5a

[Bugfix][KVConnector] Fix SimpleCPUOffloadConnector GPU->CPU store race (#46278)

Saddss committed Jun 23, 2026 at 04:08 UTC 82ede09a5a26c3b529a1e13adb68e5ecc0558790
4 files changed +244 -16
tests/v1/simple_kv_offload/test_worker.py new
+183
@@ -0,0 +1,183 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Worker-side unit tests for SimpleCPUOffloadConnector.
4 +
5 +Covers the GPU->CPU store cross-stream synchronization: the store copy must be
6 +ordered after the compute stream that writes the KV blocks, otherwise it can
7 +read partially written / stale blocks and silently corrupt the CPU cache.
8 +"""
9 +
10 +from __future__ import annotations
11 +
12 +import time
13 +
14 +import pytest
15 +import torch
16 +
17 +from vllm.platforms import current_platform
18 +
19 +if not current_platform.is_cuda_alike():
20 + pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
21 +
22 +from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
23 +from vllm.v1.simple_kv_offload.cuda_mem_ops import (
24 + CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
25 + CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
26 + build_params,
27 + pin_tensor,
28 +)
29 +from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata
30 +from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker
31 +
32 +NUM_BLOCKS = 64
33 +BLOCK_BYTES = 4096
34 +ITERS = 30
35 +# Keep the compute stream busy so the KV write lands late; this makes the
36 +# store-vs-compute race deterministic instead of timing-dependent.
37 +SLEEP_CYCLES = 50_000_000
38 +
39 +
40 +def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:
41 + gpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cuda")}
42 + cpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu")}
43 + pin_tensor(cpu["k"])
44 + low_pri, _ = torch.cuda.Stream.priority_range()
45 + backend = DmaCopyBackend()
46 + backend.init(
47 + gpu,
48 + cpu,
49 + gpu["k"].device,
50 + torch.cuda.Stream(priority=low_pri),
51 + torch.cuda.Stream(priority=low_pri),
52 + )
53 + return backend, gpu["k"], cpu["k"]
54 +
55 +
56 +def _drive_store(
57 + backend: DmaCopyBackend,
58 + gpu: torch.Tensor,
59 + cpu: torch.Tensor,
60 + *,
61 + with_barrier: bool,
62 +) -> int:
63 + """Run ITERS store cycles; return how many landed corrupted in the CPU pool.
64 +
65 + Each cycle writes a unique value on a compute stream (after a deliberate
66 + delay) and then issues the GPU->CPU store. The store is issued *after* the
67 + write in host program order, mirroring the connector's deferred-store
68 + assumption. Only the compute-done event creates a real device-side
69 + happens-before edge.
70 + """
71 + block_ids = list(range(gpu.shape[0]))
72 + compute_stream = torch.cuda.Stream()
73 + corrupt = 0
74 + for it in range(ITERS):
75 + val = (it % 126) + 1 # 1..126; distinct from the zero-initialized pool
76 + with torch.cuda.stream(compute_stream):
77 + torch.cuda._sleep(SLEEP_CYCLES)
78 + gpu.fill_(val)
79 +
80 + wait_event = None
81 + if with_barrier:
82 + wait_event = torch.Event()
83 + wait_event.record(compute_stream)
84 +
85 + store_events: list[tuple[int, torch.Event]] = []
86 + backend.launch_copy(
87 + block_ids,
88 + block_ids,
89 + is_store=True,
90 + event_idx=it,
91 + events_list=store_events,
92 + wait_event=wait_event,
93 + )
94 +
95 + deadline = time.time() + 10.0
96 + while not store_events and time.time() < deadline:
97 + time.sleep(0.0005)
98 + assert store_events, "background copy was never enqueued"
99 + store_events[0][1].synchronize()
100 +
101 + if int((cpu[:, 0].to(torch.int32) != val).sum().item()):
102 + corrupt += 1
103 + return corrupt
104 +
105 +
106 +def test_store_orders_after_compute_write():
107 + """The store must wait for the compute event; without it, it races.
108 +
109 + Asserts both directions so the test is self-validating: the no-barrier
110 + control must actually corrupt (proving the race window is exercised), and
111 + the fixed path with the compute-done event must be clean.
112 + """
113 + backend, gpu, cpu = _make_backend()
114 + try:
115 + control = _drive_store(backend, gpu, cpu, with_barrier=False)
116 + fixed = _drive_store(backend, gpu, cpu, with_barrier=True)
117 + finally:
118 + backend.shutdown()
119 +
120 + assert control > 0, (
121 + "no-barrier store did not race the compute write; the test no longer "
122 + "exercises the hazard it is meant to guard"
123 + )
124 + assert fixed == 0, f"store raced compute even with the barrier: {fixed} corrupt"
125 +
126 +
127 +class _RecordingBackend:
128 + """Captures launch_copy calls without touching the GPU."""
129 +
130 + def __init__(self) -> None:
131 + self.calls: list[dict] = []
132 +
133 + def launch_copy(
134 + self,
135 + src_blocks,
136 + dst_blocks,
137 + is_store,
138 + event_idx,
139 + events_list,
140 + wait_event=None,
141 + ) -> None:
142 + self.calls.append({"is_store": is_store, "wait_event": wait_event})
143 +
144 +
145 +def test_get_finished_passes_wait_event_for_store_only():
146 + """get_finished gates stores on a compute-done event but not loads."""
147 + worker = SimpleCPUOffloadWorker(
148 + vllm_config=None, kv_cache_config=None, cpu_capacity_bytes=0
149 + )
150 + recording = _RecordingBackend()
151 + worker._backend = recording
152 + worker._connector_metadata = SimpleCPUOffloadMetadata(
153 + load_event=0,
154 + load_gpu_blocks=[0],
155 + load_cpu_blocks=[0],
156 + store_event=1,
157 + store_gpu_blocks=[1],
158 + store_cpu_blocks=[1],
159 + )
160 +
161 + worker.get_finished(set())
162 +
163 + store_calls = [c for c in recording.calls if c["is_store"]]
164 + load_calls = [c for c in recording.calls if not c["is_store"]]
165 + assert len(store_calls) == 1
166 + assert len(load_calls) == 1
167 + assert isinstance(store_calls[0]["wait_event"], torch.Event)
168 + assert load_calls[0]["wait_event"] is None
169 +
170 +
171 +def test_build_params_src_access_order():
172 + """build_params defaults to ANY and honors an explicit STREAM override."""
173 + gpu = {"k": torch.zeros((4, 64), dtype=torch.int8, device="cuda")}
174 + cpu = {"k": torch.zeros((4, 64), dtype=torch.int8, device="cpu")}
175 + stream = torch.cuda.Stream()
176 +
177 + default = build_params(gpu, cpu, stream)
178 + assert default.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_ANY
179 +
180 + ordered = build_params(
181 + gpu, cpu, stream, src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
182 + )
183 + assert ordered.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
vllm/v1/simple_kv_offload/copy_backend.py
+38 -5
@@ -12,6 +12,8 @@ import torch
12 from vllm.logger import init_logger
13 from vllm.platforms import current_platform
14 from vllm.v1.simple_kv_offload.cuda_mem_ops import (
15 + CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
16 + CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
17 BatchMemcpyParams,
18 build_params,
19 copy_blocks,
@@ -43,8 +45,20 @@ class DmaCopyBackend:
45 self._load_stream = load_stream
46 self._store_stream = store_stream
47
46 - self._store_params = build_params(gpu_caches, cpu_caches, store_stream)
47 - self._load_params = build_params(cpu_caches, gpu_caches, load_stream)
48 + # Stores read the live KV cache -> STREAM (paired with the compute-done
49 + # wait in get_finished); loads read stable pinned host memory -> ANY.
50 + self._store_params = build_params(
51 + gpu_caches,
52 + cpu_caches,
53 + store_stream,
54 + src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
55 + )
56 + self._load_params = build_params(
57 + cpu_caches,
58 + gpu_caches,
59 + load_stream,
60 + src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
61 + )
62
63 self._queue = queue.SimpleQueue()
64 self._thread = threading.Thread(
@@ -61,11 +75,20 @@ class DmaCopyBackend:
75 is_store: bool,
76 event_idx: int,
77 events_list: list[tuple[int, torch.Event]],
78 + wait_event: torch.Event | None = None,
79 ) -> None:
80 params = self._store_params if is_store else self._load_params
81 assert params is not None and self._queue is not None
82 self._queue.put(
68 - (src_blocks, dst_blocks, params, is_store, event_idx, events_list)
83 + (
84 + src_blocks,
85 + dst_blocks,
86 + params,
87 + is_store,
88 + event_idx,
89 + events_list,
90 + wait_event,
91 + )
92 )
93
94 def shutdown(self) -> None:
@@ -89,9 +112,19 @@ class DmaCopyBackend:
112 item = q.get()
113 if item is None:
114 return
92 - src_blocks, dst_blocks, params, is_store, event_idx, events_list = item
93 - copy_blocks(src_blocks, dst_blocks, params)
115 + (
116 + src_blocks,
117 + dst_blocks,
118 + params,
119 + is_store,
120 + event_idx,
121 + events_list,
122 + wait_event,
123 + ) = item
124 stream = store_stream if is_store else load_stream
125 + if wait_event is not None:
126 + stream.wait_event(wait_event)
127 + copy_blocks(src_blocks, dst_blocks, params)
128 event = torch.Event()
129 event.record(stream)
130 events_list.append((event_idx, event))
vllm/v1/simple_kv_offload/cuda_mem_ops.py
+10 -6
@@ -13,6 +13,12 @@ from vllm.platforms import current_platform
13
14 logger = init_logger(__name__)
15
16 +# CUmemcpySrcAccessOrder values (CUDA driver API). STREAM(1): source read in
17 +# stream order, safe when the source may still be written. ANY(3): source may
18 +# be read early, only safe for a stable source (e.g. pinned host memory).
19 +CU_MEMCPY_SRC_ACCESS_ORDER_STREAM = 1
20 +CU_MEMCPY_SRC_ACCESS_ORDER_ANY = 3
21 +
22
23 def pin_tensor(tensor: torch.Tensor) -> None:
24 """Pin a CPU tensor via cudaHostRegister.
@@ -106,8 +112,8 @@ class BatchMemcpyParams(NamedTuple):
112 dst_bases: np.ndarray # [num_layers] uint64
113 bpb: np.ndarray # [num_layers] uint64 — bytes per block
114 num_layers: int
109 - # CUDA only: one attributes entry with srcAccessOrder=ANY. Unused on
110 - # ROCm (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0.
115 + # CUDA only: one attributes entry carrying srcAccessOrder. Unused on ROCm
116 + # (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0.
117 attrs: _CUmemcpyAttributes
118 attrs_idx: ctypes.c_size_t
119 # NOTE: cuMemcpyBatchAsync_v2() removed fail_idx field, but we use
@@ -120,6 +126,7 @@ def build_params(
126 src_caches: dict[str, torch.Tensor],
127 dst_caches: dict[str, torch.Tensor],
128 stream: torch.cuda.Stream,
129 + src_access_order: int = CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
130 ) -> BatchMemcpyParams:
131 global _batch_memcpy_fn
132 if _batch_memcpy_fn is None:
@@ -137,10 +144,7 @@ def build_params(
144 dst_bases.append(d.data_ptr())
145 bpb.append(s_bpb)
146
140 - # ``srcAccessOrder=3`` == CU_MEMCPY_SRC_ACCESS_ORDER_ANY /
141 - # hipMemcpySrcAccessOrderAny. See
142 - # https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f # noqa: E501
143 - attrs = _CUmemcpyAttributes(srcAccessOrder=3)
147 + attrs = _CUmemcpyAttributes(srcAccessOrder=src_access_order)
148
149 return BatchMemcpyParams(
150 src_bases=np.array(src_bases, dtype=np.uint64),
vllm/v1/simple_kv_offload/worker.py
+13 -5
@@ -57,6 +57,10 @@ class SimpleCPUOffloadWorker:
57 # Metadata for the current step
58 self._connector_metadata: SimpleCPUOffloadMetadata | None = None
59
60 + # Compute-done event recorded before each store; reused across steps
61 + # (get_finished runs once per step, copy queue is FIFO).
62 + self._store_compute_done: torch.Event | None = None
63 +
64 # Pending event index sets, populated in bind_connector_metadata
65 self._pending_load_event_indices: set[int] = set()
66 self._pending_store_event_indices: set[int] = set()
@@ -206,9 +210,11 @@ class SimpleCPUOffloadWorker:
210 ) -> tuple[set[str] | None, set[str] | None]:
211 """Submit transfers and report completed events to the scheduler.
212
209 - Called after model execution. The manager only schedules stores for
210 - blocks whose KV data is confirmed computed, so we launch both loads
211 - and stores immediately — no deferral or cross-stream sync needed.
213 + Stores (GPU->CPU) read the live KV cache, which the compute stream may
214 + still be writing under v1 overlapped execution, so they are ordered
215 + after a compute-done event recorded on the current stream. Loads
216 + (CPU->GPU) read stable pinned host memory and launch immediately. See
217 + #45704 for the bug and #39306 for the srcAccessOrder rationale.
218
219 Returns:
220 tuple of (finished_sending, finished_recving).
@@ -218,7 +224,6 @@ class SimpleCPUOffloadWorker:
224 # (1) Submit transfers
225 metadata = self._connector_metadata
226 if metadata is not None:
221 - # Launch loads (CPU->GPU).
227 if metadata.load_cpu_blocks:
228 self._backend.launch_copy(
229 metadata.load_cpu_blocks,
@@ -227,14 +232,17 @@ class SimpleCPUOffloadWorker:
232 event_idx=metadata.load_event,
233 events_list=self._load_events,
234 )
230 - # Launch stores (GPU->CPU).
235 if metadata.store_gpu_blocks:
236 + if self._store_compute_done is None:
237 + self._store_compute_done = torch.Event()
238 + self._store_compute_done.record(torch.cuda.current_stream())
239 self._backend.launch_copy(
240 metadata.store_gpu_blocks,
241 metadata.store_cpu_blocks,
242 is_store=True,
243 event_idx=metadata.store_event,
244 events_list=self._store_events,
245 + wait_event=self._store_compute_done,
246 )
247
248 # (2) Track completed transfer events