From 70514fd8995581d3d2e1183d6857e4260b69476b Mon Sep 17 00:00:00 2001 From: danyao12 Date: Wed, 14 Aug 2024 12:32:12 +0000 Subject: [PATCH] bwd rtn --- .../ck_tile/01_fmha/codegen/ops/fmha_bwd.py | 5 +- .../ck_tile/01_fmha/script/smoke_test_bwd.sh | 2 +- .../ck_tile/core/tensor/tile_elementwise.hpp | 26 +++++++ .../pipeline/block_fmha_bwd_convert_dq.hpp | 7 +- ...k_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp | 34 ++++++++-- ...a_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp | 68 +++++++++++++++---- 6 files changed, 120 insertions(+), 22 deletions(-) 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 096394c0c9..6adc953357 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -500,8 +500,9 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) -> if not cond: continue if receipt == 3: - cond = dtype in ['fp16', 'bf16'] - cond &= bias in ['no', 'alibi'] + cond = dtype in ['bf16'] + cond &= bias in ['no'] + cond &= dropout in ['no'] cond &= dpad == dvpad cond &= deterministic == "f" if not cond: 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 dbb592820e..310f23e235 100755 --- a/example/ck_tile/01_fmha/script/smoke_test_bwd.sh +++ b/example/ck_tile/01_fmha/script/smoke_test_bwd.sh @@ -9,7 +9,7 @@ export CK_REPEAT=1 COMMON_ARGS='-v=1' set -x -for prec in "fp16" "bf16" ; do +for prec in "bf16" ; do for perm in 0 1 ; do for hdim in 32 64 128 256 ; do for mode in 0 1 ; do diff --git a/include/ck_tile/core/tensor/tile_elementwise.hpp b/include/ck_tile/core/tensor/tile_elementwise.hpp index 79018b9ced..5095367485 100644 --- a/include/ck_tile/core/tensor/tile_elementwise.hpp +++ b/include/ck_tile/core/tensor/tile_elementwise.hpp @@ -227,6 +227,32 @@ CK_TILE_DEVICE auto cast_tile_pk_fp16_fp32(const InTensor& in_dstr_tensors) #endif } +template +CK_TILE_DEVICE auto cast_tile_rtn_bf16_fp32(const InTensor& in_dstr_tensors) +{ +#if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx94__) + // This API is designed to use the _pk_ serious of function + constexpr auto in_tile_dstr = InTensor::get_tile_distribution(); + + constexpr index_t thread_buffer_size = InTensor::get_thread_buffer_size(); + + auto out_dstr_tensor = make_static_distributed_tensor(in_tile_dstr); + + for(index_t i = 0; i < thread_buffer_size; i++) + { + out_dstr_tensor.get_thread_buffer().at(i) = + float_to_bf16_raw(0)>( + in_dstr_tensors.get_thread_buffer()[i]); + } + + return out_dstr_tensor; +#else + // fallback + return tile_elementwise_in(type_convert, + in_dstr_tensors); +#endif +} + #if CK_TILE_USE_SUBDWORD_TILE_CAST // this function assume either src or dst (or both) date type is under 1 dword // we pack subdword value into 1 dword to avoid compiler's default subdword behavior(which is buggy) 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 3da1104169..46ff9c6db5 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 @@ -125,7 +125,12 @@ struct BlockFmhaBwdConvertQGrad sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) { sweep_tile_span(dq_acc_spans[number<2>{}], [&](auto idx2) { constexpr auto n_i_j_idx = make_tuple(idx0, idx1, idx2); - dq_converted(n_i_j_idx) = type_convert(dq_acc[n_i_j_idx]); + if constexpr(std::is_same_v) + dq_converted(n_i_j_idx) = + float_to_bf16_raw(0)>( + dq_acc[n_i_j_idx]); + else + dq_converted(n_i_j_idx) = type_convert(dq_acc[n_i_j_idx]); }); }); }); 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 6cdd40f506..a53d0b9454 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 @@ -611,18 +611,34 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR dropout.template Run( seqlen_q_step, k_origin.at(number<0>{}), pt, randval_dram_window); } - const auto pt_gemm = [&]() { + // const auto pt_gemm = [&]() { + // if constexpr(FmhaDropout::IsDropout) + // { + // return tile_elementwise_in( + // [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); + // }, pt); + // } + // else + // { + // return cast_tile(pt); + // } + // }(); + const auto pt_dropped = [&]() { if constexpr(FmhaDropout::IsDropout) { - return tile_elementwise_in( - [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, - pt); + return tile_elementwise_in([](const auto& x) { return x > 0.f ? x : 0.f; }, pt); } else { - return cast_tile(pt); + return pt; } }(); + const auto pt_gemm = [&]() { + if constexpr(std::is_same_v) + return impl::cast_tile_rtn_bf16_fp32(pt_dropped); + else + return cast_tile(pt_dropped); + }(); // STAGE 3, P^T@OGrad^T Gemm1 auto do_block_tile = load_tile(do_dram_window); @@ -702,7 +718,13 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR auto qt_reg_tensor = load_tile(qt_lds_read_window); block_sync_lds(); - const auto dst_gemm = cast_tile(dst); + // const auto dst_gemm = cast_tile(dst); + const auto dst_gemm = [&]() { + if constexpr(std::is_same_v) + return impl::cast_tile_rtn_bf16_fp32(dst); + else + return cast_tile(dst); + }(); Policy::template SGradTFromGemm2CToGemm3A( seqlen_q_step, k_origin.at(number<0>{}), pt, randval_dram_window); } - const auto pt_gemm = [&]() { + // const auto pt_gemm = [&]() { + // if constexpr(FmhaDropout::IsDropout) + // { + // return tile_elementwise_in( + // [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); + // }, pt); + // } + // else + // { + // return cast_tile(pt); + // } + // }(); + const auto pt_dropped = [&]() { if constexpr(FmhaDropout::IsDropout) { - return tile_elementwise_in( - [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, - pt); + return tile_elementwise_in([](const auto& x) { return x > 0.f ? x : 0.f; }, pt); } else { - return cast_tile(pt); + return pt; } }(); + const auto pt_gemm = [&]() { + if constexpr(std::is_same_v) + return impl::cast_tile_rtn_bf16_fp32(pt_dropped); + else + return cast_tile(pt_dropped); + }(); // STAGE 3, P^T@OGrad^T Gemm1 Policy::template PTFromGemm0CToGemm1A(dst); + // const auto dst_gemm = cast_tile(dst); + const auto dst_gemm = [&]() { + if constexpr(std::is_same_v) + return impl::cast_tile_rtn_bf16_fp32(dst); + else + return cast_tile(dst); + }(); Policy::template SGradTFromGemm2CToGemm3A(x > 0.f ? x : 0.f); }, + // pt); + // } + // else + // { + // return cast_tile(pt); + // } + // }(); + const auto pt_dropped = [&]() { if constexpr(FmhaDropout::IsDropout) { - return tile_elementwise_in( - [](const auto& x) { return type_convert(x > 0.f ? x : 0.f); }, - pt); + return tile_elementwise_in([](const auto& x) { return x > 0.f ? x : 0.f; }, pt); } else { - return cast_tile(pt); + return pt; } }(); + const auto pt_gemm = [&]() { + if constexpr(std::is_same_v) + return impl::cast_tile_rtn_bf16_fp32(pt_dropped); + else + return cast_tile(pt_dropped); + }(); Policy::template PTFromGemm0CToGemm1A( pt_reg_tensor, pt_gemm); @@ -969,7 +1007,13 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP } // STAGE 6, SGrad^T@Q^T Gemm3 - const auto dst_gemm = cast_tile(dst); + // const auto dst_gemm = cast_tile(dst); + const auto dst_gemm = [&]() { + if constexpr(std::is_same_v) + return impl::cast_tile_rtn_bf16_fp32(dst); + else + return cast_tile(dst); + }(); Policy::template SGradTFromGemm2CToGemm3A