Skip to content

原文:第 7 章 说明:忠实翻译原网页内容,并补入与经典文献、业界系统的对照。术语首次出现给出英文锚点。

第 7 章 深度学习编译器:从通用性到专用化,再回到学习式搜索

本章在体系中的位置

这一章是第 5 章(CUDA 手工优化)的自动化延伸。第 5 章里,你在 GEMM 上亲手转动 tile 尺寸、循环顺序、寄存器分块这些旋钮,打磨一份手工内核;本章把这些旋钮连同背后的搜索空间一起交给编译器,看它如何自动完成同一件事。它也衔接第 10 章:编译器「把优化做成程序化、可搜索的过程」这套想法,会在训练并行的语境里原样重演。在原站的 16 周课程里,这一章处在「单卡优化」与「多卡训练」之间的枢纽位置;它也是原站前七章里唯一带正式参考文献的章——本版完整保留了那份书单,并做了扩充。放在 AI Infra 的坐标里,它回答的问题是:让张量计算变快这件事,该交给程序员、编译器,还是学习系统?

上一章我们看完了 GPU 上的手工优化;这一章把同一个问题交给机器。深度学习编译器是软件工程里最古老执念的一次重演:让人写靠近语义的代码,让机器把剩下的苦活干完。只是这次要自动化的不是一个表达式,而是一整片「把高层张量计算映射到低层硬件执行」的优化空间。

这一章会沿着一条时间线走完四个系统——TVM、Triton、TileLang、CUDA-L2。它们不是四个工具排排坐,而是同一个问题在四个不同时代给出的四份答案。先把这个问题钉在这里,后面每读一个系统都回来对一下:

给定一个张量工作负载和一台目标机器,应该暴露多少结构、隐藏多少结构,以及谁来搜索优化空间——程序员、编译器,还是学习型智能体?

记住这个句式。它比任何一张系统对比表都更能解释这一章在讲什么。


1. 为什么深度学习需要编译器

直觉。 一个深度学习模型是用「意义」的语言写成的。注意力、前馈网络、层归一化、卷积、激活函数、损失函数——这是你描述模型时用的词。但硬件不执行意义。硬件执行指令:把字节在各级存储之间搬来搬去、调度 warp、发出张量核指令、解析数据依赖,然后在数据没按时到达时停下来等待。

两个世界的距离。 模型层说的是「我想算什么」(what),硬件层必须决定「这件事怎么在一台具体的机器上高效地办成」(how)。这段距离不是一句语法翻译能填平的。中间隔着的是一整套决策:怎么分块、数据放哪一级存储、循环按什么顺序展开、谁跟谁融合、向量化到什么宽度、哪些线程做哪些事。把这一整套决策集合起来看,它是一个规模很大、结构很强的优化问题。深度学习编译器,就是把这个优化问题自动化的软件系统。

问题是什么时候变重要的。 不是一开始就重要的。在神经网络还只是研究原型、跑通就算赢的年代,执行效率排在正确性后面。等模型大到足够有用、进了生产环境,游戏规则就变了:光是「能跑」不够,还得跑得快、装得进内存、把硬件吃满、能适应真实部署环境的规模与约束。从这一刻起,「把张量计算映射到硬件执行」不再是工程小修,而是核心矛盾。

为什么编译器,而不是人手。 在编译器出现之前,这个映射靠的是为特定硬件手工优化的内核(hand-optimized kernel)。每新增一对「算子 × 后端」,就要有人专门写一份内核。几十上百个算子、十几个后端,乘出来是几千份内核——要写、要调、要维护、要跟随硬件换代而重写。这个支持矩阵的爆炸速度远远超过工程师队伍的增长速度。这是编译器要接管的第一个理由。

这一章的核心判断,可以压缩成一句话:

深度学习编译器,是一个把「高层张量计算映射到硬件执行」这一大规模、结构化优化问题自动化的软件系统。

四个里程碑。 这个映射在过去十年里变了好几次形状,因为工作负载变了、硬件变了、主导瓶颈也变了,于是主导的编译器设计跟着变。大致可以用四个里程碑来读这段历史:

  1. LLM 之前的编译器时代,代表是 TVM:核心问题是「很多模型 × 很多硬件后端」。
  2. LLM 时代早期的收敛,代表是 Triton:问题坍缩成「Transformer 工作负载在 NVIDIA GPU 上的高效实现」。
  3. 显式控制的回归,代表是 TileLang:内核越来越特化,黑盒自动化不够用了。
  4. 学习式搜索的新阶段,代表是 CUDA-L2:大语言模型加强化学习重新接管内核优化,但落点与过去完全不同。

这四个系统不是四个孤立工具。它们各自回答了开头那个问题的不同版本——暴露什么、隐藏什么、谁来搜索。带着这个问题往下读。


2. 编译器要解决的问题:在它出现之前,真正在优化的是什么

先把对象说清楚。 在讨论编译器之前,必须明确:真正被优化的那个东西是什么。我们读第 5 章时已经见过它的原型了。一个张量计算可以看成一个迭代空间(iteration space)——所有循环索引元组构成的集合。以矩阵乘法为例:

一种自然的实现是三层循环:

c
for (int i = 0; i < M; ++i) {
  for (int j = 0; j < N; ++j) {
    float acc = 0.0f;
    for (int k = 0; k < K; ++k) {
      acc += A[i][k] * B[k][j];
    }
    C[i][j] = acc;
  }
}

迭代空间就是集合

一个 GPU 内核,本质上就是把迭代空间映射到硬件执行层次与访存层次的一种实现方式。两个目标支配着这个映射。

2.1 并行性

第一个目标是把一个大计算拆成许多较小、可以并发运行的子任务。GPU 上的处理单元是层次化组织的:grid、block、warp、thread。并行化做得好,算术单元持续吃饱;做得差,算术单元空转。

2.2 局部性

