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分配二维计算块，降低通信-计算比率，提升效率。
  3. 实验表明，Mesh-Attention在256个GPU上实现了高达3.4倍的加速，通信量减少高达85.4%。

📝 摘要（中文）

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

🔬 方法详解

问题定义：现有分布式注意力机制，如Ring-Attention，在扩展LLM上下文窗口时面临通信瓶颈。随着模型和数据规模的增大，GPU之间的通信量急剧增加，严重影响训练效率。Ring-Attention主要采用一维的通信模式，数据局部性较差，导致大量的全局通信开销。

核心思路：Mesh-Attention的核心思路是通过二维的计算块分配策略，提高数据局部性，从而降低通信-计算比率。将计算任务划分为二维网格，每个GPU负责一个二维的tile，使得GPU可以更多地利用本地数据进行计算，减少跨GPU的数据传输。通过优化tile的形状和调度，进一步降低通信开销。

技术框架：Mesh-Attention的整体框架包括以下几个主要阶段：1) 将注意力计算任务划分为二维网格；2) 为每个GPU分配一个二维tile；3) 在tile内部进行局部注意力计算；4) 通过GPU之间的通信交换数据，完成全局注意力计算；5) 使用贪婪算法搜索最佳的tile调度方案，以最小化通信开销。

关键创新：Mesh-Attention的关键创新在于二维tile的分配策略和贪婪调度算法。与Ring-Attention的一维通信模式不同，Mesh-Attention通过二维tile提高了数据局部性，减少了全局通信量。贪婪调度算法能够在满足通信约束的条件下，高效地搜索最佳的tile调度方案，进一步优化通信效率。

关键设计：Mesh-Attention的关键设计包括：1) Tile形状的选择：不同的tile形状会影响通信-计算比率，需要根据具体的硬件和模型特点进行调整；2) 贪婪调度算法：该算法需要在搜索空间中找到最佳的tile调度方案，需要考虑GPU之间的通信带宽和计算能力；3) 通信约束：为了保证通信的正确性和效率，需要设置合理的通信约束，例如限制GPU之间的通信量和通信频率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Mesh-Attention在256个GPU上相比现有方法实现了显著的性能提升。具体来说，Mesh-Attention可以实现高达3.4倍的加速（平均2.9倍），并将通信量减少高达85.4%（平均79.0%）。此外，可扩展性实验表明，Mesh-Attention在系统规模增大时仍能保持优越的性能，有效降低大规模部署的开销。这些结果充分验证了Mesh-Attention的优势。

🎯 应用场景

Mesh-Attention具有广泛的应用前景，尤其是在需要大规模分布式训练的LLM领域。它可以加速模型训练，降低训练成本，并支持更大规模的模型和数据集。此外，该方法还可以应用于其他需要分布式计算的场景，例如图神经网络、科学计算等。Mesh-Attention的出现为解决大规模模型训练的通信瓶颈问题提供了一种新的思路。

📄 摘要（原文）

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.