reviews: 
Comments: The authors present a new algorithm, FlashAttention-2, which builds in FlashAttention to improve the efficiency of the attention algorithm when executed on GPUs. The authors focus, in particular, on maximizing the amount of time spent in "matrix-multiply" FLOPs, that is, computation that is using matrix multiplication units, which are better lent to GPU hardware given particular division of work amongst warps and loose requirements around shared memory accesses.

These optimizations include:
1. Deferring scaling of values in the online softmax computation to further reduce HBM utilization
2. Save a logsumexp for online softmax backwards rather than max and sums to reduce the memory usage.
3. The bulk of the changes present in FlashAttention-2 are related to scheduling. Parallelization over the sequence dimension results in better warp occupancy in cases where there are few attention heads or a low batch size. Further, the authors change how the backward pass shares computation in the query derivative update, which also reduce HBM utilization. Splitting the KV cache amongst thread blocks also helps to saturate memory bandwidth. Better partitioning amongst warps, overall, drives better utilization.
Strengths: - The proposed approach shows promising improvements in a performance-critical parts of end-to-end transformer computation.
- The approach is symmetrically applicable for both training and inference with broad applications in both research and production settings.
- It is important that engineering contributions such as FlashAttention-2 (which I will refer to FA-2) are part of conference literature, and the attention to detail therein is what drives the solid impact of this work rather than a special "algorithmic contribution" in a classical sense.
- The devil is in the details: careful analysis of scheduling and an intuitive approach to laying out computation drives the approach's strong results.
- The work builds on an already-strong baseline, FlashAttention.
- The paper is well-written and organized and clearly lays out the authors' contributions; in particular, the paper is quite accessible to those without a low-level background in machine learning computation or GPU programming.
Weaknesses: - The baselines benchmarked in the paper can be stronger.
  - For latency benchmarks in Figure 5, the only baselines are a FasterTransformer and PyTorch. The authors do not consider compilers
  - Do the PyTorch benchmarks in Figure 5 use CUDA Graphs? PyTorch has significant framework overhead, and CUDA Graphs can give an order of magnitude speedup for some workloads, especially latency-sensitive ones.
- The above applies more generally to the other evaluation in Section 4; FlashAttention-2 is compared to PyTorch, then implementations with Triton and Cutlass, but not with any other frameworks capable of code generation. For example: while XLA may not be memory-bandwidth-aware by default, it can still generate kernels with fused operators that significantly reduce total memory I/O.
- While not needing to resort to approximations is a significant advantage of FlashAttention-2, this could be highlighted much more in the manuscript. Section 1 discusses many alternative attention approximations -- even speculation about why these aren't used (i.e. they are riskier when researchers have limited resources and don't adapt as easily) would strengthen the exactness boon of the authors' approach.

Several improvements to writing might improve the paper:
- The constant and equation in general in Sections 4.1 (i.e. 4) is not adequately explained (why is the sequence length squared? why 4?). Clarifying these might help new readers.
- The usage of "major problem" in the first sentence of the abstract is unclear -- it's clear that scaling sequence lengths is difficult; are the authors suggesting the problem is difficult, significant, or both?
- Section 1: "However, context length increases" <-- is missing "**as** context length increases"
- Text in all of the provided diagrams can be made clearer, and the diagrams can be rendered more clearly. It is difficult to read them as is.
- Section 3.2 might more clearly explain "prefill" and "KV cache" to readers. While somewhat ubiquitously understood amongst people doing performance engineering for large-scale transformers, some clarification would help the paper flow and increase its accessibility.
- Figure 2 can be clearer with respect to rows and columns -- this is the attention matrix -- what is its size/can the axes be labeled?
Recommendation: 8: accept, good paper
Reviewer confidence: 5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness/correctness: 2 fair
Clarity: 3 good