Lightning Talk: FlexAttention + FlashAttention-4: Fast and Flexible - Driss Guessous, Meta

In this talk, TrisKus from Meta presents FlexAttention and FlashAttention V4, a new backend that significantly increases processing speeds, achieving between 1.2 to 3.2 times faster performance on Hopper and Blackwell GPUs. The speaker elaborates on the design of this new implementation, which allows for customization through user-defined score modifications and masking patterns. TrisKus explains how FlexAttention offers a flexible API that has become popular among researchers, with nearly 2,000 repositories adopting the technology. The talk goes into detail about the underlying architecture of FlashAttention, the challenges faced with performance gaps due to rapid advancements in GPU technology, and how the pipeline design addresses these issues. The session highlights the successful results of the new implementation and provides insights into future developments in the FlexAttention framework.

Hi everybody. My name is TrisKus. I work at Meta. I'm on the PyTorch team and today I'm going to be presenting on FlexAttention plus FlashAttention V4. This is a new faster back-end for FlexAttention. And here we see our lovable mascot Flexi and now he has his flash suit on. And the TLDR for this presentation is that we have a new implementation. It can achieve 1.2 to 3.2x

faster speeds and it's on a Hopper and Blackwell GPUs. Okay, let's do a quick reminder on You can think of it as a vanilla FlashAttention implementation with two hooks for customization. On the left you have score modifications and on the right you have masking patterns. And this is how a user would implement Alibi attention for say. On the right is our familiar causal attention. These modifications are

user-defined. So swapping our Alibi attention for soft capping or our causal attention for sliding window is as easy as a few changes to your PyTorch code. And you can keep going and the community has. There's been some really interesting projects and new ideas. On the left we have activation wrapper and this is a mechanism for undoing the softmax norm and replacing with an arbitrary point-wise modification. In

this case it's a ReLU. And all in all FlexAttention exposes a flexible API that has really resonated with researchers. And you don't just get to take my word for it. The adoption has been great and steadily increasing over time. There's been almost 2,000 repos that now use FlexAttention. Dozens of papers have cited FlexAttention in their work and I think this really highlights the original goal of helping

researchers explore what is possible with novel attention variants. And Flexi's got its name for its flexibility and not necessarily as raw performance. Over time we have seen the performance gap start to build. On A100 we're on par with the state-of-the-art FlashAttention implementation but on Hopper and on Blackwell we've seen this grow. So, there's really two issues. FA3 and other Soda implementations have gotten faster and faster, and

our Triton-based implementation just couldn't close the gap. And while this gap was manageable on Hopper, it's not ideal, but it's been manageable, right? There's gains of flexibility, but in sense in Blackwell, this these gains in flexibility have outweighed the cost. Um the gap has widened to a gulf. And that got us thinking that if we still wanted to provide a flexible flash attention like implementation, we were

going to have to address these performance concerns. But before we do that, I want to talk about Blackwell and why it was required to redesign our implementation to achieve Soda perf. Okay, so for those of you who do not know, this is the flash attention algorithm, loosely. Uh you can roughly think about it as a linear scan. Each CTA is assigned a chunk of queries, and they

iterate along keys and values to produce the final attention score. And basically, for Ampere and Hopper, you have one CTA that processes one chunk of queries and each scan from left to right. So, on Blackwell, what's the problem? And the problem is that Nvidia keeps making Tensor Cores faster and faster. And at least on Blackwell, on the B200, the arithmetic unit that does the exponentials couldn't keep

up. You now have the issue where your softmax takes twice as long as your gems. And so, what's the solution to this problem? And like many scheduling problems, we should probably try pipelining. And instead of having one group of queries assigned to a CTA, what if we split it in half or double it, and we had two chunks? And I want to draw out all the details.

So, this is a pretty dense animation, but hopefully we can get through it. On the top, we have the high lane, and on the bottom, we have the low lane, and these are for the Tensor Cores. They're backed by the same hardware unit, but for this pipelining, it's better to split them up. And in the middle, we have the the SFU, and And what does the softmax.

And this is the thing that now takes twice as long as the tensor cores. So this animation steps in units of tensor core time. That's what a box is. And so the step one is the top half tensor core does its gem to produce the scores the Q at K. One step later, the lower M split begins its score gem. And at the same time the upper

