File size: 5,582 Bytes
e70aa91 | 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 | """Triton circuit kernels adapted from Qwenputer by Divyansh Agrawal.
MIT license: see third_party/QWENPUTER_LICENSE.
Changes: explicit per-row active masks to freeze halted batch rows.
"""
import triton
import triton.language as tl
@triton.jit
def _ffn_kernel(x_ptr, y_ptr, D, U, rpa, ixa, wa, rpb, ixb, wb, coef, outd, active_ptr,
BLOCK: tl.constexpr, KA: tl.constexpr, KB: tl.constexpr):
pid = tl.program_id(0)
bat = tl.program_id(1)
if tl.load(active_ptr + bat) == 0:
return
u = pid * BLOCK + tl.arange(0, BLOCK)
m = u < U
xb = x_ptr + bat * D
a0 = tl.load(rpa + u, mask=m, other=0)
na = tl.load(rpa + u + 1, mask=m, other=0) - a0
ka = tl.arange(0, KA)
ma = m[:, None] & (ka[None, :] < na[:, None])
ia = tl.load(ixa + a0[:, None] + ka[None, :], mask=ma, other=0)
av = tl.sum(tl.load(wa + a0[:, None] + ka[None, :], mask=ma, other=0.0)
* tl.load(xb + ia, mask=ma, other=0.0), axis=1)
b0 = tl.load(rpb + u, mask=m, other=0)
nb = tl.load(rpb + u + 1, mask=m, other=0) - b0
kb = tl.arange(0, KB)
mb = m[:, None] & (kb[None, :] < nb[:, None])
ib = tl.load(ixb + b0[:, None] + kb[None, :], mask=mb, other=0)
bv = tl.sum(tl.load(wb + b0[:, None] + kb[None, :], mask=mb, other=0.0)
* tl.load(xb + ib, mask=mb, other=0.0), axis=1)
val = tl.load(coef + u, mask=m, other=0.0) * av * tl.maximum(bv, 0.0)
od = tl.load(outd + u, mask=m, other=0)
tl.atomic_add(y_ptr + bat * D + od, val, mask=m)
@triton.jit
def _lag_load(xb, hb, i, m, D):
"""The lagged residual: index lag * D + d (lag 0: x, lags 1..3: the history)."""
lo = i < D
return tl.where(lo, tl.load(xb + i, mask=m & lo, other=0.0),
tl.load(hb + i - D, mask=m & (i >= D), other=0.0))
@triton.jit
def _rows(xb, hb, D, rp, ix, w, base, ch, KP: tl.constexpr):
"""relu of 128 sparse rows (row base + ch) over the lagged residual."""
r0 = tl.load(rp + base + ch)
n = tl.load(rp + base + ch + 1) - r0
kk = tl.arange(0, KP)
mk = kk[None, :] < n[:, None]
i = tl.load(ix + r0[:, None] + kk[None, :], mask=mk, other=0)
acc = tl.sum(tl.load(w + r0[:, None] + kk[None, :], mask=mk, other=0.0)
* _lag_load(xb, hb, i, mk, D), axis=1)
return tl.maximum(acc, 0.0)
@triton.jit
def _scalar(xb, hb, D, rp, ix, w, vh):
lane = tl.arange(0, 32)
s0 = tl.load(rp + vh)
n = tl.load(rp + vh + 1) - s0
mk = lane < n
i = tl.load(ix + s0 + lane, mask=mk, other=0)
return tl.sum(tl.load(w + s0 + lane, mask=mk, other=0.0) * _lag_load(xb, hb, i, mk, D), axis=0)
@triton.jit
def _head_kernel(x_ptr, h_ptr, S_ptr, y_ptr, D, V,
qrp, qix, qw, krp, kix, kw, vrp, vix, vw, zrp, zix, zw,
arp, aix, aw, brp, bix, bw, bconst, odim, ow, khead, nrows,
active_ptr, KQ: tl.constexpr, KK: tl.constexpr, KV: tl.constexpr, KZ: tl.constexpr,
QHOT: tl.constexpr):
vh = tl.program_id(0)
bat = tl.program_id(1)
if tl.load(active_ptr + bat) == 0:
return
ch = tl.arange(0, 128)
xb = x_ptr + bat * D
hb = h_ptr + bat * 3 * D
h = tl.load(khead + vh)
k = _rows(xb, hb, D, krp, kix, kw, h * 128, ch, KK)
v = _rows(xb, hb, D, vrp, vix, vw, vh * 128, ch, KV)
z = _rows(xb, hb, D, zrp, zix, zw, vh * 128, ch, KZ)
q = _rows(xb, hb, D, qrp, qix, qw, h * 128, ch, KQ)
bv = _scalar(xb, hb, D, brp, bix, bw, vh) + tl.load(bconst + vh)
Sb = S_ptr + (bat * V + vh) * 16384
nr = tl.load(nrows + vh) # > 0: the head decays, in its first nr rows
if nr > 0:
av = _scalar(xb, hb, D, arp, aix, aw, vh)
fac = tl.where(av == 0.0, 1.0, tl.where(av == 110.0, 0.5, 0.0))
if fac != 1.0:
for r in range(nr):
tl.store(Sb + r * 128 + ch, tl.load(Sb + r * 128 + ch) * fac)
kmax = tl.max(k, axis=0)
krow = tl.argmax(k, axis=0)
if (kmax > 0.0) & (bv == 1.0):
tl.store(Sb + krow * 128 + ch, v)
tl.debug_barrier()
# the read: the sum of the rows whose query is set (multi-hot queries have equal values)
o = tl.zeros([128], dtype=tl.float32)
for _ in range(QHOT):
qm = tl.max(q, axis=0)
qi = tl.argmax(q, axis=0)
if qm > 0.0:
o += tl.load(Sb + qi * 128 + ch)
q = tl.where(ch == qi, 0.0, q)
ref = tl.sum(tl.where(ch < 2, o * o, 0.0), axis=0)
norm = tl.math.sqrt_rn(ref / 128.0)
out = tl.where(ref > 0.0, tl.math.div_rn(o, tl.where(ref > 0.0, norm, 1.0)), 0.0)
wv = tl.load(ow + vh * 128 + ch)
dd = tl.load(odim + vh * 128 + ch)
tl.atomic_add(y_ptr + bat * D + dd, out * z * wv, mask=wv != 0.0)
@triton.jit
def _apply_kernel(x_ptr, y_ptr, h_ptr, D, active_ptr, HIST: tl.constexpr, BLOCK: tl.constexpr):
"""x += y, y = 0; with HIST, the lag history (h: [3, D] per row, h[0] = lag 1) moves on
first, the input x becoming lag 1."""
pid = tl.program_id(0)
bat = tl.program_id(1)
if tl.load(active_ptr + bat) == 0:
return
i = pid * BLOCK + tl.arange(0, BLOCK)
m = i < D
xo = x_ptr + bat * D + i
x = tl.load(xo, mask=m, other=0.0)
if HIST:
hb = h_ptr + bat * 3 * D + i
h1 = tl.load(hb + D, mask=m, other=0.0)
h0 = tl.load(hb, mask=m, other=0.0)
tl.store(hb + 2 * D, h1, mask=m)
tl.store(hb + D, h0, mask=m)
tl.store(hb, x, mask=m)
yo = y_ptr + bat * D + i
tl.store(xo, x + tl.load(yo, mask=m, other=0.0), mask=m)
tl.store(yo, tl.zeros([BLOCK], dtype=tl.float32), mask=m)
|