AMD Iris：Triton原生多GPU通信库，以Tile级抽象实现1.79倍性能飞跃，重塑计算-通信融合范式

关键词：Iris、Triton、多 GPU 通信、计算-通信融合、对称内存抽象、tile 级编程

现代 AI 工作负载需要近乎峰值的性能以充分提取 AI 系统的效率。

然而，多 GPU 编程传统上要求开发者在性能与可编程性之间进行复杂的权衡：高性能实现通常依赖于低层 HIP/CUDA 通信库，即便实现基本的 重叠模式也需要大量工程努力；而更简单的抽象则往往牺牲性能。

• Iris: First-Class Multi-GPU Programming Experience in Triton
• https://arxiv.org/pdf/2511.12500
• 代码：https://github.com/ROCm/Iris

来自 AMD 的研究团队近日提出了一种革命性的解决方案——Iris，一个完全用 Python 和 Triton 实现的多 GPU 通信库，它消除了这种权衡。Iris 提供基于瓦片（tile）的对称内存抽象，这与 Triton 的编程模型自然对齐， 使开发者能够编写无缝交织计算与通信的单源内核。

我们的评估表明，Iris 在微基准测试中实现了接近最优的带宽利用率，并为 GEMM+All-Scatter 工作负载带来了高达 1.79 倍的性能提升，超越了 PyTorch 和 RCCL。这证明高层实现能够匹配甚至超越经过深度优化的库，同时极大地简化了多GPU编程。

这揭示了一个关键观点：实现细粒度重叠的真正障碍是抽象不匹配，而【非】硬件能力不足。当通信原语与计算存在于相同的语义空间（基于瓦片的 Triton）时，重叠模式便成为简单的扩展，而不再是艰巨的工程壮举。

一、多 GPU 编程的困境：性能与可编程性的权衡

传统的多 GPU 编程模式大多采用批量同步方式：计算内核完全执行完毕后，再启动通信内核进行数据交换。如图 1 所示，这种模式在计算与通信之间引入了明显的“气泡”（空闲等待时间），GPU 资源无法得到充分利用。

图 1：执行模型对比图。(a) 具有严格顺序阶段且由 CPU 启动内核的批量同步执行，是传统模式，CPU 主导内核启动，计算与通信间有强同步屏障，易产生空闲 “气泡”；(b) 支持以瓦片粒度动态通信且由 GPU 启动执行的细粒度重叠执行。作为 Iris 采用的模式，数据生成即通信，无需等待整阶段完成，能隐藏通信延迟，适配 AI workload 对高效多 GPU 执行的需求，如 LLM 训练中的张量并行场景

随着模型规模扩大和系统复杂度增加，细粒度重叠已成为业界共识——即在数据产生的同时就进行通信，而不是等待整个计算阶段完成。然而，在 Triton 这样的现代 GPU 编程框架中实现这一目标却【异常困难】。

二、根本挑战：通信抽象缺失

Triton 通过基于 tile 的编程模型彻底改变了优化计算内核的编写方式，自动处理内存合并、共享内存管理和内核内调度。然而，通信在 Triton 中始终是“二等公民”：

• 开发者必须依赖外部通信库（如 RCCL、NCCL），这些库在粗粒度内核边界上操作
• 通信操作对编译器不可见，无法与计算一起优化
• 现有封装方案（如 Triton-Distributed）只是将 xSHMEM 库包装为不透明字节码，无法实现真正的协同优化

三、Iris 的突破：原生 Triton 通信库

Iris 的核心创新在于完全在 Triton 内部实现通信原语，使编译器首次能够同时“看到”计算和通信操作。这一设计理念带来了根本性的优势。 需要注意的是，Iris 目前主要针对单节点内（scale-up）通过高速互联（如 AMD Infinity Fabric）紧密耦合的多 GPU 系统。多节点扩展（通过 RDMA）是未来的工作方向。

3.1 设计哲学：对称内存的现代化重塑

Iris 采用 SHMEM 的对称堆抽象，但彻底重新设计了编程接口：

