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