1. 论文基本信息

1.1. 标题

Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler

标题解析： 论文标题清晰地指出了研究的核心内容：

• Triton-distributed：这是一个新提出的系统或框架的名称，表明它是对现有 Triton 编译器的扩展，专注于“分布式” (distributed) 计算。

• Programming Overlapping Kernels：点明了该框架的核心功能——支持“重叠内核” (Overlapping Kernels) 的编程。这里的“重叠”特指计算、内存访问和通信等操作在时间上的并行执行，以隐藏通信延迟，提升系统整体效率。

• on Distributed AI Systems：明确了应用场景是“分布式人工智能系统”，特别是面向大规模模型训练和推理的场景。

• with the Triton Compiler：强调了实现的技术路径——通过扩展 Triton 编译器来实现上述功能，意味着它继承了 Triton 使用 Python 编程的高效和便捷性。

  总而言之，标题表明本文提出了一种基于 Triton 的新方法，旨在简化分布式 AI 系统中高性能“重叠内核”的开发过程。

1.2. 作者

Size Zheng, Wenlei Bao, Qi Hou, Xuegui Zheng, Jin Fang, Chenhui Huang, Tianqi (作者名无法正常显示), Haojie Duanmu, Renze Chen, Ruifan (作者名无法正常显示), Yifan Guo, Ningxin Zheng, Ziheng Jiang, Xinyi (作者名无法正常显示), Dongyang Wang, Jianxi (作者名无法正常显示), Haibin Lin, Li-Wen Chang, Liqiang Lu, Yun Liang, Jidong Zhai, Xin Liu.

隶属机构：

• ByteDance
• 清华大学 (Tsinghua University)
• 北京大学 (Peking University)
• 上海交通大学 (Shanghai Jiao Tong University)
• 浙江大学 (Zhejiang University)

背景分析： 作者团队由工业界（字节跳动）和顶尖学术界（清华、北大、上交、浙大）的研究人员组成，这种产学研结合的背景通常意味着研究工作既有前沿的学术视野，又紧密结合了工业界的实际需求和大规模应用挑战。通讯作者同样来自字节跳动和清华大学，进一步凸显了合作的紧密性。

1.3. 发表期刊/会议

本论文目前发布在 arXiv 预印本服务器上，提交于 2025 年 4 月。arXiv 是一个开放获取的学术论文发布平台，研究人员可以在同行评审前分享他们的研究成果。这表明该工作是最新研究，但尚未经过正式的同行评审和学术会议/期刊的发表流程。

1.4. 发表年份

2025年 (预印本)

1.5. 摘要

随着单芯片性能接近瓶颈，单一加速器已无法满足大语言模型的训练和推理需求，使用多加速器组成的分布式系统成为必然选择。在分布式系统中，计算、内存访问和通信是三个并发进行的核心活动。然而，在现有框架中，这三者通常在不同编程层级被独立优化，导致难以协同工作，无法充分发挥集群的性能潜力。

为应对这一挑战，本文提出了 Triton-distributed，它是对现有 Triton 编译器的扩展。Triton-distributed 是首个原生支持分布式 AI 负载重叠优化的编译器，并能广泛覆盖现有框架中的多种优化技术。

• 方法一：我们将兼容 OpenSHMEM 标准的通信原语集成到编译器中，使开发者能通过更高级的 Python 编程模型来使用它们。
• 方法二：我们展示了如何借助编译器实现计算、内存访问和通信的复杂联合优化。特别地，我们演示了如何利用重叠技术隐藏延迟，并在单节点和多节点场景下介绍了基于编译器的编程方法。
• 结果：在多达 64 个设备的测试环境中，Triton-distributed 生成的代码能充分利用异构的通信和计算资源，实现有效的重叠和高性能。在许多情况下，其性能甚至超越了手动优化的代码。此外，使用我们的编译器进行开发的难度和时间成本远低于 CUDA/C++ 等底层编程，展现了显著的生产力优势。

1.6. 原文链接

• arXiv 链接: https://arxiv.org/abs/2504.19442

• PDF 链接: https://arxiv.org/pdf/2504.19442v3.pdf

• 发布状态: 预印本 (Preprint)。

  ---

2. 整体概括

2.1. 研究背景与动机

2.1.1. 核心问题

随着大语言模型（LLM）的规模爆炸式增长，单个 AI 加速器（如 GPU）的计算和内存能力已无法满足其训练和推理的需求。因此，必须使用由成百上千个加速器组成的分布式集群。在分布式系统中，性能不仅取决于计算速度，还严重受到节点间通信的制约。理想情况下，计算、内存访问和通信这三项活动应该并行进行，即所谓的计算-通信重叠 (Computation-Communication Overlap)，以隐藏通信带来的延迟，最大化硬件利用率。

2.1.2. 现有挑战 (Gap)

然而，实现高效的计算-通信重叠在现有技术栈中极为困难，存在以下主要挑战：

