TurboGeLU A family of fast, approximate, and accurate GeLU computational kernels

TurboGeLU is a fused polynomial GeLU kernel for PyTorch (more frameworks in the future!) with custom CPU, Apple Metal (mps), and optional CUDA backends. By exploiting both mathematical symmetries and hardware-friendly code, this kernel achieves up to 11x speedups compared to PyTorch.

The formulation replaces transcendental-heavy (erf-based) GeLU evaluation with a compact piecewise polynomial that is cheap to evaluate in float32, then exposes both forward and backward kernels through torch.ops.turbogelu.* and a differentiable turbogelu.gelu(...) wrapper. The polynomial differs from the exact PyTorch variant by about 1e-5 on average, about 2 orders of magnitude more accurate than the tanh approximation built into PyTorch.

I originally developed this kernel's mathematical formulation as a part of course project at ETH Zürich, the Advanced Systems Lab. There, we developed a holistic runtime optimization of the Geometry Algebra Transformer (GATr) architecture using AVX2 intrinsics. Read the entire report here.

Motivation

GeLU is everywhere in transformer MLPs, but the standard formulations are not especially friendly to low-level kernels:

• gelu(..., approximate="none") depends on erf.
• gelu(..., approximate="tanh") still depends on tanh.
• Both are materially more expensive than fused multiply-adds, especially on single-threaded CPU code and on lightweight GPU kernels where dispatch overhead matters.

TurboGeLU targets the common case where you want a fast, standalone GeLU operator with straightforward, transparent control flow and a shared implementation strategy across accelerators which performs well across the board and requires minimum device-specific algorithmic changes. The repository focuses on:

1. A polynomial approximation generated offline
2. An Estrin-style evaluation tree with short dependency chains
3. Matched backward kernel derived from the same coefficients
4. Backend-specific implementations for CPU, Apple Metal, and CUDA

Theory

1. Clamp the tails

For sufficiently negative inputs, GeLU is extremely close to 0 and similarly, for sufficiently positive inputs, it is extremely close to x. The kernel therefore uses a piecewise form with a finite approximation window, outside of which it applies the standard ReLU operation.

TurboGeLU(x) =
  0                    x < -4.389
  0.5 x + x^2 P(x^2)  |x| <= 4.389
  x                    x > 4.389

The generated coefficients use [-4.389, 4.389] as the approximation interval. This interval was chosen based on the results of a grid search performed over the IRLS (Iterated Reweighted Least Squares) coefficient finding procedure, to give the lowest approximation error.

2. Fit an even polynomial inside the window

The script scripts/gelu_poly_codegen.py generates coefficients for a degree-17 odd approximation, represented as a degree-8 polynomial in y = x^2:

0.5 x + y P(y),  where P(y) = b1 + b2 y + ... + b9 y^8

The fitting procedure is:

1. Sample the target GeLU on [0, R].
2. Solve a least-squares problem for the coefficients.
3. Reweight the residuals with IRLS to reduce worst-case error.
4. Rescale back to the original interval.

The script fits against either one of PyTorch's GeLU implementations. For the purpose of the released kernel, I used the "exact" variant of GeLU, as the polynomial approximates it in the same sense as the tanh variant.

3. Evaluate with Estrin's scheme

The polynomial is not evaluated naively. TurboGeLU groups terms into an Estrin tree so the kernel can exploit FMAs, instruction-level parallelism (critical for CPU performance), and repeated squaring:

y  = x^2
y2 = y^2
y4 = y2^2
y8 = y4^2

a0 = b1 + b2 y
a1 = b3 + b4 y
a2 = b5 + b6 y
a3 = b7 + b8 y
h0 = a0 + a1 y2
h1 = a2 + a3 y2
P  = h0 + h1 y4 + b9 y8

or, through a graphical representation:

This is exactly the structure used in src/turbogelu/include/turbogelu/gelu_poly.h, with scalar and vectorized implementations for the available backend.

4. Differentiate the same polynomial for backward

Backward uses the derivative of the polynomial branch, with separate precomputed constants kE0..kE8 derived from kB1..kB9. Outside the approximation window, the gradient collapses to the derivative of the clamp:

• 0 for x < -4.389
• 1 for x > 4.389

That keeps forward and backward aligned and makes the backward pass just as cheap to evaluate as the forward pass.

