💻 computer science

Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants

This paper introduces Flashlight, a PyTorch-native compiler framework that automatically generates fused, FlashAttention-style kernels for arbitrary and data-dependent attention patterns without relying on static templates, thereby offering superior flexibility and competitive performance compared to existing solutions like FlexAttention.

Original authors: Bozhi You, Irene Wang, Zelal Su Mustafaoglu, Abhinav Jangda, Angélica Moreira, Roshan Dathathri, Divya Mahajan, Keshav Pingali

Published 2026-05-22

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Bozhi You, Irene Wang, Zelal Su Mustafaoglu, Abhinav Jangda, Angélica Moreira, Roshan Dathathri, Divya Mahajan, Keshav Pingali

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Traffic Jam" in AI

Imagine a Large Language Model (like the one you are talking to) is a massive factory trying to assemble a complex product (a sentence or a protein structure). The most critical part of this factory is the "Attention" department. This is where the machine decides which pieces of information are important and which to ignore.

The problem? The current way this department works is like a traffic jam.

The Old Way: The factory workers (the computer code) have to stop, write a note on a piece of paper (save data to slow memory), walk to a different station, pick up the note, and then start working again. This happens over and over. It's slow and wastes energy.
The "Flash" Fix: A few years ago, engineers invented FlashAttention. They figured out how to make the workers keep all the notes in their pockets (fast memory) and do the whole job in one continuous sprint. This was a huge speed boost, but it was like a specialized race car: it only worked on one specific type of track (standard attention). If you wanted to drive a different type of car (a new, experimental attention method), you had to build a whole new race car from scratch by hand.

The Problem: Too Many Custom Race Cars

Recently, scientists have invented many new, fancy ways to do "Attention" to make AI smarter or more efficient (like "Differential Attention" or "Gated Self-Attention").