half softmax is already running. So we're starting to see some overlap of different kinds of work. At this stage, you can see that if we focus just on the top two rows, there's one chunk of queries. Its scores have been calculated. Softmax is now finished. And so the gem for the output the P at V has all its data dependencies fully resolved. We can start computing And

meanwhile, the lower tensor cores have also finished their tensor core work. And we can can start we can start computing its softmax. And if this wasn't complicated enough, I'm going to introduce a new color. And remember flash attention is a linear scan from left to right. And the scan is done in steps. And the brownish color is for step zero. And the light blue is step one.

And you can maybe start to see the shapes of a steady state happening where we have some work from the previous step bleeding over into the next. And this is basically it. We're basically now in the end game. Our steady state pipeline will continue for our extent of our linear scans. And honestly, there's a lot of detail here. Um but you can see the shape of the

pipeline. I just I pretty much just want to focus in on the steady state. Cuz that's what we work so hard to get to. If we look at any given column, you'll see that at least the top row and the bottom row is occupied. And that's really the name of the game. And that means that we're using the tensor core at every given step of our pipeline.

And again, we bought these expensive GPUs. We want to make them to go burr. Okay. Hopefully it motivates why this very complicated pipeline we need to go a different level a lower level language to support this. And we decided to coalesce on flash attention V4. And this is implemented in Q T S L. And I want to talk about how we lower the QTSL for this pipeline.

So, we start with our user-written Python PyTorch score mod. We use PT2 to lower this user PyTorch code in a different level of IRs and extract out individual graphs. And one nice thing about FlexAttentions API is we know that score mods and mass mods are going to be pointwise. And so, there's a limited set of operations we ultimately need to support. And so, we've converted our FX

to Inductor IR. We're going down the lowering process and we now need to get into a form that QTSL can understand. And QTSL actually has a similar notion for register-backed tensors and this is where we actually apply the score and they call it tensor SSA. And so, we built an Inductor pipeline for for converting arbitrary pointwise Inductor IR expressions into tensor tensor SSA form. And this is

one of the key features of QTSL that actually made this integration possible, the way we designed it. And it's the ability for QTSL invocations either jit functions to also consume user callables. And so, the Inductor the Inductor lowering converts arbitrary user implementations into QTSL form and we pass them directly to the actual kernel. So, we had the pipeline, but we actually need to wire it up into

the kernel, right? We know there's two hook points. And this meant extending FA4 with hooks for score mods, for mass mods, implementing a block sparse implementation that also supports arbitrary loads in these callables. And with that, normal PyTorch code could drive custom behavior inside FA4. And we built this out for forward and we reuse much of the machinery for backward. But for this talk, I'm focusing on

the forward just because it's a little simpler to talk about. And okay, so numbers. I'm going to show you these numbers. Usual caveat that benchmarking is hard and it's not always portable, but I have links in our blog post at the end of this talk. You can see our full methodology and you can produce these numbers yourself. Um as you can see on Hopper, our new Flex

Flash implementation was able to greatly improve performance. And on Blackwell, the wins are even more compelling. This is again was the main motivation for this redesign. It's not perfect, right? If you look at the lower left at the document masking, you can see there actually cases where the Triton implementation is more performant. And these are for essentially memory bound sizes. But for large compute bound shapes, we're

seeing impressive speedups of up to 3x. Okay, so I hope I convinced you that if you're using Flex attention already, you should. But if you aren't, you should. And if you're on Hopper and Blackwell, you should try this new Flex Flash backend. Um the great part it's essentially a one-liner. You pass the Flash backend to your kernel options. And if you're on Blackwell, you do need to

change the default block spars from 128 to 128 to 256 by 128. And again, this is actually motivated as we saw in the pipeline from the two-stage chunking we're doing. But overall, it's the same Flex attention API you know and love. We're not done yet. Uh there's still a lot of active work. So these are two things I've been working on that I'm excited about. Left is

a deterministic backward. We've seen lately that determinism is very important for big pre-training runs. Um and then the right is CLC based work scheduling. This actually has landed since I started working on these slides. Um but this really helps performance on skewed workloads. Um yeah, okay. And then I'd love you if you want to learn more, we have the attention gym, lots of examples, full docs. Um

and on the right is the link to our paper on the PyTorch blog. Um and all the full methodology and details are there. That's it.