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

TL;DR

This paper introduces db-SP, a dual-balanced sequence parallelism method that significantly accelerates sparse attention in visual generative models by balancing workload across attention heads and blocks, leading to notable speedups.

Contribution

The paper proposes a novel sparsity-aware parallelism technique with dynamic workload balancing for sparse attention in diffusion transformers, improving inference speed.

Findings

Achieves 1.25x overall speedup over state-of-the-art methods.

Attains 1.40x attention-specific speedup.

Effectively balances workload at head and block levels dynamically.

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.