arXiv:2511.13940cs.DCcs.LG2025-11被引 10

用8个基础模块简化多GPU算子开发,提升训练效率

ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels

  • 基于8个核心原语构建统一编程模板,指导多GPU算子设计
  • 在Hopper/Blackwell上实现最高4.08倍加速,仅需50行代码
  • 适合需要高效跨GPU通信的模型训练开发者

随着模型规模扩大,跨GPU通信已成为现代AI工作负载的主要瓶颈。现有系统虽采用计算-通信重叠缓解此问题,但在异构负载和新型加速器上难以达到理论峰值性能。本文提出ParallelKittens(PK),一个极简的CUDA框架,通过八种核心原语和统一编程模板,系统性地指导多GPU算子设计。该框架基于对数据传输机制、资源调度和设计开销的全面分析,可显著降低开发复杂度。在Hopper与Blackwell架构上验证表明,仅用不足50行设备代码,PK在数据并行和张量并行任务中最高提速2.33倍,在序列并行任务中达4.08倍,在专家并行任务中提速1.22倍。

原文摘要 · Abstract (English)

Inter-GPU communication has become a major bottleneck for modern AI workloads as models scale and improvements in hardware compute throughput outpace improvements in interconnect bandwidth. Existing systems mitigate this through compute-communication overlap but often fail to meet theoretical peak performance across heterogeneous workloads and new accelerators. Instead of operator-specific techniques, we ask whether a small set of simple, reusable principles can systematically guide the design of optimal multi-GPU kernels. We present ParallelKittens (PK), a minimal CUDA framework that drastically simplifies the development of overlapped multi-GPU kernels. PK extends the ThunderKittens framework and embodies the principles of multi-GPU kernel design through eight core primitives and a unified programming template, derived from a comprehensive analysis of the factors that govern multi-GPU performance$\unicode{x2014}$data-transfer mechanisms, resource scheduling, and design overheads. We validate PK on both Hopper and Blackwell architectures. With fewer than 50 lines of device code, PK achieves up to $2.33 \times$ speedup for data- and tensor-parallel workloads, $4.08 \times$ for sequence-parallel workloads, and $1.22 \times$ for expert-parallel workloads.

多GPUCUDA并行计算算子优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。