Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference

來源原標題

Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference

作者

Tanya Lenz

原始標籤

Agentic AI / Generative AI · Data Center / Cloud · Developer Tools & Techniques · AI Inference · Inference Performance

原始連結

https://developer.nvidia.com/blog/co-designing-ai-model-attention-for-fast-interactive-long-context-inference/

摘要原文

As agentic and long-context workloads become common, the context lengths increase and attention consumes a larger share of inference time (Figure 1). Because attention now dominates that cost, how it is designed—not just how it is implemented—increasingly determines a model’s inference performance. Shaping model architecture around how GPUs execute it is the premise of AI model co-design. For a discussion of how model design choices impact both throughput and interactivity without sacrificing accuracy, see the previous post, AI Model Co-Design: Hardware-Friendly LLM Design . This post examines how group size (query heads per KV head), head dimension, and sequence length shape the performance of dense attention, where every query attends to all keys and values along the sequence length. We distill that analysis, together with how attention is parallelized across GPUs, into four practical guidelines: a co-design checklist that helps model developers raise inference throughput and interactivity on NVIDIA GPUs. Stay tuned for a post covering sparse attention. Figure 1. DeepSeek-R1 prefill time breakdown at 4K, 32K, and 128K context lengths, where the attention share rises from 18% to 85% Every analysis is grounded in two sources: analytical formulas from GEMM-shape arithmetic, and measured data from prefill and decode kernels with FP8 for both attention compute and the KV cache. Abbreviation Definition 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 (typically 64, 128, or 256) ISL Input sequence length (query tokens in prefill) KVSL Average KV cache sequence length in a decode iteration Table 1. Notations used in the equations featured in this post How are prefill and decode two different problems? Prefill processes the full prompt in parallel, producing large GEMM-M (= ISL × \(G\)) matmuls that are compute-bound. Without speculative decoding , decode generates one token at a time, producing small GEMM-M (= \(G\)) matmuls and becoming memory-bound by KV cache reads from high-bandwidth memory (HBM). Speculative decoding increases GEMM-M and can shift decode toward compute-bound. Because the query lengths, KV access, and bottleneck differ for prefill and decode (Table 2), each parameter is analyzed separately for each phase. Prefill Decode Query length Full input (ISL tokens) 1 token KV context Prompt (ISL tokens) Full KV cache (KVSL tokens) Attention GEMM-M ISL × \(G\) (large) \(G\) (small) Primary bottleneck Compute (matmul + softmax) HBM bandwidth (memory) Table 2. Prefill and decode differ in query length, KV context, attention GEMM-M, and primary bottleneck Note: With prefix caching, common in agentic and multiturn apps, a new turn may have a short ISL while attending to a large prefix cache. With a short ISL but long prefix cache, prefill behaves like decode. How arithmetic intensity governs compute- versus memory-bound behavior The roofline bounds GPU performance by compute and bandwidth ceilings, as previously explained . Arithmetic intensity determines which binds (Equation 1): Arithmetic Intensity = Total FLOPs / Total bytes accessed The ridge point marks the transition from memory-bound to compute-bound. Prefill lies well above it and is compute-bound, while decode lies below it and is memory-bound (Figure 2). Speculative decoding raises decode arithmetic intensity and can move it toward the ridge. Figure 2. Roofline model showing decode on the memory-bound ramp and prefill on the compute-bound plateau How does the FlashAttention kernel compute attention on GPU? FlashAttention computes attention without materializing the full attention matrix. It streams tiles of \(Q\), \(K\), and \(V\) from HBM to on-chip SRAM and fuses three steps into one pass: First, batched matmul (BMM1) scores queries against keys Second, online softmax normalizes the scores using a running max and sum Third, second batched matmul (BMM2) weights the values The BMMs run on Tensor Cores while the softmax exponentials run on special-function units. The BMM shapes drive the arithmetic intensity analysis that follows. Figure 3. FlashAttention kernel with BMM1 and BMM2 on Tensor Cores and fused online softmax on special-function units. Image adapted from FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness GEMM shapes Attention performance follows from the shapes of its two matmuls. Table 3 lists the per-phase (Batch, M, N, K) dimensions of BMM1 and BMM2. BMM Phase Batch M N K Meaning BMM1 Prefill PB × KH ISL × \(G\) ISL Hsz Q · Kᵀ: score queries against keys Decode DB × KH 1 × \(G\) KVSL Hsz BMM2 Prefill PB × KH ISL × \(G\) Hsz ISL Weights · V: aggregate values Decode DB × KH 1 × \(G\) Hsz KVSL Table 3. GEMM (Batch, M, N, K) shapes for BMM1 and BMM2 For decode, GEMM-M = \(G\), usually 8-16, far below a GPU tile-M of 64 or 128, limiting parallel work per tile. Larger \(G\) loads less KV per token and amortizes each load across more query heads, improving utilization. The next section quantifies this effect. Group size Group size (\(G\)) is the number of query heads that share one KV head. MHA has \(G\) = 1, GQA has \(G\) = 4, 8, 16, …, and MQA has \(G\) = QH. Arithmetic intensity as a function of \(G\). In the following formulas, “Bytes” refers to HBM bytes moved. For simplicity, assume 1 byte per element (that is, FP8 KV cache). Prefill As \(G\) grows, the 1/\(G\) vanishes and arithmetic intensity approaches 2 × ISL. At ISL = 32K, raising \(G\) from 8 to 16 improves arithmetic intensity by under 6%. In other words, prefill is dominated by ISL, not \(G\). Figure 4 confirms this, varying \(G\) from 1 (MHA) to 64 (MQA) changes prefill runtime by under 1%. Equations 2, 3, and 4: FLOPs = 4 × PB × QH × ISL² × Hsz (constant in \(G\)) Bytes = 2 × PB × KH × Hsz × ISL × (\(G\) + 1) Arithmetic Intensity = 2 × \(G\) × ISL / (\(G\) + 1) = 2 × ISL / (1 + 1/\(G\)) → 2 × ISL as \(G\) → ∞ Decode (GEMM-M = \(G\)) Doubling \(G\) doubles decode arithmetic intensity. Raising \(G\) from 1 to 8 gives an 8x gain by reducing memory traffic and improving GPU compute utilization. It is independent of KVSL: arithmetic intensity stays near 2 × \(G\), so decode remains memory-bound unless \(G\) is very large. Models such as NVIDIA Nemotron 3 adopted GQA with two KV heads, which makes decode more efficient. Equations 5, 6, and 7: FLOPs = 4 × DB × QH × KVSL × Hsz (constant in \(G\)) Bytes = 2 × DB × KH × Hsz × (\(G\) + KVSL) Arithmetic Intensity = 2 × \(G\) × KVSL / (\(G\) + KVSL) ≈ 2 × \(G\) (when KVSL ≫ \(G\)) Figure 4 shows decode runtime falls about 2x per doubling of \(G\) because halving the KV heads halves the data loaded per token. Beyond \(G\) = 16, the KVSL = 32K curve flattens. Its per-step kernel is small enough that two costs dominate: fixed setup and post-processing overhead, and the flash-decoding reduction from splitting KV across SMs to stay parallel with few KV heads. The longer KVSL = 128K kernel better amortizes these costs and continues tracking the 2x trend. Note: Speculative decoding raises effective GEMM-M to (1+\(D\)) × \(G\), where \(D\) is the number of draft tokens. Once large enough to fill compute tiles, decode moves toward compute-bound. Figure 4. Normalized runtime versus \(G\) (QH=64, Hsz=128, PB=1, DB=8; prefill ISL and decode KVSL as 32K and 128K). Prefill is flat in \(G\) (compute-bound); decode falls ~2x per doubling of \(G\) (memory-bound), from MHA (\(G\)=1) to MQA (\(G\)=64) Guideline 1: Choose \(G\) for decode efficiency and push it high. Prefill runtime is flat in \(G\), while decode arithmetic intensity ≈ 2 × \(G\), so higher \(G\) improves decode speed and GPU utilization. Speculative decoding is another lever to boost performance at a given \(G\). Head dimension Unlike group size, head dimension (Hsz) does not affect arithmetic intensity.