PyTorch 为 AMD MI350 GPU 接入 FlyDSL 后端
PyTorch Blog 报道，AMD Instinct MI350 系列 GPU 用户现在可通过现有 PyTorch 算子 API 使用 FlyDSL 内核，覆盖稠密与分组 GEMM、MXFP8 和 MXFP4 缩放 GEMM、注意力、RMSNorm 和 TopK。GEMM 可通过 torch.compile 使用该可选后端，符合条件的 RMSNorm 与 TopK 调用可在 eager 模式下自动调度，不支持的输入则继续使用现有 PyTorch 实现。FlyDSL 注意力内核由 AOTriton 0.14.2b 及以上版本提供。基于 MI355X 的官方测试显示，BF16 稠密 GEMM 测试集的几何平均加速为 1.10 倍；端到端 vLLM A/B 测试中，BF16、MXFP8 和 MXFP4 的整请求加速分别最高达到 1.13 倍、2.04 倍和 3.01 倍。结果取决于工作负载、模型和并发度，供应的报告未对其进行独立验证。
此次集成让 PyTorch 用户能够通过熟悉的 API 和调度路径使用面向 AMD 硬件优化的 FlyDSL 内核。报告中的模型级收益在低精度 MXFP8 和 MXFP4 vLLM 测试中最为突出，但实际效果取决于具体工作负载。
PyTorch users on AMD MI350-series GPUs can now use FlyDSL through existing operator APIs for dense and grouped GEMM, fused multi-head attention (FMHA/Attention), MXFP8/MXFP4 scaled GEMM, RMSNorm, and TopK, without requiring applications to call custom kernels. The optional backend supports GEMMs through torch.compile, while eligible RMSNorm and TopK calls dispatch automatically in eager mode when it is available and enabled. Unsupported inputs retain existing PyTorch implementations. FlyDSL attention kernels ship as an additional backend in AOTriton, the library behind PyTorch’s ROCm F.scaled_dot_product_attention (SDPA), starting with the 0.14b series.
Across the reported kernel-level operator suites, FlyDSL provides a 1.10x geometric-mean speedup for the 15-shape BF16 dense-GEMM suite over the faster of ATen and Triton at each shape, 1.58x and 1.68x geometric-mean speedups for MXFP8 and MXFP4, respectively, over ATen across 17 shapes per format, 1.20x geometric mean over Triton across all 24 grouped-GEMM cases, and 1.17x–3.66x over ATen across the 22 RMSNorm cases. For the 10 small-K TopK cases, the geometric mean is 4.82x over ATen with deterministic algorithms disabled. These gains are most relevant when the supported operations account for a significant share of model runtime. For attention, AOTriton 0.14.2b’s published gfx950 throughput ranges from 0.92x to 2.34x of 0.13b for forward and 0.93x to 2.47x for backward, across head dimensions and causal modes. In end-to-end vLLM A/B tests, whole-request speedup reaches 1.13x for BF16, 2.04x for MXFP8, and 3.01x for MXFP4, with the result depending on model and concurrency.
#PyTorch#AMD#FlyDSL#AMD Instinct MI350#AOTriton#ROCm#ATen#Triton#TorchInductor#vLLM