Skip to content
Published on

GPU 编译器与框架全景图 — 接收计算图、产出核函数,同一个问题,每一层给出不同的答案

🇰🇷KO🇺🇸EN🇯🇵JA🇨🇳ZH
Split View필사 모드
分享
Authors

引言 — 性能分析器指出的核函数,没有人亲手写过

用 Nsight Compute 抓一个训练步骤,会得到耗时最多的核函数名称。可如果这个名字长得像triton_poi_fused_add_mul_native_layer_norm_7,麻烦就来了。搜遍代码仓库也找不到,因为根本没有人写过它。是机器生成的。

要弄清楚究竟是什么生成了它,得先回答这些问题。这个核函数是在哪个阶段生成的?是 PyTorch 生成的,是 Triton 生成的,还是更底层的 LLVM 生成的?如果想改名字——不,如果想改代码——该动哪一层?

这个问题之所以难回答,不是因为层数多,而是因为每一层名字不同、所属组织不同、文档也各自独立,看上去彼此毫无关联。NVCC、LLVM、MLIR、Triton、Inductor、XLA、IREE、TVM,看起来就像来自完全不同的世界。

但实际上,它们解决的是同一个问题。本文先说清楚这个问题是什么,再梳理每一层切走了这个问题的哪一部分。如果说这个系列之前的文章分别讲了直接修改核函数、核函数编写的三个层次、推理引擎调优、两家厂商技术栈的差异,那么本文就是它们所处的坐标系。

版本与事实均于 2026 年 8 月 2 日直接核实自各项目的官方文档与发布说明。这个领域 6 个月就足以让格局改变,因此正文中会明确区分哪些内容已核实、哪些没有。

共同的问题 — 接收计算图,产出核函数

无论名字是什么,本文出现的这些工具输入和输出都是一样的。

输入是计算图:矩阵乘法、加法、激活函数、归一化这类节点通过张量相连而成。输出是真正在 GPU 上执行的机器码。而两者之间要做的事情,也惊人地相似。

  • 融合决定。把哪些节点合并成一个核函数。合并之后中间张量就不必写入全局内存,从而节省内存带宽。XLA 文档把融合称为"XLA 中单一最重要的优化"。
  • 布局决定。张量在内存中按什么顺序摆放。同一个运算,也会因布局不同而导致合并访存与否不同。
  • 分块决定。把大矩阵切成多大的块,分给每个线程块。块太大会导致共享内存和寄存器不够用,块太小则复用率下降。
  • 在调用库与生成代码之间做选择。这个运算,调用 cuBLAS 更快,还是直接生成核函数更快?
  • lowering是将已确定的内容落实为具体指令的阶段。

之所以会分出好几层,是因为这些决定的性质彼此不同。融合必须看到整张图才能判断,寄存器分配却只能在指令级别判断。让一个抽象层同时把两件事都做好并不容易。于是上层看图,下层看指令,中间还夹着若干中间层。

整理成表格是这样的。

层观察的单位做出的决定代表工具
计算图层整个模型融合、布局、分布式切分XLA、Inductor、IREE、TVM
核函数层单个核函数分块大小、流水线、共享内存布局Triton、CUTLASS、Helion
指令层基本块指令选择、调度、寄存器分配LLVM、ptxas
机器码层芯片实际编码SASS、AMDGCN

记住这张表再往下读,每个工具的位置就都清楚了。

最底层 — NVCC、PTX、SASS,以及 LLVM 的 GPU 后端

先从最底层看起。这是唯一一个 20 年来结构几乎没有大变化的地方。

把.cu文件丢给nvcc,它会拆成主机代码和设备代码。设备代码以两种形式产出:PTX 和 CUBIN。CUDA 13.3 编译器文档描述的流程是这样的:设备函数被编译成 PTX 或 CUBIN,再打包进 fatbinary 容器,原始源码则被转换成携带该 fatbinary 的标准 C++,交给主机编译器处理。运行时,"CUDA 运行时会检查内嵌的 fatbinary",挑出与当前 GPU 匹配的镜像。

这里需要理解的概念,是虚拟架构与实际架构的区分。

