突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

关键词计算-通信重叠块调度分布式编译器、GPU、Triton、多 GPU 工作负载

通过块级调度在单内核内实现计算与通信的深度重叠

近年来,大语言模型的规模呈指数级增长,训练这些模型需要数百甚至数千块 GPU。在多 GPU 系统中,通信已经取代计算成为主要瓶颈。即使采用 NVLink、NVSwitch 等高速互连技术,AllGather、ReduceScatter 等集合通信操作依然会占据端到端时间的很大比例

为了隐藏通信开销,现有系统主要采用内核级重叠策略:将计算内核和通信内核分配到不同的 CUDA 流中并发执行,从而实现粗粒度的计算-通信重叠。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

表 1 | 各类分布式操作相关项目的对比。该表格从粒度、计算、通信、调度、性能五个核心维度,对比了自动实现、手工实现、领域特定语言三类分布式方案与 AutoOverlap 的差异。传统自动方案仅支持内核级粒度,手工实现和 DSL 虽能实现块 / tile级优化但依赖人工开发,而 AutoOverlap 首创自动块级优化,实现计算、通信、调度全自动化,同时达到手工优化的顶级性能,填补了自动编译与细粒度优化间的技术空白。

然而,这种粗粒度的重叠存在根本性的不足,无法充分利用现代 GPU 的计算资源。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

针对这一挑战,来自加州大学圣地亚哥分校和 Meta 的研究团队提出了 AutoOverlap ——一个编译器和运行时框架,能够在单个融合内核内部实现细粒度的计算和通信重叠。

本文将对 AutoOverlap 的创新点进行解读,了解它如何通过“通信块”抽象和智能调度,为多 GPU 工作负载带来最高 4.7 倍的加速。

一、背景与相关研究

1.1 分布式编译器的现状

随着模型规模增长,张量编译器(如 TVM、TensorIR 等)开始集成多 GPU 支持。这些编译器通常采用“单 GPU 内核 + 预定义集合通信原语”的组合方式。

像 Alpa、Mercury 这样的分布式编译器会在搜索并行化策略后,构造包含 AllGather、ReduceScatter 等操作的通信 schedule,然后将计算内核和通信内核分配到不同流中,最大化内核级重叠。

然而,这种设计有一个根本限制:通信被规划为一系列完整的集体操作内核,内核的启动和完成时间是调度的基本单元所有现有的分布式编译器只关注内核层面的通信 schedule,对内核内部的细粒度重叠机会“视而不见”——例如,无法实现每分片(per-shard)或每块(per-tile)的通信与计算重叠,无法在内核内重用远程数据,也无法利用拓扑感知的流水线。

1.2 手工内核设计

另一类工作尝试超越内核级重叠,手工设计细粒度流水线。

  • Flux 在块级别融合 GEMM 和集合通信;
  • Comet 针对 MoE 模型使用共享张量抽象和 NVSHMEM 缓冲区实现词元级通信与块级计算的重叠;
  • FlashOverlap 通过轻量级就绪信号在不修改内核的情况下触发 NCCL 通信;
  • TritonDistributed 等 DSL 简化了编写重叠内核的过程。

尽管这些手工实现性能优越,但它们都依赖于专家针对特定算子和硬件平台进行手动优化,难以泛化到新的模型架构或硬件上。即使有了细粒度 DSL,发现有效的重叠策略、编码依赖关系、验证正确性的负担仍然落在程序员肩上。

1.3 通信后端多样性

现代 GPU 提供多种设备间数据传输机制,各有不同的性能和可编程性权衡,见表 2。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

