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

Update IsFulleTileInsideMask() for kUseLocal is true situtation

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-06-06 14:20:41 +00:00
parent 9e6a24010a
commit d7930cd541
2 changed files with 21 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -406,21 +406,27 @@ struct HstuAttentionFwdPipelineQRKSVS
{
if constexpr(HstuMask::kUseLocal)
{
constexpr auto s_spans = SaccBlockTileType::get_distributed_spans();
sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) {
sweep_tile_span(s_spans[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
sacc_tile.get_tile_distribution(), make_tuple(idx0, idx1));
if(!mask.IsFullTileInsideMask(q_origin.at(number<0>{}),
seqlen_k_curr,
number<kK1>{},
number<kM0>{}))
{
constexpr auto s_spans = SaccBlockTileType::get_distributed_spans();
sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) {
sweep_tile_span(s_spans[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
sacc_tile.get_tile_distribution(), make_tuple(idx0, idx1));
const auto row =
q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
const auto col = seqlen_k_curr + tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
const auto row =
q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
const auto col = seqlen_k_curr + tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
sacc_tile(i_j_idx) *= static_cast<GemmAccDataType>(
mask.IsTokenPairInsideMask(row, col));
sacc_tile(i_j_idx) *= static_cast<GemmAccDataType>(
mask.IsTokenPairInsideMask(row, col));
});
});
});
}
}
else // kUseCausal=true, kUseLocal=false
{

5
example/ck_tile/18_hstu_attention/hstu_block_masking.hpp
View File

@@ -265,10 +265,11 @@ struct HstuBlockMaskWithLocal
number<TileWidth>,
number<TileHeight>) const
{
// when local masking used, we assume all tiles need pixel-by-pixel checking
std::ignore = i_tile_top;
std::ignore = i_tile_left;
if(min_full_attn_seqlen > 0 && i_tile_top >= max_uih_len - min_full_attn_seqlen)
return true;
return false;
}
};