[KV Offload] Enable single-copy MLA layout for CPUOffloadingSpec (#50301)
Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: Cursor Agent <cursor-agent@cursor.com>
Chang Guo committed
Jul 30, 2026 at 20:20 UTC
541128bdeb623d9acee8abee53facb4e46848009
3 files changed
+119
-17
tests/v1/kv_offload/test_factory.py
+98
-5
@@ -270,9 +270,40 @@ def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkey
270
271
272
@pytest.mark.parametrize("world_size", [2, 4, 8])
273
-def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
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.
273
+def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: int):
274
+ # On shared-region (CUDA-alike) platforms the default spec now honors
275
+ # replicated layout: a single MLA copy (num_copies=1), matching tiering.
276
+ import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
277
+
278
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
279
+ worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
280
+ spec = _create_spec(
281
+ cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
282
+ worker_kv_bytes_per_block=worker_kv_bytes_per_block,
283
+ world_size=world_size,
284
+ replicated_layout=True,
285
+ )
286
+
287
+ assert isinstance(spec, CPUOffloadingSpec)
288
+ assert spec.replicated_layout is True
289
+ assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
290
+ assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block
291
+ assert spec.num_blocks == 8
292
+
293
+
294
+@pytest.mark.parametrize("world_size", [2, 4, 8])
295
+def test_cpu_spec_replicated_disabled_without_shared_region(
296
+ monkeypatch, world_size: int
297
+):
298
+ # Data-loss guard: non-CUDA-alike platforms keep a per-rank private pinned
299
+ # tensor (no shared medium), so replicated layout MUST stay off. Otherwise
300
+ # the rank-0 writer gate would ack rank>0 stores without writing, leaving
301
+ # those private buffers empty and corrupting subsequent loads.
302
+ import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
303
+
304
+ monkeypatch.setattr(
305
+ cpu_spec_module.current_platform, "is_cuda_alike", lambda: False
306
+ )
307
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
308
spec = _create_spec(
309
cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2,
@@ -288,6 +319,30 @@ def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
319
assert spec.num_blocks == 2
320
321
322
+@pytest.mark.parametrize("config_replicated", [True, False])
323
+@pytest.mark.parametrize("cuda_alike", [True, False])
324
+def test_cpu_spec_replicated_layout_truth_matrix(
325
+ monkeypatch, cuda_alike: bool, config_replicated: bool
326
+):
327
+ # replicated_layout is enabled iff the config gate passes AND the deployment
328
+ # actually allocates on the shared region (CUDA-alike).
329
+ import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
330
+
331
+ monkeypatch.setattr(
332
+ cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike
333
+ )
334
+ worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
335
+ spec = _create_spec(
336
+ cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
337
+ worker_kv_bytes_per_block=worker_kv_bytes_per_block,
338
+ world_size=4,
339
+ replicated_layout=config_replicated,
340
+ )
341
+
342
+ assert isinstance(spec, CPUOffloadingSpec)
343
+ assert spec.replicated_layout is (cuda_alike and config_replicated)
344
+
345
+
346
def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
347
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
348
@@ -321,8 +376,6 @@ def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
376
kv_caches = MagicMock()
377
spec.create_worker(kv_caches)
378
324
- # rank folds the physical device index into [0, world_size): 5 % 4 == 1.
325
- assert region_calls[0]["rank"] == 1
379
assert region_calls[0]["engine_id"] == "test-engine"
380
assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block * 4
381
assert worker_calls[0]["kv_caches"] is kv_caches
@@ -389,6 +442,46 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
442
assert worker_calls[0]["mmap_region"] is None
443
444
445
+@pytest.mark.parametrize(
446
+ ("replicated_layout", "device_index", "world_size", "expected_rank"),
447
+ [
448
+ (True, 5, 4, 0), # replicated: always slot 0
449
+ (True, 0, 4, 0), # replicated: slot 0 regardless of device
450
+ (False, 5, 4, 1), # non-replicated: 5 % 4 == 1
451
+ (False, 7, 4, 3), # non-replicated: 7 % 4 == 3
452
+ ],
453
+)
454
+def test_cpu_spec_create_worker_rank_assignment(
455
+ monkeypatch, replicated_layout, device_index, world_size, expected_rank
456
+):
457
+ import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
458
+
459
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
460
+ worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
461
+ spec = _create_spec(
462
+ cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
463
+ worker_kv_bytes_per_block=worker_kv_bytes_per_block,
464
+ world_size=world_size,
465
+ replicated_layout=replicated_layout,
466
+ )
467
+
468
+ region_calls: list[dict[str, Any]] = []
469
+
470
+ def fake_region_ctor(**kwargs):
471
+ region_calls.append(kwargs)
472
+ return MagicMock()
473
+
474
+ monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
475
+ monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", MagicMock())
476
+ monkeypatch.setattr(
477
+ cpu_spec_module.torch.accelerator, "current_device_index", lambda: device_index
478
+ )
479
+
480
+ spec.create_worker(MagicMock())
481
+
482
+ assert region_calls[0]["rank"] == expected_rank
483
+
484
+
485
def test_offloading_spec_has_replicated_layout_default():
486
spec = SingleArgExternalOffloadingSpec(_make_offloading_config())
487
assert spec.replicated_layout is False
vllm/v1/kv_offload/cpu/spec.py
+14
-11
@@ -26,7 +26,6 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
26
27
class CPUOffloadingSpec(OffloadingSpec):
28
BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
29
- SUPPORTS_REPLICATED_LAYOUT = False
29
30
@classmethod
31
def build_metric_definitions(
@@ -88,9 +87,7 @@ class CPUOffloadingSpec(OffloadingSpec):
87
self.num_blocks = 0
88
self.kv_bytes_per_chunk = 0
89
self.cpu_page_size_per_worker = 0
91
- self.replicated_layout = (
92
- config.replicated_layout and self.SUPPORTS_REPLICATED_LAYOUT
93
- )
90
+ self.replicated_layout = config.replicated_layout and self._uses_shared_region()
91
if config.worker_kv_bytes_per_block > 0 and world_size > 0:
92
num_copies = 1 if self.replicated_layout else world_size
93
kv_bytes_per_block = config.worker_kv_bytes_per_block * num_copies
@@ -144,17 +141,23 @@ class CPUOffloadingSpec(OffloadingSpec):
141
)
142
return self._manager
143
144
+ def _uses_shared_region(self) -> bool:
145
+ """Whether the worker CPU buffer is the shared mmap region (vs a private
146
+ per-rank tensor); replicated-layout dedup is gated on this being True."""
147
+ return current_platform.is_cuda_alike()
148
+
149
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
150
mmap_region: SharedOffloadRegion | None = None
151
# num_blocks == 0 would size the region to zero bytes, which cannot be
152
# 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
153
+ if self._uses_shared_region() and self.num_blocks > 0:
154
+ # Replicated layout puts all ranks on slot 0 (single MLA copy);
155
+ # otherwise each rank takes its own slot by physical device index.
156
+ if self.replicated_layout:
157
+ rank = 0
158
+ else:
159
+ world_size = self.config.parallel.world_size
160
+ rank = torch.accelerator.current_device_index() % world_size
161
mmap_region = SharedOffloadRegion(
162
engine_id=self.config.engine_id,
163
num_blocks=self.num_blocks,
vllm/v1/kv_offload/tiering/spec.py
+7
-1
@@ -76,7 +76,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
76
"""
77
78
BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
79
- SUPPORTS_REPLICATED_LAYOUT = True
79
80
@classmethod
81
@override
@@ -234,6 +233,13 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
233
234
return self._manager
235
236
+ @override
237
+ def _uses_shared_region(self) -> bool:
238
+ # Tiering always allocates on the shared region (every platform), so the
239
+ # replicated-layout gate must not be narrowed by the CPU spec's
240
+ # CUDA-alike check.
241
+ return True
242
+
243
@override
244
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
245
world_size = self.config.parallel.world_size