博客 一次搜索,两层优化:PyTorch Helion 借助 NVIDIA CompileIQ 引入编译器控制 作者: Aditya Srikanth(NVIDIA)、Rupak Roy(NVIDIA)、Abhinav Singh(NVIDIA)、Jason Ansel(Meta)、Clement Anthonioz-Blanc(Meta)和 Jongsok Choi(Meta)2026年10月19日暂无评论 /row

精选项目

随着 Helion 于去年发布,它的目标是让开发者能够用 PyTorch 原生的 Python 编写 GPU kernel,并让其强大的 自动调优器 完成过去需要手写 Triton 才能完成的工作,同时抽象掉与硬件相关的实现细节。Helion 的自动调优器会探索数千种 kernel 配置,通常能达到与专家调优代码相差几个百分点的水平,而且往往还能超越它。

现在需要回答的问题是:如果一个 kernel 已经被自动调优并优化到开发者能力的极限,还有什么可做的吗?

有的。它存在于 kernel 之下,在编译器之中。

生态系统中的每一个 kernel 自动调优器,无论是 Helion 的、Triton 的,还是 torch.compile 内部的,它们搜索的都是同样的东西:tile 形状、warp 数量、流水线阶段、内存访问策略。所有这些都是关于每个线程块执行什么工作的决策。在这些决策之下,还有第二组决策,关乎这些工作如何变成由 ptxas(NVIDIA 的 PTX 汇编器)生成的机器汇编指令。诸如寄存器分配策略、细粒度指令调度、代码生成启发式等,如今发布的所有 kernel 都沿用默认设置——不是因为默认设置最优,而是因为从来没有人去搜索和优化它们!

Helion 现在让这成为可能,而本文要讲的就是这第二个编译器专属空间包含什么,以及为什么将它和 kernel 配置空间联合搜索,能在那些已被视为最优的 kernel 上带来进一步的提升。

搜索工作由 NVIDIA 的 CompileIQ 完成,它优化 ptxas 控制参数的方式,就像自动调优器优化 tile 尺寸一样,并将结果打包为可移植的 Advanced Control File(ACF)。大多数用户不会自己去运行这项搜索:CompileIQ 会发布预先调优好的 ACF 目录,Helion 的自动调优器现在可以在它原本就探索的 kernel 配置之外,一并评估这些 ACF。

从高度优化的 kernel 出发,优化最后一公里

对比的基线是什么,这一点至关重要。如果起点是一个朴素的 kernel,那么“我们让 kernel 更快了”这种说法很容易成立,但对真实工作负载来说几乎毫无意义。

我们这里使用的 kernel 来自 Helion 仓库中的 pretuned_kernels/ ,大约 27 个 kernel,覆盖了真实模型代码中出现的各种形态:FP8 和 NVFP4 量化、RMSNorm 变体、RoPE、分组 GEMM、因果深度卷积、Gated DeltaNet 解码,以及像 Qwen3 解码层这样的融合模型区域。

这里需要注意的重要特性是:这些 kernel 并非在运行时即时自动调优。每个 kernel 都附带按架构划分的提前(ahead-of-time)启发式文件:H100 对应 _helion_aot_rms_norm_cuda_sm90.py,B200 对应 _helion_aot_rms_norm_cuda_sm100.py 等等,Helion 在运行时选择这些文件。自动调优搜索已经在离线完成,其结果被固化在内。运行该程序时,用户无需承担搜索成本即可获得调优后的性能。

这就是我们的基线:一个配置空间已经由 kernel 作者和 Helion 编译器开发者在它运行的硬件上扫描过的 kernel。如果在这之下还有提升空间,那才是 真正的 提升空间,而这正是本文关注的最后一公里。

Helion 的自动调优器搜索什么,又在哪里停下

一个 Helion kernel 就是普通的 PyTorch,外加一个新概念——tile:

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

