作者：画饼充饥
https://zhuanlan.zhihu.com/p/2040199323170951424

最近看AI infra和AI compiler的论文，一个很明显的变化是大家已经不满足于把某个算子编译的更快了。

早期深度学习编译器的叙事多半围绕operator graph，例如MatMul、Conv、Softmax、LayerNorm、Add这些算子怎么融合，layout怎么变换，schedule怎么搜索，最后怎么生成CUDA或Triton kernel。

这个范式在CNN、BERT时代很自然，但放在今天LLM、长上下文、KV cache、MoE、动态batch、分布式训练上，operator这个边界开始显得更粗粒度。

AI编译器正在重新刻画深度学习程序的执行边界。

过去一个 operator 是边界，一个 kernel launch 是边界，一个框架 runtime hook 是边界。现在这些边界都在被拆开、下沉或上移——旧的执行单位装不下现在的AI workload了。

1. 从operator到tile，编译器的单位正在下沉

过去 operator 是深度学习系统的默认单位。一个 Attention 是一个 operator，一个 LayerNorm 是一个 operator，一个 MatMul 是一个 operator。

问题是 GPU 真正在意的是数据怎么在 HBM（ High Bandwidth Memory ）、shared memory、register 之间移动，tile 怎么复用，Tensor Core 怎么喂满，warp 怎么分工，kernel 边界有没有同步bubble。

于是最近几年大量工作都在把编译器视角从 operator 下沉到 loop、block、tile、task。

TensorIR 是这条线比较早的代表。它把 tensor computation primitive 变成一等对象，让编译器能围绕 tensorized hardware primitive 进行自动优化，而不是只在传统 loop nest 上做通用变换。

Welder 则从 memory access 的角度推进了一步，用 tile-graph 显式建模 tile 级数据移动和复用，把跨算子的 memory locality 放进编译器。

Ladder 关注的是低精度和自定义数据类型，它用 tType 和扩展 tensor expression 把 custom data type 放进编译器优化空间，而不是把 dtype 当成一个简单属性。

TensorIR、Welder、Ladder 这三类工作共同说明了现代 AI 编译器需要的不是单纯 graph IR，也不是裸 CUDA，而是一个能同时表达数学结构、tile 数据流和硬件约束的中间层。

TileLang 和 ThunderKittens 则更偏向如何让人写高性能 kernel 更容易。Triton 已经把 tile programming 带到了 Python 生态，但写好 Triton kernel 依然需要很强的硬件经验。

TileLang 的思路是把 dataflow 和 scheduling space 分开，用户专注描述 tile 级数据流，thread binding、layout、tensorize、pipeline 等交给 compiler annotation 和 primitive 管理。

ThunderKittens 则用一小组面向 GPU 层级的抽象简化 AI kernel 编写，比如 warp-level tile、thread-block-level async pattern 等。它们不是要完全替代 compiler，而是在填补PyTorch 太高层、CUDA 太底层之间的开发空档。

Cypress 代表了另一个方向。现代 GPU 上 TMA、Tensor Core、warp specialization、producer-consumer pipeline 越来越复杂，用户手写同步和数据搬运的负担越来越重。

Cypress 用 task-based tensor computation 来表达 GPU 张量计算，用户写顺序语义的 task 和 mapping spec，compiler 负责插入数据移动和同步。

它和 TileLang、ThunderKittens的共同点在于GPU kernel 编程正在从写线程转向组织 tile/task。

2. 从手写专家kernel到语义融合

如果说 tile 级编程解决的是执行单位太粗，那还有一类工作解决的是另一个问题传统 fusion 规则太保守。

普通 operator fusion 最擅长处理 elementwise 链，比如relu(x)+1。但 Attention 里的 fusion 难点不是 elementwise，而是 reduction dependency。稳定 softmax 要先对一行求 max，再用这个 max 去算 exp 和 sum。

传统 loop fusion 看到后面的 sum 依赖前面的 max 最终值，就会拒绝融合。FlashAttention 的成功恰恰在于它绕开了这个限制。它不 materialize 完整 attention score 和 softmax，而是用 online softmax 在 tile 内不断更新 max、sum 和输出。

