NABLA: Neighborhood Adaptive Block-Level Attention for Efficient Video Generation

Full self-attention in video diffusion transformers scales quadratically with the spatio-temporal token count, making processing high-resolution video generation prohibitively slow and memory-intensive. We introduce NABLA, a Neighborhood-Adaptive Block-Level Attention mechanism that constructs a per-head sparse mask in three steps: 1) average-pooling queries and keys into <inline-formula> <tex-math notation="LaTeX">$N \times N$ </tex-math></inline-formula> blocks, 2) retaining the highest-probability blocks via a cumulative distribution function (CDF) threshold, and 3) optionally unioning the result with Sliding-Tile Attention (STA) to mitigate boundary artifacts. NABLA integrates seamlessly into PyTorch&#x2019;s FlexAttention without requiring custom kernels or auxiliary losses. Extensive experiments demonstrate significant acceleration for both training and inference: on the Wan 2.1 14B text-to-video model at 720p, NABLA achieves up to <inline-formula> <tex-math notation="LaTeX">$2.7\times $ </tex-math></inline-formula> speed-up in inference while matching baseline quality metrics (CLIP: <inline-formula> <tex-math notation="LaTeX">$42.06 \rightarrow 42.08$ </tex-math></inline-formula>, VBench: <inline-formula> <tex-math notation="LaTeX">$83.16 \rightarrow 83.17$ </tex-math></inline-formula>, FVD: <inline-formula> <tex-math notation="LaTeX">$68.9 \rightarrow 67.5$ </tex-math></inline-formula>). Furthermore, during pre-training of a 2B DiT at <inline-formula> <tex-math notation="LaTeX">$512^{2}$ </tex-math></inline-formula> resolution, NABLA reduces iteration time from 10.9s to 7.5s (<inline-formula> <tex-math notation="LaTeX">$1.46\times $ </tex-math></inline-formula> acceleration) while achieving lower validation loss.