PyTorch on AMD GPUs gains a FlyDSL backend: kernels map work to hardware explicitly, and vLLM requests reach up to 3.01x
Synopsis
PyTorch users on AMD MI350-series GPUs can now reach FlyDSL kernels through existing operator APIs for dense and grouped GEMM, MXFP8/MXFP4 scaled GEMM, RMSNorm, TopK, and attention via AOTriton. Those operators previously ran on ATen, Triton, or AITER ASM implementations. On MI355X operator suites, FlyDSL gives a 1.10x geometric-mean speedup over the faster ATen/Triton baseline at each of 15 BF16 dense-GEMM shapes, and 1.58x and 1.68x over ATen for MXFP8 and MXFP4. End-to-end vLLM whole-request speedups reach 1.13x for BF16, 2.04x for MXFP8, and 3.01x for MXFP4, depending on model and concurrency.
Interpretation
PyTorch applications on AMD MI350-series GPUs can use FlyDSL kernels through existing interfaces such as torch.mm, F.scaled_mm, F.grouped_mm, F.rms_norm, torch.topk, and F.scaled_dot_product_attention, without calling custom kernels. Those operators previously ran on ATen, Triton, or AITER ASM implementations, and FlyDSL's hardware mapping lived only in standalone kernels. On MI355X (gfx950), 15 BF16 NT dense-GEMM shapes give a 1.10x geometric-mean speedup over the faster ATen/Triton baseline at each shape, with 10 wins, three ties, and two losses; the losses are the two largest square-output problems in the suite.
Low-precision paths extend the gain to MXFP8 and MXFP4 scaled GEMM and to grouped GEMM that tolerates uneven and empty expert groups. Scaled GEMM must handle FP8 or packed FP4 operands, one E8M0 scale per 32 values along the reduction dimension, and CDNA4 scaled MFMA. On the 17-shape NT suite, MXFP8 reaches 1.58x and MXFP4 1.68x geometric mean over ATen; across 24 BF16 grouped-GEMM cases FlyDSL reaches 1.20x over Triton and 1.95x over ATen.
RMSNorm and TopK dispatch automatically in eager mode by input support, and unsupported inputs keep existing PyTorch implementations. These elementwise-style operators previously had no optional FlyDSL path and ran on ATen. Across 22 RMSNorm cases FlyDSL gives 1.17x to 3.66x over ATen, with 3.66x on an FP16 dimension one element above an aligned size; the 10 small-K TopK cases give a 4.82x geometric mean over ATen with deterministic algorithms disabled and 4.02x with them enabled.
Attention kernels enter F.scaled_dot_product_attention as an additional AOTriton backend, with no user-side change. SDPA previously ran on AOTriton's Triton and AITER ASM kernels, and FlyDSL does not reach it through a PyTorch-side backend. AOTriton 0.14.2b's published gfx950 throughput relative to 0.13b ranges from 0.92x to 2.34x for forward and 0.93x to 2.47x for backward, a release-to-release comparison that also includes updated Triton, AITER, and tuning databases.
Perspective
Current support targets the gfx950 architecture of AMD MI350-series GPUs, with attention kernels also shipped for RDNA4 (gfx1201) through AOTriton; operator-level numbers come from steady-state execution on MI355X, exclude compilation and initial autotuning, and do not represent end-to-end model latency.
For attention, the 0.14.2b versus 0.13b comparison also includes updated Triton, AITER, and tuning databases, so it is not an isolated FlyDSL-versus-Triton measurement; each MXFP8 and MXFP4 vLLM cell has one measured run, so small differences should not be treated as variance-qualified; TopK tie ordering can still differ from ATen in deterministic mode, and NaN payload and index selection can differ as well.