TLDR: FlashOmni is a new unified sparse attention engine for Diffusion Transformers that uses flexible sparse symbols and optimized kernels to accelerate visual synthesis. It combines feature caching and block-sparse skipping within a single framework, achieving significant speedups (up to 1.5x end-to-end) without compromising visual quality, making DiTs more efficient and easier to deploy.
Diffusion Transformers (DiTs) have made significant strides in generating high-quality visuals, from images to videos. However, their impressive capabilities come with a hefty computational cost, making them challenging to deploy efficiently, especially for large-scale or high-resolution tasks. To tackle this, researchers have explored various acceleration methods, with sparsity-based techniques being particularly popular due to their ability to speed up models without requiring extensive retraining.
The challenge with existing sparsity methods is their lack of universality. Different sparsity patterns often demand custom-built software components, known as kernels, to achieve high performance. This fragmented approach limits flexibility, increases development overhead, and makes it difficult to combine different sparsity strategies effectively.
Introducing FlashOmni: A Unified Solution
A new research paper introduces FlashOmni, a unified sparse attention engine designed to work seamlessly with any Diffusion Transformer architecture. FlashOmni aims to overcome the limitations of current sparse acceleration methods by providing a single, flexible framework that can handle a wide array of sparsity strategies.
FlashOmni’s core innovation lies in its use of “flexible sparse symbols.” These are compact 8-bit codes that standardize how different sparsity strategies are represented. This unified abstraction allows a single attention kernel to execute diverse sparse computations, from feature caching (reusing computations from previous steps) to block-sparse skipping (omitting unimportant calculations).
How FlashOmni Works: The Update-Dispatch Paradigm
FlashOmni operates on an “Update-Dispatch” paradigm, which integrates and executes two main sparsity strategies: feature caching and block-sparse skipping. In the “Update” phase, FlashOmni refreshes its sparse symbols and a feature cache based on the current step’s computations. It determines which parts of the model can be skipped or reused in subsequent steps. The “Dispatch” phase then uses these sparse symbols to accelerate attention computation over the next several timesteps. During this phase, the system intelligently skips redundant operations, significantly boosting efficiency.
The engine employs two types of sparse symbols: Sc for feature caching and Ss for block-sparse skipping. These symbols are generated by analyzing a compressed version of the attention map, identifying areas where computations can be safely reduced without compromising quality. For instance, in text-to-vision models, FlashOmni carefully avoids caching critical interactions between text and visual tokens to maintain multimodal consistency.
Optimized Sparse GEMMs for Enhanced Efficiency
Beyond its general attention kernel, FlashOmni also introduces optimized sparse General Matrix Multiplications (GEMMs) called GEMM-Q and GEMM-O. These specialized operations further eliminate redundant computations in the linear layers of attention modules, specifically during query projection and output projection. For example, if a block’s output is retrieved from the cache, FlashOmni GEMM-Q can skip the corresponding query projection. Similarly, FlashOmni GEMM-O optimizes how cached information is integrated into the final output, reducing both computational workload and memory usage.
Also Read:
- Optimizing Image Generation: HilbertA’s Efficient Sparse Attention for Diffusion Models
- Enhancing Diffusion LLM Performance with Adaptive Block Sizing
Performance and Impact
Experiments demonstrate that FlashOmni delivers impressive performance gains. Its sparse kernel design achieves near-linear speedup, closely matching the theoretical computation reduction. In attention and GEMM-Q operations, it can achieve a one-to-one acceleration with the sparsity ratio. For GEMM-O, it delivers 2.5x to 3.8x acceleration, peaking at about 87.5% of the theoretical limit.
When applied to real-world models like the Hunyuan model (33K parameters) with a multi-granularity sparsity strategy, FlashOmni enables approximately 1.5x end-to-end acceleration without degrading the visual quality of generated content. It consistently outperforms existing block-sparse skipping and feature caching methods across various quality metrics on models like FLUX and HunyuanVideo.
This unified approach not only accelerates Diffusion Transformers but also simplifies the development and deployment of sparse acceleration techniques, making high-fidelity visual generation more accessible and efficient. For more technical details, you can refer to the full research paper.
