Update the rules of hstu masking

example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp

@@ -604,7 +604,7 @@ struct HstuAttentionFwdKernel
         const auto q_dram = [&]() {
             const auto q_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                 q_ptr,
-                make_tuple(mask.max_uih_len, kargs.hdim_qk),
+                make_tuple(kargs.seqlen, kargs.hdim_qk),
                 make_tuple(kargs.seq_stride_q, 1),
                 number<HstuAttentionPipeline::kAlignmentQ>{},
                 number<1>{});
@@ -616,7 +616,7 @@ struct HstuAttentionFwdKernel
         const auto k_dram = [&]() {
             const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                 k_ptr,
-                make_tuple(mask.max_uih_len, kargs.hdim_qk),
+                make_tuple(kargs.seqlen, kargs.hdim_qk),
                 make_tuple(kargs.seq_stride_k, 1),
                 number<HstuAttentionPipeline::kAlignmentK>{},
                 number<1>{});
@@ -631,7 +631,7 @@ struct HstuAttentionFwdKernel
             {
                 const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                     v_ptr,
-                    make_tuple(mask.max_uih_len, kargs.hdim_v),
+                    make_tuple(kargs.seqlen, kargs.hdim_v),
                     make_tuple(kargs.seq_stride_v, 1),
                     number<HstuAttentionPipeline::kAlignmentV>{},
                     number<1>{});

example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline.hpp

