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

I spent the last while rebuilding Triton and Helion, the two-layer stack that powers most custom GPU kernels in modern ML, from scratch. Not as a simulator and not as a toy: the rebuild JIT-compiles Python to real tensor-core machine code, ties real Triton on memory-bound kernels, and reaches about 92% of its cold matmul throughput, in roughly 4,000 lines of Python you can read in an afternoon.

This post is the story: the architectural bet that makes a nano possible, the one layout rule that buys most of the performance, the matmul pipeline war, and the autotuner that cannot ship a wrong kernel.

Everything below is open and MIT licensed: the full from-zero explainer and the code on GitHub.

Why rebuild something that already exists

Triton (from OpenAI) changed who can write fast GPU kernels: you program a block of data instead of individual threads, and the compiler handles thread mapping, memory coalescing, shared memory and tensor cores. Helion (from PyTorch) sits one level higher, generating and autotuning Triton kernels from PyTorch-like tile code.

Both are superb tools. They are also industrial compilers with hundreds of thousands of lines, MLIR passes, and years of accumulated scheduling wisdom baked in. Reading them taught me very little, for the same reason reading a city map teaches you little about urban planning.

The better way is to rebuild the system small: keep the architecture, shrink the surface, keep the performance real, and the ideas become yours. I wanted that for the kernel stack. So the project has two parts, mirroring the production stack layer for layer:

• newt, a nano-Triton (Triton was the original genus name for newts, so a newt is literally a small triton).
• deuteron, a nano-Helion (a deuteron is a lighter atomic nucleus than a helion).

A newt kernel is a Triton kernel with the serial numbers filed off. If you know tl.*, you already know nl.*:

@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 cheat code: NVIDIA ships a compiler as a library

Triton lowers Python through its own IR, MLIR and LLVM down to PTX. A nano cannot carry that, and it turns out it does not need to. NVRTC (NVIDIA Runtime Compilation) is a full CUDA C++ compiler shipped as a shared library, callable in-process.

So newt's pipeline is: Python AST, then typed block values (every value gets a shape, dtype and layout), then a CUDA C++ source string, then NVRTC, then a cubin, loaded and launched through the raw CUDA driver API via ctypes. No nvcc subprocess, no MLIR, no wrapper packages, no dependencies beyond torch.

This is the bet that makes the whole project feasible: NVRTC does register allocation and instruction scheduling, which are the person-decades parts of a compiler. What remains is the one problem that makes Triton Triton: translating block semantics into thread semantics well. That problem fits in about 2,500 lines.

One rule for where data lives

When a kernel says "load a block of 1,024 floats", where do they go? In registers, spread across the threads of the block by a single fixed rule I call the group-cyclic layout: consecutive elements are dealt to threads in 16-byte groups, round robin.

One rule, three wins:

1. Coalescing by construction. Adjacent threads always hold adjacent memory, so every warp access collapses into one wide transaction.
2. Vectorization. Each thread owns 16 contiguous bytes, so a load can be a single 128-bit instruction. Here newt diverges from Triton on purpose: Triton proves contiguity with static analysis; newt emits a tiny runtime check per group ("are these offsets consecutive, aligned, all unmasked?") and branches to the fast path. A few integer compares against a ~500-cycle memory access is a wonderful trade, and this single trick took vector-add from 82% of Triton to parity.
3. Simplicity. Elementwise math never cares where elements live; each thread loops over its own slots.

Reductions follow the classic three-stage shape (registers, then a __shfl_xor_sync butterfly within each warp, then shared memory across warps), and broadcasts stage through a reusable shared-memory arena.

The easy half: hitting the memory roof

For memory-bound kernels, that layout rule is essentially the whole game. Softmax reads and writes each byte once; whoever saturates DRAM bandwidth wins, and coalesced plus vectorized plus fused gets you there. Any correct compiler that manages those three ties. This is the roofline model in action:

Parity is the expected result, and that is the point: it says the compiler is not leaving bandwidth on the table, and it points at where the real fight is.

The hard half: the matmul war

Matrix multiplication is compute-bound. The game flips from "save bandwidth" to "never let the tensor cores stall", and every scheduling imperfection shows up in the TFLOP/s. newt's dot went through three generations, each one a readable commit:

Generation 1: WMMA plus synchronous staging (63 TF sustained at 4096^3). NVIDIA's high-level WMMA API drives the tensor cores but hides the register layout and costs extra shared-memory round trips. Worse, the loop stalls every iteration: load tile, wait ~500 cycles, compute, repeat.

Generation 2: the cp.async ring (80 TF). Modern GPUs can copy global memory to shared memory in the background with cp.async, without the data passing through registers. The textbook use is prefetching future tiles, but a compiler tracing a user's loop cannot do that: future tiles' addresses have not been computed yet.

The fix is my favorite idea in the project: invert the problem and delay consumption instead of prefetching. Keep an N-slot ring of tiles in flight; each iteration consumes the tile staged S-1 iterations ago, then starts this iteration's background copy. Same overlap, one barrier per k-step, no loop rewriting, and the leftover tiles flush wherever the accumulator is next read. The number of slots is the same num_stages knob real Triton exposes, and the autotuner searches over it.

