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, 07:15 UTC
[Bugfix] Runtime driver check for cuMemcpyBatchAsync in swap_blocks_batch (#38919)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Itay Etelis committed
Apr 11, 2026 at 20:02 UTC
bd8bd52308ea40c0d534d17969110b03e935a382
1 file changed
+42
-30
csrc/cache_kernels.cu
+42
-30
@@ -104,37 +104,49 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
104
static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
105
static_assert(sizeof(size_t) == sizeof(int64_t));
106
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
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(
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.
131
- for (int64_t i = 0; i < n; i++) {
132
- cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
133
- reinterpret_cast<void*>(src_data[i]),
134
- static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
135
- stream);
136
- }
107
+ // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that
108
+ // binaries compiled with CUDA 12.8+ still work on older drivers, and
109
+ // we avoid the CUDA 13.0 header remapping (#define to _v2 signature).
110
+ // The function pointer is cached after the first call.
111
+ using BatchFn =
112
+ CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t,
113
+ CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream);
114
+ static BatchFn batch_fn = []() -> BatchFn {
115
+ CUdriverProcAddressQueryResult sym_status;
116
+ void* fn_ptr = nullptr;
117
+ CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080,
118
+ CU_GET_PROC_ADDRESS_DEFAULT, &sym_status);
119
+ if (res != CUDA_SUCCESS || fn_ptr == nullptr) {
120
+ return nullptr;
121
+ }
122
+ return reinterpret_cast<BatchFn>(fn_ptr);
123
+ }();
124
+
125
+ if (batch_fn != nullptr) {
126
+ CUmemcpyAttributes attr = {};
127
+ attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
128
+ size_t attrs_idx = 0;
129
+ size_t fail_idx = 0;
130
+ CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data),
131
+ reinterpret_cast<CUdeviceptr*>(src_data),
132
+ reinterpret_cast<size_t*>(size_data),
133
+ static_cast<size_t>(n), &attr, &attrs_idx, 1,
134
+ &fail_idx, static_cast<CUstream>(stream));
135
+ TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
136
+ fail_idx, " with error ", result);
137
+ } else
138
#endif
139
+ {
140
+ // Fallback for CUDA < 12.8, older drivers, and ROCm:
141
+ // individual async copies.
142
+ // cudaMemcpyDefault lets the driver infer direction from pointer types.
143
+ for (int64_t i = 0; i < n; i++) {
144
+ cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
145
+ reinterpret_cast<void*>(src_data[i]),
146
+ static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
147
+ stream);
148
+ }
149
+ }
150
}
151
152
namespace vllm {