Flash Attention and Memory Efficiency

Standard attention computes the full N x N attention matrix, requiring O(N^2) memory and multiple round-trips to GPU HBM (High Bandwidth Memory). Flash Attention (Dao et al., 2022) restructures the computation to minimize memory I/O, achieving the same result with 2-4x speedup.

The Memory Bottleneck

Standard Attention Computation

The standard implementation:

1. Compute S = QK^T (N x N matrix) — O(N^2) compute and memory
2. Compute P = softmax(S) — O(N^2) compute and memory
3. Compute O = PV — O(N^2) compute, O(Nd) memory

Problem: Step 1 and 2 require writing/reading the N x N matrix to/from HBM, which is the bottleneck.

Flash Attention Tiling

Standard Attention (HBM access):              Flash Attention (SRAM):
                                              
Q, K, V in HBM                                Q, K, V in HBM
     |                                              |
     v                                              v
Compute S = QK^T (N x N)                     Load block of Q, K, V
     |                                        into SRAM
     v                                              |
Write S to HBM                                    v
     |                                        Compute block attention
     v                                        in SRAM
Load S from HBM                                   |
     |                                        Accumulate with online
     v                                        softmax statistics
Compute P = softmax(S)                             |
     |                                              v
     v                                        Write output block
Write P to HBM                                to HBM
     |                                              |
     v                                              v
Compute O = PV                                 Done!
     |
     v
Done!

HBM Accesses: O(N^2 + Nd)                   HBM Accesses: O(N^2 d / M)

Online Softmax

def flash_attention_block(Q_block, K_block, V_block, l_prev, m_prev):
    """Compute one block of flash attention."""
    # Compute attention scores for this block
    S_block = Q_block @ K_block.T / math.sqrt(Q_block.shape[-1])
    
    # Update running max
    m_block = S_block.max(dim=-1, keepdim=True).values
    m_new = torch.maximum(m_prev, m_block)
    
    # Compute exp with corrected max
    P_block = torch.exp(S_block - m_new)
    
    # Update running sum
    l_block = P_block.sum(dim=-1, keepdim=True)
    l_new = l_prev * torch.exp(m_prev - m_new) + l_block
    
    # Update output
    O_new = (l_prev * torch.exp(m_prev - m_new) * O_prev + P_block @ V_block) / l_new
    
    return O_new, l_new, m_new

Flash Attention Variants

Flash Attention 2

Flash Attention 3

Implementation

import torch
import math

def flash_attention_forward(Q, K, V, block_size=256):
    """Simplified Flash Attention implementation."""
    batch, seq_len, num_heads, head_dim = Q.shape
    
    O = torch.zeros_like(Q)
    l = torch.zeros(batch, num_heads, seq_len, 1, device=Q.device)
    m = torch.full((batch, num_heads, seq_len, 1), float('-inf'), device=Q.device)
    
    # Process K, V in blocks
    for j in range(0, seq_len, block_size):
        K_block = K[:, j:j+block_size, :, :]
        V_block = V[:, j:j+block_size, :, :]
        
        # Process Q in blocks
        for i in range(0, seq_len, block_size):
            Q_block = Q[:, i:i+block_size, :, :]
            
            # Compute attention scores
            S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / math.sqrt(head_dim)
            
            # Online softmax update
            m_block = S_block.max(dim=-1, keepdim=True).values
            m_new = torch.maximum(m[:, i:i+block_size], m_block)
            
            P_block = torch.exp(S_block - m_new)
            l_block = P_block.sum(dim=-1, keepdim=True)
            
            # Update output
            O[:, i:i+block_size] = (
                torch.exp(m[:, i:i+block_size] - m_new) * O[:, i:i+block_size] +
                torch.matmul(P_block, V_block)
            )
            
            l[:, i:i+block_size] = (
                torch.exp(m[:, i:i+block_size] - m_new) * l[:, i:i+block_size] +
                l_block
            )
            m[:, i:i+block_size] = m_new
    
    # Normalize output
    O = O / l
    return O

Memory Comparison

Practice Exercises

1. Memory Calculation: For a 70B model with 4096 sequence length, calculate the memory savings of Flash Attention vs standard attention.

2. Block Size Analysis: How does block size affect the speed of Flash Attention? What is the optimal block size for A100 vs H100 GPUs?

3. Implementation: Implement Flash Attention for a simplified 1-layer, 1-head attention mechanism and verify it produces identical results to standard attention.

4. Profiling: Profile the HBM access patterns of standard vs Flash Attention. Where does the speedup come from?

Key Takeaways

What to Learn Next