---
license: apache-2.0
library_name: kernels
tags:
  - kernel
  - neuron
  - nki
  - trainium
  - trn2
  - bert
  - gte
  - reranker
  - embeddings
  - rope
---

# gte-encoder-neuron-kernels

NKI kernels for **NewBERT-architecture encoders** on AWS Trainium2 (trn2) with PyTorch Native
(`torch.compile(backend="neuron")`). NewBERT is the Alibaba-NLP "new-impl" modeling code used by the gte
model family: post-LayerNorm, rotate-half RoPE with NTK scaling, and a gated GELU MLP.

Each layer runs as two fused kernels instead of dozens of small ops:

| Kernel | Replaces |
|---|---|
| `attn_block` | QKV projection + RoPE + masked softmax attention + output projection + residual + LayerNorm (L <= 1024) |
| `gated_mlp_ln` | gated GELU MLP + residual + LayerNorm (weights resident in SBUF, or streamed when they do not fit) |
| `ln` / `res_ln` | (residual +) LayerNorm, using the Vector engine's `bn_stats` / `bn_aggr` (embedding LayerNorm) |

Other lengths of 768 or more (for example 1536) use `attention_cte` from the Neuron NKI library (nkilib) for
the attention core, which beats the compiler's own attention at those lengths.

Validated models:

| Model | Hidden / heads / layers | Head |
|---|---|---|
| `Alibaba-NLP/gte-multilingual-reranker-base` | 768 / 12 / 12 | sequence classification (reranker) |
| `Alibaba-NLP/gte-multilingual-base` | 768 / 12 / 12 | embeddings |
| `Alibaba-NLP/gte-base-en-v1.5` | 768 / 12 / 12 | embeddings |
| `Alibaba-NLP/gte-large-en-v1.5` | 1024 / 16 / 24 | embeddings |

## Usage

The NewBERT modeling code needs Transformers 4.x (it loads incorrect position ids and RoPE tables on 5.x).
Transformers 4.x pins `huggingface_hub < 1.0`, and `kernels` releases that can resolve kernel-type Hub repos
need `huggingface_hub >= 1.10`. So in the model's environment, load the kernels from a clone of this repo with
`get_local_kernel`:

```bash
git clone https://huggingface.co/kernels/jburtoft/gte-encoder-neuron-kernels
pip install "transformers>=4.53,<5" "kernels==0.12.3"   # 0.12.3: newest kernels compatible with hub < 1.0
```

```python
import os
os.environ.setdefault("NEURON_LOGICAL_NC_CONFIG", "2")   # read once, when the kernels are imported

from pathlib import Path
import torch
from kernels import get_local_kernel
from transformers import AutoModel

gte = get_local_kernel(Path("gte-encoder-neuron-kernels"), "gte_encoder_neuron_kernels")

model = AutoModel.from_pretrained("Alibaba-NLP/gte-large-en-v1.5", trust_remote_code=True,
                                  attn_implementation="eager").eval()
for mod in model.modules():                     # the fused MLP uses the tanh GELU (see Accuracy)
    for name, child in mod.named_children():
        if "GELU" in type(child).__name__:
            setattr(mod, name, torch.nn.GELU(approximate="tanh"))
model = model.to(torch.bfloat16)
gte.patch_model(model)                          # in place; returns the number of layers patched
model = torch.compile(model.to("neuron"), backend="neuron", dynamic=False)

out = model(input_ids=ids.to("neuron"), attention_mask=mask.to("neuron"))   # right-padded, fixed L
```

In an environment with `huggingface_hub >= 1.10` and `kernels >= 0.14`, for example when you call the ops
from your own model code, `get_kernel` works directly:
`gte = get_kernel("jburtoft/gte-encoder-neuron-kernels", version=1, trust_remote_code=True)`.

Full scripts: [`examples/rerank.py`](examples/rerank.py) (reranker) and [`examples/embed.py`](examples/embed.py)
(embeddings). Pad every batch to a fixed length (for example 128 / 256 / 512) so that each (length,
batch) shape compiles once.

`patch_model(model)` patches every `NewLayer` and the embedding LayerNorm. Shapes that the kernels do not
support fall back to the stock modules at trace time, so patching is always safe:

