【原】LLMs之FlashAttention-2：《FlashAttention-2: Faster Attention with Better Parallelism and Work Partition

地址

博客文章：FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning | Princeton NLP Group

時間

2023年7月17日

作者

谷歌學者Tri Dao，坦福大學計算機科學博士，Together.AI的首席科學家

Just within the last year, there have been several language models with much longer context than before: GPT-4 with context length 32k, MosaicML’s MPT with context length 65k, and Anthropic’s Claude with context length 100k. Emerging use cases such as long document querying and story writing have demonstrated a need for models with such long context. Scaling up the context length of Transformers is a challenge, since the attention layer at their heart has runtime and memory requirements that are quadratic in the input sequence length.

僅在過去的一年里，出現了幾種比以前上下文更長的語言模型：GPT-4的上下文長度為32k, MosaicML的MPT的上下文長度為65k, Anthropic的Claude的上下文長度為100k。新興的用例，如長文檔查詢和故事創作，已經證明了需要具有如此長上下文的模型。擴展Transformer的上下文長度是一個挑戰，因為它們核心的注意力層具有與輸入序列長度成二次方關系的運行時和內存要求。

A year ago, we released FlashAttention, a new algorithm to speed up attention and reduce its memory footprint—without any approximation. We’ve been very happy to see FlashAttention being adopted by many organizations and research labs to speed up their training & inference (see this page for a partial list). Even though FlashAttention was already 2-4x faster than optimized baselines at the time of its release, it still has quite a bit of headroom. FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40% of the theoretical maximum FLOPs/s (e.g. up to 124 TFLOPs/s on A100 GPU).

一年前，我們發布了FlashAttention，這是一種新的算法，可以加速注意力并減少其內存占用，而無需任何近似。我們很高興看到FlashAttention被許多組織和研究實驗室采用，以加快他們的訓練和推理（請查看此頁面以獲取部分列表）。盡管FlashAttention在發布時已經比優化后的基線快了2-4倍，但它仍然有相當大的提升空間。FlashAttention仍然不如優化的矩陣乘法(GEMM)操作快，僅達到理論最大FLOPs/s的25-40%(例如在A100 GPU上高達124 TFLOPs/s)。

In the past few months, we’ve been working on the next version, FlashAttention-2, that makes FlashAttention even better. Rewritten completely from scratch to use the primitives from Nvidia’s ?CUTLASS 3.x and its core library ?CuTe, FlashAttention-2 is about 2x faster than its previous version, reaching up to 230 TFLOPs/s on A100 GPUs. When used end-to-end to train GPT-style language models, we reach a training speed of up to 225 TFLOPs/s (72% model FLOP utilization). In this blogpost, we describe some of the bottlenecks of FlashAttention, and how we use better parallelism and work partitioning to get significant speedup.

FlashAttention-2 is available at: ?flash-attention

在過去的幾個月里，我們一直在研發下一個版本，即FlashAttention-2，使FlashAttention變得更加出色。FlashAttention-2完全重寫，使用Nvidia的CUTLASS 3.x和其核心庫CuTe的原語，比其前一個版本快大約2倍，在A100 GPU上最高可達230 TFLOPs/s。當用于端到端訓練類似GPT的語言模型時，我們達到了高達225 TFLOPs/s（模型FLOP利用率為72%）的訓練速度。在本文中，我們將描述FlashAttention的一些瓶頸，以及我們如何使用更好的并行性和工作分區來獲得顯著的加速。FlashAttention-2可在:flash-attention

FlashAttention is an algorithm that reorders the attention computation and leverages classical techniques (tiling, recomputation) to significantly speed it up and reduce memory usage from quadratic to linear in sequence length. Tiling means that we load blocks of inputs from HBM (GPU memory) to SRAM (fast cache), perform attention with respect to that block, and update the output in HBM. By not writing the large intermediate attention matrices to HBM, we reduce the amount of memory reads/writes, which brings 2-4x wallclock time speedup.

FlashAttention是一種重新排序注意力計算的算法，并利用經典技術(tiling、重新計算)來顯著加速它并將內存使用量從與序列長度成二次關系降低到與線性成正比的算法。tiling意味著我們從HBM（GPU內存）加載輸入塊到SRAM（快速緩存），針對該塊執行注意力計算，并在HBM中更新輸出。通過不將大型中間關注矩陣寫入HBM，我們減少了內存讀寫的數量，從而帶來2-4倍的時間加速。

