文章标题：TileLink：使用以 Tile 为中心的原语生成高效的计算-通信重叠内核
作者/机构：Size Zheng, Jin Fang, Xuegui Zheng, Qi Hou, Wenlei Bao, Ningxin Zheng, Ziheng Jiang, Dongyang Wang, Jianxi Ye, Haibin Lin, Li-Wen Chang, Xin Liu (均来自 ByteDance Seed)

A1 主要贡献

本文旨在解决大型深度学习模型分布式训练中，通信开销（占总执行开销的10%-50%）限制计算效率的核心问题。提升性能最有效的方法是计算与通信的重叠（Overlapping），现有技术主要分为两种：算子分解（operator decomposition）和核融合（kernel fusion）。算子分解易于实现，但常因内核分解过细、缓存利用率低和主机干预同步开销大而导致性能不佳。核融合性能更优，但需要深厚的硬件专业知识，开发难度大且难以跟上算法的快速迭代。

为了应对这些挑战，本文提出TileLink框架，其研究目标是通过编译技术高效地开发和生成计算-通信重叠内核。

本文的主要贡献和创新点如下：
1. 提出TileLink框架：该框架包含前端和后端。
* 前端：通过以Tile为中心（tile-centric）的原语，将通信和计算的设计空间解耦。这允许通信和计算部分可以采用不同的优化策略和分块（tiling）方法，并通过定制的屏障控制来维持生产者-消费者依赖关系，从而在不直接编写底层汇编代码的情况下，实现通信与计算核的自动化融合。
* 后端：将前端的tile-centric原语编译成底层硬件指令，并集成通信与计算内核。通过以Tile为中心的映射策略（包括形状映射、rank映射和通道映射，支持静态和动态两种方式）来确保正确的数据交换和屏障操作。
2. 通用性和灵活性：通过TileLink实现了多种重叠工作负载，包括自注意力（self-attention）、多层感知机（MLP）和混合专家模型（MoE），展示了其灵活性和通用性。
3. 高性能实现：在8卡H800 GPU上的实验表明，与非重叠基线相比，TileLink实现了1.17倍至20.76倍的加速。其性能与当前最先进的重叠库（如FLUX和RingAttention）相当或更优。在端到端模型评估中，与PyTorch相比，TileLink在8卡H800上平均加速1.32倍，在16卡H800（双节点）上平均加速1.29倍。

图1 层内并行FFN示例。
图1 层内并行FFN示例。

A3 背景知识/关键Observation/设计原则

2.1 通信与层内并行

  • 以算子为中心的通信原语：集合通信在大型模型的并行执行中非常常用。现有的库如【索引26，Nvidia collective communications library，2024】和框架如【索引1，Tensorflow: A system for large-scale machine learning，2016，OSDI】、【索引2，Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation，2024，ASPLOS】使用算子级原语来实现常见的通信模式，例如AllReduce、ReduceScatter、AllGather和All2All。这些原语为了遵循SPMD（单程序，多数据）编程模型并与其他算子无缝集成，需要在数据传输前后执行系统同步。然而，这种粗粒度的同步会导致计算单元在通信期间处于空闲状态，从而降低了计算效率。我们将这些通信原语称为以算子为中心的原语。

  • 使用以算子为中心的原语实现层内并行：对于大型模型（即基于Transformer的模型），层内并行主要应用于两个组件：注意力部分和FFN（前馈网络）部分，后者由MLP（多层感知机）层或MoE（混合专家）层组成。对于注意力部分，上下文（键和值）被分片到各个设备上。在计算之前，这些上下文分片被收集起来形成一个完整的上下文以进行自注意力计算。这种并行算法被称为序列并行【索引19，Ring attention with blockwise transformers for near-infinite context，2023】、【索引23，Efficient large-scale language model training on GPU clusters using megatron-lm，2021，SC】。

  • FFN部分的层内并行：在FFN部分，MLP或MoE中两层的权重被分片到各个设备上。首先，从不同的rank收集输入数据，然后使用相应的权重分片进行本地计算。最后，部分结果被归约并分散到相应的rank。这种算法在先前的工作【索引15，Breaking the computation and communication abstraction barrier in distributed machine learning workloads，2022，ASPLOS】、【索引23，Efficient large-scale language model training on GPU clusters using megatron-lm，2021，SC】、【索引41，Overlap communication with dependent computation via decomposition in large deep learning models，2023，ASPLOS】中常用，如图1所示，并被称为张量并行FFN。使用现有的通信库和框架，张量并行FFN表示为AllGather + GEMM（或GroupGEMM），然后是GEMM（或GroupGEMM）+ ReduceScatter。

