{"title":"11.11 Flash Attention","description":"","content":"# 11.11 Flash Attention\n\n## The Problem\n\nBy now you understand that attention is the core mechanism that makes transformers work. Every token attends to every other token, and this is what gives LLMs their ability to understand context, track references, and reason across long passages. But this power comes at a steep cost.\n\nStandard attention computes a full N x N attention matrix, where N is the sequence length. For a 4K context, that is 16 million elements. For a 32K context, that is over a billion. For a 128K context, over 16 billion. The matrix grows quadratically, and so does the memory required to store it.\n\nHere is what surprises most people: the actual arithmetic is not the bottleneck. GPUs are exceptionally fast at matrix multiplication. An A100 can do 312 TFLOPS of FP16 math per second. The real bottleneck is moving data. Every time the GPU computes that massive attention matrix, it writes it to its main memory (called HBM, or High Bandwidth Memory), then reads it back to compute softmax, then writes softmax back, then reads it again to multiply with values. Each of these read/write round-trips through slow memory dominates the total execution time.\n\nAttention is memory-bound, not compute-bound. This distinction is the entire foundation of Flash Attention.\n\nFlash Attention, introduced by Tri Dao et al. in 2022, restructures the attention computation so the full N x N matrix never leaves the GPU's fast on-chip memory. It does not approximate attention or change the math. It computes exact standard attention, just in a radically more efficient order.\n\n**What you will learn:**\n\n- Why GPU memory bandwidth, not compute, is the bottleneck in attention\n- How the GPU memory hierarchy (SRAM vs. HBM) dictates performance\n- What standard attention's memory access pattern looks like and why it is wasteful\n- How Flash Attention tiles the computation to avoid materializing the full attention matrix\n- Why recomputing attention in the backward pass is faster than saving it\n- What Flash Attention 2 and 3 improve\n- How to enable Flash Attention in PyTorch, Hugging Face, and vLLM\n\n---\n\n## GPU Memory Hierarchy\n\nTo understand why Flash Attention works, you need to understand where data lives on a GPU. There are two main levels of memory, and the gap between them is enormous.\n\n**HBM (High Bandwidth Memory):** This is the GPU's main memory. On an A100, it is 40 or 80 GB. It is what people usually mean when they say \"GPU memory\" or \"VRAM.\" HBM has high capacity but relatively slow access, around 2 TB/s bandwidth.\n\n**SRAM (Static RAM):** This is on-chip memory, sitting right next to the compute units. On an A100, each streaming multiprocessor has a small amount of shared memory, and the total across all SMs is roughly 20 MB. SRAM is tiny but extremely fast, with approximately 19 TB/s bandwidth.\n\nThat is nearly a 10x difference in bandwidth. Reading 1 GB from HBM takes about 0.5 milliseconds. Reading 1 GB from SRAM takes about 0.05 milliseconds. When your computation involves repeatedly reading and writing large intermediate results, the choice of where that data lives determines your performance far more than how fast the math runs.\n\n\n```mermaid\nflowchart TD\n subgraph GPU[\"GPU (A100)\"]\n direction TB\n SRAM[\"SRAM (On-Chip)
~20 MB
~19 TB/s bandwidth\"]:::green\n COMPUTE[\"Compute Units
312 TFLOPS (FP16)\"]:::primary\n HBM[\"HBM (Off-Chip)
40-80 GB
~2 TB/s bandwidth\"]:::orange\n end\n\n SRAM <-->|\"Fast path
~10x faster\"| COMPUTE\n HBM <-->|\"Slow path
bottleneck\"| COMPUTE\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n```\n\n\nThink of it like a kitchen. SRAM is your countertop, small but everything is within arm's reach. HBM is the pantry down the hall, it holds a lot more, but every trip to get ingredients slows you down. If your recipe requires you to walk to the pantry six times for a single dish, you are spending more time walking than cooking. A smarter chef would bring everything they need to the countertop in one trip and work from there.\n\nStandard attention is the cook who walks to the pantry for every single ingredient. Flash Attention is the one who plans ahead.\n\n---\n\n## Standard Attention: Step by Step\n\nLet's trace exactly what happens during a standard attention computation and count the memory round-trips.\n\nGiven input matrices Q (queries), K (keys), and V (values), attention computes:\n\n**Attention(Q, K, V) = softmax(Q x K^T / sqrt(d_k)) x V**\n\nHere is how a standard implementation executes this on a GPU:\n\n**Step 1:** Compute S = Q x K^T. This produces the N x N score matrix. The GPU reads Q and K from HBM, performs the matrix multiply in SRAM, and writes the result S back to HBM because it is too large to keep on-chip. That is one HBM write of O(N^2) elements.\n\n**Step 2:** Compute P = softmax(S). The GPU reads the entire N x N matrix S from HBM into SRAM (in chunks, since it does not fit all at once), computes softmax row by row, and writes the result P back to HBM. That is one HBM read and one HBM write of O(N^2) elements.\n\n**Step 3:** Compute O = P x V. The GPU reads P from HBM and V from HBM, multiplies them, and writes the output O to HBM. Another HBM read of O(N^2) for P.\n\n\n```mermaid\nflowchart LR\n subgraph Step1[\"Step 1: Scores\"]\n Q1[Q]:::primary --> MM1[\"Q × K^T\"]:::teal\n K1[K]:::primary --> MM1\n MM1 -->|\"Write O(N²)\"| S[\"S matrix
in HBM\"]:::red\n end\n\n subgraph Step2[\"Step 2: Softmax\"]\n S2[\"Read S
from HBM\"]:::red --> SM[\"softmax(S)\"]:::teal\n SM -->|\"Write O(N²)\"| P[\"P matrix
in HBM\"]:::red\n end\n\n subgraph Step3[\"Step 3: Output\"]\n P2[\"Read P
from HBM\"]:::red --> MM2[\"P × V\"]:::teal\n V1[V]:::primary --> MM2\n MM2 -->|\"Write O(N)\"| O[\"Output O
in HBM\"]:::green\n end\n\n Step1 --> Step2 --> Step3\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef teal fill:#38d9a9,stroke:#000,color:#000\n classDef red fill:#ff8787,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n```\n\n\nThe total HBM access for one attention operation is O(N^2) in three separate passes. For a sequence of 32,768 tokens, S and P each contain roughly 1 billion floating-point numbers. At FP16, that is about 2 GB per matrix. You are writing 2 GB, reading 2 GB, writing 2 GB, reading 2 GB, all through that 2 TB/s HBM pipe. The math itself finishes in microseconds, but the memory traffic takes milliseconds.\n\nThis is why naive attention is slow. The compute-to-memory ratio is terrible. The GPU's arithmetic units sit idle while waiting for data to be shuttled back and forth through HBM.\n\n---\n\n## Flash Attention: The Key Insight\n\nFlash Attention's central idea is simple to state and subtle to execute: never materialize the full N x N attention matrix in HBM. Instead, break the computation into small tiles that fit entirely in SRAM, compute partial attention within each tile, and accumulate the results.\n\nBut there is a catch, and it is the reason nobody did this before Tri Dao. Softmax is not a pointwise operation. To compute softmax for a single row, you need the maximum value across that entire row (for numerical stability) and the sum of all exponentiated values. If you are only looking at a tile of the row, you do not have the full picture. How do you compute a correct softmax without seeing the complete row?\n\nThe answer is online softmax. Instead of computing the exact max and sum upfront, Flash Attention maintains running statistics, a running maximum and a running sum of exponentials, and corrects previous partial results as new tiles reveal larger values. This is mathematically exact. There is no approximation. The final result is bit-for-bit identical to standard attention.\n\nHere is the high-level idea:\n\n\n```mermaid\nflowchart TD\n subgraph Input[\"Input Matrices\"]\n Q[\"Q blocks
Q₁, Q₂, ..., Q_m\"]:::primary\n K[\"K blocks
K₁, K₂, ..., K_n\"]:::orange\n V[\"V blocks
V₁, V₂, ..., V_n\"]:::orange\n end\n\n subgraph SRAM_Compute[\"SRAM (On-Chip Processing)\"]\n LOAD[\"Load Q_i, K_j, V_j
into SRAM\"]:::teal\n LOCAL[\"Compute local
attention scores
S_ij = Q_i × K_j^T\"]:::teal\n SOFTMAX[\"Update running
max and sum
(online softmax)\"]:::green\n ACCUM[\"Accumulate partial
output O_i\"]:::green\n end\n\n subgraph Output[\"Result\"]\n O[\"Final Output O
(exact attention)\"]:::primary\n end\n\n Q --> LOAD\n K --> LOAD\n V --> LOAD\n LOAD --> LOCAL --> SOFTMAX --> ACCUM\n ACCUM -->|\"Next tile\"| LOAD\n ACCUM -->|\"All tiles done\"| O\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef teal fill:#38d9a9,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n```\n\n\nInstead of three O(N^2) passes through HBM, Flash Attention reads Q, K, and V once from HBM and writes the final output O once. The intermediate N x N matrix is computed in tiles and never leaves SRAM. The total HBM access drops from O(N^2) to O(N^2 / M) where M is the SRAM size, and in practice this translates to massive speedups.\n\n---\n\n## The Tiling Algorithm\n\nLet's walk through how the tiling actually works. The algorithm processes attention in blocks, iterating over K/V blocks in the outer loop and Q blocks in the inner loop.\n\n**Setup:** Divide Q into blocks of size B_r rows, and K/V into blocks of size B_c rows. Choose B_r and B_c so that the tile of scores (B_r x B_c) plus the Q, K, V blocks all fit in SRAM.\n\n**The algorithm:**\n\n1. For each K/V block j (outer loop):\n\n- Load K_j and V_j from HBM into SRAM\n\n1. For each Q block i (inner loop):\n\n- Load Q_i and the current output O_i and softmax statistics (running max m_i, running sum l_i) from HBM into SRAM\n- Compute the local score tile: S_ij = Q_i x K_j^T (this is a small B_r x B_c matrix, fits in SRAM)\n- Compute the local max for this tile: m_ij = rowmax(S_ij)\n- Update the global running max: m_new = max(m_i, m_ij)\n- Compute local exponentials: P_ij = exp(S_ij - m_new)\n- Update the running sum: l_new = exp(m_i - m_new) x l_i + rowsum(P_ij)\n- Rescale previous output and add new contribution: O_i = (exp(m_i - m_new) x l_i x O_i + P_ij x V_j) / l_new\n- Write updated O_i, m_new, l_new back to HBM\n\n1. After all blocks are processed, O contains the exact attention output.\n\nHere is Python pseudocode illustrating the core logic:\n\n\n```python\nimport numpy as np\n\ndef flash_attention(Q, K, V, block_size):\n \"\"\"\n Pseudocode for Flash Attention forward pass.\n Q, K, V: (N, d) matrices\n block_size: tile size that fits in SRAM\n \"\"\"\n N, d = Q.shape\n O = np.zeros((N, d)) # output accumulator\n l = np.zeros((N, 1)) # running sum of exponentials\n m = np.full((N, 1), -np.inf) # running row max\n\n # Outer loop: iterate over K/V blocks\n for j in range(0, N, block_size):\n Kj = K[j:j+block_size] # load K block into SRAM\n Vj = V[j:j+block_size] # load V block into SRAM\n\n # Inner loop: iterate over Q blocks\n for i in range(0, N, block_size):\n Qi = Q[i:i+block_size]\n Oi = O[i:i+block_size]\n li = l[i:i+block_size]\n mi = m[i:i+block_size]\n\n # Compute local attention scores (small tile, fits in SRAM)\n Sij = Qi @ Kj.T # (block_size, block_size)\n\n # Online softmax: update running statistics\n mij = np.max(Sij, axis=1, keepdims=True)\n m_new = np.maximum(mi, mij)\n\n # Rescale previous and current exponentials\n exp_old = np.exp(mi - m_new)\n exp_new = np.exp(Sij - m_new)\n\n l_new = exp_old * li + np.sum(exp_new, axis=1, keepdims=True)\n\n # Update output: rescale old output, add new contribution\n Oi = (exp_old * li * Oi + exp_new @ Vj) / l_new\n\n # Write back updated values\n O[i:i+block_size] = Oi\n l[i:i+block_size] = l_new\n m[i:i+block_size] = m_new\n\n return O\n```\n\n\nThe critical thing to notice: at no point does this code create an N x N matrix. The largest intermediate result is `Sij`, which is `block_size x block_size`. If your block size is 256, that is a 256 x 256 tile, just 65,536 elements instead of N^2. For a 32K sequence, that is a 500,000x reduction in the size of the intermediate matrix.\n\n---\n\n## The Recomputation Trade-off\n\nStandard attention saves the full N x N attention matrix (P) during the forward pass because the backward pass needs it to compute gradients. This is the normal approach in deep learning: save intermediate results to avoid recomputing them during backpropagation.\n\nFlash Attention takes the opposite approach. It does not save P at all. During the backward pass, it recomputes the attention matrix from Q, K, and V using the same tiling strategy.\n\nThis sounds wasteful. You are literally doing the same computation twice. But here is the counterintuitive part: it is actually faster.\n\nWhy? Because attention is memory-bound. The time spent recomputing the attention tiles in SRAM is less than the time it would take to write the full N x N matrix to HBM during the forward pass and then read it back during the backward pass. You are trading cheap computation (fast SRAM arithmetic) for expensive memory access (slow HBM reads and writes).\n\nThe math works out clearly:\n\n- **Standard attention (forward + backward):** Must write and read the N x N matrix to/from HBM twice (once in forward, once in backward). Total HBM access for P alone: 4 x N^2 x bytes_per_element.\n- **Flash Attention (forward + backward):** Never writes P to HBM. Recomputes tiles in SRAM during backward pass. Total HBM access for P: zero.\n\nThe extra FLOPs from recomputation are trivial compared to the HBM bandwidth saved. On an A100, doing a 256x256 matrix multiply in SRAM takes nanoseconds. Reading 2 GB from HBM takes a millisecond. The trade-off is overwhelmingly in favor of recomputation.\n\nThis is a principle worth internalizing: on modern hardware, sometimes doing more math is faster than saving and loading results. The memory wall is that significant.\n\n---\n\n## Flash Attention 2\n\nFlash Attention 1 proved the concept, but it left performance on the table. Flash Attention 2 (also by Tri Dao, released in 2023) made several improvements that roughly doubled the speed.\n\n**Better work partitioning.** Flash Attention 1 parallelized over batch size and number of heads, but within each head, the inner loop was sequential. Flash Attention 2 also parallelizes over the sequence length dimension. On an A100 with 108 streaming multiprocessors, this means much better GPU utilization, especially for long sequences where the number of tiles along the sequence dimension is large.\n\n**Reduced non-matmul FLOPs.** Flash Attention 1 spent a surprising amount of time on operations that are not matrix multiplications: the rescaling in online softmax, the elementwise exponentials, and the bookkeeping for running statistics. Flash Attention 2 restructured the algorithm to minimize these operations. On modern GPUs, matrix multiply units (tensor cores) are vastly faster than general-purpose units, so shifting work toward matmuls and away from scalar operations is a significant win.\n\n**Swapped loop order.** Flash Attention 2 reversed the loop structure: the outer loop iterates over Q blocks, and the inner loop iterates over K/V blocks. This change reduces the amount of shared memory reads and writes, and it enables better parallelization because different Q blocks can be processed independently on different thread blocks.\n\nThe net result is that Flash Attention 2 reaches 50-73% of the A100's theoretical maximum FLOPS, compared to about 30-50% for Flash Attention 1. For a memory-bound operation, getting anywhere near peak compute utilization is remarkable.\n\n---\n\n## Flash Attention 3\n\nFlash Attention 3 (2024) targets NVIDIA's Hopper architecture (H100 and beyond) and squeezes out even more performance by exploiting hardware features that did not exist when Flash Attention 1 was designed.\n\n**Asynchronous execution.** The H100 introduced the Tensor Memory Accelerator (TMA), which can load data from HBM to SRAM asynchronously while the compute units are busy with the previous tile. Flash Attention 3 overlaps memory access with computation so that the GPU is never waiting for data. While the current tile is being processed in SRAM, the next tile is already being loaded.\n\n**FP8 support.** The H100 has native FP8 tensor cores that deliver twice the throughput of FP16. Flash Attention 3 supports FP8 attention computation, which is particularly valuable for inference where the reduced precision is acceptable. This effectively doubles the compute throughput for attention.\n\n**Warp-specialization.** Instead of having all warps (groups of 32 GPU threads) do the same thing, Flash Attention 3 assigns different roles to different warps. Some warps handle memory loading, others handle computation, and others handle softmax statistics. This producer-consumer pattern keeps all parts of the GPU busy simultaneously.\n\nThese improvements push Flash Attention 3 to 1.5-2x the speed of Flash Attention 2 on H100 GPUs.\n\n---\n\n## Impact in Practice\n\nThe numbers tell the story. Here is how Flash Attention compares to standard attention across different sequence lengths, assuming an A100 GPU with a typical transformer configuration:\n\n\n| Sequence Length | Standard Attention Memory | Flash Attention Memory | Memory Reduction | Speed Improvement |\n|----------------|--------------------------|----------------------|------------------|-------------------|\n| 2K (2,048) | ~32 MB | ~2 MB | 16x | 1.5-2x |\n| 8K (8,192) | ~512 MB | ~4 MB | 128x | 2-3x |\n| 32K (32,768) | ~8 GB | ~8 MB | 1,000x | 3-4x |\n| 128K (131,072) | ~128 GB (impossible) | ~16 MB | N/A | Enables it |\n\n\nThe memory column refers specifically to the attention matrix, not total model memory. Standard attention at 128K context would require 128 GB just for the attention matrix of a single layer, which exceeds the memory of any single GPU. Flash Attention makes it possible by never materializing that matrix.\n\nBeyond raw numbers, Flash Attention has changed what is practically achievable:\n\n- **GPT-4, Claude, Gemini** and other frontier models use Flash Attention (or equivalent fused attention kernels) to support 100K+ context windows. Without it, these context lengths would require model parallelism just for the attention computation, which is extremely expensive.\n- **Training throughput** improves by 15-25% end-to-end when Flash Attention replaces standard attention, since attention is typically 30-50% of total training time.\n- **Inference at scale** with systems like vLLM and TensorRT-LLM rely on Flash Attention to serve long-context requests without running out of memory.\n\nFlash Attention did not change the math of attention. It changed where the math happens. And that made all the difference.\n\n---\n\n## Using Flash Attention\n\nIn practice, you almost never implement Flash Attention yourself. It is a CUDA kernel that requires deep knowledge of GPU architecture to write correctly. Instead, you use it through frameworks that integrate it as a backend.\n\n**PyTorch 2.0+ (scaled_dot_product_attention)**\n\nStarting with PyTorch 2.0, the `torch.nn.functional.scaled_dot_product_attention` function automatically selects the most efficient backend, including Flash Attention when the hardware and input shapes support it.\n\n\n```python\nimport torch\nimport torch.nn.functional as F\n\n# Create Q, K, V tensors (batch, heads, seq_len, head_dim)\nbatch_size, num_heads, seq_len, head_dim = 1, 32, 8192, 128\n\nQ = torch.randn(batch_size, num_heads, seq_len, head_dim,\n device=\"cuda\", dtype=torch.float16)\nK = torch.randn(batch_size, num_heads, seq_len, head_dim,\n device=\"cuda\", dtype=torch.float16)\nV = torch.randn(batch_size, num_heads, seq_len, head_dim,\n device=\"cuda\", dtype=torch.float16)\n\n# PyTorch automatically uses Flash Attention when available\noutput = F.scaled_dot_product_attention(Q, K, V)\n\n# To explicitly check which backend is being used\nwith torch.backends.cuda.sdp_kernel(\n enable_flash=True,\n enable_math=False,\n enable_mem_efficient=False\n):\n output = F.scaled_dot_product_attention(Q, K, V)\n```\n\n\nPyTorch will use Flash Attention when: the inputs are on CUDA, the dtype is float16 or bfloat16, the head dimension is at most 256, and the GPU supports it (SM 80+ for A100, SM 90+ for H100). If Flash Attention is not available, it falls back to a memory-efficient attention implementation or standard math.\n\n**Hugging Face Transformers**\n\nHugging Face makes it even simpler. You can enable Flash Attention 2 with a single parameter when loading a model:\n\n\n```python\nfrom transformers import AutoModelForCausalLM\n\n# Load model with Flash Attention 2\nmodel = AutoModelForCausalLM.from_pretrained(\n \"meta-llama/Llama-3.1-8B-Instruct\",\n torch_dtype=torch.float16,\n device_map=\"auto\",\n attn_implementation=\"flash_attention_2\"\n)\n```\n\n\nYou need the `flash-attn` package installed (`pip install flash-attn`). Not all model architectures support it, but most popular ones do (Llama, Mistral, Falcon, GPT-NeoX, and others).\n\n**vLLM**\n\nvLLM uses Flash Attention by default for all supported models. No configuration needed:\n\n\n```python\nfrom vllm import LLM, SamplingParams\n\n# Flash Attention is used automatically\nllm = LLM(\n model=\"meta-llama/Llama-3.1-8B-Instruct\",\n dtype=\"float16\",\n max_model_len=32768 # Long context, made possible by Flash Attention\n)\n\nparams = SamplingParams(temperature=0.7, max_tokens=1024)\noutput = llm.generate(\"Explain Flash Attention in simple terms.\", params)\nprint(output[0].outputs[0].text)\n```\n\n\nvLLM also integrates PagedAttention (which you learned about in the KV Cache chapter) with Flash Attention, combining efficient memory management with efficient attention computation.\n\n---\n\n## Exercise\n\n**Exercise:** Benchmark Flash Attention vs. standard attention across sequence lengths.\n\nWrite a Python script that:\n\n1. Creates random Q, K, V tensors at three sequence lengths: 2,048, 8,192, and 32,768\n2. Runs attention using PyTorch's `scaled_dot_product_attention` with Flash Attention enabled\n3. Runs attention with Flash Attention disabled (math-only backend)\n4. For each configuration, measures:\n\n- Wall-clock time (average over 10 runs, with warmup)\n- Peak GPU memory usage (`torch.cuda.max_memory_allocated()`)\n\n1. Prints a comparison table showing the speedup and memory reduction at each sequence length\n\n**Starter code:**\n\n\n```python\nimport torch\nimport torch.nn.functional as F\nimport time\n\ndef benchmark_attention(seq_len, num_runs=10, warmup=3):\n \"\"\"Benchmark Flash vs standard attention at a given sequence length.\"\"\"\n batch, heads, dim = 1, 32, 128\n Q = torch.randn(batch, heads, seq_len, dim, device=\"cuda\", dtype=torch.float16)\n K = torch.randn(batch, heads, seq_len, dim, device=\"cuda\", dtype=torch.float16)\n V = torch.randn(batch, heads, seq_len, dim, device=\"cuda\", dtype=torch.float16)\n\n results = {}\n for name, flash, math in [\n (\"flash\", True, False),\n (\"standard\", False, True)\n ]:\n torch.cuda.reset_peak_memory_stats()\n # Warmup\n for _ in range(warmup):\n with torch.backends.cuda.sdp_kernel(\n enable_flash=flash, enable_math=math, enable_mem_efficient=False\n ):\n _ = F.scaled_dot_product_attention(Q, K, V)\n torch.cuda.synchronize()\n\n # Timed runs\n torch.cuda.reset_peak_memory_stats()\n start = time.perf_counter()\n for _ in range(num_runs):\n with torch.backends.cuda.sdp_kernel(\n enable_flash=flash, enable_math=math, enable_mem_efficient=False\n ):\n _ = F.scaled_dot_product_attention(Q, K, V)\n torch.cuda.synchronize()\n elapsed = (time.perf_counter() - start) / num_runs\n\n peak_mem = torch.cuda.max_memory_allocated() / (1024 ** 2) # MB\n results[name] = {\"time_ms\": elapsed * 1000, \"memory_mb\": peak_mem}\n\n return results\n\n# Run benchmarks\nfor seq_len in [2048, 8192, 32768]:\n r = benchmark_attention(seq_len)\n speedup = r[\"standard\"][\"time_ms\"] / r[\"flash\"][\"time_ms\"]\n mem_reduction = r[\"standard\"][\"memory_mb\"] / r[\"flash\"][\"memory_mb\"]\n print(f\"Seq {seq_len}: Flash {r['flash']['time_ms']:.1f}ms, \"\n f\"Std {r['standard']['time_ms']:.1f}ms, \"\n f\"Speedup {speedup:.1f}x, Mem reduction {mem_reduction:.1f}x\")\n```\n\n\nExpected observations: Flash Attention should be faster at all sequence lengths, with the advantage growing at longer sequences. Memory reduction should be dramatic, especially at 32K where standard attention may fail with an out-of-memory error.\n\n---\n\n## Summary\n\nFlash Attention solves the memory bottleneck in attention by restructuring the computation to avoid materializing the full N x N attention matrix in HBM.\n\nKey takeaways:\n\n- **Attention is memory-bound, not compute-bound.** The bottleneck is reading and writing the N x N matrix through slow HBM, not the matrix multiplications themselves.\n- **GPU memory hierarchy matters.** SRAM is ~10x faster than HBM but 4,000x smaller. Flash Attention keeps intermediate results in SRAM by processing attention in small tiles.\n- **Online softmax makes tiling possible.** By maintaining running max and sum statistics, Flash Attention computes exact softmax without ever seeing the full row at once.\n- **Recomputation beats saving.** Flash Attention recomputes attention during the backward pass instead of saving it, trading cheap SRAM compute for expensive HBM access. More FLOPs, less time.\n- **Flash Attention 2** doubles speed through better parallelization and reduced non-matmul operations. **Flash Attention 3** exploits H100 features like async memory access and FP8 for another 1.5-2x gain.\n- **Memory reduction is dramatic.** From O(N^2) down to O(N) in practice, enabling context lengths of 100K+ tokens that would be physically impossible with standard attention.\n- **You don't implement it yourself.** PyTorch's `scaled_dot_product_attention`, Hugging Face's `attn_implementation=\"flash_attention_2\"`, and vLLM all integrate Flash Attention as a backend.\n\nFlash Attention optimizes HOW attention is computed, making every parameter in the model work more efficiently. But what if some of the model's parameters are not needed at all? In the next chapter, we will look at pruning, which tackles that question by removing unnecessary weights entirely.","pageType":"ai-engineering"}

11.11 Flash Attention

Ashish Pratap Singh

Get Premium