用户从不说 tile 有多大。这正是关键所在:hl.tile 声明迭代空间是分块的,其余由自动调优器决定。单个 kernel 定义会隐式生成一个包含数千种 Triton 配置的搜索空间,涵盖 block_sizes、loop_orders、l2_grouping、num_warps、num_stages、索引方式(指针算术 vs. 块指针 vs. TMA 描述符)、pid_type、reduction_loops 等等。Helion 经由 TorchInductor 下沉到 Triton,编译每个候选配置,计时,并保留胜出者。

然后,胜出的 Triton kernel 被交给 ptxas,而从那一刻起往下的一切都发生在 PTXAS 的默认设置下。kernel 自动调优器的搜索空间由 kernel 语言所能表达的内容定义,而没有任何 kernel 语言能表达“以不同方式分配寄存器”。这个 编译器控制面 是一个独立的接口,由 ptxas 自身从 CUDA 13.3 起开放。

必须指出,这种探索确实很难。各种编译器控制之间存在复杂的相互作用:寄存器分配的改动会改变占用率,进而改变某个调度决策是有益还是有害。正确的设置因工作负载而异,没有可以发布的通用答案。而且每单位人力投入的回报很低,这恰恰是适合交给搜索算法的问题特征。

联合搜索:Helion 自动调优器与 PTXAS 高级控制

Helion 通过两个设置来开放底层的编译器控制空间。控制参数被打包在一个 Advanced Control File (ACF,通常以 .acf 结尾)中,这是一种在编译时应用的二进制产物,可移植、可版本化,用户可以随 kernel 一起分发。

第一个设置将已知良好的 ACF 固定为用户已选配置的一部分:

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)

第二个是更有意思的情况。它把一组 ACF 交给 Helion 自己的自动调优器,让它们作为又一个维度,与 tile 尺寸和 warp 数量一起被搜索:

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)

这种模式下有两点值得指出:

  • 空字符串会为你自动追加。Helion 会自动把“”(意为“无 ACF,使用编译器默认值”)加入所提供的列表中。因此,默认配置始终是搜索中的一个候选,自动调优器只有在某个 ACF 确实在实测时间上胜出时才会选择它。开启这一功能的最坏情况是,你花费了一些自动调优预算,最终得到的还是你原本就有的 kernel。正是这一特性,让人忍不住想在自己关心的 kernel 上试一试。
  • 顺序调优会找到局部最优;而这种联合方法不会。这就是为什么应当在一个循环中而不是两个循环中完成这件事的实质性论据。假设你先对 kernel 进行自动调优,得到 block_sizes=[128, 128] 和 num_stages=4,然后再去搜索编译器控制参数。此时,你是在一个本身是在默认控制参数下选出的 kernel 配置的条件下去搜索控制空间。但这些维度通过硬件相互耦合:块大小和阶段数决定了寄存器压力和共享内存占用,而寄存器分配策略恰恰决定了这种压力是否会导致溢出。一组在 block_sizes=[64, 64] 下会大放异彩的控制参数永远不会在那里被评估,因为 kernel 配置在编译器搜索开始之前就已经冻结了。反向顺序也有同样的问题。同时搜索两者是唯一没有这个问题的顺序,这就是为什么 autotune_search_acf 位于 Helion 的自动调优器内部,而不是作为一个后处理步骤。

代价是搜索时间。每个 ACF 都会使编译并测量的试验次数成倍增加,因此这不是一个应该对代码库中每个 kernel 都默认开启的功能。它应该针对重要的 kernel 离线运行一次,然后用 advanced_controls_file 固定结果,就像如今固定自动调优配置那样。

这些高级控制文件从何而来

既然我们已经理解了将 Helion 参数与 PTXAS 的高级控制文件(ACF)联合自动调优的好处,那么自然会问:这些 ACF 是如何生成的?这正是 NVIDIA 的 CompileIQ 登场之处。

CompileIQ 将编译器控制选择视为一个超参数优化问题。它会为目标编译器和工具包版本拉取一个精心整理的搜索空间,然后运行进化搜索:提出一组候选控制集,逐一编译,对照参考实现检查正确性,对通过者计时,让胜出者繁衍,如此重复。这一搜索过程的输出就是一个 ACF 文件。搜索算法本身没有任何 GPU 特定之处,唯一与 GPU 相关的部分完全在用户编写的目标函数中,通常只是围绕 triton.testing.do_bench 的几行代码。