However, FlashAttention still has some inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes.

然而，FlashAttention仍然存在一些低效性，這是因為在GPU上不同線程塊和線程束之間的工作分區不夠優化，導致低占用率或不必要的共享內存讀寫。

Diagram of FlashAttention forward pass: with tiling and softmax rescaling, we operate by blocks and avoid having to read/write from HBM, while obtaining the correct output with no approximation.

FlashAttention向前傳遞圖:使用tiling和softmax重新縮放，我們按塊操作，避免必須從HBM讀/寫，同時獲得正確的輸出，沒有近似。

We tweak the algorithm from FlashAttention to reduce the number of non-matmul FLOPs. This is important because modern GPUs have specialized compute units (e.g., Tensor Cores on Nvidia GPUs) that makes matmul much faster. As an example, the A100 GPU has a max theoretical throughput of 312 TFLOPs/s of FP16/BF16 matmul, but only 19.5 TFLOPs/s of non-matmul FP32. Another way to think about this is that each non-matmul FLOP is 16x more expensive than a matmul FLOP. To maintain high throughput, we want to spend as much time on matmul FLOPs as possible. We rewrite the online softmax trick used in FlashAttention to reduce the number of rescaling ops, as well as bound-checking and causal masking operations, without changing the output.

我們對FlashAttention算法進行了微調，以減少非矩陣乘法FLOPs的數量。這很重要，因為現代GPU擁有專門的計算單元（例如Nvidia GPU上的Tensor Cores），使矩陣乘法計算速度更快。例如，以A100 GPU為例，FP16/BF16的最大理論吞吐量為312 TFLOPs/s，而非矩陣FP32的最大理論吞吐量為19.5 TFLOPs/s。

另一種思考這個問題的方式是，每個非矩陣FLOP的成本是矩陣乘法FLOP的16倍。為了保持高吞吐量，我們希望盡可能多地使用矩陣乘法FLOP的時間。我們重新編寫了FlashAttention中使用的在線softmax技巧，以減少重新縮放操作的數量，以及邊界檢查和因果掩碼操作，而不改變輸出。

The first version of FlashAttention parallelizes over batch size and number of heads. We use 1 thread block to process one attention head, and there are overall (batch_size * number of heads) thread blocks. Each thread block is scheduled to run on a streaming multiprocessor (SM), and there are 108 of these SMs on an A100 GPU for example. This scheduling is efficient when this number is large (say >= 80), since we can effectively use almost all of the compute resources on the GPU.

In the case of long sequences (which usually means small batch sizes or small number of heads), to make better use of the multiprocessors on the GPU, we now additionally parallelize over the sequence length dimension. This results in significant speedup for this regime.

FlashAttention的第一個版本在批量大小和頭數上進行了并行化。我們使用1個線程塊來處理一個關注頭，總共有（批量大小*頭數）個線程塊。每個線程塊被安排在一個流多處理器（SM）上運行，例如A100 GPU上有108個這樣的SM。當這個數字很大時（例如>= 80），這種調度是有效的，因為我們可以有效地使用GPU上的幾乎所有計算資源。

對于長序列的情況（通常意味著較小的批量大小或較少的頭數），為了更好地利用GPU上的多處理器，我們現在還在序列長度維度上進行了額外的并行化。這將為這種情況帶來顯著的加速。

Even within each thread block, we also have to decide how to partition the work between different warps (a group of 32 threads working together). We typically use 4 or 8 warps per thread block, and the partitioning scheme is described below. We improve this partitioning in FlashAttention-2 to reduce the amount of synchronization and communication between different warps, resulting in less shared memory reads/writes.

即使在每個線程塊中，我們也必須決定如何在不同的warp(一組32個線程一起工作)之間劃分工作。

即使在每個線程塊內部，我們還必須決定如何在不同的線程束和線程塊之間分配工作。通常情況下，我們每個線程塊使用4或8個線程束，并且下面描述了分區方案。在FlashAttention-2中，我們改進了這種分區，以減少不同線程束之間的同步和通信量，從而減少了共享內存的讀寫。

