FlashAttention 1-2-3: The IO-Aware Revolution in LLM Inference

The Memory Wall That Standard Attention Hits

Standard multi-head attention (MHA) in PyTorch or HuggingFace Transformers computes S = Q @ K^T, then P = softmax(S), then O = P @ V. For a sequence of length N with d_head = 128 and h heads, the attention matrix S is of shape [h, N, N]. Every element of S must be written to HBM and then read back for the softmax pass. For N=8192, that's 8 GB of HBM traffic just for the attention matrix on a single layer. With 80 layers and a batch size of 1, that's 640 GB per forward pass, far exceeding the HBM bandwidth of any single GPU (H100 SXM: 3.35 TB/s, A100: 2.0 TB/s). The result: attention becomes memory-bound, not compute-bound. On an A100 80GB, a single forward pass of Llama 2 70B with N=8192 takes ~400 ms, of which ~300 ms is spent on attention HBM traffic. FlashAttention attacks this directly by never materializing the full N×N matrix in HBM. Instead, it tiles Q, K, V across the fast on-chip SRAM (192 KB per SM on A100, 256 KB on H100), computes partial attention scores, applies an online softmax trick, and writes only the final O block to HBM. The key insight: SRAM bandwidth is ~20 TB/s (on H100), 6x faster than HBM. By maximizing SRAM reuse, FlashAttention reduces HBM reads/writes from O(N^2) to O(N^2 * d / M) where M is SRAM size. For N=8192, d=128, M=192 KB, that's a ~10x reduction in HBM traffic.

# Standard attention HBM traffic estimate (single head, N=8192, d=128)
# Q, K, V reads: 3 * N * d * 2 bytes (FP16) = 6 MB
# S write+read: 2 * N^2 * 2 bytes = 268 MB
# P write+read: 2 * N^2 * 2 bytes = 268 MB
# O write: N * d * 2 bytes = 2 MB
# Total: ~544 MB HBM traffic per head per layer
# FlashAttention: ~50 MB per head per layer (10x reduction)

FlashAttention-1: Tiling, Recomputation, and Online Softmax

FlashAttention-1 (Dao et al., 2022) introduced three core ideas. First, tiling: Q, K, V are split into blocks of size B_r and B_c such that two blocks fit in SRAM. On A100 with 192 KB SRAM, typical block sizes are B_r = B_c = 128 for FP16. The outer loops over K,V blocks and the inner loop over Q blocks accumulate partial attention outputs. Second, recomputation: instead of storing the full softmax normalization statistics (the denominator sum of exponentials), FlashAttention stores only two scalars per row: the running maximum m and the running sum l. When a new K,V block arrives, the old m and l are used to rescale the previously accumulated output. This avoids writing the full softmax matrix to HBM. Third, online softmax: the algorithm maintains the invariant that at any point, the partial output O_partial = softmax_partial(Q @ K_partial^T) @ V_partial. When a new block arrives, it updates m <- max(m, new_max), rescales O_old by exp(m_old - m_new), adds the new contribution, and updates l. The final output is O / l. This requires only O(N) HBM writes for the output O, plus O(N) reads for Q, K, V. The total HBM access is O(N^2 * d / M), which for typical N and M is 5-10x less than standard attention. On an A100 80GB, FlashAttention-1 achieves 2-4x speedup over PyTorch's scaled_dot_product_attention for N=8192. For N=16384, the speedup grows to 6x because the standard attention HBM traffic scales quadratically while FlashAttention's scales sub-quadratically. The catch: FlashAttention-1 still uses a causal mask branch inside the inner loop, which adds non-matmul overhead. For causal attention (common in autoregressive LLMs), the mask forces the softmax to zero out upper-triangular entries, but the tiling still computes those entries before masking them, wasting compute.

FlashAttention-2: Reducing Non-Matmul Overhead and Causal Efficiency

