AMD 与 Meta 将 FlyDSL 接入 PyTorch，加速 MI350 系列 GPU 算子

AMD MI350 系列 GPU 上的 PyTorch 用户现在可以通过现有算子 API 使用 FlyDSL，支持稠密与分组 GEMM、融合多头注意力（FMHA/Attention）、MXFP8/MXFP4 缩放 GEMM、RMSNorm 和 TopK，无需应用调用自定义内核。该可选后端通过 torch.compile 支持 GEMM，而符合条件的 RMSNorm 和 TopK 调用在可用且启用时会在 eager 模式下自动分发。不支持的输入保留现有 PyTorch 实现。FlyDSL 注意力内核作为 AOTriton 中的额外后端发布，AOTriton 是 PyTorch ROCm F.scaled_dot_product_attention（SDPA）背后的库，从 0.14b 系列开始提供。

在所报告的各个内核级算子套件中，FlyDSL 在 15 种形状的 BF16 稠密 GEMM 套件上，相对于每个形状下 ATen 与 Triton 中更快者，提供 1.10 倍几何平均加速；在每种格式 17 种形状上，MXFP8 和 MXFP4 相对于 ATen 分别提供 1.58 倍和 1.68 倍几何平均加速；在全部 24 个分组 GEMM 用例上相对于 Triton 提供 1.20 倍几何平均加速；在 22 个 RMSNorm 用例上相对于 ATen 提供 1.17 倍至 3.66 倍加速。对于 10 个小 K 的 TopK 用例，在禁用确定性算法的情况下，几何平均为 ATen 的 4.82 倍。当受支持的操作占模型运行时间的显著比例时，这些收益最为相关。对于注意力，AOTriton 0.14.2b 公布的 gfx950 吞吐量在不同头维度和因果模式下，前向为 0.13b 的 0.92 倍至 2.34 倍，反向为 0.93 倍至 2.47 倍。在端到端 vLLM A/B 测试中，整个请求的加速比达到 BF16 的 1.13 倍、MXFP8 的 2.04 倍和 MXFP4 的 3.01 倍，结果取决于模型和并发度。

为什么将 FlyDSL 作为 PyTorch 后端

GPU DSL 选择不同的抽象层级。Triton 中的块张量编程模型让作者描述程序级瓦片，而编译器管理大部分线程、波和指令映射。FlyDSL 通过布局代数使这些映射显式化：内核作者可以在块、波、线程和值之间划分瓦片；选择平铺复制和 MFMA 原子；定义 LDS swizzle 和暂存；并放置同步或编译器调度边界。这两种方法互补。FlyDSL 适用于数据移动和指令归属属于算法一部分的架构特定模板。

这些控制直接出现在内核中：稠密/MXFP GEMM 中的 HTI 和分阶段 LDS，MXFP 中匹配的操作数缩放生命周期和缩放 MFMA，以及分组 GEMM 中的持久专家调度。FlyDSL 在 Python 中表达它们，并通过 MLIR 降低，同时保留显式硬件映射。

注意力展示了这种控制为何重要。FlyDSL 注意力内核使用双波软件流水线，将一个 KV 瓦片的 softmax VALU 工作与另一个瓦片的 MFMA 工作重叠。在最大头维度下，它们还拆分头维度：通过 LDS 在时间上拆分，并跨波拆分输出累加器。在 FlyDSL 中，两种调度都写在内核本身中，因此可以针对每种架构开发和调优，而无需更改编译器。

PyTorch 将这些专用模板转变为可叠加、可度量的后端。不支持的工作负载保留现有实现：

• Eager 执行：当可选运行时安装并启用时，符合条件的 RMSNorm 和 TopK 调用会分发到 FlyDSL；其他调用保留 ATen。
• torch.compile：TorchInductor 筛选符合条件的 Dense、Grouped 和 MXFP GEMM 候选，在支持的情况下将 FlyDSL 与 ATen 和 Triton 一起进行基准测试，为每个工作负载缓存胜出者，并运行实测最快的内核。
• 多头注意力：PyTorch 内置算子 F.scaled_dot_product_attention 调用 AOTriton，后者现在利用其离线算子调优数据库，为每个工作负载在 FlyDSL、Triton 和 AITER ASM 内核中进行选择。