2.2 通信和计算的重叠

  • 重叠技术概述：通信与计算的重叠已在先前的研究【索引5，FLUX: fast software-based communication overlap on gpus through kernel fusion，2024】、【索引6，Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning，2024，ASPLOS】、【索引15，Breaking the computation and communication abstraction barrier in distributed machine learning workloads，2022，ASPLOS】、【索引41，Overlap communication with dependent computation via decomposition in large deep learning models，2023，ASPLOS】中被广泛探索。Centauri【索引6，Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning，2024，ASPLOS】引入了一个包含模型级、层级和操作级重叠的全面三级设计空间。层内并行是这些重叠层次的基础要素。TileLink专注于层内重叠；将TileLink的技术推广到层间或模型级重叠是可行的，但超出了本文的范围。对于层内并行，实现重叠主要有两种方式：算子分解和核融合。如表1所示，我们总结了代表性研究和TileLink的特点。

  • 算子分解（Operator Decomposition）：这种方法将原始的通信和计算算子分割成更小的、细粒度的单元。这些更小的算子能够实现更精确的同步控制，允许通信算子与没有数据依赖的计算算子并行执行。算子分解的优势在于其实现简单，且与现有库和框架兼容。然而，使用更小的算子可能导致效率低下，包括L2缓存利用率低【索引39，Triton: an intermediate language and compiler for tiled neural network computations，2019，MAPL@PLDI】和资源量化效率低下【索引30，Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the GPU，2023，PPoPP】。此外，内核之间的同步需要主机干预，引入了不可忽略的开销。采用算子分解的代表性工作包括Dist-Einsum【索引41，Overlap communication with dependent computation via decomposition in large deep learning models，2023，ASPLOS】、Asynchronous Tensor Parallel PyTorch【索引2，Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation，2024，ASPLOS】和Centauri【索引6，Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning，2024，ASPLOS】。

  • 核融合（Kernel Fusion）：这种方法融合了通信和计算内核。通常，融合后的内核将一部分处理核心分配给通信任务，其余核心分配给计算任务。分配给不同任务的核心使用设备上屏障（on-device barriers）来通信执行状态。融合内核无需主机干预同步，提高了缓存利用率，并减轻了资源量化效率低下的问题，可能比算子分解方法获得更好的性能。然而，在现代加速器（如GPU）上开发融合内核存在挑战。一方面，对屏障和与硬件相关的通信指令的底层控制要求高水平的专业知识。另一方面，不当的融合设计可能因通信和计算核心之间的资源冲突而导致性能下降。因此，只有少数高度优化的库【索引5，FLUX: fast software-based communication overlap on gpus through kernel fusion，2024】、【索引31，Optimizing distributed ml communication with fused computation-collective operations，2023】或领域特定编译器【索引15，Breaking the computation and communication abstraction barrier in distributed machine learning workloads，2022，ASPLOS】支持核融合方法。

2.3 代码生成编译器

  • 代码生成编译器的发展：随着代码生成编译器【索引7，TVM: an automated end-to-end optimizing compiler for deep learning，2018，OSDI】、【索引33，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，PLDI】、【索引40，Triton: an intermediate language and compiler for tiled neural network computations，2019，MAPL@PLDI】的快速发展，为注意力或FFN生成高性能代码已变得可行。尽管之前的重叠编译器如CoCoNet【索引15，Breaking the computation and communication abstraction barrier in distributed machine learning workloads，2022，ASPLOS】和Dist-Einsum【索引41，Overlap communication with dependent computation via decomposition in large deep learning models，2023，ASPLOS】可以生成重叠内核，但它们受限于固定的重叠模式，在算子级别缺乏编程灵活性。相比之下，TileLink使用以tile为中心的原语，并为各种工作负载实现了高效的编译。

