10月23日消息，北京时间 10 月 21 日，PyTorch 官方联合 AMD 发布博文，宣布将 FlyDSL 高性能内核后端集成进 PyTorch ROCm 生态，面向 AMD Instinct MI350（gfx950）系列加速器开放。

开发者无需改动业务代码，即可在现有 API 上启用优化内核，覆盖矩阵乘法、分组 GEMM、缩放低精度矩阵运算、RMSNorm、TopK、融合多头注意力等大模型高频算子，补齐 AMD CDNA4 平台在 PyTorch 上的内核性能短板。

FlyDSL 是一套基于 Python+MLIR 的领域专用语言。和 Triton 相比，它开放了更底层的布局代数、共享内存、波前调度、硬件 MFMA 指令控制能力，适合针对 AMD 架构做极致调优；同时作为可插拔后端，和 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 倍；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 包即可开启。官方也公布了后续路线图：支持 MI450 硬件、扩展 MXFP 分组矩阵乘、适配 FlexAttention、增加算子融合、完善反向传播内核、支持 AOTInductor 部署，持续补齐训练侧能力。

对 AMD 生态开发者而言，本次集成降低了在 MI350 集群上跑 PyTorch 大模型的调优门槛，在不侵入上层代码的前提下，直接提升训练、推理吞吐，降低单位 token 算力成本。