PyTorch FlyDSL: Operator support, performance gains

1. More options in existing operator paths

Support covers standard dense matrix multiplication, which computes values without gaps, grouped matrix multiplication for processing multiple groups together, and low-precision scaled matrix multiplication. It also includes RMSNorm for normalizing values, TopK for selecting the K largest values, and multi-head attention for computing relationships between inputs.

In compiled execution, TorchInductor benchmarks eligible matrix multiplication implementations and caches the fastest candidate. Existing and new implementations compete for each workload. In eager execution, an installed and enabled optional runtime routes eligible normalization and selection operations to new kernels, code containing the computation steps to execute.

For attention, implementations are selected through the AOTriton library using pretuning data.

FlyDSL supports hardware-specific optimization by giving kernel authors direct control over data movement and computation layout.

2. Measurement ranges differ by operator

Operator-level measurements ran on Instinct MI355X (gfx950), excluding compilation and initial autotuning. Low-precision measurements started with already quantized inputs, which represent values with fewer bits. The geometric means below give equal weight to each case.

The companies reported a 1.10 times geometric-mean speedup for BF16 dense matrix multiplication over the faster of two existing implementations across 15 NT-layout shapes. NT uses a row-major left matrix and a column-major right matrix. The two baselines were ATen and Triton.

The companies reported speedups of 1.58 and 1.68 times for MXFP8 and MXFP4 scaled matrix multiplication, respectively, over the existing default implementation across 17 NT-layout shapes each.

For grouped matrix multiplication, the companies reported geometric-mean speedups of 1.20 times over Triton and 1.95 times over ATen across 24 cases.

The companies reported RMSNorm speedups of 1.17–3.66 times over ATen across 22 cases. Aligned dimensions generally saw 1.2–1.5 times speedups, with the peak coming from a specific FP16 input.

The companies reported a geometric-mean TopK speedup of 4.82 times over ATen across 10 cases with small K. This was with deterministic algorithms, procedures producing the same results for the same input, disabled. With them enabled, the speedup was 4.02 times.

3. Results for attention and model inference

Whole-model measurements used vLLM on a single MI355X. Input and output lengths were fixed at 256 and 512 tokens, with 8–256 concurrent requests. Both sides used the same serving settings and per-model caches; only the matrix multiplication candidate list differed.

For BF16 inference, the companies reported up to 1.13 times speedup for Llama-3.1-8B-Instruct and 1.06 times for Qwen3-32B when new candidates were added to the two existing backends. Llama-3.3-70B-Instruct largely stayed within measurement noise.

The companies reported the following results for low-precision 8-bit inference. The baseline was ATen using hipBLASLt, and all 18 combinations of model and concurrent request count improved.

Single MI355X (TP=1), MXFP8 A8W8, 256 input tokens and 512 output tokens, 8–256 concurrent requests, FlyDSL added versus ATen (hipBLASLt), 1 measurement per combination

Gains varied by model and concurrent request count. The companies also reported 1.25–3.01 times speedups for MXFP4 inference across 18 combinations under the same request conditions. These measurements left the final output layer unquantized and ran each combination once.

Attention results aggregate throughput relative to an earlier AOTriton version. On gfx950, the companies reported 0.92–2.34 times forward throughput and 0.93–2.47 times backward throughput for 0.14.2b versus 0.13b.

4. Requirements and performance trade-offs

Start with a recent ROCm PyTorch nightly on a gfx950-based MI350-series GPU. The optional runtime for compiled matrix multiplication is in the 0.3.x series. Below is the candidate configuration from the original example.

import torch
import torch._inductor.config as config

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

After configuration, compiling matrix multiplication with torch.compile(mode="max-autotune") makes it an autotuning candidate. Attention uses builds linked to AOTriton 0.14.2b or later, without separate packages or configuration.

Faster repeated execution comes at the cost of compilation and tuning on the first call. Low-precision inputs that do not match the required layout also incur conversion copies. Optimization becomes more valuable as supported operators account for more of the model's runtime.

The attention version comparison also includes changes to other kernels and tuning data, and the new attention kernels do not currently support FP32. Normalization backward passes use existing implementations, and the index order of tied values in TopK can differ by backend.

5. Broader support is in progress

The current integration targets the gfx950-based MI350 series. The source says work is underway on MI450 support, more low-precision and attention kernels, post-processing fusion that combines subsequent computations, and broader backward-pass support.

FlyDSL adds optimization options to PyTorch on AMD GPUs, with gains especially notable in low-precision matrix multiplication, some selection operators, and certain inference conditions.