# compute_90 = 虚拟架构,即生成 PTX 所针对的指令集世代
# sm_90      = 实际架构,面向真实芯片的二进制
nvcc -gencode arch=compute_90,code=sm_90 \
     -gencode arch=compute_90,code=compute_90 \
     kernel.cu -o kernel

# 第一行把 Hopper 用的 SASS 放进 fatbinary,第二行把 PTX 本身放进去。
# 有了第二行,未来世代的 GPU 上驱动才能用 JIT 把它救活。

按文档的说法,把 PTX 一并放进去之后,"当没有匹配当前 GPU 的二进制加载镜像时,CUDA 运行时会动态编译内嵌的 PTX 代码"。这就是前向兼容性的全部内容。这也是 PTX 不会消失的原因。PTX 看起来像是给人读的汇编,但实际上是一种分发格式,也是一份稳定的契约。SASS 的编码每一代都会变,也没有文档记录,而 PTX 有文档记录,并且保持向后兼容。

在实践中,这个区分会在某个节点真正咬人。PyTorch 2.13 的发布说明显示,CUDA 13 构建版本不再把ptxas打包进二进制里。这意味着把 PTX 转成 SASS 的工具已经不在 wheel 包里了,依赖 PTX JIT 路径的环境,或是直接调用ptxas的自定义构建脚本,都可能突然崩掉。不需要了解某一层,和那一层可以不存在,是两回事。

从 PTX 到 SASS 的最后一段路,由ptxas负责。寄存器分配和指令调度就在这个阶段发生,决定占用率的寄存器数量也是在这里确定的。核函数源码没动,只是把 CUDA 工具包的次版本号升级了一下,性能却发生了变化——这种现象几乎总是因为ptxas的调度结果变了。可以用下面的工具确认。

# 查看实际生成的 SASS
cuobjdump -sass ./kernel | head -40

# 寄存器与共享内存用量(占用率计算的输入值)
nvcc -Xptxas -v -arch=sm_90 -c kernel.cu
# ptxas info : Used 64 registers, 8192 bytes smem, 384 bytes cmem[0]

# 只提取 PTX
nvcc -ptx -arch=compute_90 kernel.cu -o kernel.ptx

不过,能生成 GPU 机器码的并不只有 CUDA C++。LLVM 里有一个叫NVPTX的后端,能把 LLVM IR lowering 成 PTX。Triton 也好,XLA 也好,其他几乎所有代码生成器最终都会经过这个后端。也就是说,PTX 既是 CUDA 的产物,同时也是 LLVM 的产物。

AMD 那一侧结构相同,只是名字不同。LLVM 的AMDGPU后端把 LLVM IR lowering 成 AMDGCN 汇编,再变成代码对象。不过 AMD 并没有直接使用上游 LLVM,而是维护着自己的一个 fork。ROCm 文档明确写着,自家编译器是"llvm/llvm-project 的一个 fork",把amdclang++列为默认编译器,并把hipcc描述为"调用 clang 或 nvcc,并传入合适的 include 与 library 选项"的驱动程序。查阅时文档记载的 llvm-project 版本是 22.0.0。

维护一个 fork 本身并不是坏信号。新芯片的支持和新功能,往往比上游发布周期需要得更快,而且 AMD 的改动有相当一部分最终会合并回上游。但这在实务中确实有影响。以 AMD GPU 为目标的第三方编译器项目,必须在自己的 LLVM 版本和 ROCm 的 LLVM 版本之间做协调,这正是对比两家厂商技术栈的文章中提到的移植摩擦的一支。

MLIR 与建立在它之上的东西 — 制造 IR 的框架,以及 Triton

从这里开始,地形变了。

LLVM IR 是为 CPU 设计的。标量与向量、基本块、SSA。它没有张量这个概念,也没有循环嵌套这个概念。用 LLVM IR 表示矩阵乘法时,循环早已被展开成低级代码,"把这个循环分块"之类的决定所需要的信息已经不存在了。

于是,2010 年代后期的每一个 ML 编译器都在重复同一件事:重新定义张量级别的 IR,重新搭建 pass 基础设施,重新写解析器、验证器和打印器,最后再重新写一遍把它 lowering 到 LLVM IR 的代码。XLA 这样做过,TVM 这样做过,每一家硬件厂商也都各自这样做过。