Neptune 这个工作亮点在于它把这个专家手写技巧抽象成了编译器 transformation。它的核心可以概括成break dependency, then repair algebraically。先故意做一个 naive fusion，让 consumer reduction 暂时读到 producer reduction 的中间值，然后自动推导一个 repair term，把旧坐标系下的 partial result 修到新坐标系下。

对 softmax 来说，这个修正项就是exp(old_max-new_max)。当这个 transformation 用在 plain attention 上，就得到类似 FlashAttention 的 rolling update，用在 decoding attention 上，就得到类似 FlashDecoding 的 split-k update。

这篇论文的启发不在于又一个 attention kernel，而在于它给了一个很清晰的研究范式是找到专家手写 kernel 中反复出现的算法结构，把它提升成带正确性条件的 compiler primitive。 FlashAttention 背后的 online softmax 是一个例子，FlashDecoding 背后的 split-k partial reduction 是一个例子。

未来 LayerNorm/RMSNorm backward、Linear + CrossEntropy、DPO loss、optimizer update，也可能有类似的专家经验编译器化空间。

FlexAttention 和 FlashInfer 则体现了 attention 生态的另一个趋势。attention 不再是一个固定算子，而是一族可编程算子。

FlexAttention 让用户用少量 PyTorch 代码表达 score modification、mask、block sparsity 等 attention 变体，再由编译器生成高性能 fused kernel。

FlashInfer 面向 LLM serving，把 KV-cache 存储异构性、block-sparse 格式、JIT attention template 和 load-balanced scheduling 结合起来，目标是在真实 serving 场景里覆盖 PagedAttention、长上下文、并行生成等动态需求。

Nautilus 则把 Neptune 往自动化方向又推了一步。

Neptune 仍然需要 high-level schedule template，而 Nautilus 试图从高层代数规格出发自动发现高级优化序列，并声称能自动发现 FlashAttention-3-like kernels。

这说明 attention compiler 的路线正在从用户写 kernel走向用户写数学定义，compiler 搜索高级语义优化。

3. Superoptimizer、Megakernel和动态执行

Superoptimizer 这两年重新火起来，但已经不是传统意义上的指令级 superoptimizer，而是 tensor program 级、甚至多层级 superoptimizer。Mirage 是代表。

它提出 µGraphs，用统一表示覆盖 GPU compute hierarchy 里的 kernel、thread block 和 thread 层级，从而同时搜索 algebraic transformations、schedule transformations 和 custom kernel generation。Mirage 还用 probabilistic equivalence verification 来检查优化后的 tensor program 是否和原始程序等价。

这个方向火起来并不奇怪，模型变体太多，shape 太多，dtype 太多，attention mask 太多，仅靠专家不可能给所有组合都手写最优 kernel。

不过，纯黑盒搜索很难成为最终答案。搜索空间大、compile time 长、dynamic shape 泛化差、浮点正确性难讲。更有前途的路线可能是语义约束下的 superoptimization。

Neptune 不是盲目搜索 FlashAttention，而是识别了 reduction repair 这个结构，Mirage 不是只在 graph 层搜索，而是把搜索空间组织成 µGraph，Nautilus 也不是单纯枚举低层 schedule，而是让高层优化、表达式重写和 tile 优化在一个 successive lowering 体系里联动。真正有影响力的优化，往往不是搜得更多，而是把搜索空间组织得更对。

Megakernel 是另一条很热的线，尤其在 LLM inference 上。decode 阶段经常是小 batch、token-by-token，kernel launch overhead、kernel boundary synchronization、tail effect 和 cold start 会被放大。

Megakernel / persistent kernel 的思路是把多个 operator 或 task 放进一个长期运行的 kernel，在 GPU 内部自己调度，从而减少 launch gap 并暴露 inter-kernel parallelism。Mirage Persistent Kernel 这类工作把 LLM inference 编译成 single megakernel，用 SM-level graph 表达依赖，做跨 operator software pipelining 和细粒度 overlap。

