"""K4 -- fused MLP (fc1 + gelu + fc2), x-prestage ("XP") in prologue mode. WHAT IT COMPUTES ---------------- ``out = gelu(x @ w1.T + b1) @ w2.T + b2``, with the rows sharded across LNC cores and the weights pre-transposed to bf16. Returns ``[R, K]`` bf16. WHY IT IS SHAPED THIS WAY ------------------------- This is the effective shipping kernel: the ``b1pre`` MLP (fc1 bias staged in the prologue so every gelu's bias is resident from t=0) with the x-prestage applied. The schedule was rebuilt from the device profile (external contribution by jlongen-aws, PR #2 on ``jburtoft/vjepa2-neuron-kernels``). Four design choices carry the performance; the instruction order and numerics are unchanged, so the output is bit-identical to the previous version: * x BEFORE weights: row-block 0's ``x`` load + transpose is issued first, then W1, then the later row blocks' ``x``, then W2. The first matmul no longer waits on ~16 MB of weights it does not need yet (the first fc1 matmul used to retire at ~57 us). * b1 in ONE DMA on the Scalar engine's HW-DGE queue: ``b1`` is host-packed to ``[128, H/128]`` (:func:`pack_b1`) and lands in a single ~16 KB DMA instead of ``n_h`` ``[128, 1]`` fp32 DMAs that trickled in and stalled each gelu (and, through its PSUM bank, the fc1 matmuls behind it). Routing it on ``engine=scalar`` HW-DGE keeps gelu's wait off the W1/W2 traffic. * row-block 0's ``x`` is transposed on the TENSOR engine (``nc_transpose``) in the prologue, while the Tensor engine is otherwise dark in the pre-K4 seam. Later row blocks keep ``dma_transpose`` (they stream under rb0's matmuls). * x prestaged in the prologue: EVERY row block's ``x`` staging is issued before the ``rb`` loop, in row-block order, so a later block's serialised DMA chain runs inside the pre-K4 seam instead of in the idle window at the rb boundary. Appending later blocks after rb=0's stages means rb=0's readiness cannot regress -- the first matmul cannot be delayed. The x staging is emitted through a module-level emitter (``_stage_xt``) so the ``x`` load, the ``nc_transpose`` and the ``dma_transpose`` each keep exactly one source site. PSUM: fc1 uses 1 bank, fc2 uses up to 4 banks; the ``y1_psum`` / ``out_psums`` tiles are declared inside their loops so the ring never accumulates onto a stale PSUM value. The b2 bias broadcast is kept (removing it regressed the sibling kernel by 23-30%). Weight DMAs stay one-per-tile: merging them into a single access-pattern DMA lost ~8 us (one completion signal per 8 MB weight). This kernel is LNC-parametric (``LNC`` from ``NEURON_LOGICAL_NC_CONFIG``) and keeps the ``min()`` remainder handling, so it is NOT restricted to exact-multiple shapes. INFERENCE ONLY: no autograd rule. """ from __future__ import annotations import math import os as _os import torch from torch_neuronx import nki_op import nki import nki.isa as nisa import nki.language as nl #: PE array partition limit, and the moving-operand free-axis limit. _P_MAX = 128 _F_MOV_MAX = 512 _BF16 = nl.bfloat16 #: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG. LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2")) _INV_SQRT2 = 1.0 / math.sqrt(2.0) _INV_SQRT_2PI = 1.0 / math.sqrt(2.0 * math.pi) def _div_ceil(n: int, d: int) -> int: """Ceiling division (never write (n + d - 1) // d inline).""" return (n + d - 1) // d def _kernel_assert(condition: bool, error_text: str): """Assert with an NKI-formatted error message.""" assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}" def _stage_xt(x_hbm, r_abs, pr, n_k, K, rb_size, tensor_transpose): """Emit one row block's x staging and return its ``n_k`` transposed bf16 tiles. ``n_rs`` loads of ``[<=128, K]`` bf16, each followed by ``n_k`` transposes into the row block's ``xt_tiles``. The transpose runs on the Tensor engine (``nc_transpose``, via a PSUM tile) when ``tensor_transpose`` is set -- used for row block 0 while the Tensor engine is dark in the pre-K4 seam -- and on the DMA engine (``dma_transpose``) otherwise. Kept as a module-level emitter so the ``x`` load and each transpose keep a single source site. Args: x_hbm: [R, K] @ HBM bf16, the MLP input. r_abs: absolute starting row of this row block. pr: number of rows in this row block (<= rb_size). n_k: number of K tiles (fc1 contraction dimension tiles). K: hidden-in dimension. rb_size: row-block width (free axis of the transposed tiles). tensor_transpose: transpose on the Tensor engine (True) or the DMA engine (False). Returns: list of ``n_k`` transposed ``[128, rb_size]`` bf16 SBUF tiles. """ n_rs = _div_ceil(pr, _P_MAX) xt_tiles = [] for _ in range(n_k): xt_tiles.append(nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf)) for rs in range(n_rs): rr0 = rs * _P_MAX prs = min(_P_MAX, pr - rr0) x_b16 = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf) nisa.dma_copy( dst=x_b16[0:prs, 0:K], src=x_hbm[r_abs + rr0:r_abs + rr0 + prs, 0:K], ) for kt in range(n_k): k0 = kt * _P_MAX pk = min(_P_MAX, K - k0) if tensor_transpose: # Transpose on the (idle) Tensor engine via PSUM, then evict to the xt tile. tp = nl.ndarray((_P_MAX, _P_MAX), dtype=_BF16, buffer=nl.psum) nisa.nc_transpose(dst=tp[0:pk, 0:prs], data=x_b16[0:prs, k0:k0 + pk], engine=nisa.engine.tensor) nisa.tensor_copy(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs], src=tp[0:pk, 0:prs]) else: nisa.dma_transpose(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs], src=x_b16[0:prs, k0:k0 + pk]) return xt_tiles @nki.jit def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm): """gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early. Args: x_hbm (nl.ndarray): [R, K] @ HBM, **bfloat16** (kernel 1 emits bf16). w1t_hbm (nl.ndarray): [K, H] @ HBM, **bf16**, == fc1.weight.T. b1_hbm (nl.ndarray): [128, H/128] @ HBM, float32, == pack_b1(fc1 bias). w2t_hbm (nl.ndarray): [H, K] @ HBM, **bf16**, == fc2.weight.T. b2_hbm (nl.ndarray): [K] @ HBM, float32. Returns: nl.ndarray: [R, K] @ HBM, **bfloat16**. """ R, K = x_hbm.shape K_w1t, H = w1t_hbm.shape H_w2t, K_w2t = w2t_hbm.shape _kernel_assert(R >= 1 and K >= 1 and H >= 1, f"R({R}) K({K}) H({H}) must be >= 1") _kernel_assert(K_w1t == K, f"w1t in-dim {K_w1t} != K {K}; expected w1t = w1.T of shape [K, H]") _kernel_assert(K_w2t == K, f"w2t out-dim {K_w2t} != K {K}; expected w2t = w2.T of shape [H, K]") _kernel_assert(H_w2t == H, f"w2t in-dim {H_w2t} != H {H}") _kernel_assert(H % _P_MAX == 0, f"H={H} must be a multiple of {_P_MAX} for the packed b1 layout") _kernel_assert(tuple(b1_hbm.shape) == (_P_MAX, H // _P_MAX), f"b1 must be pack_b1(fc1 bias): [{_P_MAX}, {H // _P_MAX}], got {tuple(b1_hbm.shape)}") _kernel_assert( R % LNC == 0, f"the {LNC}-core kernel needs R divisible by {LNC}, got R={R}. " f"program_id is a symbolic value, so a partial final row block cannot be " f"clamped; use the one-core pre-transposed kernel for uneven R", ) rows_per = R // LNC pid = nl.program_id(0) r_base = pid * rows_per out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm) n_k = _div_ceil(K, _P_MAX) # K tiles (fc1 contraction) n_h = _div_ceil(H, _P_MAX) # H tiles (fc1 out-rows / fc2 contraction) # b1 FIRST, in ONE DMA on the Scalar engine's HW-DGE queue: the packed [128, H/128] bias # lands in a single ~16 KB transfer so every gelu's bias is resident long before the first # matmul retires, and gelu's wait no longer counts the W1/W2 traffic. b1_sb[:, ht] is the # bias column for H tile ht (b1_sb[p, ht] == fc1_bias[ht*128 + p]). b1_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf) nisa.dma_copy(dst=b1_sb, src=b1_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar) # Main fused loop geometry (needed to size the x prestage). rb_size = _F_MOV_MAX n_rb = _div_ceil(rows_per, rb_size) n_ko = _div_ceil(K, _F_MOV_MAX) # x BEFORE weights: stage row block 0's x first (Tensor-engine transpose, filling the dark # pre-K4 seam), then W1, then later row blocks' x (DMA transpose, streaming under rb0's # matmuls), then W2. Appending later blocks after rb=0's stages leaves rb=0's readiness fixed. xt_sets = [None] * n_rb xt_sets[0] = _stage_xt(x_hbm, r_base, min(rb_size, rows_per), n_k, K, rb_size, tensor_transpose=True) # W1 staging (pre-transposed bf16, one DMA per K tile). w1t_tiles = [] for kt in range(n_k): k0 = kt * _P_MAX pk = min(_P_MAX, K - k0) t = nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf) nisa.dma_copy(dst=t[0:pk, :], src=w1t_hbm[k0:k0 + pk, :]) w1t_tiles.append(t) # Later row blocks' x via DMA transpose, under rb0's matmuls. for rb_pre in range(1, n_rb): r_pre = rb_pre * rb_size xt_sets[rb_pre] = _stage_xt(x_hbm, r_base + r_pre, min(rb_size, rows_per - r_pre), n_k, K, rb_size, tensor_transpose=False) # W2 staging (one DMA per H tile, so fc2 can start on the first strips). w2t_tiles = [] for ht in range(n_h): h0 = ht * _P_MAX ph = min(_P_MAX, H - h0) t = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf) nisa.dma_copy(dst=t[0:ph, :], src=w2t_hbm[h0:h0 + ph, :]) w2t_tiles.append(t) for rb in range(n_rb): r_local = rb * rb_size pr = min(rb_size, rows_per - r_local) n_rs = _div_ceil(pr, _P_MAX) r_abs = r_base + r_local # x_T bf16 tiles: consume the prologue-staged set for this row block. xt_tiles = xt_sets[rb] # ---- fc1: gt_tiles[ht] = gelu(w1 @ x.T + b1), bf16, SBUF-resident ---- gt_tiles = [] for ht in range(n_h): h0 = ht * _P_MAX ph = min(_P_MAX, H - h0) # Declared INSIDE the ht loop, so the ring never reuses a written PSUM tile. # nc_matmul's default accumulate=None auto-detects "previously written" and # would accumulate onto a stale value. y1_psum = nl.ndarray((_P_MAX, rb_size), dtype=nl.float32, buffer=nl.psum) for kt in nl.affine_range(n_k): k0 = kt * _P_MAX pk = min(_P_MAX, K - k0) nisa.nc_matmul( dst=y1_psum[0:ph, 0:pr], stationary=w1t_tiles[kt][0:pk, h0:h0 + ph], moving=xt_tiles[kt][0:pk, 0:pr], ) # b1 comes from the prologue-resident packed tile, column ht. gt = nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf) nisa.activation(dst=gt[0:ph, 0:pr], op=nl.gelu, data=y1_psum[0:ph, 0:pr], bias=b1_sb[0:ph, ht:ht + 1]) gt_tiles.append(gt) # ---- fc2: out[r,k] = sum_h g[r,h] * w2[k,h] + b2[k] ---- for ko in range(n_ko): ko0 = ko * _F_MOV_MAX nk = min(_F_MOV_MAX, K - ko0) out_psums = [] for _ in range(n_rs): out_psums.append( nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum) ) for ht in nl.affine_range(n_h): h0 = ht * _P_MAX ph = min(_P_MAX, H - h0) for rs in range(n_rs): rr0 = rs * _P_MAX prs = min(_P_MAX, pr - rr0) nisa.nc_matmul( dst=out_psums[rs][0:prs, 0:nk], stationary=gt_tiles[ht][0:ph, rr0:rr0 + prs], moving=w2t_tiles[ht][0:ph, ko0:ko0 + nk], ) # b2 broadcast, hoisted out of the rs loop (removing the broadcast itself # cost 23% and 30% in two separate attempts). b2_bc = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf) nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:nk], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, nk]], offset=ko0)) for rs in range(n_rs): rr0 = rs * _P_MAX prs = min(_P_MAX, pr - rr0) out_sb = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf) nisa.tensor_tensor(dst=out_sb[0:prs, 0:nk], data1=out_psums[rs][0:prs, 0:nk], data2=b2_bc[0:prs, 0:nk], op=nl.add) nisa.dma_copy( dst=out_hbm[r_abs + rr0:r_abs + rr0 + prs, ko0:ko0 + nk], src=out_sb[0:prs, 0:nk], ) return out_hbm def pack_b1(b1: torch.Tensor) -> torch.Tensor: """Pack fc1 bias ``[H]`` fp32 into ``[128, H/128]`` so ``b1p[p, ht] == b1[ht*128 + p]``. Lets the kernel load the whole bias in one DMA instead of ``H/128`` column DMAs. ``H`` must be a multiple of 128. """ H = b1.shape[0] assert H % _P_MAX == 0, f"pack_b1 needs H % {_P_MAX} == 0, got H={H}" return b1.reshape(-1, _P_MAX).t().contiguous() def _no_grad(*tensors) -> None: """Refuse gradient-tracking inputs -- this op registers no autograd rule.""" for t in tensors: if t is not None and t.requires_grad: raise RuntimeError( "nki_mlp_k4_xp registers no autograd rule (INFERENCE ONLY). " "Use under torch.no_grad()/eval()." ) @nki_op("vit_enc::nki_mlp_k4_xp", mutates_args={}) def nki_mlp_k4_xp( x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, w2t: torch.Tensor, b2: torch.Tensor, ) -> torch.Tensor: """Fused MLP on LNC=2, pre-transposed bf16 weights, x prestaged in the prologue. Args: x : [R, K] **bf16** on 'neuron'. ``R`` must be divisible by :data:`LNC`. w1t : [K, H] bf16 == fc1.weight.T. b1 : fc1 bias, either ``[H]`` float32 (packed on the fly) or pre-packed ``[128, H/128]`` from :func:`pack_b1`. w2t : [H, K] bf16 == fc2.weight.T. b2 : [K] float32. Returns: [R, K] bf16. """ _no_grad(x, w1t, b1, w2t, b2) if b1.dim() == 1: # accept the pre-pack [H] bias; pack it on the fly b1 = pack_b1(b1) return _mlp_k4_xp_kernel[LNC](x, w1t, b1, w2t, b2) #: Requested export alias. Same object as the incumbent-named nki_op, so the runtime #: op-name string ("vit_enc::nki_mlp_k4_xp") is unchanged whichever name is used. mlp = nki_mlp_k4_xp