Kernels:

---

aws-neuron

/

vit-encoder-neuron-kernels

Like 1

Follow

AWS Inferentia and Trainium 190

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