2.4 激励性示例

  • 示例介绍与性能对比：为了说明TileLink的优势，我们使用一个张量并行的MLP层作为激励示例。该MLP层的输入形状详见表2，对应于LLaMA-7B模型中使用的配置。此MLP层实现为AllGather + GEMM (AG + GEMM)，后接GEMM + ReduceScatter (GEMM + RS)，如图1所示。我们在表2中比较了这两部分不同技术的性能。Non-Overlap使用cuBLAS【索引24，cuBLAS，2022】和NCCL【索引26，Nvidia collective communications library，2024】，无重叠。Decomposition使用算子分解技术，性能数据来自Async-TP PyTorch【索引18，Torchtitan: One-stop pytorch native solution for production ready llm pre-training，2024】。Fusion指核融合技术，使用开源库FLUX【索引5，FLUX: fast software-based communication overlap on gpus through kernel fusion，2024】进行测量。

  • 性能与编程效率分析：一方面，我们比较了不同技术实现的性能。如表2所示，分解技术性能最低，而融合技术在AG + GEMM上取得了最好的结果。TileLink在GEMM + RS上取得了最佳性能，并且在AG + GEMM上非常接近FLUX（约99%）。这些发现表明，TileLink能够提供与先前方法相当或更好的性能。另一方面，我们比较了FLUX和TileLink所需的代码行数。FLUX涉及大约2000行CUDA代码，而TileLink仅用大约200行Python代码就实现了相似的性能，编程效率提高了约10倍。这个激励示例凸显了TileLink的显著优势。

表1 TileLink与先前工作的特性比较。
表1 TileLink与先前工作的特性比较。

表2 激励性示例。
表2 激励性示例。

表3 TileLink中的以Tile为中心的原语
表3 TileLink中的以Tile为中心的原语

A2 方法细节

3.1 解耦的设计空间

  • 采用解耦设计的原因：设计计算-通信融合内核有两种方式。一种是紧密耦合两部分的优化选择，包括tile大小、tile顺序和资源映射；另一种是解耦计算和通信内核的设计。TileLink采用后者，因为解耦的设计空间在内核设计中提供了更大的灵活性，并可能带来更好的性能。

  • 设计空间的三个子空间：我们将解耦的设计空间分为三个子空间：tile大小、tile顺序和资源映射。对于每个子空间，通信和计算组件可以做出独立的选择以优化其性能。

    • Tile大小子空间：通信和计算组件可以选择不同的tile大小。例如，如图2a所示，通信部分一次传输一个128 × 128的tile，而计算部分一次消耗一个大小为128 × 256的tile。这种tile大小的差异有助于每个组件通过与其使用的处理核心数量对齐来达到最佳性能。例如，对于一个AllGather + GEMM问题，张量大小为M × N × K，其中AllGather部分将维度M、K绑定到处理核心，GEMM部分绑定维度M、N。如果通信组件使用更多的核心，较小的tile大小将更有利，因为所有核心资源都可以被充分利用；反之，如果它使用较少的核心，较大的tile大小将更有效。
    • Tile顺序子空间：通信组件可能使用与计算组件不同的tile顺序。例如，通信可以采用各种数据传输顺序，如环形顺序、全网状全到全顺序或其他模式，而计算组件可以从任何rank开始处理数据tile。tile顺序的选择存在权衡。如果计算组件等待来自多个rank的数据tile，它在处理更大数据块时可能会获得更好的缓存效率；然而，这种方法可能导致更长的等待时间。相反，如果计算组件只等待来自单个rank的数据tile，它可以更早地开始计算，但这可能导致较低的整体计算效率。图2b展示了一个例子，其中通信使用环形顺序，而计算在每次计算迭代中等待来自两个rank的数据。
    • 资源绑定子空间：通信和计算组件可以映射到不同的单元或相同的单元，如图2c所示。如果通信组件利用复制引擎（DMA），它避免了与计算组件的资源冲突。然而，这种方法涉及主机干扰，会引入额外开销。另一方面，如果通信组件使用计算核心进行数据复制，可能会导致与计算组件的资源冲突，但消除了主机开销。这种策略特别适用于计算组件无法完全利用所有可用处理核心的场景。

图2 通信和计算的三个设计子空间示例。
图2 通信和计算的三个设计子空间示例。

图3 以Tile为中心的原语支持不同的信号控制和数据传输方向。
图3 以Tile为中心的原语支持不同的信号控制和数据传输方向。

3.2 以Tile为中心的原语

  • 引入原语的动机：解耦通信和计算的设计空间引入了同步挑战。由于两个组件使用不同的tile大小、tile顺序和资源映射，同步它们需要复杂的底层通信指令编程。例如，在GPU上，需要ld.global.acquire和red.release等指令。然而，这些指令的编程模型与代码生成编译器【索引7，TVM: an automated end-to-end optimizing compiler for deep learning，2018，OSDI】、【索引40，Triton: an intermediate language and compiler for tiled neural network computations，2019，MAPL@PLDI】的编程模型不一致，因为现有编译器缺乏对内存一致性模型的支持。

  • 原语的分类和特点：为了解决这个问题，TileLink提供了一套以tile为中心的原语。这些原语引入了内存一致性语义，并遵循编译器中使用的tile级抽象，这使它们与先前框架【索引1，Tensorflow: A system for large-scale machine learning，2016，OSDI】、【索引2，Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation，2024，ASPLOS】和库【索引26，Nvidia collective communications library，2024】提供的以算子为中心的原语有所区别。TileLink的原语总结在表3中。它们分为两组：信号原语和数据原语。每组都包含设备端原语和主机端原语。

