File size: 15,281 Bytes
c855122 7f9eb4c c855122 7f9eb4c c855122 7f9eb4c c855122 7f9eb4c c855122 7f9eb4c c855122 7f9eb4c c855122 7f9eb4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 | """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)
|