diff --git a/sgl-kernel/csrc/kvcacheio/transfer.cu b/sgl-kernel/csrc/kvcacheio/transfer.cu index 7b1e264e0..230932c15 100644 --- a/sgl-kernel/csrc/kvcacheio/transfer.cu +++ b/sgl-kernel/csrc/kvcacheio/transfer.cu @@ -806,17 +806,29 @@ inline void transfer_kv_page_first_direct_impl( } // Symbol gate: runtime may not expose cudaMemcpyBatchAsync in some environments. - using CudaMemcpyBatchAsyncFn = - cudaError_t (*)(void**, void**, size_t*, size_t, cudaMemcpyAttributes*, size_t*, size_t, size_t*, cudaStream_t); - static CudaMemcpyBatchAsyncFn cuda_memcpy_batch_async = []() { - void* symbol = dlsym(RTLD_DEFAULT, "cudaMemcpyBatchAsync"); - return reinterpret_cast(symbol); - }(); - if (cuda_memcpy_batch_async == nullptr) { + static void* cuda_memcpy_batch_async_sym = dlsym(RTLD_DEFAULT, "cudaMemcpyBatchAsync"); + if (cuda_memcpy_batch_async_sym == nullptr) { fallback_to_page_copy(); return; } + // CUDA 13.0 removed the failIdx parameter from cudaMemcpyBatchAsync. The ABI + // of the dlsym'd symbol is determined by the libcudart loaded in this process, + // not the host driver — a cu12 runtime on a cu13 driver host (common in + // containers) still exposes the 9-param v12 signature. Dispatching on the + // driver version here would segfault in that case (verified empirically). + // Use cudaRuntimeGetVersion so the signature follows the runtime. The + // runtime version is process-constant, so cache the query (static init is + // thread-safe in C++11+) to keep the KV-transfer hot path free of a redundant + // runtime API call per invocation. + static int runtime_version = 0; + static cudaError_t runtime_version_err = cudaRuntimeGetVersion(&runtime_version); + if (runtime_version_err != cudaSuccess) { + fallback_to_page_copy(); + return; + } + static const bool use_v13_signature = runtime_version >= 13000; + size_t num_copies = 0; std::vector batch_srcs; std::vector batch_dsts; @@ -916,17 +928,36 @@ inline void transfer_kv_page_first_direct_impl( TORCH_CHECK(batch_srcs.size() == num_copies, "Batch memcpy count mismatch"); if (num_copies > 0) { + cudaError_t err; size_t fail_idx = std::numeric_limits::max(); - cudaError_t err = cuda_memcpy_batch_async( - batch_dsts.data(), - batch_srcs.data(), - batch_sizes.data(), - num_copies, - &attrs, - attrs_idxs.data(), - 1, - &fail_idx, - stream); + if (use_v13_signature) { + using FnV13 = cudaError_t (*)( + void* const*, + const void* const*, + const size_t*, + size_t, + cudaMemcpyAttributes*, + size_t*, + size_t, + cudaStream_t); + auto fn = reinterpret_cast(cuda_memcpy_batch_async_sym); + err = fn( + batch_dsts.data(), batch_srcs.data(), batch_sizes.data(), num_copies, &attrs, attrs_idxs.data(), 1, stream); + } else { + using FnV12 = cudaError_t (*)( + void**, void**, size_t*, size_t, cudaMemcpyAttributes*, size_t*, size_t, size_t*, cudaStream_t); + auto fn = reinterpret_cast(cuda_memcpy_batch_async_sym); + err = + fn(batch_dsts.data(), + batch_srcs.data(), + batch_sizes.data(), + num_copies, + &attrs, + attrs_idxs.data(), + 1, + &fail_idx, + stream); + } if (err == cudaErrorNotSupported || err == cudaErrorCallRequiresNewerDriver) { fallback_to_page_copy(); return;