ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-03 05:01:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

[layernorm] hot fix (#1620)

Browse Source 
* hot fix ln

* some rename
This commit is contained in:
carlushuang
2024-11-01 11:52:50 +08:00
committed by GitHub
parent c3a4800c5f
commit 550248deec
3 changed files with 29 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
include/ck_tile/ops/layernorm2d/pipeline/layernorm2d_fwd_pipeline_one_pass.hpp
View File

@@ -111,8 +111,9 @@ struct Layernorm2dFwdPipelineOnePass
{
sweep_tile(x_resi, [&](auto idx) {
// compute x = x_resi + x
x(idx) = type_convert<YResidualDataType>(x_resi(idx)) +
type_convert<YResidualDataType>(x(idx));
auto re_ = type_convert<ComputeDataType>(x_resi(idx)) +
type_convert<ComputeDataType>(x(idx));
x(idx) = type_convert<XDataType>(re_);
});
if constexpr(kFusedAdd == Layernorm2dFusedAddEnum::PRE_ADD_STORE)
store_tile(y_residual_window, x);

10
include/ck_tile/ops/layernorm2d/pipeline/layernorm2d_fwd_pipeline_two_pass.hpp
View File

@@ -122,8 +122,9 @@ struct Layernorm2dFwdPipelineTwoPass
{
sweep_tile(x_resi, [&](auto idx) {
// compute x = x_resi + x
x(idx) = type_convert<YResidualDataType>(x_resi(idx)) +
type_convert<YResidualDataType>(x(idx));
auto re_ = type_convert<ComputeDataType>(x_resi(idx)) +
type_convert<ComputeDataType>(x(idx));
x(idx) = type_convert<XDataType>(re_);
});
if constexpr(kFusedAdd == Layernorm2dFusedAddEnum::PRE_ADD_STORE)
{
@@ -170,8 +171,9 @@ struct Layernorm2dFwdPipelineTwoPass
{
sweep_tile(x_resi, [&](auto idx) {
// compute x = x_resi + x
x(idx) = type_convert<YResidualDataType>(x_resi(idx)) +
type_convert<YResidualDataType>(x(idx));
auto re_ = type_convert<ComputeDataType>(x_resi(idx)) +
type_convert<ComputeDataType>(x(idx));
x(idx) = type_convert<XDataType>(re_);
});
}
// load gamma/beta (TODO: support no gamma/beta?)