TASP: Topology-aware Sequence Parallelism

TL;DR

TASP tackles the $O(n^2)$ self-attention bottleneck in long-context LLMs by aligning sequence parallelism with modern accelerator interconnects. It introduces topology-aware decomposition to produce mutually orthogonal ring datapaths and decomposes Ring AllGather into Multi-Ring AllGather, enabling full utilization of AlltoAll links. Empirical results across NVIDIA H100 and AMD MI300X platforms show up to $3.58\times$ speedups over Ring Attention and robust scaling in multi-node configurations, driven by improved communication efficiency. This topology- and primitive-aware approach offers a scalable path to efficient long-context inference in contemporary hardware, emphasizing the role of hardware-aware algorithm design for next-generation LLM systems.

Abstract

Long-context large language models (LLMs) face constraints due to the quadratic complexity of the self-attention mechanism. The mainstream sequence parallelism (SP) method, Ring Attention, attempts to solve this by distributing the query into multiple query chunks across accelerators and enable each Q tensor to access all KV tensors from other accelerators via the Ring AllGather communication primitive. However, it exhibits low communication efficiency, restricting its practical applicability. This inefficiency stems from the mismatch between the Ring AllGather communication primitive it adopts and the AlltoAll topology of modern accelerators. A Ring AllGather primitive is composed of iterations of ring-styled data transfer, which can only utilize a very limited fraction of an AlltoAll topology. Inspired by the Hamiltonian decomposition of complete directed graphs, we identify that modern accelerator topology can be decomposed into multiple orthogonal ring datapaths which can concurrently transfer data without interference. Based on this, we further observe that the Ring AllGather primitive can also be decomposed into the same number of concurrent ring-styled data transfer at every iteration. Based on these insights, we propose TASP, a topology-aware SP method for long-context LLMs that fully utilizes the communication capacity of modern accelerators via topology decomposition and primitive decomposition. Experimental results on both single-node and multi-node NVIDIA H100 systems and a single-node AMD MI300X system demonstrate that TASP achieves higher communication efficiency than Ring Attention on these modern accelerator topologies and achieves up to 3.58 speedup than Ring Attention and its variant Zigzag-Ring Attention. The code is available at https://github.com/infinigence/HamiltonAttention.

Figures (6)

• Figure 1: The mismatch between the ring-style data transfer of Ring AllGather and the fully-connected communication links of modern accelerators.
• Figure 2: Decomposition of 8-accelerator AlltoAll topology graph $K_8$ into 7 edge-disjoint directed Hamiltonian cycles. The decomposed Hamiltonian cycles correspond to mutually orthogonal ring datapaths that traverse all 8 accelerators.
• Figure 3: Left: The 15 Hamiltonian cycles decomposed from the topology of H100-2 via a $K_{16}-decomposition$ scheme. Right: The 8 Hamilton cycles derived from a $(8-K_8-8)^2$-decomposition scheme.
• Figure 4: Top: The Ring AllGather communication primitive used in Ring Attention and its variants. Colored circles represent accelerators 0–7. Symbols $t[0] \sim t[7]$ denote the transferred KV blocks. Bottom: The Multi-Ring AllGather communication primitive designed for $K_8-decomposition$. Each column illustrates one iteration with seven ring-styled data transfer for seven chunks of KV blocks. Symbol $t[i,j]$ represents the KV initially assigned at accelerator $j$ and circulates across all accelerators following the $i$-th ring-styled data transfer. Iterations 2–6 and ring-styled transfer 2-5 are omitted for brevity.
• Figure 5: $t_{all}$ speedup v.s. $CCR^B$. The test cases are divided into two sets (colored red and blue) to distinguish the data transfer size impact on our communication optimization. For cases with sufficient data volume to saturate the bandwidth (dots in red), TASP achieves a significant $t_{all}$ speedup compared to baselines. Furthermore, the speedup exhibits a decreasing trend with the increase of $CCR^B$, as the optimized part lies in communication.