From 62d28eb637dd8096be02a230054a74710b854ea9 Mon Sep 17 00:00:00 2001 From: ltqin Date: Mon, 8 Sep 2025 10:38:07 +0000 Subject: [PATCH] add page_block_size parameter --- example/ck_tile/01_fmha/fmha_fwd.hpp | 6 +++--- .../fmha/kernel/fmha_batch_prefill_kernel.hpp | 18 +++++++++--------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index df1e9e5699..bae132c7f5 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -428,11 +428,11 @@ struct fmha_batch_prefill_args // SGLang-style page table int32_t num_total_pages; + ck_tile::index_t page_block_size; void* kv_indptr; void* kv_page_indices; #if 0 // we assume page_block_size=1 for now void* kv_last_page_lens; - ck_tile::index_t page_block_size; #endif float scale_s; @@ -917,11 +917,11 @@ auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args) args.nhead_q, args.nhead_q / args.nhead_k, args.num_total_pages, + args.page_block_size, args.kv_indptr, args.kv_page_indices, #if 0 // we assume page_block_size=1 for now args.kv_last_page_lens, - args.page_block_size, #endif args.scale_s, args.scale_p, @@ -964,11 +964,11 @@ auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args) args.nhead_q, args.nhead_q / args.nhead_k, args.num_total_pages, + args.page_block_size, args.kv_indptr, args.kv_page_indices, #if 0 // we assume page_block_size=1 for now args.kv_last_page_lens, - args.page_block_size, #endif args.scale_s, args.scale_p, diff --git a/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp index 2850ce3379..c46c0ba3cb 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp @@ -132,13 +132,11 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel ck_tile::index_t nhead_ratio_qk; int32_t num_total_pages; + ck_tile::index_t page_block_size; const int32_t* kv_indptr; const int32_t* kv_page_indices; #if 0 // we assume page_block_size=1 for now const int32_t* kv_last_page_lens; - ck_tile::index_t page_block_size; -#else - static constexpr ck_tile::index_t page_block_size = 1; #endif float scale_s; @@ -330,11 +328,11 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel ck_tile::index_t num_head_q, ck_tile::index_t nhead_ratio_qk, int32_t num_total_pages, + ck_tile::index_t page_block_size, const void* kv_indptr, const void* kv_page_indices, #if 0 // we assume page_block_size=1 for now const void* kv_last_page_lens, - ck_tile::index_t page_block_size, #endif float scale_s, float scale_p, @@ -379,11 +377,11 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel num_head_q, nhead_ratio_qk, num_total_pages, + page_block_size, reinterpret_cast(kv_indptr), reinterpret_cast(kv_page_indices), #if 0 // we assume page_block_size=1 for now reinterpret_cast(kv_last_page_lens), - page_block_size, #endif #if CK_TILE_FMHA_FWD_FAST_EXP2 static_cast(scale_s * ck_tile::log2e_v<>), @@ -482,11 +480,11 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel ck_tile::index_t num_head_q, ck_tile::index_t nhead_ratio_qk, int32_t num_total_pages, + ck_tile::index_t page_block_size, const void* kv_indptr, const void* kv_page_indices, #if 0 // we assume page_block_size=1 for now const void* kv_last_page_lens, - ck_tile::index_t page_block_size, #endif float scale_s, float scale_p, @@ -526,11 +524,11 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel num_head_q, nhead_ratio_qk, num_total_pages, + page_block_size, reinterpret_cast(kv_indptr), reinterpret_cast(kv_page_indices), #if 0 // we assume page_block_size=1 for now reinterpret_cast(kv_last_page_lens), - page_block_size, #endif #if CK_TILE_FMHA_FWD_FAST_EXP2 static_cast(scale_s * ck_tile::log2e_v<>), @@ -756,7 +754,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel #if 0 // we assume page_block_size=1 for now kargs.seqlen_k = (num_page_blocks - 1) * kargs.page_block_size + last_page_len; #else - kargs.seqlen_k = num_page_blocks; + //Fix: + kargs.seqlen_k = num_page_blocks * kargs.page_block_size; #endif } else @@ -783,7 +782,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel #if 0 // we assume page_block_size=1 for now kargs.seqlen_k = (num_page_blocks - 1) * kargs.page_block_size + last_page_len; #else - kargs.seqlen_k = num_page_blocks; + //Fix: + kargs.seqlen_k = num_page_blocks * kargs.page_block_size; #endif }