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.
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.
9 days ago
v1: NKI kernels for RoPE ViT encoders (DINOv3, V-JEPA2) on Trainium2
