[Bugfix][ROCm] Use batch DMA for CPU KV cache loads (#49843)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Jul 27, 2026 at 04:18 UTC
394beb633b0b5b5d68aed3d2f2b5c1477e756d80
2 files changed
+16
-2
tests/v1/kv_offload/cpu/test_gpu_worker.py
+14
@@ -7,6 +7,7 @@ import uuid
7
import pytest
8
import torch
9
10
+from vllm import _custom_ops as ops
11
from vllm.platforms import current_platform
12
from vllm.utils.math_utils import round_up
13
from vllm.utils.torch_utils import set_random_seed
@@ -16,6 +17,7 @@ from vllm.v1.kv_offload.base import (
17
CanonicalKVCacheTensor,
18
GPULoadStoreSpec,
19
)
20
+from vllm.v1.kv_offload.cpu import gpu_worker
21
from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
22
from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
23
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
@@ -32,6 +34,18 @@ NUM_MAPPINGS = [3]
34
NUM_MAPPINGS_PER_GROUP = [2]
35
36
37
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test")
38
+def test_rocm_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None:
39
+ monkeypatch.setattr(gpu_worker, "HAS_TRITON", True)
40
+ monkeypatch.setattr(gpu_worker.current_platform, "is_xpu", lambda: False)
41
+ monkeypatch.setattr(gpu_worker.current_platform, "is_rocm", lambda: True)
42
+
43
+ refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=512)]]
44
+ assert gpu_worker._select_swap_blocks_fn(refs, gpu_to_cpu=False) is (
45
+ ops.swap_blocks_batch
46
+ )
47
+
48
+
49
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
50
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
51
@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
vllm/v1/kv_offload/cpu/gpu_worker.py
+2
-2
@@ -41,10 +41,10 @@ def _select_swap_blocks_fn(
41
if gpu_to_cpu:
42
return ops.swap_blocks_batch
43
# Fall back to the C++ DMA path on platforms where Triton isn't usable
44
- # (e.g. ROCm builds without Triton) or where GPU kernels cannot directly
44
+ # (e.g. ROCm host mappings) or where GPU kernels cannot directly
45
# dereference CPU pointers (XPU lacks CUDA's unified virtual address space,
46
# so the Triton kernel's tl.load(cpu_ptr) is invalid on XPU).
47
- if not HAS_TRITON or current_platform.is_xpu():
47
+ if not HAS_TRITON or current_platform.is_xpu() or current_platform.is_rocm():
48
return ops.swap_blocks_batch
49
page_sizes = [r.page_size_bytes for g in kv_cache_groups_data_refs for r in g]
50
# Triton wins only on small, 8-byte-aligned payloads.