Kernels:

---

aws-neuron

/

vit-encoder-neuron-kernels

Like 1

Follow

AWS Inferentia and Trainium 192

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