跳到正文
PyTorch:Blog· Aditya Srikanth (NVIDIA), Rupak Roy (NVIDIA), Abhinav Singh (NVIDIA), Jason Ansel (Meta), Clement Anthonioz-Blanc (Meta) and Jongsok Choi (Meta)·· 3 小时前AI 评分50

PyTorch Helion 借助 NVIDIA CompileIQ 联合搜索编译器控制:一次搜索,两层优化

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

AI 导读

PyTorch Helion 现可借助 NVIDIA CompileIQ 搜索 ptxas 编译器控制项,与原有的 tile 尺寸、warp 数等内核配置联合调优,结果打包为可移植的 Advanced Control File(ACF)。

正文 · 原文

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

Now that we’ve understood the benefits of joint autotuning of Helion parameters with PTXAS’ advanced control files (ACF), the natural question one may have is how can these ACFs be generated. That is where NVIDIA’s CompileIQ comes into the picture.

CompileIQ treats compiler-control selection as a hyperparameter optimization problem. It pulls a curated search space for the target compiler and toolkit version, then runs an evolutionary search: propose a population of candidate control sets, compile each one, check correctness against a reference, time the survivors, breed the winners, repeat. The output of this search process is an ACF file. There is nothing GPU-specific about the search algorithm itself, and the only GPU-specific parts are what is entirely in the objective function that the user writes, which is usually a few lines around triton.testing.do_bench.

For most users however, the additional search budget is unnecessary, because the searching has been done. CompileIQ publishes Booster Packs – pre-tuned ACF catalogs – on its GitHub releases. The current helion-booster-pack contains 24 ACFs tuned across a range of Helion kernels, validated against ptxas from CUDA 13.4, targeting B200. So the realistic workflow for users is: download the pack, point Helion’s autotune_search_acf at it, autotune once, pin what wins. No search space to design, no objective function to write, no compiler expertise required!

If you have a workload the pack doesn’t cover, pip install compileiq and run your own targeted search. That path is documented in NVIDIA’s CompileIQ blog post and in various CompileIQ examples.

Helion and CompileIQ: Experimental Results

This section highlights many possible ways in which one can run a CompileIQ optimization campaign on Helion kernels. As a sample case study we applied CompileIQ to the pretuned nvfp4_gemv_fp4in_kernel, a Helion W4A4 (NVFP4 weight x NVFP4 activation) decode GEMV kernel that already ships a checked-in Helion AOT config for its target shape. The question: is there performance CompileIQ can extract below Helion’s own autotuning layer ,i.e., at the ptxas compilation stage?

Environment and Methodology

Benchmarking was conducted on a GB200 system with sm_100 (Blackwell, compute capability 10.0), utilizing ptxas 13.3.73 (CUDA 13.3), CompileIQ 1.0.3, Triton 3.8.0, and an editable source installation of Helion. All timing evaluations were executed via triton.testing.do_bench with the GPU clock locked. Every observed performance improvement was re-measured independently, validated against the PyTorch reference implementation for numerical correctness, and statistically filtered using a Welch’s t-test (p < 0.05) alongside Cohen’s d (d > 0.2). Complete statistical metrics are provided alongside each result below.

Kernel Description

nvfp4_gemv_fp4in_kernel in helion/pretuned_kernels/nvfp4_gemv/nvfp4_gemv.py, computing out = (weight @ x) * alpha for a decode (one output row) GEMV. Both the packed NVFP4 weight and activation bytes are unpacked per K-group via hl.load_float4_e2m1fn_x16_to_float16, accumulated in fp16 over a 16-way static range, and scaled by a decoded E4M3 block scaler.

This kernel includes a checked-in AOT configuration specifically for the N=8192, K=28672 “down_proj” decode-GEMV shape in _helion_aot_nvfp4_gemv_cuda_sm100.py:

_CONFIG_FP4IN_DEFAULT = {
    "block_sizes": [4, 256],
    "num_warps": 2,
    "num_stages": 3,
    "range_multi_buffers": [None, True],
}