MLIR 的主张,是消灭这种重复。它不再多定义一种 IR,而是提供用于定义 IR 的基础设施。一组运算与类型的集合被称为方言,每个项目定义自己的方言,但共享 pass 管理、验证、序列化和 lowering 框架。

关键在于方言之间可以共存。同一个模块里可以混杂着抽象层次不同的运算,而 lowering 就是逐步消解这种混杂的过程。

// 同一个函数内,两种抽象层次不同的方言共存。
// linalg 保留着"这是一次矩阵乘法"的语义,
// scf/arith 则已经被 lowering 到循环和标量运算。
func.func @mm(%A: tensor<128x256xf32>, %B: tensor<256x64xf32>)
    -> tensor<128x64xf32> {
  %init = tensor.empty() : tensor<128x64xf32>
  %C = linalg.matmul
      ins(%A, %B : tensor<128x256xf32>, tensor<256x64xf32>)
      outs(%init : tensor<128x64xf32>) -> tensor<128x64xf32>
  return %C : tensor<128x64xf32>
}

linalg.matmul这一行,原样保留着"这是一次矩阵乘法"的信息。分块 pass 可以看着这个信息去切分循环。如果先把同样的东西 lowering 到 LLVM IR,这个信息就没了,之后想重建循环会困难得多。信息保留得越久,才能做出越好的决定 — 这就是 MLIR 存在的理由,浓缩成一句话。

MLIR 本身不编译任何东西。在它上面搭建什么,才是全部重点。而实际上,搭在它上面的东西已经很多了。Triton 的 IR、IREE 的 IR、XLA 的部分组件、多家硬件厂商的内部编译器,都建立在 MLIR 之上。这些名字各异的项目共享着同一个底座,这个事实对理解这片地形非常重要。

建立在 MLIR 之上的东西 — Triton 生成核函数的路径

Triton 既是一门核函数编写语言,同时也是一个基于 MLIR 的编译器。对比核函数编写三个层次的文章已经讨论过作为语言的 Triton,这里来看它作为编译器的定位。

这条流水线的定义非常清晰。PyTorch 团队在 2026 年 7 月发布的 Triton 插件扩展文档,是这样描述各个阶段的。

Python 核函数
      ↓  (AST 遍历)
TTIR    — Triton IR。以块为单位的运算。与硬件无关。
      ↓  (布局决定、线程映射)
TTGIR   — TritonGPU IR。warp 排布与共享内存已经确定。
      ↓
LLVM IR
      ↓
PTX (NVIDIA)  或  AMDGCN (AMD)

这四个阶段,把本文前面梳理的"层"结构浓缩展示了出来。TTIR 对应核函数层,TTGIR 是核函数层与指令层的边界,LLVM IR 以下则是指令层。用户大致在 TTIR 这个层级写代码,剩下的都交给编译器。这就是 Triton 提出的交易条件。

按仓库 README 的说法,受支持的后端是 NVIDIA(Compute Capability 8.0 及以上)和 AMD(ROCm 6.2 及以上),CPU 则标注为"开发中"。CPU 后端有好几路尝试在推进,但截至查阅时仍处于实验阶段,在实际工作中还不能对它抱太大期待。

2026 年这一层真正发生、值得关注的变化,是插件扩展体系。它进入了 PyTorch 分发的 pytorch-triton 3.7,PyTorch 2.13 把 Triton 版本锁定在 3.7.1。以前想加入自定义 pass 或方言,就必须 fork Triton。fork 很快会落后于上游,合并冲突不断累积,最终被困在一个拿不到新硬件支持的过时版本里。插件体系允许在运行时加载共享库,在流水线的任意位置插入 pass、关闭某个特定 pass,或者整段替换。

# 插件通过环境变量加载,import 的顺序很重要。
import os, sysconfig

dist_packages = sysconfig.get_paths()["purelib"]
os.environ["TRITON_PLUGIN_PATHS"] = os.path.join(
    dist_packages, "utlx_plugin", "libutlx.so"
)

import triton
import triton.language as tl
import utlx_plugin as tlx      # TLX 扩展运算从这里引入

