File size: 1,866 Bytes
775b8fb | 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 | """Fused Triton kernels shared by the non-causal engine.
_resid_norm_kernel x[r] += y[src[r]] * scale (optional), then RMSNorm(x[r]) * w -> bf16
_resid_ln_split_kernel x[r] += y[r] * scale, LayerNorm(x[r]) * w + b, written as [hi | lo | hi] bf16 for an fp32-accurate head GEMM
_interleave reorder SwiGLU rows [gate | up] -> gate0, up0, gate1, up1, ... (quack's gated epilogue layout)
"""
import torch
import triton
import triton.language as tl
@triton.jit
def _resid_norm_kernel(x, y, src, scale, w, hout, D: tl.constexpr, EPS: tl.constexpr, ADD: tl.constexpr):
r = tl.program_id(0)
c = tl.arange(0, D)
xr = x + r.to(tl.int64) * D
xv = tl.load(xr + c)
if ADD:
sr = tl.load(src + r).to(tl.int64)
yv = tl.load(y + sr * D + c).to(tl.float32)
xv = xv + yv * tl.load(scale + c)
tl.store(xr + c, xv)
rs = tl.rsqrt(tl.sum(xv * xv, 0) / D + EPS)
tl.store(hout + r.to(tl.int64) * D + c, (xv * rs * tl.load(w + c)).to(tl.bfloat16))
@triton.jit
def _resid_ln_split_kernel(x, y, scale, w, b, aout, D: tl.constexpr, EPS: tl.constexpr):
r = tl.program_id(0)
c = tl.arange(0, D)
xv = tl.load(x + r.to(tl.int64) * D + c)
yv = tl.load(y + r.to(tl.int64) * D + c).to(tl.float32)
xv = xv + yv * tl.load(scale + c)
mean = tl.sum(xv, 0) / D
xc = xv - mean
rs = tl.rsqrt(tl.sum(xc * xc, 0) / D + EPS)
o = xc * rs * tl.load(w + c) + tl.load(b + c)
hi = o.to(tl.bfloat16)
lo = (o - hi.to(tl.float32)).to(tl.bfloat16)
ar = aout + r.to(tl.int64) * (3 * D)
tl.store(ar + c, hi)
tl.store(ar + D + c, lo)
tl.store(ar + 2 * D + c, hi)
def _interleave(w):
"""rows [gate | up] -> gate0, up0, gate1, up1, ... (quack's gated-epilogue layout)."""
I = w.shape[0] // 2
return torch.stack([w[:I], w[I:]], 1).reshape(2 * I, *w.shape[1:]).contiguous()
|