Flash Attention & IO-Aware Algorithms

1. The Memory Bottleneck

Modern GPUs have a memory hierarchy:

SRAM (on-chip): ~20 MB, ~19 TB/s bandwidth
HBM (off-chip): ~40 GB, ~1.5 TB/s bandwidth

Standard attention requires loading the attention matrix to HBM, which is the bottleneck.

1.1 Memory Access Cost

For standard attention:

For Flash Attention:

where is the SRAM size.

2. Flash Attention Tiling Diagram

<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="400" fill="#f8fafc"/>
  
  {/* Title */}
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Flash Attention: IO-Aware Tiling</text>
  
  {/* SRAM vs HBM */}
  <g transform="translate(50, 50)">
    <text x="175" y="0" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">Memory Hierarchy</text>
    
    {/* SRAM */}
    <rect x="50" y="15" width="150" height="80" rx="8" fill="#f8fafc" stroke="#22c55e" stroke-width="2"/>
    <text x="125" y="40" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="11" font-weight="bold">SRAM</text>
    <text x="125" y="60" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">~20 MB, 19 TB/s</text>
    <text x="125" y="75" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Fast but small</text>
    
    {/* HBM */}
    <rect x="50" y="110" width="300" height="80" rx="8" fill="#f8fafc" stroke="#ec4899" stroke-width="2"/>
    <text x="200" y="135" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="11" font-weight="bold">HBM</text>
    <text x="200" y="155" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">~40 GB, 1.5 TB/s</text>
    <text x="200" y="170" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Slow but large</text>
    
    {/* Transfer arrow */}
    <line x1="125" y1="95" x2="125" y2="110" stroke="#f59e0b" stroke-width="2" marker-end="url(#382e_arrowYellow4)"/>
    <text x="140" y="103" fill="#fbbf24" font-family="monospace" font-size="8">Transfer</text>
  </g>
  
  {/* Standard Attention */}
  <g transform="translate(450, 50)">
    <text x="175" y="0" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="11" font-weight="bold">Standard Attention</text>
    
    {/* Full attention matrix */}
    <rect x="25" y="15" width="150" height="150" rx="5" fill="#f8fafc" stroke="#ec4899" stroke-width="2"/>
    <text x="100" y="90" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="10" font-weight="bold">QKᵀ</text>
    <text x="100" y="110" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">n × n matrix</text>
    <text x="100" y="125" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Must fit in HBM</text>
    <text x="100" y="140" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="8">O(n²) memory</text>
    
    {/* Arrow to HBM */}
    <line x1="100" y1="170" x2="100" y2="190" stroke="#ec4899" stroke-width="2" marker-end="url(#382e_arrowPink)"/>
    <text x="115" y="183" fill="#ec4899" font-family="monospace" font-size="8">Materialize</text>
  </g>
  
  {/* Flash Attention */}
  <g transform="translate(50, 220)">
    <text x="350" y="0" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="12" font-weight="bold">Flash Attention: Tiling</text>
    
    {/* Tiled computation */}
    <g transform="translate(0, 15)">
      {/* Tile 1 */}
      <rect x="0" y="0" width="80" height="80" rx="5" fill="#3b82f6" opacity="0.8"/>
      <text x="40" y="35" text-anchor="middle" fill="white" font-family="monospace" font-size="9">Tile 1</text>
      <text x="40" y="55" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Q₁K₁ᵀV₁</text>
      
      {/* Tile 2 */}
      <rect x="90" y="0" width="80" height="80" rx="5" fill="#a855f7" opacity="0.8"/>
      <text x="130" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="9">Tile 2</text>
      <text x="130" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Q₂K₂ᵀV₂</text>
      
      {/* Tile 3 */}
      <rect x="180" y="0" width="80" height="80" rx="5" fill="#22c55e" opacity="0.8"/>
      <text x="220" y="35" text-anchor="middle" fill="white" font-family="monospace" font-size="9">Tile 3</text>
      <text x="220" y="55" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Q₃K₃ᵀV₃</text>
      
      {/* ... */}
      <text x="290" y="40" fill="#64748b" font-family="monospace" font-size="14">.</text>
      
      {/* Tile B */}
      <rect x="310" y="0" width="80" height="80" rx="5" fill="#f59e0b" opacity="0.8"/>
      <text x="350" y="35" text-anchor="middle" fill="white" font-family="monospace" font-size="9">Tile B</text>
      <text x="350" y="55" text-anchor="middle" fill="white" font-family="monospace" font-size="8">Q_BK_BᵀV_B</text>
    </g>
    
    {/* Processing */}
    <g transform="translate(0, 100)">
      <rect width="390" height="60" rx="8" fill="#f8fafc" stroke="#22c55e" stroke-width="2"/>
      <text x="195" y="20" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">Tiled Processing</text>
      <text x="195" y="40" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">Load tile → Compute in SRAM → Write output → Next tile</text>
      <text x="195" y="55" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="8">No n×n matrix materialized in HBM!</text>
    </g>
  </g>
  
  {/* IO complexity */}
  <g transform="translate(50, 340)">
    <rect width="700" height="40" rx="5" fill="#f8fafc" stroke="#334155"/>
    <text x="350" y="25" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="10">Standard: O(n²) HBM accesses | Flash: O(n²d²/M) HBM accesses (M = SRAM size)</text>
  </g>
  
  <defs>
    <marker id="382e_arrowYellow4" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
      <polygon points="0 0, 10 3.5, 0 7" fill="#fbbf24"/>
    </marker>
    <marker id="382e_arrowPink" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
      <polygon points="0 0, 10 3.5, 0 7" fill="#ec4899"/>
    </marker>
  </defs>