图 1. 通过 PyTorch API 访问的 FlyDSL 内核。Eager RMSNorm 和 TopK 根据输入支持情况进行调度；AOTriton 选择经过调优的 SDPA 实现；TorchInductor 对符合条件的 GEMM 候选进行基准测试。

支持的功能

当前支持面向采用 gfx950 架构的 AMD MI350 系列 GPU。以下功能已在 PyTorch 中提供：

Dense GEMM 的四种布局支持接受行主序和列主序输入，包括符合条件的转置视图。在布局名称中，第一个字母描述 A，第二个字母描述 B；N 和 T 分别表示非转置和转置输入布局。Grouped GEMM 支持每组不同的行数，允许专家处理不均匀的 token 分配。其当前的输入布局要求与 Dense GEMM 的四种布局覆盖范围相互独立。RMSNorm 反向继续沿用 PyTorch 的现有实现。

MXFP 缩放 GEMM 接受 FP8 E4M3 或打包的 FP4 E2M1 操作数，沿归约维度每 32 个值配一个 E8M0 缩放（BlockWise1x32）。应用需提供量化操作数和连续、未 swizzle 的缩放张量。两种格式均使用 FP32 累加并返回 FP16 或 BF16；当前路径要求 K 为 128 的倍数。它支持可选的一维偏置，其长度为 N 且 dtype 与输出匹配；不支持快速累加。在 torch.compile 下，Inductor 的布局约束在后端选择之前运行，要求 A 和 B 沿 K 方向为单位步长——A 行主序、B 列主序——同时保留兼容的前导步长。对于不兼容的输入，Inductor 会插入布局拷贝。FlyDSL MXFP 内核及下方的测量均使用这些受约束的布局。此约束特定于缩放 GEMM：Dense torch.mm 检测符合条件的 NN/NT/TN/TT 步长，并将其布局标志传递给 FlyDSL，而不应用 MXFP 规范化。

FlyDSL 注意力内核达到了 F.scaled_dot_product_attention 的 flash 和内存高效注意力能力，但目前有一个限制：不支持 FP32。FP32 输入以及调优数据库分配到其他位置的工作负载，仍使用 AOTriton 的 Triton 或 AITER ASM 内核。AOTriton 还为 RDNA4 提供了 FlyDSL 注意力内核（gfx1201）。在 torch.compile 下，Inductor 将 SDPA 的 flash 和内存高效注意力算子保留为对 ATen 的调用，而不是为其生成代码，因此它们仍会到达 AOTriton 内核。Inductor 重写为 SDPA 的手写 matmul–softmax–matmul 模式也走同一路径。相比之下，torch.nn.attention.flex_attention 是一个单独的 API，Inductor 会将其编译为自己的 Triton 内核；它不使用 AOTriton。

每个操作都有形状和对齐要求。详细的 dense GEMM、MXFP scaled GEMM、grouped GEMM、RMSNorm 和 TopK 支持说明定义了这些边界。不支持的 eager 调用保留 ATen 行为；编译后的 GEMM 可以使用其他已启用的后端。AOTriton 0.14.2b 更新将这些内核引入 PyTorch。

内核级性能

以下算子级基准测试在 AMD Instinct MI355X（gfx950）GPU 上测量。各测试套件使用不同的基准测试设置：dense GEMM 使用图重放的 GPU 事件计时，grouped GEMM 报告稳态吞吐量，RMSNorm 和 TopK 使用重复算子调用的 GPU 事件计时。结果聚焦于稳态算子执行，不包括编译和初始自动调优；它们不代表端到端模型延迟。MXFP 测量从已量化的操作数开始，因此输入量化不在计时范围内。加速比高于 1.0 表示 FlyDSL 比指定的基线更快，几何平均值对每个采样用例赋予相同权重。每个图和标题都指定了其基线和聚合范围。

Dense GEMM：线性层工作负载

稠密矩阵乘法是 transformer 投影层和前馈层的构建模块。其维度随处理的 token 数量大幅变化，因此跨多种形状的性能对模型开发者很重要。

FlyDSL 提供全瓦片和半瓦片交错（HTI）执行路径。全瓦片路径在本地数据共享（LDS）中使用可配置的 STAGES 深度 K 瓦片环。其序言部分为环做初始填充，稳态循环将已暂存数据的计算与异步预取到可复用槽位重叠进行，最后几次迭代排空剩余阶段。关于这种生产者/消费者模式及其等待语义的分步说明，请参阅 AMD 的 Multi-Stage LDS Pipeline: Keep K Blocks in Flight。

