TASP: Topology-Aware Sequence Parallelism

标题：TASP：拓扑感知的序列并行
作者/机构：Yida Wang, Ke Hong, Xiuhong Li, Yuanchao Xu, Wenxun Wang, Guohao Dai, Yu Wang

A1 主要贡献

本文旨在解决长上下文大语言模型（LLMs）中自注意力机制的二次复杂度带来的扩展性瓶颈。现有的主流序列并行（SP）方法，如环形注意力（Ring Attention），通过将查询（Query）分块并在多个加速器之间分发，利用 Ring AllGather 通信原语使每个Q张量能访问所有KV张量。然而，这种方法的通信效率低下，限制了其实际应用。根本原因在于其采用的 Ring AllGather 通信原语与现代加速器的全连接（AlltoAll）拓扑结构之间存在不匹配，环形的数据传输方式只能利用全连接拓扑中极小部分的通信带宽。

核心问题：Ring Attention 使用的 Ring AllGather 通信原语采用环形数据传输，无法充分利用现代加速器（如NVIDIA H100, AMD MI300X）的全连接（AlltoAll）拓扑，导致通信效率低下，成为性能瓶颈，尤其是在计算通信比（CCR）低于1.0的情况下。如图1所示，当有8个加速器时，通信容量的利用率仅为1/(8-1) ≈ 14.3%。

研究目标：提出一种新颖的序列并行方法，旨在克服 Ring Attention 的通信瓶颈，通过充分利用现代加速器集群中的所有通信能力来提高效率。

创新点与主要贡献：
本文提出了 TASP (Topology-Aware Sequence Parallelism)，一种拓扑感知的序列并行方法，通过拓扑分解和原语分解来解决上述不匹配问题。
1. 方法论创新：首次引入了一种通过拓扑分解和原语分解来提升通信效率的新方法。
* 拓扑分解：受完全有向图的哈密顿分解理论启发，本文发现现代加速器的全连接拓扑可以被分解为多个相互正交（边不相交）的环形数据通路。这些通路可以并行传输数据而互不干扰。本文展示了单节点全连接拓扑（如图2所示的K8图）和多节点拓扑的分解方案。
* 原语分解：相应地，本文将 Ring AllGather 通信原语分解为多环 AllGather（Multi-Ring AllGather）。在每一轮迭代中，多个环形数据传输可以并发进行，从而与分解出的多个正交环形数据通路完美映射，实现对所有通信链路的充分利用（如图1所示）。

1. TASP 的设计与实现：

   • 成功将上述方法论应用于序列并行优化，提出了 TASP。它基于图论中的哈密顿分解，实现了数据传输与所有通信链路的完美映射。
   • 证明了 Multi-Ring AllGather 在保证分布式注意力计算正确性（可达性）和不增加额外内存开销（零拷贝）方面与 Ring AllGather 等价。
   • 针对语言模型中广泛使用的因果掩码（causal mask），提出了一种查询块放置策略（Zig-zag chunk placement），以进一步改善加速器之间的负载均衡。

2. 实验验证：

   • 在单节点和多节点的 NVIDIA H100 系统以及单节点的 AMD MI300X 系统上进行了广泛实验。结果表明，TASP 在这些现代加速器拓扑上的通信效率显著高于 Ring Attention。
   • 与 Ring Attention 及其变体 Zigzag-Ring Attention 相比，TASP 最高可实现 3.58 倍的加速。如表1所示，在CCR较低的情况下，TASP的加速效果尤为明显。

Table 1: 在8卡AMD-MI300X上，批量大小为48，不同序列长度下测量的TASP相对于Ring Attention的CCR和加速比。

A3 背景知识与动机

2.1 序列并行

长上下文LLM应用的挑战与SP技术。随着长上下文LLM应用的发展，计算开销和KV张量存储的内存需求成为巨大挑战。为解决这些问题，序列并行（SP）成为一种有前途的技术，可将工作负载有效分配到多个加速器上。SP方法主要分为两大类：Ulysses【【5，Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models，2023，arXiv】】和Ring Attention【【10，Ring attention with blockwise transformers for near-infinite context，2023，arXiv】】。

Ulysses方法的局限性。Ulysses首先沿序列长度维度分片并分发到加速器，然后通过AlltoAll操作按KV头维度对KV张量进行分区。这使得每个加速器都能保留其分配到的头的完整序列长度上的KV张量。由于每个KV头独立计算注意力，每个加速器可以为其对应的头计算完整的注意力输出。随后，通过另一次AlltoAll操作转置部分注意力输出，形成最终结果。然而，Ulysses利用KV头独立性的做法带来一个限制：其性能高度依赖于KV头的数量和各头之间的计算平衡。当前注意力机制的趋势是KV头数量减少，例如GQA（通常4或8个KV头）【【11，Cost-optimal grouped-query attention for long-context llms，2025，arXiv】】、MQA【【12，Gqa: Training generalized multi-query transformer models from multi-head checkpoints，2023，arXiv】】和MLA【【13，Multilayer attention is the amplifier of demonstration effectiveness，2025，arXiv】】（仅单个KV头）。这种减少限制了Ulysses在分布式环境中的可扩展性，因为它严重依赖足够数量的KV头来进行有效分区。此外，Ulysses很难将其通信与计算重叠，因为它们相互依赖。