表 2 | 各类 GPU 通信机制的对比。表格从硬件载体、编程方式、集体通信支持、带宽四个维度,剖析了 Copy Engine、TMA、Load/Store 三种 GPU 通信机制的核心特性与取舍。Copy Engine 带宽最高但依赖主机启动且不支持集体通信,TMA 为异步指令驱动、带宽次之,Load/Store 灵活性最强但带宽最低且占用 SM 资源。该对比为 AutoOverlap 的自适应通信后端选择提供了硬件依据,使其能根据通信场景动态匹配最优机制。

  • Copy Engine(拷贝引擎) 可以独立于 SM 运行,不消耗计算资源,但只能传输连续数据,且每次传输需要约 2-3 微秒的启动开销。
  • Tensor Memory Accelerator(TMA) 利用专用异步张量拷贝硬件,吞吐量高达 300+ GB/s,仅需少量 SM 发起指令,但目前仅支持节点内点对点通信。
  • 而普通的 Load/Store 方式最灵活,可配合 NVSHARP 实现在网规约,但消耗 SM 资源且同步性强,需要精细调度。

通信后端的多样性意味着 没有一种机制在所有场景下最优,编译器需要根据传输粒度、模式动态选择最合适的后端。

二、动机:为何需要细粒度重叠?

研究团队通过一系列微基准测试,揭示了内核级重叠的根本缺陷和细粒度重叠的潜力。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 1 | AutoOverlap 示例。左侧展示了传统内核级重叠的局限,如设备同步开销、SM 空闲时间长和通信尾延迟;右侧展示了 AutoOverlap 的创新方案,通过自适应通信后端、内核内 tile 调度和自适应块大小三大优化,将通信操作嵌入融合内核,打破传统性能瓶颈。红色数字表示细粒度重叠相比内核级重叠的直接改进,橙色数字表示 AutoOverlap 所开启的新设计空间带来的额外改进。

2.1 洞察 1:内核级重叠的局限性

内核级重叠导致 SM 利用率下降和内核启动开销

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 2 | 实验结果。(a)不同 GEMM 大小和 tile 大小下的 SM 利用率,显示 GEMM 规模减小时,内核分区会导致 SM 利用率显著下降;(b)流式 GEMM 内核与内核分区基准的性能对比,证明流式内核通过保持细粒度计算 tile,在避免额外启动开销的同时维持高利用率;(c)(d)展示了不同消息大小下各通信后端的带宽表现,量化了通信后端与传输粒度的匹配关系,表明 Copy Engine 在大消息量下带宽最优,而 TMA 和 Load/Store 各有适配场景,为后续块抽象设计提供了数据支撑。

大 GEMM 有足够多的 tile 波次来饱和 SM。但随着 GEMM 尺寸减小,tile 波次减少,最后部分填充的波次占执行时间的比例增大,导致 SM 利用率下降。将 GEMM 划分为多个子内核,强制每个启动操作于更小的形状,恰好落入了这种低利用率区域(对应图 1 中的 ①)。

图 2(b) 对比了两种变体:
* 一种是划分为多个小内核以实现重叠;
* 另一种是在单个启动内实现内部流水线(流式 GEMM)。

虽然两者执行相同算术操作,但内核划分基线因额外内核启动(①)和 SM 利用率不足(②)而导致性能大幅下降。

结论:简单地启动更多内核并非实现重叠的有效路径,我们需要一种在保持高效 GPU 执行的同时暴露内核内部并发性的机制。

2.2 洞察 2:传输粒度与后端效应

通信效率随传输粒度和后端选择剧烈变化

图 2(c) 和 2(d) 展示了不同通信后端的带宽随传输大小和 SM 数量的变化。每个后端都有不同的缩放行为:有些在中等传输大小时达到峰值带宽,有些则需要更大传输量或更多 SM 才能发挥潜力。此外,不同后端支持的模式也不同,如点对点传输与规约。

这意味着最优配置依赖于计算 tile 大小(决定结果可用的节奏)、通信传输大小(权衡延迟和带宽)以及通信后端的组合。内核级重叠无法灵活协调这些参数,因为计算和通信是独立的刚性内核。而内核内重叠可以在同一时间协调计算和通信,使 tile 产生速度与后端的“甜蜜点”对齐,选择能维持 SM 和拷贝引擎高利用率的传输粒度。

2.3 洞察 3:需要统一的块粒度抽象