</svg>

3. SRAM vs HBM

3.1 Memory Characteristics

Property	SRAM	HBM
Size	~20 MB	~40 GB
Bandwidth	19 TB/s	1.5 TB/s
Latency	~1 ns	~100 ns
Cost	High	Moderate
Location	On-chip	Off-chip

3.2 The IO Bottleneck

For standard attention:

The arithmetic intensity:

For , this is IO-bound (AI < 100).

3.3 Flash Attention IO Complexity

With MB, , :

32× reduction in IO!

4. Flash Attention Algorithm

4.1 Tiling Strategy

Divide into blocks of size and :

4.2 Online Softmax

The key innovation is online softmax computation:

4.3 Algorithm Pseudocode

Architecture Diagram

Input: Q, K, V ∈ ℝⁿˣᵈ, block sizes Br, Bc
Output: O ∈ ℝⁿˣᵈ

Initialize O = 0, l = 0, m = -∞

for j = 1 to Tc:  // Loop over K, V blocks
    Load Kj, Vj to SRAM
    
    for i = 1 to Tr:  // Loop over Q blocks
        Load Qi, Oi, li, mi to SRAM
        
        // Compute block attention
        Sij = Qi Kjᵀ / √d
        
        // Online softmax update
        mij_new = max(mi, rowmax(Sij))
        Pij = exp(Sij - mij_new)
        lij_new = exp(mi - mij_new) * li + rowsum(Pij)
        
        // Update output
        Oi_new = diag(exp(mi - mij_new)) * Oi + Pij Vj
        Oi = Oi_new / lij_new
        
        // Store back to HBM
        Store Oi, lij_new, mij_new

5. Memory Access Pattern

