ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-16 10:59:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Tiny update in IsTokenPairInsideMask()

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-04-25 15:36:58 +00:00
parent 27f7ab4f2c
commit 9996270087
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -96,16 +96,24 @@ struct HstuBlockMaskWithLocal
if(row < contextual_seqlen)
return true;
bool result = false;
if constexpr(kUseCausal)
result = (row >= col) && (row - col <= max_attn_len);
{
bool result = (row >= col) && (row - col <= max_attn_len);
if(min_full_attn_seqlen > 0)
result = result || (row >= max_uih_len - min_full_attn_seqlen);
return result;
}
else
result = abs(row - col) <= max_attn_len;
{
bool result = abs(row - col) <= max_attn_len;
if(min_full_attn_seqlen > 0)
result = result || (row >= max_uih_len - min_full_attn_seqlen);
if(min_full_attn_seqlen > 0)
result = result || (row >= max_uih_len - min_full_attn_seqlen);
return result;
return result;
}
};
};