"""Model-agnostic ViT encoder loop on the NKI kernels. A pre-norm transformer block x = x + ls1 * attn(LN1(x)) attn: q/k/v projections, RoPE on q/k, softmax, out proj x = x + ls2 * ffn(LN2(x)) ffn: GELU MLP or SwiGLU is described once per block as a :class:`BlockWeights` (plain fp32 tensors, any source), packed once by :func:`pack_block` (LayerNorm affines and LayerScale folded into the weights), and run by :class:`EncoderRunner` on a stacked ``[B*N, H]`` row buffer. Kernel choice per block (all automatic): * attention: LNC=1 -> one program over all head groups, with the residual add + LN2 in its epilogue; LNC=2 -> per-group programs whose partials the residual/LN kernel sums. * FFN: GELU -> fused MLP; SwiGLU -> weight-streaming SwiGLU. Both also do the residual add and the NEXT block's LN1 in their epilogue, so the stand-alone residual/LN kernel runs only before block 0 (and after attention at LNC=2). INFERENCE ONLY (no autograd). Constraints: hidden size % 128 == 0, head_dim in {64, 128}, tokens per image <= 2048 after padding to a multiple of 128 * LNC. """ from __future__ import annotations from dataclasses import dataclass from typing import Optional import torch import importlib as _il import os as _osr _PACK2 = _osr.environ.get("VIT_ENC_PACK2", "1") == "1" # 2 images per producer pass (N <= 256, LNC=1) _ab = _il.import_module(__package__ + ".nki_kernels.attention_block") _ab1 = _il.import_module(__package__ + ".nki_kernels.attention_block_l1") _ab1ln = _il.import_module(__package__ + ".nki_kernels.attention_block_l1_ln") from .nki_kernels.attention_block import LNC from .nki_kernels.mlp import nki_mlp_k4_xp from .nki_kernels.mlp import pack_b1 as _mlp_pack_b1 from .nki_kernels.mlp_res_ln import nki_mlp_res_ln import os as _osx _MLP_FP8 = _osx.environ.get("VIT_ENC_MLP_FP8", "0") == "1" # EXPERIMENTAL: FP8 e4m3 GELU MLP (see README) if _MLP_FP8: _mf8 = _il.import_module(__package__ + ".nki_kernels.mlp_fp8") from .nki_kernels.res_layernorm import nki_res_layernorm_fold from .nki_kernels.swiglu_stream import nki_swiglu_stream, nki_swiglu_stream_ln __all__ = ["BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC", "calibrate_fp8"] _SBUF_BUDGET = 180 * 1024 # per-partition bytes the LNC=1 resident-weights decision allows _SW_MIN_ROWS = 256 # below this many rows per core, SwiGLU runs in torch (measured faster) _L1_MAX_N = 1280 # LNC=1 all-group attention kernels verified up to this many tokens per # image; beyond it (e.g. 2048-token video clips) they exceed SBUF and # the per-group kernel + residual/LN kernel are used instead _RESIDENT_MAX_N = 1280 # LNC=2 attention: weights staged once per launch only up to this many # tokens per image (2048 tokens + resident weights exceed SBUF) @dataclass class BlockWeights: """One transformer block, as plain fp32 CPU tensors in ``nn.Linear`` layout ([out, in]).""" ln1_w: torch.Tensor ln1_b: torch.Tensor ln1_eps: float wq: torch.Tensor bq: Optional[torch.Tensor] wk: torch.Tensor bk: Optional[torch.Tensor] wv: torch.Tensor bv: Optional[torch.Tensor] wo: torch.Tensor bo: Optional[torch.Tensor] ln2_w: torch.Tensor ln2_b: torch.Tensor ln2_eps: float # GELU MLP: fc1 [I, H], fc2 [H, I]. SwiGLU: gate [I, H], up [I, H], down [H, I]. fc1_w: Optional[torch.Tensor] = None fc1_b: Optional[torch.Tensor] = None fc2_w: Optional[torch.Tensor] = None fc2_b: Optional[torch.Tensor] = None gate_w: Optional[torch.Tensor] = None gate_b: Optional[torch.Tensor] = None up_w: Optional[torch.Tensor] = None up_b: Optional[torch.Tensor] = None down_w: Optional[torch.Tensor] = None down_b: Optional[torch.Tensor] = None ls1: Optional[torch.Tensor] = None # LayerScale (None = 1) ls2: Optional[torch.Tensor] = None scale: Optional[float] = None # attention scale (None = head_dim ** -0.5) def _f(t): return None if t is None else t.detach().float().cpu() def _z(t, n): return torch.zeros(n) if t is None else t def pack_block(bw: BlockWeights, num_heads: int, rope: str, device) -> dict: """Fold and pack one block for the kernels. ``rope`` in {"half", "interleaved"}.""" wq, wk, wv, wo = _f(bw.wq), _f(bw.wk), _f(bw.wv), _f(bw.wo) H = wq.shape[1] F_ = wq.shape[0] D = F_ // num_heads gsz = 512 // D # heads per kernel program nh_p = ((num_heads + gsz - 1) // gsz) * gsz # padded head count Fp = nh_p * D ngrp = nh_p // gsz g1, t1, g2, t2 = _f(bw.ln1_w), _f(bw.ln1_b), _f(bw.ln2_w), _f(bw.ln2_b) bq, bk, bv = _z(_f(bw.bq), F_), _z(_f(bw.bk), F_), _z(_f(bw.bv), F_) ls1 = _f(bw.ls1) if bw.ls1 is not None else torch.ones(H) ls2 = _f(bw.ls2) if bw.ls2 is not None else torch.ones(H) bo = _z(_f(bw.bo), H) * ls1 wo = wo * ls1[:, None] padr = lambda w: torch.cat([w, w.new_zeros(Fp - w.shape[0], *w.shape[1:])], 0) p = {"D": D, "heads": num_heads, "nh_p": nh_p, "ngrp": ngrp, "H": H, "eps1": float(bw.ln1_eps), "eps2": float(bw.ln2_eps), "scale": float(bw.scale if bw.scale is not None else D ** -0.5)} p["wqk"], p["bqk"] = _ab.pack_weights(padr(wq * g1[None]), padr(wk * g1[None]), padr(bq + wq @ t1), padr(bk + wk @ t1), nh_p, D, rope) p["wvt"], p["bv"] = _ab.pack_v_weights(padr(wv * g1[None]), padr(bv + wv @ t1)) p["wot"] = torch.cat([wo, wo.new_zeros(H, Fp - F_)], 1).t().contiguous().to(torch.bfloat16) p["bo"] = (bo / ngrp).reshape(1, H).contiguous().to(torch.bfloat16) # per-group partials p["bo_full"] = bo.reshape(1, H).contiguous().to(torch.bfloat16) # all-group kernel if bw.gate_w is not None: # SwiGLU gw, uw, dw = _f(bw.gate_w), _f(bw.up_w), _f(bw.down_w) I = gw.shape[0] Ip = ((I + 127) // 128) * 128 pi = lambda t: torch.cat([t, t.new_zeros(Ip - t.shape[0], *t.shape[1:])], 0) p["swiglu"] = True p["s_w1t"] = pi(gw * g2[None]).t().contiguous().to(torch.bfloat16) # gate (silu) p["s_b1u"] = pi(_z(_f(bw.gate_b), I) + gw @ t2).contiguous() # [Ip] for the torch fallback p["s_b1"] = _mlp_pack_b1(p["s_b1u"]) # [128, Ip/128], one-DMA p["s_wgt"] = pi(uw * g2[None]).t().contiguous().to(torch.bfloat16) # up p["s_bgu"] = pi(_z(_f(bw.up_b), I) + uw @ t2).contiguous() # [Ip] for the torch fallback p["s_bg"] = _mlp_pack_b1(p["s_bgu"]) # [128, Ip/128], one-DMA p["s_w2t"] = pi((dw * ls2[:, None]).t().contiguous()).contiguous().to(torch.bfloat16) p["s_b2"] = (_z(_f(bw.down_b), H) * ls2).contiguous() else: # GELU MLP w1, w2 = _f(bw.fc1_w), _f(bw.fc2_w) p["swiglu"] = False p["w1t"] = (w1 * g2[None]).t().contiguous().to(torch.bfloat16) # Pre-pack fc1 bias to [128, H_mlp/128] so K4/K4+LN load it in one DMA (see mlp.pack_b1). p["b1"] = _mlp_pack_b1((_z(_f(bw.fc1_b), w1.shape[0]) + w1 @ t2).contiguous()) p["w2t"] = (w2 * ls2[:, None]).t().contiguous().to(torch.bfloat16) if _MLP_FP8: p["w1p8"], p["sw1"] = _mf8.pack_fp8_weight((w1 * g2[None]).t().contiguous()) p["w2p8"], p["sw2"] = _mf8.pack_fp8_weight((w2 * ls2[:, None]).t().contiguous()) p["g8"] = 1.0 / 32 # GELU-output scale; per-layer calibration recommended (calibrate_fp8) p["b2"] = (_z(_f(bw.fc2_b), H) * ls2).contiguous() return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in p.items()} def rope_tables(cos: torch.Tensor, sin: torch.Tensor, n_prefix: int, n_pad: int, num_heads: int, D: int, rope: str, device) -> torch.Tensor: """Kernel RoPE tables for a ``[n_prefix + n_rot, ...]`` token layout padded to ``n_pad`` rows. ``cos``/``sin``: ``[n_rot, D]`` the model's own tables for the rotated tokens (prefix tokens such as CLS/registers and padding get cos=1, sin=0). For ``rope="interleaved"`` pass the SIGNED sin table (see ``pack_tables``). """ C = torch.ones(n_pad, D) S = torch.zeros(n_pad, D) n_rot = cos.shape[0] C[n_prefix:n_prefix + n_rot] = cos.float().cpu() S[n_prefix:n_prefix + n_rot] = sin.float().cpu() gsz = 512 // D nh_p = ((num_heads + gsz - 1) // gsz) * gsz t = _ab.pack_tables(C, S, nh_p, D, rope) if _osr.environ.get("VIT_ENC_ROPE_BF16", "0") == "1": # EXPERIMENTAL bf16 RoPE (tables in bf16) t = t.to(torch.bfloat16) return t.to(device) def calibrate_fp8(runner, gelu_amax, headroom: float = 2.0): """EXPERIMENTAL (VIT_ENC_MLP_FP8=1). Set each layer's power-of-two GELU-output scale from the observed max |GELU output| per layer (list, one per block): scale = 2^floor(log2(128 / (amax * headroom))). trn2's e4m3 saturates at 240; uncalibrated activations above that turn into inf/NaN.""" import math for p, a in zip(runner.L, gelu_amax): p["g8"] = 2.0 ** math.floor(math.log2(128.0 / max(a * headroom, 1e-6))) class EncoderRunner: """Runs packed blocks on a stacked ``[B*N, H]`` fp32 residual stream (rows of image b at b*N).""" def __init__(self, packed: list): self.L = packed p0 = packed[0] self.D, self.H, self.ngrp = p0["D"], p0["H"], p0["ngrp"] self.plan = [] # (group base, grid) per LNC=2 attention call g = 0 while g < self.ngrp: gr = 2 if (LNC == 2 and self.ngrp - g >= 2) else 1 self.plan.append((g, gr)) g += gr def _res(self, mod, N): return int(mod.sbuf_estimate(self.H, self.ngrp * 512, N, self.D, True) <= _SBUF_BUDGET) def _all_group(self, N): return LNC == 1 and self.ngrp > 1 and N <= _L1_MAX_N def _attn_parts(self, z1, p, tab, n_valid, B): if self._all_group(z1.shape[0] // B): N = z1.shape[0] // B return _ab1.attention_block_l1(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], p["scale"], self.D, n_valid, B, self._res(_ab1, N)) outs = [_ab.attention_block(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo"], p["scale"], self.D, n_valid, hb, gr, B, int(z1.shape[0] // B <= _RESIDENT_MAX_N), int(_PACK2 and B % 2 == 0 and 2 * (z1.shape[0] // B) <= 512)) for hb, gr in self.plan] P = outs[0] if len(outs) == 1 else torch.cat(outs, 0) if P.shape[0] > 4: # residual/LN kernel takes <= 4 branches P = P.float().sum(0, keepdim=True).to(torch.bfloat16) return P def _ffn(self, z2, p): if p["swiglu"]: if z2.shape[0] // LNC >= _SW_MIN_ROWS: return nki_swiglu_stream(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"], p["s_w2t"], p["s_b2"]).unsqueeze(0) # torch fallback needs the UNPACKED [Ip] biases (v1.2 regression: it used the packed ones). g = torch.nn.functional.silu((z2 @ p["s_w1t"]).float() + p["s_b1u"]) u = (z2 @ p["s_wgt"]).float() + p["s_bgu"] return (((g * u).to(torch.bfloat16) @ p["s_w2t"]).float() + p["s_b2"]).to(torch.bfloat16).unsqueeze(0) if _MLP_FP8: return _mf8.nki_mlp_fp8(z2, p["w1p8"], p["b1"], p["w2p8"], p["b2"], p["sw1"], p["sw2"], p["g8"]).unsqueeze(0) return nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).unsqueeze(0) def _ffn_fused(self, z2, p, x, eps_next): if p["swiglu"]: if z2.shape[0] // LNC < _SW_MIN_ROWS: return None return nki_swiglu_stream_ln(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"], p["s_w2t"], p["s_b2"], x, eps_next) if _MLP_FP8: return _mf8.nki_mlp_fp8_res_ln(z2, p["w1p8"], p["b1"], p["w2p8"], p["b2"], x, eps_next, p["sw1"], p["sw2"], p["g8"]) return nki_mlp_res_ln(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], x, eps_next) def __call__(self, x: torch.Tensor, tab: torch.Tensor, n_valid: int, B: int) -> torch.Tensor: """x: [B*N, H] fp32 (padded rows zero). Returns the pre-final-norm residual stream [B*N, H] fp32.""" H = self.H branch = x.new_zeros(1, x.shape[0], H, dtype=torch.bfloat16) z1 = None nL = len(self.L) for li, p in enumerate(self.L): if z1 is None: x, z1 = nki_res_layernorm_fold(x, branch, p["eps1"]) N = x.shape[0] // B if self._all_group(N): x, z2 = _ab1ln.attention_block_l1_ln(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], x, p["scale"], self.D, n_valid, p["eps2"], B, self._res(_ab1ln, N), int(_PACK2 and B % 2 == 0 and 2 * N <= 512)) else: parts = self._attn_parts(z1, p, tab, n_valid, B) x, z2 = nki_res_layernorm_fold(x, parts, p["eps2"]) fused = self._ffn_fused(z2, p, x, self.L[li + 1]["eps1"]) if li + 1 < nL else None if fused is not None: x, z1 = fused branch = None else: branch = self._ffn(z2, p) z1 = None if branch is not None: x = x + branch[0].float() return x