7 ms·
FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision
- edude03 2y agoHow much is the flash attention algorithm tied to the hardware? For example, in this announcement they mention taking advantage of the async capabilities of the H100 GPUs which I assume means you won't get those speedups on non H series card. Two, the actual flash attention library requires CUDA, although the algorithm has apparently?[^0] been ported to metal. I would imagine if the algorithm was literally just a pure function it could be implemented for any GPU/ML framework? [0]: https://github.com/philipturner/metal-flash-attention https://github.com/philipturner/metal-flash-attention
- kristjansson 2y agoFlashAttention's algorithmic improvements is mostly just splitting/combining the softmax part of attention, and is itself not totally novel. The overwhelming contribution is implementing that, and all its fiddly pieces, efficiently on Nvidia hardware.
- namibj 2y agoTo clarify further, flash attention is explicitly targeting a compute engine with separate MMA and "scalar" vector execution units that allow post-processing the MMA outputs without involving memory bandwidth (though arithmetic intensity, especially relative between the MMA and the "scalar" instructions, is of concern), with a substantial amount of manually-managed L1D$ to use as sub-matrix accumulator, and a linear-in-context-length amount of "VRAM" that requires sensible arithmetic intensity to avoid being a bandwidth bottleneck (iirc in the hundreds when counting the scalar multiplies hiding in the MMA instructions). 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).
- refulgentis 2y agoClarifying: 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.
- kristjansson 2y agoI think the right analogy for FA is high-quality cache-aware BLAS kernel implementations. The algorithm(s) is (are) clever and (as you note) completely independent of hardware. However, a hardware-naive implementation is approximately worthless. Most of the value of MKL, or Accelerate, or FA is in the careful matching of the parameters and implementation of the algorithm to the capabilities of hardware it's going run on. 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. [0]: https://web.stanford.edu/class/cs345d-01/rl/olap.pdf https://web.stanford.edu/class/cs345d-01/rl/olap.pdf
- kiratp 2y ago> However, a hardware-naive implementation is approximately worthless. 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.
- saagarjha 2y agoSupporting A100 and H100 is the opposite of being hardware naive, though.
- f_devd 2y ago> How much is the flash attention algorithm tied to the hardware? 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). [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk https://hazyresearch.stanford.edu/blog/2024-05-12-tk
- kristjansson 2y agoI mean they're building an API to abstract away some of the SKU-to-SKU differences, but the broader point cuts the other way, I think: > 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?
- eyegor 2y agoWhat do you mean by "relatively universal"? This is Cuda only [0] with a promise of a rocm backend eventually. There's only one project I'm aware of that seriously tries to address the Cuda issue in ml [1]. [0] https://github.com/HazyResearch/ThunderKittens?tab=readme-ov-file#installation https://github.com/HazyResearch/ThunderKittens?tab=readme-ov... [1] https://github.com/vosen/ZLUDA https://github.com/vosen/ZLUDA
- f_devd 2y agoIf you read the article I linked they show that it's entirely based on 16x16 matrices (or "tiles") which is fairly standard across gpus.
- vhiremath4 2y agoThere are a bunch of good answers, but I wanted to succinctly say "practically, quite a bit". Here's a good little rabbit-hole example: > https://github.com/karpathy/nanoGPT/blob/master/model.py#L45 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.functional.scaled_dot_product_attention.html 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 https://arxiv.org/pdf/2307.08691 The paper for FA2 almost entirely considers itself through the hardware it's running on.
- 3abiton 2y agoTo add to the discussion, from a practical perspective, AMD hardware totally sucks and yet to have proper implementation with flash-attention-2. ROCm is moving to usable slowly, but not close to being even comparable with cuda.
- slashdave 2y agoConceptually, just a bit, practically (in terms of implementation), a lot. The standard python implementation internally compiles a kernel for your specific hardware.
- lxe 2y ago> FlashAttention-3 is optimized for Hopper GPUs (e.g. H100). How does FA3 fare for consumer GPUs such as 3090 and 4090?
- apsec112 2y agoIt's Hopper-specific, the improvements are closely tied to Hopper features like warp groups and TMA. For 4090s, you might get a speedup by using the Triton implementation of FP8 attention: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html https://triton-lang.org/main/getting-started/tutorials/06-fu...
- moffkalast 2y agoThe original flash attention (v1?) took like a year to get added to llama.cpp and only provides single digit percent VRAM savings for typical context lengths and practically no speed boost. Still nice to have, but man was this thing overhyped. I doubt v3 will do more than marginally better on the RTX 5000 series.
- apsec112 2y agoOn GPU, or on CPU/Metal? For the latter I'm not surprised, but that's because they have a totally different memory/cache hierarchy.
- moffkalast 2y agoWith CUDA offloading, I don't think it runs otherwise at all.
- Der_Einzige 2y agoThis is one of the most important improvements in all of AI, because it benefits most AI users by giving them access to more, faster, for the same hardware with little to no tradeoffs.
- snovv_crash 2y ago...for all those users with H100s.
- rfoo 2y ago... which is currently the most cost-efficient and environment-friendly way to do LLM inference [0]. [0] Small footprint time: before B100 ships; for actually large language models; for prefill only; may cause cancer in California.
- stale2002 2y agoIndeed. 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.
- andy_xor_andrew 2y agohoping an expert can answer a few Qs I have :) 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".
- zaptrem 2y ago> 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? 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.
- apsec112 2y ago1) Pretty much, it's mathematically equivalent. The only software issues are things like managing dependency versions and data formats in-memory, but Flash Attention 2 is already built into HuggingFace and other popular libraries. Flash Attention 3 probably will be soon, although it requires an H100 GPU to run 2) Flash Attention 2 added support for GQA in past version updates: https://github.com/Dao-AILab/flash-attention 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-fused-attention.html https://triton-lang.org/main/getting-started/tutorials/06-fu...
- localfirst 2y agospoiler: $xxx,xxx hardware required to run
- WanderPanda 2y agoCompiler folks: Is there any chance compilers will be able to find optimizations like FlashAttention on their own? Seems like TVM and tinygrad are working in that direction but I find it hard to believe that that would be feasible
- rfoo 2y agoNo. Think of it like a different algorithm. You just take the shape of the hardware into consideration when designing the algorithm instead of considering math only. > 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?
- dauertewigkeit 2y agoWhat's the distinction between what TVM does and FlashAttention type optimizations?
- rfoo 2y agoThere is more than layout / tile schedule in FA. For example, first, to be able to fuse all these together [0] at all, you need to "decompose" the softmax to make it combinable, which requires maintaining some extra statistics. Won't gonna repeat the math here as the original FA paper is already very clear. [0] so you can avoid materializing intermediate matrices and still being able to compute in blocks.
- FL33TW00D 2y agoGeo has explicitly stated he wants to be able to find FA in the search space of algos eventually. Actually achieving that is another matter.
- namibj 2y agoIn theory, yes, it's "just" some algebraic properties of the math used that allow for substantial reordering, and then you'd add fairly regular polyhedral loop tiling. Just expensive to do, so you'll have to cache the effort. 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!)).
- ex3ndr 2y agoI am wondering why flash attention is like 5x slower with variable masking than without it? Lack of good masking support almost zeros out the optimizations
- chillee 2y agoWhere are you seeing these benchmarks?
- cajesaw4 2y ago[flagged]
- latchkey 2y agoIf anyone wants to port this over to ROCm / AMD MI300x, reach out to me: hello@hotaisle.xyz (we won't ever spam you). Happy to donate the compute time for this work.
- saagarjha 2y agoNot trying to be rude but what is the thinking behind this offer? Why would someone do this port for…free save for access to the hardware? What’s the upside for them?
- latchkey 2y agoNot a rude question. I'm building public HPC super computers, currently focused on AMD hardware. The one I'm about to deploy is Top 150, which is a pretty good start. 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.
- gitfan86 2y agoNo one person or one company SHOULD have huge control over humanity, I agree. 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.
- latchkey 2y ago> 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. [0] https://x.com/realsharonzhou/status/1811439958277927294 https://x.com/realsharonzhou/status/1811439958277927294
- refibrillator 2y agoThe code has a comment which seems to hint that Tri Dao was working on FA3 as early as April 2022, the month after Hopper/H100 was announced. I find it mildly curious that over 2 years has elapsed before the code was released today. Perhaps it’s because now there’s better solutions in the pipeline? 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.
- nyrikki 2y agoUntil the strong exponential hypothesis is (dis-)proven, the quadratic cost is required or you have to give something up. Just the cost of exhaustive search. 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.
- MacsHeadroom 2y agoWhat makes you think that the thing you have to give up is related to model quality?
- saagarjha 2y ago> TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency. 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.
- LarsDu88 2y agoI was wondering... this post mentions that ops like sigmoid are very slow. 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?
- eigenvalue 2y agoRelu is literally just a linear function that gets clamped to zero at some point, so yes, it's much less computationally intensive than anything involving an exponential function. But I doubt you would get competitive results using such a simple activation.