Skip to content

feat: paged caching support strided cache #1505

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PanZezhong1725 merged 1 commit into from
Aug 18, 2026
Merged

feat: paged caching support strided cache #1505

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
29 changes: 23 additions & 6 deletions src/infiniop/ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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;
Expand Down Expand Up @@ -57,9 +63,9 @@ class PagedCachingKernel {
const int64_t block_offset = slot_idx % static_cast<int64_t>(_block_size);

const ptrdiff_t k_src_base = static_cast<ptrdiff_t>(token_idx) * _k_src_stride
+ static_cast<ptrdiff_t>(head_idx * _head_size);
+ static_cast<ptrdiff_t>(head_idx) * _k_src_head_stride;
const ptrdiff_t v_src_base = static_cast<ptrdiff_t>(token_idx) * _v_src_stride
+ static_cast<ptrdiff_t>(head_idx * _head_size);
+ static_cast<ptrdiff_t>(head_idx) * _v_src_head_stride;
const ptrdiff_t k_dst_base = static_cast<ptrdiff_t>(physical_block_idx) * _k_cache_block_stride
+ static_cast<ptrdiff_t>(head_idx) * _k_cache_head_stride
+ static_cast<ptrdiff_t>(block_offset) * _k_cache_slot_stride;
Expand All @@ -69,6 +75,8 @@ class PagedCachingKernel {

for (size_t d = 0; d < _head_size; ++d) {
_k_cache_gm.SetValue(k_dst_base + static_cast<ptrdiff_t>(d), _k_gm.GetValue(k_src_base + static_cast<ptrdiff_t>(d)));
}
for (size_t d = 0; d < _v_head_size; ++d) {
_v_cache_gm.SetValue(v_dst_base + static_cast<ptrdiff_t>(d), _v_gm.GetValue(v_src_base + static_cast<ptrdiff_t>(d)));
}
}
Expand All @@ -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;
Expand All @@ -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<TYPE> 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); \
Expand All @@ -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,
Expand All @@ -147,7 +163,8 @@ extern "C" infiniStatus_t paged_caching_kernel_launch(
KERNEL_NAME<<<block_dim, nullptr, stream>>>( \
k_cache, v_cache, const_cast<void *>(k), const_cast<void *>(v), \
const_cast<void *>(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); \
Expand Down
14 changes: 10 additions & 4 deletions src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -41,8 +44,8 @@ __mlu_global__ void pagedCachingKernel(
const int64_t physical_block = slot_idx / static_cast<int64_t>(block_size);
const int64_t block_offset = slot_idx - physical_block * static_cast<int64_t>(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;
Expand All @@ -53,9 +56,9 @@ __mlu_global__ void pagedCachingKernel(
char *nram_base = reinterpret_cast<char *>(((reinterpret_cast<size_t>(paged_caching_nram_buffer) + ALIGN_SIZE - 1) / ALIGN_SIZE) * ALIGN_SIZE);
Tdata *tmp = reinterpret_cast<Tdata *>(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);
}
}

Expand Down Expand Up @@ -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,
Expand Down
6 changes: 4 additions & 2 deletions src/infiniop/ops/paged_caching/cuda/kernel.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
10 changes: 9 additions & 1 deletion src/infiniop/ops/paged_caching/info.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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,
Expand Down
13 changes: 12 additions & 1 deletion src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Tdata, NUM_THREADS>(
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);
}

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
13 changes: 12 additions & 1 deletion src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Tdata, NUM_THREADS>(
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);
}

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
Loading
Loading