Helion Cuts GPU Kernel Tuning Time with Compile-Time Heuristics

Introduction

GPU kernel performance often depends on a surprisingly large configuration space. Block sizes, warp counts, pipeline stages, memory policies, loop order, and related choices can all change the result. Exhaustively testing those options may find a strong configuration, but it also makes compilation slow. Helion, PyTorch’s domain-specific language for performance-portable machine-learning kernels, has historically relied heavily on autotuning. Its new compile-time heuristics are intended to provide a faster middle ground: inspect the program structure, predict the important hardware constraints, and emit a useful configuration directly.

What the new heuristics do

• Classify the kernel by its operations. A kernel containing a matrix multiplication is sent to the matmul heuristic. A kernel without matmul but with a reduction uses the reduction heuristic, while remaining kernels use the pointwise heuristic. The matmul category is deliberately broad and can include operations such as hl.dot or torch.baddbmm, not only conventional GEMM layouts.
• Read the compiler IR and dataflow. Helion records which axes are tiled, which dimensions are configurable, and when tiles remain live during execution. A liveness analysis tracks when values are last used and estimates where intermediate data is likely to reside.
• Model hardware pressure. For candidate block sizes, the analysis estimates registers, shared memory, and, on Blackwell, tensor memory used by one CTA. These estimates feed decisions about occupancy, register spilling, and whether a configuration can launch at all.
• Balance competing knobs. More stages can hide memory latency but consume shared memory. More warps can reduce per-thread register pressure and enable larger matrix-multiply instructions, yet small workloads may benefit from more CTAs using smaller instructions. If a grid would create fewer CTAs than the GPU has SMs, the heuristic can reduce a tile to expose more parallelism.

Results and broader significance

The analysis adds only tens of milliseconds in the reported setup, a cost that is generally lost in Triton’s compilation time. The reported H100 evaluation covered more than 96 kernel-and-shape cells across nine linear-attention variants. Relative to the hand-written Triton implementation in Flash Linear Attention, the generated configuration was 1.18 times faster for forward execution and 1.12 times faster when backward execution was included. It was 2.13 times faster than Helion’s earlier no-autotune default.

A second evaluation covered more than 153 kernel-and-shape cells across six production reduction kernels from vLLM. There, the heuristic configuration was 1.27 times faster than the project’s hand-written CUDA/C++ operators and 2.2 times faster than Helion’s previous default. These figures do not mean that static rules universally beat exhaustive benchmarking; autotuning can still search a much wider space. The key benefit is a better performance-versus-wait-time tradeoff.

The design also supports both modes. With autotuning disabled, the primary heuristic supplies the default configuration. With autotuning enabled, helper heuristics produce roughly 5–20 related candidates to seed the search. For kernels with multiple matrix multiplications, Helion ranks operations by factors such as loop-carried accumulators, dynamic work, and output size before resolving conflicts over shared block-size and warp settings. This approach points toward compiler optimization that uses program structure and hardware resource models to reduce trial and error without abandoning autotuning when it is worthwhile.

Source: PyTorch Blog