1. 编程语言割裂：AI 算法开发通常使用 Python，而高性能的分布式底层优化则需要使用 CUDA/C++。这导致算法研究员和系统工程师之间存在巨大的“编程鸿沟”，需要跨语言、跨团队协作，开发效率低下。
2. 独立优化而非联合优化：在现有的深度学习框架（如 PyTorch）中，计算（由 cuBLAS 等库处理）、内存访问和通信（由 NCCL 等库处理）通常是独立优化的。这种分离使得进行细粒度的、跨操作的联合优化变得非常困难。例如，NCCL 的集合通信操作通常是阻塞的，难以与计算任务的某个特定部分（如一个 tile 的计算）进行重叠。
3. 高昂的工程成本：虽然少数顶尖团队（如 NVIDIA 的 FLUX）通过手写 CUDA/C++ 代码实现了高性能的重叠内核，但这需要极高的工程技术和大量开发时间，普通开发者难以企及。
4. 现有编译器的局限性：尽管 Triton 等编译器在单设备代码生成上取得了巨大成功，性能媲美专家手写代码，但它们普遍缺乏对分布式通信的原生支持。早期的分布式编译器则不适用于 LLM 场景或不兼容 Python 生态。

2.1.3. 论文的切入点

本文的创新切入点是：将分布式通信能力原生集成到成功的单设备编译器 Triton 中。通过扩展 Triton，作者旨在创建一个统一的、基于 Python 的编程环境，让开发者能够像编写单设备计算内核一样，直观地编写高性能的、可重叠的分布式内核。这不仅填补了现有编译器在分布式优化上的空白，还极大地降低了开发门槛。

2.2. 核心贡献/主要发现

本文最核心的贡献是提出了 Triton-distributed，一个为分布式 AI 系统设计的编译器扩展。其主要发现和贡献可以总结为以下几点：

1. 首个支持原生重叠优化的 Python 级编译器：Triton-distributed 首次让开发者能在 Python 层面实现细粒度的计算-通信重叠。它通过引入一套基于 OpenSHMEM 标准的通信原语，并将其无缝集成到 Triton 的编程模型中，从而实现了这一目标。

2. 统一的编程模型：提出了一个基于 MPMD (Multiple Programs Multiple Data) 的编程模型，其核心概念包括对称内存 (Symmetric Memory)、信号交换 (Signal Exchange) 和异步任务 (Async-Task)。这个模型为描述复杂的分布式并行行为提供了清晰的抽象。

3. 高性能代码生成：实验证明，Triton-distributed 生成的代码在多种分布式任务（如 $AllGather+GEMM$, $GEMM+ReduceScatter$, AllToAll 等）上，性能媲美甚至超越了专家手写的 CUDA/C++ 代码（如 FLUX 和 DeepEP）以及标准库（NCCL/RCCL）。在某些场景下，加速比高达 44.97倍。

4. 显著提升开发效率：相较于需要数千行 CUDA/C++ 代码的底层实现，Triton-distributed 仅需数百行 Python 代码即可完成相同功能，极大地降低了开发难度和时间成本，使普通开发者也能进行高性能分布式编程。

5. 跨平台通用性：该方法不仅在 NVIDIA GPU 上表现出色，同样在 AMD GPU 上取得了显著的性能提升，证明了其设计思想的通用性和可移植性。

   ---

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. Triton 编译器

Triton 是由 OpenAI 开发的一种基于 Python 的领域特定语言 (DSL) 和编译器，专门用于编写高效的 GPU 内核。它的核心优势在于，开发者可以用类似 Python 的高级语法来编写 GPU 计算逻辑，而 Triton 编译器会自动将其优化并编译成高性能的底层代码（如 PTX for NVIDIA GPUs）。这使得开发者无需直接编写复杂的 CUDA 或 C++ 代码，就能达到与专家级手动优化相媲美的性能。在本文之前，Triton 主要专注于单设备的内核优化。

3.1.2. 分布式 AI 系统中的并行策略

为了训练和运行大型模型，通常需要将计算和数据分布到多个设备上。常见的并行策略包括：

• 张量并行 (Tensor Parallelism, TP)：将模型中的单个大张量（如权重矩阵）切分到多个设备上。例如，在矩阵乘法 $Y = XA$ 中，可以将 $A$ 按列切分为 $[A_1, A_2]$，分别在两个设备上计算 $Y_1=XA_1$ 和 $Y_2=XA_2$，最后将结果 $Y_1$ 和 $Y_2$ 拼接起来。这个过程需要 AllGather 等通信操作。
• 专家并行 (Expert Parallelism, EP)：在 Mixture-of-Experts (MoE) 模型中，每个输入 token 只会被路由到少数几个“专家”（即小型神经网络）进行处理。专家并行将不同的专家分布在不同的设备上，token 需要通过 AllToAll 通信在设备间动态路由。

3.1.3. 计算-通信重叠 (Computation-Communication Overlap)

在分布式计算中，一个典型的步骤是：1. 通过网络/互联链路获取远程数据（通信）；2. 使用获取到的数据进行计算。如果这两个步骤串行执行，那么在通信期间，计算单元将处于闲置状态，造成资源浪费。计算-通信重叠是一种关键的性能优化技术，其目标是在进行当前数据块的计算时，预先去获取下一个数据块。这样，计算和通信在时间上并行发生，通信的延迟就被计算的时间所“隐藏”，从而提升了系统的整体吞吐量。

