PyTorch FlyDSL：算子支持，性能提升

1. 为现有算子路径增加选择

支持范围包括计算所有位置数值的普通稠密矩阵乘法、同时处理多个分组的分组矩阵乘法，以及低精度缩放矩阵乘法。还包括用于数值归一化的 RMSNorm、选取最大 K 个值的 TopK，以及计算输入之间关系的多头注意力。

在编译模式下，TorchInductor 会直接测量符合条件的矩阵乘法实现，并缓存最快的候选项。这相当于让现有实现和新实现按任务竞速后择优。在即时执行模式下，安装并启用可选运行时后，符合支持条件的归一化和选择算子会交给新内核，即包含待执行计算步骤的代码。

注意力通过 AOTriton 库，依据预先调优的数据选择实现。

FlyDSL 让内核编写者直接控制数据移动和计算布局，以支持针对不同硬件的优化。

2. 各算子的测量结果范围不同

算子层面的测量在 Instinct MI355X(gfx950) 上进行，不计编译和首次自动调优。低精度测量从已量化、即以更少比特表示数值的输入开始。下述几何平均值为各案例赋予相同权重后汇总所得。

根据厂商公告，BF16 稠密矩阵乘法在左矩阵按行优先、右矩阵按列优先的 NT 布局下，对 15 种形状的测量中，相较现有两种实现中较快者的几何平均加速比为 1.10 倍。两个对比对象为 ATen 和 Triton。

根据厂商公告，MXFP8 和 MXFP4 缩放矩阵乘法在 NT 布局下各测量了 17 种形状，相较现有默认实现分别为 1.58 倍和 1.68 倍。

根据厂商公告，分组矩阵乘法在 24 个案例中，相较 Triton 和 ATen 的几何平均加速比分别为 1.20 倍和 1.95 倍。

根据厂商公告，RMSNorm 在 22 个案例中相较 ATen 为 1.17~3.66 倍。维度对齐时大多为 1.2~1.5 倍，最高值出现在特定 FP16 输入上。

根据厂商公告，小 K 值的 TopK 在 10 个案例中，相较 ATen 的几何平均加速比为 4.82 倍。这是在关闭确定性算法、即相同输入产生相同结果的计算过程时测得的，开启时则为 4.02 倍。

3. 注意力与模型推理中的效果

整个模型的测量在单台 MI355X 上使用 vLLM 进行。固定输入 256 个 token、输出 512 个 token，将并发请求数从 8 调整到 256。对比双方的服务配置和各模型缓存相同，仅矩阵乘法候选列表不同。

根据厂商公告，在现有两个后端之外加入新候选后，BF16 推理中 Llama-3.1-8B-Instruct 最高为 1.13 倍，Qwen3-32B 最高为 1.06 倍。Llama-3.3-70B-Instruct 的变化大多处于测量噪声范围内。

厂商公布的低精度 8 比特推理结果如下。基准为使用 hipBLASLt 的 ATen，模型与并发请求数的 18 种组合均有提升。

单台 MI355X(TP=1)，MXFP8 A8W8，输入 256 个 token、输出 512 个 token，并发请求 8~256 个，相较 ATen(hipBLASLt) 添加 FlyDSL，每种组合测量 1 次

各模型的收益幅度随并发请求数而异。根据厂商公告，相同请求条件下的 MXFP4 推理在 18 种组合中也达到 1.25~3.01 倍。该测量未对最终输出层进行量化，每种组合仅运行一次。

注意力结果是相较 AOTriton 旧版本的吞吐量汇总。根据厂商公告，在 gfx950 上，0.14.2b 相较 0.13b 的前向为 0.92~2.34 倍，反向为 0.93~2.47 倍。

4. 使用条件与性能代价

可从近期 ROCm PyTorch nightly 和基于 gfx950 的 MI350 系列开始使用。用于编译模式矩阵乘法的可选运行时为 0.3.x 系列，以下是原文示例中的候选项配置部分。

import torch
import torch._inductor.config as config

config.max_autotune_gemm_backends = "ATEN,TRITON,FLYDSL"
config.flydsl_enable_autotuning = True

配置后，使用 torch.compile(mode="max-autotune") 编译矩阵乘法，即可将其纳入自动调优候选。注意力可通过链接了 AOTriton 0.14.2b 或更高版本的构建使用，无需额外软件包或配置。

重复运行速度提升的代价是首次调用需要编译和调优。若低精度输入不符合所需布局，还会增加转换复制。受支持算子占模型执行时间的比例越高，优化价值越大。

注意力版本对比还包含其他内核和调优数据的变化，目前新注意力内核不支持 FP32。归一化的反向计算使用现有实现，选取最大 K 个值时，相同值的索引顺序可能因后端而异。

5. 支持范围仍在扩展

当前集成面向 gfx950 的 MI350 系列。原文说明，后续正在推进 MI450 支持、低精度与注意力内核扩展、将后续计算合为一体的后处理融合，以及扩大反向计算支持等工作。

FlyDSL 为 AMD GPU 上的 PyTorch 增添了优化选项，收益在低精度矩阵乘法、部分选择算子及特定推理条件下尤为明显。