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

Summary

As context lengths increase in long-context and agentic workloads, attention increasingly dominates GPU inference runtime. Dense attention performance depends on architectural parameters like group size, head dimension, and sequence length, alongside GPU parallelization strategies. Prefill remains compute-bound and scales quadratically with input sequence length, showing minimal runtime sensitivity to group size variations. Conversely, decode generation is memory-bound and scales linearly with KV cache length, where doubling group size roughly halves decode execution time. To maximize throughput and interactivity on NVIDIA hardware, model architects should push group sizes higher, select head dimensions of 128 or 256, reduce effective KV state, and keep tensor parallelism less than or equal to the number of KV heads.

Takeaways

1. 1

   Decode runtime falls approximately 2x per doubling of group size because halving KV heads cuts memory traffic, whereas prefill runtime is compute-bound and varies by under 1% across group sizes from 1 to 64.

2. 2

   Prefill runtime scales quadratically with input sequence length while decode scales linearly with KV cache sequence length, making KV-cache compression, sparse attention, or hybrid architectures effective.

3. 3

   Tensor parallelism must remain less than or equal to the number of KV heads to prevent duplicate KV cache memory and bandwidth overhead across GPU ranks.