第一个使用者是 Meta 的 TLX(Triton Language Extensions)。它通过tlx.local_alloc、tlx.async_load、tlx.async_dot这类运算,在核函数内部显式处理共享内存缓冲区与异步流水线。也就是说,这是把 Triton 原本自动做出的一部分决定,重新交还给人。公开的数据显示,在 H100 的 persistent GEMM 上比 cuBLAS 快 2%~4%,在 MI350 的流水线 GEMM 上比 rocBLAS 快 12%~15%。要点在于:比厂商库更快的核函数,是从一门编译器语言里长出来的。

这里能看到一种走向。抽象层并不是只会往上走。如果说 Triton 靠隐藏线程往上走了一步,那么 TLX 就是靠重新暴露异步指令往下退了一步。抽象的高度不是由性能决定的,而是由哪个决定掌握在人手里决定的。

torch.compile 与 Inductor 到底在做什么

这是 PyTorch 用户最常碰到、却理解得最少的一层。torch.compile(model)这一行代码背后,依次运转着三个阶段。

TorchDynamo拦截 Python 字节码,从中抽取出 FX 图。纯张量运算会进入图中,而 Python 列表操作或print这类东西则会被挤到图外。这个断点就是 graph break。图一旦被切成好几段,每一段前后都要付出退回 Python 再跳出去的开销,而且段越小,融合的机会也越少。torch.compile没有达到预期加速效果的案例,大多数不是因为核函数写得差,而是因为图被切碎了。

AOTAutograd把这张图拆成前向和反向两张 ATen 级别的图,并决定哪些结果要保存、哪些要重新计算。

真正生成代码的是TorchInductor。目标是 GPU 就生成 Triton 核函数,目标是 CPU 就生成 C++ 与 OpenMP 代码。前面见到的triton_poi_fused_...这类名字,就是这么来的。名字本身携带着信息:poi是 pointwise 的缩写,fused_后面跟着的,就是被合并在一起的那串运算列表。

生成出来的代码是可以直接看到的。这是理解这一层最快的办法。

import os
os.environ["TORCH_COMPILE_DEBUG"] = "1"        # 把生成的代码留在磁盘上
os.environ["TORCH_LOGS"] = "output_code,graph_breaks"

import torch

def block(x, w, b):
    return torch.nn.functional.gelu(x @ w + b)

x = torch.randn(1024, 512, device="cuda", dtype=torch.bfloat16)
w = torch.randn(512, 512, device="cuda", dtype=torch.bfloat16)
b = torch.randn(512, device="cuda", dtype=torch.bfloat16)

compiled = torch.compile(block)
out = compiled(x, w, b)

# 单独查看 graph break 的方法
explanation = torch._dynamo.explain(block)(x, w, b)
print("图的数量:", explanation.graph_count)
print("break 数量:", explanation.graph_break_count)

打开TORCH_LOGS="output_code",Inductor 生成的 Triton 源码就会打印到标准输出。可以看到矩阵乘法通常仍然保留为 cuBLAS 调用,而 GELU 和偏置加法被合并成了一个 pointwise 核函数。那个核函数,正是前面性能分析器指出的那个无名核函数的真身。

2026 年,这一层还发生了一件重要的事。从 PyTorch 2.13 开始,Inductor 除了 Triton 之外,又多了一条CuTeDSL 路径。这是构建在 NVIDIA CuTe 之上的 DSL,根据发布说明,它是瞄准 transformer 的 GEMM 与 RMSNorm 打造的第二条高性能路径,即便"没有 Triton 也能生成质量更好的矩阵乘法代码"。同一个版本里,核函数编译也从线程池搬到了子进程池,原因是 Python 的 GIL 一直是编译时间的瓶颈。

同一个版本里还加入了torch.compiler.set_default_backend。它和torch.set_default_dtype是同样的思路:设置一个进程级别的全局默认后端,这样每次调用torch.compile就不必都带上backend=了。对于在树外构建后端的硬件厂商来说,这是一个意义不小的变化。

总结一下,torch.compile是一个计算图编译器,核函数生成的工作被委托给下面一层。这个下面一层曾经只有 Triton 一个选择,现在变成了两个。

