提出db-SP技术，解决稀疏注意力中并行计算负载不均问题。
db-SP: Accelerating Sparse Attention for Visual Generative Models with Dual-Balanced Sequence Parallelism
- 双层级分区策略，动态平衡头与块级计算负载。
- 实测相比顶尖方法平均提速1.25倍，注意力层提速1.40倍。
- 适合大规模视觉生成模型推理加速，尤其含稀疏注意力结构的DiT。
通过序列并行加速扩散变压器（DiT）推理对降低视觉生成延迟至关重要，但在使用块级稀疏注意力的模型中，因注意力头间稀疏性差异和稀疏掩码内密集块分布不均，导致严重负载不均衡。现有方法在沿头维度或块维度进行并行时均无法避免此问题。本文首次形式化定义稀疏不平衡比，并提出db-SP——一种感知稀疏性的序列并行技术。该方法采用双层级划分机制，在头与块两个层面实现近乎完美的负载均衡，且开销极低。同时，db-SP在运行时动态调整头与块维度的并行度，以应对去噪步骤和网络层间稀疏模式的演化。实验表明，db-SP在平均上相比最先进序列并行方法实现1.25倍端到端加速和1.40倍注意力层加速。代码已公开于https://github.com/thu-nics/db-SP。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。