Every time you ask an AI chatbot a question, something invisible happens in a fraction of a second. The model scans your entire conversation, figures out which words relate to which, and generates a response that (usually) makes sense.
That "figuring out which words relate to which" part? It's called attention, and it's arguably the most important computation in all of modern AI. It's also one of the most expensive.
For the past three years, a single open-source project has quietly made that computation dramatically faster for nearly every major AI model on the planet. It's called FlashAttention. And today, its creator just released the fourth major version, redesigned from scratch for NVIDIA's newest Blackwell GPUs.
Here's why you should care, even if you never touch a GPU in your life.
First up, the TL;DR
Every time you ask ChatGPT, Claude, or Gemini a question, there's an invisible piece of software doing the heavy lifting behind the scenes. It's called FlashAttention, and its creator just dropped a major new version.
Here's the deal: "Attention" is the core mechanism that lets AI models understand context; it's how a chatbot knows that "it" in your sentence refers to your dog, not your car. The problem? Attention is painfully slow and memory-hungry, especially with longer conversations. In 2022, Stanford PhD student Tri Dao figured out a clever workaround. Instead of shuffling massive amounts of data back and forth through GPU memory, he reorganized the math to keep everything in the chip's tiny but blazing-fast cache. The result: 2-4x faster training with zero loss in accuracy.
That was FlashAttention 1. Versions 2 and 3 followed, each squeezing more performance out of newer hardware. Now, FlashAttention-4 has arrived, redesigned from the ground up for NVIDIA's newest Blackwell GPUs.
Here's what matters:
- Near-maximum speed. FA4 reaches ~1,600 TFLOPS on B200 GPUs (that's 71% utilization), up to 2.7x faster than previous methods.
- New bottleneck, new fix. Blackwell's tensor cores (the math units) are so fast now that the exponential function inside softmax became the slowdown. FA4 solves this with a software workaround that splits the math between hardware and a polynomial approximation (fancy math shortcut).
- Written in Python. The whole thing is built in CuTe-DSL, so installing takes seconds instead of the hours that C++ templates used to require. You can install it right now:
pip install fa4. - PyTorch integration. FlexAttention now has an FA4 backend, delivering 1.2-3.2x speedups for custom attention variants.
Oh, and one fun detail: Tri Dao used Claude to help debug a deadlock in the code. It ran autonomously for six hours overnight, figured out part of the fix, then went down a rabbit hole convincing itself the compiler was broken. Very human behavior, honestly.
FlashAttention powers virtually every major AI model in production today. This matters because faster attention means cheaper training, longer context windows, and snappier responses. The version 4 jump is significant because Blackwell GPUs are the hardware that will define AI infrastructure for the next several years (while Rubin rolls out), and FA4 ensures that software can actually keep up.
Now, let's dive into that all a bit more in depth, shall we?
What Is Attention, and Why Is It Expensive?
The 2017 paper "Attention Is All You Need" introduced the Transformer architecture (the "T" in GPT). At its core, the Transformer uses an "attention mechanism" to figure out relationships between every word in a sequence.
Think of it this way: when you type "The cat sat on the mat because it was tired," the attention mechanism is what tells the model that "it" refers to the cat, not the mat.
The catch? Attention has to compare every word against every other word. For a conversation with 1,000 words (called "tokens" in AI), that's 1 million comparisons. For 100,000 tokens (the length of a long document), that's 10 billion. The computation scales quadratically, meaning doubling the input makes things four times more expensive. This is why longer conversations and bigger context windows cost so much money and compute.
Enter Tri Dao and FlashAttention
Tri Dao was a Stanford PhD student (advised by Christopher Ré and Stefano Ermon, friend of the pod!) when he published the original FlashAttention paper in May 2022.
The core insight was deceptively simple:
- GPUs have two types of memory: a big, slow one (HBM, which stands for High Bandwidth Memory; think of it like a warehouse) and a tiny, fast one (SRAM; think of it like the counter in front of you).
- Standard attention implementations wrote enormous intermediate results to the slow memory, then read them back.
- FlashAttention reorganized the math so everything stayed on the fast "counter" as much as possible. No approximations or quality loss; simply smarter plumbing.
The impact was immediate. FlashAttention delivered 2-4x speedups in training. Every major open-source model you've heard of (LLaMA, Falcon, MPT, and many more) adopted it. It won an ICML Outstanding Paper award.
Tri went on to become Chief Scientist at Together AI and an Assistant Professor at Princeton. He kept improving the algorithm:
- FlashAttention-2 (July 2023): Better parallelism and work partitioning. Reached up to 73% GPU utilization on A100 chips and 225 TFLOPS in end-to-end model training. About 2x faster than the original.
- FlashAttention-3 (2024): Optimized for NVIDIA's Hopper GPUs (like the H100). Exploited new hardware features including asynchronous tensor cores and warp specialization to hit 740 TFLOPS; 75% of the H100's theoretical max.
Each version was tailored to squeeze performance out of the specific hardware generation it ran on. Which brings us to today.
FlashAttention-4: Built for Blackwell
NVIDIA's Blackwell GPUs (the B200 and its variants) represent a generational leap in raw compute. But that leap created a weird new problem.
Here's the key concept: asymmetric hardware scaling. On Blackwell, the tensor cores (the units that do matrix multiplication, the bread-and-butter math of AI) got dramatically faster. From the Hopper H100 to the Blackwell B200, raw tensor core throughput jumped from 1 to 2.25 petaFLOPS. But everything else; the special function units (which handle operations like the exponential function used in softmax), and the shared memory bandwidth; stayed roughly the same.
In plain English: the engine got way more powerful, but the fuel lines didn't get any wider.
This means the bottleneck shifted. On older GPUs, the matrix multiplications were the slowdown. On Blackwell, the tensor cores blaze through the matrix math so fast that they end up waiting around for the exponential calculation inside softmax. In the forward pass (when the model generates your answer), the exponential and the matrix multiplications now take the exact same number of clock cycles: 1,024 each. In the backward pass (used during training), shared memory bandwidth is the wall.
FlashAttention-4 is a ground-up redesign that addresses both bottlenecks:
Forward pass fixes:
- Ping-pong scheduling. FA4 alternates between two sets of query data ("tiles"), so while one tile is doing matrix multiplication, the other is computing softmax exponentials. Neither unit sits idle.
- Software-emulated exponentials. Since the hardware exponential unit can't keep up, FA4 splits the work: some calculations go to the hardware's built-in exponential function, and the rest get handled by a polynomial approximation (a mathematical shortcut) running on otherwise-idle FMA (multiply-add) units. It's like opening a second checkout lane when the first one is backed up.
- Conditional rescaling. Standard online softmax requires rescaling after every step. FA4 skips rescaling when the correction would be tiny, cutting a bunch of unnecessary computation off the critical path.
Backward pass fixes:
- Tensor Memory (TMEM). Blackwell GPUs have a new on-chip scratchpad called Tensor Memory, wired directly into the tensor cores. FA4 stores intermediate results there instead of in shared memory, relieving the bandwidth bottleneck.
- 2-CTA MMA mode. Two thread blocks cooperate on a single, larger matrix multiplication tile. Each block stages half the data, roughly halving shared memory traffic.
- Deterministic mode. Important for reinforcement learning and debugging; FA4 supports fully reproducible backward passes, keeping about 85-90% of peak performance even in deterministic mode.
The result? FA4 reaches up to 1,605 TFLOPS on B200 GPUs (71% utilization). That's up to 1.3x faster than NVIDIA's cuDNN and 2.7x faster than Triton implementations.
Written in Python, Installed in Seconds
One of the most practical upgrades in FA4 might surprise you: it's written entirely in CuTe-DSL, a Python-embedded domain-specific language.
Previous FlashAttention versions relied heavily on C++ templates, which meant compilation times measured in minutes or even hours. CuTe-DSL compiles 20-30x faster. Researchers can iterate on kernel changes in seconds instead of waiting through painful build cycles.
To install: pip install fa4
That's it.
PyTorch Gets the Upgrade Too
The same day as the FA4 release, the PyTorch team announced that FlexAttention now has a FlashAttention-4 backend for Hopper and Blackwell GPUs.
FlexAttention lets researchers define custom attention variants (things like sliding window attention, ALiBi, soft-capping) in a few lines of Python. It's been widely adopted; over 1,000 repos and dozens of research papers reference it. But users kept hitting a performance ceiling. On Blackwell especially, the gap between FlexAttention's Triton implementation and a hand-optimized kernel like cuDNN had grown from a manageable gap to a chasm.
The FA4 backend closes that gap. PyTorch auto-generates CuTe-DSL score and mask modifications, then compiles them into FA4 kernels on the fly. The result: 1.2x to 3.2x speedups over the existing Triton implementation on compute-bound workloads. Researchers no longer have to choose between flexibility and performance.
The Claude Debugging Story
Here's maybe the best detail in the entire release. Tri Dao shared that he used Claude to help debug a deadlock in FA4's 2-CTA forward pass implementation.
Claude ran autonomously overnight for six hours. It figured out part of the fix, but then went down a rabbit hole convincing itself the compiler was broken. (Tri called it "very human" behavior.) From Claude's partial fix, Tri was able to resolve the hang in 10 minutes.
His takeaway? He's hoping FA5 will be written completely by AI.
Why This Matters for Everyone
You might be thinking: OK, GPU kernel optimizations. Why should I care?
Because FlashAttention is the invisible infrastructure beneath virtually every AI product you use. ChatGPT, Claude, Gemini, Llama, Mistral; they all rely on it (or implementations derived from it) for efficient attention computation. When FlashAttention gets faster, three things happen:
- Training costs drop. Higher GPU utilization means companies get more compute per dollar. That eventually translates to cheaper AI services.
- Context windows grow. The reason you can now paste entire documents into Claude or ChatGPT is, in part, because FlashAttention made long-context attention feasible. FA4 pushes that further on next-gen hardware.
- Inference gets snappier. Faster attention means faster response times, which means less waiting when you're chatting with an AI assistant.
Blackwell GPUs will define AI infrastructure for the next several years. Every major cloud provider (AWS, Google Cloud, Azure, Oracle) is deploying them. FA4 ensures the software layer can actually exploit what the hardware offers.
And the fact that a single research team; led by a researcher who started this work as a grad student just four years ago; continues to set the pace for how the entire industry computes attention? That's a story worth paying attention to. (Pun fully intended.)
