[Bugfix][ROCm][KV Offload] Use private pinned tensors for CPU KV offload (#57160)
Signed-off-by: Yuzhou Ren <ryzcqu@outlook.com>
Ren Yuzhou committed
Sep 18, 2026 at 10:35 UTC
7b942936276d59cc1912185a442b97dea8d2386c
2 files changed
+33
-18
tests/v1/kv_offload/test_factory.py
+32
-17
@@ -275,11 +275,12 @@ def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkey
275
276
@pytest.mark.parametrize("world_size", [2, 4, 8])
277
def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: int):
278
- # On shared-region (CUDA-alike) platforms the default spec now honors
278
+ # On CUDA the default spec honors
279
# replicated layout: a single MLA copy (num_copies=1), matching tiering.
280
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
281
282
monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
283
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
284
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
285
spec = _create_spec(
286
cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
@@ -295,19 +296,19 @@ def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: in
296
assert spec.num_chunks == 8
297
298
299
+@pytest.mark.parametrize("rocm", [False, True], ids=["non-cuda-alike", "rocm"])
300
@pytest.mark.parametrize("world_size", [2, 4, 8])
301
def test_cpu_spec_replicated_disabled_without_shared_region(
300
- monkeypatch, world_size: int
302
+ monkeypatch, world_size: int, rocm: bool
303
):
302
- # Data-loss guard: non-CUDA-alike platforms keep a per-rank private pinned
304
+ # Data-loss guard: platforms without shared regions keep a private pinned
305
# tensor (no shared medium), so replicated layout MUST stay off. Otherwise
306
# the rank-0 writer gate would ack rank>0 stores without writing, leaving
307
# those private buffers empty and corrupting subsequent loads.
308
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
309
308
- monkeypatch.setattr(
309
- cpu_spec_module.current_platform, "is_cuda_alike", lambda: False
310
- )
310
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm)
311
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
312
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
313
spec = _create_spec(
314
cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2,
@@ -324,17 +325,26 @@ def test_cpu_spec_replicated_disabled_without_shared_region(
325
326
327
@pytest.mark.parametrize("config_replicated", [True, False])
327
-@pytest.mark.parametrize("cuda_alike", [True, False])
328
+@pytest.mark.parametrize(
329
+ ("cuda_alike", "rocm", "shared_region"),
330
+ [(True, False, True), (True, True, False), (False, False, False)],
331
+ ids=["cuda", "rocm", "non-cuda-alike"],
332
+)
333
def test_cpu_spec_replicated_layout_truth_matrix(
329
- monkeypatch, cuda_alike: bool, config_replicated: bool
334
+ monkeypatch,
335
+ cuda_alike: bool,
336
+ rocm: bool,
337
+ shared_region: bool,
338
+ config_replicated: bool,
339
):
340
# replicated_layout is enabled iff the config gate passes AND the deployment
332
- # actually allocates on the shared region (CUDA-alike).
341
+ # actually allocates on the shared region.
342
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
343
344
monkeypatch.setattr(
345
cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike
346
)
347
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
348
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
349
spec = _create_spec(
350
cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
@@ -344,10 +354,10 @@ def test_cpu_spec_replicated_layout_truth_matrix(
354
)
355
356
assert isinstance(spec, CPUOffloadingSpec)
347
- assert spec.replicated_layout is (cuda_alike and config_replicated)
357
+ assert spec.replicated_layout is (shared_region and config_replicated)
358
359
350
-def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
360
+def test_cpu_spec_create_worker_uses_mmap_on_cuda(monkeypatch):
361
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
362
363
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
@@ -371,6 +381,7 @@ def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
381
return MagicMock()
382
383
monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
384
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
385
monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
386
monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
387
monkeypatch.setattr(
@@ -386,7 +397,10 @@ def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
397
assert worker_calls[0]["mmap_region"] is region
398
399
389
-def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch):
400
+@pytest.mark.parametrize("rocm", [False, True], ids=["non-cuda-alike", "rocm"])
401
+def test_cpu_spec_create_worker_uses_tensor_path_without_shared_region(
402
+ monkeypatch, rocm
403
+):
404
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
405
406
spec = _create_spec(worker_kv_bytes_per_block=4096, world_size=4)
@@ -403,15 +417,14 @@ def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch):
417
worker_calls.append(kwargs)
418
return MagicMock()
419
406
- monkeypatch.setattr(
407
- cpu_spec_module.current_platform, "is_cuda_alike", lambda: False
408
- )
420
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm)
421
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
422
monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
423
monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
424
425
spec.create_worker(MagicMock())
426
414
- # Non-CUDA-alike platforms keep the per-rank pinned-tensor path.
427
+ # ROCm and non-CUDA-alike platforms use the per-rank pinned-tensor path.
428
assert region_calls == []
429
assert worker_calls[0]["mmap_region"] is None
430
@@ -420,7 +433,7 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
433
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
434
435
# worker_kv_bytes_per_block=0 yields num_chunks=0; a zero-byte region cannot
423
- # be mmap'd, so even on CUDA-alike this must fall back to the tensor path.
436
+ # be mmap'd, so even on CUDA this must fall back to the tensor path.
437
spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4)
438
assert isinstance(spec, CPUOffloadingSpec)
439
assert spec.num_chunks == 0
@@ -429,6 +442,7 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
442
worker_calls: list[dict[str, Any]] = []
443
444
monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
445
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
446
monkeypatch.setattr(
447
cpu_spec_module,
448
"SharedOffloadRegion",
@@ -461,6 +475,7 @@ def test_cpu_spec_create_worker_rank_assignment(
475
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
476
477
monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
478
+ monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
479
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
480
spec = _create_spec(
481
cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
vllm/v1/kv_offload/cpu/spec.py
+1
-1
@@ -162,7 +162,7 @@ class CPUOffloadingSpec(OffloadingSpec):
162
def _uses_shared_region(self) -> bool:
163
"""Whether the worker CPU buffer is the shared mmap region (vs a private
164
per-rank tensor); replicated-layout dedup is gated on this being True."""
165
- return current_platform.is_cuda_alike()
165
+ return current_platform.is_cuda_alike() and not current_platform.is_rocm()
166
167
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
168
mmap_region: SharedOffloadRegion | None = None