Helion: a Python embedded DSL that compiles to Triton

PyTorch with tiles, one level above Triton

Helion is described as a Python embedded domain specific language for authoring machine learning kernels, designed to compile down to Triton. The readme offers two readings of what that means, and both are useful. Helion can be viewed as PyTorch with tiles, or as a higher level Triton.

The problem it solves is the gap between PyTorch, where you compose tensor operations and the compiler generates kernels, and Triton, where you write a kernel yourself and are responsible for the details that make GPU code fast or correct: how memory is indexed, which elements are masked off, how program ids map to tiles of data, how large each block is, how the reduction loop is structured. Writing those by hand is exactly what makes Triton powerful and exactly what makes it slow to develop.

Helion keeps the PyTorch programming model and moves those decisions into an autotuner. The readme is explicit about the goal, which is to raise the level of abstraction compared to Triton so that it is easier to write correct and efficient kernels while enabling more automation in the tuning process.

The name is a small joke worth knowing, since it explains the pairing. Helion refers to the nucleus of a helium-3 atom, and Triton refers to hydrogen-3. The project sits under the pytorch organisation, and the description in pyproject.toml is simply a Python embedded DSL that makes it easy to write ML kernels.

Code outside the tile loops runs on the CPU

The mechanism is easiest to understand from the example the readme gives, a minimal matrix multiplication.

import torch, helion, helion.language as hl

@helion.kernel()
def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    m, k = x.size()
    k, n = y.size()
    out = torch.empty([m, n], dtype=x.dtype, device=x.device)

    for tile_m, tile_n in hl.tile([m, n]):
        acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
        for tile_k in hl.tile(k):
            acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n])
        out[tile_m, tile_n] = acc

    return out

The split is the key. Code outside the `for` loops is standard PyTorch code executed on the CPU, and it is typically used for tasks like allocating output tensors and performing shape computations. Code inside the loops is compiled into a Triton kernel, resulting in a single GPU kernel, and a single Helion kernel is always compiled to exactly one GPU kernel.

That constraint is a guarantee rather than a coincidence. A Helion kernel maps onto one Triton kernel, so the outer `for` loop becomes the grid of the generated kernel and the inner loop becomes a loop within it.

The `hl.tile` function subdivides the iteration space, here `m` by `n`, into tiles that execute in parallel on the GPU. Tiling details such as dimensionality, tile sizes and loop ordering are determined automatically by the autotuner, and they can also be specified explicitly using the `config=` argument in `helion.kernel`.

The autotuner evaluates over a thousand Triton variants

Autotuning is not a light touch. The readme says Helion spends more time, approximately ten minutes, autotuning as it evaluates hundreds of potential Triton implementations generated from a single Helion kernel. The same passage notes that the larger search space also makes kernels more performance portable between different hardware.

The session output shows the shape of it. A search starts with a differential evolution search using a population of 40 and 20 generations, and it reports the best configuration found so far along with its minimum, median and maximum timings. By the end, the example run reports autotuning complete after searching 1520 configs.

The configuration itself is dense, and reading one tells you how much is being decided for you: block sizes for each loop dimension, loop orders, L2 groupings, range unroll factors, warp specialisation flags, number of pipeline stages, multi buffering, flattening, the number of warps, the indexing style, and the program id type. The indexing style alone appears as a choice between pointer, block pointer and TensorDescriptor styles, and the program id type changes between flat and persistent blocked mappings.

That search space is the entire value proposition of the project. A human writing Triton by hand would choose a handful of these by reasoning and measurement, which is slow, incomplete and hardware specific.

Indexing, masking and grid mapping are decided for you

Seven areas are listed as automated, and they are worth reading individually because they are the concrete difference between Helion and hand written Triton.

Tensor indexing is handled automatically: strides and indices are calculated for you, and Helion autotunes among indexing methods such as pointers, block pointers and TensorDescriptors, with support for per operation indexing strategies for fine grained control of loads and stores. Masking is mostly implicit and, the readme says, optimized away when not needed. Grid sizes and program id calculations are automatic too, and Helion autotunes multiple mappings from program ids to data tiles.

The implicit search space definition removes the need to define search configurations manually, because Helion generates configuration flags and exploration spaces. Kernel argument management is automated as well, including tensor sizes and strides, and global variables and nested closures are lifted into kernel arguments, which allows better templating.

Two more capabilities sit outside the simple automation list. Large reductions can be automatically converted into looped implementations, so a reduction that would exceed some limit becomes a loop instead of a special case. And there are automated optimisations: PID swizzling for improved L2 cache reuse, loop reordering, persistent kernel strategies, and warp specialisation choices with unrolling.

Underneath, standard PyTorch operators inside a kernel are mapped to Triton operations using TorchInductor, which is why the readme says that familiarity with PyTorch means you already know most of Helion.

