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 4, 05:00 UTC
[ROCm] Enable SimpleCPUOffloadConnector on ROCm (#40549)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Hongxia Yang committed
May 6, 2026 at 23:52 UTC
20cac26b197a2e9530a787a5c1dbf5163533b2ff
4 files changed
+90
-21
csrc/cache_kernels.cu
+27
-9
@@ -97,13 +97,13 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
97
98
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
99
100
- // Use cuMemcpyBatchAsync (CUDA 12.8+) to submit all copies in a single
101
- // driver call, amortizing per-copy submission overhead.
102
- // int64_t and CUdeviceptr/size_t are both 8 bytes on 64-bit platforms,
103
- // so we reinterpret_cast the tensor data directly to avoid copies.
104
- static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
100
+ // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a
101
+ // single driver call, amortizing per-copy submission overhead. int64_t
102
+ // and CUdeviceptr/void*/size_t are all 8 bytes on 64-bit platforms, so we
103
+ // reinterpret_cast the tensor data directly to avoid copies.
104
static_assert(sizeof(size_t) == sizeof(int64_t));
105
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
106
+ static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
107
// Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that
108
// binaries compiled with CUDA 12.8+ still work on older drivers, and
109
// we avoid the CUDA 13.0 header remapping (#define to _v2 signature).
@@ -134,12 +134,30 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
134
&fail_idx, static_cast<CUstream>(stream));
135
TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
136
fail_idx, " with error ", result);
137
- } else
137
+ return;
138
+ }
139
+#elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
140
+ // ROCm 7.1+ exposes hipMemcpyBatchAsync. The 7.2.1 implementation early-
141
+ // returns hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @
142
+ // rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so call with
143
+ // numAttrs=0.
144
+ {
145
+ hipMemcpyAttributes attr = {};
146
+ size_t attrs_idx = 0;
147
+ size_t fail_idx = 0;
148
+ hipError_t result = hipMemcpyBatchAsync(
149
+ reinterpret_cast<void**>(dst_data), reinterpret_cast<void**>(src_data),
150
+ reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
151
+ &attrs_idx, 0, &fail_idx, static_cast<hipStream_t>(stream));
152
+ TORCH_CHECK(result == hipSuccess, "hipMemcpyBatchAsync failed at index ",
153
+ fail_idx, " with error ", result);
154
+ return;
155
+ }
156
#endif
157
{
140
- // Fallback for CUDA < 12.8, older drivers, and ROCm:
141
- // individual async copies.
142
- // cudaMemcpyDefault lets the driver infer direction from pointer types.
158
+ // Fallback for CUDA < 12.8, older CUDA drivers, and ROCm < 7.1:
159
+ // individual async copies. cudaMemcpyDefault lets the driver infer
160
+ // direction from pointer types.
161
for (int64_t i = 0; i < n; i++) {
162
cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
163
reinterpret_cast<void*>(src_data[i]),
tests/v1/simple_kv_offload/test_integration.py
+2
-2
@@ -10,8 +10,8 @@ from vllm import LLM, SamplingParams, TokensPrompt
10
from vllm.config import KVTransferConfig
11
from vllm.platforms import current_platform
12
13
-if not current_platform.is_cuda():
14
- pytest.skip("Requires CUDA", allow_module_level=True)
13
+if not current_platform.is_cuda_alike():
14
+ pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
15
16
# Small models for default CI / local runs (accuracy only).
17
SMALL_MODELS = [
tests/v1/simple_kv_offload/test_scheduler.py
+7
-1
@@ -186,7 +186,13 @@ def make_request(
186
if request_id is None:
187
request_id = f"req-{_req_counter}"
188
189
- num_tokens = num_blocks * BLOCK_SIZE
189
+ # Add one extra token beyond the last full block so that
190
+ # ``max_cache_hit_length = num_tokens - 1`` (see
191
+ # KVCacheManager.get_computed_blocks) does not truncate the final
192
+ # full block: ``find_longest_cache_hit`` uses
193
+ # ``max_length // block_size`` and would otherwise drop one block
194
+ # when the prompt is an exact multiple of block_size.
195
+ num_tokens = num_blocks * BLOCK_SIZE + 1
196
start = _req_counter * 10000
197
prompt_token_ids = list(range(start, start + num_tokens))
198
sampling_params = SamplingParams(max_tokens=1)
vllm/v1/simple_kv_offload/cuda_mem_ops.py
+54
-9
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Low-level CUDA memory helpers: pinning and batch DMA transfers."""
3
+"""Low-level CUDA/HIP memory helpers: pinning and batch DMA transfers."""
4
5
import ctypes
6
from typing import Any, NamedTuple
@@ -9,6 +9,7 @@ import numpy as np
9
import torch
10
11
from vllm.logger import init_logger
12
+from vllm.platforms import current_platform
13
14
logger = init_logger(__name__)
15
@@ -39,7 +40,7 @@ class _CUmemcpyAttributes(ctypes.Structure):
40
41
42
_BATCH_MEMCPY_FUNC_TYPE = ctypes.CFUNCTYPE(
42
- ctypes.c_uint, # CUresult
43
+ ctypes.c_uint, # CUresult / hipError_t
44
ctypes.c_void_p,
45
ctypes.c_void_p,
46
ctypes.c_void_p,
@@ -56,7 +57,42 @@ _batch_memcpy_fn: Any = None
57
58
59
def _resolve_batch_memcpy():
59
- """Resolve cuMemcpyBatchAsync via cuGetProcAddress (one-time)."""
60
+ """Resolve the platform batch-memcpy entry point (one-time).
61
+
62
+ * CUDA: ``cuMemcpyBatchAsync`` via ``cuGetProcAddress`` (uses
63
+ srcAccessOrder=STREAM via one attributes entry).
64
+ * ROCm: ``hipMemcpyBatchAsync`` from libamdhip64 (ROCm 7.1+). ROCm
65
+ 7.2.1 or 7.2.2 rejects any call with ``numAttrs > 0``
66
+ (see ROCm/clr @ rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so
67
+ we call with ``numAttrs=0``.
68
+
69
+ Raises ``RuntimeError`` if the symbol is unavailable (older CUDA
70
+ driver, ROCm < 7.1, unusual install). The connector requires the
71
+ batch API.
72
+ """
73
+ if current_platform.is_rocm():
74
+ try:
75
+ lib = ctypes.CDLL("libamdhip64.so", mode=ctypes.RTLD_GLOBAL)
76
+ fn = lib.hipMemcpyBatchAsync
77
+ except (OSError, AttributeError) as e:
78
+ raise RuntimeError(
79
+ "hipMemcpyBatchAsync is unavailable in this ROCm install; "
80
+ "SimpleCPUOffloadConnector requires ROCm 7.1+."
81
+ ) from e
82
+ fn.restype = ctypes.c_uint
83
+ fn.argtypes = [
84
+ ctypes.c_void_p, # dsts
85
+ ctypes.c_void_p, # srcs
86
+ ctypes.c_void_p, # sizes
87
+ ctypes.c_size_t, # count
88
+ ctypes.c_void_p, # attrs
89
+ ctypes.c_void_p, # attrIdxs
90
+ ctypes.c_size_t, # numAttrs
91
+ ctypes.c_void_p, # failIdx
92
+ ctypes.c_void_p, # stream
93
+ ]
94
+ return fn
95
+
96
from cuda.bindings import driver as drv
97
98
err, ptr, _ = drv.cuGetProcAddress(b"cuMemcpyBatchAsync", 12080, 0)
@@ -70,6 +106,8 @@ class BatchMemcpyParams(NamedTuple):
106
dst_bases: np.ndarray # [num_layers] uint64
107
bpb: np.ndarray # [num_layers] uint64 — bytes per block
108
num_layers: int
109
+ # CUDA only: one attributes entry with srcAccessOrder=ANY. Unused on
110
+ # ROCm (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0.
111
attrs: _CUmemcpyAttributes
112
attrs_idx: ctypes.c_size_t
113
# NOTE: cuMemcpyBatchAsync_v2() removed fail_idx field, but we use
@@ -99,8 +137,10 @@ def build_params(
137
dst_bases.append(d.data_ptr())
138
bpb.append(s_bpb)
139
102
- # Refer to https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f for details. # noqa: E501
103
- attrs = _CUmemcpyAttributes(srcAccessOrder=3) # ANY
140
+ # ``srcAccessOrder=3`` == CU_MEMCPY_SRC_ACCESS_ORDER_ANY /
141
+ # hipMemcpySrcAccessOrderAny. See
142
+ # https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f # noqa: E501
143
+ attrs = _CUmemcpyAttributes(srcAccessOrder=3)
144
145
return BatchMemcpyParams(
146
src_bases=np.array(src_bases, dtype=np.uint64),
@@ -119,7 +159,7 @@ def copy_blocks(
159
dst_block_ids: list[int],
160
params: BatchMemcpyParams,
161
) -> None:
122
- """Copy blocks via cuMemcpyBatchAsync."""
162
+ """Copy blocks via cuMemcpyBatchAsync / hipMemcpyBatchAsync."""
163
n = len(src_block_ids)
164
if n == 0:
165
return
@@ -134,8 +174,13 @@ def copy_blocks(
174
params.dst_bases[:, None] + dst_ids[None, :] * params.bpb[:, None]
175
).ravel()
176
sz_all = np.repeat(params.bpb, n)
137
-
177
total = n * params.num_layers
178
+
179
+ # ROCm 7.2.1/7.2.2 rejects any call with numAttrs>0 (hipMemcpyBatchAsync
180
+ # hipamd/src/hip_memory.cpp:2819-2822); CUDA uses one attrs entry so
181
+ # srcAccessOrder is honored. attrs / attrsIdxs are ignored when
182
+ # numAttrs==0, so we pass the same values from both paths.
183
+ num_attrs = 0 if current_platform.is_rocm() else 1
184
err = _batch_memcpy_fn(
185
dst_all.ctypes.data,
186
src_all.ctypes.data,
@@ -143,11 +188,11 @@ def copy_blocks(
188
total,
189
ctypes.addressof(params.attrs),
190
ctypes.byref(params.attrs_idx),
146
- 1,
191
+ num_attrs,
192
ctypes.byref(params.fail_idx),
193
params.stream_handle,
194
)
195
if err != 0:
196
raise RuntimeError(
152
- f"cuMemcpyBatchAsync failed: err={err} failIdx={params.fail_idx.value}"
197
+ f"batch memcpy failed: err={err} failIdx={params.fail_idx.value}"
198
)