diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py index 76b9429b2e..57fdf098be 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py @@ -62,6 +62,8 @@ using fmha_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad}, {F_lse}, {F_dropout}, {F_squant}, + {F_sglang_layout}, + {F_chunked}, {F_occupancy}>; using fmha_variant_{F_idx} = ck_tile::ComposedAttention<{F_logits} * ck_tile::LOGITS_SOFT_CAP, CK_TILE_FMHA_FWD_FAST_EXP2>; @@ -98,7 +100,7 @@ using fmha_kernel_{F_idx} = ck_tile::FmhaBatchPrefillWithPagedKVCacheKernel; using trait_{F_idx} = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode},{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, - {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; + {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_sglang_layout}, {F_chunked}>; #include @@ -133,9 +135,9 @@ FMHA_FWD_API_PER_HDIM_CASE=""" {F_if} (t.hdim_q <= {F_hdim} && t.hdim_v < }} """ -FMHA_FWD_API_INNER_DISPATCH=""" {F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.has_dropout == {F_dropout}) && (t.do_fp8_static_quant == {F_squant}) && +FMHA_FWD_API_INNER_DISPATCH=""" {F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.has_dropout == {F_dropout}) && (t.do_fp8_static_quant == {F_squant}) && (t.is_sglang_layout == {F_sglang_layout}) && (t.is_chunked_prefill == {F_chunked}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ - using trait_ = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; + using trait_ = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_sglang_layout}, {F_chunked}>; return fmha_batch_prefill_(s, a); }} """ @@ -164,11 +166,13 @@ class FmhaFwdApiTrait: skpad : str dpad : str dvpad : str + cache_layout : str + chunked : str @property def name(self) -> str: return f'{self.hdim}-{self.dtype}-{self.mode}-{self.bm0}-{self.bn0}-{self.bk0}-{self.bn0}-{self.bk1}-{self.bk0max}-'+\ - f'{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.dropout}-{self.squant}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}' + f'{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.dropout}-{self.squant}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}-{self.cache_layout}-{self.chunked}' @property def scheck(self) -> str: @@ -231,6 +235,8 @@ class FmhaFwdPipeline: F_dropout : str # F_squant : str # F_mask : str # value from MASK_MAP + F_sglang_layout : str # + F_chunked : str # @property def name(self) -> str: @@ -268,6 +274,13 @@ class FmhaFwdPipeline: if self.F_squant == 't' : n += '_squant' else: n += '_nsquant' + + # if self.F_sglang_layout == 't' : n += '_sglang' + # else: n += '_vllm' + + if self.F_chunked == 't' : n += '_chunked' + else: n += '_nchunked' + return n class FmhaFwdApiPool: @@ -300,6 +313,7 @@ class FmhaFwdApiPool: F_lse=BOOL_MAP[trait.lse], F_dropout=BOOL_MAP[trait.dropout] , F_squant=BOOL_MAP[trait.squant], F_scheck=trait.scheck, F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, F_spad=BOOL_MAP[trait.spad], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad], + F_sglang_layout=BOOL_MAP[trait.cache_layout], F_chunked=BOOL_MAP[trait.chunked], F_bm0=trait.bm0, F_bn0=trait.bn0, F_bk0=trait.bk0, F_bn1=trait.bn1, F_bk1=trait.bk1, F_bk0max=trait.bk0max, F_hdim=hdim, F_dtype=FWD_DTYPE_MAP[dtype]) if_j = 'if' if j == 0 else 'else if' @@ -385,6 +399,8 @@ class FmhaFwdKernel: F_lse = BOOL_MAP[self.F_pipeline.F_lse], F_dropout = BOOL_MAP[self.F_pipeline.F_dropout], F_squant = BOOL_MAP[self.F_pipeline.F_squant], + F_sglang_layout = BOOL_MAP[self.F_pipeline.F_sglang_layout], + F_chunked = BOOL_MAP[self.F_pipeline.F_chunked], F_occupancy = self.F_tile.F_occupancy, F_pipeline_enum = PIPELINE_ENUM_MAP[self.F_pipeline.tag], F_mask = get_mask_map(self.mask_impl)[self.F_pipeline.F_mask], @@ -423,7 +439,9 @@ class FmhaFwdKernel: spad=self.F_pipeline.F_spad, skpad=self.F_pipeline.F_skpad, dpad=self.F_pipeline.F_dpad, - dvpad=self.F_pipeline.F_dvpad) + dvpad=self.F_pipeline.F_dvpad, + cache_layout=self.F_pipeline.F_sglang_layout, + chunked=self.F_pipeline.F_chunked) # TODO: design a more practical way to do it # this is current supported tile size per hdim @@ -457,31 +475,31 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl squant = 't' if dtype == 'fp8' else 'f' pipelines = [] if dtype in ['fp16', 'bf16']: - for logits, mask, bias, lse, dropout in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"]): + for logits, mask, bias, lse, dropout, chunked in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]): if hdim == 256: # if True: - pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr', 'col', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask)) + pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr', 'col', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, 'f', chunked)) # the below two is used for hdim vectorize load - pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 'f', 'f', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr', 'col', 't', 't', 'f', 'f', logits, bias, lse, dropout, squant, mask)) + pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 'f', 'f', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr', 'col', 't', 't', 'f', 'f', logits, bias, lse, dropout, squant, mask, 'f', chunked)) - pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr', 'col', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask)) + pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr', 'col', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) else: if bias == "bias": - pipelines.append(FmhaFwdPipeline('qr_async', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr_async', 'col', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask)) + pipelines.append(FmhaFwdPipeline('qr_async', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr_async', 'col', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) else: - pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask)) - pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask)) + pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) + pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) if receipt == 1 and bias != "bias": - pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask)) # TODO: cover arbitraty hdim - pipelines.append(FmhaFwdPipeline('qr', 'col', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask)) # TODO: cover arbitraty hdim + pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) # TODO: cover arbitraty hdim + pipelines.append(FmhaFwdPipeline('qr', 'col', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, 'f', chunked)) # TODO: cover arbitraty hdim elif dtype in ['fp8', 'bf8']: # no need lse/dropout kernels for logits, mask, bias in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()): diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index 1838ee5bd9..f2ae98f438 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -351,6 +351,9 @@ struct fmha_batch_prefill_args #if 0 // we assume page_block_size=1 for now void* kv_last_page_lens; ck_tile::index_t page_block_size; +#else + ck_tile::index_t page_block_size; + ck_tile::index_t block_table_batch_stride; #endif float scale_s; @@ -731,6 +734,9 @@ auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args) #if 0 // we assume page_block_size=1 for now args.kv_last_page_lens, args.page_block_size, +#else + args.page_block_size, + args.block_table_batch_stride, #endif args.scale_s, args.scale_p, @@ -837,7 +843,9 @@ template + bool kPadDv_, + bool kIsSglangLayout_=false, + bool kIsChunkedPrefill_=false> struct fmha_fwd_traits_ { static constexpr ck_tile::index_t HDim = HDim_; @@ -861,6 +869,8 @@ struct fmha_fwd_traits_ static constexpr bool kPadSK = kPadSK_; static constexpr bool kPadD = kPadD_; static constexpr bool kPadDv = kPadDv_; + static constexpr bool kIsSglangLayout = kIsSglangLayout_; + static constexpr bool kIsChunkedPrefill = kIsChunkedPrefill_; }; template @@ -995,6 +1005,8 @@ struct fmha_fwd_traits bool has_lse; bool has_dropout; bool do_fp8_static_quant; + bool is_sglang_layout = false; + bool is_chunked_prefill = false; // TODO: padding check is inside this api }; float fmha_fwd(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config&); 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 ba327ee511..5ba55b5229 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 @@ -49,6 +49,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel static constexpr bool kPadHeadDimQ = FmhaPipeline::kPadHeadDimQ; static constexpr bool kPadHeadDimV = FmhaPipeline::kPadHeadDimV; static constexpr bool kHasLogitsSoftCap = FmhaPipeline::kHasLogitsSoftCap; + static constexpr bool kIsSglangLayout = FmhaPipeline::kIsSglangLayout; + static constexpr bool kIsChunkedPrefill = FmhaPipeline::kIsChunkedPrefill; static constexpr auto BiasEnum = FmhaPipeline::BiasEnum; static constexpr bool kStoreLSE = FmhaPipeline::kStoreLSE; static constexpr bool kHasDropout = FmhaPipeline::kHasDropout; @@ -137,7 +139,10 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel 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; + ck_tile::index_t page_block_size; + ck_tile::index_t block_table_batch_stride; +// #else +// static constexpr ck_tile::index_t page_block_size = 1; #endif float scale_s; @@ -485,7 +490,10 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel 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, + ck_tile::index_t page_block_size = 1, +#else + ck_tile::index_t page_block_size, + ck_tile::index_t block_table_batch_stride, #endif float scale_s, float scale_p, @@ -530,6 +538,9 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel #if 0 // we assume page_block_size=1 for now reinterpret_cast(kv_last_page_lens), page_block_size, +#else + page_block_size, + block_table_batch_stride, #endif #if CK_TILE_FMHA_FWD_FAST_EXP2 static_cast(scale_s * ck_tile::log2e_v<>), @@ -553,7 +564,6 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel reinterpret_cast(seqstart_q_ptr), batch_stride_k, batch_stride_v}; - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { kargs.bias_ptr = bias_ptr; @@ -651,7 +661,6 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel }; const auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1); - return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch); } else @@ -711,7 +720,15 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel batch_offset_q = query_start * kargs.stride_q; - kargs.kv_page_indices += kargs.kv_indptr[i_batch]; + + if constexpr(kIsSglangLayout) + { + kargs.kv_page_indices += kargs.kv_indptr[i_batch]; + } + else + { + kargs.kv_page_indices += i_batch * kargs.block_table_batch_stride; + } if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { @@ -737,6 +754,14 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel return; } + if constexpr(kIsChunkedPrefill) + { + if(kargs.seqlen_q == 1) + { + 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 @@ -811,7 +836,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel const auto k_dram = [&]() { const auto k_dram_naive = make_naive_tensor_view( k_ptr, - make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_q), + // make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_q), + make_tuple(kargs.num_total_pages, kargs.hdim_q), make_tuple(kargs.stride_k, 1), number{}, number<1>{}); @@ -827,7 +853,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel { const auto v_dram_naive = make_naive_tensor_view( v_ptr, - make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_v), + // make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_v), + make_tuple(kargs.num_total_pages, kargs.hdim_v), make_tuple(kargs.stride_v, 1), number{}, number<1>{}); @@ -836,7 +863,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel v_dram_naive, make_tuple( make_pass_through_transform(kargs.hdim_v), - make_pass_through_transform(kargs.num_total_pages * kargs.page_block_size)), + // make_pass_through_transform(kargs.num_total_pages * kargs.page_block_size)), + make_pass_through_transform(kargs.num_total_pages)), make_tuple(sequence<1>{}, sequence<0>{}), make_tuple(sequence<0>{}, sequence<1>{})); @@ -850,7 +878,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel { const auto v_dram_naive = make_naive_tensor_view( v_ptr, - make_tuple(kargs.hdim_v, kargs.num_total_pages * kargs.page_block_size), + // make_tuple(kargs.hdim_v, kargs.num_total_pages * kargs.page_block_size), + make_tuple(kargs.hdim_v, kargs.num_total_pages), make_tuple(kargs.stride_v, 1), number{}, number<1>{}); @@ -1103,7 +1132,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel kargs.kv_page_indices, kargs.stride_k, kargs.stride_v, - dropout); + dropout, + kargs.page_block_size); } }(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp index e07cf1c94e..ec57af1180 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp @@ -6,8 +6,9 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/common/tensor_layout.hpp" #include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async_default_policy.hpp" #include "ck_tile/ops/fmha/block/block_dropout.hpp" +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async_default_policy.hpp" #include "ck_tile/ops/reduce/block/block_reduce.hpp" namespace ck_tile { @@ -64,6 +65,9 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync static constexpr bool kPadHeadDimQ = true; // support multiple of vector(like 8x) static constexpr bool kPadHeadDimV = true; // support multiple of vector(like 8x) static constexpr bool kHasLogitsSoftCap = Problem::kHasLogitsSoftCap; + // static constexpr bool kIsSglangLayout = Problem::kIsSglangLayout; + static constexpr bool kIsSglangLayout = Problem::kIsSglangLayout; + static constexpr bool kIsChunkedPrefill = Problem::kIsChunkedPrefill; static constexpr auto BiasEnum = Problem::BiasEnum; static constexpr bool kStoreLSE = Problem::kStoreLSE; static constexpr bool kHasDropout = Problem::kHasDropout; @@ -193,7 +197,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync const index_t* page_idx, const index_t stride_k, const index_t stride_v, - DropoutType& dropout) const + DropoutType& dropout, + const index_t page_block_size) const { static_assert( std::is_same_v> && @@ -322,14 +327,27 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync using KDstrEncode = typename decltype(k_dist)::DstrEncode; constexpr index_t NRepeat = KDstrEncode::hs_lengthss_[I0][I0]; statically_indexed_array k_offsets; - static_for<0, NRepeat, 1>{}([&](auto n0) { - k_offsets[n0] = page_idx[k_coord[0] + kN0 / NRepeat * n0.value] * stride_k; - }); + if constexpr(kIsSglangLayout) + { + static_for<0, NRepeat, 1>{}([&](auto n0) { + k_offsets[n0] = page_idx[k_coord[0] + kN0 / NRepeat * n0.value] * stride_k; + }); + } + else + { + static_for<0, NRepeat, 1>{}([&](auto n0) { + int32_t seqlen_k_idx_per_repeat = k_coord[0] + kN0 / NRepeat * n0.value; + int32_t i_page = seqlen_k_idx_per_repeat / page_block_size; + int32_t i_seq = seqlen_k_idx_per_repeat % page_block_size; + k_offsets[n0] = (page_idx[i_page] * page_block_size + i_seq) * stride_k; + }); + } auto k_dram_window = make_tile_scatter_gather(k_dram_block_window.get_bottom_tensor_view(), k_dram_block_window.get_window_lengths(), k_dram_block_window.get_window_origin(), k_dist, k_offsets); // K DRAM tile window for + k_dram_window.init_raw(); constexpr auto k_oob_ck = bool_constant{}; constexpr auto k_pre_np = [&]() { @@ -358,9 +376,21 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync constexpr index_t V_KRepeat = VDstrEncode::hs_lengthss_[I1][I3]; statically_indexed_array v_offsets; (void)stride_k; - static_for<0, V_KRepeat, 1>{}([&](auto k0) { - v_offsets[k0] = page_idx[v_coord[VPageIndexDim] + k0.value] * stride_v; - }); + if constexpr(kIsSglangLayout) + { + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + v_offsets[k0] = page_idx[v_coord[VPageIndexDim] + k0.value] * stride_v; + }); + } + else + { + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + int32_t seqlen_v_idx_per_repeat = v_coord[VPageIndexDim] + k0.value; + int32_t i_page = seqlen_v_idx_per_repeat / page_block_size; + int32_t i_seq = seqlen_v_idx_per_repeat % page_block_size; + v_offsets[k0] = (page_idx[i_page] * page_block_size + i_seq) * stride_v; + }); + } auto v_dram_window = make_tile_scatter_gather(v_dram_block_window_tmp.get_bottom_tensor_view(), @@ -424,9 +454,21 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync const auto bias_tile = load_tile(bias_dram_window); // load bias tile auto v_buf = load_tile(v_dram_window, number<-1>{}, bool_constant{}); - static_for<0, V_KRepeat, 1>{}([&](auto k0) { - v_offsets[k0] = page_idx[kK1 + v_coord[VPageIndexDim] + k0.value] * stride_v; - }); + if constexpr(kIsSglangLayout) + { + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + v_offsets[k0] = page_idx[kK1 + v_coord[VPageIndexDim] + k0.value] * stride_v; + }); + } + else + { + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + int32_t seqlen_v_idx_per_repeat = kK1 + v_coord[VPageIndexDim] + k0.value; + int32_t i_page = seqlen_v_idx_per_repeat / page_block_size; + int32_t i_seq = seqlen_v_idx_per_repeat % page_block_size; + v_offsets[k0] = (page_idx[i_page] * page_block_size + i_seq) * stride_v; + }); + } v_dram_window.update_page_idx(v_offsets); __builtin_amdgcn_sched_barrier(0); @@ -707,11 +749,24 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync { v_buf = load_tile( v_dram_window, number<-1>{}, bool_constant{}); // load next v_buf - static_for<0, V_KRepeat, 1>{}([&](auto k0) { - v_offsets[k0] = page_idx[kK1 * 2 + i_k1.value * kK1 + - v_coord[VPageIndexDim] + k0.value] * - stride_v; - }); + if constexpr(kIsSglangLayout) + { + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + v_offsets[k0] = page_idx[kK1 * 2 + i_k1.value * kK1 + + v_coord[VPageIndexDim] + k0.value] * + stride_v; + }); + } + else + { + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + int32_t seqlen_v_idx_per_repeat = kK1 * 2 + i_k1.value * kK1 + + v_coord[VPageIndexDim] + k0.value; + int32_t i_page = seqlen_v_idx_per_repeat / page_block_size; + int32_t i_seq = seqlen_v_idx_per_repeat % page_block_size; + v_offsets[k0] = (page_idx[i_page] * page_block_size + i_seq) * stride_v; + }); + } v_dram_window.update_page_idx(v_offsets); } block_sync_lds(); @@ -757,9 +812,22 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync move_tile_window(k_dram_block_window, {kN0, 0}); k_dram_window.set_window_origin(k_dram_block_window.get_window_origin()); - static_for<0, NRepeat, 1>{}([&](auto n0) { - k_offsets[n0] = page_idx[k_coord[0] + kN0 / NRepeat * n0.value] * stride_k; - }); + if constexpr(kIsSglangLayout) + { + static_for<0, NRepeat, 1>{}([&](auto n0) { + k_offsets[n0] = page_idx[k_coord[0] + kN0 / NRepeat * n0.value] * stride_k; + }); + } + else + { + statically_indexed_array k_offsets; + static_for<0, NRepeat, 1>{}([&](auto n0) { + int32_t seqlen_k_idx_per_repeat = k_coord[0] + kN0 / NRepeat * n0.value; + int32_t i_page = seqlen_k_idx_per_repeat / page_block_size; + int32_t i_seq = seqlen_k_idx_per_repeat % page_block_size; + k_offsets[n0] = (page_idx[i_page] * page_block_size + i_seq) * stride_k; + }); + } k_dram_window.update_page_idx(k_offsets); if constexpr(k1_loops >= 2 && LdsSeq.at(number<0>{}) == LdsSeq.at(number{})) @@ -867,7 +935,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync const index_t* page_idx, const index_t stride_k, const index_t stride_v, - DropoutType& dropout) const + DropoutType& dropout, + const index_t page_block_size) const { return operator()(q_dram_block_window_tmp, identity{}, @@ -893,7 +962,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync page_idx, stride_k, stride_v, - dropout); + dropout, + page_block_size); } }; diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp index f35c00c268..cc8d1968d2 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp @@ -53,6 +53,8 @@ struct BlockFmhaPipelineProblem static constexpr bool kPadHeadDimQ = Traits::kPadHeadDimQ; static constexpr bool kPadHeadDimV = Traits::kPadHeadDimV; static constexpr bool kHasLogitsSoftCap = Traits::kHasLogitsSoftCap; + static constexpr bool kIsSglangLayout = Traits::kIsSglangLayout; + static constexpr bool kIsChunkedPrefill = Traits::kIsChunkedPrefill; static constexpr auto BiasEnum = Traits::BiasEnum; static constexpr bool kStoreLSE = Traits::kStoreLSE; static constexpr bool kHasDropout = Traits::kHasDropout; diff --git a/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp b/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp index 4530b58d85..58a8c62c65 100644 --- a/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp +++ b/include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp @@ -19,7 +19,9 @@ template + bool kIsSglangLayout_ = false, + bool kIsChunkedPrefill_ = false, + index_t kBlockPerCu_ = -1 /* overwrite occupancy if not -1 */> struct TileFmhaTraits { static constexpr bool kPadSeqLenQ = kPadSeqLenQ_; @@ -33,6 +35,8 @@ struct TileFmhaTraits static constexpr bool kHasDropout = kHasDropout_; static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; static constexpr index_t kBlockPerCu = kBlockPerCu_; + static constexpr bool kIsSglangLayout = kIsSglangLayout_; + static constexpr bool kIsChunkedPrefill = kIsChunkedPrefill_; }; template