Rebuilding Triton and Helion from Scratch in 4,000 Lines of Python

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 is the long version of 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). It is on PyPI, so pip install nano-triton gives you both; you need torch and an NVIDIA GPU, since newt compiles the kernels at runtime with NVRTC. 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)

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. Compiled binaries are cached in memory and on disk, keyed by the compile-time constants, dtypes, num_warps and num_stages, so the second call and the second process both pay nothing. Setting NEWT_DEBUG=1 prints the generated CUDA C++, which was by a wide margin the most useful debugging feature in the project.

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. The exact mapping is:

element i  ->  thread (i / VEC) mod T,  slot (i / (T * VEC)) * VEC + i mod VEC

where T is the number of threads and VEC is how many elements fit in 16 bytes (4 floats, or 8 halves). In words: consecutive elements are handed to threads in contiguous 16-byte groups, round robin. A worked example with 8 elements, 2 threads and VEC = 2:

element:  e0 e1 e2 e3 e4 e5 e6 e7
thread:   t0 t0 t1 t1 t0 t0 t1 t1

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, falling back to a predicated scalar path when it does not. 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.

Reductions and broadcasting

nl.sum and friends have to combine values that live in different threads, and newt lowers them in three stages. Each thread first reduces its own slots in registers. Then each warp reduces its 32 lanes with a __shfl_xor_sync butterfly, where lanes exchange values directly through the register file with no memory involved, in five steps for 32 lanes. Finally one value per warp is written to shared memory and a single warp finishes the cross-warp reduction and broadcasts the result back. This is the same three-level shape Triton uses.

Broadcasting between different-sized blocks sometimes needs a value held by one thread to become visible to another. newt stages the smaller operand through a reusable shared-memory scratch area with barriers on either side. Reshapes that only add or remove size-1 dimensions are free, because they change metadata and never move data.

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, and the ordering is slightly different depending on depth: at two stages newt issues the copy first and uses two barriers, and at three or more it consumes first and gets away with a single barrier per k-step. Any tiles still in flight at the end of the loop are flushed at the point where the accumulator is first read. 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.aligned.m16n8k16 instructions, written as inline PTX. Two details did most of the work here. First, the shared-memory tiles are stored swizzled: each row's 16-byte chunks are permuted by XOR-ing the chunk index with the row index, which eliminates the bank conflicts that ldmatrix would otherwise hit, and does it without the memory overhead of padding. Second, the accumulator lives in registers with the specific per-lane layout NVIDIA documents for the mma instruction, which means newt can convert it to and from the normal group-cyclic layout whenever the kernel does elementwise math on it. That conversion is what makes a fused flash-attention kernel expressible at all. Finally, the fragment registers are double-buffered within a k-step, so the ldmatrix loads for the next step are issued before the current multiply runs. That last change, a few lines, moved the cold-start number from 96 to 110 TFLOP/s.

The remaining gap to Triton is its 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. fp32 matmul, which runs on the tensor cores as tf32, still uses the older WMMA path and sits around 45 percent of Triton; porting it to the raw PTX path is mechanical work I have not done yet.

The full numbers

One caveat before the tables: the test machine is a 110-watt laptop that thermally throttles under sustained load, so the only fair comparison is same-run columns, taken with a cooldown between suites. fp16 matmul, in TFLOP/s, with each newt row corresponding to one of the three versions above:

Memory-bound kernels, in GB/s:

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, and every tile size becomes a constexpr. You can print the generated kernel with matmul.to_newt_source(x, y, out), and the output is essentially the hand-written kernel from the tutorial, which is the most satisfying thing in the repo to look at.

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 (block sizes, num_warps, num_stages) 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. Winners are cached to disk, keyed by the kernel source, a shape bucket and the dtypes, so the next call with similar shapes launches immediately.

One detail that fell out nicely: masks propagate through the traced expressions, so a reduction on a padded tile automatically gets the right identity value (negative infinity for max, zero for sum). The layernorm example computes a correct variance on rows of length 1,500, which is not a power of two, with no explicit mask handling anywhere in the user code.

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 CI runs a GPU-free subset that checks the generated CUDA source structurally, so the compiler is exercised on every push even without a GPU on the runner.

Trying it yourself

pip install nano-triton

Then the shortest interesting thing to run is the deuteron matmul, because it prints the kernel it generated:

import torch
import deuteron as dt

@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

x = torch.randn(512, 512, device="cuda", dtype=torch.float16)
y = torch.randn(512, 512, device="cuda", dtype=torch.float16)
out = torch.empty(512, 512, device="cuda", dtype=torch.float16)

matmul(x, y, out)                          # traces, autotunes, caches, launches
print(matmul.to_newt_source(x, y, out))    # the generated newt kernel

The repository has the full example progression, from vector add up to a fused flash-attention kernel, and python -m pytest tests runs the suite (the GPU tests self-skip if there is no CUDA device).

What is supported

newt covers program_id, arange, masked load/store, the full set of arithmetic, comparison and bitwise operators with numpy-style broadcasting, where, maximum, minimum, fma, the usual transcendentals (exp, log, sqrt, rsqrt, sin, cos, tanh, erf, sigmoid and so on), full and axis reductions (sum, max, min), dot with an accumulator, dtype casts, reshape, trans and broadcast_to, atomic_add and atomic_max, and for / while / if with constexpr pruning. Dtypes run from fp16, bf16, fp32 and fp64 through the integer and unsigned types and bool. Grids go up to 3D, num_warps from 1 to 32, num_stages from 1 to 8, and @autotune / @heuristics work as they do in Triton.

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, the tf32 port to the raw PTX path, 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: arpitsinghgautam.me/nano-triton
• Code: 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.