mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-17 09:08:35 +00:00
tempsave
This commit is contained in:
aska-0096
@@ -757,20 +757,50 @@ struct FmhaFwdDecodeKernel
|
||||
|
||||
const auto make_k_dram = [&](const KDataType* data, index_t height) {
|
||||
// We don't expect K data reuse among different blocks in decode case.
|
||||
const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global,
|
||||
memory_operation_enum::set,
|
||||
amd_buffer_coherence_enum::SYSTEM_NT1>(
|
||||
data, // will update this pointer if using paged-kvcache
|
||||
make_tuple(height, kargs.hdim_q),
|
||||
make_tuple(kargs.stride_k, 1),
|
||||
number<FmhaPipeline::kAlignmentK>{},
|
||||
number<1>{});
|
||||
const auto k_dram_naive = [&]() {
|
||||
const auto k_dram_ = make_naive_tensor_view<address_space_enum::global,
|
||||
memory_operation_enum::set,
|
||||
amd_buffer_coherence_enum::SYSTEM_NT1>(
|
||||
data, // will update this pointer if using paged-kvcache
|
||||
make_tuple(height, kargs.hdim_q),
|
||||
make_tuple(kargs.stride_k, 1),
|
||||
number<FmhaPipeline::kAlignmentK>{},
|
||||
number<1>{});
|
||||
|
||||
return pad_tensor_view(
|
||||
return pad_tensor_view(
|
||||
k_dram_,
|
||||
make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
|
||||
sequence<kPadSeqLenK, kPadHeadDimQ>{});
|
||||
}();
|
||||
|
||||
const auto k_dram_unmerged = transform_tensor_view(
|
||||
k_dram_naive,
|
||||
make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
|
||||
sequence<false, kPadHeadDimQ>{});
|
||||
make_tuple(make_pass_through_transform(height),
|
||||
make_unmerge_transform(make_tuple(
|
||||
number<FmhaPipeline::kQKHeaddim / FmhaPipeline::kAlignmentK>{},
|
||||
number<FmhaPipeline::kAlignmentK>{}))),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1, 2>{}));
|
||||
|
||||
const auto k_dram_permuted = transform_tensor_view(
|
||||
k_dram_unmerged,
|
||||
make_tuple(
|
||||
make_xor_transform(make_tuple(
|
||||
height, number<FmhaPipeline::kQKHeaddim / FmhaPipeline::kAlignmentK>{})),
|
||||
make_pass_through_transform(number<FmhaPipeline::kAlignmentK>{})),
|
||||
make_tuple(sequence<0, 1>{}, sequence<2>{}),
|
||||
make_tuple(sequence<0, 1>{}, sequence<2>{}));
|
||||
|
||||
return transform_tensor_view(
|
||||
k_dram_permuted,
|
||||
make_tuple(make_pass_through_transform(height),
|
||||
make_merge_transform_v3_division_mod(make_tuple(
|
||||
number<FmhaPipeline::kQKHeaddim / FmhaPipeline::kAlignmentK>{},
|
||||
number<FmhaPipeline::kAlignmentK>{}))),
|
||||
make_tuple(sequence<0>{}, sequence<1, 2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
};
|
||||
|
||||
const auto k_dram = [&]() {
|
||||
if constexpr(kIsPagedKV)
|
||||
{
|
||||
|
||||
@@ -45,12 +45,10 @@ struct BlockFmhaFwdDecodePipelineQRKSVS
|
||||
static constexpr index_t kSubQKHeaddim = BlockFmhaShape::kSubQKHeaddim;
|
||||
static constexpr index_t kNWarp = BlockFmhaShape::Gemm0BlockWarps::at(number<1>{});
|
||||
static constexpr index_t kNXdl = BlockFmhaShape::Gemm0WarpTile::at(number<1>{});
|
||||
// using HotLoopScheduler = typename Policy::template HotLoopScheduler<Problem>;
|
||||
|
||||
static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
|
||||
|
||||
// static_assert(Problem::kPadSeqLenQ == true && Problem::kPadHeadDimQ == true &&
|
||||
// Problem::kPadHeadDimV == true);
|
||||
|
||||
static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
|
||||
static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ;
|
||||
static constexpr bool kPadSeqLenK = Problem::kPadSeqLenK;
|
||||
|
||||
@@ -198,13 +198,27 @@ struct BlockFmhaFwdDecodePipelineQRKSVSDefaultPolicy
|
||||
|
||||
constexpr index_t kKPack = GetSmemKPackK<Problem>();
|
||||
|
||||
constexpr auto k_lds_block_desc =
|
||||
make_naive_tensor_descriptor(make_tuple(number<kNPerBlock>{}, number<kKPerBlock>{}),
|
||||
make_tuple(number<kKPerBlock>{}, number<1>{}),
|
||||
number<kKPack>{},
|
||||
number<1>{});
|
||||
constexpr auto k_lds_block_desc_naive = make_naive_tensor_descriptor(
|
||||
make_tuple(number<kNPerBlock>{}, number<kKPerBlock / kKPack>{}, number<kKPack>{}),
|
||||
make_tuple(number<kKPerBlock>{}, number<kKPack>{}, number<1>{}),
|
||||
number<kKPack>{},
|
||||
number<1>{});
|
||||
|
||||
return k_lds_block_desc;
|
||||
constexpr auto k_lds_block_desc_permuted = transform_tensor_descriptor(
|
||||
k_lds_block_desc_naive,
|
||||
make_tuple(
|
||||
make_xor_transform(make_tuple(number<kNPerBlock>{}, number<kKPerBlock / kKPack>{})),
|
||||
make_pass_through_transform(number<kKPack>{})),
|
||||
make_tuple(sequence<0, 1>{}, sequence<2>{}),
|
||||
make_tuple(sequence<0, 1>{}, sequence<2>{}));
|
||||
|
||||
return transform_tensor_descriptor(
|
||||
k_lds_block_desc_permuted,
|
||||
make_tuple(
|
||||
make_pass_through_transform(number<kNPerBlock>{}),
|
||||
make_merge_transform(make_tuple(number<kKPerBlock / kKPack>{}, number<kKPack>{}))),
|
||||
make_tuple(sequence<0>{}, sequence<1, 2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
@@ -428,6 +442,14 @@ struct BlockFmhaFwdDecodePipelineQRKSVSDefaultPolicy
|
||||
return static_cast<index_t>(16 / sizeof(SDataType));
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackVT()
|
||||
{
|
||||
using VDataType = remove_cvref_t<typename Problem::VDataType>;
|
||||
// Read use ds_read_b64_tr
|
||||
return static_cast<index_t>(8 / sizeof(VDataType));
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto MakeSLdsBlockDescriptor()
|
||||
{
|
||||
@@ -527,6 +549,137 @@ struct BlockFmhaFwdDecodePipelineQRKSVSDefaultPolicy
|
||||
return max(GetSmemSizeQ<Problem>(), GetSmemSizeK<Problem>()) + GetSmemSizeS<Problem>() +
|
||||
GetSmemSizeV<Problem>();
|
||||
}
|
||||
#if 0
|
||||
template <typename Problem_>
|
||||
struct HotLoopScheduler
|
||||
{
|
||||
using Problem = Problem_;
|
||||
|
||||
template <index_t GemmStage>
|
||||
CK_TILE_DEVICE static constexpr void GemmStagedScheduler()
|
||||
{
|
||||
}
|
||||
|
||||
template <>
|
||||
CK_TILE_DEVICE constexpr void GemmStagedScheduler<0>()
|
||||
{
|
||||
// Mem: Q, LSE, OGrad, D global load, OGrad^T LDS load
|
||||
// Comp: Q x K
|
||||
constexpr index_t VMEM_READ_INST =
|
||||
Q_VMEM_READ + OGrad_VMEM_READ + LSE_VMEM_READ + D_VMEM_READ;
|
||||
constexpr index_t LDS_READ_INST = OGradT_LDS_READ;
|
||||
constexpr index_t MFMA_INST = Gemm0MFMA;
|
||||
|
||||
// Evenly distributed to relieve SQ->TA FIFO pressure
|
||||
constexpr index_t MFMA_PER_VMEM_READ = MFMA_INST / VMEM_READ_INST;
|
||||
constexpr index_t MFMA_Remainder = MFMA_INST - MFMA_PER_VMEM_READ * VMEM_READ_INST;
|
||||
// To hide instruction issue latency
|
||||
constexpr index_t LDS_READ_PER_MFMA = LDS_READ_INST / MFMA_INST;
|
||||
|
||||
static_for<0, VMEM_READ_INST, 1>{}([&](auto i) {
|
||||
ignore = i;
|
||||
__builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read
|
||||
static_for<0, MFMA_PER_VMEM_READ, 1>{}([&](auto j) {
|
||||
ignore = j;
|
||||
__builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
|
||||
__builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS read
|
||||
});
|
||||
});
|
||||
static_for<0, MFMA_Remainder, 1>{}([&](auto i) {
|
||||
ignore = i;
|
||||
__builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
|
||||
__builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS read
|
||||
});
|
||||
}
|
||||
|
||||
template <>
|
||||
CK_TILE_DEVICE constexpr void GemmStagedScheduler<1>()
|
||||
{
|
||||
// Mem: Q^T LDS load
|
||||
// Comp: OGrad x V
|
||||
constexpr index_t LDS_READ_INST = QT_LDS_READ;
|
||||
constexpr index_t MFMA_INST = Gemm1MFMA;
|
||||
|
||||
// To hide instruction issue latency
|
||||
constexpr index_t LDS_READ_PER_MFMA = LDS_READ_INST / MFMA_INST;
|
||||
|
||||
static_for<0, MFMA_INST, 1>{}([&](auto i) {
|
||||
ignore = i;
|
||||
__builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
|
||||
__builtin_amdgcn_sched_group_barrier(0x100, LDS_READ_PER_MFMA, 0); // DS read
|
||||
});
|
||||
}
|
||||
|
||||
private:
|
||||
static constexpr index_t kBlockSize = Problem::kBlockSize;
|
||||
static constexpr index_t kM0 = Problem::BlockFmhaShape::kM0;
|
||||
static constexpr index_t kN0 = Problem::BlockFmhaShape::kN0;
|
||||
static constexpr index_t kN1 = Problem::BlockFmhaShape::kN1;
|
||||
static constexpr index_t kQKHeaddim = Problem::BlockFmhaShape::kQKHeaddim;
|
||||
static constexpr index_t kVHeaddim = Problem::BlockFmhaShape::kVHeaddim;
|
||||
static constexpr index_t kK0 = Problem::BlockFmhaShape::kK0;
|
||||
static constexpr index_t kK1 = Problem::BlockFmhaShape::kK1;
|
||||
|
||||
using WarpGemm0 =
|
||||
WarpGemmMfmaDispatcher<typename Problem::QDataType,
|
||||
typename Problem::KDataType,
|
||||
typename Problem::SaccDataType,
|
||||
Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}),
|
||||
Problem::BlockFmhaShape::Gemm0WarpTile::at(number<1>{}),
|
||||
Problem::BlockFmhaShape::Gemm0WarpTile::at(number<2>{}),
|
||||
true>;
|
||||
|
||||
using WarpGemm1 =
|
||||
WarpGemmMfmaDispatcher<typename Problem::PDataType,
|
||||
typename Problem::VDataType,
|
||||
typename Problem::OaccDataType,
|
||||
Problem::BlockFmhaShape::Gemm1WarpTile::at(number<0>{}),
|
||||
Problem::BlockFmhaShape::Gemm1WarpTile::at(number<1>{}),
|
||||
Problem::BlockFmhaShape::Gemm1WarpTile::at(number<2>{}),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
WGAttrNumAccessEnum::Double>;
|
||||
|
||||
static constexpr index_t WarpGemm0M = WarpGemm0::kM;
|
||||
static constexpr index_t WarpGemm0N = WarpGemm0::kN;
|
||||
static constexpr index_t WarpGemm0K = WarpGemm0::kK;
|
||||
static constexpr index_t WarpGemm1M = WarpGemm1::kM;
|
||||
static constexpr index_t WarpGemm1N = WarpGemm1::kN;
|
||||
static constexpr index_t WarpGemm1K = WarpGemm1::kK;
|
||||
|
||||
static constexpr index_t NumOfWarps = kBlockSize / get_warp_size();
|
||||
static constexpr index_t Gemm0NWarp =
|
||||
Problem::BlockFmhaShape::Gemm0BlockWarps::at(number<1>{});
|
||||
static constexpr index_t Gemm1NWarp =
|
||||
Problem::BlockFmhaShape::Gemm1BlockWarps::at(number<1>{});
|
||||
|
||||
// Compute
|
||||
// 16*32*128/16*16*32 = 8
|
||||
static constexpr index_t Gemm0MFMA =
|
||||
kM0 * kN0 * kQKHeaddim / ( NumOfWarps * WarpGemm0M * WarpGemm0N * WarpGemm0K);
|
||||
static constexpr index_t Gemm0MFMA =
|
||||
kM0 * kN1 * kVHeaddim / ( NumOfWarps * WarpGemm1M * WarpGemm1N * WarpGemm1K);
|
||||
|
||||
// VMEM
|
||||
// 32*128/64/8 = 8
|
||||
static constexpr index_t K_VMEM_READ =
|
||||
kN0 * kQKHeaddim / kBlockSize / GetAlignmentK<Problem>();
|
||||
static constexpr index_t V_VMEM_READ =
|
||||
kN1 * kVHeaddim / kBlockSize / GetAlignmentV<Problem>();
|
||||
|
||||
// LDS Read
|
||||
static constexpr index_t K_LDS_READ =
|
||||
kN0 * kQKHeaddim / (Gemm0NWarp * get_warp_size() * GetSmemKPackK<Problem>());
|
||||
static constexpr index_t V_LDS_READ =
|
||||
kN1 * kVHeaddim / (Gemm1NWarp * get_warp_size() * GetSmemKPackVT<Problem>());
|
||||
|
||||
/* Add Shuffle S */
|
||||
|
||||
// LDS Write
|
||||
/* Add Shuffle S */
|
||||
};
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
|
||||
Reference in New Issue
Block a user