ScaleFusion: Scalable Inference of Spatial-Temporal Diffusion Transformers for High-Resolution Long Video Generation

ScaleFusion: Scalable Inference of Spatial-Temporal Diffusion Transformers for High-Resolution Long Video Generation

Speakers: Jiacheng Yang, AWS; Jun Wu, AWS; Zhen Zhang, AWS; Yida Wang, University of Toronto; Gennady Pekhimenko, University of Toronto, CentML

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

The rapid advancement of generative AI has brought about an increasing demand for high-resolution, long video generation. This talk, "ScaleFusion: Scalable Inference of Spatial-Temporal Diffusion Transformers for High-Resolution Long Video Generation," presented by Jiacheng Yang, addresses a critical bottleneck in the distributed inference of a specific class of video diffusion models: spatial-temporal diffusion transformers (ST-DiT). These models, while powerful for generating high-quality video content from text prompts, suffer from significant inference latency when scaled across multiple GPU machines due to inefficiencies in how existing distributed frameworks handle communication.

The core problem identified by the researchers from AWS, Vector Institute, University of Toronto, and CentML lies in the use of monolithic all-to-all (auto) operators. These operators, essential for data redistribution in distributed settings, hinder the crucial overlapping of communication with computation, leading to poor scalability. ScaleFusion proposes a novel approach that decouples these large all-to-all operations into smaller, independently executable slices. This fundamental insight enables the creation of new opportunities for communication-computation overlap through sophisticated inter-layer and intra-layer communication scheduling algorithms.

The impact of ScaleFusion is substantial. By strategically breaking down and rescheduling these operations, the framework achieves an impressive 14x average speedup compared to state-of-the-art distributed inference frameworks. Furthermore, it demonstrates remarkable strong scaling, achieving a 36x speedup on four AWS P4 GPU machines. This work is pivotal for democratizing the generation of high-fidelity, extended video content, making it more accessible and computationally efficient for a wide range of applications, from creative content production to synthetic data generation.

Background

▶ Watch: Introduction and problem of ST-DiT inference latency (0:00)

Video generation models have seen significant progress, with diffusion models emerging as a leading paradigm for synthesizing realistic and diverse video content from textual descriptions. The process typically begins with a user-provided text prompt, which guides the generation from a completely random noise tensor. Through multiple iterative sampling steps, the model progressively denoises this tensor, eventually passing the refined output to a Variational Autoencoder (VAE) decoder to produce the final video frames. Within this complex pipeline, the Spatial-Temporal Diffusion Transformer (ST-DiT) component is often the most computationally intensive and time-consuming part, making it a prime target for optimization.

ST-DiT models are characterized by their architecture, which is composed of multiple repetitions of spatial layers and temporal layers. As their names suggest, spatial layers compute attention mechanisms predominantly along the spatial dimensions of the video frames, while temporal layers focus on dependencies across the temporal dimension, i.e., between frames. A key challenge in high-resolution, long video generation is the sheer volume of tokens. For instance, generating a mere 32-second, 4K HD video can involve millions of tokens per layer, demanding significant computational resources. Despite this, the total model size of ST-DiT is relatively modest, typically a few gigabytes, distinguishing it from the multi-hundred-gigabyte scale of large language models (LLMs).

