Generated by Codex with GPT 5.6 Sol XHigh

Attention has become an architecture problem

The official NVIDIA Technical Blog published this post on July 31, 2026. Its central argument is that long-context inference cannot be optimized only after a model has been trained. Choices such as how many query heads share each key-value head, how wide each head is, and how attention is distributed across GPUs determine whether the hardware can execute the model efficiently. Kernel tuning still matters, but the model architecture sets the shapes, memory traffic, and parallelism limits that the kernels inherit.

The shift is visible as prompts get longer. In NVIDIA’s DeepSeek-R1 measurements, attention accounted for 18 percent of prefill time at a 4K-token context and 85 percent at 128K. That makes attention the dominant cost precisely in the agentic and multi-turn workloads where long histories, retrieved documents, tool results, and prefix caches are most valuable.

The post turns that problem into four model-design guidelines. The important contribution is not the checklist by itself, but the reasoning behind it: prefill and decode are different computational regimes, so an architectural choice that barely affects one phase can dominate the other.

Prefill and decode stress different parts of the GPU

Prefill processes all input tokens in parallel. For dense attention, each token compares with every other token, so the work grows quadratically with input length. These large matrix multiplications generally have enough computation per byte loaded to keep the GPU’s arithmetic units busy. In roofline-model terms, prefill sits above the hardware’s ridge point and is compute-bound.

Decode usually generates one token at a time. Each step performs much smaller matrix multiplications but must read the existing key-value cache for the full context from high-bandwidth memory. The amount of useful computation per byte is low, leaving decode limited by memory bandwidth rather than raw tensor-core throughput. Its per-token cost grows linearly with the key-value cache length.

This distinction also explains two important exceptions. Speculative decoding evaluates several candidate tokens together, enlarging the matrix operation and potentially moving decode toward the compute-bound regime. Prefix caching creates the reverse situation: a new turn may contain only a few fresh input tokens but attend to a large cached prefix, making that prefill operation behave more like memory-bound decode.

NVIDIA grounds the analysis in FlashAttention. Rather than materializing the full attention matrix in external memory, the kernel streams tiles of queries, keys, and values into on-chip SRAM, computes query-key scores, applies an online softmax, and immediately combines the weights with values. This reduces memory traffic, but it does not erase the underlying shape constraints. The number and dimensions of heads still determine tile utilization, while the cache size still determines how many bytes decode must stream.

Sharing key-value heads is primarily a decode optimization

The first design variable is group size: the number of query heads that share one key-value head. Standard multi-head attention gives each query head its own key and value heads. Grouped-query attention shares them within groups, while multi-query attention uses one key-value head for all query heads.

Increasing the group size leaves the total attention arithmetic nearly unchanged but reduces the key-value state that must be stored and read. NVIDIA’s derivation shows that prefill is largely insensitive to this choice once the sequence is long: increasing the group size from 8 to 16 at 32K tokens improves arithmetic intensity by less than 6 percent, and measured prefill runtime changes by under 1 percent across a much wider range.

Decode is different. When the context is much longer than the group size, decode arithmetic intensity is approximately twice the group size. Doubling the number of query heads sharing each key-value head roughly halves the cache traffic per token. In the measured kernels, decode runtime fell by about two times with each doubling until fixed overheads and the work of reducing results across streaming multiprocessors began to dominate. Longer, 128K-token caches amortized those overheads better and continued to benefit.

The practical rule is therefore to choose a high group size for decode efficiency. This is a useful example of why average FLOP counts can mislead: group size does not substantially change the mathematical work, yet it can transform runtime by changing how much state crosses the memory interface and whether the GPU’s compute tiles receive enough work.

Head width must fit the machine, not just the model

Head dimension does not change arithmetic intensity because doubling it doubles both the arithmetic and the bytes moved. It still changes performance through alignment and resource limits.

