| """NonCausalFast v3: v2 plus three changes that trade fp32 exactness (unneeded at a 30 dB target) for speed. |
| |
| * input: latent de-normalisation, post_quant_conv (1x1x1 conv) and x_embedder are one 24 -> D affine map, one GEMM; |
| * head: final LayerNorm writes bf16 once and the head is a single bf16 GEMM (v2 splits hi/lo for fp32 accuracy); |
| * everything else (window / global FA4 varlen attention, fused qkv post kernel, quack SwiGLU) is unchanged. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import triton |
| import triton.language as tl |
|
|
| from noncausal_fast import BF16, NonCausalFast, _resid_norm_kernel, qgemm, qgemm_act, _flash_attn_fwd, _nc_qkv_post_kernel |
|
|
|
|
| @triton.jit |
| def _resid_ln_bf16_kernel(x, y, scale, w, b, out, 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) |
| tl.store(out + r.to(tl.int64) * D + c, (xc * rs * tl.load(w + c) + tl.load(b + c)).to(tl.bfloat16)) |
|
|
|
|
| class NonCausalFast3(NonCausalFast): |
| def __init__(self, ref, T, H, W, *, cuda_graph: bool = True): |
| r = ref |
| dev = r.register_tokens.device |
| |
| Wc = r.post_quant_conv.weight.detach().float().reshape(24, 24) |
| bc = r.post_quant_conv.bias.detach().float() |
| We = r.x_embedder.weight.detach().float() |
| be = r.x_embedder.bias.detach().float() |
| std, mean = r.latents_std.float().to(dev), r.latents_mean.float().to(dev) |
| A = We @ Wc * std[None, :] |
| c = We @ (Wc @ mean + bc) + be |
| self._A, self._c = A.contiguous(), c.contiguous() |
| self._whead_bf = r.proj_out.weight.detach().to(BF16).contiguous() |
| self._hln = None |
| super().__init__(ref, T, H, W, cuda_graph=cuda_graph) |
|
|
| @torch.no_grad() |
| def _step(self): |
| r, T, H, W, N, M = self.r, self.T, self.H, self.W, self.N, self.M |
| if self._hln is None: |
| self._hln = torch.zeros(N, self.D, dtype=BF16, device=self.dev) |
| self._zt = torch.zeros(N, 24, device=self.dev) |
| self._zt.copy_(self.z_in[0].permute(0, 2, 3, 1).reshape(N, 24)) |
| torch.addmm(self._c, self._zt, self._A.t(), out=self.x[:N]) |
| self.x[:N].copy_(self.x[:N].to(BF16)) |
| self.x[N:].copy_(self.g0) |
| lw0 = self.lw[0] |
| _resid_norm_kernel[(M,)](self.x, self.ybuf, self.ident, lw0["s1"], lw0["n1"], self.h, D=self.D, EPS=lw0["eps1"], ADD=False, num_warps=8) |
| for li in range(self.L): |
| lw, g = self.lw[li], self.geo[self.shift[li]] |
| torch.addmm(lw["bqkv"], self.h, lw["wqkv"].t(), out=self.qkv) |
| _nc_qkv_post_kernel[(M,)](self.qkv, self.qo, self.kp, self.vp, self.kgp, self.vgp, g["cos"], g["sin"], g["qpos"], |
| g["kpos"], self.kgpos, g["gwin"], N, self.S, g["maxw"], |
| HEADS=self.heads, DH=self.dh, RD=g["rd"], EPS=lw["epsq"], num_warps=4) |
| _flash_attn_fwd(self.qo[:N], self.kp[:g["total_k"]], self.vp[:g["total_k"]], cu_seqlens_q=g["cu_q"], cu_seqlens_k=g["cu_k"], |
| max_seqlen_q=g["max_q"], max_seqlen_k=g["max_k"], out=self.obuf[:N].view(N, self.heads, self.dh)) |
| if li < self.L - 1: |
| _flash_attn_fwd(self.qo[N:], self.kgp, self.vgp, cu_seqlens_q=self.g_cuq, cu_seqlens_k=self.g_cuk, |
| max_seqlen_q=self.S, max_seqlen_k=self.g_maxk, out=self.obuf[N:].view(M - N, self.heads, self.dh)) |
| qgemm(self.obuf, lw["wo"].t(), bias=lw["bo"], out=self.ybuf) |
| _resid_norm_kernel[(M,)](self.x, self.ybuf, g["src"], lw["s1"], lw["n2"], self.h, D=self.D, EPS=lw["eps2"], ADD=True, num_warps=8) |
| qgemm_act(self.h, lw["w1"].t(), bias=lw["b1"], activation="swiglu", store_preact=False, postact_out=self.abuf) |
| torch.addmm(lw["b2"], self.abuf, lw["w2"].t(), out=self.ybuf) |
| if li < self.L - 1: |
| nxt = self.lw[li + 1] |
| _resid_norm_kernel[(M,)](self.x, self.ybuf, self.ident, lw["s2"], nxt["n1"], self.h, D=self.D, EPS=nxt["eps1"], ADD=True, num_warps=8) |
| else: |
| _resid_ln_bf16_kernel[(N,)](self.x, self.ybuf, lw["s2"], self.ln_w, self.ln_b, self._hln, D=self.D, EPS=self.ln_eps, num_warps=8) |
| torch.mm(self._hln, self._whead_bf.t(), out_dtype=torch.float32, out=self.G) |
|
|
