ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 11:07:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

add page_block_size parameter

Browse Source
This commit is contained in:
ltqin
2025-09-08 10:38:07 +00:00
parent 4cfa2c7158
commit 62d28eb637
2 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
example/ck_tile/01_fmha/fmha_fwd.hpp
View File

@@ -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,

18
include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp
View File

@@ -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<const int32_t*>(kv_indptr),
reinterpret_cast<const int32_t*>(kv_page_indices),
#if 0 // we assume page_block_size=1 for now
reinterpret_cast<const int32_t*>(kv_last_page_lens),
page_block_size,
#endif
#if CK_TILE_FMHA_FWD_FAST_EXP2
static_cast<float>(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<const int32_t*>(kv_indptr),
reinterpret_cast<const int32_t*>(kv_page_indices),
#if 0 // we assume page_block_size=1 for now
reinterpret_cast<const int32_t*>(kv_last_page_lens),
page_block_size,
#endif
#if CK_TILE_FMHA_FWD_FAST_EXP2
static_cast<float>(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
}