Every Kernel Is a Join: Automatic Multi-GPU Parallelism for AI Computations in Einsummable

2026-09-03

AI Summary

本文提出Einsummable系统，自动将AI计算分配到多GPU服务器上，通过关系连接和聚合方法优化通信成本，无需手动配置。

Abstract

Distributing an AI computation across the GPUs of a multi-GPU server is one of the central problems in systems-for-AI. We present Einsummable, a prototype system that accepts a PyTorch-like description of an AI computation and automatically distributes it across a multi-GPU server, with no device assignments, sharding annotations, or communication operations written by the programmer. Einsummable models every operation as a relational join followed by an aggregation over tensor relations, in which the tuples contain sub-tensors. Each operation exposes its possible decompositions through what we call "join-agg specs". An optimizer then selects decompositions across the whole computation to minimize a communication-cost proxy. Because it searches decompositions rather than a menu of named strategies, Einsummable discovers plans that mesh-based auto-parallelizers cannot. Each decomposed operation is implemented by synthesizing an exchange program, which is a topology-aware generalization of Volcano's exchange operator. Einsummable invokes no canned collectives: all communication and aggregation is special-purpose, derived at compile time. Despite being fully automatic, Einsummable can outperform custom-designed implementations. For example, on LLaMA transformer blocks on an eight-GPU A100 server, Einsummable achieves a geometric-mean runtime of 8.97 ms, versus 13.80 ms for hand-tuned PyTorch and 15.90 ms for vLLM.

Problem

• multi-GPU
• AI computations
• automatic distribution
• no device assignments
• sharding annotations

Innovation

• Automatic Multi-GPU Parallelism
• Join-Agg Specs
• Communication-Cost Proxy
• Exchange Program Synthesis