mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-14 02:57:45 +00:00
atomic16 base impl
This commit is contained in:
shay-li77
@@ -12,6 +12,7 @@ FWD_DTYPE_MAP = {
|
||||
|
||||
BWD_DTYPE_MAP = {
|
||||
"fp16": "FmhaBwdFp16",
|
||||
"bf16": "FmhaBwdBf16"
|
||||
}
|
||||
|
||||
MASK_IMPL = {
|
||||
|
||||
@@ -591,7 +591,7 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) ->
|
||||
continue
|
||||
for hdim_str, mode, mask, bias, dbias, dropout, spad, skpad, dpad, dvpad, deterministic, atomic32 in itertools.product(d.keys(), MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], DROPOUT_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"]):
|
||||
# for debug(xiangxli)
|
||||
if bias != 'no' or mode == 'group' or dropout!= 'no' or dpad == "t" or dvpad == "t" or deterministic == 't':
|
||||
if bias != 'no' or dropout!= 'no' or deterministic == 't':
|
||||
continue
|
||||
|
||||
tile = d[hdim_str][0]
|
||||
|
||||
@@ -123,6 +123,18 @@ auto get_elimit<FmhaBwdBf16>(ck_tile::index_t hdim_q, ck_tile::index_t hdim_v)
|
||||
return ck_tile::make_tuple(rtol, atol);
|
||||
}
|
||||
|
||||
ck_tile::index_t get_bit_ceil(const ck_tile::index_t& dim_value)
|
||||
{
|
||||
unsigned un = static_cast<unsigned>(dim_value);
|
||||
un |= un >> 1;
|
||||
un |= un >> 2;
|
||||
un |= un >> 4;
|
||||
un |= un >> 8;
|
||||
un |= un >> 16;
|
||||
un++;
|
||||
return static_cast<ck_tile::index_t>(un);
|
||||
}
|
||||
|
||||
template <typename DataTypeConfig, bool IsAtomic32=true>
|
||||
bool run(const ck_tile::ArgParser& arg_parser)
|
||||
{
|
||||
@@ -201,6 +213,12 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
bool deterministic = arg_parser.get_bool("deterministic");
|
||||
bool atomic_fp32 = arg_parser.get_bool("atomic_fp32");
|
||||
|
||||
// for dq_acc padding in atomic16
|
||||
constexpr ck_tile::index_t seqlen_dq_acc_tile_size = 16;
|
||||
const ck_tile::index_t hdim_q_pad = get_bit_ceil(hdim_q);
|
||||
const ck_tile::index_t hdim_q_dq_acc = atomic_fp32 ? hdim_q : hdim_q_pad;
|
||||
|
||||
|
||||
ck_tile::stream_config stream_config{nullptr,
|
||||
true,
|
||||
/* log_level = */ (kname ? 1 : 0),
|
||||
@@ -211,6 +229,13 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
const auto seqstart_q_host = generate_seqstarts(mode, batch, seqlen_q);
|
||||
const auto seqstart_k_host = generate_seqstarts(mode, batch, seqlen_k);
|
||||
|
||||
auto seqstart_dq_acc_host = std::vector<int32_t>(seqstart_q_host.size(), 0);
|
||||
for (int i = 0; i < batch; ++i) {
|
||||
auto cur_seqlen_q = seqstart_q_host[i+1] - seqstart_q_host[i];
|
||||
auto cur_seqlen_dq_acc = ck_tile::integer_least_multiple(cur_seqlen_q, seqlen_dq_acc_tile_size);
|
||||
seqstart_dq_acc_host[i+1] = seqstart_dq_acc_host[i] + cur_seqlen_dq_acc;
|
||||
}
|
||||
|
||||
using TypeConfig = FmhaBwdTypeConfig<DataTypeConfig>;
|
||||
|
||||
using QDataType = typename TypeConfig::QDataType;
|
||||
@@ -286,6 +311,13 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
(mode == mode_enum::batch ? seqlen_q : seqstart_q_host.back());
|
||||
const ck_tile::index_t shape_seqlen_k =
|
||||
(mode == mode_enum::batch ? seqlen_k : seqstart_k_host.back());
|
||||
const ck_tile::index_t shape_seqlen_dq_acc_batch_mode =
|
||||
atomic_fp32 ? seqlen_q : ck_tile::integer_least_multiple(seqlen_q, seqlen_dq_acc_tile_size);
|
||||
const ck_tile::index_t shape_seqlen_dq_acc_group_mode =
|
||||
atomic_fp32 ? seqstart_q_host.back() : seqstart_dq_acc_host.back();
|
||||
const ck_tile::index_t shape_seqlen_dq_acc =
|
||||
(mode == mode_enum::batch ? shape_seqlen_dq_acc_batch_mode
|
||||
: shape_seqlen_dq_acc_group_mode);
|
||||
const ck_tile::index_t kN0 = (hdim_q <= 128) ? 128 : 64;
|
||||
const ck_tile::index_t nsplits =
|
||||
deterministic ? ck_tile::integer_divide_ceil(max_seqlen_k, kN0) : 1;
|
||||
@@ -326,15 +358,13 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
use_dbias
|
||||
? get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, max_seqlen_k)
|
||||
: std::array<ck_tile::index_t, 4>{1, 1, 1, 1} /* dummy shape for simplifying code */);
|
||||
|
||||
constexpr ck_tile::index_t dq_acc_seqlen_pad_size = 100; // hyper-parameter for atomic16, need to be greater than the kernel tiling size.
|
||||
const ck_tile::index_t dq_acc_seqlen_q = IsAtomic32 ? shape_seqlen_q : shape_seqlen_q + dq_acc_seqlen_pad_size;
|
||||
|
||||
|
||||
|
||||
bool dq_acc_perm = i_perm || !atomic_fp32; // need to permute for atomic16
|
||||
ck_tile::HostTensor<QGradAccDataType> dq_acc_host(dq_acc_perm ? std::array<ck_tile::index_t,
|
||||
5>{nsplits, shape_batch, nhead, dq_acc_seqlen_q, hdim_q}
|
||||
5>{nsplits, shape_batch, nhead, shape_seqlen_dq_acc, hdim_q_dq_acc}
|
||||
: std::array<ck_tile::index_t,
|
||||
5>{nsplits, shape_batch, dq_acc_seqlen_q, nhead, hdim_q});
|
||||
5>{nsplits, shape_batch, shape_seqlen_dq_acc, nhead, hdim_q_dq_acc});
|
||||
|
||||
if(init_method == 0)
|
||||
{
|
||||
@@ -394,6 +424,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
ck_tile::DeviceMem dbias_buf(dbias_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem seqstart_dq_acc(seqstart_dq_acc_host.size() * sizeof(int32_t));
|
||||
ck_tile::DeviceMem drop_seed_buf(drop_prefs ? sizeof(uint64_t) : 0);
|
||||
ck_tile::DeviceMem drop_offset_buf(drop_prefs ? sizeof(uint64_t) : 0);
|
||||
ck_tile::DeviceMem alibi_slope_buf(alibi_slope_host.get_element_space_size_in_bytes());
|
||||
@@ -406,6 +437,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
do_buf.ToDevice(do_host.data());
|
||||
seqstart_q.ToDevice(seqstart_q_host.data());
|
||||
seqstart_k.ToDevice(seqstart_k_host.data());
|
||||
seqstart_dq_acc.ToDevice(seqstart_dq_acc_host.data());
|
||||
drop_seed_buf.ToDevice(drop_prefs ? &drop_seed : nullptr);
|
||||
drop_offset_buf.ToDevice(drop_prefs ? &drop_offset : nullptr);
|
||||
alibi_slope_buf.ToDevice(alibi_slope_host.data());
|
||||
@@ -464,7 +496,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
const ck_tile::index_t stride_dk = (i_perm ? hdim_q : nhead * hdim_q);
|
||||
const ck_tile::index_t stride_dv = (i_perm ? hdim_v : nhead * hdim_v);
|
||||
const ck_tile::index_t stride_dbias = (i_perm ? max_seqlen_k : nhead * max_seqlen_k);
|
||||
const ck_tile::index_t stride_dq_acc = (dq_acc_perm ? hdim_q : nhead * hdim_q);
|
||||
const ck_tile::index_t stride_dq_acc = (dq_acc_perm ? hdim_q_dq_acc : nhead * hdim_q_dq_acc);
|
||||
// setup nhead_stride_* arguments
|
||||
const ck_tile::index_t nhead_stride_q = (i_perm ? shape_seqlen_q * hdim_q : hdim_q);
|
||||
const ck_tile::index_t nhead_stride_k = (i_perm ? shape_seqlen_k * hdim_q : hdim_q);
|
||||
@@ -476,7 +508,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
const ck_tile::index_t nhead_stride_lsed = shape_seqlen_q;
|
||||
const ck_tile::index_t nhead_stride_dbias =
|
||||
(i_perm ? shape_seqlen_q * max_seqlen_k : max_seqlen_k);
|
||||
const ck_tile::index_t nhead_stride_dq_acc = (dq_acc_perm ? dq_acc_seqlen_q * hdim_q : hdim_q);
|
||||
const ck_tile::index_t nhead_stride_dq_acc = (dq_acc_perm ? shape_seqlen_dq_acc * hdim_q_dq_acc : hdim_q_dq_acc);
|
||||
// setup batch_stride_* arguments
|
||||
const ck_tile::index_t batch_stride_q = (nhead * shape_seqlen_q * hdim_q);
|
||||
const ck_tile::index_t batch_stride_k = (nhead_k * shape_seqlen_k * hdim_q);
|
||||
@@ -489,9 +521,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
const ck_tile::index_t batch_stride_dk = (nhead * shape_seqlen_k * hdim_q);
|
||||
const ck_tile::index_t batch_stride_dv = (nhead * shape_seqlen_k * hdim_v);
|
||||
const ck_tile::index_t batch_stride_dbias = (nhead * shape_seqlen_q * max_seqlen_k);
|
||||
const ck_tile::index_t batch_stride_dq_acc = (nhead * dq_acc_seqlen_q * hdim_q);
|
||||
const ck_tile::index_t batch_stride_dq_acc = (nhead * shape_seqlen_dq_acc * hdim_q_dq_acc);
|
||||
const ck_tile::index_t split_stride_dq_acc =
|
||||
(shape_batch * nhead * dq_acc_seqlen_q * hdim_q);
|
||||
(shape_batch * nhead * shape_seqlen_dq_acc * hdim_q_dq_acc);
|
||||
auto dq_acc_ptr = dq_acc_buf.GetDeviceBuffer();
|
||||
const auto drop_seed_offset = [&]() -> decltype(fmha_bwd_args::drop_seed_offset) {
|
||||
if(drop_prefs)
|
||||
@@ -523,6 +555,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
|
||||
seqstart_q.GetDeviceBuffer(),
|
||||
seqstart_k.GetDeviceBuffer(),
|
||||
nullptr,
|
||||
seqstart_dq_acc.GetDeviceBuffer(),
|
||||
shape_seqlen_q,
|
||||
shape_seqlen_k,
|
||||
batch,
|
||||
|
||||
@@ -92,6 +92,7 @@ struct fmha_bwd_args
|
||||
const void* seqstart_q_ptr;
|
||||
const void* seqstart_k_ptr;
|
||||
const void* seqlen_k_ptr;
|
||||
const void* seqstart_dq_acc_ptr;
|
||||
ck_tile::index_t seqlen_q;
|
||||
ck_tile::index_t seqlen_k;
|
||||
ck_tile::index_t batch;
|
||||
@@ -173,6 +174,7 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
|
||||
args.seqstart_q_ptr,
|
||||
args.seqstart_k_ptr,
|
||||
args.seqlen_k_ptr,
|
||||
args.seqstart_dq_acc_ptr,
|
||||
args.hdim_q,
|
||||
args.hdim_v,
|
||||
args.nhead_q,
|
||||
@@ -325,6 +327,7 @@ auto fmha_bwd_convert_dq_create_kargs_and_grids(fmha_bwd_args args)
|
||||
args.dq_ptr,
|
||||
args.seqstart_q_ptr,
|
||||
args.seqstart_k_ptr,
|
||||
args.seqstart_dq_acc_ptr,
|
||||
args.hdim_q,
|
||||
args.stride_dq,
|
||||
args.stride_dq_acc,
|
||||
|
||||
@@ -106,7 +106,7 @@ struct FmhaBwdDQDKDVKernel
|
||||
("o" + _TS_(kBlockPerCu) + "_") + _SS_(FmhaPipeline::name) + (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" );
|
||||
(kIsStoreRandval ? "_storerandval" : "" ) + (kIsDeterministic ? "_deterministic" : (kIsAtomic32 ? "_atomic32" : "_atomic16") );
|
||||
#undef _SS_
|
||||
#undef _TS_
|
||||
// clang-format on
|
||||
@@ -264,6 +264,11 @@ struct FmhaBwdDQDKDVKernel
|
||||
ck_tile::index_t split_stride_dq_acc = 0;
|
||||
};
|
||||
|
||||
struct FmhaBwdAtomic16GroupModeKargs
|
||||
{
|
||||
const int32_t* seqstart_dq_acc_ptr;
|
||||
};
|
||||
|
||||
struct FmhaBwdBatchModeKargs
|
||||
: FmhaBwdCommonKargs,
|
||||
std::conditional_t<BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS,
|
||||
@@ -296,7 +301,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;
|
||||
@@ -732,6 +738,7 @@ struct FmhaBwdDQDKDVKernel
|
||||
const void* seqstart_q_ptr,
|
||||
const void* seqstart_k_ptr,
|
||||
const void* seqlen_k_ptr,
|
||||
const void* seqstart_dq_acc_ptr,
|
||||
ck_tile::index_t hdim_q,
|
||||
ck_tile::index_t hdim_v,
|
||||
ck_tile::index_t num_head_q,
|
||||
@@ -803,6 +810,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)};
|
||||
@@ -858,6 +866,11 @@ struct FmhaBwdDQDKDVKernel
|
||||
kargs.split_stride_dq_acc = split_stride_dq_acc;
|
||||
}
|
||||
|
||||
if constexpr(!kIsAtomic32)
|
||||
{
|
||||
kargs.seqstart_dq_acc_ptr = reinterpret_cast<const int32_t*>(seqstart_dq_acc_ptr);
|
||||
}
|
||||
|
||||
return kargs;
|
||||
}
|
||||
|
||||
@@ -879,6 +892,7 @@ struct FmhaBwdDQDKDVKernel
|
||||
const void* seqstart_q_ptr,
|
||||
const void* seqstart_k_ptr,
|
||||
const void* seqlen_k_ptr,
|
||||
const void* seqstart_dq_acc_ptr,
|
||||
ck_tile::index_t hdim_q,
|
||||
ck_tile::index_t hdim_v,
|
||||
ck_tile::index_t num_head_q,
|
||||
@@ -928,6 +942,7 @@ struct FmhaBwdDQDKDVKernel
|
||||
seqstart_q_ptr,
|
||||
seqstart_k_ptr,
|
||||
seqlen_k_ptr,
|
||||
seqstart_dq_acc_ptr,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
num_head_q,
|
||||
@@ -980,6 +995,7 @@ struct FmhaBwdDQDKDVKernel
|
||||
const void* seqstart_q_ptr,
|
||||
const void* seqstart_k_ptr,
|
||||
const void* seqlen_k_ptr,
|
||||
const void* seqstart_dq_acc_ptr,
|
||||
ck_tile::index_t hdim_q,
|
||||
ck_tile::index_t hdim_v,
|
||||
ck_tile::index_t num_head_q,
|
||||
@@ -1029,6 +1045,7 @@ struct FmhaBwdDQDKDVKernel
|
||||
seqstart_q_ptr,
|
||||
seqstart_k_ptr,
|
||||
seqlen_k_ptr,
|
||||
seqstart_dq_acc_ptr,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
num_head_q,
|
||||
@@ -1115,13 +1132,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.seqstart_dq_acc_ptr[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)
|
||||
@@ -1331,42 +1357,69 @@ struct FmhaBwdDQDKDVKernel
|
||||
static_cast<long_index_t>(i_nhead_) * kargs.nhead_stride_dq_acc +
|
||||
batch_offset_dq_acc;
|
||||
|
||||
|
||||
auto dq_acc_dram = [&]() {
|
||||
index_t dq_acc_m = (kargs.seqlen_q + FmhaPipeline::kM0 - 1) / FmhaPipeline::kM0 * FmhaPipeline::kM0;
|
||||
constexpr auto dq_acc_n = FmhaPipeline::kQKHeaddim;
|
||||
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;
|
||||
if constexpr(kIsAtomic32)
|
||||
{
|
||||
const auto dq_acc_dram_naive =
|
||||
make_naive_tensor_view<address_space_enum::global,
|
||||
memory_operation_enum::atomic_add>(
|
||||
dq_acc_ptr,
|
||||
make_tuple(kargs.seqlen_q, kargs.hdim_q),
|
||||
make_tuple(kargs.stride_dq_acc, 1),
|
||||
number<FmhaPipeline::kAlignmentQGrad>{},
|
||||
number<1>{});
|
||||
|
||||
index_t M0 = dq_acc_m / (mfma_m1_per_lane * mfma_m_lane);
|
||||
constexpr index_t N0 = dq_acc_n / mfma_n_lane;
|
||||
return pad_tensor_view(dq_acc_dram_naive,
|
||||
make_tuple(number<FmhaPipeline::kM0>{},
|
||||
number<FmhaPipeline::kQKHeaddim>{}),
|
||||
sequence<kPadSeqLenQ, kPadHeadDimQ>{});
|
||||
}
|
||||
else
|
||||
{
|
||||
index_t dq_acc_m =
|
||||
ck_tile::integer_least_multiple(kargs.seqlen_q, FmhaPipeline::kM0);
|
||||
constexpr auto dq_acc_n = FmhaPipeline::kQKHeaddim;
|
||||
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;
|
||||
|
||||
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>{}
|
||||
);
|
||||
index_t M0 = dq_acc_m / (mfma_m1_per_lane * mfma_m_lane);
|
||||
constexpr index_t N0 = dq_acc_n / mfma_n_lane;
|
||||
|
||||
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);
|
||||
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>{}),
|
||||
@@ -1879,6 +1932,7 @@ struct FmhaBwdConvertQGradKernel
|
||||
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;
|
||||
@@ -1904,7 +1958,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
|
||||
@@ -1939,6 +1993,11 @@ struct FmhaBwdConvertQGradKernel
|
||||
ck_tile::index_t split_stride_dq_acc = 0;
|
||||
};
|
||||
|
||||
struct FmhaBwdConvertQGradAtomic16GroupModeKargs
|
||||
{
|
||||
const int32_t* seqstart_dq_acc_ptr;
|
||||
};
|
||||
|
||||
struct FmhaBwdConvertQGradBatchModeKargs
|
||||
: FmhaBwdConvertQGradCommonKargs,
|
||||
std::conditional_t<kIsDeterministic,
|
||||
@@ -1953,7 +2012,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;
|
||||
@@ -2005,6 +2067,7 @@ struct FmhaBwdConvertQGradKernel
|
||||
void* dq_ptr,
|
||||
const void* seqstart_q_ptr,
|
||||
const void* seqstart_k_ptr,
|
||||
const void* seqstart_dq_acc_ptr,
|
||||
ck_tile::index_t hdim_q,
|
||||
ck_tile::index_t stride_dq,
|
||||
ck_tile::index_t stride_dq_acc,
|
||||
@@ -2021,7 +2084,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)};
|
||||
|
||||
@@ -2030,6 +2094,11 @@ struct FmhaBwdConvertQGradKernel
|
||||
kargs.split_stride_dq_acc = split_stride_dq_acc;
|
||||
}
|
||||
|
||||
if constexpr(!kIsAtomic32)
|
||||
{
|
||||
kargs.seqstart_dq_acc_ptr = reinterpret_cast<const int32_t*>(seqstart_dq_acc_ptr);
|
||||
}
|
||||
|
||||
return kargs;
|
||||
}
|
||||
|
||||
@@ -2065,8 +2134,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.seqstart_dq_acc_ptr[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;
|
||||
@@ -2121,47 +2199,62 @@ struct FmhaBwdConvertQGradKernel
|
||||
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>{});
|
||||
index_t dq_acc_m = (kargs.seqlen_q + kM0 - 1) / kM0 * kM0;
|
||||
constexpr auto dq_acc_n = kQKHeaddim;
|
||||
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;
|
||||
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
|
||||
{
|
||||
index_t dq_acc_m = ck_tile::integer_least_multiple(kargs.seqlen_q, kM0);
|
||||
constexpr auto dq_acc_n = kQKHeaddim;
|
||||
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;
|
||||
|
||||
index_t M0 = dq_acc_m / (mfma_m1_per_lane * mfma_m_lane);
|
||||
constexpr index_t N0 = dq_acc_n / mfma_n_lane;
|
||||
index_t M0 = dq_acc_m / (mfma_m1_per_lane * mfma_m_lane);
|
||||
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_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);
|
||||
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);
|
||||
}
|
||||
}
|
||||
}();
|
||||
|
||||
|
||||
@@ -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 kPadSeqLenQ = Problem::kPadSeqLenQ;
|
||||
@@ -752,17 +753,17 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
|
||||
{
|
||||
store_tile(dq_dram_window, dq_acc);
|
||||
}
|
||||
else
|
||||
{
|
||||
if constexpr(kIsAtomic32)
|
||||
{
|
||||
update_tile(dq_dram_window, dq_acc);
|
||||
}
|
||||
else
|
||||
{
|
||||
// update_tile(dq_dram_window, cast_tile<QDataType>(dq_acc));
|
||||
if constexpr(kIsAtomic32)
|
||||
{
|
||||
update_tile(dq_dram_window, dq_acc);
|
||||
}
|
||||
else
|
||||
{
|
||||
update_tile(dq_dram_window, cast_tile<QDataType>(dq_acc));
|
||||
}
|
||||
}
|
||||
}
|
||||
move_tile_window(dq_dram_window, {kM0, 0});
|
||||
|
||||
i_total_loops += 1;
|
||||
|
||||
@@ -1035,8 +1035,6 @@ 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);
|
||||
}
|
||||
// auto wgid = blockIdx.x + blockIdx.y * gridDim.x + gridDim.x * gridDim.y * blockIdx.z;
|
||||
// auto tid = (threadIdx.z * (blockDim.x * blockDim.y)) + (threadIdx.y * blockDim.x) + threadIdx.x;
|
||||
if constexpr(kIsDeterministic)
|
||||
{
|
||||
store_tile(dq_dram_window, dq_acc);
|
||||
@@ -1050,10 +1048,6 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
|
||||
else
|
||||
{
|
||||
update_tile(dq_dram_window, cast_tile<QDataType>(dq_acc));
|
||||
// if (wgid ==0 && tid==0) {
|
||||
// printf("atomic 16 update tile \n");
|
||||
|
||||
// }
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user