Navigation
← Blog⚡ AI Infra

Attention at Scale — From FlashAttention to FlashAttention-4

· 11 min read· AI Infra Deep Dives
inferencetraininggpukv-cacheserving

TL;DR

Attention is the single most expensive operation in a transformer. It’s also where the most performance is left on the table. FlashAttention (2022) showed you could make it 2-4x faster by respecting GPU memory hierarchy — and since then, every new GPU generation has required a completely rewritten kernel to stay efficient. This isn’t a one-time optimization. It’s an ongoing arms race between hardware evolution and kernel engineering, with each generation introducing new bottlenecks and new hardware features to exploit. The stack now has three distinct layers: hardware kernels, serving-specific adaptations, and programmable abstractions.

The Problem

Here’s the fundamental tension: attention requires every token to “look at” every other token. For a sequence of N tokens, that’s N² attention scores to compute. On a 128K-context input, that’s 16 billion score computations. Per layer. Per head.

The naive implementation:

  1. Compute the full N×N attention score matrix (Q × K^T)
  2. Apply softmax
  3. Multiply by Values (scores × V)

For 128K tokens at BF16: the score matrix alone would consume 32 GB of GPU memory. An H100 has 80 GB total. You’d fill half the GPU’s memory just with intermediate scratch space for ONE layer.

Obviously nobody does this. But the alternative — chunking the computation to fit in faster on-chip memory — is where all the complexity lives.

If you’ve ever optimized code by restructuring data access to fit in CPU L1/L2 cache (loop tiling, blocking), this is the same game. Except the “cache” is 256 KB of SRAM per streaming multiprocessor, the “main memory” is 80 GB of HBM running at 3.35 TB/s, and the penalty for a “miss” (reading from HBM when you could have read from SRAM) is a 10-50x latency hit.

Why This Is a Kernel Engineering Problem, Not a Model Problem

Here’s the thing that surprised me coming from a systems background: the model architecture hasn’t changed. Attention is still Q×K^T, softmax, ×V. The math is identical since 2017.

What’s changed is how you execute that math on specific hardware. And it changes with every GPU generation — because each generation has different memory hierarchies, different instruction sets, and different bottlenecks.

This is pure systems optimization. It’s like the difference between a naive SQL query and one that’s been tuned for a specific storage engine’s buffer pool size, page layout, and I/O characteristics. Same logical operation, completely different physical execution.

The Evolution

FlashAttention (Dao, 2022) — The breakthrough: tile attention to fit in SRAM

The insight: Don’t materialize the N×N attention matrix in HBM. Instead, compute attention in tiles that fit in on-chip SRAM (shared memory), streaming through the KV sequence in blocks.

The trick that makes this possible: you can compute softmax incrementally (the “online softmax” trick). You don’t need to see all scores before normalizing — you can update a running max and sum as you process each block. This eliminates the need to store the full matrix.

Result: 2-4x speedup, memory usage from O(N²) to O(N). Made long-context training practical.

The infra analogy: it’s like switching from “load the entire table into memory, then sort” to “stream through the table in chunks, maintaining a running aggregate.” Same result, fraction of the memory, and faster because you’re working in fast cache instead of slow main memory.

FlashAttention-2 (Dao, 2023) — Better parallelism, better occupancy

Restructured the parallelism to reduce shared memory reads/writes. Achieved 50-73% of theoretical peak FLOPS on A100 (~225 TFLOPS/s).

Still left ~65% of H100 performance on the table, though — because it didn’t use any Hopper-specific hardware.

FlashAttention-3 (Shah et al., 2024) — Hopper-native: exploit the hardware

Target: H100/H200 (Hopper architecture)

The problem FA3 solves: FA2 achieves only ~35% utilization on H100 because it treats memory movement and computation as sequential steps — load a tile, compute on it, load the next. But Hopper has dedicated hardware (TMA — Tensor Memory Accelerator) that can move data in the background without using any compute threads.

The solution: warp specialization

Instead of all threads doing “load, then compute, then load, then compute” — split them into dedicated roles:

Producer warps: issue TMA loads (data movement, async, hardware-accelerated)
Consumer warps: run matmuls on tensor cores (compute)

Both run CONCURRENTLY on the same SM.

It’s the same principle as I/O threads vs compute threads in a server — one group fetches data from storage while the other group processes data already fetched. You never wait.

Results: 1.5-2x over FA2 on H100. Up to 740 TFLOPS/s (~75% utilization). Also supports FP8 (1.2 PFLOPS/s) with a block quantization scheme that keeps numerical error low.

FlashAttention-4 (Zadouri et al., March 2026) — Blackwell-native: work around lopsided hardware

Target: B200/GB200 (Blackwell architecture)

The problem FA4 solves: Blackwell doubled tensor core FLOPS compared to Hopper. But shared memory bandwidth and special function units (SFU — used for exp, log in softmax) did NOT double. The hardware scaled asymmetrically.