有效的内核内重叠需要一个粒度可调的通信单元,并能在不同通信后端上提供稳定接口。
* 这个单元的大小既要匹配 tile 产生数据的速度,也要匹配不同后端的效率点。
* 同时,它必须在高层通信调度和后端具体实现之间提供稳定边界,使得通信调度无需为每个后端重写。

因此,AutoOverlap 引入了 块(chunk)抽象 —— 一个逻辑数据块,作为通信的基本单元,具有可调粒度,并能通过统一接口映射到多样化的后端。

三、AutoOverlap 系统概述

AutoOverlap 是一个编译器和运行时框架,能够将本地编写的 Triton 内核自动转换为细粒度重叠的分布式内核。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 3 | AutoOverlap 的系统概述。该图展示了端到端的编译流程:输入为带轻量注释的本地 Triton 内核和通信调度(支持用户自定义或复用模板),经解析接口、块调优器、通信调度器等模块处理后,通过代码生成器和 tile 调度器实现编译转换,最终输出融合了 Copy Engine、TMA 或 Load/Store 等通信后端的分布式算子。整个架构无缝集成 PyTorch 分布式运行时。

输入 包括两部分:
1. 本地 Triton 内核:用户使用标准 Triton 原语编写,如同运行在单个设备上。可选注释提供轻量级调度元数据(如 tile 尺寸、tile 索引标识符、调度循环结构),但不改变内核语义。
2. 通信调度:编码期望的全局数据移动模式和设备拓扑。可以由用户通过小型 API 定义,或直接从高层分布式编译器(如 Alpa、Mercury)的并行化策略导入。

AutoOverlap 将这些输入转换为统一的块级依赖图,然后搜索细粒度调度:将 tile 映射到设备,将通信操作分配到合适的后端(拷贝引擎、TMA 或 Load/Store),插入必要的同步以遵守依赖关系。输出是一个 融合的分布式 Triton 内核,它在单个启动内部以紧密流水线方式发出异步通信和计算。从用户视角,该内核与原始本地内核具有相同签名,只需额外传入分布式运行时参数(rank、world size、mesh),即可无缝集成到现有训练代码。

四、核心创新:通信块抽象

AutoOverlap 的灵魂是 通信块(chunk) 抽象。它位于全局逻辑张量和本地计算 tile 之间,使编译器能够对称地处理通信和计算。

4.1 块的定义

一个 块(chunk) 是作为单元进行通信的逻辑数据块,包含一个或多个 tile(本地内核中的计算基本单元)。块大小在通信调度中指定逻辑传输粒度,具体实现时可能被映射到不同的物理通信模式和后端。

在代码层面,块可表示为:

chunk = Chunk(sizes=[...], layout=..., tensor=...)

4.2 基于块的通信算子

AutoOverlap 定义了两种主要通信算子:

  • 点对点(P2P)传输P2P(src_rank, dst_rank, src_chunk, dst_chunk, dependency)。若在源端定义则为推送(push),在目的端定义则为拉取(pull),这会导致后续 lowering 阶段的不同实现选择。
  • 集体通信Collective(collective_type, src_chunk, dst_chunk, ranks, dependency)。允许编译器利用后端的优化集体通信实现。

dependency 字段编码了排序约束,以 (rank, index) 元组表示,指示当前操作必须等待指定 rank 上的某个操作完成后才能开始。这允许表达复杂的通信模式,如环状交换、多阶段集体通信,通过链式依赖连接通信块。

一个通信调度(schedule)是每个 rank 上的通信操作序列:

schedule := [rank: Int, operations: List[CommOp]]

由于每个 rank 的操作可以不同,该抽象能够表达异构的通信模式。

4.3 表达力

尽管结构简单,块抽象足以捕获实践中使用的各种通信模式。图 4 展示了几个例子:

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

* (a) 和 (b) 分别以推送和拉取方式表达同一 P2P 通信。
* (c) 展示基于环的 AllGather 模式,每个 rank 以流水线方式发送和接收块,依赖关系确保操作顺序正确。
* (d) 展示基于分区的 AllReduce 模式,每个 rank 贡献一个块并在“光纤”上累加。
* (e) 展示复杂的异构交织 AllGather 模式,通过端口抽象在不同层级同时处理通信,实现多维度的细粒度流水线。

