File size: 3,822 Bytes
8e9f35a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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",
]
|