CHANGELOG and comment wording

6da9fa8 3 days ago

52.6 kB

	"""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(wqg1, wkg1, 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(wqg1, wkg1, 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' = atab0 + btab2`` and ``b' = btab1 + atab3``, so:
	* ``"half"``: ``sin`` is the PLAIN sin table (rotate-half: ``xcos + 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 bN and writing partials to result[slot, bN:(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 (8445122 =
	# 131072), mm1_masked (2420484 = 65536), exp_sb (2420482 = 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 [b32,(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)