<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="400" fill="#f8fafc"/>
  
  {/* Title */}
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Flash Attention: Memory Access Pattern</text>
  
  {/* Standard attention memory access */}
  <g transform="translate(50, 60)">
    <text x="175" y="0" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="11" font-weight="bold">Standard Attention</text>
    
    {/* HBM */}
    <rect x="0" y="15" width="350" height="150" rx="8" fill="#f8fafc" stroke="#ec4899" stroke-width="2"/>
    <text x="175" y="35" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="10">HBM Access Pattern</text>
    
    {/* Q matrix */}
    <rect x="20" y="50" width="60" height="100" rx="3" fill="#3b82f6"/>
    <text x="50" y="105" text-anchor="middle" fill="white" font-family="monospace" font-size="9">Q</text>
    
    {/* K matrix */}
    <rect x="90" y="50" width="60" height="100" rx="3" fill="#a855f7"/>
    <text x="120" y="105" text-anchor="middle" fill="white" font-family="monospace" font-size="9">K</text>
    
    {/* Attention matrix */}
    <rect x="160" y="50" width="80" height="100" rx="3" fill="#f59e0b"/>
    <text x="200" y="105" text-anchor="middle" fill="white" font-family="monospace" font-size="9">QKᵀ</text>
    
    {/* V matrix */}
    <rect x="250" y="50" width="60" height="100" rx="3" fill="#22c55e"/>
    <text x="280" y="105" text-anchor="middle" fill="white" font-family="monospace" font-size="9">V</text>
    
    {/* IO cost */}
    <text x="175" y="180" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="10">IO: O(n²) reads/writes</text>
    <text x="175" y="195" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Must materialize full n×n attention matrix</text>
  </g>
  
  {/* Flash Attention memory access */}
  <g transform="translate(450, 60)">
    <text x="175" y="0" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="11" font-weight="bold">Flash Attention</text>
    
    {/* HBM */}
    <rect x="0" y="15" width="350" height="150" rx="8" fill="#f8fafc" stroke="#22c55e" stroke-width="2"/>
    <text x="175" y="35" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">HBM Access Pattern (Tiled)</text>
    
    {/* Tiles */}
    <rect x="20" y="50" width="40" height="40" rx="3" fill="#3b82f6"/>
    <text x="40" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Q₁</text>
    
    <rect x="70" y="50" width="40" height="40" rx="3" fill="#3b82f6"/>
    <text x="90" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="7">Q₂</text>
    
    <rect x="120" y="50" width="40" height="40" rx="3" fill="#a855f7"/>
    <text x="140" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="7">K₁</text>
    
    <rect x="170" y="50" width="40" height="40" rx="3" fill="#a855f7"/>
    <text x="190" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="7">K₂</text>
    
    <rect x="220" y="50" width="40" height="40" rx="3" fill="#22c55e"/>
    <text x="240" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="7">V₁</text>
    
    <rect x="270" y="50" width="40" height="40" rx="3" fill="#22c55e"/>
    <text x="290" y="75" text-anchor="middle" fill="white" font-family="monospace" font-size="7">V₂</text>
    
    {/* Arrows showing sequential access */}
    <path d="M 40 95 L 140 95 L 240 95" fill="none" stroke="#f59e0b" stroke-width="2" marker-end="url(#382e_arrowYellow5)"/>
    
    {/* IO cost */}
    <text x="175" y="180" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">IO: O(n²d²/M) reads/writes</text>
    <text x="175" y="195" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Only load tiles to SRAM as needed</text>
  </g>
  
  {/* Bandwidth utilization */}
  <g transform="translate(50, 260)">
    <rect width="700" height="120" rx="8" fill="#f8fafc" stroke="#334155"/>
    <text x="350" y="20" text-anchor="middle" fill="#f59e0b" font-family="monospace" font-size="11" font-weight="bold">Bandwidth Utilization</text>
    
    {/* Standard */}
    <rect x="30" y="40" width="300" height="30" rx="5" fill="#ec4899" opacity="0.8"/>
    <text x="180" y="60" text-anchor="middle" fill="white" font-family="monospace" font-size="9">Standard: 25% HBM bandwidth</text>
    
    {/* Flash */}
    <rect x="30" y="80" width="650" height="30" rx="5" fill="#22c55e" opacity="0.8"/>
    <text x="355" y="100" text-anchor="middle" fill="white" font-family="monospace" font-size="9">Flash: 75% HBM bandwidth (3× improvement)</text>
  </g>
  
  <defs>
    <marker id="382e_arrowYellow5" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
      <polygon points="0 0, 10 3.5, 0 7" fill="#fbbf24"/>
    </marker>
  </defs>
</svg>

6. Tiling Block Sizes

