第九部分 · 基础设施与算力 · 第 64 章

编译器与内核

作者Changkun Ou

阅读时长约 18 分钟

框架程序距离可执行的设备代码还有好几层契约。编译器必须对守卫覆盖的每一种输入保留程序语义，同时选择布局、融合边界、分块形状、库调用和目标指令。随后，内核必须遵守承诺的形状、步幅、dtype、布局、别名、修改、数值和目标设备语义。一个结果即使很快，只要违反其中任一条件，就不是优化，而是另一个程序。

本章沿着第 63 章引出的执行路径继续向下。我们先介绍指导融合与分块的成本模型，再用 FlashAttention 说明算法如何围绕内存层级重新设计，然后把一张计算图从中间表示一路追踪到目标代码。最后几节会区分可移植性的多种含义，并解释为什么生成内核需要比普通应用代码更严格的评估者。量化算术及其专用矩阵内核仍是第 34 章的重点。

内核实现一份有边界的契约

逻辑算子说明结果意味着什么。内核是在声明域上实现该算子的一种方式。在 GPU 上，一次启动会创建线程块网格；设备把线程块调度到流式多处理器上，多组线程通道则以 warp 或 wavefront 为单位共同执行。同一个线程块中的线程可以通过共享内存与同步协作。寄存器属于单个线程，全局设备内存对整个网格可见 (Nickolls et al. 2008; NVIDIA 2026)。

启动域和内核本体同样重要。连续矩阵与它的转置视图可以有相同的逻辑形状却有不同的步幅。两个参数可能指向同一块存储。归约可能承诺特定的累加 dtype 或并列值处理规则。随机算子会推进状态，原地算子则会修改其他位置也能看见的存储。编译器可以根据这些事实特化内核，但必须在运行时守卫这些事实，或保留正确的通用路径。

由此可以区分五项责任：

层次	它看到的输入	必须保留的契约
计算图捕获	程序执行与守卫	守卫成立时，保持相同的可观察张量与状态副作用
张量中间表示（IR）	算子、形状、dtype、副作用	算子语义与合法变换
内核 IR	分块、地址计算、内存空间	完整写入、合法访问、顺序与同步
目标代码	指令或外部库调用	与选定设备及运行时兼容
运行时	缓冲区、流、启动顺序、编译缓存	生命周期、依赖关系、失败和回退行为

明确这些边界之后，性能工作才开始。

屋顶线是上界，不是秒表

屋顶线模型把算术工作量与跨越指定内存边界的流量联系起来 (Williams et al. 2009)。对一个内核，定义

I=F/Q,

其中：

$F$ 是为该内核选定的浮点运算计数；
$Q$ 是跨越测量边界传输的字节数，例如从高带宽内存（HBM）到芯片的字节数；
$I$ 是每字节浮点运算数表示的算术强度。

对同一个目标和精度采用算力上限与带宽上限，可以得到下面的约束。这里，两项分别代表算力上限与流量上限：

P_{\text{bound}}=\min\left(P_{\max}, B_{\max}I\right),

进而得到算法层面的时间下界：

T\geq \max\left(\frac{F}{P_{\max}},\frac{Q}{B_{\max}}\right).

其中：

$P_{\text{bound}}$ 是吞吐量的屋顶线上界；
$P_{\max}$ 是相关的算力上限，以每秒浮点运算数计，并且要使用与 $F$ 相同的 dtype 与指令路径；
$B_{\max}$ 是相关访问模式实际达到的带宽上限，以每秒字节数计，最好来自实测，而不是直接抄录规格表；
$T$ 是以秒计的执行时间。

这里，屋脊强度为 $I_{\text{ridge}}=P_{\max}/B_{\max}$ 。低于它时，减少流量可以抬高上界；高于它时，减少算术量或使用更快的计算路径可能有效。这是一种分类方法和乐观极限，不是运行时间预测。启动开销、依赖延迟、并行度不足、指令组合、缓存行为、未合并访问、bank 冲突、同步和占用率都可能形成更低的上限。因此，观察到的性能应该采用性能分析器给出的实际达到的带宽和算力，而不是只看名义峰值。