Ring Attention方法的瓶颈。Ring Attention沿序列长度维度对查询（Q和KV）进行分区。Q张量固定在各自的加速器上，而KV块通过Ring AllGather通信原语在所有加速器之间轮转。经过多次迭代，每个加速器通过Flash Attention【【14，Flashattention: Fast and memoryefficient exact attention with io-awareness，2022，Advances in neural information processing systems】】为其本地查询段处理所有全局KV块，从而增量计算出完整的注意力输出。原始Ring Attention的一个根本挑战是计算不平衡，特别是在广泛使用的因果掩码下，注意力的计算成本在序列长度维度上分布不均。为解决此问题，已提出多种优化方案。例如，Striped Attention【【6，Striped attention: Faster ring attention for causal transformers，2023，arXiv】】采用条带化分区策略以实现更好的负载均衡。Megatron CP【【8，Megatron Context Parallelism，2023，NVIDIA】】使用轴对称分区方案，可为因果掩码实现完美的负载均衡。Ring Attention及其因果掩码平衡变体通常采用计算-通信重叠来减轻关键的通信开销。然而，完美重叠通信开销并不容易实现。如表1所示，表示计算与通信比率的CCR在许多情况下都低于1.0。因此，整体注意力性能仍然受限于Ring AllGather原语的通信开销。因此，提高通信效率对于普遍增强Ring Attention的性能至关重要。

2.2 硬件平台的互联性

现代AI集群对互联性的高要求。随着LLM规模的增长，单个加速器已无法满足高效服务的需求。因此，部署LLM推理服务越来越依赖于利用各种并行策略（如张量并行TP和专家并行EP【【15，Mixtral of experts，2024，arXiv】】）的加速器集群。然而，这些并行策略对设备间的互联性提出了严格要求：任何一对设备都必须支持高带宽、低延迟的P2P通信，并且多个此类通信应能无干扰地并发进行，以便高效执行TP和EP中使用的AlltoAll和AllGather等集体通信操作。

全连接拓扑成为主流。为了支持这些能力，全连接拓扑已成为最先进AI加速器网络的首选，主要分为两类：全网状（full-mesh）和基于交换机（switch-based）。例如，华为的Ascend 910B NPU【【16，Ascend: a scalable and unified architecture for ubiquitous deep neural network computing: Industry track paper，2021，IEEE International Symposium on HighPerformance Computer Architecture (HPCA)】】、AMD的MI300X系列GPU【【17，Amd instinct™ mi300x: A generative ai accelerator and platform architecture，2025，IEEE Micro】】和Intel的Gaudi 3 AI加速器【【18，Intel gaudi 3 ai accelerator: Architected for gen ai training and inference，2024，IEEE Hot Chips 36 Symposium (HCS)】】都在8个设备间采用全网状拓扑。对于全网状拓扑，传统的Ring Attention只能利用n个加速器间n(n-1)个单向通信链路中的1/(n-1)。相比之下，NVIDIA的GPU采用基于交换机的全连接拓扑【【19，3.2 the a100 datacenter gpu and ampere architecture，2021，IEEE International Solid-State Circuits Conference (ISSCC)】】，与全网状方法相比，允许更灵活的带宽分配。对于更大规模的集群，多级基于交换机的架构——例如CloudMatrix384【【20，Serving large language models on huawei cloudmatrix384，2025，arXiv】】——在多达384个Ascend 910 NPU和192个CPU之间提供完全的P2P互联。

现有通信原语的局限性。这些加速器互联性的进步强调了高性能P2P通信能力既是下一代AI硬件架构的基本要求，也是关键趋势。然而，一些通信原语无法充分利用AlltoAll拓扑中的通信链路，导致通信效率低下，从而增加端到端延迟。因此，在设计和实现并行策略时，必须充分考虑硬件平台的AlltoAll互联性，以最大化通信效率。

A2 方法细节

3. TASP 的设计

TASP的设计涉及两个方面：加速器拓扑分解以揭示充分利用通信链路的潜力，以及Ring AllGather原语分解以将分块的环形数据传输工作负载映射到所有通信链路上。

首先，我们证明了现代加速器拓扑，包括单节点和多节点拓扑，可以通过哈密顿分解被分解为一组相互正交的环形数据通路。这些环形数据通路可以同时被利用而不会产生干扰。

其次，我们证明了Ring Attention中使用的Ring AllGather原语可以分解为Multi-Ring AllGather，在其中，每次迭代都会并发执行多个环形数据传输。因此，我们可以将每个环形数据传输映射到每个环形数据通路上，从而充分利用所有通信链路。

Table 2: 符号说明。

3.1 加速器拓扑分解

将拓扑建模为有向图。通过将加速器视作顶点，所有单向通信链路的集合构成了一个有向边集：

$$E = \{(R_i, R_j, c_k) \mid \text{there exists a communication link from } R_i \text{ to } R_j \text{ via cable } c_k\}.$$

