PyTorch2d agoenergy 52

One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ

image: PyTorch

With Helion’s release last year, it aimed at enabling developers to write GPU kernels in PyTorch-native Python and let its powerful autotuner do the work that used to require hand-written Triton and abstract away hardware-specific implementation details. Helion’s autotuner explores thousands of kernel configurations and routinely lands within a few percent of expert-tuned code, and oftentimes exceeding it. The question that now needs to be answered is this: If a kernel has already been autotuned and optimized to the best of a developer’s ability, is there anything left? There is. It lives below the kernel, in the compiler. Every kernel autotuner in the ecosystem whether it is Helion’s, Triton’s, the one inside torch.compile, they all search the same thing: tile shapes, warp counts, pipeline stages, memory access strategy. All of those are decisions about what work each thread block does. Underneath them sits a second set of decisions about how that work becomes machine assembly instructions generated by ptxas, NVIDIA’s PTX assembler. Things like register allocation policy, fine-grained instruction scheduling, code-generation heuristics are left at their defaults for all kernels shipping today – not because the defaults are optimal, but because nothing was searching and optimizing them! Helion now enables this, and this post is about what that second compiler-specific space contains, why searching it jointly with the kernel configuration space shows further improvements on kernels that were already deemed to be optimal. The searching is done by NVIDIA’s CompileIQ, which optimizes ptxas controls the way an autotuner optimizes tile sizes and packages the result as a portable Advanced Control File (ACF). Most users won’t run that search themselves: CompileIQ publishes pre-tuned ACF catalogs that Helion’s autotuner can now evaluate alongside the kernel configurations it already explores. Starting off with highly-optimized kernels & optimizing the last-mile It matters a great deal what baseline one compares against. A “we made kernels faster” claim is easy if the starting point is a naive kernel, and close to meaningless for real-world workloads. The kernels we use here come from the pretuned_kernels/ in the Helion repository, roughly 27 kernels covering the shapes that show up in real model code: FP8 and NVFP4 quantization, RMSNorm variants, RoPE, grouped GEMM, causal depthwise convolution, Gated DeltaNet decode, and fused model regions like a Qwen3 decode layer. The important property to note here is that these are not autotuned on the fly. Each one ships with ahead-of-time heuristic files per architecture: _helion_aot_rms_norm_cuda_sm90.py for H100, _helion_aot_rms_norm_cuda_sm100.py for B200 etc. that Helion selects at runtime. The autotuning search already happened offline, and its result is baked in. When this program is run, the user gets tuned performance with no search cost. That is our baseline: a kernel whose configuration space has already been swept, on the hardware it’s running on, by the kernel authors and the Helion compiler developers. If there’s headroom below that, it’s real headroom and that’s the last-mile focus of this post. What Helion’s autotuner searches, and where it stops A Helion kernel is ordinary PyTorch with one new idea – tiles: import torch, helion, helion.language as hl @helion.kernel() def softmax(x: torch.Tensor) -> torch.Tensor: n, _m = x.size() out = torch.empty_like(x) for tile_n in hl.tile(n): out[tile_n, :] = torch.nn.functional.softmax(x[tile_n, :], dim=1) return out The user never says how big a tile is. That’s the point: hl.tile declares that the iteration space is tiled, and the autotuner decides the rest. A single kernel definition implicitly generates a search space of thousands of Triton configurations that span block_sizes, loop_orders, l2_grouping, num_warps, num_stages, indexing (pointer arithmetic vs. block pointers vs. TMA descriptors), pid_type, reduction_loops, and more. Helion lowers through TorchInductor to Triton, compiles each candidate, times it, and keeps the winner. Then the winning Triton kernel goes to ptxas, and everything from that point down happens at PTXAS’ default settings. The kernel autotuner’s search space is defined by what the kernel language can express, and no kernel language expresses “allocate registers differently”. This compiler control surface is a separate interface, exposed by ptxas itself starting from CUDA 13.3. It is imperative to mention that this exploration is genuinely hard. There are complex interactions between various compiler controls: a change to register allocation shifts occupancy, which changes whether a scheduling decision helps or hurts. The right setting is workload-specific, and there’s no universal answer to publish. And the payoff per unit of human effort is poor, which is exactly the profile of a problem one hands off to a search algorithm. The joint search: Helion Autotuner and PTXAS’ Advanced Controls Helion exposes the underlying compiler control space through two settings. Controls are packaged in an Advanced Control File (ACF, typically ending with .acf) which is a binary artifact applied at compile time, portable and versionable, that users can ship alongside their kernel. The first setting pins a known-good ACF as part of a config the user has already chosen: softmax_tuned = helion.kernel( config=helion.Config( advanced_controls_file="h100_example1.acf", block_sizes=[4], num_warps=8, num_stages=2, reduction_loops=[512], indexing=["pointer", "pointer", "tensor_descriptor", "tensor_descriptor", "pointer", "pointer"], ), static_shapes=True, )(_softmax) The second is a more interesting case. It hands a list of ACFs to Helion’s own autotuner and lets them be searched as just another dimension, alongside tile sizes and warp counts: softmax_search = helion.kernel( autotune_search_acf=[ "acfs/helion-booster-pack/candidate_03.acf", "acfs/helion-booster-pack/candidate_11.acf", "acfs/helion-booster-pack/candidate_17.acf", ], autotune_compile_timeout=120, static_shapes=True, )(_softmax) Two things are worth pointing out in this mode: The empty string is appended for you. Helion automatically adds “”, meaning “no ACF, compiler defaults”, to the list that is supplied. The default configuration is therefore always a candidate in the search, and the autotuner will only select an ACF if it actually won on measured time. The worst case of turning this on is that you spend some autotuning budget and get back the kernel you already had. That property is what makes it tempting to try on a kernel you care about. Sequential tuning finds a local optimum; this joint method doesn’t. This is the substantive argument for doing it in one loop rather than two. Suppose you autotune the kernel first, get block_sizes=[128, 128] with num_stages=4, and then go searching for compiler controls. You are now searching the control space conditioned on a kernel config that was itself chosen under default controls. But those dimensions are coupled through the hardware: block size and stage count determine register pressure and shared-memory footprint, and register allocation policy is precisely what decides whether that pressure causes spills. A control set that would shine at block_sizes=[64, 64] never gets evaluated there, because the kernel config was frozen before the compiler search began.The reverse ordering has the same problem. Searching both together is the only ordering that doesn’t have it, which is why autotune_search_acf lives inside Helion’s autotuner rather than being a post-processing step. The cost is search time. Each ACF multiplies the number of compile-and-measure trials, so this is not something to leave on by default for every kernel in a codebase. It is something to run once, offline, on a kernel that matters, and then pin the result with advanced_controls_file, exactly the way autotuned configs are pinned today. Where do these Advanced Control Files come from N

read the original at PyTorch →