ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-06 15:54:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

Replace sentinel value before storing

Browse Source
This commit is contained in:
PoYen, Chen
2024-06-04 09:59:51 +00:00
parent 5a6b8d8606
commit 064afc69d9
1 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_combine_pipeline.hpp
View File

@@ -271,8 +271,8 @@ struct BlockFmhaFwdSplitKVCombinePipeline
}
else
{
lse_logsum(distributed_indices) =
ck_tile::log(lse_sum(distributed_indices)) + lse_max(distributed_indices);
lse_logsum(distributed_indices) = ck_tile::log(lse_sum(distributed_indices)) +
get_validated_m(lse_max(distributed_indices));
}
#if defined(PRINT_LSE_LOGSUM)
@@ -328,6 +328,18 @@ struct BlockFmhaFwdSplitKVCombinePipeline
if constexpr(kStoreLSE)
{
constexpr auto out_spans = static_distributed_tensor<
LSEDataType,
decltype(lse_logsum.get_tile_distribution())>::get_distributed_spans();
sweep_tile_span(out_spans[number<0>{}], [&](auto idx0) {
constexpr auto distributed_indices = make_tuple(idx0);
if(lse_logsum(distributed_indices) == numeric<LSEDataType>::infinity())
{
lse_logsum(distributed_indices) = -numeric<LSEDataType>::infinity();
}
});
store_tile(lse_dram_window_tmp, tile_elementwise_in(lse_element_func, lse_logsum));
}