A deep technical analysis of how attention mechanism design choices affect LLM inference performance on NVIDIA GPUs. Covers how group size (GQA/MQA), head dimension, and sequence length shape arithmetic intensity in both prefill and decode phases. Distills findings into four co-design guidelines: maximize group size for decode efficiency, use head dimensions of 128 or 256, reduce KV state through compression or hybrid architectures, and match tensor parallelism strategy to KV head count. References TensorRT-LLM's Wide EP and Helix Parallelism for scaling models with few KV heads.
Table of contents
How are prefill and decode two different problems?How arithmetic intensity governs compute- versus memory-bound behaviorHow does the FlashAttention kernel compute attention on GPU?GEMM shapesGroup sizeHead dimensionSequence lengthTensor parallelism splits attention heads across GPUsGet started co-designing AI model attention33 Impressions