AI 编译器的现实账单
过去三年,AI 编译器领域密集地涌现了一批工具和框架。Triton、OpenXLA、MLIR、TVM、Halide 改进版……每隔几个月就有一篇论文或一个新项目声称”让工程师用高层语言写出接近手写 CUDA 的性能”。
这个承诺很诱人,因为 CUDA 手写 kernel 是一件极度依赖专家经验的事。但现实是:这些工具在生产环境里的落地进展,远比 demo 和论文里展示的要慢。
承诺和现实之间的缝隙
Triton 是过去两年讨论最多的那个。OpenAI 把它开源出来时,配套的叙事是:用 Python 写 GPU kernel,性能接近 CUDA,可以让更多工程师参与算子优化。
这个叙事有真实性,但需要拆解。
Triton 确实降低了写 GPU kernel 的门槛——如果你已经理解 GPU 的内存层次、并发模型和 tile 策略,Triton 能让你更快地把想法变成代码。但如果你不理解这些,Triton 并不能帮你填平这个认知缺口。”用 Python 写”解决的是语法问题,不是算法问题。GPU 编程最硬的部分——如何划分计算 tile、如何安排共享内存、如何 hide memory latency——Triton 没有自动帮你做,只是让你用更简洁的语法手动表达。
真正在生产中用 Triton 写出高效 kernel 的团队,几乎无一例外是本来就有 CUDA 经验的人。他们得到的收益是开发效率,不是入门门槛降低。
编译器路线的结构性问题
MLIR 是 Google 在 2020 年推出的多层次中间表示框架,定位是统一 ML 编译器的基础设施。XLA 基于它重构,TensorFlow 和 JAX 都走这条路。理论上很漂亮:不同的硬件后端、不同的并行策略,都可以在统一的 IR 体系里表达和变换。
实践里的问题是调试。
MLIR 的变换链条很长——从高层 op 到底层 lowering,经过十几个 pass,最终生成目标代码。当生成的代码性能不符合预期,定位是哪一层 pass 引入了性能退化,是一件非常痛苦的事。工具链的 profiler 和 IR 可视化工具还不够成熟,出了问题很难快速定位。
相比之下,手写 CUDA kernel 加一个 Nsight Compute 的 profiler 分析,反而更直接。
这不是 MLIR 的设计缺陷,而是编译器基础设施成熟度的问题。成熟需要时间,但现在有些团队在评估的时候没有把这个时间成本算进去。
算子库路线没有消失
另一条路是 NVIDIA 的 cuDNN、cuBLAS、NCCL,以及 Flash Attention、Paged Attention 这类专项算子库。
这条路的逻辑是:对于主要的计算模式(矩阵乘法、注意力、归一化),有专家用汇编或接近汇编的 CUDA PTX 手写了极度优化的实现,直接调用就好。上层框架(PyTorch、JAX)做的事情,是把这些算子拼起来,管理内存和调度。
这条路在当下其实更务实。FlashAttention 系列的影响力已经说明了一件事:当某个算子的重要性足够高,有足够多的专家愿意在这个点上投入,手工优化的结果还是比编译器自动生成的好。
算子库路线的问题是覆盖范围。新的模型架构、新的硬件,不一定有现成的算子库覆盖。Mamba、RetNet 这类非 Transformer 架构落地时,配套算子的成熟度就明显滞后于架构本身。
两条路的关系
AI 编译器和手工算子库并不是非此即彼的关系,但它们解决的问题域有清晰的边界。
编译器路线适合的场景:模型架构在频繁变化、硬件在多样化、覆盖面比单点极致性能更重要。云厂商自研芯片(TPU、AWS Trainium、华为昇腾)选 XLA/MLIR,有很大程度是因为它们需要一个独立于 NVIDIA 的编译路径,而不完全是因为性能优于 CUDA。
算子库路线适合的场景:模型架构相对稳定、目标硬件明确、对每个算子的极致性能有要求。推理服务在主流 GPU 上跑主流 Transformer 模型,当下的最优解还是组合使用 TensorRT 和 FlashAttention。
问题在于,很多团队在选型时没有想清楚自己在哪个场景,然后在错误的框架上花了大量时间调参数。
什么时候这一层会成熟
判断一个编译器工具链是否成熟,有几个朴素的标准:文档是否足够支持独立排查问题?社区里的 bug 是否能在合理时间内得到响应?从”发现性能问题”到”定位并修复”的周期是否可接受?
Triton 在这些维度上都在快速进步,但还没有到我会放心推荐给一个没有 CUDA 背景的团队直接上生产的程度。
MLIR 生态更复杂,成熟时间线也更长。
一个比较务实的判断是:AI 编译器这一层真正成熟的时间节点,大约是它的调试体验追上传统编译器(如 LLVM + GCC)的时候。 那可能还需要三到五年。
在这之前,它们会持续在研究场景、云厂商内部和新硬件适配场景里发挥价值,但在追求生产稳定性的推理服务里,算子库 + 手动调参仍然是主流。
没有最好的工具,只有适合当前约束的工具。AI 编译器的承诺是真实的,账单也是真实的。了解两者,才能做出合理的工程决策。