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

Let Q & S, K & V reuse the LDS

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-09-06 16:35:29 +00:00
parent 2b796b5c2f
commit 4b7aa8bb90
2 changed files with 23 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp
View File

@@ -197,34 +197,23 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
__builtin_amdgcn_sched_barrier(0);
// Q tile in LDS
QDataType* q_lds_ptr =
static_cast<QDataType*>(static_cast<void*>(static_cast<char*>(smem_ptr)));
auto q_lds = make_tensor_view<address_space_enum::lds>(
QDataType* q_lds_ptr = static_cast<QDataType*>(smem_ptr);
auto q_lds = make_tensor_view<address_space_enum::lds>(
q_lds_ptr, Policy::template MakeQLdsBlockDescriptor<Problem>());
// K tile in LDS
KDataType* k_lds_ptr =
static_cast<KDataType*>(static_cast<void*>(static_cast<char*>(smem_ptr)));
KDataType* k_lds_ptr = reinterpret_cast<KDataType*>(
static_cast<char*>(smem_ptr) + max(Policy::template GetSmemSizeQ<Problem>(),
Policy::template GetSmemSizeS<Problem>()));
auto k_lds = make_tensor_view<address_space_enum::lds>(
k_lds_ptr, Policy::template MakeKLdsBlockDescriptor<Problem>());
auto k_lds_window =
make_tile_window(k_lds, make_tuple(number<kN0>{}, number<kK0>{}), {0, 0});
// V tile in LDS
auto v_lds = make_tensor_view<address_space_enum::lds>(
reinterpret_cast<VDataType*>(static_cast<char*>(smem_ptr) +
max(Policy::template GetSmemSizeQ<Problem>(),
Policy::template GetSmemSizeK<Problem>())),
Policy::template MakeVLdsBlockDescriptor<Problem>());
auto v_lds_window = make_tile_window(
v_lds, Policy::template MakeVLdsBlockDescriptor<Problem>().get_lengths(), {0, 0});
// S tile in LDS
auto s_lds = make_tensor_view<address_space_enum::lds>(
reinterpret_cast<SaccDataType*>(reinterpret_cast<char*>(smem_ptr) +
max(Policy::template GetSmemSizeQ<Problem>(),
Policy::template GetSmemSizeK<Problem>())),
Policy::template MakeSLdsBlockDescriptor<Problem>());
SaccDataType* s_lds_ptr = static_cast<SaccDataType*>(smem_ptr);
auto s_lds = make_tensor_view<address_space_enum::lds>(
s_lds_ptr, Policy::template MakeSLdsBlockDescriptor<Problem>());
auto s_write_lds_window = make_tile_window(
s_lds, Policy::template MakeSLdsBlockDescriptor<Problem>().get_lengths(), {0, 0});
auto s_read_lds_window =
@@ -233,6 +222,15 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
{0, 0},
Policy::template MakeSRegTileDistribution<Problem>());
// V tile in LDS
VDataType* v_lds_ptr = reinterpret_cast<VDataType*>(
static_cast<char*>(smem_ptr) + max(Policy::template GetSmemSizeQ<Problem>(),
Policy::template GetSmemSizeS<Problem>()));
auto v_lds = make_tensor_view<address_space_enum::lds>(
v_lds_ptr, Policy::template MakeVLdsBlockDescriptor<Problem>());
auto v_lds_window = make_tile_window(
v_lds, Policy::template MakeVLdsBlockDescriptor<Problem>().get_lengths(), {0, 0});
// Block GEMM
constexpr auto gemm_0 = Policy::template GetQKBlockGemm<Problem>();
constexpr auto gemm_1 = Policy::template GetKVBlockGemm<Problem>();
@@ -375,9 +373,6 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
// prefetch K tile
index_t i_total_loops = 0;
// ensure all the ds_read for q are done
block_sync_lds();
do
{
// STAGE 1, QK gemm
@@ -672,9 +667,6 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
});
});
// ensure all ds_read for s are done
block_sync_lds();
auto v_shuffle_tmp = make_static_distributed_tensor<VDataType>(
Policy::template MakeShuffledVRegBlockDescriptor<Problem>());
shuffle_tile(v_shuffle_tmp, v_tiles[number<0>{}]);
@@ -711,6 +703,8 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
block_sync_lds();
};
});
block_sync_lds();
} while(++i_total_loops < num_total_loop);
if constexpr(kStoreLSE)

6
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp
View File

@@ -255,8 +255,10 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVSDefaultPolicy
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
{
return max(GetSmemSizeQ<Problem>(), GetSmemSizeK<Problem>()) +
max(GetSmemSizeV<Problem>(), GetSmemSizeS<Problem>());
// in hdim128 situtation, Q and S are 16x128, K and V are 32x128 or 64x128,
// so it makes sense to let Q and S reuse LDS, K and V resuse LDS
return max(GetSmemSizeQ<Problem>(), GetSmemSizeS<Problem>()) +
max(GetSmemSizeK<Problem>(), GetSmemSizeV<Problem>());
}
};