On Blackwell, FA3’s design becomes bottlenecked on:

  1. Shared memory bandwidth — tensor cores finish matmuls faster than data can flow in/out of SMEM
  2. SFU throughput — the hardware exponential unit (used in softmax) can’t keep up with doubled matmul speed

Both are non-matmul bottlenecks. The tensor cores are fast enough — everything else is too slow.

The solutions:

TMEM (Tensor Memory): Blackwell introduced a new register-like memory (256 KB) directly accessible by tensor cores, bypassing shared memory. FA4 routes intermediate results through TMEM instead of SMEM, relieving the bandwidth bottleneck.

Software-emulated exp2(): Instead of using the hardware SFU for exponentials (which is now the bottleneck), FA4 emulates the exponential function using integer bit manipulation that runs on general CUDA cores — turning a “special function bottleneck” into work that the (now faster) general compute units can absorb. Plus conditional softmax rescaling to reduce the total number of exp calls needed.

2-CTA MMA mode: Two cooperative thread arrays share a larger tile, better utilizing the doubled tensor core throughput.

CuTe-DSL rewrite: The entire kernel is rewritten in Python-embedded CuTe-DSL instead of C++ CUTLASS templates. Compile times drop 20-30x. Same performance, dramatically faster development iteration.

Results: Up to 1613 TFLOPS/s on B200 (~71% utilization). 1.3x over cuDNN 9.13 (which already includes FA3-era optimizations).

The lesson for infra engineers: Every GPU generation shifts the bottleneck. On A100: memory bandwidth was the constraint. On H100: the constraint moved to memory-compute overlap. On Blackwell: the constraint moved to non-matmul operations and register file pressure. The kernel must be rewritten each time — there’s no “write once, run fast forever.” If you’re planning hardware refreshes, budget for kernel re-engineering.

Beyond Training: Attention Kernels for Serving

FlashAttention 1-4 optimize for training — large, contiguous tensors, uniform batch shapes, maximum throughput.

Serving has different requirements:

  • KV caches are paged (non-contiguous blocks scattered in memory)
  • Batches mix different request lengths (short chats alongside long documents)
  • Decode is the dominant phase (single-query attention against large KV caches)
  • You need to handle PagedAttention’s block table indirection

FlashInfer (MLSys 2025) — Attention kernels designed for serving

The problem: FlashAttention assumes contiguous tensors. In vLLM-style serving, the KV cache is scattered across non-contiguous pages (16-token blocks) managed by a page table. Naively, you’d gather all pages into a contiguous buffer before running attention — but that copy defeats the purpose.

Page-Asymmetric Tiling: FlashInfer tiles the query side with fixed-size tiles but adapts the KV side to the actual page structure. It streams through pages directly (using CSR-format page indices loaded into shared memory), avoiding gather copies entirely.

Handling heterogeneous batches: Uses ragged tensor format (CSR-style indptr arrays) — all requests packed contiguously without padding. A load-balanced scheduling algorithm distributes work across SMs based on actual request lengths rather than worst-case padding.

Composable attention: Decompose attention into sub-problems (shared prefix vs unique suffix) using an associative merge operator. Compute parts independently, merge results. This enables prefix sharing without duplicating KV memory.

Results: 29-69% latency reduction in serving benchmarks. Integrated into SGLang, vLLM, and MLC-Engine.

Why this matters: FlashInfer is the reason vLLM can serve paged KV caches without massive attention overhead. Without it (or something like it), PagedAttention’s fragmented layout would make attention kernels much slower than the contiguous case. It bridges the gap between the memory management layer (PagedAttention) and the compute layer (FlashAttention).

Sequence-Aware Split Heuristic (Font et al., May 2026)

The problem: Modern models use GQA/MQA (Grouped-Query / Multi-Query Attention) — fewer KV heads to save memory. But fewer heads means fewer parallel thread blocks during decode. On an H100 with 132 SMs: 8 KV heads × batch 8 = 64 blocks. That’s 68 idle SMs. Half the GPU sitting dark.

FA3/FA4’s heuristic decides whether to split sequences across multiple blocks based on sequence length alone. It says “short sequences don’t benefit from splitting.” But it ignores that the GPU is massively under-occupied due to low head count — making splitting worthwhile even for moderate sequences just to fill idle SMs.

The fix: Force parallel sequence-level splits in low-head-count regimes, even when the standard heuristic wouldn’t. The extra reduction pass is cheap compared to the utilization gain.

Result: 21-24% decode kernel efficiency recovery. For production decode (where most tokens/second throughput is determined), that’s a direct cost reduction.

The Abstraction Layer: Making Custom Attention Patterns Easy

FlexAttention (PyTorch, 2025)

The problem: Every attention variant (sliding window, ALiBi, causal, document masking, prefix-LM) previously required a handwritten CUDA kernel. Want to try a new masking pattern? Write 500 lines of Triton, debug for a week, get 70% of optimal performance.

FlexAttention’s approach: Express attention patterns in a few lines of Python:

  • score_mod: arbitrary function applied to attention scores (ALiBi bias, soft-capping, etc.)
  • mask_mod: boolean function determining which positions attend to which
  • BlockMask: precomputed block-level sparsity map (skip fully-masked blocks entirely)