4.4 从高层 IR 降级

通信调度可以由用户手动构建,也可以从高层分布式编译器 IR 自动推导。

AutoOverlap 为基于分区和基于循环的 IR 提供了前端:分析全局数据划分、识别通信点,将通信区域按所选粒度分组为块,最终得到统一的块级调度,如下面代码所示。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

集体通信操作可以直接插入调度,也可以进一步使用模板或综合算法(如 TACOS)降级为 P2P 操作。

五、基于块的代码生成

给定块级通信调度,下一步是重组本地计算,使得计算分片(tile)的执行顺序与通信块的到达和消耗对齐。AutoOverlap 通过以下步骤实现。

5.1 计算内核注释

用户在本地内核上通过 AutoOverlap API 添加轻量级注释(Annotation),暴露其分片(tiling)结构和迭代顺序。需要三类信息:
* 分片尺寸:每个分片的逻辑形状(如 GEMM 块),用于将分片映射到通信块。
* 分片索引标识符:唯一标识当前处理的分片的程序变量或元组。
* 分片调度器:推进分片索引并决定分片访问顺序的循环或控制结构。

这些注释通常可通过现有索引表达式和循环边界轻松添加,如代码 1 所示。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

5.2 依赖解析

AutoOverlap 构建一个包含通信块和计算分片的依赖图。对于每个通信块,追踪其生产者(producer)和消费者(consumer),以及通信调度中编码的显式排序约束(如流水线阶段)。对于每个计算分片,根据其分片索引和张量布局确定它读写哪些通信块。

从图中,编译器识别出遵守所有数据依赖所需的最小同步点集合。具体地,在内核中插入等待操作,使得消耗某个通信块的分片必须等到相应通信操作完成后才能启动。这种同步可以根据所选后端实现为不同的机制,但都源自同一块级依赖结构。

5.3 通信代码生成

通信调度对齐后,AutoOverlap 将抽象的块级调度 lowering 为具体通信代码。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

如图 7 所示,同一逻辑调度可以通过多种后端实现,这些后端在移动通信块的方式和流式多处理器(SM)资源的分配上有所不同。AutoOverlap 支持五种实现方式:
1. 专用拷贝引擎(copy engine)
2. 专用 SM 上的张量内存加速器(TMA)
3. 共置 SM 上的张量内存加速器(TMA)
4. 专用 SM 上的算子指令 Load/Store
5. 共置 SM 上的算子指令 Load/Store

对于每种方式,编译器根据块调度构建分片和通信步骤之间的依赖图,然后将该图 lowering 为后端特定的代码,通过构造保证所有依赖得到满足。例如:
* 当使用拷贝引擎或专用 SM 时,编译器会发出全局内存信号和内核启动,使通信相对于主计算分片异步进行;
* 当使用共置 SM 后端时,会生成共享内存屏障和索引簿记,协调同一 SM 内的通信和计算。

由于五种实现共享同一逻辑调度但暴露不同的延迟/带宽和资源权衡,它们构成了自动调优器的搜索空间:AutoOverlap 自动生成所有有效实现,测量端到端性能,并为每个算子和硬件配置选择性能最佳的后端。

5.4 分片调度器重排

通信调度与原始计算内核通常对全局张量具有不同的布局:通信根据数据移动位置将分片分组为通信块,而内核根据自身的遍历顺序将分片分组为波次。

