ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 17:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Move the dividing by max_seqlen out of f_silu to be handle outside the main-loop

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-05-18 03:21:36 +00:00
parent 27e64a682a
commit 70237d2e5c
2 changed files with 11 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp
View File

@@ -35,6 +35,8 @@ struct HstuAttentionFwdKernel
using QKVDataType = ck_tile::remove_cvref_t<typename HstuAttentionPipeline::QKVDataType>;
using BiasDataType = ck_tile::remove_cvref_t<typename HstuAttentionPipeline::BiasDataType>;
using ODataType = ck_tile::remove_cvref_t<typename HstuAttentionPipeline::ODataType>;
using GemmAccDataType =
ck_tile::remove_cvref_t<typename HstuAttentionPipeline::GemmAccDataType>;
using VLayout = ck_tile::remove_cvref_t<typename HstuAttentionPipeline::VLayout>;
@@ -740,7 +742,6 @@ struct HstuAttentionFwdKernel
bias_dram_window,
mask,
kargs.scale_s,
kargs.max_seqlen,
smem_ptr,
dropout);
}();
@@ -765,6 +766,13 @@ struct HstuAttentionFwdKernel
make_tuple(number<HstuAttentionPipeline::kM0>{}, number<HstuAttentionPipeline::kN1>{}),
{i_m0, i_n1});
tile_elementwise_inout(
[&](auto& x) {
x = x * type_convert<GemmAccDataType>(
__builtin_amdgcn_rcpf(static_cast<float>(kargs.max_seqlen)));
},
o_acc_tile);
EpiloguePipeline{}(o_dram_window, o_acc_tile);
}
};

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

@@ -133,7 +133,6 @@ struct HstuAttentionFwdPipelineQRKSVS
const OAccElementFunction& o_acc_element_func,
HstuMask mask,
float scale_s,
index_t max_seqlen, // used by silu
void* smem_ptr,
DropoutType& dropout) const
{
@@ -238,12 +237,11 @@ struct HstuAttentionFwdPipelineQRKSVS
if constexpr(std::is_same_v<CompDataType, float>)
{
x = x * __builtin_amdgcn_rcpf(neg_one - __expf(x)) *
__builtin_amdgcn_rcpf(type_convert<CompDataType>(max_seqlen));
x = x * __builtin_amdgcn_rcpf(neg_one - __expf(x));
}
else
{
x = x / (neg_one - exp(x)) / type_convert<CompDataType>(max_seqlen);
x = x / (neg_one - exp(x));
}
};
@@ -481,7 +479,6 @@ struct HstuAttentionFwdPipelineQRKSVS
const BiasDramBlockWindowTmp& bias_dram_block_window_tmp, // M0*N0 tile
HstuMask mask,
float scale_s,
index_t max_seqlen,
void* smem_ptr,
DropoutType& dropout) const
{
@@ -498,7 +495,6 @@ struct HstuAttentionFwdPipelineQRKSVS
identity{},
mask,
scale_s,
max_seqlen,
smem_ptr,
dropout);
}