PyTorch 官方联合 AMD 集成 FlyDSL 高性能内核，全面加速 MI350 系列大模型计算

SmartHey10月23日消息，北京时间10月21日，PyTorch 官方与 AMD 联合宣布，正式将 FlyDSL 高性能内核后端集成至 PyTorch ROCm 生态，首批支持 AMD Instinct MI350（gfx950）系列加速器。

开发者无需修改任何业务逻辑，即可通过现有 PyTorch API 透明启用优化内核，覆盖矩阵乘法、分组 GEMM、MXFP8/MXFP4 缩放低精度矩阵运算、RMSNorm、TopK 以及融合多头注意力（SDPA）等大模型核心算子，显著弥补 CDNA4 架构在 PyTorch 生态中的内核性能短板。

FlyDSL 是一套基于 Python + MLIR 构建的领域专用语言。相较于 Triton，它向开发者开放更底层的硬件控制能力——包括显式布局代数、共享内存管理、波前（Wavefront）调度策略及 MFMA 指令级编程接口，专为 AMD GPU 架构深度调优而设计；同时作为可插拔后端，与 ATen 和 Triton 共存：对暂不支持的输入形状，系统将自动回退至原有实现，确保全场景兼容性与稳定性。

本次发布的优化算子采用双路径调度机制：RMSNorm 与 TopK 在 Eager（即时执行）模式下全自动启用；稠密/分组矩阵乘、MXFP8/MXFP4 缩放矩阵乘则接入 torch.compile，由 TorchInductor 自动完成多后端性能评估与最优内核缓存；而缩放点积注意力（SDPA）算子依托 AOTriton 库，在 FlyDSL、Triton 与 AITER 汇编内核间动态择优，兼顾性能与泛化能力。

官方基于 MI355X 加速器开展系统性基准测试，多项关键算子实现稳定性能跃升：BF16 稠密 GEMM 相比 ATen/Triton 最优基线几何平均提速 1.10 倍；MXFP8 缩放矩阵乘平均加速 1.58 倍，MXFP4 达 1.68 倍；分组 GEMM（MoE 场景常用）相对 Triton 提升 1.20 倍；RMSNorm 加速区间为 1.17–3.66 倍；小 K 值 TopK 在关闭确定性算法时，最高几何平均加速达 4.82 倍；注意力内核借助 AOTriton 0.14.2b 更新，前向与反向吞吐量最高提升超 2 倍。

端到端推理验证已在 vLLM 框架中完成 AB 测试，覆盖 Llama、Qwen 等主流开源大模型。结果显示：BF16 精度下整体请求吞吐最高提升 1.13 倍；MXFP8 量化版本达 2.04 倍；MXFP4 量化版本更高达 3.01 倍——低精度推理收益尤为显著，对 MoE 架构与长上下文推理等高负载场景具备突出实用价值。

当前功能需搭配 ROCm 版本的 PyTorch Nightly，并额外安装 flydsl Python 包即可启用。官方同步公布了后续演进路线图：支持新一代 MI450 硬件、扩展 MXFP 分组矩阵乘支持、适配 FlexAttention 机制、增强算子融合能力、完善反向传播内核、推进 AOTInductor 生产部署支持，持续强化训练侧性能与工程落地能力。

对 AMD ROCm 生态开发者而言，此次集成大幅降低了在 MI350 集群上高效运行 PyTorch 大模型的技术门槛——无需侵入上层代码，即可获得训练与推理吞吐的实质性提升，切实降低单位 token 的算力成本。