FlashOmni: Unified Sparse Engine for DiTs

• FlashOmni is a unified sparse attention engine that standardizes heterogeneous sparsity strategies in Diffusion Transformers using a compact symbolic abstraction.
• It employs dynamic feature caching and block-sparse skipping within a single CUDA kernel to reduce computational cost without retraining.
• Empirical results demonstrate near-linear kernel speedup and preservation of visual quality, outperforming traditional isolated sparsity methods.

FlashOmni is a unified sparse attention engine for Diffusion Transformers (DiTs) that standardizes heterogeneous sparsity strategies—especially feature caching and block-sparse skipping—within a single execution framework and a single attention kernel. It is designed for arbitrary DiT architectures used in visual synthesis, where deployment is often limited by the computational cost of large models, high resolutions, and long sequences. The central claim of "FlashOmni: A Unified Sparse Attention Engine for Diffusion Transformers" is that a compact symbolic abstraction can represent multiple sparsity granularities and drive near-theoretical kernel acceleration without retraining, while preserving visual quality in practical image and video generation settings (Qiao et al., 29 Sep 2025).

1. Motivation and problem formulation

Diffusion Transformers have achieved state-of-the-art results in high-fidelity image and video synthesis, but their inference cost remains substantial. The paper identifies attention and associated projection kernels as major bottlenecks, especially for large models and long-sequence workloads such as HunyuanVideo with sequence length 33K. In response, prior sparsity-based acceleration methods have introduced mechanisms such as feature caching and block-sparse skipping, but these methods exhibit inconsistent granularity, a fragmented design space, and limited kernel generality.

FlashOmni is motivated by the observation that sparsity in DiTs has been operationalized through multiple incompatible representations. Token-level caching, block-level skipping, static masks, and dynamic masks often require different implementation logic and specialized kernels. According to the paper, this fragmentation increases engineering burden and limits universality. FlashOmni therefore targets a different level of abstraction: it is not introduced as a new sparsity heuristic, but as an execution engine that can encode and run diverse sparsity patterns within a common sparse-attention interface.

This framing is important for interpreting the system. A common misunderstanding is to treat FlashOmni as a single sparsity strategy. The paper instead presents it as a runtime unification layer for multiple structured sparsity strategies, suitable for arbitrary DiT architectures and tasks. This suggests that its primary contribution lies in systems design and kernel construction rather than in proposing a new pruning or masking criterion.

2. Sparse-symbol representation

The core abstraction in FlashOmni is a pair of 8-bit sparse symbols, denoted $\mathcal{S}_c$ and $\mathcal{S}_s$. These symbols compactly encode two different forms of structured sparsity as logical block masks. $\mathcal{S}_c$ represents feature caching by indicating which blocks are cached and can skip recomputation, while $\mathcal{S}_s$ represents block-sparse skipping by indicating which attention-matrix blocks should be skipped.

The symbols are generated during an update phase using metrics derived from a compressed blockwise attention map $\widetilde{P}$. For each block, the paper computes two quantities:

$$\mathcal{C}_{i, v\to t} = \sum_j \alpha_{j, i},\quad \mathcal{G}_{i, t\to v} = \sum_j \beta_{j, i}.$$

Here, $\mathcal{C}_{i, v\to t}$ is the Vision-to-Text Contribution, measuring how much a vision token block influences text tokens, and $\mathcal{G}_{i, t\to v}$ is the Text-to-Vision Guidance, measuring how much a vision block is controlled by text. The indices and compressed tokens are obtained by mean pooling over blocks. The selection logic described as Eqn. (1) in the paper chooses blocks with minimal combined scores, up to predefined thresholds $\tau_c$, to be cached. The resulting logical masks are then compressed into the 8-bit symbols.

This symbolization is the mechanism by which FlashOmni unifies multiple sparsity granularities. The paper emphasizes that the same representation can encode coarse feature reuse and finer blockwise attention skipping, allowing one kernel to handle mixed sparsity patterns dynamically.

Component	Encoded meaning	Execution role
$\mathcal{S}_c$	Cached blocks	Skip recomputation and reuse outputs
$\mathcal{S}_s$0	Skipped attention blocks	Omit blockwise attention submatrices
$\mathcal{S}_s$1	Cached projection bias	Reduce GEMM-$\mathcal{S}_s$2 computation

The compactness of the representation is not merely a storage detail. In the paper’s formulation, 8-bit compression is also intended to keep decoding overhead small enough that runtime sparsity remains profitable at kernel level.

3. Update-Dispatch execution model

FlashOmni organizes execution through what the paper calls an “Update-Dispatch” paradigm. The update step refreshes sparse symbols and caches using current features. The dispatch step then uses those symbols to guide efficient execution for the next $\mathcal{S}_s$3 steps.

This design separates the relatively infrequent construction of sparse metadata from the repeated execution of sparse kernels. In practical terms, sparsity decisions are amortized across subsequent denoising steps rather than recomputed at every kernel invocation. The paper presents this as a mechanism for supporting dynamic sparsity without incurring prohibitive control overhead.

A single CUDA kernel decodes sparse symbols at runtime using bitwise operations. This kernel supports dynamic switching between two execution modes. In cache-then-reuse mode, cached outputs are reused for selected blocks or tokens. In compute-on-demand mode, only the required unmasked blocks are evaluated. The kernel therefore does not hard-code a single sparse pattern; instead, it interprets the sparse symbols at runtime.

The implementation detail stressed by the paper is that decoding occurs once per block, not per element. Each cooperative thread array decodes the relevant symbol bits, stores the result in registers, and then executes the selected path. The stated purpose is to avoid branching overhead and to keep sparse control flow compatible with high-throughput attention execution.

