PyTorch Blog

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

6.9内容质量
One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ

TL;DR · AI 摘要

One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ - PyTorch Featured projects With...

核心要点

  • 主题聚焦:One Search, Two Layers: Compiler Controls Inside
  • 来源:PyTorch Blog,建议结合原文判断细节。
  • AI 分析暂不可用,本条为保底评分与摘要。
#AI#编程#后端#安全#开源
打开原文

一次搜索,两层控制:使用 NVIDIA CompileIQ 在 PyTorch Helion 中的编译器控制 - PyTorch

特色项目

随着 Helion 去年发布,其目标是让开发者能够使用原生 PyTorch Python 编写 GPU 内核,并让其强大的自动调优器完成过去需要手动编写 Triton 的工作,同时抽象出硬件相关的实现细节。Helion 的自动调优器会探索数千种内核配置,并经常将性能优化到专家调优代码的百分之几以内,有时甚至超越专家调优结果。

现在需要回答的问题是:如果一个内核已经通过自动调优和优化达到了开发者能力的最佳状态,是否还有提升空间?

有的。它存在于内核之下,即编译器中。

生态系统中所有内核自动调优器(无论是 Helion 的、Triton 的,还是 torch.compile 内置的)都在搜索相同的内容:分块形状、线程数、流水线阶段、内存访问策略。所有这些都是关于每个线程块执行什么工作的决策。在这些决策之下,还有一个关于如何将这些工作转化为 ptxas(NVIDIA 的 PTX 汇编器)生成的机器汇编指令的第二层决策。诸如寄存器分配策略、细粒度指令调度、代码生成启发式等设置,目前所有内核都使用默认值——不是因为默认值最优,而是因为没有人对这些设置进行搜索和优化!

Helion 现在实现了对这些设置的搜索,本文将探讨这一第二层编译器特定空间包含的内容,以及为何与内核配置空间联合搜索能进一步提升那些被认为已达到最优的内核性能。

搜索由 NVIDIA 的 CompileIQ 完成,它优化 ptxas 控制的方式与自动调优器优化分块大小的方式类似,并将结果封装为可移植的高级控制文件(ACF)。大多数用户不会自己运行这个搜索:CompileIQ 发布了预调优的 ACF 目录,Helion 的自动调优器现在可以将其与已探索的内核配置一同评估。

从高度优化的内核开始,优化最后一公里

基准对比的选择至关重要。如果起点是一个简单的内核,声称“我们让内核运行得更快”很容易实现,但对于实际工作负载几乎毫无意义。

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

需要注意的一个重要特性是,这些内核并非实时自动调优。每个内核都针对特定架构预置了提前生成的启发式文件:_helion_aot_rms_norm_cuda_sm90.py(用于 H100)、_helion_aot_rms_norm_cuda_sm100.py(用于 B200)等,Helion 在运行时会选择这些文件。自动调优搜索已经在离线阶段完成,结果已固化其中。当运行此程序时,用户无需搜索即可获得调优后的性能。

这就是我们的基准:一个内核的配置空间已经在其运行的硬件上被内核作者和 Helion 编译器开发者彻底探索过。如果在这一基准之下仍有提升空间,那就是真正的提升空间,也是本文关注的“最后一公里”优化重点。

Helion 自动调优器的搜索范围及其局限

Helion 内核本质上是普通的 PyTorch,但引入了一个新概念——分块(tiles):

code
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

用户从不指定分块的大小。这正是关键所在: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 的默认设置。内核自动调优器的搜索空间由内核语言所能表达的内容定义,而没有任何内核语言能表达“以不同方式分配寄存器”。这种编译器控制接口是另一个独立的接口,自 CUDA 13.3 起由 ptxas 本身直接暴露。

必须强调的是,这种探索确实非常困难。各种编译器控制参数之间存在复杂的交互:寄存器分配的更改会影响占用率,进而改变调度决策是有利还是有害。正确的设置取决于具体的工作负载,不存在放之四海而皆准的答案。而且,每单位人力投入的回报率较低,这正是将问题交给搜索算法处理的典型场景。

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

Helion 通过两个设置暴露底层的编译器控制空间。控制参数被打包在高级控制文件(ACF,通常以 .acf 结尾)中,该文件是一个二进制构件,编译时应用,具有可移植性和版本控制能力,用户可以将其与内核一起分发。

第一个设置将已知的优质 ACF 作为用户已选定配置的一部分进行固定:

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

第二种情况更有趣。它向 Helion 自动调优器提供一组 ACF 文件,允许它们与其他维度(如分块大小和线程块数量)一起被搜索:

code
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 在实际测量时间上表现更优时,自动调优器才会选择它。开启此功能的最坏情况是您消耗了一些自动调优预算,但最终得到的仍然是原本已有的内核。这种特性使得人们倾向于在关心的内核上尝试这一功能。
  • 顺序调优只能找到局部最优解;而这种联合方法则不会。这正是为何要在单个循环中进行而非两个循环的核心论点。假设您先对内核进行自动调优,得到 block_sizes=[128, 128] 和 num_stages=4,然后开始搜索编译器控制参数。此时您搜索的控制空间是基于一个本身是在默认控制参数下选择的内核配置。但这些维度通过硬件耦合:块大小和阶段数量决定了寄存器压力和共享内存占用,而寄存器分配策略正是决定这种压力是否会导致寄存器溢出的关键因素。一个在 block_sizes=[64, 64] 上表现优异的控制参数集,由于内核配置在编译器搜索开始前已被固定,因此从未在该配置下被评估。反向顺序同样存在这个问题。只有同时搜索两者才是唯一没有这个问题的顺序,这也是为何 autotune_search_acf 是 Helion 自动调优器的一部分,而非后处理步骤。

成本是搜索时间。每个 ACF 都会增加编译和测量试验的数量,因此不建议对代码库中的每个内核都默认开启此功能。正确的做法是:仅对关键内核进行一次离线搜索,然后通过 advanced_controls_file 将结果固定下来,这与当前自动调优配置的固定方式完全一致。

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

既然我们已经理解了 Helion 参数与 PTXAS 高级控制文件(ACF)联合自动调优的优势,自然会问这些 ACF 是如何生成的。这正是 NVIDIA 的 CompileIQ 发挥作用的地方。

CompileIQ 将编译器控制参数选择视为超参数优化问题。它会为特定编译器和工具包版本提取一个精选的搜索空间,然后运行进化搜索:提出一组候选控制参数集,编译每个参数集,与参考实现检查正确性,测量存活者的性能,繁殖胜出者,重复这一过程。搜索过程的输出是一个 ACF 文件。搜索算法本身并不特定于 GPU,唯一与 GPU 相关的部分完全包含在用户编写的 objective function 中,通常仅涉及 triton.testing.do_bench 几行代码。

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

如果您的工作负载不在预设包的覆盖范围内,请安装 pip install compileiq 并运行自己的定向搜索。这一路径在 NVIDIA 的 CompileIQ 博客文章和多个 CompileIQ 示例中均有详细说明。

Helion与CompileIQ:实验结果

本节展示了在Helion内核上运行CompileIQ优化活动的多种可能方式。作为案例研究,我们将CompileIQ应用于预调优的nvfp4_gemv_fp4in_kernel,这是一个Helion W4A4(NVFP4权重×NVFP4激活)解码GEMV内核,该内核已包含针对目标形状的Helion AOT配置。问题在于:CompileIQ能否在Helion自身的自动调优层之下(即ptxas编译阶段)提取性能优化?

环境与方法

基准测试在GB200系统(sm_100,Blackwell,计算能力10.0)上进行,使用ptxas 13.3.73(CUDA 13.3)、CompileIQ 1.0.3、Triton 3.8.0以及可编辑源安装的Helion。所有时序评估均通过triton.testing.do_bench执行,且GPU时钟被锁定。所有观察到的性能提升均独立重新测量,与PyTorch参考实现进行数值正确性验证,并通过Welch’s 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。通过hl.load_float4_e2m1fn_x16_to_float16,按K组解包打包的NVFP4权重和激活字节,在16路静态范围内以fp16累加,并通过解码的E4M3块缩放器进行缩放。

该内核在_helion_aot_nvfp4_gemv_cuda_sm100.py中包含针对N=8192、K=28672“down_proj”解码-GEMV形状的已提交AOT配置:

code
_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进行基准测试。任何未通过正确性验证的配置将返回无效分数,并在产生任何时序测量之前被排除。

code
def objective(config) -> float:
    acf_path = ... # 将`config`写入.acf文件,或为基线设为None
    kernel = _make_kernel(acf_path)
    got = run()
    _check(got, ref, "fp4in", N, K) # 正确性验证门
    return do_bench(run, warmup=25, rep=100, return_mode="mean")

我们初始化一个Search实例,将评分逻辑与PTXAS编译器控制的进化自动调优器耦合,然后调用.start()。它通过池大小为population的种群,在多代中扫面候选ACF,将表现最佳的控制配置传递到后续迭代中。

code
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上使用CompileIQ对PTXAS ACF进行进化搜索,基于固定的Helion配置,跨越8代、候选集池大小为15的搜索,最终获得优化的ACF控制集。该候选方案将执行时间从基准值0.0971毫秒降低至0.0928毫秒,在搜索评分循环中直接测量实现了1.047倍(4.7%)的加速比。

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

CompileIQ搜索结果验证

为验证发现,我们通过固定GPU时钟频率、增加预热和重复迭代次数,独立重新测量所有加速比。每个配置均需通过参考模型的严格正确性验证,并结合威尔科xon t检验和科恩d滤波器进行分析。这构成了我们的严格统计阈值(要求p < 0.05、cohens_d > 0.2且优化均值 < 基准均值),用于验证性能提升的有效性。

code
# 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在两者中均表现为巨大效应量,约为阈值的12.5倍,确认该ACF性能提升是真实效果而非噪声。

针对硬件指标而非运行时的调优

尽管执行速度是常见的优化目标,CompileIQ的适用范围不限于运行时。通过与Nsight Compute(ncu)的深度集成,目标函数可针对工具暴露的任何硬件指标进行优化。为确定限制内核性能的关键架构瓶颈,我们首先评估其性能分析,对GB200系统上的nvfp4_gemv_fp4in_kernel分析显示以下瓶颈:

code
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的寄存器分配控制可以直接影响的领域。

code
# objective(): do_bench(...) 变为读取 ncu 指标
return _measure_occupancy(acf_path) # ncu: sm__warps_active.avg.pct_of_peak_sustained_active

# 搜索配置:最小化运行时间 → 最大化占用率
config = SearchConfiguration(problem_type="max", generations=4, pool_size=6)

将搜索目标从运行时间切换为占用率后,仅影响两个方面:目标函数的评分方式,以及搜索的 problem_type 参数。通过在 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 缓存吞吐量。

code
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 的搜索,帕累托前沿上收敛出一个占主导地位的候选方案。与权衡曲线不同,该方案在搜索自身的评分体系中,同时在两个维度上均优于所有其他已评估方案。在全新缓存上使用发货级试验次数验证:

code
=== 运行时间 (ms,越低越好) ===
p_value=3.516e-295  cohens_d=87.024
speedup_mean=1.0689x
=== L1TEX 吞吐量 (%) (越高越好;内部取反) ===
p_value=7.123e-19  cohens_d=3.686
基线 L1TEX 吞吐量均值=81.5087%
候选 L1TEX 吞吐量均值=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。

code
def objective(config) -> tuple[float, ...]:
    """一个共享的 ACF,每个形状对应一个运行时间得分(均需最小化)"""
    acf_path = ... # 将 `config` 写入临时 .acf 文件,或 None 表示基线
    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.0944 ms | 0.0308 ms pareto_2 | 0.0917 ms | 0.0321 ms

与 down_proj 的默认基准时间 0.0972 毫秒和 15360_5120 的 0.0315 毫秒相比,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 优化的目标,通过关联的高级控制文件(ACF)从该内核实例中提取最大性能。这些配置在运行时通过 Helion 的固定自调优配置应用,具体方式取决于每个内核实现的入口点,可通过 Triton 的 ptxas_options API 或直接通过 PTXAS 命令行实现。

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

需要重点注意的是,该方法可自动发现优化点,几乎不需要编译器专业知识,且会通过您自己的正确性检查进行验证。最终结果是一个可以提交到内核代码库并在每次构建时重新应用的文件。这是理论上的几个百分比提升与实际部署的几个百分比提升之间的区别。

它还会在您最希望的位置产生复利效应。推理循环中存在高价值内核,这些内核在大规模集群中持续运行。即使在堆栈底层的几个百分比提升,只要应用于主导解码步骤的操作,其规模效应也具有重要价值。

注意事项

  • ACF 与硬件和工具包相关:为 sm_90 调优的 ACF 在 sm_100 上无效。应用不匹配的 ACF 可能导致错误结果或运行时错误,而不仅仅是代码变慢。增强包清单会声明其目标架构和工具包版本。这也是为什么搜索循环中的正确性检查不可选。
  • 搜索时间是一个非零成本:autotune_search_acf 中的每个 ACF 都会增加您的自调优试验次数。请离线对关键内核使用此功能,然后固定最佳方案。
  • 收益取决于工作负载:结果取决于工作负载、形状、架构以及内核剩余的优化空间。对您的内核而言,唯一重要的数字是您在自己的内核上测量得到的那个。

如果你从未编写过 GPU 内核:从 pretuned_kernels/ 开始。选择一个接近你需求的内核,阅读它然后直接运行。这些内核都很简洁,能够立即以调优后的速度运行,无需等待自动调优。

如果你已经在使用 Helion:下载 helion-booster-pack,将 autotune_search_acf=[...] 添加到你关心的某个内核中,然后对它进行自动调优。无 ACF 的基线版本会保留在搜索范围内,因此损失是可控的,不会超过你投入的时间。参见 Helion 的 examples/acfs/ 目录,其中包含完整的示例。

如果你的工作涉及编写内核:这里的关键信息不是工具本身。而是编译器的控制接口是一个合法的优化空间,它是可搜索的,而大多数人尚未探索这一领域。这一原则适用于任何暴露可调参数的编译器,适用于任何后端。我们希望我们能激发你去探索这个空间!

授权许可

Helion 是开源项目,由 PyTorch 基金会托管。CompileIQ 可免费使用,遵循 NVIDIA 软件许可协议(NVIDIA Software License Agreement),这是一个专有的 EULA。它生成的 ACF 文件可以自由分发,这也是为什么可以发布 booster packs,以及为什么可以在开源仓库中将 ACF 提交到内核代码库中。

链接

  • Helion https://github.com/pytorch/helion https://helionlang.com https://github.com/pytorch/helion/tree/main/examples/acfs https://github.com/pytorch/helion/tree/main/pretuned_kernels
  • CompileIQ https://github.com/NVIDIA/CompileIQ https://nvidia.github.io/CompileIQ/stable/ https://github.com/NVIDIA/CompileIQ/releases?q=booster-packs&expanded=true https://developer.nvidia.com/blog/extract-more-kernel-performance-with-nvidia-compileiq-auto-tuning

/post-content

/inner-wrap