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