3.1.4. 集合通信 (Collective Communication)

集合通信是指在一组进程（或设备）之间进行的预定义模式的数据交换。常见的操作包括：

• AllGather：每个设备都拥有数据的一部分，操作结束后，每个设备都将拥有所有设备数据的完整副本。
• ReduceScatter：与 AllGather 相反。每个设备先贡献一份数据，这些数据在所有设备间进行规约操作（如求和），然后规约结果被切分并分发回每个设备。
• AllToAll：每个设备都向其他所有设备发送一块独立的数据，并从其他所有设备接收一块独立的数据。

3.1.5. OpenSHMEM 标准

OpenSHMEM (Open Shared Memory) 是一个并行编程接口标准，它定义了一套单边通信 (One-Sided Communication) 原语。与传统的双边通信（如 MPI_Send/MPI_Recv）需要发送方和接收方显式配对不同，单边通信允许一个进程直接对另一个进程的远程内存进行读（get）或写（put）操作，而远程进程无需感知或参与。这种模型非常适合实现异步数据传输。NVIDIA 的 NVSHMEM 和 AMD 的 ROCSHMEM 是该标准在各自平台上的高性能实现。

3.2. 前人工作

1. 标准通信库 (NCCL/RCCL)：NCCL (NVIDIA Collective Communications Library) 和 RCCL (ROCm Communication Collectives Library) 是 NVIDIA 和 AMD 官方提供的、高度优化的集合通信库。它们性能极高，但其操作通常是宏观同步的。这意味着一次 AllGather 调用会阻塞，直到整个操作完成，这使得在操作内部实现细粒度的计算-通信重叠变得困难。重叠只能在更大的算子粒度上通过 CUDA Stream 等机制实现。

2. 底层优化框架 (FLUX, DeepEP)：

   • FLUX [3] 是一个专注于实现高性能重叠内核的框架，通过手写 CUDA C++ 代码，实现了如 $AllGather+GEMM$ 等融合操作，性能极高。但它的开发门槛非常高，代码复杂且难以维护。
   • DeepEP [36] 是一个为专家并行优化的系统，它实现了极低延迟的 AllToAll 内核，同样也是通过大量的底层 CUDA 编程实现的。 这些工作证明了细粒度重叠的巨大性能潜力，但也凸显了其高昂的工程成本，这正是 Triton-distributed 试图解决的问题。

3. 早期分布式编译器 (Distributed Halide, DISTAL)：

   • Distributed Halide [8] 和 DISTAL [33] 是早期探索分布式计算的编译器。然而，它们主要面向非大模型场景，并且倾向于提出专有的 DSL，而不是与 Python 生态系统对齐，因此难以满足当前 LLM 的分布式需求。

3.3. 技术演进

该领域的技术演进路径可以看作是不断追求性能和易用性平衡的过程：

1. 手动底层优化 (CUDA/C++ with NCCL)：性能好，但开发效率极低，语言割裂。
2. 单设备编译器 (Triton)：在单 GPU 上实现了 Python 编程与专家级性能的统一，但不支持分布式。
3. 专用底层框架 (FLUX, DeepEP)：在分布式场景下实现了极致性能，但牺牲了易用性和通用性。
4. 本文工作 (Triton-distributed)：试图将 Triton 的易用性扩展到分布式领域，通过编译器技术，让开发者能用 Python 实现 FLUX 等框架才能达到的高性能重叠效果，代表了编译器技术向分布式时代迈进的重要一步。

3.4. 差异化分析

Triton-distributed 与先前工作最核心的区别在于：

• 编程抽象层次：与 FLUX 和 DeepEP 等底层框架相比，Triton-distributed 在高层次的 Python 中提供了分布式编程能力。开发者无需关心底层的线程管理、内存拷贝和设备同步细节。

• 原生编译器支持：与使用 $PyTorch+NCCL$ 的方式不同，Triton-distributed 将通信原语内建于编译器中。这使得编译器能够对计算和通信进行联合优化，例如，生成依赖于通信信号的条件加载指令，实现真正的细粒度流水线。

• 通用性与生态：它建立在 Triton 这个已被广泛采用的编译器之上，天然继承了其对多种 AI 芯片的支持潜力，并且与 Python AI 生态无缝衔接。这与那些提出专有 DSL 的早期分布式编译器形成了鲜明对比。

  ---

4. 方法论

Triton-distributed 的核心思想是通过编译器提供一套高级编程抽象，将复杂的分布式协调逻辑转化为高效的底层代码。本节将深入拆解其编程模型、通信原语以及关键的重叠优化技术。

4.1. 方法原理

Triton-distributed 的方法论基于一个核心的洞察：要实现细粒度的计算-通信重叠，必须打破计算和通信之间的硬同步边界。它通过引入一个基于信号的异步编程模型来实现这一点。在这个模型中，无论是数据传输还是一个计算单元（tile）的计算，都被视为一个异步任务。这些任务的执行和依赖关系通过轻量级的信号来协调，而不是通过阻塞式的全局同步（如 barrier）。编译器负责将这种基于信号的高级描述翻译成高效的、可并行的底层硬件指令。

