db-SP: Accelerating Sparse Attention for Visual Generative Models with Dual-Balanced Sequence Parallelism

TL;DR

This work tackles the latency bottleneck of visual Diffusion Transformers by addressing workload imbalance that arises when applying block-wise sparse attention under sequence parallelism. It introduces db-SP, a dual-balanced partitioning scheme that balances head- and block-level workloads with greedy algorithms, and a sparsity-aware runtime strategy selection to switch parallel plans based on latency predictions. The approach yields notable end-to-end speedups and attention-only improvements across state-of-the-art SP methods, while keeping overheads low through plan reuse and data-locality incentives. The work provides practical improvements for real-time or multi-GPU DiT inference in high-resolution video and image generation tasks.

Abstract

Scaling Diffusion Transformer (DiT) inference via sequence parallelism is critical for reducing latency in visual generation, but is severely hampered by workload imbalance when applied to models employing block-wise sparse attention. The imbalance stems from the inherent variation in sparsity across attention heads and the irregular distribution of dense blocks within the sparse mask, when sequence parallelism is applied along the head dimension (as in Ulysses) or the block dimension (as in Ring Attention). In this paper, we formalize a sparse imbalance ratio to quantify the imbalance, and propose db-SP, a sparsity-aware sequence parallelism technique that tackles the challenge. db-SP contains a dual-level partitioning approach that achieves near-perfect workload balance at both the head and block levels with negligible overhead. Furthermore, to handle the evolving sparsity patterns across denoising steps and layers, db-SP dynamically determines the parallel degrees for the head and block dimensions at runtime. Experimental results demonstrate that db-SP delivers an end-to-end speedup of 1.25x and an attention-specific speedup of 1.40x over state-of-the-art sequence parallel methods on average. Code is available at https://github.com/thu-nics/db-SP.

Figures (12)

• Figure 1: db-SP achieves speedup via dual-balanced partitioning and parallel strategy selection. Top: Sparse attention is partitioned across GPUs by attention heads, leading to workload imbalance. Bottom: db-SP balances the workload at both head-level and block-level partitioning. Ulysses, USP, and Ring Attention are different sequence parallelism methods, where U$x$R$y$ denotes that $x$ GPUs perform Ulysses and $y$ GPUs perform Ring Attention.
• Figure 2: Comparison of sequence parallelism variants.
• Figure 3: The latency of Wan2.1-T2V-14B with PAROAttention as the GPU number increases.
• Figure 4: Dual-level workload imbalance when applying sequence parallelism to sparse attention.
• Figure 5: The overview of db-SP. For every attention computation, each parallel strategy applies its specific partitioning method to mitigate the workload imbalance, achieving a post-balancing sparse imbalance ratio close to 1. Furthermore, db-SP leverages a sparsity-aware selection mechanism that considers the parallel strategy $\bf{p}$, sparse pattern $\bf{s}$, and sparse imbalance ratio$\rho_s$ to predict and apply the optimal parallel strategy and the partitioning plan at runtime.