In systems terms, this is the main architectural unification. FlashOmni does not require one kernel for caching and another for block skipping. The same attention kernel can dispatch both behaviors, provided that the symbolic masks have been prepared during the update step.

4. Sparse attention and sparse GEMM design

The sparse attention kernel, described as Algorithm 1 in the paper, combines reuse and selective computation. If a token is cached, meaning that $\mathcal{S}_s$4 under the paper’s notation, the kernel reuses historical output. The paper notes that this reuse may include a small update, citing Taylor series expansion as in TaylorSeer. If a token is not cached, attention is computed only over unmasked blocks as specified by $\mathcal{S}_s$5, so entire attention submatrices can be skipped.

FlashOmni extends the same sparse logic to projection GEMMs. In GEMM-$\mathcal{S}_s$6, the query projection skips rows corresponding to cached tokens. The paper reports that GEMM-$\mathcal{S}_s$7 follows the sparsity ratio closely, similar to the attention kernel itself. In GEMM-$\mathcal{S}_s$8, the output projection exploits structure across multiple heads. When attention outputs for certain heads are cached across tokens, their contribution to the final output can be precomputed and stored as a bias, denoted $\mathcal{S}_s$9. This bias-caching mechanism reduces repeated projection work in subsequent steps.

The paper treats GEMM-$\mathcal{S}_c$0 as the more challenging case. Whereas sparse attention and GEMM-$\mathcal{S}_c$1 align closely with theoretical computation reduction, GEMM-$\mathcal{S}_c$2 incurs somewhat more decoding overhead because sparsity interacts with the reduction axis. Even so, the reported implementation still achieves large acceleration and approaches the theoretical ceiling.

Two aspects are significant here. First, FlashOmni covers not only attention masking but also the surrounding linear algebra that determines end-to-end attention-block cost. Second, its notion of unification is multi-granularity in a literal sense: coarse feature caching and fine block-sparse skipping are both executed within the same sparse engine rather than as separate subsystems.

5. Empirical performance and quality preservation

The paper reports near-linear speedup, closely matching the sparsity ratio, for both attention and GEMM-$\mathcal{S}_c$3 (Qiao et al., 29 Sep 2025). At 90% combined sparsity, the attention kernel reaches up to approximately $\mathcal{S}_c$4 speedup. For GEMM-$\mathcal{S}_c$5, the measured acceleration is $\mathcal{S}_c$6 to $\mathcal{S}_c$7, with a maximum peaking at about 87.5% of the theoretical limit; elsewhere in the paper, the GEMM-$\mathcal{S}_c$8 ablation is summarized as reaching 84%–93% of the theoretical maximum.

End-to-end results are necessarily smaller because attention-kernel acceleration does not eliminate the rest of the model’s runtime. Applied with a multi-granularity sparsity strategy to HunyuanVideo at sequence length 33K and about 46% sparsity, FlashOmni achieves up to $\mathcal{S}_c$9 end-to-end acceleration, while attention alone can exceed $\mathcal{S}_s$0. This gap between kernel-level and pipeline-level speedup is consistent with the paper’s systems focus: the engine substantially accelerates the targeted kernels, but total runtime still includes non-attention work.

The quality results are presented as a further empirical argument for the unified engine. Across tested image and video generative models, including FLUX.1 and HunyuanVideo, FlashOmni is reported to preserve or even improve generation quality at moderate and high sparsity settings. The paper specifically cites better SSIM, PSNR, LPIPS, and FID than baseline approaches under comparable acceleration ratios. It also states that FlashOmni consistently outperforms block-sparse skipping baselines such as DiTFastAttnV2 and SpargeAttn, as well as feature-caching baselines such as ToCa, FORA, and TaylorSeer, in terms of generation quality at the same acceleration ratios.

These empirical results support two distinct claims. The first is systems-oriented: actual runtime tracks theoretical FLOP reduction unusually closely for sparse attention and GEMM-$\mathcal{S}_s$1. The second is model-oriented: unifying sparsity strategies can improve the quality–efficiency trade-off relative to using either caching or block skipping in isolation.

6. Position within DiT sparsity research

FlashOmni is presented as a response to a fragmented sparsity ecosystem in DiT inference (Qiao et al., 29 Sep 2025). Earlier methods typically implement a single strategy in isolation, whether dynamic or static attention skipping or some form of feature caching. The paper argues that this leads to redundant engineering effort because each pattern often requires its own logic, memory treatment, and custom kernels.

Against that background, FlashOmni’s stated contribution is a unified abstraction for arbitrary structured sparsity. The paper describes this as the first such abstraction for DiT attention execution, with generic bitwise decoding enabling runtime support for mixed sparsity strategies. Its novelty therefore lies less in discovering that caching or skipping can help, and more in making them interoperable at kernel level.

Another important aspect of scope is that FlashOmni is explicitly training-free. The paper states that it requires zero changes to model weights or training and operates entirely through input-dependent dynamic sparsity at inference time. This distinguishes it from approaches that obtain efficiency through retraining, architecture modification, or sparsity-aware fine-tuning.

A plausible implication is that FlashOmni is most useful where multiple sparsity strategies already exist or are expected to coexist across tasks, resolutions, or modalities. In that setting, a general sparse execution substrate reduces the need for task-specific kernel engineering. The paper’s open-source release reinforces this interpretation by positioning FlashOmni as reusable infrastructure rather than as a model-specific optimization.

The principal limitation visible from the reported results is structural rather than negative: FlashOmni accelerates attention-centric kernels very effectively, but end-to-end speedup remains bounded by the rest of the inference stack. The reported HunyuanVideo result of about $\mathcal{S}_s$2 end-to-end acceleration, alongside larger attention-only gains, makes clear that unified sparse attention is a major but not exhaustive component of DiT systems optimization.