Kernels:

---

aws-neuron

/

vit-encoder-neuron-kernels

Like 3

Follow

AWS Inferentia and Trainium 194

"""K1/K3 -- fused residual add + affine-free LayerNorm (the "fold" kernel).

WHAT IT COMPUTES
----------------
Given a residual stream ``a`` [R, H] fp32 and a stack of 1..4 branch tensors
``b`` [B, R, H] bf16, per token row this kernel computes::

    x_new = a + sum_j b[j]                     # fp32, the untouched residual stream
    z     = (x_new - mean(x_new)) * rstd       # bf16, LayerNorm WITHOUT affine

and returns ``(x_new, z)``. gamma/beta are NOT applied here -- they have been folded
into the downstream consumers' packed weights at freeze time (K1's z feeds K2's Q/K/V
projections; K3's z feeds K4's fc1). The affine fold is exact:
``W @ (z*g + bt) + c == (W * g) @ z + (W @ bt + c)``, so scaling the consumer's input
columns by gamma and adding ``W @ bt`` to its bias reproduces the same math with the
affine removed from the norm kernel.

WHY IT IS SHAPED THIS WAY
-------------------------
* LNC = 2 shard: the kernel is launched on both NeuronCores; each owns ``R // 2`` rows
  selected by ``program_id(0)``, so the two cores never touch the same tokens.
* Every token tile has its own SBUF slot ([P, n_t, H]) and all loads are issued before any
  compute, so tile ``t+1``'s loads never wait behind tile ``t``'s stores. When the tiles of one core do not
  fit SBUF (large batch: e.g. ViT-S >= 8k rows/core, ViT-L >= ~2k), they are processed in groups that do,
  with the same pattern inside each group (v1.7; up to v1.6 these shapes failed to compile). One DMA per tile and
  operand: a single whole-shard DMA would signal completion only at the end. (External
  contribution by jlongen-aws, PR #2 on ``jburtoft/vjepa2-neuron-kernels``; replaces the
  depth-4 ring, same instructions in the same order so results are unchanged.)
* ``dge_mode.none`` on every DMA: the runtime generates the descriptors before execution,
  so no engine spends time issuing them. The default (software DGE) issued each of the ~40
  DMAs on GpSimd at ~0.75 us apiece, one after another, and the first add waited ~8 us for
  them. Hardware DGE on the Sync queue was as fast standalone, but inside the encoder it
  delayed K2/K4's loads on that queue by about as much as it saved, so it is not used.
* fp32 residual / bf16 branch: ``a`` (the residual stream) stays fp32 so ``x_new`` is
  bitwise-invariant math and can be carried forward losslessly; the branch inputs and
  the normalized output ``z`` are bf16, which is the precision boundary the downstream
  kernels consume.
* Thinned stat chain: mean and variance are computed with a shift-and-reduce chain
  (neg_shift -> copy-reduce -> neg_mu -> square-reduce -> var -> rsqrt) that keeps the
  per-lane statistics in fp32; the final normalize is a single fused ``tensor_scalar``
  (``z = (x_new + neg_mu) * rstd``) whose two per-lane operands cost the same as one op
  and whose bf16 cast rides the write. Each per-lane statistic has its own column per tile
  (``[128, n_t]``) so tile ``t+1``'s chain does not wait on tile ``t``'s.

This kernel is LNC-parametric (``LNC`` from ``NEURON_LOGICAL_NC_CONFIG``). There are no
``engine=`` pins. INFERENCE ONLY: no autograd rule is registered.
"""

from __future__ import annotations

import torch
from torch_neuronx import nki_op

import nki
import nki.isa as nisa
import nki.language as nl

#: SBUF partition count, and the token-tile height (rows processed per instruction).
_P_MAX = 128

#: The bf16 dtype used for branch inputs and the normalized output.
_BF16 = nl.bfloat16

#: Per-partition SBUF bytes one tile group's a/b/s slots may use. 200 KiB is the largest footprint measured to
#: compile (ViT-H+, 16 tiles, one branch); 210 KiB failed. Above it, tiles are processed in groups (v1.7).
_SBUF_GROUP_BUDGET = 200 * 1024

#: Logical-to-physical NeuronCore degree. Rows are split evenly across LNC cores.
import os as _os
#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))


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


