Optimizing Video Diffusion Attention on TPUs with Structured Sparsity

The Core Update

Video generation eats compute for breakfast. Denoising passes take seconds because sequence lengths explode fast. A standard 81-frame clip at 720p already pushes token counts between 50,000 and 400,000. Bump that resolution to 1440p and token counts quadruple. Because standard full attention scales quadratically, self-attention quickly swallows nearly 90% of your per-layer latency.

Google's TPU engineering team published a technical breakdown showing how to break this bottleneck on Cloud TPU v6e hardware. By mapping dynamic spatial and temporal attention masks directly to hardware tiles via JAX and Pallas Splash Attention kernels, teams can ditch dense matrix operations and retain roughly 38% of token interactions without quality loss.

Official Source: Google Announcement

Technical Impact & Mechanism

Most attention interactions in video diffusion models carry virtually zero weight. Running dense matrix multiplications across all tokens wastes cycles. But sparse attention faces an implementation problem. Pure algorithmic sparsity does not automatically translate into hardware speedups on accelerator silicon.

Hardware needs contiguous memory access. TPUs process attention through tiled matrix multiplication: slices of queries meet slices of keys. If your sparse mask only zeros out isolated values within a tile, the TPU still executes the compute instruction. You only save latency when you skip the entire tile.

The architecture relies on three structural components:

1. Head Specialization: Diffusion transformer heads fall cleanly into two categories. Spatial heads focus on local patches within immediate frames. Temporal heads focus on fixed spatial points across long timelines.
2. Sink Preservation: Both spatial and temporal passes preserve the initial frame as a visual anchor alongside a tight local interaction window.
3. Tile Classification: Pallas kernels partition the attention grid into three tile states: Full (computed normally), Boundary (masked elementwise), and Empty (bypassed completely at the scheduler level).

# Conceptual Pallas / JAX block-sparse tile execution loop
import jax.numpy as jnp
from jax.experimental import pallas as pl

def sparse_attention_tile_kernel(q_tile_ref, k_tile_ref, v_tile_ref, out_ref, tile_mask_type):
    # Skip empty blocks entirely to eliminate hardware memory traffic
    if tile_mask_type == 0:  # EMPTY TILE
        return
    
    scores = jnp.dot(q_tile_ref[...], k_tile_ref[...].T) / jnp.sqrt(128.0)
    
    if tile_mask_type == 1:  # BOUNDARY TILE
        # Apply explicit intra-tile mask for boundary conditions
        scores = jnp.where(tile_mask_ref[...], scores, -1e9)
        
    probs = jnp.exp(scores - jnp.max(scores, axis=-1, keepdims=True))
    out_ref[...] += jnp.dot(probs, v_tile_ref[...])

On TPU v6e chips running synthetic BF16 payloads with 75.6K tokens, skipping unneeded tiles eliminates major memory traffic bottlenecks. Attention compute shifts from an $O(N^2)$ tax into localized, high-throughput tile chunks.

Action Plan for Developers & Businesses

If you run custom video generation infrastructure, stop paying for full attention passes on long contexts:

1. Audit Layer Profiling: Profile your diffusion transformer heads during generation. Identify which layers exhibit distinct temporal versus spatial grouping. Do not mask blindly.
2. Align Block Sizes to Hardware Tiles: Configure your mask boundaries to match your hardware tile shapes (such as 3328 query and 1536 key/value slices). Fractional tiles kill speed gains.
3. Switch to Hardware-Aware Kernels: Drop naive boolean masking in favor of block-level sparse kernels using Pallas on TPUs or Triton on GPUs. If your kernel executes compute on masked zeros, your pipeline is burning budget.