因此，该拓扑可以表示为一个有向图 $G = (V, E)$。

3.1.1 环形数据通路定义

基于图论的定义。我们基于将加速器拓扑形式化为有向图来定义环形数据通路。考虑一组加速器，表示为 $V = \{R_0, R_1, \dots, R_{n-1}\}$，其中 $R_i$ 表示通信组中的第 $i$ 个 rank。每个加速器 $R_i$ 有多个通信端口，可以通过电缆或光纤以及可选的一组网络交换机连接到其他加速器。由于采用双工通信模式，连接 $R_i$ 和 $R_j$ 的电缆 $c_k$ 建立了两条逻辑上的单向通信链路，表示为 $(R_i, R_j, c_k)$ 和 $(R_j, R_i, c_k)$。对于通过交换机互连的加速器，每对加速器之间存在两条通信链路，但这些链路共享相同数量的聚合带宽。

环形数据通路与正交性。接下来，我们定义一个环形数据通路。图 $G$ 上的一个环形数据通路是一个子图 $D \subseteq G$，它形成一个有向哈密顿环，即一个恰好访问 $V$ 中每个顶点一次的有向环。两个环形数据通路 $D_i = (V, E_i)$ 和 $D_j = (V, E_j)$（其中 $i \neq j$）被称为正交环形数据通路，如果它们是边不相交的，即：

$E_i \cap E_j = \emptyset.$

此属性确保了一组相互正交的环形数据通路可以支持同时进行的数据传输而不会相互干扰。

3.1.2 单节点拓扑分解

应用哈密顿分解。现代单节点加速器中最常见的拓扑是AlltoAll拓扑，包括基于交换机（NVIDIA GPU）和全网状（AMD GPU）的拓扑。理论上，这两种拓扑都可以表示为一个 $n$ 阶的完全有向图 $K_n$。对于这些AlltoAll拓扑的分解，我们可以直接应用现有的将完全有向图分解为边不相交的哈密顿环的理论【【21，Edge disjoint hamiltonian cycles in k-ary n-cubes and hypercubes，2003，IEEE Transactions on Computers】、【22，Hamiltonian decompositions of prisms over cubic graphs，2004，Discrete Mathematics】、【23，On hamiltonian decompositions of complete 3- uniform hypergraphs，2024，Discrete Mathematics】、【24，Hamiltonian decomposition of complete regular multipartite digraphs，1997，Discrete Mathematics】、【25，Hamiltonian decompositions of cayley graphs on abelian groups of even order，2003，Journal of Combinatorial Theory, Series B】】（本文中称为$K_n$-分解）。

8加速器示例。如图2所示，我们以单个节点内的8个加速器（即$K_8$）为例，其分解方案通过【【26，A hamiltonian decomposition of k2m*, 2m≥8，1980，Journal of Combinatorial Theory, Series B】】提出的一个算法计算得出，该算法能在 $O(n)$ 时间复杂度内分解一个完全图 $K_n$。基于此分解， $n \times (n-1)$ 条通信链路被划分为一组 $n-1$ 个边不相交的哈密顿环。这些环随后被用作并发数据传输的相互正交的环形数据通路。

图2：8加速器AlltoAll拓扑图K8分解为7个边不相交的有向哈密顿环。分解出的哈密顿环对应于遍历所有8个加速器的相互正交的环形数据通路。

3.1.3 多节点拓扑分解

多节点系统拓扑。在多节点系统中，例如NVIDIA DGX H100 GPU集群【【27，The nvlink-network switch: Nvidia’s switch chip for high communicationbandwidth superpods，2022，IEEE Hot Chips 34 Symposium (HCS)】】，多个H100服务器（每个配备8个H100加速器）通过高带宽的Infiniband网络互连。在每个节点内部，8个加速器通过四个超高带宽的NVSwitch全连接，同时每个加速器也连接到一个IB NIC，使得所有8个加速器都能通过各自专用的400 Gb/s NIC独立且并发地进行RDMA数据传输。

朴素分解方案及其问题。为了分析这种IB连接的多节点加速器的分解，我们首先以一个2节点的H100集群为例（简写为H100-2）。由于H100-2中的任意两个加速器要么通过超高带宽（7.2 TB/s双工）的基于NVLINK的通信链路互连，要么通过高带宽（800 GB/s双工）的基于IB的通信链路互连，我们可以简单地将其拓扑建模为$K_{16}$，并使用上述的$K_n$-分解方案进行拓扑分解，如图3左侧所示。然而，这种分解方案会导致节点内基于NVLINK的通信链路带宽利用不足。在一个$K_{16}$-分解方案中，每个节点内有$7 \times 8 = 56$条通信链路，而从一个节点到另一个节点有$8 \times 8 = 64$条通信链路。因此，每个节点内通信链路的带宽约为$7200 / (56 \times 2) \approx 64$ GB/s，而每个节点间通信链路的带宽约为$800 / (64 \times 2) \approx 6.3$ GB/s，比节点内带宽低一个数量级。因此，节点内带宽的利用率受限于节点间带宽。

