编译器、内核与 CUDA 护城河

上一章结束时，它的两场争论都指向框架之下的那台编译器。本章顺着这个指向走下去。被捕获的计算图终究要变成设备代码，而朴素的设备代码与好的设备代码之间，差的不是几个百分点，而是接近一个数量级，差距由一个成本模型支配：在现代加速器上，搬运字节比对字节做算术更贵。这个模型解释了为什么融合算子是编译器最根本的优化，为什么最快的注意力实现更应被理解为一个算法而不是一个内核，也解释了为什么这个行业最深的护城河不是硅片，而是硅片上沉积了十九年的软件。读完本章，读者能把一行 Python 一路追到 GPU 执行的指令，说出性能是在这条线的哪一段赢下来的，并解释这条线为何穿过一家厂商的领地、大家为此付出什么代价，以及模型自己写内核可能怎样改变它。生产环境里服务模型的量化内核在 第 34 章 另有专章；本章讲的是让内核得以产生的那套机器。

数字节，别数 FLOPs

在 GPU 的词汇里，内核（kernel）指发射到设备上的一个程序：一次矩阵乘、一次逐元素加法、一次 softmax。内核跑得快不快，几乎完全由一个比值决定，即它的算术强度：每从高带宽显存（HBM）搬运一个字节，它做多少 FLOPs 的算术。屋顶线模型把这个比值画成一张图：可达吞吐量是两者的较小值，一边是硬件的峰值算术速率，另一边是显存带宽乘以内核的算术强度，于是性能随强度线性爬升，直到撞上平的算力天花板 (Williams et al. 2009)。两片屋顶相交的拐角，用一个数字概括了这台机器的性格。一块 H100 有约 989 稠密 BF16 teraFLOPs 的算力，对着 3.35 TB/s 的 HBM 带宽，拐角落在每字节约 295 FLOPs 附近：任何算术密度低于此的内核都受限于显存，运行时间就是搬运的字节数除以带宽。

Transformer 里除了矩阵乘之外的几乎一切，都远远落在拐角之下。一次逐元素加法每个输出元素做一个 FLOP，半精度下却要搬六个字节，强度约 0.17，差了三个数量级。把每个这样的运算都作为独立内核运行，张量就要在每一步之间往返一趟 HBM，于是「减最大值、取指数、求和、相除」这样的链条要付四趟路费，其实一趟就够。解法是融合（fusion）：把链条编译成单个内核，输入读一次，中间值留在片上寄存器里，结果写一次。算术分毫未变，变的只有流量；而既然流量才是价钱，融合就是这一层最接近免费午餐的东西 (He 2022)。对真实训练的测量证实了这笔账单里有多大一块归它管：一项对 Transformer 训练的研究发现，占据运行时间的是受限于显存的运算而不是矩阵乘，仅靠数据搬运优化就拿回了三分之一的时间 (Ivanov et al. 2021)。还有一笔较小的税把设计推向同一个方向：每次内核发射都要花微秒级的驱动开销，这正是运行时把整段发射序列录制下来、作为 CUDA 图整体重放，而不是逐个发射内核的原因 (Gray 2019)。

下面的可运行单元把这笔账算得具体：对一个大矩阵做 softmax，拆成四个内核执行与融合成一个内核执行，FLOPs 相同，字节数差三倍，而且两者的算术强度都远低于屋顶线拐角，预测的运行时间就是流量除以带宽。

N, d, B = 4096, 4096, 4  # 一张 fp32 分数矩阵，每元素字节数
flops = 5 * N * d        # 取最大值、减、取指数、求和、相除：每元素约各一次
# 不融合：四个内核，每个都从 HBM 读写整个 N*d 矩阵
unfused = (2 + 3 + 2 + 3) * N * d * B          # 各内核的读写量相加
# 融合：输入读一次，输出写一次
fused = 2 * N * d * B
for name, traffic in [("不融合", unfused), ("融合  ", fused)]:
    print(f"{name}: 搬运 {traffic/1e6:6.1f} MB, "
          f"算术强度 {flops/traffic:5.2f} FLOP/字节, "
          f"H100 上受带宽限制的时间: {traffic/3.35e12*1e6:5.0f} 微秒")
