{"id":570031,"date":"2026-08-01T08:44:25","date_gmt":"2026-08-01T08:44:25","guid":{"rendered":"https:\/\/www.newsbeep.com\/il\/570031\/"},"modified":"2026-08-01T08:44:25","modified_gmt":"2026-08-01T08:44:25","slug":"co-designing-ai-model-attention-for-fast-interactive-long-context-inference","status":"publish","type":"post","link":"https:\/\/www.newsbeep.com\/il\/570031\/","title":{"rendered":"Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference"},"content":{"rendered":"

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\u2014not just how it is implemented\u2014increasingly determines a model\u2019s 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<\/a>.\u00a0<\/p>\n

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.\u00a0<\/p>\n

\"Pie<\/p>\n

\t\tFigure 1. DeepSeek-R1 prefill time breakdown at 4K, 32K, and 128K context lengths, where the attention share rises from 18% to 85%<\/p>\n

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.\u00a0<\/p>\n

AbbreviationDefinition\u00a0PB\u00a0Prefill batch size\u00a0DB\u00a0Decode batch size\u00a0QH\u00a0Number of query heads\u00a0KH\u00a0Number of KV heads (KH = QH for MHA, KH = QH\/G for GQA, KH = 1 for MQA)\u00a0\\(G\\)\u00a0Group size = QH\/KH (query heads sharing one KV head)\u00a0Hsz\u00a0Head dimension (typically 64, 128, or 256)\u00a0ISL\u00a0Input sequence length (query tokens in prefill)\u00a0KVSL\u00a0Average KV cache sequence length in a decode iteration\u00a0Table 1. Notations used in the equations featured in this post\u00a0<\/p>\n

How are prefill and decode two different problems?\u00a0<\/a><\/p>\n

Prefill processes the full prompt in parallel, producing large GEMM-M (= ISL \u00d7 \\(G\\)) matmuls that are compute-bound. Without speculative decoding<\/a>, 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).\u00a0<\/p>\n

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. \u00a0<\/p>\n

\u00a0Prefill\u00a0Decode\u00a0Query length\u00a0Full input (ISL tokens)\u00a01 token\u00a0KV context\u00a0Prompt (ISL tokens)\u00a0Full KV cache (KVSL tokens)\u00a0Attention GEMM-M\u00a0ISL \u00d7 \\(G\\) (large)\u00a0\\(G\\) (small)\u00a0Primary bottleneck\u00a0Compute (matmul + softmax)\u00a0HBM bandwidth (memory)\u00a0Table 2. Prefill and decode differ in query length, KV context, attention GEMM-M, and primary bottleneck<\/p>\n

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.<\/p>\n

How arithmetic intensity governs compute- versus memory-bound behavior<\/a><\/p>\n

The roofline bounds GPU performance by compute and bandwidth ceilings, as previously explained<\/a>. Arithmetic intensity determines which binds (Equation 1):\u00a0<\/p>\n

Arithmetic Intensity = Total FLOPs \/ Total bytes accessed<\/p>\n

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.\u00a0<\/p>\n

\"Roofline<\/p>\n

\t\tFigure 2. Roofline model showing decode on the memory-bound ramp and prefill on the compute-bound plateau\u00a0<\/p>\n

How does the FlashAttention kernel compute attention on GPU?\u00a0<\/a><\/p>\n

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:\u00a0<\/p>\n

First, batched matmul (BMM1) scores queries against keys<\/p>\n

Second, online softmax normalizes the scores using a running max and sum<\/p>\n

Third, second batched matmul (BMM2) weights the values<\/p>\n

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.\u00a0<\/p>\n

GEMM shapes\u00a0<\/a><\/p>\n

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. \u00a0<\/p>\n

BMM\u00a0Phase\u00a0Batch\u00a0M\u00a0N\u00a0K\u00a0Meaning\u00a0BMM1\u00a0Prefill\u00a0PB \u00d7 KH\u00a0ISL \u00d7 \\(G\\)\u00a0ISL\u00a0Hsz\u00a0Q \u00b7 K\u1d40: score queries against keys\u00a0Decode\u00a0DB \u00d7 KH\u00a01 \u00d7 \\(G\\)\u00a0KVSL\u00a0Hsz\u00a0BMM2\u00a0Prefill\u00a0PB \u00d7 KH\u00a0ISL \u00d7 \\(G\\)\u00a0Hsz\u00a0ISL\u00a0Weights \u00b7 V: aggregate values\u00a0Decode\u00a0DB \u00d7 KH\u00a01 \u00d7 \\(G\\)\u00a0Hsz\u00a0KVSL\u00a0Table 3. GEMM (Batch, M, N, K) shapes for BMM1 and BMM2<\/p>\n

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.\u00a0<\/p>\n

Group size\u00a0<\/a><\/p>\n

Group size (\\(G\\)) is the number of query heads that share one KV head. MHA has \\(G\\) = 1, GQA has \\(G\\) = 4, 8, 16, \u2026, and MQA has \\(G\\) = QH.\u00a0<\/p>\n

Arithmetic intensity as a function of \\(G\\). In the following formulas, \u201cBytes\u201d refers to HBM bytes moved. For simplicity, assume 1 byte per element (that is, FP8 KV cache). \u00a0<\/p>\n

Prefill\u00a0<\/a><\/p>\n

As \\(G\\) grows, the 1\/\\(G\\) vanishes and arithmetic intensity approaches 2 \u00d7 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:<\/p>\n

FLOPs = 4 \u00d7 PB \u00d7 QH \u00d7 ISL\u00b2 \u00d7 Hsz (constant in \\(G\\))
Bytes = 2 \u00d7 PB \u00d7 KH \u00d7 Hsz \u00d7 ISL \u00d7 (\\(G\\) + 1)
Arithmetic Intensity = 2 \u00d7 \\(G\\) \u00d7 ISL \/ (\\(G\\) + 1) = 2 \u00d7 ISL \/ (1 + 1\/\\(G\\)) \u2192 2 \u00d7 ISL as \\(G\\) \u2192 \u221e\u00a0 \u00a0 \u00a0<\/p>\n

Decode (GEMM-M = \\(G\\))<\/a><\/p>\n

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 \u00d7 \\(G\\), so decode remains memory-bound unless \\(G\\) is very large. Models such as NVIDIA Nemotron 3<\/a> adopted GQA with two KV heads, which makes decode more efficient. Equations 5, 6, and 7:<\/p>\n

FLOPs = 4 \u00d7 DB \u00d7 QH \u00d7 KVSL \u00d7 Hsz (constant in \\(G\\)) \u00a0 \u00a0 \u00a0 \u00a0 \u00a0 \u00a0 \u00a0
Bytes = 2 \u00d7 DB \u00d7 KH \u00d7 Hsz \u00d7 (\\(G\\) + KVSL)
Arithmetic Intensity = 2 \u00d7 \\(G\\) \u00d7 KVSL \/ (\\(G\\) + KVSL) \u2248 2 \u00d7 \\(G\\) (when KVSL \u226b \\(G\\))<\/p>\n

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.\u00a0<\/p>\n

Note: Speculative decoding raises effective GEMM-M to (1+\\(D\\)) \u00d7 \\(G\\), where \\(D\\) is the number of draft tokens. Once large enough to fill compute tiles, decode moves toward compute-bound.\u00a0<\/p>\n

\"Two<\/p>\n

\t\tFigure 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)<\/p>\n