优化的分解方案。为此，我们提出将H100-2的拓扑建模为两个$K_8$子图，它们之间通过8条双向或16条单向的IB链路互连（表示为$(m-K_m-m)_u$-分解，在此例中$m=8$，$u=2$）。该分解方案基于将每个节点分解为一组哈密顿路径。哈密顿路径是图中的一条路径，恰好访问每个顶点一次。如图3右侧所示，每个单向节点内通信链路（灰色）的理论带宽约为64 GB/s，而每个单向节点间通信链路（蓝色）的理论带宽约为50 GB/s——仅有28%的差异，远小于采用$K_{16}$分解时超过9倍的差异。

图3：左：通过K16分解方案从H100-2拓扑中分解出的15个哈密顿环。右：通过(8-K8-8)2分解方案得到的8个哈密顿环。

方案的通用性与扩展性。接下来，我们证明2节点分解方案可以推广到任意u个节点的集群。如果我们将两个$K_8$子图分解为哈密顿路径，我们可以通过将这些路径端到端连接来获得8个哈密顿环。由于分解出的哈密顿路径同时也是一个行完备的拉丁方，在拉丁方的每一行或每一列中，元素都是互不相同的。因此，所有8个加速器都会出现在分解的哈密顿路径的起点和终点，从而利用了全部$2 \times 8 = 16$个IB NIC。

数学归纳法证明扩展性。此外，通过数学归纳法可以证明，上述分解方案可以扩展到任意数量的通过IB交换机互连的H100节点。假设我们已经使用该方案分解了一个n节点集群，我们可以通过简单的两步推导出一个n+1节点集群的分解（ranks代表加速器）：(1) 断开从rank $(8n-8) \sim (8n-1)$到rank $0 \sim 7$的通信链路；(2) 将rank $(8n-8) \sim (8n-1)$连接到rank $8n \sim (8n+7)$，并将rank $8n \sim (8n+7)$连接到rank $0 \sim 7$。

其他拓扑。其他拓扑，如超立方体，也可以通过哈密顿分解为正交的环形数据通路。然而，由于篇幅限制以及架构上对AlltoAll连接性的日益重视，本文主要关注本节分析的单节点和多节点AlltoAll拓扑。

3.2 Ring AllGather 的分解

3.2.1 问题定义

Ring Attention的构成与性质。为简单起见，我们假设对于每个查询块，其Q张量保持静止，其KV张量以环形数据传输方式发送给或从其他加速器接收。基于此假设，原生Ring Attention的公式化包含两个独立部分：一个Ring AllGather通信原语，它在每个加速器上分配一个查询块并为发送和接收这些块指定一个环形数据传输模式；以及一个块放置策略，定义了加速器上（初始的）查询块放置策略。因此，给定一个Ring AllGather原语和一个块放置策略，原生Ring Attention方案的通信和计算在每次迭代中都是明确且完全确定的。如图4所示，原生Ring Attention具有三个理想属性：

• 可达性（Accessibility）: 经过 $n-1$ 次KV传输迭代和 $n$ 次注意力计算迭代后，每个Q张量都能在其分配的加速器上本地访问到所有（最初分布在所有加速器上的）KV张量。
• 零拷贝（Zero-copy）: 由于KV张量的复制，不会消耗额外的内存。
• 负载均衡（Load balance）: 对于双向非因果注意力，只要查询块在加速器间均匀分布，计算和内存负载在所有加速器上都是平衡的。对于因果注意力，则有两种负载均衡方法，包括Striped Attention【【6，Striped attention: Faster ring attention for causal transformers，2023，arXiv】】和Zig-zag Ring Attention【【8，Megatron Context Parallelism，2023，NVIDIA】】。

分解目标。上述属性中的前两个对应于原生Ring Attention的第一个组成部分——Ring AllGather通信原语，确保了注意力计算的算法正确性。第三个属性对应于原生Ring Attention的块放置策略，并影响每次迭代的负载均衡。接下来，我们证明Ring AllGather可以被分解为Multi-Ring AllGather，同时满足前两个属性（见3.2.2节）。最后，Ring Attention的块放置策略（Striped Attention和Zig-zag Ring Attention）不能直接移植到TASP。因此，我们提出了一种用于实现完美负载均衡的Zig-zag块放置策略（见3.2.3节）。

3.2.2 Multi-Ring AllGather

Multi-Ring AllGather的定义。如图4所示，与原始Ring AllGather原语仅为每个加速器分配一个查询块相比，一个m环的Multi-Ring AllGather通信原语为每个加速器分配m个查询块，并为每个查询块指定一个特定的环形数据传输方式。m的数量可以任意选择，只要原始查询可以被均匀地分成m个查询块即可。每个查询块的环形数据传输方式也可以任意设计且互不相同。

与拓扑分解的对齐。由于前一节已证明现代加速器的拓扑可以分解为一组相互正交的环形数据通路，因此我们可以设计一个与拓扑分解百分之百对齐的Multi-Ring AllGather通信原语。例如，我们展示了一个专为8加速器AlltoAll拓扑的$K_8$-分解设计的7环Multi-Ring AllGather通信原语，如图4所示。如果环形数据传输与环形数据通路百分之百对齐，就可以实现完全的通信链路利用。

