AMD FlyDSL Joins TorchInductor: A New MLIR-Native GPU Kernel Backend

Why a New Kernel Backend Matters

Modern AI workloads increasingly rely on a wide variety of specialized GPU kernels. AMD's Liz Li and Jia Hui made a key observation in their PyTorch tech talk: no single kernel generation backend can be optimal across all operators, chips, and architectures. This means the compilation stack needs enough flexibility to choose the most appropriate kernel implementation path for each scenario.

With that motivation in mind, the AMD team integrated their in-house FlyDSL as a new backend into TorchInductor. FlyDSL is AMD's Python-native, MLIR-based GPU kernel DSL (Domain-Specific Language). It retains the ease of use of the Python ecosystem while leveraging MLIR — a modern compiler infrastructure — to unlock hardware-level optimization opportunities.

MLIR (Multi-Level Intermediate Representation) is a modern compiler infrastructure led by Google and maintained by the LLVM community. Its core design philosophy is to progressively "lower" high-level algorithm descriptions to hardware instructions through multiple layers of intermediate representation (IR). Unlike traditional compilers, MLIR allows multiple dialects to coexist within the same framework, enabling the entire transformation chain — from Python operators to GPU assembly — to be completed within a unified system while preserving optimization opportunities at every level. This contrasts with Triton's compilation path, which also generates GPU kernels but operates at a different level of abstraction and optimization granularity. Because FlyDSL is built on MLIR, it can leverage MLIR's existing pass infrastructure (such as loop transformations, vectorization, and memory hierarchy analysis) to generate code that better matches hardware characteristics — without having to reimplement those optimization routines from scratch.

Seamless Integration into the TorchInductor Compilation Pipeline

The engineering value of this work lies in how the integration was done. Rather than building an entirely separate toolchain, the AMD team plugged FlyDSL into Inductor's existing core stages — including lowering, compilation, caching, benchmarking, and the auto-tuning pipeline.

The biggest benefit of this approach is continuity in developer experience. According to the presenters, the entire process preserves the familiar Torch and Triton user experience, meaning upper-level code requires virtually no changes to switch the underlying kernel generation to FlyDSL. For the large number of models and applications already built on PyTorch, this represents a low-migration-cost path to performance gains.

Key Operators: GEMM and Attention

With this integration in place, AMD was able to bring several optimized FlyDSL kernels into PyTorch. The typical cases highlighted in the talk include GEMM (General Matrix Multiplication), Group GEMM, and Flex Attention — three of the most computationally intensive and frequently used operators in large model training and inference.

GEMM is the foundation of nearly all neural network computation. Group GEMM is especially critical in architectures like MoE (Mixture of Experts). Flex Attention directly affects the efficiency of attention computation in Transformer-based models. FlyDSL's focus on these operators is clearly targeted.

MoE (Mixture of Experts) is an architectural design that decouples model capacity from compute. The core idea is that during the forward pass, only a small subset of all "expert" sub-networks is activated for each token, allowing model parameter counts to scale significantly while keeping computation relatively manageable. Group GEMM is precisely the bottleneck in MoE computation: different experts receive different numbers of tokens, requiring a batch of small matrix multiplications of varying shapes to be executed simultaneously rather than a single large, uniform matrix multiply — creating additional challenges for kernel scheduling and hardware utilization. Leading models like DeepSeek^Ad V2/V3 use MoE architectures, making Group GEMM optimization directly valuable for real-world deployment.

Performance Numbers on MI355

The real headline is the benchmark results. On AMD MI355, FP16 GEMM achieves approximately 19% geometric mean speedup over Triton. While this may seem modest, given that GEMM is one of the largest contributors to overall compute graphs, a nearly 20% improvement translates directly into end-to-end throughput gains.

The gains for Flex Attention are even more striking. On Minimax M3 decode shapes, FlyDSL Flex Attention achieves up to 5.8× speedup over Triton; on DeepSeek V2 and V3 prefill sizes, it delivers close to 2× improvement.

It's worth noting that attention operator speedups are highly sensitive to specific shapes and scenarios — the small-batch, long-sequence characteristics of the decode phase are fundamentally different from the large-matrix computation of the prefill phase. Peak numbers like 5.8× should therefore be understood in the context of specific shapes. Even so, these results demonstrate that the MLIR-native path has a genuine advantage when handling diverse operator shapes.

Flex Attention is a programmable attention mechanism introduced in PyTorch 2.x that allows developers to express various attention variants (such as sliding window attention, causal masking, ALiBi bias, etc.) through custom score modification functions, which the compiler then fuses into efficient kernels. Decode and prefill are two phases of inference with fundamentally different characteristics: prefill processes the full input prompt, typically involving large-batch, long-sequence matrix computation and is compute-bound; decode generates tokens autoregressively one at a time, with small batch sizes and memory access patterns dominated by KV Cache growth, making it memory-bound. The dramatically different kernel requirements of these two phases explain why the same FlyDSL implementation can show speedup ratios ranging from 2× to 5.8× depending on the shape.

What This Means for AMD's Software Stack

Triton has long been considered one of the mainstream paths for writing cross-hardware kernels. With FlyDSL, AMD offers an alternative built on its own MLIR technology stack. Rather than replacing Triton, it complements it — delivering meaningfully superior performance on specific operators while maintaining a consistent PyTorch interface for upper-level code.

This reflects AMD's sustained investment in the GPU software ecosystem: by embedding optimization results directly into PyTorch's most-used compilation path, AMD lowers the barrier for developers to benefit from hardware performance gains. As this kind of work progresses, a multi-backend strategy — selecting the best backend per operator — is poised to become standard practice in high-performance AI stacks.

The presenters indicated they will share more technical details and practical experience at future events. For engineers focused on AMD GPU performance and PyTorch compilation optimization, FlyDSL is a direction worth tracking closely.