@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cd7b7c265a

[ROCm] Cpu offload for ROCm 7.13+ to align the hipMemcpyBatchAsync params and perf in 7.14x (#43018)

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

Hongxia Yang committed Aug 21, 2026 at 12:07 UTC cd7b7c265a30b8c3e5c908d59f4cb0d7dd8f8ff3
4 files changed +338 -59
csrc/libtorch_stable/cache_kernels.cu
+76 -19
@@ -16,6 +16,7 @@
16 #include <algorithm>
17 #include <cassert>
18 #include <cfloat>
19 +#include <cstdlib>
20
21 #ifdef USE_ROCM
22 #include <hip/hip_bf16.h>
@@ -76,6 +77,31 @@ void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
77 }
78 }
79
80 +namespace {
81 +// ROCm hipMemcpyBatchAsync faults for count > 8192 (MI350X/gfx950), so chunk
82 +// at that ceiling.
83 +constexpr int64_t kRocmDefaultMaxBatchDescriptors = 8192;
84 +
85 +constexpr const char* kMaxBatchDescriptorsEnv =
86 + "VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS";
87 +
88 +// Max copy descriptors per batch-memcpy call (0 = unlimited). ROCm caps and
89 +// chunks; CUDA is uncapped. The env var (>0) overrides on any platform.
90 +int64_t resolve_max_batch_descriptors() {
91 + static const int64_t cached = []() -> int64_t {
92 + const char* val = std::getenv(kMaxBatchDescriptorsEnv);
93 + const int64_t override_val = val ? std::atoll(val) : 0;
94 + if (override_val > 0) return override_val;
95 +#if defined(USE_ROCM)
96 + return kRocmDefaultMaxBatchDescriptors;
97 +#else
98 + return 0;
99 +#endif
100 + }();
101 + return cached;
102 +}
103 +} // namespace
104 +
105 void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
106 const torch::stable::Tensor& dst_ptrs,
107 const torch::stable::Tensor& sizes,
@@ -142,32 +168,63 @@ void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
168 : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
169 size_t attrs_idx = 0;
170 size_t fail_idx = 0;
145 - CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data),
146 - reinterpret_cast<CUdeviceptr*>(src_data),
147 - reinterpret_cast<size_t*>(size_data),
148 - static_cast<size_t>(n), &attr, &attrs_idx, 1,
149 - &fail_idx, static_cast<CUstream>(stream));
150 - STD_TORCH_CHECK(result == CUDA_SUCCESS,
151 - "cuMemcpyBatchAsync failed at index ", fail_idx,
152 - " with error ", result);
171 + // Uncapped on CUDA (max_desc == 0 -> single call) unless overridden by
172 + // VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS; chunk to honor the override.
173 + const int64_t max_desc = resolve_max_batch_descriptors();
174 + const int64_t step = max_desc <= 0 ? n : max_desc;
175 + for (int64_t off = 0; off < n; off += step) {
176 + const int64_t cnt = std::min(step, n - off);
177 + CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data + off),
178 + reinterpret_cast<CUdeviceptr*>(src_data + off),
179 + reinterpret_cast<size_t*>(size_data + off),
180 + static_cast<size_t>(cnt), &attr, &attrs_idx, 1,
181 + &fail_idx, static_cast<CUstream>(stream));
182 + STD_TORCH_CHECK(result == CUDA_SUCCESS,
183 + "cuMemcpyBatchAsync failed at index ", fail_idx,
184 + " with error ", result);
185 + }
186 return;
187 }
188 #elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
156 - // ROCm 7.1+ exposes hipMemcpyBatchAsync. The 7.2.1 implementation early-
157 - // returns hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @
158 - // rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so call with
159 - // numAttrs=0.
189 + // ROCm 7.1+ exposes hipMemcpyBatchAsync. ROCm 7.2.1-7.2.3 early-return
190 + // hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @ rocm-7.2.1
191 + // hipamd/src/hip_memory.cpp:2819-2822), so those releases must call with
192 + // numAttrs=0. ROCm 7.13+ accepts numAttrs > 0.
193 + // rocm-7.14+ has better performance.
194 {
195 hipMemcpyAttributes attr = {};
196 size_t attrs_idx = 0;
197 size_t fail_idx = 0;
164 - hipError_t result = hipMemcpyBatchAsync(
165 - reinterpret_cast<void**>(dst_data), reinterpret_cast<void**>(src_data),
166 - reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
167 - &attrs_idx, 0, &fail_idx, static_cast<hipStream_t>(stream));
168 - STD_TORCH_CHECK(result == hipSuccess,
169 - "hipMemcpyBatchAsync failed at index ", fail_idx,
170 - " with error ", result);
198 + size_t num_attrs = 0;
199 + #if HIP_VERSION >= 71300000
200 + static const bool runtime_accepts_attrs = []() {
201 + int runtime_version = 0;
202 + return hipRuntimeGetVersion(&runtime_version) == hipSuccess &&
203 + runtime_version >= 71300000;
204 + }();
205 + if (runtime_accepts_attrs) {
206 + attr.srcAccessOrder = is_src_access_order_any
207 + ? hipMemcpySrcAccessOrderAny
208 + : hipMemcpySrcAccessOrderStream;
209 + num_attrs = 1;
210 + }
211 + #endif
212 + // hipMemcpyBatchAsync faults (GPU memory access fault) above 8192
213 + // descriptors/call on ROCm 7.15, so chunk the batch at the resolved cap.
214 + const int64_t max_desc = resolve_max_batch_descriptors();
215 + const int64_t step = max_desc <= 0 ? n : max_desc;
216 + for (int64_t off = 0; off < n; off += step) {
217 + const int64_t cnt = std::min(step, n - off);
218 + hipError_t result = hipMemcpyBatchAsync(
219 + reinterpret_cast<void**>(dst_data + off),
220 + reinterpret_cast<void**>(src_data + off),
221 + reinterpret_cast<size_t*>(size_data + off), static_cast<size_t>(cnt),
222 + &attr, &attrs_idx, num_attrs, &fail_idx,
223 + static_cast<hipStream_t>(stream));
224 + STD_TORCH_CHECK(result == hipSuccess,
225 + "hipMemcpyBatchAsync failed at index ", fail_idx,
226 + " with error ", result);
227 + }
228 return;
229 }
230 #endif
tests/v1/simple_kv_offload/test_hip_mem_ops.py new
+145
@@ -0,0 +1,145 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the HIP/ROCm paths in vllm.v1.simple_kv_offload.cuda_mem_ops.
4 +
5 +Deliberately not skipped off ROCm: the version gate takes a plain version int
6 +and the library loader is monkeypatched, so nothing here needs a HIP runtime
7 +and the logic stays covered by the default (NVIDIA) CI that runs this
8 +directory.
9 +"""
10 +
11 +import ctypes
12 +
13 +import numpy as np
14 +import pytest
15 +
16 +from vllm.v1.simple_kv_offload import cuda_mem_ops
17 +from vllm.v1.simple_kv_offload.cuda_mem_ops import (
18 + BatchMemcpyParams,
19 + _CUmemcpyAttributes,
20 + _load_hip_runtime,
21 + _num_attrs_for_hip_version,
22 + copy_blocks,
23 +)
24 +
25 +
26 +# hipMemcpyBatchAsync rejects numAttrs > 0 on ROCm 7.2.1-7.2.3 and accepts it
27 +# on 7.13+. HIP encodes the version as major*10_000_000 + minor*100_000 + patch.
28 +@pytest.mark.parametrize(
29 + ("version", "expected"),
30 + [
31 + (0, 0), # unknown (hipRuntimeGetVersion failed) -> conservative 0
32 + (70_100_000, 0), # ROCm 7.1
33 + (70_200_000, 0), # ROCm 7.2.0
34 + (70_253_211, 0), # ROCm 7.2.2 (patch 53211)
35 + (70_300_000, 0), # ROCm 7.3
36 + (71_200_000, 0), # ROCm 7.12 -- just below the 7.13 cutoff
37 + (71_300_000, 1), # ROCm 7.13 -- first release accepting numAttrs > 0
38 + (71_326_193, 1), # ROCm 7.13 nightly (patch 26193)
39 + (71_400_000, 1), # ROCm 7.14
40 + (71_500_000, 1), # ROCm 7.15
41 + (80_000_000, 1), # ROCm 8.0
42 + ],
43 +)
44 +def test_num_attrs_for_hip_version(version: int, expected: int):
45 + assert _num_attrs_for_hip_version(version) == expected
46 +
47 +
48 +@pytest.mark.parametrize(
49 + ("available", "expected"),
50 + [
51 + # Devel install: the unversioned symlink is present and preferred.
52 + ({"libamdhip64.so", "libamdhip64.so.7"}, "libamdhip64.so"),
53 + # Runtime-only / wheel-packaged ROCm: versioned soname only.
54 + ({"libamdhip64.so.7"}, "libamdhip64.so.7"),
55 + ({"libamdhip64.so.6"}, "libamdhip64.so.6"),
56 + ],
57 +)
58 +def test_load_hip_runtime_falls_back_to_versioned_soname(
59 + monkeypatch, available: set[str], expected: str
60 +):
61 + loaded = []
62 +
63 + def fake_cdll(name, mode=0):
64 + if name not in available:
65 + raise OSError(f"{name}: cannot open shared object file")
66 + loaded.append(name)
67 + return name
68 +
69 + monkeypatch.setattr(cuda_mem_ops.ctypes, "CDLL", fake_cdll)
70 + assert _load_hip_runtime() == expected
71 + assert loaded == [expected]
72 +
73 +
74 +def test_load_hip_runtime_reports_every_candidate_when_all_fail(monkeypatch):
75 + def fake_cdll(name, mode=0):
76 + raise OSError("cannot open shared object file")
77 +
78 + monkeypatch.setattr(cuda_mem_ops.ctypes, "CDLL", fake_cdll)
79 + with pytest.raises(OSError, match="could not load the HIP runtime") as exc:
80 + _load_hip_runtime()
81 + for name in ("libamdhip64.so", "libamdhip64.so.7", "libamdhip64.so.6"):
82 + assert name in str(exc.value)
83 +
84 +
85 +def _fake_params(num_layers: int, bytes_per_block: int = 512) -> BatchMemcpyParams:
86 + """Minimal params for driving copy_blocks() without a real device.
87 +
88 + Bases are 0 and the batch-memcpy function is monkeypatched, so no memory is
89 + ever dereferenced; only the descriptor bookkeeping (count/chunking) matters.
90 + """
91 + return BatchMemcpyParams(
92 + src_bases=np.zeros(num_layers, dtype=np.uint64),
93 + dst_bases=np.zeros(num_layers, dtype=np.uint64),
94 + bpb=np.full(num_layers, bytes_per_block, dtype=np.uint64),
95 + num_layers=num_layers,
96 + attrs=_CUmemcpyAttributes(),
97 + attrs_idx=ctypes.c_size_t(0),
98 + num_attrs=0,
99 + fail_idx=ctypes.c_size_t(0),
100 + stream_handle=0,
101 + )
102 +
103 +
104 +def _record_counts(monkeypatch) -> list[int]:
105 + """Patch the batch-memcpy call to record the descriptor count per call."""
106 + counts: list[int] = []
107 +
108 + def fake_fn(dst, src, sizes, count, *rest):
109 + counts.append(count)
110 + return 0
111 +
112 + monkeypatch.setattr(cuda_mem_ops, "_batch_memcpy", (fake_fn, 0))
113 + return counts
114 +
115 +
116 +# ROCm chunks the descriptor batch at the cap (hipMemcpyBatchAsync faults above
117 +# 8192/call); CUDA leaves the cap at 0 and issues a single call.
118 +@pytest.mark.parametrize(
119 + ("max_desc", "num_layers", "num_blocks", "expected_counts"),
120 + [
121 + (8192, 3, 5000, [8192, 6808]), # 15000 desc -> 8192 + 6808
122 + (8192, 1, 8192, [8192]), # exactly at the cap -> single call
123 + (8192, 1, 8193, [8192, 1]), # one over -> spills into a second call
124 + (4096, 61, 135, [4096, 4096, 43]), # bs=1 Kimi shape (8235 desc)
125 + (0, 61, 135, [8235]), # CUDA uncapped -> single call
126 + ],
127 +)
128 +def test_copy_blocks_chunks_at_descriptor_cap(
129 + monkeypatch, max_desc, num_layers, num_blocks, expected_counts
130 +):
131 + counts = _record_counts(monkeypatch)
132 + monkeypatch.setattr(cuda_mem_ops, "_max_batch_descriptors", max_desc)
133 +
134 + ids = list(range(num_blocks))
135 + copy_blocks(ids, ids, _fake_params(num_layers))
136 +
137 + assert counts == expected_counts
138 + assert sum(counts) == num_layers * num_blocks
139 +
140 +
141 +def test_copy_blocks_noop_on_empty(monkeypatch):
142 + counts = _record_counts(monkeypatch)
143 + monkeypatch.setattr(cuda_mem_ops, "_max_batch_descriptors", 8192)
144 + copy_blocks([], [], _fake_params(num_layers=4))
145 + assert counts == []
vllm/envs.py
+7
@@ -299,6 +299,7 @@ if TYPE_CHECKING:
299 VLLM_DEBUG_MFU_METRICS: bool = False
300 VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False
301 VLLM_WEIGHT_OFFLOADING_DISABLE_UVA: bool = False
302 + VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS: int = 0
303 VLLM_WSL2_ENABLE_PIN_MEMORY: bool = False
304 VLLM_DISABLE_LOG_LOGO: bool = False
305 VLLM_LORA_DISABLE_PDL: bool = False
@@ -2059,6 +2060,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
2060 "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": lambda: bool(
2061 int(os.getenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "0"))
2062 ),
2063 + # Max descriptors per CPU-KV-offload batch-memcpy call. 0 = platform default
2064 + # (ROCm chunks at 8192, since hipMemcpyBatchAsync faults above that on
2065 + # rocm 7.14/7.15 when batch copy is optimized; CUDA uncapped). Set >0 to override.
2066 + "VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS": lambda: int(
2067 + os.getenv("VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS", "0")
2068 + ),
2069 # On WSL2 with a compatible kernel (>= 4.19.121), pinned memory is
2070 # supported but disabled by default due to a small performance regression.
2071 # Set to 1 when pinned memory or UVA is required (e.g. CPU offloading
vllm/v1/simple_kv_offload/cuda_mem_ops.py
+110 -40
@@ -8,6 +8,7 @@ from typing import Any, NamedTuple
8 import numpy as np
9 import torch
10
11 +import vllm.envs as envs
12 from vllm.logger import init_logger
13 from vllm.platforms import current_platform
14
@@ -58,27 +59,47 @@ _BATCH_MEMCPY_FUNC_TYPE = ctypes.CFUNCTYPE(
59 ctypes.c_void_p,
60 )
61
61 -# Resolved lazily on first use.
62 -_batch_memcpy_fn: Any = None
62 +# Resolved lazily on first use: (entry point, numAttrs to pass).
63 +_batch_memcpy: tuple[Any, int] | None = None
64
65 +# Max copy descriptors per batch call, resolved lazily on first use.
66 +_max_batch_descriptors: int | None = None
67
65 -def _resolve_batch_memcpy():
66 - """Resolve the platform batch-memcpy entry point (one-time).
68 +# ROCm hipMemcpyBatchAsync faults for count > 8192 (MI350X/gfx950), so we chunk
69 +# at that ceiling. Override via VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS.
70 +_ROCM_DEFAULT_MAX_BATCH_DESCRIPTORS = 8192
71
68 - * CUDA: ``cuMemcpyBatchAsync`` via ``cuGetProcAddress`` (uses
69 - srcAccessOrder=STREAM via one attributes entry).
70 - * ROCm: ``hipMemcpyBatchAsync`` from libamdhip64 (ROCm 7.1+). ROCm
71 - 7.2.1 or 7.2.2 rejects any call with ``numAttrs > 0``
72 - (see ROCm/clr @ rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so
73 - we call with ``numAttrs=0``.
72
75 - Raises ``RuntimeError`` if the symbol is unavailable (older CUDA
76 - driver, ROCm < 7.1, unusual install). The connector requires the
77 - batch API.
73 +def _resolve_max_batch_descriptors() -> int:
74 + """Max copy descriptors to pass to one batch-memcpy call (0 = unlimited).
75 +
76 + ROCm's ``hipMemcpyBatchAsync`` faults above 8192 descriptors per call, so
77 + on ROCm we cap and chunk larger transfers. CUDA's ``cuMemcpyBatchAsync``
78 + handles arbitrary counts and is left uncapped. Set
79 + ``VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS`` (>0) to override on any platform.
80 + """
81 + global _max_batch_descriptors
82 + if _max_batch_descriptors is None:
83 + override = envs.VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS
84 + if override > 0:
85 + _max_batch_descriptors = override
86 + else:
87 + _max_batch_descriptors = (
88 + _ROCM_DEFAULT_MAX_BATCH_DESCRIPTORS if current_platform.is_rocm() else 0
89 + )
90 + return _max_batch_descriptors
91 +
92 +
93 +def _resolve_batch_memcpy() -> tuple[Any, int]:
94 + """Resolve the batch-memcpy entry point and its ``numAttrs`` (one-time).
95 +
96 + CUDA uses ``cuMemcpyBatchAsync``; ROCm uses ``hipMemcpyBatchAsync``.
97 + Raises ``RuntimeError`` if the symbol is unavailable (old CUDA driver,
98 + ROCm < 7.1, unusual install).
99 """
100 if current_platform.is_rocm():
101 try:
81 - lib = ctypes.CDLL("libamdhip64.so", mode=ctypes.RTLD_GLOBAL)
102 + lib = _load_hip_runtime()
103 fn = lib.hipMemcpyBatchAsync
104 except (OSError, AttributeError) as e:
105 raise RuntimeError(
@@ -97,14 +118,55 @@ def _resolve_batch_memcpy():
118 ctypes.c_void_p, # failIdx
119 ctypes.c_void_p, # stream
120 ]
100 - return fn
121 + return fn, _rocm_num_attrs(lib)
122
123 from cuda.bindings import driver as drv
124
125 err, ptr, _ = drv.cuGetProcAddress(b"cuMemcpyBatchAsync", 12080, 0)
126 if err != drv.CUresult.CUDA_SUCCESS:
127 raise RuntimeError(f"cuGetProcAddress(cuMemcpyBatchAsync) failed: {err}")
107 - return _BATCH_MEMCPY_FUNC_TYPE(ptr)
128 + return _BATCH_MEMCPY_FUNC_TYPE(ptr), 1
129 +
130 +
131 +def _load_hip_runtime() -> ctypes.CDLL:
132 + """Load ``libamdhip64``, tolerating installs without the devel symlink.
133 +
134 + The unversioned ``libamdhip64.so`` only ships with the ROCm devel package;
135 + runtime-only and wheel-packaged ROCm installs provide just the versioned
136 + soname. ``dlopen`` returns the already-mapped library when asked for a
137 + soname the process has loaded — torch loads HIP at import — so the
138 + versioned names resolve even when they are not on the loader search path.
139 + """
140 + errors = []
141 + for name in ("libamdhip64.so", "libamdhip64.so.7", "libamdhip64.so.6"):
142 + try:
143 + return ctypes.CDLL(name, mode=ctypes.RTLD_GLOBAL)
144 + except OSError as e:
145 + errors.append(f"{name}: {e}")
146 + raise OSError("could not load the HIP runtime: " + "; ".join(errors))
147 +
148 +
149 +def _rocm_num_attrs(lib: ctypes.CDLL) -> int:
150 + """``numAttrs`` for ``hipMemcpyBatchAsync`` on the running HIP runtime."""
151 + ver = ctypes.c_int(0)
152 + try:
153 + if lib.hipRuntimeGetVersion(ctypes.byref(ver)) != 0:
154 + ver.value = 0
155 + except (OSError, AttributeError):
156 + ver.value = 0
157 + return _num_attrs_for_hip_version(ver.value)
158 +
159 +
160 +def _num_attrs_for_hip_version(version: int) -> int:
161 + """``numAttrs`` for ``hipMemcpyBatchAsync`` given a HIP runtime version int.
162 +
163 + ROCm 7.2.1-7.2.3 reject ``numAttrs > 0`` (ROCm/clr @ rocm-7.2.1
164 + hipamd/src/hip_memory.cpp:2819-2822); 7.13+ accept it. ``version`` 0
165 + (unknown) yields the conservative 0.
166 + """
167 + # HIP encodes version as major*10_000_000 + minor*100_000 + patch.
168 + major, minor = version // 10_000_000, (version // 100_000) % 100
169 + return 1 if (major, minor) >= (7, 13) else 0
170
171
172 class BatchMemcpyParams(NamedTuple):
@@ -112,10 +174,12 @@ class BatchMemcpyParams(NamedTuple):
174 dst_bases: np.ndarray # [num_layers] uint64
175 bpb: np.ndarray # [num_layers] uint64 — bytes per block
176 num_layers: int
115 - # CUDA only: one attributes entry carrying srcAccessOrder. Unused on ROCm
116 - # (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0.
177 + # One attributes entry carrying srcAccessOrder. Ignored when num_attrs is
178 + # 0, which is what ROCm runtimes older than 7.13 require (see
179 + # _num_attrs_for_hip_version).
180 attrs: _CUmemcpyAttributes
181 attrs_idx: ctypes.c_size_t
182 + num_attrs: int
183 # NOTE: cuMemcpyBatchAsync_v2() removed fail_idx field, but we use
184 # cuMemcpyBatchAsync() with fail_idx for backward compatibility
185 fail_idx: ctypes.c_size_t
@@ -128,9 +192,10 @@ def build_params(
192 stream: torch.cuda.Stream,
193 src_access_order: int = CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
194 ) -> BatchMemcpyParams:
131 - global _batch_memcpy_fn
132 - if _batch_memcpy_fn is None:
133 - _batch_memcpy_fn = _resolve_batch_memcpy()
195 + global _batch_memcpy
196 + if _batch_memcpy is None:
197 + _batch_memcpy = _resolve_batch_memcpy()
198 + _, num_attrs = _batch_memcpy
199
200 assert list(src_caches.keys()) == list(dst_caches.keys())
201 src_tensors = list(src_caches.values())
@@ -153,6 +218,7 @@ def build_params(
218 num_layers=len(src_tensors),
219 attrs=attrs,
220 attrs_idx=ctypes.c_size_t(0),
221 + num_attrs=num_attrs,
222 fail_idx=ctypes.c_size_t(0),
223 stream_handle=stream.cuda_stream,
224 )
@@ -168,6 +234,9 @@ def copy_blocks(
234 if n == 0:
235 return
236
237 + assert _batch_memcpy is not None, "build_params() must run before copy_blocks()"
238 + fn, _ = _batch_memcpy
239 +
240 src_ids = np.array(src_block_ids, dtype=np.uint64)
241 dst_ids = np.array(dst_block_ids, dtype=np.uint64)
242
@@ -180,23 +249,24 @@ def copy_blocks(
249 sz_all = np.repeat(params.bpb, n)
250 total = n * params.num_layers
251
183 - # ROCm 7.2.1/7.2.2 rejects any call with numAttrs>0 (hipMemcpyBatchAsync
184 - # hipamd/src/hip_memory.cpp:2819-2822); CUDA uses one attrs entry so
185 - # srcAccessOrder is honored. attrs / attrsIdxs are ignored when
186 - # numAttrs==0, so we pass the same values from both paths.
187 - num_attrs = 0 if current_platform.is_rocm() else 1
188 - err = _batch_memcpy_fn(
189 - dst_all.ctypes.data,
190 - src_all.ctypes.data,
191 - sz_all.ctypes.data,
192 - total,
193 - ctypes.addressof(params.attrs),
194 - ctypes.byref(params.attrs_idx),
195 - num_attrs,
196 - ctypes.byref(params.fail_idx),
197 - params.stream_handle,
198 - )
199 - if err != 0:
200 - raise RuntimeError(
201 - f"batch memcpy failed: err={err} failIdx={params.fail_idx.value}"
252 + # Chunk on ROCm: hipMemcpyBatchAsync faults above 8192 descriptors/call.
253 + # CUDA is uncapped (max_desc == 0) and issues a single call.
254 + max_desc = _resolve_max_batch_descriptors()
255 + step = total if max_desc <= 0 else max_desc
256 + for off in range(0, total, step):
257 + cnt = min(step, total - off)
258 + err = fn(
259 + dst_all[off : off + cnt].ctypes.data,
260 + src_all[off : off + cnt].ctypes.data,
261 + sz_all[off : off + cnt].ctypes.data,
262 + cnt,
263 + ctypes.addressof(params.attrs),
264 + ctypes.byref(params.attrs_idx),
265 + params.num_attrs,
266 + ctypes.byref(params.fail_idx),
267 + params.stream_handle,
268 )
269 + if err != 0:
270 + raise RuntimeError(
271 + f"batch memcpy failed: err={err} failIdx={params.fail_idx.value}"
272 + )