性质证明。然而，为确保Ring AllGather的这种分解不会产生额外的内存开销或违反注意力计算的算法正确性，我们需要证明它仍然遵循以下性质：
* 可达性证明：对于任意一个（分配给加速器$s_0$的）$Q_i$张量和任意一个$KV_j$张量，由于$KV_j$张量沿着其指定的环形数据传输路径$h_0$（一个哈密顿环）遍历所有加速器，它最终会到达加速器$s_0$，从而允许$Q_i$在本地对$KV_j$进行注意力计算。
* 零拷贝证明：每个KV张量被精确地分配给一个环形数据传输路径$h_0$，并且仅沿着$h_0$传输，确保任何时候只存在一个副本。

图4：上：Ring Attention及其变体中使用的Ring AllGather通信原语。彩色圆圈代表加速器0-7。符号t[0]~t[7]表示传输的KV块。下：为K8分解设计的多环AllGather通信原语。每一列展示了一次迭代，其中包含七个环形数据传输，用于七个KV块的分块。符号t[i, j]表示最初分配在加速器j并遵循第i个环形数据传输在所有加速器间循环的KV块。为简洁起见，省略了第2-6次迭代和第2-5个环形数据传输。

3.2.3 TASP 的 Zig-zag 块放置策略

原生块放置策略及其问题。原生 Ring Attention 采用一种简单的块放置策略：$t[i] = KV[\frac{iS}{n}, \frac{(i+1)S}{n})$（其中 S 是序列长度）。简而言之，该策略首先将所有 KV 张量均匀地分割成 n 个连续的块，并将每个块映射到一个加速器上。然而，当使用因果掩码时，这种简单的块放置策略会破坏计算负载的均衡。

因果掩码下的负载不均衡。如果将简单的块放置策略与因果掩码一起使用，在第 $i$ 次迭代中，将有 $i$ 个加速器处于空闲状态，而只有 $n-i$ 个加速器执行计算。相比之下，Megatron CP 的 Zig-Zag 块放置策略通过一个更复杂的映射，确保所有加速器在每次迭代中都执行等量的计算：

$$t[i]=KV \left[\frac{iS}{2n}, \frac{(i+1)S}{2n}\right) \cup$$

TASP的Zig-zag策略。为了实现完美负载均衡的因果注意力计算，Zig-zag Attention 的 Zig-zag 块放置策略不能直接用于 TASP。然而，通过一些必要的调整，可以推导出类似的策略。Zig-zag TASP：每个分配在加速器 $j$ 上并通过环形数据传输 $i$ 循环的 KV 块 $t[i, j]$ 被进一步分为两部分：

$$t[i,j][0] = KV\left[\frac{(2nj+i)S}{2n(n-1)}, \frac{(2nj+i+1)S}{2n(n-1)}\right]$$

$t[i, j][1] = KV[S - \frac{(2nj + i + 1)S}{2n(n-1)}, S - \frac{(2nj + i)S}{2n(n-1)})$

完美负载均衡的证明。Zig-zag TASP 完美负载均衡的证明如下：假设对于某个加速器 $r_0$，在某个迭代 $iter_0$ 中，它通过某个环形数据通路 $i_0$ 接收到最初放置在加速器 $j_0$ 的 KV 块 $t_{KV}[i_0, j_0][0]$ 和 $t_{KV}[i_0, j_0][1]$，同时 Q 块 $t_Q[i, r_0][0]$ 和 $t_Q[i, r_0][1]$ (对于所有 $i$) 保持在 $r_0$ 上。如果 $j < r_0$，那么 $t_Q[i, r_0][0]$ 和 $t_Q[i, r_0][1]$ (对于所有 $i$) 的 token 索引大于 $t_{KV}[i_0, j_0][0]$ 的 token 索引但小于 $t_{KV}[i_0, j_0][1]$ 的，使得只有 $t_{KV}[i_0, j_0][0]$ 需要被关注，这是通过 $i_0$ 传输到加速器 $r_0$ 的 KV 张量的一半。相反，如果 $j > r_0$，那么 $t_Q[i, r_0][0]$ (对于所有 $i$) 的 token 索引小于 $t_{KV}[i_0, j_0][0]$ 和 $t_{KV}[i_0, j_0][1]$ 的，而 $t_Q[i, r_0][1]$ (对于所有 $i$) 的 token 索引大于 $t_{KV}[i_0, j_0][0]$ 和 $t_{KV}[i_0, j_0][1]$ 的，使得只有 $t_Q[i, r_0][1]$ (对于所有 $i$) 需要被关注，这是查询 token 的一半。因此，对于任何加速器 $r_0$，其在任何迭代中的计算负载都是相同的，从而实现 100% 的负载均衡。

4. TASP 的实现

TASP 的实现利用了两项关键技术：（1）使用 AlltoAll 集体通信在多个环形数据通路上并发传输数据，以及（2）预计算路由表以最小化 CPU 调度开销。