@@ -362,8 +362,10 @@ struct HstuAttentionFwdPipelineQRKSVS
                     }
                     else // kUseCausal=true, kUseLocal=false
                     {
-                        if(!mask.IsFullTileInsideMask(
-                               q_origin.at(number<0>{}), seqlen_k_curr, number<kK1>{}))
+                        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) {

example/ck_tile/18_hstu_attention/hstu_block_masking.hpp

@@ -14,20 +14,34 @@ struct HstuBlockMaskWithLocal
     static constexpr bool kUseLocal = true;
     static constexpr bool IsMasking = true;
 
+    int seqlen;
     int contextual_seqlen;
-    int max_uih_len;
 
     int max_attn_len;
     int min_full_attn_seqlen;
+    int num_target;
 
-    CK_TILE_HOST_DEVICE HstuBlockMaskWithLocal(int contextual_seqlen_,
-                                               int max_uih_len_,
+    int max_uih_len;
+    int max_id;
+
+    CK_TILE_HOST_DEVICE HstuBlockMaskWithLocal(int seqlen_,
+                                               int contextual_seqlen_,
                                                int max_attn_len_,
-                                               int min_full_attn_seqlen_)
-        : contextual_seqlen(contextual_seqlen_),
-          max_uih_len(max_uih_len_),
+                                               int min_full_attn_seqlen_,
+                                               int num_target_)
+        : seqlen(seqlen_),
+          contextual_seqlen(contextual_seqlen_),
           max_attn_len(max_attn_len_),
-          min_full_attn_seqlen(min_full_attn_seqlen_){};
+          min_full_attn_seqlen(min_full_attn_seqlen_),
+          num_target(num_target_)
+    {
+        max_uih_len = seqlen - num_target;
+
+        if(contextual_seqlen > 0)
+            max_id = max_uih_len - (contextual_seqlen - 1);
+        else
+            max_id = max_uih_len;
+    };
 
     // to get the loop length along X axis, return index:[start, end), end-start=length
     // use this if need loop over X axis tile by tile (eg. seqlen_k loop-over)
@@ -36,90 +50,191 @@ struct HstuBlockMaskWithLocal
     CK_TILE_HOST_DEVICE constexpr auto
     GetTileRangeAlongX(index_t i_y, number<YTile>, number<XTile>) const
     {
-        if(i_y < contextual_seqlen)
-            return ck_tile::make_tuple(0, max_uih_len);
+        if(min_full_attn_seqlen > 0 && i_y + YTile > max_uih_len - min_full_attn_seqlen)
+        {
+            index_t x_end = min(i_y + YTile, seqlen);
+            return ck_tile::make_tuple(0, x_end);
+        };
 
         if constexpr(!kUseCausal)
         {
-            if(min_full_attn_seqlen > 0 && i_y + YTile > max_uih_len - min_full_attn_seqlen)
+            if(i_y >= contextual_seqlen)
             {
-                return ck_tile::make_tuple(0, max_uih_len);
+                if(i_y < max_uih_len)
+                {
+                    index_t x_start         = max(0, i_y - max_attn_len);
+                    index_t x_start_aligned = x_start - x_start % XTile;
+
+                    if(i_y + YTile > max_uih_len - max_attn_len)
+                    {
+                        return ck_tile::make_tuple(x_start_aligned, seqlen);
+                    }
+                    else
+                    {
+                        index_t x_end = min(i_y + YTile + max_attn_len, seqlen);
+                        return ck_tile::make_tuple(x_start_aligned, x_end);
+                    };
+                }
+                else
+                {
+                    index_t x_start = i_y - max_attn_len;
+                    index_t x_end   = seqlen;
+
+                    return ck_tile::make_tuple(x_start - x_start % XTile, x_end);
+                }
             }
             else
             {
-                index_t x_start = max(0, i_y - max_attn_len);
-                index_t x_end   = i_y + YTile + max_attn_len;
-
-                return ck_tile::make_tuple(x_start - x_start % XTile, x_end);
-            };
+                if(i_y + YTile > max_uih_len)
+                {
+                    index_t x_end = min(i_y + YTile, seqlen);
+                    return ck_tile::make_tuple(0, x_end);
+                }
+                else
+                {
+                    index_t x_end = max(i_y + YTile + max_attn_len, max_uih_len);
+                    return ck_tile::make_tuple(0, x_end);
+                };
+            }
         }
         else // kUseCausal && kUseLocal
         {
-            if(min_full_attn_seqlen > 0 && i_y + YTile > max_uih_len - min_full_attn_seqlen)
+            if(i_y >= contextual_seqlen)
             {
-                return ck_tile::make_tuple(0, max_uih_len);
+                index_t x_end = min(i_y + YTile, seqlen);
+
+                if(i_y < max_uih_len)
+                {
+                    index_t x_start = max(0, i_y - max_attn_len);
+                    return ck_tile::make_tuple(x_start - x_start % XTile, x_end);
+                }
+                else
+                {
+                    index_t x_start = max_uih_len - max_attn_len;
+                    return ck_tile::make_tuple(x_start - x_start % XTile, x_end);
+                }
             }
             else
             {
-                index_t x_start = max(0, i_y - max_attn_len);
-                index_t x_end   = min(i_y + YTile, max_uih_len);
+                index_t x_end = min(i_y + YTile, seqlen);
 
-                return ck_tile::make_tuple(x_start - x_start % XTile, x_end);
-            };
+                if(i_y + YTile > max_uih_len)
+                {
+                    return ck_tile::make_tuple(0, x_end);
+                }
+                else
+                {
+                    return ck_tile::make_tuple(0, max_uih_len);
+                };
+            }
         };
     }
 
     CK_TILE_HOST constexpr bool IsTokenPairInsideMask(int row, int col)
     {
-        if(row >= max_uih_len || col >= max_uih_len)
-            return false;
+        // row_id/col_id is clamped from physical row/col according to contextual_seqlen and
+        // max_uih_len
+        int row_id = contextual_seqlen > 0 ? max(row - contextual_seqlen + 1, 0) : row;
+        int col_id = contextual_seqlen > 0 ? max(col - contextual_seqlen + 1, 0) : col;
 
-        if(row < contextual_seqlen)
+        row_id = min(row_id, max_id);
+        col_id = min(col_id, max_id);
+
+        if(contextual_seqlen > 0 && row_id == 0 && col_id < max_id)
             return true;
 
-        bool result = false;
-        if constexpr(kUseCausal)
-            result = (row >= col) && (row - col <= max_attn_len);
-        else
-            result = abs(row - col) <= max_attn_len;
-
-        if(min_full_attn_seqlen > 0)
-            result = result || (row >= max_uih_len - min_full_attn_seqlen);
-
-        return result;
-    };
-
-    // masking codes in device don't have to compare row/col with max_uih_len, since
-    // buffer_load_xxx instruction is able to return zero for out-of-boundary access
-    CK_TILE_DEVICE constexpr int IsTokenPairInsideMask(int row, int col)
-    {
-        if(row < contextual_seqlen)
-            return 1;
-
+        // use row_id/col_id to check the dist between two q/k token pair, token pairs on the
+        // diagonal line are always considerred
         if constexpr(kUseCausal)
         {
-            bool result = (row >= col) && (row - col <= max_attn_len);
+            if(((row_id > col_id) && (row_id - col_id <= max_attn_len)) || (row == col))
+                return true;
 
-            if(min_full_attn_seqlen > 0)
-                result = result || (row >= max_uih_len - min_full_attn_seqlen);
-
-            return static_cast<int>(result);
+            if((min_full_attn_seqlen > 0) && (row_id >= max_id - min_full_attn_seqlen))
+                return true;
         }
         else
         {
-            bool result = abs(row - col) <= max_attn_len;
+            if(((row_id != col_id && abs(row_id - col_id) <= max_attn_len)) || (row == col))
+                return true;
 
-            if(min_full_attn_seqlen > 0)
-                result = result || (row >= max_uih_len - min_full_attn_seqlen);
+            if((min_full_attn_seqlen > 0) && (row >= max_id - min_full_attn_seqlen))
+                return true;
+        }
 
-            return static_cast<int>(result);
+        return false;
+    };
+
+    CK_TILE_DEVICE constexpr int IsTokenPairInsideMask(int row, int col)
+    {
+        if(contextual_seqlen > 0)
+        {
+            // row_id/col_id is clamped from physical row/col according to contextual_seqlen and
+            // max_uih_len
+            int row_id = max(row - contextual_seqlen + 1, 0);
+            int col_id = max(col - contextual_seqlen + 1, 0);
+
+            row_id = min(row_id, max_id);
+            col_id = min(col_id, max_id);
+
+            if(row_id == 0 && col_id < max_id)
+                return 1;
+
+            // use row_id/col_id to check the dist between two q/k token pair, token pairs on the
+            // diagonal line are always considerred
+            if constexpr(kUseCausal)
+            {
+                bool res1 =
+                    (((row_id > col_id) && (row_id - col_id <= max_attn_len)) || (row == col));
+                bool res2 =
+                    ((min_full_attn_seqlen > 0) && (row_id >= max_id - min_full_attn_seqlen));
+
+                return static_cast<int>(res1 || res2);
+            }
+            else
+            {
+                bool res1 = (((row_id != col_id) && (abs(row_id - col_id) <= max_attn_len)) ||
+                             (row == col));
+                bool res2 = ((min_full_attn_seqlen > 0) && (row >= max_id - min_full_attn_seqlen));
+
+                return static_cast<int>(res1 || res2);
+            }
+        }
+        else
+        {
+            // row_id/col_id is clamped from physical row/col according to contextual_seqlen and
+            // max_uih_len
+            int row_id = min(row, max_id);
+            int col_id = min(col, max_id);
+
+            // use row_id/col_id to check the dist between two q/k token pair, token pairs on the
+            // diagonal line are always considerred
+            if constexpr(kUseCausal)
+            {
+                bool res1 =
+                    (((row_id > col_id) && (row_id - col_id <= max_attn_len)) || (row == col));
+                bool res2 =
+                    ((min_full_attn_seqlen > 0) && (row_id >= max_id - min_full_attn_seqlen));
+
+                return static_cast<int>(res1 || res2);
+            }
+            else
+            {
+                bool res1 = (((row_id != col_id) && (abs(row_id - col_id) <= max_attn_len)) ||
+                             (row == col));
+                bool res2 = ((min_full_attn_seqlen > 0) && (row >= max_id - min_full_attn_seqlen));
+
+                return static_cast<int>(res1 || res2);
+            }
         }
     };
 
     // if the whole tile inside the masking area, no need for pixel-by-pixel checking
