PyTorch adds FlyDSL backend for faster operators on AMD MI350 GPUs

What is supported and how it is reached

FlyDSL covers dense and grouped GEMM, fused multi-head attention, MXFP8/MXFP4 scaled GEMM, RMSNorm and TopK. Current support targets AMD MI350-series GPUs with the gfx950 architecture, according to PyTorch.

The way each operation reaches FlyDSL depends on the execution mode:

• Eager mode: eligible RMSNorm and TopK calls dispatch to FlyDSL when the optional runtime is installed and enabled. Other calls stay on ATen.
• torch.compile: TorchInductor filters eligible dense, grouped and MXFP GEMM candidates, benchmarks FlyDSL beside ATen and Triton where supported, caches the winner for each workload and runs the fastest measured kernel.
• Attention: F.scaled_dot_product_attention calls AOTriton, which now picks among FlyDSL, Triton and AITER ASM kernels for each workload using its offline operator-tuning database.

FlyDSL attention kernels ship as an additional AOTriton backend starting with the 0.14b series. FP32 inputs are not supported by the FlyDSL attention kernels and keep AOTriton's Triton or AITER ASM kernels. AOTriton also ships FlyDSL attention kernels for RDNA4 (gfx1201).

Why FlyDSL differs from Triton

PyTorch describes the two as complementary. Triton's block-tensor model lets authors describe program-level tiles while the compiler manages much of the thread, wave and instruction mapping. FlyDSL makes those mappings explicit through layout algebra. Kernel authors can partition tiles across blocks, waves, threads and values, select tiled-copy and MFMA atoms, define LDS swizzles and staging, and place synchronization or compiler-scheduling boundaries.

FlyDSL expresses these controls in Python and lowers them through MLIR while preserving explicit hardware mapping. The post says this suits architecture-specific templates where data movement and instruction ownership are part of the algorithm.

Attention is the example given. The FlyDSL attention kernels use a dual-wave software pipeline that overlaps the softmax VALU work of one KV tile with the MFMA work of another. At the largest head dimensions they also split the head dimension, in time through LDS and across waves for the output accumulator. Because both schedules are written in the kernel itself, they can be tuned per architecture without compiler changes, according to PyTorch.

Reported performance

The kernel-level benchmarks were measured on AMD Instinct MI355X (gfx950) GPUs. PyTorch says they focus on steady-state operator execution, exclude compilation and initial autotuning, and do not represent end-to-end model latency. MXFP measurements start from already quantized operands.

• Dense GEMM: 1.10x geometric-mean speedup for the 15-shape BF16 suite over the faster of ATen and Triton at each shape.
• MXFP8 and MXFP4: 1.58x and 1.68x geometric-mean speedups over ATen across 17 shapes per format.
• Grouped GEMM: 1.20x geometric mean over Triton across all 24 cases.
• RMSNorm: 1.17x to 3.66x over ATen across 22 cases.
• TopK: 4.82x geometric mean over ATen for 10 small-K cases, with deterministic algorithms disabled.

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, depending on model and concurrency.

Limits and what it means for users

Each operation has shape and alignment requirements. MXFP scaled GEMM accepts FP8 E4M3 or packed FP4 E2M1 operands with one E8M0 scale per 32 values along the reduction dimension, and the current path requires K to be a multiple of 128. RMSNorm backward continues through PyTorch's existing implementation. Under torch.compile, flex_attention is a separate API that Inductor compiles into its own Triton kernels, so it does not use AOTriton.

For teams running PyTorch on MI350-series hardware, the change is additive: supported calls can pick up FlyDSL kernels without code changes, and unsupported ones fall back. PyTorch notes that the gains matter most when the supported operations make up a significant share of model runtime. The AOTriton 0.14.2b update brings the attention kernels to PyTorch.