Sparse attention can effectively mitigate the significant memory and throughput demands of Large Language Models (LLMs) in long contexts. Existing methods typically employ a uniform sparse attention mask, applying the same sparse pattern across different attention heads and input lengths. However, this uniform approach fails to capture the diverse attention patterns inherent in LLMs, ignoring their distinct accuracy-latency trade-offs. To address this challenge, we propose the Mixture of Attention (MoA), which automatically tailors distinct sparse attention configurations to different heads and layers. MoA constructs and navigates a search space of various attention patterns and their scaling rules relative to input sequence lengths. It profiles the model, evaluates potential configurations, and pinpoints the optimal sparse attention compression plan. MoA adapts to varying input sizes, revealing that some attention heads expand their focus to accommodate longer sequences, while other heads consistently concentrate on fixed-length local contexts. Experiments show that MoA increases the effective context length by $3.9\times$ with the same average attention span, boosting retrieval accuracy by $1.5-7.1\times$ over the uniform-attention baseline across Vicuna-{7B,13B}, and Llama3-{8B,70B} models. Moreover, MoA narrows the capability gaps between sparse and dense models, reducing the maximum relative performance drop from $9\%-36\%$ to within $5\%$ across two long-context understanding benchmarks. MoA achieves a $1.2-1.4\times$ GPU memory reduction, boosting decode throughput by $6.6-8.2\times$ and $1.7-1.9\times$ compared to FlashAttention2 and vLLM, with minimal impact on performance. Our code is available at \url{https://github.com/thu-nics/MoA}.

翻译：稀疏注意力能有效缓解大语言模型（LLMs）在长上下文场景下面临的巨大内存与吞吐量需求。现有方法通常采用统一的稀疏注意力掩码，在不同注意力头和输入长度上应用相同的稀疏模式。然而，这种统一方法无法捕捉LLMs固有的多样化注意力模式，忽视了其各自不同的精度-延迟权衡。为应对这一挑战，我们提出了注意力混合方法（MoA），它能自动为不同注意力头和网络层定制专属的稀疏注意力配置。MoA构建并探索了多种注意力模式及其随输入序列长度变化的缩放规则所组成的搜索空间。该方法通过分析模型特性、评估潜在配置，最终确定最优的稀疏注意力压缩方案。MoA能自适应变化的输入规模，研究发现：部分注意力头会扩展关注范围以适应更长序列，而其他注意力头则持续聚焦于固定长度的局部上下文。实验表明，在保持相同平均注意力范围的前提下，MoA将有效上下文长度提升至$3.9$倍，在Vicuna-{7B,13B}和Llama3-{8B,70B}系列模型上，其检索准确率较统一注意力基线提升$1.5-7.1$倍。此外，MoA显著缩小了稀疏模型与稠密模型之间的能力差距，在两个长上下文理解基准测试中，将最大相对性能下降从$9\%-36\%$缩减至$5\%$以内。相比FlashAttention2和vLLM，MoA实现了$1.2-1.4$倍的GPU内存降低，解码吞吐量分别提升$6.6-8.2$倍和$1.7-1.9$倍，且对模型性能影响极小。代码已开源：\url{https://github.com/thu-nics/MoA}。