Add Sparse VideoGen attention for Wan on TPUs

[Ravi-VK]

Video diffusion processes long sequences of video tokens across many denoising steps, making self-attention expensive. Sparse attention reduces this cost by skipping interactions that contribute little to the output.

This PR introduces an opt-in sparse attention implementation for MaxDiffusion: Sparse VideoGen (SVG) for Wan on TPUs. SVG samples a few queries to choose a spatial or temporal pattern for each head, then computes only the selected mask. Density and active steps/layers are configurable, and SVG is disabled by default.

See SVG.md for an illustrated explanation, configuration examples, and profiling instructions.

The implementation rounds sparse boundaries to hardware tiles, approximately preserving the attention-pair budget. Selected interior tiles share one mask-free kernel; only tiles touching sequence padding need validity masks.

The feature is organized into six commits:

1. Kernel primitives: compute attention over selected ranges of tiles.
2. Tile rounding and dispatch: select boundary tiles and combine main and padding-tail results.
3. Routing and layouts: choose each head’s pattern and arrange tokens for efficient execution.
4. Wan integration: connect SVG to inference with configuration controls and unsupported-mode checks.
5. Tests: cover routing, layouts, configuration, sharding, and kernel correctness.
6. Documentation: explain the method, usage, profiling, and results with illustrations.

A review follow-up adds configuration and shape guards, a static inactive-step shortcut, and regression tests.

Wan2.2 720p results on TPU v6e-8, using 81 frames and 40 denoising steps:

Policy	Estimated transformer FLOPs saved	Denoising speedup	PSNR (dB)
Conservative	≈27.2%	1.13×	26.47
Moderate	≈32.2%	1.20×	26.14
Aggressive	≈37.3%	1.28×	24.80

Timings are medians of three warm runs against same-node optimized fixed-M dense controls. PSNR uses FFmpeg’s aggregate YUV metric against dense outputs.

In an earlier evaluation using the moderate SVG policy, Wan2.2 at 720p retained 98.1% of the dense baseline’s mean VBench dimension score across a 31-prompt, 16-dimension screening subset.

SVG supports inference through the custom Ulysses/ring backends; training, Animate, external masks, periodic support, chunked Ulysses, CFG cache, and MagCache are unsupported.

[gemini-code-assist]

Code Review

This pull request implements Sparse VideoGen (SVG) attention on TPUs for Wan in MaxDiffusion, introducing modules for routing, token placement, and kernel execution. The review feedback highlights several critical issues to address: potential TypeError and AttributeError crashes in the pipeline and attention configuration logic, possible NaN propagation bugs in the attention kernels when all keys are masked out, a missing divisibility check in the Pallas partial kernel, and opportunities to optimize compilation by statically checking step_index.

[Perseus14]

Added some minor comments. PTAL

[Perseus14]

@Ravi-VK Please squash the commits