4.1 预计算路由表

最小化CPU开销。为了在分布式注意力计算期间最小化CPU开销，TASP在正向传播之前预先计算所有路由表。如算法2所述的cal_in/out_mapping函数根据给定的拓扑分解方案生成每个发送/接收的源和目标映射，从而消除了在计算密集型的注意力计算过程中进行动态路由的需求。

算法 2 计算输入/输出映射
输入：循环矩阵 looping[m][n]，方向 d ∈ {+1, -1}
输出：映射矩阵 mapping[n][n]
1: 初始化 mapping 为一个用-1填充的 n x n 矩阵
2: for i = 0 to m-1 do
3:   for j = 0 to n-1 do
4:     u ← looping[i][j]
5:     v ← looping[i][(j + d) mod n]
6:     mapping[u][v] ← i
7:   end for
8: end for
9: return mapping

4.2 通过 AlltoAll 实现并发环形数据传输

利用底层优化。与Ring Attention使用批处理的sendRecv集体操作不同，TASP可以利用单个AlltoAll集体操作来同时在所有活动的环形数据通路上传输KV块（此优化仅适用于$K_{m \times u}$-分解方案）。这种方法充分利用了通信集体库（NCCL和RCCL）中为AlltoAll集体操作实现的底层优化。

TASP正向传播算法。TASP的整个正向传播过程如算法3所示，其中每个加速器根据预先计算的路由表准备发送缓冲区，并在每次迭代中从所有其他加速器接收相应的KV块。

算法 3 TASP 正向传播
输入：q, k, v, process_group, world_size
输出：out, lse
1: // 初始化路由表: looping ← gen_hamilton_circle(world_size)
2: out_mapping ← cal_out_mapping(looping)
3: in_mapping ← cal_in_mapping(looping)
4: para_size ← world_size - 1
5: // 堆叠KV张量: this_kv ← [stack(k[i], v[i]) for i in {0, para_size}]
6: // 初始化接收缓冲区: next_kv ← [empty_like(this_kv[0]) for i in {0, para_size}]
7: for step = 0 to world_size - 1 do
8:   // 提取 K, V: k ← [this_kv[i][0]], v ← [this_kv[i][1]]
9:   if step < world_size - 1 then
10:    // 使用路由表准备发送/接收缓冲区:
11:    send_buf ← [this_kv[out_mapping[i]]]
12:    recv_buf ← [next_kv[in_mapping[i]]]
13:    All-to-All(send_buf, recv_buf)
14:  end if
15:  // 拼接KV块: kcat ← concat(k, dim=seqlen), vcat ← concat(v, dim=seqlen)
16:  // 计算注意力: output, lse_curr ← flash_attn(q, kcat, vcat)
17:  out, lse ← update_out_lse(out, lse, output, lse_curr)
18:  if step < world_size - 1 then
19:    // 等待通信完成
20:    this_kv ← next_kv, next_kv ← this_kv
21:  end if
22: end for
23: return out, lse

A4 实验环境

5.1.1 测试平台

我们在四个不同的硬件平台上评估TASP：
1. H100单节点：一个配备8块NVIDIA H100 GPU的单节点，通过NVSwitch4互连，拥有3.6TB/s的对剖带宽。
2. MI300X单节点：一个配备8块AMD MI300X GPU的单节点，采用全网状（full-mesh）拓扑，通过AMD Infinity Fabric互连，聚合带宽为896 GB/s。
3. H100双节点服务器：一个两节点的服务器，每个节点配备8块NVIDIA H100 GPU。两个节点通过16个NVIDIA Mellanox Infiniband NIC（每个节点8个NIC）互连，每个节点的聚合带宽为800 GB/s（双工）。
4. H100四节点服务器：一个四节点的服务器，其节点内和节点间配置与(3)相同。

5.1.2 基准测试

我们在三组约400个不同输入大小的测试用例上评估TASP的性能（总共1287个测试用例）。具体参数如下：
* 批量大小 (Batch size)：1到128。
* 序列长度 (Sequence length)：3k到1M。
* 注意力头数 (Attention heads)：{4, 12, 20}，头维度为64。
* 精度：BF16。

对于不同的硬件平台，我们设置了不同的最大序列长度：所有单节点平台的最大序列长度为430K；双节点和四节点平台的最大序列长度分别为490K和1M。

5.1.3 基线方法

我们使用 Ring Attention【【10，Ring attention with blockwise transformers for near-infinite context，2023，arXiv】】作为非因果掩码情况下的基线方法，使用 Zig-zag Ring Attention【【8，Megatron Context Parallelism，2023，NVIDIA】】作为因果掩码情况下的基线方法。
* Ring Attention：专为序列并行设计的基础分布式注意力实现。
* Zig-zag Ring Attention：为带因果掩码的注意力修改的分布式注意力实现，旨在解决Ring Attention中的负载不均衡问题。

