| """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 |
| w_out = G * n_och * H * 2 + H * 4 |
| 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 |
| 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 |
| 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) |
|
|
| |
| 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]) |
| |
| |
| 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_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 = 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), |
| ) |
|
|
| |
| 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]) |
| |
| |
| 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)) |
|
|
| |
| 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]) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
