Kernels:

---

aws-neuron

/

vit-encoder-neuron-kernels

Like 3

Follow

AWS Inferentia and Trainium 193

"""K2 outer: fused QKV projection + 3D RoPE + vendor attention + fused out-projection, in ONE
top-level NKI kernel, FLATTENED from ``nki_attn_block_k2v_h10`` (which forks
``nki_attn_block_k2v_dp4``). Self-contained: imports only ``nki``, ``nkilib``, ``torch``, and
``_attention_cte`` from the sibling ``attention_inner`` -- NO imports from ``experiment/`` or
the repo-root incumbent modules.

WHAT THIS COMPUTES (per core; LNC=2, this core owns 8 of the 16 heads)
======================================================================
``result[LNC, N, H] bf16`` = out-projection PARTIALS of
``attention(rope(h@Wq), rope(h@Wk), h@Wv)``. Each core computes a partial sum over its 8 heads'
512 out-projection contraction rows; K3's residual add sums the two cores' partials.

Phase 1 (producer): dma_transpose h -> x_t; project Q/K (lane-permuted) and V; apply the six
rope ``tensor_tensor`` ops; scatter K into ``k_tiles`` (SBUF), Q into ``q_priv`` (private_hbm),
V into ``v_tiles`` (SBUF).
Phase 2 (attention + out-proj): reset the allocator to reclaim the dead producer buffers, call
``_attention_cte`` per local head into feature-major SBUF O slices, then contract the fused
out-projection (tokens on partitions) and DMA token-major to ``result``.

PSUM BANK PLAN (producer): V uses bank 7 (``(0, 7*_PSUM_BANK)``); Q/K parities use banks 5/6
(``(5+par)*_PSUM_BANK``); the fused out-projection uses bank 0 (``(0, 0)``). The vendor
attention (``attention_inner``) pins its own low banks {0-3} for MM1 and {4-7} for MM2 during
its phase, which runs after the producer's high-water mark, so there is no live overlap.

THE TWO ENGINE PINS (both to Scalar), and WHY (from PERF_INVENTORY)
-------------------------------------------------------------------
  * ``:276`` V-scatter ``tensor_copy`` -> SCALAR (H9's inherited mechanism): eight copies read
    one ``vout`` before the next vsub overwrites it (a WAR chain of 8); Scalar (240 ns) beats
    GpSimd (375 ns) for that chain and sheds the copy off Vector, K2's binding engine.
  * ``:458`` QK bias-add-and-evict ``tensor_scalar`` -> SCALAR (H10's mechanism): the largest
    remaining Vector-resident relocatable site (768/core); ``data=ps`` is PSUM so GpSimd is
    illegal, and Scalar's 612 ns fee is faster than the 744 ns Vector cell it replaces, so the
    zero-lead rope chain that consumes ``proj`` cannot lengthen.
The ``:224`` bias-cast ``tensor_copy`` and ``:483`` K-scatter ``tensor_copy`` stay UNPINNED
(compiler default) -- pinning them was dropped (:224 is 0.09% of the window; :483 has no
affordable landing engine), and they MUST stay unpinned to match the shipping stream.

HOST-SIDE FOLD ALGEBRA (packing helpers below, copied flat)
-----------------------------------------------------------
Affine-free LN1 folds gamma1/beta1 into the K2 weights at freeze time:
``pack_weights(wq*g1, wk*g1, bq+wq@t1, bk+wk@t1, ...)`` and ``pack_v_weights(wv*g1, bv+wv@t1)``.
The lane permutation puts each quad-of-heads' even/odd RoPE lanes in their own 128-partition
group so RoPE is six full-width ``tensor_tensor`` ops with no transposes; it cancels in
attention's Q/K contraction (a dot product is invariant under a shared permutation of the
contraction axis), so V is natural and nothing downstream un-permutes.
LNC2 BIAS HALVING: ``pack_out_weights`` halves the out-projection bias (``bo*0.5``) so the two
cores' partials sum to exactly one bias -- halving is bitwise exact in bf16 (it only decrements
the exponent, and summing two identical bf16 values is exact).

INFERENCE ONLY: no autograd; the packed weights are detached caches.

SBUF BUDGET (bytes/lane; budget = 192 KiB = 196608 B/lane), TWO GEOMETRIES
==========================================================================
N=2048 (16-frame clips) is the SHIPPING path: nothing below is gated on, so buffer shapes,
byte layout, DMA target addresses and fill VALUES are BIT-FOR-BIT the incumbent (verified:
torch.equal(perf_out.pt) == True, max abs diff 0.0). N=8192 (64-frame clips) is a SEPARATE,
gated path (``if N > 2048``) that the N<=2048 forward never enters.

                                   |  N=2048 (16f)  |  N=8192 (64f, gated)
    ------------------------------ + -------------- + --------------------------------------
    k_tiles (8 heads resident)     |     32768      |  0  (streamed to private_hbm)
    v_tiles (8 heads resident)     |     16384      |  0  (streamed to private_hbm)
    k_stage / v_stage (1 head)     |       -        |  16384 + 8192 = 24576 (one head at a time)
    x_t   (hidden^T ring)          |     32768      |  32768 (depth-4 token-block ring, was 131072)
    tab   (rope tables ring)       |     32768      |  32768 (depth-4 token-block ring, was 131072)
    producer working set           |    ~25000      |  ~25000
    ------------------------------ + -------------- + --------------------------------------
    PHASE-1 PEAK                   |    ~124 KiB    |  ~106 KiB   (FITS)
    ------------------------------ + -------------- + --------------------------------------
    phase2_base (persistent K/V)   |     49152      |  24576 (only k_stage/v_stage persist)
    wo_sb + bo_bc + o_feat + out   |     29696      |  78848
    vendor inner (per head)        |    ~45 KiB     |  ~251 KiB (single 8192 section, n_large=4:
                                   |                |   exp_tp_sb 131072 + mm1_masked 65536 +
                                   |                |   exp_sb 32768 balloon -- see below)
    ------------------------------ + -------------- + --------------------------------------
    PHASE-2 PEAK                   |    ~125 KiB    |  ~354912 B/lane  (DOES NOT FIT -- measured
                                   |                |   on device [NCC_INLA001] 128x354912, 1.81x
                                   |                |   the 196608 budget => N>2048 is REFUSED)

WHY N>2048 IS REFUSED. At N=8192 the shipping design has TWO independent SBUF overflows
(both confirmed on device as ``[NCC_INLA001] ... out of bound 128x526944``):
  (1) PHASE 1: the four N-scaled residents k_tiles(131072)+v_tiles(65536)+x_t(131072)+
      tab(131072) = 458752 B/lane coexist -- 2.3x the whole budget.
  (2) PHASE 2: eight heads' K/V resident (196608 B/lane) is the ENTIRE budget.
The phase-1 fixes work (x_t/tab depth-4 token-block rings, bit-for-bit at N<=2048, saving
196608 B/lane; K/V streamed through private_hbm one head at a time) => phase-1 peak ~106 KiB.
But PHASE 2 does NOT fit: handing one head all 8192 keys makes the single-section flatten run
with num_large_tiles_per_section=ceil(8192/2048)=4, ballooning three inner buffers (exp_tp_sb
131072 + mm1_masked 65536 + exp_sb 32768 = 229376 B/lane) -> phase-2 peak measured 354912 B/lane,
1.81x the budget. The shrinkable phase-2 region (o_feat/out/wo/bo = 78848) is smaller than the
158304 overflow, so streaming O or shrinking o_feat cannot close the gap; re-enabling flash
sectioning fits but is numerically wrong (rel_l2 0.90, the V-JEPA2 kernels' verification notes), and
editing the inner would perturb the bit-for-bit 16-frame path. So the host wrapper
``attention_block`` raises ``NotImplementedError`` for N>2048 rather than emit a non-fitting or
wrong kernel (see the V-JEPA2 kernels' verification notes). The N<=2048 path is bit-for-bit unchanged.
"""
from __future__ import annotations

