Finally, we compute the actual output by reducing over all the splits, using the log-sum-exp to scale the contribution of each split.Īll of this is possible because the attention/softmax can be calculated iteratively.We also write 1 extra scalar per row and per split: the log-sum-exp of the attention values. We compute the attention of the query with each of these splits in parallel using FlashAttention.First, we split the keys/values in smaller chunks.Like FlashAttention, it stores very little extra data to global memory, however it fully utilizes the GPU even when the batch size is small, as long as the context length is large enough.įlash-Decoding also parallelizes across keys and values, at the cost of a small final reduction step It combines the benefits of the 2 approaches from above. Our new approach Flash-Decoding is based on FlashAttention, and adds a new parallelization dimension: the keys/values sequence length. A faster attention for decoding: Flash-Decoding In this case, the operation occupies the GPU entirely, but launches many kernels that write and read intermediate results, which is not optimal. The attention can also be done using matrix multiplication primitives - without using FlashAttention. With a batch size of 1, FlashAttention will use less than 1% of the GPU!įlashAttention parallelizes across blocks of queries and batch size only, and does not manage to occupy the entire GPU during decoding During inference, the query length is typically 1: this means that if the batch size is smaller than the number of streaming multiprocessors (SMs) on the GPU (108 for an A100), the operation will only use a small part of the GPU! This is especially the case when using long contexts, because it requires smaller batch sizes to fit in GPU memory. For training, FlashAttention parallelizes across the batch size and query length dimensions. However, these optimizations don’t apply directly to the inference case, because the bottlenecks are different. This operation has been optimized with FlashAttention (v1 and v2 recently) in the training case, where the bottleneck is the memory bandwidth to read and write the intermediate results (e.g. Multi-head attention for decodingĭuring decoding, every new token that is generated needs to attend to all previous tokens, to compute: The main idea is to load the keys and values in parallel as fast as possible, then separately rescale and combine the results to maintain the right attention outputs. We present a technique, Flash-Decoding, that significantly speeds up attention during inference, bringing up to 8x faster generation for very long sequences. This is because the amount of memory to read scales with the batch dimension, whereas it only depends on the model size for the rest of the model. When scaling on the batch size dimension, the attention can also become a bottleneck even with relatively small contexts. In this setting, attention takes a significant fraction of time during inference. As an example, most LLMs had a context length of up to 2k in 2022 (GPT-3), but we now have open-source LLMs scaling up to 32k ( Llama-2-32k), or even 100k more recently ( CodeLlama). With a longer context, LLMs can reason about longer documents, either to summarize or answer questions about them, they can keep track of longer conversations, or even process entire codebases before writing code. There are a number of important emerging use cases of LLMs that utilize a long context. This operation does not scale well with context length. Fortunately, it is possible to cache previously calculated tokens: this means that a single generation step does not depend on the context length, except for a single operation, the attention. Generating full sentences of N tokens requires N forward passes through the model. LLM inference (or “decoding”) is an iterative process: tokens are generated one at a time. As LLM applications multiply, even small efficiency gains to the generation time can have a massive impact. Some use cases are more expensive, like code auto-completion, because it runs whenever a new character is typed. Even though generating a single response can cost about $0.01 (a few seconds of an 8xA100 instance on AWS), the costs quickly add up when scaling to billions of users, who could have multiple daily interactions with such LLMs.

However, they remain massively expensive to run. Large language models (LLM) such as ChatGPT or Llama have received unprecedented attention lately. Tri Dao, Daniel Haziza, Francisco Massa, Grigory Sizov

Flashcode decoder