@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / be0c855ebd
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 04:45 UTC

[KV Offload] Unified memory layout for offloading workers (#37206)

Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

omerpaz95 committed Apr 14, 2026 at 21:33 UTC be0c855ebd6ed4a8bface6c7df6d99150fdb5d90
4 files changed +951 -55
tests/v1/kv_offload/test_cpu_gpu.py
+32 -7
@@ -2,12 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import random
4 import time
5 +import uuid
6
7 import pytest
8 import torch
9
10 from vllm.platforms import current_platform
11 from vllm.utils.torch_utils import set_random_seed
12 +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
13 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
14 from vllm.v1.kv_offload.spec import (
15 CanonicalKVCacheRef,
@@ -36,6 +38,7 @@ NUM_MAPPINGS = [3]
38 @pytest.mark.parametrize("num_tensors", NUM_TENSORS)
39 @pytest.mark.parametrize("seed", SEEDS)
40 @pytest.mark.parametrize("device", DEVICES)
41 +@pytest.mark.parametrize("use_shared_memory", [False, True])
42 @torch.inference_mode()
43 def test_transfer(
44 default_vllm_config,
@@ -48,6 +51,7 @@ def test_transfer(
51 num_tensors: int,
52 seed: int,
53 device: str,
54 + use_shared_memory: bool,
55 ) -> None:
56 set_random_seed(seed)
57
@@ -83,10 +87,24 @@ def test_transfer(
87 tensors=kv_cache_tensors,
88 group_data_refs=kv_cache_groups_data_refs,
89 )
90 +
91 + mmap_region: SharedOffloadRegion | None = None
92 + if use_shared_memory:
93 + cpu_page_size = gpu_page_size_bytes * num_tensors * block_size_factor
94 + mmap_region = SharedOffloadRegion(
95 + instance_id=str(uuid.uuid4()),
96 + total_size_bytes=num_cpu_blocks * cpu_page_size,
97 + num_blocks=num_cpu_blocks,
98 + rank=0,
99 + num_workers=1,
100 + cpu_page_size=cpu_page_size,
101 + )
102 +
103 handlers = CpuGpuOffloadingHandlers(
104 kv_caches=kv_caches,
105 block_size_factor=block_size_factor,
106 num_cpu_blocks=num_cpu_blocks,
107 + mmap_region=mmap_region,
108 )
109
110 # select block mappings
@@ -137,10 +155,8 @@ def test_transfer(
155 if finished:
156 assert finished[0].job_id == 1
157 assert finished[0].success
140 - assert (
141 - finished[0].transfer_type == ("GPU", "CPU")
142 - if gpu_to_cpu
143 - else ("CPU", "GPU")
158 + assert finished[0].transfer_type == (
159 + ("GPU", "CPU") if gpu_to_cpu else ("CPU", "GPU")
160 )
161 assert finished[0].transfer_size == (
162 len(gpu_blocks) * handler.group_block_size_in_bytes[0]
@@ -161,9 +177,9 @@ def test_transfer(
177 orig_dst_tensors,
178 ):
179 # view both GPU and CPU tensors as (n, gpu_page_size_bytes) for comparison.
164 - src_view = src_tensor.view(-1, gpu_page_size_bytes)
165 - dst_view = dst_tensor.view(-1, gpu_page_size_bytes)
166 - orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes)
180 + src_view = src_tensor.reshape(-1, gpu_page_size_bytes)
181 + dst_view = dst_tensor.reshape(-1, gpu_page_size_bytes)
182 + orig_dst_view = orig_dst_tensor.reshape(-1, gpu_page_size_bytes)
183 for dst_sub_block in range(num_dst_sub_blocks):
184 src_sub_block = dst_to_src.get(dst_sub_block)
185 if src_sub_block is not None:
@@ -171,3 +187,12 @@ def test_transfer(
187 else:
188 expected = orig_dst_view[dst_sub_block]
189 torch.testing.assert_close(dst_view[dst_sub_block].cpu(), expected.cpu())
190 +
191 + # Drop loop-variable refs so mmap_obj has no exported buffers at cleanup.
192 + del orig_tensor, tensor, src_tensor, dst_tensor, orig_dst_tensor
193 + del src_view, dst_view, orig_dst_view, expected
194 +
195 + handlers.cpu_to_gpu_handler.shutdown()
196 + handlers.gpu_to_cpu_handler.shutdown()
197 + if mmap_region:
198 + mmap_region.cleanup()
tests/v1/kv_offload/test_shared_offload_region.py new
+625
@@ -0,0 +1,625 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for SharedOffloadRegion."""
4 +
5 +import contextlib
6 +import mmap
7 +import os
8 +import threading
9 +import time
10 +import uuid
11 +
12 +import pytest
13 +
14 +from vllm.utils.system_utils import get_mp_context
15 +from vllm.v1.kv_offload.cpu.shared_offload_region import (
16 + SharedOffloadRegion,
17 + _wait_for_file_size,
18 +)
19 +
20 +PAGE_SIZE = mmap.PAGESIZE
21 +
22 +
23 +# ---------------------------------------------------------------------------
24 +# Helpers / fixtures
25 +# ---------------------------------------------------------------------------
26 +
27 +
28 +@pytest.fixture(autouse=True)
29 +def _set_spawn_method(monkeypatch):
30 + # On WSL, NVML is not compatible with fork so vLLM auto-overrides the
31 + # multiprocessing start method to 'spawn' with a warning. Set it explicitly
32 + # here so the override is a no-op and the warning is suppressed.
33 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
34 +
35 +
36 +def _make_region(
37 + instance_id: str,
38 + num_blocks: int = 4,
39 + cpu_page_size: int = PAGE_SIZE,
40 + num_workers: int = 1,
41 + rank: int = 0,
42 +) -> SharedOffloadRegion:
43 + total_size_bytes = num_blocks * num_workers * cpu_page_size
44 + assert total_size_bytes % PAGE_SIZE == 0
45 + return SharedOffloadRegion(
46 + instance_id=instance_id,
47 + total_size_bytes=total_size_bytes,
48 + num_blocks=num_blocks,
49 + rank=rank,
50 + num_workers=num_workers,
51 + cpu_page_size=cpu_page_size,
52 + )
53 +
54 +
55 +def _cleanup_file(path: str) -> None:
56 + """Best-effort file removal for test teardown."""
57 + with contextlib.suppress(FileNotFoundError):
58 + os.unlink(path)
59 +
60 +
61 +@contextlib.contextmanager
62 +def _region(instance_id: str, **kwargs):
63 + """Context manager: create one region, clean up on exit."""
64 + r = _make_region(instance_id, **kwargs)
65 + try:
66 + yield r
67 + finally:
68 + r.cleanup()
69 + _cleanup_file(r.mmap_path)
70 +
71 +
72 +@contextlib.contextmanager
73 +def _multi_region(
74 + instance_id: str,
75 + num_workers: int,
76 + num_blocks: int = 4,
77 + cpu_page_size: int = PAGE_SIZE,
78 +):
79 + """Context manager: create one SharedOffloadRegion per rank, clean up on exit."""
80 + total = num_blocks * num_workers * cpu_page_size
81 + regions = [
82 + SharedOffloadRegion(
83 + instance_id=instance_id,
84 + total_size_bytes=total,
85 + num_blocks=num_blocks,
86 + rank=rank,
87 + num_workers=num_workers,
88 + cpu_page_size=cpu_page_size,
89 + )
90 + for rank in range(num_workers)
91 + ]
92 + try:
93 + yield regions
94 + finally:
95 + for r in regions:
96 + r.cleanup()
97 + _cleanup_file(regions[0].mmap_path)
98 +
99 +
100 +def _race_construct(
101 + instance_id: str,
102 + num_workers: int,
103 + num_blocks: int = 4,
104 + cpu_page_size: int = PAGE_SIZE,
105 +) -> tuple[list[SharedOffloadRegion], list[Exception]]:
106 + """Spawn num_workers threads that all race to construct SharedOffloadRegion."""
107 + total = num_blocks * num_workers * cpu_page_size
108 + regions: list[SharedOffloadRegion | None] = [None] * num_workers
109 + errors: list[Exception] = []
110 + barrier = threading.Barrier(num_workers)
111 +
112 + def worker(rank: int) -> None:
113 + barrier.wait() # all threads start at the same instant
114 + try:
115 + regions[rank] = SharedOffloadRegion(
116 + instance_id=instance_id,
117 + total_size_bytes=total,
118 + num_blocks=num_blocks,
119 + rank=rank,
120 + num_workers=num_workers,
121 + cpu_page_size=cpu_page_size,
122 + )
123 + except Exception as e:
124 + errors.append(e)
125 +
126 + threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_workers)]
127 + for t in threads:
128 + t.start()
129 + for t in threads:
130 + t.join()
131 +
132 + return [r for r in regions if r is not None], errors
133 +
134 +
135 +def _mp_race_construct_and_write(
136 + instance_id: str,
137 + total_bytes: int,
138 + num_blocks: int,
139 + rank: int,
140 + num_workers: int,
141 + cpu_page_size: int,
142 + fill_value: int,
143 + done_queue,
144 + cleanup_queue,
145 +) -> None:
146 + """Race to construct a SharedOffloadRegion, write fill_value, then wait
147 + for the parent's cleanup signal before tearing down. The wait gives the
148 + parent a window to read the raw mmap before the creator removes the file."""
149 + try:
150 + region = SharedOffloadRegion(
151 + instance_id=instance_id,
152 + total_size_bytes=total_bytes,
153 + num_blocks=num_blocks,
154 + rank=rank,
155 + num_workers=num_workers,
156 + cpu_page_size=cpu_page_size,
157 + )
158 + t = region.create_next_view(cpu_page_size)
159 + t[:, :] = fill_value
160 + done_queue.put({"rank": rank, "error": None})
161 + cleanup_queue.get() # wait for parent's verification to finish
162 + del t # release view before cleanup to avoid BufferError
163 + region.cleanup()
164 + except Exception as e:
165 + done_queue.put({"rank": rank, "error": repr(e)})
166 +
167 +
168 +@pytest.fixture
169 +def iid():
170 + """Fresh instance ID for each test."""
171 + return str(uuid.uuid4())
172 +
173 +
174 +# ---------------------------------------------------------------------------
175 +# create_next_view — shape, stride and storage offset
176 +# ---------------------------------------------------------------------------
177 +
178 +
179 +def test_create_next_view_shape_and_stride(iid):
180 + """Returned tensor must have shape (num_blocks, tensor_page_size) and
181 + stride (row_stride, 1) where row_stride = cpu_page_size * num_workers."""
182 + with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r:
183 + t = r.create_next_view(PAGE_SIZE)
184 + assert t.shape == (4, PAGE_SIZE)
185 + # num_workers=1 → row_stride = cpu_page_size
186 + assert t.stride() == (2 * PAGE_SIZE, 1)
187 + del t
188 +
189 +
190 +def test_create_next_view_storage_offset_rank0(iid):
191 + """rank=0 worker's first tensor must start at byte 0 of the mmap."""
192 + with _region(iid, cpu_page_size=PAGE_SIZE, num_workers=2, rank=0) as r:
193 + t = r.create_next_view(PAGE_SIZE)
194 + assert t.data_ptr() == r._base.data_ptr() # storage_offset == 0
195 + del t
196 +
197 +
198 +def test_create_next_view_storage_offset_rank1(iid):
199 + """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap."""
200 + with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1):
201 + t1 = r1.create_next_view(PAGE_SIZE)
202 + assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE
203 + del t1
204 +
205 +
206 +def test_create_next_view_row_stride_with_multiple_workers(iid):
207 + """With num_workers=4, row_stride must be 4 * cpu_page_size."""
208 + with _region(iid, num_blocks=2, num_workers=4) as r:
209 + t = r.create_next_view(PAGE_SIZE)
210 + assert t.stride(0) == 4 * PAGE_SIZE
211 + del t
212 +
213 +
214 +# ---------------------------------------------------------------------------
215 +# create_next_view — cursor advancement
216 +# ---------------------------------------------------------------------------
217 +
218 +
219 +def test_create_next_view_cursor_advances(iid):
220 + """Each call to create_next_view must advance _worker_offset by tensor_page_size."""
221 + with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r:
222 + assert r._worker_offset == 0
223 + r.create_next_view(PAGE_SIZE)
224 + assert r._worker_offset == PAGE_SIZE
225 + r.create_next_view(PAGE_SIZE)
226 + assert r._worker_offset == 2 * PAGE_SIZE
227 + r.create_next_view(PAGE_SIZE)
228 + assert r._worker_offset == 3 * PAGE_SIZE # exactly at area end
229 +
230 +
231 +def test_create_next_view_exact_fill_succeeds(iid):
232 + """Allocations whose total exactly equals cpu_page_size must all succeed."""
233 + with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
234 + r.create_next_view(PAGE_SIZE) # first half
235 + r.create_next_view(PAGE_SIZE) # fills to area end — must not raise
236 +
237 +
238 +# ---------------------------------------------------------------------------
239 +# create_next_view — overflow guard
240 +# ---------------------------------------------------------------------------
241 +
242 +
243 +def test_create_next_view_single_overflow_raises(iid):
244 + """A single allocation larger than cpu_page_size must raise AssertionError."""
245 + with (
246 + _region(iid) as r,
247 + pytest.raises(AssertionError, match="exceeds worker area end"),
248 + ):
249 + r.create_next_view(PAGE_SIZE + 1)
250 +
251 +
252 +def test_create_next_view_cumulative_overflow_raises(iid):
253 + """Successive allocations that cumulatively exceed cpu_page_size must raise."""
254 + with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
255 + r.create_next_view(PAGE_SIZE) # ok — half used
256 + r.create_next_view(PAGE_SIZE) # ok — full
257 + with pytest.raises(AssertionError, match="exceeds worker area end"):
258 + r.create_next_view(1) # one byte too many
259 +
260 +
261 +def test_create_next_view_overflow_does_not_mutate_cursor(iid):
262 + """A failed create_next_view must leave _worker_offset unchanged."""
263 + with _region(iid) as r:
264 + offset_before = r._worker_offset
265 + with pytest.raises(AssertionError):
266 + r.create_next_view(PAGE_SIZE + 1)
267 + assert r._worker_offset == offset_before
268 +
269 +
270 +# ---------------------------------------------------------------------------
271 +# create_next_view — data correctness and layout
272 +# ---------------------------------------------------------------------------
273 +
274 +
275 +def test_create_next_view_write_visible_in_raw_mmap(iid):
276 + """Writes into a create_next_view view must appear at the correct raw mmap offset"""
277 + with _region(iid, num_blocks=4) as r:
278 + t = r.create_next_view(PAGE_SIZE)
279 + t[2, :] = 42 # write to block row 2
280 +
281 + raw = memoryview(r.mmap_obj)
282 + # num_workers=1 → row_stride = PAGE_SIZE; block 2 starts at byte 2*PAGE_SIZE
283 + chunk = bytes(raw[2 * PAGE_SIZE : 3 * PAGE_SIZE])
284 + assert all(b == 42 for b in chunk)
285 + del raw, t
286 +
287 +
288 +def test_create_next_view_multi_tensor_layout(iid):
289 + """Two tensors from the same worker land at consecutive byte offsets per row."""
290 + with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r:
291 + ta = r.create_next_view(PAGE_SIZE)
292 + tb = r.create_next_view(PAGE_SIZE)
293 +
294 + ta[:, :] = 1
295 + tb[:, :] = 2
296 +
297 + raw = memoryview(r.mmap_obj)
298 + for blk in range(2):
299 + row_offset = blk * 2 * PAGE_SIZE # num_workers=1
300 + assert all(b == 1 for b in raw[row_offset : row_offset + PAGE_SIZE])
301 + assert all(
302 + b == 2 for b in raw[row_offset + PAGE_SIZE : row_offset + 2 * PAGE_SIZE]
303 + )
304 + del raw, ta, tb
305 +
306 +
307 +def test_create_next_view_multiprocess_slots(iid):
308 + """Each worker process calls create_next_view and writes distinct data;
309 + the parent verifies each slot lands at the correct interleaved offset."""
310 + num_workers = 2
311 + num_blocks = 4
312 + total_bytes = num_blocks * num_workers * PAGE_SIZE
313 +
314 + ctx = get_mp_context()
315 + done_queue = ctx.Queue()
316 + cleanup_queue = ctx.Queue()
317 +
318 + # Parent is rank 0 (creator); child is rank 1 (joiner).
319 + region = SharedOffloadRegion(
320 + instance_id=iid,
321 + total_size_bytes=total_bytes,
322 + num_blocks=num_blocks,
323 + rank=0,
324 + num_workers=num_workers,
325 + cpu_page_size=PAGE_SIZE,
326 + )
327 + try:
328 + child = ctx.Process(
329 + target=_mp_race_construct_and_write,
330 + args=(
331 + iid,
332 + total_bytes,
333 + num_blocks,
334 + 1,
335 + num_workers,
336 + PAGE_SIZE,
337 + 22,
338 + done_queue,
339 + cleanup_queue,
340 + ),
341 + )
342 + child.start()
343 +
344 + t0 = region.create_next_view(PAGE_SIZE)
345 + t0[:, :] = 11
346 +
347 + result = done_queue.get(timeout=30)
348 + assert result["error"] is None, result["error"]
349 +
350 + raw = memoryview(region.mmap_obj)
351 + for blk in range(num_blocks):
352 + row_start = blk * num_workers * PAGE_SIZE
353 + w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
354 + w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
355 + assert all(b == 11 for b in w0), f"block {blk}: rank0 slot wrong"
356 + assert all(b == 22 for b in w1), f"block {blk}: rank1 slot wrong"
357 +
358 + del raw, t0 # release before finally triggers cleanup
359 + cleanup_queue.put(True)
360 + child.join(timeout=10)
361 + assert child.exitcode == 0
362 + finally:
363 + region.cleanup()
364 + _cleanup_file(region.mmap_path)
365 +
366 +
367 +def test_create_next_view_worker_isolation(iid):
368 + """Writes by worker 0 must not affect worker 1's slot and vice versa."""
369 + num_workers = 2
370 + num_blocks = 4
371 + with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions:
372 + t0 = regions[0].create_next_view(PAGE_SIZE)
373 + t1 = regions[1].create_next_view(PAGE_SIZE)
374 +
375 + t0[:, :] = 11
376 + t1[:, :] = 22
377 +
378 + raw = memoryview(regions[0].mmap_obj)
379 + for blk in range(num_blocks):
380 + row_start = blk * num_workers * PAGE_SIZE
381 + w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
382 + w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
383 + assert all(b == 11 for b in w0), f"block {blk}: worker0 slot corrupted"
384 + assert all(b == 22 for b in w1), f"block {blk}: worker1 slot corrupted"
385 + del raw, t0, t1 # release before finally triggers cleanup
386 +
387 +
388 +# ---------------------------------------------------------------------------
389 +# Constructor — creator vs joiner semantics
390 +# ---------------------------------------------------------------------------
391 +
392 +
393 +def test_creator_flag_set_on_first_open(iid):
394 + """The first worker to open the file must have _creator == True."""
395 + with _region(iid) as r:
396 + assert r._creator is True
397 +
398 +
399 +def test_joiner_flag_not_set(iid):
400 + """A second worker opening the same file must have _creator == False."""
401 + with _multi_region(iid, num_workers=2) as (r0, r1):
402 + assert r0._creator is True
403 + assert r1._creator is False
404 +
405 +
406 +def test_file_exists_after_construction(iid):
407 + """The mmap file must be present on disk after __init__ completes."""
408 + with _region(iid) as r:
409 + assert os.path.exists(r.mmap_path)
410 +
411 +
412 +def test_file_has_correct_size(iid):
413 + """The mmap file size on disk must equal total_size_bytes."""
414 + with _region(iid, num_blocks=4) as r:
415 + assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE
416 +
417 +
418 +# ---------------------------------------------------------------------------
419 +# Multi-worker race — concurrent construction
420 +# ---------------------------------------------------------------------------
421 +
422 +
423 +def test_multi_worker_race_exactly_one_creator(iid):
424 + """When N threads race to create the same region, exactly one becomes creator."""
425 + num_workers = 8
426 + regions, errors = _race_construct(iid, num_workers=num_workers)
427 + try:
428 + assert not errors, f"Workers raised: {errors}"
429 + assert len(regions) == num_workers, "Some workers failed to construct"
430 +
431 + creators = [r for r in regions if r._creator]
432 + assert len(creators) == 1, f"Expected 1 creator, got {len(creators)}"
433 + assert sum(1 for r in regions if not r._creator) == num_workers - 1, (
434 + f"Expected {num_workers - 1} non-creators, got "
435 + f"{sum(1 for r in regions if not r._creator)}"
436 + )
437 +
438 + for r in regions:
439 + assert not r.mmap_obj.closed
440 + assert r.total_size_bytes == 4 * num_workers * PAGE_SIZE
441 + finally:
442 + for r in regions:
443 + r.cleanup()
444 + _cleanup_file(regions[0].mmap_path)
445 +
446 +
447 +def test_multi_worker_race_shared_memory_visible(iid):
448 + """After a concurrent construction race, MAP_SHARED is intact across all workers."""
449 + num_workers = 4
450 + regions, errors = _race_construct(iid, num_workers=num_workers)
451 + assert not errors
452 + try:
453 + regions[0].mmap_obj[0:1] = b"\xab"
454 + for r in regions[1:]:
455 + assert memoryview(r.mmap_obj)[0:1] == b"\xab"
456 + finally:
457 + for r in regions:
458 + r.cleanup()
459 + _cleanup_file(regions[0].mmap_path)
460 +
461 +
462 +def test_multiprocess_race_construct_and_write(iid):
463 + """N processes race to construct the same SharedOffloadRegion, each writes
464 + fill_value = rank+1 into their slot; parent verifies interleaved layout."""
465 + num_workers = 4
466 + num_blocks = 3
467 + total_bytes = num_blocks * num_workers * PAGE_SIZE
468 +
469 + ctx = get_mp_context()
470 + done_queue = ctx.Queue()
471 + cleanup_queue = ctx.Queue()
472 +
473 + procs = [
474 + ctx.Process(
475 + target=_mp_race_construct_and_write,
476 + args=(
477 + iid,
478 + total_bytes,
479 + num_blocks,
480 + rank,
481 + num_workers,
482 + PAGE_SIZE,
483 + rank + 1,
484 + done_queue,
485 + cleanup_queue,
486 + ),
487 + )
488 + for rank in range(num_workers)
489 + ]
490 + for p in procs:
491 + p.start()
492 +
493 + results = {}
494 + for _ in range(num_workers):
495 + r = done_queue.get(timeout=30)
496 + results[r["rank"]] = r
497 +
498 + for rank, r in results.items():
499 + assert r["error"] is None, f"rank {rank}: {r['error']}"
500 +
501 + # Read the raw file while all workers still hold it open.
502 + mmap_path = f"/dev/shm/vllm_offload_{iid}.mmap"
503 + with open(mmap_path, "rb") as f:
504 + raw = f.read()
505 +
506 + for blk in range(num_blocks):
507 + for w in range(num_workers):
508 + slot_start = (blk * num_workers + w) * PAGE_SIZE
509 + slot = raw[slot_start : slot_start + PAGE_SIZE]
510 + expected = w + 1 # fill_value = rank + 1
511 + assert all(b == expected for b in slot), (
512 + f"block {blk}, worker {w}: expected {expected} but got wrong bytes"
513 + )
514 +
515 + # Unblock all workers to clean up.
516 + for _ in range(num_workers):
517 + cleanup_queue.put(True)
518 + for p in procs:
519 + p.join(timeout=10)
520 + assert p.exitcode == 0
521 +
522 +
523 +# ---------------------------------------------------------------------------
524 +# Cleanup
525 +# ---------------------------------------------------------------------------
526 +
527 +
528 +def test_cleanup_creator_all_effects(iid):
529 + """cleanup() on the creator closes mmap, closes fd, and removes the file."""
530 + r = _make_region(iid)
531 + path = r.mmap_path
532 + fd = r.fd
533 + mmap_obj = r.mmap_obj
534 +
535 + r.cleanup()
536 +
537 + assert mmap_obj.closed, "mmap should be closed after cleanup"
538 + assert not os.path.exists(path), "creator should remove the file"
539 + with pytest.raises(OSError):
540 + os.fstat(fd) # fd should be closed
541 +
542 +
543 +def test_cleanup_non_creator_all_effects(iid):
544 + """cleanup() on a non-creator closes mmap and fd, but leaves the file on disk."""
545 + r0 = _make_region(iid) # creator
546 + r1 = _make_region(iid) # joiner
547 + path = r0.mmap_path
548 + fd1 = r1.fd
549 + mmap_obj1 = r1.mmap_obj
550 + try:
551 + r1.cleanup()
552 +
553 + assert mmap_obj1.closed, "mmap should be closed after cleanup"
554 + assert os.path.exists(path), "non-creator must not remove the file"
555 + with pytest.raises(OSError):
556 + os.fstat(fd1) # fd should be closed
557 + finally:
558 + r0.cleanup()
559 + _cleanup_file(path)
560 +
561 +
562 +def test_cleanup_idempotent(iid):
563 + """Calling cleanup() twice must not raise any exception."""
564 + r = _make_region(iid)
565 + r.cleanup()
566 + r.cleanup() # must be a no-op
567 +
568 +
569 +def test_cleanup_after_create_next_view_releases_mmap(iid):
570 + """cleanup() must close the mmap even after create_next_view was called.
571 + create_next_view returns a view that shares storage with _base; both must be
572 + released before mmap.close() can succeed."""
573 + r = _make_region(iid)
574 + mmap_obj = r.mmap_obj
575 +
576 + t = r.create_next_view(PAGE_SIZE)
577 + del t
578 +
579 + r.cleanup()
580 +
581 + assert mmap_obj.closed, "mmap should be closed after releasing the tensor"
582 +
583 +
584 +# ---------------------------------------------------------------------------
585 +# _wait_for_file_size
586 +# ---------------------------------------------------------------------------
587 +
588 +
589 +def test_wait_for_file_size_already_large_enough(tmp_path):
590 + """_wait_for_file_size must return immediately when file is already big enough."""
591 + fd = os.open(str(tmp_path / "ready.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
592 + try:
593 + os.ftruncate(fd, PAGE_SIZE)
594 + start = time.monotonic()
595 + _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0)
596 + assert time.monotonic() - start < 0.5
597 + finally:
598 + os.close(fd)
599 +
600 +
601 +def test_wait_for_file_size_waits_for_grow(tmp_path):
602 + """_wait_for_file_size must return once a background thread grows the file."""
603 + fd = os.open(str(tmp_path / "grow.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
604 + try:
605 +
606 + def grow():
607 + time.sleep(0.05)
608 + os.ftruncate(fd, PAGE_SIZE)
609 +
610 + t = threading.Thread(target=grow)
611 + t.start()
612 + _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0) # must not raise
613 + t.join()
614 + finally:
615 + os.close(fd)
616 +
617 +
618 +def test_wait_for_file_size_timeout(tmp_path):
619 + """_wait_for_file_size must raise TimeoutError when the file never grows."""
620 + fd = os.open(str(tmp_path / "stuck.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
621 + try:
622 + with pytest.raises(TimeoutError):
623 + _wait_for_file_size(fd, PAGE_SIZE, timeout=0.1)
624 + finally:
625 + os.close(fd)
vllm/v1/kv_offload/cpu/shared_offload_region.py new
+192
@@ -0,0 +1,192 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import mmap
4 +import os
5 +import time
6 +
7 +import torch
8 +
9 +from vllm.logger import init_logger
10 +
11 +logger = init_logger(__name__)
12 +
13 +
14 +def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> None:
15 + """Spin-wait until the file reaches expected_size (creator truncated it)."""
16 + deadline = time.monotonic() + timeout
17 + while True:
18 + if os.fstat(fd).st_size >= expected_size:
19 + return
20 + if time.monotonic() > deadline:
21 + raise TimeoutError(
22 + f"Timed out waiting for mmap file to reach {expected_size} bytes"
23 + )
24 + time.sleep(0.005)
25 +
26 +
27 +class SharedOffloadRegion:
28 + """
29 + Single mmap-backed memory region shared across all workers for a
30 + vLLM instance. Workers coordinate via the filesystem: the first worker
31 + to open the file with O_EXCL becomes the creator and calls ftruncate;
32 + the rest open the existing file and wait until it reaches the expected
33 + size. Each worker then mmap()s the full file.
34 +
35 + File path: /dev/shm/vllm_offload_{instance_id}.mmap
36 + """
37 +
38 + def __init__(
39 + self,
40 + instance_id: str,
41 + total_size_bytes: int,
42 + num_blocks: int,
43 + rank: int | None,
44 + num_workers: int,
45 + cpu_page_size: int,
46 + ) -> None:
47 + self.page_size = mmap.PAGESIZE
48 +
49 + self.total_size_bytes = total_size_bytes
50 + self.mmap_path = f"/dev/shm/vllm_offload_{instance_id}.mmap"
51 + self._creator = False # set True only if this worker creates the file
52 + self.num_blocks = num_blocks
53 + self.rank = rank
54 + # interleaved-layout stride: one row = all workers' data for one block
55 + self._row_stride = cpu_page_size * num_workers
56 + if rank is not None:
57 + # byte offset to this worker's first slot within each block row
58 + self._worker_offset = rank * cpu_page_size
59 + # exclusive upper bound for this worker's area within each row
60 + self._worker_area_end = (rank + 1) * cpu_page_size
61 + try:
62 + # Exclusive create — only one worker succeeds
63 + self.fd: int | None = os.open(
64 + self.mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600
65 + )
66 + os.ftruncate(self.fd, self.total_size_bytes)
67 + self._creator = True
68 + logger.info(
69 + "Created mmap file %s (%.2f GB)",
70 + self.mmap_path,
71 + self.total_size_bytes / 1e9,
72 + )
73 + except FileExistsError:
74 + self.fd = os.open(self.mmap_path, os.O_RDWR)
75 + _wait_for_file_size(self.fd, self.total_size_bytes)
76 + logger.info("Opened existing mmap file %s", self.mmap_path)
77 +
78 + self.mmap_obj: mmap.mmap | None = mmap.mmap(
79 + self.fd,
80 + self.total_size_bytes,
81 + flags=mmap.MAP_SHARED,
82 + prot=mmap.PROT_READ | mmap.PROT_WRITE,
83 + )
84 +
85 + # MADV_POPULATE_WRITE was added in Linux 5.14 (value 23).
86 + _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
87 + if rank is not None:
88 + # Populate only this worker's pages (one slot per block row).
89 + worker_offset = rank * cpu_page_size
90 + _t0 = time.perf_counter()
91 + page_size = self.page_size
92 + for block in range(num_blocks):
93 + raw_offset = block * self._row_stride + worker_offset
94 + aligned_offset = (raw_offset // page_size) * page_size
95 + end = raw_offset + cpu_page_size
96 + aligned_length = end - aligned_offset
97 + self.mmap_obj.madvise(
98 + _MADV_POPULATE_WRITE, aligned_offset, aligned_length
99 + )
100 + logger.debug(
101 + "MADV_POPULATE_WRITE loop: %d blocks in %.3f s",
102 + num_blocks,
103 + time.perf_counter() - _t0,
104 + )
105 + else:
106 + # No rank — populate the entire shared region in one call.
107 + _t0 = time.perf_counter()
108 + self.mmap_obj.madvise(_MADV_POPULATE_WRITE, 0, self.total_size_bytes)
109 + logger.debug(
110 + "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0
111 + )
112 +
113 + self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8)
114 + self._views: list[torch.Tensor] = []
115 + self.is_pinned: bool = False
116 +
117 + def create_next_view(self, tensor_page_size: int) -> torch.Tensor:
118 + """Allocate a strided int8 view for this worker, one canonical tensor.
119 +
120 + Must be called once per canonical tensor. The full mmap layout is:
121 +
122 + worker0_block0 | worker1_block0 | ... | worker{M-1}_block0
123 + worker0_block1 | worker1_block1 | ... | worker{M-1}_block1
124 + ...
125 +
126 + Each worker_block cell is cpu_page_size bytes and holds all canonical
127 + tensors for that worker and block concatenated:
128 + [ tensor0_data | tensor1_data | ... | tensor{L-1}_data ]
129 +
130 + Consecutive rows are separated by row_stride = cpu_page_size * M.
131 +
132 + Returns an int8 tensor of shape (num_blocks, tensor_page_size) with stride
133 + (row_stride, 1). Using int8 keeps stride == bytes, so swap_blocks
134 + address arithmetic works without any dtype conversion.
135 +
136 + Args:
137 + tensor_page_size: Bytes per block for this tensor.
138 + """
139 + assert self.rank is not None
140 + new_offset = self._worker_offset + tensor_page_size
141 + assert new_offset <= self._worker_area_end, (
142 + f"Worker offset {new_offset} exceeds worker area end "
143 + f"{self._worker_area_end} (overflowed by "
144 + f"{new_offset - self._worker_area_end} bytes)"
145 + )
146 + worker_layer_view = torch.as_strided(
147 + self._base,
148 + size=(self.num_blocks, tensor_page_size),
149 + stride=(self._row_stride, 1),
150 + storage_offset=self._worker_offset,
151 + )
152 + self._worker_offset = new_offset
153 + self._views.append(worker_layer_view)
154 + return worker_layer_view
155 +
156 + def cleanup(self) -> None:
157 + if self.is_pinned and self._base is not None:
158 + base_ptr = self._base.data_ptr()
159 + result = torch.cuda.cudart().cudaHostUnregister(base_ptr)
160 + if result.value != 0:
161 + logger.warning(
162 + "cudaHostUnregister failed for rank=%d (code=%d)", self.rank, result
163 + )
164 + self.is_pinned = False
165 + # Release views before _base: each view holds a _base reference and a
166 + # direct StorageImpl reference. Freeing views first lets both refcounts
167 + # drop so the storage (which holds the mmap_obj buffer export) is freed
168 + # before mmap_obj.close() is called below.
169 + if self._views is not None:
170 + self._views.clear()
171 + self._base = None
172 + if self.mmap_obj:
173 + try:
174 + self.mmap_obj.close()
175 + except Exception:
176 + logger.warning("Failed to close mmap_obj", exc_info=True)
177 + self.mmap_obj = None
178 + if self.fd is not None:
179 + try:
180 + os.close(self.fd)
181 + except Exception:
182 + logger.warning("Failed to close fd %s", self.fd, exc_info=True)
183 + self.fd = None
184 + if self._creator and getattr(self, "mmap_path", None):
185 + try:
186 + os.unlink(self.mmap_path)
187 + logger.info("Removed mmap file %s", self.mmap_path)
188 + except Exception:
189 + logger.warning(
190 + "Failed to unlink path %s", self.mmap_path, exc_info=True
191 + )
192 + self._creator = False
vllm/v1/kv_offload/worker/cpu_gpu.py
+102 -48
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import time
4 from collections import deque
5 from dataclasses import dataclass
6
@@ -9,6 +10,7 @@ import torch
10 from vllm import _custom_ops as ops
11 from vllm.logger import init_logger
12 from vllm.utils.platform_utils import is_pin_memory_available
13 +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
14 from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec
15 from vllm.v1.kv_offload.spec import CanonicalKVCacheRef, CanonicalKVCaches
16 from vllm.v1.kv_offload.worker.worker import (
@@ -29,37 +31,76 @@ class Transfer:
31 num_bytes: int
32
33
32 -def expand_block_ids(
34 +def compute_sub_block_ptrs(
35 block_ids: np.ndarray,
36 block_size_factor: int,
37 output: np.ndarray,
38 + tensor: torch.Tensor,
39 skip_count: int = 0,
40 ):
41 """
39 - Convert a list of block IDs to a list of matching block ids,
40 - assuming each block is composed of actual block_size_factor blocks.
41 - Outputs to output tensor.
42 - The first skip_count blocks will be skipped.
43 - Note that skip_count must be less than block_size_factor.
44 -
45 - For example, if block_ids = [0, 1, 3] and block_size_factor = 4,
46 - then it yields [0, 1, 2, 3, 4, 5, 6, 7, 12, 13, 14, 15]
47 - since 0 maps to [0, 1, 2, 3]
48 - 1 maps to [4, 5, 6, 7]
49 - and 3 maps to [12, 13, 14, 15]
42 + Compute byte pointers for sub-blocks of the given block IDs.
43 +
44 + Each block in block_ids contains block_size_factor sub-blocks.
45 + The pointer for sub-block j of block b is:
46 + base_ptr + b * row_stride + j * sub_block_size
47 +
48 + where sub_block_size = tensor.shape[1] // block_size_factor (gpu page size).
49 +
50 + This handles tensors where row_stride != block_size_factor * sub_block_size
51 + (e.g. non-contiguous CPU tensors).
52 +
53 + Args:
54 + block_ids: array of block IDs at the tensor's native granularity.
55 + block_size_factor: number of sub-blocks per block.
56 + output: pre-allocated int64 array to write pointers into.
57 + tensor: the source or destination tensor.
58 + skip_count: sub-blocks to skip in the first block.
59 """
60 assert skip_count < block_size_factor
61
53 - first_range = np.arange(skip_count, block_size_factor)
54 - full_range = np.arange(0, block_size_factor)
55 -
56 - output_idx = 0
57 - for i, block_id in enumerate(block_ids):
58 - base_block_id = block_id * block_size_factor
59 - indices = first_range if i == 0 else full_range
60 - output_end_idx = output_idx + len(indices)
61 - output[output_idx:output_end_idx] = base_block_id + indices
62 - output_idx = output_end_idx
62 + num_sub_blocks = len(output)
63 + base_ptr = tensor.data_ptr()
64 + row_stride = tensor.stride(0)
65 +
66 + if block_size_factor == 1:
67 + # Fast path: 1:1 mapping, no sub-block expansion needed.
68 + output[:] = base_ptr + block_ids[:num_sub_blocks] * row_stride
69 + return
70 +
71 + # Vectorized expansion for block_size_factor > 1.
72 + assert tensor.shape[1] % block_size_factor == 0
73 + sub_block_size = tensor.shape[1] // block_size_factor
74 + sub_offsets = np.arange(block_size_factor, dtype=np.int64) * sub_block_size
75 + # (num_blocks, 1) + (1, block_size_factor) -> (num_blocks, block_size_factor)
76 + all_ptrs = (
77 + base_ptr + block_ids.astype(np.int64)[:, np.newaxis] * row_stride
78 + ) + sub_offsets[np.newaxis, :]
79 + # Flatten and apply skip_count / truncation
80 + flat = all_ptrs.ravel()
81 + output[:] = flat[skip_count : skip_count + num_sub_blocks]
82 +
83 +
84 +def pin_mmap_region(region: SharedOffloadRegion) -> None:
85 + """Register the entire mmap as CUDA pinned memory via cudaHostRegister."""
86 + rank = region.rank
87 +
88 + base_ptr = region._base.data_ptr()
89 + result = torch.cuda.cudart().cudaHostRegister(base_ptr, region.total_size_bytes, 0)
90 + if result.value != 0:
91 + logger.warning(
92 + "cudaHostRegister failed for rank=%d (code=%d) — "
93 + "transfers will still work but may be slower (unpinned DMA)",
94 + rank,
95 + result,
96 + )
97 + else:
98 + logger.debug(
99 + "cudaHostRegister rank=%d %.2f GB",
100 + rank,
101 + region.total_size_bytes / 1e9,
102 + )
103 + region.is_pinned = True
104
105
106 class SingleDirectionOffloadingHandler(OffloadingHandler):
@@ -149,13 +190,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
190 # list of CUDA events available for re-use
191 self._event_pool: list[torch.Event] = []
192
152 - # Pre-compute base pointers and block sizes for batch copies.
153 - self._src_base_ptrs = np.array(
154 - [t.data_ptr() for t in self.src_tensors], dtype=np.int64
155 - )
156 - self._dst_base_ptrs = np.array(
157 - [t.data_ptr() for t in self.dst_tensors], dtype=np.int64
158 - )
193 + # Pre-compute block sizes for batch copies.
194 self._block_size_in_bytes_arr = np.array(
195 self.tensor_block_size_in_bytes, dtype=np.int64
196 )
@@ -176,17 +211,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
211
212 assert dst_sub_block_count == src_sub_block_count - src_sub_blocks_to_skip
213
179 - src_block_ids = np.empty(dst_sub_block_count, dtype=np.int64)
180 - dst_block_ids = np.empty(dst_sub_block_count, dtype=np.int64)
181 - expand_block_ids(
182 - src_blocks,
183 - self.src_block_size_factor,
184 - src_block_ids,
185 - skip_count=src_sub_blocks_to_skip,
186 - )
187 - expand_block_ids(dst_blocks, self.dst_block_size_factor, dst_block_ids)
188 -
189 - # Build flat pointer arrays for all tensors × all block pairs.
214 num_pairs = dst_sub_block_count
215 num_tensors = len(self.src_tensors)
216 total = num_pairs * num_tensors
@@ -198,8 +222,19 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
222 for t_idx, bsz in enumerate(self._block_size_in_bytes_arr):
223 start = t_idx * num_pairs
224 end = start + num_pairs
201 - all_src[start:end] = self._src_base_ptrs[t_idx] + src_block_ids * bsz
202 - all_dst[start:end] = self._dst_base_ptrs[t_idx] + dst_block_ids * bsz
225 + compute_sub_block_ptrs(
226 + block_ids=src_blocks,
227 + block_size_factor=self.src_block_size_factor,
228 + output=all_src[start:end],
229 + tensor=self.src_tensors[t_idx],
230 + skip_count=src_sub_blocks_to_skip,
231 + )
232 + compute_sub_block_ptrs(
233 + block_ids=dst_blocks,
234 + block_size_factor=self.dst_block_size_factor,
235 + output=all_dst[start:end],
236 + tensor=self.dst_tensors[t_idx],
237 + )
238 all_sizes[start:end] = bsz
239
240 batch_src = torch.from_numpy(all_src)
@@ -281,6 +316,8 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
316 self._transfer_events.clear()
317 self._stream_pool.clear()
318 self._event_pool.clear()
319 + self.src_tensors.clear()
320 + self.dst_tensors.clear()
321
322
323 class CpuGpuOffloadingHandlers:
@@ -289,9 +326,14 @@ class CpuGpuOffloadingHandlers:
326 kv_caches: CanonicalKVCaches,
327 block_size_factor: int,
328 num_cpu_blocks: int,
329 + mmap_region: SharedOffloadRegion | None = None,
330 ):
331 pin_memory = is_pin_memory_available()
332 logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
333 + self._mmap_region = mmap_region
334 + if mmap_region is not None and pin_memory:
335 + pin_mmap_region(mmap_region)
336 +
337 gpu_tensors: list[torch.Tensor] = []
338 cpu_tensors: list[torch.Tensor] = []
339 for kv_cache_tensor in kv_caches.tensors:
@@ -300,12 +342,24 @@ class CpuGpuOffloadingHandlers:
342 (-1, gpu_page_size_bytes)
343 )
344 cpu_page_size_bytes = gpu_page_size_bytes * block_size_factor
303 - cpu_tensor = torch.zeros(
304 - (num_cpu_blocks, cpu_page_size_bytes),
305 - dtype=torch.int8,
306 - device="cpu",
307 - pin_memory=pin_memory,
308 - )
345 +
346 + if mmap_region is not None:
347 + cpu_tensor = mmap_region.create_next_view(cpu_page_size_bytes)
348 + else:
349 + t0 = time.monotonic()
350 + cpu_tensor = torch.zeros(
351 + (num_cpu_blocks, cpu_page_size_bytes),
352 + dtype=torch.int8,
353 + device="cpu",
354 + pin_memory=pin_memory,
355 + )
356 + logger.debug(
357 + "torch.zeros pinned tensor %d×%d (%.2f GB): %.3f s",
358 + num_cpu_blocks,
359 + cpu_page_size_bytes,
360 + num_cpu_blocks * cpu_page_size_bytes / 1e9,
361 + time.monotonic() - t0,
362 + )
363
364 gpu_tensors.append(gpu_tensor)
365 cpu_tensors.append(cpu_tensor)