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

Update in GetTileRangeAlongX to consider for non-causal+local_size>0 situation and add test case to test_hstu_attention.sh

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-07-25 14:53:39 +00:00
parent 43a97681b8
commit 29d3dc9662
2 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -54,13 +54,10 @@ struct HstuBlockMaskWithLocal
CK_TILE_HOST_DEVICE constexpr auto
GetTileRangeAlongX(index_t i_y, number<YTile>, number<XTile>) const
{
if constexpr(kUseCausal)
if(!is_tile_in_first_split)
{
if(!is_tile_in_first_split)
{
index_t x_end = min(i_y + YTile, seqlen);
return ck_tile::make_tuple(0, x_end);
};
index_t x_end = min(i_y + YTile, seqlen);
return ck_tile::make_tuple(0, x_end);
};
if constexpr(!kUseCausal)

3
example/ck_tile/18_hstu_attention/scripts/test_hstu_attention.sh
View File

@@ -36,5 +36,8 @@ for dtype in "fp16" "bf16"; do
## jagged causal+local+context+target
$EXE -v=1 -prec=$dtype -b=10 -jagged=1 -nhead=4 -hdim_qk=128 -hdim_v=128 -seqlens=300,300,290,280,310 -causal=1 -local_len=5 -context_len=8 -minfull_len=7 -targets=8
## jagged no-causal+local+context+target
$EXE -v=1 -prec=$dtype -b=10 -jagged=1 -nhead=4 -hdim_qk=128 -hdim_v=128 -seqlens=300,300,290,280,310 -causal=1 -local_len=5 -context_len=8 -minfull_len=7 -targets=8
set +x
done