Cost-Optimal Grouped-Query Attention for Long-Context LLMs
Authors: name: Yingfa Chen; name: Yutong Wu; name: Xu Han; name: Zhiyuan Liu; name: Maosong Sun
Published: 2025-03-12 17:50:42

Abstract: Building effective and efficient Transformer-based large language models
(LLMs) has recently become a research focus, requiring maximizing model
language capabilities and minimizing training and deployment costs. Existing
efforts have primarily described complex relationships among model performance,
parameter size, and data size, as well as searched for the optimal compute
allocation to train LLMs. However, they overlook the impacts of context length
and attention head configuration (the number of query and key-value heads in
grouped-query attention) on training and inference. In this paper, we
systematically compare models with different parameter sizes, context lengths,
and attention head configurations in terms of model performance, computational
cost, and memory cost. Then, we extend the existing scaling methods, which are
based solely on parameter size and training compute, to guide the construction
of cost-optimal LLMs during both training and inference. Our quantitative
scaling studies show that, when processing sufficiently long sequences, a
larger model with fewer attention heads can achieve a lower loss while
incurring lower computational and memory costs. Our findings provide valuable
insights for developing practical LLMs, especially in long-context processing
scenarios. We will publicly release our code and data.