Training and serving attention variants
Select an attention schedule from training, prefill, or decode requirements and the KV-cache layout
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Prefill versus decode
- IO-aware attention
- KV cache
Record three attention calls from one transformer deployment. Training may present 2048 query positions and require saved information for backward propagation. Prompt prefill may present 8192 query positions and write key-value state for future tokens. Decode may present one new query position per active request while reading thousands of earlier key and value positions from a paged cache. A profiler reports very different matrix shapes and memory traffic even though every call implements scaled dot-product attention.
The query length, retained key length, number of query heads, number of key-value heads, cache datatype, mask, and physical cache layout identify an attention workload. The recorded values determine whether the kernel has enough query work for square tiles, whether cache reads dominate, and whether address translation is required. The lesson derives kernel selection from the workload fields rather than from a generic label such as attention.
Attention performance is defined by its variant and phase. Training needs forward and backward, often dropout and saved or recomputed statistics. Prefill processes many query positions together. Decode usually has one new query per sequence and reads an ever-growing KV cache. Grouped-query and multi-query attention reduce KV heads, changing memory and communication without reducing query-head work in the same way.
A kernel name that omits these conditions is not a useful unit of comparison. The relevant shape tuple includes batch or packed token count, query length, key length, query heads, KV heads, head dimension, mask, dtype, and cache layout. Each selects a different balance of matrix efficiency, reduction, memory bandwidth, and launch overhead.
Follow one request through its lifetime. Prompt tokens first pass through prefill, which can process many query positions together. Each later decoding iteration normally adds one query position per active sequence and reads all keys and values required by that position. The unchanged equation therefore produces two different kernel regimes.
Training attention
Training processes many query positions and requires gradients. The forward pass may include dropout and save compact normalization state. The backward pass computes gradients for Q, K, and V and can recompute score blocks. Large batches and sequences provide substantial tile parallelism, but activation capacity and backward traffic are major constraints.
Packed sequences remove computation on padding tokens, but each sequence has a different valid range. The kernel must map packed token positions to sequence boundaries and prevent cross-sequence attention. Distributed context or sequence parallelism partitions the sequence dimension and introduces communication for remote keys, values, or partial results. A single-device kernel timing omits this cost.
Prefill attention
Prefill processes the prompt. The prefill phase has many queries and keys, so tiled matrix operations and FlashAttention-style reuse apply. The output produces model states for every prompt token, and the serving system also writes keys and values into the cache for later decoding. Prefill latency contributes to time to first token.
A very long prompt can occupy the GPU for enough time to delay decode work from other requests. Serving systems may divide prefill into chunks. Chunking changes the query and key shapes presented to the attention kernel and adds scheduler boundaries. Evaluate the kernel with the scheduler's real chunk sizes rather than only with the unchunked prompt length.
Decode attention
During a typical decode iteration, each active sequence contributes one new query token. The new query must attend to a cache that grows with context length. The query dimension is small, so large training tiles leave many lanes unused. The kernel more readily becomes limited by reading keys and values, following cache metadata, and combining partial reductions.
Batching several active sequences creates more parallel work, but their context lengths can differ. Some implementations partition a long sequence across several blocks and reduce partial softmax states. Others assign one block to a query-head group. The suitable mapping depends on active batch, context distribution, head dimension, number of KV heads, and cache format.
Grouped-query attention and KV reuse
In multi-head attention, each query head has its own key and value head. Grouped-query attention uses fewer key-value heads, and multi-query attention can use one key-value head for all query heads. Fewer key-value heads reduce cache size and the bytes read during decode. Query projections and output work still depend on the number of query heads.
The kernel should load one K/V head and reuse it across associated query heads when its thread mapping permits. If each query head independently rereads the same K/V data, the stored cache is smaller but the intended bandwidth reduction is not fully realized. Measure cache bytes and query-head arithmetic separately.
Paged and shared KV-cache layouts
Paged caches allocate fixed-size blocks and map logical token positions through a block table. They reduce allocation fragmentation and allow sequences to grow without one contiguous reservation. Prefix sharing lets several requests refer to the same immutable cache blocks for a common prompt.
Paged and quantized layouts add address translation and can make physical reads less contiguous. Page size determines the trade between allocation waste, metadata, and transfer granularity. Scheduling requests with nearby lengths or shared prefixes can improve locality. Kernel and scheduler results must therefore be evaluated together at a declared load and cache state.
Worked example: Grouped-query attention in decode
A model has 32 query heads but only 8 KV heads. Each KV head serves four query heads during decode.
- KV-cache bytes fall by roughly four relative to 32 independent KV heads, increasing possible context or concurrency.
- For BF16 keys and values with head dimension d and context S, compute cache bytes as 2 tensors × S × KV heads × d × 2 bytes. Compare 8 and 32 KV heads directly.
- The kernel should reuse one loaded K/V head across its four associated query heads where the tile and lane mapping permit. Otherwise the smaller stored cache is redundantly reread.
- Projection and output work still reflect query heads, so not every attention cost falls fourfold. Build separate ledgers for cache traffic and arithmetic.
- Benchmark across active batches and context lengths. At short context, launch and projection work may dominate; at long context, reduced KV traffic should become more visible.
- Add a profiler check for requested and transferred cache bytes. A smaller allocation alone does not prove that the kernel avoided redundant reads.
Conclusion: Architectural head sharing creates a systems opportunity only when the kernel schedule realizes the intended reuse and the serving cache preserves it.
Common errors
- Comparing attention implementations without a complete shape and semantic contract. Different masks, head sharing, and phases are different workloads.
- Applying a prefill-optimized kernel to decode because both call attention. Their query dimensions and reuse are fundamentally different.
- Treating cache layout as an isolated storage concern. The layout constrains kernel access and scheduler flexibility simultaneously.
- Reporting decode performance at one context length. Cache traffic and useful parallel work change throughout sequence growth.
Reference summary
Training attention processes many queries and supports backward. Prefill attention handles prompt tokens and writes cache. Decode attention has very few query tokens and reads a long, possibly paged cache. Grouped-query or multi-query attention changes K/V head count. Sliding windows, prefix sharing, quantized caches, and distributed sequence/context parallelism change layout and communication.
- Choose the kernel family from query length, key length, heads, head dimension, cache layout, and mask.
- Include scheduler behavior; a faster decode kernel may support a different batch operating point.
- Measure cache traffic and dequantization when using compressed KV.
- Treat page-table lookups and non-contiguous blocks as part of the memory-access design.
Knowledge check
Why is a training-optimized attention kernel usually a poor drop-in choice for single-token decode?
Show the answer guide
- Training attention processes many query positions and must produce forward outputs plus saved state or recomputable statistics for gradients. Its schedule can tile both query and key dimensions and exploit substantial parallel work.
- Single-token decode has one new query position per active sequence and no backward pass. Keys and values come from a growing KV cache with variable sequence lengths, so cache bandwidth, address lookup, and small-shape launch cost dominate more often.
- A training-optimized kernel can launch too many inactive query tiles, allocate unnecessary backward-related state, use tiles that exceed the one-row query shape, or fail to support paged and ragged KV layouts.
- Decode therefore needs a schedule designed for few query rows, many cached key positions, batch-level parallelism, and the service's cache layout. Mathematical equivalence does not imply schedule equivalence.
Profile one attention equation in three execution regimes
Benchmark exact attention with batch 8, 16 query heads, 4 KV heads, and head dimension 128 for: training with Q length 2048 and KV length 2048; prefill with the same lengths but no backward; and decode with Q length 1 and KV lengths 512, 2048, and 8192. Compare a general dense attention path with a decode-specific paged or KV-cache path when available.
Evidence to keep: Submit the benchmark code, output correctness checks, a result matrix with latency, tokens per second, and peak memory, plus one profiler report per regime showing launch shape and memory traffic. Explain why the training winner does or does not remain the decode winner using the observed query parallelism and KV access.
Having derived important kernels, we now choose how to author, integrate, and maintain them inside a framework and compiler ecosystem.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.