ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

hot loop schedule for page fa

Browse Source
This commit is contained in:
zanzhang
2025-04-18 16:42:26 +08:00
parent 238331dbb5
commit 3af3a368a0
1 changed files with 289 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

290
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp
View File

@@ -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<kHasDropout, BlockDropout, NullBlockDropout>;
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<Problem>();
@@ -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<PDataType>(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<kM0, kK1>{}),
v_lds_window);
block_sync_lds();
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
{
auto v_shuffle_tmp = make_static_distributed_tensor<VDataType>(
Policy::template MakeShuffledVRegBlockDescriptor<Problem>());
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<kM0, 2 * kK1>{}),
v_lds_window);
block_sync_lds();
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
{
auto v_shuffle_tmp = make_static_distributed_tensor<VDataType>(
Policy::template MakeShuffledVRegBlockDescriptor<Problem>());
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<kM0, 3 * kK1>{}),
v_lds_window);
block_sync_lds();
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
{
auto v_shuffle_tmp = make_static_distributed_tensor<VDataType>(
Policy::template MakeShuffledVRegBlockDescriptor<Problem>());
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<kM0, kN0>{}),
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);