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 5, 13:30 UTC
[Bug] Fix compile error for `swap_blocks_batch` in CUDA 13 (#38915)
Wentao Ye committed
Apr 3, 2026 at 19:56 UTC
062f1a2d706cfb631521461f6a7ad533cf69a4a8
1 file changed
+17
-8
csrc/cache_kernels.cu
+17
-8
@@ -91,9 +91,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
91
92
if (n == 0) return;
93
94
- const int64_t* src_data = src_ptrs.data_ptr<int64_t>();
95
- const int64_t* dst_data = dst_ptrs.data_ptr<int64_t>();
96
- const int64_t* size_data = sizes.data_ptr<int64_t>();
94
+ int64_t* src_data = src_ptrs.mutable_data_ptr<int64_t>();
95
+ int64_t* dst_data = dst_ptrs.mutable_data_ptr<int64_t>();
96
+ int64_t* size_data = sizes.mutable_data_ptr<int64_t>();
97
98
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
99
@@ -107,15 +107,24 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
107
CUmemcpyAttributes attr = {};
108
attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
109
size_t attrs_idx = 0;
110
+ #if defined(CUDA_VERSION) && CUDA_VERSION >= 13000
111
+ CUresult result = cuMemcpyBatchAsync(
112
+ reinterpret_cast<CUdeviceptr*>(dst_data),
113
+ reinterpret_cast<CUdeviceptr*>(src_data),
114
+ reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
115
+ &attrs_idx, 1, static_cast<CUstream>(stream));
116
+ TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed with error ",
117
+ result);
118
+ #else
119
size_t fail_idx = 0;
120
CUresult result = cuMemcpyBatchAsync(
112
- reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(dst_data)),
113
- reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(src_data)),
114
- reinterpret_cast<size_t*>(const_cast<int64_t*>(size_data)),
115
- static_cast<size_t>(n), &attr, &attrs_idx, 1, &fail_idx,
116
- static_cast<CUstream>(stream));
121
+ reinterpret_cast<CUdeviceptr*>(dst_data),
122
+ reinterpret_cast<CUdeviceptr*>(src_data),
123
+ reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
124
+ &attrs_idx, 1, &fail_idx, static_cast<CUstream>(stream));
125
TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
126
fail_idx, " with error ", result);
127
+ #endif
128
#else
129
// Fallback for CUDA < 12.8 and ROCm: individual async copies.
130
// cudaMemcpyDefault lets the driver infer direction from pointer types.