AMD GPU 上的 PyTorch 多了一条后端：FlyDSL 让算子按硬件显式排布，vLLM 整请求最高 3.01 倍
核心概要
AMD MI350 系列 GPU 上的 PyTorch 用户现在可以通过既有算子接口调用 FlyDSL 内核，覆盖稠密与分组 GEMM、MXFP8/MXFP4 缩放 GEMM、RMSNorm、TopK 以及经 AOTriton 的注意力。此前这些算子只能走 ATen 或 Triton 等实现；在 MI355X 的算子级测试中，FlyDSL 对 15 个 BF16 稠密 GEMM 形状相对逐形状更快的 ATen/Triton 取得 1.10 倍几何平均加速，MXFP8 与 MXFP4 相对 ATen 分别为 1.58 倍和 1.68 倍。端到端 vLLM 测试中整请求加速最高为 BF16 的 1.13 倍、MXFP8 的 2.04 倍和 MXFP4 的 3.01 倍，具体取决于模型与并发度。
深度剖析
AMD MI350 系列 GPU 上的 PyTorch 应用不必改写调用自定义内核，就能通过 torch.mm、F.scaled_mm、F.grouped_mm、F.rms_norm、torch.topk 和 F.scaled_dot_product_attention 这些既有接口用到 FlyDSL 内核。 此前这些算子由 ATen、Triton 或 AITER ASM 等实现承担，FlyDSL 的硬件映射只存在于独立内核中。 在 MI355X（gfx950）上，15 个 BF16 NT 稠密 GEMM 形状相对逐形状更快的 ATen/Triton 基线取得 1.10 倍几何平均加速，其中 10 例胜、3 例持平、2 例落后，落后的是该套件中两个最大的方形输出问题。
低精度路径把收益扩展到 MXFP8 与 MXFP4 缩放 GEMM，以及按专家分组、允许空组的 GEMM。 缩放 GEMM 需要处理 FP8 或打包 FP4 操作数、每 32 个归约维数值一个 E8M0 尺度，以及 CDNA4 的缩放 MFMA。 17 个形状的 NT 套件上，MXFP8 相对 ATen 为 1.58 倍、MXFP4 为 1.68 倍几何平均；24 个 BF16 分组 GEMM 用例上相对 Triton 为 1.20 倍、相对 ATen 为 1.95 倍。
RMSNorm 与 TopK 在 eager 模式下按输入条件自动分派，不支持的输入保留原有 PyTorch 实现。 这两个逐元素类算子此前没有可选的 FlyDSL 路径，只能走 ATen。 22 个 RMSNorm 用例相对 ATen 为 1.17 倍到 3.66 倍，3.66 倍出现在比对齐尺寸多一个元素的 FP16 情形；10 个小 K TopK 用例在关闭确定性算法时相对 ATen 为 4.82 倍几何平均，开启时为 4.02 倍。
注意力内核以 AOTriton 额外后端的形式进入 F.scaled_dot_product_attention，用户侧无需改动。 SDPA 此前由 AOTriton 的 Triton 与 AITER ASM 内核承担，FlyDSL 不经过 PyTorch 侧后端。 AOTriton 0.14.2b 公布的 gfx950 吞吐相对 0.13b，前向为 0.92 倍到 2.34 倍、反向为 0.93 倍到 2.47 倍，该发布间比较同时包含 Triton、AITER 与调优数据库的更新。
启示与展望
当前支持面向 AMD MI350 系列 GPU 的 gfx950 架构，注意力内核另随 AOTriton 提供 RDNA4（gfx1201）版本；算子级数字来自 MI355X 上的稳态执行，不含编译与初次自动调优，也不代表端到端模型延迟。
注意力方面 0.14.2b 对 0.13b 的比较同时包含 Triton、AITER 与调优数据库更新，并非孤立的 FlyDSL 对 Triton 测量；MXFP8 与 MXFP4 的 vLLM 每个格子只有一次测量，小差异不宜当作带方差的结果；TopK 在确定性模式下等值索引顺序仍可能与 ATen 不同，NaN 的载荷与索引选择也可能不同。