XLA、IREE、TVM — 整张图直接接收的一派

如果说 Inductor 是依附在 PyTorch 上的计算图编译器,那么还存在另一条谱系,主张独立于任何单一框架来接收计算图。

XLA资历最老,用得也最广。它是 JAX 唯一的执行路径,也是 TPU 唯一的执行路径,PyTorch 也通过 PyTorch/XLA 接入。它的输入是 HLO,框架与编译器之间的契约则由 StableHLO 承担。StableHLO 是一个承诺向后兼容 5 年、向前兼容 2 年的运算集合,正因如此,把图存下来,几年后用新的编译器打开它,才能行得通。

XLA:GPU 的代码生成方式,和本文的主题正好严丝合缝地对上了。照搬官方文档的说法,XLA:GPU 使用的是"原生(经由 LLVM 生成 PTX)发射器与 TritonIR 发射器的组合"。这分成三条路径。

  1. 常见运算,直接调用 cuBLAS、cuDNN、NCCL。
  2. 像归约或转置这类能识别出模式的运算,直接生成 LLVM IR 并 lowering 成 PTX。
  3. 涉及矩阵乘法或 softmax 的高级融合,把 HLO 融合转换成 TritonIR,选定分块参数后调用 Triton,拿回 PTX。

第三条最重要。XLA 和 Triton 不是竞争关系,而是调用关系。XLA 做计算图层的决定,核函数层的决定则交给 Triton。前面表格里划分的层次,在这里被原样实现了出来。最后,XLA 运行时会把核函数调用与库调用的序列,搬到一个叫 RuntimeIR 的自有 MLIR 方言里,再从中提取出 CUDA 图。

IREE在一个宽得多的目标范围里解决同一个问题。用 README 的话说,它是一个基于 MLIR 的编译器兼运行时,能"把 ML 模型 lowering 成统一的 IR,向上扩展到数据中心的需求,向下收缩到移动端与边缘设备的限制"。仓库的主题标签里,前端列出了 JAX、PyTorch、ONNX、TensorFlow,目标平台列出了 CUDA、ROCm、Vulkan、SPIR-V。它在 2024 年 5 月作为沙箱阶段项目加入了 LF AI & Data 基金会。

IREE 的特点在于,编译器和运行时是一起设计的。它产出的不是需要 Python 进程才能跑的对象,而是能独立运行的模块,因此同一条流水线能一路延伸到没有 Python 的嵌入式目标平台。在数据中心里遇到 IREE 的情况目前还很少见,查阅时最大的采用案例,是 AMD 把它作为自家技术栈的一部分在推广。

TVM是这片地形里位置变化最大的项目。2018 年前后,它是唯一一个真正拿得出手的开源深度学习编译器,并且凭借"用自动调优战胜厂商库"的主张吸引了大量关注。如今它是一种跨层设计:用 Relax 做计算图级别的表示,用 TensorIR 做张量级别的表示,文档写明的设计目标是"让大部分变换都可以在 Python 里定制,从而让 ML 编译器变得平易近人"。

老实说,在前沿 LLM 的训练和推理服务场景里,现在已经几乎碰不到 TVM 了。这个位置被 Inductor 和 XLA 拿走了。TVM 还活跃的地方在别处:没有厂商库的硬件、跑不起 Python 运行时的边缘目标平台,以及像 MLC LLM 那样在浏览器或移动端跑 LLM 的部署路径。"自动调优能战胜经过调优的库"这个最初的命题,在 cuBLAS 和 CUTLASS 存在的地方很难成立,而在它们不存在的地方,依然成立。

汇总成一张表是这样的。

XLAIREETVMInductor
主要前端JAX、TF、PyTorch/XLAPyTorch、ONNX、JAX、TF从多种框架导入仅限 PyTorch
计算图 IRHLO / StableHLOMLIR 方言RelaxFX → ATen
核函数生成直接生成 LLVM + 调用 TritonMLIR loweringTensorIR + 自动调优Triton、CuTeDSL、C++/OpenMP
运行时XLA 运行时、PJRT自有轻量运行时自有运行时位于 Python 进程内部
优势领域TPU、JAX、大规模训练从边缘到数据中心非主流硬件、边缘部署PyTorch 的默认路径