3.2.1 信号原语

  • 信号原语的功能与分类：信号原语设计用于管理通信和计算之间的屏障。它们包括producer(peer)_tile_notify、consumer(peer)_tile_wait和rank_notify(wait)。对于设备端原语，producer_tile_notify和consumer_tile_wait原语应用于生产者-消费者关系，例如AllGather和GEMM的tile之间。peer_tile_notify和peer_tile_wait原语主要用于不同rank上同一算子的tile，使用户能够构建各种tile顺序。对于主机端原语，rank_notify(wait)原语用于管理复制引擎和计算核心之间的屏障。当通信映射到复制引擎时，这些原语有助于控制通信和计算之间的tile顺序。图3a展示了通信和计算部分之间的信号控制。

  • Notify原语的参数：Notify原语需要mode参数或rank参数来明确通知哪些远程rank。TileLink为mode参数提供了两种选择：p2p和broadcast。p2p意味着只有一个目标rank将被通知，该rank由给定tile_id在全局张量视图中的偏移量计算得出；broadcast意味着所有rank都将被通知。

  • 内存一致性：在并行执行中，不同进程/线程执行的内存操作对其他进程/线程的可见顺序可能不一致。内存一致性模型指定了约束条件，以防止在跨进程/线程观察到的操作顺序中出现矛盾。信号原语提供严格的内存一致性语义。notify原语带有release语义，确保在producer(peer)_tile_notify和rank_notify之前发生的任何内存访问不能在这些notify原语之后执行。相反，wait原语带有acquire语义，确保在consumer(peer)_tile_wait和rank_wait之后发生的任何内存访问不能在这些wait原语之前执行。在后端编译时也必须考虑这种严格的内存一致性，这将在后面讨论。

3.2.2 数据原语

  • 数据原语的功能与模式：数据原语用于数据传输，包括tile_push(pull)_data和rank_copy_data原语。这些原语控制传输数据的资源映射和tile大小。设备端的tile_push(pull)_data原语将通信映射到处理核心，而主机端的rank_copy_data原语将通信映射到复制引擎。数据传输有两种模式——pull和push——每种模式适用于不同的同步方法。在pull模式下，生产者从所有其他rank读取数据，并使用本地屏障通知其消费者。相比之下，push模式允许生产者将本地数据写入所有其他rank，同时通知其远程消费者数据已到达。图3b说明了这两种模式之间的差异。pull和push模式的选择可能会影响性能（如FLUX【索引5，FLUX: fast software-based communication overlap on gpus through kernel fusion，2024】中所指出的），具体取决于数据形状、分块策略和可用硬件资源等因素。值得注意的是，rank_copy_data原语通过点对点复制支持两种模式，数据传输方向由源和目标指针的顺序指示。

4.1 以Tile为中心的映射

  • 映射策略概述：TileLink使用以tile为中心的映射方法将前端原语编译成底层代码。以tile为中心的映射由三个组件组成：形状映射、rank映射和通道映射。形状映射将每个tile_id与一个特定的张量形状切片关联起来。Rank映射将每个tile_id链接到一个设备rank。通道映射将每个tile_id分配给一个通信屏障。我们分别用$f_S, f_R, f_C$来表示这三种映射。根据工作负载类型的不同，应使用不同的映射函数。我们将不同的映射分为两组：静态映射和动态映射。

  • 静态映射：静态映射指的是可以在编译时确定的映射。当数据分片策略固定时，通常使用静态映射，例如张量并行的MLP和序列并行的自注意力。我们使用仿射运算来处理静态映射（$f_S, f_R, f_C$是仿射的）。例如，对于在R个rank上、每个rank有C个通道（每个rank对应C个屏障）的AllGather（pull模式）+ GEMM（问题大小为M × N × K），生产者AllGather使用tile大小$T_{mp} \times T_{np}$，并且输入张量沿M维度分片。给定生产者tile tile_idp，其形状范围、源rank和通道可以计算如下：

    公式1
    公式1

    同样，我们可以计算从消费者tile tile_idc到形状范围、rank和通道的映射。

  • 动态映射：动态映射指的是在运行时计算的映射，这对于具有动态数据分片要求的工作负载至关重要。例如，在MoE数据分片策略中，动态路由决定了数据分布，每个tile可能需要来自任何其他rank的token。在编译时无法确定从哪些rank收集数据或在哪个通道等待屏障。因此，映射必须在运行时计算。为了支持动态映射，TileLink将这些映射转换为查找表，其值可以在运行时填充，而对这些查找表的访问操作则在编译时确定。形式上，动态映射为：

    公式2
    公式2

    其中$f_{S\_low}, f_{S\_high}, f_R$和$f_C$是查找表，它们的值将在运行时由其他动态逻辑（例如，动态路由）填充。

