From 15a24c9f5ac9dc1f35e6027f695b093ece6a4554 Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Wed, 16 Sep 2026 09:18:18 +0000 Subject: [PATCH 1/8] Cap tile search VMEM limit to 32 MiB on TPU v6 --- src/maxdiffusion/generate_ltx2.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index b60f223b8..2ef7aeba9 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -105,6 +105,9 @@ def maybe_tune_block_sizes(config): vmem_limit_bytes = int( keys.get("tile_search_vmem_limit_bytes") or config.flash_block_sizes.get("vmem_limit_bytes") or 64 * 1024 * 1024 ) + if "v6" in jax.devices()[0].device_kind.lower(): + vmem_limit_bytes = min(vmem_limit_bytes, 32 * 1024 * 1024) + bench = LTX2BlockBenchmark.from_config(config, mesh, vmem_limit_bytes=vmem_limit_bytes) max_logging.log(f"[tile-search] tuning block sizes for {bench.label} before inference...") result = grid_search( From a83f0992acc5627921a50e9cf8e598027f29e954 Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Wed, 16 Sep 2026 09:21:12 +0000 Subject: [PATCH 2/8] Fix boolean mask dtype and account for CFG batch multiplier in LTX2 block benchmark --- src/maxdiffusion/utils/ltx2_block_benchmark.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/maxdiffusion/utils/ltx2_block_benchmark.py b/src/maxdiffusion/utils/ltx2_block_benchmark.py index 564c03594..f15ebaaa5 100644 --- a/src/maxdiffusion/utils/ltx2_block_benchmark.py +++ b/src/maxdiffusion/utils/ltx2_block_benchmark.py @@ -32,7 +32,7 @@ "--xla_tpu_enable_async_collective_fusion_multiple_steps=true", "--xla_tpu_overlap_compute_collective_tc=true", "--xla_enable_async_all_gather=true", - "--xla_tpu_scoped_vmem_limit_kib=65536", + "--xla_tpu_scoped_vmem_limit_kib=32768", "--xla_tpu_enable_async_all_to_all=true", "--xla_tpu_enable_all_experimental_scheduler_features=true", "--xla_tpu_enable_latency_hiding_scheduler=true", @@ -143,7 +143,13 @@ def __init__( self._width_orig = width data_shards = int(mesh.shape.get("data", 1)) * int(mesh.shape.get("fsdp", 1)) - self._batch = batch if batch is not None else max(1, data_shards) + guidance_scale = getattr(config, "guidance_scale", 1.0) + stg_scale = getattr(config, "stg_scale", 0.0) + do_cfg = guidance_scale > 1.0 + do_stg = stg_scale > 0.0 + cfg_mult = 4 if (do_cfg and do_stg) else (2 if do_cfg else 1) + default_batch = max(1, data_shards) * cfg_mult + self._batch = batch if batch is not None else default_batch self._hf_cfg = LTX2VideoTransformer3DModel.load_config(config.pretrained_model_name_or_path, subfolder="transformer") self._inputs = self._make_inputs() @@ -163,7 +169,7 @@ def tiled_seq_lens(self): return (s, s) def vmem_bytes(self): - return self._vmem + return self._vmem // max(1, self._batch) def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2): cmp = bkv_compute or bkv @@ -268,10 +274,10 @@ def _make_inputs(self): # Prompts prompt_embeds = jax.random.normal(k2, (self._batch, 1024, 3840), dtype) - prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.int32) + prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.bool_) audio_prompt_embeds = jax.random.normal(k4, (self._batch, 1024, 3840), dtype) - audio_prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.int32) + audio_prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.bool_) timestep = jnp.zeros((self._batch,), jnp.float32) repl = jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()) From 7b58086961edf8940989aa91083da4078787d53e Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Thu, 17 Sep 2026 05:57:48 +0000 Subject: [PATCH 3/8] Implement Sparse VideoGen (SVG) for LTX2 video generation - Add SVG configuration options to ltx2_video.yml and ltx2_3_video.yml - Add dynamic SVG attention dispatching in LTX2Attention for self-attention - Propagate spatiotemporal_shape, layer indices, and timestep/step index in transformer blocks - Wire SVG configuration and step index into LTX2Pipeline diffusion loops - Include SVG configuration parameters in generate_ltx2 AOT cache metadata - Update LTX2BlockBenchmark with SVG attention configuration - Add comprehensive unit tests covering SVG config propagation and attention dispatch - Update documentation in docs/svg.md --- docs/svg.md | 149 ++++++ src/maxdiffusion/configs/ltx2_3_video.yml | 18 + src/maxdiffusion/configs/ltx2_video.yml | 18 + src/maxdiffusion/generate_ltx2.py | 20 + .../models/ltx2/attention_ltx2.py | 129 +++++- .../models/ltx2/transformer_ltx2.py | 74 ++- .../pipelines/ltx2/ltx2_pipeline.py | 50 +- .../tests/ltx2/test_svg_attention_ltx2.py | 429 ++++++++++++++++++ .../ltx2/test_svg_config_propagation_ltx2.py | 209 +++++++++ .../utils/ltx2_block_benchmark.py | 7 + 10 files changed, 1083 insertions(+), 20 deletions(-) create mode 100644 docs/svg.md create mode 100644 src/maxdiffusion/tests/ltx2/test_svg_attention_ltx2.py create mode 100644 src/maxdiffusion/tests/ltx2/test_svg_config_propagation_ltx2.py diff --git a/docs/svg.md b/docs/svg.md new file mode 100644 index 000000000..d5d3cd480 --- /dev/null +++ b/docs/svg.md @@ -0,0 +1,149 @@ +# Sparse VideoGen attention on TPUs + +Video diffusion models generate a video through a sequence of denoising steps. At each step, attention lets each video token gather information from other tokens across space and time. This becomes expensive as the resolution and number of frames grow: dense attention considers every query–key pair, even though many interactions contribute very little to the output. + +Sparse spatiotemporal attention takes advantage of this structure. Instead of attending everywhere, a query attends to a smaller set of positions chosen to capture the spatial and temporal information it needs. + +## How SVG chooses where to attend + +[Sparse VideoGen (SVG)](https://arxiv.org/abs/2502.01776) observes that attention heads often favor different patterns. Spatial heads concentrate attention within a frame or nearby frames. Temporal heads concentrate attention around corresponding spatial positions across frames. These patterns let us approximate dense attention while computing fewer interactions. + +![Attention from the same query in a spatial head and a temporal head, shown across six latent frames.](images/svg/head-patterns.png) + +*Observed attention weights for the same query in two heads. The spatial head concentrates 94.8% of its attention on the query frame, while the temporal head places substantial attention near corresponding positions in other frames. Cyan squares mark the query's spatial position; `m` is the attention mass in each displayed frame. Colors use a shared logarithmic scale. These are examples of observed attention, rather than the masks themselves.* + +SVG makes the choice separately for each head: + +1. **Profile a few queries.** Compute dense attention outputs for a small sample of query tokens, using all keys. +2. **Compare two patterns.** Compute the sampled outputs under spatial and temporal masks, then measure each one's error relative to dense attention. +3. **Use the better approximation.** Select the lower-error pattern and apply sparse attention to all queries in that head. + +The choice is recomputed at each active layer and denoising step. A head does not need to keep the same assignment throughout generation. Sparsity is configurable, so the same method can trade a smaller approximation error for a larger reduction in computation. + +This implementation reimplements routing, token placement, and kernel execution for Wan and LTX2 in MaxDiffusion. The [original SVG implementation](https://github.com/svg-project/Sparse-VideoGen) provides the reference method. + +## Making sparse attention efficient on TPU + +Skipping query–key interactions only helps if the hardware can skip the corresponding work efficiently. Our implementation arranges tokens so that both spatial and temporal patterns can use the same local-band attention kernel. Spatial heads keep frame-major order; temporal heads group corresponding spatial positions across frames. Outputs are restored to their original order afterward. + +TPUs compute attention in tiles. A tile can lie entirely inside the sparse pattern, entirely outside it, or cross its boundary. + +A local attention band over a query–key tile grid, highlighting full, boundary, and skipped tiles. + +*Sparse pattern before tile rounding. Query and key indices refer to the selected token layout. Blue indicates retained interactions and gray indicates skipped interactions. The prefix anchor is omitted for clarity.* + +We round boundary tiles to either keep or skip them, approximately preserving the attention-pair budget of the original pattern. This slightly changes which interactions are retained, but lets all selected interior tiles run through one kernel without per-token sparse masking. Only tiles touching sequence padding need an additional validity mask; their outputs are combined with the main result using a numerically stable merge. + +Token placement and restoration run after the Ulysses exchange, on each device's local heads. This limits the layout work to the heads that device will actually process. + +The masks also include an optional prefix anchor. The `svg_include_first_frame` option retains the first `H × W` keys in the selected layout. For temporal heads, that prefix spans spatial positions across frames, so it does not correspond to the original first video frame. + +## Configuration and usage + +SVG is disabled by default. To enable it, set `use_svg_attention=True` and choose the densities and the steps and layers where sparsity should be active. Calls outside that interval continue to use dense attention. + +For example, these overrides select the moderate Wan2.2 policy: + +```yaml +use_svg_attention: True +svg_high_noise_density: 0.50 +svg_low_noise_density: 0.20 +svg_active_start_step: 11 +svg_active_end_step: 40 +svg_active_start_layer: 1 +svg_active_end_layer: 40 +svg_profile_query_count: 64 +svg_sample_max_row: 10000 +svg_profile_seed: 0 +svg_include_first_frame: True +``` + +Step and layer intervals are zero-based and half-open: `[11, 40)` includes steps 11 through 39. Steps refer to denoising iterations, not noise-timestep values. + +| Option | What it controls | +|---|---| +| `svg_spatial_density` | Density for single-expert Wan models; applies to either selected head pattern. | +| `svg_high_noise_density`, `svg_low_noise_density` | Separate densities for Wan2.2's two experts. | +| `svg_active_start_step`, `svg_active_end_step` | Denoising steps where SVG is enabled. | +| `svg_active_start_layer`, `svg_active_end_layer` | Transformer layers where SVG is enabled. | +| `svg_profile_query_count` | Number of queries sampled when choosing each head's pattern. | +| `svg_sample_max_row` | Limits query sampling to this prefix of the sequence. | +| `svg_profile_seed` | Seed for reproducible query sampling. | +| `svg_include_first_frame` | Enables the prefix anchor in the selected layout. | +| `svg_flash_block_sizes` | Sparse kernel tiling; an empty mapping uses `flash_block_sizes`. | + +Density controls the local-band width. The anchor and tile rounding affect the actual number of retained interactions, so density is not itself the fraction of total transformer FLOPs retained. + +The following is an example attention configuration for eight devices. Add it and the policy above to a Wan2.2 configuration used by `src/maxdiffusion/generate_wan.py`: + +```yaml +attention: ulysses_ring_custom_fixed_m +ici_data_parallelism: 2 +ici_context_parallelism: 4 +ulysses_shards: 2 +flash_block_sizes: + block_q: 6400 + block_kv: 2048 + block_kv_compute: 2048 + block_kv_compute_in: 1024 + heads_per_tile: 1 + vmem_limit_bytes: 67108864 +svg_flash_block_sizes: + block_q: 3328 + block_kv: 2816 + block_kv_compute: 256 + block_kv_compute_in: 256 + heads_per_tile: 1 + vmem_limit_bytes: 67108864 +``` + +Dense calls retain their configured ring split. SVG exchanges over the whole context axis, giving Ring2/Ulysses2 for dense calls and Ring1/Ulysses4 for sparse calls in this example. Tile sizes may need tuning for other shapes and devices. + +### Supported configurations + +SVG supports inference through the four custom Ulysses/ring attention backends. It requires matching self-attention QKV shapes, a matching video-token grid, and `heads_per_tile=1`. Heads and sequence length must divide evenly across the context shards, including any heads created by folding an unsharded batch. + +Training, Animate, external attention masks, periodic support, and chunked Ulysses are unsupported. SVG also cannot be combined with CFG cache or MagCache. + +## Performance and quality + +At 720p, SVG offers a configurable tradeoff between denoising latency and similarity to dense generation. The following Wan2.2 results use TPU v6e-8, 81 frames, and 40 denoising steps. Times are medians of three warm runs against same-node optimized fixed-M dense controls, using one prompt and seed. + +| Policy | Estimated total transformer FLOPs saved | Dense denoising | SVG denoising | Denoising speedup | PSNR (dB) | +|---|---:|---:|---:|---:|---:| +| Conservative | ≈27.2% | 153.46 s | 136.11 s | **1.13×** | **26.47** | +| Moderate | ≈32.2% | 153.37 s | 127.94 s | **1.20×** | **26.14** | +| Aggressive | ≈37.3% | 153.50 s | 119.86 s | **1.28×** | **24.80** | + +*PSNR is measured against dense outputs using FFmpeg's aggregate YUV metric.* + +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. + +## Qualitative example + +![Three rows of giraffe video frames, with dense attention on the left and SVG on the right.](images/svg/dense-vs-svg.png) + +*Dense attention (left) and SVG (right), shown at three video frames. The overall scene and subject arrangement remain similar, with visible differences in ground texture, background detail, and coat patterns.* + +This illustration is separate from the benchmark measurements above. + +## Tests and profiling + +Run the SVG tests from the repository root: + +```bash +python -m pytest -q \ + src/maxdiffusion/tests/wan/svg_attention_test.py \ + src/maxdiffusion/tests/wan/svg_balanced_rounding_test.py \ + src/maxdiffusion/tests/wan/svg_config_propagation_test.py \ + src/maxdiffusion/tests/wan/svg_head_local_test.py \ + src/maxdiffusion/tests/wan/wan_pipeline_signature_test.py \ + src/maxdiffusion/tests/ltx2/test_svg_config_propagation_ltx2.py \ + src/maxdiffusion/tests/ltx2/test_svg_attention_ltx2.py +``` + +The tests cover routing and layout, configuration propagation, schedules, sharding, unsupported configurations, and numerical agreement with attention references. Production-kernel tests cover sparse and density-one support, aligned and padded sequences, and natural and base-2 exponentials. + +For CPU semantics checks, set `JAX_PLATFORMS=cpu` and `XLA_FLAGS=--xla_force_host_platform_device_count=8`. Run the suite separately on an eight-device TPU host to exercise the compiled kernels. Tests restricted to one platform are skipped on the other. + +For profiling, enable `enable_jax_named_scopes=True` and capture a short warm denoising interval. Routing, placement, main attention, padding cleanup, merging, and restoration have named scopes, including `svg_route_profile`, `svg_layout_place`, `svg_union_main`, `svg_tail_cleanup`, `svg_lse_merge`, and `svg_layout_restore`. The `svg_kernel_c_tiles…` scope reports the fraction of physical tiles executed, which differs from attention-pair density and total transformer FLOP savings. diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml index 9c9c5432a..c9c8ff9b1 100644 --- a/src/maxdiffusion/configs/ltx2_3_video.yml +++ b/src/maxdiffusion/configs/ltx2_3_video.yml @@ -17,6 +17,24 @@ a2v_attention_kernel: 'flash' v2a_attention_kernel: 'dot_product' attention_sharding_uniform: True precision: 'bf16' +# Sparse VideoGen (SVG) Attention configuration +use_svg_attention: False +svg_spatial_density: 0.25 +svg_sample_max_row: 10000 +svg_profile_query_count: 64 +svg_profile_seed: 0 +svg_active_start_step: -1 +svg_active_end_step: -1 +svg_active_start_layer: -1 +svg_active_end_layer: -1 +svg_include_first_frame: True +svg_high_noise_density: -1.0 +svg_low_noise_density: -1.0 +# Tiling for the sparse SVG kernel only. Empty means "reuse flash_block_sizes", +# which is rarely what you want: the dense ring kernel is tuned for large kv +# compute blocks and the sparse kernel for small ones. Must stay a dict so the +# command line can override it with JSON. +svg_flash_block_sizes: {} scan_layers: True names_which_can_be_saved: [] names_which_can_be_offloaded: [] diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml index 23a4b104a..bf82085ef 100644 --- a/src/maxdiffusion/configs/ltx2_video.yml +++ b/src/maxdiffusion/configs/ltx2_video.yml @@ -17,6 +17,24 @@ a2v_attention_kernel: 'dot_product' v2a_attention_kernel: 'dot_product' attention_sharding_uniform: True precision: 'bf16' +# Sparse VideoGen (SVG) Attention configuration +use_svg_attention: False +svg_spatial_density: 0.25 +svg_sample_max_row: 10000 +svg_profile_query_count: 64 +svg_profile_seed: 0 +svg_active_start_step: -1 +svg_active_end_step: -1 +svg_active_start_layer: -1 +svg_active_end_layer: -1 +svg_include_first_frame: True +svg_high_noise_density: -1.0 +svg_low_noise_density: -1.0 +# Tiling for the sparse SVG kernel only. Empty means "reuse flash_block_sizes", +# which is rarely what you want: the dense ring kernel is tuned for large kv +# compute blocks and the sparse kernel for small ones. Must stay a dict so the +# command line can override it with JSON. +svg_flash_block_sizes: {} # For scanning transformer layers scan_layers: True diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index 2ef7aeba9..cef73a2d8 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -226,6 +226,26 @@ def ltx2_aot_metadata(config, pipeline, source_revision=None): "sharding", "weights_dtype", "activations_dtype", + "use_svg_attention", + "svg_implementation", + "svg_spatial_density", + "svg_sample_max_row", + "svg_profile_query_count", + "svg_profile_seed", + "svg_dense_layer_fraction", + "svg_dense_timestep_fraction", + "svg_active_start_step", + "svg_active_end_step", + "svg_active_start_layer", + "svg_active_end_layer", + "svg_num_train_timesteps", + "svg_num_layers", + "svg_include_first_frame", + "svg_global_stride", + "svg_global_offset", + "svg_high_noise_density", + "svg_low_noise_density", + "svg_flash_block_sizes", ] config_dict = {k: getattr(config, k, None) for k in config_keys} diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py index 3e55689eb..df2ea04c9 100644 --- a/src/maxdiffusion/models/ltx2/attention_ltx2.py +++ b/src/maxdiffusion/models/ltx2/attention_ltx2.py @@ -352,7 +352,58 @@ def __init__( use_base2_exp: bool = False, use_experimental_scheduler: bool = False, enable_jax_named_scopes: bool = False, + attention_config: Optional[dict] = None, ): + attention_config = { + "use_base2_exp": use_base2_exp, + "use_experimental_scheduler": use_experimental_scheduler, + "ulysses_shards": ulysses_shards, + "ulysses_attention_chunks": ulysses_attention_chunks, + "use_svg_attention": False, + "svg_implementation": "official_svg", + "svg_spatial_density": 0.25, + "svg_sample_max_row": 10000, + "svg_profile_query_count": 64, + "svg_profile_seed": 0, + "svg_dense_layer_fraction": 0.0, + "svg_dense_timestep_fraction": 0.0, + "svg_active_start_step": -1, + "svg_active_end_step": -1, + "svg_active_start_layer": -1, + "svg_active_end_layer": -1, + "svg_num_train_timesteps": 1000, + "svg_num_layers": 48, + "svg_include_first_frame": True, + "svg_global_stride": 0, + "svg_global_offset": 0, + "svg_high_noise_density": -1.0, + "svg_low_noise_density": -1.0, + "svg_flash_block_sizes": None, + **(attention_config or {}), + } + + self.is_self_attention = context_dim is None + self.use_svg_attention = bool(attention_config["use_svg_attention"]) and self.is_self_attention + self.svg_implementation = attention_config["svg_implementation"] + self.svg_spatial_density = attention_config["svg_spatial_density"] + self.svg_sample_max_row = attention_config["svg_sample_max_row"] + self.svg_profile_query_count = attention_config["svg_profile_query_count"] + self.svg_profile_seed = attention_config["svg_profile_seed"] + self.svg_dense_layer_fraction = attention_config["svg_dense_layer_fraction"] + self.svg_dense_timestep_fraction = attention_config["svg_dense_timestep_fraction"] + self.svg_active_start_step = attention_config["svg_active_start_step"] + self.svg_active_end_step = attention_config["svg_active_end_step"] + self.svg_active_start_layer = attention_config["svg_active_start_layer"] + self.svg_active_end_layer = attention_config["svg_active_end_layer"] + self.svg_num_train_timesteps = attention_config["svg_num_train_timesteps"] + self.svg_num_layers = attention_config["svg_num_layers"] + self.svg_include_first_frame = attention_config["svg_include_first_frame"] + self.svg_global_stride = attention_config["svg_global_stride"] + self.svg_global_offset = attention_config["svg_global_offset"] + self.svg_high_noise_density = attention_config["svg_high_noise_density"] + self.svg_low_noise_density = attention_config["svg_low_noise_density"] + self.svg_flash_block_sizes = attention_config["svg_flash_block_sizes"] + self.heads = heads self.rope_type = rope_type self.dim_head = dim_head @@ -542,10 +593,22 @@ def __call__( k_rotary_emb: Optional[Tuple[Array, Array]] = None, perturbation_mask: Optional[Array] = None, cached_kv: Optional[Tuple[Array, Array]] = None, + deterministic: bool = True, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, + svg_layer_index: Optional[int | jax.Array] = None, + svg_timestep: Optional[int | float | jax.Array] = None, + svg_step_index: Optional[int | jax.Array] = None, ) -> Array: # Determine context (Self or Cross) + is_self_attention = encoder_hidden_states is None context = encoder_hidden_states if encoder_hidden_states is not None else hidden_states + if self.use_svg_attention and is_self_attention: + if not deterministic: + raise ValueError("SVG attention supports deterministic inference only.") + if spatiotemporal_shape is None: + raise ValueError("SVG attention requires spatiotemporal_shape.") + # 1. Project and Norm with self.named_scope("QKV Projection"): query = self.to_q(hidden_states) @@ -586,8 +649,70 @@ def __call__( with self.named_scope("Attention and Output Project"): # 4. Attention - # NNXAttentionOp expects flattened input [B, S, InnerDim] for flash kernel - attn_output = self.attention_op.apply_attention(query=query, key=key, value=value, attention_mask=attention_mask) + if self.use_svg_attention and is_self_attention and spatiotemporal_shape is not None: + from maxdiffusion.models.wan.transformers import svg_attention + + is_active = svg_attention.is_svg_active( + step_index=svg_step_index, + layer_index=svg_layer_index, + timestep=svg_timestep, + start_step=self.svg_active_start_step, + end_step=self.svg_active_end_step, + start_layer=self.svg_active_start_layer, + end_layer=self.svg_active_end_layer, + dense_layer_fraction=self.svg_dense_layer_fraction, + dense_timestep_fraction=self.svg_dense_timestep_fraction, + num_train_timesteps=self.svg_num_train_timesteps, + num_layers=self.svg_num_layers, + ) + + def run_dense(_): + return self.attention_op.apply_attention( + query=query, + key=key, + value=value, + attention_mask=attention_mask, + ) + + def run_sparse_svg(_): + execution_band_width = svg_attention.svg_execution_band_width( + spatiotemporal_shape, + self.svg_spatial_density, + ) + sparse_config = { + "use_svg_attention": True, + "mask_type": "svg_spatial", + "band_width": execution_band_width, + "include_first_frame": self.svg_include_first_frame, + "global_stride": self.svg_global_stride, + "global_offset": self.svg_global_offset, + "profile_query_count": self.svg_profile_query_count, + "profile_seed": self.svg_profile_seed, + "sample_max_row": self.svg_sample_max_row, + "custom_flash_block_sizes": self.svg_flash_block_sizes, + "svg_step_index": svg_step_index, + "svg_layer_index": svg_layer_index, + "svg_timestep": svg_timestep, + } + return self.attention_op.apply_attention( + query=query, + key=key, + value=value, + attention_mask=attention_mask, + spatiotemporal_shape=spatiotemporal_shape, + sparse_config_override=sparse_config, + ) + + with self.named_scope("apply_attention"): + if isinstance(is_active, bool): + attn_output = run_sparse_svg(None) if is_active else run_dense(None) + else: + attn_output = jax.lax.cond(is_active, run_sparse_svg, run_dense, operand=None) + else: + with self.named_scope("apply_attention"): + attn_output = self.attention_op.apply_attention( + query=query, key=key, value=value, attention_mask=attention_mask + ) if perturbation_mask is not None: # value is [B, S, InnerDim] diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py index 1ac3c10c4..bdac39ca3 100644 --- a/src/maxdiffusion/models/ltx2/transformer_ltx2.py +++ b/src/maxdiffusion/models/ltx2/transformer_ltx2.py @@ -54,6 +54,9 @@ class LTX2StaticContext: audio_encoder_attention_mask: Optional[jax.Array] = None a2v_cross_attention_mask: Optional[jax.Array] = None v2a_cross_attention_mask: Optional[jax.Array] = None + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None + svg_timestep: Optional[int | float | jax.Array] = None + svg_step_index: Optional[int | jax.Array] = None @struct.dataclass @@ -63,6 +66,7 @@ class LTX2BlockContext: static: LTX2StaticContext perturbation_mask: Optional[jax.Array] = None layer_kv_cache: Optional[Mapping[str, FrozenDict]] = None + layer_index: Optional[int | jax.Array] = None def _canonicalize_attention_mask(mask: Optional[jax.Array], batch_size: int, name: str) -> Optional[jax.Array]: @@ -186,6 +190,7 @@ def __init__( use_base2_exp: bool = False, use_experimental_scheduler: bool = False, enable_jax_named_scopes: bool = False, + attention_config: Optional[dict] = None, ): self.dim = dim self.norm_eps = norm_eps @@ -230,6 +235,7 @@ def __init__( ulysses_attention_chunks=ulysses_attention_chunks, use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, + attention_config=attention_config, ) self.audio_norm1 = nnx.RMSNorm( @@ -263,6 +269,7 @@ def __init__( ulysses_attention_chunks=ulysses_attention_chunks, use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, + attention_config={"use_svg_attention": False}, ) # 2. Prompt Cross-Attention @@ -595,6 +602,10 @@ def __call__( encoder_hidden_states=None, rotary_emb=video_rotary_emb, perturbation_mask=perturbation_mask, + spatiotemporal_shape=ctx.static.spatiotemporal_shape, + svg_layer_index=ctx.layer_index, + svg_timestep=ctx.static.svg_timestep, + svg_step_index=ctx.static.svg_step_index, ) hidden_states = hidden_states + attn_hidden_states * gate_msa @@ -831,6 +842,7 @@ def __init__( use_base2_exp: bool = False, use_experimental_scheduler: bool = False, enable_jax_named_scopes: bool = False, + attention_config: Optional[dict] = None, **kwargs, ): self.spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks @@ -890,6 +902,13 @@ def __init__( self.flash_min_seq_length = flash_min_seq_length self.use_base2_exp = use_base2_exp self.use_experimental_scheduler = use_experimental_scheduler + self.attention_config = { + "use_base2_exp": use_base2_exp, + "use_experimental_scheduler": use_experimental_scheduler, + "ulysses_shards": ulysses_shards, + "ulysses_attention_chunks": ulysses_attention_chunks, + **(attention_config or {}), + } if sharding_specs is None: sharding_specs = get_sharding_specs("default", "ltx2_dit") @@ -1145,6 +1164,7 @@ def init_block(rngs): use_base2_exp=self.use_base2_exp, use_experimental_scheduler=self.use_experimental_scheduler, enable_jax_named_scopes=self.enable_jax_named_scopes, + attention_config=self.attention_config, ) if self.scan_layers: @@ -1189,6 +1209,7 @@ def init_block(rngs): use_base2_exp=self.use_base2_exp, use_experimental_scheduler=self.use_experimental_scheduler, enable_jax_named_scopes=self.enable_jax_named_scopes, + attention_config=self.attention_config, ) blocks.append(block) self.transformer_blocks = nnx.List(blocks) @@ -1399,6 +1420,7 @@ def __call__( cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, rope_cache: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, time_embed_cache: Optional[Dict[str, jax.Array]] = None, + svg_step_index: Optional[int | jax.Array] = None, ) -> Any: """ Forward pass for the full LTX2 Video/Audio Diffusion Transformer. @@ -1600,6 +1622,11 @@ def __call__( audio_encoder_hidden_states = audio_encoder_hidden_states.reshape(batch_size, -1, audio_hidden_states.shape[-1]) # 5. Run transformer blocks with self.named_scope("Transformer Blocks"): + if num_frames is not None and height is not None and width is not None: + spatiotemporal_shape = (num_frames // self.patch_size_t, height // self.patch_size, width // self.patch_size) + else: + spatiotemporal_shape = None + static_context = LTX2StaticContext( encoder_hidden_states=encoder_hidden_states, audio_encoder_hidden_states=audio_encoder_hidden_states, @@ -1620,6 +1647,9 @@ def __call__( a2v_cross_attention_mask=a2v_cross_attention_mask, v2a_cross_attention_mask=v2a_cross_attention_mask, modality_mask=modality_mask, + spatiotemporal_shape=spatiotemporal_shape, + svg_timestep=timestep, + svg_step_index=svg_step_index, ) if cached_kv is not None: @@ -1630,13 +1660,14 @@ def __call__( else: unstacked_kv_layers = [None] * self.num_layers - def apply_block_in_scan(block, hidden_states, audio_hidden_states, mask, layer_kv_cache): + def apply_block_in_scan(block, hidden_states, audio_hidden_states, mask, layer_kv_cache, layer_index=None): context = LTX2BlockContext( hidden_states=hidden_states, audio_hidden_states=audio_hidden_states, static=static_context, perturbation_mask=mask, layer_kv_cache=layer_kv_cache, + layer_index=layer_index, ) with self.named_scope("Transformer Layer"): hidden_states_out, audio_hidden_states_out = block(context) @@ -1645,24 +1676,28 @@ def apply_block_in_scan(block, hidden_states, audio_hidden_states, mask, layer_k audio_hidden_states_out.astype(audio_hidden_states.dtype), ) + layer_indices = jnp.arange(self.num_layers, dtype=jnp.int32) if perturbation_mask is None: if cached_kv is not None: def scan_fn_ltx2(carry, block_and_kv): - block, layer_kv_cache = block_and_kv + block, layer_kv_cache, layer_index = block_and_kv if isinstance(layer_kv_cache, dict): layer_kv_cache = FrozenDict(layer_kv_cache) hidden_states, audio_hidden_states, rngs_carry = carry hidden_states, audio_hidden_states = apply_block_in_scan( - block, hidden_states, audio_hidden_states, None, layer_kv_cache + block, hidden_states, audio_hidden_states, None, layer_kv_cache, layer_index=layer_index ) return (hidden_states, audio_hidden_states, rngs_carry), None else: - def scan_fn_ltx2(carry, block): + def scan_fn_ltx2(carry, block_and_idx): + block, layer_index = block_and_idx hidden_states, audio_hidden_states, rngs_carry = carry - hidden_states, audio_hidden_states = apply_block_in_scan(block, hidden_states, audio_hidden_states, None, None) + hidden_states, audio_hidden_states = apply_block_in_scan( + block, hidden_states, audio_hidden_states, None, None, layer_index=layer_index + ) return (hidden_states, audio_hidden_states, rngs_carry), None if self.scan_layers: @@ -1674,7 +1709,11 @@ def scan_fn_ltx2(carry, block): ) carry = (hidden_states, audio_hidden_states, nnx.Rngs(0)) - scan_input = (self.transformer_blocks, cached_kv) if cached_kv is not None else self.transformer_blocks + scan_input = ( + (self.transformer_blocks, cached_kv, layer_indices) + if cached_kv is not None + else (self.transformer_blocks, layer_indices) + ) (hidden_states, audio_hidden_states, _), _ = nnx.scan( rematted_scan_fn, length=self.num_layers, @@ -1685,7 +1724,7 @@ def scan_fn_ltx2(carry, block): else: for i, block in enumerate(self.transformer_blocks): hidden_states, audio_hidden_states = apply_block_in_scan( - block, hidden_states, audio_hidden_states, None, unstacked_kv_layers[i] + block, hidden_states, audio_hidden_states, None, unstacked_kv_layers[i], layer_index=i ) else: masks = jnp.ones((self.num_layers, batch_size, 1, 1), dtype=self.dtype) @@ -1697,21 +1736,23 @@ def scan_fn_ltx2(carry, block): if cached_kv is not None: def scan_fn_ltx23(carry, block_and_mask_and_kv): - block, mask, layer_kv_cache = block_and_mask_and_kv + block, mask, layer_kv_cache, layer_index = block_and_mask_and_kv if isinstance(layer_kv_cache, dict): layer_kv_cache = FrozenDict(layer_kv_cache) hidden_states, audio_hidden_states, rngs_carry = carry hidden_states, audio_hidden_states = apply_block_in_scan( - block, hidden_states, audio_hidden_states, mask, layer_kv_cache + block, hidden_states, audio_hidden_states, mask, layer_kv_cache, layer_index=layer_index ) return (hidden_states, audio_hidden_states, rngs_carry), None else: def scan_fn_ltx23(carry, block_and_mask): - block, mask = block_and_mask + block, mask, layer_index = block_and_mask hidden_states, audio_hidden_states, rngs_carry = carry - hidden_states, audio_hidden_states = apply_block_in_scan(block, hidden_states, audio_hidden_states, mask, None) + hidden_states, audio_hidden_states = apply_block_in_scan( + block, hidden_states, audio_hidden_states, mask, None, layer_index=layer_index + ) return (hidden_states, audio_hidden_states, rngs_carry), None if self.scan_layers: @@ -1723,9 +1764,9 @@ def scan_fn_ltx23(carry, block_and_mask): ) carry = (hidden_states, audio_hidden_states, nnx.Rngs(0)) scan_input = ( - (self.transformer_blocks, perturbation_mask_per_layer, cached_kv) + (self.transformer_blocks, perturbation_mask_per_layer, cached_kv, layer_indices) if cached_kv is not None - else (self.transformer_blocks, perturbation_mask_per_layer) + else (self.transformer_blocks, perturbation_mask_per_layer, layer_indices) ) (hidden_states, audio_hidden_states, _), _ = nnx.scan( rematted_scan_fn, @@ -1737,7 +1778,12 @@ def scan_fn_ltx23(carry, block_and_mask): else: for i, block in enumerate(self.transformer_blocks): hidden_states, audio_hidden_states = apply_block_in_scan( - block, hidden_states, audio_hidden_states, perturbation_mask_per_layer[i], unstacked_kv_layers[i] + block, + hidden_states, + audio_hidden_states, + perturbation_mask_per_layer[i], + unstacked_kv_layers[i], + layer_index=i, ) # 6. Output layers diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index eeaad473d..7e85e15e7 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -336,6 +336,39 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): dit_specs = get_sharding_specs(transformer_strategy, "ltx2_dit") ltx2_config["sharding_specs"] = dit_specs + high_density = float(getattr(config, "svg_high_noise_density", -1.0)) + low_density = float(getattr(config, "svg_low_noise_density", -1.0)) + expert_density = float(getattr(config, "svg_spatial_density", 0.25)) + + use_svg = bool(getattr(config, "use_svg_attention", False)) and (expert_density < 1.0) + + ltx2_config["attention_config"] = { + "use_base2_exp": getattr(config, "use_base2_exp", False), + "use_experimental_scheduler": getattr(config, "use_experimental_scheduler", False), + "ulysses_shards": getattr(config, "ulysses_shards", -1), + "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1), + "use_svg_attention": use_svg, + "svg_implementation": getattr(config, "svg_implementation", "official_svg"), + "svg_spatial_density": expert_density, + "svg_sample_max_row": getattr(config, "svg_sample_max_row", 10000), + "svg_profile_query_count": getattr(config, "svg_profile_query_count", 64), + "svg_profile_seed": getattr(config, "svg_profile_seed", 0), + "svg_dense_layer_fraction": getattr(config, "svg_dense_layer_fraction", 0.0), + "svg_dense_timestep_fraction": getattr(config, "svg_dense_timestep_fraction", 0.0), + "svg_active_start_step": getattr(config, "svg_active_start_step", -1), + "svg_active_end_step": getattr(config, "svg_active_end_step", -1), + "svg_active_start_layer": getattr(config, "svg_active_start_layer", -1), + "svg_active_end_layer": getattr(config, "svg_active_end_layer", -1), + "svg_num_train_timesteps": getattr(config, "svg_num_train_timesteps", 1000), + "svg_num_layers": getattr(config, "svg_num_layers", ltx2_config.get("num_layers", 28)), + "svg_include_first_frame": getattr(config, "svg_include_first_frame", True), + "svg_global_stride": getattr(config, "svg_global_stride", 0), + "svg_global_offset": getattr(config, "svg_global_offset", 0), + "svg_high_noise_density": high_density, + "svg_low_noise_density": low_density, + "svg_flash_block_sizes": getattr(config, "svg_flash_block_sizes", None) or None, + } + # 2. eval_shape p_model_factory = partial(create_model, ltx2_config=ltx2_config) transformer = nnx.eval_shape(p_model_factory, rngs=rngs) @@ -1832,6 +1865,10 @@ def __call__( if stg_scale > 0.0 and guidance_scale <= 1.0: raise ValueError("Spatio-temporal guidance requires guidance_scale > 1.0.") + if self.config and bool(getattr(self.config, "use_svg_attention", False)): + if getattr(self.config, "use_cfg_cache", False) or getattr(self.config, "use_magcache", False): + raise ValueError("SVG sparse attention cannot be combined with CFG cache or MagCache.") + # 2. Encode inputs (Text) t0_encode = time.perf_counter() ( @@ -2136,6 +2173,7 @@ def __call__( is_cfg_stg_mode=do_cfg and do_stg, kv_cache=kv_cache, rope_cache=rope_cache, + svg_step_index=i, ) latents_step, audio_latents_step = _select_guidance_latents( @@ -2437,6 +2475,7 @@ def transformer_forward_pass( kv_cache=None, rope_cache=None, time_embed_cache=None, + svg_step_index=None, ): """Forward pass for the transformer.""" # pylint: disable=too-many-positional-arguments,unused-argument @@ -2491,6 +2530,7 @@ def transformer_forward_pass( cached_kv=kv_cache, rope_cache=rope_cache, time_embed_cache=time_embed_cache, + svg_step_index=svg_step_index, ) return noise_pred, noise_pred_audio @@ -2598,9 +2638,9 @@ def run_diffusion_loop( def scan_body(carry, inputs): if use_kv_cache: - t, sigma_t, time_embed_cache_step = inputs + t, sigma_t, time_embed_cache_step, step_idx = inputs else: - t, sigma_t = inputs + t, sigma_t, step_idx = inputs time_embed_cache_step = None latents, audio_latents, s_state = carry @@ -2634,6 +2674,7 @@ def scan_body(carry, inputs): kv_cache=kv_cache, rope_cache=rope_cache, time_embed_cache=time_embed_cache_step, + svg_step_index=step_idx, ) latents_step, audio_latents_step = _select_guidance_latents( @@ -2680,10 +2721,11 @@ def scan_body(carry, inputs): initial_carry = (latents_jax, audio_latents_jax, scheduler_state) + step_indices = jnp.arange(len(timesteps_jax), dtype=jnp.int32) if use_kv_cache: - scan_inputs = (timesteps_jax, sigmas, time_embed_cache_full) + scan_inputs = (timesteps_jax, sigmas, time_embed_cache_full, step_indices) else: - scan_inputs = (timesteps_jax, sigmas) + scan_inputs = (timesteps_jax, sigmas, step_indices) final_carry, _ = jax.lax.scan(scan_body, initial_carry, scan_inputs) diff --git a/src/maxdiffusion/tests/ltx2/test_svg_attention_ltx2.py b/src/maxdiffusion/tests/ltx2/test_svg_attention_ltx2.py new file mode 100644 index 000000000..bec100710 --- /dev/null +++ b/src/maxdiffusion/tests/ltx2/test_svg_attention_ltx2.py @@ -0,0 +1,429 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Tests for Sparse VideoGen (SVG) attention in LTX2. +""" + +import unittest +from unittest.mock import patch, MagicMock + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh +from flax.linen import partitioning as nn_partitioning + +from maxdiffusion.models.ltx2.attention_ltx2 import LTX2Attention +from maxdiffusion.models.ltx2.transformer_ltx2 import ( + LTX2VideoTransformerBlock, + LTX2VideoTransformer3DModel, + LTX2StaticContext, + LTX2BlockContext, +) +from maxdiffusion.models.wan.transformers.svg_attention import is_svg_active + + +class LTX2SVGAttentionTest(unittest.TestCase): + + def setUp(self): + devices = np.array(jax.devices()[:1]).reshape((1, 1)) + self.mesh = Mesh(devices, ("data", "fsdp")) + self.rngs = nnx.Rngs(0) + self.logical_axis_rules = ( + ("activation_batch", ("data", "fsdp")), + ("activation_length", None), + ("activation_embed", None), + ) + + def test_is_svg_active_ltx2_boundaries(self): + """Verifies that SVG active predicate correctly checks step, layer, and density.""" + # Active range: step in [10, 30), layer in [1, 28) + # Outside active steps: step 5, 30, 35 -> False + self.assertFalse( + is_svg_active( + layer_index=5, + step_index=5, + timestep=500.0, + start_step=10, + end_step=30, + start_layer=1, + end_layer=28, + ) + ) + self.assertFalse( + is_svg_active( + layer_index=5, + step_index=30, + timestep=200.0, + start_step=10, + end_step=30, + start_layer=1, + end_layer=28, + ) + ) + + # Outside active layers: layer 0, layer 28 -> False + self.assertFalse( + is_svg_active( + layer_index=0, + step_index=15, + timestep=500.0, + start_step=10, + end_step=30, + start_layer=1, + end_layer=28, + ) + ) + self.assertFalse( + is_svg_active( + layer_index=28, + step_index=15, + timestep=500.0, + start_step=10, + end_step=30, + start_layer=1, + end_layer=28, + ) + ) + + # Inside active steps & layers: step 15, layer 5 -> True + self.assertTrue( + is_svg_active( + layer_index=5, + step_index=15, + timestep=500.0, + start_step=10, + end_step=30, + start_layer=1, + end_layer=28, + ) + ) + + # Dynamic execution under JIT with JAX Array inputs + @jax.jit + def check_dynamic(step, layer): + return is_svg_active( + layer_index=layer, + step_index=step, + timestep=500.0, + start_step=10, + end_step=30, + start_layer=1, + end_layer=28, + ) + + self.assertTrue(bool(check_dynamic(jnp.int32(15), jnp.int32(5)))) + self.assertFalse(bool(check_dynamic(jnp.int32(5), jnp.int32(5)))) + self.assertFalse(bool(check_dynamic(jnp.int32(15), jnp.int32(0)))) + + def test_ltx2_attention_forward_dense_and_sparse_dispatch(self): + """Tests LTX2Attention routing under inactive and active SVG steps.""" + B = 1 + F, H, W = 4, 8, 8 + seq_len = F * H * W + dim = 128 + num_heads = 4 + head_dim = dim // num_heads + + attention_config = { + "use_svg_attention": True, + "svg_spatial_density": 0.25, + "svg_active_start_step": 2, + "svg_active_end_step": 10, + "svg_active_start_layer": 0, + "svg_active_end_layer": 10, + "svg_sample_max_row": 100, + "svg_profile_query_count": 16, + } + + with self.mesh, nn_partitioning.axis_rules(self.logical_axis_rules): + attn = LTX2Attention( + rngs=self.rngs, + query_dim=dim, + context_dim=None, + heads=num_heads, + dim_head=head_dim, + attention_kernel="dot_product", + mesh=self.mesh, + attention_config=attention_config, + ) + + hidden_states = jax.random.normal(jax.random.key(1), (B, seq_len, dim), dtype=jnp.float32) + + # 1. Inactive step (step 0 < active_start_step 2) -> runs dense path successfully + out_dense = attn( + hidden_states=hidden_states, + spatiotemporal_shape=(F, H, W), + svg_layer_index=0, + svg_timestep=jnp.array([100.0]), + svg_step_index=0, + ) + self.assertEqual(out_dense.shape, (B, seq_len, dim)) + self.assertTrue(jnp.all(jnp.isfinite(out_dense))) + + # 2. Active step (step 3 in [2, 10)) -> triggers SVG branch which requires custom Ulysses backend + with self.assertRaisesRegex(ValueError, "Head-local SVG requires a custom Ulysses attention backend"): + attn( + hidden_states=hidden_states, + spatiotemporal_shape=(F, H, W), + svg_layer_index=0, + svg_timestep=jnp.array([100.0]), + svg_step_index=3, + ) + + # 3. Verify mock dispatch receives SVG sparse_config_override and spatiotemporal_shape + mock_apply = MagicMock(return_value=jnp.zeros((B, seq_len, dim), dtype=jnp.float32)) + with patch.object(attn.attention_op, "apply_attention", mock_apply): + _ = attn( + hidden_states=hidden_states, + spatiotemporal_shape=(F, H, W), + svg_layer_index=0, + svg_timestep=jnp.array([100.0]), + svg_step_index=3, + ) + mock_apply.assert_called_once() + _, kwargs = mock_apply.call_args + self.assertEqual(kwargs.get("spatiotemporal_shape"), (F, H, W)) + sp_cfg = kwargs.get("sparse_config_override") + self.assertIsNotNone(sp_cfg) + self.assertTrue(sp_cfg.get("use_svg_attention")) + + def test_ltx2_transformer_block_forward_svg_dispatch(self): + """Tests LTX2VideoTransformerBlock forward pass routing with SVG.""" + B = 1 + F, H, W = 4, 8, 8 + seq_len = F * H * W + audio_seq_len = 16 + dim = 64 + audio_dim = 64 + num_heads = 2 + head_dim = dim // num_heads + + attention_config = { + "use_svg_attention": True, + "svg_spatial_density": 0.25, + "svg_active_start_step": 2, + "svg_active_end_step": 10, + "svg_active_start_layer": 0, + "svg_active_end_layer": 10, + "svg_sample_max_row": 100, + "svg_profile_query_count": 16, + } + + with self.mesh, nn_partitioning.axis_rules(self.logical_axis_rules): + block = LTX2VideoTransformerBlock( + rngs=self.rngs, + dim=dim, + num_attention_heads=num_heads, + attention_head_dim=head_dim, + cross_attention_dim=dim, + audio_dim=audio_dim, + audio_num_attention_heads=num_heads, + audio_attention_head_dim=head_dim, + audio_cross_attention_dim=audio_dim, + attention_kernel="dot_product", + a2v_attention_kernel="dot_product", + v2a_attention_kernel="dot_product", + mesh=self.mesh, + attention_config=attention_config, + ) + + hidden_states = jax.random.normal(jax.random.key(1), (B, seq_len, dim), dtype=jnp.float32) + audio_hidden_states = jax.random.normal(jax.random.key(2), (B, audio_seq_len, audio_dim), dtype=jnp.float32) + encoder_hidden_states = jax.random.normal(jax.random.key(3), (B, 16, dim), dtype=jnp.float32) + audio_encoder_hidden_states = jax.random.normal(jax.random.key(4), (B, 16, audio_dim), dtype=jnp.float32) + + # Inactive step 0 -> dense execution succeeds + static_ctx_inactive = LTX2StaticContext( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + temb=jnp.zeros((B, 6 * dim)), + temb_audio=jnp.zeros((B, 6 * audio_dim)), + temb_ca_scale_shift=jnp.zeros((B, 4 * dim)), + temb_ca_audio_scale_shift=jnp.zeros((B, 4 * audio_dim)), + temb_ca_gate=jnp.zeros((B, 1 * dim)), + temb_ca_audio_gate=jnp.zeros((B, 1 * audio_dim)), + spatiotemporal_shape=(F, H, W), + svg_timestep=jnp.array([100.0]), + svg_step_index=0, + ) + block_ctx_inactive = LTX2BlockContext( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + static=static_ctx_inactive, + layer_index=0, + ) + out_h, out_a = block(block_ctx_inactive) + self.assertEqual(out_h.shape, (B, seq_len, dim)) + self.assertEqual(out_a.shape, (B, audio_seq_len, audio_dim)) + + # Active step 3 -> triggers SVG + static_ctx_active = LTX2StaticContext( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + temb=jnp.zeros((B, 6 * dim)), + temb_audio=jnp.zeros((B, 6 * audio_dim)), + temb_ca_scale_shift=jnp.zeros((B, 4 * dim)), + temb_ca_audio_scale_shift=jnp.zeros((B, 4 * audio_dim)), + temb_ca_gate=jnp.zeros((B, 1 * dim)), + temb_ca_audio_gate=jnp.zeros((B, 1 * audio_dim)), + spatiotemporal_shape=(F, H, W), + svg_timestep=jnp.array([100.0]), + svg_step_index=3, + ) + block_ctx_active = LTX2BlockContext( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + static=static_ctx_active, + layer_index=0, + ) + with self.assertRaisesRegex(ValueError, "Head-local SVG requires a custom Ulysses attention backend"): + block(block_ctx_active) + + def test_ltx2_model_full_forward_with_svg(self): + """Tests LTX2VideoTransformer3DModel full forward pass with SVG configuration.""" + B = 1 + F, H, W = 2, 8, 8 + seq_len = F * H * W + audio_seq_len = 16 + in_channels = 8 + out_channels = 8 + audio_in_channels = 4 + num_heads = 2 + head_dim = 32 + + # Step range [5, 15) so step 0 is inactive and completes full dense pass, while step 6 triggers SVG + attention_config = { + "use_svg_attention": True, + "svg_spatial_density": 0.25, + "svg_active_start_step": 5, + "svg_active_end_step": 15, + "svg_active_start_layer": 0, + "svg_active_end_layer": 10, + "svg_sample_max_row": 100, + "svg_profile_query_count": 16, + } + + with self.mesh, nn_partitioning.axis_rules(self.logical_axis_rules): + # Non-scanned blocks: step 0 statically resolves inactive and executes dense + model_unscanned = LTX2VideoTransformer3DModel( + rngs=nnx.Rngs(0), + in_channels=in_channels, + out_channels=out_channels, + patch_size=1, + patch_size_t=1, + num_attention_heads=num_heads, + attention_head_dim=head_dim, + cross_attention_dim=num_heads * head_dim, + caption_channels=16, + audio_in_channels=audio_in_channels, + audio_out_channels=audio_in_channels, + audio_num_attention_heads=num_heads, + audio_attention_head_dim=head_dim, + audio_cross_attention_dim=num_heads * head_dim, + num_layers=2, + mesh=self.mesh, + attention_kernel="dot_product", + a2v_attention_kernel="dot_product", + v2a_attention_kernel="dot_product", + scan_layers=False, + attention_config=attention_config, + ) + + hidden_states = jax.random.normal(jax.random.key(10), (B, seq_len, in_channels), dtype=jnp.float32) + audio_hidden_states = jax.random.normal(jax.random.key(11), (B, audio_seq_len, audio_in_channels), dtype=jnp.float32) + timestep = jnp.array([1.0]) + encoder_hidden_states = jax.random.normal(jax.random.key(12), (B, 16, 16), dtype=jnp.float32) + audio_encoder_hidden_states = jax.random.normal(jax.random.key(13), (B, 16, 16), dtype=jnp.float32) + + # 1. Inactive step (svg_step_index=0) -> executes full dense pass successfully + output = model_unscanned( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + timestep=timestep, + num_frames=F, + height=H, + width=W, + audio_num_frames=audio_seq_len, + svg_step_index=0, + return_dict=True, + ) + + self.assertEqual(output["sample"].shape, (B, seq_len, out_channels)) + self.assertEqual(output["audio_sample"].shape, (B, audio_seq_len, audio_in_channels)) + self.assertTrue(jnp.all(jnp.isfinite(output["sample"]))) + self.assertTrue(jnp.all(jnp.isfinite(output["audio_sample"]))) + + # 2. Active step (svg_step_index=6) -> triggers SVG routing and checks backend + with self.assertRaisesRegex(ValueError, "Head-local SVG requires a custom Ulysses attention backend"): + model_unscanned( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + timestep=timestep, + num_frames=F, + height=H, + width=W, + audio_num_frames=audio_seq_len, + svg_step_index=6, + return_dict=True, + ) + + # 3. Scanned blocks: fail-closed validation when backend is not a custom Ulysses kernel + model_scanned = LTX2VideoTransformer3DModel( + rngs=nnx.Rngs(0), + in_channels=in_channels, + out_channels=out_channels, + patch_size=1, + patch_size_t=1, + num_attention_heads=num_heads, + attention_head_dim=head_dim, + cross_attention_dim=num_heads * head_dim, + caption_channels=16, + audio_in_channels=audio_in_channels, + audio_out_channels=audio_in_channels, + audio_num_attention_heads=num_heads, + audio_attention_head_dim=head_dim, + audio_cross_attention_dim=num_heads * head_dim, + num_layers=2, + mesh=self.mesh, + attention_kernel="dot_product", + a2v_attention_kernel="dot_product", + v2a_attention_kernel="dot_product", + scan_layers=True, + attention_config=attention_config, + ) + with self.assertRaisesRegex(ValueError, "Head-local SVG requires a custom Ulysses attention backend"): + model_scanned( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + timestep=timestep, + num_frames=F, + height=H, + width=W, + audio_num_frames=audio_seq_len, + svg_step_index=6, + return_dict=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/ltx2/test_svg_config_propagation_ltx2.py b/src/maxdiffusion/tests/ltx2/test_svg_config_propagation_ltx2.py new file mode 100644 index 000000000..0a7d668df --- /dev/null +++ b/src/maxdiffusion/tests/ltx2/test_svg_config_propagation_ltx2.py @@ -0,0 +1,209 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Configuration propagation tests for Sparse VideoGen attention in LTX2. +""" + +from types import SimpleNamespace +import unittest +from unittest.mock import MagicMock + +from flax import nnx +import jax +import numpy as np +from jax.sharding import Mesh + +from maxdiffusion.pipelines.ltx2.ltx2_pipeline import create_sharded_logical_transformer, LTX2Pipeline + + +class LTX2SVGConfigPropagationTest(unittest.TestCase): + + def _get_test_ltx2_config(self): + return { + "num_layers": 2, + "num_attention_heads": 2, + "attention_head_dim": 32, + "cross_attention_dim": 64, + "in_channels": 16, + "out_channels": 16, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 32, + "audio_cross_attention_dim": 64, + "caption_channels": 32, + "patch_size": 1, + "patch_size_t": 1, + "pos_embed_max_pos": 16, + "base_height": 32, + "base_width": 32, + "audio_pos_embed_max_pos": 16, + "audio_sampling_rate": 16000, + "audio_hop_length": 160, + "audio_scale_factor": 1.0, + } + + def test_svg_config_propagation_through_transformer_construction(self): + test_ltx2_config = self._get_test_ltx2_config() + devices = np.array(jax.devices()[:1]).reshape((1, 1)) + mesh = Mesh(devices, ("data", "fsdp")) + rngs = nnx.Rngs(0) + + cfg = SimpleNamespace( + use_svg_attention=True, + svg_spatial_density=0.25, + svg_active_start_step=8, + svg_active_end_step=30, + svg_active_start_layer=1, + svg_active_end_layer=28, + precision="DEFAULT", + flash_block_sizes={}, + activations_dtype="bfloat16", + weights_dtype="bfloat16", + attention="dot_product", + a2v_attention_kernel="dot_product", + v2a_attention_kernel="dot_product", + remat_policy="none", + names_which_can_be_saved=[], + names_which_can_be_offloaded=[], + flash_min_seq_length=0, + dropout=0.0, + scan_layers=False, + enable_jax_named_scopes=False, + use_base2_exp=False, + use_experimental_scheduler=False, + logical_axis_rules=(), + ) + + m = create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=rngs, + config=cfg, + restored_checkpoint={"ltx2_config": dict(test_ltx2_config), "ltx2_state": {}}, + subfolder="", + ) + + # Video self-attention (attn1) must have SVG enabled + first_block = m.transformer_blocks[0] + self.assertTrue(first_block.attn1.use_svg_attention) + self.assertEqual(first_block.attn1.svg_spatial_density, 0.25) + self.assertEqual(first_block.attn1.svg_active_start_step, 8) + self.assertEqual(first_block.attn1.svg_active_end_step, 30) + self.assertEqual(first_block.attn1.svg_active_start_layer, 1) + self.assertEqual(first_block.attn1.svg_active_end_layer, 28) + + # Audio self-attention (audio_attn1) and cross attentions must remain dense + self.assertFalse(first_block.audio_attn1.use_svg_attention) + self.assertFalse(first_block.attn2.use_svg_attention) + self.assertFalse(first_block.audio_attn2.use_svg_attention) + self.assertFalse(first_block.audio_to_video_attn.use_svg_attention) + self.assertFalse(first_block.video_to_audio_attn.use_svg_attention) + + def test_svg_block_sizes_are_independent_of_the_dense_block_sizes(self): + test_ltx2_config = self._get_test_ltx2_config() + devices = np.array(jax.devices()[:1]).reshape((1, 1)) + mesh = Mesh(devices, ("data", "fsdp")) + sparse_blocks = {"block_q": 3328, "block_kv": 2816, "block_kv_compute": 256, "block_kv_compute_in": 256} + common = { + "use_svg_attention": True, + "svg_spatial_density": 0.25, + "precision": "DEFAULT", + "activations_dtype": "bfloat16", + "weights_dtype": "bfloat16", + "attention": "dot_product", + "a2v_attention_kernel": "dot_product", + "v2a_attention_kernel": "dot_product", + "remat_policy": "none", + "names_which_can_be_saved": [], + "names_which_can_be_offloaded": [], + "flash_min_seq_length": 0, + "dropout": 0.0, + "scan_layers": False, + "enable_jax_named_scopes": False, + "use_base2_exp": False, + "use_experimental_scheduler": False, + "logical_axis_rules": (), + } + + def build(**extra): + return create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=nnx.Rngs(0), + config=SimpleNamespace(**common, **extra), + restored_checkpoint={"ltx2_config": dict(test_ltx2_config), "ltx2_state": {}}, + subfolder="", + ) + + tuned = build(flash_block_sizes={}, svg_flash_block_sizes=sparse_blocks) + self.assertEqual(tuned.transformer_blocks[0].attn1.svg_flash_block_sizes, sparse_blocks) + + default = build(flash_block_sizes={}, svg_flash_block_sizes={}) + self.assertIsNone(default.transformer_blocks[0].attn1.svg_flash_block_sizes) + + absent = build(flash_block_sizes={}) + self.assertIsNone(absent.transformer_blocks[0].attn1.svg_flash_block_sizes) + + def test_svg_disabled_when_use_svg_attention_is_false(self): + test_ltx2_config = self._get_test_ltx2_config() + devices = np.array(jax.devices()[:1]).reshape((1, 1)) + mesh = Mesh(devices, ("data", "fsdp")) + + cfg_dense = SimpleNamespace( + use_svg_attention=False, + precision="DEFAULT", + flash_block_sizes={}, + activations_dtype="bfloat16", + weights_dtype="bfloat16", + attention="dot_product", + a2v_attention_kernel="dot_product", + v2a_attention_kernel="dot_product", + remat_policy="none", + names_which_can_be_saved=[], + names_which_can_be_offloaded=[], + flash_min_seq_length=0, + dropout=0.0, + scan_layers=False, + enable_jax_named_scopes=False, + use_base2_exp=False, + use_experimental_scheduler=False, + logical_axis_rules=(), + ) + m_dense = create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=nnx.Rngs(0), + config=cfg_dense, + restored_checkpoint={"ltx2_config": dict(test_ltx2_config), "ltx2_state": {}}, + subfolder="", + ) + self.assertFalse(m_dense.transformer_blocks[0].attn1.use_svg_attention) + + def test_cfg_and_magcache_incompatibility_validation(self): + pipeline = LTX2Pipeline.__new__(LTX2Pipeline) + pipeline.check_inputs = MagicMock() + pipeline.config = SimpleNamespace(use_svg_attention=True, use_cfg_cache=True, use_magcache=False) + + with self.assertRaisesRegex(ValueError, "SVG sparse attention cannot be combined with CFG cache or MagCache"): + pipeline(prompt="test prompt") + + pipeline.config = SimpleNamespace(use_svg_attention=True, use_cfg_cache=False, use_magcache=True) + with self.assertRaisesRegex(ValueError, "SVG sparse attention cannot be combined with CFG cache or MagCache"): + pipeline(prompt="test prompt") + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/utils/ltx2_block_benchmark.py b/src/maxdiffusion/utils/ltx2_block_benchmark.py index f15ebaaa5..673e6a1e7 100644 --- a/src/maxdiffusion/utils/ltx2_block_benchmark.py +++ b/src/maxdiffusion/utils/ltx2_block_benchmark.py @@ -256,6 +256,13 @@ def _build_model(self, bq, bkv, cmp): remat_policy=getattr(c, "remat_policy", "NONE"), scan_layers=False, num_layers=1, + attention_config={ + "use_base2_exp": getattr(c, "use_base2_exp", False), + "use_experimental_scheduler": getattr(c, "use_experimental_scheduler", False), + "ulysses_shards": getattr(c, "ulysses_shards", -1), + "ulysses_attention_chunks": getattr(c, "ulysses_attention_chunks", 1), + "use_svg_attention": getattr(c, "use_svg_attention", False), + }, ) model = LTX2VideoTransformer3DModel(**ltx2_config, rngs=nnx.Rngs(params=0)) gd, state, rest = nnx.split(model, nnx.Param, ...) From 51f51a6b0391240a865ba1f06c16e81c97b87d88 Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Fri, 18 Sep 2026 06:23:49 +0000 Subject: [PATCH 4/8] Fix custom block size fallback and enable framewise VAE tiling for LTX-2 --- src/maxdiffusion/configs/ltx2_3_video.yml | 2 ++ src/maxdiffusion/configs/ltx2_video.yml | 2 ++ src/maxdiffusion/models/attention_flax.py | 11 +++++++---- src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py | 2 ++ 4 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml index c9c8ff9b1..8c01da9d2 100644 --- a/src/maxdiffusion/configs/ltx2_3_video.yml +++ b/src/maxdiffusion/configs/ltx2_3_video.yml @@ -143,6 +143,8 @@ profiler_steps: 5 enable_jax_named_scopes: False replicate_vae: False +enable_vae_tiling: False +enable_vae_slicing: False run_text_encoder_on_tpu: False # Dynamically disables VAE slicing and distributes the batch dimension to avoid HBM OOM for larger batch sizes. diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml index bf82085ef..4778e285c 100644 --- a/src/maxdiffusion/configs/ltx2_video.yml +++ b/src/maxdiffusion/configs/ltx2_video.yml @@ -149,6 +149,8 @@ enable_jax_named_scopes: False replicate_vae: False use_bwe: False +enable_vae_tiling: False +enable_vae_slicing: False run_text_encoder_on_tpu: False # Dynamically disables VAE slicing and distributes the batch dimension to avoid HBM OOM for larger batch sizes. diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 3db542586..de6712235 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -381,7 +381,7 @@ def _extract_custom_block_sizes(flash_block_sizes): bq = 4864 bkv = 1024 bkv_compute = 1024 - bkv_compute_in = 1024 + bkv_compute_in = None heads_per_tile = 1 vmem_limit_bytes = None if flash_block_sizes is not None: @@ -390,14 +390,14 @@ def _extract_custom_block_sizes(flash_block_sizes): bq = get("block_q", None) or bq bkv = get("block_kv", None) or bkv bkv_compute = get("block_kv_compute", None) or bkv_compute - bkv_compute_in = get("block_kv_compute_in", None) or bkv_compute_in + bkv_compute_in = get("block_kv_compute_in", None) heads_per_tile = get("heads_per_tile", None) or heads_per_tile vmem_limit_bytes = get("vmem_limit_bytes", None) or vmem_limit_bytes else: bq = getattr(flash_block_sizes, "block_q", None) or bq bkv = getattr(flash_block_sizes, "block_kv", None) or bkv bkv_compute = getattr(flash_block_sizes, "block_kv_compute", None) or bkv_compute - bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", None) or bkv_compute_in + bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", None) heads_per_tile = getattr(flash_block_sizes, "heads_per_tile", None) or heads_per_tile vmem_limit_bytes = getattr(flash_block_sizes, "vmem_limit_bytes", None) or vmem_limit_bytes # A BlockSizes object carries heads_per_tile=None when the config dict omitted @@ -405,7 +405,10 @@ def _extract_custom_block_sizes(flash_block_sizes): # to 1 (the custom-kernel default) to keep the `heads_per_tile > 1` guards safe. if heads_per_tile is None: heads_per_tile = 1 - bkv_compute_in = min(bkv_compute, bkv_compute_in) + if bkv_compute_in is None: + bkv_compute_in = bkv_compute + else: + bkv_compute_in = min(bkv_compute, bkv_compute_in) return bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index 7e85e15e7..fefba2ee4 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -591,9 +591,11 @@ def enable_vae_tiling(self): if hasattr(self.vae, "enable_tiling"): self.vae.enable_tiling() self.vae.use_tiling = True + self.vae.use_framewise_decoding = True def disable_vae_tiling(self): self.vae.use_tiling = False + self.vae.use_framewise_decoding = False @classmethod def load_tokenizer(cls, config: HyperParameters): From c9fbddb2872513506e90cb427c0ce1be7f06a524 Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Mon, 28 Sep 2026 09:52:17 +0000 Subject: [PATCH 5/8] ltx2_block_benchmark: set scoped VMEM limit per TPU generation xla_tpu_scoped_vmem_limit_kib was hardcoded to 32768 (TPU v6e) in LIBTPU_INIT_ARGS. Pick it from jax.devices()[0].device_kind instead (v6e: 32 MiB, v7x: 64 MiB) and pass it as a jax.jit compile option, since LIBTPU_INIT_ARGS is only read when jax.devices() initializes libtpu. nnx.jit has no compiler_options, so _forward now uses jax.jit over nnx.split state. --- .../utils/ltx2_block_benchmark.py | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/src/maxdiffusion/utils/ltx2_block_benchmark.py b/src/maxdiffusion/utils/ltx2_block_benchmark.py index 673e6a1e7..ab306880e 100644 --- a/src/maxdiffusion/utils/ltx2_block_benchmark.py +++ b/src/maxdiffusion/utils/ltx2_block_benchmark.py @@ -32,7 +32,6 @@ "--xla_tpu_enable_async_collective_fusion_multiple_steps=true", "--xla_tpu_overlap_compute_collective_tc=true", "--xla_enable_async_all_gather=true", - "--xla_tpu_scoped_vmem_limit_kib=32768", "--xla_tpu_enable_async_all_to_all=true", "--xla_tpu_enable_all_experimental_scheduler_features=true", "--xla_tpu_enable_latency_hiding_scheduler=true", @@ -75,9 +74,18 @@ def tiled_seq_len(full_seq: int, attention: str, context_shards: int, ulysses_sh return local_tiled_seq_len(full_seq, attention, context_shards, ulysses_shards) -@functools.partial(nnx.jit, static_argnames=("num_frames", "height", "width")) +# xla_tpu_scoped_vmem_limit_kib per jax device_kind. Passed as a compile option +# (not via LIBTPU_INIT_ARGS) because jax.devices() initializes libtpu, after +# which LIBTPU_INIT_ARGS is no longer read. +_SCOPED_VMEM_LIMIT_KIB = { + "TPU v6 lite": 32768, + "TPU7x": 65536, +} + + def _forward( - model, + graphdef, + state, latents, timestep, prompt_embeds, @@ -90,6 +98,7 @@ def _forward( height, width, ): + model = nnx.merge(graphdef, state) return model( hidden_states=latents, audio_hidden_states=audio_latents, @@ -134,6 +143,10 @@ def __init__( self._ulysses_shards = int(getattr(config, "ulysses_shards", 1) or 1) self._vmem = int(vmem_limit_bytes) self.label = f"ltx2/{self._attention}/u{self._ulysses_shards}" + device_kind = jax.devices()[0].device_kind + self._compiler_options = { + "xla_tpu_scoped_vmem_limit_kib": str(_SCOPED_VMEM_LIMIT_KIB[device_kind]), + } self._lf = (num_frames - 1) // 8 + 1 self._lh, self._lw = height // 32, width // 32 @@ -176,6 +189,12 @@ def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2): try: with self._mesh: model = self._build_model(bq, bkv, cmp) + graphdef, state = nnx.split(model) + forward = jax.jit( + functools.partial(_forward, graphdef), + static_argnames=("num_frames", "height", "width"), + compiler_options=self._compiler_options, + ) ( latents, timestep, @@ -187,8 +206,8 @@ def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2): ) = self._inputs with self._mesh, nn_partitioning.axis_rules(self._rules): mean, std, times, compile_ms = time_callable( - lambda: _forward( - model, + lambda: forward( + state, latents, timestep, prompt_embeds, From 49163e95987f012362efd74f94169d7822916aad Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Mon, 28 Sep 2026 10:03:29 +0000 Subject: [PATCH 6/8] ltx2: fall back to defaults when optional SVG config values are None --- .../pipelines/ltx2/ltx2_pipeline.py | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index fefba2ee4..46a7e529d 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -336,9 +336,14 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): dit_specs = get_sharding_specs(transformer_strategy, "ltx2_dit") ltx2_config["sharding_specs"] = dit_specs - high_density = float(getattr(config, "svg_high_noise_density", -1.0)) - low_density = float(getattr(config, "svg_low_noise_density", -1.0)) - expert_density = float(getattr(config, "svg_spatial_density", 0.25)) + def _cfg(name, default): + """Returns config., falling back to default if missing or None.""" + value = getattr(config, name, None) + return default if value is None else value + + high_density = float(_cfg("svg_high_noise_density", -1.0)) + low_density = float(_cfg("svg_low_noise_density", -1.0)) + expert_density = float(_cfg("svg_spatial_density", 0.25)) use_svg = bool(getattr(config, "use_svg_attention", False)) and (expert_density < 1.0) @@ -355,12 +360,12 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): "svg_profile_seed": getattr(config, "svg_profile_seed", 0), "svg_dense_layer_fraction": getattr(config, "svg_dense_layer_fraction", 0.0), "svg_dense_timestep_fraction": getattr(config, "svg_dense_timestep_fraction", 0.0), - "svg_active_start_step": getattr(config, "svg_active_start_step", -1), - "svg_active_end_step": getattr(config, "svg_active_end_step", -1), - "svg_active_start_layer": getattr(config, "svg_active_start_layer", -1), - "svg_active_end_layer": getattr(config, "svg_active_end_layer", -1), - "svg_num_train_timesteps": getattr(config, "svg_num_train_timesteps", 1000), - "svg_num_layers": getattr(config, "svg_num_layers", ltx2_config.get("num_layers", 28)), + "svg_active_start_step": _cfg("svg_active_start_step", -1), + "svg_active_end_step": _cfg("svg_active_end_step", -1), + "svg_active_start_layer": _cfg("svg_active_start_layer", -1), + "svg_active_end_layer": _cfg("svg_active_end_layer", -1), + "svg_num_train_timesteps": _cfg("svg_num_train_timesteps", 1000), + "svg_num_layers": _cfg("svg_num_layers", ltx2_config.get("num_layers", 28)), "svg_include_first_frame": getattr(config, "svg_include_first_frame", True), "svg_global_stride": getattr(config, "svg_global_stride", 0), "svg_global_offset": getattr(config, "svg_global_offset", 0), From e47f86b3715df805a0bd2538b97b9cc1d0a8afa9 Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Mon, 28 Sep 2026 11:31:23 +0000 Subject: [PATCH 7/8] linter fixes --- src/maxdiffusion/aot_cache.py | 9 +++------ src/maxdiffusion/models/ltx2/attention_ltx2.py | 4 +--- src/maxdiffusion/pipelines/wan/wan_pipeline.py | 10 ++-------- src/maxdiffusion/tests/aot_cache_test.py | 4 +--- src/maxdiffusion/tests/wan/svg_attention_test.py | 12 +++--------- 5 files changed, 10 insertions(+), 29 deletions(-) diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py index 0e7cc344e..1c27c2c92 100644 --- a/src/maxdiffusion/aot_cache.py +++ b/src/maxdiffusion/aot_cache.py @@ -131,9 +131,8 @@ def extract_svg_meta(config: Any, pipeline: Any = None) -> dict[str, Any]: t = getattr(pipeline, attr, None) if t is not None: t_cfg = getattr(t, "config", None) - attn_cfg = ( - getattr(t_cfg, "attention_config", None) - or (t_cfg.get("attention_config") if isinstance(t_cfg, dict) else None) + attn_cfg = getattr(t_cfg, "attention_config", None) or ( + t_cfg.get("attention_config") if isinstance(t_cfg, dict) else None ) if isinstance(attn_cfg, dict): meta[f"{attr}_svg_config"] = json.dumps( @@ -162,9 +161,7 @@ def _dynamic_signature(args: tuple, kwargs: dict) -> str: Array leaves contribute shape/dtype; non-array leaves (python scalars, None flags) contribute an address-stripped repr. """ - leaves_with_paths = jax.tree_util.tree_flatten_with_path( - (args, kwargs), is_leaf=_is_graphdef - )[0] + leaves_with_paths = jax.tree_util.tree_flatten_with_path((args, kwargs), is_leaf=_is_graphdef)[0] parts = [] for path, leaf in leaves_with_paths: if _is_graphdef(leaf): diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py index df2ea04c9..16e32c727 100644 --- a/src/maxdiffusion/models/ltx2/attention_ltx2.py +++ b/src/maxdiffusion/models/ltx2/attention_ltx2.py @@ -710,9 +710,7 @@ def run_sparse_svg(_): attn_output = jax.lax.cond(is_active, run_sparse_svg, run_dense, operand=None) else: with self.named_scope("apply_attention"): - attn_output = self.attention_op.apply_attention( - query=query, key=key, value=value, attention_mask=attention_mask - ) + attn_output = self.attention_op.apply_attention(query=query, key=key, value=value, attention_mask=attention_mask) if perturbation_mask is not None: # value is [B, S, InnerDim] diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index fd673ca30..4006e2fd1 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -1422,10 +1422,7 @@ def _has_svg_enabled(obj: Any) -> bool: return True cfg = getattr(obj, "config", None) if cfg is not None: - attn_cfg = ( - getattr(cfg, "attention_config", None) - or (cfg.get("attention_config") if isinstance(cfg, dict) else None) - ) + attn_cfg = getattr(cfg, "attention_config", None) or (cfg.get("attention_config") if isinstance(cfg, dict) else None) if isinstance(attn_cfg, dict) and bool(attn_cfg.get("use_svg_attention", False)): return True if hasattr(obj, "attributes") and hasattr(obj, "nodes"): @@ -1434,10 +1431,7 @@ def _has_svg_enabled(obj: Any) -> bool: return True if k == "config": val = getattr(v, "value", None) - ac = ( - getattr(val, "attention_config", None) - or (val.get("attention_config") if isinstance(val, dict) else None) - ) + ac = getattr(val, "attention_config", None) or (val.get("attention_config") if isinstance(val, dict) else None) if isinstance(ac, dict) and bool(ac.get("use_svg_attention", False)): return True return False diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py index 25ddd1ca4..a0a43170d 100644 --- a/src/maxdiffusion/tests/aot_cache_test.py +++ b/src/maxdiffusion/tests/aot_cache_test.py @@ -237,9 +237,7 @@ def test_extract_svg_meta_changes_fingerprint_on_svg_changes(self): def test_step_array_preserves_dynamic_signature_across_steps(self): sigs = { - aot_cache._dynamic_signature( - (self._a,), {"svg_step_index": jnp.asarray(step, dtype=jnp.int32)} - ) + aot_cache._dynamic_signature((self._a,), {"svg_step_index": jnp.asarray(step, dtype=jnp.int32)}) for step in range(40) } self.assertEqual(len(sigs), 1) diff --git a/src/maxdiffusion/tests/wan/svg_attention_test.py b/src/maxdiffusion/tests/wan/svg_attention_test.py index 638d03ac0..1be65b03e 100644 --- a/src/maxdiffusion/tests/wan/svg_attention_test.py +++ b/src/maxdiffusion/tests/wan/svg_attention_test.py @@ -598,21 +598,15 @@ def test_partially_specified_schedules(self): # 3. Fully specified schedule: steps [11, 40), layers [1, 40) self.assertIs( - svg_attention.is_svg_active( - step_index=11, layer_index=0, start_step=11, end_step=40, start_layer=1, end_layer=40 - ), + svg_attention.is_svg_active(step_index=11, layer_index=0, start_step=11, end_step=40, start_layer=1, end_layer=40), False, ) self.assertIs( - svg_attention.is_svg_active( - step_index=10, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40 - ), + svg_attention.is_svg_active(step_index=10, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40), False, ) self.assertIs( - svg_attention.is_svg_active( - step_index=11, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40 - ), + svg_attention.is_svg_active(step_index=11, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40), True, ) From 65f1722032a89c65a279abae3b5b037c42495ed2 Mon Sep 17 00:00:00 2001 From: Jitendra Jalwaniya Date: Tue, 29 Sep 2026 06:46:48 +0000 Subject: [PATCH 8/8] ltx2: address PR review feedback on SVG config, AOT step index, and block benchmark - Wrap svg_step_index in jnp.asarray(i, dtype=jnp.int32) in unrolled loop to avoid per-step AOT recompilation - Guard TPU scoped VMEM compiler option and compute per-device local_batch in LTX2BlockBenchmark - Use _cfg fallback uniformly for all optional SVG config attributes in create_sharded_logical_transformer - Move svg_attention import to top level in attention_ltx2.py - Add LTX-2 config example and update svg_spatial_density description in docs/svg.md --- docs/svg.md | 18 ++++++++++++++++- .../models/ltx2/attention_ltx2.py | 3 +-- .../pipelines/ltx2/ltx2_pipeline.py | 20 +++++++++---------- .../utils/ltx2_block_benchmark.py | 14 +++++++++---- 4 files changed, 38 insertions(+), 17 deletions(-) diff --git a/docs/svg.md b/docs/svg.md index d5d3cd480..88b260d7e 100644 --- a/docs/svg.md +++ b/docs/svg.md @@ -58,11 +58,27 @@ svg_profile_seed: 0 svg_include_first_frame: True ``` +For LTX-2 (`src/maxdiffusion/configs/ltx2_video.yml` or `ltx2_3_video.yml`), use `svg_spatial_density`: + +```yaml +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 +``` + Step and layer intervals are zero-based and half-open: `[11, 40)` includes steps 11 through 39. Steps refer to denoising iterations, not noise-timestep values. | Option | What it controls | |---|---| -| `svg_spatial_density` | Density for single-expert Wan models; applies to either selected head pattern. | +| `svg_spatial_density` | Density for single-expert Wan models and LTX-2; applies to either selected head pattern. | | `svg_high_noise_density`, `svg_low_noise_density` | Separate densities for Wan2.2's two experts. | | `svg_active_start_step`, `svg_active_end_step` | Denoising steps where SVG is enabled. | | `svg_active_start_layer`, `svg_active_end_layer` | Transformer layers where SVG is enabled. | diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py index 16e32c727..65db95e5a 100644 --- a/src/maxdiffusion/models/ltx2/attention_ltx2.py +++ b/src/maxdiffusion/models/ltx2/attention_ltx2.py @@ -22,6 +22,7 @@ import jax.numpy as jnp from ... import common_types from ..attention_flax import NNXAttentionOp +from ..wan.transformers import svg_attention from .logical_sharding_ltx2 import get_sharding_specs, LTX2DiTShardingSpecs Array = common_types.Array @@ -650,8 +651,6 @@ def __call__( with self.named_scope("Attention and Output Project"): # 4. Attention if self.use_svg_attention and is_self_attention and spatiotemporal_shape is not None: - from maxdiffusion.models.wan.transformers import svg_attention - is_active = svg_attention.is_svg_active( step_index=svg_step_index, layer_index=svg_layer_index, diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index 46a7e529d..cd41c1da3 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -353,22 +353,22 @@ def _cfg(name, default): "ulysses_shards": getattr(config, "ulysses_shards", -1), "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1), "use_svg_attention": use_svg, - "svg_implementation": getattr(config, "svg_implementation", "official_svg"), + "svg_implementation": _cfg("svg_implementation", "official_svg"), "svg_spatial_density": expert_density, - "svg_sample_max_row": getattr(config, "svg_sample_max_row", 10000), - "svg_profile_query_count": getattr(config, "svg_profile_query_count", 64), - "svg_profile_seed": getattr(config, "svg_profile_seed", 0), - "svg_dense_layer_fraction": getattr(config, "svg_dense_layer_fraction", 0.0), - "svg_dense_timestep_fraction": getattr(config, "svg_dense_timestep_fraction", 0.0), + "svg_sample_max_row": _cfg("svg_sample_max_row", 10000), + "svg_profile_query_count": _cfg("svg_profile_query_count", 64), + "svg_profile_seed": _cfg("svg_profile_seed", 0), + "svg_dense_layer_fraction": _cfg("svg_dense_layer_fraction", 0.0), + "svg_dense_timestep_fraction": _cfg("svg_dense_timestep_fraction", 0.0), "svg_active_start_step": _cfg("svg_active_start_step", -1), "svg_active_end_step": _cfg("svg_active_end_step", -1), "svg_active_start_layer": _cfg("svg_active_start_layer", -1), "svg_active_end_layer": _cfg("svg_active_end_layer", -1), "svg_num_train_timesteps": _cfg("svg_num_train_timesteps", 1000), "svg_num_layers": _cfg("svg_num_layers", ltx2_config.get("num_layers", 28)), - "svg_include_first_frame": getattr(config, "svg_include_first_frame", True), - "svg_global_stride": getattr(config, "svg_global_stride", 0), - "svg_global_offset": getattr(config, "svg_global_offset", 0), + "svg_include_first_frame": _cfg("svg_include_first_frame", True), + "svg_global_stride": _cfg("svg_global_stride", 0), + "svg_global_offset": _cfg("svg_global_offset", 0), "svg_high_noise_density": high_density, "svg_low_noise_density": low_density, "svg_flash_block_sizes": getattr(config, "svg_flash_block_sizes", None) or None, @@ -2180,7 +2180,7 @@ def __call__( is_cfg_stg_mode=do_cfg and do_stg, kv_cache=kv_cache, rope_cache=rope_cache, - svg_step_index=i, + svg_step_index=jnp.asarray(i, dtype=jnp.int32), ) latents_step, audio_latents_step = _select_guidance_latents( diff --git a/src/maxdiffusion/utils/ltx2_block_benchmark.py b/src/maxdiffusion/utils/ltx2_block_benchmark.py index ab306880e..7281664d1 100644 --- a/src/maxdiffusion/utils/ltx2_block_benchmark.py +++ b/src/maxdiffusion/utils/ltx2_block_benchmark.py @@ -144,9 +144,13 @@ def __init__( self._vmem = int(vmem_limit_bytes) self.label = f"ltx2/{self._attention}/u{self._ulysses_shards}" device_kind = jax.devices()[0].device_kind - self._compiler_options = { - "xla_tpu_scoped_vmem_limit_kib": str(_SCOPED_VMEM_LIMIT_KIB[device_kind]), - } + if "tpu" in device_kind.lower(): + scoped_vmem = _SCOPED_VMEM_LIMIT_KIB.get(device_kind, 65536) + self._compiler_options = { + "xla_tpu_scoped_vmem_limit_kib": str(scoped_vmem), + } + else: + self._compiler_options = {} self._lf = (num_frames - 1) // 8 + 1 self._lh, self._lw = height // 32, width // 32 @@ -182,7 +186,9 @@ def tiled_seq_lens(self): return (s, s) def vmem_bytes(self): - return self._vmem // max(1, self._batch) + data_shards = int(self._mesh.shape.get("data", 1)) * int(self._mesh.shape.get("fsdp", 1)) + local_batch = max(1, self._batch // max(1, data_shards)) + return self._vmem // local_batch def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2): cmp = bkv_compute or bkv