Mesh-Attention: A New Communication-Efficient Distributed Attention with Improved Data Locality

📄 arXiv: 2512.20968v1 📥 PDF

作者: Sirui Chen, Jingji Chen, Siqi Zhu, Ziheng Jiang, Yanghua Peng, Xuehai Qian

分类: cs.DC, cs.AI

发布日期: 2025-12-24

💡 一句话要点

提出Mesh-Attention，通过优化数据局部性提升分布式Attention的通信效率，加速LLM训练。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 分布式注意力 通信效率 大型语言模型 GPU集群 数据局部性

📋 核心要点

  1. 现有Ring-Attention方法在扩展LLM上下文窗口时，由于通信量过大，可扩展性受限，成为瓶颈。
  2. Mesh-Attention通过为每个GPU分配二维计算块，降低通信-计算比率，优化数据局部性，提升分布式Attention效率。
  3. 实验表明，Mesh-Attention在256个GPU上实现了高达3.4倍的加速，并显著降低了通信量，验证了其优势。

📝 摘要（中文）

分布式注意力机制是扩展大型语言模型(LLMs)上下文窗口的关键。现有最优方法Ring-Attention因通信量过大而面临可扩展性限制。本文提出了一种新的分布式注意力算法Mesh-Attention，通过基于矩阵的新模型重新思考分布式注意力的设计空间。该方法为每个GPU分配一个二维计算块（而非一维行或列），通过降低通信-计算(CommCom)比率来实现更高的效率。该方法将Ring-Attention作为特例，并允许通过不同的块形状调整CommCom比率。重要的是，我们提出了一种贪婪算法，可以在块内有效地搜索调度空间，同时限制GPU之间的有效通信。理论分析表明，与其他现有算法相比，Mesh-Attention具有更低的通信复杂度和良好的可扩展性。实验结果表明，在256个GPU上，Mesh-Attention可以实现高达3.4倍的加速（平均2.9倍），并将通信量减少高达85.4%（平均79.0%）。可扩展性结果进一步表明，Mesh-Attention在系统扩展时保持卓越的性能，从而大大减少了大规模部署中的开销。结果令人信服地证实了Mesh-Attention的优势。

🔬 方法详解

问题定义：论文旨在解决分布式注意力机制在大规模语言模型训练中，因通信量过大导致的可扩展性瓶颈问题。现有的Ring-Attention方法虽然能够实现分布式计算，但其通信模式导致了较高的通信开销，限制了其在大规模GPU集群上的应用效果。

核心思路：Mesh-Attention的核心思路是通过改变数据划分和计算分配方式，优化通信模式，从而降低通信开销。具体来说，它将计算任务划分为二维的块，并将这些块分配给GPU进行计算。通过合理安排块的形状和GPU之间的通信方式，可以显著降低通信-计算比率，提高整体效率。

技术框架：Mesh-Attention的技术框架主要包括以下几个阶段：1) 数据划分：将注意力计算所需的输入数据划分为二维的块。2) 任务分配：将这些块分配给不同的GPU进行计算。3) 通信调度：设计一种高效的通信调度策略，使得GPU之间能够以最小的通信量完成数据交换。4) 结果聚合：将各个GPU的计算结果进行聚合，得到最终的注意力输出。

关键创新：Mesh-Attention的关键创新在于其二维块划分和通信调度策略。与Ring-Attention的一维划分方式不同，二维划分能够更好地利用GPU的计算资源，降低通信开销。此外，论文还提出了一种贪婪算法，用于在块内搜索最佳的调度方案，以确保GPU之间的有效通信。

关键设计：Mesh-Attention的关键设计包括：1) 块的形状：块的形状会影响通信-计算比率，需要根据具体的硬件和任务特点进行调整。2) 通信调度策略：论文提出了一种贪婪算法，用于搜索最佳的通信调度方案。该算法考虑了GPU之间的通信带宽和延迟等因素，以最小化通信开销。3) 损失函数：论文没有特别提到损失函数的设计，但可以推测其使用了标准的注意力机制损失函数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Mesh-Attention在256个GPU上实现了高达3.4倍的加速（平均2.9倍），并将通信量减少高达85.4%（平均79.0%）。与Ring-Attention相比，Mesh-Attention在可扩展性方面表现出更强的优势，能够更好地适应大规模GPU集群。这些结果充分证明了Mesh-Attention在分布式注意力机制方面的优越性。

🎯 应用场景

Mesh-Attention在大型语言模型训练、图像识别、语音识别等领域具有广泛的应用前景。通过降低分布式注意力机制的通信开销，可以加速模型的训练过程，提高模型的性能。此外，该方法还可以应用于其他需要大规模分布式计算的场景，例如科学计算、金融分析等。未来，Mesh-Attention有望成为大规模AI模型训练的重要技术支撑。

📄 摘要（原文）

Distributed attention is a fundamental problem for scaling context window for Large Language Models (LLMs). The state-of-the-art method, Ring-Attention, suffers from scalability limitations due to its excessive communication traffic. This paper proposes a new distributed attention algorithm, Mesh-Attention, by rethinking the design space of distributed attention with a new matrix-based model. Our method assigns a two-dimensional tile -- rather than one-dimensional row or column -- of computation blocks to each GPU to achieve higher efficiency through lower communication-computation (CommCom) ratio. The general approach covers Ring-Attention as a special case, and allows the tuning of CommCom ratio with different tile shapes. Importantly, we propose a greedy algorithm that can efficiently search the scheduling space within the tile with restrictions that ensure efficient communication among GPUs. The theoretical analysis shows that Mesh-Attention leads to a much lower communication complexity and exhibits good scalability comparing to other current algorithms. Our extensive experiment results show that Mesh-Attention can achieve up to 3.4x speedup (2.9x on average) and reduce the communication volume by up to 85.4% (79.0% on average) on 256 GPUs. Our scalability results further demonstrate that Mesh-Attention sustains superior performance as the system scales, substantially reducing overhead in large-scale deployments. The results convincingly confirm the advantage of Mesh-Attention.