From 36dc973cbfb9c35688cc2799a61591d80ce25d13 Mon Sep 17 00:00:00 2001 From: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Date: Sat, 28 Feb 2026 03:43:14 +0800 Subject: [PATCH] [HiCache] refactor page_first_direct io kernel (#18113) Co-authored-by: hzh0425 --- sgl-kernel/csrc/kvcacheio/transfer.cu | 191 +++++++++++++++++++++++--- sgl-kernel/tests/test_kvcacheio.py | 78 ++++++----- 2 files changed, 216 insertions(+), 53 deletions(-) diff --git a/sgl-kernel/csrc/kvcacheio/transfer.cu b/sgl-kernel/csrc/kvcacheio/transfer.cu index c1f37dfc6..7b1e264e0 100644 --- a/sgl-kernel/csrc/kvcacheio/transfer.cu +++ b/sgl-kernel/csrc/kvcacheio/transfer.cu @@ -1,10 +1,14 @@ #include #include #include +#include #include +#include +#include #ifndef USE_ROCM +#include #define WARP_SIZE 32 #include "pytorch_extension_utils.h" #else @@ -743,48 +747,195 @@ inline void transfer_kv_page_first_direct_impl( auto src_indices_cpu = src_indices.cpu(); auto dst_indices_cpu = dst_indices.cpu(); const int64_t num_pages = src_indices_cpu.size(0) / page_size; + int64_t* src_indices_ptr = src_indices_cpu.data_ptr(); + int64_t* dst_indices_ptr = dst_indices_cpu.data_ptr(); + + auto fallback_to_page_copy = [&]() { + if constexpr (IsLf2Pf) { + const bool is_mla = dst_ptrs.size() == 1; + const int64_t num_layers = is_mla ? src_ptrs.size() : src_ptrs.size() / 2; + for (const auto i : c10::irange(num_pages)) { + const int64_t s_index = src_indices_ptr[i * page_size]; + const int64_t d_index = dst_indices_ptr[i * page_size] / page_size; + for (int64_t j = 0; j < num_layers; ++j) { + transfer_page_direct( + src_ptrs[j], dst_ptrs[0].select(0, d_index).select(0, start_layer_id + j), s_index, 0, page_size); + if (!is_mla) { + transfer_page_direct( + src_ptrs[j + num_layers], + dst_ptrs[1].select(0, d_index).select(0, start_layer_id + j), + s_index, + 0, + page_size); + } + } + } + } else { + const bool is_mla = src_ptrs.size() == 1; + const int64_t num_layers = is_mla ? dst_ptrs.size() : dst_ptrs.size() / 2; + for (const auto i : c10::irange(num_pages)) { + const int64_t s_index = src_indices_ptr[i * page_size] / page_size; + const int64_t d_index = dst_indices_ptr[i * page_size]; + for (int64_t j = 0; j < num_layers; ++j) { + transfer_page_direct( + src_ptrs[0].select(0, s_index).select(0, start_layer_id + j), dst_ptrs[j], 0, d_index, page_size); + if (!is_mla) { + transfer_page_direct( + src_ptrs[1].select(0, s_index).select(0, start_layer_id + j), + dst_ptrs[j + num_layers], + 0, + d_index, + page_size); + } + } + } + } + }; + +#if defined(USE_ROCM) || !defined(CUDA_VERSION) || CUDA_VERSION < 12080 + fallback_to_page_copy(); + return; + +#else + // Driver capability gate: only use cudaMemcpyBatchAsync on CUDA 12.8+ drivers. + int driver_version = 0; + cudaError_t driver_version_err = cudaDriverGetVersion(&driver_version); + if (driver_version_err != cudaSuccess || driver_version < 12080) { + fallback_to_page_copy(); + return; + } + + // 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) { + fallback_to_page_copy(); + return; + } + + size_t num_copies = 0; + std::vector batch_srcs; + std::vector batch_dsts; + std::vector batch_sizes; + std::vector attrs_idxs(1, 0); + cudaMemcpyAttributes attrs{}; + const int device_id = at::cuda::current_device(); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + auto append_copy = [&](void* src, void* dst, size_t size_bytes) { + batch_srcs.push_back(src); + batch_dsts.push_back(dst); + batch_sizes.push_back(size_bytes); + }; if constexpr (IsLf2Pf) { const bool is_mla = dst_ptrs.size() == 1; const int64_t num_layers = is_mla ? src_ptrs.size() : src_ptrs.size() / 2; + const int64_t dst_stride0 = dst_ptrs[0].stride(0); + const int64_t dst_stride1 = dst_ptrs[0].stride(1); + const int64_t src_stride0 = src_ptrs[0].stride(0); + const int64_t elem_size = dst_ptrs[0].element_size(); + const int64_t copy_size_bytes = page_size * src_stride0 * elem_size; + attrs.srcAccessOrder = cudaMemcpySrcAccessOrderStream; + attrs.srcLocHint.type = cudaMemLocationTypeDevice; + attrs.srcLocHint.id = device_id; + attrs.dstLocHint.type = cudaMemLocationTypeHost; + attrs.dstLocHint.id = 0; + attrs.flags = 0; + + num_copies = static_cast(num_pages) * static_cast(num_layers) * static_cast(is_mla ? 1 : 2); + batch_srcs.reserve(num_copies); + batch_dsts.reserve(num_copies); + batch_sizes.reserve(num_copies); + for (const auto i : c10::irange(num_pages)) { - auto s_index = src_indices_cpu[i * page_size].item(); - auto d_index = dst_indices_cpu[i * page_size].item() / page_size; + auto s_index = src_indices_ptr[i * page_size]; + auto d_index = dst_indices_ptr[i * page_size] / page_size; + for (int64_t j = 0; j < num_layers; ++j) { - transfer_page_direct( - src_ptrs[j], dst_ptrs[0].select(0, d_index).select(0, start_layer_id + j), s_index, 0, page_size); + const char* src_k_ptr = static_cast(src_ptrs[j].data_ptr()) + s_index * src_stride0 * elem_size; + char* dst_k_ptr = static_cast(dst_ptrs[0].data_ptr()) + d_index * dst_stride0 * elem_size + + (start_layer_id + j) * dst_stride1 * elem_size; + append_copy(const_cast(src_k_ptr), dst_k_ptr, copy_size_bytes); + if (!is_mla) { - transfer_page_direct( - src_ptrs[j + num_layers], - dst_ptrs[1].select(0, d_index).select(0, start_layer_id + j), - s_index, - 0, - page_size); + const char* src_v_ptr = + static_cast(src_ptrs[j + num_layers].data_ptr()) + s_index * src_stride0 * elem_size; + char* dst_v_ptr = static_cast(dst_ptrs[1].data_ptr()) + d_index * dst_stride0 * elem_size + + (start_layer_id + j) * dst_stride1 * elem_size; + append_copy(const_cast(src_v_ptr), dst_v_ptr, copy_size_bytes); } } } + } else { const bool is_mla = src_ptrs.size() == 1; const int64_t num_layers = is_mla ? dst_ptrs.size() : dst_ptrs.size() / 2; + const int64_t src_stride0 = src_ptrs[0].stride(0); + const int64_t src_stride1 = src_ptrs[0].stride(1); + const int64_t dst_stride0 = dst_ptrs[0].stride(0); + const int64_t elem_size = src_ptrs[0].element_size(); + const int64_t copy_size_bytes = page_size * dst_stride0 * elem_size; + attrs.srcAccessOrder = cudaMemcpySrcAccessOrderStream; + attrs.srcLocHint.type = cudaMemLocationTypeHost; + attrs.srcLocHint.id = 0; + attrs.dstLocHint.type = cudaMemLocationTypeDevice; + attrs.dstLocHint.id = device_id; + attrs.flags = 0; + + num_copies = static_cast(num_pages) * static_cast(num_layers) * static_cast(is_mla ? 1 : 2); + batch_srcs.reserve(num_copies); + batch_dsts.reserve(num_copies); + batch_sizes.reserve(num_copies); + for (const auto i : c10::irange(num_pages)) { - auto s_index = src_indices_cpu[i * page_size].item() / page_size; - auto d_index = dst_indices_cpu[i * page_size].item(); + auto s_index = src_indices_ptr[i * page_size] / page_size; + auto d_index = dst_indices_ptr[i * page_size]; + for (int64_t j = 0; j < num_layers; ++j) { - transfer_page_direct( - src_ptrs[0].select(0, s_index).select(0, start_layer_id + j), dst_ptrs[j], 0, d_index, page_size); + const char* src_k_ptr = static_cast(src_ptrs[0].data_ptr()) + s_index * src_stride0 * elem_size + + (start_layer_id + j) * src_stride1 * elem_size; + char* dst_k_ptr = static_cast(dst_ptrs[j].data_ptr()) + d_index * dst_stride0 * elem_size; + append_copy(const_cast(src_k_ptr), dst_k_ptr, copy_size_bytes); + if (!is_mla) { - transfer_page_direct( - src_ptrs[1].select(0, s_index).select(0, start_layer_id + j), - dst_ptrs[j + num_layers], - 0, - d_index, - page_size); + const char* src_v_ptr = static_cast(src_ptrs[1].data_ptr()) + s_index * src_stride0 * elem_size + + (start_layer_id + j) * src_stride1 * elem_size; + char* dst_v_ptr = static_cast(dst_ptrs[j + num_layers].data_ptr()) + d_index * dst_stride0 * elem_size; + append_copy(const_cast(src_v_ptr), dst_v_ptr, copy_size_bytes); } } } } + + TORCH_CHECK(batch_srcs.size() == num_copies, "Batch memcpy count mismatch"); + if (num_copies > 0) { + 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 (err == cudaErrorNotSupported || err == cudaErrorCallRequiresNewerDriver) { + fallback_to_page_copy(); + return; + } + if (err != cudaSuccess) { + TORCH_CHECK(false, "cudaMemcpyBatchAsync failed. failIdx=", fail_idx, " error=", cudaGetErrorString(err)); + } + } +#endif } void transfer_kv_per_layer_direct_pf_lf( diff --git a/sgl-kernel/tests/test_kvcacheio.py b/sgl-kernel/tests/test_kvcacheio.py index 5ba1c85a6..16fc7e826 100644 --- a/sgl-kernel/tests/test_kvcacheio.py +++ b/sgl-kernel/tests/test_kvcacheio.py @@ -317,6 +317,7 @@ def test_transfer_kv_pf_direct( torch.set_default_dtype(dtype) device = "cuda" torch.cuda.manual_seed(42) + test_stream = torch.cuda.Stream() num_layers = 4 @@ -356,13 +357,16 @@ def test_transfer_kv_pf_direct( dst_pool_direct = torch.zeros_like(dst_pool_ref) torch.cuda.synchronize() - transfer_kv_all_layer_direct_lf_pf( - src_pool_ptrs, - [dst_pool_direct], - src_indices_host, - dst_indices_host, - page_size, - ) + with torch.cuda.stream(test_stream): + transfer_kv_all_layer_direct_lf_pf( + src_pool_ptrs, + [dst_pool_direct], + src_indices_host, + dst_indices_host, + page_size, + ) + test_stream.synchronize() + for i in range(num_layers): ref_copy_with_indices_pf_direct( src_pool, @@ -393,13 +397,16 @@ def test_transfer_kv_pf_direct( dst_v_pool_direct = torch.zeros_like(dst_v_pool_ref) torch.cuda.synchronize() - transfer_kv_all_layer_direct_lf_pf( - src_k_pool_ptrs + src_v_pool_ptrs, - [dst_k_pool_direct, dst_v_pool_direct], - src_indices_host, - dst_indices_host, - page_size, - ) + with torch.cuda.stream(test_stream): + transfer_kv_all_layer_direct_lf_pf( + src_k_pool_ptrs + src_v_pool_ptrs, + [dst_k_pool_direct, dst_v_pool_direct], + src_indices_host, + dst_indices_host, + page_size, + ) + test_stream.synchronize() + for i in range(num_layers): ref_copy_with_indices_pf_direct( src_k_pool, @@ -435,14 +442,17 @@ def test_transfer_kv_pf_direct( dst_pool_direct_ptrs = [dst_pool_direct[i] for i in range(num_layers)] torch.cuda.synchronize() - transfer_kv_per_layer_direct_pf_lf( - [src_pool], - [dst_pool_direct_ptrs[layer_idx_to_test]], - src_indices_host, - dst_indices_host, - layer_idx_to_test, - page_size, - ) + with torch.cuda.stream(test_stream): + transfer_kv_per_layer_direct_pf_lf( + [src_pool], + [dst_pool_direct_ptrs[layer_idx_to_test]], + src_indices_host, + dst_indices_host, + layer_idx_to_test, + page_size, + ) + test_stream.synchronize() + ref_copy_with_indices_pf_direct( src_pool, dst_pool_ref, @@ -473,17 +483,19 @@ def test_transfer_kv_pf_direct( dst_v_pool_direct_ptrs = [dst_v_pool_direct[i] for i in range(num_layers)] torch.cuda.synchronize() - transfer_kv_per_layer_direct_pf_lf( - [src_k_pool, src_v_pool], - [ - dst_k_pool_direct_ptrs[layer_idx_to_test], - dst_v_pool_direct_ptrs[layer_idx_to_test], - ], - src_indices_host, - dst_indices_host, - layer_idx_to_test, - page_size, - ) + with torch.cuda.stream(test_stream): + transfer_kv_per_layer_direct_pf_lf( + [src_k_pool, src_v_pool], + [ + dst_k_pool_direct_ptrs[layer_idx_to_test], + dst_v_pool_direct_ptrs[layer_idx_to_test], + ], + src_indices_host, + dst_indices_host, + layer_idx_to_test, + page_size, + ) + test_stream.synchronize() ref_copy_with_indices_pf_direct( src_k_pool,