jburtoft/gelu-erf-neuron

Exact-erf GELU NKI Kernel for AWS Trainium/Inferentia

Version 1.0.0: Fused exact-erf GELU forward + backward for PyTorch Native on Neuron. The transcendental (erf) is computed by a single hardware activation-table instruction on the Scalar/activation engine, replacing the compiler's multi-engine rational-erf lowering. Drop-in replacement for nn.GELU() (default, exact/erf), with a custom NKI backward. Integrates under torch.compile(backend="neuron") at 1 graph / 0 graph breaks.

What ships

Two NKI kernels + a functional wrapper + an nn.Module. All figures below are measured on trn2.3xlarge, PyTorch Native Beta 4 (SDK 2.31, NKI 0.5.0), unless stated otherwise.

Component	Purpose
gelu_fwd	Exact (erf) GELU forward. [R, C] -> [R, C]. One nl.gelu activation-table op per 128-row tile.
gelu_bwd	Exact GELU backward. Needs the saved pre-activation x and incoming grad. One nl.gelu_dx gate (Scalar engine) + one Vector multiply per tile.
gelu_erf	Functional API: torch.autograd.Function wiring gelu_fwd (forward) and gelu_bwd (backward).
GELUErf (nn.Module)	Drop-in replacement for nn.GELU().

GELUErf is exact erf-GELU (not the tanh approximation). It applies to any model whose MLP uses nn.GELU() with default args (e.g. ViT-family, BERT-family). It was developed and measured on the DINOv3 ViT-L/16 MLP (hidden width 4096).

Quick start

Drop-in module replacement

import torch, torch.nn as nn
from kernels import get_kernel

# Personal (non-org) repo: pass trust_remote_code=True, and revision or version.
k = get_kernel("jburtoft/gelu-erf-neuron", trust_remote_code=True, revision="v1")

# Swap nn.GELU() for the NKI kernel in every MLP block of your model:
def swap_gelu(module):
    n = 0
    for m in module.modules():
        if hasattr(m, "act") and isinstance(m.act, nn.GELU):
            m.act = k.GELUErf()
            n += 1
    return n

model = model.to(torch.bfloat16).to("neuron")
swap_gelu(model)

compiled = torch.compile(model, backend="neuron", fullgraph=False, dynamic=False)
# forward + backward + optimizer step as usual; gradients flow through gelu_bwd

Functional API

import torch
from kernels import get_kernel

k = get_kernel("jburtoft/gelu-erf-neuron")

x = torch.randn(3216, 4096, device="neuron", dtype=torch.bfloat16, requires_grad=True)
y = k.gelu_erf(x)      # forward via gelu_fwd
y.sum().backward()     # backward via gelu_bwd

Raw kernels

from kernels import get_kernel
k = get_kernel("jburtoft/gelu-erf-neuron")

y  = k.gelu_fwd(x2d)             # x2d: [R, C] on device
dx = k.gelu_bwd(x2d, grad_out)   # x2d = saved pre-activation, grad_out = dL/dy

Convention: gelu_fwd/gelu_bwd take 2-D [R, C] device tensors (partition dim tiled by 128, C = hidden free dim). The gelu_erf / GELUErf wrappers flatten leading dims to [R, C] and restore afterwards.

Correctness (measured, Beta 4 trn2.3xlarge, NKI 0.5.0)

Reference: CPU FP32 exact-erf GELU (torch.nn.functional.gelu, default) and its autograd derivative.

Kernel	dtype	cos_sim	max_abs
gelu_fwd	FP32	1.000000	2.1e-6
gelu_bwd	FP32	1.000000	3.1e-6
gelu_fwd	BF16	0.999998	3.1e-2
gelu_bwd	BF16	0.999995	4.0e-2

The FP32 cos_sim of 1.000000 confirms nl.gelu implements exact erf-GELU and nl.gelu_dx its exact derivative — no tanh substitution. Integrated into all 24 DINOv3 ViT-L MLP blocks, the 20-step training loss curve matched the native-nn.GELU baseline within 2.3e-3.

Performance (measured)

