"""Weight-streaming SwiGLU FFN, with an optional fused residual + next-layer LN epilogue. out = (silu(x @ w1.T + b1) * (x @ wg.T + bg)) @ w2.T + b2 (SwiGLU FFN) WHY STREAMING. ``swiglu.py`` keeps gate+up+down weights SBUF-resident (3 x K x I bf16). At ViT-H+ (K=1280, I=3416) that is ~205 KiB/partition, over the ~192 KiB budget, so the compiler spills. Here the intermediate dim I is processed in chunks of CI columns: only one chunk's gate/up columns and down rows are live (double-buffered by the scheduler: chunk c+1's DMA has no dependency on chunk c's matmuls), and the down-projection is accumulated across chunks in an fp32 SBUF accumulator with rows on partitions. LOOP ORDER. Rows (per core) are processed in super-blocks of S rows (host-chosen, <=1024). Within a super-block the chunk loop is OUTER, so each weight chunk is DMA'd once per super-block and reused by every 512-row slice in it -- weight traffic per row falls as S grows. FUSED EPILOGUE (``fuse=1``). The accumulator is initialised from the fp32 residual ``a``; after the last chunk: x_new = acc + b2 (fp32), z = LN_no_affine(x_new; eps) (bf16, same stat chain as res_layernorm.py). Returns (x_new, z). ``fuse=0`` returns bf16 ``acc + b2`` like ``swiglu.py``. Rows are sharded across LNC cores (R % (LNC*128) == 0). INFERENCE ONLY. """ from __future__ import annotations import torch from torch_neuronx import nki_op import nki import nki.isa as nisa import nki.language as nl from .mlp import pack_b1 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")) _P_MAX = 128 _F_MOV_MAX = 512 _BF16 = nl.bfloat16 def _div_ceil(n: int, d: int) -> int: return (n + d - 1) // d def _kernel_assert(cond: bool, msg: str): assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}" def plan(K: int, I: int, rows_per_core: int, budget: int = 176 * 1024): """Host-side choice of (S, CI) from a per-partition SBUF estimate (bytes).""" n_k = _div_ceil(K, _P_MAX) best = None for CI in (1024, 512, 256, 128): for S in (1024, 512, 256, 128): if S > _div_ceil(rows_per_core, _P_MAX) * _P_MAX and S != 128: continue nh = CI // _P_MAX w = 2 * (2 * n_k * CI * 2 + nh * K * 2) # 2 chunks in flight acc = (S // _P_MAX) * K * 4 # fp32 accumulator rows xt = n_k * S * 2 + (S // _P_MAX) * K * 2 # x^T + staging rows g = nh * _F_MOV_MAX * (2 + 4) # g bf16 + silu fp32 per tile ep = 2 * K * 2 + K * 4 # z tiles + b2 broadcast tot = w + acc + xt + g + ep + _div_ceil(I, _P_MAX) * 8 + 4096 if tot <= budget: n_sb = _div_ceil(rows_per_core, S) key = (-n_sb, CI, -S) # fewest weight re-streams, then big CI, then tight S if best is None or key > best[0]: best = (key, S, CI, tot) _kernel_assert(best is not None, f"no SBUF plan for K={K} I={I} rows={rows_per_core}") return best[1], best[2], best[3] def _stage_xt(x_hbm, r_abs, S, n_k, K): xt = [] for _ in range(n_k): xt.append(nl.ndarray((_P_MAX, S), dtype=_BF16, buffer=nl.sbuf)) for rs in range(S // _P_MAX): rr0 = rs * _P_MAX xb = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf) nisa.dma_copy(dst=xb[0:_P_MAX, 0:K], src=x_hbm[r_abs + rr0:r_abs + rr0 + _P_MAX, 0:K]) for kt in range(n_k): k0 = kt * _P_MAX pk = min(_P_MAX, K - k0) nisa.dma_transpose(dst=xt[kt][0:pk, rr0:rr0 + _P_MAX], src=xb[0:_P_MAX, k0:k0 + pk]) return xt def _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI, fuse): R, K = x_hbm.shape K1, I = w1t_hbm.shape _kernel_assert(K1 == K and tuple(wgt_hbm.shape) == (K, I), "w1t/wgt must be [K, I]") _kernel_assert(tuple(w2t_hbm.shape) == (I, K), "w2t must be [I, K]") _kernel_assert(R % (LNC * _P_MAX) == 0, f"R={R} must be a multiple of LNC*128") rows_per = R // LNC _kernel_assert(S % _P_MAX == 0, f"S={S} must be a multiple of 128") _kernel_assert(CI % _P_MAX == 0, "CI must be a multiple of 128") pid = nl.program_id(0) r_base = pid * rows_per n_k = _div_ceil(K, _P_MAX) n_h = _div_ceil(I, _P_MAX) n_c = _div_ceil(I, CI) n_ko = _div_ceil(K, _F_MOV_MAX) n_sb = _div_ceil(rows_per, S) if fuse: xnew_hbm = nl.ndarray((R, K), dtype=nl.float32, buffer=nl.shared_hbm) z_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm) else: out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm) # biases (tiny, resident). b1/bg are host-packed to [128, I/128] (pack_bias; b[p, ht] == # bias[ht*128 + p]) and each lands in ONE DMA on the Scalar HW-DGE queue -- replaces 2 * I/128 # [128, 1] DMAs (~54 per core at ViT-H+) that stalled the first silu/gate ops. b2 broadcast over # partitions. _kernel_assert(I % _P_MAX == 0, "I must be padded to a multiple of 128 (pack_bias)") _kernel_assert(tuple(b1_hbm.shape) == (_P_MAX, n_h) and tuple(bg_hbm.shape) == (_P_MAX, n_h), f"b1/bg must be pack_bias(...): [{_P_MAX}, {n_h}]") b1_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf) nisa.dma_copy(dst=b1_sb, src=b1_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar) bg_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf) nisa.dma_copy(dst=bg_sb, src=bg_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar) b2_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf) nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:K], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0)) inv_k = 1.0 / float(K) neg_inv_k = -inv_k for sb in range(n_sb): r_sb = r_base + sb * S Ssb = min(S, rows_per - sb * S) # ragged last super-block (multiple of 128) n_rs = Ssb // _P_MAX xt = _stage_xt(x_hbm, r_sb, Ssb, n_k, K) acc = [] for rs in range(n_rs): t = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf) if fuse: nisa.dma_copy(dst=t[0:_P_MAX, 0:K], src=a_hbm[r_sb + rs * _P_MAX:r_sb + (rs + 1) * _P_MAX, 0:K]) acc.append(t) for c in range(n_c): i0 = c * CI ci = min(CI, I - i0) nh_c = _div_ceil(ci, _P_MAX) # ---- this chunk's weights (gate/up columns i0:i0+ci, down rows i0:i0+ci) ---- w1c = [] wgc = [] for kt in range(n_k): k0 = kt * _P_MAX pk = min(_P_MAX, K - k0) t1 = nl.ndarray((_P_MAX, CI), dtype=_BF16, buffer=nl.sbuf) nisa.dma_copy(dst=t1[0:pk, 0:ci], src=w1t_hbm[k0:k0 + pk, i0:i0 + ci]) w1c.append(t1) tg = nl.ndarray((_P_MAX, CI), dtype=_BF16, buffer=nl.sbuf) nisa.dma_copy(dst=tg[0:pk, 0:ci], src=wgt_hbm[k0:k0 + pk, i0:i0 + ci]) wgc.append(tg) w2c = [] for hl in range(nh_c): h0 = i0 + hl * _P_MAX ph = min(_P_MAX, I - h0) t2 = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf) nisa.dma_copy(dst=t2[0:ph, 0:K], src=w2t_hbm[h0:h0 + ph, 0:K]) w2c.append(t2) for rb in range(_div_ceil(Ssb, _F_MOV_MAX)): q0 = rb * _F_MOV_MAX rw = min(_F_MOV_MAX, Ssb - q0) # ---- gate/up for this chunk x this 512-row slice -> g [ci, rw] bf16 ---- g_t = [] for hl in range(nh_c): hg = (i0 // _P_MAX) + hl ph = min(_P_MAX, I - (i0 + hl * _P_MAX)) c0 = hl * _P_MAX y1 = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum) for kt in nl.affine_range(n_k): pk = min(_P_MAX, K - kt * _P_MAX) nisa.nc_matmul(dst=y1[0:ph, 0:rw], stationary=w1c[kt][0:pk, c0:c0 + ph], moving=xt[kt][0:pk, q0:q0 + rw]) yg = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum) for kt in nl.affine_range(n_k): pk = min(_P_MAX, K - kt * _P_MAX) nisa.nc_matmul(dst=yg[0:ph, 0:rw], stationary=wgc[kt][0:pk, c0:c0 + ph], moving=xt[kt][0:pk, q0:q0 + rw]) sa = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf) nisa.activation(dst=sa[0:ph, 0:rw], op=nl.silu, data=y1[0:ph, 0:rw], bias=b1_sb[0:ph, hg:hg + 1]) g = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf) # g = (yg + bg) * silu(y1 + b1) -- one fused op, PSUM operand nisa.scalar_tensor_tensor(dst=g[0:ph, 0:rw], data=yg[0:ph, 0:rw], op0=nl.add, operand0=bg_sb[0:ph, hg:hg + 1], op1=nl.multiply, operand1=sa[0:ph, 0:rw]) g_t.append(g) # ---- down-projection partial for this chunk, accumulated into acc (fp32 SBUF) ---- for rs_l in range(rw // _P_MAX): rs = (q0 // _P_MAX) + rs_l for ko in range(n_ko): ko0 = ko * _F_MOV_MAX nk = min(_F_MOV_MAX, K - ko0) ps = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum) for hl in nl.affine_range(nh_c): ph = min(_P_MAX, I - (i0 + hl * _P_MAX)) nisa.nc_matmul(dst=ps[0:_P_MAX, 0:nk], stationary=g_t[hl][0:ph, rs_l * _P_MAX:(rs_l + 1) * _P_MAX], moving=w2c[hl][0:ph, ko0:ko0 + nk]) if c == 0 and not fuse: nisa.tensor_copy(dst=acc[rs][0:_P_MAX, ko0:ko0 + nk], src=ps[0:_P_MAX, 0:nk]) else: nisa.tensor_tensor(dst=acc[rs][0:_P_MAX, ko0:ko0 + nk], data1=acc[rs][0:_P_MAX, ko0:ko0 + nk], data2=ps[0:_P_MAX, 0:nk], op=nl.add) # ---- epilogue per 128-row subtile ---- for rs in range(n_rs): r0 = r_sb + rs * _P_MAX row = acc[rs] nisa.tensor_tensor(dst=row[0:_P_MAX, 0:K], data1=row[0:_P_MAX, 0:K], data2=b2_bc[0:_P_MAX, 0:K], op=nl.add) zt = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf) if fuse: nisa.dma_copy(dst=xnew_hbm[r0:r0 + _P_MAX, 0:K], src=row[0:_P_MAX, 0:K]) neg_shift = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) acc_s = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) neg_mu = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) var = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) rstd = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) nisa.tensor_scalar(dst=neg_shift, data=row[0:_P_MAX, 0:1], op0=nl.multiply, operand0=-1.0) nisa.activation_reduce(dst=zt, op=nl.copy, data=row[0:_P_MAX, 0:K], reduce_op=nl.add, reduce_res=acc_s, bias=neg_shift, scale=1.0) nisa.tensor_scalar(dst=neg_mu, data=acc_s, op0=nl.multiply, operand0=neg_inv_k, op1=nl.add, operand1=neg_shift) nisa.activation_reduce(dst=zt, op=nl.square, data=row[0:_P_MAX, 0:K], reduce_op=nl.add, reduce_res=acc_s, bias=neg_mu, scale=1.0) nisa.tensor_scalar(dst=var, data=acc_s, op0=nl.multiply, operand0=inv_k, op1=nl.add, operand1=eps) nisa.activation(dst=rstd, op=nl.rsqrt, data=var) nisa.tensor_scalar(dst=zt, data=row[0:_P_MAX, 0:K], op0=nl.add, operand0=neg_mu, op1=nl.multiply, operand1=rstd) nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:K], src=zt) else: nisa.tensor_copy(dst=zt, src=row[0:_P_MAX, 0:K]) nisa.dma_copy(dst=out_hbm[r0:r0 + _P_MAX, 0:K], src=zt) if fuse: return xnew_hbm, z_hbm return out_hbm @nki.jit def _swiglu_stream_kernel(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, S, CI): return _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, None, 0.0, S, CI, 0) @nki.jit def _swiglu_stream_ln_kernel(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI): return _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI, 1) def pack_bias(b: torch.Tensor) -> torch.Tensor: """Pack an fp32 bias ``[I]`` (I % 128 == 0) to ``[128, I/128]`` so ``b_p[p, t] == b[t*128 + p]``.""" if b.dim() == 2: return b assert b.shape[0] % _P_MAX == 0, f"pack_bias needs I % {_P_MAX} == 0, got {b.shape[0]}" return b.reshape(-1, _P_MAX).t().contiguous() def _no_grad(*ts): for t in ts: if t is not None and t.requires_grad: raise RuntimeError("swiglu_stream is inference-only; use torch.no_grad().") def _plan_for(x, w1t): S, CI, _ = plan(x.shape[1], w1t.shape[1], x.shape[0] // LNC) S = int(_os.environ.get("VIT_ENC_SW_S", S)) CI = int(_os.environ.get("VIT_ENC_SW_CI", CI)) return S, CI @nki_op("vit_enc::nki_swiglu_stream", mutates_args={}) def nki_swiglu_stream(x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, wgt: torch.Tensor, bg: torch.Tensor, w2t: torch.Tensor, b2: torch.Tensor) -> torch.Tensor: """x [R,K] bf16; w1t/wgt [K,I] bf16; b1/bg [I] f32; w2t [I,K] bf16; b2 [K] f32 -> [R,K] bf16.""" _no_grad(x, w1t, b1, wgt, bg, w2t, b2) S, CI = _plan_for(x, w1t) return _swiglu_stream_kernel[LNC](x, w1t, pack_bias(b1), wgt, pack_bias(bg), w2t, b2, S, CI) @nki_op("vit_enc::nki_swiglu_stream_ln", mutates_args={}) def nki_swiglu_stream_ln(x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, wgt: torch.Tensor, bg: torch.Tensor, w2t: torch.Tensor, b2: torch.Tensor, a: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: """SwiGLU(x) + residual a, then affine-free LN. Returns (x_new [R,K] fp32, z [R,K] bf16).""" _no_grad(x, w1t, b1, wgt, bg, w2t, b2, a) S, CI = _plan_for(x, w1t) return _swiglu_stream_ln_kernel[LNC](x, w1t, pack_bias(b1), wgt, pack_bias(bg), w2t, b2, a, eps, S, CI)