FlashAttention-2 (Dao, 2023) made two critical improvements. First, it reordered the tiling loops: instead of iterating over Q blocks in the inner loop, it iterates over K,V blocks in the inner loop. This allows the softmax reduction to be done once per Q block rather than once per (Q, K) block pair, reducing the number of online softmax rescaling operations from O(N^2 / (B_r * B_c)) to O(N / B_r). Second, it eliminated the causal mask branch entirely. Instead of computing the full S = Q @ K^T and then masking, FlashAttention-2 simply skips the computation of S entries that are in the upper triangle. This is done by adjusting the block sizes: for each Q block, only K,V blocks with indices less than or equal to the Q block index are loaded. This halves the compute for causal attention at large N. The result: FlashAttention-2 is ~2x faster than FlashAttention-1 for causal attention on A100, and ~1.5x faster on H100. In practice, this means Llama 2 70B with 32K context runs at 18 tokens/sec on a single H100 with FlashAttention-2, vs 12 tokens/sec with FlashAttention-1. The non-matmul overhead drops from ~30% of total attention time to ~10%. FlashAttention-2 also introduced better warp-level partitioning: each warp handles one row of the Q block, and the softmax reduction is done across warps using shared memory. This reduces shared memory bank conflicts and improves occupancy. The implementation in the flash-attn repository (v2.6.1) is the default backend for PyTorch's torch.nn.functional.scaled_dot_product_attention as of PyTorch 2.5. Libraries like vLLM, TensorRT-LLM, and llama.cpp all integrate FlashAttention-2 kernels directly. For example, vLLM uses FlashAttention-2 in its PagedAttention implementation to handle KV cache blocks efficiently, achieving 95% HBM utilization on H100.

// Pseudocode for FlashAttention-2 causal: skip upper-triangular blocks
for (int q_block = 0; q_block < num_q_blocks; q_block++) {
    for (int kv_block = 0; kv_block <= q_block; kv_block++) {
        // Load Q_block from HBM to SRAM (once per q_block)
        // Load K_block, V_block from HBM to SRAM
        // Compute S = Q_block @ K_block^T (only lower-triangular entries)
        // Online softmax update
        // Accumulate O_block
    }
    // Write O_block to HBM
}

FlashAttention-3: FP8, Async Warp Overlap, and Hopper-Specific Optimizations

FlashAttention-3 (Shah et al., 2024) targets the NVIDIA Hopper architecture (H100, H200, B200) and leverages three new hardware features. First, FP8 tensor cores: H100's fourth-gen tensor cores support FP8 input with FP32 accumulation at 2x the throughput of FP16 (197 TFLOPS vs 989 TFLOPS for sparse FP8). FlashAttention-3 quantizes Q and K to FP8 on the fly, while keeping V and O in FP16 to maintain precision. This doubles the matmul throughput for the Q @ K^T step. Second, asynchronous warp-level overlap: FlashAttention-3 uses H100's asynchronous copy (cp.async) to load the next K,V block into SRAM while the current block's matmul is still executing. This hides HBM latency almost perfectly. The kernel is structured as a producer-consumer pipeline: two SRAM buffers are used, one for the current block's compute, one for the next block's prefetch. Third, FlashAttention-3 employs warp specialization: one set of warps handles data movement (loading tiles, writing results), while another set handles computation. This reduces warp stalls and improves occupancy. On H100 SXM, FlashAttention-3 achieves 1.5-2x speedup over FlashAttention-2 for N=16384 with FP16, and up to 3x with FP8. For N=32768, the speedup is even larger because the quadratic HBM traffic of standard attention becomes prohibitive. A key benchmark: Llama 3.1 405B with 128K context on 8x H100s using TensorRT-LLM with FlashAttention-3 achieves 24 tokens/sec, compared to 16 tokens/sec with FlashAttention-2. The FP8 mode reduces KV cache memory by 2x (FP8 vs FP16), which is critical for long-context inference. However, FlashAttention-3 requires Hopper GPUs, it will not run on A100 or earlier architectures. For AMD MI300X, the equivalent is the ROCm-compatible FlashAttention fork (rocm/flash-attention), which uses MFMA instructions but lacks FP8 tensor core support (MI300X has FP8 but the software stack is immature).

