[PR 1/2] 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.

This is PR 1/2 (model side). It depends on #493 (pyink formatting fix on main). The config, pipeline, AOT metadata and docs wiring are in #498.

Changes in this PR:

• attention_ltx2.py: LTX2Attention accepts an attention_config dict with SVG settings and dispatches video self-attention to the SVG kernel (or dense via jax.lax.cond) based on the active step/layer window.
• transformer_ltx2.py: plumbs spatiotemporal_shape, svg_timestep, svg_step_index and per-layer layer_index through LTX2StaticContext/LTX2BlockContext (scanned and unscanned paths). Only attn1 (video self-attention) gets SVG; audio_attn1 is forced dense.
• tests/ltx2/test_svg_attention_ltx2.py: new unit tests.

---

VABench Evaluation: SVG vs. Dense Attention

The end-to-end results below require both this PR and #498.

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 TPU v6e-8 (8 chips), followed by a 15-dimension VABench evaluation across 8× NVIDIA A100-80GB GPUs. Each prompt was generated once per configuration:

• Dense Baseline: use_svg_attention=False, attention=ulysses_custom
• SVG Ulysses: use_svg_attention=True, svg_spatial_density=0.25, attention=ulysses_custom, svg_active_* left at defaults (SVG active on all steps and layers)

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 comparable overall quality: most metrics are on par or slightly higher, with a small drop in judged visual realism (-1.78%):

• 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). We observed this with video heads routed through SVG while audio and cross-modal attention stay dense; since each prompt was generated once, we have not yet measured seed-to-seed variance for this metric.
• 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 most 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%]); visual realism is slightly lower (4.55 vs. 4.47 [-1.78%]).
  • 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%

---

Testing

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

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

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

[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 into the LTX2 model. It introduces SVG configuration parameters, updates the attention layer to route between dense and sparse SVG attention based on active steps and layers, and propagates the necessary spatiotemporal and step metadata through the transformer blocks and static/block contexts. Additionally, comprehensive unit tests are added to verify the SVG activation boundaries, dispatch routing, and full model forward passes. There are no review comments, so no additional feedback is provided.

[Perseus14]

Nice work wiring SVG into LTX-2 and running the full 778-prompt VABench eval! The context plumbing through both the scanned and unscanned transformer paths is clean, and keeping audio/cross-attention dense makes sense.

I left inline comments on a few things to tighten up before merging:

1. Deduplicating SVG config/dispatch with Wan (attention_ltx2.py vs FlaxWanAttention in attention_flax.py) so we don't maintain two copies of the same ~100 lines.
2. Passing num_layers from the model into attention_config instead of hardcoding svg_num_layers: 48.
3. Cleaning up unused or ignored attention_config keys so callers aren't surprised when a key has no effect.
4. Adding one numerical test on TPU (ulysses_custom) and a check that audio_attn1 stays dense.

Also a quick note on the PR description:

• Since each prompt was generated with a single seed, the -31.6% LatentSync and +4.9% QA deltas likely include seed-to-seed variance (sparse attention is an approximation of dense). I'd frame those as "comparable / on par with dense" unless we have multi-seed numbers.
• Since #493 is already merged into main, you can remove the "depends on #493" note.

[Perseus14]

Please squash the commits! @jitendra-jalwaniya

[jitendra-jalwaniya]

Please squash the commits! @jitendra-jalwaniya

Done

[Perseus14]

LGTM!