import torch
from nki.framework.kernel import wrap_nki_to_torch_hop
from torch_neuronx import nki_op

import nki
import nki.isa as nisa
import nki.language as nl
import os as _osr
_ROPE_BF16 = _osr.environ.get("VIT_ENC_ROPE_BF16", "0") == "1"
_ROPE_DT = nl.bfloat16 if _ROPE_BF16 else nl.float32
from nki.isa import engine

from nkilib.core.utils.modular_allocator import ModularAllocator

# The vendor attention, flattened to this model's reached path with the h8/h10 pins inlined.
# The N<=2048 path calls THIS single-section flatten; the N>2048 path is REFUSED (see below).
# seqlen_k=8192 < the vendor's _FLASH_ATTENTION_THRESHOLD (10240), so num_sections==1 -- the
# SAME single-section, non-flash code path the verified 16-frame configuration uses (no
# multi-section online-softmax combine, which is numerically wrong here: the V-JEPA2 kernels' verification notes).
# BUT single-section at 8192 does NOT fit: the inner's mm1_masked / exp_sb / exp_tp_sb are
# allocated with num_free_tiles carrying num_large_tiles_per_section = ceil(8192/2048) = 4
# (attention_inner.py:658-701), so those three buffers BALLOON with total_seqlen_k -- they are
# NOT reused across the large-tile loop. exp_tp_sb 131072 + mm1_masked 65536 + exp_sb 32768 =
# 229376 B/lane, driving phase-2 peak to ~354912 B/lane (1.81x the 196608 budget; measured on
# device [NCC_INLA001] 128x354912). See the SBUF table and refusal rationale below.
from .attention_inner import _attention_cte

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"))
_SPLIT = 2  # heads are always split in two halves (8 heads per program)
_P_MAX = 128
_PSUM_FMAX = 512
_PSUM_BANK = 2048          # bytes per partition per bank
_K_TILE = 512              # vendor k_cache_sbuf tiles are (d, 512)
_V_TILE = 128              # vendor v_cache_sbuf tiles are (128, d)


def _kernel_assert(cond: bool, msg: str):
    assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}"


# ==========================================================================================
# Host-side packing helpers (copied flat from nki_qk_fused_2core / nki_attn_block_k2v_dp4).
# All are frozen-weight work, cached once; INFERENCE ONLY.
# ==========================================================================================
def lane_permutation(heads: int, D: int, rope: str = "half") -> torch.Tensor:
    """Column order that puts each quad-of-heads' RoPE pair halves in their own 128-row group.

    RoPE rotates pairs of lanes ``(a, b)`` of each head. ``rope`` picks the pairing:
      * ``"half"``        -- (j, j + D/2), the rotate-half convention (DINOv3, EVA-02, HF Llama style).
      * ``"interleaved"`` -- (2j, 2j + 1), adjacent lanes (V-JEPA2 / GPT-NeoX-interleaved style).
    The projection's output columns are reordered so group ``2a`` holds the FIRST lane of every
    pair for heads ``4a..4a+3`` (at D=64; 2 groups per 4 heads at D=128) and group ``2a+1`` the
    SECOND lane, so RoPE is six full-width elementwise ops with no transposes. The permutation
    cancels in the Q/K contraction, so nothing downstream un-permutes.

    Returns ``perm`` of length ``heads*D`` s.t. ``w.t()[:, perm]`` is the reordered weight.
    """
    _kernel_assert(rope in ("half", "interleaved"), f"rope must be 'half' or 'interleaved', got {rope!r}")
    half = D // 2
    quad = _P_MAX // half
    _kernel_assert(quad * half == _P_MAX, f"D={D} gives half={half}, needs quad*half=128")
    _kernel_assert(heads % quad == 0, f"heads={heads} must be divisible by quad={quad}")
    perm = torch.empty(heads * D, dtype=torch.long)
    for hd in range(heads):
        a, b = hd // quad, hd % quad
        for j in range(half):
            lo, hi = (j, half + j) if rope == "half" else (2 * j, 2 * j + 1)
            perm[(2 * a) * _P_MAX + b * half + j] = hd * D + lo
            perm[(2 * a + 1) * _P_MAX + b * half + j] = hd * D + hi
    return perm


def pack_weights(wq: torch.Tensor, wk: torch.Tensor,
                 bq: torch.Tensor | None, bk: torch.Tensor | None,
                 heads: int, D: int, rope: str = "half"):
    """``(w_packed [2,H,H] bf16, b_packed [2,128,H/128] fp32)`` -- transposed and column-permuted.

    FOLD: in the affine-free-LN1 freeze this is called on gamma-folded inputs
    ``pack_weights(wq*g1, wk*g1, bq+wq@t1, bk+wk@t1, heads, D)``. ``wq``/``wk`` are torch
    ``Linear.weight`` ``[H_out, H_in]``; the kernel wants ``[H_in, H_out]`` with ``H_out``
    permuted, i.e. ``w.t()[:, perm]``. The bias is packed by :func:`pack_qk_bias` so each program
    loads all of its 128-row groups in one DMA per q/k.
    """
    H = wq.shape[0]
    perm = lane_permutation(heads, D, rope).to(wq.device)
    w = torch.stack([wq.detach().t()[:, perm].contiguous(),
                     wk.detach().t()[:, perm].contiguous()], 0)
    zero = torch.zeros(H, dtype=wq.dtype, device=wq.device)
    b = torch.stack([(zero if bq is None else bq.detach())[perm].contiguous(),
                     (zero if bk is None else bk.detach())[perm].contiguous()], 0)
    return w.to(torch.bfloat16).contiguous(), pack_qk_bias(b.to(torch.bfloat16).float())