4.2 内存一致性的编译

  • 编译挑战与解决方案：在后端编译中，具有内存一致性语义的前端原语被编译成相应的设备指令，如ld.global.acquire和red.release。然而，直接翻译这些原语不足以保证内存一致性。对于大多数计算内核，会应用多级流水线来增强加载-计算平衡并提高整体性能。将原始程序编译成多级版本需要算子重排序，在此期间，一些内存访问操作可能会意外地被重排到TileLink原语之前或之后。为了解决这个问题，TileLink在其原语与其后的加载/存储操作之间强制执行严格的数据依赖关系，以便其原语可以被流水线传递正确地重排序和展开。

4.3 其他编译优化

  • 单设备优化策略：除了上述技术，TileLink还利用了单设备优化的策略来实现高性能，这在先前的工作【索引7，TVM: an automated end-to-end optimizing compiler for deep learning，2018，OSDI】、【索引40，Triton: an intermediate language and compiler for tiled neural network computations，2019，MAPL@PLDI】中已有深入研究。这些优化主要包括两个方面：内存优化和流水线优化。内存优化涉及为计算自动分配片上寄存器缓冲和共享内存缓冲。对全局缓冲的数据访问被合并，对共享内存的访问模式被转换以避免bank冲突。流水线优化涉及重新排列数据加载/存储操作和计算，以形成多级流水线。本地数据复制被映射到专用的异步引擎，例如GPU的张量内存加速器（TMA）。计算被映射到高性能单元，例如GPU的张量核心单元。

图4 使用TileLink的GEMM+RS重叠内核。
图4 使用TileLink的GEMM+RS重叠内核。

  • 引言：为了展示TileLink的灵活性和通用性，我们介绍了如何为GEMM + 环形ReduceScatter、AllGather + MoE以及AllGather KV + 自注意力设计重叠内核。这三个例子具有代表性，因为它们利用了不同的tile顺序（环形和全对全）、不同的映射（静态和动态）以及不同的硬件资源（设备和主机）。

  • GEMM + 环形ReduceScatter：图4展示了GEMM + 环形ReduceScatter内核的伪代码。计算和通信都使用SM（流式多处理器），在这个例子中我们为通信分配了20个SM（见第1行）。生产者GEMM将部分输出存储在本地张量中，并使用producer_tile_notify通知其消费者（第9行）。消费者ReduceScatter在第16行等待其生产者。一旦来自生产者的数据准备就绪，消费者内核执行本地归约（第20行）并将部分结果传递给其前一个rank（第24行）。对等rank之间的信号控制分别在第19行和第26行使用peer_tile_wait和peer_tile_notify原语进行管理。这个例子使用静态映射，并演示了如何在两个方向上编程通信：生产者-消费者和对等体。

  • AllGather + MoE：图5展示了AllGather + MoE的伪代码。同样，计算和通信都使用SM，我们为通信分配了20个SM（见第1行）。请注意，MoE需要动态路由（输入中的topk_ids）来为每个token选择专家，这需要动态映射。我们使用table来表示形状映射、rank映射和通道映射的查找表。所有涉及的原语都应将table作为参数，以便TileLink可以使用动态映射生成正确的代码。此外，load原语也需要table，因为它使用table中的形状映射来收集正确的token（第11行）和相应token的topk_ids（第12行），这些是当前tile所需要的。

  • AllGather KV + 自注意力：图6展示了AllGather KV + 自注意力（序列并行）的伪代码。在这个例子中，通信使用复制引擎。我们使用主机原语来触发复制引擎。通信和计算在两个不同的流上运行。通信使用rank_copy_data原语完成，通信部分的tile大小简单地将KVCache序列长度（S）除以总rank数（WORLD_SIZE）。对于计算部分，tile大小是不同的。以tile为中心的映射被用来保证通信和计算部分之间正确的屏障操作。

  • 总结：这些例子表明，得益于我们以tile为中心的原语和映射，TileLink在重叠内核设计方面非常灵活，并减少了编程工作量。