The PyTorch compiler fuses these into the FlashAttention inner loop — no materializing the full attention matrix, no kernel engineering required.

Performance: ~90% of handwritten FlashAttention2 kernels. Automatic backward pass. Block sparsity gives ~2x speedup on causal masks by skipping empty blocks.

Why this matters for serving: As model architectures diversify (sliding window, hybrid attention patterns, prefix caching with different mask shapes), the serving stack needs to handle them without per-architecture kernel engineering. FlexAttention makes new attention variants a config change, not a kernel rewrite.

DeepSeek MLA — Compressing the problem away

A different approach entirely: instead of making attention faster on the same data, reduce the data.

MLA down-projects Key/Value into a low-rank latent vector during prefill, then up-projects on-the-fly during decode. The KV cache shrinks to ~1/6th of GQA size. Less data to attend over = less memory movement = faster attention, even with a simpler kernel.

This is architectural optimization vs. kernel optimization — and in practice you need both.

TornadoAttention (ICLR 2026) — Sparsity through reordering

For video models (Diffusion Transformers): video attention has natural spatio-temporal locality (nearby frames and spatial regions attend to each other more). TornadoAttention discovers this structure offline, then physically reorders token layout in memory to place related tokens contiguously along the matrix diagonal.

Result: sparse attention patterns become dense, contiguous tiles — no dynamic GPU thread divergence, no scatter-gather. Hardware sparse acceleration “for free.” Up to 70% of DiT runtime is attention, so this matters enormously for video generation.

The infra analogy: it’s like reordering rows in a database table to match common query access patterns (clustered index). The logical data is the same, but physical layout determines I/O efficiency.

The Full Stack (How It All Fits Together)

┌─────────────────────────────────────────────────────────┐
│  Layer 3: Abstractions                                   │
│  FlexAttention, MLA, TornadoAttention                   │
│  "Express what you want, compile to fast kernels"        │
├─────────────────────────────────────────────────────────┤
│  Layer 2: Serving Kernels                                │
│  FlashInfer, Sequence-Aware Split                        │
│  "Handle paged KV, heterogeneous batches, low occupancy" │
├─────────────────────────────────────────────────────────┤
│  Layer 1: Hardware Kernels                               │
│  FlashAttention-2 (A100) → FA3 (H100) → FA4 (B200)     │
│  "Maximum throughput on specific silicon"                 │
└─────────────────────────────────────────────────────────┘

Each layer solves a different problem:

  • Layer 1 squeezes maximum FLOPS from the hardware (rewrites per GPU generation)
  • Layer 2 adapts to serving-specific memory layouts and workload patterns
  • Layer 3 lets you change attention logic without touching kernels

They stack: FlexAttention compiles down to FlashInfer-style tiled kernels, which use FA3/FA4-level hardware primitives underneath.

What’s Still Unsolved

The “rewrite every generation” problem. FA2 → FA3 → FA4 are each nearly complete rewrites. Every new GPU architecture shifts the bottleneck (SMEM bandwidth, SFU throughput, register pressure) requiring new strategies. There’s no portable abstraction that generates optimal kernels across hardware — CuTe-DSL helps with development speed but doesn’t eliminate the redesign.

Decode occupancy on large models. The Sequence-Aware Split fixes one case (low head count), but decode occupancy remains poor when KV cache per request is large relative to SRAM capacity. The fundamental issue: decode is one query against a large KV cache — inherently low arithmetic intensity. Making GPUs efficient at this requires architectural change (more memory bandwidth) not just kernel tricks.

Heterogeneous attention in one batch. Some requests use sliding window, others use full causal, others have custom masks. Batching them together means either: (a) run separate kernels per mask type (loses batching efficiency), or (b) use FlexAttention’s block mask (adds overhead per-block). No solution is overhead-free.

Cross-layer KV sharing and attention. Some architectures share KV caches across layers or use cross-attention between layers. Current kernel stacks assume per-layer, per-head independence. Violating that assumption requires rethinking the tiling strategy.

References

  1. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022) — https://arxiv.org/abs/2205.14135
  2. FlashAttention-2: Faster Attention with Better Parallelism (Dao, 2023) — https://arxiv.org/abs/2307.08691
  3. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (Shah et al., 2024) — https://arxiv.org/abs/2407.08608
  4. FlashAttention-4: Hardware-Adaptive Attention for Blackwell GPUs (Zadouri et al., 2026) — https://arxiv.org/abs/2603.05451
  5. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving (Li et al., MLSys 2025) — https://arxiv.org/abs/2501.01005
  6. Sequence-Aware Split Heuristic for Efficient LLM Decoding (Font et al., 2026) — https://arxiv.org/abs/2604.00028
  7. FlexAttention (PyTorch / Meta AI, 2025) — https://pytorch.org/blog/flexattention/
  8. TornadoAttention: Efficient Attention for Video Diffusion Transformers (Zhang, Tang et al., ICLR 2026) — https://openreview.net/forum?id=fx3Ht2gYTR