Transformer 的性能剖析说明了这种区别为什么重要。Ivanov 等人发现，张量收缩与内存受限算子都在 BERT 运行时间中占有显著份额；减少数据移动确实加速了他们测量的工作负载，但这并不意味着每个 Transformer 或每个算子都受带宽限制 (Ivanov et al. 2021)。

下面的可运行示例只计算一个刻意简化的案例。它在 FP32 向量上计算 $y=\max(ax+b,0)$ ，其中 $x$ 是输入向量， $a$ 和 $b$ 都是标量。未融合版本会写出两个完整中间结果；融合版本则把两个标量中间值保留在同一个内核里。计算还假定每个实体化向量都完整读写 HBM，而且没有可利用的跨内核缓存保留。它统计的是最小算法流量，不是实测设备流量或时间。

elements = 16 * 1024 * 1024
bytes_per_element = 4  # FP32
operations_per_element = 3  # multiply, add, maximum

operations = elements * operations_per_element
unfused_bytes = 3 * 2 * elements * bytes_per_element
fused_bytes = 2 * elements * bytes_per_element

for name, traffic in [("unfused", unfused_bytes), ("fused", fused_bytes)]:
    intensity = operations / traffic
    print(f"{name:7s}: {traffic / 1e6:5.1f} MB minimum traffic, "
          f"{intensity:.3f} operations/byte")

print(f"same arithmetic: {operations} operations")
print(f"traffic reduction: {unfused_bytes / fused_bytes:.1f}x")

融合是合法性与资源决策

融合把生产者与消费者操作合进一个内核，使中间结果不必在全局内存中实体化。它还可以减少启动次数，并暴露代数简化机会。只有当合并后的内核仍然合法，而且适合目标硬件时，这些收益才真正存在。

编译器首先要证明融合会保留依赖关系和可观察副作用。修改与别名关系可能让看似无用的中间结果被其他位置观察到。随机数状态、宿主回调、集合通信、原子操作和异常都会限制重排。归约尤其需要谨慎，因为新的归约顺序会改变舍入、NaN 传播和确定性行为。即使只是逐点融合，也可能引入融合乘加指令或重新结合运算，因此数学等价并不意味着逐位一致。

随后才轮到成本决策。融合会延长值的存活区间，增加寄存器压力或共享内存需求。若数值溢出到本地或全局内存，流量会以代价更高的形式回来。更大的内核可能降低占用率，减少调度自由度，重复计算，或阻止系统调用高度优化的库。因此，过度融合确实可能发生。实用的编译器会估算减少的字节数、启动次数、并行度和资源用量，并对不确定的选择进行剖析，而不是把融合当作无条件规则 (Ansel et al. 2024)。

分块把工作映射到有限资源

分块从更大的计算中选取一小块工作，并在处理器附近复用它的操作数。以矩阵乘法为例，一个线程块可以把 $A$ 与 $B$ 的块从全局内存载入共享内存，完成同步，在寄存器中执行许多乘加，然后沿归约维度重复。相邻线程通道应采用合并的全局内存访问，让硬件用尽量少的事务完成服务。

每一种分块形状都是一组耦合选择。更大的块能产生更多复用，却会消耗更多寄存器和共享内存。让更多线程块同时驻留可以隐藏延迟，但最大占用率本身并不是目标：如果每个块完成更多有用工作，占用率较低的内核仍然可能更快。bank 冲突、同步频率、尾部掩码、张量单元对齐和流水线阶段数也会改变结果。当形状或架构变化时，编译器通常用启发式方法缩小候选范围，再通过自动调优实测有限的一组调度方案 (Chen et al. 2018; Tillet et al. 2019; Triton Project 2026)。

自动调优结果是带条件的产物。某个分块方案在一种形状、dtype、步幅模式、设备、驱动或库版本上胜出，换一个条件就可能落败。缓存键必须包含选择所依赖的每一项事实；如果内核会修改输入，还必须在各次试验之间恢复输入。

