FlashAttention Implementation › GPU Memory Hierarchy
Language AI HandbookAll chapters FlashAttention Implementation › GPU Memory HierarchyOne section · Sections in this chapter2 / 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
Before understanding FlashAttention, you need to understand where bottlenecks occur on GPUs. Modern deep learning is often described as "compute-bound," but attention computation is "memory-bound." That classification determines whether arithmetic or memory traffic is the limiting factor.
A computation is compute-bound when the limiting factor is the number of arithmetic operations the processor can perform. A computation is memory-bound when the limiting factor is how quickly data can be transferred to and from memory. Attention is memory-bound because we spend more time moving data than computing with it.
GPUs have a hierarchical memory structure with three main levels, each with different speeds and capacities:
def gpu_memory_specs():
"""
Approximate specifications for NVIDIA A100 GPU memory hierarchy.
These numbers illustrate the dramatic differences between memory levels.
"""
specs = {
"HBM (High Bandwidth Memory)": {
"capacity_gb": 80,
"bandwidth_tb_s": 2.0,
"latency_ns": 200,
},
"L2 Cache": {
"capacity_gb": 0.040, # 40 MB
"bandwidth_tb_s": 4.0,
"latency_ns": 50,
},
"SRAM (Registers + Shared Memory)": {
"capacity_gb": 0.000192, # 192 KB per SM, ~108 SMs
"bandwidth_tb_s": 19.0,
"latency_ns": 5,
},
}
return specs
specs = gpu_memory_specs()NVIDIA A100 GPU Memory Hierarchy: Level Capacity Bandwidth Latency ---------------------------------------------------------------------------------- HBM (High Bandwidth Memory) 80 GB 2.0 TB/s 200 ns L2 Cache 40.0 MB 4.0 TB/s 50 ns SRAM (Registers + Shared Memory) 0.2 MB 19.0 TB/s 5 ns
The numbers reveal a fundamental tension. HBM (the main GPU memory where tensors live) has enormous capacity but relatively slow access. SRAM (registers and shared memory) is much faster but tiny. The speed difference is nearly 10x, while the capacity difference is over 400,000x. Standard attention computation repeatedly reads and writes to HBM, suffering the slow access speeds. FlashAttention restructures the computation to keep data in SRAM as long as possible.
Why Standard Attention is Memory-Bound
Consider what happens during standard attention computation:
- Load Q, K, V from HBM to compute
- Write S back to HBM (this is an matrix)
- Load S from HBM to compute softmax
- Write attention weights P back to HBM
- Load P and V from HBM to compute output
- Write output back to HBM
Each step involves transfers to and from the slow HBM. The matrices S and P are particularly problematic: for a sequence of 4096 tokens with 32-bit floats, each matrix is 64 MB. With 32 attention heads, that's 2 GB just for intermediate storage, per layer, per batch element.
def memory_transfers_standard_attention(n, d, dtype_bytes=4):
"""
Calculate memory transferred in standard attention.
Args:
n: Sequence length
d: Model dimension
dtype_bytes: Bytes per element (4 for fp32, 2 for fp16)
Returns:
Dictionary with memory transfer statistics
"""
# Input matrices Q, K, V: each n x d
input_read = 3 * n * d * dtype_bytes
# Score matrix S = QK^T: n x n, written then read for softmax
score_write = n * n * dtype_bytes
score_read = n * n * dtype_bytes
# Attention weights P after softmax: n x n, written then read
weights_write = n * n * dtype_bytes
weights_read = n * n * dtype_bytes
# V read again for P @ V
v_read = n * d * dtype_bytes
# Output: n x d
output_write = n * d * dtype_bytes
total = (
input_read
+ score_write
+ score_read
+ weights_write
+ weights_read
+ v_read
+ output_write
)
return {
"input_read_mb": input_read / 1e6,
"intermediate_mb": (
score_write + score_read + weights_write + weights_read
)
/ 1e6,
"output_mb": output_write / 1e6,
"total_mb": total / 1e6,
}
# Example: typical transformer settings
transfers = memory_transfers_standard_attention(n=4096, d=768, dtype_bytes=4)Memory Transfers in Standard Attention (n=4096, d=768, fp32): Input Q, K, V: 37.7 MB Intermediate S, P: 268.4 MB Output: 12.6 MB Total transfers: 331.4 MB The n×n intermediate matrices dominate memory traffic.
The intermediate matrices consume over 128 MB of memory transfers, dwarfing the input and output. This is where FlashAttention intervenes: by never materializing the full matrices, it eliminates most of this memory traffic.