对于大型 GEMM，调优后的 HTI 配置使用 256 × 256 输出瓦片和八波工作组（512 线程，Wave64）。HTI 沿 M 拆分 A、沿 N 拆分 B，在寄存器中维护四个 128 × 128 输出象限累加器。在每个 K 瓦片内，A0 馈送 C00/C01，A1 馈送 C10/C11，B0 馈送 C00/C10，B1 馈送 C01/C11。较小的形状可以在更窄的全瓦片和 HTI 配置之间自动调优。

HTI 使用两个 LDS 阶段成对处理连续的 K 块。在稳态下，预取会逐步复用存放 t 和 t+1 的阶段来存放 t+2 和 t+3。复用按半块进行：一旦某个操作数半块旧 LDS 内容所需的读取完成且满足同步要求，其区域就可以重新填充，同时矩阵融合乘加（MFMA）计算继续在寄存器驻留的片段上进行。在同一计算单元（CU）内，相位偏移的屏障调度将工作组的两个四波组错开，使一个组中的 MFMA 计算与另一组中的操作数加载和异步预取重叠。两个组在这些活动之间交替，将半块缓冲区复用与错开执行相结合，而不是等待整个块的计算完成。

图 2. HTI 波作用域与相位错开执行，以两组各四个 Wave64 波展示。每个网格单元是一个 16 × 16 输出块。每组中的波数可通过块配置调整，受内核约束限制。

在 15 个 BF16 NT 形状上，FlyDSL 在每个形状上相比更快的 ATen/Triton 基线实现了 1.10 倍的几何平均加速。增益在较小和中等规模问题上最为显著：例如 M × N × K = 64 × 4096 × 4096 提升了 1.33 倍。

图 3. MI355 上的 BF16 NT GEMM，相对于每个形状上更快的 ATen/Triton 基线。展示了全部 15 个案例：使用 ±1% 平局带，10 胜、3 平、2 负。

随后对当前全布局内核在相同形状上的回归检查报告称，相比更快的原始基线实现了 1.12 倍的总体加速。图 3 保留了较早的逐形状测量结果，因为后续检查仅发布了总体结果。

本套件中两个最大的方形输出问题略微有利于 ATen。保持启用多个后端可让 PyTorch 为每个工作负载做出该选择。这些结果表征的是 NT GEMM；NN、TN 和 TT 也受支持，其性能取决于工作负载。

MXFP 缩放 GEMM：扩展到低精度工作负载

MXFP 通过 FP8 操作数处理、打包 FP4 存储、E8M0 缩放暂存和 CDNA4 缩放 MFMA，扩展了共享的 gfx950 全块和 HTI GEMM 实现。

在 HTI 中，A/B 缓冲区按 K 块推进，而单独的缩放槽位存放跨多个 K 块的数据块。每个 MFMA 使用寄存器驻留的操作数及其 K 位置对应的缩放。在数据块边界处，等待和屏障协调另一个缩放槽位的复用，以便在当前数据块被消费时预取下一个数据块。数据块长度取决于块和工作组几何结构；下面的四块数据块是一个示例，不是固定的格式要求。全块路径则使用与 A/B 相同的阶段索引，将每个 K 块的缩放暂存在单独的 LDS 缓冲区中。

图 4. MXFP HTI 缓冲，以每个缩放数据块四个 K 块展示。A/B 读取槽位按块交替；缩放读取槽位按数据块交替。在每个数据块边界处，下一个缩放数据块被预取到另一个槽位中。

在 MI355X 上，17 形状 NT 套件显示 MXFP8 相比 ATen 实现了 1.58 倍的几何平均加速，MXFP4 相比 ATen 实现了 1.68 倍的加速。

图 5. MI355X 上的 NT 缩放 GEMM。每个面板展示了全部 17 个案例及其相对于 ATen 的几何平均值。

在所有报告的案例中，FlyDSL 在两种格式下的吞吐量均超过 ATen。MXFP8 的加速比范围为 1.31 倍至 2.23 倍，最大增益出现在 32 × 4096 × 4096。MXFP4 的范围为 1.31 倍至 2.96 倍，峰值出现在 32 × 14336 × 4096。MXFP 自动调优路径在两个后端均启用且适用时，会将 FlyDSL 与 ATen 一起进行基准测试。测量覆盖 NT 布局。