FlashAttention 围绕 IO 改写算法

单个注意力头的缩放点积注意力可以写成下面的形式，各符号紧接着定义：

S=QK^\mathsf{T}/\sqrt{d}+C, \qquad P_{ij}=\frac{\exp(S_{ij}-m_i)}{\ell_i}, \qquad O=PV,

其中

m_i=\max_{1\leq j\leq N}S_{ij}, \qquad \ell_i=\sum_{j=1}^{N}\exp(S_{ij}-m_i).

各符号含义如下：

$Q,K\in\mathbb{R}^{N\times d}$ 是查询矩阵与键矩阵；
$V\in\mathbb{R}^{N\times d_v}$ 是值矩阵；
$N$ 是序列长度， $d$ 是查询或键的头宽度， $d_v$ 是值宽度；
$C\in(\mathbb{R}\cup\{-\infty\})^{N\times N}$ 是加性偏置或掩码，其中 $-\infty$ 表示禁止访问某个键；
$S$ 是缩放后的分数矩阵， $m_i$ 是查询行 $i$ 的最大值， $\ell_i>0$ 是该行的归一化因子；
$P$ 是逐行 softmax 矩阵， $O$ 是注意力输出， $\mathsf{T}$ 表示转置；
假定每个查询行至少有一个允许访问的键。

直接实现可能在 HBM 中实体化完整的 $N\times N$ 分数矩阵和概率矩阵。FlashAttention 不这样做，它让一个查询行或查询块依次处理多个键值块，只在高速内存中保留分块与运行统计量 (Milakov and Gimelshein 2018; Dao et al. 2022)。对固定查询行 $i$ ，初始化 $m_0=-\infty$ 、 $\ell_0=0$ 和 $o_0=0\in\mathbb{R}^{d_v}$ 。对键值块 $\mathcal{B}_t$ ，各项递推关系为

\begin{aligned} b_t &= \max_{j\in\mathcal{B}_t} S_{ij},\\ m_t &= \max(m_{t-1},b_t),\\ \alpha_t &= \exp(m_{t-1}-m_t),\\ \ell_t &= \alpha_t\ell_{t-1} +\sum_{j\in\mathcal{B}_t}\exp(S_{ij}-m_t),\\ o_t &= \alpha_t o_{t-1} +\sum_{j\in\mathcal{B}_t}\exp(S_{ij}-m_t)v_j. \end{aligned}

处理完 $T$ 个块后，

O_i=o_T/\ell_T.

其中：

$\mathcal{B}_1,\ldots,\mathcal{B}_T$ 把允许访问的键索引分区， $t$ 是块索引；
$b_t$ 是当前键值块中的最大分数；
$m_t$ 是运行最大值，最大值变化时， $\alpha_t$ 会把先前贡献重新缩放；
$\ell_t$ 是运行归一化因子；
$o_t\in\mathbb{R}^{d_v}$ 是运行中的值加权分子；
$v_j$ 是 $V$ 的第 $j$ 行， $O_i$ 是第 $i$ 个输出行。

在精确算术中，这套递推保留了与完整 softmax 相同的分子和分母。“精确注意力”指的是同一个数学函数，不是逐位一致，因为分块改变了浮点求值顺序。算法仍然执行二次复杂度的注意力算术。它的收益在于不会在 HBM 中实体化完整的二次中间结果，而且反向传播可以重算局部分数块和概率块，而不必保存它们。

在原论文采用的两级内存模型中，当前向算法满足 $d\leq M\leq Nd$ 时，它在慢速内存中的标量字传输次数为

H_{\text{FA}}=\Theta\left(\frac{N^2d^2}{M}\right)

其中， $H_{\text{FA}}$ 是传输次数， $M$ 是以标量字计的高速内存容量。这是一项 IO 结论，不是 FLOP 数或字节数；实际字节数还取决于存储 dtype。后续工作在 $M\geq d^2$ 时证明了常数因子内相匹配的逐点下界，并发现阈值以下存在另一种情形 (Saha and Ye 2024)。因此，不能声称原始算法对所有可能的 SRAM 容量都最优。

