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):

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