下图（原文 Figure 2）展示了 Triton-distributed 的编译流程。用户在 Python 中分别编写计算和通信逻辑。通信部分使用 OpenSHMEM 原语，与一个预编译的 bitcode 库一起被降级到 LLVM IR。计算部分则通过 Triton 原有的编译栈降级到 LLVM IR。最终，两者被合并生成目标硬件（NVIDIA 或 AMD）的后端代码。

该图像是示意图，展示了用户分类及我们的编译工作流。其中，用户分为四类：power user、ordinary user、fresh hand 和我们的目标用户。编译流程包含从用户编写的通信和计算脚本，到生成 LLVM IR 及相关库的多个步骤，直观呈现了 Triton-distributed 编译器的运行机制。

4.2. 核心方法详解

4.2.1. 编程模型：MPMD

Triton-distributed 遵循 MPMD (Multiple Programs Multiple Data) 模型，允许计算和通信任务并行运行并相互协作。该模型包含三个核心概念：

1. 对称内存 (Symmetric Memory)：

   • 定义：集群中的每个设备（rank）都在全局作用域内分配一个相同大小的内存缓冲区。
   • 特性：每个缓冲区的地址空间是独立的，不存在一个全局统一的虚拟地址空间（UVA）。因此，一个设备不能通过简单的指针解引用来直接访问另一个设备的远程内存。访问远程数据必须使用特定的通信原语。

2. 信号交换 (Signal Exchange)：

   • 定义：信号是位于对称内存中的一个数据对象（通常是一个整数），用于任务间的同步和通信。
   • 操作：有一组固定的原子操作来操纵信号的值，例如 set（设置值）、add（增加值）和 wait_until（等待信号满足特定条件，通常通过自旋锁实现）。信号是实现细粒度协调的关键机制。

3. 异步任务 (Async-Task)：

   • 定义：所有操作，如数据传输和计算，都被视为可以并行运行的异步任务。

   • 实现：在 GPU 上，这可以通过多 CUDA Stream 或多线程来实现。多 Stream 依赖运行时调度，而多线程则利用硬件的并行单元。任务间的同步完全依赖于信号。

     下图（原文 Figure 3）直观地解释了这三个概念。在一个 2 节点、每节点 2 个 rank 的例子中，左侧展示了每个 rank 的对称内存和信号交换。右侧的时间线显示，每个 rank 上有三个任务（节点间 P2P 传输、节点内 P2P 传输、计算）并行运行。计算内核中的 tile 执行顺序经过精心安排（swizzling），确保计算任务总是在其所需数据到达后立即开始，从而避免等待。

     该图像是图表，展示了在两个 H800 节点上执行 AllGather GEMM 的相对性能。不同编程框架的性能比较包括 PyTorch+NCCL、FLUX 和 Triton-distributed，图中显示了各自的性能值。

4.2.2. 通信原语

Triton-distributed 提供了一套丰富的通信原语，分为 OpenSHMEM 兼容原语和为特定优化设计的非 OpenSHMEM 原语。下表（原文 Table 1）列出了这些原语。

关键原语解释：

• wait 和 consume_token：这对原语是实现计算与通信在内核内部重叠的关键。在计算内核（如 GEMM）中，当需要使用一块远程数据时，代码会先调用 wait 等待一个与该数据块关联的信号。wait 操作返回一个 token。然后，在实际加载该数据之前，调用 consume_token。编译器会利用这个 token 建立数据依赖关系，确保加载操作只有在信号到达后才会执行，从而实现精确的流水线式同步。

4.2.3. 编程范例：跨节点重叠 AllGather GEMM

下图（原文 Figure 4）展示了如何使用 Triton-distributed 编写一个跨节点的重叠 AllGather GEMM。

• 通信部分（左侧）：使用不同的线程块 (threadblock) 并行执行节点内数据分发和节点间数据传输。

• 计算部分（右上）：复用标准的 Triton GEMM 实现，仅在加载数据前插入 wait 和 consume_token 两个原语。GEMM 中的每个计算 tile 会等待其对应的数据块信号，从而将数据依赖的等待时间与其他 tile 的计算时间重叠起来。

• Host 端代码（右下）：分配对称内存，并在不同的 CUDA Stream 上启动通信和计算内核，实现宏观上的并行。

  该图像是一个条形图，展示了在2个H800x8节点上进行GEMM ReduceScatter的相对性能。不同颜色的条形代表了不同的实现方式，包括PyTorch+NCCL、FLUX和Triton-distributed，相比之下，Triton-distributed展现出最佳性能。

4.2.4. 关键重叠优化技术

Triton-distributed 实现了多种先进的重叠优化技术，覆盖了从单节点到多节点，从高带宽到低延迟的各种场景。

1. 节点内 AllGather (使用 Copy Engine)