分组 GEMM：专家间工作负载不均

在混合专家模型中，不同专家可能接收到数量差异极大的 token。分组 GEMM 将这些矩阵乘法一起处理，同时适应不均匀——有时甚至为空——的分组。PyTorch 将输入表示为拼接的激活值 A[sum(M_g), K]、分组权重 B[G, K, N]，以及标记每个专家行边界的偏移量。

持久化网格不为每个专家启动一个内核，而是为每个工作组分配一个跨累积分组偏移量的跨步 tile 序列，从而自然地跳过空分组。计算裸片感知的 swizzle 将工作分布到 GPU 的八个加速器复合裸片（XCD）上，并带有 N 主序回退方案，使并发工作组保持在可复用的 B tile 上；HTI 配置仍可供自动调优使用。

图 6. 分组 GEMM 调度。图中显示一个 N-tile 列：各专家向累积流贡献 ceil(M_g / BM) 个 tile，而零行专家不贡献任何 tile。持久化工作组从 blockIdx.x 开始并按网格大小推进，因此它们可以跨越专家边界而无需为每个专家启动。

报告的 24 个案例按所测试的工作负载维度划分。在统一形状 G × M × K × N 中，G 是专家数量，M 是每个专家的 token 数，K 和 N 是归约和输出维度：

• 统一分组（14 个案例）：所有专家使用相同的 M，而 G、M、K 和 N 在常见的分组 GEMM 形状中变化。
• 投影维度（5 个案例）：G = 8 和 M = 512 固定，而 K 和 N 覆盖 MoE 上/下投影以及一个小 N 案例。
• 不规则专家负载（5 个案例）：G = 8 和 K = N = 4096 固定，而每个标签列出八个逐专家 M 值，包括空专家和高度不均衡的专家。

在所有 24 个案例中，FlyDSL 相对 Triton 达到 1.20 倍几何平均加速比，相对 ATen 达到 1.95 倍。统一分组分别达到 1.21 倍和 2.12 倍，FlyDSL 在 14 个案例中的 11 个中是最快的后端。投影维度案例达到 1.23 倍和 1.73 倍；不规则案例达到 1.15 倍和 1.73 倍，FlyDSL 在全部五个案例中获胜。

图 7. gfx950 上全部 24 个 BF16 分组 GEMM 案例，分为三个面板。每个条形将 FlyDSL 与该案例中更快的 ATen/Triton 结果进行比较；每个面板以其几何平均值结束。

逐案例视图使形状敏感性变得明确：投影扫描包含一个 Triton 更快的小 N 案例和一个与 ATen 打平的宽 N 案例，而全部五个不规则专家负载案例均有利于 FlyDSL。

RMSNorm：跨隐藏维度的增益

RMSNorm 对激活值进行归一化和缩放，这是许多 transformer 模型中重复的步骤。FlyDSL 加速前向操作，同时保留正常的 PyTorch API 和反向行为。

内核为每一行分配一个线程块（CTA），并从归一化维度中选择 256、512 或 1024 个线程，分别对应 4、8 或 16 个 Wave64。它用 128 位拷贝加载可向量化的主体部分，用标量尾部处理剩余元素，并将已加载的输入保留在寄存器中。FP32 平方和首先在每个 Wave64 内通过 shuffle 操作进行归约，然后通过一个小的 LDS 缓冲区将每个 wave 的一个部分和合并起来。驻留值被复用来应用 rsqrt 和权重，并返回反向传播所需的 rstd。

图 8. RMSNorm 调度。一个线程块将一行加载到寄存器中，根据 N 产生 4/8/16 个 Wave64 部分和，通过 LDS 将它们合并，然后复用驻留值来写出归一化输出和反向 rstd。

对于所测量的对齐隐藏维度，相对于 ATen 的加速约为 1.2 倍至 1.5 倍。比对齐尺寸多一个元素的维度——例如 4097 而非 4096——显示出更大的增益，达到 3.66 倍。

图 9. MI355X 上报告的全部 22 个 RMSNorm 用例。标签标识 dtype 和 M × N，其中 M 是行数，N 是归一化维度。两个面板使用相同的加速比刻度。

