diff --git a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py index e89774156a..9e15a822ef 100644 --- a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py +++ b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py @@ -12,6 +12,7 @@ FWD_DTYPE_MAP = { BWD_DTYPE_MAP = { "fp16": "FmhaBwdFp16", + "bf16": "FmhaBwdBf16" } MASK_IMPL = { 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 4027342323..9525299bc1 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -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] diff --git a/example/ck_tile/01_fmha/fmha_bwd.cpp b/example/ck_tile/01_fmha/fmha_bwd.cpp index 9e538b68ee..647e3e4094 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.cpp +++ b/example/ck_tile/01_fmha/fmha_bwd.cpp @@ -123,6 +123,18 @@ 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) +{ + unsigned un = static_cast(dim_value); + un |= un >> 1; + un |= un >> 2; + un |= un >> 4; + un |= un >> 8; + un |= un >> 16; + un++; + return static_cast(un); +} + template 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(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; 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{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 dq_acc_host(dq_acc_perm ? std::array{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{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, diff --git a/example/ck_tile/01_fmha/fmha_bwd.hpp b/example/ck_tile/01_fmha/fmha_bwd.hpp index 632753977d..60f316ee7c 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd.hpp @@ -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, diff --git a/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp index f01aedc829..bfadfbf379 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp @@ -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" : "_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>, std::conditional_t>, std::conditional_t>, - std::conditional_t> + std::conditional_t>, + std::conditional_t> { 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(seqstart_q_ptr), reinterpret_cast(seqstart_k_ptr), reinterpret_cast(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(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(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( + dq_acc_ptr, + make_tuple(kargs.seqlen_q, kargs.hdim_q), + make_tuple(kargs.stride_dq_acc, 1), + number{}, + 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{}, + number{}), + sequence{}); + } + 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{}, number{}, number{}, number{}, number{}), - make_tuple(number{}, number{}, number{}, number{}, number{}, number<1>{}), - number{}, - 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{}, number{}, number{})), - make_merge_transform(make_tuple(number{}, number{}))), - make_tuple(sequence<0, 3, 2, 5>{}, sequence<1, 4>{}), - make_tuple(sequence<0>{}, sequence<1>{}) - ); - return - make_tensor_view( - dq_acc_ptr, - q_grad_dram_desc); + const auto q_grad_dram_desc_0 = make_naive_tensor_descriptor( + make_tuple(M0, + number{}, + number{}, + number{}, + number{}, + number{}), + make_tuple(number{}, + number{}, + number{}, + number{}, + number{}, + number<1>{}), + number{}, + 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{}, + number{}, + number{})), + make_merge_transform( + make_tuple(number{}, number{}))), + make_tuple(sequence<0, 3, 2, 5>{}, sequence<1, 4>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + return make_tensor_view( + dq_acc_ptr, q_grad_dram_desc); + } }(); - - + return make_tile_window( dq_acc_dram, make_tuple(number{}, number{}), @@ -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 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> + FmhaBwdConvertQGradEmptyKargs<0>>, + std::conditional_t> { 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(seqstart_q_ptr), reinterpret_cast(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(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(kargs.dq_acc_ptr) + static_cast(i_nhead_) * (kargs.nhead_stride_dq_acc) + batch_offset_dq_acc; + if constexpr(kIsAtomic32) + { - // auto dq_acc_dram_naive = make_naive_tensor_view( - // dq_acc_ptr, - // make_tuple(kargs.seqlen_q, kargs.hdim_q), - // make_tuple(kargs.stride_dq_acc, 1), - // number{}, - // number<1>{}); - // return pad_tensor_view(dq_acc_dram_naive, - // make_tuple(number{}, number{}), - // sequence{}); - 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( + dq_acc_ptr, + make_tuple(kargs.seqlen_q, kargs.hdim_q), + make_tuple(kargs.stride_dq_acc, 1), + number{}, + number<1>{}); + return pad_tensor_view(dq_acc_dram_naive, + make_tuple(number{}, number{}), + sequence{}); + } + 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{}, number{}, number{}, number{}, number{}), - make_tuple(number{}, number{}, number{}, number{}, number{}, number<1>{}), - number{}, - number<1>{} - ); + const auto q_grad_dram_desc_0 = make_naive_tensor_descriptor( + make_tuple(M0, + number{}, + number{}, + number{}, + number{}, + number{}), + make_tuple(number{}, + number{}, + number{}, + number{}, + number{}, + number<1>{}), + number{}, + 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{}, number{}, number{})), - make_merge_transform(make_tuple(number{}, number{}))), - make_tuple(sequence<0, 3, 2, 5>{}, sequence<1, 4>{}), - make_tuple(sequence<0>{}, sequence<1>{}) - ); - return - make_tensor_view( - 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{}, + number{}, + number{})), + make_merge_transform(make_tuple(number{}, number{}))), + make_tuple(sequence<0, 3, 2, 5>{}, sequence<1, 4>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + return make_tensor_view(dq_acc_ptr, + q_grad_dram_desc); + } } }(); 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 3a9d9c5acb..e68d42e8ab 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 @@ -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(dq_acc)); + if constexpr(kIsAtomic32) + { + update_tile(dq_dram_window, dq_acc); + } + else + { + update_tile(dq_dram_window, cast_tile(dq_acc)); + } } - } move_tile_window(dq_dram_window, {kM0, 0}); i_total_loops += 1; 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 ad8f42794c..334bf41108 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 @@ -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(dq_acc)); - // if (wgid ==0 && tid==0) { - // printf("atomic 16 update tile \n"); - - // } } }