Blog Tags GitHub
Back to blog
communication efficient-attention parallelism-sp

SwiftFusion: Scalable Sequence Parallelism for Distributed Inference of Diffusion Transformers on GPUs

面向扩散Transformer分布式推理的可扩展序列并行引擎

SwiftFusion: Scalable Sequence Parallelism for Distributed Inference of Diffusion Transformers on GPUs

一、论文概述

项目内容
标题SwiftFusion: Scalable Sequence Parallelism for Distributed Inference of Diffusion Transformers on GPUs
作者Jiacheng Yang, Jun Wu, Yaoyao Ding, Zhiying Xu, Yida Wang, Gennady Pekhimenko
机构University of Toronto & Vector Institute, Amazon, NVIDIA
会议ACM CAIS ‘26 (Conference on AI and Agentic Systems)
论文arXiv:2601.20273
领域cs.DC, cs.CV

二、核心思想

问题定义

Diffusion Transformers (DiTs) 在高质量图像和视频生成中广泛应用。随着分辨率和视频时长增加，单GPU推理变得低效。

DiT推理示意图

内存瓶颈：CogVideoX模型权重仅12 GiB，但生成10秒768×1360视频会导致单A100 GPU (40 GiB) OOM。

现有序列并行技术的局限：

技术通信模式问题
Ring Attention环形KV传输通信量不随GPU数增加而减少
Ulysses AttentionAll-to-All不支持与计算重叠，受头数限制
USP (统一SP)Ring (机间) + Ulysses (机内)机间通信开销大，可扩展性差

SP技术对比

三大挑战

带宽差异与延迟分解

挑战说明影响
机间/机内带宽差异NVSwitch带宽远高于InfiniBand/EFAUSP使用Ring进行机间通信，通信量不减少
All-to-All不可重叠现有实现将All-to-All视为原子操作机间通信成为瓶颈
双侧通信开销NCCL等库需要严格同步同步和计算开销随GPU数增加

同步开销

解决方案概述

SwiftFusion提出三个关键创新：

创新说明效果
拓扑感知序列并行Ulysses用于机间，Ring用于机内减少机间通信量
Torus Attention分块All-to-All与计算重叠隐藏机间通信延迟
单侧通信实现使用NVSHMEM替代NCCL减少同步和计算开销

三、技术架构

拓扑感知通信调度

拓扑感知SP

核心思想：将N×M个GPU组织为2D设备网格，形状为Pu×PrP_u × P_r。

  • Ulysses Attention (PuP_u)：用于机间通信，通信量随机器数增加而减少
  • Ring Attention (PrP_r)：用于机内通信，利用NVSwitch高带宽

并行度计算：

  • 当H=NH = N时：Pu=NP_u = N，Pr=MP_r = M
  • 当H≠NH \neq N时：Pu=gcd⁡(NM,H)P_u = \gcd(NM, H)，Pr=NM/PuP_r = NM/P_u

与USP对比：SwiftFusion在大多数情况下产生更少的机间通信量（详见Appendix D证明）。

Torus Attention

Torus Attention

核心观察：All-to-All操作前后存在静态元素（当前rank的head索引元素），可直接用于计算。

分阶段执行：

1. Pull Q阶段 (NN个阶段)：

  • 阶段kk：GPU tt使用Q(t−k+1)%N,tQ_{(t-k+1)\%N,t}计算注意力，同时发送Qt,(t+k)%NQ_{t,(t+k)\%N}
  • 通信与计算重叠

2. Pull KV阶段 (N−1N-1个阶段)：

  • 阶段kk：GPU tt使用接收到的QQ和K(t−k)%N,tK_{(t-k)\%N,t}、V(t−k)%N,tV_{(t-k)\%N,t}计算注意力
  • 同时发送下一轮需要的KV

3. Push O阶段 (N−1N-1个阶段)：

  • 计算Ot,tO_{t,t}时，同时发送Oi,tO_{i,t}到其他GPU
  • 输出张量通信与计算重叠

单侧通信实现

Algorithm 1：SwiftFusion单侧通信伪代码

