[Bugfix][KV Offload] Centralize shared mmap cleanup in CPU worker (#51622)
Signed-off-by: Alex <jihui.huang@daocloud.io>
AlexHuang committed
Aug 11, 2026 at 18:22 UTC
0fec3d652babd3a4a2919ff7b19a35701b298426
2 files changed
+220
-13
tests/v1/kv_offload/cpu/test_gpu_worker.py
+173
-2
@@ -1,8 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import logging
4
import random
5
import time
6
import uuid
7
+from unittest.mock import MagicMock
8
9
import pytest
10
import torch
@@ -46,6 +48,177 @@ def test_rocm_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None:
48
)
49
50
51
+def test_worker_shutdown_releases_region_and_runs_both_handlers() -> None:
52
+ """Both directions drain before the worker releases its shared region."""
53
+ worker = CPUOffloadingWorker.__new__(CPUOffloadingWorker)
54
+ calls: list[str] = []
55
+
56
+ def record_store_shutdown() -> bool:
57
+ calls.append("store")
58
+ return True
59
+
60
+ def record_load_shutdown() -> bool:
61
+ calls.append("load")
62
+ return True
63
+
64
+ store_handler = MagicMock()
65
+ load_handler = MagicMock()
66
+ store_handler.shutdown.side_effect = record_store_shutdown
67
+ load_handler.shutdown.side_effect = record_load_shutdown
68
+
69
+ def record_region_cleanup(**_: bool) -> bool:
70
+ calls.append("region")
71
+ return True
72
+
73
+ mmap_region = MagicMock()
74
+ mmap_region.cleanup.side_effect = record_region_cleanup
75
+ worker._store_handler = store_handler
76
+ worker._load_handler = load_handler
77
+ worker._mmap_region = mmap_region
78
+
79
+ worker.shutdown()
80
+
81
+ assert calls == ["store", "load", "region"]
82
+ store_handler.shutdown.assert_called_once_with()
83
+ load_handler.shutdown.assert_called_once_with()
84
+ mmap_region.cleanup.assert_called_once_with()
85
+ assert worker._mmap_region is None
86
+
87
+
88
+@pytest.mark.parametrize("failing_handler", ["store", "load"])
89
+def test_worker_logs_handler_error_and_cleans_region(
90
+ caplog_vllm, monkeypatch: pytest.MonkeyPatch, failing_handler
91
+) -> None:
92
+ worker = CPUOffloadingWorker.__new__(CPUOffloadingWorker)
93
+ calls: list[str] = []
94
+ store_handler = MagicMock()
95
+ load_handler = MagicMock()
96
+ if failing_handler == "store":
97
+
98
+ def fail_store_shutdown() -> bool:
99
+ calls.append("store")
100
+ raise RuntimeError("transfer did not drain")
101
+
102
+ store_handler.shutdown.side_effect = fail_store_shutdown
103
+ else:
104
+
105
+ def fail_load_shutdown() -> bool:
106
+ calls.append("load")
107
+ raise RuntimeError("transfer did not drain")
108
+
109
+ load_handler.shutdown.side_effect = fail_load_shutdown
110
+
111
+ def record_other_shutdown() -> bool:
112
+ calls.append("load" if failing_handler == "store" else "store")
113
+ return True
114
+
115
+ if failing_handler == "store":
116
+ load_handler.shutdown.side_effect = record_other_shutdown
117
+ else:
118
+ store_handler.shutdown.side_effect = record_other_shutdown
119
+
120
+ def record_device_sync() -> None:
121
+ calls.append("sync")
122
+
123
+ monkeypatch.setattr(torch.accelerator, "synchronize", record_device_sync)
124
+
125
+ def record_region_cleanup() -> None:
126
+ calls.append("region")
127
+
128
+ mmap_region = MagicMock()
129
+ mmap_region.cleanup.side_effect = record_region_cleanup
130
+ worker._store_handler = store_handler
131
+ worker._load_handler = load_handler
132
+ worker._mmap_region = mmap_region
133
+
134
+ with caplog_vllm.at_level(
135
+ logging.ERROR, logger="vllm.v1.kv_offload.cpu.gpu_worker"
136
+ ):
137
+ worker.shutdown()
138
+
139
+ other = "load" if failing_handler == "store" else "store"
140
+ getattr(worker, f"_{other}_handler").shutdown.assert_called_once_with()
141
+ mmap_region.cleanup.assert_called_once_with()
142
+ assert worker._mmap_region is None
143
+ assert (
144
+ f"Failed to shut down {failing_handler} offloading handler" in caplog_vllm.text
145
+ )
146
+ assert calls[-2:] == ["sync", "region"]
147
+
148
+
149
+def test_handler_shutdown_skips_transfers_after_event_sync_failure() -> None:
150
+ handler = gpu_worker.SingleDirectionOffloadingHandler.__new__(
151
+ gpu_worker.SingleDirectionOffloadingHandler
152
+ )
153
+ failed_event = MagicMock()
154
+ failed_event.synchronize.side_effect = RuntimeError("device lost")
155
+ skipped_event = MagicMock()
156
+ handler._transfers = gpu_worker.deque(
157
+ [
158
+ MagicMock(end_event=failed_event),
159
+ MagicMock(end_event=skipped_event),
160
+ ]
161
+ )
162
+ handler._transfer_events = {1: failed_event, 2: skipped_event}
163
+ handler._stream_pool = [MagicMock()]
164
+ handler._event_pool = [MagicMock()]
165
+ handler._buffer_pool = [(MagicMock(), MagicMock(), MagicMock())]
166
+ handler.src_tensors = [MagicMock()]
167
+ handler.dst_tensors = [MagicMock()]
168
+
169
+ with pytest.raises(RuntimeError, match="device lost"):
170
+ handler.shutdown()
171
+ failed_event.synchronize.assert_called_once_with()
172
+ skipped_event.synchronize.assert_not_called()
173
+ assert not handler._transfers
174
+ assert not handler._transfer_events
175
+ assert not handler._stream_pool
176
+ assert not handler._event_pool
177
+ assert not handler._buffer_pool
178
+ assert not handler.src_tensors
179
+ assert not handler.dst_tensors
180
+
181
+
182
+@pytest.mark.parametrize("device_sync_fails", [False, True])
183
+def test_worker_syncs_before_cleanup_after_handler_failure(
184
+ caplog_vllm, monkeypatch: pytest.MonkeyPatch, device_sync_fails: bool
185
+) -> None:
186
+ worker = CPUOffloadingWorker.__new__(CPUOffloadingWorker)
187
+ calls: list[str] = []
188
+ store_handler = MagicMock()
189
+
190
+ def fail_store_shutdown() -> None:
191
+ raise RuntimeError("device lost")
192
+
193
+ store_handler.shutdown.side_effect = fail_store_shutdown
194
+ load_handler = MagicMock()
195
+
196
+ def record_device_sync() -> None:
197
+ calls.append("sync")
198
+ if device_sync_fails:
199
+ raise RuntimeError("device lost")
200
+
201
+ def record_region_cleanup() -> None:
202
+ calls.append("region")
203
+
204
+ monkeypatch.setattr(torch.accelerator, "synchronize", record_device_sync)
205
+ mmap_region = MagicMock()
206
+ mmap_region.cleanup.side_effect = record_region_cleanup
207
+ worker._store_handler = store_handler
208
+ worker._load_handler = load_handler
209
+ worker._mmap_region = mmap_region
210
+
211
+ with caplog_vllm.at_level(
212
+ logging.WARNING, logger="vllm.v1.kv_offload.cpu.gpu_worker"
213
+ ):
214
+ worker.shutdown()
215
+
216
+ assert calls == ["sync", "region"]
217
+ mmap_region.cleanup.assert_called_once_with()
218
+ if device_sync_fails:
219
+ assert "Device sync before mmap cleanup failed" in caplog_vllm.text
220
+
221
+
222
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
223
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
224
@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
@@ -229,8 +402,6 @@ def test_transfer(
402
del src_view, dst_view, orig_dst_view, expected
403
404
worker.shutdown()
232
- if mmap_region:
233
- mmap_region.cleanup()
405
406
407
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
vllm/v1/kv_offload/cpu/gpu_worker.py
+47
-11
@@ -248,7 +248,6 @@ class SingleDirectionOffloadingHandler:
248
blocks_per_chunk: int,
249
layer_refs_per_group: list[list[CanonicalKVCacheRef]],
250
gpu_to_cpu: bool,
251
- mmap_region: SharedOffloadRegion | None = None,
251
canonical_layout: bool = False,
252
):
253
"""
@@ -326,8 +325,6 @@ class SingleDirectionOffloadingHandler:
325
self._scratch_bases_src = np.empty(num_scratch_blocks, dtype=np.uint64)
326
self._scratch_bases_dst = np.empty(num_scratch_blocks, dtype=np.uint64)
327
329
- # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)
330
- self._mmap_region = mmap_region
328
# job_id -> event
329
self._transfer_events: dict[int, torch.Event] = {}
330
# queue of transfers (job_id, stream, event)
@@ -711,18 +708,31 @@ class SingleDirectionOffloadingHandler:
708
event.synchronize()
709
710
def shutdown(self) -> None:
711
+ """Drain this direction and release its transfer-side resources."""
712
+ sync_error: Exception | None = None
713
while self._transfers:
715
- transfer = self._transfers.popleft()
716
- transfer.end_event.synchronize()
714
+ transfer = self._transfers[0]
715
+ try:
716
+ transfer.end_event.synchronize()
717
+ except Exception as e:
718
+ logger.exception(
719
+ "Failed to synchronize transfer end event; "
720
+ "skipping %d remaining transfers",
721
+ len(self._transfers) - 1,
722
+ )
723
+ self._transfers.clear()
724
+ sync_error = e
725
+ break
726
+ self._transfers.popleft()
727
+
728
self._transfer_events.clear()
729
self._stream_pool.clear()
730
self._event_pool.clear()
731
self._buffer_pool.clear()
732
self.src_tensors.clear()
733
self.dst_tensors.clear()
723
- if self._mmap_region is not None:
724
- self._mmap_region.cleanup()
725
- self._mmap_region = None
734
+ if sync_error is not None:
735
+ raise sync_error
736
737
738
class CPUOffloadingWorker(OffloadingWorker):
@@ -742,6 +752,10 @@ class CPUOffloadingWorker(OffloadingWorker):
752
canonical_layout: bool = False,
753
):
754
assert not canonical_layout or mmap_region is not None
755
+ # The caller owns mmap_region until this constructor returns. After a
756
+ # successful construction, the worker is the sole owner and releases
757
+ # it after both transfer directions have stopped.
758
+ self._mmap_region = mmap_region
759
pin_memory = PIN_MEMORY
760
logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
761
if mmap_region is not None and pin_memory:
@@ -794,7 +808,6 @@ class CPUOffloadingWorker(OffloadingWorker):
808
blocks_per_chunk=blocks_per_chunk,
809
layer_refs_per_group=kv_caches.group_data_refs,
810
gpu_to_cpu=True,
797
- mmap_region=mmap_region,
811
canonical_layout=canonical_layout,
812
)
813
@@ -827,5 +840,28 @@ class CPUOffloadingWorker(OffloadingWorker):
840
self._load_handler.wait(job_ids)
841
842
def shutdown(self) -> None:
830
- self._store_handler.shutdown()
831
- self._load_handler.shutdown()
843
+ handler_failed = False
844
+ try:
845
+ self._store_handler.shutdown()
846
+ except Exception:
847
+ logger.exception("Failed to shut down store offloading handler")
848
+ handler_failed = True
849
+
850
+ try:
851
+ self._load_handler.shutdown()
852
+ except Exception:
853
+ logger.exception("Failed to shut down load offloading handler")
854
+ handler_failed = True
855
+
856
+ if self._mmap_region is not None:
857
+ if handler_failed:
858
+ try:
859
+ torch.accelerator.synchronize()
860
+ except Exception:
861
+ logger.warning(
862
+ "Device sync before mmap cleanup failed; "
863
+ "proceeding with cleanup anyway",
864
+ exc_info=True,
865
+ )
866
+ self._mmap_region.cleanup()
867
+ self._mmap_region = None