| --- |
| license: apache-2.0 |
| library_name: kernels |
| tags: |
| - kernel |
| - neuron |
| - nki |
| - trainium |
| - trn2 |
| - vision |
| - vit |
| - dinov3 |
| - vjepa2 |
| - rope |
| - swiglu |
| --- |
| |
| # vit-encoder-neuron-kernels |
|
|
| NKI kernels that run the whole transformer stack of **RoPE vision-transformer encoders** on AWS |
| Trainium2 (trn2). Each block is 2-3 coarse kernels instead of dozens of small ops: attention plus |
| the residual add and LayerNorm, and the MLP / SwiGLU FFN plus the next block's LayerNorm. The |
| residual stream stays in fp32 between kernels. |
|
|
| Drop-in layers are included for two Hugging Face Transformers models: |
|
|
| | Layer | Replaces | Models | |
| |---|---|---| |
| | `NeuronDINOv3ViTModel` | `DINOv3ViTModel.forward` | DINOv3 ViT-S/16, S+/16, B/16, L/16, H+/16 | |
| | `NeuronVJEPA2Encoder` | `VJEPA2Encoder.forward` | V-JEPA2 ViT-L (video), up to 2048 tokens per clip | |
|
|
| The building blocks (`BlockWeights` / `pack_block` / `EncoderRunner`) are model-agnostic. Any |
| pre-norm ViT with the shape constraints below can use them, but only the two models above |
| have been validated. |
|
|
| ## Performance (trn2.3xlarge, bf16, `torch.compile(backend="neuron")`) |
|
|
| Measured 2026-10 on PyTorch Native Beta 6 (torch 2.13.0, torch-neuronx 2.13.3, NKI 0.7.0b1). |
| Real Hugging Face weights. "Stock" is the same model compiled with |
| `torch.compile(backend="neuron")` without these kernels: DINOv3 cast to bf16, V-JEPA2 as its fp32 |
| checkpoint (stock V-JEPA2 is the encoder as shipped, and its accuracy is in the table below). |
|
|
| **Per instance** (one process per logical core, all cores busy): |
|
|
| | Model | Resolution | Stock (img/s) | **These kernels (img/s)** | Speedup | |
| |---|---|---|---|---| |
| | DINOv3 ViT-B/16 | 224 | 2,901.5 | **4,105.2** | 1.41x | |
| | DINOv3 ViT-L/16 | 224 | 1,061.2 | **1,767.5** | 1.67x | |
| | DINOv3 ViT-L/16 | 448 | 201.1 | **491.9** | 2.45x | |
| | DINOv3 ViT-H+/16 | 224 | 400.0 | **626.6** | 1.57x | |
| | DINOv3 ViT-H+/16 | 448 | 106.8 | **200.3** | 1.88x | |
| | V-JEPA2 ViT-L, 16 frames x 256 px | - | 36.8 clips/s | **157.3 clips/s** | 4.27x | |
|
|
| Configuration: LNC=1 (8 logical cores) for every row except V-JEPA2 stock, whose best was LNC=2 |
| (4 cores). DINOv3 kernels run at batch 8 (224 px) or batch 2 (448 px). DINOv3 stock runs at batch 4 |
| (224 px) or batch 1 (448 px), its best measured. V-JEPA2 runs at batch 1 per core for both. |
|
|
| **Single core, latency-oriented** (LNC=2, batch 1): |
|
|
| | Model | Resolution | Stock | **These kernels** | Speedup | |
| |---|---|---|---|---| |
| | DINOv3 ViT-B/16 | 224 | 251.0 img/s | **510.7 img/s** | 2.03x | |
| | DINOv3 ViT-L/16 | 224 | 97.1 img/s | **220.3 img/s** (4.5 ms) | 2.27x | |
| | DINOv3 ViT-L/16 | 448 | 25.1 img/s | **83.0 img/s** | 3.31x | |
| | DINOv3 ViT-H+/16 | 224 | 21.4 img/s | **88.2 img/s** | 4.12x | |
| | DINOv3 ViT-H+/16 | 448 | 17.7 img/s | **32.4 img/s** | 1.83x | |
| | V-JEPA2 ViT-L, 16 frames x 256 px | - | 101.5 ms/clip | **25.8 ms/clip** | 3.93x | |
|
|
| The gain grows with resolution: attention cost grows with token count, and these kernels keep |
| attention and its surroundings on chip. At 224 px with large batches the gap narrows. For example, |
| stock ViT-L/16 at LNC=2 batch 4 is 252 img/s and the kernels give 321 img/s (1.27x). |
|
|
| ## Accuracy |
|
|
| Every result above passes a per-image check against the fp32 CPU model (rel_l2 <= 5e-2). |
| On DINOv3 the kernels are **closer to fp32 than stock bf16**, because the residual stream and |
| LayerNorm statistics stay in fp32: |
| |
| | Model | Stock bf16 rel_l2 vs fp32 | These kernels | |
| |---|---|---| |
| | DINOv3 ViT-B/16 | 1.4-1.8e-2 | 0.9-1.7e-2 | |
| | DINOv3 ViT-L/16 | 1.3-1.6e-2 | 0.8-1.0e-2 | |
| | DINOv3 ViT-H+/16 | 1.0-1.7e-2 | 0.6-1.0e-2 | |
| | V-JEPA2 ViT-L | 1.4e-2 (fp32 model, compiler-chosen precision) | 3.9e-2 (bf16 matmuls) | |
|
|
| For V-JEPA2 the stock path compiles the fp32 checkpoint and keeps more of it in higher precision, |
| so it is closer to fp32 than the kernels (which run every matmul in bf16). |
|
|
| ## Quick start |
|
|
| ```python |
| import types, torch |
| from kernels import get_kernel |
| from transformers import AutoModel |
| |
| k = get_kernel("aws-neuron/vit-encoder-neuron-kernels", version=1, trust_remote_code=True) |
| |
| model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m", |
| dtype=torch.float32).eval().to("neuron") |
| model.forward = types.MethodType(k.NeuronDINOv3ViTModel.forward, model) |
| model = torch.compile(model, backend="neuron") |
| |
| with torch.no_grad(): |
| out = model(pixel_values.to("neuron")) # [B, 3, H, W] fp32, any batch size |
| feats = out.last_hidden_state # [B, 1 + 4 registers + patches, 1024] |
| cls = out.pooler_output |
| ``` |
|
|
| V-JEPA2: |
|
|
| ```python |
| 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") |
| emb = enc(clip.to("neuron")).last_hidden_state # clip [B, 16, 3, 256, 256] -> [B, 2048, 1024] |
| ``` |
|
|
| Notes: |
| * Load the model in **fp32**. The layer folds LayerNorm affines and LayerScale into the weights |
| once, on the first device forward, and casts the weights to bf16 itself. |
| * **`torch.compile(backend="neuron")` is part of the configuration.** The layers set |
| `can_torch_compile = True`. |
| * **Set `NEURON_LOGICAL_NC_CONFIG` before importing**, and to the same value the process runs |
| with: `2` (default) or `1`. The kernels are built for that core configuration. For LNC=1 also pass |
| `NEURON_CC_FLAGS=--lnc=1`. |
| * `NKI_ENABLE_TRACE_CACHE=0` is recommended when you edit kernels. The cross-process trace cache |
| can otherwise reuse a stale NEFF. |
| * On a vLLM-Neuron venv (CUDA-tagged torch), `import libtorch_neuronx_lite` before `get_kernel` |
| so the `torch-neuron` variant is selected. On PyTorch Native this is not needed. |
| |
| **Environment switches** (defaults are the measured best): `VIT_ENC_PACK2=1` packs two images per attention producer |
| pass at LNC=1 when the batch is even and the images are <= 256 tokens (224 px); `VIT_ENC_ROPE_BF16=1` (experimental) |
| runs RoPE in bf16; `VIT_ENC_WB_FUSED=1` (experimental) fuses the attention normalisation; `VIT_ENC_MM1_WIDE=0` |
| restores the v1.3 QK^T PSUM layout. |
| |
| **Experimental FP8 MLP** (v1.4, GELU-MLP models only): set `VIT_ENC_MLP_FP8=1` before importing to run the MLP |
| matmuls in FP8 e4m3 (`double_row`). That is +7-10% end to end on DINOv3 ViT-L/16, at ~7x the feature error of the |
| default bf16 path (CLS rel_l2 5.8e-2 vs 0.84e-2; imagenette k-NN predictions unchanged). Calibrate the per-layer |
| activation scale with `calibrate_fp8(runner, gelu_amax)`, because trn2's e4m3 saturates at 240. Off by default. |
| |
| Which LNC to use: **LNC=1** (8 logical cores, one process per core) gives the most throughput per |
| instance. **LNC=2** (4 larger cores) gives the lowest latency per request. |
| |
| ## Kernels |
| |
| | Kernel | What it does | |
| |---|---| |
| | `attention_block` | LN'd hidden -> Q/K/V projection + RoPE + attention + output projection, one 8-head (D=64) or 4-head (D=128) group per program; any number of images per launch; padded keys skipped | |
| | `attention_block_l1` | LNC=1: one program over **all** head groups, one output (no per-group partials) | |
| | `attention_block_l1_ln` | `attention_block_l1` plus residual add and LayerNorm in the epilogue | |
| | `res_layernorm_fold` | residual add (sum of up to 4 branches) + affine-free LayerNorm | |
| | `mlp` | GELU MLP (fc1 -> GELU -> fc2) | |
| | `mlp_res_ln` | `mlp` plus residual add and the next block's LayerNorm in the epilogue | |
| | `swiglu_stream` | SwiGLU FFN that streams the weights in chunks of the intermediate dimension (so large FFNs fit on chip) | |
| | `swiglu_stream_ln` | `swiglu_stream` plus residual add and the next block's LayerNorm | |
|
|
| Host-side packing (`pack_weights`, `pack_tables`, `lane_permutation`, `pack_block`) supports both |
| RoPE pairings via `rope=`: |
| * `"half"`: pairs lane `j` with `j + D/2` (rotate-half; DINOv3, EVA-02, HF Llama style). |
| * `"interleaved"`: pairs lane `2j` with `2j + 1` (V-JEPA2 style). |
|
|
| ## Supported shapes and limits |
|
|
| * **Bidirectional (encoder) self-attention with RoPE.** No causal masks, KV cache, cross-attention, |
| or learned / absolute position embeddings. A model without RoPE could pass identity tables (cos=1, |
| sin=0), but that is untested. |
| * **Pre-norm blocks with standard LayerNorm** (not RMSNorm), optional LayerScale. |
| * hidden size a multiple of 128; head_dim 64 or 128; any head count (padded to a multiple of 8 |
| at D=64 or 4 at D=128). |
| * **<= 2048 tokens per image** after padding to a multiple of 128 x LNC (e.g. 720 x 720 px at |
| patch 16). Longer inputs fall back to the stock forward. |
| * FFN: GELU MLP (`nl.gelu`, the erf form) or SwiGLU (SiLU gate). |
| * Inference only (no autograd). The layers refuse in training mode. |
| * ViT-7B-sized models (hidden 4096) are not covered by the drop-in layer. They need tensor parallelism. |
| The streaming SwiGLU kernel is slower than the compiler's FFN at that width. |
| * Validated end to end on DINOv3 ViT-B/L/H+ and V-JEPA2 ViT-L only. |
| |
| ## Requirements |
| |
| * trn2 (verified on trn2.3xlarge, LNC=1 and LNC=2). |
| * PyTorch Native (TorchNeuron) with NKI 0.7 and `nkilib`. Verified on Beta 6 (torch 2.13.0, |
| torch-neuronx 2.13.3, NKI 0.7.0b1, neuronx-cc 2.0.404056). |
| * Transformers with DINOv3 / V-JEPA2 (verified on 4.57.6); `kernels` for `get_kernel`. |
|
|
| ## Repository layout |
|
|
| ``` |
| build/torch-neuron/ |
| ├── __init__.py # public API |
| ├── metadata.json |
| ├── layers.py # NeuronDINOv3ViTModel, NeuronVJEPA2Encoder |
| ├── encoder.py # BlockWeights, pack_block, rope_tables, EncoderRunner |
| └── nki_kernels/ |
| ├── attention_block.py # attention (per head group) + host-side packing |
| ├── attention_block_l1.py # attention, all head groups (LNC=1) |
| ├── attention_block_l1_ln.py # ... + residual + LayerNorm epilogue |
| ├── attention_inner.py # modified copy of nkilib attention_cte (vendor attention core) |
| ├── res_layernorm.py # residual + affine-free LayerNorm |
| ├── mlp.py # GELU MLP |
| ├── mlp_res_ln.py # GELU MLP + residual + next LayerNorm |
| └── swiglu_stream.py # weight-streaming SwiGLU (+ residual/LayerNorm variant) |
| ``` |
|
|
| ## Attribution and license |
|
|
| Apache-2.0 (see `LICENSE`, `NOTICE`). The attention, MLP and residual/LayerNorm kernels derive from |
| the V-JEPA2 encoder kernels in |
| [`aws-neuron/vjepa2-neuron-kernels`](https://huggingface.co/kernels/aws-neuron/vjepa2-neuron-kernels), |
| developed by Antonio Mena ([@menaman123](https://huggingface.co/menaman123)). This repo generalizes |
| them: both RoPE pairings, any head count, head_dim 128, batching, padded-key skipping, an LNC=1 |
| all-group variant, and the fused residual/LayerNorm epilogues. `swiglu_stream.py` is new. |
| `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. |
|
|
| DINOv3 and V-JEPA2 model weights are distributed under their own licenses. |
|
|