• 目标：在单个节点内部的多个 GPU 之间高效传输数据，同时不占用计算核心（SMs）。
• 方法：利用 GPU 内置的专用 DMA 引擎 (Copy Engine) 进行数据传输。
• 实现：
  • Push 模式 (算法 1)：每个 rank 主动将自己的数据块“推送”到所有其他 rank 的对称内存中，然后发送信号通知对方。优点是少一次同步，但无法控制数据到达顺序。
  • Pull 模式 (算法 2)：每个 rank 先将自己的数据复制到对称内存的指定位置，进行一次 barrier 确保所有 rank 都完成了本地复制。然后，每个 rank 主动从其他所有 rank 的对称内存中“拉取”所需数据。优点是数据到达顺序可控，但多一次同步开销。

2. 跨节点低延迟 AllGather

• 挑战：对于推理等延迟敏感场景，传统循环式的数据传输会导致“偏斜” (skew)，即使消息很小，总延迟也会被拉长。

• 方法：

  • 节点内广播：使用 NVIDIA PTX 中的 multimem_st 指令。该指令可以一次性将数据广播到节点内所有其他 GPU 的内存中，非常高效。
  • 节点间传输：采用低延迟协议 (Low-latency Protocol, LL)。该协议将数据和标志位打包成一个 8 字节的数据块，利用 GPU 硬件对 8 字节的跨设备原子读写支持，直接发送数据。接收方通过自旋锁检查标志位来判断数据是否到达。该协议速度快，但会使消息体积翻倍，适合小消息场景。

• 时间线对比：下图（原文 Figure 5）展示了优化前后的时间线。优化后的低延迟 AllGather 将延迟从约 25μs 降低到了 13.5μs。

  该图像是图表，包含三部分展示了分布式闪存解码的性能。左侧为强缩放延迟图，显示不同节点数（1、2、4节点）在不同全局KV长度下的延迟；中间为弱缩放带宽图，展示不同数量GPU在相同本地KV长度下的带宽表现；右侧为强缩放带宽图，展示不同GPU数量在不同全局KV长度下的带宽变化。这些数据展示了系统的性能特征和资源利用情况。

3. 跨节点 ReduceScatter (异构通信)

• 目标：实现 GEMM 计算和 ReduceScatter 通信的重叠。

• 分解：该操作被分解为三个阶段：节点内 scatter、本地 reduction、节点间 P2P 通信。

• 资源调度：为了最小化对主计算（GEMM）的影响，采用精细的资源划分策略：

  • 节点内 scatter 使用 Copy Engine，不占用 SM。
  • 节点间 P2P 通信仅需 1 个 SM。
  • 本地 reduction 分配最少数量的 SM（例如，在 H800 上只需不到 15 个 SM），确保其带宽能匹配通信速率，实现完美重叠。

• 时间线：下图（原文 Figure 9）展示了 GEMM（生产者）与 ReduceScatter 各阶段（消费者）的重叠时间线。

  该图像是图表，展示了在 L20 GPU 上低延迟 AllGather 的性能。图中显示了不同消息大小下各种方法的延迟，包括 NVSHMEM 32bit、NVSHMEM 64bit 和 Triton-distributed。性能随着消息大小的变化而变化，在 8xL20 和 16xL20 的设置中进行比较。

4. Swizzling 优化 (计算顺序重排)

• 核心思想：调整计算内核中 tile 的处理顺序 (swizzling)，使其与通信的数据到达顺序相匹配。这能最大化数据局部性（缓存命中率）并最小化等待时间。

• 硬件拓扑感知：swizzling 策略需要根据硬件的互联拓扑进行设计。下图（原文 Figure 6）展示了 NVIDIA H800 (通过 NVSwitch) 和 AMD MI308X (全连接网状) 的不同拓扑。

  该图像是图表，展示了在不同 GPU 数量下，Triton-distributed 和 DeepEP 的 AllToAll 调度与合并相对性能的比较。左侧图表为 AllToAll Dispatch，右侧为 AllToAll Combine，均显示出二者在不同设备上的性能差异。

• NVIDIA 上的 AllGather GEMM Swizzling (原文 Figure 7)：由于 NVSwitch 提供了任意两 GPU 间的高带宽，swizzling 策略是在每一步中，让每个 rank 从另一个特定的 rank 收集数据。计算从本地数据开始，然后依次使用从 rank 1, rank 2... 收集来的数据。

  该图像是图表，展示了在单节点 8xMI308X 上进行的 AllGather GEMM 性能相对比较。图中绿色柱状表示使用 PyTorch+RCCL 的性能，橙色柱状表示 Triton-distributed 的性能。提供了不同配置下的相对性能评估。

• AMD 上的 AllGather GEMM Swizzling (原文 Figure 8)：由于 AMD GPU 间是点对点直连，为充分利用总带宽，swizzling 策略是在每一步中，让每个 rank 同时从所有其他 rank 收集一小部分数据（sub-chunk），然后进行计算。

  该图像是图表，展示了在单节点 8xMI308X 上 GEMM ReduceScatter 的相对性能。图中比较了使用 PyTorch+RCCL 和 Triton-distributed 的性能表现，结果显示在不同输入参数下，两种方法的性能差异。

