From 667ed5825c7735313b5d2f91340dce704cbfa366 Mon Sep 17 00:00:00 2001 From: ltqin Date: Fri, 12 Sep 2025 08:59:05 +0000 Subject: [PATCH] add kv_offset_array_transform to batch async for page size 16 --- .../01_fmha/codegen/ops/fmha_batch_prefill.py | 3 +- ..._batch_prefill_pipeline_qr_ks_vs_async.hpp | 168 ++++++++++++++---- .../pipeline/block_fmha_pipeline_problem.hpp | 1 + 3 files changed, 136 insertions(+), 36 deletions(-) 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 b6194805d1..9e3f1220e7 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 @@ -63,6 +63,7 @@ using fmha_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad}, {F_dropout}, {F_squant}, {F_occupancy}, + false, {F_sglang_layout}>; using fmha_variant_{F_idx} = ck_tile::ComposedAttention<{F_logits} * ck_tile::LOGITS_SOFT_CAP, CK_TILE_FMHA_FWD_FAST_EXP2>; @@ -517,8 +518,6 @@ class KernelComponentFactory: for logits, mask, bias, lse, dropout, sglang in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]): pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, sglang)) pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, sglang)) - # 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)) else: assert False return pipelines 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 6398bf316e..f534c6f47e 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 @@ -12,6 +12,57 @@ #include "ck_tile/ops/reduce/block/block_reduce.hpp" namespace ck_tile { +template +CK_TILE_HOST_DEVICE void kv_offset_array_transform(const index_t* page_vec, + const index_t& stride_kv, + const CoordVecType& coord_vec, + OffsetVecType& kv_offset_vec) +{ + const index_t& thread_coord_start = coord_vec[kCoordAxis]; + if constexpr(kIsSglangLayout) + { + static_for<0, kLoopCount, 1>{}([&](auto k0) { + kv_offset_vec[k0] = + page_vec[thread_coord_start + kLoopStart + kLoopStride * k0.value] * stride_kv; + }); + } + else + { + constexpr index_t kPageMask = (1 << kPageShiftSize) - 1; + if constexpr(kIsKcache) + { + // for k_offset_vec + static_for<0, kLoopCount, 1>{}([&](auto k0) { + constexpr index_t kPageId = kLoopStride * k0.value >> kPageShiftSize; + const index_t page_offset = + (thread_coord_start + kLoopStride * k0.value) & kPageMask; + kv_offset_vec[k0] = + ((page_vec[kPageId] << kPageShiftSize) + page_offset) * stride_kv; + }); + } + else + { + // for v_offset_vec + const index_t lane0_start = __builtin_amdgcn_readfirstlane(thread_coord_start); + const index_t lane0_page_id = (lane0_start + kLoopStart) >> kPageShiftSize; + const index_t page_loc = page_vec[lane0_page_id] << kPageShiftSize; + static_for<0, kLoopCount, 1>{}([&](auto k0) { + const index_t page_offset = + (thread_coord_start + kLoopStart + k0.value) & kPageMask; + kv_offset_vec[k0] = (page_loc + page_offset) * stride_kv; + }); + } + } +} // a variation of qr/ks/vs, where we use async copy to load k (potentially v in the future) template {}; - static constexpr auto I1 = number<1>{}; - static constexpr auto I2 = number<2>{}; - static constexpr auto I3 = number<3>{}; + static constexpr index_t kM0 = BlockFmhaShape::kM0; + static constexpr index_t kN0 = BlockFmhaShape::kN0; + static constexpr index_t kK0 = BlockFmhaShape::kK0; + static constexpr index_t kN1 = BlockFmhaShape::kN1; + static constexpr index_t kK1 = BlockFmhaShape::kK1; + static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim; + static constexpr index_t kSubQKHeaddim = BlockFmhaShape::kSubQKHeaddim; + static constexpr index_t kPageBlockSize = 16; + static constexpr index_t kPageShiftSize = 4; + static constexpr auto I0 = number<0>{}; + static constexpr auto I1 = number<1>{}; + static constexpr auto I2 = number<2>{}; + static constexpr auto I3 = number<3>{}; static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!"); @@ -68,6 +121,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync static constexpr auto BiasEnum = Problem::BiasEnum; static constexpr bool kStoreLSE = Problem::kStoreLSE; static constexpr bool kHasDropout = Problem::kHasDropout; + static constexpr bool kIsSglangLayout = Problem::kIsSglangLayout; static_assert((CK_TILE_FMHA_FWD_FAST_EXP2 && (kHasLogitsSoftCap && Problem::BiasEnum == BlockAttentionBiasEnum::NO_BIAS || @@ -323,9 +377,17 @@ 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; - }); + kv_offset_array_transform, + decltype(k_coord), + 0, + kPageBlockSize, + kPageShiftSize, + 0, + NRepeat, + kN0 / NRepeat, + kIsSglangLayout, + true>(page_idx, stride_k, k_coord, k_offsets); + 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(), @@ -358,10 +420,16 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync using VDstrEncode = typename decltype(v_dist)::DstrEncode; 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; - }); + kv_offset_array_transform, + decltype(v_coord), + VPageIndexDim, + kPageBlockSize, + kPageShiftSize, + 0, + V_KRepeat, + 1, + kIsSglangLayout, + false>(page_idx, stride_v, v_coord, v_offsets); auto v_dram_window = make_tile_scatter_gather(v_dram_block_window_tmp.get_bottom_tensor_view(), @@ -425,9 +493,16 @@ 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; - }); + kv_offset_array_transform, + decltype(v_coord), + VPageIndexDim, + kPageBlockSize, + kPageShiftSize, + kK1, + V_KRepeat, + 1, + kIsSglangLayout, + false>(page_idx, stride_v, v_coord, v_offsets); v_dram_window.update_page_idx(v_offsets); __builtin_amdgcn_sched_barrier(0); @@ -594,10 +669,16 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync {0, kK1}); // will have scratch if move this right after load_tile(v_dram)... 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 + v_coord[VPageIndexDim] + k0.value] * stride_v; - }); + kv_offset_array_transform, + decltype(v_coord), + VPageIndexDim, + kPageBlockSize, + kPageShiftSize, + 2 * kK1, + V_KRepeat, + 1, + kIsSglangLayout, + false>(page_idx, stride_v, v_coord, v_offsets); v_dram_window.update_page_idx(v_offsets); } __builtin_amdgcn_sched_barrier(0); @@ -725,11 +806,16 @@ 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; - }); + kv_offset_array_transform, + decltype(v_coord), + VPageIndexDim, + kPageBlockSize, + kPageShiftSize, + (2 + i_k1.value) * kK1, + V_KRepeat, + 1, + kIsSglangLayout, + false>(page_idx, stride_v, v_coord, v_offsets); v_dram_window.update_page_idx(v_offsets); } block_sync_lds(); @@ -770,14 +856,28 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync i_total_loops++; if(i_total_loops < num_total_loop) { - page_idx += kN0; + if constexpr(kIsSglangLayout) + { + page_idx += kN0; + } + else + { + page_idx += kN0 / kPageBlockSize; + } // move K tile windows 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; - }); + kv_offset_array_transform, + decltype(k_coord), + 0, + kPageBlockSize, + kPageShiftSize, + 0, + NRepeat, + kN0 / NRepeat, + kIsSglangLayout, + true>(page_idx, stride_k, k_coord, k_offsets); k_dram_window.update_page_idx(k_offsets); if constexpr(k1_loops >= 2 && LdsSeq.at(number<0>{}) == LdsSeq.at(number{})) 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 86ac713b6f..e301bcfa10 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 @@ -61,6 +61,7 @@ struct BlockFmhaPipelineProblem static constexpr bool kHasDropout = Traits::kHasDropout; static constexpr bool kDoFp8StaticQuant = Traits::kDoFp8StaticQuant; static constexpr index_t kBlockPerCu = Traits::kBlockPerCu; + static constexpr bool kIsSglangLayout = Traits::kIsSglangLayout; }; template