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

PyTorch Helion 借助 NVIDIA CompileIQ 联合搜索编译器控制与内核配置

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

AI 导读

PyTorch Helion 现可通过 NVIDIA CompileIQ 搜索 ptxas 编译器控制空间,并将其与 tile 大小、warp 数等内核配置联合调优,在已离线调优的约 27 个预调内核上进一步提速。

正文 · AI 翻译

去年发布的 Helion 旨在让开发者用 PyTorch 原生的 Python 编写 GPU kernel,并让其强大的 autotuner 完成过去需要手写 Triton 才能完成的工作,同时抽象掉硬件相关的实现细节。Helion 的 autotuner 会探索数千种 kernel 配置,通常能达到专家调优代码的几个百分点以内,而且常常超越它。

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

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

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

Helion 现在实现了这一点,而本文要讲的是第二个编译器特定空间包含什么,以及为什么将其与 kernel 配置空间联合搜索,能在已被视为最优的 kernel 上展现出进一步的提升。

搜索由 NVIDIA 的 CompileIQ 完成,它像 autotuner 优化 tile 大小那样优化 ptxas 控制项,并将结果打包为可移植的高级控制文件(ACF)。大多数用户不会自己运行该搜索:CompileIQ 发布预调优的 ACF 目录,Helion 的 autotuner 现在可以将其与它已经探索的 kernel 配置一起评估。

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

与什么基线比较非常重要。如果起点是朴素 kernel,“我们让 kernel 更快了”的说法很容易成立,但对真实工作负载几乎毫无意义。

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

这里需要注意的重要特性是,这些 kernel 并非在运行时进行 autotune。每个 kernel 都附带按架构预先生成的启发式文件:用于 H100 的 _helion_aot_rms_norm_cuda_sm90.py、用于 B200 的 _helion_aot_rms_norm_cuda_sm100.py 等,Helion 在运行时选择它们。autotuning 搜索已经在离线完成,其结果已固化其中。运行该程序时,用户获得调优后的性能,且没有搜索成本。

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

Helion 的 autotuner 搜索什么,以及它在哪里停止

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 声明迭代空间被分块,其余由自动调优器决定。单个内核定义会隐式生成一个包含数千种 Triton 配置的搜索空间,涵盖 block_sizes、loop_orders、l2_grouping、num_warps、num_stages、索引方式(指针算术 vs. 块指针 vs. TMA 描述符)、pid_type、reduction_loops 等。Helion 通过 TorchInductor 下沉到 Triton,编译每个候选,计时,并保留胜出者。

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

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

联合搜索:Helion Autotuner 与 PTXAS 高级控制

Helion 通过两个设置暴露底层的编译器控制空间。控制项打包在高级控制文件(ACF,通常以 .acf 结尾)中,这是一个在编译时应用的二进制产物,可移植、可版本化,用户可以随内核一起分发。

第一个设置将已知良好的 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 做自动调优,得到 block_sizes=[128, 128] 和 num_stages=4,然后再去搜索编译器控制参数。此时你是在一个 kernel 配置的条件下搜索控制空间,而这个 kernel 配置本身又是在默认控制参数下选出来的。但这些维度通过硬件相互耦合:block size 和 stage 数量决定了寄存器压力和共享内存占用,而寄存器分配策略恰恰决定了这种压力是否会导致溢出。一组在 block_sizes=[64, 64] 下本会表现出色的控制参数永远不会在那里被评估,因为 kernel 配置在编译器搜索开始之前就已经冻结了。反过来排序也有同样的问题。只有把两者放在一起搜索,才不存在这个问题,这也是为什么 autotune_search_acf 位于 Helion 的自动调优器内部,而不是作为一个后处理步骤。

代价是搜索时间。每个 ACF 都会成倍增加编译和测量的试验次数,所以这不是一个可以对代码库中每个 kernel 都默认开启的东西。它应该针对重要的 kernel 在离线状态下运行一次,然后用 advanced_controls_file 固定结果,就像今天固定自动调优配置那样。

这些 Advanced Control Files 从何而来

既然我们已经理解了将 Helion 参数与 PTXAS 的 advanced control files(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 内核上运行 CompileIQ 优化活动的多种可能方式。作为示例案例研究,我们将 CompileIQ 应用于预调优的 nvfp4_gemv_fp4in_kernel,这是一个 Helion W4A4(NVFP4 权重 x NVFP4 激活)解码 GEMV 内核,其目标形状已附带一份已检入的 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)进行统计过滤。完整的统计指标随下方每项结果一并提供。

内核描述

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 块缩放器进行缩放。

该内核包含一份专门针对 N=8192、K=28672 的“down_proj”解码 GEMV 形状的已检入 AOT 配置,位于 _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],
}

以运行时作为 CompileIQ 目标函数

以运行时作为 CompileIQ 目标函数,使我们能够直接优化执行速度。目标函数接收一个候选 ACF 配置,将其应用于我们固定的 Helion 内核并进行评估:内核运行一次,其输出对照 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) 集成,使目标函数能够针对该工具暴露的任何硬件指标。为了确定限制内核性能的关键架构瓶颈,我们首先评估其性能剖析,并在 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.

性能剖析显示该内核受寄存器限制。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 优化

虽然早期的评估集中在单一内核形状 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 成功同时在两个目标维度上超越默认基线。

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

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

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

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

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

它还会在你期望的地方产生复合效应。推理循环中存在高价值内核,它们在大规模集群上持续运行。即使是堆栈底部的几个百分点,应用于主导解码步骤的操作,在那个规模上也是有价值的。

注意事项

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

轮到你来试试了!

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

如果你已经在使用 Helion:下载 helion-booster-pack,在你关心的某个内核中加入 autotune_search_acf=[…],然后对它进行自动调优。无 ACF 的基线仍保留在搜索中,因此代价上限就是你投入的时间。参见 Helion 的 examples/acfs/ 获取完整示例。

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

许可

Helion 是开源的,由 PyTorch 基金会托管。CompileIQ 可免费使用,依据 NVIDIA 软件许可协议(一份专有 EULA)分发。它生成的 ACF 文件可以自由分发,因此 booster pack 可以发布,ACF 也可以与内核一起提交到开源仓库中。

链接

文章 One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ 首次发表于 PyTorch。

来源:PyTorch Blog · pytorch.org