---

5. 实验设置

5.1. 硬件平台

实验在多个硬件集群上进行，以验证 Triton-distributed 的性能和通用性：

• NVIDIA GPU 集群：
  • H800 GPU：用于测试高性能的节点内和跨节点内核，集群规模从 8 GPUs (1节点) 扩展到 64 GPUs (8节点)。
  • L20 GPU：用于测试仅支持 PCIe 通信的场景，验证低延迟内核在非 NVLink 环境下的性能。
• AMD GPU 集群：
  • MI308X GPU：用于测试节点内内核，验证该方法在 AMD 平台上的可移植性和性能。

5.2. 评估任务

为了全面评估 Triton-distributed 的能力，论文设计了一系列涵盖不同并行策略和通信模式的优化内核。下表（原文 Table 3）总结了所有测试的内核。

5.3. 评估指标

实验主要使用以下指标来衡量性能：

1. 延迟 (Latency)
   • 概念定义：指完成一次操作所需的总时间，单位通常是毫秒 (ms) 或微秒 (μs)。延迟越低，性能越好。这个指标在衡量推理等对响应速度要求高的场景中尤为重要。
   • 数学公式：通常直接通过计时器测量得到，即 $T_{end} - T_{start}$。
2. 带宽 (Bandwidth)
   • 概念定义：指单位时间内传输的数据量，单位通常是 GB/s (Gigabytes per second) 或 TB/s (Terabytes per second)。带宽越高，表示数据传输效率越高。这个指标在衡量训练等需要处理大量数据的场景中非常关键。
   • 数学公式： $$\text{Bandwidth} = \frac{\text{Total Data Size}}{\text{Latency}}$$
   • 符号解释：
     • Total Data Size：操作中传输的总数据量。
     • Latency：完成该操作所花费的时间。
3. 加速比 (Speedup)
   • 概念定义：指优化后方法的性能相对于基线方法的性能提升倍数。它直观地展示了优化的效果。
   • 数学公式： $$\text{Speedup} = \frac{\text{Latency}_{\text{baseline}}}{\text{Latency}_{\text{ours}}}$$
   • 符号解释：
     • $\text{Latency}_{\text{baseline}}$：基线方法的执行延迟。
     • $\text{Latency}_{\text{ours}}$：Triton-distributed 方法的执行延迟。

5.4. 对比基线

论文将 Triton-distributed 与当前领域中最具代表性的方法进行了比较：

• PyTorch+NCCL/RCCL：这是目前最主流的、基于框架和标准库的实现方式。它代表了普通开发者能够轻易达到的性能基准。

• FLUX：一个由 NVIDIA 研究人员开发的、专门用于实现高性能重叠内核的底层 CUDA C++ 框架。它代表了专家级手动优化的性能上限。

• DeepEP：一个为 MoE 模型专家并行优化的系统，其手写的 AllToAll 内核是目前已知的性能最佳实现之一。

• NVSHMEM：NVIDIA 官方的 OpenSHMEM 实现，用于对比低延迟通信原语的性能。

  ---

6. 实验结果与分析

本章节将详细解析论文中的核心实验结果，展示 Triton-distributed 在不同场景下的性能表现。

6.1. 节点内内核性能 (NVIDIA H800)

6.1.1. AllGather + GEMM (AG+GEMM)

• 结果分析：下图（原文 Figure 11）展示了在 8 卡 H800 节点上，Triton-distributed 与 $PyTorch+NCCL$ 和 FLUX 的性能对比。尽管 Triton 生成的 GEMM 内核性能略低于 cuBLAS 或 CUTLASS（约 95%），但得益于更高效的计算-通信重叠，Triton-distributed 的端到端性能全面超越了另外两者。

• 量化数据：平均而言，Triton-distributed 相比 $PyTorch+NCCL$ 实现了 1.42倍 的加速，相比专家级优化的 FLUX 实现了 1.09倍 的加速。这证明了编译器生成的重叠策略优于手动优化。

  该图像是示意图，展示了用户分类及我们的编译工作流。其中，用户分为四类：power user、ordinary user、fresh hand 和我们的目标用户。编译流程包含从用户编写的通信和计算脚本，到生成 LLVM IR 及相关库的多个步骤，直观呈现了 Triton-distributed 编译器的运行机制。

6.1.2. GEMM + ReduceScatter (GEMM+RS)

• 结果分析：下图（原文 Figure 12）展示了 GEMM+RS 的性能。Triton-distributed 采用将 scatter 操作放在独立 stream 上异步执行的策略，而 FLUX 则是将 scatter 融合进 GEMM 内核。实验结果表明，Triton-distributed 的策略更优，取得了最好的性能。

• 量化数据：平均而言，Triton-distributed 相比 $PyTorch+NCCL$ 加速 1.28倍，相比 FLUX 加速 1.30倍。

  该图像是示意图，展示了对称内存、信号交换和异步任务的工作机制。在 NODE 0 的 rank 0 和 rank 1 中，通过信号交换实现通信和计算的协调，显现对称内存的使用。此外，图中的关系图详细展示了不同节点间的计算和通信任务的安排。

