{"id":3146,"date":"2026-07-31T22:16:00","date_gmt":"2026-07-31T22:16:00","guid":{"rendered":"https:\/\/futurenews24.com\/index.php\/2026\/07\/31\/co-designing-ai-model-attention-for-fast-interactive-long-context-inference\/"},"modified":"2026-08-01T17:59:46","modified_gmt":"2026-08-01T17:59:46","slug":"co-designing-ai-model-attention-for-fast-interactive-long-context-inference","status":"publish","type":"post","link":"https:\/\/futurenews24.com\/index.php\/2026\/07\/31\/co-designing-ai-model-attention-for-fast-interactive-long-context-inference\/","title":{"rendered":"Co-Designing AI Mannequin Consideration for Quick, Interactive Lengthy-Context Inference"},"content":{"rendered":"
\n<\/p>\n
\n
As agentic and long-context workloads grow to be widespread, the context lengths improve and a spotlight consumes a bigger share of inference time (Determine 1). As a result of consideration now dominates that price, how it’s designed\u2014not simply how it’s applied\u2014more and more determines a mannequin\u2019s inference efficiency. Shaping mannequin structure round how GPUs execute it’s the premise of AI mannequin co-design. For a dialogue of how mannequin design selections affect each throughput and interactivity with out sacrificing accuracy, see the earlier submit, AI Mannequin Co-Design: {Hardware}-Pleasant LLM Design.\u00a0<\/p>\n
This submit examines how group measurement (question heads per KV head), head dimension, and sequence size form the efficiency of dense consideration, the place each question attends to all keys and values alongside the sequence size. We distill that evaluation, along with how consideration is parallelized throughout GPUs, into 4 sensible pointers: a co-design guidelines that helps mannequin builders elevate inference throughput and interactivity on NVIDIA GPUs. Keep tuned for a submit protecting sparse consideration.\u00a0<\/p>\n
\nDetermine 1. DeepSeek-R1 prefill time breakdown at 4K, 32K, and 128K context lengths, the place the eye share rises from 18% to 85%<\/figcaption><\/figure>\n<\/div>\n
Each evaluation is grounded in two sources: analytical formulation from GEMM-shape arithmetic, and measured knowledge from prefill and decode kernels with FP8 for each consideration compute and the KV cache.\u00a0<\/p>\nAbbreviationDefinition\u00a0PB\u00a0Prefill batch measurement\u00a0DB\u00a0Decode batch measurement\u00a0QH\u00a0Variety of question heads\u00a0KH\u00a0Variety of KV heads (KH = QH for MHA, KH = QH\/G for GQA, KH = 1 for MQA)\u00a0(G)\u00a0Group measurement = QH\/KH (question heads sharing one KV head)\u00a0Hsz\u00a0Head dimension (sometimes 64, 128, or 256)\u00a0ISL\u00a0Enter sequence size (question tokens in prefill)\u00a0KVSL\u00a0Common KV cache sequence size in a decode iteration\u00a0Desk 1. Notations used within the equations featured on this submit\u00a0<\/figcaption><\/figure>\n
How are prefill and decode two completely different issues?\u00a0<\/h2>\n
Prefill processes the complete immediate in parallel, producing massive GEMM-M (= ISL \u00d7 (G)) matmuls which are compute-bound. With out speculative decoding, decode generates one token at a time, producing small GEMM-M (= (G)) matmuls and turning into memory-bound by KV cache reads from high-bandwidth reminiscence (HBM).\u00a0<\/p>\n
Speculative decoding will increase GEMM-M and may shift decode towards compute-bound. As a result of the question lengths, KV entry, and bottleneck differ for prefill and decode (Desk 2), every parameter is analyzed individually for every part. \u00a0<\/p>\n\u00a0Prefill\u00a0Decode\u00a0Question size\u00a0Full enter (ISL tokens)\u00a01 token\u00a0KV context\u00a0Immediate (ISL tokens)\u00a0Full KV cache (KVSL tokens)\u00a0Consideration GEMM-M\u00a0ISL \u00d7 (G) (massive)\u00a0(G) (small)\u00a0Major bottleneck\u00a0Compute (matmul + softmax)\u00a0HBM bandwidth (reminiscence)\u00a0Desk 2. Prefill and decode differ in question size, KV context, consideration GEMM-M, and first bottleneck<\/figcaption><\/figure>\n
Be aware: With prefix caching, widespread in agentic and multiturn apps, a brand new flip could have a brief ISL whereas attending to a big prefix cache. With a brief ISL however lengthy prefix cache, prefill behaves like decode.<\/p>\n
How arithmetic depth governs compute- versus memory-bound habits<\/h2>\n
The roofline bounds GPU efficiency by compute and bandwidth ceilings, as beforehand defined. Arithmetic depth determines which binds (Equation 1):\u00a0<\/p>\n
The ridge level marks the transition from memory-bound to compute-bound. Prefill lies properly above it and is compute-bound, whereas decode lies under it and is memory-bound (Determine 2). Speculative decoding raises decode arithmetic depth and may transfer it towards the ridge.\u00a0<\/p>\n
\nDetermine 2. Roofline mannequin exhibiting decode on the memory-bound ramp and prefill on the compute-bound plateau\u00a0<\/figcaption><\/figure>\n<\/div>\n
How does the FlashAttention kernel compute consideration on GPU?\u00a0<\/h2>\n
FlashAttention computes consideration with out materializing the complete consideration matrix. It streams tiles of (Q), (Ok), and (V) from HBM to on-chip SRAM and fuses three steps into one cross:\u00a0<\/p>\n
First, batched matmul (BMM1) scores queries in opposition to keys<\/p>\n
Second, on-line softmax normalizes the scores utilizing a working max and sum<\/p>\n
Third, second batched matmul (BMM2) weights the values<\/p>\n
The BMMs run on Tensor Cores whereas the softmax exponentials run on special-function items. The BMM shapes drive the arithmetic depth evaluation that follows.\u00a0<\/p>\n
GEMM shapes\u00a0<\/h2>\n
Consideration efficiency follows from the shapes of its two matmuls. Desk 3 lists the per-phase (Batch, M, N, Ok) dimensions of BMM1 and BMM2. \u00a0<\/p>\nBMM\u00a0Section\u00a0Batch\u00a0M\u00a0N\u00a0Ok\u00a0Which means\u00a0BMM1\u00a0Prefill\u00a0PB \u00d7 KH\u00a0ISL \u00d7 (G)\u00a0ISL\u00a0Hsz\u00a0Q \u00b7 K\u1d40: rating queries in opposition to keys\u00a0Decode\u00a0DB \u00d7 KH\u00a01 \u00d7 (G)\u00a0KVSL\u00a0Hsz\u00a0BMM2\u00a0Prefill\u00a0PB \u00d7 KH\u00a0ISL \u00d7 (G)\u00a0Hsz\u00a0ISL\u00a0Weights \u00b7 V: combination values\u00a0Decode\u00a0DB \u00d7 KH\u00a01 \u00d7 (G)\u00a0Hsz\u00a0KVSL\u00a0Desk 3. GEMM (Batch, M, N, Ok) shapes for BMM1 and BMM2<\/figcaption><\/figure>\n
For decode, GEMM-M = (G), normally 8-16, far under a GPU tile-M of 64 or 128, limiting parallel work per tile. Bigger (G) masses much less KV per token and amortizes every load throughout extra question heads, enhancing utilization. The following part quantifies this impact.\u00a0<\/p>\n
Group measurement\u00a0<\/h2>\n
Group measurement ((G)) is the variety of question 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 depth as a perform of (G). Within the following formulation, \u201cBytes\u201d refers to HBM bytes moved. For simplicity, assume 1 byte per ingredient (that’s, FP8 KV cache). \u00a0<\/p>\n
Prefill\u00a0<\/h3>\n
As (G) grows, the 1\/(G) vanishes and arithmetic depth approaches 2 \u00d7 ISL. At ISL = 32K, elevating (G) from 8 to 16 improves arithmetic depth by beneath 6%. In different phrases, prefill is dominated by ISL, not (G). Determine 4 confirms this, various (G) from 1 (MHA) to 64 (MQA) adjustments prefill runtime by beneath 1%. Equations 2, 3, and 4:<\/p>\n
Doubling (G) doubles decode arithmetic depth. Elevating (G) from 1 to eight offers an 8x achieve by decreasing reminiscence visitors and enhancing GPU compute utilization. It’s unbiased of KVSL: arithmetic depth stays close to 2 \u00d7 (G), so decode stays memory-bound until (G) may be very massive. Fashions akin to NVIDIA Nemotron 3 adopted GQA with two KV heads, which makes decode extra environment friendly. Equations 5, 6, and seven:<\/p>\n
Determine 4 exhibits decode runtime falls about 2x per doubling of (G) as a result of halving the KV heads halves the info loaded per token. Past (G) = 16, the KVSL = 32K curve flattens. Its per-step kernel is sufficiently small that two prices dominate: mounted setup and post-processing overhead, and the flash-decoding discount from splitting KV throughout SMs to remain parallel with few KV heads. The longer KVSL = 128K kernel higher amortizes these prices and continues monitoring the 2x pattern.\u00a0<\/p>\n
Be aware: Speculative decoding raises efficient GEMM-M to (1+(D)) \u00d7 (G), the place (D) is the variety of draft tokens. As soon as massive sufficient to fill compute tiles, decode strikes towards compute-bound.\u00a0<\/p>\n
\nDetermine 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)<\/figcaption><\/figure>\n<\/div>\n
Guideline 1: Select (G) for decode effectivity and push it excessive. Prefill runtime is flat in (G), whereas decode arithmetic depth \u2248 2 \u00d7 (G), so greater (G) improves decode velocity and GPU utilization. Speculative decoding is one other lever to spice up efficiency at a given (G).<\/p>\n
Head dimension\u00a0\u00a0<\/h2>\n
Not like group measurement, head dimension (Hsz) doesn’t have an effect on arithmetic depth. Doubling Hsz doubles each FLOPs (Equations 2 and 5) and bytes (Equations 3 and 6), leaving their ratio unchanged. But Determine 5 exhibits runtime growing with Hsz as a result of consideration kernels carry out three forms of work that scale in a different way with Hsz.\u00a0<\/p>\n
Matmul: Grows with Hsz, however in aligned steps. The earlier submit recommends mannequin dimensions which are multiples of 128 to align with GPU tile sizes and cache-line widths. {A partially} crammed tile prices as a lot as a full tile, so Hsz = 64 pays for 128. Hsz \u2265 512 pushes near the tensor reminiscence (TMEM) capability restrict. That makes 128 and 256 the environment friendly selections.\u00a0<\/p>\n
Reminiscence (the KV state): Additionally grows with Hsz. Because the GPU strikes knowledge in 128-byte items, reminiscence entry, like matmul, is most effective when Hsz is a number of of 128.\u00a0<\/p>\n
Softmax is completely different: Its price is unbiased of Hsz as a result of it operates on attention-score matrix (question tokens \u00d7 keys), which has no head dimension. Equations 8 and 9:\u00a0<\/p>\n
The stability between these determines Hsz price in every part (Determine 5).\u00a0<\/p>\n
\nDetermine 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<\/figcaption><\/figure>\n<\/div>\n
Prefill is compute-bound (matmul plus softmax). As Hsz grows, matmul FLOPs develop whereas softmax stays mounted. If prefill had been pure matmul, doubling Hsz would double runtime; however the mounted softmax doesn’t scale, so runtime rises lower than Hsz does. Determine 5 confirms this: prefill climbs with Hsz however slower than Hsz grows. A wider Hsz amortizes softmax, shifting extra kernel time to matmul and making prefill much less softmax-bound.\u00a0<\/p>\n
Decode is memory-bound (streaming the KV cache). Bigger Hsz will increase KV bytes per token (Equation 6), so runtime ought to scale with Hsz. Determine 5 confirms this, although barely sublinearly as a result of setup, post-processing, and flash-decoding discount overheads don\u2019t scale with Hsz and weigh extra on the shorter KVSL = 32K kernel.\u00a0<\/p>\n
Guideline 2: Use Hsz of 128 or 256. Hsz doesn\u2019t change arithmetic depth however should align with the {hardware}. Usually, Hsz = 64 nonetheless pays for a 128-wide tile, whereas Hsz \u2265 512 pushes near the TMEM capability restrict. That makes 128 and 256 the candy spot.\u00a0<\/p>\n
Sequence size\u00a0<\/h2>\n
Sequence size (ISL \/ KVSL) impacts prefill and decode in a different way, as a result of it enters every part by a special variable: prefill processes all ISL enter tokens collectively, whereas every decode step reads a KV cache of size KVSL. The 2 subsequently scale at completely different charges (Determine 6).\u00a0<\/p>\n
Prefill scales quadratically\u00a0<\/h3>\n
Prefill performs ISL\u00b2 work (each token attends to each token), whereas KV visitors grows solely proportional to ISL (Eq. 2, 3). Arithmetic depth subsequently rises linearly with ISL, maintaining prefill properly above the ridge level and compute-bound. Doubling ISL ought to roughly quadruple runtime (Determine 6). At quick ISL, scaling is under 4x as a result of mounted setup and post-processing overheads dominate; quadratic scaling seems as soon as ISL is massive sufficient to amortize them.\u00a0<\/p>\n
Decode scales linearly\u00a0<\/h3>\n
Every step reads the complete KV cache to generate one token, so bytes develop with KVSL whereas per-step work stays small (Equations 5 and 6). Arithmetic depth stays close to 2 \u00d7 (G), properly under the ridge level, so decode stays memory-bound in any respect lengths. Doubling KVSL ought to double runtime, which Determine 6 confirms. At quick KVSL, scaling is under 2x as a result of setup, post-processing, and flash-decoding discount overheads don\u2019t scale with KVSL. Their share shrinks as KVSL grows.\u00a0<\/p>\n
\nDetermine 6. Normalized runtime versus sequence size (QH=64, (G)=32, Hsz=256, PB=1, DB=8). Prefill scales as O(n\u00b2) with ISL; decode as O(n) with KVSL; each scale under the perfect fee at quick lengths<\/figcaption><\/figure>\n<\/div>\n
Guideline 3: Scale back efficient KV state the place doable. The use case units sequence size, however the associated fee is uneven: prefill grows as ISL\u00b2, whereas decode grows linearly with KVSL. Scale back efficient KV state by KV-cache compression, sparse or sliding-window consideration, or hybrid mannequin architectures (like Nemotron 3) the place just some layers carry the rising international KV state.\u00a0<\/p>\n
Tensor parallelism splits consideration heads throughout GPUs \u00a0<\/h2>\n
Tensor parallelism (TP) splits the eye heads throughout GPUs, giving every GPU QH\/TP question heads and KH\/TP KV heads. It shards heads, not tokens. In Desk 3, solely the batch dimension, containing KH\/TP, shrinks; the per-GPU GEMM form and arithmetic depth stay unchanged. \u00a0<\/p>\n
TP has a sensible restrict: KV heads should divide evenly throughout GPUs. As soon as TP > KH, a bunch\u2019s question heads span a number of ranks, every requiring a replica of the shared KV head. This duplicates KV state, including reminiscence and bandwidth overhead with out profit (Determine 7). Thus, hold TP \u2264 KH so every GPU owns at the very least one full group: one KV head and its (G) question heads.\u00a0<\/p>\n
Fashions with few KV heads (Nemotron 3 with two, for instance), shortly exhaust TP as a result of the cache can’t be sharded under one KV head per GPU with out duplication. Consideration should then scale in a different way: Consideration Information Parallelism (ADP) shards requests, whereas KV Parallelism (KVP) shards long-sequence KV caches throughout GPUs. The FFN scales individually with Knowledgeable Parallelism (EP).\u00a0<\/p>\n
TensorRT-LLM combines these as Huge EP (ADP for consideration plus EP for the FFN), and Helix Parallelism (KVP for consideration plus EP for the FFN). In each circumstances, KH determines environment friendly scaling.\u00a0<\/p>\n
\nDetermine 7. Consideration sharding methods: when TP > KH, the KV cache is duplicated, including reminiscence and bandwidth overhead<\/figcaption><\/figure>\n<\/div>\n
Guideline 4: Let KH decide the parallelism technique. Preserve TP \u2264 KH so every GPU has a full KV head. Fashions with few KV heads (MQA with 1, or GQA with 2) exhaust TP shortly and are higher served by ADP or KVP for consideration plus EP for the MoE FFN (applied as Huge EP and Helix Parallelism in TensorRT-LLM).\u00a0<\/p>\n
Get began co-designing AI mannequin consideration<\/h2>\n
Use the 4 pointers summarized under as a mannequin design guidelines to get began co-designing AI mannequin consideration. These selections can enhance GPU utilization, inference velocity, throughput, and interactivity on the identical {hardware}.\u00a0\u00a0<\/p>\n
Guideline 1: Select the group measurement ((G)) for decode and push it excessive. Prefill just isn’t delicate to group measurement.\u00a0<\/p>\n
Guideline 2: Use a head dimension (Hsz) = 128 or 256 to align with the GPU tiles and 128-byte transfers whereas staying inside TMEM finances. A bigger head additionally hides softmax in prefill.\u00a0<\/p>\n
Guideline 3: Scale back efficient KV state by KV-cache compression, sparse or sliding-window consideration, or hybrid fashions.\u00a0<\/p>\n
Guideline 4: Match parallelism to KV head depend (KH). Preserve TP \u2264 KH and scale a number of KH fashions with Huge EP and Helix Parallelism.\u00a0<\/p>\n
Acknowledgments\u00a0<\/h3>\n
This submit is an NVIDIA cross-team effort. We’re grateful to Timmy Liu, Jatin Mitra, Tiyasa Mitra, Bhargava Gopireddy, Brian Pharris, Julien Demouth, and Eduardo Alvarez for his or her assist.<\/p>\n<\/div>\n
![\"Pie](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/attention-prefill-time-ai-model.webp\")
![\"Roofline](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/roofline-plot-prefill-decode.webp\")
![\"Two](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/normalized-runtime-versus-g.webp\")
![\"Two](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/normalized-runtime-versus-hsz.webp\")
![\"](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/normalized-runtime-versus-sequence-length.webp\")
![\"A](\"https:\/\/developer-blogs.nvidia.com\/wp-content\/uploads\/2026\/07\/atention-sharding-strategies.webp\")