FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Figure 6c reports backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 64.

Figure 6a reports backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 64.

Table 1 reports per-GPU training throughput for GPT3-2.7B at 2k context on 8×A100 GPUs for the baseline without FlashAttention, FlashAttention, and FlashAttention-2.

Figure 7c reports forward-plus-backward attention throughput on H100 80GB SXM5; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 64.

The H100 measurements use the unmodified implementation and no H100-specific TMA or fourth-generation Tensor Core instructions. The authors forecast, rather than measure, a further 1.5–2× speedup from those features and leave that optimization to future work.

Table 1 reports per-GPU training throughput for GPT3-1.3B at 2k context on 8×A100 GPUs for the baseline without FlashAttention, FlashAttention, and FlashAttention-2.

Figure 5c reports forward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 64.

Figure 7b reports forward-plus-backward attention throughput on H100 80GB SXM5; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 128.

Figure 4b reports forward-plus-backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 128.

Figure 4d reports forward-plus-backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 128.

Figure 5b reports forward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 128.

Table 1 reports per-GPU training throughput for GPT3-1.3B at 8k context on 8×A100 GPUs for the baseline without FlashAttention, FlashAttention, and FlashAttention-2.

Figure 4a reports forward-plus-backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 64.

Figure 5d reports forward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 128.

Figure 6d reports backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 128.

The A100 results are summarized as reaching up to 230 TFLOPs/s, 73% of theoretical maximum in forward, and 63% of theoretical maximum in backward.

Figure 5a reports forward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 64.

Running the same implementation on H100 without H100-specific instructions is described in the prose as reaching up to 335 TFLOPs/s.

Figure 4c reports forward-plus-backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 64.

Across the A100 attention benchmarks, the paper summarizes FlashAttention-2 as 1.7–3.0× faster than FlashAttention, 1.3–2.5× faster than the Triton implementation, and 3–10× faster than standard PyTorch attention. A more operation-specific summary says 1.3–1.5× over Triton in forward and around 2× in backward, and around 2× over FlashAttention and xFormers overall.

The reported end-to-end TFLOPs/s uses 6 × sequence_length × parameter_count + 12 × layer_count × hidden_dimension × sequence_length². The paper notes that the attention term could be halved for causal masking, but deliberately does not halve it to remain consistent with Megatron-LM and prior reporting.

Figure 7d reports forward-plus-backward attention throughput on H100 80GB SXM5; PyTorch is marked OOM at sequence length 16k. The panel covers causal attention with head dimension 128.

For 1.3B- and 2.7B-parameter GPT-style models at 2k and 8k context, FlashAttention-2 is summarized as up to 1.3× faster than FlashAttention and 2.8× faster than a baseline without FlashAttention, reaching 225 TFLOPs/s/GPU and 72% model-FLOPs utilization.

Figure 6b reports backward attention throughput on A100 80GB SXM4; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 128.

Table 1 reports per-GPU training throughput for GPT3-2.7B at 8k context on 8×A100 GPUs for the baseline without FlashAttention, FlashAttention, and FlashAttention-2.

Figure 7a reports forward-plus-backward attention throughput on H100 80GB SXM5; PyTorch is marked OOM at sequence length 16k. The panel covers non-causal attention with head dimension 64.

For causal attention, skipping blocks wholly above the causal diagonal is reported to make the kernel around 1.7–1.8× faster than attention without a causal mask; only one square block per row needs elementwise causal masking.