FlashAttention Implementation › FlashAttention-2 Improvements
Language AI HandbookAll chapters FlashAttention Implementation › FlashAttention-2 ImprovementsOne section · Sections in this chapter5 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
FlashAttention-2, released in 2023, builds on the original algorithm with several key optimizations that improve hardware utilization from ~50-70% to ~70-90% of theoretical peak performance.
Reduced Non-Matrix-Multiply FLOPs
The original FlashAttention spent significant time on non-matrix-multiply operations: rescaling outputs, updating statistics, and coordinating between iterations. FlashAttention-2 restructures the algorithm to minimize these overheads.
The key change is delaying the rescaling of output blocks. Instead of rescaling after each K-V block, FlashAttention-2 accumulates unnormalized outputs and applies a single rescaling at the end. This requires careful bookkeeping but reduces the number of scalar operations significantly.
def compare_non_matmul_ops(n, B_r, B_c):
"""
Compare non-matrix-multiply operations between FlashAttention versions.
"""
num_blocks = n // B_c # Number of K-V blocks per query block
# FlashAttention-1: rescale after each block
# For each K-V block: update m, l, rescale output
flash1_rescales = num_blocks * B_r # One rescale per query per block
flash1_stats_updates = num_blocks * B_r * 2 # m and l updates
# FlashAttention-2: single rescale at the end
flash2_rescales = B_r # One rescale per query at the end
flash2_stats_updates = num_blocks * B_r * 2 # Same stats tracking
return {
"flash1_rescales": flash1_rescales,
"flash2_rescales": flash2_rescales,
"rescale_reduction": flash1_rescales / flash2_rescales,
"num_kv_blocks": num_blocks,
}
# Example calculation
ops = compare_non_matmul_ops(n=4096, B_r=64, B_c=64)Non-MatMul Operations Comparison (n=4096, block=64): K-V blocks processed: 64 FlashAttention-1 rescales: 4,096 FlashAttention-2 rescales: 64 Reduction: 64x fewer rescale operations
With 64 K-V blocks to process, FlashAttention-1 performs 4,096 rescale operations (one per query per block), while FlashAttention-2 performs only 64 (one per query at the end). This 64x reduction in rescaling operations translates directly to improved throughput, particularly when matrix multiplies are already saturating the GPU's compute units.
Better Work Partitioning
FlashAttention-2 improves how work is distributed across thread blocks and warps:
- Parallelism over batch and heads: Each thread block handles one (batch, head) pair, maximizing parallelism across the GPU
- Within-block parallelism: Different warps handle different parts of the Q block, with careful synchronization
- Sequence parallelism: For very long sequences, work can be split across multiple blocks even for a single head
Improved Memory Access
FlashAttention-2 optimizes memory access patterns in several ways:
- Swapping loops: The inner and outer loops are reordered so that the outer loop iterates over Q blocks and the inner loop over K-V blocks. This reduces the number of times Q blocks are loaded from HBM.
- Pipelining: Memory loads are overlapped with computation. While computing attention for the current K-V block, the next block is being loaded from HBM.
- Warp-level primitives: Instead of synchronizing entire thread blocks, FlashAttention-2 uses warp-level operations that are faster and more efficient.
def memory_access_comparison(n, d, B_r, B_c, dtype_bytes=2):
"""
Compare HBM access patterns between FlashAttention versions.
"""
num_q_blocks = n // B_r
num_kv_blocks = n // B_c
# FlashAttention-1: outer loop over K-V, inner over Q
# Q loaded num_kv_blocks times
flash1_q_loads = num_kv_blocks * n * d * dtype_bytes
flash1_kv_loads = (
n * d * 2 * dtype_bytes
) # K and V each once per Q block pass
# FlashAttention-2: outer loop over Q, inner over K-V
# Q loaded once per Q block
flash2_q_loads = n * d * dtype_bytes # Just once total
flash2_kv_loads = (
num_q_blocks * n * d * 2 * dtype_bytes
) # K, V for each Q block
return {
"flash1_total_mb": (flash1_q_loads + flash1_kv_loads) / 1e6,
"flash2_total_mb": (flash2_q_loads + flash2_kv_loads) / 1e6,
"flash2_q_reduction": flash1_q_loads / flash2_q_loads,
}
access = memory_access_comparison(n=4096, d=64, B_r=64, B_c=64)Memory Access Patterns (n=4096, d=64, block=64): FlashAttention-1 total HBM reads: 34.6 MB FlashAttention-2 total HBM reads: 67.6 MB Q matrix read reduction: 64x fewer Q loads
By swapping the loop order, FlashAttention-2 reads the Q matrix 64x fewer times from HBM. This is the dominant source of the memory access reduction. Since memory bandwidth is the bottleneck, reducing redundant reads directly improves throughput.
Performance Gains
These optimizations combine to give FlashAttention-2 a 2x speedup over FlashAttention on many workloads:
def theoretical_speedup():
"""
Approximate speedups from FlashAttention to FlashAttention-2.
Based on published benchmarks.
"""
# Speedups vary by sequence length and head dimension
benchmarks = [
{
"config": "n=1K, d=64",
"flash1_tflops": 124,
"flash2_tflops": 185,
"speedup": 1.49,
},
{
"config": "n=2K, d=64",
"flash1_tflops": 128,
"flash2_tflops": 210,
"speedup": 1.64,
},
{
"config": "n=4K, d=64",
"flash1_tflops": 115,
"flash2_tflops": 220,
"speedup": 1.91,
},
{
"config": "n=8K, d=64",
"flash1_tflops": 109,
"flash2_tflops": 230,
"speedup": 2.11,
},
{
"config": "n=16K, d=64",
"flash1_tflops": 102,
"flash2_tflops": 235,
"speedup": 2.30,
},
]
return benchmarks
benchmarks = theoretical_speedup()FlashAttention-2 vs FlashAttention-1 Performance (A100): Configuration FA-1 (TFLOPS) FA-2 (TFLOPS) Speedup ------------------------------------------------------------ n=1K, d=64 124 185 1.49x n=2K, d=64 128 210 1.64x n=4K, d=64 115 220 1.91x n=8K, d=64 109 230 2.11x n=16K, d=64 102 235 2.30x
The benchmarks show a clear trend: FlashAttention-2's advantage grows with sequence length. At 1K tokens, the speedup is a modest 1.49x. At 16K tokens, it reaches 2.30x. This pattern occurs because longer sequences have more K-V blocks, amplifying the benefits of reduced rescaling operations and optimized memory access. FlashAttention-2 also achieves higher absolute TFLOPS, indicating better GPU utilization.
