| """CTA-local routing-mask helpers shared by the CuTe architecture backends.""" |
|
|
| import cutlass |
| import cutlass.cute as cute |
| from cutlass import Float32, Int32, const_expr |
| from cutlass._mlir.dialects import llvm |
| from cutlass.cutlass_dsl import T, dsl_user_op |
|
|
|
|
| @dsl_user_op |
| def sol_attn_bfind_b32( |
| value: Int32, |
| *, |
| loc=None, |
| ip=None, |
| ) -> Int32: |
| return Int32( |
| llvm.inline_asm( |
| T.i32(), |
| [Int32(value).ir_value(loc=loc, ip=ip)], |
| "bfind.u32 $0, $1;", |
| "=r,r", |
| has_side_effects=False, |
| is_align_stack=False, |
| ) |
| ) |
|
|
|
|
| @dsl_user_op |
| def sol_attn_popc_b32( |
| value: Int32, |
| *, |
| loc=None, |
| ip=None, |
| ) -> Int32: |
| return Int32( |
| llvm.inline_asm( |
| T.i32(), |
| [Int32(value).ir_value(loc=loc, ip=ip)], |
| "popc.b32 $0, $1;", |
| "=r,r", |
| has_side_effects=False, |
| is_align_stack=False, |
| ) |
| ) |
|
|
|
|
| @cute.jit |
| def _mask_word( |
| mask0: Int32, |
| mask1: Int32, |
| mask2: Int32, |
| mask3: Int32, |
| word: Int32, |
| ) -> Int32: |
| result = mask0 |
| if word == Int32(1): |
| result = mask1 |
| if word == Int32(2): |
| result = mask2 |
| if word == Int32(3): |
| result = mask3 |
| return result |
|
|
|
|
| @cute.jit |
| def _test_exact_bit( |
| mask0: Int32, |
| mask1: Int32, |
| mask2: Int32, |
| mask3: Int32, |
| offset: Int32, |
| ) -> cutlass.Boolean: |
| word = offset // Int32(32) |
| bit = offset - word * Int32(32) |
| return ( |
| _mask_word(mask0, mask1, mask2, mask3, word) |
| & (Int32(1) << bit) |
| ) != Int32(0) |
|
|
|
|
| @cute.jit |
| def sol_attn_test_exact_bit_limited_words( |
| mask0: Int32, |
| mask1: Int32, |
| mask2: Int32, |
| mask3: Int32, |
| offset: Int32, |
| group_words: cutlass.Constexpr[int], |
| ) -> cutlass.Boolean: |
| bit = offset & Int32(31) |
| if const_expr(group_words == 1): |
| return (mask0 & (Int32(1) << bit)) != Int32(0) |
| if const_expr(group_words == 2): |
| word = mask0 |
| if offset >= Int32(32): |
| word = mask1 |
| return (word & (Int32(1) << bit)) != Int32(0) |
| if const_expr(group_words == 3): |
| index = offset // Int32(32) |
| word = mask0 |
| if index == Int32(1): |
| word = mask1 |
| if index == Int32(2): |
| word = mask2 |
| return (word & (Int32(1) << bit)) != Int32(0) |
| return _test_exact_bit(mask0, mask1, mask2, mask3, offset) |
|
|
|
|
| @cute.jit |
| def sol_attn_set_exact_bit( |
| mask0: Int32, |
| mask1: Int32, |
| mask2: Int32, |
| mask3: Int32, |
| offset: Int32, |
| ): |
| word = offset // Int32(32) |
| bit_value = Int32(1) << (offset - word * Int32(32)) |
| if word == Int32(0): |
| mask0 = mask0 | bit_value |
| if word == Int32(1): |
| mask1 = mask1 | bit_value |
| if word == Int32(2): |
| mask2 = mask2 | bit_value |
| if word == Int32(3): |
| mask3 = mask3 | bit_value |
| return mask0, mask1, mask2, mask3 |
|
|
|
|
| @cute.jit |
| def sol_attn_route_is_exact( |
| q_block: Int32, |
| kv_block: Int32, |
| column_mean: Float32, |
| threshold: Float32, |
| valid: cutlass.Boolean, |
| ) -> cutlass.Boolean: |
| distance = q_block - kv_block |
| if distance < Int32(0): |
| distance = Int32(0) - distance |
| return ((column_mean > threshold) or distance <= Int32(1)) and valid |
|
|
|
|
| @cute.jit |
| def sol_attn_mask_word_constexpr( |
| mask0: Int32, |
| mask1: Int32, |
| mask2: Int32, |
| mask3: Int32, |
| word: cutlass.Constexpr[int], |
| ) -> Int32: |
| if const_expr(word == 0): |
| return mask0 |
| if const_expr(word == 1): |
| return mask1 |
| if const_expr(word == 2): |
| return mask2 |
| return mask3 |
|
|
|
|
| __all__ = [ |
| "sol_attn_bfind_b32", |
| "sol_attn_mask_word_constexpr", |
| "sol_attn_popc_b32", |
| "sol_attn_route_is_exact", |
| "sol_attn_set_exact_bit", |
| "sol_attn_test_exact_bit_limited_words", |
| ] |
|
|
