From 5944a63f1148167f3cb7f3d00cd5a237be7d295a Mon Sep 17 00:00:00 2001
From: Qianfeng Zhang <Qianfeng.Zhang@amd.com>
Date: Fri, 6 Jun 2025 14:20:41 +0000
Subject: [PATCH] Update IsFulleTileInsideMask() for kUseLocal is true
 situtation

---
 .../hstu_attention_fwd_pipeline.hpp           | 30 +++++++++++--------
 .../18_hstu_attention/hstu_block_masking.hpp  |  5 ++--
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline.hpp
index 0b9ffec66b..dacdea67e1 100644
--- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline.hpp
+++ b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline.hpp
@@ -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
                     {
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 16b080227d..3ae5924f6a 100644
--- a/example/ck_tile/18_hstu_attention/hstu_block_masking.hpp
+++ b/example/ck_tile/18_hstu_attention/hstu_block_masking.hpp
@@ -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;
     }
 };