@nki.jit
def _res_layernorm_fold_kernel(a_hbm, b_hbm, eps):
    """Emit the fold kernel body: ``(a + sum(b), LN_no_affine(a + sum(b)))``, token-major.

    gamma/beta live in the consumers' packed weights, so this kernel takes no affine
    parameters. Rows are sharded across ``LNC`` cores; every token tile is loaded up front
    into its own SBUF slot (all loads issued before any compute), summed in fp32, normalized
    with the thinned stat chain, and the normalized result is cast to bf16 on write. Every DMA
    uses ``dge_mode.none`` so the runtime pre-generates the descriptors.

    Args:
        a_hbm (nl.ndarray): [R, H] @ HBM float32 -- the residual stream.
        b_hbm (nl.ndarray): [B, R, H] @ HBM bf16 -- B branch tensors (1 <= B <= 4), summed in.
        eps   (float):      LayerNorm epsilon.

    Returns:
        tuple: ``(x_new [R, H] float32, z [R, H] bf16)`` both in shared HBM.
    """
    R, H = a_hbm.shape
    n_b, Rb, Hb = b_hbm.shape

    _kernel_assert(R >= 1 and H >= 1, f"R({R}) H({H}) must be >= 1")
    _kernel_assert(Rb == R and Hb == H, f"b {(Rb, Hb)} must match a {(R, H)}")
    _kernel_assert(1 <= n_b <= 4, f"b must stack 1..4 branches, got {n_b}")
    _kernel_assert(
        H <= 16384,
        f"H={H} is one instruction's free axis here; no chunking is implemented.",
    )
    _kernel_assert(
        R % (LNC * _P_MAX) == 0,
        f"need R divisible by LNC*{_P_MAX} = {LNC * _P_MAX}, got R={R}",
    )

    rows_per = R // LNC
    n_t = rows_per // _P_MAX
    pid = nl.program_id(0)
    r_base = pid * rows_per

    x_new_hbm = nl.ndarray((R, H), dtype=nl.float32, buffer=nl.shared_hbm)
    z_hbm = nl.ndarray((R, H), dtype=_BF16, buffer=nl.shared_hbm)

    inv_h = 1.0 / float(H)
    neg_inv_h = -inv_h

    # Tile GROUPS (v1.7): the kernel stages every tile of a group in SBUF at once (loads first, see the
    # module doc), so the group size G is capped by the SBUF budget. Measured footprint per tile and
    # partition is the a (fp32), b (bf16 x n_b) and s (fp32) slots, H * (8 + 2 n_b) bytes -- the compiler
    # reuses z's slot (ViT-H+ 16 tiles = 200 KiB compiles; ViT-S 56 tiles = 210 KiB does not).
    # G = all tiles when they fit (the <= v1.6 schedule, unchanged), else the largest G under
    # _SBUF_GROUP_BUDGET. Same per-tile instructions in the same order -> bit-identical results.
    per_tile = H * (8 + 2 * n_b)
    G = max(1, min(n_t, _SBUF_GROUP_BUDGET // per_tile))
    n_g = (n_t + G - 1) // G

    # Scratch destination for the reduce activations (its full tile is discarded; only
    # the reduce result is consumed).
    sink = nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf)

    for g in range(n_g):
        t0 = g * G
        gt = min(G, n_t - t0)
        # One SBUF slot per tile of this group; all of the group's loads are issued before any compute or store.
        # a_sb[:, t, :] is tile t's residual, b_sb[:, j, t, :] is branch j of tile t.
        a_sb = nl.ndarray((_P_MAX, G, H), dtype=nl.float32, buffer=nl.sbuf)
        b_sb = nl.ndarray((_P_MAX, n_b, G, H), dtype=_BF16, buffer=nl.sbuf)
        for t in range(gt):
            r0 = r_base + (t0 + t) * _P_MAX
            nisa.dma_copy(dst=a_sb[:, t, :], src=a_hbm[r0:r0 + _P_MAX, 0:H],
                          dge_mode=nisa.dge_mode.none)
            for j in range(n_b):
                nisa.dma_copy(dst=b_sb[:, j, t, :], src=b_hbm[j, r0:r0 + _P_MAX, 0:H],
                              dge_mode=nisa.dge_mode.none)

        s_sb = nl.ndarray((_P_MAX, G, H), dtype=nl.float32, buffer=nl.sbuf)
        z_sb = nl.ndarray((_P_MAX, G, H), dtype=_BF16, buffer=nl.sbuf)

        # Per-lane statistics, one column per tile so tile t+1's chain doesn't wait on tile t's.
        acc = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
        acc2 = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
        neg_shift = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
        neg_mu = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
        var = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
        rstd = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)

        for t in range(gt):
            r0 = r_base + (t0 + t) * _P_MAX
            s = s_sb[:, t, :]
            c = slice(t, t + 1)

            # ---- x_new = a + sum_j b[j] in fp32, then write the residual stream out ----
            nisa.tensor_tensor(dst=s, data1=a_sb[:, t, :], data2=b_sb[:, 0, t, :], op=nl.add)
            for j in range(1, n_b):
                nisa.tensor_tensor(dst=s, data1=s, data2=b_sb[:, j, t, :], op=nl.add)

            nisa.dma_copy(dst=x_new_hbm[r0:r0 + _P_MAX, 0:H], src=s, dge_mode=nisa.dge_mode.none)

            # ---- thinned stat chain: mean via shift-and-reduce, variance, rstd ----
            nisa.tensor_scalar(dst=neg_shift[:, c], data=s_sb[:, t, 0:1],
                               op0=nl.multiply, operand0=-1.0)
            nisa.activation_reduce(dst=sink, op=nl.copy, data=s,
                                   reduce_op=nl.add, reduce_res=acc[:, c],
                                   bias=neg_shift[:, c], scale=1.0)
            nisa.tensor_scalar(dst=neg_mu[:, c], data=acc[:, c], op0=nl.multiply, operand0=neg_inv_h,
                               op1=nl.add, operand1=neg_shift[:, c])
            nisa.activation_reduce(dst=sink, op=nl.square, data=s,
                                   reduce_op=nl.add, reduce_res=acc2[:, c], bias=neg_mu[:, c], scale=1.0)
            nisa.tensor_scalar(dst=var[:, c], data=acc2[:, c], op0=nl.multiply, operand0=inv_h,
                               op1=nl.add, operand1=eps)
            nisa.activation(dst=rstd[:, c], op=nl.rsqrt, data=var[:, c])

            # ---- one-op affine-free apply: z = (x_new + neg_mu) * rstd, bf16 cast on write.
            # Both operands are per-lane vectors; the fused two-op tensor_scalar costs one op.
            nisa.tensor_scalar(dst=z_sb[:, t, :], data=s,
                               op0=nl.add, operand0=neg_mu[:, c],
                               op1=nl.multiply, operand1=rstd[:, c])

            nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:H], src=z_sb[:, t, :],
                          dge_mode=nisa.dge_mode.none)

    return x_new_hbm, z_hbm