Package import and usage

import torch
import turbogelu

x = torch.randn(1_048_576, device="cpu", dtype=torch.float32, requires_grad=True)
y = turbogelu.gelu(x)

Exported entry points:

• turbogelu.gelu(x): autograd-aware public API
• turbogelu.gelu_op(x): raw forward op
• turbogelu.gelu_backward_op(grad_out, x): raw backward op

Performance Benchmarks

Benchmarks were produced by scripts/benchmark_gelu.py.

Setup used by the script:

• 1,048,576 random elements plus 15 edge-case values
• 50 timed batches
• 100 GeLU calls per timed batch
• single-threaded CPU timing
• explicit mps / cuda synchronization so GPU timings measure execution rather than simple kernel dispatch/enqueue latency
• Disclaimer: I evaluated only on hardware I have access to, for now. In the future, I am aiming to evaluate on high-bandwidth HBMx-based Nvidia GPUs to verify the memory/compute bounds of these kernels.

Apple M1 Pro CPU

Forward	Backward

TurboGeLU is noticeably faster than both PyTorch baselines on this single-threaded CPU benchmark. The biggest gains are against torch approx, but it also consistently outperforms exact GeLU in both passes and all tested dtypes.

direction	dtype	method	median us/op	p5	p95	speedup vs turbo
forward	float32	torch exact	1256.41	1243.24	1323.10	3.27x
forward	float32	torch approx	4350.62	4340.33	4385.41	11.33x
forward	float32	turbogelu	383.86	378.99	391.11	-
forward	float16	torch exact	1199.46	1183.75	1239.50	2.79x
forward	float16	torch approx	4460.40	4449.67	4622.82	10.37x
forward	float16	turbogelu	430.19	428.63	441.14	-
forward	bfloat16	torch exact	1382.43	1356.36	1432.85	1.68x
forward	bfloat16	torch approx	4704.08	4698.19	4739.10	5.73x
forward	bfloat16	turbogelu	820.61	804.14	857.27	-
backward	float32	torch exact	2393.74	2355.08	2469.22	5.31x
backward	float32	torch approx	4761.72	4729.56	4898.81	10.56x
backward	float32	turbogelu	450.76	438.93	463.28	-
backward	float16	torch exact	2377.02	2363.06	2405.23	4.84x
backward	float16	torch approx	4492.40	4482.40	4738.43	9.14x
backward	float16	turbogelu	491.56	486.09	509.50	-
backward	bfloat16	torch exact	2597.50	2589.88	2624.62	2.90x
backward	bfloat16	torch approx	4708.50	4696.37	4872.91	5.26x
backward	bfloat16	turbogelu	894.39	886.13	911.18	-

Apple M1 Pro GPU (Metal Performance Shaders)

Forward	Backward

On Apple's GPU backend, my polynomial still wins, but the margin is smaller than on CPU. The strongest gains show up in float16 and bfloat16, while float32 remains a modest improvement over both PyTorch implementations.

direction	dtype	method	median us/op	p5	p95	speedup vs turbo
forward	float32	torch exact	34.91	34.26	37.24	1.17x
forward	float32	torch approx	32.19	31.79	34.89	1.07x
forward	float32	turbogelu	29.94	29.21	33.28	-
forward	float16	torch exact	32.32	31.75	32.94	1.61x
forward	float16	torch approx	27.22	26.81	29.33	1.35x
forward	float16	turbogelu	20.13	19.23	22.89	-
forward	bfloat16	torch exact	36.50	36.04	37.19	1.51x
forward	bfloat16	torch approx	30.08	29.64	31.52	1.25x
forward	bfloat16	turbogelu	24.10	23.19	24.63	-
backward	float32	torch exact	44.33	43.63	48.47	1.12x
backward	float32	torch approx	42.06	41.62	43.39	1.06x
backward	float32	turbogelu	39.73	38.78	40.49	-
backward	float16	torch exact	39.15	38.53	40.55	1.75x
backward	float16	torch approx	34.71	33.27	39.26	1.55x
backward	float16	turbogelu	22.35	21.94	24.23	-
backward	bfloat16	torch exact	44.17	43.67	46.09	1.91x
backward	bfloat16	torch approx	35.28	34.77	36.94	1.53x
backward	bfloat16	turbogelu	23.09	22.49	26.30	-

