V-JEPA2 ViT-L Encoder NKI Kernels for AWS Trainium2

Superseded. For new work, use aws-neuron/vit-encoder-neuron-kernels. It contains a maintained NeuronVJEPA2Encoder built from these kernels. Compared with this repo it adds batching, LNC=1 (8 logical cores), fused residual + LayerNorm epilogues and later attention optimizations, and it covers DINOv3 and XLM-RoBERTa as well. Its README has the current V-JEPA2 performance numbers. On the same SDK 2.33 instance it is faster than this repo: 23.8 vs 26.7 ms/clip single-core, and 159.6 clips/s per instance (3.59x over stock). This repo is kept unchanged, so existing version=1 pins keep working. It receives no new features.

Four NKI kernels replace the per-layer torch ops of the V-JEPA2 ViT-L video encoder on trn2, at 16 frames, batch 1 and LNC=2 only. The kernels are unchanged since v1 (2026-09-25).

Quick start

import types
import torch
from kernels import get_kernel
from transformers import AutoModel

k = get_kernel("aws-neuron/vjepa2-neuron-kernels", version=1, trust_remote_code=True)

model = AutoModel.from_pretrained("facebook/vjepa2-vitl-fpc64-256", dtype=torch.float32).eval().to("neuron")
model.encoder.forward = types.MethodType(k.NeuronVJEPA2Encoder.forward, model.encoder)

enc = torch.compile(model.encoder, backend="neuron")    # required; the numbers below are of this graph
with torch.no_grad():
    embeddings = enc(your_clip.to("neuron")).last_hidden_state   # [1, 16, 3, 256, 256] fp32 -> [1, 2048, 1024]

Set NKI_ENABLE_TRACE_CACHE=0 in the environment. The runtime enables the cross-process trace cache by default, and a stale NEFF silently captures a different instruction stream.

Use the forward graft, not transformers.KernelConfig. transformers does not register VJEPA2Encoder as a hub-kernel layer, so a KernelConfig({"VJEPA2Encoder": ...}) mapping fails with Layer VJEPA2Encoder is not registered in the model. This was checked on 5.13.1–5.22.2; earlier revisions of this README showed KernelConfig.

SDK 2.33 install. The public SDK 2.33 DLAMI's torch_neuronx_2_14_0 venv does not include transformers or kernels. One way to add them without touching its torch: pip install --target ~/xpkgs --no-deps "transformers==5.22.2" "kernels>=0.17.2,<0.18" tokenizers safetensors regex typer-slim accelerate, then export PYTHONPATH=~/xpkgs. transformers 5.22.2 rejects kernels 0.18.x.

Performance and correctness (trn2.3xlarge, LNC=2, batch 1, 16 frames x 256 px)

• How times were measured: device time is neuron-explorer's total_time for the steady-state encoder NEFF, and MFU is its mfu_estimated_percent.
• SDK 2.33 wall-clock: with a device sync every iteration, 84.68 vs 27.95 ms (3.03x). Eager mode is 33.6 ms.
• Why the speedup shrank on 2.33: the stock compiled encoder got 21% faster, and these kernels got 11% slower. They were not re-tuned for 2.33.
• Earlier README: it listed torch-neuronx 2.11.3 for the pre-release stack. The Beta 5 build is 2.12.3.
• Bitwise identity: on Beta 5, the output was bitwise identical (torch.equal) to the verified reference implementation, under torch.compile. This was not re-checked on 2.33, because that reference is not public.
• Run-to-run nondeterminism: the vendor attention inner varies run-to-run at a few bf16 positions on some layers. fp64 agreement is unchanged. Bitwise comparisons must use a fixed captured sample.
• Seeded test inputs: build them on CPU and then copy them to the device. torch.rand(..., device="neuron") uses a different random stream from the CPU.

Known limitations

• 16 frames / 2048 tokens / batch 1 / 256x256 / LNC=2 only. 64-frame clips (8192 tokens) need the attention kernel's flash-sectioning path, which this integration does not support. Its SBUF buffers need 347 KiB/lane, against a 192 KiB budget. Off-geometry input falls back to the stock encoder.
• Inference only. The layer refuses to run in training mode.
• The first device forward performs the one-time weight fold eagerly, outside torch.compile. Expect two compiles before steady state.

Repository layout

build/torch-neuron/
├── __init__.py            # public API re-exports
├── metadata.json          # kernels-library metadata (backend: neuron)
├── layers.py              # NeuronVJEPA2Encoder (forward graft target)
└── nki_kernels/
    ├── res_layernorm.py   # K1/K3
    ├── attention_block.py # K2 outer (QKV + RoPE + out-proj) and weight packing
    ├── attention_inner.py # K2 inner: MODIFIED copy of the Neuron SDK's nkilib attention_cte
    ├── mlp.py             # K4
    └── rope_tables.py     # host-side 3D RoPE tables

Attribution and license

Apache-2.0 (see LICENSE, NOTICE). nki_kernels/attention_inner.py is a modified copy of the AWS Neuron SDK's nkilib core/attention/attention_cte.py (Copyright Amazon.com, Inc. or its affiliates, Apache-2.0). The file header lists the modifications: four engine-placement pins, worth about 1.19x end-to-end, and removal of code paths this configuration cannot reach. The other three kernels, the layer glue and the packing are original work.

Contact

Kernels developed by Antonio Mena (@menaman123).