-    template <index_t TileWidth>
-    CK_TILE_DEVICE constexpr bool
-    IsFullTileInsideMask(index_t i_tile_top, index_t i_tile_left, number<TileWidth>) const
+    template <index_t TileWidth, index_t TileHeight>
+    CK_TILE_DEVICE constexpr bool IsFullTileInsideMask(index_t i_tile_top,
+                                                       index_t i_tile_left,
+                                                       number<TileWidth>,
+                                                       number<TileHeight>) const
     {
         // when local masking used, we assume all tiles need pixel-by-pixel checking
         std::ignore = i_tile_top;
@@ -135,11 +250,23 @@ struct HstuBlockMaskNoLocal
     static constexpr bool kUseLocal = false;
     static constexpr bool IsMasking = kUseCausal;
 
+    int seqlen;
     int contextual_seqlen;
-    int max_uih_len;
+    int num_target;
 
-    CK_TILE_HOST_DEVICE HstuBlockMaskNoLocal(int contextual_seqlen_, int max_uih_len_)
-        : contextual_seqlen(contextual_seqlen_), max_uih_len(max_uih_len_){};
+    int max_uih_len;
+    int max_id;
+
+    CK_TILE_HOST_DEVICE HstuBlockMaskNoLocal(int seqlen_, int contextual_seqlen_, int num_target_)
+        : seqlen(seqlen_), contextual_seqlen(contextual_seqlen_), num_target(num_target_)
+    {
+        max_uih_len = seqlen - num_target;
+
+        if(contextual_seqlen > 0)
+            max_id = max_uih_len - (contextual_seqlen - 1);
+        else
+            max_id = max_uih_len;
+    };
 
     // to get the loop length along X axis, return index:[start, end), end-start=length
     // use this if need loop over X axis tile by tile (eg. seqlen_k loop-over)
@@ -150,71 +277,135 @@ struct HstuBlockMaskNoLocal
     {
         if constexpr(!IsMasking)
         {
-            return ck_tile::make_tuple(0, max_uih_len);
+            return ck_tile::make_tuple(0, seqlen);
         }
         else
         {
+            index_t x_end = min(i_y + YTile, seqlen);
+
             if(i_y < contextual_seqlen)
-                return ck_tile::make_tuple(0, max_uih_len);
-
-            index_t x_end = min(i_y + YTile, max_uih_len); // for lower-triangular masking, x <= y
-
-            return ck_tile::make_tuple(0, x_end);
+            {
+                if(i_y + YTile > max_uih_len)
+                {
+                    return ck_tile::make_tuple(0, x_end);
+                }
+                else
+                {
+                    return ck_tile::make_tuple(0, max_uih_len);
+                };
+            }
+            else
+            {
+                return ck_tile::make_tuple(0, x_end);
+            };
         };
     }
 
     CK_TILE_HOST constexpr bool IsTokenPairInsideMask(int row, int col)
     {
-        if(row >= max_uih_len || col >= max_uih_len)
-            return false;
+        // row_id/col_id is clamped from physical row/col according to contextual_seqlen and
+        // max_uih_len
+        int row_id = contextual_seqlen > 0 ? max(row - contextual_seqlen + 1, 0) : row;
+        int col_id = contextual_seqlen > 0 ? max(col - contextual_seqlen + 1, 0) : col;
 
-        if(row < contextual_seqlen)
+        row_id = min(row_id, max_id);
+        col_id = min(col_id, max_id);
+
+        if(contextual_seqlen > 0 && row_id == 0 && col_id < max_id)
             return true;
 
+        // use row_id/col_id to check the dist between two q/k token pair, token pairs on the
+        // diagonal line are always considerred
         if constexpr(IsMasking)
         {
-            bool result = (row >= col);
-
-            return result;
+            return (row_id > col_id) || (row == col);
         }
-
-        return true;
+        else
+        {
+            return (row_id != col_id) || (row == col);
+        };
     };
 
-    // masking codes in device don't have to compare row/col with max_uih_len, since
-    // buffer_load_xxx instruction is able to return zero for out-of-boundary access
     CK_TILE_DEVICE constexpr int IsTokenPairInsideMask(int row, int col)
     {
-        if(row < contextual_seqlen)
-            return 1;
-
-        if constexpr(IsMasking)
+        if(contextual_seqlen > 0)
         {
-            bool result = (row >= col);
+            // row_id/col_id is clamped from physical row/col according to contextual_seqlen and
+            // max_uih_len
+            int row_id = max(row - contextual_seqlen + 1, 0);
+            int col_id = max(col - contextual_seqlen + 1, 0);
 
-            return static_cast<int>(result);
+            row_id = min(row_id, max_id);
+            col_id = min(col_id, max_id);
+
+            if(row_id == 0 && col_id < max_id)
+                return 1;
+
+            // use row_id/col_id to check the dist between two q/k token pair, token pairs on the
+            // diagonal line are always considerred
+            if constexpr(IsMasking)
+            {
+                bool res = ((row_id > col_id) || (row == col));
+
+                return static_cast<int>(res);
+            }
+            else
+            {
+                bool res = ((row_id != col_id) || (row == col));
+
+                return static_cast<int>(res);
+            };
         }
+        else
+        {
+            // row_id/col_id is clamped from physical row/col according to contextual_seqlen and
+            // max_uih_len
+            int row_id = min(row, max_id);
+            int col_id = min(col, max_id);
 
-        return 1;
+            // use row_id/col_id to check the dist between two q/k token pair, token pairs on the
+            // diagonal line are always considerred
+            if constexpr(IsMasking)
+            {
+                bool res = ((row_id > col_id) || (row == col));
+
+                return static_cast<int>(res);
+            }
+            else
+            {
+                bool res = ((row_id != col_id) || (row == col));
+
+                return static_cast<int>(res);
+            };
+        }
     };
 
     // if the whole tile inside the masking area, no need for pixel-by-pixel checking
-    template <index_t TileWidth>
-    CK_TILE_DEVICE constexpr bool
-    IsFullTileInsideMask(index_t i_tile_top, index_t i_tile_left, number<TileWidth>) const
+    template <index_t TileWidth, index_t TileHeight>
+    CK_TILE_DEVICE constexpr bool IsFullTileInsideMask(index_t i_tile_top,
+                                                       index_t i_tile_left,
+                                                       number<TileWidth>,
+                                                       number<TileHeight>) const
     {
         if constexpr(kUseCausal)
         {
-            index_t i_tile_right = i_tile_left + TileWidth;
+            index_t i_tile_right  = i_tile_left + TileWidth;
+            index_t i_tile_bottom = i_tile_top + TileHeight;
 
-            if(i_tile_right > i_tile_top)
+            if(i_tile_right > i_tile_top ||
+               (i_tile_bottom > max_uih_len && i_tile_right > max_uih_len))
                 return false;
 
             return true;
         }
         else
         {
-            // need further check kPadSeqLenK in the masking context
+            index_t i_tile_right  = i_tile_left + TileWidth;
+            index_t i_tile_bottom = i_tile_top + TileHeight;
+
+            if(i_tile_bottom > max_uih_len && i_tile_right > max_uih_len)
+                return false;
+
             return true;
         };
     }
@@ -235,29 +426,15 @@ CK_TILE_HOST_DEVICE constexpr auto make_hstu_block_mask_with_local(int seqlen_,
                                                                    int max_attn_len_,
                                                                    int min_full_attn_seqlen_)
 {
-    auto max_uih_len_ = [&]() {
-        if(contextual_seqlen_ > 0)
-            return seqlen_ - (contextual_seqlen_ - 1) - num_target;
-        else
-            return seqlen_ - num_target;
-    }();
-
     return HstuBlockMaskType{
-        contextual_seqlen_, max_uih_len_, max_attn_len_, min_full_attn_seqlen_};
+        seqlen_, contextual_seqlen_, max_attn_len_, min_full_attn_seqlen_, num_target};
 };
 
 template <typename HstuBlockMaskType>
 CK_TILE_HOST_DEVICE constexpr auto
 make_hstu_block_mask_without_local(int seqlen_, int contextual_seqlen_, int num_target)
 {
-    auto max_uih_len_ = [&]() {
-        if(contextual_seqlen_ > 0)
-            return seqlen_ - (contextual_seqlen_ - 1) - num_target;
-        else
-            return seqlen_ - num_target;
-    }();
-
-    return HstuBlockMaskType{contextual_seqlen_, max_uih_len_};
+    return HstuBlockMaskType{seqlen_, contextual_seqlen_, num_target};
 };
 
 } // namespace ck_tile