ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

bwd rtn

Browse Source
This commit is contained in:
danyao12
2024-08-14 12:32:12 +00:00
parent 17c97f5814
commit 70514fd899
6 changed files with 120 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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:

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

@@ -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

26
include/ck_tile/core/tensor/tile_elementwise.hpp
View File

@@ -227,6 +227,32 @@ CK_TILE_DEVICE auto cast_tile_pk_fp16_fp32(const InTensor& in_dstr_tensors)
#endif
}
template <typename OutDataType, typename InTensor>
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<OutDataType>(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<static_cast<bf16_rounding_mode>(0)>(
in_dstr_tensors.get_thread_buffer()[i]);
}
return out_dstr_tensor;
#else
// fallback
return tile_elementwise_in(type_convert<OutDataType, typename InTensor::DataType>,
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)

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

@@ -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<QGradDataType>(dq_acc[n_i_j_idx]);
if constexpr(std::is_same_v<QGradDataType, bf16_t>)
dq_converted(n_i_j_idx) =
float_to_bf16_raw<static_cast<bf16_rounding_mode>(0)>(
dq_acc[n_i_j_idx]);
else
dq_converted(n_i_j_idx) = type_convert<QGradDataType>(dq_acc[n_i_j_idx]);
});
});
});

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

@@ -611,18 +611,34 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
dropout.template Run<decltype(gemm_0), RandValOutputDataType>(
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<GemmDataType>(x > 0.f ? x : 0.f);
// }, pt);
// }
// else
// {
// return cast_tile<GemmDataType>(pt);
// }
// }();
const auto pt_dropped = [&]() {
if constexpr(FmhaDropout::IsDropout)
{
return tile_elementwise_in(
[](const auto& x) { return type_convert<GemmDataType>(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<GemmDataType>(pt);
return pt;
}
}();
const auto pt_gemm = [&]() {
if constexpr(std::is_same_v<GemmDataType, bf16_t>)
return impl::cast_tile_rtn_bf16_fp32<GemmDataType>(pt_dropped);
else
return cast_tile<GemmDataType>(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<GemmDataType>(dst);
// const auto dst_gemm = cast_tile<GemmDataType>(dst);
const auto dst_gemm = [&]() {
if constexpr(std::is_same_v<GemmDataType, bf16_t>)
return impl::cast_tile_rtn_bf16_fp32<GemmDataType>(dst);
else
return cast_tile<GemmDataType>(dst);
}();
Policy::template SGradTFromGemm2CToGemm3A<Problem,
decltype(dst_reg_tensor),

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

@@ -647,18 +647,34 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
dropout.template Run<decltype(gemm_0), RandValOutputDataType>(
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<GemmDataType>(x > 0.f ? x : 0.f);
// }, pt);
// }
// else
// {
// return cast_tile<GemmDataType>(pt);
// }
// }();
const auto pt_dropped = [&]() {
if constexpr(FmhaDropout::IsDropout)
{
return tile_elementwise_in(
[](const auto& x) { return type_convert<GemmDataType>(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<GemmDataType>(pt);
return pt;
}
}();
const auto pt_gemm = [&]() {
if constexpr(std::is_same_v<GemmDataType, bf16_t>)
return impl::cast_tile_rtn_bf16_fp32<GemmDataType>(pt_dropped);
else
return cast_tile<GemmDataType>(pt_dropped);
}();
// STAGE 3, P^T@OGrad^T Gemm1
Policy::template PTFromGemm0CToGemm1A<Problem,
@@ -733,7 +749,13 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
}
// STAGE 6, SGrad^T@Q^T Gemm3
const auto dst_gemm = cast_tile<GemmDataType>(dst);
// const auto dst_gemm = cast_tile<GemmDataType>(dst);
const auto dst_gemm = [&]() {
if constexpr(std::is_same_v<GemmDataType, bf16_t>)
return impl::cast_tile_rtn_bf16_fp32<GemmDataType>(dst);
else
return cast_tile<GemmDataType>(dst);
}();
Policy::template SGradTFromGemm2CToGemm3A<Problem,
decltype(dst_reg_tensor),
@@ -900,18 +922,34 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
}
// STAGE 3, P^T@OGrad^T Gemm1
const auto pt_gemm = [&]() {
// const auto pt_gemm = [&]() {
// if constexpr(FmhaDropout::IsDropout)
// {
// return tile_elementwise_in(
// [](const auto& x) { return type_convert<GemmDataType>(x > 0.f ? x : 0.f); },
// pt);
// }
// else
// {
// return cast_tile<GemmDataType>(pt);
// }
// }();
const auto pt_dropped = [&]() {
if constexpr(FmhaDropout::IsDropout)
{
return tile_elementwise_in(
[](const auto& x) { return type_convert<GemmDataType>(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<GemmDataType>(pt);
return pt;
}
}();
const auto pt_gemm = [&]() {
if constexpr(std::is_same_v<GemmDataType, bf16_t>)
return impl::cast_tile_rtn_bf16_fp32<GemmDataType>(pt_dropped);
else
return cast_tile<GemmDataType>(pt_dropped);
}();
Policy::template PTFromGemm0CToGemm1A<Problem, decltype(pt_reg_tensor), decltype(pt_gemm)>(
pt_reg_tensor, pt_gemm);
@@ -969,7 +1007,13 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP
}
// STAGE 6, SGrad^T@Q^T Gemm3
const auto dst_gemm = cast_tile<GemmDataType>(dst);
// const auto dst_gemm = cast_tile<GemmDataType>(dst);
const auto dst_gemm = [&]() {
if constexpr(std::is_same_v<GemmDataType, bf16_t>)
return impl::cast_tile_rtn_bf16_fp32<GemmDataType>(dst);
else
return cast_tile<GemmDataType>(dst);
}();
Policy::template SGradTFromGemm2CToGemm3A<Problem,
decltype(dst_reg_tensor),