| --- |
| license: apache-2.0 |
| library_name: kernels |
| tags: |
| - kernel |
| - neuron |
| - nki |
| - trainium |
| - trn2 |
| - vjepa2 |
| - video |
| - vision |
| - attention |
| - encoder |
| --- |
| |
| # V-JEPA2 ViT-L Encoder NKI Kernels for AWS Trainium2 |
|
|
| **Version 1 (2026-09-25).** Four NKI kernels replacing the per-layer torch ops of the |
| [V-JEPA2](https://huggingface.co/facebook/vjepa2-vitl-fpc64-256) ViT-L video encoder on trn2: |
| **4.35x over the stock torch path (104.2 → 23.9 ms/clip), MFU 10.1% → 44.1%** (neuron-explorer |
| `mfu_estimated_percent`, device window). Output is **bitwise identical** to the reference |
| implementation this was verified against, and rel_l2 4.2e-2 vs the fp32 torch encoder (the |
| model's own bf16-matmul sensitivity). |
| |
| > **See also: [`aws-neuron/vit-encoder-neuron-kernels`](https://huggingface.co/kernels/aws-neuron/vit-encoder-neuron-kernels)**, |
| > a generalized version of these kernels. It adds batching (any batch size), LNC=1 (8-core) support, |
| > fused residual + LayerNorm epilogues, a SwiGLU FFN kernel, both RoPE pairings, and a DINOv3 layer |
| > alongside a V-JEPA2 one. This repo is unchanged and keeps its verified 16-frame / batch-1 path. |
| |
| ## What ships |
| |
| | Component | Purpose | |
| |---|---| |
| | `NeuronVJEPA2Encoder` (layer) | KernelConfig-mappable drop-in for `VJEPA2Encoder`: the 4-kernel layer loop, lazy one-time weight fold, stock fallback off-device/off-geometry | |
| | `res_layernorm_fold` | K1/K3: residual add + affine-free LayerNorm (fp32 stream in, bf16 normalized out); the LN affines are folded into the neighboring matmuls | |
| | `attention_block` | K2: fused QKV projection (LN-gain folded, lane-permuted) + 3D RoPE + flash-attention inner + fused out-projection → [2, N, H] per-core partials (LNC=2) | |
| | `mlp` | K4: fc1 + GELU + fc2, bf16 matmuls, fp32 boundary, prologue-prestaged row blocks | |
| | `pack_*`, `build_rope_tables` | host-side weight folding (used internally by the layer; exported for direct kernel callers) | |
|
|
| ## Quick start |
|
|
| ```python |
| import torch |
| from transformers import AutoModel, KernelConfig |
| |
| kernel_config = KernelConfig({ |
| "VJEPA2Encoder": ( |
| "aws-neuron/vjepa2-neuron-kernels:NeuronVJEPA2Encoder", |
| {"version": 1, "trust_remote_code": True}, |
| ), |
| }) |
| model = AutoModel.from_pretrained( |
| "facebook/vjepa2-vitl-fpc64-256", |
| dtype=torch.float32, |
| kernel_config=kernel_config, |
| device_map="neuron", |
| ).eval() |
| |
| enc = torch.compile(model.encoder, backend="neuron") # the published numbers are of this graph |
| video = your_clip # [1, 16, 3, 256, 256] fp32 |
| with torch.no_grad(): |
| embeddings = enc(video).last_hidden_state # [1, 2048, 1024] |
| ``` |
|
|
| > **`torch.compile(backend="neuron")` is part of the configuration.** Eager dispatch adds |
| > ~5.5 ms/forward. The layer sets `can_torch_compile`, so the kernels library will not silently |
| > swap it out under compile. |
|
|
| > **Resolving via `get_kernel` / `kernels`:** on a vLLM (0.24) venv where `torch` is a |
| > CUDA-tagged build, `import libtorch_neuronx_lite` (or `torch_xla`) **before** `get_kernel(...)` |
| > so `torch.neuron` is registered and the `torch-neuron` variant is selected. On a PyTorch Native |
| > container (`torch.neuron` already registered) this import is unnecessary. |
| |
| ## Performance (trn2.3xlarge, verified 2026-09-25) |
| |
| | | stock torch encoder | these kernels | |
| |---|---|---| |
| | throughput, compiled graph | 104.18 ms/clip | **23.94 ms/clip (4.35x)** | |
| | device MFU (`mfu_estimated_percent`) | 10.1% | **44.08%** | |
| | end-to-end serving, preprocess on device + double-buffered uint8 H2D | — | 24.59 ms/clip = **42.8%** | |
| | end-to-end single stream, nothing overlapped | — | 29.86 ms = 35.7% | |
| |
| MFU is neuron-explorer's own field: model FLOPs (1.6557e12, equal to the closed-form |
| QK+PV+projections+MLP+patch count to 0.000%) over 157.3 TFLOP/s (one logical NeuronCore = |
| 2 physical cores, bf16) times the window. One logical core of four: a single stream uses ~11% |
| of the chip; four concurrent streams reach ~44% chip-wide. |
| |
| Where the remaining 2.3x to physics goes (measured, not estimated): the softmax exp runs on the |
| Scalar engine at 1 element/lane/cycle (+2.6 ms over physics — unreachable without different |
| hardware), instruction issue overhead (~3.8 ms), and the queueing of three engines co-bound |
| within 0.5% (~3.7 ms). The exact-math kernel-level candidate classes are enumerated to |
| exhaustion; details in the source repository's cycle records. |
| |
| ## Correctness |
| |
| * Full-forward output **bitwise identical** (`torch.equal`) to the verified reference |
| implementation, including through this repo's KernelConfig-style forward graft under |
| `torch.compile`. |
| * rel_l2 vs the fp32 torch encoder: 4.2028e-2 — identical to the stock-vs-fp32 difference |
| (both paths run bf16 matmuls; neither is "more approximate" than the other). |
| * K1/K3's GpSimd `rsqrt` is measurably *closer* to fp64 than the torch path's (6.9e-8 vs |
| 3.9e-5 relative). |
| * The vendor attention inner is run-to-run nondeterministic on some layers at a few bf16 |
| positions (fp64 agreement unchanged) — an upstream kernel property, reported to the Neuron |
| team. Bitwise comparisons must use a fixed captured sample. |
|
|
| ## Requirements |
|
|
| * **Hardware**: trn2 (verified on trn2.3xlarge, LNC=2). |
| * **SDK**: neuronx-cc 2.27.2878.0; torch-neuronx 2.11.3 (PyTorch Native); nki 0.6.0; |
| nkilib (bundled with the SDK — `ModularAllocator` is imported from it). |
| * **Transformers** `>= 5.14` (tuple-form `KernelConfig`); `kernels >= 0.17`. |
| * `NKI_ENABLE_TRACE_CACHE=0` in the environment (the runtime defaults the cross-process trace |
| cache on; a stale NEFF silently captures a different instruction stream). |
|
|
| ## Repository layout |
|
|
| ``` |
| build/torch-neuron/ |
| ├── __init__.py # public API re-exports |
| ├── metadata.json # kernels-library metadata (backend: neuron) |
| ├── layers.py # NeuronVJEPA2Encoder (KernelConfig 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 |
| ``` |
|
|
| ## Known limitations |
|
|
| * **16 frames / 2048 tokens / batch 1 / 256x256 only.** The model's native 64-frame clip |
| (8192 tokens) is NOT supported: three of the vendor attention kernel's SBUF buffers scale |
| with keys-per-section (347 KiB/lane needed vs the 192 KiB budget, device-confirmed), so 8192 |
| keys require the flash-sectioning path, whose cross-section combine this integration does not |
| yet reproduce correctly. Off-geometry inputs fall back to the stock implementation. |
| * **Inference only** — no autograd rule; the layer refuses in training mode. |
| * The first device forward performs the one-time weight fold eagerly (~seconds); it is |
| deliberately kept outside `torch.compile` (compiled folding rounds differently and breaks |
| bitwise parity with the verified reference). |
|
|
|
|
| ## 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 ~1.19x end-to-end, and removal of code paths unreachable for this configuration). |
| The other three kernels, the layer glue and the packing are original work. |
| |
| ## Contact |
| |
| Kernels developed by Antonio Mena ([@menaman123](https://huggingface.co/menaman123)). |
