"""Attention for LNC=1: ONE program over ALL head groups. The 2-core kernel in ``attention_block.py`` owns one 8-head group per program. At LNC=1 that means one launch per group per layer (2 for ViT-L, 3 for ViT-H+), each re-transposing x, each re-loading the rope tables, each writing an [N, H] bf16 out-projection PARTIAL to HBM that the LN-fold kernel then sums. This kernel, per image (``nimg`` images stacked on the row axis, one launch): * transposes x once and loads the rope tables once (shared by every group: tables are quad-tiled per head, identical for all heads); * for each group: projects Q/K/V + RoPE (same producer as attention_block), runs the vendor attention per head, and writes O feature-major into a PERSISTENT per-group SBUF slab; * after all groups: ONE out-projection contracting over ALL features (G*4 chunks accumulate in one fp32 PSUM tile), adds the FULL bias once, writes ONE [N, H] result. No partials, no host sum; the LN-fold kernel receives a single branch. ``resident=1``: Q/K/V/O weights, biases and rope tables are staged ONCE per launch (bottom of SBUF) and reused by every image. Host decides by an SBUF estimate (does not fit at N>=1024 for ViT-L); ``resident=0`` restages each group's weights per image. N<=2048 only (the shipping single-section vendor path). Numerics: identical producer and vendor attention; the out-projection now accumulates all groups in fp32 PSUM instead of summing bf16 per-group partials, so results differ from attention_block by bf16 rounding of the partials. """ 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 from .attention_inner import _attention_cte from .attention_block import pack_qk_bias as _pack_qk_bias _P_MAX = 128 _PSUM_FMAX = 512 _PSUM_BANK = 2048 _K_TILE = 512 _V_TILE = 128 def _kernel_assert(cond: bool, msg: str): assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}" def sbuf_estimate(H: int, Fp: int, N: int, D: int, resident: bool) -> int: """Conservative bytes/lane at the kernel's peak (host-side decision helper, not exact).""" G = Fp // _PSUM_FMAX n_chunks = H // _P_MAX n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX n_ktiles = (N + _K_TILE - 1) // _K_TILE lh = _PSUM_FMAX // D n_och = (lh * D) // _P_MAX w_grp = 3 * n_chunks * _PSUM_FMAX * 2 + _PSUM_FMAX * 4 + 64 # wqk + wv + bv_bc + biases w_out = G * n_och * H * 2 + H * 4 # wo + bo_bc tab = 4 * n_tt * _PSUM_FMAX * 4 o_feat = G * n_och * N * 2 x_t = n_chunks * n_tt * _PSUM_FMAX * 2 kv = lh * (n_ktiles * _K_TILE * 2 + (N // _V_TILE) * D * 2) work = 2 * _PSUM_FMAX * 2 * 2 + 2 * _PSUM_FMAX * 4 + _PSUM_FMAX * 4 inner = 48 * 1024 resid = (G * w_grp + w_out + tab) if resident else 0 base = resid + o_feat + x_t + (0 if resident else tab) p1 = kv + (0 if resident else w_grp) + work p2 = kv + inner p3 = (0 if resident else w_out) + 2048 return base + max(p1, p2, p3) @nki.jit def _attention_block_l1_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm, wot_hbm, bo_hbm, scale, head_dim, n_valid, nimg=1, resident=0): NT, H = h_hbm.shape N = NT // nimg Fp = wqk_hbm.shape[2] D = head_dim half = D // 2 quad = _P_MAX // half local_heads = _PSUM_FMAX // D local_quads = local_heads // quad G = Fp // _PSUM_FMAX # head groups, ALL handled by this one program n_chunks = H // _P_MAX n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX n_ktiles = (N + _K_TILE - 1) // _K_TILE n_vtiles = N // _V_TILE n_ochunks = (local_heads * D) // _P_MAX # feature chunks per group v_per_tt = _PSUM_FMAX // _V_TILE _kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk") _kernel_assert(bqk_hbm.shape == (2, _P_MAX, Fp // _P_MAX), "bqk must be pack_qk_bias(...)") _kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab") _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, "H must be a multiple of 128") _kernel_assert(quad * half == _P_MAX, "D must give quad*half=128") _kernel_assert(Fp % _PSUM_FMAX == 0, "padded features must be a multiple of 512") _kernel_assert(N % _P_MAX == 0, "N must be a multiple of 128") _kernel_assert(N <= _PSUM_FMAX * 4, "N>2048 not supported") _kernel_assert(NT == nimg * N, "rows must be nimg x N") result = nl.ndarray((1, NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm) # ---------------- resident (per-launch invariant) staging ---------------- res_top = 0 r_w = [] r_wv = [] r_b = [] r_bv = [] if resident: ra = ModularAllocator(initial_address=0) for g in range(G): col0 = g * _PSUM_FMAX 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=w_sb_[qk][c], src=wqk_hbm[qk, c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_quads * 2 * _P_MAX]) 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=wv_sb_[c], src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + 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). 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=b_sb_[qk], src=bqk_hbm[qk, 0:_P_MAX, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + local_quads * 2], dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar) bv_bc_ = ra.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=col0)) r_w.append(w_sb_) r_wv.append(wv_sb_) r_b.append(b_sb_) r_bv.append(bv_bc_) 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 = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16, block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks]) for oc in range(G * n_ochunks): nisa.dma_copy(dst=r_wo[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H]) r_bo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4) nisa.dma_copy(dst=r_bo, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0)) res_top = ra.get_current_address() n_grp = N // _P_MAX for img in range(nimg): rofs = img * N * H allocator = ModularAllocator(initial_address=res_top) # O for ALL groups stays in SBUF until the single out-projection. o_feat = [] for g in range(G): of_ = allocator.alloc_sbuf_tensor(shape=(_P_MAX, N), dtype=nl.bfloat16, block_dim=[n_ochunks], num_free_tiles=[n_ochunks], align_to=32) o_feat.append(of_) # x^T: shared by every group (transposed ONCE per image). x_t = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16, block_dim=[n_chunks, n_tt], num_free_tiles=[n_chunks, n_tt], align_to=32) for s in range(n_tt): 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, N - s * _PSUM_FMAX)]], offset=0), src=h_hbm.ap([[H, min(_PSUM_FMAX, N - s * _PSUM_FMAX)], [1, 1], [1, 1], [1, _P_MAX]], offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX), ) # rope tables: shared by every group. if resident: tab = r_tab else: tab = allocator.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=tab[t_i][s][:, 0:_tw], src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw]) img_base = allocator.get_current_address() for g in range(G): allocator.set_current_address(img_base) col0 = g * _PSUM_FMAX k_tiles = [] v_tiles = [] for _lh in range(local_heads): kt_ = allocator.alloc_sbuf_tensor(shape=(D, _K_TILE), dtype=nl.bfloat16, block_dim=[n_ktiles], num_free_tiles=[n_ktiles], align_to=32) k_tiles.append(kt_) vt_ = allocator.alloc_sbuf_tensor(shape=(_V_TILE, D), dtype=nl.bfloat16, block_dim=[n_vtiles], num_free_tiles=[n_vtiles]) v_tiles.append(vt_) phase2_base = allocator.get_current_address() if resident: w_sb = r_w[g] wv_sb = r_wv[g] b_sb = r_b[g] bv_bc = r_bv[g] else: 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, col0:col0 + 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). 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, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + 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=col0)) 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_priv = [] for _lh in range(local_heads): q_priv.append(nl.ndarray((1, D, N), dtype=nl.bfloat16, buffer=nl.private_hbm)) # ---- producer (identical to attention_block's N<=2048 path) ---- for tt in range(n_tt): t0 = tt * _PSUM_FMAX tw = min(_PSUM_FMAX, N - t0) 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): nisa.tensor_copy(dst=v_tiles[lh][tt * v_per_tt + vsub], src=vout[:, lh * D:(lh + 1) * D], engine=engine.scalar) for qk in range(2): for j in range(local_quads): for par in range(2): gq = j * 2 + par cols = gq * _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)) nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw], op0=nl.add, operand0=b_sb[qk][:, gq:gq + 1]) 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) for b in range(quad): lh = j * quad + b r0 = b * half if qk == 1: for par2 in range(2): nisa.tensor_copy( dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half, 0:tw], src=roped[par2][r0:r0 + half, 0:tw]) else: for par2 in range(2): nisa.dma_copy( dst=q_priv[lh][0, par2 * half:(par2 + 1) * half, t0:t0 + tw], src=roped[par2][r0:r0 + half, 0:tw]) # ---- attention for this group's heads -> o_feat[g] (SBUF, feature-major) ---- allocator.set_current_address(phase2_base) o_sb = [] for _lh in range(local_heads): oc = (_lh * D) // _P_MAX rr = (_lh * D) % _P_MAX grp = [] for _g in range(n_grp): grp.append(o_feat[g][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): _attention_cte( q_priv[lh], None, None, scale=scale, causal_mask=False, tp_q=False, tp_k=False, tp_out=True, cache_softmax=False, k_cache_sbuf=k_tiles[lh], v_cache_sbuf=v_tiles[lh], out_o_hbm=o_sb[lh], init_sbuf_addr=attn_base, seqlen_k_valid=n_valid, ) allocator.set_current_address(attn_base) # ---- ONE out-projection over ALL G*n_ochunks feature chunks, full bias once ---- allocator.set_current_address(img_base) if resident: wo_sb = r_wo bo_bc = r_bo else: wo_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16, block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks]) for oc in range(G * n_ochunks): nisa.dma_copy(dst=wo_sb[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H]) 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)) 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(n_grp): for ob in range(n_ob): ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX) ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum, address=(0, 0)) for g in range(G): for oc in range(n_ochunks): nisa.nc_matmul( dst=ps_out[:, 0:ow], stationary=o_feat[g][oc][:, tb * _P_MAX:(tb + 1) * _P_MAX], moving=wo_sb[g * n_ochunks + oc][:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow], accumulate=((g > 0) or (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=rofs + tb * _P_MAX * H + ob * _PSUM_FMAX), src=out_sb[:, 0:ow]) return result _ATTN_L1 = wrap_nki_to_torch_hop(_attention_block_l1_kernel[1]) @nki_op("vit_enc::attention_block_l1", mutates_args={}) def attention_block_l1( 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, nimg: int = 1, resident: int = 0, ) -> torch.Tensor: """All head groups in one program. Returns ONE [1, nimg*N, H] bf16 result (full bias).""" for t in (h, wqk, bqk, tab, wvt, bv, wot, bo): if t is not None and t.requires_grad: raise RuntimeError("attention_block_l1 is inference-only; use torch.no_grad().") if h.shape[0] // nimg > 2048: raise NotImplementedError("attention_block_l1 supports N<=2048 only") if bqk.dim() == 3 and bqk.shape[1] != 128: bqk = _pack_qk_bias(bqk) return _ATTN_L1(h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, nimg, resident)