// FlashAttention-3 kernel structure (simplified)
__global__ void flash_attn_3(
    half* Q, half* K, half* V, half* O,
    int N, int d) {
    __shared__ half Q_sram[BLOCK_SIZE][D];
    __shared__ half K_sram[BLOCK_SIZE][D];
    __shared__ half V_sram[BLOCK_SIZE][D];
    __shared__ half K_next[BLOCK_SIZE][D];
    __shared__ half V_next[BLOCK_SIZE][D];
    
    // Producer warp: load first K,V block
    cp_async(&K_sram, &K[0], BLOCK_SIZE * D * sizeof(half));
    cp_async(&V_sram, &V[0], BLOCK_SIZE * D * sizeof(half));
    cp_async_wait();
    
    for (int kv = 0; kv < N / BLOCK_SIZE; kv++) {
        // Producer: prefetch next block asynchronously
        if (kv + 1 < N / BLOCK_SIZE) {
            cp_async(&K_next, &K[(kv+1)*BLOCK_SIZE*D], ...);
            cp_async(&V_next, &V[(kv+1)*BLOCK_SIZE*D], ...);
        }
        // Consumer warps: compute S = Q @ K^T (FP8)
        // Online softmax + accumulate
        // Wait for prefetch to complete before next iteration
        cp_async_wait();
        swap(K_sram, K_next);
        swap(V_sram, V_next);
    }
}

Real-World Performance Numbers: H100, RTX 5090, MI300X, and M3 Ultra

Let's ground this in concrete benchmarks. On a single H100 SXM (3.35 TB/s HBM3, 989 TFLOPS FP8), FlashAttention-3 with FP8 achieves 450 TFLOPs (45% utilization) for N=16384, d=128. The same kernel in FP16 achieves 250 TFLOPs (25% utilization). Standard attention in PyTorch 2.5 achieves only 40 TFLOPs for the same configuration, a 11x gap. On an RTX 5090 (expected 1.8 TB/s GDDR7, 1800 AI TOPS FP8), FlashAttention-2 should deliver ~200 TFLOPs for N=8192, enabling 70B models at 4-bit quantization to run with 32K context at ~10 tokens/sec. On AMD MI300X (5.2 TB/s HBM3, 1307 TFLOPS FP8), the ROCm FlashAttention fork (based on FlashAttention-2) achieves ~180 TFLOPs for N=8192, about 60% of H100 performance due to less mature software and smaller SRAM (192 KB vs 256 KB). On Apple M3 Ultra (800 GB/s unified memory, 64-core GPU), the MLX framework's FlashAttention implementation (based on FlashAttention-1) achieves ~80 TFLOPs for N=8192, limited by Metal's lack of warp-level primitives. For context, running Llama 3.1 70B at 4-bit with 128K context requires ~70 GB of KV cache (FP8). On an H200 with 141 GB HBM3e, this fits with room for weights. Without FlashAttention, the same model would need ~140 GB of KV cache (FP16) and would be HBM-bandwidth-bound at ~5 tokens/sec. With FlashAttention-3, it runs at ~18 tokens/sec. The key takeaway: FlashAttention is not optional for long-context inference, it's a requirement.

Integration in Major Inference Engines: vLLM, TensorRT-LLM, llama.cpp, ExLlamaV2