1. 纯 Python 和 Triton 实现：不依赖任何外部通信库，仅使用 PyTorch 生态系统和 HIP 运行时 API
2. 熟悉的 C++/HIP 内存模型：使用 acquire/release 语义和内存作用域，无需学习新概念
3. 基于值和指针的双重 API：既支持细粒度的寄存器到远程内存传输，也支持粗粒度的缓冲区间复制

表 1：多 GPU 通信方法对比表。该表对比多 GPU 通信方法的关键维度，区分 “为 Triton 构建” 与 “基于包装器” 两类方案。从方法、编程模型、编译器可见性等 7 个维度对比 Iris 与 Triton-Distributed、PyTorch 等。Iris 采用原生 Triton 实现，支持 C++/HIP 内存模型，编译器可全量优化；包装器方法依赖 xSHMEM，接口为 C 风格，编译器可见性有限，难以协同优化计算与通信

3.2 核心技术：基于 tile 的对称内存 API

Iris 的核心抽象是基于 tile 的对称内存，这与 Triton 的计算模型自然对齐。 每个 GPU 分配一个对称堆，所有对等 GPU 都可以直接访问彼此的内存区域，无需通过主机 CPU 中转。

指针转换机制：实现透明的远程内存访问

图 2：Iris 地址转换机制。该图展示了本地指针通过三步转换为远程地址的过程，揭示了对称堆（SymHeap）与地址映射的核心逻辑。

Iris 的核心创新之一是其透明的指针转换机制，它使跨 GPU 的远程内存访问如同访问本地内存一样简单。该机制基于 HIP 的进程间通信（IPC）能力，无需开发者手动管理复杂的地址映射。

其转换过程分为三个清晰步骤：
1. 计算偏移量：确定本地指针在其所属对称堆内的偏移（例如，0x448）。
2. 叠加基地址：将此偏移量添加到目标 GPU 的对称堆基地址（例如，远程 rank 的基地址 0xFFFCABC0）。
3. 类型转换：将计算结果转换回目标 GPU 上可用的正确指针类型。

通过这一机制，Iris 将跨 GPU 内存访问的复杂性完全抽象，显著简化了分布式 GPU 编程模型。

简洁直观的编程接口

Iris 提供了层次清晰的 API，涵盖主机端（CPU）和设备端（GPU）操作。

表 2：Iris 主机端核心 API。涵盖库初始化、对称堆管理、张量创建及序列生成等功能。

主机端 API 负责全局管理与协调，核心功能包括：
* 基础设施：初始化（构建对称堆、交换 IPC 句柄）、全局同步（barrier）。
* 张量创建：支持 zeros、ones 等与 PyTorch 兼容的操作，所有张量自动分配在对称堆中，天然支持跨 GPU 访问。
* 序列生成：如 arange，简化多 GPU 环境下的数据初始化。

表 3：Iris 设备端核心 API。分为内存操作、原子操作及地址转换三类。

设备端 API 供 Triton 内核调用，实现高性能计算与通信：
* 内存操作：包括加载/存储（值级）、内存拷贝（指针级），并支持非阻塞访问。
* 原子操作：覆盖加、减、位运算等，可配置内存序（如 acquire/release）与同步范围（块/设备/系统）。
* 地址转换：内部使用的核心函数，是所有跨 GPU 内存访问的基础。所有 API 均需传入堆基地址以适配 Triton 的编程范式。

四、计算-通信重叠模式的分类学

Iris 支持从简单到复杂的多种计算-通信重叠模式，旨在最大化硬件利用率。图 3 展示了其完整的分类体系。

图 3：Iris 支持的计算-通信重叠模式分类法。主要分为未融合（Unfused）与融合（Fused）两大类。

4.1 未融合模式：清晰的关注点分离

未融合模式将计算与通信任务分离到不同的内核或流中执行。
1. 批量同步：最简模式，计算和通信在单独的内核中严格顺序执行。
2. 生产者-消费者（流并发）：利用多个异步流，将计算内核和通信内核并发执行，通过资源分区实现重叠。

图 4：未融合模式执行时间线示意图。左图为批量同步模式，存在明显的 GPU 空闲间隙；右图为流并发模式，通过分配部分计算单元（CU）专门进行通信，实现了部分重叠。

4.2 融合模式：极致的性能优化

