| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Low-level CUDA/HIP memory helpers: pinning and batch DMA transfers.""" |
| 4 | |
| 5 | import ctypes |
| 6 | from typing import Any, NamedTuple |
| 7 | |
| 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 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | # CUmemcpySrcAccessOrder values (CUDA driver API). STREAM(1): source read in |
| 18 | # stream order, safe when the source may still be written. ANY(3): source may |
| 19 | # be read early, only safe for a stable source (e.g. pinned host memory). |
| 20 | CU_MEMCPY_SRC_ACCESS_ORDER_STREAM = 1 |
| 21 | CU_MEMCPY_SRC_ACCESS_ORDER_ANY = 3 |
| 22 | |
| 23 | |
| 24 | def pin_tensor(tensor: torch.Tensor) -> None: |
| 25 | """Pin a CPU tensor via cudaHostRegister. |
| 26 | |
| 27 | This bypasses PyTorch's CUDACachingHostAllocator which rounds |
| 28 | every ``pin_memory=True`` allocation up to the next power of 2 |
| 29 | (e.g. 100 GB becomes 128 GB). |
| 30 | """ |
| 31 | err = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), tensor.nbytes, 0) |
| 32 | if err.value != 0: |
| 33 | raise RuntimeError(f"cudaHostRegister failed: {err}") |
| 34 | |
| 35 | |
| 36 | class _CUmemLocation(ctypes.Structure): |
| 37 | _fields_ = [("type", ctypes.c_uint), ("id", ctypes.c_int)] |
| 38 | |
| 39 | |
| 40 | class _CUmemcpyAttributes(ctypes.Structure): |
| 41 | _fields_ = [ |
| 42 | ("srcAccessOrder", ctypes.c_uint), |
| 43 | ("srcLocHint", _CUmemLocation), |
| 44 | ("dstLocHint", _CUmemLocation), |
| 45 | ("flags", ctypes.c_uint), |
| 46 | ] |
| 47 | |
| 48 | |
| 49 | _BATCH_MEMCPY_FUNC_TYPE = ctypes.CFUNCTYPE( |
| 50 | ctypes.c_uint, # CUresult / hipError_t |
| 51 | ctypes.c_void_p, |
| 52 | ctypes.c_void_p, |
| 53 | ctypes.c_void_p, |
| 54 | ctypes.c_size_t, |
| 55 | ctypes.c_void_p, |
| 56 | ctypes.c_void_p, |
| 57 | ctypes.c_size_t, |
| 58 | ctypes.c_void_p, |
| 59 | ctypes.c_void_p, |
| 60 | ) |
| 61 | |
| 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 | |
| 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 | |
| 72 | |
| 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: |
| 102 | lib = _load_hip_runtime() |
| 103 | fn = lib.hipMemcpyBatchAsync |
| 104 | except (OSError, AttributeError) as e: |
| 105 | raise RuntimeError( |
| 106 | "hipMemcpyBatchAsync is unavailable in this ROCm install; " |
| 107 | "SimpleCPUOffloadConnector requires ROCm 7.1+." |
| 108 | ) from e |
| 109 | fn.restype = ctypes.c_uint |
| 110 | fn.argtypes = [ |
| 111 | ctypes.c_void_p, # dsts |
| 112 | ctypes.c_void_p, # srcs |
| 113 | ctypes.c_void_p, # sizes |
| 114 | ctypes.c_size_t, # count |
| 115 | ctypes.c_void_p, # attrs |
| 116 | ctypes.c_void_p, # attrIdxs |
| 117 | ctypes.c_size_t, # numAttrs |
| 118 | ctypes.c_void_p, # failIdx |
| 119 | ctypes.c_void_p, # stream |
| 120 | ] |
| 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}") |
| 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): |
| 173 | src_bases: np.ndarray # [num_layers] uint64 — data_ptr per layer |
| 174 | dst_bases: np.ndarray # [num_layers] uint64 |
| 175 | bpb: np.ndarray # [num_layers] uint64 — bytes per block |
| 176 | num_layers: int |
| 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 |
| 186 | stream_handle: int # raw cudaStream_t / CUstream |
| 187 | |
| 188 | |
| 189 | def build_params( |
| 190 | src_caches: dict[str, torch.Tensor], |
| 191 | dst_caches: dict[str, torch.Tensor], |
| 192 | stream: torch.cuda.Stream, |
| 193 | src_access_order: int = CU_MEMCPY_SRC_ACCESS_ORDER_ANY, |
| 194 | ) -> BatchMemcpyParams: |
| 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()) |
| 202 | dst_tensors = list(dst_caches.values()) |
| 203 | |
| 204 | src_bases, dst_bases, bpb = [], [], [] |
| 205 | for s, d in zip(src_tensors, dst_tensors): |
| 206 | s_bpb = s.stride(0) * s.element_size() |
| 207 | assert s_bpb == d.stride(0) * d.element_size() |
| 208 | src_bases.append(s.data_ptr()) |
| 209 | dst_bases.append(d.data_ptr()) |
| 210 | bpb.append(s_bpb) |
| 211 | |
| 212 | attrs = _CUmemcpyAttributes(srcAccessOrder=src_access_order) |
| 213 | |
| 214 | return BatchMemcpyParams( |
| 215 | src_bases=np.array(src_bases, dtype=np.uint64), |
| 216 | dst_bases=np.array(dst_bases, dtype=np.uint64), |
| 217 | bpb=np.array(bpb, dtype=np.uint64), |
| 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 | ) |
| 225 | |
| 226 | |
| 227 | def copy_blocks( |
| 228 | src_block_ids: list[int], |
| 229 | dst_block_ids: list[int], |
| 230 | params: BatchMemcpyParams, |
| 231 | ) -> None: |
| 232 | """Copy blocks via cuMemcpyBatchAsync / hipMemcpyBatchAsync.""" |
| 233 | n = len(src_block_ids) |
| 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 | |
| 243 | src_all = ( |
| 244 | params.src_bases[:, None] + src_ids[None, :] * params.bpb[:, None] |
| 245 | ).ravel() |
| 246 | dst_all = ( |
| 247 | params.dst_bases[:, None] + dst_ids[None, :] * params.bpb[:, None] |
| 248 | ).ravel() |
| 249 | sz_all = np.repeat(params.bpb, n) |
| 250 | total = n * params.num_layers |
| 251 | |
| 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 | ) |