| # Changelog |
|
|
| ## v1.9 (2026-10) |
|
|
| * **New: padding masks on the kernel path for `NeuronXLMRobertaModel`.** v1.8 sent any batch with a padding |
| `attention_mask` to the stock forward. The attention kernel now takes an optional per-image additive key bias |
| (`attention_block_l1_ln_post_kb`: 0 for attended keys, -1e9 for masked keys), and the layer turns the mask into that |
| bias, so variable-length batches stay on the kernels. `VIT_ENC_MASK_KERNEL=0` restores the v1.8 fallback. |
|
|
| XLM-RoBERTa-Large, LNC=1, 8 processes on one trn2 device, a log-normal length mix padded to the sequence length |
| (mean valid length 53 / 89 / 191), inferences/s: |
|
|
| | seq | stock with mask (v1.8 fallback) | v1.9 kernels with mask | speedup | v1.9 kernels, no mask | |
| |---|---|---|---|---| |
| | 128 | 956.2 | **2,923.0** | 3.06x | 3,177.0 | |
| | 256 | 464.5 | **1,899.2** | 4.09x | 1,922.7 | |
| | 512 | 213.0 | **894.4** | 4.20x | 962.9 | |
|
|
| Accuracy vs the fp32 CPU model on the same mask (valid positions): rel_l2 0.005-0.055 for right padding, half |
| padding, length-1 sequences and interior holes (the unmasked v1.8 path gives 0.007-0.051 at the same shapes). |
| Masked positions are computed like the stock model does (their keys still attend to valid keys). Every query row |
| is still computed, so padding is not skipped: throughput follows the padded length. |
| * **No change for existing models:** 28/28 regression cases (DINOv3 S/S+/B/L/H+ at 224/448, V-JEPA2, LNC=1 and LNC=2) |
| give output bitwise identical to v1.8 at the same speed (within ±0.15%). Without the new argument the kernels emit |
| the v1.8 instruction stream. Verified on the public SDK 2.33 container |
| (`public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314-sdk2.33.0-ubuntu26.04`: torch 2.14.0, torch-neuronx 2.14.3, |
| neuronx-cc 2.28.7631, NKI 0.7.0) with transformers 4.57.6. |
| |
| ## v1.8 (2026-10) |
| |
| * **New: post-LayerNorm encoders (BERT / RoBERTa / XLM-R).** `NeuronXLMRobertaModel` is a drop-in forward for |
| `XLMRobertaModel`, `RobertaModel` and `BertModel` (GELU, absolute positions, no RoPE). It runs on the same kernels via |
| two new options: |
| - `attention_block_l1_ln_post` / `mlp_res_ln_post`: the existing fused attention and MLP epilogues also emit the |
| post-LN residual `LN(x_new) * gamma + beta` in fp32, computed from the mean and rstd they already have. |
| - `norope=1` on the attention kernel skips the rotation (and its table loads) for models without RoPE. |
| |
| XLM-RoBERTa-Large, LNC=1, 8 processes (one per logical core) on one trn2.3xlarge, fully padded batches (inferences/s): |
|
|
| | seq | stock `torch.compile` | v1.8 `NeuronXLMRobertaModel` (`attention_mask=None`) | speedup | rel_l2 vs fp32 (stock / kernels) | |
| |---|---|---|---|---| |
| | 128 | 2,092.7 | **3,144.8** | 1.50x | 0.031 / 0.0066 | |
| | 256 | 944.4 (BS 16) | **1,869.5** | 1.98x | 0.051 / 0.0147 | |
| | 512 | 397.1 (BS 8) | **935.1** | 2.35x | 0.018 / 0.0143 | |
| |
| Batch per core: 16 / 8 / 4. Passing an all-ones `attention_mask` costs 3-4% (a host-side check that no position is |
| masked: 3,022.7 / 1,789.7 / 902.8). A padded mask falls back to the stock forward (checked: rel_l2 vs fp32 1.2e-6). |
| |
| Measured on the SDK 2.33 DLAMI (`torch_neuronx_2_14_0` venv: torch 2.14.0, torch-neuronx 2.14.3.0, neuronx-cc 2.28.7631, |
| NKI 0.7.0) with transformers 4.57.6. The kernel path needs LNC=1, a sequence length that is a multiple of 128 (up to 2048) |
| and no padding mask (all positions attended). Anything else falls back to the stock forward. |
| * **Fix: models with an odd number of head groups failed to compile at LNC=2 on neuronx-cc 2.28** (`NCC_EVRF066`: a |
| one-program attention grid on an LNC=2 build). That covers ViT-S/16, ViT-S+/16 and ViT-H+/16. At LNC=2 the host packing now pads |
| an odd group count with one all-zero head group, so every attention launch uses a two-program grid. The padded heads have |
| zero V and out-projection weights and add exactly 0. Models with an even group count (ViT-B, ViT-L, V-JEPA2) are unchanged. |
| Single core, LNC=2: ViT-S/16 224 BS4 1,148.8 img/s (rel_l2 0.0109), ViT-S+/16 1,097.9 (0.0148), ViT-H+/16 224 BS1/BS4 |
| 87.8 / 130.9, 448 BS1/BS2 32.6 / 33.9 (rel_l2 0.0059-0.0080). The ViT-H+ LNC=2 path is 1.18-1.34x faster than its LNC=1 |
| fallback on the same core. |
| * **No change for existing models at existing shapes:** in the 28-case regression (DINOv3 S/S+/B/L/H+ at 224/448 and |
| V-JEPA2 at batch 1/2, LNC=1 and LNC=2), every case v1.7 runs gives output bitwise identical to v1.7, at the same speed |
| (within ±1%). That is 22 cases. The other 6 are the odd-group LNC=2 cases above, which v1.7 cannot compile on this |
| toolchain. Verified on the SDK 2.33 DLAMI. |
|
|
| ## v1.7 (2026-10) |
|
|
| * **Fix: `res_layernorm_fold` failed to compile at large batch** (`NCC_IBIR229`, state buffer allocation failed). The |
| kernel staged every row tile of a core in SBUF at once. Above about 200 KiB per partition that overflowed: |
| - ViT-S/16 at LNC=1, batch 32 at 224 px or batch 8 at 448 px; |
| - ViT-B/16 at LNC=1, batch 16 at 224 px; |
| - ViT-L/16 at LNC=1, batch 16 at 224 px; |
| - ViT-H+/16 at LNC=1, batch 16 at 224 px. |
|
|
| Each failure killed the worker process. Tiles are now processed in groups that fit, with the same per-tile |
| instructions in the same order. Every shape that compiled before keeps its schedule and is **bit-identical** to v1.6 |
| (checked at 9 shapes, LNC=1 and LNC=2). The newly compiling shapes match a float64 reference |
| (z rel_l2 1.66e-3, the same as before). |
| |
| Per instance (LNC=1, 8 cores) at the newly reachable batch sizes: |
| |
| | model | batch | img/s | best before | |
| |---|---|---|---| |
| | ViT-S/16 224 px | 32 / 64 | 8,205.4 / 8,009.5 | 9,422.6 (batch 16) | |
| | ViT-S/16 448 px | 8 / 16 | 2,176.4 / 2,192.3 | 2,486.1 (batch 4) | |
| | ViT-B/16 224 px | 16 / 32 | 4,775.4 / 4,679.0 | 4,737.8 (batch 8) | |
| | ViT-L/16 224 px | 16 / 32 | 1,892.3 / 1,866.8 | 1,905.8 (batch 8) | |
| | ViT-H+/16 224 px | 16 | 737.2 | 742.7 (batch 8) | |
| |
| None of these larger batch sizes beat the best smaller one (ViT-B 16 is +0.8%, within noise), so the recommended |
| batch sizes are unchanged. The fix means a user-chosen batch no longer crashes. |
| |
| **No cost at the recommended batch sizes:** grouping turns on only when the tiles do not fit. v1.6 vs v1.7 per instance, |
| 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%, |
| ViT-L 448 BS 2 -0.03% (LNC=1); ViT-L 224 BS 8 LNC=2 +0.12%. |
| |
| ## v1.6 (2026-10) |
| |
| * **Two images per producer pass in `attention_block` too** (new `pack2` argument; the encoder enables it for even |
| batches up to 256 tokens per image, e.g. 224 px). This kernel serves every model at LNC=2 and single-head-group models |
| such as ViT-S/16 at LNC=1. v1.5 had packing only in the LNC=1 all-group kernel. Q/K/V projections and RoPE run once |
| over 2N columns; attention stays per image. Output bit-identical to v1.5 (same rel_l2 vs fp32 on all 20 cases). |
| Kernel device cycles, ViT-L 224 px LNC=2: batch 4 -8.4% (resident weights) / -12.0% (restaged), batch 2 -8.4%. |
| Per instance at 224 px, batch 8: |
|
|
| | model | config | v1.5 img/s | v1.6 img/s | change | |
| |---|---|---|---|---| |
| | ViT-S/16 | LNC=1 DP=8 | 8,592.8 | **9,365.6** | +9.0% | |
| | ViT-S/16 | LNC=2 DP=4 | 4,667.6 | 5,061.6 | +8.4% | |
| | ViT-B/16 | LNC=2 DP=4 | 3,578.3 | 3,856.8 | +7.8% | |
| | ViT-L/16 | LNC=2 DP=4 | 1,514.8 | 1,589.7 | +4.9% | |
| | ViT-H+/16 | LNC=2 DP=4 | 507.3 | 554.0 | +9.2% | |
|
|
| Single core, batch 4: +5.9% to +10.8%. Batch 1 and 448 px unchanged (packing off). `VIT_ENC_PACK2=0` disables it. |
|
|
| ## v1.5 (2026-10) |
|
|
| * **Two images per producer pass** (LNC=1 all-group attention, `attention_block_l1_ln`, new `pack2` argument, on by |
| default from the encoder when the batch is even and two images fit one 512-token tile, i.e. up to 256 tokens such as |
| 224 px). The two images' Q/K/V projections and RoPE run as one pass over 2N columns; attention stays per image. |
| Kernel device cycles, ViT-L 224 px, batch 4: -6.2% (resident weights) / -14.8% (restaged), batch 2: -8.5%. |
| End to end (real HF weights, LNC=1): per instance (8 cores, batch 8) ViT-L/16 **1,838 -> 1,906 img/s (+3.7%)**, |
| ViT-H+/16 **672 -> 743 (+10.6%)**, ViT-B/16 **4,390 -> 4,738 (+7.9%)**. Single core batch 4: ViT-L +4.6%, ViT-H+ |
| +12.6%, ViT-B +4.1%. Output bit-identical to v1.4 (same rel_l2 vs fp32 on all 12 cases). `VIT_ENC_PACK2=0` disables it. |
| A fused-attention-block technique (two images per producer pass). |
| * **Fix: inputs of 128 tokens or fewer** failed to compile in every attention kernel (`index 1 out of bounds`). They |
| now use a non-pipelined single-Q-group path (rel_l2 vs a float64 reference 1.9-2.5e-3). Longer inputs are |
| unchanged (bit-identical). |
| * **EXPERIMENTAL bf16 RoPE** (`VIT_ENC_ROPE_BF16=1`, off by default): RoPE tables and products in bf16. +0.6% to +3.3% |
| end to end on DINOv3; rel_l2 vs fp32 grows by up to 1.16x (V-JEPA2: 0.039 -> 0.041). |
| * **EXPERIMENTAL fused attention write-back** (`VIT_ENC_WB_FUSED=1`, off by default): one Vector multiply normalises and |
| casts O. -3.4% at 256 tokens / LNC=2, but +6.6% at 2048 tokens (Vector is the attention bottleneck), so off. |
|
|
| ## v1.4 (2026-10) |
|
|
| * **QK^T PSUM banks:** for 1024 or fewer tokens per image (two key tiles or fewer), the QK^T matmul now rotates over |
| five PSUM banks across Q groups instead of reusing four, so it no longer waits for the Vector row-max evict of the |
| tile four back. Attention kernel device cycles vs v1.3: -2.1% (N=256 and N=1024, LNC=2), -1.1% (N=1024, LNC=1), |
| -4.7% (all-group kernel, N=256, LNC=1). At 2048 tokens the old layout is kept (the rotation collided with other banks |
| and measured +7%). `VIT_ENC_MM1_WIDE=0` restores the v1.3 layout. |
| End to end (single core, real HF weights) vs v1.3: DINOv3 ViT-L/16 -0.5% to +3.7%, ViT-H+/16 0.0% to +2.6%. |
| V-JEPA2 (2048 tokens): 23.80 ms/clip vs 23.81, output identical to v1.3. All 24 DINOv3 cases pass (rel_l2 vs fp32 |
| 0.56-1.75e-2). |
| * **EXPERIMENTAL FP8 MLP** (`VIT_ENC_MLP_FP8=1`, GELU-MLP models only, off by default): fc1 and fc2 run as e4m3 |
| `double_row` matmuls, with power-of-two weight scaling and a per-layer GELU-output scale (`calibrate_fp8`). Measured |
| on DINOv3 ViT-L/16, single core: +7.3% (224 px, LNC=2) to +10.5% (448 px, LNC=1) end to end. On imagenette |
| (real weights), k-NN and linear-probe predictions were identical to fp32. But the CLS feature error rises to |
| rel_l2 5.8e-2 (bf16 path: 0.84e-2). **Calibrate before use**: trn2's e4m3 saturates at 240, and DINOv3's first |
| block produces GELU outputs near 2,000. Without calibration the default scale (1/32) is used, which is safe for |
| DINOv3 but coarse. |
| * Outputs are deterministic per compiled graph. Across separate compilations they can differ at bf16-ulp level |
| (rel ~1e-5), as in v1.3. |
| |
| ## v1.3 (2026-10) |
| |
| Attention-path schedule improvements, plus a SwiGLU fallback fix. Output of every attention kernel is |
| **bitwise identical** to v1.2. Verified on trn2.3xlarge at LNC=1 and LNC=2, PyTorch Native Beta 6. |
| |
| Device-cycle A/B, v1.2 vs v1.3 (median of 3 captures, same instance, container and inputs): |
| |
| | Kernel / shape | LNC | change | |
| |---|---|---| |
| | `attention_block`, N=1024, B=2 | 2 | -1.9% | |
| | `attention_block`, N=2048, B=1 (V-JEPA2 clip) | 2 | **-5.2%** | |
| | `attention_block`, N=256, B=8 | 2 | -2.0% | |
| | `attention_block`, N=1024, B=2 | 1 | -3.3% | |
| | `attention_block_l1`, N=256, B=4 | 1 | **-11.1%** | |
| | `attention_block_l1_ln`, N=256, B=4 | 1 | **-10.2%** | |
| | `attention_block_l1_ln`, N=1024, B=2 | 1 | -4.5% | |
|
|
| End to end (real Hugging Face weights, single core, vs v1.0 on the same harness): DINOv3 ViT-L/16 +1.0% to |
| +12.7%, ViT-H+/16 +0.6% to +17.3%. Largest gains are at 224 px and batch 1, where attention's fixed costs |
| dominate. V-JEPA2 (16 frames, LNC=2): 24.78 -> 24.48 ms/clip vs v1.2, output identical to v1.2. |
| All 16 DINOv3 cases pass rel_l2 <= 5e-2 vs fp32 (0.56-0.98e-2). |
| |
| Changes: |
| |
| * **Q/K bias in one DMA.** `pack_weights` now returns the Q/K bias packed as `[2, 128, Fp/128]` fp32 |
| (`pack_qk_bias`, exported). Each attention program loads its groups with one DMA per Q/K on the Scalar |
| HW-DGE queue, replacing one `[128, 1]` DMA plus one cast per 128-row group. The values are rounded to bf16 |
| first, which keeps the output bit-identical. The kernels still accept the legacy `[2, Fp, 1]` layout. |
| * **Dead work removed in the attention core** (`attention_inner.py`, single-section path, N <= 2048): |
| the `mm2_sb` and `exp_partial_sum` zero-fills (never observed), and the `tensor_reduce` over the chained |
| exp-sum columns (the sum is already in one column; it is now read directly). |
| * **Chained row-max reduce**: the QK^T row-max accumulates across the 512-key tiles of a large tile and is |
| read out once, instead of once per tile followed by a `tensor_reduce`. Pattern from the upstream nkilib |
| `attention_cte` row-max chaining, re-implemented for this fork. It adds -3.5% at N=2048 and is neutral at |
| N <= 1024 (fewer key tiles to chain). |
| * **Fix (v1.2 regression):** the SwiGLU torch fallback (taken below 256 rows per core, e.g. ViT-H+/16 at |
| LNC=2, batch 1, 224 px) used the packed biases and failed to compile. It now uses unpacked copies. The |
| kernel path was unaffected. |
|
|
| ## v1.2 (2026-10) |
|
|
| Applied the v1.1 packed-bias one-DMA technique to the SwiGLU FFN kernels (`swiglu_stream`, |
| `swiglu_stream_ln`). The silu bias `b1` and the up-projection bias `bg` are now host-packed to |
| `[128, I/128]` (`pack_b1`) and each loaded in a single DMA on the Scalar engine's HW-DGE queue, |
| replacing `I/128` `[128, 1]` column DMAs apiece. Output is **bitwise identical** to v1.1 at LNC=1 and |
| LNC=2 (`torch.equal`, verified on trn2.3xlarge). Device cycles (same stack, A/B): `swiglu_stream` |
| -2.5% (reproducible); `swiglu_stream_ln` within measurement noise (the LN epilogue dominates and the |
| bias DMAs already overlap compute). The wrappers still accept an unpacked `[I]` bias (packed on the |
| fly). Affects only the gated-MLP (SwiGLU) models; the GELU-MLP path is unchanged from v1.1. |
|
|
| ## v1.1 (2026-10) |
|
|
| Reworked DMA/engine schedule for `mlp` (K4), `mlp_res_ln` (K4 + next LN) and `res_layernorm_fold` |
| (K1/K3). Output is **bitwise identical** to v1 at both LNC=1 and LNC=2. Based on an external |
| contribution by [jlongen-aws](https://huggingface.co/jlongen-aws) (PR #2 on |
| `jburtoft/vjepa2-neuron-kernels`), re-implemented here with the LNC-parametric path and the `min()` |
| remainder handling preserved. |
|
|
| Verified on trn2.3xlarge (A/B on one consistent stack: old vs new compiled and run in the same |
| container + inputs, device cycles from `neuron-explorer show-session`). LNC=2, V-JEPA2 ViT-L shape |
| (R=2048, H=1024, MLP hidden 4096): |
|
|
| * K4 `mlp`: -7.5% device cycles (reproducible, -7.8% +/- ~1% over 3 captures). |
| * K4+LN `mlp_res_ln`: -6.8%. |
| * K1 `res_layernorm_fold` (1 branch): -16.2%; K3 (2 branches): -18.3%. |
| * Full 24-layer encoder (synthetic ViT-L weights, includes the unchanged attention): -3.7% device |
| cycles, output `torch.equal` to v1. |
| * Real-weight V-JEPA2 (`facebook/vjepa2-vitl-fpc64-256`, 16 frames x 256 px, LNC=2): output bitwise |
| identical to v1 (`new_vs_v1_rel` 0.0), rel_l2 0.0389 vs fp32 CPU (unchanged from v1), -4.0% ms/clip |
| (speedup vs stock fp32 3.19x -> 3.33x). |
| * Per-kernel output is `torch.equal` to v1 at LNC=1 and LNC=2. |
| |
| Changes: |
| |
| * `mlp` / `mlp_res_ln`: fc1 bias is host-packed to `[128, H/128]` (`pack_b1`, now exported) and |
| loaded in one DMA on the Scalar engine's HW-DGE queue instead of `H/128` column DMAs; row block 0's |
| `x` is staged and transposed on the (idle) Tensor engine before the weights, later row blocks stream |
| under row block 0's matmuls. `mlp` still accepts an unpacked `[H]` bias (packed on the fly). |
| * `res_layernorm_fold`: every DMA uses `dge_mode.none` (the runtime pre-generates descriptors) and all |
| token tiles are made SBUF-resident before any compute, replacing the depth-4 ring. Per-tile statistic |
| columns (`[128, n_t]`) so one tile's stat chain does not wait on the previous tile's. |
| * `pack_block` now stores the GELU fc1 bias pre-packed; `pack_b1` is exported from `nki_kernels`. |
|
|
| The compute instructions and their order are unchanged, so results are bit-identical to v1. |
|
|
| ## v1 (2026-10) |
|
|
| First release. |
|
|
| * Drop-in layers for Hugging Face Transformers: `NeuronDINOv3ViTModel` (DINOv3 ViT-S/16, S+/16, B/16, |
| L/16, H+/16) and `NeuronVJEPA2Encoder` (V-JEPA2 ViT-L, any batch size, up to 2048 tokens per clip). |
| * Model-agnostic encoder loop: `BlockWeights`, `pack_block`, `rope_tables`, `EncoderRunner`. |
| * Kernels: `attention_block` (per head group; batched; padded keys skipped; head_dim 64 or 128; |
| any head count), `attention_block_l1` and `attention_block_l1_ln` (LNC=1, all head groups, optional |
| residual + LayerNorm epilogue), `res_layernorm_fold`, `mlp` and `mlp_res_ln` (GELU MLP, optional |
| residual + next LayerNorm epilogue), `swiglu_stream` and `swiglu_stream_ln` (weight-streaming |
| SwiGLU FFN). |
| * RoPE pairing selectable at packing time: `rope="half"` (rotate-half) or `rope="interleaved"`. |
| * Verified on trn2.3xlarge, PyTorch Native Beta 6 (torch 2.13.0, torch-neuronx 2.13.3, NKI 0.7.0b1), |
| Transformers 4.57.6, LNC=1 and LNC=2. Every measured case passes rel_l2 <= 5e-2 vs the fp32 CPU model. |
| |
| Known limitations: encoder (bidirectional) attention with RoPE only; LayerNorm (not RMSNorm); |
| <= 2048 tokens per image; inference only; hidden size 4096 (ViT-7B) not covered by the drop-in layer. |
| |