为了公平比较，所有方法（包括TASP）都在PyTorch【【29，Pytorch: An imperative style, high-performance deep learning library，2019，Advances in neural information processing systems】】中实现，并使用相同的Flash Attention【【14，Flashattention: Fast and memoryefficient exact attention with io-awareness，2022，Advances in neural information processing systems】】实现。我们采用了Ring Attention和Zig-zag Ring Attention的开源代码【【30，ring-flash-attention，2024，https://github.com/zhuzilin/ring-flash-attention】】进行基线比较 。

A4 实验结果

5.2 单节点评估

我们在代表交换机拓扑的8卡NVIDIA H100服务器（H100-1）和代表全网状拓扑的8卡AMD MI300X服务器（MI300X-1）上评估了TASP的单节点性能。

核心发现：如图5所示，当通信数据量足以饱和加速器互连带宽时（红色数据点），TASP在总时间（tall）上取得了显著的加速。
* 在 H100-1 上，对于全掩码和因果掩码情况，平均加速比分别为 1.05倍 和 1.08倍，最大加速比达到 2.31倍。
* 在 MI300X-1 上，对于全掩码和因果掩码情况，平均加速比分别为 1.57倍 和 1.68倍，最大加速比达到 3.58倍。
* 在通信时间（tcomm）方面，TASP在MI300X-1上最高实现了4.72倍的加速，在H100-1上最高实现了2.87倍的加速。

影响因素分析：
* CCRB的影响：基线方法的计算通信比（CCRB）越低，表明其越受通信限制，TASP的加速效果越明显。
* 掩码模式的影响：因果掩码的注意力计算量理论上是全掩码的一半，导致CCRB更低，因此TASP在因果掩码测试用例上取得了更高的加速比。
* 硬件平台的影响：TASP在MI300X-1上的性能优化远超H100-1，这不仅因为MI300X-1的节点内通信带宽较低使其在多数情况下更受通信限制，也与我们对全网状拓扑与交换机拓扑的理论分析相符。

5.3 多节点评估

我们在H100 GPU组成的双节点（H100-2）和四节点（H100-4）服务器上进行了多节点评估。我们比较了两种不同的拓扑分解方案：$K_{m \times u}$-分解 和 $(m-K_m-m)_u$-分解。

核心发现：如图6所示，在通信数据量充足的情况下（红色数据点）：
* 在 H100-2 上，使用 $K_{8 \times 2}$-分解的平均加速比为 1.43倍，而使用 $(8-K_8-8)_2$-分解的平均加速比为 1.20倍。
* 在 H100-4 上，使用 $K_{8 \times 4}$-分解的平均加速比为 1.27倍，而使用 $(8-K_8-8)_4$-分解的平均加速比为 1.20倍。

分解方案分析：
* 理论上，$(m-K_m-m)_u$-分解因其拓扑感知的通信链路利用而优于$K_{m \times u}$-分解。然而，实验在双节点上得出了相反的结果。这是因为$K_{m \times u}$-分解方案可以直接调用NCCL/RCCL中高度优化的AlltoAll API，而$(m-K_m-m)_u$-分解方案需要使用批处理的SendRecv进行自定义通信，缺乏底层优化导致效率下降。
* 可扩展性：$(m-K_m-m)_u$-分解方案的可扩展性要好得多。当从H100-2扩展到H100-4时，其加速比几乎保持不变，而$K_{m \times u}$-分解的加速比则显著下降。

A7 补充细节

5.4 开销分析

通信开销。TASP采用复杂的通信原语（AlltoAll集体操作）。基于$(m-K_m-m)_u$分解的TASP所采用的批处理SendRecv也比原生Ring Attention中采用的批处理SendRecv更复杂（m个并发SendRecv vs 2个）。虽然这种方法能有效利用所有正交环形数据通路，但它也带来了更高的运行时开销。因此，在通信量较小的场景下（例如批量大小和序列长度较小），这种开销会降低整体性能，如图5和图6中的蓝色数据点所示。然而，在其他场景中（约占所有测试用例的3/4），这种开销与通信延迟相比可以忽略不计。

计算开销。TASP的额外计算开销主要源于其对KV Cache的细粒度划分。例如，在一个基于$K_8$分解的TASP方案中，KV Cache被分成56个块（每个加速器7个块），而原生Ring Attention只需要8个块。尽管这些KV块的路由和映射可以预先计算，但增加的块数与原生Ring Attention相比，引入了稍高的CPU开销。幸运的是，这种CPU开销不影响整体执行时间，因为TASP的核函数启动延迟可以被GPU计算完全重叠——这得益于没有GPU-CPU同步。

GPU计算开销。除了CPU开销，我们当前TASP的实现与原生Ring Attention相比也引入了轻微的GPU计算开销，因为所有KV Cache块必须在每个rank上拼接起来，然后才能调用flash-attn核函数进行注意力计算。经验测量表明，这种额外开销是适度的，导致总计算时间(tcomp)仅减慢1%–5%。尽管如此，我们注意到这种开销可以通过核函数优化技术进一步最小化。

7. 局限性与讨论

与CCR的强相关性。我们提出的算法在长序列多加速器推理中的有效性与计算通信比（CCR）密切相关：当该比率较低时，我们以通信为中心的优化针对主要瓶颈（通信），从而带来显著的性能提升；然而，在计算密集型的注意力场景中——特别是当CCR超过1.5时——这些增益会消失。

