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

atomic16 base impl

Browse Source 
formatting code

fix compile error

fix conflict

use global_atomic_pk_add instr

remove redundant modifications

formatting code

remove seqstart_dq_acc in varlen mode

formatting code
This commit is contained in:
shay-li77
2025-08-02 00:16:37 +08:00
parent 33418b201f
commit 5be2aae20e
16 changed files with 603 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
include/ck_tile/core/arch/amd_buffer_addressing.hpp
View File

@@ -790,6 +790,34 @@ struct buffer_atomic_add_if<bf16_t, 2, pre_nop>
}
};
template <bool pre_nop>
struct buffer_atomic_add_if<fp16_t, 2, pre_nop>
{
template <typename T>
CK_TILE_DEVICE void operator()(const T& value,
int32x4_t res /*buffer resource*/,
index_t v_offset,
index_t /*s_offset*/,
index_t i_offset /*max 0xFFF*/,
index_t flag = 1)
{
static_assert(sizeof(T) == 4);
auto save_exec = __builtin_amdgcn_read_exec();
using mbuf_t = float;
asm volatile("v_cmpx_le_u32 exec, 1, %4\n"
"global_atomic_pk_add_f16 %0, %1, %2 offset:%3\n"
"s_mov_b64 exec %5"
:
: "v"(v_offset),
"v"(bit_cast<mbuf_t>(value)),
"s"(res.xy),
"n"(i_offset),
"v"(flag),
"s"(save_exec)
: "memory");
}
};
template <typename scalar_type, index_t N, bool pre_nop = false>
struct buffer_atomic_add;
@@ -813,6 +841,26 @@ struct buffer_atomic_add<bf16_t, 2, pre_nop>
}
};
template <bool pre_nop>
struct buffer_atomic_add<fp16_t, 2, pre_nop>
{
template <typename T>
CK_TILE_DEVICE void operator()(const T& value,
int32x4_t res /*buffer resource*/,
index_t v_offset,
index_t /*s_offset*/,
index_t i_offset /*max 0xFFF*/,
index_t /*flag = 1*/)
{
static_assert(sizeof(T) == 4);
using mbuf_t = float;
asm volatile("global_atomic_pk_add_f16 %0, %1, %2 offset:%3"
:
: "v"(v_offset), "v"(bit_cast<mbuf_t>(value)), "s"(res.xy), "n"(i_offset)
: "memory");
}
};
namespace impl {
// below type indicate the data type used for buffer load inline asm
// clang-format off

48
include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp
View File

@@ -658,6 +658,34 @@ struct buffer_atomic_add_if<bf16_t, 2, pre_nop>
}
};
template <bool pre_nop>
struct buffer_atomic_add_if<fp16_t, 2, pre_nop>
{
template <typename T>
CK_TILE_DEVICE void operator()(const T& value,
int32x4_t res /*buffer resource*/,
index_t v_offset,
index_t /*s_offset*/,
index_t i_offset /*max 0xFFF*/,
index_t flag = 1)
{
static_assert(sizeof(T) == 4);
auto save_exec = __builtin_amdgcn_read_exec();
using mbuf_t = float;
asm volatile("v_cmpx_le_u32 exec, 1, %4\n"
"global_atomic_pk_add_f16 %0, %1, %2 offset:%3\n"
"s_mov_b64 exec %5"
:
: "v"(v_offset),
"v"(bit_cast<mbuf_t>(value)),
"s"(res.xy),
"n"(i_offset),
"v"(flag),
"s"(save_exec)
: "memory");
}
};
template <typename scalar_type, index_t N, bool pre_nop = false>
struct buffer_atomic_add;
@@ -681,6 +709,26 @@ struct buffer_atomic_add<bf16_t, 2, pre_nop>
}
};
template <bool pre_nop>
struct buffer_atomic_add<fp16_t, 2, pre_nop>
{
template <typename T>
CK_TILE_DEVICE void operator()(const T& value,
int32x4_t res /*buffer resource*/,
index_t v_offset,
index_t /*s_offset*/,
index_t i_offset /*max 0xFFF*/,
index_t /*flag = 1*/)
{
static_assert(sizeof(T) == 4);
using mbuf_t = float;
asm volatile("global_atomic_pk_add_f16 %0, %1, %2 offset:%3"
:
: "v"(v_offset), "v"(bit_cast<mbuf_t>(value)), "s"(res.xy), "n"(i_offset)
: "memory");
}
};
namespace impl {
// below type indicate the data type used for buffer load inline asm
// clang-format off

2
include/ck_tile/core/tensor/tensor_view.hpp
View File

@@ -455,7 +455,7 @@ CK_TILE_HOST_DEVICE constexpr auto make_tensor_view(DataType* __restrict__ p,
auto buffer_view =
make_buffer_view<BufferAddressSpace, Coherence>(p, desc.get_element_space_size());
return tensor_view<decltype(buffer_view), decltype(desc)>{buffer_view, desc};
return tensor_view<decltype(buffer_view), decltype(desc), DstInMemOp>{buffer_view, desc};
}
template <address_space_enum BufferAddressSpace = address_space_enum::generic,

2
include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp
View File

@@ -1143,7 +1143,7 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN1>{}),
{i_m0, i_n1});
EpiloguePipeline{}(o_dram_window, o_acc_tile);
EpiloguePipeline{}(o_dram_window, o_acc_tile, nullptr);
}
};