图5 使用TileLink的AG + MoE重叠内核。
图5 使用TileLink的AG + MoE重叠内核。

图6 AG KV + self-attention overlapping Kernel.
图6 AG KV + self-attention overlapping Kernel.

6. 实现

  • 实现细节与编译流程：TileLink是基于Triton【索引40，Triton: an intermediate language and compiler for tiled neural network computations，2019，MAPL@PLDI】用Python实现的。我们通过在Python级别实现以tile为中心的原语来扩展Triton的语言特性，而以tile为中心的映射机制则通过Python抽象语法树（AST）转换实现。当前的实现可以轻松适应其他编译器框架，如TVM【索引7，TVM: an automated end-to-end optimizing compiler for deep learning，2018，OSDI】和MLIR【索引17，MLIR: A compiler infrastructure for the end of moore’s law，2020】。

  • 编译与运行时：如图7所示，编译器接收一个纯Python程序作为输入，该程序结合了TileLink的原语和Triton的原生原语。提供了一个特殊的参数BlockChannel，作为计算和通信的以tile为中心的映射上下文。BlockChannel参数封装了分布式映射元数据，包括当前进程rank、总world size、同步屏障配置以及生产者/消费者块关系。Python程序被解析成一个AST，并翻译成Triton IR。在翻译过程中，BlockChannel参数被分解，以使用嵌入的元数据构建以tile为中心的映射。TileLink的原语被转换成Triton的ElementwiseInlineAsmOp。然后，Triton IR被降低到Triton GPU IR和一个由TileLink引入的新的分布式IR（Distributed IR）。这个分布式IR用于将通过ElementwiseInlineAsmOp表达的特殊指令翻译成LLVM IR，后者再被编译成用于NVIDIA GPU的PTX。通过将LLVM IR翻译成目标特定的底层汇编，可以支持其他后端架构。在运行时，使用NVSHMEM【索引27，NVSHMEM，2025】来初始化分布式执行环境和分配共享内存。生成的代码在所有进程上启动，以执行并发的计算和通信，完成后再进行适当的共享内存释放。

图7 TileLink的编译和运行时。
图7 TileLink的编译和运行时。

表4 基准测试形状。S是序列长度，H是隐藏维度长度，I是中间大小，E是专家数量。
表4 基准测试形状。S是序列长度，H是隐藏维度长度，I是中间大小，E是专家数量。

A4 实验环境

  • 基准测试：实验使用了三个基准：MLP层、MoE层和自注意力层。
  • 输入配置：输入形状如论文表4所示，这些配置来源于真实工作负载，如LLaMA【索引10，The llama 3 herd of models，2024】、Gemma【索引36，Gemma 2: Improving open language models at a practical size，2024】和Qwen【索引42，Qwen2 technical report，2024】。
  • 硬件配置：
    • GPU：单节点8xH800 GPU集群，多节点评估使用2个节点共16xH800 GPU。
  • 软件配置与基线：
    • TileLink实现：基于Triton，使用Python实现。
    • 基线方法：
      • 非重叠：cuBLAS+NCCL。
      • 算子分解法：Async-TP PyTorch (由于Centauri【索引6】和Dist-Einsum【索引41】未公开)。
      • 核融合法：FLUX【索引5】(由于CoCoNet【索引15】代码已弃用)。
    • 所有基线都使用了一致的并行配置。

A4 实验结果

