FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Summary

FLASH ATTENTION sped up attention on GPUs but achieved only 35% utilization on H100 GPUs. To address this, FLASH ATTENTION-3 introduces three techniques for Hopper GPUs: warp-specialization to overlap computation and data movement by assigning warps to producer and consumer roles; interleaving operations to combine block-wise matrix multiplication (matmul) and softmax operations for asynchronous execution; and support for FP8 precision using block quantization and incoherent processing to leverage hardware capabilities. These techniques result in a 1.5-2.0× speedup with FP16, up to 1.2 PFLOPs/s with FP8, and 2.6× lower numerical error than baseline FP8 attention, improving computational efficiency and hardware utilization on modern GPUs.

Strengths

The strengths of the paper are that it is well presented and shows good results. The paper addresses the important topic of accelerating Transformer modules by presenting an incremental improvement over previous work, specifically FLASHATTENTION-2, and applies these advancements to the new H100 GPU. The paper demonstrates how the enhanced methods leverage the capabilities of the H100 GPU to achieve significant performance gains, showcasing advancements in computational speed and efficiency. This progression highlights the continuous evolution in optimizing Transformer architectures for cutting-edge hardware, emphasizing the relevance and impact of the research in the field of deep learning and hardware acceleration.

Weaknesses

The work primarily focuses on the H100 GPU, demonstrating the performance enhancements of FLASHATTENTION-3 for this hardware. However, evaluating how FLASHATTENTION-3 performs across other GPU families would increase the work's impact, showing its versatility and potential for wider adoption in different computing environments.

Questions

A few question and notes are listed below: How would the methods are applicable to other GPU devices? Is speed is slower for flashattention3 in small sequence len only seen in H100? In 4.1: '4·seqlen2 ·head dimension·number of heads'. would be better to format as equation

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations and, if applicable, potential negative societal impact of their work

Summary

The paper builds on the existing work on FlashAttention, and introduces an advanced method to accelerate the attention mechanism in Transformer model. It specifically targets the newer GPU architectures, like NVIDIA H100, by exploiting asynchrony in Tensor Cores and Tensor Memory Accelerators. There are three techniques that are proposed: 1) producer-consumer asynchrony across warps; 2) compute asynchrony within consumer warps; 3) hardware accelerated low precision GEMM. The extensive empirical results show impressive performance, with significant speedups compared to state-of-the-art methods, and with reduced numerical errors.

Strengths

The authors present novel techniques that exploit hardware features like asynchrony and low-precision computation. The performance gains are impressive, in speed and accuracy. The empirical validation includes comprehensive benchmarks and comparisons with previous methods. It is nice to see the open source commitment.

Weaknesses

The proposed methods are highly complex, and may be difficult to implement and understand for practitioners. The techniques are tailored to NVIDIA Hopper GPUs, and may have limited applicability to other hardware achitectures.

Questions

The backward pass algorithm is only included in the appendix. It would be useful if it could be included in the main paper. Can you comment if FlashAttention-3 could be used on older GPU architectures that lack some of the advanced features of the Hopper GPUs?

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

Yes.

Summary

This is a system paper that introduces FlashAttention-3, an optimized version of FlashAttention for NVIDIA's SM90+ GPUs. The key contributions (summarized by the authors) are: - Taking advantage of warp specialization in NVIDIA SM90+ and designing a producer-consumer computation paradigm to allow better intra-CTA overlapping between producers and consumers; - Designing 2-stage software pipeline to allow better overlapping within consumers; - Designing on-the-fly layout transformation to support FP8 WGMMA and using incoherent processing to reduce quantization error by 3x.

Strengths

- The empirical results of this paper are very strong, and are expected to generate immediate impact to the community after code release (promised by the authors). - The presentation of this paper is clear. It is easy for people who are familiar with NVIDIA SM8x but not experts in SM90+ to follow the writing and essence of the optimizations. I learned a lot from the descriptions related to the warp specialization design. - This is the first (to-be) open-sourced system paper that discusses the challenges and solutions related to FP8 attention, which has been implemented and close-sourced by NVIDIA for a long time. The insights related to FP8 such as layout transformation and reducing quantization error are very helpful for the readers.

Weaknesses