228
include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp
View File

@@ -71,15 +71,20 @@ struct FmhaBwdDQDKDVKernel
static constexpr bool kHasDropout = FmhaDropout::IsDropout;
static constexpr bool kIsStoreRandval = FmhaDropout::IsStoreRandval;
static constexpr bool kIsDeterministic = FmhaPipeline::kIsDeterministic;
static constexpr bool kIsAtomic32 = FmhaPipeline::kIsAtomic32;
static constexpr bool kUseTrLoad = FmhaPipeline::kUseTrLoad;
static constexpr index_t kMaxSeqLenQ = FmhaPipeline::BlockFmhaShape::kMaxSeqLenQ;
static_assert(kUseQrQtrDorPipeline == (kMaxSeqLenQ != 0));
static_assert(!kUseTrLoad || kIsAtomic32);
static_assert(!kIsDeterministic || kIsAtomic32);
#if defined(__gfx950__)
static constexpr bool kIsAvialable = true;
#else
static constexpr bool kIsAvialable = !kUseTrLoad;
#endif
using QGradAccDataType = std::conditional_t<kIsAtomic32, AccDataType, QDataType>;
// clang-format off
template <typename T> struct t2s;
template <> struct t2s<ck_tile::fp16_t> { static constexpr const char * name = "fp16"; };
@@ -116,7 +121,7 @@ struct FmhaBwdDQDKDVKernel
("o" + _TS_(kBlockPerCu)) + (pn.empty() ? "_npad" : "_" + pn) +
(BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("_nbias") : (_SS_("_") + BlockAttentionBiasEnumToStr<BiasEnum>::name)) +
(kHasBiasGrad ? "_dbias" : "_ndbias") + (kHasMask ? "_" + _SS_(FmhaMask::name) : "_nmask") + (kHasDropout ? "_dropout" : "_ndropout" ) +
(kIsStoreRandval ? "_storerandval" : "" ) + (kIsDeterministic ? "_deterministic" : "_ndeterministic" ) + (kUseTrLoad ? "_trload" : "_ntrload");
(kIsStoreRandval ? "_storerandval" : "" ) + (kIsDeterministic ? "_deterministic" : (kIsAtomic32 ? "_atomic32" : "_atomic16")) + (kUseTrLoad ? "_trload" : "_ntrload");
#undef _SS_
#undef _TS_
// clang-format on
@@ -274,6 +279,11 @@ struct FmhaBwdDQDKDVKernel
ck_tile::index_t split_stride_dq_acc = 0;
};
struct FmhaBwdAtomic16GroupModeKargs
{
ck_tile::index_t max_seqlen_q_aligned = 0;
};
struct FmhaBwdBatchModeKargs
: FmhaBwdCommonKargs,
std::conditional_t<BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS,
@@ -306,7 +316,8 @@ struct FmhaBwdDQDKDVKernel
std::conditional_t<kHasBiasGrad, FmhaBwdCommonBiasGradKargs, FmhaBwdEmptyKargs<1>>,
std::conditional_t<kHasMask, FmhaBwdMaskKargs, FmhaBwdEmptyKargs<2>>,
std::conditional_t<kHasDropout, FmhaBwdCommonDropoutKargs, FmhaBwdEmptyKargs<3>>,
std::conditional_t<kIsDeterministic, FmhaBwdDeterministicKargs, FmhaBwdEmptyKargs<4>>
std::conditional_t<kIsDeterministic, FmhaBwdDeterministicKargs, FmhaBwdEmptyKargs<4>>,
std::conditional_t<!kIsAtomic32, FmhaBwdAtomic16GroupModeKargs, FmhaBwdEmptyKargs<5>>
{
const int32_t* seqstart_q_ptr;
const int32_t* seqstart_k_ptr;
@@ -518,6 +529,7 @@ struct FmhaBwdDQDKDVKernel
const void* seqstart_q_ptr,
const void* seqstart_k_ptr,
const void* seqlen_k_ptr,
ck_tile::index_t max_seqlen_q_aligned,
ck_tile::index_t hdim_q,
ck_tile::index_t hdim_v,
ck_tile::index_t num_head_q,
@@ -589,6 +601,7 @@ struct FmhaBwdDQDKDVKernel
{}, // placeholder for mask
{}, // placeholder for dropout
{}, // placeholder for deterministic
{}, // placeholder for atomic16
reinterpret_cast<const int32_t*>(seqstart_q_ptr),
reinterpret_cast<const int32_t*>(seqstart_k_ptr),
reinterpret_cast<const int32_t*>(seqlen_k_ptr)};
@@ -644,6 +657,11 @@ struct FmhaBwdDQDKDVKernel
kargs.split_stride_dq_acc = split_stride_dq_acc;
}
if constexpr(!kIsAtomic32)
{
kargs.max_seqlen_q_aligned = max_seqlen_q_aligned;
}
return kargs;
}
@@ -707,13 +725,22 @@ struct FmhaBwdDQDKDVKernel
// get starting offset for each batch
const long_index_t query_start = kargs.seqstart_q_ptr[i_batch];
const long_index_t key_start = kargs.seqstart_k_ptr[i_batch];
long_index_t dq_acc_start = 0;
if constexpr(kIsAtomic32)
{
dq_acc_start = kargs.seqstart_q_ptr[i_batch];
}
else
{
dq_acc_start = kargs.max_seqlen_q_aligned * i_batch;
}
batch_offset_q = query_start * kargs.stride_q;
batch_offset_k = key_start * kargs.stride_k;
batch_offset_v = key_start * kargs.stride_v;
batch_offset_do = query_start * kargs.stride_do;
batch_offset_lsed = query_start;
batch_offset_dq_acc = query_start * kargs.stride_dq_acc;
batch_offset_dq_acc = dq_acc_start * kargs.stride_dq_acc;
batch_offset_dk = key_start * kargs.stride_dk;
batch_offset_dv = key_start * kargs.stride_dv;
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
@@ -879,7 +906,9 @@ struct FmhaBwdDQDKDVKernel
auto dq_dram_window = [&, i_tile_n_ = i_tile_n, i_nhead_ = i_nhead]() {
constexpr bool kUseKSplit = !kUseQrQtrDorPipeline && kIsDeterministic;
using DType = std::conditional_t<kUseQrQtrDorPipeline, QGradDataType, AccDataType>;
using DType = std::
conditional_t<kUseQrQtrDorPipeline || !kIsAtomic32, QGradDataType, AccDataType>;
auto dq_acc_ptr = reinterpret_cast<DType*>(kargs.dq_acc_ptr) + [&]() {
if constexpr(kUseKSplit)
@@ -893,17 +922,71 @@ struct FmhaBwdDQDKDVKernel
constexpr auto DstInMemOp = conditional_expr<kUseKSplit>(
memory_operation_enum::set, memory_operation_enum::atomic_add);
const auto dq_acc_dram_naive =
make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
dq_acc_ptr,
make_tuple(kargs.seqlen_q, kargs.hdim_q),
make_tuple(kargs.stride_dq_acc, 1),
number<FmhaPipeline::kAlignmentQGrad>{},
number<1>{});
const auto dq_acc_dram = pad_tensor_view(
dq_acc_dram_naive,
make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kQKHeaddim>{}),
sequence<false, kPadHeadDimQ>{});
auto dq_acc_dram = [&]() {
if constexpr(kIsAtomic32)
{
const auto dq_acc_dram_naive =
make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
dq_acc_ptr,
make_tuple(kargs.seqlen_q, kargs.hdim_q),
make_tuple(kargs.stride_dq_acc, 1),
number<FmhaPipeline::kAlignmentQGrad>{},
number<1>{});
return pad_tensor_view(
dq_acc_dram_naive,
make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kQKHeaddim>{}),
sequence<false, kPadHeadDimQ>{});
}
else
{
constexpr index_t m_pack = 2; // dword alignment for atomic 16 instr.
constexpr index_t mfma_m1_per_lane = 4;
constexpr index_t m1_pack_num = mfma_m1_per_lane / m_pack;
constexpr index_t mfma_n_lane = FmhaPipeline::kGemm4WarpN;
constexpr index_t mfma_m_lane = get_warp_size() / mfma_n_lane;
constexpr index_t m_align_size = mfma_m1_per_lane * mfma_m_lane;
static_assert(
FmhaPipeline::kM0 % m_align_size == 0,
"tiling size in the m direction must be divisible by the m align size.");
index_t M0 = (kargs.seqlen_q + FmhaPipeline::kM0 - 1) / m_align_size;
constexpr auto dq_acc_n = FmhaPipeline::kQKHeaddim;
constexpr index_t N0 = dq_acc_n / mfma_n_lane;
const auto q_grad_dram_desc_0 = make_naive_tensor_descriptor(
make_tuple(M0,
number<N0>{},
number<m1_pack_num>{},
number<mfma_m_lane>{},
number<mfma_n_lane>{},
number<m_pack>{}),
make_tuple(number<dq_acc_n * mfma_m1_per_lane * mfma_m_lane>{},
number<mfma_n_lane * mfma_m1_per_lane * mfma_m_lane>{},
number<mfma_m_lane * mfma_n_lane * m_pack>{},
number<mfma_n_lane * m_pack>{},
number<m_pack>{},
number<1>{}),
number<m_pack>{},
number<1>{});
const auto q_grad_dram_desc = transform_tensor_descriptor(
q_grad_dram_desc_0,
make_tuple(
make_merge_transform(make_tuple(M0,
number<mfma_m_lane>{},
number<m1_pack_num>{},
number<m_pack>{})),
make_merge_transform(make_tuple(number<N0>{}, number<mfma_n_lane>{}))),
make_tuple(sequence<0, 3, 2, 5>{}, sequence<1, 4>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tensor_view<address_space_enum::global,
memory_operation_enum::atomic_add>(dq_acc_ptr,
q_grad_dram_desc);
}
}();
return make_tile_window(
dq_acc_dram,
make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kQKHeaddim>{}),
@@ -1430,14 +1513,18 @@ struct FmhaBwdConvertQGradKernel
static constexpr ck_tile::index_t kM0 = FmhaBwdConvertQGrad::kM0;
static constexpr ck_tile::index_t kN0 = FmhaBwdConvertQGrad::kN0;
static constexpr ck_tile::index_t kQKHeaddim = FmhaBwdConvertQGrad::kQKHeaddim;
static constexpr ck_tile::index_t kGemm4WarpN = FmhaBwdConvertQGrad::kGemm4WarpN;
using AccDataType = ck_tile::remove_cvref_t<typename FmhaBwdConvertQGrad::AccDataType>;
using QGradDataType = ck_tile::remove_cvref_t<typename FmhaBwdConvertQGrad::QGradDataType>;
using QGradAccDataType =
ck_tile::remove_cvref_t<typename FmhaBwdConvertQGrad::QGradAccDataType>;
static constexpr bool kIsGroupMode = FmhaBwdConvertQGrad::kIsGroupMode;
static constexpr bool kPadSeqLenQ = FmhaBwdConvertQGrad::kPadSeqLenQ;
static constexpr bool kPadHeadDimQ = FmhaBwdConvertQGrad::kPadHeadDimQ;
static constexpr bool kIsDeterministic = FmhaBwdConvertQGrad::kIsDeterministic;
static constexpr bool kIsAtomic32 = FmhaBwdConvertQGrad::kIsAtomic32;
// clang-format off
template <typename T> struct t2s;
@@ -1463,7 +1550,7 @@ struct FmhaBwdConvertQGradKernel
+ "b" + _TS_(kM0) + "x" + _TS_(kN0) + "_"
+ (kIsGroupMode ? "group" : "batch") + "_"
+ ("o" + _TS_(kBlockPerCu)) + (pn.empty() ? "_npad" : "_" + pn)
+ (kIsDeterministic ? "_deterministic" : "_ndeterministic") ;
+ (kIsDeterministic ? "_deterministic" : (kIsAtomic32 ? "_atomic32" : "_atomic16")) ;
#undef _SS_
#undef _TS_
// clang-format on
@@ -1498,6 +1585,11 @@ struct FmhaBwdConvertQGradKernel
ck_tile::index_t split_stride_dq_acc = 0;
};
struct FmhaBwdConvertQGradAtomic16GroupModeKargs
{
ck_tile::index_t max_seqlen_q_aligned = 0;
};
struct FmhaBwdConvertQGradBatchModeKargs
: FmhaBwdConvertQGradCommonKargs,
std::conditional_t<kIsDeterministic,
@@ -1512,7 +1604,10 @@ struct FmhaBwdConvertQGradKernel
: FmhaBwdConvertQGradCommonKargs,
std::conditional_t<kIsDeterministic,
FmhaBwdConvertQGradDeterministicKargs,
FmhaBwdConvertQGradEmptyKargs<0>>
FmhaBwdConvertQGradEmptyKargs<0>>,
std::conditional_t<!kIsAtomic32,
FmhaBwdConvertQGradAtomic16GroupModeKargs,
FmhaBwdConvertQGradEmptyKargs<1>>
{
const int32_t* seqstart_q_ptr;
const int32_t* seqstart_k_ptr;
@@ -1564,6 +1659,7 @@ struct FmhaBwdConvertQGradKernel
void* dq_ptr,
const void* seqstart_q_ptr,
const void* seqstart_k_ptr,
ck_tile::index_t max_seqlen_q_aligned,
ck_tile::index_t hdim_q,
ck_tile::index_t stride_dq,
ck_tile::index_t stride_dq_acc,
@@ -1580,7 +1676,8 @@ struct FmhaBwdConvertQGradKernel
stride_dq_acc,
nhead_stride_dq,
nhead_stride_dq_acc},
{},
{}, // placeholder for deterministic
{}, // placeholder for atomic16
reinterpret_cast<const int32_t*>(seqstart_q_ptr),
reinterpret_cast<const int32_t*>(seqstart_k_ptr)};
@@ -1589,6 +1686,11 @@ struct FmhaBwdConvertQGradKernel
kargs.split_stride_dq_acc = split_stride_dq_acc;
}
if constexpr(!kIsAtomic32)
{
kargs.max_seqlen_q_aligned = max_seqlen_q_aligned;
}
return kargs;
}
@@ -1624,8 +1726,17 @@ struct FmhaBwdConvertQGradKernel
{
// get starting offset for each batch
const long_index_t query_start = kargs.seqstart_q_ptr[i_batch];
batch_offset_dq = query_start * kargs.stride_dq;
batch_offset_dq_acc = query_start * kargs.stride_dq_acc;
long_index_t dq_acc_start = 0;
if constexpr(kIsAtomic32)
{
dq_acc_start = kargs.seqstart_q_ptr[i_batch];
}
else
{
dq_acc_start = kargs.max_seqlen_q_aligned * i_batch;
}
batch_offset_dq = query_start * kargs.stride_dq;
batch_offset_dq_acc = dq_acc_start * kargs.stride_dq_acc;
// get real # queries & # keys under group mode
const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch;
@@ -1676,20 +1787,75 @@ struct FmhaBwdConvertQGradKernel
}
else
{
const AccDataType* dq_acc_ptr =
reinterpret_cast<const AccDataType*>(kargs.dq_acc_ptr) +
const QGradAccDataType* dq_acc_ptr =
reinterpret_cast<const QGradAccDataType*>(kargs.dq_acc_ptr) +
static_cast<long_index_t>(i_nhead_) * (kargs.nhead_stride_dq_acc) +
batch_offset_dq_acc;
if constexpr(kIsAtomic32)
{
auto dq_acc_dram_naive = make_naive_tensor_view<address_space_enum::global>(
dq_acc_ptr,
make_tuple(kargs.seqlen_q, kargs.hdim_q),
make_tuple(kargs.stride_dq_acc, 1),
number<FmhaBwdConvertQGrad::kAlignmentQGradAcc>{},
number<1>{});
return pad_tensor_view(dq_acc_dram_naive,
make_tuple(number<kM0>{}, number<kQKHeaddim>{}),
sequence<kPadSeqLenQ, kPadHeadDimQ>{});
auto dq_acc_dram_naive = make_naive_tensor_view<address_space_enum::global>(
dq_acc_ptr,
make_tuple(kargs.seqlen_q, kargs.hdim_q),
make_tuple(kargs.stride_dq_acc, 1),
number<FmhaBwdConvertQGrad::kAlignmentQGradAcc>{},
number<1>{});
return pad_tensor_view(dq_acc_dram_naive,
make_tuple(number<kM0>{}, number<kQKHeaddim>{}),
sequence<kPadSeqLenQ, kPadHeadDimQ>{});
}
else
{
constexpr index_t m_pack = 2; // dword alignment for atomic 16 instr.
constexpr index_t mfma_m1_per_lane = 4;
constexpr index_t m1_pack_num = mfma_m1_per_lane / m_pack;
constexpr index_t mfma_n_lane = kGemm4WarpN;
constexpr index_t mfma_m_lane = get_warp_size() / mfma_n_lane;
constexpr index_t m_align_size = mfma_m1_per_lane * mfma_m_lane;
static_assert(
kM0 % m_align_size == 0,
"tiling size in the m direction must be divisible by the m align size.");
index_t M0 = (kargs.seqlen_q + m_align_size - 1) / m_align_size;
constexpr auto dq_acc_n = kQKHeaddim;
constexpr index_t N0 = dq_acc_n / mfma_n_lane;
const auto q_grad_dram_desc_0 = make_naive_tensor_descriptor(
make_tuple(M0,
number<N0>{},
number<m1_pack_num>{},
number<mfma_m_lane>{},
number<mfma_n_lane>{},
number<m_pack>{}),
make_tuple(number<dq_acc_n * mfma_m1_per_lane * mfma_m_lane>{},
number<mfma_n_lane * mfma_m1_per_lane * mfma_m_lane>{},
number<mfma_m_lane * mfma_n_lane * m_pack>{},
number<mfma_n_lane * m_pack>{},
number<m_pack>{},
number<1>{}),
number<m_pack>{},
number<1>{});
const auto q_grad_dram_desc = transform_tensor_descriptor(
q_grad_dram_desc_0,
make_tuple(
make_merge_transform(make_tuple(M0,
number<mfma_m_lane>{},
number<m1_pack_num>{},
number<m_pack>{})),
make_merge_transform(make_tuple(number<N0>{}, number<mfma_n_lane>{}))),
make_tuple(sequence<0, 3, 2, 5>{}, sequence<1, 4>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
auto dq_acc_dram_view = make_tensor_view<address_space_enum::global,
memory_operation_enum::atomic_add>(
dq_acc_ptr, q_grad_dram_desc);
return pad_tensor_view(
dq_acc_dram_view,
make_tuple(number<kM0>{}, number<kQKHeaddim>{}),
sequence<kPadSeqLenQ, false>{}); // we have already padded the dram buffer
// in headdim direction
}
}
}();

38
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp
View File

@@ -20,11 +20,15 @@ struct BlockFmhaBwdConvertQGrad
static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
static constexpr index_t kBlockSize = Problem::kBlockSize;
static constexpr index_t kQKHeaddim = Problem::kQKHeaddim;
static constexpr index_t kGemm4WarpN = Problem::kGemm4WarpN;
static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ;
static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
static constexpr bool kIsAtomic32 = Problem::kIsAtomic32;
using QGradAccDataType = std::conditional_t<kIsAtomic32, AccDataType, QGradDataType>;
static constexpr index_t kAlignmentQGradAcc =
kPadHeadDimQ ? 1 : Policy::template GetAlignmentPostQGradAcc<Problem>();
@@ -40,7 +44,7 @@ struct BlockFmhaBwdConvertQGrad
QGradDramBlockWindowTmp& dq_dram_block_window_tmp) const
{
static_assert(
std::is_same_v<AccDataType,
std::is_same_v<QGradAccDataType,
remove_cvref_t<typename QGradAccDramBlockWindowTmp::DataType>> &&
std::is_same_v<QGradDataType,
remove_cvref_t<typename QGradDramBlockWindowTmp::DataType>>,
@@ -48,16 +52,32 @@ struct BlockFmhaBwdConvertQGrad
static_assert(kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}], "wrong!");
auto dq_acc_dram_window =
make_tile_window(dq_acc_dram_block_window_tmp.get_bottom_tensor_view(),
dq_acc_dram_block_window_tmp.get_window_lengths(),
dq_acc_dram_block_window_tmp.get_window_origin(),
Policy::template MakePostQGradDramTileDistribution<Problem>());
if constexpr(kIsAtomic32)
{
auto dq_acc_dram_window =
make_tile_window(dq_acc_dram_block_window_tmp.get_bottom_tensor_view(),
dq_acc_dram_block_window_tmp.get_window_lengths(),
dq_acc_dram_block_window_tmp.get_window_origin(),
Policy::template MakePostQGradDramTileDistribution<Problem>());
auto dq_acc = load_tile(dq_acc_dram_window);
const auto dq = cast_tile<QGradDataType>(dq_acc);
auto dq_acc = load_tile(dq_acc_dram_window);
const auto dq = cast_tile<QGradDataType>(dq_acc);
store_tile(dq_dram_block_window_tmp, dq);
store_tile(dq_dram_block_window_tmp, dq);
}
else
{
auto dq_acc_dram_window = make_tile_window(
dq_acc_dram_block_window_tmp.get_bottom_tensor_view(),
dq_acc_dram_block_window_tmp.get_window_lengths(),
dq_acc_dram_block_window_tmp.get_window_origin(),
Policy::template MakePostQGradAccAtomic16DramTileDistribution<Problem>());
auto shuffled_dq = make_static_distributed_tensor<QGradDataType>(
Policy::template MakePostQGradAtomic16DramTileDistribution<Problem>());
auto dq_acc = load_tile(dq_acc_dram_window);
shuffle_tile(shuffled_dq, dq_acc);
store_tile(dq_dram_block_window_tmp, shuffled_dq);
}
}
// Reduce + Convert