单层性能

  • MLP层：

    • 实验内容：分别评估了MLP的两个主要部分（AG+GEMM和GEMM+RS）以及完整的MLP层。
    • 实验结果（图8）：
      • AG + GEMM：TileLink实现了1.27倍加速（vs cuBLAS+NCCL），达到了FLUX性能的94.5%。Async-TP PyTorch由于GEMM分解过小和主机同步开销大而未能提速。
      • GEMM + RS：TileLink表现最佳，相比cuBLAS+NCCL、Async-TP PyTorch和FLUX分别取得了1.25倍、2.22倍和1.28倍的加速。这得益于其解耦的设计空间，使得GEMM和ReduceScatter可以各自优化。
      • 完整MLP：TileLink性能与FLUX相当（101.4%），比cuBLAS+NCCL快1.24倍。
    • 结论：TileLink能以更少的代码量（约200行Python vs FLUX的2000行CUDA）达到与SOTA融合库相当的性能。

  • MoE层：

    • 实验内容：评估了MoE层的两个复杂部分（AG+Gather+GroupGEMM 和 GroupGEMM+Scatter+Topk Reduce+RS）。MoE层需要动态映射，现有库如FLUX和Async-TP PyTorch不支持。
    • 实验结果（图9）：
      • 第一部分：TileLink通过融合gather-scatter并重叠通信与计算，性能平均比vLLM高1.51倍，比cuBLAS基线快9.82倍以上。
      • 第二部分：TileLink平均比vLLM快1.31倍，比CUTLASS+NCCL快10.56倍。TileLink成功重叠了三个内核，展示了其原语的通用性。
      • 完整MoE：TileLink平均比vLLM快1.14倍，最大相比cuBLAS+NCCL加速20.76倍。
    • 结论：TileLink的灵活性和动态映射能力使其能够有效优化复杂的MoE层，性能远超现有实现。

  • 自注意力层：

    • 实验内容：使用TileLink实现Flash-Attention，并重叠AllGather和Flash-Attention的计算，测试了16k到128k的不同序列长度。
    • 实验结果（图10）：
      • 在所有序列长度上，TileLink都表现出一致的加速效果。平均比PyTorch（Torch）快5.04倍，比RingAttention（RingAttn）快1.97倍。
      • 通过计算重叠率（Overlap Ratio），结果显示TileLink平均能有效隐藏43.9%的通信开销。
    • 结论：TileLink能有效加速长序列自注意力计算，并显著隐藏通信延迟。

端到端评估

  • 实验内容：将TileLink集成到PyTorch中，在H800集群上评估了8种不同LLM（5个密集模型，3个MoE模型）的端到端性能，包括单节点（8卡）和双节点（16卡）场景。
  • 实验结果（图11）：
    • 单节点（8xH800）：TileLink平均比PyTorch基线快1.32倍。其中，密集模型平均加速1.20倍，MoE模型平均加速1.54倍。端到端加速低于单层MoE加速，因为模型中MLP层仍占约50%的执行时间。
    • 双节点（16xH800）：采用节点间数据并行、节点内张量并行。整体平均加速1.29倍，与单节点结果相似，略低是由于额外的跨节点通信开销。
  • 结论：TileLink在端到端模型训练中表现出稳定且显著的性能提升，效果可扩展到多节点环境。

图8 8xH800上MLP层（AG+GEMM和GEMM+RS）的性能结果。
图8 8xH800上MLP层（AG+GEMM和GEMM+RS）的性能结果。

图9 8xH800上MoE层（AG + Gather + GroupGEMM和GroupGEMM + Scatter + Reduce + RS）的性能结果。
图9 8xH800上MoE层（AG + Gather + GroupGEMM和GroupGEMM + Scatter + Reduce + RS）的性能结果。

图10 8xH800上自注意力层的性能结果和重叠率。
图10 8xH800上自注意力层的性能结果和重叠率。

公式 重叠率
公式 重叠率

图11 8xH800和16xH800上端到端模型的性能结果。
图11 8xH800和16xH800上端到端模型的性能结果。

A5 结论

本文提出了TileLink，一个用于生成高性能计算-通信重叠内核的框架。为了解决大型DNN模型在分布式系统部署中通信与计算重叠的关键问题，TileLink通过一套以tile为中心的原语来提高开发效率，并通过以tile为中心的映射来自动生成底层代码。这种方法克服了现有重叠研究中性能次优或开发困难的缺点。实验结果表明，与非重叠基线相比，TileLink实现了1.17倍至20.76倍的显著加速，其性能与最先进的重叠库相当。这证明了TileLink在提高分布式深度学习系统效率方面的有效性和潜力。