第二个目标是降低数据移动的成本。算术和存储在物理上是分开的:寄存器快但很小,共享内存和 L1 也快但容量有限,HBM 大却远得多、慢得多。一个高性能内核的关键,是让可复用的数据尽量长时间地待在层级的高处。这也是为什么一大堆看起来五花八门的优化技巧,最后都能归成同一句话:每搬一个字节,都要做更多有用计算。

2.3 一整套规范化旋钮

第 5 章讲的那些优化手段不是孤立的小技巧,它们是张量内核优化的标准自由度。值得把它们摆在一起看,因为它们正是后面所有编译器——TVM、Triton、TileLang、CUDA-L2——都在操作的那组旋钮:

  • 平铺(tiling):把迭代空间切成更小的块。GEMM 里可以把输出空间切成 的 tile,再把归约维度按选定块大小切分。平铺同时改善局部性(块能装进共享内存或寄存器)和并行度(不同块分给不同 block 或 warp)。
  • 内存放置(memory placement):决定中间数据住在哪一级——全局内存、共享内存、寄存器,还是某种专门化的片上存储。GEMM 里 A、B 的 tile 通常先搬进共享内存,供块内许多线程复用。
  • 循环置换(loop permutation):改变迭代顺序。数学结果不变,但数据重用模式大变。把 换成 ,内层循环复用的值就变了,局部性跟着变。
  • 融合与切分(fusion / fission):融合把多个算子或循环合成一个更大的计算,中间结果不必物化回全局内存,内存流量大减;切分则把一个融合体拆小,有时能改善并行度、降低寄存器压力或简化调度。
  • 循环展开(unrolling):把多次迭代展开成直线代码,减少循环控制开销、提升指令级并行,代价是代码体积和寄存器压力变大。
  • 线程映射(thread/block mapping):决定哪块迭代空间交给哪个 block、warp、thread。这个选择强烈影响访存合并(coalescing)、数据重用、occupancy 和同步模式。
  • 同步(synchronization):数据被共享时,同步是正确性的前提,但它本身有开销。同步太多毁掉吞吐,同步太少破坏正确性。

把七个旋钮写成一个元组,就是编译器要搜索的对象:

每个分量几乎都是离散的:tile 尺寸从有限集合里选,循环顺序是排列,展开因子是整数,线程映射和同步方式是类别选择。真实内核的变量远不止这七个——常常超过十个,有时超过二十个。这就是后面所有故事的开端。


3. GPU 内核优化是一个结构化的离散搜索空间

核心抽象。 这一节要建立的抽象,是把上面所有选择合在一起看:一个内核的优化从来不是一个决策,而是一组决策的组合——而且这组决策构成一个结构化的离散搜索空间(structured discrete search space)

直觉。 试试把它想成做饭。一道菜的口味不是由一个变量决定的,而是火候、盐量、时间、下料顺序一起决定的;更关键的是它们互相纠缠——盐放多了火就得小一点,菜切得薄时间就得短。单个变量都不复杂,放在一起就成了一个巨大的、互相牵制的配置空间。内核优化正是这样。

耦合才是难点。 搜索空间大,还不算最难;最难的是变量之间强耦合:

  • tile 变大可能提升数据复用,却把共享内存撑爆;
  • 展开更多可能提升指令级并行,却把寄存器压力顶破;
  • 换个循环顺序可能改善局部性,却毁掉向量化;
  • 更激进的融合可能减少 HBM 流量,却压低 occupancy;
  • 更好的线程映射可能改善合并访问,却在别处引入 warp 发散。

换句话说,性能不是独立决策的简单叠加,而是一张交互面(interaction surface)。这让最坏情形下的优化问题变得组合式、难以精确求解。编译器实践上只能用启发式、剪枝、整数规划、混合整数规划、代价模型(cost model)、规则系统、自动调优(autotuning),或学习式搜索来逼近。

停下来想想

上面这些耦合例子没有一个是「理论上的」。第 5 章里你手工调 GEMM 时遇到的每一个「调好这个、坏了那个」,都是这里列出的交互项之一。编译器存在的全部理由,就是自动处理这种耦合。

为什么要「栈」。 这个观察直接解释了一个设计事实:一整个巨型优化问题一次性解决是不可能的,系统必须拆成层。这就是下一节的主题。


4. 为什么编译器必须以「栈」的形式存在

核心思想。 一个计算需要很多种表示。栈顶的表示对人有意义,栈底的表示对硬件有意义,中间的每种表示保留一类不同的结构。某些优化需要看到全局依赖,另一些优化需要看到循环、缓冲区和访存的细节——没有任何单一表示能同时满足两者。

类比。 想想一栋楼。建筑师画概念图,结构工程师画结构图,施工队看钢筋明细表。同一栋楼、好几套图纸,每套服务于一类决策,说其中任何一张是「真正的图纸」都没有意义。编译器栈同理。

一个典型的深度学习编译器栈,包含五层:

  1. 模型层(model layer)——神经网络的函数结构;
  2. 图 IR(graph IR)——算子与数据依赖构成的图;
  3. 张量 IR(tensor IR)——感知循环与访存的张量程序;
  4. 硬件 IR(hardware IR)——面向具体设备的执行表示;
  5. 运行时(runtime)——动态启动、内存与协调决策。

这不是官僚主义,是分解(decomposition)。 模型层说「算什么函数」,图 IR 说「哪些算子依赖哪些算子」,张量 IR 说「循环结构和访存长什么样」,硬件 IR 说「工作怎么映射到机器」,运行时说「编译好的计划此刻怎么被启动和协调」。性能来自在每一层保留并利用正确的结构——这五层各自存在的理由,就是不同优化需要不同类型的可见性。


5. 历史转变:LLM 之前与 LLM 之后

5.1 LLM 之前:处处是多样性

情境。 大约 2017 到 2020 年,深度学习处处是多样性。模型有 CNN、RNN、LSTM、GAN、混合架构、早期注意力模型;算子种类繁多、形态各异。硬件同样碎片化:NVIDIA GPU、Google TPU、ARM CPU、FPGA、移动端加速器、华为昇腾(Ascend)、Graphcore IPU 并存。

核心问题因此是:

怎么让一个模型写一次,到处高效运行?