FlashAttention is not a standalone library, it's a kernel that must be integrated into the broader inference stack. vLLM (v0.6.0+) uses FlashAttention-2 as the default attention backend for all GPU types. It integrates with PagedAttention by treating each KV cache page as a block that FlashAttention can tile over. This allows vLLM to serve Llama 3.1 405B with 128K context at 24 tokens/sec on 8x H100s (as of April 2026). TensorRT-LLM (v0.14.0) uses FlashAttention-3 for Hopper GPUs and FlashAttention-2 for Ampere. It also supports FP8 KV cache via FlashAttention-3, reducing memory by 2x. The TensorRT-LLM implementation is hand-tuned for each GPU architecture, achieving 95% of theoretical roofline on H100. llama.cpp uses a custom implementation called ggml_flash_attn that is based on FlashAttention-1 but optimized for CPU and GPU backends (CUDA, Metal, Vulkan). For CPU inference on Sapphire Rapids with AVX-512, ggml_flash_attn achieves 2x speedup over naive attention for N=4096. ExLlamaV2 (v0.2.0+) uses FlashAttention-2 for its GPTQ/AWQ inference, with support for 4-bit KV cache (via quantization) combined with FlashAttention tiling. This allows 70B models at 4-bit to run with 64K context on a single RTX 4090 (24 GB VRAM) at ~8 tokens/sec. For Apple Silicon, MLX (v0.18.0) provides mlx.core.fast.attention which implements FlashAttention-1 with Metal performance optimizations. On M3 Ultra (192 GB unified memory), this enables running Llama 3.1 405B (4-bit) with 32K context at ~3 tokens/sec, slow but feasible. The common pattern: every inference engine has adopted FlashAttention because the memory savings are too large to ignore.

# Using FlashAttention in vLLM (v0.6.0+)
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Meta-Llama-3.1-405B",
    tensor_parallel_size=8,
    max_model_len=131072,
    dtype="bfloat16",
    kv_cache_dtype="fp8",  # Enables FlashAttention-3 FP8 path
)

params = SamplingParams(temperature=0.7, max_tokens=1024)
output = llm.generate("Explain FlashAttention in one paragraph.", params)
print(output[0].outputs[0].text)

The Math Behind the Online Softmax Trick