FlashAttention-2 在论文报告的 A100 配置上改进了工作划分，并减少非矩阵运算 (Dao 2024)。FlashAttention-3 在报告的 H100 结果中使用了 Hopper 特有的异步复制和矩阵指令 (Shah et al. 2024)。这些案例体现了一条普遍规律：算法层面的 IO 节省可以跨目标迁移，最佳调度通常仍与目标有关。

降级在增加决策的同时保留语义

编译器采用多层中间表示，因为没有单一层级适合做出所有决策。高层张量 IR 仍能表达捕获计算图中的广播、归约、副作用与符号维度。低层内核 IR 则能表达分块、地址计算、内存空间、屏障和目标指令。渐进式降级会明确写出每一项新决策，同时保留检查决策所需的信息。Halide 奠定了算法与调度分离这一重要思想；TVM 加入张量计算图优化和测量引导的调度搜索；MLIR 则为定义与组合多层 IR 提供可复用基础设施 (Ragan-Kelley et al. 2013; Chen et al. 2018; Lattner et al. 2021)。

图 64.1. 一条通用降级流水线。可移植张量语义逐步收窄为目标特定的调度与产物；生成内核和外部库调用都是合法的叶节点。

这张图刻意保持通用。编译器可以把高层操作保留为库调用，而不是生成内核。矩阵乘法、卷积和集合通信经常走这条路径，因为厂商库或框架库已经包含许多面向具体目标的算法。其他区域则降级到内核 IR，让融合、布局和资源选择变得显式。

PyTorch 的守卫路径

对 torch.compile 而言，TorchDynamo 从 Python 执行中捕获带守卫的 FX 计算图片段。训练还会加入 AOTAutograd，用它捕获并划分前向与反向计算图，同时应用分解和函数化。随后，TorchInductor 调度循环级 IR，并可生成 Triton GPU 内核、C++ CPU 内核、模板内核，或调用外部算子 (Ansel et al. 2024; PyTorch Contributors 2026)。计算图中断会把不受支持的工作交还即时执行。守卫失败则可能选中另一条编译缓存记录、触发重新编译，或转入更通用的路径。

动态形状并不会消除特化。符号维度与守卫可以让某些编译区域覆盖多种输入，但数值、秩、步幅和目标特性仍可能需要不同代码。因此，编译延迟、守卫失败、计算图中断和缓存增长都属于生产成本模型，而不是前端无关紧要的细节。

JAX 与 OpenXLA

JAX 把特化的 Python 函数追踪为 jaxpr。降级过程通过 StableHLO 把程序送入 XLA 的优化与代码生成流水线。XLA 会完成与目标有关的融合、布局分配、缓冲区规划、集合通信处理，以及内核或库选择，然后由运行时载入目标可执行程序 (Frostig et al. 2018; OpenXLA Project 2026; OpenXLA Project 2026)。

StableHLO 定义可移植张量语义，并为可移植产物规定兼容窗口。它不会固定物理布局、调度、每个消费者上的数值精度或性能，这些决策属于消费者与目标。JAX 也会根据形状、dtype 和静态参数等事实做特化，因此它的编译缓存需要与其他分阶段系统同等严肃的运行管理。

布局与自动调优都是编译的一部分

逻辑形状并不能决定物理布局。编译器可以选择哪个维度连续，加入填充或分块，把缓冲区放进不同内存空间，或把分块元素分配给不同线程通道与寄存器。有些转置只改元数据，另一些会移动整个张量。库调用和矩阵指令会施加布局约束，在两种原本都合法的布局之间转换，可能抵消快速内核带来的收益。

调度选择同样具有条件性。Triton 把程序员的工作单位从标量线程提升为分块程序，编译器则负责合并访问、向量化、共享内存分配和指令选择等任务 (Tillet et al. 2019; Triton Project 2026)。作者仍需选择网格、掩码、块大小、warp 数或阶段数，有时还要选择目标特定的机制。TVM 风格的学习成本模型与经验自动调优会在这些选择中搜索，但不会让某一种调度普遍最优。