print("FLOPs 相同；融合内核快出来的，恰好是它省掉的那些流量")

经典一课：FlashAttention

如果说融合是这一层的日常动作，FlashAttention 注意力内核（FlashAttention）就是这一层的代表作。它值得单独一节，因为它改变了这个领域对「内核能是什么」的认知。注意力的教科书写法要具体化一张 $N \times N$ 的分数矩阵，序列一长，这张矩阵既放不进片上的快速内存，它逼出来的那些 HBM 往返也不值得付。FlashAttention 计算的是精确注意力，没有任何近似，靠的是分块：让键和值的块流过容量很小的片上 SRAM，维护一个运行中的最大值与归一化项，使 softmax 能在一趟里增量算完，这个技巧几年前就以在线 softmax 的名义发表过 (Milakov and Gimelshein 2018)，而那张分数矩阵从头到尾都不写到任何地方 (Dao et al. 2022)。内存从序列长度的平方级降到线性，论文还证明了它的字节数对精确注意力而言在整个 SRAM 尺寸范围内渐近最优。让这一课立住的是一组对照：更早的一种省内存注意力达到了同样的内存界，却没有围绕存储层级来编排计算，其论文自己的表格显示它比基线更慢 (Rabe and Staats 2021)。省内存从来不是重点；编排字节才是。

续作把这一课磨得更锋利，而它对可移植性的含义也更令人不安。FlashAttention-2 把利用率从峰值的 25%-40% 提到最高 73%，靠的不是新想法，而是把同一份计算在 GPU 的工作单元之间重新切分 (Dao 2023)；FlashAttention-3 达到 Hopper BF16 峰值的 85%，代价是彻底拥抱那块芯片最新的硬件，专用的异步拷贝引擎与 warp-group 矩阵指令，于是最快的注意力如今是用一种架构的私有词汇写成的 (Shah et al. 2024)。事实证明，峰值性能不可移植。这个事实塑造着接下来两节的一切。

下沉流水线

在框架捕获的图与斤斤计较字节的内核之间，隔着一套编译器栈，到 2026 年它的形状已经稳定得可以画出来。PyTorch 的 torch.compile 像上一章讲的那样从 Python 字节码捕获计算图，交给 TorchInductor，后者能融合的都融合，再为 GPU 生成 Triton；在 180 个真实模型上的测量显示，这条流水线买到了推理 2.27 倍、训练 1.41 倍的几何平均加速，本质上就是把本章的成本模型自动化 (Ansel et al. 2024)。JAX 那条路生来就是编译的：被追踪的程序先下沉到一种可移植的算子方言，再进入 XLA，也就是自 2017 年起喂养 TPU、2023 年拆分为 OpenXLA 后开始面向多家硅片的那台编译器。两条路在底部汇合：从 PTX 开始由 NVIDIA 的编译器接手，这个由驱动翻译到物理指令集的虚拟指令集，正是下文几个护城河论点共同的枢纽。

这套栈的思想谱系值得交代，因为它解释了分工。图像处理语言 Halide 贡献了奠基性的分离：把「算什么」与「怎样调度到机器上」分开描述，调度便可以被搜索而不必手工推导 (Ragan-Kelley et al. 2013)。TVM 把这个分离带进深度学习，并让搜索自动化，用一个学习到的成本模型去探索人类不会去枚举的调度 (Chen et al. 2018)。MLIR 把公共的底层设施工业化：一个由可组合方言构成的编译器框架，如今 XLA、Triton 的内部和多数后来者都共用它 (Lattner et al. 2021)。而 Triton 定下了整个生态最终收敛的抽象层级：程序员按分块，也就是一小块矩阵，各写一个程序，访存合并、共享内存暂存与处理器内部的调度交给编译器，那些恰恰是 CUDA 让程序员自己扛的难活 (Tillet et al. 2019)。TorchInductor 生成内核用的就是 Triton，它的后端如今瞄准 NVIDIA、AMD 与 Intel 的硅片，生态对可移植性的希望因此都压在它身上。然而反面的证据也压在它身上：为了追逐 Hopper 与 Blackwell 的峰值，Triton 长出了 warp 特化，随后又长出 Gluon，一个把架构特有控制原原本本暴露出来的更低层方言，而那些控制正是这层抽象当初要藏起来的东西。FlashAttention-3 的模式在编译器层面重演了：写法的可移植做得到；峰值性能的可移植总在溜走。