这是 TVM 最合理的历史语境。那时代的关键词叫通用性(generality):编译器必须同时应对很多模型和很多后端。

5.2 LLM 时代:收敛,然后是压力

收敛。 2020 年之后局面变了。工作负载向 Transformer 收敛,硬件向 NVIDIA GPU 和 CUDA 生态收敛。这在某种意义上大幅简化了编译器问题——不再需要为所有人优化一切了。但简化换来了新压力:Transformer 训练和推理贵到每一个百分点的效率都直接换算成基础设施成本,模型开发者再也无法无视底层性能。这是 Triton 出场时的世界。

特化。 然后问题又变了。Transformer 衍生的模型和推理内核越来越专门化:FlashAttention 的各种变体、分组查询注意力(GQA)、MLA、量化内核、融合 epilogue、依赖特定硬件的流水线技巧。在这个阶段,黑盒编译抽象开始不够用了。这是 TileLang 出场时的世界。

再自动化。 当暴露出来的优化空间对人类又变得太大、太细、太依赖硬件细节,新的回答出现了:与其把更多旋钮塞回人手里,不如用 LLM + 搜索 + 强化学习去探索那片离散空间。这是 CUDA-L2 以及更广的 LLM-for-AI-infra 方向所处的世界。

一句话概括这段历史:问题变了四次形状,编译器的形状就跟着变四次。


6. TVM:LLM 之前的深度学习编译器

6.1 TVM 试图解决的历史问题

情境。 TVM 是 pre-LLM 时代最典型的深度学习编译器,由华盛顿大学、上海交大、CMU 等多家机构合作完成,2018 年发表 OSDI 论文(Chen 等人)。它的动机清晰且有历史意义。

当时的 TensorFlow、Caffe、早期 PyTorch 高度依赖为特定硬件手工优化的内核。这导致支持矩阵爆炸:很多算子 × 很多后端。全部手工组合去支持,昂贵、脆弱、演进缓慢——算子每多一个、后端每多一个,维护负担都是相乘的。

TVM 的回答:构建一个统一的编译器栈,从高层模型出发,穿过中间表示,为多样的硬件后端自动生成高效代码。它的奠基论文把这件事定义为深度学习模型的端到端优化问题,而不只是局部的代码生成问题。

6.2 TVM 的栈式组织:图、张量、代码生成、运行时

TVM 把编译问题明确拆成几个层面:

  • 图层:推理整个模型的结构——有哪些算子、怎么连接、什么可以融合、哪些常量可以折叠、内存怎么规划;
  • 张量程序层:推理单个算子的实际实现——循环、内存作用域、平铺、并行化、向量化、张量化、缓存;
  • 后端层:把调度好的张量程序降到目标特定代码,如 LLVM IR、CUDA、ROCm 或其他后端;
  • 运行时:处理模块加载、内存、启动与协调。

历史上,TVM 的图级 IR 是 Relay;后来的版本越来越强调 Relax,用于图级抽象和动态形状支持。张量侧从早期的 Tensor Expression(TE) 演进到更显式、更便于变换的 TensorIR(TIR)。当前 TVM 文档把 Relax 定位为 ML 模型的图抽象,把 TensorIR 定位为「循环、缓冲区与硬件相关调度结构」的张量程序抽象。

6.3 compute-schedule 分离:TVM 的核心思想

TVM 最核心的思想是 compute-schedule 分离。 compute 定义算子的数学逻辑,schedule 定义这个计算怎么被执行。矩阵乘法的 TE 风格定义,概念上是这样的:

python
k = te.reduce_axis((0, K), name="k")
C = te.compute(
    (M, N),
    lambda i, j: te.sum(A[i, k] * B[k, j], axis=k)
)

这定义了「要算什么」,还没说平铺、循环顺序、内存作用域、向量化、线程怎么组织。schedule 负责说「怎么做」:

python
s = te.create_schedule(C.op)
i, j = s[C].op.axis
io, ii = s[C].split(i, factor=32)
jo, ji = s[C].split(j, factor=32)
s[C].reorder(io, jo, ii, ji)
s[C].parallel(io)
s[C].vectorize(ji)

这个六行小例子同时演示了好几个深刻思想:

  • split 做平铺;
  • reorder 改变循环顺序,以改善局部性和向量化的可行性;
  • parallel 暴露并行工作;
  • vectorize 把最内层连续迭代映射到 SIMD 或向量化执行。

TensorIR 文档把这个风格描述为对张量程序做的一系列变换:一个 TensorIR 程序包含缓冲区、循环嵌套和计算块;调度就是调用循环重排、切分、张量化等变换原语,把程序推向高效实现。

这个分离是现代 ML 编译器最重要的思想之一。 它把正确性和性能解耦:算子语义保持稳定,调度空间可以被大规模自动探索。就像菜谱把「做什么菜」和「用什么火候」分开写,然后让一个自动化的厨师把火候调到最优。

6.4 图级优化:为什么不只是「生成内核」

在详细调度开始之前,TVM 可以在图级做硬件无关的变换:算子融合、常量折叠、死代码消除、布局更改、内存规划。其中算子融合(operator fusion)最值得展开。考虑

不融合时,MatMulAddGELU 是三个独立算子,中间张量被写回全局内存再读回来。融合后,这些操作可以并成一个更大的内核,中间值留在寄存器或缓存里,而不是物化在 HBM 中。第 5 章的 Conv2D → BiasAdd → ReLU → BatchNorm 例子讲的是同一件事:卷积计算密集,后面的 BiasAdd、ReLU、推理期 BatchNorm 大多是逐元素、访存受限的,把它们融合既省内存流量,又省内核启动开销。

合法且有利可图的融合需要四个条件:

  1. 单向数据依赖;
  2. 计算能合成一个循环嵌套;
  3. 语义等价;
  4. 访问模式兼容,同一中间结果能留在片上。

6.5 张量级优化:结构化的 CUDA 优化空间