Running the kernel means paying for the search once

Calling the kernel looks like ordinary PyTorch code.

out = matmul(torch.randn([2048, 2048], device="cuda"),
             torch.randn([2048, 2048], device="cuda"))

The first time a kernel runs, Helion initiates autotuning, and that is where the ten minutes go. Subsequent runs of the same kernel reuse the configuration that was chosen.

The readme notes that the best config can be hardcoded to skip autotuning, using the `config=` argument with a `helion.Config` object, and the same `config=` argument is how you pin tiling details explicitly. For a deployment pipeline this matters: autotuning belongs in a benchmarking step, and the selected configuration should be committed alongside the kernel.

The repository gives you the material to do that without starting from scratch. There is a `pretuned_kernels/` directory at the top level, an `examples/` directory holding kernels such as attention, softmax variants, fp8 GEMM and flex attention, a `benchmarks/` directory, and a `notebooks/` directory. The readme also links a Colab notebook for the softmax example and an AMD DevCloud entry, which are the fastest ways to try the language without local hardware setup.

For installation, the readme points to the documentation site rather than giving commands. pyproject.toml shows the packaging details: the project name is `helion`, it requires Python 3.10 or newer, and the build backend is hatchling with hatch-vcs, so the version is derived from the git tags rather than written by hand.

Where Helion is the wrong tool

Three boundaries are worth weighing before you commit.

The first is the autotuning cost. Roughly ten minutes on the first run of every new kernel is a real expense in a test suite or a CI pipeline that runs on every commit, and the readme's answer is to hardcode the configuration. That is a workflow change, not a free default.

The second is the supported operation set. Helion supports pointwise operations such as add and sigmoid, reductions such as sum and softmax, views, and matrix multiplication operations. Arbitrary function calls within a kernel are supported, but must be traceable with `make_fx`. That is the constraint that decides whether a particular custom operation can move into a Helion kernel at all, and the failure mode is a traceability error rather than a slow kernel.

The third is the backend story. Helion compiles to Triton, and the optional dependency groups show how far the surface extends. There are `cute` extras pinning `nvidia-cutlass-dsl[cu13]==4.7.0` together with `apache-tvm-ffi`, and `pallas-tpu` and `pallas-cpu` extras built on `jax` and `absl-py`. The comment in pyproject.toml is blunt about the cute backend: Helion integrates tightly with the cutlass DSL internal APIs, each upstream release tends to break the cute backend, and the pin tracked in `.github/ci_commit_pins/nvidia_cutlass_dsl.txt` is the single source of truth. That is an explicit statement that the extra backends are pinned tightly and break on upstream releases.

There is also the question of hardware. The Dockerfile is built on `nvidia/cuda:13.2.0-cudnn-devel-ubuntu24.04` and installs torch and triton from the PyTorch nightly index for cu130, so the reference environment is a recent NVIDIA CUDA stack.

How Helion differs from writing Triton by hand

The alternative for a kernel author is Triton itself, or the existing collection of hand tuned PyTorch kernels, or TorchInductor generating code for you without you writing a kernel at all.

Against hand written Triton the difference is who makes the decisions. In Triton you choose block sizes, write the index arithmetic with strides and masks, decide how program ids map to tiles, and select the number of warps and stages, then measure and iterate. In Helion you write the tiled loop in PyTorch syntax and the autotuner makes those choices by evaluating hundreds of generated variants. The price is the ten minute search and the loss of direct control over the generated code.

Against letting a compiler generate the kernel, Helion occupies the middle. Inductor generates kernels for you but does not give you a place to express an algorithm the way a Helion tile loop does, and it cannot fuse a reduction with a specific access pattern the way you can inside a single tile loop. Helion's claim is that it keeps the expression of the algorithm while removing the mechanical parts of the kernel.

The portability argument follows from the same place. Because the search space is wide and the tuner evaluates it on the actual device, the chosen configuration is tuned for the hardware in front of it, which the readme presents as making kernels more performance portable between different hardware. The flip side is that the best configuration differs per device, so a kernel tuned on one GPU is not the same kernel as the one tuned on another.

Editorial conclusion

Helion suits kernel authors who know PyTorch and want tiled GPU code without hand tuning strides, masks and program id mapping, and it does not suit teams that need a kernel to run immediately, that depend on operations outside the supported set, or that must target backends other than Triton without extra work. Before adopting it, check that your operations are pointwise, reductions, views or matrix multiplications, budget roughly ten minutes of autotuning per new kernel, and pin the `nvidia-cutlass-dsl[cu13]==4.7.0` extra to the same commit the CI tracks, since each upstream release breaks the cute backend.

Frequently asked questions

Official sources

1. Issues
2. License: BSD-3-Clause
3. pytorch/helion on GitHub
4. README
5. Releases