方法细节中的引用汇总

  • 引用 [5] FLUX: fast software-based communication overlap on gpus through kernel fusion (2024)

    • 引用段落: A3 - 2.2 Communication and Computation Overlapping
    • 原文描述: 在介绍核融合技术时，引用FLUX作为高度优化的库的代表。
    • 引用段落: A2 - 3.2.2 Data Primitives
    • 原文描述: 在讨论数据原语的pull和push模式时，指出FLUX曾提到这两种模式的选择会影响性能。

  • 引用 [6] Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning (2024, ASPLOS)

    • 引用段落: A3 - 2.2 Communication and Computation Overlapping
    • 原文描述: 在概述重叠技术时，引用Centauri引入了模型级、层级和操作级的三级设计空间。在介绍算子分解技术时，将其列为代表性工作。

  • 引用 [7] TVM: an automated end-to-end optimizing compiler for deep learning (2018, OSDI)

    • 引用段落: A3 - 2.3 Code Generation Compilers
    • 原文描述: 引用TVM作为代码生成编译器的代表，说明了这类编译器的发展使得生成高性能代码成为可能。
    • 引用段落: A2 - 4.3 Other Compilation Optimizations
    • 原文描述: 在讨论单设备优化时，指出这些策略在TVM等先前工作中已有深入研究。
    • 引用段落: A2 - 6. 实现
    • 原文描述: 提及TileLink的当前实现可以轻松适配到TVM等其他编译器框架。

  • 引用 [15] Breaking the computation and communication abstraction barrier in distributed machine learning workloads (2022, ASPLOS)

    • 引用段落: A3 - 2.1 Communication and Intra-Layer Parallelism
    • 原文描述: 在介绍张量并行FFN算法时，引用CoCoNet作为使用该算法的先前工作之一。
    • 引用段落: A3 - 2.2 Communication and Computation Overlapping
    • 原文描述: 在介绍核融合技术时，引用CoCoNet作为支持核融合的领域特定编译器。
    • 引用段落: A3 - 2.3 Code Generation Compilers
    • 原文描述: 提及CoCoNet是能生成重叠核的编译器，但受限于固定的重叠模式。

  • 引用 [18] Torchtitan: One-stop pytorch native solution for production ready llm pre-training (2024)

    • 引用段落: A3 - 2.4 Motivational Example
    • 原文描述: 在激励示例的性能对比中，Decomposition方法的数据来源于Async-TP PyTorch。

  • 引用 [19] Ring attention with blockwise transformers for near-infinite context (2023)

    • 引用段落: A3 - 2.1 Communication and Intra-Layer Parallelism
    • 原文描述: 引用RingAttention来解释序列并行（sequence-parallel）算法。

  • 引用 [23] Efficient large-scale language model training on GPU clusters using megatron-lm (2021, SC)

    • 引用段落: A3 - 2.1 Communication and Intra-Layer Parallelism
    • 原文描述: 引用Megatron-LM来解释序列并行和张量并行FFN算法。

  • 引用 [26] Nvidia collective communications library (2024)

    • 引用段落: A3 - 2.1 Communication and Intra-Layer Parallelism
    • 原文描述: 在介绍以算子为中心的通信原语时，引用NCCL作为提供这些原语的现有库。
    • 引用段落: A2 - 3.2 Tile-Centric Primitives
    • 原文描述: 指出TileLink的原语与NCCL等库提供的以算子为中心的原语不同。

  • 引用 [27] NVSHMEM (2025)

    • 引用段落: A2 - 6. 实现
    • 原文描述: 提及在运行时使用NVSHMEM来初始化分布式执行环境。

  • 引用 [30] Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the GPU (2023, PPoPP)

    • 引用段落: A3 - 2.2 Communication and Computation Overlapping
    • 原文描述: 在解释算子分解的缺点时，引用Stream-k来说明小算子可能导致资源量化效率低下。

  • 引用 [31] Optimizing distributed ml communication with fused computation-collective operations (2023)

    • 引用段落: A3 - 2.2 Communication and Computation Overlapping
    • 原文描述: 引用此文献作为支持核融合方法的高度优化的库之一。

  • 引用 [33] Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines (2013, PLDI)

    • 引用段落: A3 - 2.3 Code Generation Compilers
    • 原文描述: 引用Halide作为代码生成编译器的代表。

  • 引用 [39, 40] Triton: an intermediate language and compiler for tiled neural network computations (2019, MAPL@PLDI)

    • 引用段落: A3 - 2.2 Communication and Computation Overlapping
    • 原文描述: 引用Triton来说明小算子可能导致L2缓存利用率低。
    • 引用段落: A3 - 2.3 Code Generation Compilers & A2 - 3.2 Tile-Centric Primitives & A2 - 4.3 Other Compilation Optimizations & A2 - 6. 实现
    • 原文描述: 多次引用Triton，作为代码生成编译器的代表、现有编译器不支持内存一致性模型的例子、单设备优化已有研究的来源，以及TileLink的实现基础。

  • 引用 [41] Overlap communication with dependent computation via decomposition in large deep learning models (2023, ASPLOS)

    • 引用段落: A3 - 2.1 Communication and Intra-Layer Parallelism
    • 原文描述: 在介绍张量并行FFN算法时，引用Dist-Einsum作为使用该算法的先前工作。
    • 引用段落: A3 - 2.2 Communication and Computation Overlapping
    • 原文描述: 引用Dist-Einsum作为采用算子分解方法的代表性工作。
    • 引用段落: A3 - 2.3 Code Generation Compilers
    • 原文描述: 提及Dist-Einsum是能生成重叠核的编译器，但受限于固定的重叠模式。