Ryzen 7 5700X CPU

Forward	Backward

The x86 CPU results tell the same story as Apple M1 CPU, TurboGeLU clearly outperforms both PyTorch paths, and the advantage is especially large against the tanh approximation.

direction	dtype	method	median us/op	p5	p95	speedup vs turbo
forward	float32	torch exact	612.11	611.59	619.79	2.47x
forward	float32	torch approx	2511.82	2494.47	2549.99	10.13x
forward	float32	turbogelu	248.00	241.03	248.59	-
forward	float16	torch exact	1133.07	1129.77	1139.41	3.98x
forward	float16	torch approx	2511.42	2511.07	2513.07	8.83x
forward	float16	turbogelu	284.56	284.48	288.93	-
forward	bfloat16	torch exact	1219.17	1218.88	1227.39	3.22x
forward	bfloat16	torch approx	2577.12	2576.61	2579.31	6.80x
forward	bfloat16	turbogelu	378.94	378.76	380.73	-
backward	float32	torch exact	858.82	852.79	862.12	3.26x
backward	float32	torch approx	2524.10	2514.69	2527.87	9.58x
backward	float32	turbogelu	263.37	263.17	263.85	-
backward	float16	torch exact	1678.91	1677.10	1681.87	5.17x
backward	float16	torch approx	2391.74	2391.18	2395.15	7.37x
backward	float16	turbogelu	324.71	323.29	325.05	-
backward	bfloat16	torch exact	1802.09	1796.33	1807.05	4.19x
backward	bfloat16	torch approx	2549.67	2549.08	2551.45	5.92x
backward	bfloat16	turbogelu	430.35	430.18	433.72	-

RTX 3070 Ti CUDA

Forward	Backward

CUDA is the outlier in the current benchmark set, TurboGeLU is effectively on par with PyTorch and occasionally slightly behind the fastest baseline in low precision. That suggests the current CUDA path is not yet the main win of the project. Based on some preliminary measurements, it seems that both Nvidia's low-level erf implementation and my kernel are memory bound essentially sitting on the GPU roofline, and as such, lower runtime cannot be achieved.

direction	dtype	method	median us/op	p5	p95	speedup vs turbo
forward	float32	torch exact	17.23	17.21	17.26	1.00x
forward	float32	torch approx	17.23	17.21	17.26	1.00x
forward	float32	turbogelu	17.17	17.13	17.23	-
forward	float16	torch exact	5.36	5.34	5.38	1.00x
forward	float16	torch approx	5.08	5.07	5.11	0.95x
forward	float16	turbogelu	5.36	5.30	5.42	-
forward	bfloat16	torch exact	5.55	5.53	5.61	1.03x
forward	bfloat16	torch approx	4.95	4.94	4.96	0.92x
forward	bfloat16	turbogelu	5.39	5.30	5.45	-
backward	float32	torch exact	24.49	24.44	24.53	1.00x
backward	float32	torch approx	24.56	24.50	24.64	1.00x
backward	float32	turbogelu	24.52	24.49	24.56	-
backward	float16	torch exact	12.24	12.20	12.28	1.00x
backward	float16	torch approx	12.22	12.19	12.26	1.00x
backward	float16	turbogelu	12.19	12.16	12.22	-
backward	bfloat16	torch exact	12.29	12.26	12.34	1.00x
backward	bfloat16	torch approx	12.17	12.14	12.21	0.99x
backward	bfloat16	turbogelu	12.29	12.23	12.36	-

Numerical Comparison Benchmark

To verify numerical correctness, and its consequent long-term effects on training consistency/stability, a minimalist example of an MLP trained on MNIST is available in here.

From the training dynamics, it is clear that TurboGeLU performs essentially identically to standard PyTorch (exact) implementation.

