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

arXiv:2512.20968v2 Announce Type: replace-cross Abstract: 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.

Source: arXiv cs.AI | 2026-08-11