6.1.3. MoE 相关内核 (AG+MoE & MoE+RS)

• 结果分析：对于 MoE 内核，$PyTorch+NCCL$ 的基线实现使用了 Python 循环来处理 GroupGEMM，效率极低，因此 Triton-distributed 取得了巨大的性能优势。
• 量化数据：
  • 对于 AG+MoE-intra，平均加速比高达 44.97倍。
  • 对于 MoE+RS-intra，平均加速比也达到了 15.55倍。 以下是原文 Table 4 和 Table 5 的部分数据，展示了绝对性能值（单位：ms）。

Table 4: AllGather MoE 性能 (ms)

Table 5: MoE ReduceScatter 性能 (ms)

6.2. 跨节点内核性能 (NVIDIA H800)

6.2.1. AllGather + GEMM (AG+GEMM)

• 结果分析：在 16 卡 H800 (2节点) 环境下，如下图（原文 Figure 13）所示，Triton-distributed 的性能依然稳定超越 $PyTorch+NCCL$，并且与专家级手写的 FLUX 性能非常接近。

• 量化数据：相比 $PyTorch+NCCL$ 平均加速 1.33倍，达到了 FLUX 性能的 95.60%。考虑到开发效率的巨大提升，这是一个非常出色的结果。

  该图像是代码示例，展示了实现跨节点的AllGather GEMM操作的过程。代码主要分为两部分：生产者和消费者。生产者部分负责协调节点间的数据传输，使用信号等待机制来确保数据的同步；而消费者部分则进行矩阵乘法计算，并利用信号机制获取数据。整个结构清晰，反映了在分布式AI系统中进行优化计算和通信的编程方法。

6.2.2. GEMM + ReduceScatter (GEMM+RS)

• 结果分析：在 16 卡 H800 (2节点) 环境下，如下图（原文 Figure 14）所示，Triton-distributed 的性能同样优于 $PyTorch+NCCL$，并接近 FLUX。

• 量化数据：相比 $PyTorch+NCCL$ 平均加速 1.42倍，达到了 FLUX 性能的 96.36%。

  该图像是图表，展示了基线 AllGather 和优化后的 AllGather 的时间线对比。左侧为基线 AllGather 的时间，最坏情况约为 25 微秒；右侧为优化后的 AllGather，稳定时间约为 13.5 微秒。通过优化，显著降低了延迟。

6.2.3. 分布式 Flash Decoding

• 背景：Flash Decoding 是用于加速长文本生成的核心技术。作者首次将其扩展到分布式多设备场景。

• 结果分析：下图（原文 Figure 15）展示了分布式 Flash Decoding 的性能。

  • 弱扩展 (Weak Scaling)：保持每个 GPU 的 KV 缓存长度不变，增加 GPU 数量。结果显示，即使扩展到 32 个 GPU，每个 GPU 的 HBM 带宽依然能保持在较高水平（1.7 TB/s），证明了良好的扩展性。
  • 强扩展 (Strong Scaling)：保持全局 KV 缓存长度不变，增加 GPU 数量。结果显示，对于超长序列（如 1M），使用更多 GPU 能有效降低解码延迟。

• 结论：Triton-distributed 的低延迟 AllGather 内核是实现高效分布式解码的关键，为未来处理超长上下文的推理模型铺平了道路。

  该图像是图表，展示了H800 NVLink拓扑（左）与AMD MI308X拓扑（右）的对比。H800 NVLink的点对点带宽约为200 GB/s，而AMD MI308X的点对点带宽为50 GB/s，需使用所有7条链路才能达到350 GB/s。

6.2.4. 低延迟 AllToAll

• 背景：AllToAll 是 MoE 模型专家并行的性能瓶颈。DeepEP 的手写 CUDA 内核是该领域的标杆。

• 结果分析：下图（原文 Figure 16）显示，Triton-distributed 使用数百行 Python 代码实现的 AllToAll 内核，在 8 到 64 GPU 的规模上，性能与 DeepEP 相当甚至更优。

  • AllToAll Dispatch 平均加速 1.18倍。
  • AllToAll Combine 平均加速 1.44倍。

• 原因分析：Triton-distributed 的性能优势部分来源于更简单的内存管理逻辑（分配大块 buffer 而非 DeepEP 复杂的内存队列控制）。作者也指出，在更大规模（如 128 GPU）上，DeepEP 因使用了更具扩展性的 IBGDA 技术而表现更佳，这是未来工作的改进方向。

  该图像是示意图，展示了在Nvidia GPU上进行内部节点AllGather的Swizzle示例，假设有4个rank。图中展示了四个步骤，每个rank如何利用本地数据和收集的数据进行计算和汇总。

6.3. 节点内内核性能 (AMD MI308X)

