arXiv:2509.25401cs.LGcs.AI2025-09被引 1

统一稀疏注意力引擎,让扩散模型加速更高效

FlashOmni: A Unified Sparse Attention Engine for Diffusion Transformers

  • 用通用符号表示多种稀疏模式,一套内核适配所有架构
  • 注意力与GEMM计算加速达2.5至3.8倍,接近理论极限
  • 适用于多粒度稀疏策略,适合部署大模型的性能优化

多模态扩散变压器(DiTs)在图像生成方面表现卓越,但部署受限于高计算开销。现有稀疏加速方法因稀疏模式多样,常需定制内核,缺乏通用性。本文提出FlashOmni,一种兼容任意DiT架构的统一稀疏注意力引擎。它引入灵活的稀疏符号,标准化多种稀疏策略(如特征缓存、块稀疏跳过),使不同稀疏计算可在单一注意力内核中执行。同时,设计优化的稀疏GEMM,利用稀疏符号消除冗余计算,进一步提升效率。实验表明,FlashOmni在注意力和GEMM-Q上实现近线性加速,速度提升与稀疏比几乎一致(1:1);在GEMM-O上加速达2.5×–3.8×,最高达到理论极限的87.5%。结合多粒度稀疏策略,使Hunyuan模型(33K参数)获得约1.5×端到端加速,且不降低视觉质量。

原文摘要 · Abstract (English)

Multi-Modal Diffusion Transformers (DiTs) demonstrate exceptional capabilities in visual synthesis, yet their deployment remains constrained by substantial computational demands. To alleviate this bottleneck, many sparsity-based acceleration methods have been proposed. However, their diverse sparsity patterns often require customized kernels for high-performance inference, limiting universality. We propose FlashOmni, a unified sparse attention engine compatible with arbitrary DiT architectures. FlashOmni introduces flexible sparse symbols to standardize the representation of a wide range of sparsity strategies, such as feature caching and block-sparse skipping. This unified abstraction enables the execution of diverse sparse computations within a single attention kernel. In addition, FlashOmni designs optimized sparse GEMMs for attention blocks, leveraging sparse symbols to eliminate redundant computations and further improve efficiency. Experiments demonstrate that FlashOmni delivers near-linear, closely matching the sparsity ratio speedup (1:1) in attention and GEMM-$Q$, and achieves 2.5$\times$-3.8$\times$ acceleration in GEMM-$O$ (max peaking at about 87.5% of the theoretical limit). Applied with a multi-granularity sparsity strategy, it enables the Hunyuan model (33K) to achieve about 1.5$\times$ end-to-end acceleration without degrading visual quality.

扩散模型稀疏计算注意力机制性能优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。