Changelog

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.