Author(s): Kailash Thiyagarajan Originally published on Towards AI. Accelerating AI: A Deep Dive into Flash Attention and Its ImpactsImage Generated by AuthorIntroductionTransformers, introduced in the groundbreaking paper Attention Is All You Need, have revolutionized artificial intelligence, particularly in natural language processing and image classification. At the core of this success is the attention mechanism, which enables models to dynamically focus on relevant parts of the input. However, as transformers grow larger and deeper, the attention mechanism faces significant computational bottlenecks, especially with long input sequences.ProblemThe self-attention module in transformers has a time and memory complexity that scales quadratically with sequence length, making it challenging to handle long contexts. While methods like sparse and low-rank approximations aim to reduce computational costs, they often overlook memory access overheads, limiting practical speedups. As modern GPUs have advanced in compute speed more than memory speed, memory access remains a critical bottleneck. IO-aware algorithms have optimized memory-bound tasks in other fields, but deep learning frameworks like PyTorch and TensorFlow lack fine-grained memory control.OverviewIn this article, we will explore the FlashAttention mechanism and its approach to addressing these challenges. We will examine GPU requirements and demonstrate its implementation with a short code example.Key-Value Attention: The Backbone of TransformersHaving outlined the challenges faced by traditional attention mechanisms, we now turn our focus to the core component that underpins these models: key-value attention. This mechanism is crucial for enabling transformers to efficiently process and prioritize information within input sequences. Lets delve into how key-value attention operatesThe process involves three matrices:Query (Q): What were looking for.Key (K): Where to find it.Value (V): The information we care about.For simplicity, The attention mechanism can be broken down into three steps:Compute the similarity between queries and keys to generate a score matrix (S).Apply the softmax function to turn scores into probabilities.Multiply the probabilities with the values to get the final output.For a detailed explanation, refer to my previous article on Improving LLM Efficiency. Each step involves reading and writing data to high-bandwidth memory (HBM), which can slow down the attention process.Image Source: HuggingFaceThe GPU Memory Pyramid: A Balancing ActA GPU has three main types of memory, and each has different strengths and weaknesses. Think about them as layers in a pyramid, with very marked trade-offs between speed and capacity:Image Source: Flash Attention arxivImage Source: AuthorThe faster the memory, the smaller it gets, creating a classic computing trade-off. While SRAM is ideal for speed, its limited capacity means we often lean on HBM for larger tasks.The Problem with Moving Everything to SRAMWhy not simply perform all these operations in SRAM, the fastest memory layer? While it sounds like a great idea in theory, there are some significant hurdles:Size Limitations:The score matrix S alone can be massive, especially for long sequences or high-dimensional embeddings. A single computation could easily exceed the 20 MB capacity of SRAM.Complexity of Access:Even if we could split data into smaller chunks to fit into SRAM, the frequent movement of these chunks in and out of SRAM would lead to inefficiencies, defeating the purpose of using fast memory.Energy Costs:SRAM is designed for speed, not for handling the repeated, large-scale memory operations required by attention mechanisms. Constantly managing this flow would drain computational resources.As a result, we rely on HBM for these operations, despite its slower speed and higher latency compared to SRAM.Enter Flash Attention: A Smarter Way to WorkFlash Attention changes the game by rethinking how attention computations are performed. Instead of relying on HBM for storing large intermediate results, Flash Attention cleverly restructures the process to make the most of SRAMs speed.Heres how it works:Tiling the Computation:Tiling is a technique used to optimize matrix multiplication by breaking down large matrices into smaller sub-matrices or tiles. This approach enhances performance by improving cache usage and reducing memory bandwidth requirements.Image Source : Tiled Matrix MultiplicationSteps in TilingDivide Matrices into Tiles: The large matrices (Share A and Share B) are divided into smaller blocks or tiles. In this example, each matrix is divided into four smaller 22 tiles.Multiply Tiles: Each tile from Share A is multiplied with the corresponding tile from Share B. This multiplication is done independently for each pair of tiles, resulting in a temporary result (Temp).Accumulate Results: The temporary results from each tile multiplication are accumulated into the final result matrix. This is done by adding the Temp results to the corresponding positions in the Result matrix.Repeat for All Tiles:The process is repeated for all combinations of tiles until the entire matrix multiplication is complete.In Simple terms, Imagine you have two big grids of numbers (Share A and Share B), and you want to multiply them to get a new grid (Result). Doing this all at once can be slow and use a lot of memory.Tiling is like cutting these big grids into smaller squares (tiles). You multiply each small square from Share A with a matching square from Share B to get a small result (Temp). Then, you add up all these small results to get your final big grid (Result).This method is faster because:It works with small pieces at a time, which fits better in computers fast memory.It can do many small multiplications at the same time if you have multiple processors.Overall, tiling makes the multiplication of large matrices more efficient by breaking the task into smaller, more manageable pieces.2. Online Softmax:The term online in online softmax refers to the process of computing the softmax in a streaming or incremental manner, rather than computing it all at once. This is particularly useful for handling large sequences that might not fit into memory if processed in a single batch.Chunking: The input sequence is divided into smaller chunks that can be processed independently. This reduces memory usage because you only need to keep a small part of the sequence in memory at any given time.Incremental Computation: As each chunk is processed, the softmax is computed for that chunk, and intermediate results are stored. This allows the algorithm to build up the final result incrementally.Numerical Stability: Online softmax can be designed to maintain numerical stability by carefully managing the range of values during the computation. This is crucial because exponentiating large numbers can lead to overflow, and small numbers can lead to underflow.Efficiency: By processing data in chunks and only keeping necessary intermediate results, online softmax can significantly reduce the computational overhead and memory footprint, making it suitable for large-scale applications.In Simple terms, Imagine you have a long list of numbers, and you want to convert them into a list of probabilities. Normally, youd take all the numbers at once, do the math (exponentiation and division), and get your probabilities. But if your list is really long, this can be slow and use a lot of memory.Online softmax is like breaking that long list into smaller, manageable pieces. You do the math on each piece one at a time, and then combine the results. This way, you never have to deal with the whole list at once, which saves time and memory.In the context of FlashAttention, this means you can handle really long sequences of data more efficiently, which is great for tasks like natural language processing where you often deal with long texts.3. Weighted Sum of Values:Apply Attention Weights: Use the probabilities from the online softmax to weight the corresponding value vectors.Sum Weighted Values:For each query, compute the weighted sum of the value vectors. This gives the final attention output for each query.In Simple terms, Think of it like using the probabilities to decide how much of each piece of information (value vectors) to include in the final answer. You combine these weighted pieces to get the result for each part of your input.The result? Faster computations, lower energy costs, and the ability to handle larger models without hitting memory bottlenecks. Putting it all together , here is the representation from the original paperImage Source: Flash Attention arxivWith a clear understanding of Flash Attention, lets now take a closer look at its next evolution: Flash Attention v2.Diving into Flash Attention v2:Flash Attention v2 is an improved version of the original Flash Attention algorithm, designed to further optimize the memory and computational efficiency of transformer models. It introduces advanced techniques for multi-query and grouped attention, making it suitable for both inference and training at scale. By efficiently leveraging SRAM through better tiling and streamlining operations, Flash Attention v2 minimizes memory bottlenecks and improves throughput, especially for large models.Key GPU Requirements of Flash AttentionTensor Core Support: Flash Attention heavily relies on Tensor Cores to perform efficient mixed-precision computations such as FP16 or BF16. Tensor Cores were introduced in the NVIDIA Volta architecture (V100) and have been improved in subsequent generations.Warp-Level Primitives: Flash Attention leverages warp-level parallelism in GPUs, which is optimized in NVIDIA GPUs starting from Turing (T4) and beyond.Image: AuthorBF16 is Generally Optimized for Training/Inference:Larger Dynamic Range: Matches FP32 for handling extreme values, crucial for stable computations in attention mechanisms.Efficiency Without Accuracy Loss: Reduces precision minimally while maintaining model accuracy, enabling faster computations.Hardware Support: Modern GPUs (e.g., NVIDIA A100) are optimized for BF16, enhancing throughput for training and inference.Image FP16 vs BF16FP16 vs. BF16 RepresentationFP16 (16-bit Floating Point): Uses 1 bit for the sign, 5 bits for the exponent, and 10 bits for the mantissa. It has a smaller range and precision but is efficient for neural network computations.BF16 (Brain Floating Point): Uses the same 8-bit exponent as FP32 but reduces the mantissa to 7 bits. It provides a larger range and better compatibility with FP32 while maintaining efficiency.Example: Evaluating Flash Attention with Metas Llama ModelNow that weve covered the fundamentals of Flash Attention and its evolution to Flash Attention v2, lets shift gears to see it in action. To truly understand its performance benefits, well walk through an example using the Meta-LLaMA model on the Orca Math Word Problems dataset. This practical demonstration will highlight how Flash Attention v2 improves efficiency and scalability, especially when leveraging mixed precision (FP16/BF16) for inference. So, lets dive into the code to explore this in more detail.Follow this with the code block for seamless engagement.The dataset is limited to 10,000 entries and pre-trained model for faster experimentation. The evaluation focuses on runtime performance for generating text outputs.Heres the code:The performance gains from using FlashAttention with BF32 compared to other mechanisms are more noticeable with larger model sizes and higher batch sizes. These conditions allow the optimizations in memory access and computation to be fully leveraged. Additionally, factors such as sequence length, input data characteristics, and hardware capabilities play a crucial role in realizing these benefits. By optimizing these parameters, you can better exploit the efficiencies offered by FlashAttention in large-scale applications.I hope you found this article useful, and if you did, consider giving claps. #llm #flash-attn #optimization #kv-cacheReferences:FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTransformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention arearxiv.orgJoin thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming asponsor. Published via Towards AI