Accelerating PyTorch on AMD GPUs with FlyDSL

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. Why FlyDSL as a PyTorch Backend GPU DSLs choose different abstraction levels. The block-tensor programming model in Triton 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. The two approaches are complementary. FlyDSL is useful for architecture-specific templates where data movement and instruction ownership are part of the algorithm. These controls appear directly in the kernels: HTI and staged LDS in Dense/MXFP GEMM, matched operand-scale lifetimes and scaled MFMA in MXFP, and persistent expert scheduling in Grouped GEMM. FlyDSL expresses them in Python and lowers them through MLIR while preserving explicit hardware mapping. Attention shows why this control matters. The FlyDSL attention kernels use a dual-wave software pipeline, which 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. In FlyDSL, both schedules are written in the kernel itself, so they can be developed and tuned per architecture without changes to the compiler. PyTorch turns those specialized templates into an additive, measurable backend. Unsupported workloads retain existing implementations: Eager execution: eligible RMSNorm and TopK calls dispatch to FlyDSL when the optional runtime is installed and enabled; other calls retain 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. Multi-head attention: PyTorch’s built-in operator 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. Figure 1. FlyDSL kernels reached through PyTorch APIs. Eager RMSNorm and TopK dispatch by input support; AOTriton selects a tuned SDPA implementation; and TorchInductor benchmarks eligible GEMM candidates. Supported Features Current support targets AMD MI350-series GPUs with the gfx950 architecture. The following features are available in PyTorch: Operation Mode PyTorch API Input dtype Layout / shape Output / options Dense GEMM Compile torch.mm FP16, BF16 Static 2D; NN/NT/TN/TT Same dtype MXFP scaled GEMM Compile F.scaled_mm MXFP8, MXFP4 Static 2D; A row-major, B column-major FP16/BF16; block scales Grouped GEMM Compile F.grouped_mm FP16, BF16 Ragged 2D A; grouped 3D B Uneven/empty groups RMSNorm Eager F.rms_norm FP16, BF16, FP32 Contiguous; 1-D norm shape Forward TopK Eager torch.topk FP32 Contiguous; last dimension Largest, sorted; functional/ out= Attention (via AOTriton) Eager, compile F.scaled_dot_product_attention FP16, BF16 D innermost; head dim 32–512 Forward and backward Dense GEMM’s four-layout support accepts row-major and column-major inputs, including eligible transpose views. In the layout names, the first letter describes A and the second describes B; N and T denote non-transposed and transposed input layouts. Grouped GEMM supports different row counts per group, allowing experts to process uneven token assignments. Its current input-layout requirements are separate from dense GEMM’s four-layout coverage. RMSNorm backward continues through PyTorch’s existing implementation. MXFP scaled GEMM accepts FP8 E4M3 or packed FP4 E2M1 operands, with one E8M0 scale per 32 values along the reduction dimension ( BlockWise1x32 ). Applications supply the quantized operands and contiguous, unswizzled scale tensors. Both formats use FP32 accumulation and return FP16 or BF16; the current path requires K to be a multiple of 128. It supports an optional one-dimensional bias whose length is N and whose dtype matches the output; fast accumulation is not supported. Under torch.compile , Inductor’s layout constraint runs before backend selection and requires unit stride along K for A and B—row-major A and column-major B—while preserving compatible leading strides. For incompatible inputs, Inductor inserts a layout copy. The FlyDSL MXFP kernel and the measurements below use these constrained layouts. This constraint is specific to scaled GEMM: Dense torch.mm detects eligible NN/NT/TN/TT strides and passes their layout flags to FlyDSL without applying the MXFP canonicalization. The FlyDSL attention kernels match the flash and memory-efficient attention capabilities of F.scaled_dot_product_attention , with one current limit: FP32 is not supported. FP32 inputs, and workloads that the tuning database assigns elsewhere, keep AOTriton’s Triton or AITER ASM kernels. AOTriton also ships FlyDSL attention kernels for RDNA4 ( gfx1201 ). Under torch.compile , Inductor keeps SDPA’s flash and memory-efficient attention operators as calls to ATen rather than generating code for them, so they still reach the AOTriton kernels. Hand-written matmul–softmax–matmul patterns that Inductor rewrites into SDPA follow the same path. By contrast, torch.nn.attention.flex_attention is a separate API that Inductor compiles into its own Triton kernels; it does not use AOTriton. Each operation has shape and alignment requirements. The detailed dense GEMM , MXFP scaled GEMM , grouped GEMM , RMSNorm , and TopK support descriptions define those boundaries. Unsupported eager calls retain ATen behavior; compiled GEMMs can use the other enabled backends. The AOTriton 0.14.2b update brings these kernels to PyTorch. Kernel-Level Performance The following operator-level benchmarks were measured on AMD Instinct MI355X ( gfx950 ) GPUs. The suites use different benchmark setups: dense GEMM uses GPU-event timing of graph replay, grouped GEMM reports steady-state throughput, and RMSNorm and TopK use GPU-event timing of repeated operator calls. The results focus on steady-state operator execution, excluding compilation and initial autotuning; they do not represent end-to-end model latency. MXFP measurements start from already quantized operands, so input quantiza