然而对大多数用户来说,额外的搜索预算是不必要的,因为搜索已经完成了。CompileIQ 在其 GitHub releases 上发布了 Booster Packs——预先调优的 ACF 目录。当前的 helion-booster-pack 包含 24 个 ACF,针对一系列 Helion kernel 进行了调优,已对照 CUDA 13.4 的 ptxas 验证,面向 B200。因此,用户现实的工作流程是:下载该包,将 Helion 的 autotune_search_acf 指向它,自动调优一次,固定胜出者。无需设计搜索空间,无需编写目标函数,无需编译器专业知识!

如果你有该包未覆盖的工作负载,可以 pip install compileiq 并运行你自己的针对性搜索。该路径记录在 NVIDIA 的 CompileIQ 博客文章 以及各种 CompileIQ 示例.

中。Helion 与 CompileIQ:实验结果

本节重点介绍在 Helion kernel 上运行 CompileIQ 优化活动的多种可能方式。作为示例案例研究,我们将 CompileIQ 应用于预调优的 nvfp4_gemv_fp4in_kernel,这是一个 Helion W4A4(NVFP4 权重 × NVFP4 激活)解码 GEMV kernel,已经为其目标形状附带了一个已检入的 Helion AOT 配置。问题是:CompileIQ 能否在 Helion 自身的自动调优层之下,即在 ptxas 编译阶段,挖掘出性能?

环境与方法

基准测试在配备 sm_100(Blackwell,计算能力 10.0)的 GB200 系统上进行,使用 ptxas 13.3.73(CUDA 13.3)、CompileIQ 1.0.3、Triton 3.8.0,以及 Helion 的可编辑源码安装。所有计时评估均通过 triton.testing.do_bench 在 GPU 时钟锁定的情况下执行。每一项观察到的性能提升都经过独立重新测量,对照 PyTorch 参考实现验证数值正确性,并使用 Welch t 检验(p < 0.05)以及 Cohen's d(d > 0.2)进行统计过滤。完整的统计指标随下方每项结果一并提供。

Kernel 描述

helion/pretuned_kernels/nvfp4_gemv/nvfp4_gemv.py 中的 nvfp4_gemv_fp4in_kernel,针对解码(单输出行)GEMV 计算 out = (weight @ x) * alpha。打包的 NVFP4 权重和激活字节均通过 hl.load_float4_e2m1fn_x16_to_float16 按 K 组解包,在 fp16 中沿 16 路静态范围累加,并由解码后的 E4M3 块缩放因子缩放。

该 kernel 包含一个专门针对 N=8192、K=28672 的“down_proj”解码 GEMV 形状、已检入 _helion_aot_nvfp4_gemv_cuda_sm100.py 的 AOT 配置:

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

使用运行时作为 CompileIQ 目标函数

使用运行时作为 CompileIQ 目标函数,让我们可以直接优化执行速度。目标函数接收一个候选 ACF 配置,将其应用于我们固定的 Helion kernel 并进行评估:kernel 运行一次,其输出对照 PyTorch 参考实现进行检查,只有通过正确性检查的配置才会用 triton.testing.do_bench 进行基准测试。任何未通过正确性检查的配置都会返回无效分数,并在贡献任何计时测量之前就被从搜索中剔除。

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")

我们初始化一个 Search 实例,将我们的评分例程与针对 PTXAS 编译器控制的进化式自动调优器耦合,然后调用 .start()。它在 pool_size 大小的种群中跨代扫描候选 ACF,将表现最好的控制配置繁殖到后续每一次迭代中。

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)

在 GB200 上针对我们固定的 Helion 配置,使用 CompileIQ 跨 PTXAS ACF 执行进化搜索,跨越 8 代、种群规模为 15 个候选,得到了一个优化后的 ACF 控制集。该候选将执行时间从 0.0971 ms 的基线降低到 0.0928 ms,在搜索的评分循环中直接测得 1.047 倍(4.7%)的加速。

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