51
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp
View File

@@ -38,15 +38,16 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
static constexpr index_t kBlockSize = Problem::kBlockSize;
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 kK1 = BlockFmhaShape::kK1;
static constexpr index_t kK2 = BlockFmhaShape::kK2;
static constexpr index_t kK3 = BlockFmhaShape::kK3;
static constexpr index_t kK4 = BlockFmhaShape::kK4;
static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim;
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 kK1 = BlockFmhaShape::kK1;
static constexpr index_t kK2 = BlockFmhaShape::kK2;
static constexpr index_t kK3 = BlockFmhaShape::kK3;
static constexpr index_t kK4 = BlockFmhaShape::kK4;
static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim;
static constexpr index_t kGemm4WarpN = BlockFmhaShape::Gemm0WarpTile::at(ck_tile::number<1>{});
static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
@@ -54,6 +55,7 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
static constexpr auto BiasEnum = Problem::BiasEnum;
static constexpr bool kHasBiasGrad = Problem::kHasBiasGrad;
static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
static constexpr bool kIsAtomic32 = Problem::kIsAtomic32;
static constexpr bool kUseTrLoad = Problem::kUseTrLoad;
static_assert(!kUseTrLoad, "This pipeline does not use trload!");
@@ -468,14 +470,26 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
{0, 0},
Policy::template MakeShuffledBiasTileDistribution<Problem>());
// ----------------------------Loop write out------------------------------//
auto dq_dram_window = make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
dq_dram_block_window_tmp.get_window_lengths(),
{seqlen_q_start, 0});
using SPBlockTileType = decltype(gemm_0.MakeCBlockTile());
using SPGradBlockTileType = decltype(gemm_2.MakeCBlockTile());
using QGradBlockTileType = decltype(gemm_4.MakeCBlockTile());
// ----------------------------Loop write out------------------------------//
auto dq_dram_window = [&]() {
if constexpr(kIsAtomic32)
{
return make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
dq_dram_block_window_tmp.get_window_lengths(),
{seqlen_q_start, 0});
}
else
{
return make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
dq_dram_block_window_tmp.get_window_lengths(),
{seqlen_q_start, 0},
decltype(cast_tile<QGradDataType>(
QGradBlockTileType{}))::get_tile_distribution());
}
}();
index_t i_total_loops = 0;
index_t seqlen_q_step = seqlen_q_start;
@@ -750,7 +764,14 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
}
else
{
update_tile(dq_dram_window, dq_acc);
if constexpr(kIsAtomic32)
{
update_tile(dq_dram_window, dq_acc);
}
else
{
update_tile(dq_dram_window, cast_tile<QGradDataType>(dq_acc));
}
}
move_tile_window(dq_dram_window, {kM0, 0});

