attention_block gains the pack2 producer packing v1.5 added to attention_block_l1_ln: for even batches up to
256 tokens per image the Q/K/V projections and RoPE run once over 2N columns, attention stays per image.
Output bit-identical to v1.5. Per instance at 224 px, batch 8: ViT-S/16 LNC=1 +9.0%; LNC=2 ViT-S +8.4%,
ViT-B +7.8%, ViT-L +4.9%, ViT-H+ +9.2%.
examples: use aws-neuron/vit-encoder-neuron-kernels repo id (docs only)
attention_block_l1_ln packs two images into one Q/K/V projection + RoPE pass (pack2, default on for even batches
at 224 px): per instance ViT-L/16 +3.7%, ViT-H+/16 +10.6%, ViT-B/16 +7.9%, output bit-identical to v1.4.
Inputs of 128 tokens or fewer now use a single-Q-group path (previously failed to compile).
Experimental opt-ins: bf16 RoPE (VIT_ENC_ROPE_BF16=1), fused attention write-back (VIT_ENC_WB_FUSED=1).
Packing technique from the BERT-reranker fused attention block.
QK^T rotates over five PSUM banks across Q groups when a group has <=2 key tiles, so it no longer waits on
the Vector row-max evict: attention device cycles -1.1% to -4.7% vs v1.3; 2048-token layout unchanged.
FP8 e4m3 double_row GELU MLP behind VIT_ENC_MLP_FP8=1 (off by default) with weight and per-layer activation
scaling. Verified on trn2.3xlarge: all 24 DINOv3 cases and V-JEPA2 pass.
Q/K bias host-packed [2,128,Fp/128] and loaded in one Scalar HW-DGE DMA per Q/K (pack_qk_bias);
dead zero-fills and the chained exp-sum tensor_reduce removed in the single-section attention core;
row-max chained across key tiles. Attention kernels -1.7% to -11.1% device cycles on trn2.3xlarge,
output torch.equal to v1.2 at LNC=1 and LNC=2. Fixes a v1.2 regression where the SwiGLU torch
fallback (rows/core < 256) used packed biases and failed to compile.
5 days ago
v1.2: packed-bias one-DMA for SwiGLU (swiglu_stream/_ln)
Silu bias b1 and up bias bg host-packed to [128, I/128] and loaded in one DMA each on the
Scalar HW-DGE queue (same technique as v1.1 K4). Bitwise-identical to v1.1 at LNC=1 and LNC=2;
swiglu_stream -2.5% device cycles, swiglu_stream_ln within noise. GELU path unchanged.
Bitwise-identical to v1 at LNC=1 and LNC=2. Packed fc1 bias loaded in one DMA on the
Scalar HW-DGE queue; row-block-0 x transposed on the Tensor engine before the weights;
res_layernorm uses dge_mode.none with all token tiles SBUF-resident. Based on an external
contribution by jlongen-aws (PR #2 on jburtoft/vjepa2-neuron-kernels), re-implemented with
the LNC-parametric path preserved. Verified on trn2.3xlarge: real-weight V-JEPA2 output
bitwise identical to v1, -4.0% ms/clip.
8 days ago
v1: NKI kernels for RoPE ViT encoders (DINOv3, V-JEPA2) on Trainium2
