From 124539e12375ea7966b0bfa51ad0d8a9c8b3d709 Mon Sep 17 00:00:00 2001 From: Qianfeng Zhang Date: Tue, 13 May 2025 10:37:19 +0000 Subject: [PATCH] Update the rules of hstu masking --- .../hstu_attention_fwd_kernel.hpp | 6 +- .../hstu_attention_fwd_pipeline.hpp | 6 +- .../18_hstu_attention/hstu_block_masking.hpp | 385 +++++++++++++----- 3 files changed, 288 insertions(+), 109 deletions(-) diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp index f290bac92d..ac1462ae47 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp +++ b/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( 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{}, number<1>{}); @@ -616,7 +616,7 @@ struct HstuAttentionFwdKernel const auto k_dram = [&]() { const auto k_dram_naive = make_naive_tensor_view( 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{}, number<1>{}); @@ -631,7 +631,7 @@ struct HstuAttentionFwdKernel { const auto v_dram_naive = make_naive_tensor_view( 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{}, number<1>{}); 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 51c1068bdb..c79732b780 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 @@ -362,8 +362,10 @@ struct HstuAttentionFwdPipelineQRKSVS } else // kUseCausal=true, kUseLocal=false { - if(!mask.IsFullTileInsideMask( - q_origin.at(number<0>{}), seqlen_k_curr, number{})) + if(!mask.IsFullTileInsideMask(q_origin.at(number<0>{}), + seqlen_k_curr, + number{}, + number{})) { constexpr auto s_spans = SaccBlockTileType::get_distributed_spans(); sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) { 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 3eb2199502..22e3349105 100644 --- a/example/ck_tile/18_hstu_attention/hstu_block_masking.hpp +++ b/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, number) 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(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(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(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(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(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(res1 || res2); + } } }; // if the whole tile inside the masking area, no need for pixel-by-pixel checking - template - CK_TILE_DEVICE constexpr bool - IsFullTileInsideMask(index_t i_tile_top, index_t i_tile_left, number) const + template + CK_TILE_DEVICE constexpr bool IsFullTileInsideMask(index_t i_tile_top, + index_t i_tile_left, + number, + number) 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(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(res); + } + else + { + bool res = ((row_id != col_id) || (row == col)); + + return static_cast(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(res); + } + else + { + bool res = ((row_id != col_id) || (row == col)); + + return static_cast(res); + }; + } }; // if the whole tile inside the masking area, no need for pixel-by-pixel checking - template - CK_TILE_DEVICE constexpr bool - IsFullTileInsideMask(index_t i_tile_top, index_t i_tile_left, number) const + template + CK_TILE_DEVICE constexpr bool IsFullTileInsideMask(index_t i_tile_top, + index_t i_tile_left, + number, + number) 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 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