File size: 17,401 Bytes

# 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.