NVIDIA recommends head dimensions of 128 or 256 on the hardware it measured. A 64-wide head can still occupy a 128-wide tensor-core tile, paying for capacity it does not use. Dimensions aligned to 128-byte memory transactions also make cache access more efficient. At the other extreme, dimensions of 512 or more approach tensor-memory capacity limits and make scheduling harder.

The cost composition differs by phase. During prefill, both matrix multiplication and softmax contribute to runtime. Matrix work grows with head width, while softmax operates on the token-to-token score matrix and does not. A wider head therefore increases prefill time less than proportionally because it amortizes the fixed softmax component. During decode, widening the head enlarges every key and value vector in the cache, so the memory-bound runtime tracks the additional bytes more directly.

This is model-hardware co-design in a concrete form. A mathematically reasonable dimension can be operationally inefficient if it leaves accelerator tiles partly empty, misaligns transfers, or consumes scarce on-chip storage. Those effects are difficult to recover with a better kernel after the architecture is fixed.

Context length makes cache management a first-class concern

Dense-attention prefill grows with the square of input length because all token pairs interact. Decode grows linearly because each new token reads the accumulated cache once. Both curves eventually dominate fixed kernel overhead, but they point to different remedies.

For prefill, long inputs create an unavoidable compute problem unless the model stops attending densely to every prior token. For decode, the expanding key-value state becomes a bandwidth and capacity problem. NVIDIA therefore recommends reducing the effective cache through compression, sparse or sliding-window attention, or hybrid architectures in which only some layers maintain global attention state.

The word “effective” matters. Applications may need to preserve a long logical history without keeping a full-resolution key and value vector for every token in every layer. Model architects can trade among global reach, local detail, cache size, and layer structure. Serving systems can add prefix reuse or speculative decoding, but the largest gains may require changing what the model stores and attends to in the first place.

Key-value heads also set the multi-GPU scaling limit

Tensor parallelism usually divides attention heads across GPUs. It reduces the number of heads assigned to each device, but it does not change the per-head matrix shape or arithmetic intensity. Efficient sharding therefore depends on the number of key-value heads.

Once the tensor-parallel degree exceeds that count, the GPUs can no longer each receive a complete, distinct key-value head. The system must duplicate shared key-value state across ranks, adding memory capacity and bandwidth costs without eliminating attention work. NVIDIA’s rule is to keep tensor parallelism at or below the key-value head count.

That creates a deliberate tension. A high group size improves single-request decode by reducing the number of key-value heads, but it also removes head-level parallelism. Models with only one or two key-value heads quickly exhaust conventional tensor parallelism. They must scale attention along another axis.

The post points to Attention Data Parallelism, which distributes requests, and KV Parallelism, which partitions a long cache across devices. Feed-forward mixture-of-experts layers can be scaled separately with expert parallelism. TensorRT-LLM combines these strategies in Wide EP, which pairs attention data parallelism with expert parallelism, and Helix Parallelism, which pairs key-value parallelism with expert parallelism. The architecture of one layer therefore constrains the distributed topology of the whole serving system.

The broader engineering takeaway

The four recommendations form a coherent design method: use a large query-to-key-value group for decode efficiency, choose 128- or 256-wide heads that fit GPU execution units, reduce the effective key-value state, and let the remaining key-value head count determine the parallelism strategy.

More broadly, the post shows why inference performance should be modeled before model architecture is frozen. A team’s latency and throughput targets can be translated into anticipated prompt lengths, decode batch sizes, cache footprints, arithmetic intensity, tile shapes, and sharding limits. That analysis exposes conflicts early—for example, when reducing cache traffic also reduces tensor-parallel opportunities—and encourages the model and serving stack to be designed together.

Long-context capability is not merely the ability to accept more tokens. It is the ability to process and revisit those tokens within an acceptable latency, memory, and cost envelope. As attention takes over the runtime, the boundary between model research and systems engineering disappears: the model defines the workload, and the machine’s constraints should help define the model.