The Bottleneck: To make these new methods fast, you usually need a human expert to hand-code a custom "race car" (a specialized computer kernel) for each one. This takes forever.
The Middle Ground (FlexAttention): A tool called FlexAttention tried to fix this by giving builders a set of Lego templates. If your new car fit the template, it would work fast. But if your car was weird or unique (didn't fit the template), FlexAttention couldn't help you, and you were back to the slow, manual coding.

The Solution: FLASHLIGHT (The Universal Factory Upgrade)

Enter FLASHLIGHT. The authors describe it as a compiler-native framework. In simple terms, it's a smart "auto-pilot" built directly into the PyTorch software (the language most AI researchers use).

Here is how FLASHLIGHT works, using an analogy:

1. The "Smart Chef" vs. The "Recipe Book"

Old Way: You have a cookbook (the code). If you want to make a standard soup, the cookbook tells you to chop, boil, and serve. If you want a weird soup, you have to hire a chef to invent a new recipe and write a new cookbook page.
FlexAttention: The cookbook has a "Special Soup Section." If your soup fits the "Special Soup" template, it's fast. If not, you're stuck.
FLASHLIGHT: FLASHLIGHT is like a super-smart kitchen robot that watches you cook. It doesn't care if you are making a standard soup or a weird, experimental soup. It looks at your actions (the code), realizes you are doing a lot of chopping and boiling, and says, "Hey, I can combine all these steps into one super-efficient motion so you don't have to walk back and forth to the fridge."

It automatically rewrites your code to be as fast as a hand-crafted race car, without you needing to know how to build the car or fit a template.

2. How It Does the Magic (The Three Tricks)

The paper explains three main "tricks" FLASHLIGHT uses to speed things up:

Trick A: The "Demotion" (Fusing Steps)
Imagine you are baking a cake. Step 1 is mixing flour. Step 2 is adding eggs. Usually, you mix the flour, put the bowl down, pick up a new bowl, and add eggs.
FLASHLIGHT says, "Why stop? Keep the flour in the bowl and just add the eggs right there." It merges separate steps into one continuous flow, so data doesn't have to travel back and forth.
Trick B: The "Math Magic" (Algebraic Transformation)
Sometimes, to be safe, you have to do a calculation twice (like checking the temperature, then checking it again to be sure). This is slow.
FLASHLIGHT uses a mathematical trick (called a "homomorphism") to realize that you can do both checks at the same time while you are cooking, rather than stopping to do them separately. It's like checking the oven temperature while you are stirring the batter, rather than stopping to look at the oven, then going back to stir.
Trick C: The "Tiling" (Packing the Truck)
Imagine you are moving houses. If you carry one chair at a time, it takes forever. If you pack a whole truck (a "tile") with furniture and drive it once, it's fast.
FLASHLIGHT is smart about how it packs the truck. It figures out the perfect size for the truck based on what you are moving. If a piece of furniture is small, it fits it into a corner so you don't waste space. It ensures the computer's "truck" (memory) is always full and moving efficiently.

What Did They Prove?

The researchers tested FLASHLIGHT on powerful computer chips (NVIDIA H100 and A100).

For Standard Tasks: When they used the "standard" attention methods that FlexAttention could already handle, FLASHLIGHT was just as fast or even faster.
For Weird/New Tasks: When they tried attention methods that FlexAttention couldn't handle (like the ones used in AlphaFold, a famous protein-folding AI), FLASHLIGHT was 5 times faster than the standard, unoptimized code.
Real-World Test: They tested it on a real protein-folding model (AlphaFold). It made the model run 6% to 9% faster overall.

The Bottom Line

FLASHLIGHT is a tool that lets AI developers write code in the normal, easy way they always have, but it automatically turns that code into a high-speed, super-efficient engine.

Before: You had to choose between "Easy to write but slow" OR "Fast but hard to write."
With FLASHLIGHT: You get "Easy to write AND Fast."

It removes the need for experts to hand-code special speed-ups for every new idea, allowing scientists to experiment with new types of AI attention models without worrying about the code running too slowly. It is currently open-source, meaning anyone can use it right now as a fork of PyTorch.

Technical Summary: FLASHLIGHT

Problem Statement

Optimizing attention mechanisms is critical for accelerating the training and inference of Large Language Models (LLMs). While FlashAttention successfully improved performance for standard attention by leveraging tiling and kernel fusion to reduce memory bandwidth and launch overhead, it is a hand-optimized kernel specific to vanilla attention. It cannot be applied to newer, more complex attention variants (e.g., differential attention, row/column-wise gated self-attention in AlphaFold's Evoformer, Invariant Point Attention, or Rectified Sparse Attention) without significant manual engineering effort.

Existing solutions like FlexAttention and FlashInfer attempt to bridge this gap by providing static, higher-order templates or just-in-time compilation for specific subsets of attention variants. However, these approaches suffer from limited expressiveness: they require programmers to re-express their models to fit predefined templates. Consequently, many novel attention formulations that do not fit these templates remain constrained by slower, unfused PyTorch implementations, hindering their adoption. Furthermore, the standard PyTorch compiler stack (TorchInductor) currently lacks key optimizations such as reduction fusion and complex operator fusion across memory boundaries, limiting the performance of generated code for these variants.

Methodology

The paper introduces FLASHLIGHT, a compiler-native framework within the PyTorch ecosystem designed to automatically generate fused, FlashAttention-style kernels for arbitrary attention-based programs. Unlike previous approaches, FLASHLIGHT does not rely on static templates or predefined kernel specializations. Instead, it integrates directly into the torch.compile workflow, analyzing and transforming standard PyTorch code to discover and fuse constituent operations.

FLASHLIGHT extends the TorchInductor intermediate representation (IR) and optimization passes through three primary mechanisms:

Unified Reduction IR: FLASHLIGHT models tensor contractions (like General Matrix Multiplication, GEMM) as generalized reductions within a unified IR that distinguishes between parallel dimensions ( $p$ -dimensions) and reduction dimensions ( $r$ -dimensions). This allows matrix multiplications to participate fully in the end-to-end fusion engine, removing the artificial boundary that typically isolates GEMM operations from surrounding computations.
Structural Fusion with Dimension Demotion: The framework introduces a fusion rule that allows a parallel loop ( $p$ -dimension) from a producer kernel to be "demoted" into a reduction loop ( $r$ -dimension) within a fused consumer kernel. This transformation trades parallelism for the elimination of high-latency materialization of intermediate tensors to global memory, enabling operations like the max() in softmax to be fused directly with the preceding matrix multiplication ( $QK^\top$ ).
Semantic Fusion with Algebraic Transformation: To fuse state-dependent operations like softmax, FLASHLIGHT leverages algebraic properties (specifically homomorphisms) to transform multi-pass algorithms into single-pass online algorithms. For instance, it converts the stable softmax (which requires two passes: finding the max, then summing exponentials) into an online softmax algorithm. This is achieved by maintaining running accumulators for the maximum and the sum, dynamically rescaling the sum whenever the running maximum changes using the property $e^{x-y} = e^x / e^y$ .
Tiling-Aware Dimension Elimination and Logical Grids: The framework refines loop execution by iterating over contiguous tiles. It introduces logical grid dimensions to decouple the problem's structure from physical GPU hardware constraints. This allows for flexible, per-dimension tile sizes and enables the elimination of small dimensions that fit entirely within a single tile, unlocking fusion opportunities for complex multi-stage computations (e.g., fusing $softmax(QK^\top)V$ ) that were previously impossible for the compiler.

Key Contributions

Kernel-Level Fusion: FLASHLIGHT implements a mechanism to fuse compatible compute blocks (e.g., matmul + softmax, matmul + matmul) automatically when parallelism permits, exploiting performance gains without requiring manual kernel writing.
Generality via Compiler-Driven Approach: It supports both standard and emerging attention mechanisms through a unified compiler-driven approach, removing the need for static templates or custom kernels.
Performance with Flexibility: It produces kernels with competitive or superior performance to existing template-based solutions (like FlexAttention) while preserving the full flexibility of native PyTorch code.

Results

The authors evaluated FLASHLIGHT on NVIDIA H100 and A100 GPUs across various attention variants:

FlexAttention-Supported Variants: For variants supported by FlexAttention (Vanilla, ALiBi, Softcap, Causal, Sliding Window, etc.), FLASHLIGHT-generated code is competitive with or faster than FlexAttention. Specifically, for score_mod variants, FLASHLIGHT is up to 1.48× faster due to simpler fused kernels that avoid the overhead of handling block masks. For block_mask variants, while FlexAttention's kernel execution is faster due to sparsity skipping, FLASHLIGHT avoids the significant overhead of constructing sparse block masks at runtime.
Complex Variants (Unsupported by FlexAttention): For attention variants not expressible in FlexAttention's template (Differential Attention, AlphaFold's Evoformer row/column-wise gated self-attention), FLASHLIGHT significantly outperforms standard torch.compile.
- Differential Attention: Shows higher speedups on H100 compared to A100.
- Evoformer: FLASHLIGHT achieves speedups of 5× or more compared to torch.compile on both GPU architectures.
End-to-End Inference:
- AlphaFold2: FLASHLIGHT improves inference latency by 6% to 9% for row/column-wise gated self-attention layers.
- LLaMa-3.2-1B: In a vLLM system using a Mooncake conversation trace, FLASHLIGHT outperforms FlexAttention for Softcap attention. For Causal attention (a block-mask variant), FlexAttention remains slightly more performant due to its optimized sparse block mask handling, though FLASHLIGHT avoids out-of-memory errors that occur with standard torch.compile due to intermediate tensor materialization.

Significance and Claims

The paper claims that FLASHLIGHT represents a significant shift in how attention variants are optimized, moving the burden from manual engineering to compiler automation. By enabling the automatic generation of fused, FlashAttention-style kernels from standard PyTorch code, it allows AI scientists to rapidly explore new attention models without sacrificing performance or scalability.

The authors emphasize that FLASHLIGHT is open source and available as a fork of PyTorch. They position the framework as a general-purpose solution that not only supports all variants expressible in the FlexAttention model but also handles more general, data-dependent formulations that exceed FlexAttention's capabilities. While the current evaluation focuses on inference, the authors note that the underlying extensions and passes are general-purpose and could be extended to support training workloads and optimize for block sparsity in future work.