对齐维度与差一维度之间的划分说明了在解读峰值结果时输入形状为何重要。3.66 倍的加速适用于某个特定的 FP16 用例；对齐维度上的增益则较为温和。

TopK：面向小 K 和大 K 的更快选择

TopK 从每一行中选出得分最高的元素。请求的输出数量会改变工作负载：选择少量元素与选择数百个元素是不同的。FlyDSL 为这些情形提供了专门的路径。对于小的固定 K，一个 gfx950 CTA 包含两个 Wave64 并处理两行——每行一个 wave。每个 lane 读取 128 位数据块，一次对一个组进行双调排序，并将每个组折叠进一个驻留寄存器的局部 top-K。随后蝶形 shuffle 跨 wave 合并这些候选，lane 0 写出最终的 K 个对。对于更大的 K，四趟 8 位基数排序找到第 K 个阈值，仅收集存活的候选，并对一个从 K 向上取整到下一个 2 的幂的缓冲区进行双调排序，而不是对整行排序。

这里的“确定性开启”是指 torch.use_deterministic_algorithms(True)。基数排序路径随后使用前缀和槽位来保持有限值的并列顺序；在确定性关闭时，它使用原子槽位分配，因此相等值的索引可能不同。寄存器内核本身在两种模式下都是可复现的，尽管其并列排序可能与 ATen 不同。NaN 被规范化为一个序数，因此即使在确定性模式下，确切的 NaN 载荷和索引选择也可能与 ATen 不同。

图 10. TopK 调度。寄存器路径每个 CTA 处理两行，每行一个 Wave64，然后蝶形合并局部 top-K 候选。基数排序路径每行使用一个 CTA 和四趟字节处理来确定阈值，然后仅对存活的候选进行排序。

对于小的 K = {2, 4, 8, 16}，FlyDSL 在禁用确定性算法时相对于 ATen 实现了 4.82 倍的几何平均加速，在启用时实现了 4.02 倍。对于较大的采样 K 区间，几何平均加速范围为 1.40 倍至 1.97 倍。

图 11. MI355X 上的 FP32 TopK，在相同确定性设置下与 ATen 比较。柱状图按内核系列和 K 区间汇总了全部 33 个报告的用例。

小 K 的寄存器路径显示出最大的平均增益。基数选择路径将收益扩展到更大的 K 范围，尽管个别情况接近持平。对于需要检查索引的应用，不同后端对相等值的并列项排序可能不同；可复现性并不保证并列索引完全一致。

注意：通过 AOTriton 的 FMHA 内核

与上述算子不同，SDPA 并不通过 PyTorch 侧后端到达 FlyDSL。AOTriton 将 FlyDSL 内核编译进自己的库，然后与它的 Triton 和 AITER ASM 内核一起调度。前向是一个内核；反向使用独立的 dQ 和 dK/dV 内核。

对于最高 256 的头维度，前向内核使用双波软件流水线。每个工作组有两个 KV 分块在途，一个分块的 softmax 在 VALU 上运行，而另一个分块的 QK 和 PV 乘积在 MFMA 单元上运行。在头维度 512 时，两个常驻 K/V 分块将需要约 272 KB 的 LDS，而限制为 160 KB，因此单独的宽头主体一次处理一个分块。它两次拆分头维度：在时间上，通过流水线化的 D 切片将 K 和 V 暂存到相同的 LDS 缓冲区；在波之间，每个波累积输出的一部分，使累加器能放入寄存器。

反向复用前向构建块，而不是推导新的 lane 映射。dK/dV 是前向循环的转置：K 和 V 留在寄存器中，而 Q 和 dO 流经 LDS，转置后的操作数使用前向的 LDS 转置读取（ds_read_b64_tr_b16）。在 dQ 中，三个 GEMM 中的每一个都匹配前向已经发出的一个。K 分块暂存一次并以两种方向读取，这使头维度 512 保持在 LDS 限制内。

图 12. 头维度最高 256 的双波前向流水线。在一个波内，每个矩阵核心步骤与相邻 KV 分块的 softmax 重叠；在整个工作组中，两个波组相差一个阶段运行，因此一组加载 K 和 V，而另一组进行计算。

图 13. 头维度 512 的宽头前向主体。K 和 V 通过两个 LDS 缓冲区以 256 列的一半暂存，两个波计算相同的分数，而每个波累积输出的一半。