Guideline 1: Choose \\(G\\) for decode efficiency and push it high. Prefill runtime is flat in \\(G\\), while decode arithmetic intensity \u2248 2 \u00d7 \\(G\\), so higher \\(G\\) improves decode speed and GPU utilization. Speculative decoding is another lever to boost performance at a given \\(G\\).<\/p>\n

Head dimension\u00a0\u00a0<\/a><\/p>\n

Unlike group size, head dimension (Hsz) does not affect arithmetic intensity. Doubling Hsz doubles both FLOPs (Equations 2 and 5) and bytes (Equations 3 and 6), leaving their ratio unchanged. Yet Figure 5 shows runtime increasing with Hsz because attention kernels perform three types of work that scale differently with Hsz.\u00a0<\/p>\n

Matmul: Grows with Hsz, but in aligned steps. The previous post<\/a> recommends model dimensions that are multiples of 128 to align with GPU tile sizes and cache-line widths. A partially filled tile costs as much as a full tile, so Hsz = 64 pays for 128. Hsz \u2265 512 pushes close to the tensor memory (TMEM) capacity limit. That makes 128 and 256 the efficient choices.\u00a0<\/p>\n

Memory (the KV state): Also grows with Hsz. Since the GPU moves data in 128-byte units, memory access, like matmul, is most efficient when Hsz is multiple of 128.\u00a0<\/p>\n

