FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision
together.ai
together.ai
Tri’s publication history has been leaning toward SSM and Mamba style architectures recently. Unlike Flash Attention which has quadratic time complexity wrt sequence length, these latest algorithms are subquadratic. Thus they do much less computation, instead of just doing it more efficiently a la Flash Attention.
Dao and Gu published a really long paper this year which demonstrated (among other things) how Mamba/SSM can be formulated such that it’s amenable to acceleration using the same hardware primitives that Transformers benefit from.
As (dis-)proving SETH will resolve the P vs NP problem, I wouldn't hold my breath.
The question is if a particular use case can accept those costs.
This v3 with async might for once be so tied to Hopper that it's not trivially portable to another platform that has the mentioned hardware blocks (AFAIK every AMD GCN card that can do compute shaders would qualify, though they do lack a specialized MMA unit).
Given the question: "How much is the flash attention algorithm tied to the hardware?"
The answer is 0.
ex. you can find generic flash attention recently added in llama.cpp and ONNX (MS needed it for Phi-3, needed for Recall).
On the side, novelty, I have no direct knowledge on, IMHO, asking that question would devolve the way novelty arguments do in any field: there's always someone else who can claim they did 80% of $X via $X-1, therefore, $X is by and large not novel. Ad infinitum.
I definitely don't mean to take away from Tri/FA by mentioning novelty - I'm just repeating from paper, which refers back to algebraic aggregates[0] in its discussion of their tiled softmax.
This isn’t true when there is one vendor that’s 90% of the market and 2 maybe 3 generations of hardware to consider. Support A100, H100 and you are supporting most of the current market.
The original FA, almost none.
For the latest versions depends on your abstraction, ThunderKittens[0] provides about the same speed up over FA2 (1.3x-2x%) as the article but relatively universal across GPUs. For any new hardware there may be hardware specific features that make it edge out more performance; usually vendors will adopt any new features that seems to beat them, but you do get fragmented API/libraries (which is already true for CUDA).
> In fact, more broadly we believe we should really reorient our ideas of AI around what maps well onto the hardware. How big should a recurrent state be? As big can fit onto an SM. How dense should the compute be? No less so than what the hardware demands. An important future direction of this work for us is to use our learnings about the hardware to help us design the AI to match.
The value is in adapting the implementation (either manually at write-time or programmatically at run-time) to the specifics of the hardware.
Also, great line:
> And we ask: if your matrix multiply is smaller than 16x16, are you sure what you’re doing is AI?
[0] https://github.com/HazyResearch/ThunderKittens?tab=readme-ov...
> https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
Karpathy's nanoGPT calling flash attention by checking if torch.nn.functional.scaled_dot_product_attention exists
> https://pytorch.org/docs/stable/generated/torch.nn.functiona...
Looking at the docs, in reality, most of the time you want this to call out to FA2 which optimizes the kernals on the device to split ops on the Softmax of the triangular matrix as well as reduce moving unnecessary batches of floating point numbers back and forth from the GPU to the CPU.
> https://arxiv.org/pdf/2307.08691
The paper for FA2 almost entirely considers itself through the hardware it's running on.
> Seems like TVM
Fair enough, though technically they are still about different things but it's indeed very close, but
> and tinygrad
?????? what gives you this impression?
[0] so you can avoid materializing intermediate matrices and still being able to compute in blocks.
The area of e-graph optimizers seems well-suited to this, btw. It's not really deployed outside of some niche tooling though, as it's a big paradigm shift in optimizer pass handling (e.g., doesn't work well with chairs classic call graphs, so control flow needs to be massively revamped to deploy e-graphs outside/across basic blocks and for loops (break and return not supported!)).
I would like to understand why you say e-graph would need control-flow to be revamped.
Do you have anything I could read on it ?
Happy to chat more btw, feel free to hit me up on discord.
I'm not sure what the state of the art in compiler optimisation is with regard to data positioning and targeting maximum processor usage
There was a video on optimisation a while back that showed small optimisations caused increases in speed that were insignificant when compared to the speed variance induced by the memory layout that the optimisation (or even a random change) caused.
While that talk was more focused on getting a signal past the noise. That noise itself is an artifact of compilers being not particularly good at handling a much simpler form of the problem you describe.
CPU and memory architectures are complex when caches and access patterns impact upon speed.
When you add in GPU architectures to the mix I think you might be in fairly uncharted territory.
Maybe one day.
Of course since we are in the field of AI there is also the question of could a sufficiently smart AI do this. It depends on the value of sufficient.
I would like to think that an extremely high level test for an AI model could be to give it something like micrograd and tell it to produce something with the same interface that outperforms torch.
We're not even in the ballpark of being able to do that yet, but it will be interesting when and if that happens.
Happy to donate the compute time for this work.
The goal is to encourage a developer flywheel. The more developers working with AMD hardware, the more hardware that is needed, the more hardware I can justify buying, the bigger my super computers get.
Nvidia has been doing the flywheel for years and it has clearly worked. Why not do the same for AMD? As I said in another thread, anyone who thinks that there should be a single provider for all of AI compute needs, will be on the wrong side of history.
But practically speaking this is a unique time in history of technology because there are quick feedback loops that cause that flywheel you mentioned to be a insurmountable first mover advantage.
I'm staking my career and business on you being wrong about the insurmountable part. This is just the beginning of a long road and I'm not the only one who believes this. My partnership with Dell, Advizex and a huge soon to be announced datacenter company, isn't small beans.
Much like how I didn't know how the internet would look when I first joined in 1991. But, what I can see very clearly from my decades of experience in the tech field, is that history is repeating itself with what is happening in AI.
As I'm also prone to say... this isn't a football match where one team needs to "beat" the other. It really is enough to have multiple players in the market and nothing more than that. In fact, I'm more than happy to deploy any type of compute that my customers want me to deploy for them, including Nvidia.
Even Lamini, whom were previously AMD only, just announced [0] that they are partnering with Nvidia. Their software will run equally well on any system. Why? Because it builds a simple bridge from one platform to the next. Reminds me of the Java "write once, run anywhere" slogan. It actually worked pretty well.
I'm saying NVDA uses their own chips to help build more chips, AND they are intricately involved in the buildout of the 100B data centers and intricately involved in TSMC roadmaps. That with the combination of huge profits that are increasing create even more advantages over competitors.
Obviously this doesn't go on forever, NVDA will never have 100T of profit in a quarter. Years from now the feedback loops will have diminishing returns and there will be commodity AI systems eventually.
> Years from now the feedback loops will have diminishing returns and there will be commodity AI systems eventually.
Maybe, but the cat is out of the bag. Before it was a question of Moore's law and speed, but nobody talks about that anymore... all they talk about is that the need for raw compute (not even the fastest), is officially boundless.
How does FA3 fare for consumer GPUs such as 3090 and 4090?
My understanding was that while it frees up registers it more importantly lets the hardware handle address generation, which can become a bottleneck as other operations around it become faster.
[0] Small footprint time: before B100 ships; for actually large language models; for prefill only; may cause cancer in California.
Anyone who is doing anything important or at scale would be at least renting those, or even using an abstracted service that is on top of another service.
Those cost savings allow people to train things for cheaper, causing those cost savings to benefit almost everyone doing important stuff in the space.
Is FlashAttention simply a drop-in replacement for the attention operation in an LLM? Can it be used anywhere that an "attention" operation is used? Or does a LLM need to be trained specially to use FA?
How does FA relate to attention strategies like GQA (grouped query attention) or sliding-window attention? Are they orthogonal concepts? Or you need a specific FA implementation for each strategy?
Recently llama.cpp added flash attention support - does this just mean they started consuming a flash attention-provided CUDA kernel or something?
lastly, in this post, they compare FlashAttention to Triton. I thought Triton was like an abstraction layer? Couldn't FA be implemented in Triton? I just don't really get what it means to say "FlashAttention vs. Triton".
Yes
> How does FA relate to attention strategies like GQA (grouped query attention) or sliding-window attention? Are they orthogonal concepts? Or you need a specific FA implementation for each strategy?
Flash Attention is a way of calculating the Softmax(QK^T)V part of attention, whereas GQA is a way of calculating the Q, K, and V matricies. Sliding window attention (less sure about this, there are a bunch of windowed attention techniques) change the attention mask (the thing that controls which queries can attend to which keys).
> Recently llama.cpp added flash attention support - does this just mean they started consuming a flash attention-provided CUDA kernel or something?
I don't use llama.cpp but that sounds about right.
> lastly, in this post, they compare FlashAttention to Triton. I thought Triton was like an abstraction layer? Couldn't FA be implemented in Triton? I just don't really get what it means to say "FlashAttention vs. Triton".
They're talking about a previous Flash Attention implementation written in Triton.
2) Flash Attention 2 added support for GQA in past version updates:
https://github.com/Dao-AILab/flash-attention
3) They're comparing this implementation of Flash Attention (which is written in raw CUDA C++) to the Triton implementation of a similar algorithm (which is written in Triton): https://triton-lang.org/main/getting-started/tutorials/06-fu...
A lot of modern LLMs use activation functions with sigmoid or soft max like SiLU, Swish, and SOLU.
Does Relu take less of a performance hit, and if so, maybe it'd be better to go back to good old relu?