The core innovation that enables FlashAttention is the online softmax algorithm. Standard softmax for a row vector x of length N computes m = max(x_i), then s_i = exp(x_i - m), then l = sum(s_i), then p_i = s_i / l. This requires two passes over x: one to find m, one to compute s and l. In FlashAttention, we process x in blocks. Let x^(1) and x^(2) be two blocks. For the first block, we compute m1 = max(x^(1)), s1 = exp(x^(1) - m1), l1 = sum(s1). The partial output O1 = s1 @ V1 (where V1 is the corresponding V block). For the second block, we compute m2 = max(x^(2)), then the combined max m = max(m1, m2). We then rescale the first block's statistics: s1' = exp(x^(1) - m) = s1 * exp(m1 - m), l1' = l1 * exp(m1 - m). For the second block: s2' = exp(x^(2) - m), l2' = sum(s2'). Then the new total l = l1' + l2'. The new output O = (O1 * exp(m1 - m) + s2' @ V2) / l. This requires only two scalars per row (m and l) to be kept in registers, not the full N-length softmax vector. The rescaling factor exp(m1 - m) is a single float per row. This is the key to avoiding O(N^2) HBM writes. The algorithm is numerically stable because it always subtracts the maximum exponent, preventing overflow. The error introduced by the online update is on the order of machine epsilon (1e-7 for FP32), negligible for transformer training and inference. In FlashAttention-3, the online softmax is further optimized by using FP32 accumulation for the exponentials even when the inputs are FP8, maintaining precision.

When FlashAttention Doesn't Help: Short Sequences and Batch Size > 1

FlashAttention's benefits are most pronounced for long sequences (N > 2048) and batch size 1 (or small batch). For short sequences (N < 512), the overhead of tiling and the online softmax rescaling can outweigh the HBM savings. On an H100, for N=256, FlashAttention-2 is only ~1.1x faster than standard attention because the attention matrix fits entirely in SRAM anyway. For large batch sizes (B > 8), the Q, K, V matrices become large, and the tiling overhead increases because each query in the batch must be processed independently. FlashAttention's speedup drops from 10x at B=1, N=16384 to ~2x at B=16, N=16384. This is because the HBM traffic for standard attention scales as O(B * N^2), but FlashAttention's traffic scales as O(B * N^2 / M) where M is SRAM per SM. With B=16, the Q matrix alone is 16 * 16384 * 128 * 2 bytes = 64 MB, which doesn't fit in SRAM (256 KB). So FlashAttention must tile over the batch dimension as well, increasing the number of kernel launches and reducing efficiency. For training, where batch sizes are typically 32-128, FlashAttention still helps but the speedup is modest (1.5-2x). For inference with continuous batching (vLLM style), the effective batch size is often 1-4, making FlashAttention ideal. Another limitation: FlashAttention requires the head dimension d to be a multiple of 128 (for FP16) or 64 (for FP8). Models with non-standard head dimensions (e.g., d=96 in some GPT variants) require padding or custom kernels, which reduces efficiency. Finally, FlashAttention-3's FP8 mode requires hardware support for FP8 tensor cores (H100, B200, MI300X). On older GPUs (A100, RTX 4090), FP8 is not available, so FlashAttention-3 falls back to FP16, losing the 2x throughput advantage.

# Benchmark: FlashAttention vs standard attention (H100, FP16, N=8192)
import torch
from flash_attn import flash_attn_func

q = torch.randn(1, 32, 8192, 128, device='cuda', dtype=torch.half)
k = torch.randn(1, 32, 8192, 128, device='cuda', dtype=torch.half)
v = torch.randn(1, 32, 8192, 128, device='cuda', dtype=torch.half)

# Standard attention
import time
start = time.time()
out_std = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
torch.cuda.synchronize()
print(f"Standard: {time.time() - start:.4f}s")

# FlashAttention-2
start = time.time()
out_fa2 = flash_attn_func(q, k, v, causal=True)
torch.cuda.synchronize()
print(f"FlashAttention-2: {time.time() - start:.4f}s")

Beyond FlashAttention: Future Directions and Alternatives

FlashAttention is not the end of the story. Several research directions aim to push attention efficiency further. First, FlashDecoding (2023) optimizes the decoding phase (where N=1 for query but KV cache is long) by parallelizing over the KV cache length using a separate kernel that computes partial softmax and then reduces. This is orthogonal to FlashAttention and is used in vLLM and TensorRT-LLM. Second, PagedAttention (2022) manages the KV cache as fixed-size blocks, enabling virtual memory-like paging and reducing fragmentation. FlashAttention integrates with PagedAttention by treating each page as a tile. Third, state space models (Mamba, Mamba-2) replace attention entirely with a recurrent scan that is O(N) in memory and compute. Mamba-2 achieves 5x speedup over FlashAttention-2 on long sequences (N=65536) on H100, but at the cost of quality on certain tasks (e.g., long-range retrieval). Fourth, sparse attention methods (e.g., BigBird, Longformer, SparseGPT) reduce the attention matrix to O(N log N) entries, but require custom hardware or software support for irregular sparsity. FlashAttention-3's FP8 mode is a step toward mixed-precision attention, and future GPUs (B200 with 2x HBM3e bandwidth) will further reduce the memory bottleneck. On the software side, the Triton language (OpenAI) allows writing custom FlashAttention kernels that are portable across NVIDIA and AMD GPUs. The Triton implementation of FlashAttention-2 (in the Triton tutorial) achieves ~80% of the performance of the hand-tuned CUDA version, and is used in projects like MLC-LLM for AMD GPUs. For NPUs (e.g., Intel Gaudi 3, Qualcomm Cloud AI 100), FlashAttention is not yet available due to the lack of warp-level programming models. Instead, these NPUs rely on tiling in the host compiler (OpenVINO, DirectML) to approximate FlashAttention's behavior. The gap between GPU and NPU attention efficiency remains large, typically 2-4x in favor of GPUs for long-context inference.

Further reading

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• FlashAttention-3: Fast and Accurate Attention with Asynchronous Overlap and FP8
• vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention
• NVIDIA TensorRT-LLM Documentation