但 megakernel 不是简单地kernel 越大越好。

真实 LLM serving 有 dynamic batching、paged KV cache、不同请求长度、speculative decoding 分支、data-dependent computation。静态 megakernel 很难处理这些动态性。

Event Tensor 的出发点正是，近期 megakernel 技术可以消除 launch gaps，但面对 dynamic shapes 和 data-dependent computation 仍然困难，它用 Event Tensor 显式编码 tiled tasks 之间的依赖，使 dynamic megakernel 能支持 shape-dependent 和 data-dependent dynamism，并在此基础上生成 persistent kernels。

Infera 站在 serving system 侧给出了另一个版本的答案。

它不是把所有东西都塞进一个 persistent megakernel，而是在编译期把大 operator partition 成 tiles/micro-kernels，并生成多版本 kernel，运行时再根据 GPU 状态、任务优先级和 kernel 属性进行动态调度。

它还强调 compile strategy 和 scheduling strategy 必须匹配，否则 kernel 单独最优并不等于系统最优。

4. 分布式编译器，通信也在进入IR

训练侧的趋势和推理侧不一样。推理侧强调 kernel launch、dynamic serving 和 attention/KV cache，而训练侧最大的系统瓶颈之一仍然是通信。DDP、Tensor Parallel、Pipeline Parallel、ZeRO/FSDP、MoE expert parallel 都会引入通信。

过去这些通信往往藏在框架实现里，DDP reducer hook、FSDP all-gather hook、ZeRO bucket、Megatron-LM 手写 all-reduce overlap。通信如果不在图里，编译器就只能优化 compute graph，看不见真正的执行瓶颈。

Concerto 的贡献就是把这个问题抽象成编译器问题。

它先 trace PyTorch function 得到 FX graph，再根据 parallel method 转成包含 computation operators 和 communication operators 的 ConcertoIR，之后把调度建模成 resource-constrained project scheduling problem，并通过 auto-decomposition 给 critical communication 创造 overlap 机会。

Triton-distributed 和 TileLink 则把这个问题进一步推进到 kernel 编程层。

Triton-distributed 扩展 Triton，把 OpenSHMEM-compliant communication primitives 集成到 compiler 中，让程序员能用高层 Python/Triton 风格写分布式 GPU kernels，并在单机和多机上做 computation、memory access、communication 的联合优化。它的目标不是简单替代 NCCL，而是让 fine-grained overlap 进入 compiler-generated kernel 的表达能力。

TileLink 的角度更像是给重叠 kernel 一个 tile-centric 编程模型。传统做 compute-communication overlap 有两条路。operator decomposition 容易做但性能可能差，communication kernel 和 compute kernel 手工融合性能好但很难写。

TileLink 把 computation 和 communication 的设计空间解耦，再用 tile-centric primitives 把二者连接起来，后端把这些 primitives 翻译成低层通信指令并生成 overlapped kernels。

这说明分布式 AI 编译器也在经历和单卡 kernel compiler 类似的抽象演进。

• 第一阶段是 runtime 里手工调 bucket、prefetch、overlap；
• 第二阶段是 Concerto/DeepCompile 式的 compute-communication graph；
• 第三阶段是 Triton-distributed/TileLink 式的 distributed kernel programming，让跨 GPU 通信成为 kernel 内部的可编译对象。

这个变化很重要，因为未来 MoE、tensor parallel、sequence/context parallel 都越来越需要 tile/chunk 级别的通信计算交错，单纯靠框架 hook 很难做到足够细。

5. 低精度、硬件异步和编译器越来越像执行系统

还有一条趋势容易被忽略：dtype、硬件异步能力和编译器正在深度耦合。

FlashAttention-3 在 Hopper 上利用 Tensor Cores 和 TMA 的 asynchrony，通过 warp specialization overlap 数据搬运和计算、interleave block-wise matmul 和 softmax，并结合 FP8 low precision，在 H100 上相对 FlashAttention-2 有 1.5–2.0× 的 FP16 加速，FP8 路径接近 1.2 PFLOPs/s。

