[ROCm] Enable SimpleCPUOffloadConnector on ROCm

csrc/cache_kernels.cu

@@ -18,6 +18,7 @@
 #endif
 
 #include <algorithm>
+#include <atomic>
 #include <cassert>
 #include <cfloat>
 
@@ -97,13 +98,13 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
 
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-  // Use cuMemcpyBatchAsync (CUDA 12.8+) to submit all copies in a single
-  // driver call, amortizing per-copy submission overhead.
-  // int64_t and CUdeviceptr/size_t are both 8 bytes on 64-bit platforms,
-  // so we reinterpret_cast the tensor data directly to avoid copies.
-  static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
+  // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a
+  // single driver call, amortizing per-copy submission overhead. int64_t
+  // and CUdeviceptr/void*/size_t are all 8 bytes on 64-bit platforms, so we
+  // reinterpret_cast the tensor data directly to avoid copies.
   static_assert(sizeof(size_t) == sizeof(int64_t));
 #if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
+  static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
   // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that
   // binaries compiled with CUDA 12.8+ still work on older drivers, and
   // we avoid the CUDA 13.0 header remapping (#define to _v2 signature).
@@ -134,12 +135,37 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
                                &fail_idx, static_cast<CUstream>(stream));
     TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
                 fail_idx, " with error ", result);
