File size: 3,046 Bytes
df13683 | 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 | # REVISED FROM
# https://github.com/shawntan/stickbreaking-attention/blob/main/stickbreaking_attention/sb_varlen/softplus.py
import triton
from triton import language as tl
from fla.utils import IS_NVIDIA
def _generate_softplus(num_pack):
template = """
.reg .pred p;
setp.gt.f32 p, ${in_reg}, 20.;
@p mov.f32 ${out_reg}, ${in_reg};
@!p mul.f32 ${out_reg}, ${in_reg}, 1.4426950408889634;
@!p ex2.approx.ftz.f32 ${out_reg}, ${out_reg};
@!p add.f32 ${out_reg}, ${out_reg}, 1.0;
@!p lg2.approx.ftz.f32 ${out_reg}, ${out_reg};
@!p mul.f32 ${out_reg}, ${out_reg}, 0.6931471805599453;
"""
out_str = ""
for i in range(num_pack):
inner_str = template.format(out_reg=i, in_reg=i + num_pack)
out_str += "{" + inner_str + "}\n"
# flatten out because torch.compile doesn't like newlines
out_str = " ".join(out_str.split("\n"))
return out_str
def _generate_softplus2(num_pack):
template = """
.reg .pred p;
setp.gt.f32 p, ${in_reg}, 15.;
@p mov.f32 ${out_reg}, ${in_reg};
@!p ex2.approx.ftz.f32 ${out_reg}, ${in_reg};
@!p add.f32 ${out_reg}, ${out_reg}, 1.0;
@!p lg2.approx.ftz.f32 ${out_reg}, ${out_reg};
"""
out_str = ""
for i in range(num_pack):
inner_str = template.format(out_reg=i, in_reg=i + num_pack)
out_str += "{" + inner_str + "}\n"
# flatten out because torch.compile doesn't like newlines
out_str = " ".join(out_str.split("\n"))
return out_str
def _generate_constraints(num_pack):
return ",".join("=r" for i in range(num_pack)) + "," + ",".join("r" for i in range(num_pack))
_NUM_REG = 1
s_softplus: tl.constexpr = tl.constexpr(_generate_softplus(_NUM_REG))
s_softplus2: tl.constexpr = tl.constexpr(_generate_softplus2(_NUM_REG))
s_constraints: tl.constexpr = tl.constexpr(_generate_constraints(_NUM_REG))
NUM_REG: tl.constexpr = tl.constexpr(_NUM_REG)
@triton.jit
def softplus_nv(x):
# equivalent to:
# return tl.where(x < 20.0, tl.math.log(1 + tl.math.exp(x)), x)
return tl.inline_asm_elementwise(
asm=s_softplus,
constraints=s_constraints,
pack=NUM_REG,
args=[
x,
],
dtype=tl.float32,
is_pure=True,
)
@triton.jit
def softplus_triton(x):
return tl.where(x < 20.0, tl.math.log(1 + tl.math.exp(x)), x)
@triton.jit
def softplus2_nv(x):
# equivalent to:
# return tl.where(x < 15.0, tl.math.log2(1 + tl.math.exp2(x)), x)
return tl.inline_asm_elementwise(
asm=s_softplus2,
constraints=s_constraints,
pack=NUM_REG,
args=[
x,
],
dtype=tl.float32,
is_pure=True,
)
@triton.jit
def softplus2_triton(x):
return tl.where(x < 15.0, tl.math.log2(1 + tl.math.exp2(x)), x)
if IS_NVIDIA:
softplus = softplus_nv
softplus2 = softplus2_nv
else:
softplus = softplus_triton
softplus2 = softplus2_triton
|