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

formatting code

Browse Source
This commit is contained in:
shay-li77
2025-08-14 02:01:06 +08:00
parent 777c6e8f53
commit 178e0b71c9
11 changed files with 60 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
example/ck_tile/01_fmha/CMakeLists.txt
View File

@@ -138,7 +138,6 @@ endif()
# Allow comparing floating points directly in order to check sentinel values
list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-float-equal)
list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-float-equal)
# list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS --save-temps)
target_compile_options(${EXAMPLE_FMHA_FWD} PRIVATE ${EXAMPLE_FMHA_FWD_COMPILE_OPTIONS})
target_compile_options(${EXAMPLE_FMHA_BWD} PRIVATE ${EXAMPLE_FMHA_BWD_COMPILE_OPTIONS})

16
example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py
View File

@@ -286,11 +286,8 @@ class FmhaBwdDQDKDVApiTrait:
else : return '!t.is_deterministic && !t.is_atomic_fp32'
def get_kernel_group(self) -> Tuple[str]:
# kernel_group = [FMHA_BWD_CONVERT_DQ]
kernel_group = [FMHA_BWD_DOT_DO_O, FMHA_BWD_DQ_DQ_DV, FMHA_BWD_CONVERT_DQ]
# if self.deterministic == 't' or self.atomic32 == 't':
# kernel_group.append(FMHA_BWD_CONVERT_DQ)
return tuple(kernel_group)
return tuple(kernel_group)
class FmhaBwdApiPool:
def __init__(self, mask_impl):
@@ -568,10 +565,10 @@ class FmhaBwdDQDKDVKernel:
def get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype : str) -> Optional[dict]:
if dtype == 'fp16' or dtype == 'bf16':
return {
# '32' : [FmhaBwdDQDKDVTileSize( 32, 128, 32, 32, 32, 32, 64, 32, 32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
# "kr_ktr_vr_iglp", "kr_ktr_vr"],
# '64' : [FmhaBwdDQDKDVTileSize( 32, 128, 64, 32, 64, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
# "kr_ktr_vr_iglp", "kr_ktr_vr"],
'32' : [FmhaBwdDQDKDVTileSize( 32, 128, 32, 32, 32, 32, 64, 32, 32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
"kr_ktr_vr_iglp", "kr_ktr_vr"],
'64' : [FmhaBwdDQDKDVTileSize( 32, 128, 64, 32, 64, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
"kr_ktr_vr_iglp", "kr_ktr_vr"],
'128' : [FmhaBwdDQDKDVTileSize( 16, 128, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
"kr_ktr_vr_iglp", "kr_ktr_vr"],
'256' : [FmhaBwdDQDKDVTileSize( 16, 64, 256, 16, 256, 16, 32, 256, 256, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
@@ -590,9 +587,6 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) ->
if d == None:
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 dropout!= 'no' or deterministic == 't':
continue
tile = d[hdim_str][0]
ppl = d[hdim_str][1]

36
example/ck_tile/01_fmha/fmha_bwd.cpp
View File

@@ -95,7 +95,7 @@ auto create_args(int argc, char* argv[])
"0",
"if set to 1 will use multi-buffer reduction strategy for dq, atomic opeartion "
"will not be used")
.insert("atomic_fp32", "1", "if set to 0 will use atomic fp16/bf16(w/o convert_dq kernel)");
.insert("atomic_fp32", "1", "if set to 0 will use atomic fp16/bf16");
bool result = arg_parser.parse(argc, argv);
return std::make_tuple(result, arg_parser);
@@ -123,7 +123,7 @@ 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)
ck_tile::index_t get_bit_ceil(const ck_tile::index_t dim_value)
{
unsigned un = static_cast<unsigned>(dim_value);
un |= un >> 1;
@@ -215,9 +215,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
// 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;
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,
@@ -230,10 +229,12 @@ bool run(const ck_tile::ArgParser& arg_parser)
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;
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>;
@@ -524,7 +525,6 @@ bool run(const ck_tile::ArgParser& arg_parser)
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 * 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)
{
@@ -551,7 +551,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
dk_buf.GetDeviceBuffer(),
dv_buf.GetDeviceBuffer(),
dbias_buf.GetDeviceBuffer(),
dq_acc_ptr,
dq_acc_buf.GetDeviceBuffer(),
seqstart_q.GetDeviceBuffer(),
seqstart_k.GetDeviceBuffer(),
nullptr,
@@ -1022,22 +1022,26 @@ int main(int argc, char* argv[])
return -1;
const std::string data_type = arg_parser.get_str("prec");
const bool atomic_fp32 = arg_parser.get_bool("atomic_fp32");
const bool atomic_fp32 = arg_parser.get_bool("atomic_fp32");
if(data_type == "fp16")
{
if (atomic_fp32) {
if(atomic_fp32)
{
return run<FmhaBwdFp16>(arg_parser) ? 0 : -2;
}
else {
else
{
return run<FmhaBwdFp16, false>(arg_parser) ? 0 : -2;
}
}
else if(data_type == "bf16")
{
if (atomic_fp32) {
if(atomic_fp32)
{
return run<FmhaBwdBf16>(arg_parser) ? 0 : -2;
}
else {
else
{
return run<FmhaBwdBf16, false>(arg_parser) ? 0 : -2;
}
}

2
example/ck_tile/01_fmha/fmha_bwd.hpp
View File

@@ -423,7 +423,7 @@ template <ck_tile::index_t HDim_,
bool kPadS_,
bool kPadD_,
bool kIsDeterministic_,
bool kAtomic32_>
bool kAtomic32_ = true>
struct fmha_bwd_convert_dq_traits_
{
static constexpr ck_tile::index_t HDim = HDim_;

14
example/ck_tile/01_fmha/script/smoke_test_bwd.sh
View File

@@ -16,13 +16,14 @@ for bias in "n" "a" ; do
for dbias in 0 ; do
for p_drop in 0.0 0.2 ; do
for deterministic in 0 ; do
for atomic_fp32 in 0 1 ; do
$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -deterministic=$deterministic -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -deterministic=$deterministic -atomic_fp32=$atomic_fp32 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
done
done
@@ -32,4 +33,5 @@ done
done
done
done
done
set +x

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

@@ -267,7 +267,7 @@ struct buffer_view<address_space_enum::global,
CK_TILE_HOST_DEVICE constexpr buffer_view(T* p_data, BufferSizeType buffer_size)
: p_data_{p_data},
buffer_size_{buffer_size},
buffer_size_{buffer_size / PackedSize},
cached_buf_res_{0},
invalid_element_value_{0}
{

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

@@ -64,9 +64,9 @@ 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 kIsAtomic32 = FmhaPipeline::kIsAtomic32;
using QGradAccDataType = std::conditional_t<kIsAtomic32, AccDataType, QDataType>;
using QGradAccDataType = std::conditional_t<kIsAtomic32, AccDataType, QDataType>;
// clang-format off
template <typename T> struct t2s;
@@ -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" : (kIsAtomic32 ? "_atomic32" : "_atomic16") );
(kIsStoreRandval ? "_storerandval" : "" ) + (kIsDeterministic ? "_deterministic" : (kIsAtomic32 ? "_atomic32" : "_atomic16"));
#undef _SS_
#undef _TS_
// clang-format on
@@ -1132,7 +1132,7 @@ 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;
long_index_t dq_acc_start = 0;
if constexpr(kIsAtomic32)
{
dq_acc_start = kargs.seqstart_q_ptr[i_batch];

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

@@ -28,7 +28,7 @@ struct BlockFmhaBwdConvertQGrad
static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
static constexpr bool kIsAtomic32 = Problem::kIsAtomic32;
using QGradAccDataType = std::conditional_t<kIsAtomic32, AccDataType, QGradDataType>;
using QGradAccDataType = std::conditional_t<kIsAtomic32, AccDataType, QGradDataType>;
static constexpr index_t kAlignmentQGradAcc =
kPadHeadDimQ ? 1 : Policy::template GetAlignmentPostQGradAcc<Problem>();
@@ -51,28 +51,30 @@ struct BlockFmhaBwdConvertQGrad
"wrong!");
static_assert(kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}], "wrong!");
if constexpr(kIsAtomic32) {
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>());
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);
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>(
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);
auto dq_acc = load_tile(dq_acc_dram_window);
shuffle_tile(shuffled_dq, dq_acc);
store_tile(dq_dram_block_window_tmp, shuffled_dq);
}

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

@@ -69,7 +69,6 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
kPadHeadDimV ? 1 : Policy::template GetAlignmentV<Problem>();
static constexpr index_t kAlignmentOGrad =
kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad<Problem>();
// static constexpr index_t kAlignmentQGrad = kPadHeadDimQ ? 1 : Policy::template GetAlignmentQGrad<Problem>();
static constexpr index_t kAlignmentQGrad = 1;
static constexpr index_t kAlignmentKGrad =
kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad<Problem>();

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

@@ -69,8 +69,7 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
kPadHeadDimV ? 1 : Policy::template GetAlignmentV<Problem>();
static constexpr index_t kAlignmentOGrad =
kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad<Problem>();
static constexpr index_t kAlignmentQGrad = kPadHeadDimQ ? 1 : Policy::template GetAlignmentQGrad<Problem>();
// static constexpr index_t kAlignmentQGrad = 1;
static constexpr index_t kAlignmentQGrad = 1;
static constexpr index_t kAlignmentKGrad =
kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad<Problem>();
static constexpr index_t kAlignmentVGrad =
@@ -475,7 +474,6 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
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});
// auto dq_dram_window = dq_dram_block_window_tmp;
using SPBlockTileType = decltype(gemm_0.MakeCBlockTile());
using SPGradBlockTileType = decltype(gemm_2.MakeCBlockTile());

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

@@ -286,25 +286,6 @@ struct BlockFmhaBwdPipelineDefaultPolicy
return kVecLoad;
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentQGrad()
{
using QGradDataType = remove_cvref_t<typename Problem::QGradDataType>;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t kMNPerBlock = Problem::BlockFmhaShape::kM0;
constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kQKHeaddim;
constexpr index_t kMaxVecLoad = 16 / sizeof(QGradDataType);
constexpr index_t kMinVecLoad = 4 / sizeof(QGradDataType);
constexpr index_t total_pixels = kMNPerBlock * kKPerBlock / kBlockSize;
constexpr index_t kVecLoad = ((total_pixels / kMaxVecLoad) >= kMinVecLoad)
? kMaxVecLoad
: (total_pixels / kMinVecLoad);
return kVecLoad;
}
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentKGrad()
{
@@ -631,13 +612,12 @@ struct BlockFmhaBwdPipelineDefaultPolicy
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 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<>,
@@ -657,13 +637,12 @@ struct BlockFmhaBwdPipelineDefaultPolicy
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 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<>,
@@ -737,27 +716,6 @@ struct BlockFmhaBwdPipelineDefaultPolicy
return 16 / sizeof(GemmDataType);
}
// template <index_t mSize, index_t nSize, index_t mPack>
// CK_TILE_HOST_DEVICE static constexpr auto MakeQGradDramBlockDescriptor()
// {
// constexpr auto q_grad_dram_desc_0 = make_naive_tensor_descriptor(
// make_tuple(number<mSize / mPack>{}, number<nSize>{}, number<mPack>{}),
// make_tuple(number<nSize * mPack>{}, number<mPack>{}, number<1>{}),
// number<mPack>{},
// number<1>{}
// );
// constexpr auto q_grad_dram_desc = transform_tensor_descriptor(
// q_grad_dram_desc_0,
// make_tuple(
// make_merge_transform(make_tuple(number<mSize / mPack>{}, number<mPack>{})),
// make_pass_through_transform(number<nSize>{})),
// make_tuple(sequence<0, 2>{}, sequence<1>{}),
// make_tuple(sequence<0>{}, sequence<1>{})
// );
// return q_grad_dram_desc;
// }
template <index_t MNPerBlock, index_t KPerBlock, index_t KPack>
CK_TILE_HOST_DEVICE static constexpr auto MakeXLdsBlockDescriptor()
{