Paper: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR-2024)

Reviewer comment (Reproducibility):
How do the "tweaks" in Section 4.1.1 help reduce non-matrixmul FLOPs? I know a fair amount about tensor cores, but it wasn't obvious to me.

Author response:
Q: How does the tweaks reduce non-matmul FLOPs.

A: The amount of matmul FLOPs stay the same (e.g. Q @ K^T and attn_prob @ V). However, FlashAttention incurs more vector flops to rescale the output (equation at bottom of page 4): extra divisions by l^2 on both terms. FlashAttention2 gets rid of this division (top equation at top of page 5).

In addition, masking (non-matmul operation) is only done to 1 block and not every blocks, again reducing non-matmul FLOPs so that most of the time is spent on matmul FLOPs accelerated by tensor cores.

These changes generally reduce register pressure as an added benefit.

Reviewer comment (Evaluation):
In Figure 2, does an uncolored square mean no computation? Does the backward pass for a given worker start right away or do workers need to synchronize between forward and backward pass? Do you not need to compute the combined result for the forward pass before you can start the backward pass? If you do need to wait, then how can one achieve greater than 50% use of peak performance if roughly half the compute cycles are spent waiting for the longest running forward/backward pass thread block to complete? If you don't need to wait, why not?

Author response:
Q: Figure 2, does an uncolored square mean no computation?

A: Yes, we skip those blocks since they don't contribute to the outputs (i.e. masked out) any way.

Q: backward pass synchronization with forward pass.

A: The backward pass is done separately from the forward pass, so there's no overlap between them. As an example, in the model forward pass, the attention and MLP forward pass are performed, then we calculate a loss function, which then gives a gradient that gets propagated to the MLP backward pass and then attention backward pass.

Reviewer comment (Evaluation):
I'm not sure how to relate Figure 3 to Algorithm 1 (i.e., which lines it is meant to illustrate). From the two paragraphs above Figure 3 I get it there are two potential sources of reduced execution time: fewer shared memory accesses and fewer synchronizations (__syncthreads, I assume). Unclear which of those matters most and why given that shared memory accesses proceed about as fast as register file accesses and synchronization with a thread block is low overhead.

Author response:
Q: Figure 3.

A: Thanks for this clarification, we will add this to the paper. Figure 3 corresponds to line 8 and 10 in Algorithm 1: matmul Q K^T and matmul P V. Figure 3 illustrates how to map these matrix multiplications to 4 warps in the same threadblock.

Reviewer comment (Experiments):
Why is FlashAttention (version 1) missing in Figure 5?

Author response:
Q: FlashAttention in Figure 5.

A: For decoding with long sequences, FlashAttention is actually slower than Pytorch due to lack of parallelism. As a result FlashAttention was only used for prompt processing (prefill) and not decoding (KV cache). FlashAttention-2 now speeds up both prompt processing (prefill) and decoding.

We will add this note to the paper.

Reviewer comment (Evaluation):
As someone who knows GPUs well, I would have liked to see more performance counter data to backup the claims of the sources of performance improvements. I understand space is limited in the main text, but in checking the supplemental material, while it is great to see all the code, there appeared to be no PDF providing additional data or details. Including one might have helped.

Author response:
Q: performance counter data.

A: Thanks for this great suggestion. We will add this to the paper to help clarify the effectiveness of our approaches.

Reviewer comment (Experiments):
The baselines benchmarked in the paper can be stronger.

Author response:
Q: Figure 5 baselines.

A: Here we measure just the attention decoding operation, not the end-to-end latency. We agree that CUDA graph helps end-to-end generation speed, especially for small models where the CPU overhead can be large.
We have benchmarked the kernel generated by torch.compile which is faster than Pytorch eager but still not as fast as the hand-written kernel from FasterTransformer. FA2 decoding kernel is faster still.

Reviewer comment (Experiments):
For latency benchmarks in Figure 5, the only baselines are a FasterTransformer and PyTorch. The authors do not consider compilers.

Author response:
Q: Comparison to compilers.

A: Compilers can generally perform fusion. However, optimizations that require mathematical rewriting of the same expression (while maintaining numerical stability) are generally harder for compilers. Jax uses the powerful XLA compiler, but currently the best implementation of attention on TPUs is a version of FA2 implemented in Pallas, where the programmer still specifies some details (which elements should be in HBM vs SRAM) leaving the Pallas compiler to generate code to invoke the systolic array for matmuls or vector units for other computation.

Reviewer comment (Experiments):
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.

Author response:
Q: Figure 5 baselines.

A: Here we measure just the attention decoding operation, not the end-to-end latency. We agree that CUDA graph helps end-to-end generation speed, especially for small models where the CPU overhead can be large.
We have benchmarked the kernel generated by torch.compile which is faster than Pytorch eager but still not as fast as the hand-written kernel from FasterTransformer. FA2 decoding kernel is faster still.

Reviewer comment (Experiments):
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.

Author response:
Q: gap between different implementations, e.g. xformers, Cutlass, Triton.

xformers uses cutlass to implement a similar algorithm to FA, while Triton generates ptx directly. The differences in speed are due to low-level implementation: e.g. the Triton compiler only deals with power-of-2 block sizes, while with cutlass one can use non-power-of-2 block sizes. The triton compiler will use async memory copy automatically, while with Cutlass one has to spend more effort to use these hardware features.

Reviewer comment (Evaluation):
The performance gains are relatively marginal, especially when compared to those of the original FlashAttention over the baseline.

Author response:
Attention is now one of the most optimized components as one of the core layers of Transformers training on massive clusters. FlashAttention-2 speeding up attention by up to 2x can directly benefit these large training runs (costing millions of dollars). It can also enable new capabilities such as long context models (for books and code, high-res images, audio signal). It also make it easier to deploy models locally on desktop GPUs or enable research into these models due to faster iteration speed.

Reviewer comment (Experiments):
The absence of an ablation study makes it difficult to pinpoint the exact sources of efficiency.

Author response:
Thank you for the helpful suggestions on the ablation study to improve the work. We have included such a study in the common response.

Reviewer comment (Experiments):
Perhaps the most significant omission in this paper is the lack of an ablation study. This makes it challenging to discern the contributions of individual optimizations. Among the proposed optimizations, which one has the highest impact?

Author response:
Thank you for the helpful suggestions on the ablation study to improve the work. We have included such a study in the common response.

Reviewer comment (Experiments):
In Section 3.1, is the technique of skipping blocks for "causal masking" also applied to FlashAttention? As the authors mention, this technique can be applied to both FlashAttention and FlashAttention-2, and I am curious about how the application of this technique would affect the performance gap between FlashAttention and FlashAttention-2 if it had not been applied to FlashAttention.

Author response:
Yes, in our benchmarks, we also apply the same technique to FlashAttention. The difference in speed is due to our 3 contributions.

Reviewer comment (Experiments):
In Section 3.3, what is the performance impact of "tuning block sizes"? Was the same level of parameter tuning effort applied to FlashAttention? My question concerns the extent to which the performance gains over FlashAttention can be attributed to algorithmic improvements versus additional tuning effort, with the latter being less significant.

Author response:
This was done at the very end to squeeze out a little more performance, on the order of 5%. It was not done with FlashAttention since its implementation was much more rigid and the block sizes were hard-coded. We do not expect block size tuning to play a major role here.