From ee259a89247dfa4229cd50781e49ece54c0dfbb8 Mon Sep 17 00:00:00 2001 From: Qianfeng Zhang Date: Fri, 18 Apr 2025 15:37:49 +0000 Subject: [PATCH] Fix the GetTileRangeAlongX() to align with the hstu masking definition when both causal=true and local=true --- example/ck_tile/18_hstu_attention/hstu_block_masking.hpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/example/ck_tile/18_hstu_attention/hstu_block_masking.hpp b/example/ck_tile/18_hstu_attention/hstu_block_masking.hpp index 99ce3e0447..3c39b332a6 100644 --- a/example/ck_tile/18_hstu_attention/hstu_block_masking.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_block_masking.hpp @@ -61,9 +61,10 @@ struct HstuBlockMaskWithLocal } else { - index_t x_end = i_y + YTile + max_attn_len; + index_t x_start = max(0, i_y - max_attn_len); + index_t x_end = min(i_y + YTile, max_uih_len); - return ck_tile::make_tuple(0, x_end); + return ck_tile::make_tuple(x_start, x_end); }; }; }