Einsummable: Automatic Multi-GPU Parallelism via Join-Agg Specs
Authors: Zhimin Ding, Chen-Kuan Liao, Chima Adiole, Brianna Barrow, Fangzhou Du, Yu Hsiao, Ge Huang, Yicheng Jin, Ismail Syed, Chris Jermaine Affiliation: Rice University arXiv: 2609.03905（2026-09-04，cs.DC；PVLDB 风格预印） Venue: 预印本 / VLDB 风格 PDF: arXiv PDF
单机多 GPU（NVLink 域可达 72）上，intra-operator 并行常靠手写 TP/SP/DP 或 mesh 注解自动并行。Einsummable 把每个算子建模为 张量关系上的 join + aggregation，用 join-agg specs 枚举合法分解，再用通信字节代理做 DP 选计划，最后合成 exchange program（拓扑感知，不依赖罐头 NCCL）。对照 LLM Distributed Training Collectives：集体不是库调用，而是编译出来的专用交换。
动机
- FlexFlow / GSPMD / Alpa 等在固定 mesh 词汇上搜；表达不了 3D matmul、packed-sequence 边界切分、GQA 头继承等。
- 通信开销可使「加卡变慢」；需要逻辑分解与物理路由分离。
- 目标：程序员写 PyTorch 风格图，编译器自动选分解与数据移动。
方案
- 抽象算子 → join-agg specs： 每个算子按可切维导出 (p) 个 function specs（(p=) GPU 数）。
- 逻辑优化： 按通信代理（repartition + join 输入量 + aggregation 副本）DP 选整图分解；可恢复经典 3D matmul。
- 物理优化： 为每个算子合成 ExProg（聚合+多播+中继），硬件模拟器成本；编成 CUDA graph。
- 内核： cuBLAS / FlashAttention 风格 / Triton 元素与分解核。
效果（仅论文数字）
平台： DGX A100 8×40GB NVSwitch；另测 DGX V100 cube-mesh。FP16；warmup 2 + timed 10。
| 工作负载 | Einsummable | 对照 |
|---|---|---|
| LLaMA-scale transformer block，5 负载几何均值（8 GPU） | 8.97 ms | PyTorch 13.65；vLLM 14.87 |
| 单序列 128K tokens | 143 ms | PyTorch 506；vLLM 498（约 3.5×） |
| 矩阵链 A–E | 与 JAX/PyTorch 基本持平（差约 ≤6%） | — |
| 随机计划 vs 优化器（transformer） | 中位随机 22.8 ms（约 2.5× 更差） | — |
| V100 优化 vs naive exchange | transformer 几何均值约 +5.6%；bushy chain 最高 1.41× | A100 上几乎无差 |
128K 计划叙述：自动得到类似 DeepSpeed-Ulysses（attention 内按 KV 头切、外围按 token 切）。通信代理与 runtime Pearson r 多在 0.72–0.92。
与 wiki 的关系
- LLM Distributed Training Collectives — 专用交换 vs 罐头 AllReduce/A2A
- BASP — 手写 Ulysses 子组优化；本文自动发现同类分解
- NVLink NVSwitch Scale-Up Fabric — 单机 NVSwitch 域是默认假设
- Synchronization Tax — 集体墙钟含同步税；本文用模拟器估交换时间
- Parallelism Transition Point — 推理侧 DP/TP 切换直觉对照
开放问题
- 多机慢互联下是否仍以 intra-op 为主，还是必须接 pipeline？
- 计算代价（形状敏感 kernel）进逻辑成本模型后排名是否变化？
Citations
[1] arXiv PDF — Ding et al., arXiv:2609.03905 [2] einsummable-multi-gpu-parallelism.md — ingest stub