PyTorch's FlyDSL Boosts AMD MI350 GPU Speed by 3x With Low-Precision Math

PyTorch has added FlyDSL as an optional backend for AMD Instinct MI350-series data-center GPUs. Applications can reach specialized gfx950 kernels through standard PyTorch operators, covering dense and grouped general matrix multiplication (GEMM), scaled low-precision GEMM, RMSNorm, TopK, and attention. Unsupported operations continue through existing ATen or Triton paths.

FlyDSL is AMD's open-source Python domain-specific language for GPU kernels. Its PyTorch integration matters most for workloads using MXFP8 and MXFP4, block-scaled 8-bit and 4-bit matrix formats, where the published MI355X benchmarks show larger gains than BF16. Applications keep their existing operator calls, so enabling the backend requires no custom-kernel integration or model-code changes.

Standard operators, three dispatch paths

PyTorch exposes FlyDSL through its existing eager, compilation, and attention systems. The selection mechanism depends on the operator:

A ROCm nightly environment requires FlyDSL 0.3.x:

python -m pip install 'flydsl>=0.3.0,<0.4'

The following TorchInductor settings add FlyDSL to the GEMM autotuning pool:

import torch._inductor.config as config

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

These settings allow Inductor to benchmark FlyDSL during compilation. They do not force every GEMM onto the backend; ATen or Triton can still win for a given shape, layout, and data type. Autotuning adds compilation work on the first encounter and reuses the cached choice afterward.

Hardware control moves into Python

Triton presents kernel authors with a block-tensor abstraction and delegates many hardware choices to the compiler. FlyDSL exposes lower-level scheduling controls: authors can partition tiles across thread blocks, wavefronts, threads, and values; choose tiled-copy and matrix fused multiply-add instructions; define Local Data Share layouts and swizzles; stage buffers; and place synchronization or compiler-scheduling boundaries.

Those controls let a kernel encode MI350-specific decisions about data movement, instruction ownership, and reuse of on-chip shared memory. Waves are groups of GPU threads, MFMA instructions perform matrix multiply-accumulate work, and LDS is AMD's on-chip shared-memory region.

FlyDSL's attention kernel uses a dual-wave software pipeline that overlaps vector softmax work for one key-value tile with matrix instructions for another. At a head dimension of 512, the kernel also partitions the dimension over time through LDS and distributes the output accumulator across waves. That schedule targets the execution and memory behavior of gfx950 directly.

Low precision leads the kernel results

The published operator benchmarks ran on an AMD Instinct MI355X. Each result uses the stated baseline, and geometric means combine performance across multiple shapes:

FlyDSL exceeded ATen throughput in every reported low-precision case. MXFP8 speedups ranged from 1.31x to 2.23x, while MXFP4 ranged from 1.31x to 2.96x. RMSNorm gained roughly 1.2x to 1.5x on aligned hidden dimensions; an off-by-one dimension such as 4097 reached 3.66x because ATen handled the remaining tail less efficiently. The TopK result was measured with deterministic algorithms disabled.

Quantization widens the serving gains

The end-to-end tests used vllm bench serve on one MI355X with Llama-3.1-8B, Qwen3-32B, and Llama-3.3-70B. The benchmark covered six concurrency settings spanning 8 to 256 and changed the TorchInductor GEMM candidate list between runs.

MXFP8 improved all 18 combinations of model and concurrency, using ATen as the baseline. MXFP4 produced the largest full-request result, reaching 3.01x for Qwen3-32B. The smaller BF16 changes show that faster kernels translate unevenly to serving throughput when other work, including scheduling, communication, and memory movement, occupies more of the request.

The support boundary remains narrow

• GPU architecture: The PyTorch backend described here targets gfx950 in the MI350 series. It does not cover MI300 or consumer RDNA GPUs. AOTriton separately ships FlyDSL attention kernels for RDNA4 gfx1201.
• MXFP shapes and scales: Scaled GEMM requires K to be a multiple of 128 and uses BlockWise1x32 E8M0 scales.
• MXFP layouts: TorchInductor requires row-major A and column-major B, inserting a layout copy when an input does not match.
• Attention types: FlyDSL attention lacks FP32 input support, so those calls remain on Triton or AITER paths.
• Grouped GEMM: Grouped operations have layout constraints separate from the four layouts supported by dense GEMM.
• TopK ties: Tie ordering can differ from ATen in deterministic mode because FlyDSL canonicalizes NaNs to one ordinal value.
• Benchmark scope: The figures come from project-published MI355X tests. Deployment results will depend on model shapes, prompt and output lengths, concurrency, layouts, and compilation-cache behavior.

The published roadmap includes MI450-series support, MXFP8 grouped GEMM, a backend for FlexAttention, fused GEMM epilogues, and wider backward-pass coverage for training.

MI350 deployments using vLLM and MXFP8 or MXFP4 have the clearest evidence for enabling FlyDSL as an Inductor candidate. BF16 workloads show smaller and more model-dependent gains, while unsupported shapes retain the existing ATen and Triton fallbacks.