For each block, FlashAttention splits K and V across 4 warps while keeping Q accessible by all warps. This is referred to as the “sliced-K” scheme. However, this is inefficient since all warps need to write their intermediate results out to shared memory, synchronize, then add up the intermediate results. These shared memory reads/writes slow down the forward pass in FlashAttention.

In FlashAttention-2, we instead split Q across 4 warps while keeping K and V accessible by all warps. After each warp performs matrix multiply to get a slice of Q K^T, they just need to multiply with the shared slice of V to get their corresponding slice of the output. There is no need for communication between warps. The reduction in shared memory reads/writes yields speedup.

對于每個塊，FlashAttention將K和V分配給4個線程束，同時讓Q對所有線程束可訪問。這被稱為“切片-K”方案。然而，這是低效的，因為所有線程束都需要將其中間結果寫入共享內存，進行同步，然后將中間結果相加。這些共享內存的讀寫減慢了FlashAttention的向前傳遞。

在FlashAttention-2中，我們改為將Q分配給4個線程束，同時讓K和V對所有線程束可訪問。每個線程束執行矩陣乘法以獲取Q K^T的一個切片，然后它們只需將其與共享切片的V相乘，即可獲得相應的輸出切片。不需要在線程束之間進行通信。減少共享內存的讀寫帶來了加速。

FlashAttention only supported head dimensions up to 128, which works for most models but a few were left out. FlashAttention-2 now supports head dimension up to 256, which means that models such as GPT-J, CodeGen and CodeGen2, and StableDiffusion 1.x can use FlashAttention-2 to get speedup and memory saving.

This new version also supports multi-query attention (MQA) as well as grouped-query attention (GQA). These are variants of attention where multiple heads of query attend to the same head of key and value, in order to reduce the size of KV cache during inference and can lead to significantly higher inference throughput.

FlashAttention僅支持頭維度高達128，這適用于大多數模型，但有些模型不能使用。FlashAttention-2現在支持頭維度高達256，這意味著模型如GPT-J、CodeGen和CodeGen2以及StableDiffusion 1.x可以使用FlashAttention-2來獲得加速和節省內存。

這個新版本還支持多查詢注意力（MQA）以及分組查詢注意力（GQA）。這些是注意力的變種，其中多個查詢頭注意力相同的鍵頭和值頭，以減小推斷過程中的KV緩存大小，可以顯著提高推斷吞吐量。

We measure the runtime of different attention methods on an A100 80GB SXM4 GPU for different settings (without / with causal mask, head dimension 64 or 128). We see that FlashAttention-2 is around 2x faster than FlashAttention (as well as its other implementations in the xformers library and in Triton). Compared to a standard attention implementation in PyTorch, FlashAttention-2 can be up to 9x faster.

我們在A100 80GB SXM4 GPU上對不同設置（無/有因果掩碼，頭維度64或128）上的不同關注方法的運行時進行了測量。我們發現FlashAttention-2比FlashAttention快大約2倍（以及xformers庫和Triton中的其他實現）。與PyTorch中的標準關注實現相比，FlashAttention-2的速度可以提高多達9倍。

Attention forward + backward speed on A100 GPU

Just running the same implementation on H100 GPUs (using no special instructions to make use of new hardware features such as TMA and 4th-gen Tensor Cores), we obtain up to 335 TFLOPs/s.

Attention forward + backward speed on H100 GPU

只要在H100 GPU上運行相同的實現（不使用特殊指令來利用新的硬件特性，如TMA和第四代張量核心），我們可以獲得多達335 TFLOPs/s的吞吐量。

When used to train a GPT-style model end-to-end, FlashAttention-2 helps achieve up to 225 TFLOPs/s on A100 GPU (72% model FLOPs utilization). This is a 1.3x end-to-end speedup over an already very optimized model with FlashAttention.

當用于端到端訓練GPT風格的模型時，FlashAttention-2可以在A100 GPU上實現高達225 TFLOPs/s的吞吐量（模型FLOP利用率為72%）。這相當于在已經經過優化的模型中使用FlashAttention獲得了1.3倍的端到端加速。

*Baseline: Megatron-LM without FlashAttention. Megatron-LM now has an option to use FlashAttention. We plan to integrate FlashAttention-2 to Megatron-LM in the near future.

