v1.6: two images per producer pass in attention_block (LNC=2 path; LNC=1 single-group models)

a827ebd 2 days ago

14 kB

	"""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 ``[BN, H]`` fp32 residual stream (rows of image b at bN)."""

	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: [BN, H] fp32 (padded rows zero). Returns the pre-final-norm residual stream [BN, 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