From 3af3a368a04db6dd915718d7bc9d35d37ed5de61 Mon Sep 17 00:00:00 2001 From: zanzhang Date: Fri, 18 Apr 2025 16:42:26 +0800 Subject: [PATCH] hot loop schedule for page fa --- .../pipeline/block_fmha_pipeline_qr_ks_vs.hpp | 290 +++++++++++++++++- 1 file changed, 289 insertions(+), 1 deletion(-) diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp index 2af044d539..c1f7673f08 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp @@ -11,6 +11,9 @@ #include "ck_tile/core/tensor/tile_distribution.hpp" #include "ck_tile/core/tensor/tile_scatter_gather.hpp" +#define qk_scheduler +#define kv_scheduler + namespace ck_tile { // This pipeline is qkv all located in LDS @@ -115,6 +118,34 @@ struct BlockFmhaPipelineQRKSVS using DropoutType = std::conditional_t; + CK_TILE_HOST_DEVICE static constexpr auto HotLoopScheduler() + { + __builtin_amdgcn_sched_group_barrier(0x100, 2, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + } + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return Policy::template GetSmemSize(); @@ -337,6 +368,12 @@ struct BlockFmhaPipelineQRKSVS 0); // prevent from messing up the order of global loads } const auto bias_tile = load_tile(bias_dram_window); // load bias tile +#ifdef qk_scheduler + __builtin_amdgcn_sched_group_barrier(0x020, 2, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x200, 2, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x020, 4, 0); // VMEM read + block_sync_lds(); +#endif if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { __builtin_amdgcn_sched_barrier( @@ -359,6 +396,9 @@ struct BlockFmhaPipelineQRKSVS k_lds_window, tile_elementwise_in(k_element_func, k_block_tile)); // LDS write i + 1 k_block_tile = load_tile(k_dram_window); // global read i + 2 +#ifdef qk_scheduler + HotLoopScheduler(); +#endif }); } @@ -379,6 +419,26 @@ struct BlockFmhaPipelineQRKSVS block_sync_lds(); store_tile(k_lds_window, tile_elementwise_in(k_element_func, k_block_tile)); +#ifdef qk_scheduler + __builtin_amdgcn_sched_group_barrier(0x100, 2, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA +#endif block_sync_lds(); gemm_0(s_acc, @@ -571,13 +631,219 @@ struct BlockFmhaPipelineQRKSVS tile_elementwise_in(v_element_func, v_prefetch)); // store the prefetch } move_tile_window(v_dram_window, {0, kK1}); - +#ifdef kv_scheduler + __builtin_amdgcn_sched_group_barrier(0x200, 4, 0); // DS write +#endif const auto p = cast_tile(tile_elementwise_in(p_compute_element_func, p_compute)); // STAGE 3, KV gemm if constexpr(k1_loops > 1) { + +#ifdef kv_scheduler + { + const auto v = load_tile(v_dram_window); // load next v + + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + v_offsets[k0] = page_idx[kK1 * 2 + v_coord[VPageIndexDim] + k0.value] * stride_v; + }); + v_dram_window.update_page_idx(v_offsets); + + block_sync_lds(); + gemm_1(o_acc, + get_slice_tile( + p, sequence<0, 0>{}, sequence{}), + v_lds_window); + block_sync_lds(); + if constexpr(std::is_same_v) + { + auto v_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledVRegBlockDescriptor()); + shuffle_tile(v_shuffle_tmp, v); + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, + v_shuffle_tmp)); // store the prefetch + } + else + { + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, v)); // store next v + } + move_tile_window(v_dram_window, {0, kK1}); + { + __builtin_amdgcn_sched_group_barrier(0x020, 3, 0); // VMEM read + + __builtin_amdgcn_sched_group_barrier(0x100, 2, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + + block_sync_lds(); + } + } + + { + const auto v = load_tile(v_dram_window); // load next v + + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + v_offsets[k0] = page_idx[kK1 * 2 + kK1 + v_coord[VPageIndexDim] + k0.value] * stride_v; + }); + v_dram_window.update_page_idx(v_offsets); + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + + block_sync_lds(); + gemm_1(o_acc, + get_slice_tile( + p, sequence<0, kK1>{}, sequence{}), + v_lds_window); + block_sync_lds(); + if constexpr(std::is_same_v) + { + auto v_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledVRegBlockDescriptor()); + shuffle_tile(v_shuffle_tmp, v); + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, + v_shuffle_tmp)); // store the prefetch + } + else + { + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, v)); // store next v + } + move_tile_window(v_dram_window, {0, kK1}); + { + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + block_sync_lds(); + } + } + + { + const auto v = load_tile(v_dram_window); // load next v + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + static_for<0, V_KRepeat, 1>{}([&](auto k0) { + v_offsets[k0] = page_idx[kK1 * 2 + 2 * kK1 + v_coord[VPageIndexDim] + k0.value] * stride_v; + }); + v_dram_window.update_page_idx(v_offsets); + + block_sync_lds(); + gemm_1(o_acc, + get_slice_tile( + p, sequence<0, 2 * kK1>{}, sequence{}), + v_lds_window); + block_sync_lds(); + if constexpr(std::is_same_v) + { + auto v_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledVRegBlockDescriptor()); + shuffle_tile(v_shuffle_tmp, v); + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, + v_shuffle_tmp)); // store the prefetch + } + else + { + store_tile(v_lds_window, + tile_elementwise_in(v_element_func, v)); // store next v + } + move_tile_window(v_dram_window, {0, kK1}); + { + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA 1 + block_sync_lds(); + } + + } +#else static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) { const auto v = load_tile(v_dram_window); // load next v @@ -608,6 +874,7 @@ struct BlockFmhaPipelineQRKSVS } move_tile_window(v_dram_window, {0, kK1}); }); +#endif } // move K tile windows move_tile_window(k_dram_block_window, {kN0, 0}); @@ -618,6 +885,27 @@ struct BlockFmhaPipelineQRKSVS get_slice_tile(p, sequence<0, (k1_loops - 1) * kK1>{}, sequence{}), v_lds_window); block_sync_lds(); +#ifdef kv_scheduler + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read 2 + __builtin_amdgcn_sched_group_barrier(0x008, 2, 0); // MFMA 1 +#endif } page_idx += kN0; } while(++i_total_loops < num_total_loop);