*基線：Megatron-LM沒有使用FlashAttention。Megatron-LM現在可以選擇使用FlashAttention。我們計劃在不久的將來將FlashAttention-2集成到Megatron-LM中。

FlashAttention-2 is 2x faster than FlashAttention, which means that e.g. we can train models with 16k longer context for the same price as previously training a 8k context model. We’re excited about how this can be used to understand long books and reports, high resolution images, audio and video. FlashAttention-2 will also speed up training, finetuning, and inference of existing models.

In the near future, we plan to collaborate with folks to make FlashAttention widely applicable in different kinds of devices (e.g. H100 GPUs, AMD GPUs), as well as new data types such as fp8. As an immediate next step, we plan to optimize FlashAttention-2 for H100 GPUs to use new hardware features (TMA, 4th-gen Tensor Cores, fp8). Combining the low-level optimizations in FlashAttention-2 with high-level algorithmic changes (e.g. local, dilated, block-sparse attention) could allow us to train AI models with much longer context. We’re also excited to work with compiler researchers to make these optimization techniques easily programmable.

FlashAttention-2比FlashAttention快2倍，這意味著，例如，我們可以用與以前訓練8k上下文模型相同的價格來訓練具有16k更長上下文的模型。我們很高興看到它可以用來理解長篇書籍和報告、高分辨率圖像、音頻和視頻。FlashAttention-2還將加速現有模型的訓練、微調和推理。

在不久的將來，我們計劃與其他人合作，使FlashAttention在不同類型的設備（例如H100 GPU、AMD GPU）上廣泛適用，以及新的數據類型，如fp8。下一步，，我們計劃優化FlashAttention-2以適應H100 GPU，以利用新的硬件特性（TMA、第四代張量內核、fp8）。將FlashAttention-2中的低級優化與高級算法更改（例如局部、擴張、塊稀疏關注）結合起來，可以使我們能夠訓練具有更長上下文的AI模型。我們也很高興與編譯器研究人員合作，使這些優化技術易于編程。

We thank Phil Tillet and Daniel Haziza, who have implemented versions of FlashAttention in Triton and the xformers library. FlashAttention-2 was motivated by exchange of ideas between different ways that attention could be implemented. We are grateful to the Nvidia CUTLASS team (especially Vijay Thakkar, Haicheng Wu, and Andrew Kerr) for their CUTLASS library, in particular the CUTLASS 3.x release, which provides clean abstractions and powerful building blocks for the implementation of FlashAttention-2. We thank Driss Guessous for integrating FlashAttention to PyTorch. FlashAttention-2 has benefited from helpful discussions with Phil Wang, Markus Rabe, James Bradbury, Young-Jun Ko, Julien Launay, Daniel Hesslow, Micha?l Benesty, Horace He, Ashish Vaswani, and Erich Elsen. Thanks to Stanford CRFM and Stanford NLP for the compute support. We thank Dan Fu and Christopher Ré for their collaboration, constructive feedback, and constant encouragement on this line of work of designing hardware-efficient algorithms. We thank Albert Gu and Beidi Chen for their helpful suggestions on early drafts of the FlashAttention-2 technical report.

我們要感謝Phil Tillet和Daniel Haziza，他們在Triton和xformers庫中實現了FlashAttention的版本。FlashAttention-2的靈感來自于不同注意力執行方式之間的思想交流。我們非常感謝Nvidia CUTLASS團隊(特別是Vijay Thakkar, Haicheng Wu和Andrew Kerr)的CUTLASS庫，特別是CUTLASS 3。它為FlashAttention-2的實現提供了清晰的抽象和強大的構建塊。我們感謝Driss Guessous將FlashAttention集成到PyTorch。flashatten2得益于與Phil Wang, Markus Rabe, James Bradbury, Young-Jun Ko, Julien Launay, Daniel Hesslow, Micha?l Benesty, Horace He, Ashish Vaswani和Erich Elsen的有益討論。感謝斯坦福CRFM和斯坦福NLP提供的計算支持。我們感謝Dan Fu和Christopher r<s:1>在設計硬件高效算法方面的合作、建設性的反饋和不斷的鼓勵。我們感謝Albert Gu和Beidi Chen在FlashAttention-2技術報告的早期草稿中提出的有益建議。