融合模式将计算与通信逻辑写入同一个内核，以消除内核启动开销并实现更细粒度的重叠。
1. 顺序融合：在同一内核中，先完成一个计算单元（如Tile）的计算，随后立即执行其对应的通信操作。
2. 工作组分化：在同一内核内，将不同的 GPU 工作组（Workgroup）专门分配给计算任务或通信任务，实现硬件层面的并行。

图 5：融合模式下的 GEMM+All-Scatter 执行时间线。左图为顺序融合，计算与通信交替进行；右图为工作组分化，计算与通信工作组并行工作。

示例：GEMM + All-Scatter 的顺序融合实现

以下代码片段展示了“顺序融合模式”，体现了 Iris 在 Tile 粒度进行融合的设计思想。

与未融合模式不同，该代码在 GEMM 主循环中，每个 Tile 计算完成后立即通过 iris.store API 将结果分发到远程 GPU。这避免了先将结果写入全局内存再读取的额外带宽开销。
* 编译器提示（如 tl.max_contiguous）帮助 Triton 优化内存访问模式。
* 全局偏移量计算精确匹配 All-Scatter 的通信语义，确保数据正确拼接。

示例：GEMM + All-Scatter 的工作组分化实现

以下代码展示了“工作组分化模式”，这是实现极致性能重叠的关键。

代码将 GPU 的计算单元划分为两组：GEMM_SMS 负责计算，COMM_SMS 负责通信。两者通过原子操作（atomic_cas）配合 release/acquire 内存序进行同步，确保数据就绪后通信立即启动。
* 持久化循环让工作组持续处理任务，消除了内核重复启动的开销。
* iris.put 指针式通信支持跨 GPU 内存的直接访问。

该模式有效解决了未融合模式中的资源闲置问题，在特定问题规模上实现了高达 2.5 倍的性能加速。

五、性能表现：超越传统优化库

在 8×AMD Instinct™ MI300X GPU 集群上的评估表明，Iris 在保持编程简洁性的同时，带来了显著的性能提升。

图 8：Iris（融合GEMM+All-Scatter）与 RCCL（GEMM+AllGather）在不同问题规模和 GPU 数量下的性能对比。

测试中，Iris 平均性能达到传统 RCCL 方法的 1.21 倍，在最佳场景（8192×3584×14336，8卡）下甚至达到 2.5 倍。RCCL 因采用粗粒度的批量同步模式而存在大量计算空闲，而 Iris 通过细粒度的计算-通信重叠有效隐藏了通信延迟。

5.1 微基准测试：接近理想的带宽利用率

Iris 在基础通信原语上实现了极高的效率。

图 6：加载、存储及原子操作的带宽利用率热图。

测试显示，Iris 的加载/存储操作带宽利用率高达 95%-99%，原子操作也超过 90%。这验证了其原生 Triton 实现相较于传统包装器方案（如基于 xSHMEM 的方案）的更低开销优势。

图 7：全加载/全存储操作的带宽随缓冲区大小变化曲线。

随着缓冲区增大，Iris 的性能迅速逼近硬件峰值带宽。当缓冲区达到 2GB 时，带宽利用率可达 99%，表明其特别适合大数据量的 AI 训练负载。

5.2 实际工作负载：显著加速

综合性能测试证实，Iris 能够将微基准测试中的高效转化为实际复杂工作负载的端到端性能飞跃，为分布式高性能计算提供了新的高效范式。

对于 GEMM+All-Scatter 工作负载，Iris 相比 PyTorch 和 RCCL 实现了最高 1.79 倍的加速，平均加速比为 1.21 倍：
* 生产者-消费者模式：在 8192×3584×14336 问题形状上实现最高 2.5 倍加速。
* 顺序融合模式：在 8192×4608×36864 问题大小上实现 1.8 倍（4 GPU）和 1.5 倍（8 GPU）加速。

下图展示了不同问题形状下 GEMM 和通信的细粒度分解，突出了 Iris 通过灵活调度实现的重叠优势。

图 9：8192×4608×36864 矩阵规模性能分解图。该规模下 Iris 与 PyTorch、RCCL 的 GEMM（深色）和通信（浅色）耗时分解，Iris 因适配问题规模实现双环节提速，瓦片抽象支持编译时调参。Iris 的 GEMM 通过编译时调整瓦片大小减少计算冗余，通信因细粒度传输降低延迟，整体比 PyTorch+RCCL 快 20%。而 RCCL 的 AllGather 因依赖 CPU 启动内核，存在额外调度开销，凸显 Iris 瓦片抽象的灵活性。