70
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp
View File

@@ -38,15 +38,16 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
static constexpr index_t kBlockSize = Problem::kBlockSize;
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 kK1 = BlockFmhaShape::kK1;
static constexpr index_t kK2 = BlockFmhaShape::kK2;
static constexpr index_t kK3 = BlockFmhaShape::kK3;
static constexpr index_t kK4 = BlockFmhaShape::kK4;
static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim;
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 kK1 = BlockFmhaShape::kK1;
static constexpr index_t kK2 = BlockFmhaShape::kK2;
static constexpr index_t kK3 = BlockFmhaShape::kK3;
static constexpr index_t kK4 = BlockFmhaShape::kK4;
static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
static constexpr index_t kVHeaddim = BlockFmhaShape::kVHeaddim;
static constexpr index_t kGemm4WarpN = BlockFmhaShape::Gemm0WarpTile::at(ck_tile::number<1>{});
static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
@@ -54,6 +55,7 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
static constexpr auto BiasEnum = Problem::BiasEnum;
static constexpr bool kHasBiasGrad = Problem::kHasBiasGrad;
static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
static constexpr bool kIsAtomic32 = Problem::kIsAtomic32;
static constexpr bool kUseTrLoad = Problem::kUseTrLoad;
static_assert(!kUseTrLoad, "This pipeline does not use trload!");
@@ -467,14 +469,26 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
{0, 0},
Policy::template MakeShuffledBiasTileDistribution<Problem>());
// ----------------------------Loop write out------------------------------//
auto dq_dram_window = make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
dq_dram_block_window_tmp.get_window_lengths(),
{seqlen_q_start, 0});
using SPBlockTileType = decltype(gemm_0.MakeCBlockTile());
using SPGradBlockTileType = decltype(gemm_2.MakeCBlockTile());
using QGradBlockTileType = decltype(gemm_4.MakeCBlockTile());
// ----------------------------Loop write out------------------------------//
auto dq_dram_window = [&]() {
if constexpr(kIsAtomic32)
{
return make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
dq_dram_block_window_tmp.get_window_lengths(),
{seqlen_q_start, 0});
}
else
{
return make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
dq_dram_block_window_tmp.get_window_lengths(),
{seqlen_q_start, 0},
decltype(cast_tile<QGradDataType>(
QGradBlockTileType{}))::get_tile_distribution());
}
}();
index_t i_total_loops = 0;
index_t seqlen_q_step = seqlen_q_start;
@@ -792,8 +806,20 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
}
else
{
update_tile(dq_dram_window, dq_acc);
if constexpr(kIsAtomic32)
{
update_tile(dq_dram_window, dq_acc);
}
else
{
buffer_store_fence();
update_tile_raw(dq_dram_window,
cast_tile<QGradDataType>(dq_acc),
number<-1>{},
bool_constant<false>{});
}
}
move_tile_window(dq_dram_window, {kM0, 0});
i_total_loops += 1;
@@ -1027,14 +1053,24 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc);
tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc);
}
if constexpr(kIsDeterministic)
{
store_tile(dq_dram_window, dq_acc);
}
else
{
update_tile(dq_dram_window, dq_acc);
if constexpr(kIsAtomic32)
{
update_tile(dq_dram_window, dq_acc);
}
else
{
buffer_store_fence();
update_tile_raw(dq_dram_window,
cast_tile<QGradDataType>(dq_acc),
number<-1>{},
bool_constant<false>{});
}
}
return make_tuple(dk_acc, dv_acc);

