From f09473f0fab1853ca99f09ba9d128926aa3d8afe Mon Sep 17 00:00:00 2001 From: PanZezhong Date: Tue, 18 Aug 2026 00:36:05 +0000 Subject: [PATCH] feat: paged caching support strided cache --- .../ascend/paged_caching_ascend.cc | 3 + .../ascend/paged_caching_ascend.h | 3 + .../ascend/paged_caching_ascend_kernel.cpp | 29 +++++++-- .../paged_caching/bang/paged_caching_bang.mlu | 14 +++-- .../ops/paged_caching/cuda/kernel.cuh | 6 +- src/infiniop/ops/paged_caching/info.h | 10 ++- .../metax/paged_caching_metax.maca | 13 +++- .../moore/paged_caching_moore.mu | 13 +++- .../nvidia/paged_caching_nvidia.cu | 12 ++++ test/infinicore/ops/paged_caching.py | 63 ++++++++++++++++--- 10 files changed, 141 insertions(+), 25 deletions(-) diff --git a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc index d505ecc55..ff4e59faa 100644 --- a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc +++ b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc @@ -47,9 +47,12 @@ infiniStatus_t Descriptor::calculate( _info.num_tokens, _info.num_kv_heads, _info.head_size, + _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, + _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, diff --git a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.h b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.h index 9edb1250f..62d0d0690 100644 --- a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.h +++ b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.h @@ -13,9 +13,12 @@ extern "C" infiniStatus_t paged_caching_kernel_launch( size_t num_tokens, size_t num_kv_heads, size_t head_size, + size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, ptrdiff_t v_src_stride, + ptrdiff_t k_src_head_stride, + ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, diff --git a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp index c200201ac..b655d8994 100644 --- a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp +++ b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp @@ -15,9 +15,12 @@ class PagedCachingKernel { GM_ADDR slot_mapping, size_t num_kv_heads, size_t head_size, + size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, ptrdiff_t v_src_stride, + ptrdiff_t k_src_head_stride, + ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, @@ -26,9 +29,12 @@ class PagedCachingKernel { ptrdiff_t v_cache_slot_stride) { _num_kv_heads = num_kv_heads; _head_size = head_size; + _v_head_size = v_head_size; _block_size = block_size; _k_src_stride = k_src_stride; _v_src_stride = v_src_stride; + _k_src_head_stride = k_src_head_stride; + _v_src_head_stride = v_src_head_stride; _k_cache_block_stride = k_cache_block_stride; _v_cache_block_stride = v_cache_block_stride; _k_cache_head_stride = k_cache_head_stride; @@ -57,9 +63,9 @@ class PagedCachingKernel { const int64_t block_offset = slot_idx % static_cast(_block_size); const ptrdiff_t k_src_base = static_cast(token_idx) * _k_src_stride - + static_cast(head_idx * _head_size); + + static_cast(head_idx) * _k_src_head_stride; const ptrdiff_t v_src_base = static_cast(token_idx) * _v_src_stride - + static_cast(head_idx * _head_size); + + static_cast(head_idx) * _v_src_head_stride; const ptrdiff_t k_dst_base = static_cast(physical_block_idx) * _k_cache_block_stride + static_cast(head_idx) * _k_cache_head_stride + static_cast(block_offset) * _k_cache_slot_stride; @@ -69,6 +75,8 @@ class PagedCachingKernel { for (size_t d = 0; d < _head_size; ++d) { _k_cache_gm.SetValue(k_dst_base + static_cast(d), _k_gm.GetValue(k_src_base + static_cast(d))); + } + for (size_t d = 0; d < _v_head_size; ++d) { _v_cache_gm.SetValue(v_dst_base + static_cast(d), _v_gm.GetValue(v_src_base + static_cast(d))); } } @@ -82,9 +90,12 @@ class PagedCachingKernel { size_t _num_kv_heads; size_t _head_size; + size_t _v_head_size; size_t _block_size; ptrdiff_t _k_src_stride; ptrdiff_t _v_src_stride; + ptrdiff_t _k_src_head_stride; + ptrdiff_t _v_src_head_stride; ptrdiff_t _k_cache_block_stride; ptrdiff_t _v_cache_block_stride; ptrdiff_t _k_cache_head_stride; @@ -97,14 +108,16 @@ class PagedCachingKernel { extern "C" __global__ __aicore__ void KERNEL_NAME( \ GM_ADDR k_cache, GM_ADDR v_cache, GM_ADDR k, GM_ADDR v, \ GM_ADDR slot_mapping, size_t num_kv_heads, size_t head_size, \ - size_t block_size, ptrdiff_t k_src_stride, \ - ptrdiff_t v_src_stride, ptrdiff_t k_cache_block_stride, \ + size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, \ + ptrdiff_t v_src_stride, ptrdiff_t k_src_head_stride, \ + ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, \ ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, \ ptrdiff_t v_cache_head_stride, ptrdiff_t k_cache_slot_stride, \ ptrdiff_t v_cache_slot_stride) { \ PagedCachingKernel op; \ op.init(k_cache, v_cache, k, v, slot_mapping, num_kv_heads, \ - head_size, block_size, k_src_stride, v_src_stride, \ + head_size, v_head_size, block_size, k_src_stride, \ + v_src_stride, k_src_head_stride, v_src_head_stride, \ k_cache_block_stride, v_cache_block_stride, \ k_cache_head_stride, v_cache_head_stride, \ k_cache_slot_stride, v_cache_slot_stride); \ @@ -127,9 +140,12 @@ extern "C" infiniStatus_t paged_caching_kernel_launch( size_t num_tokens, size_t num_kv_heads, size_t head_size, + size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, ptrdiff_t v_src_stride, + ptrdiff_t k_src_head_stride, + ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, @@ -147,7 +163,8 @@ extern "C" infiniStatus_t paged_caching_kernel_launch( KERNEL_NAME<<>>( \ k_cache, v_cache, const_cast(k), const_cast(v), \ const_cast(slot_mapping), num_kv_heads, head_size, \ - block_size, k_src_stride, v_src_stride, \ + v_head_size, block_size, k_src_stride, v_src_stride, \ + k_src_head_stride, v_src_head_stride, \ k_cache_block_stride, v_cache_block_stride, \ k_cache_head_stride, v_cache_head_stride, \ k_cache_slot_stride, v_cache_slot_stride); \ diff --git a/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu b/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu index bbe57a216..92fc38dd8 100644 --- a/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu +++ b/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu @@ -17,9 +17,12 @@ __mlu_global__ void pagedCachingKernel( size_t num_tokens, size_t num_kv_heads, size_t head_size, + size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, ptrdiff_t v_src_stride, + ptrdiff_t k_src_head_stride, + ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, @@ -41,8 +44,8 @@ __mlu_global__ void pagedCachingKernel( const int64_t physical_block = slot_idx / static_cast(block_size); const int64_t block_offset = slot_idx - physical_block * static_cast(block_size); - const Tdata *k_src = k + token_idx * k_src_stride + head_idx * head_size; - const Tdata *v_src = v + token_idx * v_src_stride + head_idx * head_size; + const Tdata *k_src = k + token_idx * k_src_stride + head_idx * k_src_head_stride; + const Tdata *v_src = v + token_idx * v_src_stride + head_idx * v_src_head_stride; Tdata *k_dst = k_cache + physical_block * k_cache_block_stride + head_idx * k_cache_head_stride + block_offset * k_cache_slot_stride; @@ -53,9 +56,9 @@ __mlu_global__ void pagedCachingKernel( char *nram_base = reinterpret_cast(((reinterpret_cast(paged_caching_nram_buffer) + ALIGN_SIZE - 1) / ALIGN_SIZE) * ALIGN_SIZE); Tdata *tmp = reinterpret_cast(nram_base); __memcpy(tmp, k_src, head_size * sizeof(Tdata), GDRAM2NRAM); - __memcpy(tmp + head_size, v_src, head_size * sizeof(Tdata), GDRAM2NRAM); + __memcpy(tmp + head_size, v_src, v_head_size * sizeof(Tdata), GDRAM2NRAM); __memcpy(k_dst, tmp, head_size * sizeof(Tdata), NRAM2GDRAM); - __memcpy(v_dst, tmp + head_size, head_size * sizeof(Tdata), NRAM2GDRAM); + __memcpy(v_dst, tmp + head_size, v_head_size * sizeof(Tdata), NRAM2GDRAM); } } @@ -85,9 +88,12 @@ infiniStatus_t launchPagedCaching( info.num_tokens, info.num_kv_heads, info.head_size, + info.v_head_size, info.block_size, info.k_src_stride, info.v_src_stride, + info.k_src_head_stride, + info.v_src_head_stride, info.k_cache_block_stride, info.v_cache_block_stride, info.k_cache_head_stride, diff --git a/src/infiniop/ops/paged_caching/cuda/kernel.cuh b/src/infiniop/ops/paged_caching/cuda/kernel.cuh index a1b2106ce..c2ac7d096 100644 --- a/src/infiniop/ops/paged_caching/cuda/kernel.cuh +++ b/src/infiniop/ops/paged_caching/cuda/kernel.cuh @@ -38,6 +38,8 @@ __device__ void pagedCachingKernel( // ----- Stride Information ----- const ptrdiff_t k_src_stride, // Stride between tokens in the source K tensor const ptrdiff_t v_src_stride, // Stride between tokens in the source V tensor + const ptrdiff_t k_src_head_stride, // Stride between heads in the source K tensor + const ptrdiff_t v_src_head_stride, // Stride between heads in the source V tensor const ptrdiff_t k_cache_block_stride, // Stride between blocks in the K cache pool const ptrdiff_t v_cache_block_stride, // Stride between blocks in the V cache pool const ptrdiff_t k_cache_head_stride, // Stride between heads in the K cache pool @@ -66,8 +68,8 @@ __device__ void pagedCachingKernel( const int64_t block_offset = slot_idx % block_size; // Calculate base pointers for source and destination for this specific token. - const Tdata *k_src_head_ptr = k_ptr + token_idx * k_src_stride + head_idx * head_size; - const Tdata *v_src_head_ptr = v_ptr + token_idx * v_src_stride + head_idx * v_head_size; + const Tdata *k_src_head_ptr = k_ptr + token_idx * k_src_stride + head_idx * k_src_head_stride; + const Tdata *v_src_head_ptr = v_ptr + token_idx * v_src_stride + head_idx * v_src_head_stride; // Destination pointer calculation assumes a [num_blocks, block_size, num_heads, head_size] layout. // We point to the beginning of the memory region for this token's slot. diff --git a/src/infiniop/ops/paged_caching/info.h b/src/infiniop/ops/paged_caching/info.h index e5cf1e4fa..8a6808390 100644 --- a/src/infiniop/ops/paged_caching/info.h +++ b/src/infiniop/ops/paged_caching/info.h @@ -25,6 +25,8 @@ class PagedCachingInfo { // --- Strides for Memory Layout --- ptrdiff_t k_src_stride; ptrdiff_t v_src_stride; + ptrdiff_t k_src_head_stride; + ptrdiff_t v_src_head_stride; ptrdiff_t k_cache_block_stride; ptrdiff_t v_cache_block_stride; ptrdiff_t k_cache_head_stride; @@ -75,13 +77,17 @@ class PagedCachingInfo { if (v_cache_shape[0] != k_cache_shape[0] || v_cache_shape[1] != num_kv_heads || v_cache_shape[2] != block_size) { return INFINI_STATUS_BAD_TENSOR_SHAPE; } - if (k_cache_shape[1] != num_kv_heads || k_cache_shape[3] != head_size || v_cache_shape[3] != v_head_size) { + // The V cache may be wider than the incoming V. This is used by models + // that pad V to the Q/K head width for the attention backend. + if (k_cache_shape[1] != num_kv_heads || k_cache_shape[3] != head_size || v_cache_shape[3] < v_head_size) { return INFINI_STATUS_BAD_TENSOR_SHAPE; } // --- Extract strides for memory access --- ptrdiff_t k_src_stride = k_desc->stride(0); ptrdiff_t v_src_stride = v_desc->stride(0); + ptrdiff_t k_src_head_stride = k_desc->stride(1); + ptrdiff_t v_src_head_stride = v_desc->stride(1); ptrdiff_t k_cache_block_stride = k_cache_desc->stride(0); ptrdiff_t v_cache_block_stride = v_cache_desc->stride(0); ptrdiff_t k_cache_head_stride = k_cache_desc->stride(1); @@ -98,6 +104,8 @@ class PagedCachingInfo { block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, diff --git a/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca b/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca index 82e3ff210..a14f5073a 100644 --- a/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca +++ b/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca @@ -10,12 +10,14 @@ INFINIOP_METAX_KERNEL pagedCaching( const int64_t *slot_mapping, const size_t head_size, const size_t v_head_size, const size_t block_size, const ptrdiff_t k_src_stride, const ptrdiff_t v_src_stride, + const ptrdiff_t k_src_head_stride, const ptrdiff_t v_src_head_stride, const ptrdiff_t k_cache_block_stride, const ptrdiff_t v_cache_block_stride, const ptrdiff_t k_cache_head_stride, const ptrdiff_t v_cache_head_stride, const ptrdiff_t k_cache_slot_stride, const ptrdiff_t v_cache_slot_stride) { op::paged_caching::cuda::pagedCachingKernel( k_cache, v_cache, k, v, slot_mapping, head_size, v_head_size, - block_size, k_src_stride, v_src_stride, + block_size, k_src_stride, v_src_stride, + k_src_head_stride, v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, v_cache_head_stride, k_cache_slot_stride, v_cache_slot_stride); } @@ -61,6 +63,7 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, const void *slot_mapping, size_t num_tokens, size_t num_kv_heads, size_t head_size, size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, ptrdiff_t v_src_stride, + ptrdiff_t k_src_head_stride, ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, ptrdiff_t v_cache_head_stride, ptrdiff_t k_cache_slot_stride, ptrdiff_t v_cache_slot_stride, @@ -88,6 +91,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -107,6 +112,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -126,6 +133,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -157,6 +166,7 @@ infiniStatus_t Descriptor::calculate( _info, k_cache, v_cache, _info.dtype, k, v, slot_mapping, _info.num_tokens, _info.num_kv_heads, _info.head_size, _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, _info.v_cache_head_stride, _info.k_cache_slot_stride, _info.v_cache_slot_stride, @@ -166,6 +176,7 @@ infiniStatus_t Descriptor::calculate( _info, k_cache, v_cache, _info.dtype, k, v, slot_mapping, _info.num_tokens, _info.num_kv_heads, _info.head_size, _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, _info.v_cache_head_stride, _info.k_cache_slot_stride, _info.v_cache_slot_stride, diff --git a/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu b/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu index bb6be3489..579e7cf05 100644 --- a/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu +++ b/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu @@ -10,12 +10,14 @@ INFINIOP_MOORE_KERNEL pagedCaching( const int64_t *slot_mapping, const size_t head_size, const size_t v_head_size, const size_t block_size, const ptrdiff_t k_src_stride, const ptrdiff_t v_src_stride, + const ptrdiff_t k_src_head_stride, const ptrdiff_t v_src_head_stride, const ptrdiff_t k_cache_block_stride, const ptrdiff_t v_cache_block_stride, const ptrdiff_t k_cache_head_stride, const ptrdiff_t v_cache_head_stride, const ptrdiff_t k_cache_slot_stride, const ptrdiff_t v_cache_slot_stride) { op::paged_caching::cuda::pagedCachingKernel( k_cache, v_cache, k, v, slot_mapping, head_size, v_head_size, - block_size, k_src_stride, v_src_stride, + block_size, k_src_stride, v_src_stride, + k_src_head_stride, v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, v_cache_head_stride, k_cache_slot_stride, v_cache_slot_stride); } @@ -61,6 +63,7 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, const void *slot_mapping, size_t num_tokens, size_t num_kv_heads, size_t head_size, size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, ptrdiff_t v_src_stride, + ptrdiff_t k_src_head_stride, ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, ptrdiff_t v_cache_head_stride, ptrdiff_t k_cache_slot_stride, ptrdiff_t v_cache_slot_stride, @@ -88,6 +91,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -107,6 +112,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -126,6 +133,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -156,6 +165,7 @@ infiniStatus_t Descriptor::calculate( _info, k_cache, v_cache, _info.dtype, k, v, slot_mapping, _info.num_tokens, _info.num_kv_heads, _info.head_size, _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, _info.v_cache_head_stride, _info.k_cache_slot_stride, _info.v_cache_slot_stride, @@ -165,6 +175,7 @@ infiniStatus_t Descriptor::calculate( _info, k_cache, v_cache, _info.dtype, k, v, slot_mapping, _info.num_tokens, _info.num_kv_heads, _info.head_size, _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, _info.v_cache_head_stride, _info.k_cache_slot_stride, _info.v_cache_slot_stride, diff --git a/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu b/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu index a0924d605..202c02b23 100644 --- a/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu +++ b/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu @@ -10,12 +10,14 @@ INFINIOP_CUDA_KERNEL pagedCaching( const int64_t *slot_mapping, const size_t head_size, const size_t v_head_size, const size_t block_size, const ptrdiff_t k_src_stride, const ptrdiff_t v_src_stride, + const ptrdiff_t k_src_head_stride, const ptrdiff_t v_src_head_stride, const ptrdiff_t k_cache_block_stride, const ptrdiff_t v_cache_block_stride, const ptrdiff_t k_cache_head_stride, const ptrdiff_t v_cache_head_stride, const ptrdiff_t k_cache_slot_stride, const ptrdiff_t v_cache_slot_stride) { op::paged_caching::cuda::pagedCachingKernel( k_cache, v_cache, k, v, slot_mapping, head_size, v_head_size, block_size, k_src_stride, v_src_stride, + k_src_head_stride, v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, v_cache_head_stride, k_cache_slot_stride, v_cache_slot_stride); } @@ -61,6 +63,7 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, const void *slot_mapping, size_t num_tokens, size_t num_kv_heads, size_t head_size, size_t v_head_size, size_t block_size, ptrdiff_t k_src_stride, ptrdiff_t v_src_stride, + ptrdiff_t k_src_head_stride, ptrdiff_t v_src_head_stride, ptrdiff_t k_cache_block_stride, ptrdiff_t v_cache_block_stride, ptrdiff_t k_cache_head_stride, ptrdiff_t v_cache_head_stride, ptrdiff_t k_cache_slot_stride, ptrdiff_t v_cache_slot_stride, @@ -88,6 +91,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -107,6 +112,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -126,6 +133,8 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, block_size, k_src_stride, v_src_stride, + k_src_head_stride, + v_src_head_stride, k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, @@ -156,6 +165,7 @@ infiniStatus_t Descriptor::calculate( _info, k_cache, v_cache, _info.dtype, k, v, slot_mapping, _info.num_tokens, _info.num_kv_heads, _info.head_size, _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, _info.v_cache_head_stride, _info.k_cache_slot_stride, _info.v_cache_slot_stride, @@ -165,6 +175,7 @@ infiniStatus_t Descriptor::calculate( _info, k_cache, v_cache, _info.dtype, k, v, slot_mapping, _info.num_tokens, _info.num_kv_heads, _info.head_size, _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, _info.v_cache_head_stride, _info.k_cache_slot_stride, _info.v_cache_slot_stride, @@ -174,6 +185,7 @@ infiniStatus_t Descriptor::calculate( _info, k_cache, v_cache, _info.dtype, k, v, slot_mapping, _info.num_tokens, _info.num_kv_heads, _info.head_size, _info.v_head_size, _info.block_size, _info.k_src_stride, _info.v_src_stride, + _info.k_src_head_stride, _info.v_src_head_stride, _info.k_cache_block_stride, _info.v_cache_block_stride, _info.k_cache_head_stride, _info.v_cache_head_stride, _info.k_cache_slot_stride, _info.v_cache_slot_stride, diff --git a/test/infinicore/ops/paged_caching.py b/test/infinicore/ops/paged_caching.py index 9c6b07303..fc547a739 100644 --- a/test/infinicore/ops/paged_caching.py +++ b/test/infinicore/ops/paged_caching.py @@ -18,7 +18,9 @@ # Operator-specific configuration # ============================================================================== -# Test cases format: (num_seqs, max_seq_len, num_kv_heads, head_size, block_size, permute_dim_1_2[, value_size]) +# Test cases format: (num_seqs, max_seq_len, num_kv_heads, head_size, +# block_size, permute_dim_1_2[, value_size[, value_cache_size +# [, value_head_stride]]]) _TEST_CASES_DATA = [ (1, 128, 8, 128, 16, False), (1, 128, 8, 128, 16, True), @@ -31,6 +33,10 @@ # New DeepSeek MLA wrapper case: verifies cache writes when key and # value head sizes differ. (1, 32, 1, 576, 16, False, 512), + # Kimi K3 MLA case: V is narrowed from a fused [K_nope, V] tensor, so + # consecutive V heads are separated by the fused projection width. Its + # 128 real values are copied into a cache padded to the 192-wide Q/K head. + (1, 32, 4, 192, 16, True, 128, 192, 256), ] # Tolerance configuration @@ -78,10 +84,10 @@ def ref_paged_caching( if permute_dim_1_2: k_cache_ref[block_idx, block_offset, :, :] = key_token - v_cache_ref[block_idx, block_offset, :, :] = value_token + v_cache_ref[block_idx, block_offset, :, : value.shape[-1]] = value_token else: k_cache_ref[block_idx, :, block_offset, :] = key_token - v_cache_ref[block_idx, :, block_offset, :] = value_token + v_cache_ref[block_idx, :, block_offset, : value.shape[-1]] = value_token return k_cache_ref, v_cache_ref @@ -103,6 +109,32 @@ def parse_test_cases(): permute_dim_1_2, ) = case value_size = head_size + value_cache_size = value_size + value_head_stride = value_size + elif len(case) == 7: + ( + num_seqs, + max_seq_len, + num_kv_heads, + head_size, + block_size, + permute_dim_1_2, + value_size, + ) = case + value_cache_size = value_size + value_head_stride = value_size + elif len(case) == 8: + ( + num_seqs, + max_seq_len, + num_kv_heads, + head_size, + block_size, + permute_dim_1_2, + value_size, + value_cache_size, + ) = case + value_head_stride = value_size else: ( num_seqs, @@ -112,6 +144,8 @@ def parse_test_cases(): block_size, permute_dim_1_2, value_size, + value_cache_size, + value_head_stride, ) = case num_blocks = 4096 # A reasonably large cache pool for testing @@ -131,9 +165,9 @@ def parse_test_cases(): current_slot += length.item() # Ensure we don't exceed the total number of slots in the cache - assert ( - current_slot <= num_blocks * block_size - ), "Not enough blocks in the cache pool for this test case" + assert current_slot <= num_blocks * block_size, ( + "Not enough blocks in the cache pool for this test case" + ) slot_mapping = torch.tensor(slot_mapping_list, dtype=torch.int64) @@ -143,10 +177,10 @@ def parse_test_cases(): k_shape = (ntok, num_kv_heads, head_size) v_shape = (ntok, num_kv_heads, value_size) k_cache_shape = (num_blocks, num_kv_heads, block_size, head_size) - v_cache_shape = (num_blocks, num_kv_heads, block_size, value_size) + v_cache_shape = (num_blocks, num_kv_heads, block_size, value_cache_size) if permute_dim_1_2: k_cache_shape = (num_blocks, block_size, num_kv_heads, head_size) - v_cache_shape = (num_blocks, block_size, num_kv_heads, value_size) + v_cache_shape = (num_blocks, block_size, num_kv_heads, value_cache_size) # Generate test cases for all data types for dtype in _TENSOR_DTYPES: @@ -154,7 +188,12 @@ def parse_test_cases(): # Create typed tensor specs k_spec = TensorSpec.from_tensor(k_shape, None, dtype) - v_spec = TensorSpec.from_tensor(v_shape, None, dtype) + v_strides = ( + num_kv_heads * value_head_stride, + value_head_stride, + 1, + ) + v_spec = TensorSpec.from_tensor(v_shape, v_strides, dtype) k_cache_spec = TensorSpec.from_tensor( k_cache_shape, None, dtype, init_mode=TensorInitializer.ZEROS ) @@ -168,7 +207,11 @@ def parse_test_cases(): dtype=infinicore.int64, ) - for comparison_target in [0, 1] if value_size != head_size else [0]: + for comparison_target in ( + [0, 1] + if value_size != head_size or value_cache_size != head_size + else [0] + ): test_cases.append( TestCase( inputs=[