To efficiently handle the massive token count and accelerate inference in a distributed environment, sequence parallelism is a common strategy. In this approach, the initial noise tensor, typically represented with dimensions like Batch (B), Temporal (T), Spatial (S), Head (H), and Head Dimension (D) (BTSHD), is sharded across multiple GPUs, often along the spatial dimension. This means different portions of the video frame data reside on different GPUs. The execution of ST-DiT layers under this distributed scheme presents a critical problem: while temporal layers can operate largely independently without inter-GPU communication (as temporal attention doesn't require spatially sharded data to be re-aggregated), spatial layers do. When a spatial layer needs to compute attention, a query (Q) tensor on one GPU might need to interact with key (K) and value (V) tensors residing on other GPUs to form the complete attention score matrix.

Prior state-of-the-art distributed inference frameworks address this inter-GPU dependency using a blocking all-to-all (auto) operator. Before a spatial layer, this operator transposes the data from its initial spatial-distributed form to a temporal-distributed form. This redistribution ensures that all necessary Q, K, and V components for a given spatial attention computation are co-located on the same GPU, allowing the attention calculation to proceed correctly. After the spatial layer completes, another all-to-all operator is typically invoked to transpose the data back to its original spatial-distributed form, preparing it for subsequent layers or the VAE decoder. While functional, this monolithic, blocking nature of the all-to-all operator introduces substantial communication overhead. As demonstrated by the researchers, on just two AWS P4 GPU machines, this overhead can account for as much as 34% of the total inference time, severely limiting the scalability of ST-DiT models across multiple GPU systems. This fundamental inefficiency necessitates a more granular and intelligent approach to communication scheduling.

Key Findings

▶ Watch: ST-DiT execution, large tokens, and sequence parallelism (1:58)

The research behind ScaleFusion identifies and addresses a critical limitation in the distributed inference of Spatial-Temporal Diffusion Transformers (ST-DiT). The primary findings are centered on overcoming the communication bottlenecks imposed by conventional distributed frameworks and introducing novel scheduling mechanisms.

Firstly, the team pinpointed that monolithic all-to-all (auto) operators are the main culprits hindering efficient communication-computation overlap in existing distributed ST-DiT inference. These operators, designed for data redistribution in sequence parallelism, act as blocking operations, preventing computations from proceeding while data is being exchanged across GPUs. This leads to significant idle time and poor scaling performance, particularly as the number of GPU machines increases.

Secondly, a pivotal discovery made by the ScaleFusion team is that the all-to-all operator is intrinsically decomposable. Rather than treating it as an indivisible, large-scale collective operation on an entire tensor, they recognized that it can be broken down into multiple smaller, independent all-to-all operations, each applied to a "slice" of the tensor. This decomposition fundamentally changes how data can be moved and processed in a distributed setting.

Building upon this insight, the researchers developed two synergistic communication scheduling algorithms: intra-layer communication scheduling and inter-layer communication scheduling. The intra-layer approach leverages the decomposability of the all-to-all operator to overlap its execution for one data slice with the computation of a previously processed slice within the same layer. The inter-layer scheduling further refines this by anticipating data availability, moving parts of the all-to-all operations for subsequent layers to overlap with computations of the current layer, effectively minimizing or eliminating communication "bubbles" that persist even with intra-layer scheduling.

These algorithmic innovations yielded impressive performance gains. ScaleFusion achieved an outstanding 14x average speedup compared to existing state-of-the-art distributed inference frameworks for ST-DiT models. This dramatic acceleration significantly reduces the time required to generate high-resolution, long videos. Furthermore, the framework demonstrated exceptional strong scaling, achieving a 36x speedup when scaling from a single P4 GPU machine to four AWS P4 GPU machines (totaling 32 GPUs). This near-linear scaling indicates the effectiveness of ScaleFusion's approach in distributing the workload efficiently across an increasing number of hardware resources. The researchers also observed that the benefits of ScaleFusion become even more pronounced with larger and more demanding video generation workloads, suggesting its robustness for future, higher-fidelity applications.

Technical Deep Dive

▶ Watch: Prior work's all-to-all operators cause communication overhead (2:50)

ScaleFusion's technical innovation revolves around a deep understanding of the ST-DiT architecture and the inefficiencies of distributed collective operations, specifically the all-to-all (auto) operator, within current frameworks.

The ST-DiT model is structured as a sequence of alternating spatial and temporal transformer layers. For an input noise tensor with shape B (batch size), T (temporal dimension), S (spatial dimension), H (number of attention heads), and D (head dimension), which is typically sharded along the spatial dimension (S) for sequence parallelism:

• Spatial Layer: This layer computes attention primarily across the spatial dimension. When the tensor is sharded spatially, different parts of the spatial dimension reside on different GPUs. To compute attention correctly, query (Q), key (K), and value (V) tensors from different GPUs might need to interact. In prior work, this necessitated a global data redistribution.
• Temporal Layer: This layer computes attention along the temporal dimension. Since the tensor is sharded spatially, and temporal attention does not inherently require data from different spatial shards to interact for a given temporal slice, these layers can often proceed without explicit inter-GPU communication if the data is in the correct distributed form.

Prior Work's Approach and its Limitations:

In existing state-of-the-art frameworks (e.g., DSP, as mentioned in the talk), the execution for a spatial layer proceeds as follows:

1. Input: A noise tensor is distributed across GPUs, sharded along the spatial dimension.
2. Pre-Spatial Layer All-to-All: Before the spatial layer's computation, a large, blocking all-to-all operator is invoked. This operator transposes the data from its spatial-distributed form to a temporal-distributed form. The purpose is to ensure that for any given spatial attention computation, all necessary Q, K, and V components are co-located on the same GPU. This is crucial because, without it, a Q tensor on GPU 1 might need KV tensors from GPU 4 to form the complete attention score.
3. Spatial Layer Computation: Once the data is in the temporal-distributed form, the spatial attention can be computed efficiently on each GPU without further communication.
4. Post-Spatial Layer All-to-All: After the spatial layer, another large, blocking all-to-all operator is used to transpose the data back to its original spatial-distributed form. This is often necessary for subsequent temporal layers or the VAE decoder, which might expect the data in a specific distributed layout to avoid communication.

The critical issue with this approach is that these all-to-all operators are monolithic and blocking. They introduce significant communication overheads, creating "bubbles" or periods where GPUs are idle, waiting for data transfer to complete before computation can resume. This non-overlapped communication can consume up to 34% of the total inference time on just two P4 GPU machines.

ScaleFusion's Core Insight: Decoupling the All-to-All Operator:

The fundamental breakthrough of ScaleFusion is the realization that a large all-to-all operator, applied to a large tensor, can be decoupled into multiple smaller, independent all-to-all operations. Instead of processing the entire tensor at once, the tensor can be conceptually divided into N "slices." Each slice can then undergo its all-to-all transformation independently. This fine-grained control over data movement is the cornerstone for enabling communication-computation overlap.

Intra-Layer Communication Scheduling:

With the all-to-all operator decoupled, ScaleFusion introduces intra-layer scheduling. Consider the process around a spatial layer:

1. Sliced All-to-All: Instead of one large all-to-all before the spatial layer, ScaleFusion performs N smaller all-to-all operations, one for each slice.
2. Overlap Opportunity: As soon as the first slice (slice 0) completes its all-to-all operation, its data can be sent to the spatial layer for computation. While slice 0 is being computed, slice 1 can simultaneously undergo its all-to-all operation. This creates a pipeline: All-to-All(slice_i) overlaps with Computation(slice_{i-1}).

This significantly reduces the idle time compared to the monolithic approach. However, a "bubble" still exists: the very first all-to-all operation (for slice 0) cannot be overlapped with any prior computation within the same layer, as it's the first operation. The researchers found that increasing the number of slices (N) reduces this initial bubble but eventually leads to performance degradation beyond approximately 16 slices due to the increased overhead of managing more, smaller operations and potential re-synchronization costs.

Inter-Layer Communication Scheduling:

To address the remaining "bubble" from the first slice's non-overlapped communication and further optimize the pipeline, ScaleFusion proposes inter-layer communication scheduling. This technique leverages the fact that computations within a layer often complete asynchronously for different parts of the tensor.

1. Anticipatory Communication: For example, the "back slices" of a tensor being processed by a spatial layer might finish computation slightly earlier than the "front slices."
2. Proactive All-to-All: Instead of waiting for the entire tensor output from the spatial layer to be ready before initiating the next all-to-all (e.g., the one that transposes data back for the temporal layer), ScaleFusion starts the all-to-all operation for the already-completed slices of the spatial layer's output.
3. Extended Overlap: This effectively moves part of the subsequent all-to-all operation to overlap with the ongoing computation of the current layer (the spatial layer). This "chaining" mechanism can be applied not just between spatial and temporal layers but also temporal and spatial layers, progressively reducing or even eliminating all communication bubbles throughout the entire ST-DiT execution.

By combining these intra-layer and inter-layer scheduling algorithms, ScaleFusion transforms the execution flow from a sequence of blocking communication and computation stages into a highly pipelined process where data movement and processing occur concurrently, maximizing GPU utilization and dramatically reducing inference latency.

Experimental Setup & Results

▶ Watch: ScaleFusion's key insight: decouple all-to-all operators (4:08)

The evaluation of ScaleFusion focused on demonstrating its effectiveness in accelerating the distributed inference of ST-DiT (Spatial-Temporal Diffusion Transformer) models, particularly for high-resolution, long video generation workloads. The experiments were meticulously designed to highlight the framework's advantages over existing state-of-the-art approaches.

Model and Workloads:

The core model under test was the ST-DiT, which is central to video diffusion processes. The workloads represented demanding scenarios, such as generating 32-second, 4K HD videos. These workloads are characterized by an extremely high number of tokens per layer, often reaching millions, posing significant challenges for distributed systems due to the large data volume and communication requirements.

Hardware:

All experiments were conducted on AWS P4 GPU machines. These instances are equipped with NVIDIA A100 GPUs, providing substantial computational power and high-bandwidth interconnects, making them suitable for evaluating distributed ML workloads. The strong scaling experiments specifically scaled the number of machines from one to four P4 machines, which translates to a total of 32 GPUs (each P4 machine typically houses 8 GPUs).

Baselines:

The performance of ScaleFusion was compared against two primary baselines:

1. Ulysses Attention (OS): This represents a common approach used for LLM inference, likely a basic form of sequence parallelism that might not include advanced communication optimizations for ST-DiT. It serves as a general reference point for standard distributed attention mechanisms.
2. State-of-the-Art (DSP): This baseline represents the most advanced prior work in distributed ST-DiT inference. It utilizes the monolithic all-to-all (auto) operators as described in the background, which, while functional, suffer from significant communication overheads due to their blocking nature.

ScaleFusion Variants:

To illustrate the individual and combined contributions of its novel scheduling algorithms, ScaleFusion was evaluated in two configurations:

1. SO (ScaleFusion Intra-Layer Only): This variant incorporated only the intra-layer communication scheduling algorithm, demonstrating the benefits of decoupling the all-to-all operator and overlapping communication with computation within the same layer.
2. SF (ScaleFusion Full): This represents the complete ScaleFusion framework, leveraging both the intra-layer and inter-layer communication scheduling algorithms. This variant aims to achieve maximum overlap by addressing the residual communication bubbles.

Metrics:

The primary metrics for evaluation included:

• Inference Latency: The total time taken to generate a video.
• Speedup: The ratio of the baseline inference time to ScaleFusion's inference time.
• Strong Scaling: How the inference time reduces as more computational resources (GPU machines) are added, for a fixed problem size.

Headline Numbers and Key Results:

The experimental results vividly showcased the superiority of ScaleFusion:

• Communication Overhead: The "State-of-the-Art (DSP)" baseline exhibited substantial communication overhead, accounting for as much as 34% of the total inference time on just two P4 GPU machines. This highlighted the severe bottleneck that ScaleFusion aimed to address.
• Overlap Efficiency: ScaleFusion (SF) demonstrated an almost complete overlap of communication overhead with computation. This was visually represented in performance graphs, showing a dramatic reduction in idle time compared to the baselines.
• Linear Scaling: The framework achieved nearly linear scaling when moving from one to four P4 machines. This is a critical indicator of efficient distributed resource utilization.
• Average Speedup: ScaleFusion delivered an impressive 14x average speedup when compared to state-of-the-art frameworks (DSP). This translates directly to a significant reduction in the time required for video generation.
• Strong Scaling Speedup: For the specific configuration of four AWS P4 GPU machines, ScaleFusion achieved a remarkable 36x strong scaling speedup. This indicates that as more resources are added, ScaleFusion can leverage them exceptionally well to reduce the inference time for a fixed video generation task.
• Workload Adaptability: The researchers noted that ScaleFusion's algorithms perform even better with larger video generation workloads. This suggests that the benefits of fine-grained scheduling become more pronounced as the scale and complexity of the task increase, making it future-proof for even higher-resolution or longer video requirements.

Scaling Beyond Four Machines:

An important discussion point in the Q&A session addressed scaling beyond four P4 machines. The presenter acknowledged that while resource limitations prevented testing on more than four machines, at some point, with a significantly larger number of GPUs, the system might transition from being computation-bottlenecked to communication-bottlenecked. In such a scenario, further optimizations might need to focus on the underlying communication fabric itself, rather than just scheduling techniques. However, the speaker emphasized that four P4 machines (32 GPUs) is still considered a "very reasonable setup" for current video generation workloads, making ScaleFusion highly relevant and useful in its current scope.

Practical Implications

▶ Watch: Applying ScaleFusion to overlap communication and computation (4:40)

ScaleFusion introduces significant practical implications across various facets of the AI/ML ecosystem, particularly for those involved in developing, deploying, and managing generative video models.

For practitioners and model builders, ScaleFusion directly translates into the ability to generate high-resolution, long videos with drastically reduced inference latency. This is crucial for applications demanding real-time or near real-time video synthesis, such as interactive content creation, personalized media experiences, or rapid prototyping in film and animation. The 14x average speedup means that tasks that previously took minutes could now complete in seconds, fundamentally changing user experience and creative workflows. It also enables the generation of higher quality and longer duration videos that might have been prohibitively expensive or slow with prior methods.

Infrastructure teams and ML deployers stand to benefit from the enhanced scalability and efficiency of distributed ST-DiT inference. The 36x strong scaling on four P4 GPU machines demonstrates that ScaleFusion can effectively utilize increasing hardware resources to meet growing demands without hitting early performance ceilings. This means better utilization of expensive GPU clusters, potentially reducing operational costs for large-scale video generation services. The insights gained from ScaleFusion regarding fine-grained control over collective operations (like all-to-all) provide a blueprint for optimizing other complex distributed transformer architectures, extending beyond just video diffusion models.

However, the adoption of ScaleFusion also comes with certain tradeoffs and limitations:

• Increased Algorithmic Complexity: Implementing ScaleFusion's inter-layer and intra-layer scheduling algorithms requires a deeper understanding of the model's execution graph and fine-grained control over communication primitives. This is more complex than simply relying on off-the-shelf, monolithic collective operations provided by standard deep learning frameworks.
• Tuning and Optimization: The effectiveness of intra-layer scheduling depends on carefully selecting the optimal number of slices (N). Too few slices limit overlap opportunities, while too many can introduce overhead from managing a larger number of smaller operations, potentially leading to performance degradation. This necessitates careful profiling and tuning for specific model configurations and hardware setups.
• Hardware and Network Considerations: While ScaleFusion effectively mitigates communication bottlenecks in the tested configurations, the talk acknowledges that at much larger scales (e.g., beyond 4 P4 machines/32 GPUs), the system might eventually become inherently communication-bottlenecked. At this point, the limiting factor would shift from scheduling inefficiencies to the raw bandwidth and latency of the underlying network interconnects (e.g., NVLink, InfiniBand). Future optimizations at such scales might require hardware-aware communication primitives or network topology-aware scheduling.
• Resource Availability: The research team explicitly stated that resource limitations prevented testing on more than four P4 machines. While the theoretical benefits are expected to continue, the practical validation of strong scaling on larger clusters remains an area for future work.

Despite these considerations, ScaleFusion represents a significant leap forward in making high-resolution, long video generation practical and efficient. It highlights the critical importance of moving beyond coarse-grained distributed operations towards a more intelligent, anticipatory, and finely-tuned approach to communication and computation overlap in the era of large-scale generative AI models.

Key Takeaways

• Existing distributed inference frameworks for Spatial-Temporal Diffusion Transformers (ST-DiT) suffer from significant latency due to monolithic all-to-all (auto) operators that prevent effective communication-computation overlap.
• ScaleFusion's core innovation is the discovery that the all-to-all operator can be decoupled into multiple smaller, independent slices, creating novel opportunities for pipelined execution.
• The framework employs two synergistic scheduling algorithms: intra-layer communication scheduling overlaps communication of one slice with computation of a previous slice, while inter-layer communication scheduling further reduces "bubbles" by proactively initiating subsequent layer's communication based on early completion of current layer's computation.
• ScaleFusion achieves a remarkable 14x average speedup compared to state-of-the-art frameworks and an impressive 36x strong scaling speedup on four AWS P4 GPU machines for high-resolution long video generation.
• The benefits of ScaleFusion become even more pronounced with larger and more demanding video generation workloads, indicating its robustness and scalability for future applications.
• While effective up to 32 GPUs, scaling to significantly larger GPU clusters might eventually shift the bottleneck to underlying communication hardware, necessitating further specialized optimizations.

About the Speaker(s)

The research presented on ScaleFusion is a collaborative effort involving experts from leading institutions and industry. The talk was delivered by Jiacheng Yang, who is associated with AWS. The co-authors of this work include Jun Wu and Zhen Zhang from AWS, Yida Wang from the University of Toronto, and Gennady Pekhimenko from the University of Toronto and CentML. This diverse team brings together expertise from cloud computing, academic research, and AI optimization, highlighting the interdisciplinary nature of addressing complex challenges in scalable machine learning.