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

add kv_last_page_lens to kernel

Browse Source
This commit is contained in:
ltqin
2025-09-13 06:43:35 +00:00
parent 667ed5825c
commit bb7ff2978c
2 changed files with 17 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -431,7 +431,7 @@ struct fmha_batch_prefill_args
ck_tile::index_t page_block_size;
void* kv_indptr;
void* kv_page_indices;
#if 0 // we assume page_block_size=1 for now
#if 1 // we assume page_block_size=1 for now
void* kv_last_page_lens;
#endif
@@ -920,7 +920,7 @@ auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args)
args.page_block_size,
args.kv_indptr,
args.kv_page_indices,
#if 0 // we assume page_block_size=1 for now
#if 1 // we assume page_block_size=1 for now
args.kv_last_page_lens,
#endif
args.scale_s,
@@ -967,7 +967,7 @@ auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args)
args.page_block_size,
args.kv_indptr,
args.kv_page_indices,
#if 0 // we assume page_block_size=1 for now
#if 1 // we assume page_block_size=1 for now
args.kv_last_page_lens,
#endif
args.scale_s,

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

@@ -54,6 +54,7 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
static constexpr bool kStoreLSE = FmhaPipeline::kStoreLSE;
static constexpr bool kHasDropout = FmhaPipeline::kHasDropout;
static constexpr bool kDoFp8StaticQuant = FmhaPipeline::Problem::kDoFp8StaticQuant;
static constexpr bool kIsSglangLayout = FmhaPipeline::kIsSglangLayout;
using AttentionVariant = ck_tile::remove_cvref_t<typename FmhaPipeline::AttentionVariant>;
using FmhaMask = ck_tile::remove_cvref_t<typename FmhaPipeline::FmhaMask>;
static constexpr bool kHasMask = FmhaMask::IsMasking;
@@ -135,9 +136,7 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
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;
#endif
float scale_s;
@@ -331,9 +330,7 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
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,
#endif
float scale_s,
float scale_p,
float scale_o,
@@ -380,9 +377,7 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
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),
#endif
#if CK_TILE_FMHA_FWD_FAST_EXP2
static_cast<float>(scale_s * ck_tile::log2e_v<>),
#else
@@ -483,9 +478,7 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
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,
#endif
float scale_s,
float scale_p,
float scale_o,
@@ -527,9 +520,7 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
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),
#endif
#if CK_TILE_FMHA_FWD_FAST_EXP2
static_cast<float>(scale_s * ck_tile::log2e_v<>),
#else
@@ -715,9 +706,11 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
long_index_t batch_offset_o = 0;
const int32_t num_page_blocks = kargs.kv_indptr[i_batch + 1] - kargs.kv_indptr[i_batch];
#if 0 // we assume page_block_size=1 for now
const int32_t last_page_len = kargs.kv_last_page_lens[i_batch];
#endif
int32_t last_page_len = 0;
if constexpr(!kIsSglangLayout)
last_page_len = kargs.kv_last_page_lens[i_batch];
if constexpr(kIsGroupMode)
{
// get starting offset for each batch
@@ -751,12 +744,10 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
return;
}
#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
//Fix:
kargs.seqlen_k = num_page_blocks * kargs.page_block_size;
#endif
if constexpr(kIsSglangLayout)
kargs.seqlen_k = num_page_blocks;
else
kargs.seqlen_k = (num_page_blocks - 1) * kargs.page_block_size + last_page_len;
}
else
{
@@ -779,12 +770,10 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
}
batch_offset_o = static_cast<long_index_t>(i_batch) * kargs.batch_stride_o;
#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
//Fix:
kargs.seqlen_k = num_page_blocks * kargs.page_block_size;
#endif
if constexpr(kIsSglangLayout)
kargs.seqlen_k = num_page_blocks;
else
kargs.seqlen_k = (num_page_blocks - 1) * kargs.page_block_size + last_page_len;
}
// for simplicity, batch stride we just modify the pointer