CompileIQ 搜索结果的验证

为确认我们的发现,每个加速都在固定 GPU 时钟频率下,使用更多的预热和重复迭代独立重新测量。每个配置都要经过针对参考模型的严格正确性验证,以及 Welch t 检验和 Cohen's d 过滤。这作为我们严格的统计阈值(要求 p < 0.05、cohens_d > 0.2 且 optimized.mean < baseline.mean),在信任这些数字之前验证性能提升。

# 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

两次执行在实验方差范围内几乎完全相同:d ≈ 2.5 表明效应量巨大,轻松超过我们 0.2 的统计阈值,验证了这一 ACF 性能提升是真实的。

两次运行在实验方差范围内几乎完全相同,且两者中 d ≈ 2.5 都是巨大的效应量,接近我们 d > 0.2 门槛的 12.5 倍,确认这一 ACF 性能提升是真实的,而非噪声。

以硬件指标而非运行时为目标的调优

虽然执行速度是常见的优化目标,但 CompileIQ 并不局限于运行时。直接与 Nsight Compute(ncu)集成,使目标函数能够以该工具暴露的任何硬件指标为目标。为确定限制 kernel 性能的关键架构瓶颈,我们首先评估其性能剖析,并在 GB200 系统上分析 nvfp4_gemv_fp4in_kernel,突出显示以下瓶颈:

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.

性能剖析显示该 kernel 受寄存器限制。NCU 自身的成本模型将 21% 的“预计加速”机会归因于占用率,而这一上限由寄存器数量决定,这正是 PTXAS 的寄存器分配控制可以直接影响的。

将目标从运行时切换为占用率只涉及两件事:目标函数评分的内容,以及搜索的 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)

将搜索目标指向实际达到的占用率本身,跨越 4 代 × 种群规模 6,而非运行时,找到了一个候选,使代理指标提升了 +0.16pp(实际达到的占用率从 21.60% 提升到 21.76%,p = 9.2e-23,d = 13.17),并且经独立重新测量,转化为约 1.75-2.0% 的适度但统计上稳健的运行时收益(p = 3.6e-116,d = 0.78)。

多目标调优

CompileIQ 还可以执行多目标调优,其中评分例程通过返回元组而非标量浮点数,同时沿多个维度评估候选。自动调优过程不是孤立出单个最佳表现者,而是沿帕累托前沿描绘候选,识别在所有目标指标上均未被支配的控制集。这产生了一个统一的 ACF,同时针对执行运行时和 L1TEX 缓存吞吐量进行联合优化。

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 代 × 种群规模 15 收敛到帕累托前沿上的单个主导候选。与权衡曲线不同,这个单点在搜索自身的评分中同时在两个轴上击败了所有其他被评估的候选。在全新缓存上以生产级试验次数验证它:

=== 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%

多目标调优找到了一个 ACF,经验证带来 6.9% 的运行时加速,同时伴随 4.5pp 的 L1TEX 吞吐量提升(从 81.51% 提升到 85.96%)。两项结果均具有统计显著性(p ≪ 0.05,Cohen's d ≫ 0.2),且两项指标之间没有相互权衡。

跨多个输入形状的联合 ACF 优化

虽然此前的评估集中在单一 kernel 形状 N=8192、K=28672,但真实世界的模型部署要求覆盖多种配置。nvfp4_gemv 实现附带若干预调优的 Helion AOT 配置,包括针对 N=15360、K=5120 的专用设置。CompileIQ 不是为每种形状维护不同的控制文件,而是支持统一搜索,以发现一个同时最大化两种形状性能的单一联合 ACF。

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)

在两种形状上执行 8 代、种群规模为 15 的搜索,得到了包含 4 个最优候选的帕累托前沿:

候选 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

与 down_proj 的默认基准时间 0.0972 ms 和 15360_5120 的 0.0315 ms 相比,pareto_0 脱颖而出,成为在两个目标张量维度上始终超越基线的配置。