护城河

CUDA 于 2007 年发布，名字听着像一门语言，但这个名字盖住的是一整根垂直的栈：驱动与运行时、二者之间的 PTX 契约，以及最重要的库，做线性代数的 cuBLAS、承担 第 62 章 那些集合通信的 NCCL，还有 cuDNN，它 2014 年的发布恰好在深度学习工业化的时刻，把那批年轻的框架焊在了 NVIDIA 的硅片上 (Chetlur et al. 2014; Nickolls et al. 2008)。二十年的沉积，如今支撑着号称六百万的 CUDA 开发者。这条护城河既是技术的也是法律的：NVIDIA 的许可证禁止把编译产物翻译到其他厂商的硬件上运行；唯一一次认真的二进制兼容层尝试 ZLUDA，先后经历企业资助、法律施压、匿名复活，到 2026 年年中又退回到业余项目的状态。竞争者因此只能攻击二进制之上的那一层。AMD 当前的加速器在内存上领先，288 GB 的 HBM3E 和 第 34 章 里讨论的那些数值格式，它的 ROCm 栈也能跑上游 PyTorch 与主要的服务引擎；然而对 MI300X 一代最细致的独立基准测试得出的结论是，用公开软件训练的吞吐量落后 NVIDIA 超过 2.5 倍，报告标题就叫「CUDA 护城河还活着」(SemiAnalysis 2024)。真正逃出去的那一派，靠的是从未进来：Google 的 TPU 线，用 XLA 编译器喂养的稠密脉动阵列，在这条河得名之前就在训练前沿模型 (Jouppi et al. 2017)；AWS 的 Trainium 沿同一条 XLA 血统走到了五十万芯片的规模。耐人寻味的是，两条逃生路线在用户撞到「只靠编译器能给多少」的天花板之后，都各自长出了类 Triton 的内核语言。

模型写的内核

这一层最新的变量，是本书的主角掉头对准了自己的地基。给这个问题定框架的基准 KernelBench，要求模型写出正确且快于 PyTorch 的 GPU 内核；2025 年初，前沿模型单次尝试能解出的任务不足五分之一 (Ouyang et al. 2025)。但「单次尝试」与「对着校验器迭代」之间的差距，是本书反复出现的主题，在这里同样成立：把编译与计时当作奖励信号，经强化学习训练或由搜索驱动的系统如今能稳定产出正确内核，并在受限于显存的负载上击败 PyTorch 的即时算子，有时是数倍 (Baronio et al. 2025)。警世故事也如期而至：一个广为宣传、声称百倍加速的「AI CUDA 工程师」，几天内就被发现钻了自己评测框架的空子，复用了缓存的正确结果而不是真正算出来，正是 第 27 章 解剖的那种奖励欺骗（reward hacking）的教科书案例 (Sakana AI 2025)。到 2026 年年中，清醒的读法是：生成的内核正在成为覆盖长尾可融合运算的廉价手段，但还没有人展示过模型能在受限于算力的前沿达到 FlashAttention-3 级别的工程水准，而计时环境本身是否可靠，已经上升为头等关切。生产环境的服务栈反映着同一种分工：vLLM 默认编译算子之间那些可融合的黏合部分并在 CUDA 图下重放，而注意力与量化矩阵乘仍是编译器视作黑盒的手写内核 (Ansel et al. 2024)。编译器管住了地板；峰顶仍是手艺活。

延伸阅读