什么时候该了解哪一层

层多,不代表都得懂。只需要在真正有必要时才往下探。而"必要"的信号到底是什么,是本文最实用的部分。

你正遇到的症状该下探的层具体该做的事
训练就是慢还没到任何一层先上性能分析器。数据加载器和同步点是大多数问题的根源
开了 torch.compile 却没变快计算图层用 torch._dynamo.explain 检查 graph break
编译本身耗时太长计算图层检查是否存在动态 shape 引发的重新编译,考虑 mode="reduce-overhead"
分析结果里出现了生成的核函数名字,但不知道它是什么核函数层用 TORCH_LOGS="output_code" 读 Triton 源码
Inductor 生成的核函数比库还慢核函数层直接用 Triton 编写,或替换成自定义 op
Triton 核函数怎么都达不到理论带宽附近核函数与指令的边界用 TTGIR dump 检查布局与共享内存排布
怀疑发生了寄存器溢出指令层用 -Xptxas -v 或 ncu 检查寄存器数量
只是升级了工具包版本,性能就变了指令层对比 SASS 反汇编
需要把框架接到新硬件上所有层从 MLIR 方言与 PJRT 插件设计开始

根据经验,90% 的从业者到上面这两行就结束了。消灭一个 graph break,通常比手工优化核函数带来的收益更大。往下探得越深,能拿到的收益越大,但花费的时间也会以更陡的斜率增长。

如果想直接查看 Triton IR,dump 的方法是这样的。

# 把每个阶段的 IR 都保留成文件
export TRITON_KERNEL_DUMP=1
export TRITON_DUMP_DIR=/tmp/triton_ir
python train.py

ls /tmp/triton_ir/*/
# *.ttir     TTIR   — 以块为单位,与硬件无关
# *.ttgir    TTGIR  — 布局与 warp 排布已经确定
# *.llir     LLVM IR
# *.ptx      PTX
# *.cubin    包含 SASS 的最终二进制

把同一个核函数的这五个文件并排打开看一遍,会比读完这整篇文章更快建立理解。从上往下走,什么被决定、什么消失了,都会看得一清二楚。

最后,把这次调查中未能核实的内容列在这里。

  • Triton 稳定发布版本的号码和日期。仓库 README 没有写明版本号,能确认的只有 PyTorch 锁定的 3.7.1。Triton 自身的发布编号,和 PyTorch 分发的 pytorch-triton 编号,未必一致。
  • IREE 最新发布标签的号码和日期。README 里没有显示,也没有另行核实。
  • Apache TVM 近期的发布节奏和活跃程度。项目文档确认了设计方向,但没有直接查阅发布历史。上文关于采用现状的描述,是一种判断,并没有统计数据支撑。
  • Helion。确认了它是 PyTorch 在 2026 年力推的核函数 DSL,并且 TPU 后端会编译到 Pallas,但在 NVIDIA 与 AMD 的路径上它与 Triton 是什么关系,没能在参考的文档中得到确认。

结语 — 层的划分依据是决定权,不是性能

把本文这张地图折成一句话,就是:层与层之间划分的标准不是速度,而是谁掌握着哪个决定。

计算图层拿走了融合与布局。作为代价,它放弃了单个核函数的细节。核函数层拿走了分块与流水线。作为代价,它看不到整张图。指令层拿走了寄存器与调度。作为代价,它不知道自己算的到底是什么。这就是为什么没有任何一层能取代另一层——这与 30 年前 CPU 编译器经历的层次分化,结构上是一样的。

所以,每当出现一个新项目的名字,该问的问题也是固定的。不是问它哪里新,而是问它从谁手里拿走了哪个决定。Triton 从人手里拿走了线程映射。TLX 又从编译器手里把异步指令抢了回来。Inductor 从用户手里拿走了融合。MLIR 从每个项目手里拿走了 IR 基础设施。

光靠这一个问题,基本就能确定一个新出现的框架该放在这张地图的什么位置。而这,比追着每 6 个月就换一批的名字跑,更经得起时间。

参考资料

Discuss on Twitter • View on GitHub