可移植性不止一种

声称某个“可移植内核”时，必须说清边界：

可移植性主张	它意味着什么	它不保证什么
源码可移植性	多个后端都能接受同一种源语言	功能覆盖相同，或生成代码完全一致
语义可移植性	各后端实现声明的算子行为	浮点结果逐位一致
产物可移植性	序列化 IR 或二进制文件能在承诺的兼容窗口内载入	能够移出该产物生态系统
性能可移植性	同一种实现能在多个目标上保持高效	由源码或语义可移植性自然推出
运行可移植性	工具、部署、可观测性和回退在每个目标上都能工作	由一次成功的内核启动自然推出

PTX 是 NVIDIA 的虚拟指令集，不是通用 GPU IR。CUDA 设备代码可以发布为 PTX、目标特定的 cubin，或包含多种镜像的 fat binary。PTX 可以提前汇编，也可以由驱动即时编译；cubin 的兼容范围更窄 (NVIDIA 2026)。AMD 后端可能通过 AMD LLVM 降级到 HSACO 产物，TPU 后端则生成自己的可执行程序。这些分支可以共享高层张量语义，却不必共享最终调度或二进制文件。

Triton 也遵循同样的区别。它的分块编程模型可以减少受支持目标之间的源码重复，但后端覆盖、可用指令、布局选择和调优配置仍会不同。性能可移植性往往需要目标特定的调度，而且随着硬件和编译器变化，这套调度必须重新验证。

生态系统也是目标的一部分

CUDA 最初是一套编程与执行模型，但应用实际依赖更大的契约：驱动程序与运行时、编译器与产物格式、数学库与通信库、性能分析器与调试器、框架集成，以及围绕这些组件积累的测试与运行知识 (Nickolls et al. 2008; Chetlur et al. 2014)。其他目标也有对应的软件栈。因此，只比较内核语言或峰值算术速度，会漏掉很大一部分切换成本。

库尤其重要。它们的契约包括支持的 dtype 与布局、工作区限制、算法选择、确定性、流语义、集合通信顺序、ABI 兼容性和目标代际。迁移工作负载意味着为每一项已依赖的契约找到等价行为，再用实际形状分布验证数值结果和性能。源代码翻译可以减少工作量，却不能免除验证。

这也是一项研究约束。“硬件彩票”描述了一个现象：容易映射到现有硬件与软件的想法，可以用更低成本、更快完成实验 (Hooker 2021)。一种新架构表现不佳，可能是因为缺少融合扫描、稀疏操作或布局感知的编译器路径，而不是数学思想本身薄弱。内核可用性会影响哪些模型设计能够得到可信评估。

生成内核需要一名挑剔的评估者

KernelBench 用 250 个 PyTorch 工作负载和一项同时要求正确性与速度的指标，让生成内核评估变得具体 (Ouyang et al. 2025)。对任务集 $\mathcal{D}$ ，该指标的一种写法是

\operatorname{fast}_p = \frac{1}{|\mathcal{D}|} \sum_{i\in\mathcal{D}} \mathbf{1}\!\left[ \operatorname{correct}(k_i) \land t(k_i)<\frac{t(b_i)}{p} \right].

其中：

$\mathcal{D}$ 是基准任务集， $|\mathcal{D}|$ 是其大小；
$k_i$ 是任务 $i$ 的生成内核， $b_i$ 是其基线；
$t(\cdot)$ 是基准规程下经过同步的执行时间；
$p>0$ 是要求达到的加速阈值；
$\operatorname{correct}(k_i)$ 是基准给出的语义判定；
条件成立时 $\mathbf{1}[\cdot]$ 为一，否则为零。

