AI 编译器正在重写执行边界
青稞AI · wechat · 2026-07-18
这篇文章系统回顾了 AI 编译器与推理/训练基础设施的演进，核心观点是：**传统以 operator 为边界的编译范式已经不够用了**，现在的工作正在把执行单位下沉到 tile/task、把语义融合做深、把通信也纳入 IR。主要内容包括： ### 1) 编译边界从 operator 下沉到 tile/task 作者认为，LLM、长上下文、KV cache、MoE、动态 batch、分布式训练等工作负载，使得“一个算子一个边界”的抽象越来越粗。 文中提到的代表方向包括： - **TensorIR**：把 tensor computation primitive 变成一等对象。 - **Welder**：从 memory access 角度显式建模 tile 级数据移动与复用。 - **Ladder**：把低精度和自定义 datatype 纳入优化空间。 - **TileLang / ThunderKittens**：让开发者更容易写高性能 GPU kernel。 - **Cypress**：用 task-based tensor computation 描述 GPU 计算，交给编译器处理数据搬运和同步。 ### 2) 从普通 fusion 走向“语义融合” 文章指出，传统 operator fusion 对 attention 这类带 reduction 依赖的模式不够用。FlashAttention 的关键是用 online softmax 绕开 materialize 全量 score 的限制。 它重点介绍了 **Neptune**： - 先做一个“看似不合法”的 naive fusion； - 再通过 algebraic repair 自动修正中间 reduction； - 例如 softmax 里会用到 `exp(old_max - new_max)` 这类修正项； - 该思路既能覆盖 plain attention，也能覆盖 decoding attention，接近 FlashAttention / FlashDecoding 的专家技巧。 同时还提到： - **FlexAttention**：让用户用少量 PyTorch 代码表达 mask、score modification、blocksparsity 等变体。 - **FlashInfer**：面向 LLM serving，把 KV cache、block-sparse 格式、JIT attention template 和负载均衡调度结合起来。 - **Nautilus**：尝试从高层代数规格自动发现类似 FlashAttention-3 的优化序列。 ### 3) superoptimizer、megakernel 与动态执行 文中把最近一波 superoptimizer 重新火起来的原因归结为：模型变体、shape、dtype、attention mask 太多，专家无法手写所有最优 kernel。 代表工作包括： - **Mirage**：用 `µGraphs` 统一表示 kernel / threadblock / thread 层级，并搜索代数变换、调度变换和自定义 kernel 生成。 - **megakernel / persistent kernel**：把多个 operator 或 task 放到一个长期运行的 kernel 中，减少 launch overhead。 - **EventTensor**：显式编码 task 之间的依赖，让 dynamic megakernel 能处理 shape-dependent 和 data-dependent 动态性。 - **Infera**：在编译期把大 operator 切成 tiles/micro-kernels，再结合运行时状态做动态调度。 作者强调：真正有效的优化不只是“搜得更多”，而是**把搜索空间组织得更对**。 ### 4) 分布式编译器开始把通信纳入 IR 训练侧的瓶颈往往是通信，而不是纯计算。文章总结了把通信显式化的路线： - **Concerto**：把 PyTorch trace 成 FX graph，再构造成包含计算和通信算子的 IR，用资源受限项目调度去安排 overlap。 - **Triton-distributed**：把通信原语直接集成进 Triton 风格的高层编程模型。 - **TileLink**：用 tile-centric 的方式表达 computation-communication overlap。 ### 5) 总结 作者最终的判断是： - AI 编译器正在从“算子优化”走向“语义与系统协同优化”； - attention、megakernel、通信 overlap 都在被编译器重新抽象； - 未来真正有影响力的方向，可能不是某个单点 kernel，而是**把专家经验抽象成可组合的编译原语**。