* fused attention block: L in {128, 256, 384, 512, 768, 1024}; batch divisible by the number of logical
  cores, or (L / 128) divisible by it (an odd batch at LNC=2 splits the queries across the two cores);
* MLP and LayerNorm kernels: batch x L divisible by 128 x LNC;
* other L >= 768: stock attention projections + nkilib `attention_cte` + the NKI LayerNorm.

The raw kernels (`attn_block_kernel`, `gated_mlp_ln_kernel`, `gated_mlp_ln_stream_kernel`,
`res_ln_kernel`) and the host-side weight packer (`pack_attn_block`) are exported too, for use from other
frameworks or under `nki.simulate`.

## Performance (trn2.3xlarge, bf16)

Measured 2026-10 on PyTorch Native Beta 6 (torch 2.13.0, torch-neuronx 2.13.3, NKI 0.7.0b1) with real
Hugging Face weights. "Stock" is the same model, also with the tanh GELU, compiled with
`torch.compile(backend="neuron")` without these kernels, in the same container on the same instance.

**Whole instance** (one process per logical core, all cores busy at once; best LNC for each column):

| Model | L x batch | Stock (seq/s) | **These kernels (seq/s)** | Speedup |
|---|---|---|---|---|
| gte-multilingual-reranker-base | 128 x 32 | 8,463 | **12,427** | 1.47x |
| gte-multilingual-reranker-base | 256 x 32 | 3,864 | **5,865** | 1.52x |
| gte-base-en-v1.5 | 128 x 32 | 8,407 | **12,397** | 1.47x |
| gte-base-en-v1.5 | 256 x 32 | 3,879 | **5,831** | 1.50x |
| gte-large-en-v1.5 | 128 x 32 | 2,645 | **3,427** | 1.30x |
| gte-large-en-v1.5 | 256 x 32 | 1,278 | **1,622** | 1.27x |

**One logical core** (single process, LNC=2, ms per batch, median):

| Model | L x batch | Stock | **These kernels** | Speedup |
|---|---|---|---|---|
| gte-base-en-v1.5 | 128 x 32 | 16.3 | **10.1** | 1.62x |
| gte-base-en-v1.5 | 256 x 32 | 35.8 | **20.2** | 1.77x |
| gte-base-en-v1.5 | 512 x 16 | 39.4 | **22.8** | 1.73x |
| gte-base-en-v1.5 | 512 x 1 | 4.0 | **3.2** | 1.24x |
| gte-large-en-v1.5 | 128 x 32 | 44.8 | **30.5** | 1.47x |
| gte-large-en-v1.5 | 256 x 32 | 104.4 | **61.1** | 1.71x |
| gte-large-en-v1.5 | 512 x 16 | 111.7 | **68.0** | 1.64x |
| gte-large-en-v1.5 | 512 x 1 | 9.2 | **7.9** | 1.17x |

**Long sequences** (single process, ms per batch, median; v1.0.0 used `attention_cte` + unfused projections at
these lengths, v1.1.0 runs the fused attention block):

| Model | LNC | L x batch | Stock | v1.0.0 | **v1.1.0** | vs stock |
|---|---|---|---|---|---|---|
| gte-multilingual-reranker-base | 2 | 768 x 8 | 40.0 | 33.3 | **20.1** | 1.99x |
| gte-multilingual-reranker-base | 2 | 1024 x 8 | 54.7 | 46.6 | **29.3** | 1.87x |
| gte-multilingual-reranker-base | 2 | 1024 x 1 | 8.0 | 7.7 | **5.7** | 1.41x |
| gte-multilingual-reranker-base | 1 | 1024 x 8 | 107.2 | 82.9 | **54.0** | 1.98x |
| gte-large-en-v1.5 | 2 | 768 x 8 | 108.0 | 95.3 | **62.2** | 1.74x |
| gte-large-en-v1.5 | 2 | 1024 x 8 | 155.7 | 127.9 | **87.3** | 1.78x |
| gte-large-en-v1.5 | 2 | 1024 x 1 | 21.5 | 18.8 | **15.1** | 1.43x |
| gte-large-en-v1.5 | 1 | 1024 x 8 | 307.5 | 242.0 | **168.7** | 1.82x |

