Google Speeds Up Video Diffusion on TPUs with Optimized Sparse Attention

Video diffusion models require many denoising steps to generate smooth videos, and at scale, the computational bottleneck often lies in self-attention operations within those steps. As resolutions climb from 720p to 1440p, attention's share of processing time can surge from 55% to 88%, making it the prime target for optimization.

The Challenge: Attention at Scale#

For an 81-frame video at 720p resolution, the sequence length can range from 50,000 to 400,000 tokens—and quadrupling resolution quadruples that sequence length. Since full attention computes every interaction between tokens in quadratic time, higher resolutions quickly become prohibitively expensive to process. However, research into video diffusion models reveals that attention doesn't distribute evenly. Many query-key interactions carry negligible importance, creating opportunities to skip unnecessary computations.

Sparse Attention Patterns#

Google researchers built on work from Sparse VideoGen (SVG), which recognizes that attention heads specialize into two categories: spatial heads that look mostly at patches in the same or nearby frames, and temporal heads that track motion in narrow spatial regions across many frames. This structure is far from random—it follows regular geometric patterns that can be exploited. The key insight is that certain positions in a frame connect heavily to one another, while others barely interact, allowing algorithms to dynamically route each head to either a spatial or temporal attention mask based on sampling.

Efficient Kernel Implementation#

Understanding the pattern is one thing; executing it efficiently on hardware is another. Google developed a series of optimizations for its JAX and Pallas Splash Attention kernels on TPU v6e chips:

• Tile-aware masking: Instead of evaluating the attention mask at every position within every visited tile, the team classified tiles as "full" (containing only retained pairs) or "boundary" (containing both), giving each a specialized execution path. This dropped latency from 96 ms to 54 ms.

• Tile size tuning: Smaller tiles follow the mask boundary more closely, but they reduce compute efficiency. Testing revealed that the latency sweet spot lay at a query tile size of 3,328—not the configuration with minimal boundary work.