2
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp
View File

@@ -54,8 +54,10 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadKRKTRVR
static constexpr auto BiasEnum = Problem::BiasEnum;
static constexpr bool kHasBiasGrad = Problem::kHasBiasGrad;
static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
static constexpr bool kIsAtomic32 = Problem::kIsAtomic32;
static constexpr bool kUseTrLoad = Problem::kUseTrLoad;
static_assert(kUseTrLoad, "This pipeline uses trload!");
static_assert(kIsAtomic32, "This pipeline does not use atomic16!");
// last dimension vector length used to create tensor view(and decide buffer_load vector length)
// ... together with tensor distribution. tensor dist should able to overwrite this

2
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp
View File

@@ -56,8 +56,10 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR
static constexpr auto BiasEnum = Problem::BiasEnum;
static constexpr bool kHasBiasGrad = Problem::kHasBiasGrad;
static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
static constexpr bool kIsAtomic32 = Problem::kIsAtomic32;
static constexpr bool kUseTrLoad = Problem::kUseTrLoad;
static_assert(kUseTrLoad, "This pipeline uses trload!");
static_assert(kIsAtomic32, "This pipeline does not use atomic16!");
// last dimension vector length used to create tensor view(and decide buffer_load vector length)
// ... together with tensor distribution. tensor dist should able to overwrite this

