Add compact fixed-page CPU KV offload for packed heterogeneous caches

csrc/cache.h

@@ -13,7 +13,8 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
 void swap_blocks_batch(const torch::Tensor& src_ptrs,
                        const torch::Tensor& dst_ptrs,
                        const torch::Tensor& sizes,
-                       bool is_src_access_order_any);
+                       bool is_src_access_order_any,
+                       bool use_batch_api = true);
 
 void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
                        torch::Tensor& key_cache, torch::Tensor& value_cache,

csrc/libtorch_stable/cache_kernels.cu

@@ -79,7 +79,8 @@ void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
 void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
                        const torch::stable::Tensor& dst_ptrs,
                        const torch::stable::Tensor& sizes,
-                       bool is_src_access_order_any) {
+                       bool is_src_access_order_any,
+                       bool use_batch_api) {
   STD_TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU");
   STD_TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU");
   STD_TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU");
@@ -102,85 +103,89 @@ void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
 
   const cudaStream_t stream = get_current_cuda_stream();
 
-  // 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).
-  // The function pointer is cached after the first call.
-  using BatchFn =
-      CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t,
-                   CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream);
-  static BatchFn batch_fn = []() -> BatchFn {
-    CUdriverProcAddressQueryResult sym_status;
-    void* fn_ptr = nullptr;
-    CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080,
-                                    CU_GET_PROC_ADDRESS_DEFAULT, &sym_status);
-    if (res != CUDA_SUCCESS || fn_ptr == nullptr) {
-      return nullptr;
+  // When use_batch_api is false, skip cuMemcpyBatchAsync/hipMemcpyBatchAsync
+  // entirely and use the per-descriptor cudaMemcpyAsync fallback. This is
+  // the durable fix for compact-transfer cuMemcpyBatchAsync segfaults with
+  // very large descriptor counts.
+  if (use_batch_api) {
+    // 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).
+    // The function pointer is cached after the first call.
+    using BatchFn =
+        CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t,
+                     CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream);
+    static BatchFn batch_fn = []() -> BatchFn {
+      CUdriverProcAddressQueryResult sym_status;
+      void* fn_ptr = nullptr;
+      CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080,
+                                      CU_GET_PROC_ADDRESS_DEFAULT, &sym_status);
+      if (res != CUDA_SUCCESS || fn_ptr == nullptr) {
+        return nullptr;
+      }
+      return reinterpret_cast<BatchFn>(fn_ptr);
+    }();
+
+    // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 /
+    // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy
+    // fallback below, which is correct on any stream. Real and per-thread-default
+    // streams take the batch fast path.
+    const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy;
+    if (batch_fn != nullptr && usable_stream) {
+      CUmemcpyAttributes attr = {};
+      // ANY lets the DMA engine prefetch source bytes out of stream order,
+      // which is only safe when no GPU stream is concurrently writing the
+      // source.
+      attr.srcAccessOrder = is_src_access_order_any
+                                ? CU_MEMCPY_SRC_ACCESS_ORDER_ANY
+                                : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
+      size_t attrs_idx = 0;
+      size_t fail_idx = 0;
+      CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data),
+                                  reinterpret_cast<CUdeviceptr*>(src_data),
+                                  reinterpret_cast<size_t*>(size_data),
+                                  static_cast<size_t>(n), &attr, &attrs_idx, 1,
+                                  &fail_idx, static_cast<CUstream>(stream));
+      STD_TORCH_CHECK(result == CUDA_SUCCESS,
+                      "cuMemcpyBatchAsync failed at index ", fail_idx,
+                      " with error ", result);
+      return;
     }
-    return reinterpret_cast<BatchFn>(fn_ptr);
-  }();
-
-  // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 /
-  // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy
-  // fallback below, which is correct on any stream. Real and per-thread-default
-  // streams take the batch fast path.
-  const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy;
-  if (batch_fn != nullptr && usable_stream) {
-    CUmemcpyAttributes attr = {};
-    // ANY lets the DMA engine prefetch source bytes out of stream order,
-    // which is only safe when no GPU stream is concurrently writing the
-    // source.
-    attr.srcAccessOrder = is_src_access_order_any
-                              ? CU_MEMCPY_SRC_ACCESS_ORDER_ANY
-                              : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
-    size_t attrs_idx = 0;
-    size_t fail_idx = 0;
-    CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data),
-                               reinterpret_cast<CUdeviceptr*>(src_data),
-                               reinterpret_cast<size_t*>(size_data),
-                               static_cast<size_t>(n), &attr, &attrs_idx, 1,
-                               &fail_idx, static_cast<CUstream>(stream));
-    STD_TORCH_CHECK(result == CUDA_SUCCESS,
-                    "cuMemcpyBatchAsync failed at index ", fail_idx,
-                    " with error ", result);
-    return;
-  }
-#elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
-  // ROCm 7.1+ exposes hipMemcpyBatchAsync. The 7.2.1 implementation early-
-  // returns hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @
-  // rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so call with
-  // numAttrs=0.
-  {
-    hipMemcpyAttributes attr = {};
-    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, 0, &fail_idx, static_cast<hipStream_t>(stream));
-    STD_TORCH_CHECK(result == hipSuccess,
-                    "hipMemcpyBatchAsync failed at index ", fail_idx,
-                    " with error ", result);
-    return;
-  }
-#endif
-  {
-    // Fallback for CUDA < 12.8, older CUDA drivers, and ROCm < 7.1:
-    // 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]),
-                      static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
-                      stream);
+  #elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
+    // ROCm 7.1+ exposes hipMemcpyBatchAsync. The 7.2.1 implementation early-
+    // returns hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @
+    // rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so call with
+    // numAttrs=0.
+    {
+      hipMemcpyAttributes attr = {};
+      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, 0, &fail_idx, static_cast<hipStream_t>(stream));
+      STD_TORCH_CHECK(result == hipSuccess,
+                      "hipMemcpyBatchAsync failed at index ", fail_idx,
+                      " with error ", result);
+      return;
     }
+  #endif
+  }
+  // Fallback for CUDA < 12.8, older CUDA drivers, ROCm < 7.1, or when
+  // use_batch_api is false: 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]),
+                    static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
+                    stream);
   }
 }
 

csrc/libtorch_stable/ops.h

@@ -461,10 +461,14 @@ void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
                  const torch::stable::Tensor& block_mapping);
 
 // Batch swap: submit all block copies in a single driver call.
+// When use_batch_api is false, skip cuMemcpyBatchAsync/hipMemcpyBatchAsync
+// and use per-descriptor cudaMemcpyAsync fallback (safe for very large
+// descriptor counts that trigger driver batch-API segfaults).
 void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
                        const torch::stable::Tensor& dst_ptrs,
                        const torch::stable::Tensor& sizes,
-                       bool is_src_access_order_any);
+                       bool is_src_access_order_any,
+                       bool use_batch_api = true);
 
 void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
                        torch::stable::Tensor& key_cache,

csrc/libtorch_stable/torch_bindings.cpp

@@ -782,10 +782,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
       "            int block_size_in_bytes, Tensor block_mapping) -> ()");
 
   // Batch swap: submit all block copies in a single driver call.
+  // When use_batch_api is false, skip cuMemcpyBatchAsync/hipMemcpyBatchAsync
+  // and use per-descriptor cudaMemcpyAsync fallback.
   ops.def(
       "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
       "                  Tensor sizes,"
-      "                  bool is_src_access_order_any=False) -> ()");
+      "                  bool is_src_access_order_any=False,"
+      "                  bool use_batch_api=True) -> ()");
 
   // Reshape the key and value tensors and cache them.
   ops.def(