先前的工作通过显式地在通信和计算之间插入数据重排来解决这个不匹配,这会增加额外的全局内存流量和同步。AutoOverlap 则不同:它保持通信块就位,转而 在内核内部对分片调度器进行重排(swizzling)

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 6 | Tile调度器swizzling | (a)计算和通信自然遵循不同的tile / 块布局,造成错位;(b)现有方案通过在两条路径间插入显式数据重排来解决该问题;(c)Syncopate(AutoOverlap 核心模块)则重写tile调度器以遵循块顺序,并应用块内混洗优化局部性,无需额外数据移动即可使计算与通信进度对齐。总得来说,传统方案(b)通过数据重排解决布局错位,却引入额外内存开销和同步延迟;而 AutoOverlap(c)创新性地通过调度器重排实现对齐 —— 在块到达后立即调度对应tile计算,并通过块内混洗保持缓存局部性。这种 “调度适配而非数据适配” 的思路,彻底消除了数据重排的性能损耗,是实现高效细粒度重叠的关键创新。

如图 6 所示,Syncopate 通过重排波次顺序,确保每个数据块一旦到达即被计算消费,并应用块内 swizzling 以保持局部性访问顺序。这种基于块的 tile 调度纯粹通过调度逻辑实现了计算与通信进度的对齐,无需引入额外的数据重排内核。

六、通信中心的自动调优

块抽象为通信中心的自动调优提供了自然的操作空间。由于块恰好位于全局通信调度与本地 tile 调度器的边界,调整块级参数会同时影响数据在 rank 间的移动方式和每个内核内部的计算顺序。AutoOverlap 暴露了一个高层搜索空间,其控制旋钮直接作用于重叠效率和资源共享。

6.1 块间调优

在块间层面,调优器调整每个逻辑传输的 块大小、形状和分裂因子

  • 大块:通常在拷贝引擎和 TMA 上能实现更高的有效带宽,但会降低重叠的粒度。
  • 小块:能实现更细粒度的流水线,但会带来更高的每块开销和更多的同步操作。

不同算子和模型规模倾向于不同的权衡点。例如,通信密集型的 A2A-GEMM 和 GEMM-AR 受益于能平衡带宽与重叠需求的中等分裂因子。调优器在硬件特定约束下(如拷贝引擎和 TMA 的最小有效传输尺寸、内存对齐规则)搜索该空间,并剪枝违反硬件限制的配置。

6.2 块内调优

在块内层面,调优器同时调整 计算 tile 配置每个块的通信后端实现方式

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 7 | 通信后端选择。(a)通过 Copy Engine 进行通信并配合全局内存信号;(b)由专用 SM 驱动传输;(c)通信与计算共置于同一 SM。

对于固定的逻辑调度,编译器可以将每个传输实例化为图 7 中的任一后端,并在适用时调整分配给通信的 SM 数量。有些调度适合使用 TMA 进行节点内张量移动,而另一些则适合用 Load/Store 处理小的、规约密集型分片。同时,自动调优器探索不同的 tile 尺寸和块内执行顺序,使计算波次更好地与所选块布局对齐,从而提高局部性并避免过长的通信尾部延迟。

6.3 调优空间的可管理性

关键在于,所有这些决策都基于同一块级依赖图。改变通信后端、SM 分配或 tile 执行顺序无需重新推导全局通信调度;AutoOverlap 复用现有调度,并重新生成遵守相同依赖关系的后端特定代码。这种逻辑调度与物理实现的分离,使得搜索空间既丰富又易于管理。后续的消融研究也表明,合理但次优的配置可能轻易导致超过 2 倍的性能损失,而 AutoOverlap 找到的调优配置始终与计算、通信和硬件利用率之间的最佳平衡点保持一致。

七、实验评估

AutoOverlap 在配备 8 块 NVIDIA H100 GPU(通过 NVLink 连接)的服务器上进行了全面评估。测试负载包括代表 LLM 工作负载的 GEMM 变体(AG-GEMM、GEMM-RS、GEMM-AR)和注意力变体(Head-Parallel、Sequence-Parallel、RingAttention),算子形状源自 Llama-3 和 Qwen 模型。

7.1 性能对比

