[KV Offload] Move CPUOffloadingSpec onto SharedOffloadRegion (#50094)
Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: Cursor Agent <cursor-agent@cursor.com>
Chang Guo committed
Jul 29, 2026 at 03:05 UTC
542a8fad6da059a0fd4a8b0a03b4476e9ec9b66a
2 files changed
+132
-5
tests/v1/kv_offload/test_factory.py
+111
-4
@@ -136,8 +136,12 @@ def test_create_cpu_offloading_spec():
136
137
138
def test_cpu_spec_sizes_normalized_worker_layout():
139
+ # The CPU spec now rounds the offloaded row up to the mmap page size
140
+ # (matching the shared region), so kv_bytes_per_chunk picks up padding
141
+ # while cpu_page_size_per_worker stays the un-padded per-worker slot.
142
+ alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
143
spec = _create_spec(
140
- cpu_bytes_to_use=1920,
144
+ cpu_bytes_to_use=alignment * 3,
145
worker_kv_bytes_per_block=16,
146
blocks_per_chunk=2,
147
world_size=6,
@@ -147,8 +151,8 @@ def test_cpu_spec_sizes_normalized_worker_layout():
151
152
assert isinstance(spec, CPUOffloadingSpec)
153
assert spec.cpu_page_size_per_worker == 32
150
- assert spec.kv_bytes_per_chunk == 192
151
- assert spec.num_blocks == 10
154
+ assert spec.kv_bytes_per_chunk == alignment
155
+ assert spec.num_blocks == 3
156
157
158
def test_cpu_spec_zero_worker_bytes_produces_empty_cache():
@@ -267,7 +271,9 @@ def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkey
271
272
@pytest.mark.parametrize("world_size", [2, 4, 8])
273
def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
270
- worker_kv_bytes_per_block = 4096
274
+ # Use a page-multiple per-worker block so the row size is unaffected by the
275
+ # page-alignment rounding regardless of the host mmap page size.
276
+ worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
277
spec = _create_spec(
278
cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2,
279
worker_kv_bytes_per_block=worker_kv_bytes_per_block,
@@ -282,6 +288,107 @@ def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
288
assert spec.num_blocks == 2
289
290
291
+def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
292
+ import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
293
+
294
+ worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
295
+ spec = _create_spec(
296
+ cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
297
+ worker_kv_bytes_per_block=worker_kv_bytes_per_block,
298
+ world_size=4,
299
+ )
300
+ assert isinstance(spec, CPUOffloadingSpec)
301
+
302
+ region = MagicMock()
303
+ region_calls: list[dict[str, Any]] = []
304
+ worker_calls: list[dict[str, Any]] = []
305
+
306
+ def fake_region_ctor(**kwargs):
307
+ region_calls.append(kwargs)
308
+ return region
309
+
310
+ def fake_worker_ctor(**kwargs):
311
+ worker_calls.append(kwargs)
312
+ return MagicMock()
313
+
314
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
315
+ monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
316
+ monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
317
+ monkeypatch.setattr(
318
+ cpu_spec_module.torch.accelerator, "current_device_index", lambda: 5
319
+ )
320
+
321
+ kv_caches = MagicMock()
322
+ spec.create_worker(kv_caches)
323
+
324
+ # rank folds the physical device index into [0, world_size): 5 % 4 == 1.
325
+ assert region_calls[0]["rank"] == 1
326
+ assert region_calls[0]["engine_id"] == "test-engine"
327
+ assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block * 4
328
+ assert worker_calls[0]["kv_caches"] is kv_caches
329
+ assert worker_calls[0]["mmap_region"] is region
330
+
331
+
332
+def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch):
333
+ import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
334
+
335
+ spec = _create_spec(worker_kv_bytes_per_block=4096, world_size=4)
336
+ assert isinstance(spec, CPUOffloadingSpec)
337
+
338
+ region_calls: list[dict[str, Any]] = []
339
+ worker_calls: list[dict[str, Any]] = []
340
+
341
+ def fake_region_ctor(**kwargs):
342
+ region_calls.append(kwargs)
343
+ return MagicMock()
344
+
345
+ def fake_worker_ctor(**kwargs):
346
+ worker_calls.append(kwargs)
347
+ return MagicMock()
348
+
349
+ monkeypatch.setattr(
350
+ cpu_spec_module.current_platform, "is_cuda_alike", lambda: False
351
+ )
352
+ monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
353
+ monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
354
+
355
+ spec.create_worker(MagicMock())
356
+
357
+ # Non-CUDA-alike platforms keep the per-rank pinned-tensor path.
358
+ assert region_calls == []
359
+ assert worker_calls[0]["mmap_region"] is None
360
+
361
+
362
+def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
363
+ import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
364
+
365
+ # worker_kv_bytes_per_block=0 yields num_blocks=0; a zero-byte region cannot
366
+ # be mmap'd, so even on CUDA-alike this must fall back to the tensor path.
367
+ spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4)
368
+ assert isinstance(spec, CPUOffloadingSpec)
369
+ assert spec.num_blocks == 0
370
+
371
+ region_calls: list[dict[str, Any]] = []
372
+ worker_calls: list[dict[str, Any]] = []
373
+
374
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
375
+ monkeypatch.setattr(
376
+ cpu_spec_module,
377
+ "SharedOffloadRegion",
378
+ lambda **kwargs: region_calls.append(kwargs),
379
+ )
380
+ monkeypatch.setattr(
381
+ cpu_spec_module,
382
+ "CPUOffloadingWorker",
383
+ lambda **kwargs: worker_calls.append(kwargs),
384
+ )
385
+
386
+ spec.create_worker(MagicMock())
387
+
388
+ assert region_calls == []
389
+ assert worker_calls[0]["mmap_region"] is None
390
+
391
+
392
def test_offloading_spec_has_replicated_layout_default():
393
spec = SingleArgExternalOffloadingSpec(_make_offloading_config())
394
assert spec.replicated_layout is False
vllm/v1/kv_offload/cpu/spec.py
+21
-1
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import Any
4
5
+import torch
6
from typing_extensions import override
7
8
from vllm.platforms import current_platform
@@ -20,10 +21,11 @@ from vllm.v1.kv_offload.config import OffloadingConfig
21
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
22
from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
23
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
24
+from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
25
26
27
class CPUOffloadingSpec(OffloadingSpec):
26
- BLOCK_SIZE_ALIGNMENT = 1
28
+ BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
29
SUPPORTS_REPLICATED_LAYOUT = False
30
31
@classmethod
@@ -143,10 +145,28 @@ class CPUOffloadingSpec(OffloadingSpec):
145
return self._manager
146
147
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
148
+ mmap_region: SharedOffloadRegion | None = None
149
+ # num_blocks == 0 would size the region to zero bytes, which cannot be
150
+ # mmap'd; fall back to the tensor path (empty tensors) as before.
151
+ if current_platform.is_cuda_alike() and self.num_blocks > 0:
152
+ # Back each worker's CPU buffer with a private slot in a single
153
+ # shared mmap region instead of a per-rank pinned tensor. Fold the
154
+ # global physical device index into this replica's
155
+ # [0, world_size) slot range.
156
+ world_size = self.config.parallel.world_size
157
+ rank = torch.accelerator.current_device_index() % world_size
158
+ mmap_region = SharedOffloadRegion(
159
+ engine_id=self.config.engine_id,
160
+ num_blocks=self.num_blocks,
161
+ rank=rank,
162
+ kv_bytes_per_block=self.kv_bytes_per_chunk,
163
+ cpu_page_size=self.cpu_page_size_per_worker,
164
+ )
165
return CPUOffloadingWorker(
166
kv_caches=kv_caches,
167
blocks_per_chunk=self.blocks_per_chunk,
168
num_cpu_blocks=self.num_blocks,
169
+ mmap_region=mmap_region,
170
)
171
172
@override