Hugging Face's logo

Kernels:
aws-neuron
/
vit-encoder-neuron-kernels

Kernel card Files
xet
 Community
vit-encoder-neuron-kernels / build /torch-neuron /nki_kernels /swiglu_stream.py
jburtoft's picture
jburtoft
v1.2: packed-bias one-DMA for SwiGLU (swiglu_stream/_ln)
7f9eb4c
Raw History Blame
15.3 kB
"""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)