Generation 3: raw PTX, swizzle and fragment ping-pong (82 TF sustained, 110 TF cold). WMMA went away: fp16/bf16 dots now compile to raw ldmatrix and mma.sync.m16n8k16 inline PTX over shared memory whose 16-byte chunks are XOR-permuted by row index, which kills bank conflicts without padding. Because NVIDIA documents the accumulator's per-lane register layout, newt can convert fragments to and from the normal layout whenever the kernel does elementwise math on them; that conversion is what makes a fused flash-attention kernel expressible.

The final piece: within one k-step, the fragment loads and the mma math form their own dependency chain. Double-buffering the fragments, so step k+1's ldmatrix issues before step k's math, took the cold-start number from 96 to 110 TFLOP/s.

The honest summary: 76-83% of Triton sustained on my thermally limited 110 W laptop, and 109.7 TFLOP/s cold, about 92% of Triton's own cold numbers. The remaining gap is Triton's finest-grained machinery (per-iteration address strength reduction and warp specialization), known, documented, and out of scope for a nano on purpose. Benchmarks on a laptop deserve their caveats spelled out: the numbers above come from same-run columns with cooldown between suites, because that is the only fair comparison on a machine that throttles.

deuteron: kernels that check themselves

One level up, deuteron is ~700 lines doing to newt what Helion does to Triton. You write tile-level code with zero kernel details:

@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 the AST: the outer dt.tile loop becomes the launch grid, indexing becomes pointer math plus boundary masks, @ becomes a tensor-core dot fused into the accumulator. The output is a newt kernel source string in which every tile size is a compile-time constant, and printing it is the most satisfying demo in the repo, because it looks exactly like the hand-written tutorial kernel.

The design decision worth stealing is Helion's oracle trick, replicated here: the same function also runs as plain PyTorch, which gives ground-truth outputs for free. During autotuning, every candidate config runs on cloned inputs and its output is compared against the oracle; a config that compiles, runs, and computes garbage is rejected before it is ever timed. The autotuner cannot ship a wrong kernel, by construction.

Trust before speed

A compiler that is fast and wrong is worse than useless. The suite sits at 176 tests comparing every operation against PyTorch references, and the tricky corners (sync hazards, layout algebra, swizzle consistency, pipeline wait-counting) got hammered with hundreds of small targeted GPU programs. Whatever broke became a regression test.

The single most valuable debugging feature cost one line: NEWT_DEBUG=1 prints the generated CUDA source. Being able to read what the compiler wrote turns most bug hunts into text diffs.

What I took away from this

Some of these I expected going in. Most I did not.

• NVRTC changes what a small compiler can be. Register allocation and instruction scheduling are the parts that take person-decades, and it turns out you can just rent them. The lines you actually write should go into the semantics gap.
• One good layout rule quietly replaces a whole subsystem. Once elements are dealt out group-cyclically, coalescing and vectorization stop being analysis problems.
• A runtime check costing a few integer ops is a fair price for skipping static analysis, as long as it guards something as slow as a DRAM access.
• Deferred consumption was the best idea in the project. I could not prefetch tiles whose addresses did not exist yet, but delaying the math by S-1 iterations bought the same overlap without touching the user's loop.
• Getting softmax to tie Triton felt like a milestone until I understood it is a floor: any correct compiler saturates the bus there. Matmul is where you find out whether your scheduling is any good.
• The eager oracle is the reason I trust my own autotuner. Filter configs on correctness first and the search can be as aggressive as you like.
• On a laptop, a benchmark number without its thermal context is a lie. I ended up trusting only same-run columns, with cooldown between suites.

Go read it

The whole thing is built to be read: a from-zero explainer that assumes no GPU background, a README with the architecture diagrams, benchmarks with the methodology spelled out, and a git history where each compiler stage is one self-contained commit.

• Site: arpitsinghgautam.me/nano-triton
• Code: github.com/arpitsinghgautam/nano-triton

If you build something on top of it, or spot a scheduling trick I left on the table, I want to hear about it.

Credits and further reading

This project stands on other people's work, and the reading list is half the value of the write-up:

• Philippe Tillet and the Triton team at OpenAI: Triton defined the block programming model that newt copies; the MAPL 2019 paper is still the best short read on why the block level is the right level.
• 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 compute-roof fight; if you enjoyed the matmul section, read it next.
• Horace He: Making Deep Learning Go Brrrr From First Principles is the cleanest telling of the memory-bound vs compute-bound framing used throughout this post.
• Sasha Rush: GPU Puzzles and Triton Puzzles, the hands-on way to learn everything above.
• Tri Dao: FlashAttention is where fused attention comes from; the repo's fused-attention example is a homage.
• Mark Saroufim and the GPU MODE community: the lectures and Discord where kernel knowledge actually circulates.
• Andrej Karpathy: nanoGPT and llm.c set the tone for minimal rebuilds that keep real performance.
• NVIDIA docs: the PTX ISA sections on mma and ldmatrix document the exact register mappings newt relies on, and NVRTC is the library that makes a 4,000-line compiler possible at all.