def pack_qk_bias(b: torch.Tensor) -> torch.Tensor:
    """``[2, Fp]`` (or legacy ``[2, Fp, 1]``) permuted Q/K bias -> ``[2, 128, Fp/128]`` fp32.

    ``out[qk, p, g] == b[qk, g*128 + p]``: column ``g`` is the per-partition bias of 128-row group
    ``g``, so a program loads all of its groups with one DMA per q/k. Values are first rounded to
    bf16 (the precision the v1 kernel loaded and cast), which keeps the output bit-identical to v1.
    """
    if b.dim() == 3 and b.shape[1] == _P_MAX:
        return b.float().contiguous()
    if b.dim() == 3:
        b = b.squeeze(-1)
    Fp = b.shape[1]
    _kernel_assert(Fp % _P_MAX == 0, f"bias length {Fp} must be a multiple of 128")
    return b.to(torch.bfloat16).float().reshape(2, Fp // _P_MAX, _P_MAX).transpose(1, 2).contiguous()


def pack_tables(cos: torch.Tensor, sin: torch.Tensor, heads: int, D: int, rope: str = "half"):
    """``tab [4, 128, N] fp32`` RoPE tables, quad-tiled to match ``lane_permutation(..., rope)``.

    Inputs are the model's own ``[N, D]`` cos / sin tables (shared by every head; prefix tokens that
    take no rotation: cos=1, sin=0). With ``(a, b)`` a lane pair the kernel computes
    ``a' = a*tab0 + b*tab2`` and ``b' = b*tab1 + a*tab3``, so:
      * ``"half"``:        ``sin`` is the PLAIN sin table (rotate-half: ``x*cos + rotate_half(x)*sin``);
                           tab = (cos_lo, cos_hi, -sin_lo, +sin_hi).
      * ``"interleaved"``: ``sin`` is the SIGNED sin table (``-sin`` on even lanes, ``+sin`` on odd);
                           tab = (cos_even, cos_odd, sinsigned_even, sinsigned_odd).
    """
    _kernel_assert(rope in ("half", "interleaved"), f"rope must be 'half' or 'interleaved', got {rope!r}")
    half = D // 2
    quad = _P_MAX // half
    if rope == "half":
        parts = [cos[:, :half], cos[:, half:], -sin[:, :half], sin[:, half:]]
    else:
        parts = [cos[:, 0::2], cos[:, 1::2], sin[:, 0::2], sin[:, 1::2]]
    tab = torch.stack([p.t().contiguous().repeat(quad, 1) for p in parts], 0)
    return tab.to(torch.float32).contiguous()


def pack_v_weights(wv: torch.Tensor, bv: torch.Tensor | None):
    """``(wv_t [H, H] bf16, bv [1, H] bf16)``. V is NOT lane-permuted (natural order).

    FOLD: in the affine-free-LN1 freeze this is called on gamma-folded inputs
    ``pack_v_weights(wv*g1, bv+wv@t1)``.
    """
    H = wv.shape[0]
    zero = torch.zeros(H, dtype=wv.dtype, device=wv.device)
    return (wv.detach().t().contiguous().to(torch.bfloat16),
            (bv.detach() if bv is not None else zero).reshape(1, H).contiguous().to(torch.bfloat16))


def pack_out_weights(wo, bo):
    """``(wo_t [H, H] bf16, bo/2 [1, H] bf16)`` for the fused out projection.

    ``wo`` is ``attn.proj.weight`` [H, H] (attention features -> hidden). The kernel contracts
    over attention features, so it wants those on partitions: ``wo.T`` = [features, hidden].

    LNC2 BIAS HALVING. Both cores add ``bo/2`` and K3's partial sum restores ``bo``. Adding the
    full bias on both would double it; adding it on one core only would need ``program_id`` as an
    arithmetic operand (it is symbolic -- neither a compile-time constant nor a [P,1] tile).
    Halving is bitwise exact in bf16: it only decrements the exponent, and summing two identical
    bf16 values is exact.
    """
    import torch as _t
    Hh = wo.shape[0]
    zero = _t.zeros(Hh, dtype=wo.dtype, device=wo.device)
    half = (bo.detach() if bo is not None else zero) * 0.5
    return (wo.detach().t().contiguous().to(_t.bfloat16),
            half.reshape(1, Hh).contiguous().to(_t.bfloat16))


# ==========================================================================================
# THE KERNEL. Flattened from _attn_block_k2v_h10_dp_kernel (verbatim dp4 body with the two
# engine= keywords hardcoded to engine.scalar at :276 and :458; :224 and :483 left unpinned).
# ==========================================================================================
@nki.jit
def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
                            wot_hbm, bo_hbm, scale, head_dim, n_valid, hb, grid, nimg=1, resident=0, pack2=0):
    """out[16, N, D] bf16 = attention(rope(h@Wq), rope(h@Wk), h@Wv), one kernel (per-core partials).

    BATCHED (nimg>1): h_hbm is the STACKED [nimg*N, H] buffer (images back to back, N rows each).
    The per-image body below runs once per image inside this ONE launch, reading image b's rows at
    row offset b*N and writing partials to result[slot, b*N:(b+1)*N]. ONE launch, no host-side
    slice/cat glue.
    resident=0: the loop wraps the whole body, so weights/biases/rope tables are re-DMA'd per image.
    resident=1 (N<=2048): Q/K/V/O weights, biases and rope tables are staged ONCE at the bottom of
    SBUF before the image loop (~38 KiB/lane ViT-L + tables); each image's buffers are allocated
    above them. Same math, bit-identical outputs."""
    NT, H = h_hbm.shape
    N = NT // nimg
    Fp = wqk_hbm.shape[2]                 # padded attention features = 512 * n_groups
    D = head_dim
    half = D // 2
    quad = _P_MAX // half                 # heads per 128-partition parity group = 4
    heads = Fp // D                       # padded head count (multiple of 8)
    local_heads = _PSUM_FMAX // D         # heads per program: 8 at D=64, 4 at D=128
    local_quads = local_heads // quad     # 2
    n_chunks = H // _P_MAX                # 8 contraction chunks
    n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX   # token tiles of 512 (last may be ragged)
    n_ktiles = (N + _K_TILE - 1) // _K_TILE     # K tiles / head (last may be partially filled)
    n_vtiles = N // _V_TILE               # 16 V tiles / head
    n_ochunks = (local_heads * D) // _P_MAX   # feature chunks of the out projection
    v_per_tt = _PSUM_FMAX // _V_TILE      # 4 V tiles per 512-token tile

    # N>2048 gate (see the phase-1/phase-2 note below). Computed early so the K/V residents can
    # be allocated as HBM (streamed) for the big path.
    _big = N > _PSUM_FMAX * 4

    # N>2048 IS NOT SUPPORTED (guarded at the host wrapper `attention_block` below, which raises
    # NotImplementedError before this kernel is ever traced -- NKI kernels cannot `raise`). The
    # phase-1 memory fixes (x_t/tab depth-4 rings, K/V streamed through private_hbm one head at a
    # time) DO fit -- phase-1 peak ~106 KiB/lane. But phase 2 does NOT: handing one head all 8192
    # keys makes the flattened single-section _attention_cte run with num_large_tiles_per_section
    # =ceil(8192/2048)=4, which balloons three N-scaled inner buffers -- exp_tp_sb (8*4*4*512*2 =
    # 131072), mm1_masked (2*4*2048*4 = 65536), exp_sb (2*4*2048*2 = 32768) = 229376 B/lane for
    # three buffers alone. Measured on device (2026-09-25, neuronx-cc 2.25): [NCC_INLA001] out of
    # bound 128x354912 -> phase-2 peak 354912 B/lane, 1.81x the 196608 budget. The shrinkable
    # phase-2 region (o_feat 65536 + out 1024 + wo/bo 12288 = 78848) is smaller than the 158304
    # overflow, so streaming O or shrinking o_feat CANNOT close the gap -- the inner's N-scaled
    # buffers are the wall. Re-enabling flash sectioning (2048-key sections) shrinks the inner but
    # activates the vendor's multi-section online-softmax combine, a path the shipping config never
    # exercises and which produced rel_l2 0.90 (wrong). See the V-JEPA2 kernels' verification notes.

    _kernel_assert(wqk_hbm.shape == (2, H, Fp), f"wqk must be [2,{H},{Fp}]")
    _kernel_assert(bqk_hbm.shape == (2, _P_MAX, Fp // _P_MAX), f"bqk must be pack_qk_bias(...): [2,128,{Fp // _P_MAX}]")
    _kernel_assert(tab_hbm.shape == (4, _P_MAX, N), f"tab must be [4,{_P_MAX},{N}]")
    _kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
    _kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
    _kernel_assert(H % _P_MAX == 0, f"H={H} must be a multiple of 128")
    _kernel_assert(quad * half == _P_MAX, f"D={D} must give 4 heads per group")
    _kernel_assert(Fp % _PSUM_FMAX == 0, "padded features must be a multiple of 512")
    _kernel_assert(local_heads * D == _PSUM_FMAX,
                   "inline V projection wants all local heads in one 512-wide stream")
    _kernel_assert(N % _P_MAX == 0, f"N={N} must be a multiple of 128")
    _kernel_assert(NT == nimg * N, f"rows {NT} must be nimg({nimg}) x N")
    _kernel_assert(str(h_hbm.dtype) == str(nl.bfloat16), "h must be bf16")

    # [LNC, N, H] PARTIALS. The out projection contracts over all H features == all 16 heads,
    # and this core owns 8, so what it computes is a partial sum; K3 adds the two.
    result = nl.ndarray((grid, NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)

    # Each program handles one 8-head group: group = hb + program_id. grid==1 -> single program.
    if grid == 1:
        pid = hb
        out_slot = 0
    else:
        out_slot = nl.program_id(0)       # SYMBOLIC: index arithmetic only
        pid = hb + out_slot
    head0 = pid * local_heads
    col0 = pid * (local_heads * D)        # first permuted qk weight column
    vcol0 = pid * (local_heads * D)       # first natural v weight column (same offset)
    gcol0 = pid * ((local_heads * D) // _P_MAX)   # first 128-row bias group (multiply only: pid may be symbolic)

    # ---- RESIDENT MODE (resident=1, N<=2048 only): stage the per-launch-invariant inputs ONCE,
    # at the BOTTOM of SBUF, before the image loop. Everything per image is then allocated above
    # `res_top`. Invariant: Q/K/V/O weights, their biases, and the rope tables (same for every image).
    # Cost ~38 KiB/lane (ViT-L) + tables; per-image phase peaks were <=106 KiB, budget 192 KiB.
    _res = bool(resident) and not (N > _PSUM_FMAX * 4)
    res_top = 0
    if _res:
        ra = ModularAllocator(initial_address=0)
        r_w_sb = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2 * _P_MAX), dtype=nl.bfloat16,
                                      block_dim=[2, n_chunks], num_free_tiles=[2, n_chunks])
        for qk in range(2):
            for c in range(n_chunks):
                nisa.dma_copy(dst=r_w_sb[qk][c],
                              src=wqk_hbm[qk, c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_quads * 2 * _P_MAX])
        r_wv_sb = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_heads * D), dtype=nl.bfloat16,
                                       block_dim=[n_chunks], num_free_tiles=[n_chunks])
        for c in range(n_chunks):
            nisa.dma_copy(dst=r_wv_sb[c], src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, vcol0:vcol0 + local_heads * D])
        # Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
        # ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
        r_b_sb = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
            block_dim=[2], num_free_tiles=[2], align_to=4)
        for qk in range(2):
            nisa.dma_copy(dst=r_b_sb[qk], src=bqk_hbm[qk, 0:_P_MAX, gcol0:gcol0 + local_quads * 2],
                          dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
        r_bv_bc = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
        nisa.dma_copy(dst=r_bv_bc, src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=vcol0))
        r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
                                     block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
        for s in range(n_tt):
            for t_i in range(4):
                _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
                nisa.dma_copy(dst=r_tab[t_i][s][:, 0:_tw], src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
        r_wo_sb = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
                                       block_dim=[n_ochunks], num_free_tiles=[n_ochunks])
        for oc in range(n_ochunks):
            nisa.dma_copy(dst=r_wo_sb[oc], src=wot_hbm[col0 + oc * _P_MAX:col0 + (oc + 1) * _P_MAX, 0:H])
        r_bo_bc = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
        nisa.dma_copy(dst=r_bo_bc, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
        res_top = ra.get_current_address()

    # v1.6 (pack2): two images share one producer pass when 2N <= 512 (e.g. 224 px): their rows are contiguous,
    # so one 512-token x^T tile holds both and the Q/K/V projections + RoPE stream 2N columns per matmul.
    # Attention runs per image (separate _attention_cte calls on per-image K/V/Q slices); the out-projection
    # covers both images' rows. Same arithmetic per element -> bit-identical to the unpacked kernel.
    PK = 2 if (pack2 and nimg % 2 == 0 and 2 * N <= _PSUM_FMAX) else 1
    NP = N * PK
    n_ttp = (NP + _PSUM_FMAX - 1) // _PSUM_FMAX
    n_ktp = (NP + _K_TILE - 1) // _K_TILE
    n_vtp = NP // _V_TILE
    n_grp1 = N // _P_MAX
    for img in range(nimg // PK):        # unrolled at trace time (python int)
        rofs = img * NP * H               # element offset of this image (pair)'s rows in h / result
        allocator = ModularAllocator(initial_address=res_top)

        # ---- persistent consumer tiles (written by producer, read by attention) ---------
        # N<=2048: all 8 heads' K/V resident in SBUF (bit-for-bit shipping path). N>2048: K/V for
        # all 8 heads (196608 B/lane) is the ENTIRE 192 KiB budget, so they cannot stay resident
        # alongside x_t/tab/working set. STREAM them: the producer scatters each head's K/V to
        # private_hbm, and phase-2 loads ONE head's K/V into small resident staging rings just
        # before its attention call (attention runs per head and writes O to HBM, so only one
        # head's K/V need be live at a time).
        k_tiles = []
        v_tiles = []
        k_hbm = []
        v_hbm = []
        if _big:
            for _lh in range(local_heads):
                k_hbm.append(nl.ndarray((n_ktiles, D, _K_TILE), dtype=nl.bfloat16,
                                        buffer=nl.private_hbm))
                v_hbm.append(nl.ndarray((n_vtiles, _V_TILE, D), dtype=nl.bfloat16,
                                        buffer=nl.private_hbm))
            # ONE head's resident staging rings (reused across heads in phase 2).
            k_stage = allocator.alloc_sbuf_tensor(
                shape=(D, _K_TILE), dtype=nl.bfloat16,
                block_dim=[n_ktiles], num_free_tiles=[n_ktiles], align_to=32)
            v_stage = allocator.alloc_sbuf_tensor(
                shape=(_V_TILE, D), dtype=nl.bfloat16,
                block_dim=[n_vtiles], num_free_tiles=[n_vtiles])
        else:
            for _lh in range(local_heads):
                kt_ = allocator.alloc_sbuf_tensor(
                    shape=(D, _K_TILE), dtype=nl.bfloat16,
                    block_dim=[n_ktp], num_free_tiles=[n_ktp], align_to=32)
                k_tiles.append(kt_)
                vt_ = allocator.alloc_sbuf_tensor(
                    shape=(_V_TILE, D), dtype=nl.bfloat16,
                    block_dim=[n_vtp], num_free_tiles=[n_vtp])
                v_tiles.append(vt_)
        # HIGH WATER MARK for phase 2: k_tiles/v_tiles are the only phase-1 products attention
        # needs, so every buffer above this address is dead once the producer loop finishes.
        phase2_base = allocator.get_current_address()

        # ------------------------------------------------------------------------------------
        # N>2048 GATE (64-frame clips, N=8192). At N<=2048 every ring below keeps its full-N
        # depth, so the buffer shape, byte layout, DMA target addresses and fill VALUES are
        # bit-for-bit identical to the shipping 16-frame path (see the SBUF budget table in the
        # module docstring), and the attention call goes to the pinned single-section flatten
        # ``_attention_cte`` (all engine pins intact). At N>2048 two independent overflows appear:
        #   (1) PHASE 1: the four N-scaled residents (k_tiles+v_tiles+x_t+tab = 458752 B/lane)
        #       coexist. FIX: x_t and tab become depth-XT_RING token-block rings (a rolling 4-tile
        #       window filled inside the per-token-tile loop; a pure copy, identical values).
        #   (2) PHASE 2: eight heads' K/V resident (196608 B/lane) is the whole budget. FIX: stream
        #       K/V through private_hbm and stage ONE head at a time into k_stage/v_stage. The
        #       attention call is the SAME single-section flatten ``_attention_cte`` the N<=2048 path
        #       uses (seqlen_k=8192 < 10240 threshold => num_sections==1, no multi-section flash
        #       combine); its large-tile loop iterates 4x over _LARGE_TILE_SZ-sized buffers.
        # The N<=2048 path enters NEITHER branch and is bit-for-bit unchanged.
        _big = N > _PSUM_FMAX * 4                 # N>2048  -> the 64-frame geometry
        XT_RING = min(n_tt, 4) if _big else n_ttp  # token-block ring depth for x_t / tab

        # ---- resident producer inputs ----------------------------------------------------
        # x^T: hidden on partitions, tokens on the free axis, as a per-512-token-block ring.
        # At N<=2048 XT_RING==n_tt so the ring holds every block (== the old (128,N) buffer,
        # same bytes, same addresses); at N=8192 it holds a rolling window of 4 blocks. The fill
        # is a pure dma_transpose (no arithmetic), so any fill ORDER yields identical values.
        x_t = allocator.alloc_sbuf_tensor(
            shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16,
            block_dim=[n_chunks, n_ttp], num_free_tiles=[n_chunks, XT_RING], align_to=32)

        if not _big:
            # N<=2048: fill the whole buffer up front, exactly as the shipping path did (pure
            # dma_transpose; block-indexed AP lands at the same addresses as the old (128,N) buffer).
            for s in range(n_ttp):
                for c in range(n_chunks):
                    nisa.dma_transpose(
                        dst=x_t[c][s].ap([[_PSUM_FMAX, _P_MAX], [1, 1], [1, 1],
                                          [1, min(_PSUM_FMAX, NP - s * _PSUM_FMAX)]], offset=0),
                        src=h_hbm.ap([[H, min(_PSUM_FMAX, NP - s * _PSUM_FMAX)], [1, 1], [1, 1], [1, _P_MAX]],
                                     offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX),
                    )

        if _res:
            w_sb, wv_sb, b_sb, bv_bc = r_w_sb, r_wv_sb, r_b_sb, r_bv_bc
        else:
            # this core's Q/K weight columns (lane-permuted) and V columns (natural)
            w_sb = allocator.alloc_sbuf_tensor(
                shape=(_P_MAX, local_quads * 2 * _P_MAX), dtype=nl.bfloat16,
                block_dim=[2, n_chunks], num_free_tiles=[2, n_chunks])
            for qk in range(2):
                for c in range(n_chunks):
                    nisa.dma_copy(dst=w_sb[qk][c],
                                  src=wqk_hbm[qk, c * _P_MAX:(c + 1) * _P_MAX,
                                              col0:col0 + local_quads * 2 * _P_MAX])
            wv_sb = allocator.alloc_sbuf_tensor(
                shape=(_P_MAX, local_heads * D), dtype=nl.bfloat16,
                block_dim=[n_chunks], num_free_tiles=[n_chunks])
            for c in range(n_chunks):
                nisa.dma_copy(dst=wv_sb[c],
                              src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX,
                                          vcol0:vcol0 + local_heads * D])

            # biases: qk as (128,1) per group; v broadcast to (128, 8*64) once
            # Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
            # ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
            b_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
                block_dim=[2], num_free_tiles=[2], align_to=4)
            for qk in range(2):
                nisa.dma_copy(dst=b_sb[qk], src=bqk_hbm[qk, 0:_P_MAX, gcol0:gcol0 + local_quads * 2],
                              dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
            bv_bc = allocator.alloc_sbuf_tensor(
                shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
            nisa.dma_copy(dst=bv_bc,
                          src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]],
                                        offset=vcol0))

        # rope tables, resident (quad-tiled on the host: identical rows for every quad). Token-block
        # ring, same gating as x_t: at N<=2048 TAB_RING==n_tt holds every block (== old (128,N)
        # buffer, bit-for-bit); at N=8192 a rolling 4-block window filled inside the tt loop. Pure
        # dma_copy -> identical values regardless of fill order.
        if _res and PK == 1:
            tab = r_tab
        elif PK == 2:
            # pack2: the same per-image tables side by side in one 512-wide tile
            tab = allocator.alloc_sbuf_tensor(
                shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
                block_dim=[4, 1], num_free_tiles=[4, 1], align_to=4)
            for t_i in range(4):
                for sub in range(2):
                    nisa.dma_copy(dst=tab[t_i][0][:, sub * N:(sub + 1) * N], src=tab_hbm[t_i, 0:_P_MAX, 0:N])
        else:
            tab = allocator.alloc_sbuf_tensor(
                shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
                block_dim=[4, n_tt], num_free_tiles=[4, XT_RING], align_to=4)

            if not _big:
                for s in range(n_tt):
                    for t_i in range(4):
                        _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
                        nisa.dma_copy(dst=tab[t_i][s][:, 0:_tw],
                                      src=tab_hbm[t_i, 0:_P_MAX,
                                                  s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])

        # ---- per-tile working set (reused across iterations; same tensor objects) --------
        proj = allocator.alloc_sbuf_tensor(
            shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16, block_dim=[2],
            num_free_tiles=[2])
        acc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
                                          align_to=4)
        tmp = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
                                          align_to=4)
        roped = allocator.alloc_sbuf_tensor(
            shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16, block_dim=[2],
            num_free_tiles=[2])
        vout = allocator.alloc_sbuf_tensor(
            shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)

        # Q staging via private_hbm, one per LOCAL head (vendor segmented-kernel pattern)
        q_priv = []
        for _lh in range(local_heads):
            qs_ = []
            for _sub in range(PK):
                qs_.append(nl.ndarray((1, D, N), dtype=nl.bfloat16, buffer=nl.private_hbm))
            q_priv.append(qs_)

        # ---- producer: per 512-token tile, straight through the pipeline ----------------
        for tt in range(n_ttp):
            t0 = tt * _PSUM_FMAX
            tw = min(_PSUM_FMAX, NP - t0)         # ragged last token tile (N % 512 != 0)

            # N>2048: fill this token-block's x_t / tab slice NOW (rolling ring window). At
            # N<=2048 the whole buffer was already filled up front (bit-for-bit), so skip.
            if _big:
                for c in range(n_chunks):
                    nisa.dma_transpose(
                        dst=x_t[c][tt].ap([[_PSUM_FMAX, _P_MAX], [1, 1], [1, 1], [1, _PSUM_FMAX]],
                                          offset=0),
                        src=h_hbm.ap([[H, _PSUM_FMAX], [1, 1], [1, 1], [1, _P_MAX]],
                                     offset=rofs + t0 * H + c * _P_MAX),
                    )
                for t_i in range(4):
                    nisa.dma_copy(dst=tab[t_i][tt],
                                  src=tab_hbm[t_i, 0:_P_MAX, t0:t0 + _PSUM_FMAX])

            # V for all 8 local heads of this token tile.
            for vsub in range(tw // _V_TILE):
                ps_v = nl.ndarray((_V_TILE, local_heads * D), dtype=nl.float32,
                                  buffer=nl.psum, address=(0, 7 * _PSUM_BANK))
                for c in range(n_chunks):
                    nisa.nc_matmul(
                        dst=ps_v,
                        stationary=x_t[c][tt][:, vsub * _V_TILE:(vsub + 1) * _V_TILE],
                        moving=wv_sb[c],
                        accumulate=(c > 0),
                    )
                nisa.tensor_tensor(dst=vout, data1=ps_v, data2=bv_bc, op=nl.add)
                for lh in range(local_heads):
                    if _big:
                        # N>2048: stream V to HBM (SBUF->HBM dma from vout directly). The pinned
                        # :276 tensor_copy is the N<=2048 shipping site and is left untouched below.
                        nisa.dma_copy(
                            dst=v_hbm[lh][tt * v_per_tt + vsub],
                            src=vout[:, lh * D:(lh + 1) * D])
                        continue
                    # ======== SITE :276 -- pinned to Scalar. WAR chain of 8 copies read one vout
                    # before the next vsub overwrites it; Scalar (240 ns) beats GpSimd (375 ns) and
                    # sheds the copy off Vector, the binding engine. ========
                    nisa.tensor_copy(
                        dst=v_tiles[lh][tt * v_per_tt + vsub],
                        src=vout[:, lh * D:(lh + 1) * D],
                        engine=engine.scalar)

            # Q and K for both local quads, both parities, this token tile.
            for qk in range(2):
                for j in range(local_quads):
                    for par in range(2):
                        g = j * 2 + par
                        cols = g * _P_MAX
                        ps = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32,
                                        buffer=nl.psum,
                                        address=(0, (5 + par) * _PSUM_BANK))
                        for c in range(n_chunks):
                            nisa.nc_matmul(
                                dst=ps[:, 0:tw],
                                stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
                                moving=x_t[c][tt][:, 0:tw],
                                accumulate=(c > 0),
                            )
                        # ======== SITE :458 -- pinned to Scalar. Largest remaining Vector-resident
                        # relocatable site (768/core); data=ps is PSUM so GpSimd is illegal, and
                        # Scalar's 612 ns fee is faster than the 744 ns Vector cell, so the
                        # zero-lead rope chain consuming proj cannot lengthen. ========
                        # NKI 0.7: a Scalar-engine pin here fails ISA opcode_engine_check on newer compilers
                        # newer compilers -> unpinned (compiler picks the engine).
                        nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
                                           op0=nl.add, operand0=b_sb[qk][:, g:g + 1])

                    # rope: the six ops of nki_rope, same order, tables sliced per tile
                    nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
                                       data2=tab[0][tt][:, 0:tw], op=nl.multiply)
                    nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
                                       data2=tab[2][tt][:, 0:tw], op=nl.multiply)
                    nisa.tensor_tensor(dst=roped[0][:, 0:tw], data1=acc[:, 0:tw], data2=tmp[:, 0:tw], op=nl.add)
                    nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[1][:, 0:tw],
                                       data2=tab[1][tt][:, 0:tw], op=nl.multiply)
                    nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[0][:, 0:tw],
                                       data2=tab[3][tt][:, 0:tw], op=nl.multiply)
                    nisa.tensor_tensor(dst=roped[1][:, 0:tw], data1=acc[:, 0:tw], data2=tmp[:, 0:tw], op=nl.add)

                    # scatter: head b of the quad owns rows [b*32,(b+1)*32)
                    for b in range(quad):
                        lh = j * quad + b
                        r0 = b * half
                        if qk == 1:
                            for par2 in range(2):
                                if _big:
                                    # N>2048: stream K to HBM (SBUF->HBM dma). The :483 tensor_copy
                                    # below is the N<=2048 shipping site, left untouched.
                                    nisa.dma_copy(
                                        dst=k_hbm[lh][tt][par2 * half:(par2 + 1) * half,
                                                          0:_K_TILE],
                                        src=roped[par2][r0:r0 + half, 0:_PSUM_FMAX])
                                    continue
                                # SITE :483 -- K-scatter COPY, left UNPINNED (no affordable landing
                                # engine: it is a WAR chain against the next iteration's rope write).
                                nisa.tensor_copy(
                                    dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half,
                                                        0:tw],
                                    src=roped[par2][r0:r0 + half, 0:tw])
                        elif PK == 1:
                            for par2 in range(2):
                                nisa.dma_copy(
                                    dst=q_priv[lh][0][0, par2 * half:(par2 + 1) * half,
                                                      t0:t0 + tw],
                                    src=roped[par2][r0:r0 + half, 0:tw])
                        else:
                            for par2 in range(2):
                                for sub in range(2):
                                    nisa.dma_copy(
                                        dst=q_priv[lh][sub][0, par2 * half:(par2 + 1) * half, 0:N],
                                        src=roped[par2][r0:r0 + half, sub * N:(sub + 1) * N])

        # ---- PHASE 2: reclaim the dead producer buffers. NO ALIASING. -------------------
        allocator.set_current_address(phase2_base)

        if _res:
            wo_sb, bo_bc = r_wo_sb, r_bo_bc
        else:
            # ---- out-projection weights: THIS core's feature rows, all H output columns -------
            wo_sb = allocator.alloc_sbuf_tensor(
                shape=(_P_MAX, H), dtype=nl.bfloat16,
                block_dim=[n_ochunks], num_free_tiles=[n_ochunks])
            for oc in range(n_ochunks):
                nisa.dma_copy(dst=wo_sb[oc],
                              src=wot_hbm[col0 + oc * _P_MAX:col0 + (oc + 1) * _P_MAX, 0:H])
            # bias already HALVED on the host (pack_out_weights) so the two cores' partials sum to
            # exactly one bias. Broadcast across partitions (output columns are on the free axis).
            bo_bc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
            nisa.dma_copy(dst=bo_bc,
                          src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))

        # ---- O STAYS IN SBUF, ALLOCATED EXACTLY ONCE. tp_out=True gives [D, 128 tokens] per Q
        # group (feature-major), which the projection contracts over, so O never transposes.
        # The list handed to the vendor is SLICES into o_feat, so _write_back_o_impl copies
        # straight into final position. ------------------------------------------------------
        n_grp = NP // _P_MAX
        o_feat = allocator.alloc_sbuf_tensor(
            shape=(_P_MAX, NP), dtype=nl.bfloat16,
            block_dim=[n_ochunks], num_free_tiles=[n_ochunks], align_to=32)
        o_sb = []
        for _lh in range(local_heads):
            oc = (_lh * D) // _P_MAX
            rr = (_lh * D) % _P_MAX
            # explicit loop, NOT a comprehension: the NKI tracer rejects comprehensions in a kernel.
            grp = []
            for _g in range(n_grp):
                grp.append(o_feat[oc][rr:rr + D, _g * _P_MAX:(_g + 1) * _P_MAX])
            o_sb.append(grp)

        attn_base = allocator.get_current_address()
        for lh in range(local_heads):
            if _big:
                # N>2048: load THIS head's K/V from HBM into the resident staging rings (only one
                # head live at a time -> K/V cost is one head's 24576 B/lane, not eight heads'
                # 196608). Then run the SAME single-section flatten the N<=2048 path uses: seqlen_k
                # =8192 < 10240 threshold => num_sections==1, no multi-section flash combine. The
                # inner's large-tile loop iterates 4x over _LARGE_TILE_SZ-sized buffers; O is written
                # straight into this head's o_feat SBUF slices (o_sb[lh]), exactly as N<=2048.
                for _i in range(n_ktiles):
                    nisa.dma_copy(dst=k_stage[_i], src=k_hbm[lh][_i])
                for _i in range(n_vtiles):
                    nisa.dma_copy(dst=v_stage[_i], src=v_hbm[lh][_i])
                # 1D python-list views [seq_tile] (the flatten indexes k_sb[k_tile_idx_in_section]
                # and v_sb[v_tile_idx] as flat lists -- same convention as k_tiles[lh]/v_tiles[lh]).
                k1d = []
                for _i in range(n_ktiles):
                    k1d.append(k_stage[_i])
                v1d = []
                for _i in range(n_vtiles):
                    v1d.append(v_stage[_i])
                _attention_cte(
                    q_priv[lh][0], None, None,
                    scale=scale, causal_mask=False,
                    tp_q=False, tp_k=False, tp_out=True,
                    cache_softmax=False,
                    k_cache_sbuf=k1d,
                    v_cache_sbuf=v1d,
                    out_o_hbm=o_sb[lh],
                    init_sbuf_addr=attn_base,
                    seqlen_k_valid=n_valid,
                )
            else:
                # N<=2048: the pinned single-section flatten (h8/h10 engine pins intact, bit-for-bit).
                for sub in range(PK):
                    if PK == 1:
                        k_l = k_tiles[lh]
                        v_l = v_tiles[lh]
                        o_l = o_sb[lh]
                    else:
                        # pack2: this image's K columns / V tiles / O slices of the shared buffers
                        k_l = [k_tiles[lh][0][:, sub * N:(sub + 1) * N]]
                        v_l = []
                        for _v in range(n_grp1):
                            v_l.append(v_tiles[lh][sub * n_grp1 + _v])
                        o_l = []
                        for _o in range(n_grp1):
                            o_l.append(o_sb[lh][sub * n_grp1 + _o])
                    _attention_cte(
                        q_priv[lh][sub], None, None,
                        scale=scale, causal_mask=False,
                        tp_q=False, tp_k=False, tp_out=True,
                        cache_softmax=False,
                        k_cache_sbuf=k_l,
                        v_cache_sbuf=v_l,
                        out_o_hbm=o_l,
                        init_sbuf_addr=attn_base,
                        seqlen_k_valid=n_valid,
                    )
                    allocator.set_current_address(attn_base)
            allocator.set_current_address(attn_base)

        # ---- PARTIAL out projection: tokens on partitions, output columns on the free axis
        # (chosen by the HBM write: result is token-major, so a row of consecutive output columns
        # is a 1024-byte contiguous run). 128 calls/layer/core (16 token blocks x 2 col blocks x
        # 4 chunks). --------------------------------------------------------------------------
        out_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX),
                                             dtype=nl.bfloat16, align_to=32)
        n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
        for tb in range(NP // _P_MAX):
            for ob in range(n_ob):
                ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX)   # ragged last column block (H%512!=0)
                ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
                                    address=(0, 0))
                for oc in range(n_ochunks):
                    nisa.nc_matmul(
                        dst=ps_out[:, 0:ow],
                        stationary=o_feat[oc][:, tb * _P_MAX:(tb + 1) * _P_MAX],
                        moving=wo_sb[oc][:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
                        accumulate=(oc > 0),
                    )
                nisa.tensor_tensor(dst=out_sb[:, 0:ow], data1=ps_out[:, 0:ow],
                                   data2=bo_bc[:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
                                   op=nl.add)
                nisa.dma_copy(
                    dst=result.ap(pattern=[[H, _P_MAX], [1, ow]],
                                  offset=out_slot * NT * H + rofs + tb * _P_MAX * H + ob * _PSUM_FMAX),
                    src=out_sb[:, 0:ow])
    return result


_ATTN = {g: wrap_nki_to_torch_hop(_attention_block_kernel[g]) for g in sorted({1, LNC})}


def _no_grad(*tensors) -> None:
    for t in tensors:
        if t is not None and t.requires_grad:
            raise RuntimeError(
                "attention_block registers no autograd rule and its packed weights are "
                "detached caches. Use under torch.no_grad()/eval().")


@nki_op("vit_enc::attention_block", mutates_args={})
def attention_block(
    h: torch.Tensor,
    wqk: torch.Tensor,
    bqk: torch.Tensor,
    tab: torch.Tensor,
    wvt: torch.Tensor,
    bv: torch.Tensor,
    wot: torch.Tensor,
    bo: torch.Tensor,
    scale: float,
    head_dim: int,
    n_valid: int,
    hb: int = 0,
    grid: int = LNC,
    nimg: int = 1,
    resident: int = 0,
    pack2: int = 0,
) -> torch.Tensor:
    """Fused norm1'd-hidden -> QKV+rope -> attention -> fused out-projection PARTIALS.

    Returns ``[LNC=2, N, H]`` bf16 per-core out-projection partials; K3's residual add sums the
    two cores. The two Scalar engine pins (V-scatter :276, QK bias-evict :458) are hardcoded.
    """
    _no_grad(h, wqk, bqk, tab, wvt, bv, wot, bo)
    # N>2048 (e.g. 64-frame clips) is NOT supported. The phase-1 memory fixes fit (~106 KiB/lane),
    # but the single-section vendor inner overflows phase 2 (measured 354912 B/lane vs the 196608
    # budget: exp_tp_sb/mm1_masked/exp_sb scale with num_large_tiles_per_section=ceil(N/2048)=4).
    # The shrinkable phase-2 region (78848 B/lane) is smaller than the 158304 overflow, so
    # streaming O / shrinking o_feat cannot close the gap; re-enabling flash sectioning fits but
    # is numerically wrong (rel_l2 0.90). Refuse here rather than emit a non-fitting or wrong
    # kernel. See the V-JEPA2 kernels' verification notes. NKI kernels cannot `raise`, so the guard lives here.
    N = h.shape[0] // nimg
    if N > 2048:
        raise NotImplementedError(
            f"attention_block N={N} (>2048, e.g. 64-frame clips) is not supported: the "
            f"single-section vendor inner does not fit SBUF (phase-2 peak measured 354912 B/lane "
            f"vs the 196608 budget). Streaming O or shrinking o_feat cannot close the 158304 "
            f"B/lane gap, and flash sectioning is numerically wrong (rel_l2 0.90). See "
            f"the V-JEPA2 kernels' verification notes. Use N<=2048 (up to 16-frame clips) -- the bit-for-bit "
            f"shipping path."
        )
    if bqk.dim() == 3 and bqk.shape[1] != 128:
        bqk = pack_qk_bias(bqk)
    return _ATTN[grid](h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, hb, grid, nimg, resident, pack2)