SVG implementation for LTX 2

[jitendra-jalwaniya]

Overview

This PR extends Sparse VideoGen (SVG) spatiotemporal attention support to LTX-2 (LTX2) video generation models on Cloud TPUs, building on the custom Ulysses/ring SVG kernel infrastructure introduced for Wan (PR #480).

Self-attention in LTX-2 transformer blocks dynamically profiles query tokens to choose between spatial and temporal attention patterns per head, skipping unneeded query–key interactions while executing through hardware-aligned local-band kernels on TPU. Sparse attention is opt-in (use_svg_attention: True), disabled by default, and configurable across denoising steps, layers, and sparsity densities. Audio self-attention and cross-modal attention remain dense to preserve temporal and semantic grounding.

In addition, several changes in this PR ensure robust support for LTX-2 inference on both TPU v6 (Trillium) and TPU v7, including capping tile search VMEM limits to 32 MiB on TPU v6, correcting custom block size fallbacks (bkv_compute_in), enabling framewise VAE tiling, and accounting for CFG/STG batch scaling in block benchmarks.

---

TPU v6 and TPU v7 Compatibility

This PR includes dedicated improvements to guarantee smooth execution across both TPU v6 and v7 architectures:

• Tile Search VMEM Cap on TPU v6: In generate_ltx2.py, dynamically caps tile_search_vmem_limit_bytes to 32 MiB on TPU v6 devices to prevent VMEM allocations from spilling into host memory during automatic block-size tuning.
• Block Size Fallback Fix: In attention_flax.py, corrected _extract_custom_block_sizes so that omitted block_kv_compute_in values default directly to bkv_compute, preventing unoptimized compute tiles on TPU hardware.
• Framewise VAE Tiling: Enabled use_framewise_decoding when VAE tiling is activated to prevent HBM OOMs during long video decoding.
• CFG/STG Batch Accounting: In ltx2_block_benchmark.py, properly scaled batch sizes (2× for CFG, 4× for CFG + STG) and normalized per-batch VMEM limits to match true runtime inference workloads.

---

VABench Evaluation: SVG vs. Dense Attention

We evaluated SVG against dense attention on the Full VABench Benchmark suite (778 prompts across all 24 Easy/Hard bundles and 7 content categories) for LTX-2 synchronized text-to-audio-video (T2AV) generation at long sequence length (768 × 1280 × 241 frames, $N = 29,760$ video tokens, 10.04s @ 24 fps video + 24 kHz PCM audio) on 8× TPU v6e cores, followed by a 15-dimension VABench evaluation across 8× NVIDIA A100-80GB GPUs:

• Dense Baseline: use_svg_attention=False, attention=ulysses_custom
• SVG Ulysses: use_svg_attention=True, svg_spatial_density=0.25, attention=ulysses_custom

1. TPU v6e-8 Generation Performance (778 Videos @ 768 × 1280 × 241)

Metric	Dense Baseline	SVG Ulysses	Delta (SVG vs. Dense)
Denoising Time / Video (40 steps)	72.00 s	61.62 s	-10.38 s (-14.42% / 1.17× speedup)
Per-Step Denoising Latency	1.800 s / step	1.541 s / step	-0.259 s / step
Total Inference / Video (end-to-end)	102.94 s	91.23 s	-11.71 s (-11.38%)
Benchmark Wall Time (778 Videos)	~22.24 hours	~19.72 hours	-2.52 hours saved

2. 15-Dimension VABench Quality Highlights

Enabling SVG yields faster generation with zero quality regression, while improving audio-video synchronization and cross-modal alignment:

• Audio-Video Synchronization & Lip-Sync:
  • Synchformer Temporal Desynchronization (second_desync $\downarrow$): Reduced from 0.6743 s $\rightarrow$ 0.6404 s (-5.03% better sync overall), with strong gains on Animals (-18.25%), Music (-12.24%), Virtual Worlds (-10.97%), and Synchronous Physical Sounds (-6.12%).
  • LatentSync Lip-Sync Error (second_lsa $\downarrow$): Reduced by -31.60% overall (0.9813 $\rightarrow$ 0.6712), and by -39.15% on Human Sounds (1.3270 $\rightarrow$ 0.8075), showing that routing video heads through SVG while preserving dense audio and cross-modal attention significantly sharpens temporal alignment.
• Cross-Modal Alignment: Higher alignment across all three embedding models: ImageBind-Huge (+1.91%), ViCLIP-L (+1.82%), and LAION-CLAP (+1.79%).
• Audio Aesthetic Quality: AudioBox Aesthetics improved by +1.67% (3.5645 $\rightarrow$ 3.6241), with DNSMOS (+0.55%) and NISQA (+0.39%) on par.
• Multimodal Judge & Fine-Grained QA (Qwen2.5-Omni-7B):
  • Parity on multimodal judge criteria (alignment: 4.47 vs. 4.45 [-0.60%], audio realism: 3.94 vs. 3.91 [-0.82%], expressiveness: 4.25 vs. 4.23 [-0.45%]).
  • Notable accuracy gains on multi-turn question answering: Audio QA (+4.90%) and Visual QA (+4.64%).

Full 15-Dimension VABench Comparison Table (778 Prompts)

Module	Dimension	Metric / Criterion	Direction	Dense Baseline	SVG Ulysses	Absolute Delta	Relative Change
M1: Audio Quality & Aesthetics	first_dnsmos	Microsoft DNSMOS (sig_bak_ovr + p808)	$\uparrow$	1.6538	1.6629	+0.0091	+0.55%
	first_nisqa	NISQA v2 Speech/Audio Naturalness MOS	$\uparrow$	1.5264	1.5324	+0.0060	+0.39%
	first_audiobox	Meta AudioBox Aesthetics	$\uparrow$	3.5645	3.6241	+0.0596	+1.67%
M2: Cross-Modal Sync & Alignment	second_viclip	ViCLIP-L Text-Video Similarity	$\uparrow$	0.1918	0.1953	+0.0035	+1.82%
	second_clap	LAION-CLAP Text-Audio Similarity	$\uparrow$	0.3835	0.3904	+0.0069	+1.79%
	second_imagebind	Meta ImageBind-Huge AV Alignment	$\uparrow$	0.2144	0.2185	+0.0041	+1.91%
	second_desync	Synchformer Temporal Desync Offset (s)	$\downarrow$	0.6743 s	0.6404 s	-0.0339 s	-5.03% (Better Sync)
	second_lsa	LatentSync Lip-Sync Distance	$\downarrow$	0.9813	0.6712	-0.3101	-31.60% (Better Lip-Sync)
M3: Multimodal Judge (Qwen2.5-Omni-7B)	third_alignment	AV Semantic & Temporal Alignment (1–5)	$\uparrow$	4.4743	4.4473	-0.0270	-0.60% (Parity)
	third_audio_reality	Acoustic Realism & Fidelity (1–5)	$\uparrow$	3.9383	3.9062	-0.0321	-0.82% (Parity)
	third_visual_reality	Visual Realism & Coherence (1–5)	$\uparrow$	4.5476	4.4666	-0.0810	-1.78%
	third_expressiveness	Emotional & Dynamic Expressiveness (1–5)	$\uparrow$	4.2468	4.2275	-0.0193	-0.45% (Parity)
	third_artistry	Audiovisual Aesthetic Quality (1–5)	$\uparrow$	3.6272	3.6478	+0.0206	+0.57%
	Module 3 Mean	Mean Multimodal Judge Score (1–5)	$\uparrow$	4.1668	4.1391	-0.0278	-0.67% (Parity)
M4: Multi-Turn Question Answering	fourth_qa_audio	Audio QA Accuracy (0–1)	$\uparrow$	0.6438	0.6754	+0.0315	+4.90%
	fourth_qa_vision	Visual QA Accuracy (0–1)	$\uparrow$	0.6088	0.6370	+0.0283	+4.64%
	Module 4 Mean	Mean Multi-Modal QA Accuracy (0–1)	$\uparrow$	0.6263	0.6562	+0.0299	+4.77%

---

Configuration Example

To enable SVG on LTX-2, add the following to ltx2_video.yml or override via CLI flags:

attention: ulysses_ring_custom_fixed_m
use_svg_attention: True
svg_spatial_density: 0.20
svg_active_start_step: 10
svg_active_end_step: 35
svg_active_start_layer: 1
svg_active_end_layer: 28
svg_profile_query_count: 64
svg_sample_max_row: 10000
svg_profile_seed: 0
svg_include_first_frame: True

---

Testing

Run the LTX-2 SVG unit tests from the repository root:

python -m pytest -q \
  src/maxdiffusion/tests/ltx2/test_svg_config_propagation_ltx2.py \
  src/maxdiffusion/tests/ltx2/test_svg_attention_ltx2.py

All existing Wan and LTX-2 unit tests continue to pass.

Note: This PR also includes minor pyink==23.10.0 formatting fixes in aot_cache.py, aot_cache_test.py, wan_pipeline.py, and svg_attention_test.py (introduced in #480) because the CI Linter workflow runs pyink src/maxdiffusion --check across the
entire repo rather than only changed files.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[gemini-code-assist]

Code Review

This pull request integrates Sparse VideoGen (SVG) attention for LTX2 models on TPUs, including configuration propagation, routing logic, pipeline updates, and comprehensive unit tests. The feedback identifies potential TypeError crashes in ltx2_pipeline.py if optional SVG density, step, or layer parameters are explicitly configured as None, and provides robust suggestions to safely handle these cases.

[mbohlool]

This change is very big, would it be possible to break it into smaller PRs. we have strict guidelines for huge diffs (that probably creating with AI help).