From 3e6281d0aad083211f6348d967556c9633ecd77c Mon Sep 17 00:00:00 2001 From: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Date: Sun, 26 Oct 2025 15:53:50 +0800 Subject: [PATCH] [HiCache]Page head layout IO kernel (#11615) --- sgl-kernel/csrc/common_extension.cc | 10 + sgl-kernel/csrc/kvcacheio/transfer.cu | 274 ++++++++++++++++++++-- sgl-kernel/include/sgl_kernel_ops.h | 30 +++ sgl-kernel/python/sgl_kernel/kvcacheio.py | 86 ++++++- sgl-kernel/tests/test_kvcacheio.py | 205 ++++++++++++++++ 5 files changed, 574 insertions(+), 31 deletions(-) diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index e597e3111..c9d6f47fb 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -370,6 +370,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "transfer_kv_per_layer_pf_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor " "dst_indices, int layer_id, int item_size, int src_layout_dim, int block_quota, int num_warps_per_block) -> ()"); m.impl("transfer_kv_per_layer_pf_lf", torch::kCUDA, &transfer_kv_per_layer_pf_lf); + m.def( + "transfer_kv_per_layer_ph_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor " + "dst_indices, int layer_id, int item_size, int src_layout_dim, int page_size, int head_num, int block_quota, int " + "num_warps_per_block) -> ()"); + m.impl("transfer_kv_per_layer_ph_lf", torch::kCUDA, &transfer_kv_per_layer_ph_lf); m.def( "transfer_kv_all_layer(Tensor src_k_layers, Tensor dst_k_layers, Tensor src_v_layers, Tensor dst_v_layers, " "Tensor src_indices, Tensor dst_indices, int item_size, int num_layers, int block_quota, int " @@ -380,6 +385,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int block_quota, int " "num_warps_per_block) -> ()"); m.impl("transfer_kv_all_layer_lf_pf", torch::kCUDA, &transfer_kv_all_layer_lf_pf); + m.def( + "transfer_kv_all_layer_lf_ph(Tensor src_k_layers, Tensor dst_k, Tensor src_v_layers, Tensor dst_v, " + "Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int page_size, int " + "head_num, int block_quota, int num_warps_per_block) -> ()"); + m.impl("transfer_kv_all_layer_lf_ph", torch::kCUDA, &transfer_kv_all_layer_lf_ph); m.def( "transfer_kv_per_layer_mla(Tensor src, Tensor dst, Tensor src_indices, Tensor dst_indices, int item_size, int " "block_quota, int num_warps_per_block) -> ()"); diff --git a/sgl-kernel/csrc/kvcacheio/transfer.cu b/sgl-kernel/csrc/kvcacheio/transfer.cu index bca9f326c..c1f37dfc6 100644 --- a/sgl-kernel/csrc/kvcacheio/transfer.cu +++ b/sgl-kernel/csrc/kvcacheio/transfer.cu @@ -68,6 +68,140 @@ __device__ __forceinline__ T* get_global_offset_lf_tbl( return reinterpret_cast(layer_base_tbl[layer_id]) + page_id * item_size_bytes; } +template +__device__ __forceinline__ T* get_global_offset_per_head_lf( + T* base, + const uintptr_t* __restrict__ /*unused*/, + int64_t layer_id, + int64_t layer_dim, + int64_t page_id, + int64_t item_size_bytes, + int64_t head_id, + int64_t head_num, + int64_t /*unused*/) { + // layer first offset func per head + return base + layer_id * layer_dim + page_id * item_size_bytes + item_size_bytes / head_num * head_id; +} + +template +__device__ __forceinline__ T* get_global_offset_per_head_lf_tbl( + T* /*unused*/, + const uintptr_t* __restrict__ layer_base_tbl, + int64_t layer_id, + int64_t /*unused*/, + int64_t page_id, + int64_t item_size_bytes, + int64_t head_id, + int64_t head_num, + int64_t /*unused*/) { + return reinterpret_cast(layer_base_tbl[layer_id]) + page_id * item_size_bytes + + item_size_bytes / head_num * head_id; +} + +template +__device__ __forceinline__ T* get_global_offset_ph( + T* base, + const uintptr_t* __restrict__ /*unused*/, + int64_t layer_id, + int64_t page_dim, + int64_t page_id, + int64_t item_size_bytes, + int64_t head_id, + int64_t head_num, + int64_t page_size) { + // page head layout: [page_num, head_num, page_size, layer_num, head_dim] + return base + page_id / page_size * page_size * page_dim + // page_num dimension offset + page_dim / head_num * head_id * page_size + // head_num dimension offset + page_id % page_size * page_dim / head_num + // page_size dimension offset + layer_id * item_size_bytes / head_num; // layer_num dimension offset +} + +template +__global__ void transfer_page_head_kernel_impl( + const void* __restrict__ src_k, + void* __restrict__ dst_k, + const void* __restrict__ src_v, + void* __restrict__ dst_v, + const int64_t* __restrict__ src_indices, + const int64_t* __restrict__ dst_indices, + int64_t start_layer_id, + int64_t num_layers_to_process, + int64_t num_items, + int64_t items_per_warp, + int64_t item_size_bytes, + int64_t src_layout_dim, + int64_t dst_layout_dim, + const uintptr_t* __restrict__ src_k_layer_tbl, + const uintptr_t* __restrict__ dst_k_layer_tbl, + const uintptr_t* __restrict__ src_v_layer_tbl, + const uintptr_t* __restrict__ dst_v_layer_tbl, + const int64_t page_size, + const int64_t head_num) { + int32_t tid = blockIdx.x * blockDim.x + threadIdx.x; + int32_t lane_id = tid % WARP_SIZE; + int32_t warp_id = tid / WARP_SIZE; + const int64_t head_size_bytes = item_size_bytes / head_num; + + for (int i = 0; i < items_per_warp; ++i) { + int64_t item_id = warp_id * items_per_warp + i; + if (item_id >= num_items) { + break; + } + const int64_t src_page_id = src_indices[item_id]; + const int64_t dst_page_id = dst_indices[item_id]; + + // Loop over layers if necessary + for (int64_t layer_id = start_layer_id; layer_id < start_layer_id + num_layers_to_process; ++layer_id) { + // For page head layout, the cache of each head in the token is discontinuous, need to loop + for (int64_t head_id = 0; head_id < head_num; ++head_id) { + const char* src_k_ptr = SrcOffsetFn( + static_cast(src_k), + src_k_layer_tbl, + layer_id, + src_layout_dim, + src_page_id, + item_size_bytes, + head_id, + head_num, + page_size); + char* dst_k_ptr = DstOffsetFn( + static_cast(dst_k), + dst_k_layer_tbl, + layer_id, + dst_layout_dim, + dst_page_id, + item_size_bytes, + head_id, + head_num, + page_size); + transfer_item_warp(lane_id, src_k_ptr, dst_k_ptr, head_size_bytes); + + const char* src_v_ptr = SrcOffsetFn( + static_cast(src_v), + src_v_layer_tbl, + layer_id, + src_layout_dim, + src_page_id, + item_size_bytes, + head_id, + head_num, + page_size); + char* dst_v_ptr = DstOffsetFn( + static_cast(dst_v), + dst_v_layer_tbl, + layer_id, + dst_layout_dim, + dst_page_id, + item_size_bytes, + head_id, + head_num, + page_size); + transfer_item_warp(lane_id, src_v_ptr, dst_v_ptr, head_size_bytes); + } + } + } +} + template __global__ void transfer_kernel_impl( const void* __restrict__ src_k, @@ -118,7 +252,7 @@ __global__ void transfer_kernel_impl( } } -template +template void transfer_kv_launcher( const at::Tensor& src_k, at::Tensor& dst_k, @@ -136,7 +270,9 @@ void transfer_kv_launcher( const at::Tensor& src_v_layers, const at::Tensor& dst_v_layers, int64_t block_quota, - int64_t num_warps_per_block) { + int64_t num_warps_per_block, + const int64_t page_size = 16, + const int64_t head_num = 1) { TORCH_CHECK(src_indices.is_cuda(), "Source indices must be a CUDA tensor"); TORCH_CHECK(dst_indices.is_cuda(), "Destination indices must be a CUDA tensor"); TORCH_CHECK(src_indices.scalar_type() == at::kLong, "Source indices must be of type long"); @@ -161,24 +297,47 @@ void transfer_kv_launcher( const uintptr_t* dst_v_tbl_ptr = IsMLA || !dst_v_layers.defined() ? nullptr : dst_v_layers.data_ptr(); cudaStream_t torch_current_stream = at::cuda::getCurrentCUDAStream(); - transfer_kernel_impl<<>>( - src_k_ptr, - dst_k_ptr, - src_v_ptr, - dst_v_ptr, - src_indices.data_ptr(), - dst_indices.data_ptr(), - start_layer_id, - num_layers_to_process, - num_items, - items_per_warp, - item_size, - src_layout_dim, - dst_layout_dim, - src_k_tbl_ptr, - dst_k_tbl_ptr, - src_v_tbl_ptr, - dst_v_tbl_ptr); + if constexpr (PageHeadLayout) { + transfer_page_head_kernel_impl<<>>( + src_k_ptr, + dst_k_ptr, + src_v_ptr, + dst_v_ptr, + src_indices.data_ptr(), + dst_indices.data_ptr(), + start_layer_id, + num_layers_to_process, + num_items, + items_per_warp, + item_size, + src_layout_dim, + dst_layout_dim, + src_k_tbl_ptr, + dst_k_tbl_ptr, + src_v_tbl_ptr, + dst_v_tbl_ptr, + page_size, + head_num); + } else { + transfer_kernel_impl<<>>( + src_k_ptr, + dst_k_ptr, + src_v_ptr, + dst_v_ptr, + src_indices.data_ptr(), + dst_indices.data_ptr(), + start_layer_id, + num_layers_to_process, + num_items, + items_per_warp, + item_size, + src_layout_dim, + dst_layout_dim, + src_k_tbl_ptr, + dst_k_tbl_ptr, + src_v_tbl_ptr, + dst_v_tbl_ptr); + } C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -246,6 +405,43 @@ void transfer_kv_per_layer_pf_lf( num_warps_per_block); } +void transfer_kv_per_layer_ph_lf( + const at::Tensor src_k, + at::Tensor dst_k, + const at::Tensor src_v, + at::Tensor dst_v, + const at::Tensor src_indices, + const at::Tensor dst_indices, + int64_t layer_id, + int64_t item_size, + int64_t src_layout_dim, + int64_t page_size, + int64_t head_num, + int64_t block_quota, + int64_t num_warps_per_block) { + at::Tensor empty; + transfer_kv_launcher, get_global_offset_per_head_lf, false, true>( + src_k, + dst_k, + src_v, + dst_v, + src_indices, + dst_indices, + layer_id, + 1, + item_size, + src_layout_dim, + 0, + empty, + empty, + empty, + empty, + block_quota, + num_warps_per_block, + page_size, + head_num); +} + void transfer_kv_all_layer( const at::Tensor src_k_layers, const at::Tensor dst_k_layers, @@ -313,6 +509,44 @@ void transfer_kv_all_layer_lf_pf( num_warps_per_block); } +void transfer_kv_all_layer_lf_ph( + const at::Tensor src_k_layers, + at::Tensor dst_k, + const at::Tensor src_v_layers, + at::Tensor dst_v, + const at::Tensor src_indices, + const at::Tensor dst_indices, + int64_t item_size, + int64_t dst_layout_dim, + int64_t num_layers, + int64_t page_size, + int64_t head_num, + int64_t block_quota, + int64_t num_warps_per_block) { + TORCH_CHECK(num_layers == src_k_layers.size(0), "Number of layers in source k tensor does not match num_layers"); + at::Tensor empty; + transfer_kv_launcher, get_global_offset_ph, false, true>( + empty, + dst_k, + empty, + dst_v, + src_indices, + dst_indices, + 0, + num_layers, + item_size, + 0, + dst_layout_dim, + src_k_layers, + empty, + src_v_layers, + empty, + block_quota, + num_warps_per_block, + page_size, + head_num); +} + void transfer_kv_per_layer_mla( const at::Tensor src, at::Tensor dst, diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 6be8af703..10fd29c5e 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -562,6 +562,21 @@ void transfer_kv_per_layer_pf_lf( int64_t block_quota, int64_t num_warps_per_block); +void transfer_kv_per_layer_ph_lf( + const at::Tensor src_k, + at::Tensor dst_k, + const at::Tensor src_v, + at::Tensor dst_v, + const at::Tensor src_indices, + const at::Tensor dst_indices, + int64_t layer_id, + int64_t item_size, + int64_t src_layout_dim, + int64_t page_size, + int64_t head_num, + int64_t block_quota, + int64_t num_warps_per_block); + void transfer_kv_all_layer( const at::Tensor src_k_layers, const at::Tensor dst_k_layers, @@ -587,6 +602,21 @@ void transfer_kv_all_layer_lf_pf( int64_t block_quota, int64_t num_warps_per_block); +void transfer_kv_all_layer_lf_ph( + const at::Tensor src_k_layers, + at::Tensor dst_k, + const at::Tensor src_v_layers, + at::Tensor dst_v, + const at::Tensor src_indices, + const at::Tensor dst_indices, + int64_t item_size, + int64_t dst_layout_dim, + int64_t num_layers, + int64_t page_size, + int64_t head_num, + int64_t block_quota, + int64_t num_warps_per_block); + void transfer_kv_per_layer_mla( const at::Tensor src, at::Tensor dst, diff --git a/sgl-kernel/python/sgl_kernel/kvcacheio.py b/sgl-kernel/python/sgl_kernel/kvcacheio.py index 5714b6a0d..908a3efa9 100644 --- a/sgl-kernel/python/sgl_kernel/kvcacheio.py +++ b/sgl-kernel/python/sgl_kernel/kvcacheio.py @@ -21,7 +21,7 @@ def transfer_kv_per_layer( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_per_layer( + torch.ops.sgl_kernel.transfer_kv_per_layer.default( src_k, dst_k, src_v, @@ -47,7 +47,7 @@ def transfer_kv_per_layer_pf_lf( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_per_layer_pf_lf( + torch.ops.sgl_kernel.transfer_kv_per_layer_pf_lf.default( src_k, dst_k, src_v, @@ -62,6 +62,38 @@ def transfer_kv_per_layer_pf_lf( ) +def transfer_kv_per_layer_ph_lf( + src_k: torch.Tensor, + dst_k: torch.Tensor, + src_v: torch.Tensor, + dst_v: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + layer_id: int, + item_size: int, + src_layout_dim: int, + page_size: int, + head_num: int, + block_quota: int = 2, + num_warps_per_block: int = 16 if _is_hip else 32, +): + torch.ops.sgl_kernel.transfer_kv_per_layer_ph_lf.default( + src_k, + dst_k, + src_v, + dst_v, + src_indices, + dst_indices, + layer_id, + item_size, + src_layout_dim, + page_size, + head_num, + block_quota, + num_warps_per_block, + ) + + def transfer_kv_all_layer( src_k_layers: torch.Tensor, dst_k_layers: torch.Tensor, @@ -74,7 +106,7 @@ def transfer_kv_all_layer( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_all_layer( + torch.ops.sgl_kernel.transfer_kv_all_layer.default( src_k_layers, dst_k_layers, src_v_layers, @@ -101,7 +133,7 @@ def transfer_kv_all_layer_lf_pf( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_all_layer_lf_pf( + torch.ops.sgl_kernel.transfer_kv_all_layer_lf_pf.default( src_k_layers, dst_k, src_v_layers, @@ -116,6 +148,38 @@ def transfer_kv_all_layer_lf_pf( ) +def transfer_kv_all_layer_lf_ph( + src_k_layers: torch.Tensor, + dst_k: torch.Tensor, + src_v_layers: torch.Tensor, + dst_v: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + item_size: int, + dst_layout_dim: int, + num_layers: int, + page_size: int, + head_num: int, + block_quota: int = 2, + num_warps_per_block: int = 16 if _is_hip else 32, +): + torch.ops.sgl_kernel.transfer_kv_all_layer_lf_ph.default( + src_k_layers, + dst_k, + src_v_layers, + dst_v, + src_indices, + dst_indices, + item_size, + dst_layout_dim, + num_layers, + page_size, + head_num, + block_quota, + num_warps_per_block, + ) + + def transfer_kv_direct( src_layers: List[torch.Tensor], dst_layers: List[torch.Tensor], @@ -123,7 +187,7 @@ def transfer_kv_direct( dst_indices: torch.Tensor, page_size: int, ): - torch.ops.sgl_kernel.transfer_kv_direct( + torch.ops.sgl_kernel.transfer_kv_direct.default( src_layers, dst_layers, src_indices, dst_indices, page_size ) @@ -136,7 +200,7 @@ def transfer_kv_per_layer_direct_pf_lf( layer_id: int, page_size: int, ): - torch.ops.sgl_kernel.transfer_kv_per_layer_direct_pf_lf( + torch.ops.sgl_kernel.transfer_kv_per_layer_direct_pf_lf.default( src_ptrs, dst_ptrs, src_indices, dst_indices, layer_id, page_size ) @@ -148,7 +212,7 @@ def transfer_kv_all_layer_direct_lf_pf( dst_indices: torch.Tensor, page_size: int, ): - torch.ops.sgl_kernel.transfer_kv_all_layer_direct_lf_pf( + torch.ops.sgl_kernel.transfer_kv_all_layer_direct_lf_pf.default( src_ptrs, dst_ptrs, src_indices, dst_indices, page_size ) @@ -162,7 +226,7 @@ def transfer_kv_per_layer_mla( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_per_layer_mla( + torch.ops.sgl_kernel.transfer_kv_per_layer_mla.default( src, dst, src_indices, @@ -184,7 +248,7 @@ def transfer_kv_per_layer_mla_pf_lf( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_per_layer_mla_pf_lf( + torch.ops.sgl_kernel.transfer_kv_per_layer_mla_pf_lf.default( src, dst, src_indices, @@ -207,7 +271,7 @@ def transfer_kv_all_layer_mla( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_all_layer_mla( + torch.ops.sgl_kernel.transfer_kv_all_layer_mla.default( src_layers, dst_layers, src_indices, @@ -230,7 +294,7 @@ def transfer_kv_all_layer_mla_lf_pf( block_quota: int = 2, num_warps_per_block: int = 16 if _is_hip else 32, ): - torch.ops.sgl_kernel.transfer_kv_all_layer_mla_lf_pf( + torch.ops.sgl_kernel.transfer_kv_all_layer_mla_lf_pf.default( src_layers, dst, src_indices, diff --git a/sgl-kernel/tests/test_kvcacheio.py b/sgl-kernel/tests/test_kvcacheio.py index 07fcc2413..6328cc651 100644 --- a/sgl-kernel/tests/test_kvcacheio.py +++ b/sgl-kernel/tests/test_kvcacheio.py @@ -3,11 +3,13 @@ import torch from sgl_kernel.kvcacheio import ( transfer_kv_all_layer, transfer_kv_all_layer_direct_lf_pf, + transfer_kv_all_layer_lf_ph, transfer_kv_all_layer_mla, transfer_kv_direct, transfer_kv_per_layer, transfer_kv_per_layer_direct_pf_lf, transfer_kv_per_layer_mla, + transfer_kv_per_layer_ph_lf, ) @@ -30,6 +32,32 @@ def ref_copy_with_indices_pf_direct( ][layer_id].to(dst_pool.device) +def ref_copy_with_indices_page_head( + src_pool, + dst_pool, + src_indices, + dst_indices, + page_size, + layer_id, + head_num, + lf_to_ph=False, +): + if lf_to_ph: + for head_id in range(head_num): + for i in range(0, len(src_indices)): + dst_pool[dst_indices[i] // page_size][head_id][ + dst_indices[i] % page_size + ][layer_id] = src_pool[layer_id][src_indices[i]][head_id].to( + dst_pool.device + ) + else: + for head_id in range(head_num): + for i in range(0, len(src_indices)): + dst_pool[layer_id][dst_indices[i]][head_id] = src_pool[ + src_indices[i] // page_size + ][head_id][src_indices[i] % page_size][layer_id].to(dst_pool.device) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("num_items_to_transfer", [1, 128, 1024]) @pytest.mark.parametrize("page_size", [1, 16, 64]) @@ -481,5 +509,182 @@ def test_transfer_kv_pf_direct( torch.set_default_dtype(original_dtype) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("num_items_to_transfer", [256, 1024]) +@pytest.mark.parametrize("page_size", [16, 64, 128]) +@pytest.mark.parametrize("item_size", [1024]) +@pytest.mark.parametrize("head_num", [8, 16]) +@pytest.mark.parametrize("total_items_in_pool", [4096]) +@pytest.mark.parametrize("lf_to_ph", [False, True]) +def test_transfer_kv_page_head( + dtype: torch.dtype, + num_items_to_transfer: int, + page_size: int, + item_size: int, + head_num: int, + total_items_in_pool: int, + lf_to_ph: bool, +): + original_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + device = "cuda" + torch.cuda.manual_seed(42) + + num_layers = 4 + + total_pages_in_pool = total_items_in_pool // page_size + num_pages_to_transfer = num_items_to_transfer // page_size + if num_pages_to_transfer == 0: + torch.set_default_dtype(original_dtype) + return + + assert item_size % head_num == 0 + head_dim = item_size // head_num + + page_indices = torch.randperm(total_pages_in_pool, dtype=torch.int64) + src_indices_host = torch.cat( + [ + torch.arange(p * page_size, (p + 1) * page_size) + for p in page_indices[:num_pages_to_transfer] + ] + ) + src_indices_device = src_indices_host.to(device) + dst_indices_host = torch.cat( + [ + torch.arange(p * page_size, (p + 1) * page_size) + for p in page_indices[num_pages_to_transfer : 2 * num_pages_to_transfer] + ] + ) + dst_indices_device = dst_indices_host.to(device) + + # We will test the per-layer function on the first layer (index 0) of the pool. + layer_idx_to_test = 0 + + if lf_to_ph: + src_k_pool = torch.randn( + num_layers, total_items_in_pool, head_num, head_dim + ).to(device) + src_v_pool = torch.randn( + num_layers, total_items_in_pool, head_num, head_dim + ).to(device) + src_k_pool_ptrs = [src_k_pool[i] for i in range(num_layers)] + src_k_pool_ptrs = torch.tensor( + [x.data_ptr() for x in src_k_pool_ptrs], + dtype=torch.uint64, + device=device, + ) + src_v_pool_ptrs = [src_v_pool[i] for i in range(num_layers)] + src_v_pool_ptrs = torch.tensor( + [x.data_ptr() for x in src_v_pool_ptrs], + dtype=torch.uint64, + device=device, + ) + + dst_k_pool_ref = torch.zeros( + total_pages_in_pool, head_num, page_size, num_layers, head_dim + ).pin_memory() + dst_v_pool_ref = torch.zeros_like(dst_k_pool_ref).pin_memory() + + dst_k_pool_kernel = torch.zeros_like(dst_k_pool_ref).pin_memory() + dst_v_pool_kernel = torch.zeros_like(dst_v_pool_ref).pin_memory() + torch.cuda.synchronize() + + transfer_kv_all_layer_lf_ph( + src_k_pool_ptrs, + dst_k_pool_kernel, + src_v_pool_ptrs, + dst_v_pool_kernel, + src_indices_device, + dst_indices_device, + item_size * dtype.itemsize, + item_size * num_layers * dtype.itemsize, + num_layers, + page_size, + head_num, + ) + torch.cuda.synchronize() + + for i in range(num_layers): + ref_copy_with_indices_page_head( + src_k_pool, + dst_k_pool_ref, + src_indices_device, + dst_indices_host, + page_size, + i, + head_num, + lf_to_ph=True, + ) + ref_copy_with_indices_page_head( + src_v_pool, + dst_v_pool_ref, + src_indices_device, + dst_indices_host, + page_size, + i, + head_num, + lf_to_ph=True, + ) + torch.cuda.synchronize() + torch.testing.assert_close(dst_k_pool_kernel, dst_k_pool_ref) + torch.testing.assert_close(dst_v_pool_kernel, dst_v_pool_ref) + else: + src_k_pool = torch.randn( + total_pages_in_pool, head_num, page_size, num_layers, head_dim + ).pin_memory() + src_v_pool = torch.randn( + total_pages_in_pool, head_num, page_size, num_layers, head_dim + ).pin_memory() + + dst_k_pool_ref = torch.zeros( + num_layers, total_items_in_pool, head_num, head_dim + ).to(device) + dst_v_pool_ref = torch.zeros_like(dst_k_pool_ref) + dst_k_pool_kernel = torch.zeros_like(dst_k_pool_ref) + dst_v_pool_kernel = torch.zeros_like(dst_v_pool_ref) + dst_k_pool_kernel_ptrs = [dst_k_pool_kernel[i] for i in range(num_layers)] + dst_v_pool_kernel_ptrs = [dst_v_pool_kernel[i] for i in range(num_layers)] + torch.cuda.synchronize() + + transfer_kv_per_layer_ph_lf( + src_k_pool, + dst_k_pool_kernel_ptrs[layer_idx_to_test], + src_v_pool, + dst_v_pool_kernel_ptrs[layer_idx_to_test], + src_indices_device, + dst_indices_device, + layer_idx_to_test, + item_size * dtype.itemsize, + item_size * num_layers * dtype.itemsize, + page_size, + head_num, + ) + + ref_copy_with_indices_page_head( + src_k_pool, + dst_k_pool_ref, + src_indices_host, + dst_indices_device, + page_size, + layer_idx_to_test, + head_num, + lf_to_ph=False, + ) + ref_copy_with_indices_page_head( + src_v_pool, + dst_v_pool_ref, + src_indices_host, + dst_indices_device, + page_size, + layer_idx_to_test, + head_num, + lf_to_ph=False, + ) + torch.cuda.synchronize() + torch.testing.assert_close(dst_k_pool_kernel, dst_k_pool_ref) + torch.testing.assert_close(dst_v_pool_kernel, dst_v_pool_ref) + torch.set_default_dtype(original_dtype) + + if __name__ == "__main__": pytest.main([__file__])