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

Tiny codes simplification in pipeline

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-04-18 08:22:11 +00:00
parent ca1ae84fc6
commit f12a47218f
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline.hpp
View File

@@ -349,25 +349,25 @@ struct HstuAttentionFwdPipelineQRKSVS
move_tile_window(v_dram_window, {0, kK1});
});
s_acc = tile_elementwise_in(s_acc_element_func, s_acc);
// STAGE 2, scale_s, add bias, mask, siLU
if constexpr(kHasBias)
{
s_acc = tile_elementwise_in(s_acc_element_func, s_acc);
tile_elementwise_inout([&scale_s](auto& x) { x = x * scale_s; }, s_acc);
tile_elementwise_inout(
[&](auto& x, const auto& y) {
x += type_convert<GemmAccDataType>(bias_element_func(y));
[&scale_s, &bias_element_func](auto& x, const auto& y) {
x = x * scale_s + type_convert<GemmAccDataType>(bias_element_func(y));
},
s_acc,
bias_tile);
move_tile_window(bias_dram_window, {0, kN0});
}
else
{
s_acc = tile_elementwise_in(s_acc_element_func, s_acc);
tile_elementwise_inout([&scale_s](auto& x) { x = x * scale_s; }, s_acc);
}
move_tile_window(bias_dram_window, {0, kN0});
if constexpr(HstuMask::IsMasking)
{
const auto k_origin = k_dram_block_window.get_window_origin();