Softmax is different: Its cost is independent of Hsz because it operates on attention-score matrix (query tokens \u00d7 keys), which has no head dimension. Equations 8 and 9:\u00a0<\/p>\n

Softmax ops (prefill) \u2248 PB \u00d7 QH \u00d7 ISL\u00b2
Softmax ops (decode) \u2248 DB \u00d7 QH \u00d7 KVSL<\/p>\n

The balance between these determines Hsz cost in each phase (Figure 5).\u00a0<\/p>\n

\"Two<\/p>\n

\t\tFigure 5. Normalized runtime versus Hsz (QH=64, \\(G\\)=32, PB=1, DB=8; prefill ISL and decode KVSL as 32K and 128K). Runtime rises with Hsz\u00a0<\/p>\n

Prefill is compute-bound (matmul plus softmax). As Hsz grows, matmul FLOPs grow while softmax stays fixed. If prefill were pure matmul, doubling Hsz would double runtime; but the fixed softmax does not scale, so runtime rises less than Hsz does. Figure 5 confirms this: prefill climbs with Hsz but slower than Hsz grows. A wider Hsz amortizes softmax, shifting more kernel time to matmul and making prefill less softmax-bound.\u00a0<\/p>\n

Decode is memory-bound (streaming the KV cache). Larger Hsz increases KV bytes per token (Equation 6), so runtime should scale with Hsz. Figure 5 confirms this, though slightly sublinearly because setup, post-processing, and flash-decoding reduction overheads don\u2019t scale with Hsz and weigh more on the shorter KVSL = 32K kernel.\u00a0<\/p>\n

Guideline 2: Use Hsz of 128 or 256. Hsz doesn\u2019t change arithmetic intensity but must align with the hardware<\/a>. Typically, Hsz = 64 still pays for a 128-wide tile, while Hsz \u2265 512 pushes close to the TMEM capacity limit. That makes 128 and 256 the sweet spot.\u00a0<\/p>\n

Sequence length\u00a0<\/a><\/p>\n

Sequence length (ISL \/ KVSL) affects prefill and decode differently, because it enters each phase through a different variable: prefill processes all ISL input tokens together, while each decode step reads a KV cache of length KVSL. The two therefore scale at different rates (Figure 6).\u00a0<\/p>\n

Prefill scales quadratically\u00a0<\/a><\/p>\n

Prefill performs ISL\u00b2 work (every token attends to every token), while KV traffic grows only proportional to ISL (Eq. 2, 3). Arithmetic intensity therefore rises linearly with ISL, keeping prefill well above the ridge point and compute-bound. Doubling ISL should roughly quadruple runtime (Figure 6). At short ISL, scaling is below 4x because fixed setup and post-processing overheads dominate; quadratic scaling appears once ISL is large enough to amortize them.\u00a0<\/p>\n

Decode scales linearly\u00a0<\/a><\/p>\n

Each step reads the full KV cache to generate one token, so bytes grow with KVSL while per-step work remains small (Equations 5 and 6). Arithmetic intensity stays near 2 \u00d7 \\(G\\), well below the ridge point, so decode remains memory-bound at all lengths. Doubling KVSL should double runtime, which Figure 6 confirms. At short KVSL, scaling is below 2x because setup, post-processing, and flash-decoding reduction overheads don\u2019t scale with KVSL. Their share shrinks as KVSL grows.\u00a0<\/p>\n

\"<\/p>\n

\t\tFigure 6. Normalized runtime versus sequence length (QH=64, \\(G\\)=32, Hsz=256, PB=1, DB=8). Prefill scales as O(n\u00b2) with ISL; decode as O(n) with KVSL; both scale below the ideal rate at short lengths<\/p>\n

Guideline 3: Reduce effective KV state where possible. The use case sets sequence length, but the cost is asymmetric: prefill grows as ISL\u00b2, while decode grows linearly with KVSL. Reduce effective KV state through KV-cache compression, sparse or sliding-window attention, or hybrid model architectures (like Nemotron 3) where only some layers carry the growing global KV state.\u00a0<\/p>\n