AOTriton 0.14.2b 版本按头维度、方向和因果模式发布了 AOTriton 0.13b 和 0.14.2b 的吞吐量。对于下面显示的 gfx950 聚合数据，前向加速范围从 0.92x 到 2.34x，反向加速从 0.93x 到 2.47x。

图 14. AOTriton 0.14.2b 与 0.13b 在 gfx950 上的对比，使用每个头维度已发布的吞吐量聚合数据。此版本间比较还包括更新的 Triton、AITER 和调优数据库，因此它不是孤立的 FlyDSL 与 Triton 测量。

端到端 vLLM 推理

我们在一个 MI355X（TP=1）上使用 vllm bench serve，搭配 Llama-3.1-8B-Instruct、Qwen3-32B 和 Llama-3.3-70B-Instruct。每个请求有 256 个 token 的输入和固定的 512 个 token 输出；并发范围从 8 到 256。基线和处理使用相同的服务配置和每模型 KV 缓存，仅更改了 Inductor GEMM 候选列表。图中报告的是整个请求的加速比。

BF16。将 FlyDSL 添加到 ATen/Triton 基线后，Llama-3.1-8B 在低到中等并发下产生最明显的增益，峰值达到 1.13x；Qwen3-32B 峰值为 1.06x，而 Llama-3.3-70B 基本保持在 1.79% 的测量噪声底限内。BF16 结果报告的是至少两次 ABBA 交错重复的中位数。

图 15. 在单块 MI355X 上的 BF16 整请求加速比。基线启用 ATen 和 Triton；处理组额外加入 FlyDSL。高于 1x 的数值表示 FlyDSL 更优。

MXFP8 A8W8。以 ATen（hipBLASLt）为基线，FlyDSL 在全部 18 个模型-并发组合上均有提升：Qwen3-32B 为 1.40x–2.03x，Llama-3.1-8B 为 1.04x–1.92x，Llama-3.3-70B 为 1.38x–2.04x。

图 16. 在单块 MI355X 上的 MXFP8 A8W8 整请求加速比。基线为 ATen；处理组额外加入 FlyDSL。高于 1x 的数值表示 FlyDSL 更优。

MXFP4 A4W4。在 lm_head 保持未量化的条件下，全部 18 个组合同样均有提升：Qwen3-32B 为 1.53x–3.01x，Llama-3.1-8B 为 1.25x–1.85x，Llama-3.3-70B 为 1.50x–2.82x。MXFP8 和 MXFP4 每个单元格仅有一次测量运行，因此不应将微小差异视为具备方差统计意义的结果。

图 17. 在单块 MI355X 上的 MXFP4 A4W4 整请求加速比。基线为 ATen；处理组额外加入 FlyDSL。高于 1x 的数值表示 FlyDSL 更优。

在 PyTorch 中试用

对于 MI350 系列 GPU，请使用较新的 ROCm PyTorch nightly。本文假定该版本已包含上述所有 FlyDSL 集成，包括 MXFP 缩放 GEMM。然后从受支持的 0.3.x 系列安装可选运行时：

python -m pip install --upgrade "flydsl>=0.3.0,<0.4"

以下示例需要 gfx950 GPU；ROCm 构建使用 PyTorch 的 "cuda" 设备字符串。

稠密与分组 GEMM

将 FlyDSL 启用为候选实现，并让 torch.compile 选择具体实现：

import torch
import torch._inductor.config as config

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

a = torch.randn(128, 4096, device="cuda", dtype=torch.bfloat16)
b = torch.randn(4096, 14336, device="cuda", dtype=torch.bfloat16)

@torch.compile(mode="max-autotune")
def mm(a, b):
    return a @ b

out = mm(a, b)

此示例使用 NN 布局。符合条件的转置输入无需额外标志即可启用 NT、TN 和 TT。相同的设置也适用于 F.grouped_mm。

首次调用包含编译和调优；后续调用会复用已选定的实现。启用 FlyDSL 会为 PyTorch 增加一个候选实现，因此仍可能选中 ATen 或 Triton。当工作负载与 GPU 图捕获不兼容时，请使用 max-autotune-no-cudagraphs。这些 Inductor 配置控制项对版本敏感，可能在不同 nightly 构建之间发生变化。

MXFP8 与 MXFP4 缩放 GEMM