执行与性能分析反馈可以改进结果，多轮训练也已在固定基准划分上显示收益 (Baronio et al. 2025)。这些结果并不能证明生成内核普遍可靠。一个内核可能只在少数形状上通过测试，却只写入输出的一部分；也可能依赖狭窄输入分布、修改输入、调用参考路径，或利用分配器与同步状态。后续审计已经表明，更广的隐藏分布与更严格的基线可能逆转表面上的胜利 (Zhang et al. 2026)。这与第 27 章讨论的是同一个评估者问题，只是这里又加入了异步硬件和可变内存。

图 64.2. 生成内核只有通过彼此独立的语义门、编译器门和性能门，才能被接受。失败会回到生成器或已知回退路径，绝不会直接进入生产环境。

生产评估至少应覆盖以下六项：

契约表面。 固定支持的秩、形状、dtype 与数学模式、设备架构、对齐、布局、别名、修改、工作区、流、错误与确定性行为。
隐藏语义。 使用新鲜的随机输入与输出、边界形状、零尺寸与奇数尺寸、广播情况、非连续步幅、未对齐、抵消、并列值、有符号零，以及支持范围内的 NaN 与无穷值；容差应独立地从 dtype 与归约长度推导。
内存与并发。 检测越界访问、未初始化数据、不完整写入、数据竞争、错误同步、非法别名，以及并发流、计算图重放和不同分配器历史下的失败。
评测框架完整性。 隔离候选实现，禁止调用参考实现、monkey-patching、隐藏回退、网络或文件访问，以及复用参考实现的分配状态。扰动每个输入，确保输出确实依赖它。
测量。 分开统计编译时间、自动调优、冷启动延迟、纯内核时间与端到端时间。采用预热、设备同步、随机配对顺序、多次重复与中位数，并提供有用的分位数或置信区间。
代表性比较。 对齐 dtype、快速数学设置、布局、工作区、计算图模式和硬件。在真实工作负载分布上比较生产基线，而不只是即时执行参考实现。

运行编译器与内核的边界

可靠的发布流程会把生成产物与它的假设视为一个整体：

记录计算图或算子模式、守卫、副作用与别名模型、数值容差和回退。
固定目标架构、驱动程序、运行时、数学库、编译器、工具链版本、标志和自动调优状态。
建立语义门，覆盖即时执行与编译执行的输出对照、相关梯度、边界形状、步幅、dtype、修改和并发。
建立编译器门，覆盖计算图中断、守卫失败、重新编译、编译缓存增长、sanitizer 发现和目标代码生成。
在代表性工作负载分布上建立性能门，测量启动次数、实际带宽或算力、占用率、寄存器、共享内存、溢出、工作区和峰值内存。
把冷启动编译与自动调优和热执行分开报告。
对产物进行灰度验证，一旦发生数值、编译器或延迟回归，自动回滚到已知回退路径。
计算图、目标、驱动、库、编译器或工作负载分布发生变化时，重新运行所有门。

下层约束

编译器无法恢复框架丢弃的信息，无法证明程序没有暴露的别名事实，也无法凭空造出目标不支持的指令。反过来，局部很快的内核仍可能输掉端到端性能：它的布局迫使系统转置，工作区增加内存压力，动态形状让编译缓存碎片化，或者它的启动打断了更大的融合区域。因此，编译器与内核性能必须在消费它们的计算图和服务中测量，并与第 62 章给出的带宽、内存和互连限制比较。

下一章第 65 章会沿运行时间向上移动一层：从生成可执行工作，转向让一套长期运行的训练系统持续获得数据、保持可观察，并能从故障中恢复。

争议所在

尚无定论的问题，不是编译器能否生成有用内核，它们已经能做到。问题在于，多少目标特定信息可以藏在可移植接口之后，同时又不牺牲重要性能。高层 IR 与分块内核语言提高了源码和语义可移植性；新的设备指令、布局规则、库契约和异步流水线，却不断制造新的目标特定调度工作。生成内核系统或许能减少这部分劳动，但收益始终取决于评估者、基线、形状分布、硬件和搜索预算。因此，值得检验的主张必须具体：对一项明确的工作负载与契约，把编译与运行成本计入之后，这套工具链能否保留语义，并击败维护最好的回退实现？