End-to-end (DINOv3 ViT-L/16 training, all 24 MLP blocks swapped)

Same-session A/B, torch.compile(backend="neuron"), BF16:

Config	native nn.GELU	GELUErf (NKI)	Delta
Single-core, BS=16 (Beta 4)	28.58 img/s	30.30 img/s	+6.0%
FSDP2 DP=4, per-rank BS=32 (Beta 4)	112.6 img/s	121.9 img/s	+8.3%
Single-core, BS=16 (Beta 5, SDK 2.32)	—	—	+4.4%

Integration: torch._dynamo.explain reports 1 graph, 0 graph breaks.

Standalone kernel profile ([3216, 4096], BF16, Beta 4)

Instruction-count comparison of the compiled op:

	native nn.GELU lowering	gelu_fwd (NKI)
MATMUL instrs	1664	0
RECIPROCAL instrs	256	0
Vector TENSOR_TENSOR instrs	2816	0
ACTIVATE instrs	3584	26
Vector-engine compute time	2,501,177 ns	1,047 ns

The native lowering implements erf as a rational/polynomial expansion across the Tensor, Vector, and Scalar engines; the NKI kernel is one nl.gelu activation-table lookup per 128-row tile.

Engine-active share of the integrated training step, native → GELUErf: forward Vector-engine 75.3% → 63.1%; backward Vector-engine 49.7% → 36.1%. Forward MFU 10.0% → 10.7%; backward MFU 10.7% → 11.6%.

Standalone kernel bound (profiled)

Metric	gelu_fwd	gelu_bwd
total_time	166.8 µs	218.5 µs
dma_active_time_percent	90.9%	95.0%
MBU	44.1%	50.5%
Scalar ACTIVATE time	93.9 µs (hidden)	93.9 µs (hidden)
Vector TENSOR_TENSOR time	—	59.6 µs (hidden)
effective DMA throughput	316 GB/s	362 GB/s

Both kernels are DMA-bound: the activation compute is fully overlapped under DMA, and all transfers run on a single DMA queue. Measured double-buffer, static-unroll, wider-free-dim, and dge_mode=hwdge variants were each bit-identical to the baseline (max_abs vs baseline = 0.0) and moved the end-to-end result by −0.08% (single-core) / +0.3% (replica-DP=4) — within noise.

Requirements

• Hardware: trn2 (measured on trn2.3xlarge).
• SDK: PyTorch Native Beta 4 (NKI 0.5.0, SDK 2.31). Also runs on Beta 5 (NKI 0.6.0, SDK 2.32) with no kernel changes.
• PyTorch: 2.11.0 (bundled with Beta 4).
• Kernels library: pip install kernels.

Repository layout

build/torch-neuron/
├── __init__.py          # public API: GELUErf, gelu_erf, gelu_fwd, gelu_bwd
├── metadata.json        # HF kernels library metadata
├── layers.py            # GELUErf nn.Module (drop-in for nn.GELU)
├── autograd.py          # _FusedGELUErfFn autograd.Function + gelu_erf functional
└── nki_kernels/
    ├── __init__.py
    └── gelu_erf.py       # gelu_fwd + gelu_bwd NKI kernels

Known limitations

• DMA-bound at the measured shape. [3216, 4096] BF16: both kernels sit at the single-DMA-queue throughput ceiling (~316–362 GB/s) with compute fully hidden. Measured multi-buffering and DMA-routing variants produced no end-to-end change on this SDK/shape.
• End-to-end gain is modest and workload-dependent. The measured +6.0% single-core / +8.3% FSDP2 on DINOv3 ViT-L reflects GELU being one op among the MLP matmuls the compiler already schedules well; the gain will differ on other models/shapes and should be re-measured for your workload.
• Exact-erf GELU only. For nn.GELU(approximate="tanh") this kernel is not a drop-in.

Attribution

Original NKI kernels, developed and measured on DINOv3 ViT-L/16 training on trn2.3xlarge. Uses NKI 0.5.0 nl.gelu / nl.gelu_dx activation-table ops. Published under Apache 2.0.

License

Apache 2.0.