使用上述 GEMM 配置，在已量化的操作数及其块缩放上编译一次调用。此封装使用该集成所支持的公开 F.scaled_mm API：

import torch.nn.functional as F
from torch.nn.functional import ScalingType, SwizzleType

@torch.compile(mode="max-autotune")
def mxfp_mm(a, b, scale_a, scale_b):
    return F.scaled_mm(
        a,
        b,
        scale_a,
        ScalingType.BlockWise1x32,
        scale_b,
        ScalingType.BlockWise1x32,
        SwizzleType.NO_SWIZZLE,
        SwizzleType.NO_SWIZZLE,
        output_dtype=torch.bfloat16,
    )

对于 MXFP8 NT 调用，传入行主序的 a[M, K] 和列主序的 b[K, N]，dtype 为 torch.float8_e4m3fn。两个缩放均为 dtype torch.float8_e8m0fnu，连续形状分别为 [M, K // 32] 和 [N, K // 32]。对于 MXFP4，使用打包的 torch.float4_e2m1fn_x2 操作数，存储形状为 [M, K // 2] 和 [K // 2, N]；缩放形状使用逻辑形状 K。该封装以相同的后端设置支持两种格式。

Eager RMSNorm 与 TopK

符合条件的 eager 调用会自动使用 FlyDSL。后端控制器便于与 ATen 进行对比：

import torch
import torch.nn.functional as F
import torch.backends.python_native as pn

assert torch.version.hip is not None
assert torch.cuda.get_device_properties(0).gcnArchName.split(":")[0] == "gfx950"
assert pn.flydsl.available

x = torch.randn(8192, 4096, device="cuda", dtype=torch.bfloat16)
weight = torch.ones(4096, device="cuda", dtype=torch.bfloat16)

y = F.rms_norm(x, (4096,), weight, eps=1e-5)
with pn.flydsl.disabled():
    y_aten = F.rms_norm(x, (4096,), weight, eps=1e-5)

torch.testing.assert_close(y, y_aten)

相同的 pn.flydsl.disabled() 上下文管理器适用于符合条件的 torch.topk 调用。实际分派取决于操作的形状、dtype 和布局。

注意力

SDPA 无需 FlyDSL 包或配置。它只需要链接到 AOTriton 0.14.2b 或更高版本的 PyTorch 构建，该版本已捆绑预编译内核：

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

q, k, v = (torch.randn(4, 16, 4096, 128, device="cuda", dtype=torch.bfloat16,
                       requires_grad=True) for _ in range(3))

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
out.sum().backward()

AOTriton 会为每个工作负载选择 FlyDSL 或其他内核。

后续计划

后续步骤将扩展可从 FlyDSL 受益的工作负载范围：

• 更广泛的硬件覆盖：将 PyTorch 集成和优化内核集扩展到 AMD Instinct MI450 系列 GPU。
• 更多低精度工作负载与注意力：为更多推理工作负载添加 MXFP8 分组 GEMM 和 FlexAttention 内核。
• FlyDSL 自身的注意力：将 AOTriton 注意力内核移入 FlyDSL 的 kernels/attention/，使 AOTriton 从上游构建它们，而不是使用 vendored 副本。同时添加感知序列长度的注意力调度。
• FlyDSL 的 FlexAttention：在 torch.compile 下为 torch.nn.attention.flex_attention 添加 FlyDSL 后端，与 Inductor 的 Triton 内核并存。
• 融合与训练：支持 GEMM epilogue 融合并扩大反向传播覆盖范围。
• 部署与验证：增加 AOTInductor 支持并扩展端到端模型基准测试。

这些扩展仍在进行中，不属于本文所呈现的当前支持范围和性能结果。

结论

当前的集成将 FlyDSL 内核开发同时连接到 eager 算子和编译后的 GEMM，为 PyTorch 用户提供了一种从优化后的 AMD 内核中获益的实用方式。通过 AOTriton，它还能在用户侧无需任何改动的情况下覆盖 F.scaled_dot_product_attention。vLLM 的结果表明，这些内核选择如何转化为模型层面的收益，尤其是对于低精度 MXFP8 和 MXFP4 工作负载。请在你的工作负载上试用，并通过 PyTorch issue tracker 分享结果或功能请求。

文章 Accelerating PyTorch on AMD GPUs with FlyDSL 首次发布于 PyTorch。