Triton and Helion are how most custom GPU kernels get written today. Triton, from OpenAI, lets you write a kernel as a Python function over a block of data and handles the thread mapping, memory coalescing, shared memory and tensor cores for you. Helion, from PyTorch, sits one level higher and generates and tunes Triton kernels from tile-level code. Both are excellent, and both are large: Triton alone is hundreds of thousands of lines built on MLIR and LLVM.
I wanted to understand how that stack actually works, and reading the source was not getting me there. The reliable way to understand a big system is to rebuild a small one, so that is what I did: a working Triton and a working Helion, small enough to read in an afternoon but fast enough to benchmark honestly against the originals.
The result is about 4,000 lines of Python with no dependencies beyond torch. Memory-bound kernels like softmax and layernorm match Triton at the memory-bandwidth ceiling; fp16 matmul reaches 76 to 83 percent of Triton's tensor-core throughput sustained, and 110 TFLOP/s cold. This post covers how it works and where the performance came from.
The Triton layer is called newt (a newt is a small triton) and the Helion layer is called deuteron (a lighter nucleus than a helion). If you have written a Triton kernel, newt will look familiar. Swap tl for nl and most of it runs unchanged:
@newt.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: nl.constexpr):
pid = nl.program_id(0)
offs = pid * BLOCK + nl.arange(0, BLOCK)
mask = offs < n
x = nl.load(x_ptr + offs, mask=mask)
y = nl.load(y_ptr + offs, mask=mask)
nl.store(out_ptr + offs, x + y, mask=mask)
The package is on PyPI: pip install nano-triton gives you both newt and deuteron. It needs torch and an NVIDIA GPU, since newt compiles the kernels at runtime with NVRTC.
Compiling to the GPU without MLIR or LLVM
The first question was how to get from Python to something the GPU can run, without rebuilding the parts of Triton I had no interest in reimplementing. Triton lowers its IR through MLIR and LLVM to PTX, and that machinery is most of the compiler and none of the ideas I wanted to learn.
The way around it is NVRTC, NVIDIA's runtime compiler. It is a full CUDA C++ compiler shipped as a library: you pass it a string of CUDA C++ and get back a compiled binary, in-process, with no temporary files and no subprocess. This is the decision the whole project rests on. NVRTC handles register allocation and instruction scheduling, which are the hardest and most time-consuming parts of a backend, so newt does not have to.
The pipeline is therefore short. newt parses the Python function into a syntax tree, assigns every value a shape, dtype and layout, emits CUDA C++, compiles it with NVRTC, and launches the result through the CUDA driver API using ctypes.
That leaves one real problem, and it is the one that makes Triton worth studying: turning block-level semantics into thread-level code.
The register layout
When a kernel loads "a block of 1,024 floats," those values have to be distributed across the threads of the block, and the distribution largely determines performance. newt uses a single rule, the group-cyclic layout: elements are handed to threads in contiguous 16-byte groups, round robin.
This one rule gives three properties without any further work:
- Coalescing. Adjacent threads hold adjacent memory, so a warp-wide load becomes one wide memory transaction instead of up to 32 scattered ones.
- Vectorization. Each thread owns a contiguous 16 bytes, which is exactly one 128-bit load or store.
- Simplicity. Elementwise operations never need to know where an element lives; each thread iterates over its own slots.
There is one deliberate difference from Triton. Triton proves that accesses are contiguous through static analysis at compile time. Writing that analysis is a substantial amount of work, so newt takes a different route: it emits a small runtime check (are these offsets consecutive, aligned and unmasked?) and branches to the vectorized path when the check passes. A few integer comparisons in front of a 500-cycle memory access is a good trade, and it is what brought vector-add from about 82 percent of Triton to parity.
Memory-bound kernels
Memory-bound kernels were the easy part. Softmax, layernorm and elementwise operations touch each byte once, so performance is bounded entirely by memory bandwidth. Once the layout is coalesced and vectorized and the passes are fused into a single kernel, there is nothing left to optimize.
newt matches Triton on all of them. This is worth stating plainly rather than celebrating: parity here is the expected result, not an achievement. Any correct compiler that coalesces, vectorizes and fuses lands in the same place, because the bandwidth ceiling is a hard limit. The workload that actually distinguishes compilers is matrix multiplication.
Matmul
Matmul is compute-bound: the arithmetic dominates the memory traffic, so the objective changes from conserving bandwidth to keeping the tensor cores busy. Every scheduling inefficiency shows up directly in the achieved TFLOP/s. newt's matmul went through three versions.
Version 1: WMMA with synchronous staging. The first implementation used WMMA, NVIDIA's high-level tensor-core API, and staged tiles into shared memory synchronously. At 4096-cubed it reached 63 TFLOP/s against Triton's 119 on the same GPU. The bottleneck was clear in hindsight: the loop stalled every iteration, loading a tile, waiting roughly 500 cycles for it to arrive, then computing. The tensor cores were idle most of the time.
Version 2: a cp.async pipeline. The standard fix is cp.async, which copies from global to shared memory in the background while computation proceeds, so you can prefetch the next tile while working on the current one. This ran into a real obstacle. newt generates code while tracing the user's loop, and the address of the next tile does not exist yet; it is produced later by code that has not executed. You cannot prefetch an address you cannot compute.
The solution was to invert the problem. Instead of prefetching the next tile, newt keeps a ring of tiles in flight and consumes the one it started several iterations earlier. The overlap is identical, but every address involved has already been computed, because it belongs to a past iteration, and this required no change to the user's loop. The depth of the ring, num_stages, is exposed as a tuning parameter, exactly as in Triton. This brought matmul to about 80 TFLOP/s.
Version 3: raw PTX, swizzling and fragment double-buffering. The last version replaced WMMA with the underlying ldmatrix and mma.sync instructions, written as inline PTX. Shared-memory tiles are stored swizzled, with each row's 16-byte chunks permuted by XOR-ing with the row index, which eliminates bank conflicts without the memory overhead of padding. Finally, the fragment registers are double-buffered within a k-step so that the loads for the next step are issued before the current multiply. That last change moved the cold-start number from 96 to 110 TFLOP/s.
The final numbers, stated honestly: 76 to 83 percent of Triton sustained, and about 92 percent when both run cold. The remaining gap is Triton's most specialized machinery, which I did not implement: strength reduction of the address arithmetic across iterations, and warp specialization into producer and consumer roles. These are well understood and deliberately out of scope; past a point they stop illustrating the ideas and become an engineering project of their own. One caveat on the measurements: the test machine is a 110-watt laptop that thermally throttles under sustained load, so the only fair comparison is same-run columns with a cooldown between suites, which is how these were taken.
deuteron: the autotuning layer
newt is the compiler; deuteron is the layer on top that removes the remaining kernel details, the way Helion does for Triton. You write tiles, with no program ids, offsets, masks or block sizes:
@dt.kernel
def matmul(x, y, out):
for tile_m, tile_n in dt.tile([x.shape[0], y.shape[1]]):
acc = dt.zeros([tile_m, tile_n], dtype=dt.float32)
for tile_k in dt.tile(x.shape[1]):
acc += x[tile_m, tile_k] @ y[tile_k, tile_n]
out[tile_m, tile_n] = acc
deuteron traces this. The outer loop becomes the launch grid, tensor indexing becomes pointer arithmetic and boundary masks, and @ becomes a fused tensor-core dot. It generates newt source, which you can print; the output is essentially the hand-written kernel from the tutorial.
The important design decision, taken from Helion, is how it tunes. The same function also runs as ordinary PyTorch, which gives a correct reference for free. During autotuning, each candidate configuration is run on cloned inputs and checked against that reference before it is timed. A configuration that compiles and runs but produces the wrong result is discarded before its performance is ever measured, so the tuner cannot select a fast but incorrect kernel.
Testing
A fast compiler that is occasionally wrong is not useful, so correctness got the majority of the effort. There are 176 tests comparing every operation against PyTorch, along with a large number of small targeted kernels aimed at the components most likely to fail subtly: synchronization, the layout arithmetic, and the swizzle. Every bug that was found became a regression test.
The most useful debugging feature was a single line: setting NEWT_DEBUG=1 prints the generated CUDA C++. Being able to read exactly what the compiler produced turned most debugging into reading a diff.
What I left out
Most of the difficulty in a project like this is deciding what not to build. NVRTC meant no register allocator. The group-cyclic layout meant no contiguity analysis. Deferred consumption meant no rewriting of the user's loop. The parts left unimplemented are the last few percent of matmul scheduling described above, plus in-kernel random number generation, device-side printing, non-NVIDIA backends, and fp8. None of them changes the point the project is meant to demonstrate: the modern GPU kernel stack, from tile-level Python down to tensor-core machine code, fits in about 4,000 lines once the essential problems are separated from the incidental ones.
Everything is open and MIT licensed, including a from-zero explainer for readers without a GPU background, the full benchmarks with their caveats, and a commit history where each stage of the matmul work is a separate commit.
- Install:
pip install nano-triton - Site: https://arpitsinghgautam.me/nano-triton/
- Code: https://github.com/arpitsinghgautam/nano-triton
If you build something on top of it, or find a scheduling improvement I missed, I would be glad to hear about it.
Credits and further reading
This project builds directly on other people's work, and the reading list is half the value of the writeup.
- Philippe Tillet and the Triton team at OpenAI: Triton defined the block programming model that newt copies, and the MAPL 2019 paper is still the best short read on why the block level is the right level to program at.
- Jason Ansel and the PyTorch compiler team: Helion is the tile-level layer deuteron miniaturizes, and the correctness-first autotuning idea is theirs.
- Simon Boehm: How to Optimize a CUDA Matmul Kernel is the classic worklog of the same compute-bound problem; if you want the CUDA-level version of the matmul section, read it next.
- Horace He: Making Deep Learning Go Brrrr From First Principles is the clearest explanation of the memory-bound versus compute-bound distinction used throughout this post.
- Sasha Rush: GPU Puzzles and Triton Puzzles, the hands-on way to learn the material above.
- Tri Dao: FlashAttention is the origin of fused attention; the fused-attention example in the repo is a homage to it.
- Mark Saroufim and the GPU MODE community: the lectures and Discord where much of this knowledge circulates.
- Andrej Karpathy: nanoGPT and llm.c set the precedent for small rebuilds that keep real performance.
- NVIDIA's documentation: the PTX ISA pages on mma and ldmatrix document the exact register mappings newt depends on, and NVRTC is the library that makes a 4,000-line compiler possible.
Top comments (0)