GEMM 算子结果:图 8 显示,AutoOverlap 在所有 GEMM 配置中均达到或接近最佳性能曲线。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 8 | GEMM 算子性能对比。ThunderKittens 仅支持 8 块 GPU;不支持 4 块 GPU 配置。当两种配置均不支持时,相应柱状图省略。该图基于 Llama-3 和 Qwen 模型的真实算子形状,在 4 块和 8 块 H100 GPU 上进行基准测试:对于 AG-GEMM 和 GEMM-RS 等成熟算子,AutoOverlap 性能接近手工优化的 ThunderKittens 和 TritonDistributed(4 块 GPU 平均达最优基准的 99.8%);而在 GEMM-AR 及大模型配置下,其凭借动态块调优和后端适配能力实现反超(8 块 GPU 达 104%)。结果证明自动编译方案可媲美专家手工优化,且具备更优的扩展性。

在常见且高度优化的 AG-GEMM 和 GEMM-RS 场景中,手动优化内核(如 ThunderKittens、TritonDistributed、AsyncTP、Flux)已接近硬件极限。AutoOverlap 基本匹配它们的峰值吞吐量:在 4 GPU 上平均达到最佳基线的 99.8%,在 8 GPU 上达到 104%。对于 GEMM-AR,AutoOverlap 在 7B/8B 形状上略低于 TritonDistributed,但在更大模型配置上扩展性更好并成为领先内核。这表明 AutoOverlap 的通用编译流程能够复现专家手工设计的高度调优重叠模式。

注意力算子结果:图 9 展示了类似趋势。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 9 | AutoOverlap 优化的算子与现有最优基准的性能对比。该图聚焦注意力算子的关键场景:在中等序列长度的 HP-Attention 场景下,AutoOverlap 与手工优化内核性能持平;而随着序列长度增加(达 128K)或采用通信密集型的 Ring-Attention 时,基准内核性能显著下降,AutoOverlap 则通过灵活调整块大小和通信后端,维持高 TFLOPS 输出。这一优势源于其对通信延迟的精准隐藏,尤其适配大模型长上下文训练推理的核心需求。

在标准的 Head-Parallel 注意力下,AutoOverlap 性能紧追最佳手动实现。当问题难度增加(如 RingAttention、长序列、8 GPU 运行)时,AutoOverlap 的优势开始扩大。它能够重塑数据块、重新平衡计算与通信负载,并在序列长度和并行策略变化时选择不同后端,从而在通信最密集的 RingAttention 设置中仍能交付最佳性能。

集成现有编译器结果:图 10 展示了 AutoOverlap 与现有自动分布式编译器(Domino、Alpa、Mercury)的集成效果。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 10 | 对由高层分布式编译器搜索得到的基于分区的中间表示和基于循环的中间表示进行降级评估,并与 AutoOverlap 集成。该图验证了 AutoOverlap 的兼容性与增强能力:将 Domino、Alpa 的分区式 IR 和 Mercury 的循环式 IR 转换为块级调度后,集成 AutoOverlap 的方案在 4 块和 8 块 H100 GPU 上均实现 latency 降低。这表明 AutoOverlap 可作为现有分布式编译器的增强模块,在保留其全局并行策略的基础上,通过细粒度重叠优化进一步挖掘性能潜力,且无需重构原有编译框架。

对于每个编译器,保持其并行化策略和通信调度不变,将其转换为 AutoOverlap 的块表示,然后由 AutoOverlap 生成细粒度重叠内核。在 4 和 8 GPU 配置上,集成 AutoOverlap 始终比这些系统的原生实现降低了端到端算子延迟,说明块级内核内重叠在它们的全局并行化决策之上提供了一个额外的、有效的优化维度。

7.2 消融与敏感性研究

图 11 展示了关键设计选择的消融研究,验证了 AutoOverlap 调优组件的必要性。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 11:AutoOverlap 自动调优设计空间的消融与敏感性研究。(a)通信后端选择 (b)块拆分调优 (c)SM 数量调优 (d)块内tile调度。整体来说,该图通过消融实验量化了各设计维度的重要性:(a)证明 Copy Engine 和 Intra-SM TMA 是最优后端选择,后端不当会导致性能损失超 50%;(b)显示块拆分因子存在最优区间(2-3 拆分,约 128MB),过粗或过细均会影响性能;(c)表明 SM 分配需与模型大小匹配,405B 模型最优 SM 数量高于 70B;(d)则验证tile调度顺序对性能影响显著,最优顺序可使 TFLOPS 翻倍。这些结果凸显了自动调优框架的必要性。

