@ycombinator:在我们的最新 YC Paper Club 上,研究人员和开发者介绍了多 GPU 内核、每瓦特智能、异构……
摘要
Y Combinator 举办了一场 Paper Club,研究人员展示了多 GPU 内核优化的创新,包括 ParallelKittens,这是一个简化重叠多 GPU 内核开发并在多种工作负载下实现显著加速的 CUDA 框架。
查看缓存全文
缓存时间: 2026/07/30 03:46
在我们的最新一期 YC Paper Club 中,研究人员和开发者围绕多GPU内核、每瓦特智能、异构推理等主题展开了分享。感谢以下演讲者:0:00 – @FrancoisChauba1:芯片与内核专业化的理由 7:16 – @stuart_sul:Parallel Kittens - 多GPU AI内核的系统化与实用简化 (https://arxiv.org/abs/2511.13940) 21:29 – @JonSaadFalcon:每瓦特智能 - 衡量本地与云端AI的智能效率 (https://arxiv.org/abs/2511.07885) 31:05 – @MarkSaroufim:当AI开始编写系统代码 47:04 – Misha Smelyanskiy:为什么AI推理需要异构硬件 1:04:33 – @shacklettbp:构建一个完全在GPU上运行的高吞吐量游戏引擎 (https://madrona-engine.github.io/shacklett_siggraph23.pdf…) — # 多GPU AI内核的系统化与实用简化 来源:https://arxiv.org/html/2511.13940 ###### 摘要 随着模型规模的扩大以及硬件计算吞吐量的提升远超互连带宽的改进,GPU间通信已成为现代AI工作负载的主要瓶颈。现有系统通过计算与通信重叠来缓解这一问题,但在异构工作负载和新加速器上往往无法达到理论峰值性能。我们不采用特定于算子技术,而是探究是否可以用一小套简单、可复用的原则系统性地指导最优多GPU内核的设计。我们提出ParallelKittens(PK),一个极简的CUDA框架,极大简化了重叠式多GPU内核的开发。PK扩展了ThunderKittens框架,并通过八个核心原语和一个统一编程模板体现了多GPU内核设计的原则,这些原则源于对影响多GPU性能的因素(数据传输机制、资源调度和设计开销)的全面分析。我们在Hopper和Blackwell架构上验证了PK。使用不到50行设备代码,PK在数据并行和张量并行工作负载上实现了高达2.33倍的加速,序列并行工作负载上实现4.08倍加速,专家并行工作负载上实现1.22倍加速。 ## 1 引言 见图1:我们研究了高性能多GPU内核的原则,并引入了ParallelKittens(PK),一个封装这些原则的、有主见的编程原语集合。左侧显示了GPU内存层次结构及相应的PK抽象(第3.2.1节),右侧显示了PK程序模板及其关键的多GPU内核组件(第3.2.3节)。几年前,GPU计算利用率通常受GPU内部内存访问的限制。然而,像FlashAttention[4]这样的IO感知算法、支持算子高效映射到硬件的领域特定语言(DSL)[29,18,27]以及AI模型的持续扩展,使得GPU间通信成为主要的剩余瓶颈。即使使用高速互连如NVLink[15]和计算友好阶段如预填充,通信在大型语言模型(LLM)工作负载中也可能占用超过50%的执行时间,导致GPU计算空闲[3]。问题因通信硬件的相对缓慢改进而加剧:从Nvidia A100[17]到B200[20],BF16张量核心性能提升了7.2倍,高带宽内存(HBM)带宽提升了5.1倍,而节点内通信(NVLink)仅提升了3倍,节点间通信(PCIe/InfiniBand)仅提升了2倍。为缓解通信开销,先前方法将GPU间通信与常见算子(如通用矩阵乘法(GEMM)、注意力机制和混合专家(MoE)层)的GPU内计算重叠[13,3,32,35,31,1]。这些方法减少了数据并行、张量并行、序列并行和专家并行[25,12]中的非重叠通信时间,这些是跨多个GPU分布行业规模训练和推理的常见策略。然而,先前工作要么(i)依赖于针对特定AI算子的定制内核和复杂的低级原语(例如CUTLASS、NVSHMEM、Linux IPC),要么(ii)采用编译器方法,但无法适应新加速器——偶尔生成的内核甚至比非重叠基线更慢,或者(iii)使用现成库,导致性能比手工调优实现慢多达4.08倍。随着硬件向统一多GPU系统转变——以Nvidia从NVL72到NVL144(2026年)和NVL576(2027年)的路线图为例[21]——我们需要简单、通用的原则和编程原语来实现峰值性能的多GPU操作。在这项工作中,我们确定了设计高效多GPU内核的三个关键原则,并详细分析了每个原则(第3.1节)。1. 传输机制。GPU间网络依赖三种机制——拷贝引擎、张量内存加速器(TMA)和寄存器级指令——它们在最大带宽、有效消息粒度、支持的功能和计算占用率方面有所不同。理解这些权衡并选择正确的机制对于峰值性能至关重要。例如,拷贝引擎达到最高效率(理论最大值的81%),但需要大消息(≥256 MB)才能饱和。TMA仅用2 KB消息即可达到近峰值吞吐量(74%)(图2)。寄存器级指令在128 B粒度下高效运行,但需要约76个流多处理器(SM)才能饱和带宽(70%),而TMA只需15个(图3)。然而,只有寄存器级指令支持网内归约。现有系统未能捕捉这些权衡;例如,Triton Distributed、Flux和CUTLASS依赖拷贝引擎进行节点内all-gather GEMM,在较小矩阵尺寸上变得比非重叠基线更慢(图7)。2. 调度。计算和通信工作在各SM之间的分配必须根据工作负载特征进行选择。我们将SM间和SM内重叠确定为两种主要调度策略,在计算利用率和通信多功能性之间进行权衡。当计算和通信粒度对齐时,SM内重叠是首选;例如,在GEMM reduce-scatter中,SM内重叠比SM间重叠性能提高1.2倍。相比之下,SM间重叠可以实现通信模式,从而显著减少传输大小。例如,通过SM间重叠利用网内归约,GEMM all-reduce实现了3.62倍的性能提升(图5),而all-gather GEMM实现了1.57倍的性能提升。没有先前工作探索过两种调度策略;现有方法要么依赖单一类型,要么完全省略设备侧重叠,从而无法泛化(例如,将Flux的SM内重叠设计应用于GEMM all-reduce会导致上述减速)。3. 设计开销。广泛使用的通信库(例如NCCL、NVSHMEM)封装了设计选择——特别是在同步和缓冲方面——这些选择偏向简单性而非性能。我们表明,先前库中的选择可能导致纯通信内核(例如all-reduce)性能损失超过1.7倍,通信延迟高达4.5倍。通过采用允许用户显式控制内存分配和同步的设计,这些开销可以大幅减少。基于这些见解,我们引入ParallelKittens(PK),一个扩展了ThunderKittens(TK)框架[27]的、有主见的C++嵌入式编程原语集合(第3.2节)。PK仅暴露每种功能最高效的传输机制(例如,点对点通信使用TMA,网内加速使用寄存器操作),提供极简的同步原语和一个通用程序模板,简化实现SM间和SM内重叠调度,并提供对性能关键组件(例如NVLink传输)的完全控制,同时抽象掉非必需的多GPU复杂性(例如进程间通信和虚拟内存交换)。我们在Hopper和Blackwell架构上的多种并行AI工作负载(包括数据、张量、序列和专家并行,即融合并行GEMM、分布式注意力变体和MoE)上验证了PK。与最强基线相比,PK在数据并行和张量并行工作负载上实现了高达2.33倍的计算吞吐量(FLOP/s),序列并行工作负载上实现4.08倍,专家并行工作负载上实现1.22倍,有效将非重叠通信时间分别降至1%、9%和15%。PK匹配了最强手工优化内核(Flux、Comet、CUTLASS)的性能,在不同问题规模上比基于编译器的方法(Triton Distributed)快1.07–5.63倍,比基于通信库的方法(xDiT、YunChang)快1.01–4.08倍。每个PK内核在原始单GPU GEMM或注意力内核之外,仅需不到50行额外设备代码。PK的完整实现(包括其内核)已完全开源,目前正在Cursor用于大规模内部训练。总之,我们的贡献包括:- •对多GPU编程的详细分析,将性能分解为可解释的因素(传输机制、调度策略和设计开销),并通过微基准测试验证每个因素。- •ParallelKittens,一个极简的多GPU原语集合和一个统一编程模板,扩展了熟悉的ThunderKittens框架。- •使用ParallelKittens构建的内核,匹配或超越手工优化内核的性能,同时大幅降低代码复杂度。 ## 2 背景 在本节中,我们提供现代数据中心级GPU的背景,并回顾先前在多GPU AI内核优化方面的工作。111除非另有说明,我们以Nvidia HGX H100[30]平台(配备8×H100 80GB SXM GPU、第4代NVLink/NVSwitch和第5代PCIe)作为运行示例;然而,这些原则可扩展到其他现代平台(例如Blackwell架构)和硬件供应商(例如AMD)。 ### 2.1 GPU架构 GPU内核从HBM加载数据,执行计算,并将结果写回HBM。多GPU内核将工作负载分布到多个GPU上,并访问所有设备的内存。##### GPU层次结构。GPU内核在数百个流多处理器(SM)上并行执行数万个硬件线程。远离SM的内存提供更大容量,但延迟更高。每个SM包含64 KB寄存器,专用于单个线程,每个时钟周期可访问。线程组织为线程块,每个线程块在单个分配的SM上执行。线程块中的线程通过227 KB共享内存(SMEM)通信,这是一种每SM的片上SRAM,提供高达33 TB/s的带宽。所有线程共享一个50 MB L2缓存(约12 TB/s),连接到80 GB HBM(3 TB/s)。线程还可以通过NVLink(450 GB/s单向)访问对等GPU HBM,从而实现多GPU内核开发。##### GPU网络。多GPU系统依赖层次化的互连。PCIe(64 GB/s)是CPU到GPU(例如内核启动、主机发起传输)和多节点通信(通过InfiniBand/TCP)的通道。NVLink(450 GB/s)提供GPU之间和到NVSwitch的点对点连接;NVSwitch将所有NVLink端点互连成非阻塞结构,实现GPU到GPU的完全通信。NVSwitch还支持网内、设备外加速,用于多播和归约。除非另有说明,本文中的所有GPU间通信均通过NVLink/NVSwitch进行。##### 执行重叠。GPU包含各种执行单元,专门用于不同的计算、内存和通信操作。计算方面,张量核心执行分块矩阵乘法,而CUDA核心处理逐元素算术。内存方面,张量内存加速器(TMA)在SMEM和HBM之间执行批量数据传输,可由单个线程异步调用。另外,每个GPU的拷贝引擎(专用DMA单元)独立于SM移动大型连续设备内存区域,由主机调用。在SM内,线程可以同时向不同执行单元发出指令。因此,实现最佳性能依赖于有效地重叠它们的使用,以隐藏非关键操作并最大化关键操作的吞吐量。我们区分SM间重叠(整个SM几乎完全专用于计算、内存或通信任务)和SM内重叠(同一SM内的不同线程束或线程同时驱动计算、内存或GPU间流量)。这些资源以不同速率饱和,为各种重叠策略创造了机会。 ### 2.2 相关工作 我们受到大量加速多GPU AI工作负载的工作的启发。##### 算子特定内核。许多先前工作通过重叠计算和通信手工调优特定AI算子,例如TP-Async[9]、Flux[3]、Ring Attention[13]、DeepEP[5]、Comet[31]、FlashDMoE[1]以及来自CUTLASS[28]的一些分布式GEMM内核。这些方法采用的技术范围从将主机触发的拷贝与设备内核重叠,到高度优化的设备端调度器和设备发起的通信。虽然这些系统为特定目标提供了强大性能,但它们要求复杂的实现,并且提供的可复用抽象有限。例如,FlashDMoE仅针对TF32精度优化,BF16/FP16支持在发布五个月后仍在开发中。相比之下,PK提炼出适用于各种工作负载的通用原则,实现了具有竞争力的速度。
相似文章
@levidiamode: GPU编程的第163/365天 - 今天看几个不同的agentic GPU内核优化系统。我最感兴趣的两个是…
一条推文讨论了两种agentic GPU内核优化系统:@dogacel0的Auto GPU Kernel和@songhan_mit实验室的Kernel Design Agents,两者均在MLSys Sparse Attention FlashInfer比赛中获胜。该帖子突出了使用子代理和Claude技能进行GPU编程的不同方法。
并行编程的禅意:内核的姿态
一篇博客文章,通过HipKittens论文的视角探索并行编程概念,重点介绍了AMD GPU上重叠计算与内存移动的八波乒乓调度,并与禅宗原则进行了哲学类比。
一个可定制的编译器,用于为AI模型生成高效的融合GPU内核 [P]
作者介绍了一款用 Python 编写、高度可定制且易于修改的 ML 编译器。该编译器通过多级 IR 流水线将 LLMs 转换为优化的 CUDA 内核,在特定操作上实现了与 PyTorch 相当甚至更优的性能。文章详细阐述了该编译器的优化过程、降级规则以及用于生成高效融合 GPU 内核的 CLI 用法。
最后,衷心感谢这个了不起的团队:@jcz42, Arjun, Driss, @tensorcore, @yoonrkim 和 @tri_dao!PDF: https://a…
CODA 引入了一种 GPU 内核抽象,将 transformer 计算重写为 GEMM-plus-epilogue 程序,减少内存受限操作,提高训练效率。
@lauriewired:真有意思。三星几天前几乎悄无声息地发布了一篇论文:在CXL内存池上卸载KV缓存。R…
三星发表了一篇关于通过CXL内存池卸载KV缓存的论文,展示了即使使用早期版本的CXL硬件,GPU也能像使用真实DRAM一样高效地获取数据,使得该方法易于复现。