Using Runtime as CompileIQ Objective Function

Using runtime as the CompileIQ objective function lets us optimize execution speed directly. The objective function takes a candidate ACF configuration, applies it to our fixed Helion kernel, and evaluates it: the kernel runs once, its output is checked against the PyTorch reference implementation, and only configurations that pass correctness are then benchmarked with triton.testing.do_bench. Any configuration that fails correctness returns an invalid score and is dropped from the search before it ever contributes a timing measurement.

def objective(config) -> float:
    acf_path = ... # write `config` to a .acf, or None for baseline
    kernel = _make_kernel(acf_path)
    got = run()
    _check(got, ref, "fp4in", N, K) # correctness gate
    return do_bench(run, warmup=25, rep=100, return_mode="mean")

We initialize a Search instance that couples our scoring routine to an evolutionary autotuner over PTXAS compiler controls, then invoke .start(). It sweeps candidate ACFs across generations generations of a pool_size-sized population, breeding the top-performing control configurations forward into each subsequent iteration.

config = SearchConfiguration(problem_type="min", generations=8, pool_size=15)
tuner = Search(
    objective_function=objective,
    search_space=PtxasSearchSpace(version=cuda_version, variant="default"),
    search_config=config,
    dump_results="nvfp4_gemv_results.csv",
    worker_type=WorkerTypes.ISOLATED, # crashing ACF candidates shouldn't kill the search
)

with gpu_benchmark_mode(clock_mhz=target_clock_mhz, raise_on_failure=False):
    results = tuner.start(task_timeout=60)

Executing an evolutionary search with CompileIQ across PTXAS ACFs, spanning 8 generations with a pool size of 15 candidates over our fixed Helion configuration on GB200 yielded an optimized ACF control set. This candidate reduced execution time from a baseline of 0.0971 ms down to 0.0928 ms, delivering a 1.047x (4.7%) speedup as measured directly within the search’s scoring loop.

sample baseline run: 0.09713 ms
Best runtime (CompileIQ Search): 0.0928 ms
Speedup: 1.047x

Validation of CompileIQ Search Results

To confirm our findings, every speedup is re-measured independently using increased warmup and repetition iterations under fixed GPU clock frequencies. Each configuration undergoes strict correctness verification against the reference model alongside a Welch’s t-test and Cohen’s d filter. This serves as our rigorous statistical threshold (requiring p < 0.05, cohens_d > 0.2, and optimized.mean < baseline.mean) to verify performance gains before trusting the numbers.

# run 1
baseline:  n=1744 mean=0.09752ms std=0.00218
optimized: n=2018 mean=0.09301ms std=0.00133
speedup_mean:   1.0485
p_value:  0.0
cohens_d: 2.501
significant: True

# run 2 (independent re-run)
baseline:  n=1757 mean=0.09744ms std=0.00216
optimized: n=2028 mean=0.09299ms std=0.00130
speedup_mean:   1.0478
p_value:  0.0
cohens_d: 2.489
significant: True

The two executions remain virtually identical within experimental variance: with d ≈ 2.5 demonstrating a massive effect size that easily surpasses our statistical threshold of 0.2 , verifying that this ACF performance gain is genuine.

The two runs are virtually identical within experimental variance, and d ≈ 2.5 in both is a massive effect size nearly 12.5x past our d > 0.2 gate, confirming this ACF performance gain is genuine, not noise.

Tuning against a hardware metric instead of runtime

While execution speed is a common optimization goal, CompileIQ is not limited to runtime alone. Integrating directly with Nsight Compute (ncu) enables the objective function to target any hardware metric exposed by the tool. To determine the key architectural bottleneck limiting kernel performance, we start by evaluating its profile and analyzing nvfp4_gemv_fp4in_kernel on a GB200 system highlights the following bottlenecks:

Section: GPU Speed of Light Throughput
Memory Throughput               %       78.99
L1/TEX Cache Throughput         %       89.35
Compute (SM) Throughput         %       37.29

Section: Occupancy
Theoretical Occupancy                   %    25.0
Achieved Occupancy                      %    21.63

  Est. Speedup: 21.01%
  ... theoretical occupancy (25.0%) is limited by the number of required
  registers.

Profiling shows the kernel is register-limited. NCU’s own cost model attributes a 21% ‘Est. Speedup’ opportunity to occupancy, and this ceiling is set by register count, something PTXAS’s register-allocation controls can directly influence.

Switching the objective from runtime to occupancy only touches two things: what the objective function scores, and the search’s problem_type:

# objective(): do_bench(...) becomes an ncu metric read
return _measure_occupancy(acf_path) # ncu: sm__warps_active.avg.pct_of_peak_sustained_active

# search config: minimize runtime -> maximize occupancy
config = SearchConfiguration(problem_type="max", generations=4, pool_size=6)

Pointing the search at achieved occupancy itself across 4 generations × pool size 6, rather than runtime, found a candidate that moved the proxy metric by +0.16pp (21.60% → 21.76% achieved occupancy, p = 9.2e-23, d = 13.17) and, independently re-measured, translated into a modest but statistically solid runtime win of ~1.75-2.0% (p = 3.6e-116, d = 0.78).

Multi-Objective Tuning

CompileIQ can also perform multi-objective tuning, where the scoring routine evaluates candidates along multiple dimensions simultaneously by returning a tuple rather than a scalar float. Instead of isolating a single top performer, the autotuning process maps out candidates along the Pareto frontier, identifying control sets that are not dominated across all target metrics. This yields a unified ACF optimized jointly for execution runtime alongside L1TEX cache throughput.

def objective(config) -> tuple[float, float]:
    runtime_ms = do_bench(run, warmup=25, rep=100, return_mode="mean")
    l1_throughput = _measure_l1throughput(acf_path)
    return runtime_ms, -l1_throughput

8 generations × pool size 15 converged to a single dominant candidate on the Pareto front. Unlike a tradeoff curve, this single point beats every other evaluated candidate on both axes simultaneously in the search’s own scoring. Validating it on a fresh cache at shipping-grade trial counts:

=== RUNTIME (ms, lower is better) ===
p_value=3.516e-295  cohens_d=87.024
speedup_mean=1.0689x
=== L1TEX THROUGHPUT (%, higher is better; internally negated) ===
p_value=7.123e-19  cohens_d=3.686
baseline L1TEX THROUGHPUT mean=81.5087%
candidate L1TEX THROUGHPUT mean=85.9620%

Multi-objective tuning found an ACF delivering a validated 6.9% runtime speedup, accompanied by a simultaneous 4.5pp L1TEX throughput gain (81.51% → 85.96%). Both results are statistically significant (p ≪ 0.05, Cohen’s d ≫ 0.2), with neither metric being traded off against the other.

Joint ACF Optimization Across Multiple Input Shapes

While earlier evaluations focused on a single kernel shape, N=8192, K=28672, real-world model deployment demands coverage across multiple configurations. The nvfp4_gemv implementation ships with several pre-tuned Helion AOT configurations, including a dedicated setup for N=15360, K=5120. Instead of maintaining distinct control files per shape, CompileIQ enables a unified search to discover a single, joint ACF that maximizes performance across both shapes simultaneously.

def objective(config) -> tuple[float, ...]:
    """One shared ACF, one runtime score per shape (both minimize)."""
    acf_path = ... # write `config` to a temp .acf, or None for baseline
    return tuple(_bench_shape(acf_path, shape) for shape in SHAPES)

Executing an 8-generation search with a pool size of 15 across both shapes yielded a Pareto front containing 4 optimal candidates:

Candidate down_proj 15360_5120
pareto_0 0.0927 ms 0.0309 ms
pareto_1 0.0927 ms 0.0309 ms
pareto_2 0.0944 ms 0.0308 ms
pareto_3 0.0917 ms 0.0321 ms

