Cost-Optimal Grouped-Query Attention for Long-Context Modeling

Grouped-Query Attention (GQA) is a widely adopted strategy for reducing the computational cost of attention layers in large language models (LLMs). However, current GQA configurations are often suboptimal because they overlook how context length influences inference cost. Since inference cost grows with context length, the most cost-efficient GQA configuration should vary accordingly. In this work, we analyze the relationship among context length, model size, GQA configuration, and model loss, and introduce two innovations: (1) we decouple the total head size from the hidden size, enabling more flexible control over attention FLOPs; and (2) we jointly optimize the model size and the GQA configuration to arrive at a better allocation of inference resources between attention layers and other components. Our analysis reveals that commonly used GQA configurations are highly suboptimal for longcontext scenarios. Moreover, we propose a recipe for deriving cost-optimal GQA configurations. Our results show that for long-context scenarios, one should use fewer attention heads while scaling up the model size. Configurations selected by our recipe can reduce both memory usage and FLOPs by more than 50% compared to Llama-3's GQA, with no degradation in model capabilities. Our findings offer valuable insights for designing efficient longcontext LLMs. 1 Memory (GB) -57.8% -50.8% FLOPs (1e15) Llama-3 GQA Ours Llama-3 GQA Ours Change 1: Result Change 2: Time-variant memory/FLOPs Time-invariant memory/FLOPs Time-variant memory/FLOPs Time-invariant memory/FLOPs Faster but worse Faster and better KV Cache Attention FLOPs Adjustable Fixed Decoupling hidden size and head number KV Cache Attention FLOPs