ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels
Stuart H. Sul, Simran Arora, Benjamin F. Spector, Christopher Ré
cs.DC, cs.LG
2025-11-18
ParallelKittens 把多 GPU 内核的重叠抽象成 8 条原语和一套模板,几十行设备代码在 Hopper 和 Blackwell 上跑出最高 4.08 倍加速,通信空窗压到 1%。
模型越做越大,跑在多张卡上时,卡与卡之间的数据搬运反而成了瓶颈。从 A100 到 B200,BF16 张量核心的算力涨了 7.2 倍,HBM 带宽涨了 5.1 倍,可卡间 NVLink 只快了 3 倍,跨节点的 PCIe 和 InfiniBand 更只快了 2 倍。在大模型负载里,通信吃掉一半以上的运行时间并不稀奇,算力在那干等。
标准解法是让通信和计算重叠:一张卡在等数据的时候,先用已经在卡上的数据做计算。但现有的重叠方案各有硬伤。手工内核(Flux、FlashDMoE、CUTLASS)能把单个算子压到峰值,代价是代码高度定制,依赖一堆底层原语。编译器路线(Triton Distributed、TileLink)有时生成的内核比不重叠的基线还慢,而且换一代硬件就水土不服。现成库(NCCL、NVSHMEM)要双向同步和中间缓存,比手工调优慢最多 4.08 倍。ParallelKittens(下称 PK)想问的是:与其每个算子都重新写一套,能不能用一小撮可复用的原则系统性地指导多 GPU 内核的设计。
PK 建在 ThunderKittens 之上。ThunderKittens 本来是个单卡、基于 tile 的 CUDA DSL,PK 给它加了 8 条新原语,全部按 tile 粒度(16×16 到约 256×256)运作,而且由设备端主动发起,不经 host。这 8 条原语分三组:点对点 tile 传输(storeasync、storeaddasync)、利用网络内归约加速的集合通信(reduce、allreduce)、以及靠 multicast 同步的 barrier(signal、signalall、wait、barrier)。
原语背后是三条设计原则。
第一,按消息大小挑传输机制。Copy engine 能跑到峰值带宽的约 81%,但只有消息大到 256 MB 以上才吃得满;TMA(Hopper 的张量内存加速单元)从 2 KB 就能跑到约 74%;寄存器级指令 128 B 也很高效,但要占用 3 到 5 倍的 SM 才能跑满,而且它是唯一支持网络内归约的机制。
第二,按粒度选调度方式。SM 内重叠在同一个 SM 里把线程切成计算池和通信池,同步开销约 64 ns;SM 间重叠把整块 SM 专用化,能解锁网络内归约和更好的 L2 复用,但同步开销约 832 ns。怎么选取决于计算和通信的粒度能不能对得上。
第三,把设计开销砍掉。NCCL 那套要双向同步,PK 改成单向写进预分配的缓冲;不做中间拷贝,访问对端内存时不加多余的 syncthreads。
原语之上是一套编程模板(Load-Compute-Store-Communicate),配四个 worker:loader、storer、consumer、communicator。模板自动处理共享内存布局、barrier 设置和 SM/warp 划分。开发者描述「搬什么、算什么」,模板把重叠的接线接好。
在 8×H100(Hopper,第四代 NVLink 450 GB/s)和 B200(Blackwell,第五代 900 GB/s)上验证,CUDA 12.6。
| 工作负载 | 对照基线 | PK 加速 |
| 数据/张量并行 GEMM+通信 | Flux、CUTLASS、Triton Distributed | 最高 2.33 倍 |
| 序列并行 Ring Attention | xDiT | 1.07–4.08 倍 |
| DeepSpeed-Ulysses | YunChang | 1.01–1.39 倍 |
| 专家并行 MoE dispatch | Comet | 0.92–1.22 倍 |
当规约轴足够大时,数据并行内核里不重叠的通信时间压到总时间的 1%;序列并行的 Ring Attention 压到 9%。每个内核的通信部分都是在原单卡内核之外加了不到 50 行设备代码(MoE 那个更少,不到 40 行)。在非连续内存布局上,PK 对细粒度集合通信也跑赢 NCCL,因为 NCCL 对不连续的分块要先重排再拷。
需要诚实指出:加速幅度是区间,因为高度依赖问题大小和哪个基线算公平。MoE 上 PK 和 Comet 基本打平(0.92–1.22 倍),个别 GEMM 配置下比 Flux 还略慢。这框架的价值是让短程序够到帕累托前沿,不是对每个基线的全面碾压。
跑多卡训练或推理的团队大多不写自己的集合内核,而是直接调 NCCL、把开销吃下去。PK 论证的是:NCCL 之上那一层(手工的重叠内核)可以做得便宜,几十行代码加一个已经把硬接线都处理好的模板。对于占延迟大头的算子(attention、MoE dispatch、GEMM 前后的 all-gather/reduce-scatter),GPU 的空闲时间就藏在这儿。这套框架瞄准的是节点内扩展(NVL72→NVL144→NVL576),正好是英伟达在堆密度的方向,所以相关性是当下的。不过它是个扩展 ThunderKittens 的研究框架,不是拿来即用的库,对从业者最直接的收获是那三条原则,不是一行能粘进 Makefile 的东西。
最大的局限作者说得直接:PK 只管节点内通信。多节点集群还得靠 MPI/NCCL 走 fabric,跨节点扩展被列为未来工作。
此外,所有加速都是区间值,低的那头(对 Comet 0.92 倍、对 YunChang 1.01 倍)意味着 PK 并不严格强于现有的手工内核,有些配置打平甚至略输。验证只覆盖 Hopper 和 Blackwell,AMD MI 系列和 TPU 没碰,而且三条原则相当依赖英伟达特有的硬件(TMA、copy engine、NVSwitch multicast)。PK 还要求开发者熟悉 ThunderKittens 这个本身相当底层的 DSL,所谓「50 行」只算设备代码,不算周边仍要写的脚手架。