Compared to the default benchmark times of 0.0972 ms for down_proj and 0.0315 ms for 15360_5120, pareto_0 emerges as the standout configuration that consistently outpaces the baseline across both target tensor dimensions.

To confirm these performance gains, we execute an independent validation of pareto_0 on each individual shape. Both configurations comfortably pass our statistical validation (1.05x speedup with a p-value of 1.1e-224 and Cohen’s d of 67.71 for down_proj, and 1.01x speedup with a p-value of 6.2e-73 and a d of 4.54 on the 15360_5120 shape), far exceeding our benchmark requirements and proving these gains represent genuine performance improvements rather than system noise. A single joint ACF successfully outperforms the default baseline across both target dimensions simultaneously.

Applicability of this workflow for end-to-end model tuning

While the focus of this entire post is about optimizing the underlying compiler for extracting more performance from individual kernels, it is important to zoom out and highlight how a workflow like this could be utilized for end-to-end tuning of models on PyTorch.

The first step is to profile and characterize the most important kernels that are part of the model which you wish to optimize. Then, each of these kernels can then be a target of CompileIQ optimization, with an associated Advanced Control File to extract maximum performance from that instantiation of the kernel. These get applied during runtime using Helion’s pinned autotune configurations, using Triton’s ptxas_options API or directly via the PTXAS command line depending on the entry point of each kernel’s implementation.

We want to be upfront about the expected gains: On kernels that are already highly optimized, compiler-control tuning has shown runtime improvements on the order of 2-7% which may still mean a lot for critical kernels in top models. Less optimized kernels have achieved double-digit performance gains, offering greater potential for meaningful end-to-end improvements.

The important points to note are that it is found automatically, and requires little to no compiler expertise, it is verified against your own correctness check, and the result is a file you can commit next to the kernel and reapply on every build. It is the difference between a few percent that exists in principle and a few percent you actually ship.

It also compounds exactly where you’d want it to. There are high-value kernels inside inference loops that run continuously across large fleets. Even a few percent at the bottom of the stack, applied to the operations that dominate a decode step, is valuable at that scale.

Caveats

  • ACFs are hardware- and toolkit-specific: An ACF tuned for sm_90 is not valid on sm_100. Applying a mismatched ACF can produce incorrect results or runtime errors, not just slower code. The booster pack manifest states its target architecture and toolkit version. This is also why the correctness check in the search loop is not optional.
  • Search-time is a nontrivial cost: Every ACF in autotune_search_acf multiplies your autotuning trials. Use it offline on kernels that matter, then pin the winner.
  • Gains are workload-specific: The result depends on the workload, the shapes, the architecture, and how much headroom the kernel has left. The only number that matters for your kernel is the one you measure on your kernel.

Your turn to try this out!

If you’ve never written a GPU kernel: start at pretuned_kernels/. Pick something close to what you need, read it and then run it. The kernels are short, and they run at tuned speed immediately, with no autotuning wait.

If you already use Helion: download the helion-booster-pack, add autotune_search_acf=[…] to one kernel you care about, and autotune it. The no-ACF baseline stays in the search, so the downside is bounded at the time you spend. See Helion’s examples/acfs/ for a worked example.

If you work on kernels for a living: the takeaway here is not the tool itself. It is that the compiler’s control surface is a legitimate optimization space, it is searchable, and most people are not yet searching it. That is true of any compiler that exposes tunable controls, on any backend. We’re hopeful that we’ve inspired you to explore this space!

Licensing

Helion is open source and hosted by the PyTorch Foundation. CompileIQ is free to use and it is distributed under the NVIDIA Software License Agreement, a proprietary EULA. The ACF files it produces are freely distributable, which is why booster packs can be published and why an ACF can be committed alongside a kernel in an open repository.

Links

The post One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ appeared first on PyTorch.

来源:PyTorch:Blog · pytorch.org