[Results, Benchmarking] - There has already been a paper from Colfax (arXiv 2312.11918) available online more than half a year ago describing how to implement efficient FlashAttention on Hopper GPUs. It will be better if the authors could cite this paper and compare the performance of FlashAttention-3 and this paper. - NVIDIA's TensorRT-LLM has a close-sourced implementation of FP8 attention. The result comparison with this library is missing. - While the authors mention RMSE reduction from incoherence processing using simulated data, there is no evaluation on how this step improves real LLM inference workloads (such as Wikitext perplexity, MMLU accuracy, etc.) - How many end-to-end LLM inference / training speedup can we obtain by switching to this FP8 attention? For inference, I am interested in a TensorRT-LLM-like implementation with FP8 GEMM + FP16 attention (FlashAttention-3) VS FP8 GEMM + FP16 attention (FlashAttention-2) VS FP8 GEMM + FP8 FlashAttention-3. [Methodology, Novelty] - While all the methodologies mentioned in the paper make sense to me, it is important to know that some of them occurred in previous literature / open source implementation on the same workload. For example, FlashInfer from University of Washington implements within-consumer software pipelining on SM80+ devices. Incoherent processing is also not new, and has been adopted in related work such as QUIP#, QuaRot, etc. - The authors mentioned that a lot of operations can be fused into the preceding RoPE operation **with no additional overhead**. For example, transposing V, applying Hadamard transformation, block quantization, etc. While I agree that intuitively computation operations could be fused with memory bound Ops, it is also important to notice that the CUDA cores of H100 are not that strong. I want to see real measured numbers in end-to-end workloads that justify the authors' claims on "no additional overhead".

Questions

Please address my comments in the Weaknesses section. Additional questions related to ablation studies: - Will it be also possible to open source the benchmarking code in Table 2? This will be super helpful for system researchers to understand the impact of warp specialization and 2-stage computation pipeline on the end-to-end throughput in different workloads. - Please comment why your implementation achieves performance advantage over OpenAI Triton ones. What is the most important design space extension that leads to this improvement? - Please comment on how the register re-allocation feature provided by NVIDIA SM90+ impacts the performance. Are there any important design spaces enabled by this feature which would otherwise lead to register spilling to local memory in SM8x? - Please comment on whether GEMM-SoftMax 2-stage pipelining bring about performance improvement on SM8x (it's related to the question above. If there are a lot of register spills then I think this is also a specialized optimization for SM90+).

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors have adequately addressed the limitations of this paper. There is no potential negative societal impact.

Summary

This paper presents FlashAttention-3, which speedup the commonly-used attention operator on Hopper GPUs. The paper proposes to leverage the asynchronous execution of the Tensor Cores and Tensor Memory Acceleator to better utilize the GPU hardware. Specifically, the paper proposes three techniques: (1) It overlaps the data movement and GEMM computation by warp specialization; (2) It proposes to overlap the GEMM and softmax operations by interleaving these operations; and (3) It leverages FP8 tensor core by block quantization. The experimental results show that FlashAttention-3 can achieve superior speedup by 1.5-2.0x with FP16 and can outperform vendor libraries.

Strengths

* It optimze attention operator on Hopper architecture, which is a timely problem and is of great importance; * It provides new insights for optimizing for Hopper architecture; * It achieves superior speedup compared to prior work and comparable speedup to vendor libraries;

Weaknesses

* The method introduces a few hyper-parameters such as number of pipeline stages, block sizes $B_r$ and $B_c$. However, the paper has never discussed how to tune or select this parameters in practice. * The paper does not evaluate with the end-to-end LLM inference task, which makes it obscure to see its real performance gain in practice.

Questions

Thanks for submitting the excellent paper to NeurIPS. The paper is in general well-written and the ideas are novel and insightful. However, I would like to make the following comments for further polishing the paper: The paper has mentioned a few hyperparameters. However, how do you tune these hyperparameters to achieve the best performance? Do you use grid search or simply heuristics? This is important because the inputs to language models often have varying sequence lengths, and tuning each of them is not practical. It would be great to see a systematic way to find these hyperparameters. The paper achieves superior speedup on the single flash attention operator. However, given the existence of FlashAttention-2, how much percentage does the attention operation takes in the context of whole LLM inference or training pipeline? This could help better understanding the positioning of this paper.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

* The paper is only specifically optimized for the Hopper architecture and may not applies to other or future GPUs.

Rebuttal Acknowledgment

Thanks for the clarifications. I have read the rebuttals.