| """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) |
| 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() |
| |
| 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) |
|
|
|
|
