arXiv:2512.20968cs.DCcs.AI2025-12被引 2

提出Mesh-Attention,让大模型长文本推理更高效且通信开销更低。

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

  • 用二维分块方式分配计算任务,平衡查询与键值局部性。
  • 在256张GPU上处理百万级令牌时,速度比环形注意力快17倍。
  • 适合大规模分布式语言模型训练,尤其对长序列场景优化显著。

分布式注意力对扩展大语言模型至长上下文至关重要,但现有方法或并行度受限,或通信开销高。Ulysses采用高效的全对全通信,但无法超越注意力头数量;而环形注意力虽无此限制,却带来每卡高通信量。分析表明,环形注意力将每张卡分配完整行的键值交互矩阵,保留查询局部性但牺牲键值局部性,导致每卡接收近全部键值分区,通信量随序列长度线性增长。本文提出Mesh-Attention,将每张卡分配一个二维块(tile),仅收集部分查询和键值分区,实现查询与键值局部性的均衡,从而在渐进意义上降低通信复杂度,且不限制并行度。此外,结合键值分块旋转(KVR)、贪心调度与拓扑感知映射,进一步平衡流量、重叠通信与计算,并减少低带宽链路的传输。在最多256张GPU和100万令牌序列上的实验显示,相比环形注意力、USP(Ulysses阶8)和StarTrail,平均提速分别为17.00倍、2.62倍和2.94倍(峰值达32.27倍、4.36倍和4.59倍),且随着GPU数量和序列长度增加仍保持优势,大幅降低大规模下的通信开销。

原文摘要 · Abstract (English)

Distributed attention is essential for scaling large language models (LLMs) to long contexts, yet existing methods either have limited parallelism or incur high communication costs. Ulysses uses efficient all-to-all communication but cannot scale beyond the number of attention heads, whereas Ring-Attention removes this limit at the cost of high per-GPU communication. In our analytical Q-KV interaction matrix model, Ring-Attention assigns each GPU an entire row, preserving Q locality while sacrificing KV locality. Each GPU therefore receives nearly all KV partitions, and its communication grows linearly with sequence length. We present Mesh-Attention, which instead assigns each GPU a 2D tile to balance Q and KV locality, so that it collects only subsets of Q and KV partitions. This gives Mesh-Attention asymptotically lower communication complexity than Ring-Attention without limiting parallelism. Mesh-Attention further uses KV Partition Rotation (KVR), greedy scheduling, and topology-aware GPU mapping to balance traffic, overlap communication with computation, and reduce traffic over low-bandwidth links. Across experiments on up to 256 GPUs and 1M-token sequences, Mesh-Attention achieves average speedups of 17.00x, 2.62x and 2.94x (up to 32.27x, 4.36x and 4.59x) over Ring-Attention, USP (Ulysses degree 8) and StarTrail, respectively. It maintains this performance advantage as GPU count and sequence length increase while substantially reducing communication overhead at scale.

分布式注意力机制大模型通信优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。