SAFE-DiT: Semantics-Aware Fast-path Execution for High-Resolution Diffusion Transformers

2026-06-28

AI Summary

This work identifies and quantifies a previously overlooked performance bottleneck in high-resolution diffusion Transformer inference—termed Mask-Induced Dispatch Tax (MIDT)—caused by spatially redundant attention masks that disable the fast path of FlashAttention, leading to substantial system overhead. To address this, the authors propose SAFE, a training-free, semantic-aware fast-path execution framework that preserves semantically critical structures while enabling efficient adaptive inference. SAFE integrates source-verified mask pruning, prompt-guided token partitioning, selective state updates, and periodic context refreshing. Fully compatible with PyTorch’s SDPA stack, SAFE achieves end-to-end speedups of 2.69× and 5.09× at resolutions of 1024² and 2560², respectively, reduces peak memory consumption from 94.1 GB to 27.9 GB, and enables generation at ultra-high resolution (3072²).

Abstract

High-resolution Diffusion Transformer (DiT) inference contains substantial spatial redundancy, but many spatially adaptive implementations encode regional computation as attention masks, which can inadvertently move scaled dot-product attention (SDPA) away from FlashAttention fast paths. We identify this avoidable systems bottleneck as Mask-Induced Dispatch Tax (MIDT) and show that it grows with latent sequence length. We introduce SAFE-DiT, a training-free Semantics-Aware Fast-path Execution framework that separates exact mask elision from approximation-based spatial scheduling. SAFE-DiT removes only provenance-certified image self-attention masks that induce a row-wise constant shift in attention logits, preserves semantics-bearing masks such as text-padding masks, and realizes spatial adaptation through prompt-conditioned token partitioning, selective state updates with global context, and periodic context refresh. We call this acceleration-only configuration SAFE-Core and report sensitivity-weighted classifier-free guidance separately as SAFE-DiT+SW. On the evaluated PyTorch SDPA stack, redundant masks make long-sequence attention $4.1\times$ to $5.8\times$ slower than the mask-free path. On Lumina-Next, SAFE-DiT achieves $2.69\times$ end-to-end acceleration at $1024^2$ resolution and $5.09\times$ at $2560^2$, reduces peak memory at $2560^2$ from 94.1 to 27.9 GB, and enables $3072^2$ generation when dense inference runs out of memory. Paired metrics, component ablations, and a blinded human study support visual non-inferiority of SAFE-Core to the dense fast-path baseline, while SAFE-DiT+SW provides a separate prompt-alignment operating point without reintroducing spatial self-attention masks. Code is available at https://github.com/xuanhuayin/SAFE-DiT.

Problem

• Diffusion Transformer
• spatial redundancy
• attention masks
• FlashAttention
• Mask-Induced Dispatch Tax

Innovation

• Diffusion Transformer
• FlashAttention
• spatial redundancy
• mask elision
• high-resolution generation