Tensor parallelism splits attention heads across GPUs \u00a0<\/a><\/p>\n

Tensor parallelism (TP) splits the attention heads across GPUs, giving each GPU QH\/TP query heads and KH\/TP KV heads. It shards heads, not tokens. In Table 3, only the batch dimension, containing KH\/TP, shrinks; the per-GPU GEMM shape and arithmetic intensity remain unchanged. \u00a0<\/p>\n

TP has a practical limit: KV heads must divide evenly across GPUs. Once TP > KH, a group\u2019s query heads span multiple ranks, each requiring a copy of the shared KV head. This duplicates KV state, adding memory and bandwidth overhead without benefit (Figure 7). Thus, keep TP \u2264 KH so each GPU owns at least one complete group: one KV head and its \\(G\\) query heads.\u00a0<\/p>\n

Models with few KV heads (Nemotron 3 with two, for example), quickly exhaust TP because the cache cannot be sharded below one KV head per GPU without duplication. Attention must then scale differently: Attention Data Parallelism (ADP) shards requests, while KV Parallelism (KVP) shards long-sequence KV caches across GPUs. The FFN scales separately with Expert Parallelism (EP).\u00a0<\/p>\n

TensorRT-LLM<\/a> combines these as Wide EP<\/a> (ADP for attention plus EP for the FFN), and Helix Parallelism<\/a> (KVP for attention plus EP for the FFN). In both cases, KH determines efficient scaling.\u00a0<\/p>\n

\"A<\/p>\n

\t\tFigure 7. Attention sharding strategies: when TP > KH, the KV cache is duplicated, adding memory and bandwidth overhead<\/p>\n

Guideline 4: Let KH determine the parallelism strategy. Keep TP \u2264 KH so each GPU has a full KV head. Models with few KV heads (MQA with 1, or GQA with 2) exhaust TP quickly and are better served by ADP or KVP for attention plus EP for the MoE FFN (implemented as Wide EP<\/a> and Helix Parallelism<\/a> in TensorRT-LLM).\u00a0<\/p>\n

Get started co-designing AI model attention<\/a><\/p>\n

Use the four guidelines summarized below as a model design checklist to get started co-designing AI model attention. These choices can improve GPU utilization, inference speed, throughput, and interactivity on the same hardware.\u00a0\u00a0<\/p>\n

Guideline 1: Choose the group size (\\(G\\)) for decode and push it high. Prefill is not sensitive to group size.\u00a0<\/p>\n

Guideline 2: Use a head dimension (Hsz) = 128 or 256 to align with the GPU tiles and 128-byte transfers while staying within TMEM budget. A larger head also hides softmax in prefill.\u00a0<\/p>\n

Guideline 3: Reduce effective KV state through KV-cache compression, sparse or sliding-window attention, or hybrid models.\u00a0<\/p>\n

Guideline 4: Match parallelism to KV head count (KH). Keep TP \u2264 KH and scale a few KH models with Wide EP<\/a> and Helix Parallelism<\/a>.\u00a0<\/p>\n

Acknowledgments\u00a0<\/a><\/p>\n

This post is an NVIDIA cross-team effort. We are grateful to Timmy Liu, Jatin Mitra, Tiyasa Mitra, Bhargava Gopireddy, Brian Pharris, Julien Demouth, and Eduardo Alvarez for their help.<\/p>\n","protected":false},"excerpt":{"rendered":"As agentic and long-context workloads become common, the context lengths increase and attention consumes a larger share of…\n","protected":false},"author":2,"featured_media":570032,"comment_status":"","ping_status":"","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[20],"tags":[345,343,344,85,46,125],"class_list":["post-570031","post","type-post","status-publish","format-standard","has-post-thumbnail","category-artificial-intelligence","tag-ai","tag-artificial-intelligence","tag-artificialintelligence","tag-il","tag-israel","tag-technology"],"_links":{"self":[{"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/posts\/570031","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/comments?post=570031"}],"version-history":[{"count":0,"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/posts\/570031\/revisions"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/media\/570032"}],"wp:attachment":[{"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/media?parent=570031"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/categories?post=570031"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.newsbeep.com\/il\/wp-json\/wp\/v2\/tags?post=570031"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}