epoch 01 | exact train/test acc 0.8861/0.9453 | turbo train/test acc 0.8861/0.9453
epoch 02 | exact train/test acc 0.9538/0.9621 | turbo train/test acc 0.9538/0.9621
epoch 03 | exact train/test acc 0.9675/0.9694 | turbo train/test acc 0.9675/0.9694
epoch 04 | exact train/test acc 0.9759/0.9732 | turbo train/test acc 0.9759/0.9732
epoch 05 | exact train/test acc 0.9810/0.9765 | turbo train/test acc 0.9810/0.9765
epoch 06 | exact train/test acc 0.9841/0.9741 | turbo train/test acc 0.9841/0.9741
epoch 07 | exact train/test acc 0.9875/0.9773 | turbo train/test acc 0.9875/0.9773
epoch 08 | exact train/test acc 0.9905/0.9760 | turbo train/test acc 0.9905/0.9760
epoch 09 | exact train/test acc 0.9917/0.9787 | turbo train/test acc 0.9917/0.9787
epoch 10 | exact train/test acc 0.9944/0.9794 | turbo train/test acc 0.9944/0.9794
epoch 11 | exact train/test acc 0.9947/0.9805 | turbo train/test acc 0.9947/0.9805
epoch 12 | exact train/test acc 0.9957/0.9779 | turbo train/test acc 0.9957/0.9779
epoch 13 | exact train/test acc 0.9965/0.9783 | turbo train/test acc 0.9965/0.9783
epoch 14 | exact train/test acc 0.9974/0.9811 | turbo train/test acc 0.9974/0.9811
epoch 15 | exact train/test acc 0.9981/0.9799 | turbo train/test acc 0.9981/0.9797
epoch 16 | exact train/test acc 0.9978/0.9804 | turbo train/test acc 0.9976/0.9818
epoch 17 | exact train/test acc 0.9973/0.9782 | turbo train/test acc 0.9971/0.9780
epoch 18 | exact train/test acc 0.9980/0.9796 | turbo train/test acc 0.9990/0.9801
epoch 19 | exact train/test acc 0.9975/0.9781 | turbo train/test acc 0.9984/0.9801
epoch 20 | exact train/test acc 0.9982/0.9779 | turbo train/test acc 0.9978/0.9803
epoch 21 | exact train/test acc 0.9982/0.9786 | turbo train/test acc 0.9974/0.9756
epoch 22 | exact train/test acc 0.9987/0.9792 | turbo train/test acc 0.9981/0.9812
epoch 23 | exact train/test acc 0.9976/0.9798 | turbo train/test acc 0.9983/0.9788
epoch 24 | exact train/test acc 0.9984/0.9782 | turbo train/test acc 0.9981/0.9802
epoch 25 | exact train/test acc 0.9988/0.9822 | turbo train/test acc 0.9990/0.9808
epoch 26 | exact train/test acc 0.9975/0.9811 | turbo train/test acc 0.9984/0.9795
epoch 27 | exact train/test acc 0.9992/0.9823 | turbo train/test acc 0.9989/0.9811
epoch 28 | exact train/test acc 1.0000/0.9833 | turbo train/test acc 0.9992/0.9821
epoch 29 | exact train/test acc 1.0000/0.9830 | turbo train/test acc 0.9996/0.9795
epoch 30 | exact train/test acc 1.0000/0.9833 | turbo train/test acc 0.9996/0.9825
epoch 31 | exact train/test acc 1.0000/0.9834 | turbo train/test acc 1.0000/0.9833
epoch 32 | exact train/test acc 1.0000/0.9831 | turbo train/test acc 1.0000/0.9828
epoch 33 | exact train/test acc 1.0000/0.9831 | turbo train/test acc 1.0000/0.9830
epoch 34 | exact train/test acc 1.0000/0.9829 | turbo train/test acc 1.0000/0.9830
epoch 35 | exact train/test acc 1.0000/0.9829 | turbo train/test acc 1.0000/0.9830

Reproducing the polynomial and benchmarks

To compile the CPU kernels, standard apple-clang (xcode command line tools) or clang/gcc (linux) are sufficient. MPS kernels for Mac are compiled using PyTorch's built in runtime, while for CUDA kernels the CUDA Toolkit must be present.

Install locally:

pip install -e .

Regenerate coefficients and the approximation plot:

python scripts/gelu_poly_codegen.py --range 4.389 --iters 35 --points 20000 --gelu_mode tanh

Run the benchmark harness:

python scripts/benchmark_gelu.py --direction forward
python scripts/benchmark_gelu.py --direction backward

The benchmark script writes a PNG plot to benchmarks/ and prints the markdown timing table used above.

TODOs

Currently, some work remains to be done to further validate and test the performance improvements of my kernel on latest Nvidia hardware.

Ideally, one should run the benchmarks on an HBMx-based GPU.