pytorch/helion

Helion – A Python‑embedded DSL for writing GPU kernels

What it is – Helion lets you write kernel code (the tiny programs that run on a GPU) directly in Python, using a syntax that feels like ordinary PyTorch. Under the hood the code is transformed into a Triton kernel, a low‑level language that compiles to fast GPU machine code.

Why it matters – Writing efficient GPU kernels is hard: you have to decide tile sizes, grid dimensions, memory‑access patterns, and many other low‑level knobs. Helion automates most of that work. A single Helion function is compiled to one GPU kernel, but before the first launch it runs an autotuner that tries hundreds of Triton variants and picks the fastest. The result is a kernel that is both easy to author and close to hand‑tuned performance.

Key features

• Python‑first API – Define a kernel with @helion.kernel() and use normal PyTorch tensors inside. The outer Python code runs on the CPU (allocation, shape checks), the inner loops become GPU code.
• Automatic tiling & grid selection – hl.tile splits the iteration space; Helion decides tile dimensions, loop ordering, and grid size.
• Rich autotuning – Searches a large configuration space (block sizes, loop orders, indexing strategies, warp specialization, etc.) and can cache the best config for later runs. You can also supply your own configs to skip the expensive search.
• Memory‑access strategies – Chooses between pointer‑based, block‑pointer, or TensorDescriptor (TMA) indexing automatically, with the ability to override.
• Support for reductions, masking, and persistent kernels – Handles common patterns like softmax, sum, or large reductions without extra code.
• Back‑ends – Primarily Triton, but also experimental TileIR (for NVIDIA Blackwell GPUs) and a CuTe backend (CUDA 13+).
• Debug & dev tools – Print generated Triton code, emit reproducible scripts, and run kernels in an interpreter mode for step‑by‑step debugging.

Typical workflow

1. Write a kernel function using @helion.kernel() and hl.tile.
2. Call the function with CUDA tensors. The first call triggers autotuning (≈10 min for a matrix‑multiply example).
3. Optionally extract the best helion.Config and hard‑code it to avoid future tuning.
4. Deploy the kernel with the saved config for predictable start‑up time.

Who should use it – Researchers and engineers who need custom GPU kernels but don’t want to hand‑write Triton or CUDA. It fits nicely into PyTorch codebases and can accelerate workloads where existing libraries lack a specialized implementation.

Installation – pip install helion (requires Python 3.10‑3.14, PyTorch ≥ 2.9, Triton ≥ 3.5). Development setup uses the uv tool and optional extras for the CuTe backend.

Community – Discussions happen on the GPU MODE Discord (#helion). The project follows a governance model and is BSD‑licensed.