这里的重点不是又优化了 Attention，而是现代 GPU 的能力已经要求 kernel 同时处理异步拷贝、warp specialization、Tensor Core scheduling、低精度数据路径和数值误差。

Ladder 说明 dtype 本身也在变成编译器对象。

低精度不再只是fp16/bf16/int8这种 tensor 属性，而是一套计算协议：storage type、scale、conversion、accumulation、layout、hardware instruction mapping。Ladder 用 tType 和扩展 tensor expression 把 custom data types 作为 first-class citizen，目标是在不断演进的数据类型和固定硬件 precision format 之间架桥。

这也解释了为什么 TileLang、Cypress、ThunderKittens 这些工作会出现。硬件越来越强，但也越来越难编程。

写高性能 kernel 不再只是选 tile size，而是要组织 producer/consumer pipeline、异步搬运、warp specialization、register/shared memory 资源、Tensor Core 指令、低精度转换。AI 编译器如果还停留在传统 loop schedule，就很难覆盖这些复杂性。

6. 总结

如果总结起来AI编译器最火的是什么，回答可以是megakernel、superoptimizer、tile programming、attention compiler、distributed compiler。但这些词只是表面，深层的共同点是传统抽象边界正在失效。

Operator 边界失效，所以有 Welder、Neptune、TileLang、ThunderKittens。

Kernel launch 边界失效，所以有 megakernel、persistent kernel、Event Tensor。

Framework hook 边界失效，所以有 Concerto、Triton-distributed、TileLink。

固定 attention 算子边界失效，所以有 FlexAttention、FlashInfer、Neptune、Nautilus。

固定 dtype 边界失效，所以有 Ladder 和 FP8/INT4/自定义格式相关编译工作。

Compile-time schedule 和 runtime schedule 边界失效，所以有 Infera 这类 compiler-scheduler co-design。

这给后续研究一个很清晰的启发。**不要只追又快了多少的局部优化，而要问一个更本质的问题：现有系统里还有什么重要对象没有进入 IR？ **

过去几年，通信进入了 IR，tile 进入了 IR，dtype 进入了 IR，event 进入了 IR，attention mask 进入了 IR。

训练里还有很多东西没有被很好地抽象，比如 backward、RNG、optimizer state、saved tensor、loss function、post-training preference objective、MoE routing。

谁能把这些对象变成可编译、可验证、可生成高性能实现的 IR object，可能就会有机会做出好的 AI compiler 工作。

参考文献

[1] https://doi.org/10.1145/3575693.3576933
[2] https://docs.pytorch.org/assets/pytorch2-2.pdf
[3] https://www.usenix.org/system/files/osdi23-shi.pdf
[4] https://www.usenix.org/system/files/osdi24-wang-lei.pdf
[5] https://arxiv.org/abs/2504.17577
[6] https://arxiv.org/abs/2410.20399
[7] https://arxiv.org/abs/2504.07004
[8] https://arxiv.org/abs/2205.14135
[9] https://arxiv.org/abs/2307.08691
[10] https://arxiv.org/abs/2407.08608
[11] https://arxiv.org/abs/2412.05496
[12] https://arxiv.org/abs/2501.01005
[13] https://arxiv.org/abs/2510.08726
[14] https://arxiv.org/abs/2604.14825
[15] https://arxiv.org/abs/2405.05751
[16] https://arxiv.org/abs/2512.22219
[17] https://arxiv.org/abs/2604.13327
[18] https://dl.acm.org/doi/10.1145/3767295.3769392
[19] https://dl.acm.org/doi/pdf/10.1145/3669940.3707223
[20] https://arxiv.org/abs/2504.09983
[21] https://arxiv.org/abs/2503.20313
[22] https://arxiv.org/abs/2504.19442
[23] https://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf
[24] https://www.usenix.org/system/files/osdi22-zhu.pdf
[25] https://arxiv.org/abs/2006.06762