AMD FlyDSL接入TorchInductor：MLIR原生GPU内核新后端

一个新的内核后端为何重要

现代AI工作负载越来越依赖种类繁多的专用GPU内核。来自AMD的Liz Li与Jia Hui在PyTorch相关技术分享中给出了一个核心判断：没有任何单一的内核生成后端能够在所有算子、芯片和架构上都做到最优。这意味着编译栈需要具备足够的灵活性，能够按场景选择最合适的内核实现路径。

正是基于这一出发点，AMD团队将自研的FlyDSL作为新后端集成进TorchInductor。FlyDSL是AMD推出的Python原生、基于MLIR的GPU内核DSL（领域特定语言）。它既保留了Python生态的易用性，又借助MLIR这一现代编译器基础设施，提供了更接近硬件的优化空间。

MLIR（Multi-Level Intermediate Representation）是由Google主导、LLVM社区维护的现代编译器基础设施，其核心设计理念是通过多层次的中间表示（IR）将高层算法描述逐步"下降"到硬件指令。与传统编译器不同，MLIR允许在同一框架内共存多套方言（Dialect），使得从Python算子到GPU汇编的整条变换链可以在统一框架下完成，同时保留每一层的优化机会。这与Triton的编译路径形成对比——Triton同样生成GPU内核，但其抽象层次和优化粒度有所不同。FlyDSL基于MLIR意味着它能够利用MLIR已有的Pass基础设施（如循环变换、向量化、内存层级分析）来生成更贴近硬件特性的代码，而无需从头实现这些优化逻辑。

无缝嵌入TorchInductor编译流水线

这项工作的工程价值在于集成方式。AMD团队并没有另起炉灶构建一套独立工具链，而是把FlyDSL接入了Inductor现有的核心环节——包括lowering（算子下降）、编译、缓存、benchmarking以及auto-tuning自动调优流水线。

这种做法最大的好处是开发者体验的连续性。按照演讲者的说法，整个过程保留了用户熟悉的Torch与Triton使用体验，意味着上层代码几乎不需要改动，就能在底层切换到FlyDSL生成的内核。对于已经构建在PyTorch之上的大量模型和应用而言，这是一种低迁移成本的性能升级路径。

覆盖GEMM与注意力的关键算子

有了这一集成基础，AMD得以把多个经过优化的FlyDSL内核带入PyTorch。演讲中列举的典型案例包括GEMM（通用矩阵乘法）、Group GEMM以及Flex Attention——这三类恰好是大模型训练与推理中的高频、重计算算子。

GEMM是几乎所有神经网络计算的基石，Group GEMM在MoE（专家混合）等结构中尤为关键，而Flex Attention则直接影响Transformer类模型的注意力计算效率。FlyDSL把优化重心放在这些算子上，针对性明确。

MoE（Mixture of Experts，专家混合）是一种将模型容量与计算量解耦的架构设计。其核心思想是在前向传播时，对每个token只激活全部"专家"子网络中的少数几个，从而在参数规模大幅扩展的同时保持计算量相对可控。Group GEMM正是MoE计算的瓶颈所在：不同专家接收的token数量不同，导致需要同时执行一批形状各异的小矩阵乘法，而非一个规整的大矩阵乘，这给内核调度和硬件利用率带来额外挑战。DeepSeek^广告 V2/V3等当前主流大模型均采用MoE结构，因此Group GEMM的优化对实际部署场景具有直接价值。

在MI355上的性能数据

真正值得关注的是实测数字。在AMD MI355上，FP16 GEMM相比Triton实现了约19%的几何平均加速。这一数字看似温和，但考虑到GEMM是整个计算图中占比最大的算子之一，接近两成的提升会直接传导到端到端吞吐。

Flex Attention的收益则更为显著。在Minimax M3的decode形状下，FlyDSL Flex Attention相对Triton最高可达5.8倍加速；在DeepSeek V2与V3的prefill尺寸上，也实现了接近2倍的提升。

需要说明的是，注意力算子的加速比往往对具体的shape和场景高度敏感——decode阶段的小batch、长序列特征与prefill阶段的大矩阵计算有本质差异，因此5.8倍这样的峰值数字应放在特定形状语境下理解。即便如此，这些数据表明MLIR原生路径在应对多样化算子形状时确有优势。

Flex Attention是PyTorch 2.x引入的一种可编程注意力机制，允许开发者通过自定义score modification函数来表达多种注意力变体（如滑动窗口注意力、因果掩码、ALiBi偏置等），编译器再将其融合为高效内核。Decode与Prefill是推理过程中两个特性截然不同的阶段：Prefill处理完整的输入提示，通常涉及大batch与长序列的矩阵计算，计算密集；Decode则逐token自回归生成，每步只产生一个新token，batch小、序列维度以KV Cache形式增长，访存密集。两种阶段对内核的要求差异悬殊，这也解释了为何同一FlyDSL实现在不同形状下加速比跨度可以从2倍到5.8倍如此之大。

对AMD软件栈的意义

长期以来，Triton被视为跨硬件内核编写的主流路径之一，而AMD通过FlyDSL给出了一条基于自家MLIR技术栈的替代方案。它既不排斥Triton，又在特定算子上提供了明显超越Triton的性能，同时保持了PyTorch上层接口的一致性。

这反映出AMD在GPU软件生态上的持续投入：通过把优化成果直接嵌入PyTorch最常用的编译路径，降低开发者享受硬件红利的门槛。随着这类工作的推进，多后端共存、按算子择优的编译策略有望成为高性能AI栈的常态。

演讲者表示将在后续活动中分享更多技术细节与实践经验。对于关注AMD GPU性能、PyTorch编译优化的工程师而言，FlyDSL是一个值得持续跟踪的方向。