Guide

Core optimizations

FlashDecoding and split-KV

Decode often has one query and a long KV cache. Split the KV across thread blocks, then reduce, so the GPU has parallel work again.

56 of 138 · 2 min read

Training and prefill attend over many query tokens at once, so n_q is large. Decode often has n_q = 1, or a small batch of queries, and a long KV cache. The original FlashAttention tiling assumes enough parallel work along the query axis. At decode that axis disappears, and a naive FlashAttention kernel underuses the GPU.

Concept

Prefill processes the whole prompt as one block of queries. Decode adds one new query token per step and reads the full cached K and V. Split-KV (also called FlashDecoding) cuts the KV length across thread blocks, then combines the partial softmax state.

The big idea

Prefill is a wide classroom taking one exam. Decode is one student against a very long answer key. FlashAttention's original seating plan needs many students. Split-KV tears the answer key into sections, grades each section in parallel, then merges the scores.

The picture breaks if you force one kernel for both phases. Prefill wants classic tiling over large Q. Decode with long KV wants parallelism along the sequence.

Walking through it

Each split keeps its own online-softmax numerator and denominator. A final reduction rescales those partials the same way online softmax rescales tiles. The math stays exact.

This matters for long-context serving at small batch sizes: exactly when interactive decode feels slow even though "FlashAttention is on." Flash-decoding and paged KV help utilization. They do not remove the O(n) read of the cache.

Serving engines switch backends by phase. Production stacks also need paged or ragged KV layouts, GQA head grouping, and fused append-KV paths. FlashInfer and vendor libraries package those variants. When you benchmark, pin the backend or log which one ran. People say "FlashAttention" when they mean FA2, FA3, FlashInfer, or a TensorRT kernel.

Common mistakes

  • Assuming FlashAttention alone fixes batch-1 decode at huge context. Prefill gets a large win. Decode still wants split-KV and careful KV memory.
  • Forcing one kernel for prefill and decode. You leave performance on the table.
  • Comparing engines without logging the attention backend. The name on the slide may not be the kernel that ran.

Check yourself

  1. Why does original FlashAttention tiling starve the GPU at decode? Answer: It parallelizes along the query axis, and decode often has n_q = 1.
  2. What do you split instead, and what must you do after? Answer: Split the KV length across blocks, then reduce the partial softmax statistics.

Key takeaways

  • Prefill has many queries. Decode has one query and a long KV cache.
  • Split-KV restores parallelism along the sequence axis.
  • The reduction is online softmax across splits, so the result stays exact.
  • Engines should pick different backends for prefill and decode.

Further learning

Next, FlashAttention-2 and 3 cover what changed in later kernels. For the serving side, see long-context serving.

FlashAttention

Why standard attention is HBM-bound, how tiling plus online softmax computes exact attention in on-chip SRAM, and what FlashAttention-2/3 and decode variants change.

Serving long contexts

Million-token prompts: quadratic prefill attention, linear KV growth, context parallelism, offloading, sparse/hybrid architectures, and how providers price the pain.

Prefill vs decode

The two phases of LLM inference. Prefill is compute-bound and sets TTFT. Decode is memory-bound and sets TPOT. They interfere when they share a GPU.