gte-large's MLP weights (1024 x 8192 + 4096 x 1024 per layer) do not fit in SBUF next to the working
set, so its MLP kernel streams them. With every core busy, that streaming competes for HBM bandwidth,
which is why its whole-instance gain is lower than its single-core gain.

## Accuracy

Checked against the stock Hugging Face model on CPU in fp32 with the exact GELU, on 64-72 real-sentence
inputs per shape with ragged lengths (6 tokens up to the padded length), at L = 128 / 256 / 512 / 768 / 1024,
LNC=1 and LNC=2 (768 / 1024 measured on the reranker and gte-large). "Stock" is stock bf16 PyTorch Native on the same inputs.

| Model | Metric | Stock bf16 | These kernels |
|---|---|---|---|
| gte-multilingual-reranker-base | score Pearson | 0.99997-0.99999 | 0.99997-0.99998 |
| gte-multilingual-reranker-base | top-1 passage per query = fp32 | same as stock | same as stock |
| gte-multilingual-base | CLS cosine, mean | 0.99984-0.99987 | 0.99992-0.99993 |
| gte-base-en-v1.5 | CLS cosine, mean | 0.99991-0.99992 | 0.99993 |
| gte-large-en-v1.5 | CLS cosine, mean | 0.99986-0.99991 | 0.99990-0.99992 |
| all embedding models | retrieval top-1 = fp32 | 0.88-1.00 | 0.88-1.00 (never below stock) |

The kernels keep the residual sum and the LayerNorm statistics in fp32, so on average they are slightly
closer to fp32 than stock bf16. The worst single input is on par with stock: the lowest per-text cosine
was 0.99900-0.99991 with the kernels against 0.99834-0.99984 for stock. On gte-large it trailed stock's
worst by up to 2.3e-4 in some configurations. Padding is exact: a sequence's output is bit-identical
whatever the other sequences in its batch contain.

**GELU.** The MLP kernel computes `gelu_tanh` (the tanh approximation), while the models specify the exact
(erf) GELU. All accuracy numbers above are against the exact-GELU fp32 model, so they include this
difference. The usage snippet swaps the stock GELU module for the tanh version too, so that fallback shapes
compute the same function.

## Constraints

* Inference only (eval mode, no dropout, no `head_mask`).
* bf16 activations and weights; hidden size a multiple of 128; head_dim 64; number of heads a multiple of 4.
* Right padding only, with a standard 0/1 `attention_mask`.
* Set `NEURON_LOGICAL_NC_CONFIG` before importing the kernels.
* Requires the NewBERT `NewModel` eager path: `unpad_inputs=False` and `use_memory_efficient_attention=False`
  (the defaults of all four validated models).

## Tests

`tests/` has simulator tests (`pytest -m sim`, kernels under `nki.simulate` against NumPy fp32 references),
a CPU test of the weight packing against the Hugging Face layer (`-m cpu`), and on-device accuracy tests
(`-m neuron`). See [`tests/README.md`](tests/README.md). A simulator pass is not device evidence: the
simulator does not model SBUF or PSUM capacity, and some DMA access patterns behave differently on
hardware.

## Requirements

* trn2, 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 4.x (verified on 4.53.3) with `trust_remote_code=True` for the NewBERT modeling code.
  Transformers 5.x does not load these models correctly.
* `kernels`: 0.12.3 with `get_local_kernel` next to Transformers 4.x (verified), or `get_kernel(...,
  version=1, trust_remote_code=True)` with kernels 0.17.2 / huggingface_hub 2.x (verified on trn2).

## Repository layout

```
build/torch-neuron/
├── __init__.py            # public API
├── metadata.json
├── torch_ops.py           # torch custom ops (gte_enc::*), NkiLN / NkiMlpLN / NkiAttnBlock, patch_model
├── attn_block_kernel.py   # fused attention block
├── attn_block_pack.py     # host-side weight packing (NumPy)
├── gated_mlp_kernel.py    # fused gated MLP + residual + LayerNorm (resident and streaming)
└── res_ln_kernel.py       # (residual +) LayerNorm
examples/                  # rerank.py, embed.py
tests/                     # simulator, CPU and device tests
```

## License

Apache-2.0. See [LICENSE](LICENSE) and [NOTICE](NOTICE).