From bb7ff2978c2c11895b75d6a5dc9ef4ba0f8d60b5 Mon Sep 17 00:00:00 2001 From: ltqin Date: Sat, 13 Sep 2025 06:43:35 +0000 Subject: [PATCH] add kv_last_page_lens to kernel --- example/ck_tile/01_fmha/fmha_fwd.hpp | 6 +-- .../fmha/kernel/fmha_batch_prefill_kernel.hpp | 39 +++++++------------ 2 files changed, 17 insertions(+), 28 deletions(-) diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index 698b94b74f..fa11e497f3 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -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, 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 c46c0ba3cb..590346877c 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 @@ -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; using FmhaMask = ck_tile::remove_cvref_t; 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(kv_indptr), reinterpret_cast(kv_page_indices), -#if 0 // we assume page_block_size=1 for now reinterpret_cast(kv_last_page_lens), -#endif #if CK_TILE_FMHA_FWD_FAST_EXP2 static_cast(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(kv_indptr), reinterpret_cast(kv_page_indices), -#if 0 // we assume page_block_size=1 for now reinterpret_cast(kv_last_page_lens), -#endif #if CK_TILE_FMHA_FWD_FAST_EXP2 static_cast(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(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