6.1 Block Size Selection

The optimal block sizes and must satisfy:

For typical values (, MB):

6.2 Block Size Trade-offs

	SRAM Usage	IO Cost	Quality
32	Low	Higher	Slightly lower
64	Medium	Medium	Good
128	High	Lower	Better
256	Very high	Lowest	Best

6.3 Practical Block Sizes

Flash Attention typically uses:

or
or
Trade-off between IO and numerical precision

7. Backward Pass

7.1 Recomputation Strategy

Instead of storing the attention matrix, recompute it during the backward pass:

7.2 Memory Savings

Standard backward pass:

Store: (attention matrix)
Memory:

Flash backward pass:

Store: (softmax statistics)
Memory:

7.3 Recomputation Cost

The recomputation adds:

But saves:

Trade-off: 33% more compute for linear memory!

8. Flash Attention v2 (Dao, 2023)

8.1 Improvements over v1

Better parallelism: Parallelize over batch and head dimensions
Reduced non-matmul FLOPs: ~50% reduction
Better work partitioning: Between warps

8.2 FLOP Reduction

Flash v1:

Flash v2:

8.3 Speedup

Sequence Length	Flash v1 Speedup	Flash v2 Speedup
1024	2.0×	2.5×
4096	3.5×	5.0×
16384	5.0×	8.0×
65536	6.0×	10.0×

9. Memory Efficiency

9.1 Memory Comparison

Method	Forward Memory	Backward Memory
Standard
Flash v1
Flash v2

9.2 Practical Memory Usage

For , , batch size 32:

Method	Forward	Backward	Total
Standard	2.1 GB	4.2 GB	6.3 GB
Flash	0.5 GB	1.0 GB	1.5 GB

9.3 Scaling to Longer Sequences

With Flash Attention:

: ~3 GB (fits in most GPUs)
: ~6 GB (requires A100 40GB)
: ~12 GB (requires A100 80GB)
: ~24 GB (requires multiple GPUs)

10. Implementation Considerations

10.1 Numerical Precision

Flash Attention uses online softmax which can accumulate numerical errors. Mitigations:

Periodic rescaling
Double precision for accumulators
Careful ordering of operations

10.2 Hardware Compatibility

Flash Attention is optimized for:

NVIDIA GPUs with Tensor Cores (A100, H100)
High bandwidth memory (HBM2e, HBM3)
Specific warp sizes (32 threads)

10.3 Integration with Frameworks

# PyTorch implementation using Flash Attention
from flash_attn import flash_attn_func

# Standard attention
output = torch.softmax(Q @ K.T / sqrt(d), dim=-1) @ V

# Flash Attention
output = flash_attn_func(Q, K, V, causal=True)

Flash Attention has become the de facto standard for efficient Transformer training, enabling models with tens of thousands of context length on modern hardware.

Flash Attention & IO-Aware Algorithms

Flash Attention & IO-Aware Algorithms

1. The Memory Bottleneck

1.1 Memory Access Cost

2. Flash Attention Tiling Diagram

3. SRAM vs HBM

3.1 Memory Characteristics

3.2 The IO Bottleneck

3.3 Flash Attention IO Complexity

4. Flash Attention Algorithm

4.1 Tiling Strategy

4.2 Online Softmax

4.3 Algorithm Pseudocode

5. Memory Access Pattern

6. Tiling Block Sizes

6.1 Block Size Selection

6.2 Block Size Trade-offs

6.3 Practical Block Sizes

7. Backward Pass

7.1 Recomputation Strategy

7.2 Memory Savings

7.3 Recomputation Cost

8. Flash Attention v2 (Dao, 2023)

8.1 Improvements over v1

8.2 FLOP Reduction

8.3 Speedup

9. Memory Efficiency

9.1 Memory Comparison

9.2 Practical Memory Usage

9.3 Scaling to Longer Sequences

10. Implementation Considerations

10.1 Numerical Precision

10.2 Hardware Compatibility

10.3 Integration with Frameworks

Need Expert AI/ML Premium Help?