通信后端与 SM 分配
图 11(a) 显示,同一逻辑调度在不同后端的性能差异巨大。对于 GEMM-RS 和 AG-GEMM,拷贝引擎和 SM 内 TMA 后端达到最高 TFLOPS,而纯 CUDA Load/Store 实现的吞吐量则低得多。最佳与最差后端之间的性能差距,堪比图 8 中 AutoOverlap 与某些基线的差距,证实了后端选择的重要性。

图 11(c) 显示,对于固定后端,分配给通信的 SM 数量存在明显“甜蜜点”。分配太少会浪费链路带宽,太多则会饿死主计算内核。最优 SM 数量随模型规模变化,AutoOverlap 能够自动为每个算子和硬件选择近优值。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 8 | GEMM 算子性能对比。ThunderKittens 仅支持 8 块 GPU;不支持 4 块 GPU 配置。当两种配置均不支持时,相应柱状图省略。

块大小调优
图 11(b) 展示了改变分裂因子(块数)对 A2A-GEMM 和 GEMM-AR 性能的影响。
* 较小的分裂因子对应大块,单次传输效率高但重叠机会少;
* 较大的分裂因子增加重叠但带来每块开销。

性能在中等分裂(如 2-3,约 128MB)处达到峰值,而手工方便选择的单一大块或按 rank 划分一次等简单策略远非最优。这直接反映了块级代码生成中的权衡:手动无法可靠选择“好”的块大小,而 AutoOverlap 自动搜索并与算子的计算/通信平衡匹配。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 11:AutoOverlap 自动调优设计空间的消融与敏感性研究。(a)通信后端选择 (b)块拆分调优 (c)SM 数量调优 (d)块内tile调度。

块内 tile 调度
图 11(d) 展示了不同块内 tile 访问顺序带来的性能差异。通过改变 M、N、K 维度上的 tile 尺寸和流水线阶段,保持语义不变但改变局部性和负载平衡。TFLOPS 的广泛分布显示,仅 tile 顺序就能引起超过 2 倍的性能差异,验证了 tile 调度器 swizzling 的重要性。高性能调度集中于那些使 tile 波次与通信块顺序对齐的顺序,而低性能调度则反复访问 tile,破坏局部性并产生长尾未完成工作。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 11:AutoOverlap 自动调优设计空间的消融与敏感性研究。(a)通信后端选择 (b)块拆分调优 (c)SM 数量调优 (d)块内tile调度。

八、结论与展望

AutoOverlap 通过引入 通信块抽象,弥合了通信规划与细粒度重叠之间的抽象鸿沟。它将通信调度与分块的 GPU 计算统一起来,通过自动对齐 tile 调度与通信进度,并在异构后端间选择,将内核内重叠从手工优化转变为通用编译器能力。AutoOverlap 能够与现有分布式编译器无缝集成,补充其全局并行化决策,提供更深入的优化机会。

展望未来,块抽象可以扩展到更多样的通信模式(如稀疏通信、异步集体)、更多硬件平台(如 AMD、Intel GPU),并可能与自动并行化搜索相结合,实现端到端的完全自动化重叠优化。

突破多GPU通信瓶颈:AutoOverlap实现块级细粒度计算-通信重叠,最高加速4.7倍

图 11:AutoOverlap 自动调优设计空间的消融与敏感性研究。(a)通信后端选择 (b)块拆分调优 (c)SM 数量调优 (d)块内tile调度。


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

本文来自网络搜集,不代表鲸林向海立场,如有侵权,联系删除。转载请注明出处:https://www.itsolotime.com/archives/22315

(0)
上一篇 11小时前
下一篇 5小时前

相关推荐