一旦进入 TensorIR,程序就是感知循环与访存的了。TensorIR 文档强调:张量程序显式表示缓冲区、循环和块,这个抽象被设计成能暴露线程化、访存和专用指令等硬件加速选项。这正是 TVM 的调度空间与结构化 CUDA 优化重合的地方。

典型的张量级调度原语包括:

  • 循环切分 / 平铺(split / tile)
  • 循环重排(reorder)
  • 并行化 / 绑定(parallel / bind)
  • 向量化(vectorize)
  • 缓存(cache_readcache_write
  • 张量化(tensorize)
  • 软件流水(software pipelining)

一句话理解 TensorIR:它是把 CUDA 优化空间做成了结构化的符号编码。 底层自由度还是第 2 节那七种,但表示形式是为系统性变换设计的,而不是为手写代码设计的。

6.6 AutoTVM 与 MetaSchedule:为什么必须自动搜索

手工写 schedule 无法扩展,搜索空间涨得太快。TVM 的回答先是 AutoTVM,后来是 MetaSchedule

AutoTVM 用人工定义的调度模板加机器学习引导的搜索。基本循环是:

  1. 用参数化模板定义调度空间;
  2. 生成候选调度;
  3. 在真实设备上编译运行;
  4. 测量延迟;
  5. 训练或更新代价模型来预测性能;
  6. 用代价模型引导进一步探索。

这就是经典的自动调优闭环:测量、预测、搜索、反馈。

MetaSchedule 把这个想法推广了一层。2021 年的 TVM RFC 把它描述为 TensorIR 上的概率调度 DSL,统一了 AutoTVM 和 AutoScheduler 风格的方法,直接在 TIR 的调度原语上操作,每一层的自动化基础设施都可定制——RFC 明确称它是 TVM 的第三代自动调度系统。关键的转变是:MetaSchedule 在 IR 变换空间上搜索,而不是在人写的模板上搜索。 当张量化、循环切分这类新调度原语出现时,搜索系统能直接吸收它们,而不必为每个算子重写模板。

6.7 TVM 为什么重要

TVM 重要,因为它改变了机器学习系统的心智模型。以前,优化内核是每个框架、每个后端手工写死的固定工件;TVM 把优化变成可编程、可组合、部分可学习的编译过程。

它的强项恰好是 pre-LLM 时代需要的:

  • 跨很多模型的图级推理;
  • 跨很多算子的张量级调度;
  • 很多后端;
  • 用自动调优扛住组合爆炸。

它的弱点后来也越看越清楚:广泛的通用性带来复杂性。栈变深,跨层效应难以推理,调试困难——在一层做的优化可能伤到另一层。这与其说是 TVM 的偶然缺点,不如说是「广义编译器问题」本身的代价。


7. Triton:LLM 时代的第一代编译器

7.1 Triton 为什么会出现

Triton 属于一个与 TVM 明显不同的世界。到了 LLM 时代早期,问题已经简化:主导工作负载收敛到 Transformer 类算子,主导平台收敛到 NVIDIA GPU。在这个前提下,广泛的硬件通用性不再是第一优先。什么才是?让主要在 Python 里工作的模型开发者,能够写出高性能的自定义 GPU 内核,而不必成为专职 CUDA 专家。

Triton 的设计目标因此不是「到处优化一切」,而是更接近:

给 Python 用户一条路,写出接近 CUDA 性能的快速自定义 GPU 内核,同时隐藏大部分底层机制。

Triton 文档就是这个定位:教程主打简短高性能的矩阵乘、融合 softmax、融合注意力、grouped GEMM、persistent matmul,都用一门 Python 化的 DSL 写,比裸 CUDA 简洁得多。Triton 由 Tillet 等人在 2019 年提出(MAPL 论文),OpenAI 在 2021 年开源了 triton-lang,之后成为 LLM 生态里被引用最多的内核语言之一。

7.2 Triton 的程序模型

Triton 的核心心智模型是块级程序(block-level program)。一个 Triton 内核用 @triton.jit 装饰;你写的代码仿佛一个程序实例在操作一个数据块。底层程序实例到 CUDA 执行层次的映射,由 Triton 完成。

课程课件里的对比把这种抽象讲得很准:

  • CUDA 直接暴露全局/共享/寄存器内存;Triton 大量自动管理内存。
  • CUDA 显式暴露 block、warp、thread;Triton 把块/程序作为用户可见的主要单元。
  • CUDA 要求手工使用张量核、手工向量化;Triton 自动完成大部分。

7.3 @triton.jitprogram_idarangemask

即使在一个最简单的向量加例子里,也能看到 Triton 的核心构件:

  • @triton.jit:标记函数做 JIT 编译;
  • tl.program_id(axis):沿某个启动维度返回程序实例编号;
  • tl.arange:在块内生成向量化偏移;
  • tl.load(..., mask=...)tl.store(..., mask=...):带掩码的访存,处理边界。

7.4 怎么读 Triton 代码

一个 Triton 内核概念上是这样的:

python
@triton.jit
def kernel(x_ptr, y_ptr, z_ptr, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < N
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    z = x + y
    tl.store(z_ptr + offsets, z, mask=mask)

这样读它:

  • 一个程序实例处理 BLOCK 个元素的一块数据;
  • tl.arange 生成块内的 lane 偏移向量;
  • mask 保护边界;
  • 你操作的是向量,不是显式的 CUDA 线程。

这是最关键的抽象跃迁:Triton 把线程级并行藏进了向量风格的块计算里。 概念上,Triton 的流程是:分析 Python 内核 → 按块大小这类运行时常量做动态特化 → 经过 LLVM / MLIR / PTX 相关机制下降 → 缓存编译产物 → 启动内核。动态特化和缓存不是边角功能,它们正是「Triton 好用又快速」的核心原因。

7.5 Triton 用于 GEMM 与注意力

Triton 的矩阵乘教程把主要思想说得很清楚:分块 GEMM 算法中,每一对双重嵌套的块迭代由一个专用 Triton 程序实例处理,块加载用显式指针算术,性能通过块大小、warp 数等配置参数自动调优;文档还强调程序重排能改善 L2 命中率。

融合 softmax 教程展示了 Triton 的另一个关键优势:带宽受限的算子从融合中获益——数据加载一次、片上处理、写回一次,而不是在 HBM 里物化多个中间张量。

融合注意力教程和 persistent matmul 教程则展示了下一级:Triton 的表述力足够实现 SOTA 注意力和更高级的 matmul 变体,同时仍然暴露块大小、warp 数、stage 数等参数上的自动调优。

7.6 为什么在 LLM 时代早期这么好用

课程里有个微妙但重要的点:Triton 不是魔法。 它特别好用,是因为问题暂时被简化了。如果主导工作负载族是 Transformer 类,主导目标是 NVIDIA CUDA GPU,那么编译器可以高度特化。在这种条件下,暴露一个又小又舒服的编程接口、同时仍然给出强性能,是可能的。

这正是为什么对 GEMM、FlashAttention 这类关键算子,Triton 常常能给出「精心手写 CUDA 性能的 80–90%」,而开发复杂度低得多。

7.7 Triton 的局限

Triton 的强项是它的抽象,弱点也是。一旦内核变得不规则或高度定制,黑盒调度就开始撞天花板。程序员无法直接控制这些在困难内核里变得重要的细节:

  • 共享内存的精确布局;
  • swizzle 模式;
  • warp 特化(warp specialization);
  • 细粒度线程绑定;
  • 显式多级流水线;
  • 某些复杂融合内核。

Triton 的优雅来自收缩优化空间。当工作负载族再次变宽,这种收缩就变成限制。


8. TileLang:当控制旋钮回归

8.1 TileLang 为什么会出现

TileLang 出现在 Triton 的简洁不再够用的地方。随着 LLM 内核越来越特化、GPU 本身越来越复杂,模型开发者和系统研究者需要更显式的控制。挑战不再是「写一个快速自定义内核」,而是「写一个性能取决于布局、流水线、线程绑定、张量核映射和形状相关调度的微妙选择的内核」。

TileLang 的回答不是抛弃抽象,而是把抽象往下移一层:

仍然以 tile 为主要概念单位,但把围绕 tile 的调度和内存机制多暴露出来。

8.2 tile 是一等公民

TileLang 文档明确说:tile 是其编程模型的心脏——一块有形状的数据,可以被一个 warp、一个线程块或其他等价并行单元拥有和操作。在 T.Kernel 中,执行上下文定义块索引和线程数,用户可见的显式原语把 tile 缓冲区放进共享内存这样的物理存储空间。

这与 Triton 的决定性区别:Triton 里你写块级逻辑,把很多调度细节留给编译器;TileLang 里 tile 仍然是推理单位,但你可以围绕它显式地塑造更多调度。

8.3 TileLang 显式暴露了什么

TileLang 构建在 TVM 风格的基础设施之上,把数据流与调度空间明确解耦。用户可见的原语包括:

  • T.Kernel(...):内核启动上下文;
  • T.alloc_shared(...):共享内存分配;
  • T.alloc_fragment(...):寄存器 fragment 分配;
  • T.copy(...)T.async_copy(...):跨内存作用域的显式数据搬移;
  • T.Pipelined(...):沿循环维表达软件流水;
  • T.Parallel(...):并行拷贝或循环并行;
  • T.gemm(...) 及相关的计算原语;
  • layout / swizzle 工具:内存布局控制。

TileLang 文档把 DSL 按这几类组织:数据搬移、计算原语、控制辅助、诊断,以及内存屏障、warp 组操作这类高级操作。

8.4 三层控制:从舒适区到专家区

TileLang 的矩阵乘教程描述了多个控制层级。Level 1 相对舒适、高层;Level 2 暴露线程块、共享内存、fragment、流水线和 tile 级原语;Level 3 逼近手写 CUDA/HIP 式的显式度,面向需要控制线程级行为甚至内联 PTX 级细节的专家。

这个概念上很重要:TileLang 不是「Triton 加更多旋钮」,它试图横跨一个更宽的可编程性—性能权衡区间。

8.5 内存布局:关键例子

课程特别强调,TileLang 暴露的最重要东西之一是内存布局(memory layout)。在困难内核里,tile 在共享内存或寄存器 fragment 里的精确排布极其重要:bank 冲突、对齐、张量核 fragment 布局、合并访问,全取决于这些细节。

TileLang 因此允许你控制:

  • tile 在共享内存里的排布;
  • 避免 bank 冲突的 swizzle 变换;
  • padding 与对齐;
  • 面向张量核友好访问的 fragment 布局。

文档里有专门的 layout / swizzle 工具,包括 GEMM fragment 辅助布局和线性布局;数据搬移原语也显式地在 Global、Shared、Fragment 三个作用域之间中转数据。

8.6 流水线与 warp 特化

TileLang 比 Triton 更直接地暴露流水线控制。matmul 教程明确用 T.Pipelined(...) 表达沿 K 维的软件流水;更广的变换文档里还有 warp 特化的生产者—消费者 pass,把可被改写的流水循环重写为带显式屏障同步的生产者分支和消费者分支。

这很重要,因为很多高性能内核靠精心构造的流水来重叠数据搬移和计算。Triton 通过 num_stages 这类参数间接暴露一部分;TileLang 给你更直接的影响力——流水深度、缓冲方式,甚至 warp 特化的结构。

8.7 为什么 TileLang 强,也让人痛

TileLang 强,因为它重新打开了 Triton 收缩掉的那部分优化空间。这正是它能在困难内核上超过 Triton 的原因:

  • 复杂注意力变体;
  • 量化内核;
  • 特殊内存布局的内核;
  • 跨硬件调度场景;
  • 需要显式控制片上中转的融合内核。

但这也是它痛的原因:开发者必须显式地推理内存布局、线程绑定、流水和 tensorization,可用性成本是真实的。

一个诚实的判断

课程在这里有一个强烈的观点。从人因工程的角度看,把这些底层旋钮重新塞回模型开发者手里,是深不可取的——编译器的意义本来就在于自动化这类搜索。TileLang 是对真实问题的有效工程回应,但更深一层说,它也是「自动化还不够好」的征兆。


9. CUDA-L2:用学习式搜索击败人类启发式

9.1 CUDA-L2 在概念上为什么重要

CUDA-L2 不只是又一个内核优化器,它代表一个不同的想法:

与其把更多优化旋钮交还给人,不如训练一个系统,让它比人更有效地搜索这个巨大的离散内核设计空间。

论文聚焦 A100 上的半精度 GEMM(HGEMM)。它把大语言模型和强化学习结合起来,在 1000 组 配置上优化内核,并在自己的评测设置下报告平均性能超过了包括 cuBLAS 和 cuBLASLt 在内的强基线。

这为什么重要?两条理由。第一,GEMM 是 LLM 的典范工作负载。第二,cuBLAS 和 cuBLASLt 本身是多年专家工程的产物。打败它们,是「搜索空间里仍藏着高质量解,只是人类要么没想到、要么没有系统化地去找」的强证据。

9.2 CUDA-L2 由哪些部件组成

按论文,CUDA-L2 在更早的 CUDA-L1 之上加了几个关键成分:

  1. 在更多样 CUDA 代码上的持续预训练;
  2. 多阶段 RL 训练,从较宽的内核优化走向 matmul 特化;
  3. 上下文里加入更丰富的 Nsight Compute profiling 指标;
  4. 检索增强上下文,注入基础模型里没有充分编码的架构知识或优化经验。

奖励是执行速度。候选内核被编译、运行、在真实硬件上测量。搜索因此扎根于真实执行行为,而不是只靠静态推理。

9.3 为什么学习式搜索有机会打败专家

课程的解释一针见血:人类不搜索整个空间,人类按规则搜索。 专家通常带着很强的先验:

  • tile 尺寸应该整除张量维度;
  • padding 是多余的浪费;
  • 预取深度通常很小;
  • 某个 swizzle 是「标准答案」;
  • 双缓冲通常好——或者通常不值;
  • 中间临时布局多半必要。

这些规则不是非理性的,它们是压缩过的经验。但问题在于:它们在搜索开始之前就剪掉了大量可能性。 CUDA-L2 能赢,不是因为它有神秘直觉,而是因为它愿意尝试人类会早早拒绝的选项。

9.4 CUDA-L2 找到的几个代表性案例

课程幻灯片总结了六个案例,把上述论点讲得非常生动。

案例 1:用 padding 解锁更优 tile。 人类常坚持 tile 尺寸干净地整除矩阵维度。CUDA-L2 找到了一些配置:padding 把维度略微撑大,从而解锁一个好得多的 tile。幻灯片里的例子:tile 160 不整除 8192,只需约 1.6% 的 padding 开销,却比 128、256 这些更「自然」的选择性能好得多。论文把这个现象作为关键例子报告。教训很微妙:padding 的局部成本,可能远小于利用率和寄存器压力上的整体收益。

案例 2:深层预取,而不是 K+1。 标准内核通常在 K 维只预取一块。CUDA-L2 发现更深的 lookahead(比如 K+4)在 A100 上能更好地隐藏延迟——异步拷贝延迟足够高,浅预取并不总是够用。这是架构敏感型优化的完美例子:延迟隐藏、缓冲成本和形状依赖之间的交互,人类手工很难推理,因而不太会激进地探索深预取。

案例 3:直接从寄存器做 128 位宽写进共享内存。 人类常引入一个临时中转张量,在写共享内存前先重打包或重排数据。CUDA-L2 发现,当形状、对齐和顺序已经有利时,中间中转是多余的:直接从寄存器对共享内存做宽 128 位 store,内存指令更少、epilogue 流水更简单。主题是:人类设计常常带着「以防万一」的保守中转,学习式搜索能发现这些步骤何时其实是冗余的。

案例 4:寄存器 fragment 双缓冲,但只在划算的时候。 双缓冲是经典技术:一个缓冲区喂 MMA 计算时,另一个预取下一块。CUDA-L2 既学会了怎么用双缓冲,也学会了什么时候不用——幻灯片明确说,寄存器压力会溢出时它会避开。这正是人类概念上理解、但难以跨很多形状精确优化的权衡:延迟隐藏的收益 vs 寄存器压力压低 occupancy 的损失。

案例 5:错开的 A/B 预取调度。 与其总在计算前把 prefetch-A 和 prefetch-B 一起发出,CUDA-L2 发现把两者在计算周围错开,能改善指令级并行、平滑内存流量。这又是那种收益取决于精确指令混合和 warp 调度行为的底层调度模式。

案例 6:形状特定的 CTA swizzle。 CUDA-L2 对小、中、大 GEMM 分别调不同的 CTA 块 swizzle 以改善 L2 复用:小 GEMM 常不用 swizzle,中等用小步长 swizzle,大尺寸用大步长变体。这是为什么学习式搜索有吸引力的漂亮例子:可能根本不存在单一「最佳」swizzle 策略,最佳选择取决于形状和缓存行为。

9.5 CUDA-L2 教给我们的三件事

第一,即使 GEMM 这样成熟的领域,内核优化空间也远未被穷尽。 第二,人类内核工程受限制的不只是知识,还有搜索策略——人类带着强先验和有限的耐心搜索,学习系统可以搜得更广、更系统。第三,面对 TileLang 这类系统暴露的复杂性,长期正解也许不是继续让人承担更多控制,而是让机器去搜索这些控制。


10. 更广的新方向:LLM 接管 AI 基础设施

CUDA-L2 属于一个更大的运动。课程结尾的幻灯片把近年工作正确分成几个簇:

  • 特化的后训练内核模型;
  • 代理式(agentic)迭代搜索系统;
  • 硬件感知 / profiling 感知的优化;
  • 演化式或种群式搜索;
  • 跨硬件扩展;
  • 基准与评测基础设施。

这些工作共享的信念是:AI 基础设施本身,正在成为 AI 驱动自动化的目标。 编译器不再只是手写的优化系统,它越来越变成一个环境——学习型智能体在其中提出候选、测量、修订、积累优化策略。

这是历史性转变。TVM 用机器学习引导自动调优;CUDA-L2 和相关系统走得更远——把底层优化本身当作推理与搜索问题,让基于 LLM 的系统参与进来。

课程里那句话说得传神:

用魔法打败魔法。

说得更冷静一点:用机器学习来优化机器学习的执行。


11. 把四个系统放在一起看

现在可以把四个里程碑摆在一起比较。

TVM

最典型的编译器中心型系统。它假设问题宽、异构、多层次,所以建一个丰富的 IR 栈,把 compute 和 schedule 分离,用图级优化加张量级自动调优覆盖大空间。

Triton

程序员中心、编译器辅助的系统,面向一个更窄的区间。它假设工作负载与硬件的配对足够受限,因此一个更简单的 DSL 就能暴露正确的块级抽象,同时保持高性能。

TileLang

恢复控制型系统。它保留 tile 抽象,但重新暴露内存布局、流水线、线程绑定和张量化——因为对高度定制的内核来说,Triton 的收缩抽象不再够用。

CUDA-L2

学习式搜索型系统。它的判断是:优化空间对人来说太大、太碎、太依赖形状与硬件,所以让 LLM + RL 系统直接用执行时奖励和 profiling 信号搜索它。

这不止是四个工具,而是「自动化应该住在哪一层」的四种回答。 依次看:TVM 把自动化放在编译器里,Triton 放在 DSL 的简洁里,TileLang 退回给人类,CUDA-L2 放在学习系统里。


12. 这一章真正的收获

本章最深的教训,不在某个具体编译机制,而在下面这几条。

第一,一个计算不只有一种表示,它必须有很多种表示。 模型层保留语义与设计意图;图层保留全局依赖结构;张量层保留循环、缓冲区与可调度的结构;硬件层保留执行与资源约束;运行时保留动态协调与启动行为。不同的优化只在合适的结构可见时才可能发生——这就是编译器栈存在的原因。

第二,内核优化本质上是一个结构化离散搜索问题。 这解释了为什么手工 CUDA 优化那么难;为什么编译器栈要分层;为什么 compute-schedule 分离有意义;为什么 Triton 能暂时简化问题;为什么 TileLang 不得不重新暴露更多控制;为什么 CUDA-L2 这样的 LLM + RL 系统如今如此有吸引力。所有现象都是同一个根因的不同投影。

第三,这是历史的教训。 深度学习编译器会演化,不是因为编译器研究者口味多变,而是因为世界的形状变了:

  • 世界异构时,需要宽泛的编译器抽象;
  • 世界收敛时,需要高生产力的专用 DSL;
  • 特化制造新复杂时,需要更细粒度的控制;
  • 细粒度控制超出人力的系统管理能力时,又回到学习式搜索的自动化。

这个循环还没结束。 下一代的 ML 系统很可能更用力地朝一个综合走:保留编译器的结构化分层视角,同时让学习型智能体自动搜索设计空间里困难的部分。这不仅是技术方向,也是务实方向——优化空间对大脑来说已经太大了。赢的系统会是那些保留结构、暴露正确抽象、并在人的推理变得脆弱处自动化搜索的系统。


13. 最新进展

自 CUDA-L2(2025 年 12 月)以来,这个领域扩张得很快,工作不再是单一主线,而是分成几个比较清晰的分支。读这片地形最实用的方式:特化的后训练模型、代理式搜索、硬件/profiling 感知搜索、演化式搜索,以及新的基准与评测基础设施。一篇近期的综述几乎就按这个方式组织,是很好的伞形参考(这里的分簇是我个人的读法,和那篇综述基本一致)。

第一个分支:为内核生成本身训练的特化 LLM。 CUDA-L2 之前的关键前驱是 AutoTriton(面向 Triton 生成,SFT + RL + 执行感知奖励)和 CUDA-L1(用对比式 RL 做 CUDA 优化,在 KernelBench 上报告了显著加速)。CUDA-L2 把这条线在 HGEMM 上再推一步。2026 年的工作则聚焦让 RL 训练更稳定、更不容易被刷:Dr. Kernel 引入 KernelGYM、多轮 RL、profiling 奖励和 TRLOO,以降低回合级信用分配的偏差;CUDA Agent 用合成数据、技能增强环境和稳定训练技巧扩展代理式 RL;DRTriton 为 Triton 内核做大规模合成数据 RL。大趋势是:社区不再把内核编码当作普通代码生成,它正在成为自己的后训练任务域——有专门的环境、奖励和数据管线。

第二个分支:不做重训练、强调闭环迭代的代理式搜索。 很多系统把内核优化当成一个闭环程序合成问题:生成候选、编译、测试正确性、profile、变异、保留胜者。2025 年早期的例子有 CUDA-LLM(对 CUDA 内核做迭代精化循环)和 NVIDIA 的 DeepSeek-R1 + 推理时扩展实验(把注意力内核生成当作推理时的搜索问题)。2026 年的系统更系统化:AutoKernel 对 PyTorch 瓶颈跑自主优化循环,带强正确性测试;StitchCUDA 用 planner/coder/verifier 代理做端到端 CUDA 生成;KernelFalcon 用深层代理、层次化分解和持久内存,在 Meta 的博客里报告 KernelBench 100% 正确率。关键转变是从「LLM 写一个内核」变成「LLM 编排一个长程搜索过程」。

第三个分支:硬件感知与 profiling 感知搜索。 这条线想弥合「泛泛的代码推理」和「人类性能工程师真正做的——从瓶颈、计数器和硬件结构推理」之间的差距。SwizzlePerf 是个强例子:把架构细节、访问模式、profiling 日志和历史反思喂给 LLM,报告在 GEMM 上恢复了专家级 swizzle 决策,并改善 L2 命中率和内核套件速度。PRAGMA 把细粒度 profiling 直接整合进推理循环。KernelBand 把内核优化建模为层次化多臂老虎机问题,用 profiling 引导探索/利用、减少浪费的试错。更近的一篇论文(《Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidance》)把同一想法再推一步:代理工作的抽象本身很重要,系统一旦接近硬件极限就该停止搜索。这个分支整体在走向瓶颈感知的搜索,而不是盲目变异。

第四个分支:演化与种群式搜索,这可能是 CUDA-L2 之后增长最快的子领域。动机是内核优化的地形不规则、非单调、充满局部最优,维持一个多样化的候选种群比单链局部精化更好。GPU Kernel Scientist 用演化流程做 AMD MI300/HIP 内核;cuPilot 引入策略协调的多代理演化,把「策略」当作中间语义表示,加 roofline 引导的提示;K-Search 加入共同演化的世界模型,让规划和代码实例化部分解耦,报告在 FlashInfer 的 GQA、MLA、MoE 复杂内核上有显著收益;KernelFoundry 用 MAP-Elites 质量多样性搜索、元提示演化和硬件感知参数调优;AVO 让变异算子本身代理化,在 Blackwell 注意力内核上报告超过 cuDNN 和 FlashAttention-4;Kernel-Smith 把演化代理与从演化轨迹提炼的后训练配方结合,报告了整体最优的 KernelBench 性能,还迁移到 MetaX 的 MACA 后端和 SGLang、LMDeploy 等真实集成。共同模式是:搜索不再是「多采样几个补全」,而变成内核程序上的结构化种群动力学。

第五个分支,也很重要:异构硬件与非 CUDA 扩展。 CUDA-L2 之后的一个大问题是:这个领域会不会永远 NVIDIA/CUDA 中心。现在答案越来越像「不会」。KernelEvolve 面向 NVIDIA GPU、AMD GPU、Meta 加速器和 CPU 的异构生产环境,明确横跨 Triton 和 CuTe 等多种编程抽象。KernelCraft 为带新 ISA 的新兴硬件提供近裸机内核生成的基准。CuTeGen 把 LLM 驱动生成推进到 CuTe/CUTLASS 风格的设计空间。更早的 QiMeng 系列(QiMeng-GEMM、QiMeng-TensorOp、QiMeng-Attention)都用 LLM 引导搜索跨目标生成硬件感知算子,是最早把「LLM + 结构化算子优化」从投机变成具体工作的系统之一。

还有一个重要进展是更好的基准与评测纪律。 KernelBench 成了「模型能否从 PyTorch 工作负载生成快速且正确的内核」的核心公共基准。更新的工作开始攻击评测的弱点:KernelCraft 朝新兴硬件和近裸机 ISA 扩展;多篇论文强调稳健的正确性测试、多轮日志、抵抗 reward hacking;综述明确把评测稳健性和泛化列为领域主要开放问题之一。这很重要,因为很多报告的战绩严重取决于比较对象是谁——eager PyTorch、torch.compile、cuBLAS 启发式,还是完全自动调优的厂商库——以及是否允许长搜索预算。

小结。 CUDA-L2 证明了 LLM + RL 能在狭窄的设定里打败被高度优化的矩阵内核;下一波把这个洞见推广成完整的代理式系统——组合规划、profiling、演化式搜索,有时加 RL 后训练。最强的新系统在收敛到同一个配方:领域专用抽象、严格正确性检查、硬件反馈、种群或多轮轨迹上的结构化搜索、对过去成功编辑的记忆。剩下的主要瓶颈是样本效率、可靠的奖励设计、对强基线的公平评测,以及跨硬件泛化。


14. 结语

如果你只从这一章带走一样东西,带走这个。

读任何一个深度学习编译器时,问自己五个问题。 它面向哪种问题区间——很多模型很多后端,还是一个主导平台上的窄而高价值的工作负载?它工作在哪一层表示——图 IR、张量 IR、调度 IR、硬件 IR、运行时,还是几层的混合?那一层里哪些结构可见——全局依赖、循环嵌套、内存作用域、线程层次,还是流水线形状?它把哪些优化旋钮暴露给谁——程序员、调度语言、代价模型,还是学习型智能体?它的搜索策略是什么——手写启发式、模板、代价模型引导、演化、强化学习,还是 LLM 代理?能回答这五个问题,你就通常已经理解了一个系统的本质。

把这个框架套回本章的四个系统,四格就填满了:TVM 面向多模型多后端,工作在丰富 IR 栈上,把旋钮暴露给调度语言和代价模型,用模板加自动调优搜索。Triton 面向收敛后的 Transformer/NVIDIA 世界,工作在块级 DSL,把大多数旋钮藏起来,用参数自动调优搜索。TileLang 面向特化内核,在 tile 抽象下把布局、流水、线程绑定还给人类,搜索策略是手工专家经验。CUDA-L2 面向同样的特化内核,但让 LLM + RL 系统去搜,旋钮暴露给学习智能体。四格填完,历史自己就浮现出来了。

最后的总结不必再重复每个系统。 深度学习编译器存在,是因为高层张量计算必须映射到底层硬件执行,而这个映射是一个大规模、结构化的优化问题。优化旋钮——平铺、内存放置、循环顺序、融合、展开、线程映射、同步——构成一个离散而耦合的搜索空间,大到无法在大规模场景下靠人完成。编译器栈存在,因为不同层保留不同结构,不同优化需要不同可见性。TVM 是 LLM 之前宽泛的答案;Triton 是 LLM 时代早期高生产力的答案;TileLang 是特化让自动化不够用时、控制回归的答案;CUDA-L2 是让学习系统接管搜索的答案。放在一起,它们讲的是一个连贯的故事:AI 系统的未来不在任何单一编译器抽象里,而在表示、优化和日益自动化的搜索三者之间的分层综合里。

这个循环还没走完,而且大概率不会走完。工作负载还会变,硬件还会变,自动化与人工控制的平衡还会继续摆动。这正是这门课里最值得带走的习惯:看一个系统,不要只看它今天的界面,要看它是对哪个世界的形状做出的回答。


延伸阅读

核心编译器系统

Triton

TileLang

CUDA-L2 与学习式内核优化

相邻语境