diff --git a/example/ck_tile/01_fmha/CMakeLists.txt b/example/ck_tile/01_fmha/CMakeLists.txt index 8170717d52..1b004ec100 100644 --- a/example/ck_tile/01_fmha/CMakeLists.txt +++ b/example/ck_tile/01_fmha/CMakeLists.txt @@ -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}) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py index 9525299bc1..a92ce4924c 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -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] diff --git a/example/ck_tile/01_fmha/fmha_bwd.cpp b/example/ck_tile/01_fmha/fmha_bwd.cpp index 647e3e4094..02bc745185 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.cpp +++ b/example/ck_tile/01_fmha/fmha_bwd.cpp @@ -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(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(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(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; @@ -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(arg_parser) ? 0 : -2; } - else { + else + { return run(arg_parser) ? 0 : -2; } } else if(data_type == "bf16") { - if (atomic_fp32) { + if(atomic_fp32) + { return run(arg_parser) ? 0 : -2; } - else { + else + { return run(arg_parser) ? 0 : -2; } } diff --git a/example/ck_tile/01_fmha/fmha_bwd.hpp b/example/ck_tile/01_fmha/fmha_bwd.hpp index 60f316ee7c..cc34f27848 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd.hpp @@ -423,7 +423,7 @@ template + bool kAtomic32_ = true> struct fmha_bwd_convert_dq_traits_ { static constexpr ck_tile::index_t HDim = HDim_; diff --git a/example/ck_tile/01_fmha/script/smoke_test_bwd.sh b/example/ck_tile/01_fmha/script/smoke_test_bwd.sh index 5ba3425e26..2cc35edbd8 100755 --- a/example/ck_tile/01_fmha/script/smoke_test_bwd.sh +++ b/example/ck_tile/01_fmha/script/smoke_test_bwd.sh @@ -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 diff --git a/include/ck_tile/core/tensor/buffer_view.hpp b/include/ck_tile/core/tensor/buffer_view.hpp index a31aa3dd7e..13b038bc48 100644 --- a/include/ck_tile/core/tensor/buffer_view.hpp +++ b/include/ck_tile/core/tensor/buffer_view.hpp @@ -267,7 +267,7 @@ struct buffer_view; + using QGradAccDataType = std::conditional_t; // clang-format off template 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::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]; diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp index d4e45e6e02..fad323f344 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp @@ -28,7 +28,7 @@ struct BlockFmhaBwdConvertQGrad static constexpr bool kIsDeterministic = Problem::kIsDeterministic; static constexpr bool kIsAtomic32 = Problem::kIsAtomic32; - using QGradAccDataType = std::conditional_t; + using QGradAccDataType = std::conditional_t; static constexpr index_t kAlignmentQGradAcc = kPadHeadDimQ ? 1 : Policy::template GetAlignmentPostQGradAcc(); @@ -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()); + dq_acc_dram_block_window_tmp.get_window_lengths(), + dq_acc_dram_block_window_tmp.get_window_origin(), + Policy::template MakePostQGradDramTileDistribution()); auto dq_acc = load_tile(dq_acc_dram_window); const auto dq = cast_tile(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()); - auto shuffled_dq = make_static_distributed_tensor( + 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()); + auto shuffled_dq = make_static_distributed_tensor( Policy::template MakePostQGradAtomic16DramTileDistribution()); - 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); } diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp index e68d42e8ab..4ebf1c0109 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp @@ -69,7 +69,6 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); static constexpr index_t kAlignmentOGrad = kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); - // static constexpr index_t kAlignmentQGrad = kPadHeadDimQ ? 1 : Policy::template GetAlignmentQGrad(); static constexpr index_t kAlignmentQGrad = 1; static constexpr index_t kAlignmentKGrad = kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp index 334bf41108..b92c4a41f7 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp @@ -69,8 +69,7 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP kPadHeadDimV ? 1 : Policy::template GetAlignmentV(); static constexpr index_t kAlignmentOGrad = kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad(); - static constexpr index_t kAlignmentQGrad = kPadHeadDimQ ? 1 : Policy::template GetAlignmentQGrad(); - // static constexpr index_t kAlignmentQGrad = 1; + static constexpr index_t kAlignmentQGrad = 1; static constexpr index_t kAlignmentKGrad = kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad(); 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()); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp index 870a50693d..a4eeaa6c66 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp @@ -286,25 +286,6 @@ struct BlockFmhaBwdPipelineDefaultPolicy return kVecLoad; } - template - CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentQGrad() - { - using QGradDataType = remove_cvref_t; - 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 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, @@ -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, @@ -737,27 +716,6 @@ struct BlockFmhaBwdPipelineDefaultPolicy return 16 / sizeof(GemmDataType); } - // template - // CK_TILE_HOST_DEVICE static constexpr auto MakeQGradDramBlockDescriptor() - // { - // constexpr auto q_grad_dram_desc_0 = make_naive_tensor_descriptor( - // make_tuple(number{}, number{}, number{}), - // make_tuple(number{}, number{}, number<1>{}), - // number{}, - // number<1>{} - // ); - - // constexpr auto q_grad_dram_desc = transform_tensor_descriptor( - // q_grad_dram_desc_0, - // make_tuple( - // make_merge_transform(make_tuple(number{}, number{})), - // make_pass_through_transform(number{})), - // make_tuple(sequence<0, 2>{}, sequence<1>{}), - // make_tuple(sequence<0>{}, sequence<1>{}) - // ); - // return q_grad_dram_desc; - // } - template CK_TILE_HOST_DEVICE static constexpr auto MakeXLdsBlockDescriptor() {