算法有效性的解释。这并不表明算法本身存在缺陷，而是说明这种分布式注意力方法的整体延迟是由多种因素共同决定的。当其他因素的影响超过通信开销时，面向通信的优化的有效性自然会减弱，这也阐明了该算法在与计算导向的优化（如稀疏注意力）协同工作时的优势。

A5 结论

本文提出了TASP，一种新颖的序列并行（SP）方法，它利用完全有向图的哈密顿分解来优化长上下文LLM的注意力预填充（prefill）阶段。为了克服现有方法的局限性，TASP采用了两步优化：拓扑分解和原语分解，以充分利用AI加速器原生的互连拓扑。通过实验评估，TASP在单节点和多节点拓扑上均实现了比Ring Attention（及Zig-zag Ring Attention）更高的通信效率，突显了在设计高性能SP方法时，拓扑感知优化的至关重要性。

方法细节中的引用汇总

1. [5] Sam Ade Jacobs, et al. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023.

   • 引用章节：2.1 序列并行
   • 引用描述：作为序列并行方法的两大类别之一被提及，并详细阐述了其通过AlltoAll操作按KV头维度分区的方法及其可扩展性局限。

2. [6] William Brandon, et al. Striped attention: Faster ring attention for causal transformers. arXiv preprint arXiv:2311.09431, 2023.

   • 引用章节：2.1 序列并行，3.2.1 问题定义
   • 引用描述：作为Ring Attention的优化变体被提及，它采用条带化分区策略来解决因果掩码下的负载不均衡问题。

3. [8] NVIDIA. Megatron Context Parallelism, 2023.

   • 引用章节：2.1 序列并行，3.2.1 问题定义，3.2.3 TASP的Zig-zag块放置策略
   • 引用描述：作为Ring Attention的另一种负载均衡优化方法（也称为Zig-zag Ring Attention）被提及，并详细介绍了其为实现完美负载均衡而设计的轴对称分区方案。

4. [9] Rajeev Thakur, et al. Optimization of collective communication operations in mpich. The International Journal of High Performance Computing Applications, 19(1):49–66, 2005.

   • 引用章节：1 Introduction
   • 引用描述：指出Ring Attention中使用的Ring AllGather原语是MPICH中AllGather原语最古老的实现之一。

5. [10] Hao Liu, et al. Ring attention with blockwise transformers for near-infinite context, 2023.

   • 引用章节：2.1 序列并行
   • 引用描述：作为主流的序列并行方法被详细介绍，说明其通过Ring AllGather原语在加速器间轮转KV块。

6. [11-13]

   • 引用章节：2.1 序列并行
   • 引用描述：用于说明当前注意力机制（如GQA, MQA, MLA）的趋势是减少KV头数量，这限制了Ulysses方法的可扩展性。

7. [14] Tri Dao, et al. Flashattention: Fast and memoryefficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022.

   • 引用章节：2.1 序列并行
   • 引用描述：被提及为Ring Attention中用于增量计算注意力输出的高效实现。

8. [16-20]

   • 引用章节：2.2 硬件平台的互联性
   • 引用描述：用于举例说明现代AI加速器（如Ascend 910B, MI300X, Gaudi 3, NVIDIA A100）普遍采用全连接拓扑（全网状或基于交换机）。

9. [21-25]

   • 引用章节：3.1.2 单节点拓扑分解
   • 引用描述：引用了将完全有向图分解为边不相交的哈密顿循环的现有图论工作，作为TASP拓扑分解的理论基础。

10. [26] Timothy W Tillson. A hamiltonian decomposition of k2m*, 2m≥8. Journal of Combinatorial Theory, Series B, 29(1):68–74, 1980.

    • 引用章节：3.1.2 单节点拓扑分解
    • 引用描述：引用了一个在O(n)时间复杂度内分解完全图$K_n$的特定算法，该算法被用于计算图2中的分解方案。

11. [27] Alexander Ishii and Ryan Wells. The nvlink-network switch: Nvidia’s switch chip for high communicationbandwidth superpods. In 2022 IEEE Hot Chips 34 Symposium (HCS), pages 1–23, 2022.

    • 引用章节：3.1.3 多节点拓扑分解
    • 引用描述：用于描述NVIDIA DGX H100 GPU集群的硬件互连架构。

12. [28] Mohammad Shoeybi, et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.

    • 引用章节：5.1.3 基线方法 (引用了Zig-zag Ring Attention的来源，尽管文中直接称其为Megatron CP)
    • 引用描述：作为Zig-zag Ring Attention的实现来源之一被引用。

13. [29] Adam Paszke, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.

    • 引用章节：5.1.3 基线方法
    • 引用描述：说明所有方法（包括TASP和基线）都是在PyTorch中实现的。

14. [30] Zilin Zhu. ring-flash-attention. https://github. com/zhuzilin/ring-flash-attention, 2024.

    • 引用章节：5.1.3 基线方法
    • 引用描述：说明基线比较采用了此开源代码库。