50
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp
View File

@@ -741,6 +741,56 @@ struct BlockFmhaBwdPipelineDefaultPolicy
return dstr;
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakePostQGradAccAtomic16DramTileDistribution()
{
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t kMPerBlock = Problem::kM0;
constexpr index_t kNPerBlock = Problem::kQKHeaddim;
constexpr index_t mPack = 2; // for b16
constexpr index_t M1 = mPack;
constexpr index_t M0 = kMPerBlock / M1;
constexpr index_t N0 = kBlockSize / get_warp_size();
constexpr index_t N1 = get_warp_size() / M0;
constexpr index_t N2 = kNPerBlock / (N0 * N1);
return make_static_tile_distribution(
tile_distribution_encoding<sequence<>,
tuple<sequence<M0, M1>, sequence<N0, N1, N2>>,
tuple<sequence<2>, sequence<1, 2>>,
tuple<sequence<0>, sequence<0, 1>>,
sequence<2, 1>,
sequence<2, 1>>{});
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakePostQGradAtomic16DramTileDistribution()
{
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t kMPerBlock = Problem::kM0;
constexpr index_t kNPerBlock = Problem::kQKHeaddim;
constexpr index_t mPack = 2; // for b16
constexpr index_t M1 = mPack;
constexpr index_t M0 = kMPerBlock / M1;
constexpr index_t N0 = kBlockSize / get_warp_size();
constexpr index_t N1 = get_warp_size() / M0;
constexpr index_t N2 = kNPerBlock / (N0 * N1);
return make_static_tile_distribution(
tile_distribution_encoding<sequence<>,
tuple<sequence<M0, M1>, sequence<N0, N1, N2>>,
tuple<sequence<2>, sequence<1, 2>>,
tuple<sequence<0>, sequence<0, 1>>,
sequence<1, 2>,
sequence<1, 2>>{});
}
// these are for lds
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackQ()

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

@@ -25,6 +25,7 @@ template <typename QDataType_,
typename BlockFmhaShape_,
bool kIsGroupMode_,
bool kIsDeterministic_,
bool kIsAtomic32_,
typename FmhaMask_,
typename FmhaDropout_,
bool kUseTrLoad_,
@@ -54,6 +55,7 @@ struct BlockFmhaBwdPipelineProblem
static constexpr index_t kBlockSize = BlockFmhaShape::NumWarps * get_warp_size();
static constexpr bool kIsGroupMode = kIsGroupMode_;
static constexpr bool kIsDeterministic = kIsDeterministic_;
static constexpr bool kIsAtomic32 = kIsAtomic32_;
static constexpr bool kUseTrLoad = kUseTrLoad_;
// attributes from traits
@@ -99,8 +101,10 @@ template <typename AccDataType_,
index_t kM0_,
index_t kN0_,
index_t kQKHeaddim_,
index_t kGemm4WarpN_,
bool kIsGroupMode_,
bool kIsDeterministic_,
bool kIsAtomic32_,
typename Traits_>
struct BlockFmhaBwdConvertQGradPipelineProblem
{
@@ -115,8 +119,10 @@ struct BlockFmhaBwdConvertQGradPipelineProblem
static constexpr index_t kM0 = kM0_;
static constexpr index_t kN0 = kN0_;
static constexpr index_t kQKHeaddim = kQKHeaddim_;
static constexpr index_t kGemm4WarpN = kGemm4WarpN_;
static constexpr bool kIsGroupMode = kIsGroupMode_;
static constexpr bool kIsDeterministic = kIsDeterministic_;
static constexpr bool kIsAtomic32 = kIsAtomic32_;
// attributes from traits
static constexpr bool kPadSeqLenQ = Traits::kPadSeqLenQ;