为确认这些性能提升,我们在每个单独形状上对 pareto_0 执行独立验证。两个配置都轻松通过我们的统计验证(down_proj 上 1.05 倍加速,p 值为 1.1e-224,Cohen's d 为 67.71;15360_5120 形状上 1.01 倍加速,p 值为 6.2e-73,d 为 4.54),远超我们的基准要求,证明这些提升代表真实的性能改进而非系统噪声。单个联合 ACF 成功同时在两个目标维度上超越默认基线。

此工作流用于端到端模型调优的适用性

虽然整篇文章的重点是优化底层编译器以从单个 kernel 中提取更多性能,但重要的是拉远视角,强调这样的工作流如何能用于在 PyTorch 上进行模型的端到端调优。

第一步是对你希望优化的模型中最重要的 kernel 进行剖析和特征描述。然后,这些 kernel 中的每一个都可以成为 CompileIQ 优化的目标,并配有相应的 Advanced Control File,以便从该 kernel 实例中提取最大性能。这些会在运行时通过 Helion 固定的自动调优配置、使用 Triton 的 ptxas_options API 或直接通过 PTXAS 命令行(取决于每个 kernel 实现的入口点)来应用。

我们希望坦诚说明预期收益:在已经高度优化的 kernel 上,编译器控制调优已显示出约 2-7% 的运行时提升,这对顶级模型中的关键 kernel 而言可能仍然意义重大。优化程度较低的 kernel 则实现了两位数的性能提升,为有意义的端到端改进提供了更大潜力。

需要注意的要点是:它是自动找到的,几乎不需要编译器专业知识;它会根据你自己的正确性检查进行验证;结果是一个文件,你可以将其提交到 kernel 旁边,并在每次构建时重新应用。这正是理论上存在的几个百分点与你实际交付的几个百分点之间的差别。

它还会在你最希望的地方产生复利效应。推理循环中存在高价值 kernel,它们会在大型集群上持续运行。即使是底层栈的几个百分点,应用到主导解码步骤的操作上,在这样的规模下也很有价值。

注意事项

  • ACF 与硬件和工具包特定相关: 为 sm_90 调优的 ACF 在 sm_100 上无效。 应用不匹配的 ACF 可能产生错误结果或运行时错误,而不仅仅是代码变慢。booster pack 清单会说明其目标架构和工具包版本。这也是搜索循环中的正确性检查并非可选的原因。
  • 搜索时间是一项不可忽视的成本: autotune_search_acf 中的每个 ACF 都会成倍增加你的自动调优试验次数。请在离线状态下对重要的 kernel 使用它,然后固定胜出者。
  • 收益因工作负载而异: 结果取决于工作负载、形状、架构以及 kernel 剩余的优化空间。对你的 kernel 而言,唯一重要的数字就是你在自己的 kernel 上测得的数字。

轮到你来试试了!

如果你从未编写过 GPU kernel:请从 pretuned_kernels/开始。挑选一个接近你需求的内容,阅读它,然后运行它。这些 kernel 很短,而且会立即以调优后的速度运行,无需等待自动调优。

如果你已经在使用 Helion:请下载 helion-booster-pack,将 autotune_search_acf=[…] 添加到你关心的某个 kernel 中,然后对其进行自动调优。无 ACF 的基线仍保留在搜索中,因此下行风险仅限于你所花费的时间。请参阅 Helion 的 examples/acfs/ 获取完整示例。

如果你以编写 kernel 为生:这里的要点不是工具本身。而是编译器的控制面是一个正当的优化空间,它是可搜索的,而大多数人尚未开始搜索它。这对任何暴露可调控制项的编译器、在任何后端上都成立。我们希望已经启发了你去探索这个空间!

许可

Helion 是开源项目,由 PyTorch 基金会托管。CompileIQ 可免费使用,并根据 NVIDIA 软件许可协议(一项专有 EULA)分发。它生成的 ACF 文件可自由分发,这就是 booster pack 可以发布、以及 ACF 可以与 kernel 一起提交到开放仓库的原因。

链接

/row /row