RainFusion2.0: Temporal-Spatial Awareness and Hardware-Efficient Block-wise Sparse Attention

TL;DR

Diffusion Transformer models incur prohibitive attention costs in video/image generation, hindering practical deployment. RainFusion2.0 introduces an online adaptive, hardware-efficient sparse attention framework that combines blockwise representative tokens, 3D window permutation, and a First Frame Sink to maintain video quality while reducing computation. The approach achieves up to 80% sparsity with end-to-end speedups of 1.5–1.8× across multiple diffusion-based generators and demonstrates generality across heterogeneous hardware such as NPUs/ASICs. This work enables broader deployment of DiT-based generators on diverse devices and suggests integration with quantization or distillation for further acceleration.

Abstract

In video and image generation tasks, Diffusion Transformer (DiT) models incur extremely high computational costs due to attention mechanisms, which limits their practical applications. Furthermore, with hardware advancements, a wide range of devices besides graphics processing unit (GPU), such as application-specific integrated circuit (ASIC), have been increasingly adopted for model inference. Sparse attention, which leverages the inherent sparsity of attention by skipping computations for insignificant tokens, is an effective approach to mitigate computational costs. However, existing sparse attention methods have two critical limitations: the overhead of sparse pattern prediction and the lack of hardware generality, as most of these methods are designed for GPU. To address these challenges, this study proposes RainFusion2.0, which aims to develop an online adaptive, hardware-efficient, and low-overhead sparse attention mechanism to accelerate both video and image generative models, with robust performance across diverse hardware platforms. Key technical insights include: (1) leveraging block-wise mean values as representative tokens for sparse mask prediction; (2) implementing spatiotemporal-aware token permutation; and (3) introducing a first-frame sink mechanism specifically designed for video generation scenarios. Experimental results demonstrate that RainFusion2.0 can achieve 80% sparsity while achieving an end-to-end speedup of 1.5~1.8x without compromising video quality. Moreover, RainFusion2.0 demonstrates effectiveness across various generative models and validates its generalization across diverse hardware platforms.

Figures (3)

• Figure 1: Workflow of RainFusion2.0
• Figure 2: Results of RainFusion on Diffusion Models. HunyuanVideo1.5 and Wan2.2 generate 720p videos under two configurations: full attention and RainFusion with 80% sparsity. Qwen-image-edit generates 1024×1024 images using RainFusion with 60% sparsity.
• Figure 3: Experimental results on the Wan2.2 dataset. As shown in Subfigure (b), the video generated by RainFusion (80% sparsity, without 3D permutation) is overall comparable to that of full attention (Subfigure (a)). However, two prominent spurious rock artifacts emerge in the bottom-right corner of the video frame. In contrast, these visual artifacts are completely eliminated when 3D permutation is integrated into RainFusion (Subfigure (c)).