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

arXiv: https://arxiv.org/abs/2503.09579 · HTML: https://mlsanity.com/paper/2503.09579/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 39 references that are in the map.
- First posted on arXiv 12 Mar 2025.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization](https://mlsanity.com/paper/2401.18079/) (2024)
- [Deep Learning Scaling is Predictable, Empirically](https://mlsanity.com/paper/1712.00409/) (2017)
- [Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs](https://mlsanity.com/paper/2310.01801/) (2023)
- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [LLMLingua-2: Data Distillation for Efficient and Faithful Task-Agnostic Prompt Compression](https://mlsanity.com/paper/2403.12968/) (2024)
- [Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Laws](https://mlsanity.com/paper/2401.00448/) (2023)
- [A Survey on Large Language Model Acceleration based on KV Cache Management](https://mlsanity.com/paper/2412.19442/) (2024)
- [Small Language Models: Survey, Measurements, and Insights](https://mlsanity.com/paper/2409.15790/) (2024)
- [Keep the Cost Down: A Review on Methods to Optimize LLM' s KV-Cache Consumption](https://mlsanity.com/paper/2407.18003/) (2024)
- [Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers](https://mlsanity.com/paper/2406.16747/) (2024)
- [KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches](https://mlsanity.com/paper/2407.01527/) (2024)
- [LoCoCo: Dropping In Convolutions for Long Context Compression](https://mlsanity.com/paper/2406.05317/) (2024)

## Built on by (most cited first, 4 of 4)

- [Hybrid Linear Attention Done Right: Efficient Distillation and Effective Architectures for Extremely Long Contexts](https://mlsanity.com/paper/2601.22156/) (2026)
- [A Theory of Inference Compute Scaling: Reasoning through Directed Stochastic Skill Search](https://mlsanity.com/paper/2507.00004/) (2025)
- [Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection](https://mlsanity.com/paper/2603.04427/) (2026)
- [Entropy Meets Importance: A Unified Head Importance-Entropy Score for Stable and Efficient Transformer Pruning](https://mlsanity.com/paper/2510.13832/) (2025)

## Abstract (arXiv)

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 also 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 long-context scenarios. More importantly, 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 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 long-context LLMs. The code is available at https://www.github.com/THUNLP/cost-optimal-gqa .