-  } else
+    return;
+  }
+#elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
+  // HIP 7.1+ exposes hipMemcpyBatchAsync with the same signature. ROCm 7.2
+  // ships the symbol as a stub that returns hipErrorNotSupported; detect
+  // that at call time, clear the sticky last error (otherwise subsequent
+  // torch ops surface the stale NotSupported), and fall through to the
+  // per-op loop for the remainder of this process. This auto-activates
+  // when a future ROCm release wires the real implementation.
+  static std::atomic<bool> batch_supported{true};
+  if (batch_supported.load(std::memory_order_relaxed)) {
+    hipMemcpyAttributes attr = {};
+    attr.srcAccessOrder = hipMemcpySrcAccessOrderStream;
+    size_t attrs_idx = 0;
+    size_t fail_idx = 0;
+    hipError_t result = hipMemcpyBatchAsync(
+        reinterpret_cast<void**>(dst_data), reinterpret_cast<void**>(src_data),
+        reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
+        &attrs_idx, 1, &fail_idx, static_cast<hipStream_t>(stream));
+    if (result == hipSuccess) return;
+    TORCH_CHECK(result == hipErrorNotSupported,
+                "hipMemcpyBatchAsync failed at index ", fail_idx,
+                " with error ", result);
+    batch_supported.store(false, std::memory_order_relaxed);
+    (void)hipGetLastError();  // clear the sticky stub error
+  }
 #endif
   {
-    // Fallback for CUDA < 12.8, older drivers, and ROCm:
-    // individual async copies.
-    // cudaMemcpyDefault lets the driver infer direction from pointer types.
+    // Fallback for CUDA < 12.8, older drivers, and ROCm without a working
+    // batch impl: individual async copies. cudaMemcpyDefault lets the
+    // driver infer direction from pointer types.
     for (int64_t i = 0; i < n; i++) {
       cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
                       reinterpret_cast<void*>(src_data[i]),

tests/v1/simple_kv_offload/test_integration.py

@@ -10,8 +10,8 @@
 from vllm.config import KVTransferConfig
 from vllm.platforms import current_platform
 
-if not current_platform.is_cuda():
-    pytest.skip("Requires CUDA", allow_module_level=True)
+if not current_platform.is_cuda_alike():
+    pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
 
 # Small models for default CI / local runs (accuracy only).
 SMALL_MODELS = [
@@ -37,6 +37,7 @@ def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM:
     )
     return LLM(
         model=model,
+        gpu_memory_utilization=0.7,
         kv_cache_memory_bytes=40 << 30,  # 40 GiB
         disable_hybrid_kv_cache_manager=False,
         enable_prefix_caching=True,

tests/v1/simple_kv_offload/test_scheduler.py

@@ -186,7 +186,13 @@ def make_request(
     if request_id is None:
         request_id = f"req-{_req_counter}"
 
-    num_tokens = num_blocks * BLOCK_SIZE
+    # Add one extra token beyond the last full block so that
+    # ``max_cache_hit_length = num_tokens - 1`` (see
+    # KVCacheManager.get_computed_blocks) does not truncate the final
+    # full block: ``find_longest_cache_hit`` uses
+    # ``max_length // block_size`` and would otherwise drop one block
+    # when the prompt is an exact multiple of block_size.
+    num_tokens = num_blocks * BLOCK_SIZE + 1
     start = _req_counter * 10000
     prompt_token_ids = list(range(start, start + num_tokens))
     sampling_params = SamplingParams(max_tokens=1)

vllm/v1/simple_kv_offload/cuda_mem_ops.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Low-level CUDA memory helpers: pinning and batch DMA transfers."""
+"""Low-level CUDA/HIP memory helpers: pinning and batch DMA transfers."""
 
 import ctypes
 from typing import Any, NamedTuple
@@ -9,12 +9,22 @@
 import torch
 
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
+# hipError_t / CUresult value returned when a symbol is exported but the
+# underlying implementation is a stub (seen on ROCm 7.2 for
+# ``hipMemcpyBatchAsync``).
+_ERR_NOT_SUPPORTED_HIP = 801
+_ERR_NOT_SUPPORTED_CUDA = 801  # CUDA_ERROR_NOT_SUPPORTED
+
+# hipMemcpyKind / cudaMemcpyKind
+_MEMCPY_DEFAULT = 4
+
 
 def pin_tensor(tensor: torch.Tensor) -> None:
-    """Pin a CPU tensor via cudaHostRegister.
+    """Pin a CPU tensor via cudaHostRegister / hipHostRegister.
 
     This bypasses PyTorch's CUDACachingHostAllocator which rounds
     every ``pin_memory=True`` allocation up to the next power of 2
@@ -25,6 +35,8 @@ def pin_tensor(tensor: torch.Tensor) -> None:
         raise RuntimeError(f"cudaHostRegister failed: {err}")
 
 
+# NOTE: ``CUmemcpyAttributes`` and ``hipMemcpyAttributes`` share the same
+# layout in ROCm 7.x, so a single ctypes struct definition works for both.
 class _CUmemLocation(ctypes.Structure):
     _fields_ = [("type", ctypes.c_uint), ("id", ctypes.c_int)]
 
@@ -39,7 +51,7 @@ class _CUmemcpyAttributes(ctypes.Structure):
 
 
 _BATCH_MEMCPY_FUNC_TYPE = ctypes.CFUNCTYPE(
-    ctypes.c_uint,  # CUresult
+    ctypes.c_uint,  # CUresult / hipError_t
     ctypes.c_void_p,
     ctypes.c_void_p,
     ctypes.c_void_p,
@@ -53,10 +65,41 @@ class _CUmemcpyAttributes(ctypes.Structure):
 
 # Resolved lazily on first use.
 _batch_memcpy_fn: Any = None
+# ``hipMemcpyAsync`` / ``cudaMemcpyAsync`` fallback; resolved lazily only if
+# the batch API returns NotSupported.
+_memcpy_async_fn: Any = None
+# Flips to False after we observe NotSupported from the batch API so
+# subsequent calls skip it.
+_batch_memcpy_supported: bool = True
 
 
 def _resolve_batch_memcpy():
-    """Resolve cuMemcpyBatchAsync via cuGetProcAddress (one-time)."""
+    """Resolve the platform batch-memcpy entry point (one-time).
+
+    * CUDA: ``cuMemcpyBatchAsync`` via ``cuGetProcAddress``.
+    * ROCm: ``hipMemcpyBatchAsync`` from libamdhip64 (ROCm 7.1+).
+
+    NOTE: ROCm 7.2 ships the symbol but the implementation returns
+    ``hipErrorNotSupported``; in that case ``copy_blocks`` falls back to
+    per-block ``hipMemcpyAsync`` via ``_resolve_memcpy_async``.
+    """
+    if current_platform.is_rocm():
+        lib = ctypes.CDLL("libamdhip64.so", mode=ctypes.RTLD_GLOBAL)
+        fn = lib.hipMemcpyBatchAsync
+        fn.restype = ctypes.c_uint
+        fn.argtypes = [
+            ctypes.c_void_p,  # dsts
+            ctypes.c_void_p,  # srcs
+            ctypes.c_void_p,  # sizes
+            ctypes.c_size_t,  # count
+            ctypes.c_void_p,  # attrs
+            ctypes.c_void_p,  # attrIdxs
+            ctypes.c_size_t,  # numAttrs
+            ctypes.c_void_p,  # failIdx
+            ctypes.c_void_p,  # stream
+        ]
+        return fn
+
     from cuda.bindings import driver as drv
 
     err, ptr, _ = drv.cuGetProcAddress(b"cuMemcpyBatchAsync", 12080, 0)
@@ -65,6 +108,38 @@ def _resolve_batch_memcpy():
     return _BATCH_MEMCPY_FUNC_TYPE(ptr)
 
 
+def _resolve_memcpy_async():
+    """Resolve per-op ``hipMemcpyAsync`` / ``cudaMemcpyAsync`` (ROCm fallback)."""
+    lib_name = "libamdhip64.so" if current_platform.is_rocm() else "libcudart.so"
+    sym = "hipMemcpyAsync" if current_platform.is_rocm() else "cudaMemcpyAsync"
+    lib = ctypes.CDLL(lib_name, mode=ctypes.RTLD_GLOBAL)
+    fn = getattr(lib, sym)
+    fn.restype = ctypes.c_uint
+    fn.argtypes = [
+        ctypes.c_void_p,  # dst
+        ctypes.c_void_p,  # src
+        ctypes.c_size_t,  # sizeBytes
+        ctypes.c_int,  # kind (hipMemcpyKind / cudaMemcpyKind)
+        ctypes.c_void_p,  # stream
+    ]
+    return fn
+
+
+def _clear_last_error() -> None:
+    """Clear the sticky last error on the current HIP/CUDA context.
+
+    Needed because ``hipMemcpyBatchAsync`` on ROCm 7.2 returns a stub
+    ``hipErrorNotSupported`` that remains sticky — subsequent torch ops
+    on the device would otherwise surface the stale error.
+    """
+    lib_name = "libamdhip64.so" if current_platform.is_rocm() else "libcudart.so"
+    sym = "hipGetLastError" if current_platform.is_rocm() else "cudaGetLastError"
+    lib = ctypes.CDLL(lib_name, mode=ctypes.RTLD_GLOBAL)
+    fn = getattr(lib, sym)
+    fn.restype = ctypes.c_uint
+    fn()
+
+
 class BatchMemcpyParams(NamedTuple):
     src_bases: np.ndarray  # [num_layers] uint64 — data_ptr per layer
     dst_bases: np.ndarray  # [num_layers] uint64
@@ -75,7 +150,7 @@ class BatchMemcpyParams(NamedTuple):
     # NOTE: cuMemcpyBatchAsync_v2() removed fail_idx field, but we use
     # cuMemcpyBatchAsync() with fail_idx for backward compatibility
     fail_idx: ctypes.c_size_t
-    stream_handle: int  # raw cudaStream_t / CUstream
+    stream_handle: int  # raw cudaStream_t / CUstream / hipStream_t
 
 
 def build_params(
@@ -99,8 +174,10 @@ def build_params(
         dst_bases.append(d.data_ptr())
         bpb.append(s_bpb)
 
-    # Refer to https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f for details.  # noqa: E501
-    attrs = _CUmemcpyAttributes(srcAccessOrder=3)  # ANY
+    # ``srcAccessOrder=3`` == CU_MEMCPY_SRC_ACCESS_ORDER_ANY /
+    # hipMemcpySrcAccessOrderAny. See
+    # https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f  # noqa: E501
+    attrs = _CUmemcpyAttributes(srcAccessOrder=3)
 
     return BatchMemcpyParams(
         src_bases=np.array(src_bases, dtype=np.uint64),
@@ -114,12 +191,62 @@ def build_params(
     )
 
 
+def _copy_blocks_batch(
+    n: int,
+    params: BatchMemcpyParams,
+    src_all: np.ndarray,
+    dst_all: np.ndarray,
+    sz_all: np.ndarray,
+) -> int:
+    """Call the batch memcpy API. Returns the driver error code."""
+    total = n * params.num_layers
+    return _batch_memcpy_fn(
+        dst_all.ctypes.data,
+        src_all.ctypes.data,
+        sz_all.ctypes.data,
+        total,
+        ctypes.addressof(params.attrs),
+        ctypes.byref(params.attrs_idx),
+        1,
+        ctypes.byref(params.fail_idx),
+        params.stream_handle,
+    )
+
+
+def _copy_blocks_per_op(
+    params: BatchMemcpyParams,
+    src_all: np.ndarray,
+    dst_all: np.ndarray,
+    sz_all: np.ndarray,
+) -> None:
+    """Fallback: issue one ``hipMemcpyAsync`` per (block, layer) on the stream.
+
+    Used when the batch API returns NotSupported (e.g. ROCm 7.2 stub).
+    """
+    global _memcpy_async_fn
+    if _memcpy_async_fn is None:
+        _memcpy_async_fn = _resolve_memcpy_async()
+    stream = params.stream_handle
+    for dst_ptr, src_ptr, sz in zip(dst_all, src_all, sz_all):
+        err = _memcpy_async_fn(
+            int(dst_ptr), int(src_ptr), int(sz), _MEMCPY_DEFAULT, stream
+        )
+        if err != 0:
+            raise RuntimeError(f"per-op memcpy async failed: err={err}")
+
+
 def copy_blocks(
     src_block_ids: list[int],
     dst_block_ids: list[int],
     params: BatchMemcpyParams,
 ) -> None:
-    """Copy blocks via cuMemcpyBatchAsync."""
+    """Copy blocks via cuMemcpyBatchAsync / hipMemcpyBatchAsync.
+
+    Falls back to per-op ``hipMemcpyAsync`` if the batch API returns
+    NotSupported on the current platform (ROCm 7.2 ships the symbol but
+    the implementation is a stub).
+    """
+    global _batch_memcpy_supported
     n = len(src_block_ids)
     if n == 0:
         return
@@ -135,19 +262,22 @@ def copy_blocks(
     ).ravel()
     sz_all = np.repeat(params.bpb, n)
 
-    total = n * params.num_layers
-    err = _batch_memcpy_fn(
-        dst_all.ctypes.data,
-        src_all.ctypes.data,
-        sz_all.ctypes.data,
-        total,
-        ctypes.addressof(params.attrs),
-        ctypes.byref(params.attrs_idx),
-        1,
-        ctypes.byref(params.fail_idx),
-        params.stream_handle,
-    )
-    if err != 0:
-        raise RuntimeError(
-            f"cuMemcpyBatchAsync failed: err={err} failIdx={params.fail_idx.value}"
-        )
+    if _batch_memcpy_supported:
+        err = _copy_blocks_batch(n, params, src_all, dst_all, sz_all)
+        if err == 0:
+            return
+        if err == _ERR_NOT_SUPPORTED_HIP:
+            logger.warning(
+                "Batch memcpy API returned NotSupported; falling back to "
+                "per-op async memcpy for the remainder of this process."
+            )
+            _batch_memcpy_supported = False
+            # ROCm 7.2's stub leaves the error sticky on the context; clear
+            # it so subsequent device work doesn't surface the stale error.
+            _clear_last_error()
+        else:
+            raise RuntimeError(
+                f"batch memcpy failed: err={err} failIdx={params.fail_idx.value}"
+            )
+
+    _copy_blocks_per_op(params, src_all, dst_all, sz_all)