Designing Attention for Fast, Interactive Long-Context Inference on NVIDIA GPUs

As agentic and long‑context workloads become common, attention now consumes an increasing share of inference time. Because attention has become the dominant cost, architectural design choices—not only implementation—play a growing role in a model’s inference performance. NVIDIA examined how group size (G), head dimension (Hsz), sequence length, and GPU parallelism shape dense attention performance, and distilled the analysis into practical guidance for improving throughput and interactivity on NVIDIA GPUs.

The analysis uses two foundations: analytic GEMM‑shape arithmetic and measured data from prefill and decode kernels using FP8 for attention compute and the KV cache.

Notation and definitions

• PB: Prefill batch size
• DB: Decode batch size
• QH: number of query heads
• KH: number of KV heads (KH = QH for MHA, KH = QH/G for GQA, KH = 1 for MQA)
• G: group size = QH/KH (query heads sharing one KV head)
• Hsz: head dimension (commonly 64, 128, or 256)
• ISL: input sequence length (number of query tokens in prefill)
• KVSL: average KV cache sequence length during a decode iteration

Prefill vs decode: fundamentally different bottlenecks

• Prefill processes the entire prompt in parallel, producing large GEMM‑M (= ISL × G) matmuls and is typically compute‑bound (matmul + softmax).
• Decode (without speculative decoding) generates one token at a time, producing small GEMM‑M (= G) matmuls and is memory‑bound by KV cache reads from high‑bandwidth memory (HBM). Speculative decoding can raise GEMM‑M and shift decode toward compute‑bound.

Because query lengths, KV access patterns and bottlenecks differ, the parameters must be analyzed separately for each phase.

Arithmetic intensity and the roofline

The roofline model bounds GPU performance by compute and bandwidth ceilings. Arithmetic intensity (Total FLOPs / Total bytes accessed) determines whether a kernel is compute‑ or memory‑bound. Prefill sits well above the ridge point and is compute‑bound; decode is below it and memory‑bound. Speculative decoding increases decode’s arithmetic intensity and can move it toward the ridge.

FlashAttention kernel overview

FlashAttention avoids materializing the full attention matrix by streaming tiles of Q, K and V from HBM into on‑chip SRAM and fusing three steps into one pass:

1. BMM1: batched matmul scores Q against K
2. online softmax: normalize scores using a running max and sum
3. BMM2: batched matmul to weight V The BMMs run on Tensor Cores; softmax exponentials run on special‑function units. The shapes of these GEMMs drive the arithmetic intensity analysis.

GEMM shapes in prefill and decode

Per‑phase GEMM (Batch, M, N, K) dimensions for BMM1 and BMM2 are:

• BMM1 prefill: Batch = PB × KH, M = ISL × G, N = ISL, K = Hsz
• BMM1 decode: Batch = DB × KH, M = 1 × G, N = KVSL, K = Hsz
• BMM2 prefill: Batch = PB × KH, M = ISL × G, N = Hsz, K = ISL
• BMM2 decode: Batch = DB × KH, M = 1 × G, N = Hsz, K = KVSL

Note: decode GEMM‑M = G is often 8–16, much smaller than GPU tile‑M (64 or 128), which limits parallel work per tile. Increasing G amortizes KV loads across more query heads and improves utilization.

Group size (G)

Group size equals the number of query heads sharing a KV head (MHA: G=1, GQA: G in {4,8,16,…}, MQA: G=QH).

• Prefill: as G increases the 1/G term vanishes and arithmetic intensity approaches 2×ISL. Consequently prefill runtime is dominated by ISL and is largely insensitive to G (empirically, changing G from 1 to 64 changed prefill runtime by <1% in an example).

• Decode: arithmetic intensity ≈ 2×G, so doubling G roughly doubles decode arithmetic intensity. Raising G from 1 to 8 can yield an ~8× improvement in memory traffic efficiency and GPU utilization. In practice, measured decode runtime falls by about 2× per doubling of G until fixed overheads and flash‑decoding parallelization effects dominate; very large KVSL values amortize those overheads and better track the ideal trend.

Guideline 1: Choose G to optimize decode efficiency and push it high. Prefill is not sensitive to G. Speculative decoding is another lever to increase effective GEMM‑M.

Head dimension (Hsz)

Hsz does not change arithmetic intensity, because doubling Hsz doubles both FLOPs and bytes. Nonetheless runtime grows with Hsz because the attention kernel performs three types of work that scale differently:

• Matmul: increases with Hsz but in alignment‑sensitive steps; tile alignment (multiples of 128) matters.
• Memory (KV state): grows with Hsz and benefits from 128‑byte transfer alignment.
• Softmax: independent of Hsz (operates on the query×key score matrix).

Consequently, Hsz choices matter for practical runtime. Hsz = 64 often pays for a 128‑wide tile; Hsz ≥ 512 approaches TMEM capacity limits. Empirical and hardware considerations make Hsz = 128 or 256 the most efficient options.

Guideline 2: Use Hsz = 128 or 256 to align with GPU tile sizes and 128‑byte transfers while staying within TMEM budget.

Sequence length (ISL / KVSL)

Sequence length affects prefill and decode asymmetrically:

• Prefill: does O(ISL²) work (each token attends to each token). Arithmetic intensity grows with ISL, keeping prefill compute‑bound; doubling ISL roughly quadruples runtime (ignoring fixed overheads at short ISL).
• Decode: each step reads the full KV cache, so per‑step bytes grow linearly with KVSL while per‑step work stays small. Arithmetic intensity remains near 2×G and decode stays memory‑bound; doubling KVSL roughly doubles runtime (again, short KVSL sees fixed overheads weaken ideal scaling).

Guideline 3: Reduce effective KV state where possible—via KV‑cache compression, sparse or sliding‑window attention, or hybrid architectures (e.g., Nemotron 3) where only a subset of layers accumulate global KV state.

Parallelism: let KH determine strategy

Tensor parallelism (TP) shards attention across GPUs by splitting heads: each GPU gets QH/TP query heads and KH/TP KV heads. TP reduces the batch dimension (KH/TP) but leaves per‑GPU GEMM shapes and arithmetic intensity unchanged.

Practical limit: KH must divide across GPUs. If TP > KH, a KV head spans multiple ranks and the KV cache is duplicated across ranks, adding memory and bandwidth overhead without benefit. Keep TP ≤ KH so each GPU owns at least one full group (one KV head and its G query heads).

Models with few KV heads (e.g., Nemotron 3 with two, MQA with 1) quickly exhaust TP; for those, better options are Attention Data Parallelism (ADP) or KV Parallelism (KVP) for attention plus Expert Parallelism (EP) for the FFN. TensorRT‑LLM implements Wide EP (ADP + EP) and Helix Parallelism (KVP + EP) as practical combinations.

Guideline 4: Match parallelism to KH. Keep TP ≤ KH; for small KH prefer ADP/KVP for attention and EP for the MoE/FFN.

Co‑design checklist (four guidelines)

1. Pick G for decode and push it high—prefill is insensitive to G.
2. Use head dimension Hsz = 128 or 256 to align with GPU tiles and 128‑byte transfers while staying within TMEM budget.
3. Reduce effective KV state (compression, sparse/sliding attention, or hybrid models).
4. Let KH determine parallelism: keep TP ≤ KH; scale few‑KH models with Wide EP or Helix Parallelism where appropriate.

Acknowledgments

This work is a cross‑team NVIDIA effort. NVIDIA thanks Timmy Liu, Jatin Mitra, Tiyasa Mitra, Bhargava Gopireddy, Brian Pharris, Julien Demouth and Eduardo Alvarez for their contributions.