Kernels · Softmax Attention

How Flash Attention Speeds Up And by how much

Attention waits on HBM. Keep scores in registers and that wait goes away. Profiled on one B200.

Where this kernel is used

Default softmax kernel in PyTorch SDPA, vLLM, SGLang — Qwen3, Llama, Gemma, and other GQA models. Not linear attention.

1.How

FLOPs stay the same. The speedup is skipping the score-matrix round-trip through HBM.

Reminder — one head:

The big matrix is scores
\[ \mathrm{scores}=QK^{\top} \] \[ O=\mathrm{softmax}(\mathrm{scores}/\sqrt{d})\,V \]

scores: queries × keys  ·  grows with sequence²

Naive · scores in HBM

queries · keys
↓ write
full score matrix
↓ read · softmax
weights · values
output

Flash · scores in registers

query tile stays on-chip
↓ walk key tiles
one score tile
↓ fold in · discard
running max, sum, output
output

The wait is the red matrix.

Naive: every score written.

Flash: one live tile. The rest never exist.

in HBM live tile already folded away

2.How much

Qwen3.8-27B, one softmax layer, batch 1, 24 query heads, float32 scores. Flash only runs on these 16 GQA layers — the other 48 are Gated DeltaNet. Move sequence length.

2,048

HBM · off-chip DRAM

Entire score matrix

384 MB

Naive writes it here, then reads it back.

Registers · on-chip

One score tile

32 KB

Flash overwrites this every key tile. Size does not grow.

Naive score matrix 384 MB written to HBM, then read back
Flash extra for scores 0 the 32 KB tile is overwritten, not allocated in HBM
Extra HBM traffic 768 MB write scores + read scores
Flash score traffic 0 queries, keys, values, output still move
Naive scores
384 MB
Flash tile
32 KB

At 8k tokens one softmax layer is 6 GB of scores. Sixteen such layers in this model.

3.The tile walk

One query, four keys. Already-scaled scores \([1,\ 3,\ 0,\ 2]\). Values \([10,\ 20,\ 30,\ 40]\). Same softmax as writing the full row — one score at a time.

Filled = live score. Dashed = already folded into the running output.

Q stays on-chip. K and V stream. The score slot is overwritten.

Query
stays
reused all tiles
Key
loaded this tile
Value
loaded this tile
Score
registers, then gone

Only the query is reused. Scores never accumulate. After the last key, \(o \approx 24.2\) — the same mix as a full-row softmax.

The fold-in is online softmax: keep a running max, a running sum of exponentials, and a running weighted mix of values. A new tile updates those three. If its max is bigger, rescale the old sum and mix first, then add. You never need the rest of the row.

