v1.7: res_layernorm_fold processes row tiles in groups when they do not fit SBUF
Browse filesLarge batches (e.g. ViT-S/16 LNC=1 batch 32, ViT-B/L/H+ LNC=1 batch 16) failed to compile with NCC_IBIR229.
Shapes that compiled before keep the same schedule: bit-identical output and the same speed (within 0.12%).
- CHANGELOG.md +31 -0
- build/torch-neuron/nki_kernels/res_layernorm.py +76 -57
CHANGELOG.md
CHANGED
|
@@ -1,5 +1,36 @@
|
|
| 1 |
# Changelog
|
| 2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
## v1.6 (2026-10)
|
| 4 |
|
| 5 |
* **Two images per producer pass in `attention_block` too** (new `pack2` argument; the encoder enables it for even
|
|
|
|
| 1 |
# Changelog
|
| 2 |
|
| 3 |
+
## v1.7 (2026-10)
|
| 4 |
+
|
| 5 |
+
* **Fix: `res_layernorm_fold` failed to compile at large batch** (`NCC_IBIR229`, state buffer allocation failed). The
|
| 6 |
+
kernel staged every row tile of a core in SBUF at once. Above about 200 KiB per partition that overflowed:
|
| 7 |
+
- ViT-S/16 at LNC=1, batch 32 at 224 px or batch 8 at 448 px;
|
| 8 |
+
- ViT-B/16 at LNC=1, batch 16 at 224 px;
|
| 9 |
+
- ViT-L/16 at LNC=1, batch 16 at 224 px;
|
| 10 |
+
- ViT-H+/16 at LNC=1, batch 16 at 224 px.
|
| 11 |
+
|
| 12 |
+
Each failure killed the worker process. Tiles are now processed in groups that fit, with the same per-tile
|
| 13 |
+
instructions in the same order. Every shape that compiled before keeps its schedule and is **bit-identical** to v1.6
|
| 14 |
+
(checked at 9 shapes, LNC=1 and LNC=2). The newly compiling shapes match a float64 reference
|
| 15 |
+
(z rel_l2 1.66e-3, the same as before).
|
| 16 |
+
|
| 17 |
+
Per instance (LNC=1, 8 cores) at the newly reachable batch sizes:
|
| 18 |
+
|
| 19 |
+
| model | batch | img/s | best before |
|
| 20 |
+
|---|---|---|---|
|
| 21 |
+
| ViT-S/16 224 px | 32 / 64 | 8,205.4 / 8,009.5 | 9,422.6 (batch 16) |
|
| 22 |
+
| ViT-S/16 448 px | 8 / 16 | 2,176.4 / 2,192.3 | 2,486.1 (batch 4) |
|
| 23 |
+
| ViT-B/16 224 px | 16 / 32 | 4,775.4 / 4,679.0 | 4,737.8 (batch 8) |
|
| 24 |
+
| ViT-L/16 224 px | 16 / 32 | 1,892.3 / 1,866.8 | 1,905.8 (batch 8) |
|
| 25 |
+
| ViT-H+/16 224 px | 16 | 737.2 | 742.7 (batch 8) |
|
| 26 |
+
|
| 27 |
+
None of these larger batch sizes beat the best smaller one (ViT-B 16 is +0.8%, within noise), so the recommended
|
| 28 |
+
batch sizes are unchanged. The fix means a user-chosen batch no longer crashes.
|
| 29 |
+
|
| 30 |
+
**No cost at the recommended batch sizes:** grouping turns on only when the tiles do not fit. v1.6 vs v1.7 per instance,
|
| 31 |
+
same instance, 2 runs each: ViT-S 224 BS 16 -0.02%, ViT-B 224 BS 8 -0.12%, ViT-L 224 BS 8 +0.04%, ViT-H+ 224 BS 8 +0.01%,
|
| 32 |
+
ViT-L 448 BS 2 -0.03% (LNC=1); ViT-L 224 BS 8 LNC=2 +0.12%.
|
| 33 |
+
|
| 34 |
## v1.6 (2026-10)
|
| 35 |
|
| 36 |
* **Two images per producer pass in `attention_block` too** (new `pack2` argument; the encoder enables it for even
|
build/torch-neuron/nki_kernels/res_layernorm.py
CHANGED
|
@@ -20,7 +20,9 @@ WHY IT IS SHAPED THIS WAY
|
|
| 20 |
* LNC = 2 shard: the kernel is launched on both NeuronCores; each owns ``R // 2`` rows
|
| 21 |
selected by ``program_id(0)``, so the two cores never touch the same tokens.
|
| 22 |
* Every token tile has its own SBUF slot ([P, n_t, H]) and all loads are issued before any
|
| 23 |
-
compute, so tile ``t+1``'s loads never wait behind tile ``t``'s stores.
|
|
|
|
|
|
|
| 24 |
operand: a single whole-shard DMA would signal completion only at the end. (External
|
| 25 |
contribution by jlongen-aws, PR #2 on ``jburtoft/vjepa2-neuron-kernels``; replaces the
|
| 26 |
depth-4 ring, same instructions in the same order so results are unchanged.)
|
|
@@ -59,6 +61,10 @@ _P_MAX = 128
|
|
| 59 |
#: The bf16 dtype used for branch inputs and the normalized output.
|
| 60 |
_BF16 = nl.bfloat16
|
| 61 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 62 |
#: Logical-to-physical NeuronCore degree. Rows are split evenly across LNC cores.
|
| 63 |
import os as _os
|
| 64 |
#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
|
|
@@ -114,67 +120,80 @@ def _res_layernorm_fold_kernel(a_hbm, b_hbm, eps):
|
|
| 114 |
inv_h = 1.0 / float(H)
|
| 115 |
neg_inv_h = -inv_h
|
| 116 |
|
| 117 |
-
#
|
| 118 |
-
#
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
|
| 123 |
-
|
| 124 |
-
|
| 125 |
-
|
| 126 |
-
nisa.dma_copy(dst=b_sb[:, j, t, :], src=b_hbm[j, r0:r0 + _P_MAX, 0:H],
|
| 127 |
-
dge_mode=nisa.dge_mode.none)
|
| 128 |
-
|
| 129 |
-
s_sb = nl.ndarray((_P_MAX, n_t, H), dtype=nl.float32, buffer=nl.sbuf)
|
| 130 |
-
z_sb = nl.ndarray((_P_MAX, n_t, H), dtype=_BF16, buffer=nl.sbuf)
|
| 131 |
|
| 132 |
# Scratch destination for the reduce activations (its full tile is discarded; only
|
| 133 |
# the reduce result is consumed).
|
| 134 |
sink = nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf)
|
| 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 |
return x_new_hbm, z_hbm
|
| 180 |
|
|
|
|
| 20 |
* LNC = 2 shard: the kernel is launched on both NeuronCores; each owns ``R // 2`` rows
|
| 21 |
selected by ``program_id(0)``, so the two cores never touch the same tokens.
|
| 22 |
* Every token tile has its own SBUF slot ([P, n_t, H]) and all loads are issued before any
|
| 23 |
+
compute, so tile ``t+1``'s loads never wait behind tile ``t``'s stores. When the tiles of one core do not
|
| 24 |
+
fit SBUF (large batch: e.g. ViT-S >= 8k rows/core, ViT-L >= ~2k), they are processed in groups that do,
|
| 25 |
+
with the same pattern inside each group (v1.7; up to v1.6 these shapes failed to compile). One DMA per tile and
|
| 26 |
operand: a single whole-shard DMA would signal completion only at the end. (External
|
| 27 |
contribution by jlongen-aws, PR #2 on ``jburtoft/vjepa2-neuron-kernels``; replaces the
|
| 28 |
depth-4 ring, same instructions in the same order so results are unchanged.)
|
|
|
|
| 61 |
#: The bf16 dtype used for branch inputs and the normalized output.
|
| 62 |
_BF16 = nl.bfloat16
|
| 63 |
|
| 64 |
+
#: Per-partition SBUF bytes one tile group's a/b/s slots may use. 200 KiB is the largest footprint measured to
|
| 65 |
+
#: compile (ViT-H+, 16 tiles, one branch); 210 KiB failed. Above it, tiles are processed in groups (v1.7).
|
| 66 |
+
_SBUF_GROUP_BUDGET = 200 * 1024
|
| 67 |
+
|
| 68 |
#: Logical-to-physical NeuronCore degree. Rows are split evenly across LNC cores.
|
| 69 |
import os as _os
|
| 70 |
#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
|
|
|
|
| 120 |
inv_h = 1.0 / float(H)
|
| 121 |
neg_inv_h = -inv_h
|
| 122 |
|
| 123 |
+
# Tile GROUPS (v1.7): the kernel stages every tile of a group in SBUF at once (loads first, see the
|
| 124 |
+
# module doc), so the group size G is capped by the SBUF budget. Measured footprint per tile and
|
| 125 |
+
# partition is the a (fp32), b (bf16 x n_b) and s (fp32) slots, H * (8 + 2 n_b) bytes -- the compiler
|
| 126 |
+
# reuses z's slot (ViT-H+ 16 tiles = 200 KiB compiles; ViT-S 56 tiles = 210 KiB does not).
|
| 127 |
+
# G = all tiles when they fit (the <= v1.6 schedule, unchanged), else the largest G under
|
| 128 |
+
# _SBUF_GROUP_BUDGET. Same per-tile instructions in the same order -> bit-identical results.
|
| 129 |
+
per_tile = H * (8 + 2 * n_b)
|
| 130 |
+
G = max(1, min(n_t, _SBUF_GROUP_BUDGET // per_tile))
|
| 131 |
+
n_g = (n_t + G - 1) // G
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 132 |
|
| 133 |
# Scratch destination for the reduce activations (its full tile is discarded; only
|
| 134 |
# the reduce result is consumed).
|
| 135 |
sink = nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf)
|
| 136 |
|
| 137 |
+
for g in range(n_g):
|
| 138 |
+
t0 = g * G
|
| 139 |
+
gt = min(G, n_t - t0)
|
| 140 |
+
# One SBUF slot per tile of this group; all of the group's loads are issued before any compute or store.
|
| 141 |
+
# a_sb[:, t, :] is tile t's residual, b_sb[:, j, t, :] is branch j of tile t.
|
| 142 |
+
a_sb = nl.ndarray((_P_MAX, G, H), dtype=nl.float32, buffer=nl.sbuf)
|
| 143 |
+
b_sb = nl.ndarray((_P_MAX, n_b, G, H), dtype=_BF16, buffer=nl.sbuf)
|
| 144 |
+
for t in range(gt):
|
| 145 |
+
r0 = r_base + (t0 + t) * _P_MAX
|
| 146 |
+
nisa.dma_copy(dst=a_sb[:, t, :], src=a_hbm[r0:r0 + _P_MAX, 0:H],
|
| 147 |
+
dge_mode=nisa.dge_mode.none)
|
| 148 |
+
for j in range(n_b):
|
| 149 |
+
nisa.dma_copy(dst=b_sb[:, j, t, :], src=b_hbm[j, r0:r0 + _P_MAX, 0:H],
|
| 150 |
+
dge_mode=nisa.dge_mode.none)
|
| 151 |
+
|
| 152 |
+
s_sb = nl.ndarray((_P_MAX, G, H), dtype=nl.float32, buffer=nl.sbuf)
|
| 153 |
+
z_sb = nl.ndarray((_P_MAX, G, H), dtype=_BF16, buffer=nl.sbuf)
|
| 154 |
+
|
| 155 |
+
# Per-lane statistics, one column per tile so tile t+1's chain doesn't wait on tile t's.
|
| 156 |
+
acc = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
|
| 157 |
+
acc2 = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
|
| 158 |
+
neg_shift = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
|
| 159 |
+
neg_mu = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
|
| 160 |
+
var = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
|
| 161 |
+
rstd = nl.ndarray((_P_MAX, G), dtype=nl.float32, buffer=nl.sbuf)
|
| 162 |
+
|
| 163 |
+
for t in range(gt):
|
| 164 |
+
r0 = r_base + (t0 + t) * _P_MAX
|
| 165 |
+
s = s_sb[:, t, :]
|
| 166 |
+
c = slice(t, t + 1)
|
| 167 |
+
|
| 168 |
+
# ---- x_new = a + sum_j b[j] in fp32, then write the residual stream out ----
|
| 169 |
+
nisa.tensor_tensor(dst=s, data1=a_sb[:, t, :], data2=b_sb[:, 0, t, :], op=nl.add)
|
| 170 |
+
for j in range(1, n_b):
|
| 171 |
+
nisa.tensor_tensor(dst=s, data1=s, data2=b_sb[:, j, t, :], op=nl.add)
|
| 172 |
+
|
| 173 |
+
nisa.dma_copy(dst=x_new_hbm[r0:r0 + _P_MAX, 0:H], src=s, dge_mode=nisa.dge_mode.none)
|
| 174 |
+
|
| 175 |
+
# ---- thinned stat chain: mean via shift-and-reduce, variance, rstd ----
|
| 176 |
+
nisa.tensor_scalar(dst=neg_shift[:, c], data=s_sb[:, t, 0:1],
|
| 177 |
+
op0=nl.multiply, operand0=-1.0)
|
| 178 |
+
nisa.activation_reduce(dst=sink, op=nl.copy, data=s,
|
| 179 |
+
reduce_op=nl.add, reduce_res=acc[:, c],
|
| 180 |
+
bias=neg_shift[:, c], scale=1.0)
|
| 181 |
+
nisa.tensor_scalar(dst=neg_mu[:, c], data=acc[:, c], op0=nl.multiply, operand0=neg_inv_h,
|
| 182 |
+
op1=nl.add, operand1=neg_shift[:, c])
|
| 183 |
+
nisa.activation_reduce(dst=sink, op=nl.square, data=s,
|
| 184 |
+
reduce_op=nl.add, reduce_res=acc2[:, c], bias=neg_mu[:, c], scale=1.0)
|
| 185 |
+
nisa.tensor_scalar(dst=var[:, c], data=acc2[:, c], op0=nl.multiply, operand0=inv_h,
|
| 186 |
+
op1=nl.add, operand1=eps)
|
| 187 |
+
nisa.activation(dst=rstd[:, c], op=nl.rsqrt, data=var[:, c])
|
| 188 |
+
|
| 189 |
+
# ---- one-op affine-free apply: z = (x_new + neg_mu) * rstd, bf16 cast on write.
|
| 190 |
+
# Both operands are per-lane vectors; the fused two-op tensor_scalar costs one op.
|
| 191 |
+
nisa.tensor_scalar(dst=z_sb[:, t, :], data=s,
|
| 192 |
+
op0=nl.add, operand0=neg_mu[:, c],
|
| 193 |
+
op1=nl.multiply, operand1=rstd[:, c])
|
| 194 |
+
|
| 195 |
+
nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:H], src=z_sb[:, t, :],
|
| 196 |
+
dge_mode=nisa.dge_mode.none)
|
| 197 |
|
| 198 |
return x_new_hbm, z_hbm
|
| 199 |
|