• 结果分析：为了验证跨平台能力，作者在 AMD MI308X GPU 上也进行了测试。如下图（原文 Figure 17 和 18）所示，即使 Triton 生成的 GEMM 性能略低于 AMD 官方的 rocBLAS，Triton-distributed 凭借其卓越的重叠优化能力，最终的端到端性能依然超越了 $PyTorch+RCCL$ 基线。
• 量化数据：

  • 对于 $AG+GEMM$，平均加速 1.09倍。

  • 对于 $GEMM+RS$，平均加速 1.16倍。

    该图像是图表，展示了在 AMD GPU 上进行 intra-node AllGather 操作的示例，假设有 4 个 rank。图中分为四个步骤，分别为从不同 rank 收集数据并进行计算，清晰地描述了每个步骤的数据流动和计算过程。

    该图像是图表，展示了在116个SM上执行的GEMM内核（生产者）与其他计算流的时间线。包括了设备间数据复制、局部归约内核及跨节点p2p内核的操作，显示了不同流之间的交互与重叠。

---

7. 总结与思考

7.1. 结论总结

本文成功地解决了分布式 AI 系统中高性能内核开发的一大痛点——即在保持高开发效率的同时实现极致的性能。通过提出 Triton-distributed，作者将分布式通信能力原生、深度地集成到 Triton 编译器中，实现了以下关键贡献：

1. 统一了编程语言：开发者现在可以在统一的 Python 环境中，同时编写计算逻辑和复杂的分布式通信与同步逻辑，彻底消除了 Python 与 CUDA/C++ 之间的编程鸿沟。
2. 实现了专家级性能：实验结果有力地证明，由 Triton-distributed 编译器生成的代码，其性能不仅远超标准库，甚至能够媲美乃至超越由领域专家耗费大量精力手写的底层优化代码。
3. 极大地提升了生产力：将原本需要数千行复杂 CUDA 代码才能实现的功能，简化为数百行直观的 Python 代码，这无疑将极大地加速分布式 AI 系统的研发和迭代速度。
4. 展示了良好的通用性：该方法不仅在 NVIDIA 平台上表现优异，也在 AMD 平台上取得了成功，并支持 NVLink 和 PCIe 等不同互联方式，展示了其作为基础架构的巨大潜力。

7.2. 局限性与未来工作

尽管 Triton-distributed 取得了显著的成功，论文作者也坦诚地指出了其当前存在的局限性和未来的研究方向：

• 特定内核的扩展性问题：MoE+RS-inter 内核在从单节点扩展到多节点时性能扩展不佳，表明需要为该场景设计一个专用的 ReduceScatter 内核。
• 底层通信库依赖：AllToAll 内核在超大规模集群（>64 GPU）上的性能受限于 NVSHMEM bitcode 库当前不支持 IBGDA。未来的工作需要等待或推动底层库的支持以进一步提升扩展性。
• 自动化程度：当前的编程模型虽然大大简化了开发，但仍然需要开发者手动设计 swizzling 策略和资源划分方案。这是一个需要深厚领域知识的任务。将这些手动策略自动化，让编译器能够根据计算图和硬件拓扑自动推导出最优的执行计划，是未来一个重要的研究方向（论文中提到的 TileLink [40] 可能就是朝这个方向的探索）。

7.3. 个人启发与批判

这篇论文带来了深刻的启发，并对 AI 系统领域的发展具有重要意义。

• 启发：

  1. 编译器的力量：该工作是编译器技术赋能系统优化的又一典范。它表明，通过提升编程抽象层次并让编译器处理底层的复杂性，可以在不牺牲性能的前提下，极大地解放生产力。这对于整个 AI 社区 democratize 高性能计算至关重要。
  2. 软硬件协同设计的未来：论文中对 swizzling 策略的设计，深度依赖于对 NVIDIA NVSwitch 和 AMD 全连接拓扑的理解。这启示我们，未来的 AI 系统软件设计必须更加“硬件感知”，而编译器是连接上层算法与底层硬件细节的最佳桥梁。
  3. 基础架构的价值：Triton-distributed 提供的是一种底层能力，它为上层应用的创新提供了坚实的基础。例如，高效的分布式 Flash Decoding 将可能催生出能够处理更长上下文、具备更强推理能力的新模型。

• 批判性思考与潜在改进：

  1. 编程模型的学习曲线：尽管基于 Python，但 Triton-distributed 引入的对称内存、信号、异步任务等概念，以及手动设计通信模式和 swizzling 的要求，对于没有分布式编程背景的 AI 算法开发者来说，仍然存在一定的学习曲线。
  2. 调试与可维护性：异步和分布式程序的调试是众所周知的难题。虽然论文未详细讨论，但在实际应用中，如何为 Triton-distributed 提供有效的调试工具和可观测性，将是决定其能否被广泛采纳的关键因素。
  3. 从“能力”到“自动化”的飞跃：目前的工作提供了一种强大的“能力”，让开发者可以做到以前很难做到的事。但终极目标应该是“自动化”，即开发者只需描述计算逻辑，编译器就能自动完成所有分布式并行化和优化。本文是通往该目标的关键一步，但前路依然漫长。未来的研究可以在算子融合、自动并行策略选择、动态资源调度等方面进行更深入的探索。