图 10：8192×3584×14336 矩阵规模性能分解图。该规模下 Iris 与 PyTorch、RCCL 的耗时分解，凸显生产者-消费者/工作组专用化调度的优势——通过 CU 分区将通信开销隐藏于 GEMM 计算中。8 卡场景下，Iris 非融合生产者-消费者模式达 2.5 倍提速，通信耗时几乎被 GEMM 覆盖；PyTorch+RCCL 因同步屏障，通信时 GPU 空闲占比超 40%，无法利用计算时间掩盖延迟，体现 Iris 重叠策略的价值。

图 11：8192×4608×36864 矩阵规模融合顺序模式性能分解图。该规模下 Iris 融合顺序模式的优势展示——因通信开销占比低（输出瓦片小、K 值大），可集中资源给 GEMM，提升整体效率。4 卡时 Iris 比 PyTorch+RCCL 快 1.8 倍，8 卡快 1.5 倍。因 K 值极大（36864），GEMM 占总耗时超 80%，通信开销仅 20%，融合后无需内核切换，资源集中用于计算，避免传统模式的内核启动与同步开销。

关键洞察是：性能提升主要归功于 Iris 能够实现瓦片（tile）粒度同步，从而在计算和通信之间实现细粒度重叠，这与 PyTorch 和 RCCL 的刚性批量同步编程模型形成鲜明对比。

六、相关工作：从封装到原生

Iris 的设计哲学与现有方法有本质区别：
1. xSHMEM 库（NVSHMEM、rocSHMEM）：基于 OpenSHMEM 规范，但其 API 与 Triton 的瓦片中心模型不匹配。
2. Flux：直接在 CUDA 和 CUTLASS 中实现，性能优化但开发周期长。
3. Triton-Distributed：将 xSHMEM 库封装在 Triton 特定的 Python API 后，通信层对编译器仍然不透明。
4. PyTorch 对称内存：提供框架级支持，但仍依赖供应商提供的通信后端。

Iris 的独特之处在于：完全在 Triton 内部实现通信原语，使编译器能够跨计算和通信边界进行整体优化。

七、意义与未来方向

Iris 的成功证明了一个关键观点：细粒度重叠的主要障碍是抽象不匹配，而不是硬件能力不足。当通信原语与计算处于相同的语义空间（基于瓦片的 Triton）时，重叠模式成为简单的扩展，而非艰巨的工程努力。

Iris 的 1.79 倍峰值加速反映了编程模型的质的转变：同步粒度从内核边界转移到瓦片级别。更重要的是，Iris 使各种重叠模式变得易于实现——从批量同步到融合顺序，再到生产者-消费者和工作组分划——通常只需要在同一 Triton 内核中添加几行代码。

未来工作将集中于：
1. 多节点扩展：通过 RDMA 支持扩展到多节点设置。
2. 更多融合模式：探索 Gluon 中的波前分划和工作队列。
3. 编译器卸载：利用 Triton 优化整个计算-通信管道的能力。

八、结语

在 AI 计算需求持续爆炸式增长的背景下，Iris 代表了一种范式转变：高性能多 GPU 编程不再需要牺牲可编程性和开发效率。通过将通信提升为 Triton 中的“一等公民”，Iris 为开发者提供了同时获得最佳性能和最佳开发体验的路径。

对于那些厌倦了在复杂底层代码和性能妥协之间做选择的 AI 工程师和研究者来说，Iris 提供了一个令人信服的新选择：用 Python 和 Triton 的简洁性，实现传统上只有 HIP/CUDA 专家才能达到的性能水平。

这项工作的开源实现在：https://github.com/ROCm/Iris，为社区采用和进一步创新提供了坚实基础。

注：Iris 是 AMD 更广泛的软件生态系统优化的一部分，与 RCCL、rocSHMEM 和 ROCm 编译器工具链协同工作，为 AMD Instinct 加速器提供端到端的高性能计算解决方案。

关注“鲸栖”小程序，掌握最新AI资讯