function SwiftFusion(Q, K, V, T, t, U, u, R, r):
    # 1. Ulysses Attention: 机内All-to-All
    ScatterPush({Q_t,:, K_t,:, V_t,:}, ..., (t, :, r))
    BarrierAll()

    # 2. 预取所有Pull操作
    for k = 1 to N-1:
        E_k^Q = GatherPull(Q_t':, ..., (t', :, r))
        E_k^KV = GatherPull(K_t':, V_t':, ..., (t', :, r))

    # 3. Pull Q阶段: 重叠通信与计算
    RingAttn(Q_t,:, K_t,:, V_t,:)
    for k = 1 to N-1:
        Wait(E_k^Q)
        RingAttn(Q_t':, K_t,:, V_t,:)

    # 4. Pull KV阶段
    for k = 1 to N-1:
        Wait(E_k^KV); Barrier(R)
        RingAttn(Q:{\t},:, K_t':, V_t':)

    # 5. Push O阶段
    for k = 1 to N-1:
        ScatterPush(O_t':, ..., (t', :, r))
    RingAttn(Q_t,:, K:{\t},:, V:{\t},:)

    return O

关键优化：

  • 预取所有Pull操作，最大化重叠
  • 仅需机内同步+两次机间同步（层开始和结束）
  • 最多一个QQ、KK、VV、OO的缓冲副本

四、核心创新

创新点说明理论/实验依据
拓扑感知调度Ulysses机间 + Ring机内机间通信量减少近NN倍
Torus Attention分块All-to-All重叠隐藏机间通信延迟
单侧通信NVSHMEM替代NCCL消除每传输同步开销
静态元素利用All-to-All前后不变元素直接开始计算

五、实验结果

实验设置

配置详情
硬件4× AWS p4de.24xlarge, 每台8× A100 (40GB), NVSwitch互联
网络AWS EFA 400 Gbps
软件CUDA 12.8.0, PyTorch 2.8.0, NCCL 2.27.3, NVSHMEM 3.4.5
模型Flux (12B, 图像), CogVideoX (5B, 视频)
工作负载3072×3072/4096×4096图像, 20s/40s 768×1360视频
基线USP (Unified Sequence Parallelism)

端到端性能

端到端性能

最优分布式配置：

对比加速比说明
SwiftFusion vs USP1.35× 平均 (最高1.77×)所有关键创新协同工作
TAS vs USP1.27× 平均 (最高1.64×)仅拓扑感知调度

关键发现：

  • 2台机器时TAS性能不如USP（通信量相同但不可重叠）
  • 3-4台机器时SwiftFusion优势明显
  • 内存消耗无额外开销

其他分布式配置

各种配置

方法平均加速比最高加速比
TAS vs USP1.47×2.54×
SFU vs USP1.61×3.11×

观察：

  • 增加Ulysses度数UU通常可获得更好加速
  • TAS中U12R2U12R2优于U24R1U24R1（避免过大All-to-All开销）

逐层性能分析

不同头维度 (DD)：

头维度平均加速比
D=32D=321.12×
D=64D=641.28×
D=128D=1281.32×

不同序列长度：

序列长度平均加速比
96K1.59×
128K1.27×
160K1.07×
192K1.14×

不同批量大小：

批量大小平均加速比
所有1.13× - 1.27×

趋势：序列越长，计算量二次增长，通信优化收益相对减少。

六、相关工作

方法类别代表方法特点与SwiftFusion的区别
分布式DiT推理USP, DistriFusion, PipeFusionRing/Ulysses组合或有损通信隐藏未优化拓扑感知和重叠
通信-计算重叠FLUX, DeepEP, Comet, ScaleFusion手动或编译器优化针对LLM或特定架构，非通用DiT
编译器方法TileLink, Triton-Distributed, Mercury自动重叠难以发现Torus Attention等专用技术

七、总结

核心贡献

  1. 拓扑感知序列并行：首次提出Ulysses机间 + Ring机内的调度策略
  2. Torus Attention：新型分布式注意力算法，实现All-to-All与计算重叠
  3. 单侧通信统一实现：使用NVSHMEM统一Torus/Ulysses/Ring Attention
  4. 显著性能提升：平均1.35×加速（最高1.77×）

技术影响

  • 可扩展性：支持多机多卡高效扩展
  • 内存效率：无额外内存开销
  • 通用性：适用于图像和视频生成任务
  • 生产就绪：基于PyTorch和NVSHMEM实现

局限性

  1. 2机器场景：拓扑感知调度在2机器时无优势
  2. 头数限制：Ulysses Attention需要HH可被PuP_u整除
  3. 序列长度：超长序列（>160K）时计算主导，收益减少
  4. 硬件依赖：需要NVSwitch和NVSHMEM支持

八、参考资源

  • 论文：arXiv:2601.20273
  • 会议：ACM CAIS ‘26
  • 相关项目：USP, FlashAttention, NVSHMEM
  • 代码库：未公开