Table of Contents
Fetching ...

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

Jiacheng Yang, Jun Wu, Yaoyao Ding, Zhiying Xu, Yida Wang, Gennady Pekhimenko

TL;DR

This work targets scalable distributed inference for Diffusion Transformers (DiTs), where activation sizes scale with sequence length and single-GPU inference becomes impractical. StreamFusion delivers a topology-aware sequence parallelism framework that combines inter-machine Ulysses Attention with intra-machine Ring Attention, introduces Torus Attention to overlap inter-machine all-to-all transfers with computation, and uses one-sided NVSHMEM-based communication to minimize synchronization overhead. The approach yields consistent latency reductions, achieving an average speedup of $1.35\\times$ (up to $1.77\\times$) over the previous best USP across image and video workloads, with further gains depending on workload and configuration. Overall, StreamFusion enables scalable, low-latency DiT inference on multi-GPU clusters, accelerating high-resolution image synthesis and long-duration video generation in distributed settings.

Abstract

Diffusion Transformers (DiTs) have gained increasing adoption in high-quality image and video generation. As demand for higher-resolution images and longer videos increases, single-GPU inference becomes inefficient due to increased latency and large activation sizes. Current frameworks employ sequence parallelism (SP) techniques such as Ulysses Attention and Ring Attention to scale inference. However, these implementations have three primary limitations: (1) suboptimal communication patterns for network topologies on modern GPU machines, (2) latency bottlenecks from all-to-all operations in inter-machine communication, and (3) GPU sender-receiver synchronization and computation overheads from using two-sided communication libraries. To address these issues, we present StreamFusion, a topology-aware efficient DiT serving engine. StreamFusion incorporates three key innovations: (1) a topology-aware sequence parallelism technique that accounts for inter- and intra-machine bandwidth differences, (2) Torus Attention, a novel SP technique enabling overlapping of inter-machine all-to-all operations with computation, and (3) a one-sided communication implementation that minimizes GPU sender-receiver synchronization and computation overheads. Our experiments demonstrate that StreamFusion outperforms the state-of-the-art approach by an average of $1.35\times$ (up to $1.77\times$).

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

TL;DR

This work targets scalable distributed inference for Diffusion Transformers (DiTs), where activation sizes scale with sequence length and single-GPU inference becomes impractical. StreamFusion delivers a topology-aware sequence parallelism framework that combines inter-machine Ulysses Attention with intra-machine Ring Attention, introduces Torus Attention to overlap inter-machine all-to-all transfers with computation, and uses one-sided NVSHMEM-based communication to minimize synchronization overhead. The approach yields consistent latency reductions, achieving an average speedup of (up to ) over the previous best USP across image and video workloads, with further gains depending on workload and configuration. Overall, StreamFusion enables scalable, low-latency DiT inference on multi-GPU clusters, accelerating high-resolution image synthesis and long-duration video generation in distributed settings.

Abstract

Diffusion Transformers (DiTs) have gained increasing adoption in high-quality image and video generation. As demand for higher-resolution images and longer videos increases, single-GPU inference becomes inefficient due to increased latency and large activation sizes. Current frameworks employ sequence parallelism (SP) techniques such as Ulysses Attention and Ring Attention to scale inference. However, these implementations have three primary limitations: (1) suboptimal communication patterns for network topologies on modern GPU machines, (2) latency bottlenecks from all-to-all operations in inter-machine communication, and (3) GPU sender-receiver synchronization and computation overheads from using two-sided communication libraries. To address these issues, we present StreamFusion, a topology-aware efficient DiT serving engine. StreamFusion incorporates three key innovations: (1) a topology-aware sequence parallelism technique that accounts for inter- and intra-machine bandwidth differences, (2) Torus Attention, a novel SP technique enabling overlapping of inter-machine all-to-all operations with computation, and (3) a one-sided communication implementation that minimizes GPU sender-receiver synchronization and computation overheads. Our experiments demonstrate that StreamFusion outperforms the state-of-the-art approach by an average of (up to ).
Paper Structure (23 sections, 1 theorem, 11 equations, 12 figures, 2 algorithms)

This paper contains 23 sections, 1 theorem, 11 equations, 12 figures, 2 algorithms.

Key Result

Lemma 4.1

Let $\text{V}_\text{diff} = \frac{\text{V}_\text{USP} - \text{V}_\text{SFU\xspace}}{BLHD/N} = \frac{4N}{P_u^2} - \frac{4M+ 6N}{P_u} - \frac{2P_u}{M} + 2N + 6$, then $\text{V}_\text{diff} \geq 0$ when $2 \leq M \leq P_u \leq N$.

Figures (12)

  • Figure 1: Generating images with diffusion transformers
  • Figure 2: Illustration of commonly used SP techniques for DiT inference.
  • Figure 3: (a) Differences in aggregated bandwidth between intra- and inter-machine network; and (b) Latency breakdown of the state-of-the-art SP technique USP usp.
  • Figure 4: Implicit synchronization in each step of the Ring Attention due to two-sided communication APIs, i.e, ncclSendRecv.
  • Figure 5: Topology-aware sequence parallelism in StreamFusion.
  • ...and 7 more figures

Theorems & Definitions (2)

  • Lemma 4.1
  • proof