[Bugfix][KV Offload] Submit CPU stores on no-forward steps (#56621)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Sep 13, 2026 at 01:56 UTC
410f6da5c4bb62010728502035bee1b5f0eab2ac
2 files changed
+110
-4
tests/v1/simple_kv_offload/test_worker.py
+100
-3
@@ -10,6 +10,7 @@ read partially written / stale blocks and silently corrupt the CPU cache.
10
from __future__ import annotations
11
12
import time
13
+from contextlib import nullcontext
14
from unittest.mock import MagicMock
15
16
import pytest
@@ -21,7 +22,13 @@ if not current_platform.is_cuda_alike():
22
pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
23
24
from tests.v1.attention.utils import dense_kv_cache_tensor, dense_kv_cache_views
25
+from tests.v1.kv_connector.unit.test_kv_connector_lifecycle import (
26
+ _make_empty_scheduler_output,
27
+)
28
from vllm.config import CacheConfig
29
+from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (
30
+ SimpleCPUOffloadConnector,
31
+)
32
from vllm.v1.core.kv_cache_utils import (
33
get_kv_cache_config_from_groups,
34
is_kv_cache_spec_uniform,
@@ -41,11 +48,12 @@ from vllm.v1.simple_kv_offload.cuda_mem_ops import (
48
CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
49
CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
50
build_params,
44
- pin_tensor,
51
)
52
from vllm.v1.simple_kv_offload.disk_backend import DiskBackend
53
from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata
54
from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker
55
+from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
56
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
57
from vllm.v1.worker.utils import allocate_kv_cache
58
59
NUM_BLOCKS = 64
@@ -58,8 +66,11 @@ SLEEP_CYCLES = 50_000_000
66
67
def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:
68
gpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cuda")}
61
- cpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu")}
62
- pin_tensor(cpu["k"])
69
+ cpu = {
70
+ "k": torch.zeros(
71
+ (NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu", pin_memory=True
72
+ )
73
+ }
74
low_pri, _ = torch.cuda.Stream.priority_range()
75
backend = DmaCopyBackend()
76
backend.init(
@@ -72,6 +83,92 @@ def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:
83
return backend, gpu["k"], cpu["k"]
84
85
86
+@pytest.mark.parametrize("use_v2", [False, True])
87
+def test_no_forward_step_completes_cpu_store(monkeypatch, use_v2):
88
+ """A finished request's final store must drain without another model step."""
89
+ backend, gpu, cpu = _make_backend()
90
+ worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0)
91
+ worker._backend = backend
92
+ connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector)
93
+ connector.worker_handler = worker
94
+ output = _make_empty_scheduler_output()
95
+ output.kv_connector_metadata = SimpleCPUOffloadMetadata(
96
+ store_event=0, store_gpu_blocks=[0], store_cpu_blocks=[0]
97
+ )
98
+ module = "vllm.v1.worker.kv_connector_model_runner_mixin"
99
+ monkeypatch.setattr(f"{module}.get_kv_transfer_group", lambda: connector)
100
+ monkeypatch.setattr(f"{module}.set_forward_context", lambda *a: nullcontext())
101
+ monkeypatch.setattr(f"{module}.get_forward_context", lambda: None)
102
+ try:
103
+ gpu.fill_(91)
104
+ if use_v2:
105
+ active = ActiveKVConnector.__new__(ActiveKVConnector)
106
+ active.kv_connector = connector
107
+ active._disabled = False
108
+ active._pending_load_start = False
109
+ module = "vllm.v1.worker.gpu.kv_connector"
110
+ monkeypatch.setattr(f"{module}.is_forward_context_available", lambda: True)
111
+ monkeypatch.setattr(f"{module}.get_forward_context", lambda: None)
112
+ result = active.no_forward(output)
113
+ else:
114
+ result = KVConnectorModelRunnerMixin.kv_connector_no_forward(output, None)
115
+ completion = (
116
+ result.kv_connector_output.kv_connector_worker_meta
117
+ if result.kv_connector_output is not None
118
+ else None
119
+ )
120
+ deadline = time.monotonic() + 5
121
+ while completion is None and time.monotonic() < deadline:
122
+ connector.get_finished(set())
123
+ completion = connector.build_connector_worker_meta()
124
+ time.sleep(0.001)
125
+ assert completion is not None, "CPU store never completed on the empty step"
126
+ assert completion.completed_store_events == {0: 1}
127
+ assert torch.equal(cpu[0], gpu[0].cpu())
128
+ finally:
129
+ backend.shutdown()
130
+
131
+
132
+def test_deferred_store_waits_for_draft_forward(monkeypatch):
133
+ """Target completion must not submit the store before draft finalization."""
134
+ backend, gpu, cpu = _make_backend()
135
+ worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0)
136
+ worker._backend = backend
137
+ connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector)
138
+ connector.worker_handler = worker
139
+ output = _make_empty_scheduler_output()
140
+ output.kv_connector_metadata = SimpleCPUOffloadMetadata(
141
+ store_event=0, store_gpu_blocks=[0], store_cpu_blocks=[0]
142
+ )
143
+ module = "vllm.v1.worker.kv_connector_model_runner_mixin"
144
+ monkeypatch.setattr(f"{module}.get_kv_transfer_group", lambda: connector)
145
+ monkeypatch.setattr(f"{module}.has_kv_transfer_group", lambda: True)
146
+ monkeypatch.setattr(f"{module}.get_forward_context", lambda: None)
147
+ launch = MagicMock(wraps=backend.launch_copy)
148
+ monkeypatch.setattr(backend, "launch_copy", launch)
149
+ try:
150
+ with KVConnectorModelRunnerMixin._get_kv_connector_output(
151
+ output, defer_finalize=True
152
+ ):
153
+ gpu.fill_(17)
154
+ launch.assert_not_called()
155
+ gpu.fill_(91)
156
+ KVConnectorModelRunnerMixin.finalize_kv_connector()
157
+ assert launch.call_count == 1
158
+ completion = None
159
+ deadline = time.monotonic() + 5
160
+ while completion is None and time.monotonic() < deadline:
161
+ connector.get_finished(set())
162
+ completion = connector.build_connector_worker_meta()
163
+ time.sleep(0.001)
164
+ assert completion is not None
165
+ assert completion.completed_store_events == {0: 1}
166
+ assert torch.equal(cpu[0], gpu[0].cpu())
167
+ assert launch.call_count == 1
168
+ finally:
169
+ backend.shutdown()
170
+
171
+
172
def _drive_store(
173
backend: DmaCopyBackend,
174
gpu: torch.Tensor,
vllm/v1/simple_kv_offload/worker.py
+10
-1
@@ -67,6 +67,7 @@ class SimpleCPUOffloadWorker:
67
68
# Metadata for the current step
69
self._connector_metadata: SimpleCPUOffloadMetadata | None = None
70
+ self._store_submitted = False
71
72
# Compute-done event recorded before each store; reused across steps
73
# (get_finished runs once per step, copy queue is FIFO).
@@ -244,12 +245,15 @@ class SimpleCPUOffloadWorker:
245
246
def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None:
247
self._connector_metadata = metadata
248
+ self._store_submitted = False
249
if metadata.load_event >= 0:
250
self._pending_load_event_indices.add(metadata.load_event)
251
if metadata.store_event >= 0:
252
self._pending_store_event_indices.add(metadata.store_event)
253
254
def clear_connector_metadata(self) -> None:
255
+ # No-forward steps skip the normal wait_for_save hook.
256
+ self.wait_for_save()
257
self._connector_metadata = None
258
259
def start_load_kv(self) -> None:
@@ -278,7 +282,11 @@ class SimpleCPUOffloadWorker:
282
#45704 for the bug and #39306 for the srcAccessOrder rationale.
283
"""
284
metadata = self._connector_metadata
281
- if metadata is not None and metadata.store_gpu_blocks:
285
+ if (
286
+ metadata is not None
287
+ and metadata.store_gpu_blocks
288
+ and not self._store_submitted
289
+ ):
290
backend = self._backend
291
assert backend is not None
292
if self._store_compute_done is None:
@@ -292,6 +300,7 @@ class SimpleCPUOffloadWorker:
300
events_list=self._store_events,
301
wait_event=self._store_compute_done,
302
)
303
+ self._store_submitted = True
304
305
def get_finished(
306
self, finished_req_ids: set[str]