def _no_grad(*tensors: torch.Tensor) -> None:
    """Refuse gradient-tracking inputs -- this op registers no autograd rule."""
    for t in tensors:
        if t.requires_grad:
            raise RuntimeError(
                "nki_res_layernorm_fold registers no autograd rule. Use under "
                "torch.no_grad()/eval()."
            )


@nki_op("vit_enc::nki_res_layernorm_fold", mutates_args={})
def nki_res_layernorm_fold(
    a: torch.Tensor,
    b: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fused residual add + affine-free LayerNorm on LNC=2 (gamma/beta folded away).

    The consumers MUST have been packed with folded weights; feeding ``z`` to unfolded
    weights silently drops the affine.

    Args:
        a:   [R, H] float32 on 'neuron'. ``R`` divisible by ``LNC * 128``.
        b:   [B, R, H] bf16 on 'neuron'. B branch tensors (1 <= B <= 4), summed in.
        eps: LayerNorm epsilon.

    Returns:
        ``(x_new [R, H] float32, z [R, H] bf16)``.
    """
    _no_grad(a, b)
    return _res_layernorm_fold_kernel[LNC](a, b, eps)


#: Requested export alias. Same object as the incumbent-named nki_op, so the runtime
#: op-name string ("vit_enc::nki_res_layernorm_fold") is unchanged whichever name is used.
res_layernorm_fold = nki_res_layernorm_fold