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:

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:

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:

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

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 (reranker) and 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):

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

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):

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.

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. 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 and NOTICE.