Online softmax, one update
\[ \begin{aligned} m' &= \max(m,x)\\ \ell &\leftarrow e^{m-m'}\ell + e^{x-m'}\\ O &\leftarrow e^{m-m'}O + e^{x-m'}v \end{aligned} \]

\(x\) this score, \(v\) this value. \(m\) running max, \(\ell\) running exp-sum, \(O\) running mix. Then \(o = O/\ell\). Flash does this per key-tile. The algebra is in the guides below.

Milakov & Gimelshein, 2018 — the original online-softmax paper Ceyxasm — FlashAttention, with the three-step merge Datta — why tiling plus that merge is still exact
query tile
key tiles →

Same loop on a real block: one query tile, many key tiles.

Mental model

Speedup is HBM trips, not FLOPs. Q is reused. K and V stream. Scores are a tile in registers, folded in with online softmax. Hopper and Blackwell only change the multiplier for the same walk.

4.Performance

First, the clock. Lower is faster. Every dot is the mean latency for the same Qwen3.8-27B full-attention shape on one B200.

We compare four implementations: an eager PyTorch baseline, PyTorch scaled dot-product attention (SDPA), FlashAttention-4, and FlashInfer. FlashInfer is tested in both phases: its fused single-prefill kernel for prefill and its tensor-core single-decode kernel for decode.

Model dimensions that matter here: hidden size 5,120; 64 layers, of which 16 are full attention; 24 query heads; 4 key/value heads; head dimension 256; six query heads share each key/value head.

Prefill latency

NaivePyTorch SDPA FlashAttention-4FlashInfer

CUDA-event mean latency in milliseconds · logarithmic y-axis · bfloat16 · batch 1 · 24 query heads · 4 key/value heads · head dimension 256. Hover for the exact value.

Measured prefill gain

At 4k tokens, FlashAttention-4 takes 0.190 ms: 27.8× faster than the eager materializing baseline at 5.291 ms, and 3.4× faster than PyTorch SDPA at 0.639 ms. At 8k, it is 0.599 ms versus 2.279 ms for SDPA — 3.8× faster.

What gives prefill the speedup?

1
Materializing attention → fused attention

Stop sending the full score matrix through HBM

The eager baseline launches 12 GPU kernels, explicitly expands grouped key/value heads, and writes the growing score matrix out before reading it back. At 4k, 91.5% of its instrumented time is outside the two matrix multiplies — in expansion, masking, conversion, softmax, and layout work. FlashAttention fuses the walk from Section 3 into one kernel; together those avoided costs produce the measured 27.8× gain over this diagnostic baseline.

2
Fused SDPA → Blackwell FlashAttention-4

Use fewer on-chip instructions to do the same work

SDPA is already fused, so both kernels avoid the full score matrix. At 4k they move about the same HBM bytes — 79 MB for FA4 and 84 MB for SDPA — and neither reaches 4% of peak HBM throughput. The extra 3.4× gain is therefore on-chip, not another HBM saving.

Nsight Compute summary showing twelve separate eager prefill kernels

The eager baseline is visibly twelve kernels. Score conversion, masking, standalone softmax, output conversion, and grouped-head expansion consume 91.5% of its profiled duration; the two matrix multiplies take only 0.60 ms.

Nsight Compute Summary showing one fused FlashAttention-4 prefill kernel row

The comparison is now literal: the eager call above has twelve result rows; fused FlashAttention-4 has one. Every profiler section below describes this one complete launch.

The profiler shows what changed. At 4k, these rows are FA4 / SDPA:

Instructions per cycle1.28 / 1.27same issue rate
Executed instructions61M / 186M3.06× fewer
Nsight duration0.36 / 1.08 ms2.98× shorter

The issue rate is the same; the amount issued is not. FlashAttention-4's B200 path keeps tensor-core accumulators in Tensor Memory and uses asynchronous matrix operations. Its ordinary load/store pipeline is active for 0.43% of elapsed cycles, versus 22.17% for SDPA. Fewer staging instructions produce nearly the same 3× ratio as the measured time.

What the profiler sections show

Speed of Light: fewer cycles, not a higher clock

Nsight Compute GPU Speed of Light section for FlashAttention-4 prefill

Speed of Light reports each hardware path as a percentage of its peak while the kernel runs.

Compute Workload: one-third the instruction stream

Nsight Compute Compute Workload Analysis for FlashAttention-4 prefill

The left chart counts cycles occupied; the right counts issued instructions. Tensor Memory keeps dedicated hardware busy with comparatively few instructions.

Why Blackwell-native matters

Tensor Memory is a dedicated ~256 KB on-chip store per streaming multiprocessor for fifth-generation tensor cores. A thread can launch a large asynchronous matrix multiply whose accumulators land there instead of being distributed through thread registers. The profile links that path to fewer instructions; an otherwise-identical non-Tensor-Memory kernel would be needed to isolate its exact contribution.

Memory Workload: HBM is not the remaining prefill gap

Nsight Compute Memory Workload Analysis for FlashAttention-4 prefill

This section separates byte rate, cache hits, request bandwidth, memory-instruction pipes, and spilling.

Warp State: longer waits can still win

Nsight Compute Warp State Statistics for FlashAttention-4 prefill

The bars divide the average cycles between issued warp instructions by the reason the next instruction could not issue.

Launch and Occupancy: prefill has enough blocks

Nsight Compute Occupancy section for PyTorch SDPA prefill

Occupancy is resident warps as a percentage of the hardware maximum. It is a latency-hiding resource, not a performance score.

Decode is a different amount of work

Decode latency

NaivePyTorch SDPAFlashAttention-4 FlashInfer

Each dot is the CUDA-event mean latency for one new query token over the existing key/value cache. Hover for the exact value.

At a 32k cache, FlashInfer takes 0.039 ms; SDPA takes 0.054 ms; FA4 takes 0.307 ms; the eager baseline takes 0.862 ms.

The cause is parallelism: FA4 launches 48 blocks for 148 SMs. FlashInfer splits the cache across 296 blocks, enough for two waves across the chip. One-token decode is too small for FA4's large prefill-shaped work unit.

Nsight Compute summary showing FlashInfer decode main and merge kernels

FlashInfer decode is two launches: a 296-block attention kernel (45.63 µs in Nsight) and a 24-block merge (11.17 µs). SDPA uses a 256-block main kernel plus a 6-block combine.

What the measurements say

Prefill: FlashAttention wins first by eliminating the score-matrix round-trip, then FA4 gains again from Blackwell Tensor Memory and fewer staging instructions. Decode: use a kernel that exposes enough parallel work; the large FA4 prefill kernel underfills the GPU.

References

How to cite this post

Dong, S. (2026). How Flash Attention Speeds Up — and By How Much.
https://simondong1.github.io/flash-attention.html