Fast by Default: Compile-Time Heuristics for Helion Kernels

TL;DR Helion now uses compile-time heuristics that analyze a kernel’s IR and dataflow to pick fast settings without slow autotuning. It classifies the kernel as matmul, reduction, or pointwise and chooses block sizes, warps, and stages, giving near-autotuned speed for free. Introduction Helion, PyTorch’s domain-specific language (DSL) for performance portable machine learning kernels, relies heavily on autotuning for its performance. Autotuning time can be a major bottleneck: for example, a previous team tuning a single paged-attention kernel across 72 shapes spent 10 hours in Helion’s quick mode, and 25 at full effort. To address these issues, we’ve developed a set of heuristics that read a kernel’s IR and dataflow at compile time and emit a config directly, without any autotuning . We classify kernels into broad families based on the operations they contain – matmul, reduction, or pointwise – and apply a family-specific heuristic to each. These families are intentionally general: the matmul heuristic, for example, looks for any kernel containing a matmul operation such as hl.dot or torch.baddbmm, not just standard m-n-k matrix multiplications. These are designed to be generic and apply broadly to custom kernels that users may write. Highlights: Autotuner Heuristics. Our heuristics cover a large set of kernels, including GEMM, reductions, and even linear attention. Linear attention. Over 96 kernel-and-shape cells spanning nine linear-attention variants, the computed config beats the hand-written Triton in Flash Linear Attention (FLA) by 1.18x in the forward pass, and by 1.12x with backward included on H100. Overall the heuristic is 2.13x faster than Helion’s previous no-autotune default on H100. vLLM reduction kernels. Over 153 kernel-and-shape cells across six production kernels on H100, the computed config is 1.27x faster than vLLM’s own handwritten CUDA/C++ operators, and 2.2x faster than Helion’s previous default. The analysis itself is cheap. Compiling a kernel with the heuristics on versus off differs negligibly (tens of milliseconds), which is lost in the noise of Triton compilation. Our Approach: Decide the Config at Compile Time A single Helion kernel can lower into thousands of different Triton kernels. One definition covers block sizes, loop orders, indexing modes and pipelining strategies, but the autotuner has to choose which of those thousands to emit, by compiling and benchmarking many candidates. And this autotuning cost is per kernel and per shape , which multiplies the problem the more kernels, shapes, and hardware backends you want to support. We compute the config at compile time instead of searching for it at tuning time. We read the device IR and record structural facts about it: which axes are tiled, which block sizes are tunable, which tiles are live at which part of the kernel, etc. A heuristic reads those facts and emits a config. The point here is not to come up with a better config than the autotuner, which has the benefit of benchmarking potentially thousands of configs before selecting one. It is to add a point in the tradeoff space between performance and wait time. Our heuristics serve two purposes. A primary heuristic emits the default config used when autotuning is disabled. When autotuning is enabled, helper heuristics generate roughly 5–20 related or alternative configs that seed the autotuner. The goals are different: the default config should perform well on its own, while the seeds should give the autotuner several promising starting points to speed up autotuning. In this blog, we’ll touch on both, but focus on the primary heuristic that gets emitted with autotuning off. How the Heuristics Work A heuristic sets the values for various config knobs, which must be tuned carefully to achieve high performance and are normally picked by the autotuner. These include block sizes for each tiled axis, num_warps, num_stages, the L2 grouping, load eviction policies, and more. We pick a heuristic by scanning the device IR for the operations a kernel contains. Any matmul routes to the matmul heuristic regardless of what else is in the kernel; a reduction with no matmul routes to the reduction heuristic; everything else receives the pointwise heuristic. If a kernel’s extents are not all known at compile time, no heuristic fires and Helion falls back to the original default. Our heuristics all lean on the same underlying analysis, which estimates resource usage of the kernel. At any step of the kernel, the analysis determines: given a candidate set of block sizes, how many registers, shared memory, and — on Blackwell — how much tensor memory, one CTA needs at that moment. To do this, we run a liveness analysis on the compiler IR and walk each graph to find where every value is last used, and guess where it will be placed in memory. This is a hard problem: we are trying to predict what Triton will keep in registers, what it will stage through shared memory, what it will materialize at all rather than fuse away. Sometimes it is easier to predict (e.g., we assume the output of an hl.dot will be placed in TMEM on B200), but the estimate is approximate. Using this we can estimate things that matter for performance, such as: Occupancy : how many CTAs fit on an SM at once. Whichever space runs out first decides it, e.g., a CTA needing a tenth of the SM’s registers and a tenth of its shared memory, but six tenths of its tensor memory, gets one CTA per SM, since two CTAs would go over the tensor memory budget by 20% (0.6 x 2 = 1.2). Register pressure: Peak live registers spread over the threads we intend to launch tells us whether the assembler will be forced to spill. These considerations inform the selection of different config knobs. We’ll illustrate this for the matmul heuristic : block_size: If the launch produces fewer CTAs than the GPU has SMs, we shrink block size on a grid axis to produce more of them. num_stages: More stages hide more load latency, so we want as many as shared memory allows. But stages can also accidentally lose occupancy: six of them might fit while eating the CTA’s entire shared-memory budget, leaving one CTA per SM, where three stages and two CTAs per SM is the better trade. num_warps. Two things pull on this: Which instruction Triton emits. Below four warps, or with an M block under 64, Triton uses the smaller synchronous MMA; otherwise it uses the big high-throughput one (wgmma on Hopper, tcgen05 on Blackwell). Usually we want the big one. But when there is little matmul work sometimes the better answer is many CTAs per SM running small MMAs that overlap each other’s loads. So we count dynamic matmul work (essentially counting the number of times the hl.dot will be executed dynamically) and choose based on that. Register pressure. Warps divide the tile across more threads, so more warps means less pressure on each one. If the peak live state at the current count would spill, that argues for increasing warps. Feasibility Check. A final pessimistic pass over the whole kernel: add up tensor memory, shared memory and registers, and confirm the config can actually launch. The adjustments above already respect these budgets, so this is mostly defensive. In the case where there are multiple matmuls in one kernel (for example for attention and linear attention), we run the matmul heuristic on each hl.dot and use those as the starting point. They may disagree: two dots can share one block-size knob — one dot’s N is frequently the next dot’s K — and there is only one num_warps for the whole kernel. To resolve these tradeoffs, we rank each dot by estimated importance: priority goes to any dot with a loop-carried accumulator, since that accumulator stays resident for the whole loop and its tile sets the kernel’s footprint; then its dynamic matmul work; then its own output area. Wherever two dots disagree about a knob, the higher-ranked one wins (for example num_warps is set by the top-ranked dot alone). The reduct