From 716cc72002e27f5fc14b172651bf7afcd8d15c51 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Fri, 24 Jul 2026 22:51:32 +0000 Subject: [PATCH 01/10] LTX2 and LTX2.3 improvements and Tile-size auto-tuning - Fix IndivisibleError for 5D latents when scan_layers=False - Fix KV cache in run_diffusion_loop --- src/maxdiffusion/configs/ltx2_3_video.yml | 28 +- src/maxdiffusion/configs/ltx2_video.yml | 26 +- src/maxdiffusion/configs/ltx_video.yml | 2 + src/maxdiffusion/generate_ltx2.py | 74 +- src/maxdiffusion/models/attention_flax.py | 4 +- .../models/ltx2/attention_ltx2.py | 100 ++- src/maxdiffusion/models/ltx2/ltx2_utils.py | 140 +++- .../text_encoders/torchax_text_encoder.py | 20 +- .../models/ltx2/transformer_ltx2.py | 558 +++++++++----- .../pipelines/ltx2/ltx2_pipeline.py | 702 ++++++++++++------ .../tests/ltx2/test_pipeline_ltx2.py | 5 +- .../tests/ltx2/test_transformer_ltx2.py | 5 +- .../utils/ltx2_block_benchmark.py | 323 ++++++++ 13 files changed, 1542 insertions(+), 445 deletions(-) create mode 100644 src/maxdiffusion/utils/ltx2_block_benchmark.py diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml index a6275cfae..c4f536bff 100644 --- a/src/maxdiffusion/configs/ltx2_3_video.yml +++ b/src/maxdiffusion/configs/ltx2_3_video.yml @@ -1,7 +1,13 @@ #hardware hardware: 'tpu' skip_jax_distributed_system: False -attention: 'flash' +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +use_base2_exp: False +use_experimental_scheduler: False +# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. +ulysses_shards: -1 +# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. +ulysses_attention_chunks: 1 a2v_attention_kernel: 'flash' v2a_attention_kernel: 'dot_product' attention_sharding_uniform: True @@ -106,6 +112,10 @@ enable_ondemand_xprof: True skip_first_n_steps_for_profiler: 0 profiler_steps: 5 +# Enable JAX named scopes for detailed profiling and debugging +# When enabled, adds named scopes around key operations in transformer and attention layers +enable_jax_named_scopes: False + replicate_vae: False run_text_encoder_on_tpu: False @@ -173,4 +183,18 @@ upsampler_temporal_patch_size: 1 upsampler_adain_factor: 0.0 upsampler_tone_map_compression_ratio: 0.0 upsampler_rational_spatial_scale: 2.0 -upsampler_output_type: "pil" \ No newline at end of file +upsampler_output_type: "pil" + +aot_cache_dir: '' +converted_weights_dir: '' + + +# Tile-size auto-tuning. When enable_tile_search: True, generate_ltx2 runs a fast one-DiT-block +# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites +# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op). +enable_tile_search: False +tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep) +tile_search_iters: 10 +tile_search_out: '' # dir for the results CSV; '' -> print only +use_kv_cache: False +cross_attn_mod: False diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml index 8ddfaba1b..856b4626e 100644 --- a/src/maxdiffusion/configs/ltx2_video.yml +++ b/src/maxdiffusion/configs/ltx2_video.yml @@ -1,7 +1,13 @@ #hardware hardware: 'tpu' skip_jax_distributed_system: False -attention: 'flash' +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +use_base2_exp: False +use_experimental_scheduler: False +# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. +ulysses_shards: -1 +# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. +ulysses_attention_chunks: 1 a2v_attention_kernel: 'dot_product' v2a_attention_kernel: 'dot_product' attention_sharding_uniform: True @@ -111,6 +117,10 @@ enable_ondemand_xprof: True skip_first_n_steps_for_profiler: 0 profiler_steps: 5 +# Enable JAX named scopes for detailed profiling and debugging +# When enabled, adds named scopes around key operations in transformer and attention layers +enable_jax_named_scopes: False + replicate_vae: False use_bwe: False @@ -171,3 +181,17 @@ upsampler_adain_factor: 0.0 upsampler_tone_map_compression_ratio: 0.0 upsampler_rational_spatial_scale: 2.0 upsampler_output_type: "pil" + +aot_cache_dir: '' +converted_weights_dir: '' + + +# Tile-size auto-tuning. When enable_tile_search: True, generate_ltx2 runs a fast one-DiT-block +# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites +# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op). +enable_tile_search: False +tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep) +tile_search_iters: 10 +tile_search_out: '' # dir for the results CSV; '' -> print only +use_kv_cache: False +cross_attn_mod: False diff --git a/src/maxdiffusion/configs/ltx_video.yml b/src/maxdiffusion/configs/ltx_video.yml index d70154e0e..6dc09b5d4 100644 --- a/src/maxdiffusion/configs/ltx_video.yml +++ b/src/maxdiffusion/configs/ltx_video.yml @@ -109,3 +109,5 @@ enable_single_replica_ckpt_restoring: False enable_ml_diagnostics: False profiler_gcs_path: "" enable_ondemand_xprof: False +use_kv_cache: False +cross_attn_mod: False diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index 0445913ef..6c355bd2a 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -19,7 +19,7 @@ import os import subprocess from maxdiffusion.checkpointing.ltx2_checkpointer import LTX2Checkpointer -from maxdiffusion import pyconfig, max_logging, max_utils +from maxdiffusion import aot_cache, pyconfig, max_logging, max_utils from absl import app from google.cloud import storage from google.api_core.exceptions import GoogleAPIError @@ -113,7 +113,48 @@ def call_pipeline(config, pipeline, prompt, negative_prompt): return out +def maybe_tune_block_sizes(config): + """If enable_tile_search, run a fast one-DiT-block tile-size grid search and overwrite + flash_block_sizes' block_q/block_kv/block_kv_compute with the winner IN PLACE. + """ + keys = config.get_keys() + val = keys.get("enable_tile_search", False) + if str(val).lower() not in ("true", "1", "yes"): + return + from maxdiffusion.utils.tile_size_grid_search import grid_search + from maxdiffusion.utils.ltx2_block_benchmark import LTX2BlockBenchmark + + mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes) + bench = LTX2BlockBenchmark.from_config(config, mesh) + max_logging.log(f"[tile-search] tuning block sizes for {bench.label} before inference...") + result = grid_search( + bench, + mode=keys.get("tile_search_mode", "smart"), + iters=keys.get("tile_search_iters", 10), + out_dir=(keys.get("tile_search_out", "") or None), + log=max_logging.log, + ) + if result.best is None: + max_logging.log("[tile-search] no config succeeded; keeping configured flash_block_sizes") + return + fbs = dict(config.flash_block_sizes) + fbs.update({ + "block_q": result.best.bq, + "block_kv": result.best.bkv, + "block_kv_compute": result.best.bkv_compute, + "block_kv_compute_in": result.best.bkv_compute, + }) + config.get_keys()["flash_block_sizes"] = fbs + max_logging.log( + f"[tile-search] using block_q={result.best.bq} block_kv={result.best.bkv} " + f"(block-bench {result.best.mean_ms:.2f} ms)" + ) + + def run(config, pipeline=None, filename_prefix="", commit_hash=None): + if pipeline is None: + maybe_tune_block_sizes(config) + writer = max_utils.initialize_summary_writer(config) if jax.process_index() == 0 and writer: max_logging.log(f"TensorBoard logs will be written to: {config.tensorboard_dir}") @@ -189,14 +230,37 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): original_enable_mld = config.get_keys().get("enable_ml_diagnostics", False) original_num_steps = config.get_keys().get("num_inference_steps", 40) + # Per-shape AOT executable cache + aot_cache.install( + getattr(config, "aot_cache_dir", ""), + meta={ + "model": config.pretrained_model_name_or_path, + "attention": getattr(config, "attention", ""), + "flash_block_sizes": str(getattr(config, "flash_block_sizes", "")), + "mesh_shape": str(pipeline.mesh.shape) if pipeline and hasattr(pipeline, "mesh") and pipeline.mesh else "", + "weights_dtype": str(getattr(config, "weights_dtype", "bfloat16")), + "activations_dtype": str(getattr(config, "activations_dtype", "bfloat16")), + "scan_layers": str(getattr(config, "scan_layers", True)), + "jax": jax.__version__, + }, + mesh=pipeline.mesh if pipeline else None, + ) + aot_cache.wait_for_loads() + # --------------------------------------------------------- # Run 1: Warmup Compilation (Original steps, NO profiling) # --------------------------------------------------------- config.get_keys()["enable_profiler"] = False config.get_keys()["enable_ml_diagnostics"] = False + warmup_steps = min(2, original_num_steps) + config.get_keys()["num_inference_steps"] = warmup_steps + + max_logging.log(f"🚀 Starting warmup compilation pass ({warmup_steps} steps)...") + with aot_cache.warmup_mode(): + _ = call_pipeline(config, pipeline, prompt, negative_prompt) - max_logging.log(f"🚀 Starting warmup compilation pass ({original_num_steps} steps)...") - _ = call_pipeline(config, pipeline, prompt, negative_prompt) + aot_cache.save_pending() + config.get_keys()["num_inference_steps"] = original_num_steps compile_time = time.perf_counter() - s0 max_logging.log(f"compile_time: {compile_time}") @@ -237,7 +301,9 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # Export videos for i in range(len(videos)): - video_path = f"{filename_prefix}ltx2_output_{getattr(config, 'seed', 0)}_{i}.mp4" + model_name = getattr(config, "model_name", "ltx2") or "ltx2" + model_name_prefix = model_name.replace(".", "_") + video_path = f"{filename_prefix}{model_name_prefix}_output_{getattr(config, 'seed', 0)}_{i}.mp4" audio_i = audios[i] if audios is not None else None audio_format = getattr(config, "audio_format", "s16") diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index c56d8c64c..a76da7dea 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -1914,7 +1914,7 @@ def __init__( self, mesh: Mesh, attention_kernel: str, - scale: int, + scale: float, heads: int, dim_head: int, use_memory_efficient_attention: bool = False, @@ -2009,7 +2009,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask class AttentionOp(nn.Module): mesh: Mesh attention_kernel: str - scale: int + scale: float heads: int dim_head: int use_memory_efficient_attention: bool = False diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py index 214690c98..730f81771 100644 --- a/src/maxdiffusion/models/ltx2/attention_ltx2.py +++ b/src/maxdiffusion/models/ltx2/attention_ltx2.py @@ -88,13 +88,8 @@ def apply_split_rotary_emb(x: Array, freqs: Tuple[Array, Array]) -> Array: first_x = split_x[..., 0, :] second_x = split_x[..., 1, :] - cos_u = jnp.expand_dims(cos, axis=-2) - sin_u = jnp.expand_dims(sin, axis=-2) - - out = split_x * cos_u - - out_first = out[..., 0, :] - second_x * sin_u.squeeze(-2) - out_second = out[..., 1, :] + first_x * sin_u.squeeze(-2) + out_first = first_x * cos - second_x * sin + out_second = second_x * cos + first_x * sin out = jnp.stack([out_first, out_second], axis=-2) out = out.reshape(*out.shape[:-2], last_dim) @@ -176,12 +171,6 @@ def prepare_video_coords( patch_ends = grid + patch_size_delta # Combine start and end coordinates - latent_coords = jnp.stack([grid, patch_ends], axis=-1) # [3, N_F, N_H, N_W, 2] - latent_coords = latent_coords.transpose(1, 2, 3, 0, 4) # [N_F, N_H, N_W, 3, 2] - latent_coords = latent_coords.reshape(-1, 3, 2) # [num_patches, 3, 2] - latent_coords = jnp.expand_dims(latent_coords, 0) # [1, num_patches, 3, 2] - latent_coords = jnp.tile(latent_coords, (batch_size, 1, 1, 1)) # [B, num_patches, 3, 2] - latent_coords = jnp.stack([grid, patch_ends], axis=-1) # [3, N_F, N_H, N_W, 2] latent_coords = latent_coords.reshape(3, -1, 2) # [3, num_patches, 2] latent_coords = jnp.expand_dims(latent_coords, 0) # [1, 3, num_patches, 2] @@ -257,7 +246,10 @@ def __call__(self, coords: Array) -> Tuple[Array, Array]: # 2. Fractions if self.modality == "video": - max_positions = jnp.array((self.base_num_frames, self.base_height, self.base_width), dtype=coords.dtype) + max_positions = jnp.array( + (self.base_num_frames, self.base_height, self.base_width), + dtype=coords.dtype, + ) elif self.modality == "audio": max_positions = jnp.array((self.base_num_frames,), dtype=coords.dtype) @@ -352,6 +344,8 @@ def __init__( flash_min_seq_length: int = 4096, sharding_specs: Optional[LTX2DiTShardingSpecs] = None, gated_attn: bool = False, + ulysses_shards: int = -1, + ulysses_attention_chunks: int = 1, ): self.heads = heads self.rope_type = rope_type @@ -392,10 +386,22 @@ def __init__( # Handle Self vs Cross Attention input dims kv_dim = context_dim if context_dim is not None else query_dim self.to_k = nnx.Linear( - kv_dim, self.inner_dim, use_bias=bias, kernel_init=qkv_kernel_init, bias_init=qkv_bias_init, rngs=rngs, dtype=dtype + kv_dim, + self.inner_dim, + use_bias=bias, + kernel_init=qkv_kernel_init, + bias_init=qkv_bias_init, + rngs=rngs, + dtype=dtype, ) self.to_v = nnx.Linear( - kv_dim, self.inner_dim, use_bias=bias, kernel_init=qkv_kernel_init, bias_init=qkv_bias_init, rngs=rngs, dtype=dtype + kv_dim, + self.inner_dim, + use_bias=bias, + kernel_init=qkv_kernel_init, + bias_init=qkv_bias_init, + rngs=rngs, + dtype=dtype, ) # 3. Normalization (Applied to full inner_dim, NOT per-head) @@ -445,6 +451,43 @@ def __init__( dtype=dtype, ) + is_self_attention = context_dim is None + if not is_self_attention: + if attention_kernel in ( + "tokamax_ring", + "tokamax_ring_custom", + "ulysses_ring", + ): + attention_kernel = "tokamax_flash" # do not use ring attention for cross attention + if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir"): + attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention + + axis_names_q = ( + common_types.BATCH, + common_types.CROSS_ATTN_HEAD, + common_types.CROSS_ATTN_Q_LENGTH, + common_types.D_KV, + ) + axis_names_kv = ( + common_types.BATCH, + common_types.CROSS_ATTN_HEAD, + common_types.CROSS_ATTN_KV_LENGTH, + common_types.D_KV, + ) + else: + axis_names_q = ( + common_types.BATCH, + common_types.SELF_ATTN_HEAD, + common_types.SELF_ATTN_Q_LENGTH, + common_types.D_KV, + ) + axis_names_kv = ( + common_types.BATCH, + common_types.SELF_ATTN_HEAD, + common_types.SELF_ATTN_KV_LENGTH, + common_types.D_KV, + ) + self.attention_op = NNXAttentionOp( mesh=mesh, attention_kernel=attention_kernel, @@ -452,12 +495,20 @@ def __init__( heads=heads, dim_head=dim_head, dtype=dtype, - axis_names_q=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV), - axis_names_kv=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV), + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, flash_block_sizes=flash_block_sizes, flash_min_seq_length=flash_min_seq_length, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) + def compute_kv(self, context: Array) -> Tuple[Array, Array]: + key = self.to_k(context) + value = self.to_v(context) + key = self.norm_k(key) + return key, value + def __call__( self, hidden_states: Array, @@ -466,6 +517,7 @@ def __call__( rotary_emb: Optional[Tuple[Array, Array]] = None, k_rotary_emb: Optional[Tuple[Array, Array]] = None, perturbation_mask: Optional[Array] = None, + cached_kv: Optional[Tuple[Array, Array]] = None, ) -> Array: # Determine context (Self or Cross) context = encoder_hidden_states if encoder_hidden_states is not None else hidden_states @@ -473,19 +525,23 @@ def __call__( # 1. Project and Norm with jax.named_scope("QKV Projection"): query = self.to_q(hidden_states) - key = self.to_k(context) - value = self.to_v(context) + if cached_kv is not None: + key, value = cached_kv + else: + key = self.to_k(context) + value = self.to_v(context) with jax.named_scope("QKV Norm"): query = self.norm_q(query) - key = self.norm_k(key) + if cached_kv is None: + key = self.norm_k(key) # 3. Apply RoPE to tensors of shape [B, S, InnerDim] # Frequencies are shape [B, S, InnerDim] # 3. Apply RoPE with jax.named_scope("Apply RoPE"): if rotary_emb is not None: - if hasattr(self, "rope_type") and self.rope_type == "split": + if self.rope_type == "split": # Split RoPE: passing full freqs [B, H, S, D//2] # apply_split_rotary_emb handles reshaping query/key diff --git a/src/maxdiffusion/models/ltx2/ltx2_utils.py b/src/maxdiffusion/models/ltx2/ltx2_utils.py index 7cf7fa761..31819cecc 100644 --- a/src/maxdiffusion/models/ltx2/ltx2_utils.py +++ b/src/maxdiffusion/models/ltx2/ltx2_utils.py @@ -17,6 +17,7 @@ import json from typing import Optional import torch +import numpy as np import jax import jax.numpy as jnp from maxdiffusion import max_logging @@ -184,6 +185,14 @@ def load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device, fi return tensors +def _torch_tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray: + import ml_dtypes + + if tensor.dtype == torch.bfloat16: + return tensor.view(torch.uint16).numpy().view(ml_dtypes.bfloat16) + return tensor.numpy() + + def load_transformer_weights( pretrained_model_name_or_path: str, eval_shapes: dict, @@ -193,38 +202,125 @@ def load_transformer_weights( scan_layers: bool = True, subfolder: str = "transformer", ): + import threading + import concurrent.futures + import time + device = jax.local_devices(backend=device)[0] max_logging.log(f"Load and port {pretrained_model_name_or_path} {subfolder} on {device}") - with jax.default_device(device): - # Support sharded loading - tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device) + index_file = "diffusion_pytorch_model.safetensors.index.json" + try: + index_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=index_file) + with open(index_path, "r") as f: + index_data = json.load(f) + weight_map = index_data["weight_map"] + shards = sorted(set(weight_map.values())) - flax_state_dict = {} - cpu = jax.local_devices(backend="cpu")[0] - flattened_dict = flatten_dict(eval_shapes) + def resolve_shard_path(model_file): + return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=model_file) - random_flax_state_dict = {} - for key in flattened_dict: - random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key] + except EntryNotFoundError: + shards = ["diffusion_pytorch_model.safetensors"] - for pt_key, tensor in tensors.items(): - renamed_pt_key = rename_key(pt_key) - renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key) + def resolve_shard_path(model_file): + try: + return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=model_file) + except EntryNotFoundError: + return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename="diffusion_pytorch_model.bin") - pt_tuple_key = tuple(renamed_pt_key.split(".")) + t_start = time.perf_counter() - flax_key, flax_tensor = get_key_and_value( - pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers, num_layers - ) + flattened_dict = flatten_dict(eval_shapes) + random_flax_state_dict = {} + for key in flattened_dict: + random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key] - flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu) + flax_state_dict = {} + dict_lock = threading.Lock() + + def convert_chunk(ckpt_shard_path, chunk_keys): + if ckpt_shard_path.endswith(".safetensors"): + with safe_open(ckpt_shard_path, framework="pt") as f: + for pt_key in chunk_keys: + tensor = _torch_tensor_to_numpy(f.get_tensor(pt_key)) + process_tensor(pt_key, tensor) + else: + loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu") + for pt_key in chunk_keys: + tensor = _torch_tensor_to_numpy(loaded_state_dict[pt_key]) + process_tensor(pt_key, tensor) + + def process_tensor(pt_key, tensor): + renamed_pt_key = rename_key(pt_key) + renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key) + pt_tuple_key = tuple(renamed_pt_key.split(".")) + + block_index = None + if len(pt_tuple_key) > 0 and "transformer_blocks_" in pt_tuple_key[0]: + import re + + m = re.match(r"transformer_blocks_(\d+)", pt_tuple_key[0]) + if m: + if scan_layers: + block_index = int(m.group(1)) + pt_tuple_key = ("transformer_blocks",) + pt_tuple_key[1:] + else: + # For nnx.List, NNX uses string indices ('0', '1', etc.) + pt_tuple_key = ("transformer_blocks", m.group(1)) + pt_tuple_key[1:] - validate_flax_state_dict(eval_shapes, flax_state_dict) - flax_state_dict = unflatten_dict(flax_state_dict) - del tensors - jax.clear_caches() - return flax_state_dict + flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, scan_layers) + flax_key_str = [str(k) for k in flax_key] + if "scale_shift_table" in flax_key_str and flax_key_str[-1] in ["kernel", "weight"]: + flax_key_str.pop() + flax_key = tuple(flax_key_str) + # Convert string indices back to integers only if they were natively numeric keys + # But wait, nnx.List expects strings. _tuple_str_to_int converts '0' to 0. + # If the logical state sharding has '0', _tuple_str_to_int will break the match! + # Let's conditionally skip _tuple_str_to_int for nnx.List indices. + + # Actually, Flax traverse_util flattens nnx.List keys to strings in nnx, but _tuple_str_to_int + # makes them integers. Let's just use _tuple_str_to_int but format it back to string if needed. + flax_key = _tuple_str_to_int(flax_key) + + # We must ensure that if scan_layers=False, the index is a string to match NNX List! + if not scan_layers and len(flax_key) > 1 and flax_key[0] == "transformer_blocks" and isinstance(flax_key[1], int): + flax_key = ("transformer_blocks", str(flax_key[1])) + flax_key[2:] + + if block_index is not None: + with dict_lock: + stacked = flax_state_dict.get(flax_key) + if stacked is None: + stacked = np.empty((num_layers,) + flax_tensor.shape, dtype=flax_tensor.dtype) + flax_state_dict[flax_key] = stacked + stacked[block_index] = flax_tensor + else: + value = np.array(flax_tensor, dtype=flax_tensor.dtype, copy=True, order="C") + with dict_lock: + flax_state_dict[flax_key] = value + + chunk_size = 32 + tasks = [] + for model_file in shards: + ckpt_shard_path = resolve_shard_path(model_file) + if ckpt_shard_path.endswith(".safetensors"): + with safe_open(ckpt_shard_path, framework="pt") as f: + shard_keys = list(f.keys()) + else: + loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu") + shard_keys = list(loaded_state_dict.keys()) + for i in range(0, len(shard_keys), chunk_size): + tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) + + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: + futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks] + for future in concurrent.futures.as_completed(futures): + future.result() + + validate_flax_state_dict(eval_shapes, flax_state_dict) + flax_state_dict = unflatten_dict(flax_state_dict) + max_logging.log(f"Converted weights in {time.perf_counter() - t_start:.1f}s") + return flax_state_dict def load_vae_weights( diff --git a/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py b/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py index 882c65a25..e18cf90c4 100644 --- a/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py +++ b/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py @@ -20,11 +20,9 @@ import jax from torchax import interop, default_env -# --- Monkeypatch transformers masking_utils to avoid torchax integer tracing bug --- +import contextlib import transformers.masking_utils -_orig_sliding_window_overlay = transformers.masking_utils.sliding_window_overlay - def _patched_sliding_window_overlay(sliding_window: int): # pylint: disable=unused-argument @@ -44,6 +42,16 @@ def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: return inner_mask +@contextlib.contextmanager +def patch_sliding_window_overlay(): + orig = transformers.masking_utils.sliding_window_overlay + transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay + try: + yield + finally: + transformers.masking_utils.sliding_window_overlay = orig + + class TorchaxGemma3TextEncoder(interop.JittableModule): """ A jittable Torchax module for wrapping the HuggingFace PyTorch @@ -57,8 +65,7 @@ def __call__( self, input_ids: jax.Array, attention_mask: jax.Array, output_hidden_states: bool = True ) -> Tuple[jax.Array, ...]: # Dynamically patch transformers.masking_utils only during the duration of this call - transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay - try: + with patch_sliding_window_overlay(): with default_env(): input_ids = interop.torch_view(input_ids) attention_mask = interop.torch_view(attention_mask) @@ -72,9 +79,6 @@ def __call__( output_hidden_states=output_hidden_states, ) return interop.jax_view(output) - finally: - # Restore original behavior to prevent side effects on other potential models in same env - transformers.masking_utils.sliding_window_overlay = _orig_sliding_window_overlay @staticmethod def _forward_inner(model, input_ids, attention_mask, output_hidden_states=True): diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py index f2511ed07..3ff086915 100644 --- a/src/maxdiffusion/models/ltx2/transformer_ltx2.py +++ b/src/maxdiffusion/models/ltx2/transformer_ltx2.py @@ -27,6 +27,34 @@ from maxdiffusion.configuration_utils import ConfigMixin, register_to_config from maxdiffusion.common_types import BlockSizes from .logical_sharding_ltx2 import get_sharding_specs, LTX2DiTShardingSpecs +from flax import struct + + +@struct.dataclass +class LTX2BlockContext: + hidden_states: jax.Array + audio_hidden_states: jax.Array + encoder_hidden_states: jax.Array + audio_encoder_hidden_states: jax.Array + temb: jax.Array + temb_audio: jax.Array + temb_ca_scale_shift: jax.Array + temb_ca_audio_scale_shift: jax.Array + temb_ca_gate: jax.Array + temb_ca_audio_gate: jax.Array + temb_prompt: Optional[jax.Array] = None + temb_prompt_audio: Optional[jax.Array] = None + modality_mask: Optional[jax.Array] = None + video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + ca_video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + ca_audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + encoder_attention_mask: Optional[jax.Array] = None + audio_encoder_attention_mask: Optional[jax.Array] = None + a2v_cross_attention_mask: Optional[jax.Array] = None + v2a_cross_attention_mask: Optional[jax.Array] = None + perturbation_mask: Optional[jax.Array] = None + cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None class LTX2AdaLayerNormSingle(nnx.Module): @@ -89,7 +117,12 @@ def __call__( resolution = added_cond_kwargs.get("resolution", None) aspect_ratio = added_cond_kwargs.get("aspect_ratio", None) - embedded_timestep = self.emb(timestep, resolution=resolution, aspect_ratio=aspect_ratio, hidden_dtype=hidden_dtype) + embedded_timestep = self.emb( + timestep, + resolution=resolution, + aspect_ratio=aspect_ratio, + hidden_dtype=hidden_dtype, + ) return self.linear(self.silu(embedded_timestep)), embedded_timestep @@ -128,6 +161,8 @@ def __init__( flash_min_seq_length: int = 4096, sharding_specs: Optional[LTX2DiTShardingSpecs] = None, perturbed_attn: bool = False, + ulysses_shards: int = -1, + ulysses_attention_chunks: int = 1, ): self.dim = dim self.norm_eps = norm_eps @@ -166,6 +201,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) self.audio_norm1 = nnx.RMSNorm( @@ -194,6 +231,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) # 2. Prompt Cross-Attention @@ -223,6 +262,8 @@ def __init__( flash_block_sizes=flash_block_sizes, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) self.audio_norm2 = nnx.RMSNorm( @@ -252,6 +293,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention @@ -279,9 +322,11 @@ def __init__( attention_kernel=a2v_attention_kernel, rope_type=rope_type, flash_block_sizes=flash_block_sizes, - flash_min_seq_length=0, + flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) self.video_to_audio_norm = nnx.RMSNorm( @@ -311,6 +356,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) # 4. Feed Forward @@ -361,69 +408,99 @@ def __init__( self.scale_shift_table = nnx.Param( nnx.with_partitioning( - lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(self.dim), table_sharding + lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(self.dim), + table_sharding, )(k1, (table_size, self.dim)) ) if self.cross_attn_mod: self.prompt_scale_shift_table = nnx.Param( nnx.with_partitioning( - lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(self.dim), table_sharding + lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(self.dim), + table_sharding, )(k5, (2, self.dim)) ) self.audio_scale_shift_table = nnx.Param( nnx.with_partitioning( - lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(audio_dim), table_sharding + lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(audio_dim), + table_sharding, )(k2, (table_size, audio_dim)) ) self.video_a2v_cross_attn_scale_shift_table = nnx.Param( - nnx.with_partitioning(lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype), table_sharding)( - k3, (5, self.dim) - ) + nnx.with_partitioning( + lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype), + table_sharding, + )(k3, (5, self.dim)) ) self.audio_a2v_cross_attn_scale_shift_table = nnx.Param( - nnx.with_partitioning(lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype), table_sharding)( - k4, (5, audio_dim) - ) + nnx.with_partitioning( + lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype), + table_sharding, + )(k4, (5, audio_dim)) ) if self.cross_attn_mod: self.audio_prompt_scale_shift_table = nnx.Param( nnx.with_partitioning( - lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(audio_dim), table_sharding + lambda key, shape: jax.random.normal(key, shape, dtype=weights_dtype) / jnp.sqrt(audio_dim), + table_sharding, )(k6, (2, audio_dim)) ) + def compute_kv(self, encoder_hidden_states: jax.Array, audio_encoder_hidden_states: jax.Array): + text_k, text_v = self.attn2.compute_kv(encoder_hidden_states) + audio_text_k, audio_text_v = self.audio_attn2.compute_kv(audio_encoder_hidden_states) + return {"attn2": (text_k, text_v), "audio_attn2": (audio_text_k, audio_text_v)} + def __call__( self, - hidden_states: jax.Array, # Video - audio_hidden_states: jax.Array, # Audio - encoder_hidden_states: jax.Array, # Context (Text) - audio_encoder_hidden_states: jax.Array, # Audio Context - # Timestep embeddings for AdaLN - temb: jax.Array, - temb_audio: jax.Array, - temb_ca_scale_shift: jax.Array, - temb_ca_audio_scale_shift: jax.Array, - temb_ca_gate: jax.Array, - temb_ca_audio_gate: jax.Array, - temb_prompt: Optional[jax.Array] = None, - temb_prompt_audio: Optional[jax.Array] = None, - modality_mask: Optional[jax.Array] = None, - # RoPE - video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - ca_video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - ca_audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - encoder_attention_mask: Optional[jax.Array] = None, - audio_encoder_attention_mask: Optional[jax.Array] = None, - a2v_cross_attention_mask: Optional[jax.Array] = None, - v2a_cross_attention_mask: Optional[jax.Array] = None, - perturbation_mask: Optional[jax.Array] = None, + ctx: "LTX2BlockContext", ) -> Tuple[jax.Array, jax.Array]: + """ + Forward pass of the LTX2 video/audio transformer block. + + This block handles complex multi-modal attention including: + - Video Self-Attention (video -> video) + - Audio Self-Attention (audio -> audio) + - Video Cross-Attention (video -> text caption) + - Audio Cross-Attention (audio -> text caption) + - Video-to-Audio Cross-Attention + - Audio-to-Video Cross-Attention + + Args: + ctx: An `LTX2BlockContext` object containing all hidden states, timestep + embeddings, attention masks, rotary embeddings, and modulation + parameters needed for this layer's forward pass. + + Returns: + A tuple of `(output_hidden_states, output_audio_hidden_states)`. + """ + hidden_states = ctx.hidden_states + audio_hidden_states = ctx.audio_hidden_states + encoder_hidden_states = ctx.encoder_hidden_states + audio_encoder_hidden_states = ctx.audio_encoder_hidden_states + temb = ctx.temb + temb_audio = ctx.temb_audio + temb_ca_scale_shift = ctx.temb_ca_scale_shift + temb_ca_audio_scale_shift = ctx.temb_ca_audio_scale_shift + temb_ca_gate = ctx.temb_ca_gate + temb_ca_audio_gate = ctx.temb_ca_audio_gate + temb_prompt = ctx.temb_prompt + temb_prompt_audio = ctx.temb_prompt_audio + modality_mask = ctx.modality_mask + video_rotary_emb = ctx.video_rotary_emb + audio_rotary_emb = ctx.audio_rotary_emb + ca_video_rotary_emb = ctx.ca_video_rotary_emb + ca_audio_rotary_emb = ctx.ca_audio_rotary_emb + encoder_attention_mask = ctx.encoder_attention_mask + audio_encoder_attention_mask = ctx.audio_encoder_attention_mask + a2v_cross_attention_mask = ctx.a2v_cross_attention_mask + v2a_cross_attention_mask = ctx.v2a_cross_attention_mask + perturbation_mask = ctx.perturbation_mask + cached_kv = ctx.cached_kv batch_size = hidden_states.shape[0] axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) @@ -515,11 +592,13 @@ def __call__( scale_text_kv = prompt_ada_values[:, :, 1, :] encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv + attn2_kv = cached_kv.get("attn2") if cached_kv is not None else None attn_hidden_states = self.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, rotary_emb=None, attention_mask=encoder_attention_mask, + cached_kv=attn2_kv, ) if getattr(self, "cross_attn_mod", False): attn_hidden_states = attn_hidden_states * gate_q @@ -537,11 +616,13 @@ def __call__( audio_scale_text_kv = audio_prompt_ada_values[:, :, 1, :] audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv + audio_attn2_kv = cached_kv.get("audio_attn2") if cached_kv is not None else None attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, encoder_hidden_states=audio_encoder_hidden_states, rotary_emb=None, attention_mask=audio_encoder_attention_mask, + cached_kv=audio_attn2_kv, ) if getattr(self, "cross_attn_mod", False): attn_audio_hidden_states = attn_audio_hidden_states * audio_gate_q @@ -695,6 +776,8 @@ def __init__( use_prompt_embeddings: bool = True, perturbed_attn: bool = False, spatio_temporal_guidance_blocks: Tuple[int, ...] = (), + ulysses_shards: int = -1, + ulysses_attention_chunks: int = 1, **kwargs, ): self.spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks @@ -760,6 +843,9 @@ def __init__( inner_dim = self.num_attention_heads * self.attention_head_dim audio_inner_dim = self.audio_num_attention_heads * self.audio_attention_head_dim + self.inner_dim = inner_dim + self.audio_inner_dim = audio_inner_dim + # 1. Patchification input projections self.proj_in = nnx.Linear( self.in_channels, @@ -879,17 +965,19 @@ def __init__( # 3. Output Layer Scale/Shift Modulation parameters param_rng = rngs.params() + audio_param_rng = rngs.params() table_sharding = self.sharding_specs.scale_shift_table self.scale_shift_table = nnx.Param( nnx.with_partitioning( - lambda key, shape: jax.random.normal(key, shape, dtype=self.weights_dtype) / jnp.sqrt(inner_dim), table_sharding + lambda key, shape: jax.random.normal(key, shape, dtype=self.weights_dtype) / jnp.sqrt(inner_dim), + table_sharding, )(param_rng, (2, inner_dim)) ) self.audio_scale_shift_table = nnx.Param( nnx.with_partitioning( lambda key, shape: jax.random.normal(key, shape, dtype=self.weights_dtype) / jnp.sqrt(audio_inner_dim), table_sharding, - )(param_rng, (2, audio_inner_dim)) + )(audio_param_rng, (2, audio_inner_dim)) ) # 4. Rotary Positional Embeddings (RoPE) @@ -956,7 +1044,12 @@ def __init__( # 5. Transformer Blocks @nnx.split_rngs(splits=self.num_layers) - @nnx.vmap(in_axes=0, out_axes=0, axis_size=self.num_layers, transform_metadata={nnx.PARTITION_NAME: "layers"}) + @nnx.vmap( + in_axes=0, + out_axes=0, + axis_size=self.num_layers, + transform_metadata={nnx.PARTITION_NAME: "layers"}, + ) def init_block(rngs): return LTX2VideoTransformerBlock( rngs=rngs, @@ -990,6 +1083,8 @@ def init_block(rngs): flash_block_sizes=flash_block_sizes, flash_min_seq_length=self.flash_min_seq_length, perturbed_attn=self.perturbed_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) if self.scan_layers: @@ -999,6 +1094,7 @@ def init_block(rngs): for _ in range(self.num_layers): block = LTX2VideoTransformerBlock( rngs=rngs, + sharding_specs=self.sharding_specs, dim=inner_dim, num_attention_heads=self.num_attention_heads, attention_head_dim=self.attention_head_dim, @@ -1028,6 +1124,8 @@ def init_block(rngs): flash_block_sizes=flash_block_sizes, flash_min_seq_length=self.flash_min_seq_length, perturbed_attn=self.perturbed_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) blocks.append(block) self.transformer_blocks = nnx.List(blocks) @@ -1035,7 +1133,13 @@ def init_block(rngs): # 6. Output layers self.gradient_checkpoint = GradientCheckpointType.from_str(remat_policy) self.norm_out = nnx.LayerNorm( - inner_dim, epsilon=1e-6, use_scale=False, use_bias=False, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32 + inner_dim, + epsilon=1e-6, + use_scale=False, + use_bias=False, + rngs=rngs, + dtype=jnp.float32, + param_dtype=jnp.float32, ) self.proj_out = nnx.Linear( inner_dim, @@ -1048,7 +1152,13 @@ def init_block(rngs): ) self.audio_norm_out = nnx.LayerNorm( - audio_inner_dim, epsilon=1e-6, use_scale=False, use_bias=False, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32 + audio_inner_dim, + epsilon=1e-6, + use_scale=False, + use_bias=False, + rngs=rngs, + dtype=jnp.float32, + param_dtype=jnp.float32, ) self.audio_proj_out = nnx.Linear( audio_inner_dim, @@ -1060,6 +1170,71 @@ def init_block(rngs): bias_init=nnx.with_partitioning(nnx.initializers.zeros, self.sharding_specs.out_embed_bias), ) + def compute_kv_cache( + self, + encoder_hidden_states: jax.Array, + audio_encoder_hidden_states: jax.Array, + num_frames: int, + height: int, + width: int, + fps: float, + audio_num_frames: int, + ): + batch_size = encoder_hidden_states.shape[0] + if self.use_prompt_embeddings and self.caption_projection is not None: + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states) + + encoder_hidden_states = encoder_hidden_states.reshape(batch_size, -1, self.inner_dim) + audio_encoder_hidden_states = audio_encoder_hidden_states.reshape(batch_size, -1, self.audio_inner_dim) + + kv_cache = None + if not self.cross_attn_mod: + if self.scan_layers: + + @nnx.vmap( + in_axes=(0, None, None), + out_axes=0, + transform_metadata={nnx.PARTITION_NAME: "layers"}, + ) + def _compute_kv(block, enc_states, audio_enc_states): + return block.compute_kv(enc_states, audio_enc_states) + + kv_cache = _compute_kv( + self.transformer_blocks, + encoder_hidden_states, + audio_encoder_hidden_states, + ) + else: + kv_cache_list = [] + for block in self.transformer_blocks: + kv_cache_list.append(block.compute_kv(encoder_hidden_states, audio_encoder_hidden_states)) + keys = kv_cache_list[0].keys() + kv_cache = {} + for k in keys: + k_list = [d[k][0] for d in kv_cache_list] + v_list = [d[k][1] for d in kv_cache_list] + kv_cache[k] = (jnp.stack(k_list, axis=0), jnp.stack(v_list, axis=0)) + + # RoPE pre-computation + video_coords = self.rope.prepare_video_coords(batch_size, num_frames, height, width, fps=fps) + audio_coords = self.audio_rope.prepare_audio_coords(batch_size, audio_num_frames) + + video_rotary_emb = self.rope(video_coords) + audio_rotary_emb = self.audio_rope(audio_coords) + + video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :]) + audio_cross_attn_rotary_emb = self.cross_attn_audio_rope(audio_coords[:, 0:1, :]) + + rope_cache = { + "video_rotary_emb": video_rotary_emb, + "audio_rotary_emb": audio_rotary_emb, + "video_cross_attn_rotary_emb": video_cross_attn_rotary_emb, + "audio_cross_attn_rotary_emb": audio_cross_attn_rotary_emb, + } + + return kv_cache, rope_cache, encoder_hidden_states, audio_encoder_hidden_states + def __call__( self, hidden_states: jax.Array, @@ -1084,10 +1259,48 @@ def __call__( modality_mask: Optional[jax.Array] = None, return_dict: bool = True, perturbation_mask: Optional[jax.Array] = None, + cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, + rope_cache: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, ) -> Any: + """ + Forward pass for the full LTX2 Video/Audio Diffusion Transformer. + + Args: + hidden_states: Video latent patches of shape `(batch, seq_len, in_channels)`. + audio_hidden_states: Audio latent patches of shape `(batch, audio_seq_len, audio_in_channels)`. + encoder_hidden_states: Text embeddings for video generation. + audio_encoder_hidden_states: Text embeddings for audio generation. + timestep: Timestep array for video diffusion. + audio_timestep: Optional timestep array for audio diffusion. If None, uses `timestep`. + sigma: Optional noise scale for video (for flow matching). + audio_sigma: Optional noise scale for audio. + encoder_attention_mask: Mask for video text embeddings. + audio_encoder_attention_mask: Mask for audio text embeddings. + num_frames: Number of video frames. + height: Height of the video frames. + width: Width of the video frames. + fps: Frames per second. + audio_num_frames: Number of audio frames. + video_coords: Optional pre-computed 3D coordinates for video RoPE. + audio_coords: Optional pre-computed 1D coordinates for audio RoPE. + attention_kwargs: Additional kwargs for the attention mechanisms. + use_cross_timestep: Whether to use a cross-modal timestep interaction. + modality_mask: Mask indicating which modality to drop/keep. + return_dict: If True, returns a dictionary. Otherwise, returns a tuple. + perturbation_mask: Optional mask for perturbing attention. + + Returns: + Output dict containing `sample` (video) and `audio_sample` (audio). + """ # Determine timestep for audio. audio_timestep = audio_timestep if audio_timestep is not None else timestep + a2v_cross_attention_mask = None + v2a_cross_attention_mask = None + if attention_kwargs is not None: + a2v_cross_attention_mask = attention_kwargs.get("a2v_cross_attention_mask", None) + v2a_cross_attention_mask = attention_kwargs.get("v2a_cross_attention_mask", None) + if self.attention_kernel == "dot_product": if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: encoder_attention_mask = (1 - encoder_attention_mask.astype(self.dtype)) * -10000.0 @@ -1097,20 +1310,34 @@ def __call__( audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.astype(self.dtype)) * -10000.0 audio_encoder_attention_mask = jnp.expand_dims(audio_encoder_attention_mask, axis=1) + if a2v_cross_attention_mask is not None and a2v_cross_attention_mask.ndim == 2: + a2v_cross_attention_mask = (1 - a2v_cross_attention_mask.astype(self.dtype)) * -10000.0 + a2v_cross_attention_mask = jnp.expand_dims(a2v_cross_attention_mask, axis=1) + + if v2a_cross_attention_mask is not None and v2a_cross_attention_mask.ndim == 2: + v2a_cross_attention_mask = (1 - v2a_cross_attention_mask.astype(self.dtype)) * -10000.0 + v2a_cross_attention_mask = jnp.expand_dims(v2a_cross_attention_mask, axis=1) + batch_size = hidden_states.shape[0] # 1. Prepare RoPE positional embeddings with jax.named_scope("RoPE Preparation"): - if video_coords is None: - video_coords = self.rope.prepare_video_coords(batch_size, num_frames, height, width, fps=fps) - if audio_coords is None: - audio_coords = self.audio_rope.prepare_audio_coords(batch_size, audio_num_frames) + if rope_cache is not None: + video_rotary_emb = rope_cache["video_rotary_emb"] + audio_rotary_emb = rope_cache["audio_rotary_emb"] + video_cross_attn_rotary_emb = rope_cache["video_cross_attn_rotary_emb"] + audio_cross_attn_rotary_emb = rope_cache["audio_cross_attn_rotary_emb"] + else: + if video_coords is None: + video_coords = self.rope.prepare_video_coords(batch_size, num_frames, height, width, fps=fps) + if audio_coords is None: + audio_coords = self.audio_rope.prepare_audio_coords(batch_size, audio_num_frames) - video_rotary_emb = self.rope(video_coords) - audio_rotary_emb = self.audio_rope(audio_coords) + video_rotary_emb = self.rope(video_coords) + audio_rotary_emb = self.audio_rope(audio_coords) - video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :]) - audio_cross_attn_rotary_emb = self.cross_attn_audio_rope(audio_coords[:, 0:1, :]) + video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :]) + audio_cross_attn_rotary_emb = self.cross_attn_audio_rope(audio_coords[:, 0:1, :]) # 2. Patchify input projections with jax.named_scope("Input Projection"): @@ -1152,9 +1379,8 @@ def __call__( temb_prompt_audio = None if use_cross_timestep: - assert ( - sigma is not None and audio_sigma is not None - ), "sigma and audio_sigma must be provided when use_cross_timestep is True" + if sigma is None or audio_sigma is None: + raise ValueError("sigma and audio_sigma must be provided when use_cross_timestep is True") video_ca_timestep = audio_sigma.flatten() audio_ca_timestep = sigma.flatten() else: @@ -1187,7 +1413,7 @@ def __call__( ) audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.reshape(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1]) - if self.use_prompt_embeddings and self.caption_projection is not None: + if self.use_prompt_embeddings and self.caption_projection is not None and rope_cache is None: encoder_hidden_states = self.caption_projection(encoder_hidden_states) audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states) @@ -1195,38 +1421,60 @@ def __call__( audio_encoder_hidden_states = audio_encoder_hidden_states.reshape(batch_size, -1, audio_hidden_states.shape[-1]) # 5. Run transformer blocks with jax.named_scope("Transformer Blocks"): + base_context = LTX2BlockContext( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + temb=temb, + temb_audio=temb_audio, + temb_ca_scale_shift=video_cross_attn_scale_shift, + temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, + temb_ca_gate=video_cross_attn_a2v_gate, + temb_ca_audio_gate=audio_cross_attn_v2a_gate, + temb_prompt=temb_prompt, + temb_prompt_audio=temb_prompt_audio, + video_rotary_emb=video_rotary_emb, + audio_rotary_emb=audio_rotary_emb, + ca_video_rotary_emb=video_cross_attn_rotary_emb, + ca_audio_rotary_emb=audio_cross_attn_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=audio_encoder_attention_mask, + a2v_cross_attention_mask=a2v_cross_attention_mask, + v2a_cross_attention_mask=v2a_cross_attention_mask, + modality_mask=modality_mask, + cached_kv=cached_kv, + ) + + def apply_block(block, context: LTX2BlockContext, mask, block_cached_kv=None): + orig_perturbation_mask = context.perturbation_mask + orig_cached_kv = context.cached_kv + context = context.replace(perturbation_mask=mask, cached_kv=block_cached_kv) + with jax.named_scope("Transformer Layer"): + hidden_states_out, audio_hidden_states_out = block(context) + context = context.replace( + hidden_states=hidden_states_out.astype(context.hidden_states.dtype), + audio_hidden_states=audio_hidden_states_out.astype(context.audio_hidden_states.dtype), + perturbation_mask=orig_perturbation_mask, + cached_kv=orig_cached_kv, + ) + return context + if perturbation_mask is None: - # Fast-path: No perturbation masking (standard LTX-2 or disabled STG) - def scan_fn_ltx2(carry, block): - hidden_states, audio_hidden_states, rngs_carry = carry - with jax.named_scope("Transformer Layer"): - hidden_states_out, audio_hidden_states_out = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=None, - modality_mask=modality_mask, - ) - return ( - hidden_states_out.astype(hidden_states.dtype), - audio_hidden_states_out.astype(audio_hidden_states.dtype), - rngs_carry, - ), None + if cached_kv is not None: + + def scan_fn_ltx2(carry, block_and_kv): + block, layer_kv_cache = block_and_kv + context, rngs_carry = carry + context = apply_block(block, context, None, layer_kv_cache) + return (context, rngs_carry), None + + else: + + def scan_fn_ltx2(carry, block): + context, rngs_carry = carry + context = apply_block(block, context, None, None) + return (context, rngs_carry), None if self.scan_layers: rematted_scan_fn = self.gradient_checkpoint.apply( @@ -1235,77 +1483,53 @@ def scan_fn_ltx2(carry, block): self.names_which_can_be_offloaded, prevent_cse=not self.scan_layers, ) - carry = (hidden_states, audio_hidden_states, nnx.Rngs(0)) - (hidden_states, audio_hidden_states, _), _ = nnx.scan( + carry = (base_context, nnx.Rngs(0)) + + if cached_kv is not None: + scan_input = (self.transformer_blocks, cached_kv) + else: + scan_input = self.transformer_blocks + + (final_context, _), _ = nnx.scan( rematted_scan_fn, length=self.num_layers, in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0), transform_metadata={nnx.PARTITION_NAME: "layers"}, - )(carry, self.transformer_blocks) + )(carry, scan_input) + hidden_states = final_context.hidden_states + audio_hidden_states = final_context.audio_hidden_states else: - for block in self.transformer_blocks: - hidden_states, audio_hidden_states = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=None, - modality_mask=modality_mask, - ) + current_context = base_context + for i, block in enumerate(self.transformer_blocks): + layer_kv_cache = None + if cached_kv is not None: + layer_kv_cache = jax.tree.map(lambda x: x[i], cached_kv) + current_context = apply_block(block, current_context, None, layer_kv_cache) + hidden_states = current_context.hidden_states + audio_hidden_states = current_context.audio_hidden_states else: - # Slow-path: Dynamic perturbation masking (LTX-2.3 STG enabled) masks = jnp.ones((self.num_layers, batch_size, 1, 1), dtype=self.dtype) for i in self.spatio_temporal_guidance_blocks: if i < self.num_layers: masks = masks.at[i].set(perturbation_mask) perturbation_mask_per_layer = masks - def scan_fn_ltx23(carry, block_and_mask): - block, mask = block_and_mask - hidden_states, audio_hidden_states, rngs_carry = carry - with jax.named_scope("Transformer Layer"): - hidden_states_out, audio_hidden_states_out = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=mask, - modality_mask=modality_mask, - ) - return ( - hidden_states_out.astype(hidden_states.dtype), - audio_hidden_states_out.astype(audio_hidden_states.dtype), - rngs_carry, - ), None + 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 + context, rngs_carry = carry + context = apply_block(block, context, mask, layer_kv_cache) + return (context, rngs_carry), None + + else: + + def scan_fn_ltx23(carry, block_and_mask): + block, mask = block_and_mask + context, rngs_carry = carry + context = apply_block(block, context, mask, None) + return (context, rngs_carry), None if self.scan_layers: rematted_scan_fn = self.gradient_checkpoint.apply( @@ -1314,39 +1538,39 @@ def scan_fn_ltx23(carry, block_and_mask): self.names_which_can_be_offloaded, prevent_cse=not self.scan_layers, ) - carry = (hidden_states, audio_hidden_states, nnx.Rngs(0)) - (hidden_states, audio_hidden_states, _), _ = nnx.scan( + carry = (base_context, nnx.Rngs(0)) + + if cached_kv is not None: + scan_input = ( + self.transformer_blocks, + perturbation_mask_per_layer, + cached_kv, + ) + else: + scan_input = ( + self.transformer_blocks, + perturbation_mask_per_layer, + ) + + (final_context, _), _ = nnx.scan( rematted_scan_fn, length=self.num_layers, in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0), transform_metadata={nnx.PARTITION_NAME: "layers"}, - )(carry, (self.transformer_blocks, perturbation_mask_per_layer)) + )(carry, scan_input) + hidden_states = final_context.hidden_states + audio_hidden_states = final_context.audio_hidden_states else: + current_context = base_context for i, block in enumerate(self.transformer_blocks): mask = perturbation_mask_per_layer[i] if perturbation_mask_per_layer is not None else None - hidden_states, audio_hidden_states = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=mask, - modality_mask=modality_mask, - ) + layer_kv_cache = None + if cached_kv is not None: + layer_kv_cache = jax.tree.map(lambda x: x[i], cached_kv) + current_context = apply_block(block, current_context, mask, layer_kv_cache) + hidden_states = current_context.hidden_states + audio_hidden_states = current_context.audio_hidden_states # 6. Output layers with jax.named_scope("Output Projection & Norm"): diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index 19069b6f9..4eb8ce2ef 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -56,9 +56,15 @@ from ...pyconfig import HyperParameters from ... import max_logging from ... import max_utils +from ... import aot_cache from ...max_utils import get_precision, device_put_replicated, get_flash_block_sizes +@partial(jax.jit, static_argnums=(1,)) +def _enforce_layout(x, axes): + return jax.lax.with_sharding_constraint(x, axes) + + TORCH_DTYPE_MAP = { "bfloat16": torch.bfloat16, "float16": torch.float16, @@ -81,6 +87,8 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): """ std_text = jnp.std(noise_pred_text, axis=list(range(1, noise_pred_text.ndim)), keepdims=True) std_cfg = jnp.std(noise_cfg, axis=list(range(1, noise_cfg.ndim)), keepdims=True) + # Prevent division by zero + std_cfg = jnp.maximum(std_cfg, 1e-5) # rescale the results from guidance (fixes overexposure) noise_pred_rescaled = noise_cfg * (std_text / std_cfg) # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images @@ -148,6 +156,8 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): ltx2_config["precision"] = get_precision(config) ltx2_config["flash_block_sizes"] = get_flash_block_sizes(config) ltx2_config["flash_min_seq_length"] = getattr(config, "flash_min_seq_length", 4096) + ltx2_config["ulysses_shards"] = getattr(config, "ulysses_shards", -1) + ltx2_config["ulysses_attention_chunks"] = getattr(config, "ulysses_attention_chunks", 1) ltx2_config["remat_policy"] = config.remat_policy ltx2_config["names_which_can_be_saved"] = config.names_which_can_be_saved ltx2_config["names_which_can_be_offloaded"] = config.names_which_can_be_offloaded @@ -186,13 +196,64 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): ) params = jax.tree_util.tree_map_with_path( - lambda path, x: cast_with_exclusion(path, x, dtype_to_cast=config.weights_dtype), params + lambda path, x: cast_with_exclusion(path, x, dtype_to_cast=config.weights_dtype), + params, ) + + t_put_start = time.perf_counter() + put_specs = [] for path, val in flax.traverse_util.flatten_dict(params).items(): if restored_checkpoint: path = path[:-1] sharding = logical_state_sharding[path].get_value() - state[path].set_value(device_put_replicated(val, sharding)) + put_specs.append((path, val, sharding)) + + if jax.process_count() == 1: + n_devices = mesh.devices.size + dim0_sharding = NamedSharding(mesh, P(mesh.axis_names)) + + def stages_via_ici(val, sharding) -> bool: + return ( + sharding.is_fully_replicated + and val.ndim > 0 + and val.shape[0] % n_devices == 0 + and val.nbytes >= 1 << 26 # 64MB: below this, staging overhead wins + ) + + staged_ids = [i for i, (_, val, sharding) in enumerate(put_specs) if stages_via_ici(val, sharding)] + direct_ids = [i for i in range(len(put_specs)) if i not in set(staged_ids)] + + put_arrays = [None] * len(put_specs) + if staged_ids: + staged = jax.device_put([put_specs[i][1] for i in staged_ids], [dim0_sharding] * len(staged_ids)) + replicate_fn = jax.jit(lambda xs: xs, out_shardings=[put_specs[i][2] for i in staged_ids]) + for i, replicated in zip(staged_ids, replicate_fn(staged)): + put_arrays[i] = replicated + if direct_ids: + for i, put_array in zip( + direct_ids, + jax.device_put( + [put_specs[i][1] for i in direct_ids], + [put_specs[i][2] for i in direct_ids], + ), + ): + put_arrays[i] = put_array + for (path, _, _), put_array in zip(put_specs, put_arrays): + state[path].set_value(put_array) + else: + for path, val, sharding in put_specs: + try: + state[path].set_value(device_put_replicated(val, sharding)) + except Exception as e: + max_logging.log(f"Failed to device_put_replicated {path}: {e}") + max_logging.log(f"Trying to use process_allgather for {path}") + val_on_host = jax.experimental.multihost_utils.process_allgather(val, tiled=True) + state[path].set_value(device_put_replicated(val_on_host, sharding)) + del val_on_host + + jax.block_until_ready([state[path].value for path, _, _ in put_specs]) + max_logging.log(f"Transformer {subfolder or 'transformer'} weights on device in {time.perf_counter() - t_put_start:.1f}s") + state = nnx.from_flat_state(state) transformer = nnx.merge(graphdef, state, rest_of_state) @@ -364,7 +425,13 @@ def load_text_encoder(cls, config: HyperParameters): return text_encoder @classmethod - def load_connectors(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters): + def load_connectors( + cls, + devices_array: np.array, + mesh: Mesh, + rngs: nnx.Rngs, + config: HyperParameters, + ): max_logging.log("Loading Connectors...") def create_model(rngs: nnx.Rngs, config: HyperParameters): @@ -411,7 +478,13 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters): return connectors @classmethod - def load_vae(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters): + def load_vae( + cls, + devices_array: np.array, + mesh: Mesh, + rngs: nnx.Rngs, + config: HyperParameters, + ): max_logging.log("Loading Video VAE...") def create_model(rngs: nnx.Rngs, config: HyperParameters): @@ -461,7 +534,13 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters): return vae @classmethod - def load_audio_vae(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters): + def load_audio_vae( + cls, + devices_array: np.array, + mesh: Mesh, + rngs: nnx.Rngs, + config: HyperParameters, + ): max_logging.log("Loading Audio VAE...") def create_model(rngs: nnx.Rngs, config: HyperParameters): @@ -527,7 +606,13 @@ def load_transformer( return transformer @classmethod - def load_vocoder(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters): + def load_vocoder( + cls, + devices_array: np.array, + mesh: Mesh, + rngs: nnx.Rngs, + config: HyperParameters, + ): max_logging.log("Loading Vocoder...") # Determine correct source path based on which vocoder we are loading @@ -584,7 +669,13 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters, model_path: str, use_b return vocoder @classmethod - def load_upsampler(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters): + def load_upsampler( + cls, + devices_array: np.array, + mesh: Mesh, + rngs: nnx.Rngs, + config: HyperParameters, + ): """ LTX2LatentUpsamplerModel is a flax.linen.Module, so we do not use nnx.eval_shape or nnx.split. Instead, we return the instantiated Module and the explicitly loaded parameters to be used @@ -615,7 +706,9 @@ def load_upsampler(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, con spatial_upsample = upsampler_config.get("spatial_upsample", True) temporal_upsample = upsampler_config.get("temporal_upsample", False) rational_spatial_scale = getattr( - config, "upsampler_rational_spatial_scale", upsampler_config.get("rational_spatial_scale", 2.0) + config, + "upsampler_rational_spatial_scale", + upsampler_config.get("rational_spatial_scale", 2.0), ) mid_channels = upsampler_config.get("mid_channels", 1024) @@ -657,7 +750,11 @@ def load_upsampler(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, con # Load weights from disk. Evaluating eval_shapes=None returns the raw checkpoint dict. params = load_upsampler_weights( - config.upsampler_model_path, eval_shapes=None, device="cpu", subfolder=subfolder, filename=filename + config.upsampler_model_path, + eval_shapes=None, + device="cpu", + subfolder=subfolder, + filename=filename, ) if hasattr(config, "weights_dtype"): @@ -714,27 +811,45 @@ def _create_common_components(cls, config: HyperParameters, vae_only=False): @classmethod def _load_and_init( - cls, config: HyperParameters, restored_checkpoint, vae_only=False, load_transformer=True, load_upsampler=False + cls, + config: HyperParameters, + restored_checkpoint, + vae_only=False, + load_transformer=True, + load_upsampler=False, ): - components = cls._create_common_components(config, vae_only) + import concurrent.futures - transformer = None - if load_transformer: - max_logging.log("Loading Transformer...") - transformer = cls.load_transformer( - devices_array=components["devices_array"], - mesh=components["mesh"], - rngs=components["rngs"], - config=config, - restored_checkpoint=restored_checkpoint, - ) + common_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + common_future = common_executor.submit(cls._create_common_components, config, vae_only) + transformer = None latent_upsampler = None latent_upsampler_params = None - if load_upsampler: - latent_upsampler, latent_upsampler_params = cls.load_upsampler( - devices_array=components["devices_array"], mesh=components["mesh"], rngs=components["rngs"], config=config - ) + + if load_transformer or load_upsampler: + # Need mesh and devices to load these + devices_array = max_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + rngs = nnx.Rngs(jax.random.key(config.seed)) + + if load_transformer: + max_logging.log("Loading Transformer...") + transformer = cls.load_transformer( + devices_array=devices_array, + mesh=mesh, + rngs=rngs, + config=config, + restored_checkpoint=restored_checkpoint, + ) + + if load_upsampler: + latent_upsampler, latent_upsampler_params = cls.load_upsampler( + devices_array=devices_array, mesh=mesh, rngs=rngs, config=config + ) + + components = common_future.result() + common_executor.shutdown() pipeline = cls( scheduler=components["scheduler"], @@ -755,13 +870,24 @@ def _load_and_init( return pipeline, pipeline.transformer @classmethod - def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transformer=True, load_upsampler=False): + def from_pretrained( + cls, + config: HyperParameters, + vae_only=False, + load_transformer=True, + load_upsampler=False, + ): pipeline, _ = cls._load_and_init(config, None, vae_only, load_transformer, load_upsampler) return pipeline @classmethod def from_checkpoint( - cls, config: HyperParameters, restored_checkpoint, vae_only=False, load_transformer=True, load_upsampler=False + cls, + config: HyperParameters, + restored_checkpoint, + vae_only=False, + load_transformer=True, + load_upsampler=False, ): pipeline, _ = cls._load_and_init(config, restored_checkpoint, vae_only, load_transformer, load_upsampler) return pipeline @@ -855,6 +981,8 @@ def _get_gemma_prompt_embeds( prompt = [p.strip() for p in prompt] + target_dtype = dtype if dtype is not None else jnp.bfloat16 + if self.text_encoder is not None: run_text_encoder_on_tpu = getattr(self.config, "run_text_encoder_on_tpu", False) if hasattr(self, "config") else False if run_text_encoder_on_tpu: @@ -872,30 +1000,33 @@ def _get_gemma_prompt_embeds( # Distribute the batch dimension across available TPUs to prevent Softmax OOM # (reduces 512MB allocation down to 64MB per TPU for batch size 16) - devices = np.array(jax.devices()) - num_shards = 1 - for i in range(len(devices), 0, -1): - if text_input_ids.shape[0] % i == 0: - num_shards = i - break - - if num_shards > 1: - mesh = Mesh(devices[:num_shards], axis_names=("batch",)) - sharding = NamedSharding(mesh, P("batch")) + if hasattr(self, "mesh") and self.mesh is not None: + data_axis = self.mesh.axis_names[0] + sharding = NamedSharding(self.mesh, P(data_axis)) text_input_ids = jax.device_put(text_input_ids, sharding) prompt_attention_mask = jax.device_put(prompt_attention_mask, sharding) + else: + devices = np.array(jax.devices()) + num_shards = 1 + for i in range(len(devices), 0, -1): + if text_input_ids.shape[0] % i == 0: + num_shards = i + break + + if num_shards > 1: + mesh = Mesh(devices[:num_shards], axis_names=("batch",)) + sharding = NamedSharding(mesh, P("batch")) + text_input_ids = jax.device_put(text_input_ids, sharding) + prompt_attention_mask = jax.device_put(prompt_attention_mask, sharding) # Torchax wrapper returns tuple of hidden states natively text_encoder_hidden_states = self.text_encoder( - input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, ) - prompt_embeds_list = [] - # Iterate instead of stacking eagerly to avoid 5.7+ GB HBM allocations outside JIT - for state in text_encoder_hidden_states: - prompt_embeds_list.append(state.astype(jnp.bfloat16)) - - prompt_embeds = prompt_embeds_list + prompt_embeds = jax.tree.map(lambda x: x.astype(target_dtype), list(text_encoder_hidden_states)) del text_encoder_hidden_states # Free memory prompt_attention_mask = prompt_attention_mask.astype(jnp.bool_) @@ -917,31 +1048,27 @@ def _get_gemma_prompt_embeds( with torch.no_grad(): text_encoder_outputs = self.text_encoder( - input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, ) text_encoder_hidden_states = text_encoder_outputs.hidden_states del text_encoder_outputs # Free memory - prompt_embeds_list = [] - # Iterate instead of stacking eagerly to avoid 5.7+ GB HBM allocations outside JIT - for state in text_encoder_hidden_states: - state_np = state.cpu().to(torch.float32).numpy() - prompt_embeds_list.append(jnp.array(state_np, dtype=jnp.bfloat16)) - - prompt_embeds = prompt_embeds_list + prompt_embeds = jax.tree.map( + lambda state: jnp.array(state.cpu().to(torch.float32).numpy(), dtype=target_dtype), + list(text_encoder_hidden_states), + ) del text_encoder_hidden_states # Free PyTorch tensor memory - prompt_attention_mask = jnp.array(prompt_attention_mask.cpu().to(torch.float32).numpy(), dtype=jnp.bool_) + prompt_attention_mask = jnp.array( + prompt_attention_mask.cpu().to(torch.float32).numpy(), + dtype=jnp.bool_, + ) else: raise ValueError("`text_encoder` is required to encode prompts.") - if dtype is not None: - if isinstance(prompt_embeds, list): - prompt_embeds = [state.astype(dtype) for state in prompt_embeds] - else: - prompt_embeds = prompt_embeds.astype(dtype) - if isinstance(prompt_embeds, list): _, seq_len, _ = prompt_embeds[0].shape prompt_embeds = [ @@ -1030,7 +1157,10 @@ def encode_prompt( f" {type(prompt)}." ) - negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + ( + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self._get_gemma_prompt_embeds( prompt=negative_prompt, num_videos_per_prompt=num_videos_per_prompt, max_sequence_length=max_sequence_length, @@ -1038,7 +1168,12 @@ def encode_prompt( dtype=dtype, ) - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + return ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) def check_inputs( self, @@ -1117,7 +1252,12 @@ def _pack_latents(latents: jax.Array, patch_size: int = 1, patch_size_t: int = 1 @staticmethod def _unpack_latents( - latents: jax.Array, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 + latents: jax.Array, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, ) -> jax.Array: batch_size = latents.shape[0] # latents: (Batch, SeqLen, Channels*Patches) @@ -1136,7 +1276,10 @@ def _unpack_latents( @staticmethod def _normalize_latents( - latents: jax.Array, latents_mean: jax.Array, latents_std: jax.Array, scaling_factor: float = 1.0 + latents: jax.Array, + latents_mean: jax.Array, + latents_std: jax.Array, + scaling_factor: float = 1.0, ) -> jax.Array: latents_mean = latents_mean.reshape(1, -1, 1, 1, 1).astype(latents.dtype) latents_std = latents_std.reshape(1, -1, 1, 1, 1).astype(latents.dtype) @@ -1145,7 +1288,10 @@ def _normalize_latents( @staticmethod def _denormalize_latents( - latents: jax.Array, latents_mean: jax.Array, latents_std: jax.Array, scaling_factor: float = 1.0 + latents: jax.Array, + latents_mean: jax.Array, + latents_std: jax.Array, + scaling_factor: float = 1.0, ) -> jax.Array: latents_mean = latents_mean.reshape(1, -1, 1, 1, 1).astype(latents.dtype) latents_std = latents_std.reshape(1, -1, 1, 1, 1).astype(latents.dtype) @@ -1169,7 +1315,8 @@ def _decode_latents_to_video( # Replicate VAE weights graphdef, state = nnx.split(self.vae) state = jax.tree_util.tree_map( - lambda x: jax.lax.with_sharding_constraint(x, replicated_sharding) if isinstance(x, jax.Array) else x, state + lambda x: jax.lax.with_sharding_constraint(x, replicated_sharding) if isinstance(x, jax.Array) else x, + state, ) self.vae = nnx.merge(graphdef, state) except Exception as e: # pylint: disable=broad-exception-caught @@ -1235,6 +1382,9 @@ def _create_noised_state(latents: jax.Array, noise_scale: float, generator: Opti else: # Fallback or expect noise to be handled otherwise? # pipeline prepare_latents typically generates noise. + max_logging.log( + "WARNING: No PRNG generator provided. Falling back to deterministic zero-seed noise (jax.random.key(0))." + ) noise = jax.random.normal(jax.random.key(0), latents.shape, dtype=latents.dtype) # Default fallback noised_latents = noise_scale * noise + (1 - noise_scale) * latents @@ -1242,13 +1392,22 @@ def _create_noised_state(latents: jax.Array, noise_scale: float, generator: Opti @staticmethod def _pack_audio_latents( - latents: jax.Array, patch_size: Optional[int] = None, patch_size_t: Optional[int] = None + latents: jax.Array, + patch_size: Optional[int] = None, + patch_size_t: Optional[int] = None, ) -> jax.Array: if patch_size is not None and patch_size_t is not None: batch_size, _, latent_length, latent_mel_bins = latents.shape post_patch_latent_length = latent_length // patch_size_t post_patch_mel_bins = latent_mel_bins // patch_size - latents = latents.reshape(batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size) + latents = latents.reshape( + batch_size, + -1, + post_patch_latent_length, + patch_size_t, + post_patch_mel_bins, + patch_size, + ) # Permute to (Batch, T', F', C, p_t, p) latents = latents.transpose(0, 2, 4, 1, 3, 5) latents = latents.reshape(batch_size, post_patch_latent_length * post_patch_mel_bins, -1) @@ -1280,9 +1439,19 @@ def _unpack_audio_latents( # latents: (Batch, Seq, Dim) # Pack: (B, C, L, F) -> (B, C, L', pt, F', p) -> (B, C, L', pt, F', p) -> (B, L', F', C, pt, p) -> (B, L', F', C*pt*p) # Unpack: (B, L'*F', C*pt*p) -> (B, L', F', C, pt, p) -> (B, C, L', pt, F', p) -> (B, C, L'*pt, F'*p) - latents = latents.reshape(batch_size, -1, num_mel_bins // patch_size, num_channels * patch_size_t * patch_size) latents = latents.reshape( - batch_size, latent_length // patch_size_t, num_mel_bins // patch_size, num_channels, patch_size_t, patch_size + batch_size, + -1, + num_mel_bins // patch_size, + num_channels * patch_size_t * patch_size, + ) + latents = latents.reshape( + batch_size, + latent_length // patch_size_t, + num_mel_bins // patch_size, + num_channels, + patch_size_t, + patch_size, ) latents = latents.transpose(0, 3, 1, 4, 2, 5).reshape(batch_size, num_channels, latent_length, num_mel_bins) # Wait, reshape order needs to match pack? @@ -1321,7 +1490,11 @@ def prepare_latents( # The packing and unpacking mechanisms expect (B, C, T, H, W). latents = latents.transpose(0, 4, 1, 2, 3) - latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) + latents = self._pack_latents( + latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) if latents.ndim != 3: raise ValueError("Unexpected latents shape") latents = self._create_noised_state(latents, noise_scale, generator) @@ -1337,7 +1510,11 @@ def prepare_latents( generator = jax.random.key(seed) latents = jax.random.normal(generator, shape, dtype=dtype or jnp.float32) - latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) + latents = self._pack_latents( + latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) return latents def prepare_audio_latents( @@ -1356,7 +1533,9 @@ def prepare_audio_latents( if latents.ndim == 4: # (Batch, Channels, Length, Mel) -> Pack latents = self._pack_audio_latents( - latents, getattr(self.audio_vae.config, "patch_size", None), getattr(self.audio_vae.config, "patch_size_t", None) + latents, + getattr(self.audio_vae.config, "patch_size", None), + getattr(self.audio_vae.config, "patch_size_t", None), ) if latents.ndim != 3: raise ValueError("Unexpected audio latents shape") @@ -1376,7 +1555,9 @@ def prepare_audio_latents( latents = jax.random.normal(generator, shape, dtype=dtype or jnp.float32) latents = self._pack_audio_latents( - latents, getattr(self.audio_vae.config, "patch_size", None), getattr(self.audio_vae.config, "patch_size_t", None) + latents, + getattr(self.audio_vae.config, "patch_size", None), + getattr(self.audio_vae.config, "patch_size_t", None), ) return latents @@ -1431,12 +1612,23 @@ def __call__( t0_init = time.perf_counter() # 1. Check inputs self.check_inputs( - prompt, height, width, prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask + prompt, + height, + width, + prompt_embeds, + negative_prompt_embeds, + prompt_attention_mask, + negative_prompt_attention_mask, ) # 2. Encode inputs (Text) t0_encode = time.perf_counter() - prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt( + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( prompt, negative_prompt, do_classifier_free_guidance=guidance_scale > 1.0, @@ -1548,7 +1740,13 @@ def __call__( ] else: prompt_embeds_jax = jnp.concatenate( - [negative_prompt_embeds_jax, prompt_embeds_jax, prompt_embeds_jax, prompt_embeds_jax], axis=0 + [ + negative_prompt_embeds_jax, + prompt_embeds_jax, + prompt_embeds_jax, + prompt_embeds_jax, + ], + axis=0, ) prompt_attention_mask_jax = jnp.concatenate( @@ -1607,7 +1805,10 @@ def __call__( t0_connectors = time.perf_counter() with jax.named_scope("connectors_pass"): video_embeds, audio_embeds, new_attention_mask = self._run_connectors( - connectors_graphdef, connectors_state, prompt_embeds_jax, prompt_attention_mask_jax.astype(jnp.bool_) + connectors_graphdef, + connectors_state, + prompt_embeds_jax, + prompt_attention_mask_jax.astype(jnp.bool_), ) video_embeds = video_embeds.block_until_ready() connectors_time = time.perf_counter() - t0_connectors @@ -1617,12 +1818,14 @@ def __call__( audio_embeds_sharded = audio_embeds if not self.transformer.scan_layers: - activation_axes = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) - activation_axes_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) - spec = NamedSharding(self.mesh, P(*activation_axes)) - spec_audio = NamedSharding(self.mesh, P(*activation_axes_audio)) - video_embeds_sharded = jax.device_put(video_embeds, spec) - audio_embeds_sharded = jax.device_put(audio_embeds, spec_audio) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + activation_axes = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) + activation_axes_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) + + latents_jax = _enforce_layout(latents_jax, P(*activation_axes)) + audio_latents_jax = _enforce_layout(audio_latents_jax, P(*activation_axes_audio)) + video_embeds_sharded = _enforce_layout(video_embeds_sharded, P(*activation_axes)) + audio_embeds_sharded = _enforce_layout(audio_embeds_sharded, P(*activation_axes_audio)) timesteps_jax = jnp.array(timesteps, dtype=jnp.float32) t0_denoise = time.perf_counter() @@ -1657,6 +1860,9 @@ def __call__( self.scheduler.step, tuple(tuple(rule) if isinstance(rule, list) else rule for rule in self.config.logical_axis_rules), use_cross_timestep=use_cross_timestep, + do_cfg=do_cfg, + do_stg=do_stg, + use_kv_cache=getattr(self.config, "use_kv_cache", False), ) else: # Old Python loop path @@ -1668,8 +1874,13 @@ def __call__( audio_latents_jax_sharded = audio_latents_jax if not self.transformer.scan_layers: - activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) - activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + activation_axis_names = nn.logical_to_mesh_axes(( + "activation_batch", + "activation_length", + "activation_embed", + )) + activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) latents_jax_sharded = jax.lax.with_sharding_constraint(latents_jax, activation_axis_names) audio_latents_jax_sharded = jax.lax.with_sharding_constraint(audio_latents_jax, activation_axis_names_audio) @@ -1716,7 +1927,12 @@ def convert_to_vel(lat, x0, sig): return (lat - x0) / sig if do_cfg and do_stg: - noise_pred_uncond, noise_pred_text, noise_pred_perturb, noise_pred_isolated = jnp.split(noise_pred, 4, axis=0) + ( + noise_pred_uncond, + noise_pred_text, + noise_pred_perturb, + noise_pred_isolated, + ) = jnp.split(noise_pred, 4, axis=0) x0_uncond = convert_to_x0(latents_step, noise_pred_uncond, sigma_t) x0_text = convert_to_x0(latents_step, noise_pred_text, sigma_t) x0_perturb = convert_to_x0(latents_step, noise_pred_perturb, sigma_t) @@ -1732,9 +1948,12 @@ def convert_to_vel(lat, x0, sig): noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t) # Audio - noise_pred_audio_uncond, noise_pred_audio_text, noise_pred_audio_perturb, noise_pred_audio_isolated = jnp.split( - noise_pred_audio, 4, axis=0 - ) + ( + noise_pred_audio_uncond, + noise_pred_audio_text, + noise_pred_audio_perturb, + noise_pred_audio_isolated, + ) = jnp.split(noise_pred_audio, 4, axis=0) x0_audio_uncond = convert_to_x0(audio_latents_step, noise_pred_audio_uncond, sigma_t) x0_audio_text = convert_to_x0(audio_latents_step, noise_pred_audio_text, sigma_t) x0_audio_perturb = convert_to_x0(audio_latents_step, noise_pred_audio_perturb, sigma_t) @@ -1752,7 +1971,9 @@ def convert_to_vel(lat, x0, sig): x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta if audio_guidance_rescale > 0: x0_audio_combined = rescale_noise_cfg( - x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale + x0_audio_combined, + x0_audio_text, + guidance_rescale=audio_guidance_rescale, ) noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) @@ -1764,7 +1985,11 @@ def convert_to_vel(lat, x0, sig): latents_step, _ = self.scheduler.step(scheduler_state, noise_pred, t, latents_step, return_dict=False) audio_latents_step, _ = self.scheduler.step( - scheduler_state, noise_pred_audio, t, audio_latents_step, return_dict=False + scheduler_state, + noise_pred_audio, + t, + audio_latents_step, + return_dict=False, ) if do_cfg and do_stg: @@ -1804,7 +2029,10 @@ def convert_to_vel(lat, x0, sig): self.transformer_temporal_patch_size, ) latents = self._denormalize_latents( - latents, self.vae.latents_mean.value, self.vae.latents_std.value, self.vae.config.scaling_factor + latents, + self.vae.latents_mean.value, + self.vae.latents_std.value, + self.vae.config.scaling_factor, ) # VAE expects channels last (B, T, H, W, C) but unpack returns (B, C, T, H, W) @@ -1848,14 +2076,19 @@ def convert_to_vel(lat, x0, sig): # Denormalize and Unpack Audio (Order important: Denorm THEN Unpack) audio_latents = self._denormalize_audio_latents( - audio_latents_jax, self.audio_vae.latents_mean.value, self.audio_vae.latents_std.value + audio_latents_jax, + self.audio_vae.latents_mean.value, + self.audio_vae.latents_std.value, ) num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64 latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio audio_latents = self._unpack_audio_latents( - audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins, num_channels=audio_channels + audio_latents, + audio_num_frames, + num_mel_bins=latent_mel_bins, + num_channels=audio_channels, ) # Audio VAE expects channels last (B, T, F, C) but unpack returns (B, C, T, F) @@ -1956,7 +2189,7 @@ def convert_to_vel(lat, x0, sig): @partial( - jax.jit, + aot_cache.cached_jit, static_argnames=( "latent_num_frames", "latent_height", @@ -1988,6 +2221,8 @@ def transformer_forward_pass( audio_sigma=None, use_cross_timestep=False, is_cfg_stg_mode: bool = False, + kv_cache=None, + rope_cache=None, ): """Forward pass for the transformer.""" # pylint: disable=too-many-positional-arguments,unused-argument @@ -2039,22 +2274,16 @@ def transformer_forward_pass( return_dict=False, perturbation_mask=perturbation_mask, use_cross_timestep=use_cross_timestep, + cached_kv=kv_cache, + rope_cache=rope_cache, ) return noise_pred, noise_pred_audio @partial( - jax.jit, + aot_cache.cached_jit, static_argnames=( - "guidance_scale", - "stg_scale", - "modality_scale", - "guidance_rescale", - "audio_guidance_scale", - "audio_stg_scale", - "audio_modality_scale", - "audio_guidance_rescale", "latent_num_frames", "latent_height", "latent_width", @@ -2065,6 +2294,9 @@ def transformer_forward_pass( "scheduler_step", "logical_axis_rules", "use_cross_timestep", + "do_cfg", + "do_stg", + "use_kv_cache", ), ) def run_diffusion_loop( @@ -2097,14 +2329,35 @@ def run_diffusion_loop( logical_axis_rules, perturbation_mask=None, use_cross_timestep=False, + do_cfg=False, + do_stg=False, + use_kv_cache=False, + cached_kv=None, + rope_cache=None, ): """Runs the diffusion loop.""" # pylint: disable=too-many-positional-arguments latents_jax = latents_jax.astype(jnp.float32) audio_latents_jax = audio_latents_jax.astype(jnp.float32) - do_cfg = guidance_scale > 1.0 - do_stg = stg_scale > 0.0 + transformer = nnx.merge(graphdef, state) + kv_cache = None + rope_cache = None + if use_kv_cache: + ( + kv_cache, + rope_cache, + video_embeds_sharded, + audio_embeds_sharded, + ) = transformer.compute_kv_cache( + video_embeds_sharded, + audio_embeds_sharded, + latent_num_frames, + latent_height, + latent_width, + fps, + audio_num_frames, + ) # Helper functions matching Diffusers Delta formulation def convert_to_x0(lat, vel, sigma_t): @@ -2113,137 +2366,156 @@ def convert_to_x0(lat, vel, sigma_t): def convert_to_vel(lat, x0, sigma_t): return (lat - x0) / sigma_t + if not scan_layers: + with nn_partitioning.axis_rules(logical_axis_rules): + activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) + activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) + latents_jax = jax.lax.with_sharding_constraint(latents_jax, activation_axis_names) + audio_latents_jax = jax.lax.with_sharding_constraint(audio_latents_jax, activation_axis_names_audio) + video_embeds_sharded = jax.lax.with_sharding_constraint(video_embeds_sharded, activation_axis_names) + audio_embeds_sharded = jax.lax.with_sharding_constraint(audio_embeds_sharded, activation_axis_names_audio) + def scan_body(carry, inputs): t, sigma_t = inputs latents, audio_latents, s_state = carry - with nn_partitioning.axis_rules(logical_axis_rules): - latents_sharded = latents - audio_latents_sharded = audio_latents - - if not scan_layers: - activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) - latents_sharded = jax.lax.with_sharding_constraint(latents, activation_axis_names) - audio_latents_sharded = jax.lax.with_sharding_constraint(audio_latents, activation_axis_names) - - # Forward Pass - noise_pred, noise_pred_audio = transformer_forward_pass( - graphdef, - state, - latents_sharded, - audio_latents_sharded, - t, - video_embeds_sharded, - audio_embeds_sharded, - new_attention_mask, - new_attention_mask, - latent_num_frames=latent_num_frames, - latent_height=latent_height, - latent_width=latent_width, - audio_num_frames=audio_num_frames, - fps=fps, - global_batch_size=batch_size, - sigma=t, - audio_sigma=t, - use_cross_timestep=use_cross_timestep, - is_cfg_stg_mode=do_cfg and do_stg, - ) - - # Extract latents_step based on stacking strategy - if do_cfg and do_stg: - latents_step = latents[batch_size : 2 * batch_size] - audio_latents_step = audio_latents[batch_size : 2 * batch_size] - elif do_cfg: - latents_step = latents[batch_size:] - audio_latents_step = audio_latents[batch_size:] - else: - latents_step = latents - audio_latents_step = audio_latents - - # Apply Diffusers STG + CFG + Modality Delta Logic - if do_cfg and do_stg: - noise_pred_uncond, noise_pred_text, noise_pred_perturb, noise_pred_isolated = jnp.split(noise_pred, 4, axis=0) - - # Convert to x0 - x0_uncond = convert_to_x0(latents_step, noise_pred_uncond, sigma_t) - x0_text = convert_to_x0(latents_step, noise_pred_text, sigma_t) - x0_perturb = convert_to_x0(latents_step, noise_pred_perturb, sigma_t) - x0_isolated = convert_to_x0(latents_step, noise_pred_isolated, sigma_t) - - # Delta formulation - cfg_delta = (guidance_scale - 1) * (x0_text - x0_uncond) - stg_delta = stg_scale * (x0_text - x0_perturb) - video_modality_delta = (modality_scale - 1) * (x0_text - x0_isolated) + latents_sharded = latents + audio_latents_sharded = audio_latents + video_embeds_sharded_constrained = video_embeds_sharded + audio_embeds_sharded_constrained = audio_embeds_sharded + + # Forward Pass + noise_pred, noise_pred_audio = transformer_forward_pass( + graphdef, + state, + latents_sharded, + audio_latents_sharded, + t, + video_embeds_sharded_constrained, + audio_embeds_sharded_constrained, + new_attention_mask, + new_attention_mask, + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + audio_num_frames=audio_num_frames, + fps=fps, + global_batch_size=batch_size, + sigma=t, + audio_sigma=t, + use_cross_timestep=use_cross_timestep, + is_cfg_stg_mode=do_cfg and do_stg, + kv_cache=kv_cache, + rope_cache=rope_cache, + ) - x0_combined = x0_text + cfg_delta + stg_delta + video_modality_delta + # Extract latents_step based on stacking strategy + if do_cfg and do_stg: + latents_step = latents[batch_size : 2 * batch_size] + audio_latents_step = audio_latents[batch_size : 2 * batch_size] + elif do_cfg: + latents_step = latents[batch_size:] + audio_latents_step = audio_latents[batch_size:] + else: + latents_step = latents + audio_latents_step = audio_latents - if guidance_rescale > 0: - x0_combined = rescale_noise_cfg(x0_combined, x0_text, guidance_rescale=guidance_rescale) + # Apply Diffusers STG + CFG + Modality Delta Logic + if do_cfg and do_stg: + ( + noise_pred_uncond, + noise_pred_text, + noise_pred_perturb, + noise_pred_isolated, + ) = jnp.split(noise_pred, 4, axis=0) + + # Convert to x0 + x0_uncond = convert_to_x0(latents_step, noise_pred_uncond, sigma_t) + x0_text = convert_to_x0(latents_step, noise_pred_text, sigma_t) + x0_perturb = convert_to_x0(latents_step, noise_pred_perturb, sigma_t) + x0_isolated = convert_to_x0(latents_step, noise_pred_isolated, sigma_t) + + # Delta formulation + cfg_delta = (guidance_scale - 1) * (x0_text - x0_uncond) + stg_delta = stg_scale * (x0_text - x0_perturb) + video_modality_delta = (modality_scale - 1) * (x0_text - x0_isolated) + + x0_combined = x0_text + cfg_delta + stg_delta + video_modality_delta + + x0_combined = jax.lax.cond( + guidance_rescale > 0, + lambda: rescale_noise_cfg(x0_combined, x0_text, guidance_rescale=guidance_rescale), + lambda: x0_combined, + ) - noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t) + noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t) - # Audio guidance - noise_pred_audio_uncond, noise_pred_audio_text, noise_pred_audio_perturb, noise_pred_audio_isolated = jnp.split( - noise_pred_audio, 4, axis=0 - ) + # Audio guidance + ( + noise_pred_audio_uncond, + noise_pred_audio_text, + noise_pred_audio_perturb, + noise_pred_audio_isolated, + ) = jnp.split(noise_pred_audio, 4, axis=0) - x0_audio_uncond = convert_to_x0(audio_latents_step, noise_pred_audio_uncond, sigma_t) - x0_audio_text = convert_to_x0(audio_latents_step, noise_pred_audio_text, sigma_t) - x0_audio_perturb = convert_to_x0(audio_latents_step, noise_pred_audio_perturb, sigma_t) - x0_audio_isolated = convert_to_x0(audio_latents_step, noise_pred_audio_isolated, sigma_t) + x0_audio_uncond = convert_to_x0(audio_latents_step, noise_pred_audio_uncond, sigma_t) + x0_audio_text = convert_to_x0(audio_latents_step, noise_pred_audio_text, sigma_t) + x0_audio_perturb = convert_to_x0(audio_latents_step, noise_pred_audio_perturb, sigma_t) + x0_audio_isolated = convert_to_x0(audio_latents_step, noise_pred_audio_isolated, sigma_t) - cfg_audio_delta = (audio_guidance_scale - 1 if audio_guidance_scale is not None else guidance_scale - 1) * ( - x0_audio_text - x0_audio_uncond - ) - stg_audio_delta = (audio_stg_scale if audio_stg_scale is not None else stg_scale) * ( - x0_audio_text - x0_audio_perturb - ) - audio_modality_delta = (audio_modality_scale - 1 if audio_modality_scale is not None else modality_scale - 1) * ( - x0_audio_text - x0_audio_isolated - ) + cfg_audio_delta = (audio_guidance_scale - 1 if audio_guidance_scale is not None else guidance_scale - 1) * ( + x0_audio_text - x0_audio_uncond + ) + stg_audio_delta = (audio_stg_scale if audio_stg_scale is not None else stg_scale) * (x0_audio_text - x0_audio_perturb) + audio_modality_delta = (audio_modality_scale - 1 if audio_modality_scale is not None else modality_scale - 1) * ( + x0_audio_text - x0_audio_isolated + ) - x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta + x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta - if audio_guidance_rescale > 0: - x0_audio_combined = rescale_noise_cfg(x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale) + x0_audio_combined = jax.lax.cond( + audio_guidance_rescale > 0, + lambda: rescale_noise_cfg( + x0_audio_combined, + x0_audio_text, + guidance_rescale=audio_guidance_rescale, + ), + lambda: x0_audio_combined, + ) - noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) + noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) - # ... (Standard CFG paths can be added here, but for brevity and since LTX2.3 runs with STG this handles the core logic) - elif do_cfg: - noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + # ... (Standard CFG paths can be added here, but for brevity and since LTX2.3 runs with STG this handles the core logic) + elif do_cfg: + noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0) - noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond) + noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0) + noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond) - # Step scheduler - latents_step, _ = scheduler_step(s_state, noise_pred, t, latents_step, return_dict=False) - latents_step = latents_step.astype(latents.dtype) + # Step scheduler + latents_step, _ = scheduler_step(s_state, noise_pred, t, latents_step, return_dict=False) + latents_step = latents_step.astype(latents.dtype) - audio_latents_step, _ = scheduler_step(s_state, noise_pred_audio, t, audio_latents_step, return_dict=False) - audio_latents_step = audio_latents_step.astype(audio_latents.dtype) + audio_latents_step, _ = scheduler_step(s_state, noise_pred_audio, t, audio_latents_step, return_dict=False) + audio_latents_step = audio_latents_step.astype(audio_latents.dtype) - if do_cfg and do_stg: - latents_next = jnp.concatenate([latents_step] * 4, axis=0) - audio_latents_next = jnp.concatenate([audio_latents_step] * 4, axis=0) - elif do_cfg: - latents_next = jnp.concatenate([latents_step] * 2, axis=0) - audio_latents_next = jnp.concatenate([audio_latents_step] * 2, axis=0) - else: - latents_next = latents_step - audio_latents_next = audio_latents_step + if do_cfg and do_stg: + latents_next = jnp.concatenate([latents_step] * 4, axis=0) + audio_latents_next = jnp.concatenate([audio_latents_step] * 4, axis=0) + elif do_cfg: + latents_next = jnp.concatenate([latents_step] * 2, axis=0) + audio_latents_next = jnp.concatenate([audio_latents_step] * 2, axis=0) + else: + latents_next = latents_step + audio_latents_next = audio_latents_step - new_carry = (latents_next, audio_latents_next, s_state) - return new_carry, None + new_carry = (latents_next, audio_latents_next, s_state) + return new_carry, None initial_carry = (latents_jax, audio_latents_jax, scheduler_state) scan_inputs = (timesteps_jax, sigmas) - final_carry, _ = nnx.scan( - scan_body, - in_axes=(nnx.Carry, 0), - out_axes=(nnx.Carry, 0), - )(initial_carry, scan_inputs) + final_carry, _ = jax.lax.scan(scan_body, initial_carry, scan_inputs) return final_carry[0], final_carry[1] diff --git a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py index 831699474..4e4d86634 100644 --- a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py @@ -188,13 +188,16 @@ def test_encode_prompt_no_cfg(self, list_embed_mock): self.assertIsNone(n_e) self.assertIsNone(n_a) + @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.Mesh") + @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.max_utils.create_device_mesh") @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.LTX2Pipeline.load_transformer") @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.LTX2Pipeline._create_common_components") @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.LTX2Pipeline.quantize_transformer") - def test_load_and_init(self, mock_quantize, mock_create_common, mock_load_transformer): + def test_load_and_init(self, mock_quantize, mock_create_common, mock_load_transformer, mock_create_device_mesh, mock_Mesh): """Test that pipeline loading correctly wires all the dependencies down to __init__.""" mock_config = MagicMock() mock_mesh = MagicMock() + mock_Mesh.return_value = mock_mesh mock_common = { "vae": MagicMock(), diff --git a/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py b/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py index f8c5b8d62..281b82784 100644 --- a/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py @@ -25,6 +25,7 @@ from maxdiffusion import pyconfig from maxdiffusion.max_utils import create_device_mesh from maxdiffusion.models.ltx2.transformer_ltx2 import ( + LTX2BlockContext, LTX2VideoTransformerBlock, LTX2VideoTransformer3DModel, LTX2AdaLayerNormSingle, @@ -203,7 +204,7 @@ def test_ltx2_transformer_block(self): temb_ca_gate = jnp.zeros((batch_size, 1 * dim)) temb_ca_audio_gate = jnp.zeros((batch_size, 1 * audio_dim)) - output_hidden, output_audio = block( + ctx = LTX2BlockContext( hidden_states=hidden_states, audio_hidden_states=audio_hidden_states, encoder_hidden_states=encoder_hidden_states, @@ -216,6 +217,8 @@ def test_ltx2_transformer_block(self): temb_ca_audio_gate=temb_ca_audio_gate, ) + output_hidden, output_audio = block(ctx) + self.assertEqual(output_hidden.shape, hidden_states.shape) self.assertEqual(output_audio.shape, audio_hidden_states.shape) diff --git a/src/maxdiffusion/utils/ltx2_block_benchmark.py b/src/maxdiffusion/utils/ltx2_block_benchmark.py new file mode 100644 index 000000000..3aefd7ee4 --- /dev/null +++ b/src/maxdiffusion/utils/ltx2_block_benchmark.py @@ -0,0 +1,323 @@ +""" +LTX2 implementation of the model-agnostic `BlockBenchmark` (see tile_size_grid_search.py). +""" +import os +import functools + +os.environ.setdefault( + "LIBTPU_INIT_ARGS", + " ".join([ + "--xla_tpu_dvfs_p_state=7", + "--xla_tpu_spmd_rng_bit_generator_unsafe=true", + "--xla_tpu_enable_dot_strength_reduction=true", + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true", + "--xla_enable_async_collective_permute=true", + "--xla_tpu_enable_async_collective_fusion=true", + "--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_enable_async_all_to_all=true", + "--xla_tpu_enable_all_experimental_scheduler_features=true", + "--xla_tpu_enable_latency_hiding_scheduler=true", + "--xla_tpu_enable_megacore_fusion=true", + ]), +) +os.environ.setdefault("JAX_DEFAULT_MATMUL_PRECISION", "bfloat16") +os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.95") + +import jax +import jax.numpy as jnp +from flax import linen as nn +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxdiffusion import max_logging, max_utils, pyconfig +from maxdiffusion.models.ltx2.transformer_ltx2 import LTX2VideoTransformer3DModel +from maxdiffusion.utils.tile_size_grid_search import ( + BenchResult, + BlockBenchmark, + grid_search, + time_callable, +) + +_IN_CHANNELS = 128 # LTX2 uses 128 channels in latent space +_VAE_T, _VAE_S = ( + 8, + 32, +) # LTX2 spatial/temporal compression: patch_size_t=1, patch_size=1 (it actually operates on 32x32 compressed) Wait, LTX2 spatial downsampling is 32x32 in pixel space, so in latent space it's 1x1 patch? No, let's just use the latent shape. + + +def latent_seq_len(num_frames: int, height: int, width: int) -> int: + lf = (num_frames - 1) // 8 + 1 + lh, lw = height // 32, width // 32 + return lf * lh * lw + + +def tiled_seq_len(full_seq: int, attention: str, context_shards: int, ulysses_shards: int) -> int: + _RING_VARIANTS = {"ulysses_ring_custom", "ulysses_ring_custom_bidir", "tokamax_ring", "ring"} + if attention in _RING_VARIANTS and ulysses_shards >= 1 and context_shards >= 1: + ring_shards = max(1, context_shards // max(1, ulysses_shards)) + return full_seq // ring_shards + return full_seq + + +@functools.partial(nnx.jit, static_argnames=("num_frames", "height", "width")) +def _forward( + model, + latents, + timestep, + prompt_embeds, + prompt_attention_mask, + audio_latents, + audio_prompt_embeds, + audio_prompt_attention_mask, + *, + num_frames, + height, + width, +): + return model( + hidden_states=latents, + audio_hidden_states=audio_latents, + encoder_hidden_states=prompt_embeds, + audio_encoder_hidden_states=audio_prompt_embeds, + timestep=timestep, + audio_timestep=None, + sigma=None, + audio_sigma=None, + encoder_attention_mask=prompt_attention_mask, + audio_encoder_attention_mask=audio_prompt_attention_mask, + num_frames=num_frames, + height=height, + width=width, + fps=24.0, + audio_num_frames=128, + video_coords=None, + audio_coords=None, + attention_kwargs=None, + use_cross_timestep=False, + ) + + +class LTX2BlockBenchmark(BlockBenchmark): + + def __init__( + self, + config, + mesh, + *, + num_frames, + height, + width, + batch=None, + vmem_limit_bytes=64 * 1024 * 1024, + ): + self._config = config + self._mesh = mesh + self._rules = config.logical_axis_rules + self._attention = getattr(config, "attention", "flash") + self._context_shards = int(mesh.shape.get("context", 1)) + 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}" + + self._lf = (num_frames - 1) // 8 + 1 + self._lh, self._lw = height // 32, width // 32 + self._full_seq = latent_seq_len(num_frames, height, width) + self._num_frames_orig = num_frames + self._height_orig = height + 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) + self._hf_cfg = LTX2VideoTransformer3DModel.load_config(config.pretrained_model_name_or_path, subfolder="transformer") + self._inputs = self._make_inputs() + + @classmethod + def from_config(cls, config, mesh, **overrides): + return cls( + config, + mesh, + num_frames=config.num_frames, + height=config.height, + width=config.width, + **overrides, + ) + + def tiled_seq_lens(self): + s = tiled_seq_len(self._full_seq, self._attention, self._context_shards, self._ulysses_shards) + return (s, s) + + def vmem_bytes(self): + return self._vmem + + def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2): + cmp = bkv_compute or bkv + try: + with self._mesh: + model = self._build_model(bq, bkv, cmp) + ( + latents, + timestep, + prompt_embeds, + prompt_attention_mask, + audio_latents, + audio_prompt_embeds, + audio_prompt_attention_mask, + ) = self._inputs + with self._mesh, nn_partitioning.axis_rules(self._rules): + mean, std, times, compile_ms = time_callable( + lambda: _forward( + model, + latents, + timestep, + prompt_embeds, + prompt_attention_mask, + audio_latents, + audio_prompt_embeds, + audio_prompt_attention_mask, + num_frames=self._lf, + height=self._lh, + width=self._lw, + ), + iters=iters, + warmup=warmup, + sync=jax.block_until_ready, + ) + return BenchResult( + bq, + bkv, + cmp, + "ok", + mean_ms=mean, + std_ms=std, + times_ms=times, + compile_ms=compile_ms, + ) + except Exception as e: + import traceback + + traceback.print_exc() + msg = str(e) + oom = any(t in msg for t in ("RESOURCE_EXHAUSTED", "out of memory", "Mosaic", "VMEM")) + return BenchResult(bq, bkv, cmp, "oom" if oom else "error", detail=msg[:200]) + + def _flash_block_sizes(self, bq, bkv, cmp): + from maxdiffusion.max_utils import CustomFlashBlockSizes + + return CustomFlashBlockSizes( + block_q=bq, + block_kv=bkv, + block_kv_compute=cmp, + block_kv_compute_in=cmp, + heads_per_tile=1, + vmem_limit_bytes=self._vmem, + ) + + def _build_model(self, bq, bkv, cmp): + c = self._config + ltx2_config = dict(self._hf_cfg) + if ltx2_config.get("activation_fn") == "gelu-approximate": + ltx2_config["activation_fn"] = "gelu" + ltx2_config.update( + mesh=self._mesh, + dtype=c.activations_dtype, + weights_dtype=c.weights_dtype, + attention_kernel=c.attention, + a2v_attention_kernel=getattr(c, "a2v_attention_kernel", "flash"), + v2a_attention_kernel=getattr(c, "v2a_attention_kernel", "dot_product"), + precision=max_utils.get_precision(c), + flash_block_sizes=self._flash_block_sizes(bq, bkv, cmp), + flash_min_seq_length=getattr(c, "flash_min_seq_length", 4096), + ulysses_shards=getattr(c, "ulysses_shards", -1), + ulysses_attention_chunks=getattr(c, "ulysses_attention_chunks", 1), + remat_policy=getattr(c, "remat_policy", "NONE"), + scan_layers=False, + num_layers=1, + ) + model = LTX2VideoTransformer3DModel(**ltx2_config, rngs=nnx.Rngs(params=0)) + gd, state, rest = nnx.split(model, nnx.Param, ...) + shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), self._mesh, self._rules) + return nnx.merge(gd, jax.device_put(state, shardings), rest) + + def _make_inputs(self): + dtype = self._config.activations_dtype + k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4) + seq_len = self._lf * self._lh * self._lw + latents = jax.random.normal(k1, (self._batch, seq_len, _IN_CHANNELS), dtype) + + # Audio latents + audio_seq_len = 128 + audio_latents = jax.random.normal(k3, (self._batch, audio_seq_len, _IN_CHANNELS), dtype) + + # Prompts + prompt_embeds = jax.random.normal(k2, (self._batch, 1024, 3840), dtype) + prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.int32) + + audio_prompt_embeds = jax.random.normal(k4, (self._batch, 1024, 3840), dtype) + audio_prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.int32) + + timestep = jnp.zeros((self._batch,), jnp.float32) + repl = jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()) + return ( + jax.device_put(latents, repl), + jax.device_put(timestep, repl), + jax.device_put(prompt_embeds, repl), + jax.device_put(prompt_attention_mask, repl), + jax.device_put(audio_latents, repl), + jax.device_put(audio_prompt_embeds, repl), + jax.device_put(audio_prompt_attention_mask, repl), + ) + + +def _build_cli_config(argv_yaml, attention, ulysses_shards, num_frames, height, width): + pyconfig.initialize([ + "ltx2_block_benchmark", + argv_yaml, + f"attention={attention}", + f"ulysses_shards={ulysses_shards}", + "skip_jax_distributed_system=True", + "weights_dtype=bfloat16", + "activations_dtype=bfloat16", + "per_device_batch_size=1", + f"num_frames={num_frames}", + f"height={height}", + f"width={width}", + "ici_data_parallelism=1", + "ici_fsdp_parallelism=1", + f"ici_context_parallelism={jax.device_count()}", + "ici_tensor_parallelism=1", + "allow_split_physical_axes=True", + "use_base2_exp=true", + "use_experimental_scheduler=true", + "enable_jax_named_scopes=false", + ]) + return pyconfig.config + + +def main(): + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("config", help="Path to yaml config (e.g. configs/ltx2_video.yml)") + parser.add_argument("--attention", default="ulysses_ring_custom") + parser.add_argument("--ulysses-shards", type=int, default=1) + parser.add_argument("--num-frames", type=int, default=161) + parser.add_argument("--height", type=int, default=512) + parser.add_argument("--width", type=int, default=768) + parser.add_argument("--smart-search", action="store_true") + parser.add_argument("--full-search", action="store_true") + parser.add_argument("--out-dir", default="") + args = parser.parse_args() + + config = _build_cli_config(args.config, args.attention, args.ulysses_shards, args.num_frames, args.height, args.width) + mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes) + bench = LTX2BlockBenchmark.from_config(config, mesh) + + mode = "full" if args.full_search else "smart" + grid_search(bench, mode=mode, out_dir=args.out_dir or None, log=max_logging.log) + + +if __name__ == "__main__": + main() From 65c601d8e4879886bd0006c1e81244b3e340b672 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sat, 25 Jul 2026 12:22:45 +0530 Subject: [PATCH 02/10] fix: use full step count for warmup when scan_diffusion_loop=True jax.lax.scan bakes iteration count into array shapes, so a 2-step warmup compiles a different XLA program than the N-step real run, causing a wasteful double compilation. Only reduce warmup steps when scan_diffusion_loop=False (Python loop), where per-step JIT compilation is independent of total iteration count. --- src/maxdiffusion/generate_ltx2.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index 6c355bd2a..1e727c61d 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -252,7 +252,18 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # --------------------------------------------------------- config.get_keys()["enable_profiler"] = False config.get_keys()["enable_ml_diagnostics"] = False - warmup_steps = min(2, original_num_steps) + + # When scan_diffusion_loop=True the entire denoising loop is compiled as a + # single jax.lax.scan whose iteration count is baked into the XLA program via + # array shapes. A 2-step warmup would compile a *different* program than the + # real N-step run, forcing a wasteful second compilation. Only reduce warmup + # steps when scan_diffusion_loop=False (Python loop), where compilation + # happens per-step and is independent of the total iteration count. + scan_diffusion_loop = getattr(config, "scan_diffusion_loop", True) + if scan_diffusion_loop: + warmup_steps = original_num_steps + else: + warmup_steps = min(2, original_num_steps) config.get_keys()["num_inference_steps"] = warmup_steps max_logging.log(f"🚀 Starting warmup compilation pass ({warmup_steps} steps)...") From 0f8e422f783827d03d7b3b322a14b47342d719e6 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sat, 25 Jul 2026 17:43:57 +0530 Subject: [PATCH 03/10] fix: reject incompatible LTX2 KV cache --- .../models/ltx2/transformer_ltx2.py | 70 ++++++------ .../pipelines/ltx2/ltx2_pipeline.py | 2 +- .../tests/ltx2/test_transformer_ltx2.py | 104 ++++++++++++++++++ 3 files changed, 142 insertions(+), 34 deletions(-) diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py index 3ff086915..976abcb75 100644 --- a/src/maxdiffusion/models/ltx2/transformer_ltx2.py +++ b/src/maxdiffusion/models/ltx2/transformer_ltx2.py @@ -532,7 +532,7 @@ def __call__( scale_mlp = ada_values[:, :, 4, :] gate_mlp = ada_values[:, :, 5, :] - if getattr(self, "cross_attn_mod", False): + if self.cross_attn_mod: shift_q = ada_values[:, :, 6, :] scale_q = ada_values[:, :, 7, :] gate_q = ada_values[:, :, 8, :] @@ -563,7 +563,7 @@ def __call__( audio_scale_mlp = audio_ada_values[:, :, 4, :] audio_gate_mlp = audio_ada_values[:, :, 5, :] - if getattr(self, "cross_attn_mod", False): + if self.cross_attn_mod: audio_shift_q = audio_ada_values[:, :, 6, :] audio_scale_q = audio_ada_values[:, :, 7, :] audio_gate_q = audio_ada_values[:, :, 8, :] @@ -581,10 +581,10 @@ def __call__( # 2. Video and Audio Cross-Attention with the text embeddings norm_hidden_states = self.norm2(hidden_states) - if getattr(self, "cross_attn_mod", False): + if self.cross_attn_mod: norm_hidden_states = norm_hidden_states * (1 + scale_q) + shift_q - if getattr(self, "cross_attn_mod", False) and temb_prompt is not None: + if self.cross_attn_mod and temb_prompt is not None: prompt_table_reshaped = jnp.expand_dims(self.prompt_scale_shift_table, axis=(0, 1)) temb_prompt_reshaped = temb_prompt.reshape(batch_size, 1, 2, -1) prompt_ada_values = prompt_table_reshaped + temb_prompt_reshaped @@ -600,15 +600,15 @@ def __call__( attention_mask=encoder_attention_mask, cached_kv=attn2_kv, ) - if getattr(self, "cross_attn_mod", False): + if self.cross_attn_mod: attn_hidden_states = attn_hidden_states * gate_q hidden_states = hidden_states + attn_hidden_states norm_audio_hidden_states = self.audio_norm2(audio_hidden_states) - if getattr(self, "cross_attn_mod", False): + if self.cross_attn_mod: norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_q) + audio_shift_q - if getattr(self, "cross_attn_mod", False) and temb_prompt_audio is not None: + if self.cross_attn_mod and temb_prompt_audio is not None: audio_prompt_table_reshaped = jnp.expand_dims(self.audio_prompt_scale_shift_table, axis=(0, 1)) temb_prompt_audio_reshaped = temb_prompt_audio.reshape(batch_size, 1, 2, -1) audio_prompt_ada_values = audio_prompt_table_reshaped + temb_prompt_audio_reshaped @@ -624,7 +624,7 @@ def __call__( attention_mask=audio_encoder_attention_mask, cached_kv=audio_attn2_kv, ) - if getattr(self, "cross_attn_mod", False): + if self.cross_attn_mod: attn_audio_hidden_states = attn_audio_hidden_states * audio_gate_q audio_hidden_states = audio_hidden_states + attn_audio_hidden_states @@ -1180,6 +1180,12 @@ def compute_kv_cache( fps: float, audio_num_frames: int, ): + if self.cross_attn_mod: + raise ValueError( + "KV caching is incompatible with cross_attn_mod=True because prompt embeddings are modulated by the " + "timestep. Disable use_kv_cache for this model." + ) + batch_size = encoder_hidden_states.shape[0] if self.use_prompt_embeddings and self.caption_projection is not None: encoder_hidden_states = self.caption_projection(encoder_hidden_states) @@ -1188,33 +1194,31 @@ def compute_kv_cache( encoder_hidden_states = encoder_hidden_states.reshape(batch_size, -1, self.inner_dim) audio_encoder_hidden_states = audio_encoder_hidden_states.reshape(batch_size, -1, self.audio_inner_dim) - kv_cache = None - if not self.cross_attn_mod: - if self.scan_layers: + if self.scan_layers: - @nnx.vmap( - in_axes=(0, None, None), - out_axes=0, - transform_metadata={nnx.PARTITION_NAME: "layers"}, - ) - def _compute_kv(block, enc_states, audio_enc_states): - return block.compute_kv(enc_states, audio_enc_states) + @nnx.vmap( + in_axes=(0, None, None), + out_axes=0, + transform_metadata={nnx.PARTITION_NAME: "layers"}, + ) + def _compute_kv(block, enc_states, audio_enc_states): + return block.compute_kv(enc_states, audio_enc_states) - kv_cache = _compute_kv( - self.transformer_blocks, - encoder_hidden_states, - audio_encoder_hidden_states, - ) - else: - kv_cache_list = [] - for block in self.transformer_blocks: - kv_cache_list.append(block.compute_kv(encoder_hidden_states, audio_encoder_hidden_states)) - keys = kv_cache_list[0].keys() - kv_cache = {} - for k in keys: - k_list = [d[k][0] for d in kv_cache_list] - v_list = [d[k][1] for d in kv_cache_list] - kv_cache[k] = (jnp.stack(k_list, axis=0), jnp.stack(v_list, axis=0)) + kv_cache = _compute_kv( + self.transformer_blocks, + encoder_hidden_states, + audio_encoder_hidden_states, + ) + else: + kv_cache_list = [] + for block in self.transformer_blocks: + kv_cache_list.append(block.compute_kv(encoder_hidden_states, audio_encoder_hidden_states)) + keys = kv_cache_list[0].keys() + kv_cache = {} + for k in keys: + k_list = [d[k][0] for d in kv_cache_list] + v_list = [d[k][1] for d in kv_cache_list] + kv_cache[k] = (jnp.stack(k_list, axis=0), jnp.stack(v_list, axis=0)) # RoPE pre-computation video_coords = self.rope.prepare_video_coords(batch_size, num_frames, height, width, fps=fps) diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index 4eb8ce2ef..d3c92bef1 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -2485,7 +2485,7 @@ def scan_body(carry, inputs): noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) - # ... (Standard CFG paths can be added here, but for brevity and since LTX2.3 runs with STG this handles the core logic) + # Standard CFG path when spatio-temporal guidance is disabled. elif do_cfg: noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) diff --git a/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py b/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py index 281b82784..3e3e17051 100644 --- a/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py @@ -138,6 +138,110 @@ def test_ltx2_rope_split(self): self.assertEqual(cos.shape, (1, 32, 10, 16)) self.assertEqual(sin.shape, (1, 32, 10, 16)) + def test_video_coords_use_channel_first_layout(self): + """Video coordinate axes remain [batch, axis, sequence, start_or_end].""" + rope = LTX2RotaryPosEmbed( + dim=self.dim, + patch_size=self.patch_size, + patch_size_t=self.patch_size_t, + base_num_frames=8, + base_height=32, + base_width=32, + modality="video", + ) + + coords = rope.prepare_video_coords(batch_size=2, num_frames=2, height=3, width=4, fps=24.0) + + self.assertEqual(coords.shape, (2, 3, 24, 2)) + + def test_kv_cache_rejects_timestep_modulated_prompt_embeddings(self): + """Caching must not bypass the per-timestep LTX2.3 prompt modulation.""" + model = Mock(cross_attn_mod=True) + + with self.assertRaisesRegex(ValueError, "KV caching is incompatible with cross_attn_mod=True"): + LTX2VideoTransformer3DModel.compute_kv_cache( + model, + encoder_hidden_states=None, + audio_encoder_hidden_states=None, + num_frames=1, + height=1, + width=1, + fps=24.0, + audio_num_frames=1, + ) + + def test_kv_cache_matches_uncached_output_with_and_without_layer_scan(self): + """Layer-scanned KV inputs must be equivalent to normal cross-attention projections.""" + batch_size = 1 + num_frames, height, width = 2, 2, 2 + audio_num_frames = 4 + video_dim, audio_dim, caption_dim = 64, 64, 32 + hidden_states = jnp.ones((batch_size, num_frames * height * width, 8), dtype=jnp.float32) + audio_hidden_states = jnp.ones((batch_size, audio_num_frames, 4), dtype=jnp.float32) + encoder_hidden_states = jnp.ones((batch_size, 3, caption_dim), dtype=jnp.float32) + audio_encoder_hidden_states = jnp.ones((batch_size, 3, caption_dim), dtype=jnp.float32) + attention_mask = jnp.ones((batch_size, 3), dtype=jnp.float32) + + for scan_layers in (False, True): + with self.subTest(scan_layers=scan_layers), self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + model = LTX2VideoTransformer3DModel( + rngs=nnx.Rngs(jax.random.key(0)), + in_channels=8, + out_channels=8, + num_attention_heads=4, + attention_head_dim=16, + cross_attention_dim=video_dim, + caption_channels=caption_dim, + audio_in_channels=4, + audio_out_channels=4, + audio_num_attention_heads=4, + audio_attention_head_dim=16, + audio_cross_attention_dim=audio_dim, + num_layers=2, + mesh=self.mesh, + scan_layers=scan_layers, + attention_kernel="dot_product", + a2v_attention_kernel="dot_product", + v2a_attention_kernel="dot_product", + ) + + common_args = dict( + 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=jnp.array([1.0]), + num_frames=num_frames, + height=height, + width=width, + audio_num_frames=audio_num_frames, + encoder_attention_mask=attention_mask, + audio_encoder_attention_mask=attention_mask, + return_dict=True, + ) + uncached_output = model(**common_args) + kv_cache, rope_cache, cached_video_embeds, cached_audio_embeds = model.compute_kv_cache( + encoder_hidden_states, + audio_encoder_hidden_states, + num_frames, + height, + width, + 24.0, + audio_num_frames, + ) + cached_output = model( + **{ + **common_args, + "encoder_hidden_states": cached_video_embeds, + "audio_encoder_hidden_states": cached_audio_embeds, + "cached_kv": kv_cache, + "rope_cache": rope_cache, + } + ) + + self.assertTrue(bool(jnp.allclose(uncached_output["sample"], cached_output["sample"]))) + self.assertTrue(bool(jnp.allclose(uncached_output["audio_sample"], cached_output["audio_sample"]))) + def test_ltx2_ada_layer_norm_single(self): """Tests LTX2AdaLayerNormSingle initialization and execution.""" key = jax.random.key(0) From 55ee9a8518e975f01fef21f30e19c831246c1959 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sat, 25 Jul 2026 18:09:56 +0530 Subject: [PATCH 04/10] fix: complete LTX2 KV cache and named scopes --- src/maxdiffusion/configs/ltx2_3_video.yml | 1 - src/maxdiffusion/configs/ltx2_video.yml | 1 - src/maxdiffusion/configs/ltx_video.yml | 2 - .../models/ltx2/attention_ltx2.py | 14 +++-- .../models/ltx2/transformer_ltx2.py | 39 +++++++++---- .../pipelines/ltx2/ltx2_pipeline.py | 56 +++++++++++++++++++ 6 files changed, 95 insertions(+), 18 deletions(-) diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml index c4f536bff..f9271ef12 100644 --- a/src/maxdiffusion/configs/ltx2_3_video.yml +++ b/src/maxdiffusion/configs/ltx2_3_video.yml @@ -197,4 +197,3 @@ tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' ( tile_search_iters: 10 tile_search_out: '' # dir for the results CSV; '' -> print only use_kv_cache: False -cross_attn_mod: False diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml index 856b4626e..6be535010 100644 --- a/src/maxdiffusion/configs/ltx2_video.yml +++ b/src/maxdiffusion/configs/ltx2_video.yml @@ -194,4 +194,3 @@ tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' ( tile_search_iters: 10 tile_search_out: '' # dir for the results CSV; '' -> print only use_kv_cache: False -cross_attn_mod: False diff --git a/src/maxdiffusion/configs/ltx_video.yml b/src/maxdiffusion/configs/ltx_video.yml index 6dc09b5d4..d70154e0e 100644 --- a/src/maxdiffusion/configs/ltx_video.yml +++ b/src/maxdiffusion/configs/ltx_video.yml @@ -109,5 +109,3 @@ enable_single_replica_ckpt_restoring: False enable_ml_diagnostics: False profiler_gcs_path: "" enable_ondemand_xprof: False -use_kv_cache: False -cross_attn_mod: False diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py index 730f81771..601999cf4 100644 --- a/src/maxdiffusion/models/ltx2/attention_ltx2.py +++ b/src/maxdiffusion/models/ltx2/attention_ltx2.py @@ -14,6 +14,7 @@ limitations under the License. """ +import contextlib from typing import Optional, Tuple from flax import nnx import jax @@ -346,12 +347,14 @@ def __init__( gated_attn: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + enable_jax_named_scopes: bool = False, ): self.heads = heads self.rope_type = rope_type self.dim_head = dim_head self.inner_dim = dim_head * heads self.dropout_rate = dropout + self.enable_jax_named_scopes = enable_jax_named_scopes if sharding_specs is None: specs = get_sharding_specs("default", "ltx2_dit") @@ -509,6 +512,9 @@ def compute_kv(self, context: Array) -> Tuple[Array, Array]: key = self.norm_k(key) return key, value + def named_scope(self, name: str): + return jax.named_scope(name) if self.enable_jax_named_scopes else contextlib.nullcontext() + def __call__( self, hidden_states: Array, @@ -523,7 +529,7 @@ def __call__( context = encoder_hidden_states if encoder_hidden_states is not None else hidden_states # 1. Project and Norm - with jax.named_scope("QKV Projection"): + with self.named_scope("QKV Projection"): query = self.to_q(hidden_states) if cached_kv is not None: key, value = cached_kv @@ -531,7 +537,7 @@ def __call__( key = self.to_k(context) value = self.to_v(context) - with jax.named_scope("QKV Norm"): + with self.named_scope("QKV Norm"): query = self.norm_q(query) if cached_kv is None: key = self.norm_k(key) @@ -539,7 +545,7 @@ def __call__( # 3. Apply RoPE to tensors of shape [B, S, InnerDim] # Frequencies are shape [B, S, InnerDim] # 3. Apply RoPE - with jax.named_scope("Apply RoPE"): + with self.named_scope("Apply RoPE"): if rotary_emb is not None: if self.rope_type == "split": # Split RoPE: passing full freqs [B, H, S, D//2] @@ -560,7 +566,7 @@ def __call__( elif encoder_hidden_states is None: key = apply_rotary_emb(key, rotary_emb) - with jax.named_scope("Attention and Output Project"): + 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) diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py index 976abcb75..374038bb6 100644 --- a/src/maxdiffusion/models/ltx2/transformer_ltx2.py +++ b/src/maxdiffusion/models/ltx2/transformer_ltx2.py @@ -14,6 +14,7 @@ limitations under the License. """ +import contextlib from typing import Optional, Tuple, Any, Dict import jax import jax.numpy as jnp @@ -163,12 +164,14 @@ def __init__( perturbed_attn: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + enable_jax_named_scopes: bool = False, ): self.dim = dim self.norm_eps = norm_eps self.norm_elementwise_affine = norm_elementwise_affine self.attention_kernel = attention_kernel self.perturbed_attn = perturbed_attn + self.enable_jax_named_scopes = enable_jax_named_scopes if sharding_specs is None: sharding_specs = get_sharding_specs("default", "ltx2_dit") @@ -186,6 +189,7 @@ def __init__( ) self.attn1 = LTX2Attention( rngs=rngs, + enable_jax_named_scopes=enable_jax_named_scopes, query_dim=self.dim, heads=num_attention_heads, dim_head=attention_head_dim, @@ -216,6 +220,7 @@ def __init__( ) self.audio_attn1 = LTX2Attention( rngs=rngs, + enable_jax_named_scopes=enable_jax_named_scopes, query_dim=audio_dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, @@ -247,6 +252,7 @@ def __init__( ) self.attn2 = LTX2Attention( rngs=rngs, + enable_jax_named_scopes=enable_jax_named_scopes, query_dim=dim, context_dim=cross_attention_dim, heads=num_attention_heads, @@ -277,6 +283,7 @@ def __init__( ) self.audio_attn2 = LTX2Attention( rngs=rngs, + enable_jax_named_scopes=enable_jax_named_scopes, query_dim=audio_dim, context_dim=audio_cross_attention_dim, heads=audio_num_attention_heads, @@ -309,6 +316,7 @@ def __init__( ) self.audio_to_video_attn = LTX2Attention( rngs=rngs, + enable_jax_named_scopes=enable_jax_named_scopes, query_dim=dim, context_dim=audio_dim, heads=audio_num_attention_heads, @@ -340,6 +348,7 @@ def __init__( ) self.video_to_audio_attn = LTX2Attention( rngs=rngs, + enable_jax_named_scopes=enable_jax_named_scopes, query_dim=audio_dim, context_dim=dim, heads=audio_num_attention_heads, @@ -455,6 +464,9 @@ def compute_kv(self, encoder_hidden_states: jax.Array, audio_encoder_hidden_stat audio_text_k, audio_text_v = self.audio_attn2.compute_kv(audio_encoder_hidden_states) return {"attn2": (text_k, text_v), "audio_attn2": (audio_text_k, audio_text_v)} + def named_scope(self, name: str): + return jax.named_scope(name) if self.enable_jax_named_scopes else contextlib.nullcontext() + def __call__( self, ctx: "LTX2BlockContext", @@ -539,7 +551,7 @@ def __call__( norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa - with jax.named_scope("Video Self-Attention"): + with self.named_scope("Video Self-Attention"): attn_hidden_states = self.attn1( hidden_states=norm_hidden_states, encoder_hidden_states=None, @@ -570,7 +582,7 @@ def __call__( norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa - with jax.named_scope("Audio Self-Attention"): + with self.named_scope("Audio Self-Attention"): attn_audio_hidden_states = self.audio_attn1( hidden_states=norm_audio_hidden_states, encoder_hidden_states=None, @@ -671,7 +683,7 @@ def __call__( mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale) + video_a2v_ca_shift mod_norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_a2v_ca_scale) + audio_a2v_ca_shift - with jax.named_scope("Audio-to-Video Cross-Attention"): + with self.named_scope("Audio-to-Video Cross-Attention"): a2v_attn_hidden_states = self.audio_to_video_attn( mod_norm_hidden_states, encoder_hidden_states=mod_norm_audio_hidden_states, @@ -687,7 +699,7 @@ def __call__( mod_norm_hidden_states_v2a = norm_hidden_states * (1 + video_v2a_ca_scale) + video_v2a_ca_shift mod_norm_audio_hidden_states_v2a = norm_audio_hidden_states * (1 + audio_v2a_ca_scale) + audio_v2a_ca_shift - with jax.named_scope("Video-to-Audio Cross-Attention"): + with self.named_scope("Video-to-Audio Cross-Attention"): v2a_attn_hidden_states = self.video_to_audio_attn( mod_norm_audio_hidden_states_v2a, encoder_hidden_states=mod_norm_hidden_states_v2a, @@ -778,6 +790,7 @@ def __init__( spatio_temporal_guidance_blocks: Tuple[int, ...] = (), ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + enable_jax_named_scopes: bool = False, **kwargs, ): self.spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks @@ -830,6 +843,7 @@ def __init__( self.gated_attn = gated_attn self.cross_attn_mod = cross_attn_mod self.perturbed_attn = perturbed_attn + self.enable_jax_named_scopes = enable_jax_named_scopes self.a2v_attention_kernel = a2v_attention_kernel self.v2a_attention_kernel = v2a_attention_kernel self.flash_min_seq_length = flash_min_seq_length @@ -1085,6 +1099,7 @@ def init_block(rngs): perturbed_attn=self.perturbed_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + enable_jax_named_scopes=self.enable_jax_named_scopes, ) if self.scan_layers: @@ -1126,6 +1141,7 @@ def init_block(rngs): perturbed_attn=self.perturbed_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + enable_jax_named_scopes=self.enable_jax_named_scopes, ) blocks.append(block) self.transformer_blocks = nnx.List(blocks) @@ -1170,6 +1186,9 @@ def init_block(rngs): bias_init=nnx.with_partitioning(nnx.initializers.zeros, self.sharding_specs.out_embed_bias), ) + def named_scope(self, name: str): + return jax.named_scope(name) if self.enable_jax_named_scopes else contextlib.nullcontext() + def compute_kv_cache( self, encoder_hidden_states: jax.Array, @@ -1325,7 +1344,7 @@ def __call__( batch_size = hidden_states.shape[0] # 1. Prepare RoPE positional embeddings - with jax.named_scope("RoPE Preparation"): + with self.named_scope("RoPE Preparation"): if rope_cache is not None: video_rotary_emb = rope_cache["video_rotary_emb"] audio_rotary_emb = rope_cache["audio_rotary_emb"] @@ -1344,12 +1363,12 @@ def __call__( audio_cross_attn_rotary_emb = self.cross_attn_audio_rope(audio_coords[:, 0:1, :]) # 2. Patchify input projections - with jax.named_scope("Input Projection"): + with self.named_scope("Input Projection"): hidden_states = self.proj_in(hidden_states) audio_hidden_states = self.audio_proj_in(audio_hidden_states) # 3. Prepare timestep embeddings and modulation parameters - with jax.named_scope("Timestep and Caption Projection"): + with self.named_scope("Timestep and Caption Projection"): timestep_cross_attn_gate_scale_factor = self.cross_attn_timestep_scale_multiplier / self.timestep_scale_multiplier temb, embedded_timestep = self.time_embed( @@ -1424,7 +1443,7 @@ def __call__( encoder_hidden_states = encoder_hidden_states.reshape(batch_size, -1, hidden_states.shape[-1]) audio_encoder_hidden_states = audio_encoder_hidden_states.reshape(batch_size, -1, audio_hidden_states.shape[-1]) # 5. Run transformer blocks - with jax.named_scope("Transformer Blocks"): + with self.named_scope("Transformer Blocks"): base_context = LTX2BlockContext( hidden_states=hidden_states, audio_hidden_states=audio_hidden_states, @@ -1454,7 +1473,7 @@ def apply_block(block, context: LTX2BlockContext, mask, block_cached_kv=None): orig_perturbation_mask = context.perturbation_mask orig_cached_kv = context.cached_kv context = context.replace(perturbation_mask=mask, cached_kv=block_cached_kv) - with jax.named_scope("Transformer Layer"): + with self.named_scope("Transformer Layer"): hidden_states_out, audio_hidden_states_out = block(context) context = context.replace( hidden_states=hidden_states_out.astype(context.hidden_states.dtype), @@ -1577,7 +1596,7 @@ def scan_fn_ltx23(carry, block_and_mask): audio_hidden_states = current_context.audio_hidden_states # 6. Output layers - with jax.named_scope("Output Projection & Norm"): + with self.named_scope("Output Projection & Norm"): scale_shift_values = jnp.expand_dims(self.scale_shift_table, axis=(0, 1)) + jnp.expand_dims(embedded_timestep, axis=2) shift = scale_shift_values[:, :, 0, :] scale = scale_shift_values[:, :, 1, :] diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index d3c92bef1..db45059d7 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -158,6 +158,7 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): ltx2_config["flash_min_seq_length"] = getattr(config, "flash_min_seq_length", 4096) ltx2_config["ulysses_shards"] = getattr(config, "ulysses_shards", -1) ltx2_config["ulysses_attention_chunks"] = getattr(config, "ulysses_attention_chunks", 1) + ltx2_config["enable_jax_named_scopes"] = getattr(config, "enable_jax_named_scopes", False) ltx2_config["remat_policy"] = config.remat_policy ltx2_config["names_which_can_be_saved"] = config.names_which_can_be_saved ltx2_config["names_which_can_be_offloaded"] = config.names_which_can_be_offloaded @@ -1866,6 +1867,25 @@ def __call__( ) else: # Old Python loop path + kv_cache = None + rope_cache = None + if getattr(self.config, "use_kv_cache", False): + ( + kv_cache, + rope_cache, + video_embeds_sharded, + audio_embeds_sharded, + ) = transformer_kv_cache( + graphdef, + state, + video_embeds_sharded, + audio_embeds_sharded, + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + audio_num_frames=audio_num_frames, + fps=frame_rate, + ) for i in range(len(timesteps_jax)): t = timesteps_jax[i] sigma_t = sigmas[i] @@ -1904,6 +1924,8 @@ def __call__( audio_sigma=t, use_cross_timestep=use_cross_timestep, is_cfg_stg_mode=do_cfg and do_stg, + kv_cache=kv_cache, + rope_cache=rope_cache, ) # Extract latents_step based on stacking strategy @@ -2188,6 +2210,40 @@ def convert_to_vel(lat, x0, sig): return LTX2PipelineOutput(frames=video, audio=audio, timings=timings) +@partial( + aot_cache.cached_jit, + static_argnames=( + "latent_num_frames", + "latent_height", + "latent_width", + "audio_num_frames", + "fps", + ), +) +def transformer_kv_cache( + graphdef, + state, + encoder_hidden_states, + audio_encoder_hidden_states, + latent_num_frames, + latent_height, + latent_width, + audio_num_frames, + fps, +): + """Precomputes cross-attention K/V and RoPE caches for a diffusion run.""" + transformer = nnx.merge(graphdef, state) + return transformer.compute_kv_cache( + encoder_hidden_states, + audio_encoder_hidden_states, + latent_num_frames, + latent_height, + latent_width, + fps, + audio_num_frames, + ) + + @partial( aot_cache.cached_jit, static_argnames=( From 55fa53cf5be4659acda0fa0bb6121ab230d7fc3e Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sat, 25 Jul 2026 18:18:53 +0530 Subject: [PATCH 05/10] fix: strengthen LTX2 AOT cache key --- src/maxdiffusion/generate_ltx2.py | 81 ++++++++++++++++++++++++++----- 1 file changed, 70 insertions(+), 11 deletions(-) diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index 1e727c61d..637a6ee4d 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -13,6 +13,7 @@ # limitations under the License. from typing import Sequence +import json import jax import jax.numpy as jnp import time @@ -151,6 +152,73 @@ def maybe_tune_block_sizes(config): ) +def _canonical_aot_value(value): + return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str) + + +def ltx2_aot_metadata(config, pipeline): + """Returns every graph- and topology-shaping input to LTX2 AOT caching. + + We deliberately exclude values such as model weights and RNG state: they are + executable inputs, not compilation inputs. `use_kv_cache` is also excluded + because it is a static argument of `run_diffusion_loop` and therefore part + of that executable's per-call signature. + """ + transformer_config = dict(getattr(pipeline.transformer, "config", {})) + for key in ( + "rngs", + "mesh", + "dtype", + "weights_dtype", + "precision", + "flash_block_sizes", + "flash_min_seq_length", + "scan_layers", + "attention_kernel", + "a2v_attention_kernel", + "v2a_attention_kernel", + "ulysses_shards", + "ulysses_attention_chunks", + "remat_policy", + "names_which_can_be_saved", + "names_which_can_be_offloaded", + "sharding_specs", + "enable_jax_named_scopes", + ): + transformer_config.pop(key, None) + + device = jax.devices()[0] + return { + "model": config.pretrained_model_name_or_path, + "transformer_architecture": _canonical_aot_value(transformer_config), + "attention": getattr(config, "attention", ""), + "a2v_attention_kernel": getattr(config, "a2v_attention_kernel", "flash"), + "v2a_attention_kernel": getattr(config, "v2a_attention_kernel", "dot_product"), + "flash_block_sizes": _canonical_aot_value(getattr(config, "flash_block_sizes", {})), + "flash_min_seq_length": str(getattr(config, "flash_min_seq_length", 4096)), + "ulysses_shards": str(getattr(config, "ulysses_shards", -1)), + "ulysses_attention_chunks": str(getattr(config, "ulysses_attention_chunks", 1)), + "scan_layers": str(getattr(config, "scan_layers", True)), + "scan_diffusion_loop": str(getattr(config, "scan_diffusion_loop", True)), + "remat_policy": str(getattr(config, "remat_policy", "NONE")), + "spatio_temporal_guidance_blocks": _canonical_aot_value( + getattr(config, "spatio_temporal_guidance_blocks", ()) + ), + "enable_jax_named_scopes": str(getattr(config, "enable_jax_named_scopes", False)), + "logical_axis_rules": _canonical_aot_value(getattr(config, "logical_axis_rules", ())), + "sharding": _canonical_aot_value(getattr(config, "sharding", {})), + "mesh_shape": str(pipeline.mesh.shape), + "mesh_axes": _canonical_aot_value(pipeline.mesh.axis_names), + "backend": jax.default_backend(), + "device_kind": getattr(device, "device_kind", ""), + "process_count": str(jax.process_count()), + "weights_dtype": str(getattr(config, "weights_dtype", "bfloat16")), + "activations_dtype": str(getattr(config, "activations_dtype", "bfloat16")), + "jax": jax.__version__, + "jaxlib": getattr(jax.lib, "__version__", ""), + } + + def run(config, pipeline=None, filename_prefix="", commit_hash=None): if pipeline is None: maybe_tune_block_sizes(config) @@ -233,17 +301,8 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # Per-shape AOT executable cache aot_cache.install( getattr(config, "aot_cache_dir", ""), - meta={ - "model": config.pretrained_model_name_or_path, - "attention": getattr(config, "attention", ""), - "flash_block_sizes": str(getattr(config, "flash_block_sizes", "")), - "mesh_shape": str(pipeline.mesh.shape) if pipeline and hasattr(pipeline, "mesh") and pipeline.mesh else "", - "weights_dtype": str(getattr(config, "weights_dtype", "bfloat16")), - "activations_dtype": str(getattr(config, "activations_dtype", "bfloat16")), - "scan_layers": str(getattr(config, "scan_layers", True)), - "jax": jax.__version__, - }, - mesh=pipeline.mesh if pipeline else None, + meta=ltx2_aot_metadata(config, pipeline), + mesh=pipeline.mesh, ) aot_cache.wait_for_loads() From 92dcf7504ea7c6ef408f05cdf0a86f0d6e28d2fc Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sat, 25 Jul 2026 20:05:22 +0530 Subject: [PATCH 06/10] refactor: deduplicate LTX2 guidance paths --- .../models/ltx2/transformer_ltx2.py | 10 +- .../pipelines/ltx2/ltx2_pipeline.py | 355 +++++++++--------- .../tests/ltx2/test_pipeline_ltx2.py | 79 +++- 3 files changed, 252 insertions(+), 192 deletions(-) diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py index 374038bb6..98f21d278 100644 --- a/src/maxdiffusion/models/ltx2/transformer_ltx2.py +++ b/src/maxdiffusion/models/ltx2/transformer_ltx2.py @@ -55,7 +55,6 @@ class LTX2BlockContext: a2v_cross_attention_mask: Optional[jax.Array] = None v2a_cross_attention_mask: Optional[jax.Array] = None perturbation_mask: Optional[jax.Array] = None - cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None class LTX2AdaLayerNormSingle(nnx.Module): @@ -470,6 +469,7 @@ def named_scope(self, name: str): def __call__( self, ctx: "LTX2BlockContext", + cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, ) -> Tuple[jax.Array, jax.Array]: """ Forward pass of the LTX2 video/audio transformer block. @@ -512,7 +512,6 @@ def __call__( a2v_cross_attention_mask = ctx.a2v_cross_attention_mask v2a_cross_attention_mask = ctx.v2a_cross_attention_mask perturbation_mask = ctx.perturbation_mask - cached_kv = ctx.cached_kv batch_size = hidden_states.shape[0] axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) @@ -1466,20 +1465,17 @@ def __call__( a2v_cross_attention_mask=a2v_cross_attention_mask, v2a_cross_attention_mask=v2a_cross_attention_mask, modality_mask=modality_mask, - cached_kv=cached_kv, ) def apply_block(block, context: LTX2BlockContext, mask, block_cached_kv=None): orig_perturbation_mask = context.perturbation_mask - orig_cached_kv = context.cached_kv - context = context.replace(perturbation_mask=mask, cached_kv=block_cached_kv) + context = context.replace(perturbation_mask=mask) with self.named_scope("Transformer Layer"): - hidden_states_out, audio_hidden_states_out = block(context) + hidden_states_out, audio_hidden_states_out = block(context, cached_kv=block_cached_kv) context = context.replace( hidden_states=hidden_states_out.astype(context.hidden_states.dtype), audio_hidden_states=audio_hidden_states_out.astype(context.audio_hidden_states.dtype), perturbation_mask=orig_perturbation_mask, - cached_kv=orig_cached_kv, ) return context diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index db45059d7..fbf820d75 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -96,6 +96,113 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): return noise_cfg +def _select_guidance_latents(latents, audio_latents, batch_size, do_cfg, do_stg): + """Returns the conditional latent slice used for the scheduler update.""" + if do_cfg and do_stg: + return latents[batch_size : 2 * batch_size], audio_latents[batch_size : 2 * batch_size] + if do_cfg: + return latents[batch_size:], audio_latents[batch_size:] + return latents, audio_latents + + +def _restore_guidance_latents(latents, audio_latents, do_cfg, do_stg): + """Restores the batch layout expected by the next denoising iteration.""" + if do_cfg and do_stg: + return jnp.concatenate([latents] * 4, axis=0), jnp.concatenate([audio_latents] * 4, axis=0) + if do_cfg: + return jnp.concatenate([latents] * 2, axis=0), jnp.concatenate([audio_latents] * 2, axis=0) + return latents, audio_latents + + +def _apply_ltx2_guidance( + noise_pred, + noise_pred_audio, + latents, + audio_latents, + sigma_t, + guidance_scale, + stg_scale, + modality_scale, + guidance_rescale, + audio_guidance_scale, + audio_stg_scale, + audio_modality_scale, + audio_guidance_rescale, + do_cfg, + do_stg, + use_lax_cond=False, +): + """Applies the shared LTX2 CFG/STG delta formulation to video and audio.""" + if not do_cfg: + return noise_pred, noise_pred_audio + + if not do_stg: + noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0) + noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond) + return noise_pred, noise_pred_audio + + def convert_to_x0(latents_to_convert, velocity): + return latents_to_convert - velocity * sigma_t + + def convert_to_vel(latents_to_convert, x0): + return (latents_to_convert - x0) / sigma_t + + def maybe_rescale(noise, text_noise, scale): + if use_lax_cond: + return jax.lax.cond( + scale > 0, + lambda: rescale_noise_cfg(noise, text_noise, guidance_rescale=scale), + lambda: noise, + ) + if scale > 0: + return rescale_noise_cfg(noise, text_noise, guidance_rescale=scale) + return noise + + ( + noise_pred_uncond, + noise_pred_text, + noise_pred_perturb, + noise_pred_isolated, + ) = jnp.split(noise_pred, 4, axis=0) + x0_uncond = convert_to_x0(latents, noise_pred_uncond) + x0_text = convert_to_x0(latents, noise_pred_text) + x0_perturb = convert_to_x0(latents, noise_pred_perturb) + x0_isolated = convert_to_x0(latents, noise_pred_isolated) + + cfg_delta = (guidance_scale - 1) * (x0_text - x0_uncond) + stg_delta = stg_scale * (x0_text - x0_perturb) + modality_delta = (modality_scale - 1) * (x0_text - x0_isolated) + x0_combined = maybe_rescale(x0_text + cfg_delta + stg_delta + modality_delta, x0_text, guidance_rescale) + noise_pred = convert_to_vel(latents, x0_combined) + + ( + noise_pred_audio_uncond, + noise_pred_audio_text, + noise_pred_audio_perturb, + noise_pred_audio_isolated, + ) = jnp.split(noise_pred_audio, 4, axis=0) + x0_audio_uncond = convert_to_x0(audio_latents, noise_pred_audio_uncond) + x0_audio_text = convert_to_x0(audio_latents, noise_pred_audio_text) + x0_audio_perturb = convert_to_x0(audio_latents, noise_pred_audio_perturb) + x0_audio_isolated = convert_to_x0(audio_latents, noise_pred_audio_isolated) + + audio_guidance_scale = guidance_scale if audio_guidance_scale is None else audio_guidance_scale + audio_stg_scale = stg_scale if audio_stg_scale is None else audio_stg_scale + audio_modality_scale = modality_scale if audio_modality_scale is None else audio_modality_scale + cfg_audio_delta = (audio_guidance_scale - 1) * (x0_audio_text - x0_audio_uncond) + stg_audio_delta = audio_stg_scale * (x0_audio_text - x0_audio_perturb) + audio_modality_delta = (audio_modality_scale - 1) * (x0_audio_text - x0_audio_isolated) + x0_audio_combined = maybe_rescale( + x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta, + x0_audio_text, + audio_guidance_rescale, + ) + noise_pred_audio = convert_to_vel(audio_latents, x0_audio_combined) + return noise_pred, noise_pred_audio + + logger = logging.get_logger(__name__) @@ -1621,6 +1728,8 @@ def __call__( prompt_attention_mask, negative_prompt_attention_mask, ) + if stg_scale > 0.0 and guidance_scale <= 1.0: + raise ValueError("Spatio-temporal guidance requires guidance_scale > 1.0.") # 2. Encode inputs (Text) t0_encode = time.perf_counter() @@ -1928,82 +2037,31 @@ def __call__( rope_cache=rope_cache, ) - # Extract latents_step based on stacking strategy - if do_cfg and do_stg: - latents_step = latents_jax[batch_size : 2 * batch_size] - audio_latents_step = audio_latents_jax[batch_size : 2 * batch_size] - elif do_cfg: - latents_step = latents_jax[batch_size:] - audio_latents_step = audio_latents_jax[batch_size:] - elif do_stg: - latents_step = latents_jax[:batch_size] - audio_latents_step = audio_latents_jax[:batch_size] - else: - latents_step = latents_jax - audio_latents_step = audio_latents_jax - - def convert_to_x0(lat, vel, sig): - return lat - vel * sig - - def convert_to_vel(lat, x0, sig): - return (lat - x0) / sig - - if do_cfg and do_stg: - ( - noise_pred_uncond, - noise_pred_text, - noise_pred_perturb, - noise_pred_isolated, - ) = jnp.split(noise_pred, 4, axis=0) - x0_uncond = convert_to_x0(latents_step, noise_pred_uncond, sigma_t) - x0_text = convert_to_x0(latents_step, noise_pred_text, sigma_t) - x0_perturb = convert_to_x0(latents_step, noise_pred_perturb, sigma_t) - x0_isolated = convert_to_x0(latents_step, noise_pred_isolated, sigma_t) - - cfg_delta = (guidance_scale - 1) * (x0_text - x0_uncond) - stg_delta = stg_scale * (x0_text - x0_perturb) - video_modality_delta = (modality_scale - 1) * (x0_text - x0_isolated) - - x0_combined = x0_text + cfg_delta + stg_delta + video_modality_delta - if guidance_rescale > 0: - x0_combined = rescale_noise_cfg(x0_combined, x0_text, guidance_rescale=guidance_rescale) - noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t) - - # Audio - ( - noise_pred_audio_uncond, - noise_pred_audio_text, - noise_pred_audio_perturb, - noise_pred_audio_isolated, - ) = jnp.split(noise_pred_audio, 4, axis=0) - x0_audio_uncond = convert_to_x0(audio_latents_step, noise_pred_audio_uncond, sigma_t) - x0_audio_text = convert_to_x0(audio_latents_step, noise_pred_audio_text, sigma_t) - x0_audio_perturb = convert_to_x0(audio_latents_step, noise_pred_audio_perturb, sigma_t) - x0_audio_isolated = convert_to_x0(audio_latents_step, noise_pred_audio_isolated, sigma_t) - - cfg_audio_delta = (audio_guidance_scale - 1 if audio_guidance_scale is not None else guidance_scale - 1) * ( - x0_audio_text - x0_audio_uncond - ) - stg_audio_delta = (audio_stg_scale if audio_stg_scale is not None else stg_scale) * ( - x0_audio_text - x0_audio_perturb - ) - audio_modality_delta = (audio_modality_scale - 1 if audio_modality_scale is not None else modality_scale - 1) * ( - x0_audio_text - x0_audio_isolated - ) - x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta - if audio_guidance_rescale > 0: - x0_audio_combined = rescale_noise_cfg( - x0_audio_combined, - x0_audio_text, - guidance_rescale=audio_guidance_rescale, - ) - noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) - - elif do_cfg: - noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0) - noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond) + latents_step, audio_latents_step = _select_guidance_latents( + latents_jax, + audio_latents_jax, + batch_size, + do_cfg, + do_stg, + ) + noise_pred, noise_pred_audio = _apply_ltx2_guidance( + noise_pred, + noise_pred_audio, + latents_step, + audio_latents_step, + sigma_t, + guidance_scale, + stg_scale, + modality_scale, + guidance_rescale, + audio_guidance_scale, + audio_stg_scale, + audio_modality_scale, + audio_guidance_rescale, + do_cfg, + do_stg, + use_lax_cond=False, + ) latents_step, _ = self.scheduler.step(scheduler_state, noise_pred, t, latents_step, return_dict=False) audio_latents_step, _ = self.scheduler.step( @@ -2014,15 +2072,12 @@ def convert_to_vel(lat, x0, sig): return_dict=False, ) - if do_cfg and do_stg: - latents_jax = jnp.concatenate([latents_step] * 4, axis=0) - audio_latents_jax = jnp.concatenate([audio_latents_step] * 4, axis=0) - elif do_cfg: - latents_jax = jnp.concatenate([latents_step] * 2, axis=0) - audio_latents_jax = jnp.concatenate([audio_latents_step] * 2, axis=0) - else: - latents_jax = latents_step - audio_latents_jax = audio_latents_step + latents_jax, audio_latents_jax = _restore_guidance_latents( + latents_step, + audio_latents_step, + do_cfg, + do_stg, + ) jax.block_until_ready(latents_jax) denoise_time = time.perf_counter() - t0_denoise @@ -2415,13 +2470,6 @@ def run_diffusion_loop( audio_num_frames, ) - # Helper functions matching Diffusers Delta formulation - def convert_to_x0(lat, vel, sigma_t): - return lat - vel * sigma_t - - def convert_to_vel(lat, x0, sigma_t): - return (lat - x0) / sigma_t - if not scan_layers: with nn_partitioning.axis_rules(logical_axis_rules): activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) @@ -2465,89 +2513,31 @@ def scan_body(carry, inputs): rope_cache=rope_cache, ) - # Extract latents_step based on stacking strategy - if do_cfg and do_stg: - latents_step = latents[batch_size : 2 * batch_size] - audio_latents_step = audio_latents[batch_size : 2 * batch_size] - elif do_cfg: - latents_step = latents[batch_size:] - audio_latents_step = audio_latents[batch_size:] - else: - latents_step = latents - audio_latents_step = audio_latents - - # Apply Diffusers STG + CFG + Modality Delta Logic - if do_cfg and do_stg: - ( - noise_pred_uncond, - noise_pred_text, - noise_pred_perturb, - noise_pred_isolated, - ) = jnp.split(noise_pred, 4, axis=0) - - # Convert to x0 - x0_uncond = convert_to_x0(latents_step, noise_pred_uncond, sigma_t) - x0_text = convert_to_x0(latents_step, noise_pred_text, sigma_t) - x0_perturb = convert_to_x0(latents_step, noise_pred_perturb, sigma_t) - x0_isolated = convert_to_x0(latents_step, noise_pred_isolated, sigma_t) - - # Delta formulation - cfg_delta = (guidance_scale - 1) * (x0_text - x0_uncond) - stg_delta = stg_scale * (x0_text - x0_perturb) - video_modality_delta = (modality_scale - 1) * (x0_text - x0_isolated) - - x0_combined = x0_text + cfg_delta + stg_delta + video_modality_delta - - x0_combined = jax.lax.cond( - guidance_rescale > 0, - lambda: rescale_noise_cfg(x0_combined, x0_text, guidance_rescale=guidance_rescale), - lambda: x0_combined, - ) - - noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t) - - # Audio guidance - ( - noise_pred_audio_uncond, - noise_pred_audio_text, - noise_pred_audio_perturb, - noise_pred_audio_isolated, - ) = jnp.split(noise_pred_audio, 4, axis=0) - - x0_audio_uncond = convert_to_x0(audio_latents_step, noise_pred_audio_uncond, sigma_t) - x0_audio_text = convert_to_x0(audio_latents_step, noise_pred_audio_text, sigma_t) - x0_audio_perturb = convert_to_x0(audio_latents_step, noise_pred_audio_perturb, sigma_t) - x0_audio_isolated = convert_to_x0(audio_latents_step, noise_pred_audio_isolated, sigma_t) - - cfg_audio_delta = (audio_guidance_scale - 1 if audio_guidance_scale is not None else guidance_scale - 1) * ( - x0_audio_text - x0_audio_uncond - ) - stg_audio_delta = (audio_stg_scale if audio_stg_scale is not None else stg_scale) * (x0_audio_text - x0_audio_perturb) - audio_modality_delta = (audio_modality_scale - 1 if audio_modality_scale is not None else modality_scale - 1) * ( - x0_audio_text - x0_audio_isolated - ) - - x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta - - x0_audio_combined = jax.lax.cond( - audio_guidance_rescale > 0, - lambda: rescale_noise_cfg( - x0_audio_combined, - x0_audio_text, - guidance_rescale=audio_guidance_rescale, - ), - lambda: x0_audio_combined, - ) - - noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) - - # Standard CFG path when spatio-temporal guidance is disabled. - elif do_cfg: - noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0) - noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond) + latents_step, audio_latents_step = _select_guidance_latents( + latents, + audio_latents, + batch_size, + do_cfg, + do_stg, + ) + noise_pred, noise_pred_audio = _apply_ltx2_guidance( + noise_pred, + noise_pred_audio, + latents_step, + audio_latents_step, + sigma_t, + guidance_scale, + stg_scale, + modality_scale, + guidance_rescale, + audio_guidance_scale, + audio_stg_scale, + audio_modality_scale, + audio_guidance_rescale, + do_cfg, + do_stg, + use_lax_cond=True, + ) # Step scheduler latents_step, _ = scheduler_step(s_state, noise_pred, t, latents_step, return_dict=False) @@ -2556,15 +2546,12 @@ def scan_body(carry, inputs): audio_latents_step, _ = scheduler_step(s_state, noise_pred_audio, t, audio_latents_step, return_dict=False) audio_latents_step = audio_latents_step.astype(audio_latents.dtype) - if do_cfg and do_stg: - latents_next = jnp.concatenate([latents_step] * 4, axis=0) - audio_latents_next = jnp.concatenate([audio_latents_step] * 4, axis=0) - elif do_cfg: - latents_next = jnp.concatenate([latents_step] * 2, axis=0) - audio_latents_next = jnp.concatenate([audio_latents_step] * 2, axis=0) - else: - latents_next = latents_step - audio_latents_next = audio_latents_step + latents_next, audio_latents_next = _restore_guidance_latents( + latents_step, + audio_latents_step, + do_cfg, + do_stg, + ) new_carry = (latents_next, audio_latents_next, s_state) return new_carry, None diff --git a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py index 4e4d86634..348e85438 100644 --- a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py @@ -16,10 +16,18 @@ import unittest from unittest.mock import MagicMock, patch +import jax import jax.numpy as jnp import numpy as np -from maxdiffusion.pipelines.ltx2.ltx2_pipeline import LTX2Pipeline, calculate_shift, rescale_noise_cfg +from maxdiffusion.pipelines.ltx2.ltx2_pipeline import ( + LTX2Pipeline, + _apply_ltx2_guidance, + _restore_guidance_latents, + _select_guidance_latents, + calculate_shift, + rescale_noise_cfg, +) class LTX2PipelineTest(unittest.TestCase): @@ -55,6 +63,75 @@ def test_rescale_noise_cfg(self): rescaled_0 = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0) np.testing.assert_allclose(rescaled_0, noise_cfg, rtol=1e-5) + def test_guidance_helpers(self): + latents = jnp.array([[10.0]]) + audio_latents = jnp.array([[20.0]]) + + # Standard CFG uses the two-way [unconditional, conditional] layout. + cfg_noise, cfg_audio_noise = _apply_ltx2_guidance( + jnp.array([[1.0], [3.0]]), + jnp.array([[2.0], [4.0]]), + latents, + audio_latents, + sigma_t=2.0, + guidance_scale=2.0, + stg_scale=0.0, + modality_scale=1.0, + guidance_rescale=0.0, + audio_guidance_scale=None, + audio_stg_scale=None, + audio_modality_scale=None, + audio_guidance_rescale=0.0, + do_cfg=True, + do_stg=False, + ) + np.testing.assert_allclose(cfg_noise, jnp.array([[5.0]])) + np.testing.assert_allclose(cfg_audio_noise, jnp.array([[6.0]])) + + # CFG+STG returns the same delta formulation in eager and traced execution. + stg_noise = jnp.array([[1.0], [3.0], [4.0], [5.0]]) + apply_stg = jax.jit( + lambda rescale: _apply_ltx2_guidance( + stg_noise, + stg_noise, + latents, + latents, + sigma_t=2.0, + guidance_scale=2.0, + stg_scale=1.0, + modality_scale=2.0, + guidance_rescale=rescale, + audio_guidance_scale=None, + audio_stg_scale=None, + audio_modality_scale=None, + audio_guidance_rescale=rescale, + do_cfg=True, + do_stg=True, + use_lax_cond=True, + ) + ) + guided_noise, guided_audio_noise = apply_stg(jnp.array(0.0)) + np.testing.assert_allclose(guided_noise, jnp.array([[2.0]])) + np.testing.assert_allclose(guided_audio_noise, jnp.array([[2.0]])) + + stacked_latents = jnp.concatenate([latents] * 4, axis=0) + stacked_audio_latents = jnp.concatenate([audio_latents] * 4, axis=0) + selected_latents, selected_audio_latents = _select_guidance_latents( + stacked_latents, + stacked_audio_latents, + batch_size=1, + do_cfg=True, + do_stg=True, + ) + restored_latents, restored_audio_latents = _restore_guidance_latents( + selected_latents, + selected_audio_latents, + do_cfg=True, + do_stg=True, + ) + np.testing.assert_array_equal(restored_latents, stacked_latents) + np.testing.assert_array_equal(restored_audio_latents, stacked_audio_latents) + def test_pipeline_init(self): """Test LTX2Pipeline initialization and property extraction.""" mock_vae = MagicMock() From af55325008befffd9d8d4064d76580e6f477124c Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Mon, 27 Jul 2026 12:18:46 +0530 Subject: [PATCH 07/10] Fix LTX2 TPU correctness and cache safety --- LTX2_BRANCH_REVIEW_AND_IMPLEMENTATION_PLAN.md | 445 ++++++++++++++++++ src/maxdiffusion/aot_cache.py | 30 +- src/maxdiffusion/configs/ltx2_3_video.yml | 19 +- src/maxdiffusion/configs/ltx2_video.yml | 19 +- src/maxdiffusion/generate_ltx2.py | 122 ++++- src/maxdiffusion/max_utils.py | 40 ++ src/maxdiffusion/models/attention_flax.py | 305 +++++++++--- .../models/ltx2/attention_ltx2.py | 29 +- src/maxdiffusion/models/ltx2/ltx2_utils.py | 137 ++++-- .../models/ltx2/transformer_ltx2.py | 82 +++- .../pipelines/ltx2/ltx2_pipeline.py | 104 ++-- src/maxdiffusion/pyconfig.py | 13 +- src/maxdiffusion/tests/aot_cache_test.py | 13 + src/maxdiffusion/tests/attention_test.py | 111 ++++- .../tests/ltx2/test_attention_ltx2.py | 244 ++++++++++ .../tests/ltx2/test_generate_ltx2_aot.py | 183 +++++++ .../tests/ltx2/test_ltx2_utils_loader.py | 167 +++++++ .../tests/ltx2/test_pipeline_ltx2.py | 39 ++ .../tests/tile_size_grid_search_test.py | 36 ++ .../utils/ltx2_block_benchmark.py | 26 +- .../utils/tile_size_grid_search.py | 109 ++++- 21 files changed, 2050 insertions(+), 223 deletions(-) create mode 100644 LTX2_BRANCH_REVIEW_AND_IMPLEMENTATION_PLAN.md create mode 100644 src/maxdiffusion/tests/ltx2/test_generate_ltx2_aot.py create mode 100644 src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py diff --git a/LTX2_BRANCH_REVIEW_AND_IMPLEMENTATION_PLAN.md b/LTX2_BRANCH_REVIEW_AND_IMPLEMENTATION_PLAN.md new file mode 100644 index 000000000..148e06942 --- /dev/null +++ b/LTX2_BRANCH_REVIEW_AND_IMPLEMENTATION_PLAN.md @@ -0,0 +1,445 @@ +# LTX2 Branch Review and Implementation Plan + +Date: 2026-07-26 + +Branch: `ltx2-improvements` + +Review base: local `origin/main` at merge base `bda41e1572c82da0dc6129ade09f2fc3155b5420` + +## Executive Summary + +The branch contains useful improvements, particularly the scan-loop warmup fix, KV-cache compatibility guard, guidance-path deduplication, and new KV-cache parity tests. It is not ready to merge yet. + +The main blockers are: + +1. AOT executables can be reused after the underlying MaxDiffusion source changes. +2. The `.bin` checkpoint fallback can load dozens of full checkpoint copies concurrently. +3. Tile auto-tuning fails with the shipped `attention: flash` configuration. +4. Different JAX processes can choose different tile-search winners. +5. Ring-configured cross-attention can perform softmax over only the local KV shard. + +Several medium-severity sharding, masking, state-management, and validation problems should also be fixed before the new paths are treated as production-ready. + +## Implementation Status + +Updated: 2026-07-27 + +The fixes in phases 1–6 are now implemented in the working tree: + +- Checkpoint loading no longer reloads `.bin` files per chunk, and scanned parameters require exact layer coverage. +- AOT cache identity includes source/build revision, an ABI, and graph-shaping options; dirty or unversioned runs disable persistent caching. +- Attention masks remain canonical per-example keep masks until backend dispatch. Dot, Splash, Ulysses, and ring paths now receive the correct representation. +- Local-flash cross-attention keeps KV global, including independently configured a2v/v2a kernels and remapped ring/custom variants. +- Tile tuning uses production block-size construction, production-style padded ring lengths, a shared VMEM budget, and one multi-host winner. +- TPU text batches select divisible mesh axes; VAE slicing changes are temporary and restored on failures. +- Unsupported/unused configuration entries and whitespace issues were cleaned up. + +Additional integration fixes found during implementation: + +- Removed an unsafe conditional multi-host parameter `process_allgather` fallback that could deadlock or expand parameter shapes. +- Explicitly reject tile tuning with an already-built pipeline, where static block sizes cannot be changed. +- Reset AOT global state when caching is disabled after a prior cached run. + +Static verification passes: + +- `git diff --check origin/main` +- `python -m compileall -q src/maxdiffusion` + +Runtime unit tests and Ruff could not run in the desktop Python because the project dependencies are not installed (`packaging`, JAX, Flax, NumPy, PyTorch, pytest, and Ruff are absent). TPU numerical parity and multi-device qualification therefore remain merge gates. + +## Review Findings + +### P1: AOT Cache Does Not Include the Source Revision + +Location: [`src/maxdiffusion/generate_ltx2.py`](src/maxdiffusion/generate_ltx2.py#L159) + +`ltx2_aot_metadata` fingerprints model configuration, topology, attention settings, and JAX versions, but it does not include the MaxDiffusion source revision or an explicit cache ABI. + +This is already observable within this branch: the AOT metadata was strengthened in one commit, and a later commit materially changed cached diffusion and guidance code without changing that metadata. If the shapes and configuration remain unchanged, an executable generated from the earlier code remains eligible for loading. + +Impact: + +- Old HLO can execute silently after a code upgrade. +- Fixes in Python/JAX source may appear to have no effect. +- The result may be numerically wrong without retracing or producing an error. + +Recommended fix: + +- Add an explicit `LTX2_AOT_ABI_VERSION`. +- Include the repository commit or package build revision in the metadata. +- Treat dirty working trees distinctly, or disable disk-cache loading for dirty source. +- Log the complete fingerprint inputs when installing the cache. + +### P1: `.bin` Loading Can Create Dozens of Full Checkpoint Copies + +Locations: + +- [`src/maxdiffusion/models/ltx2/ltx2_utils.py`](src/maxdiffusion/models/ltx2/ltx2_utils.py#L242) +- [`src/maxdiffusion/models/ltx2/ltx2_utils.py`](src/maxdiffusion/models/ltx2/ltx2_utils.py#L302) + +The fallback path first loads the entire `.bin` state dictionary to enumerate its keys. It then creates 32-key tasks, and every worker calls `torch.load` on the same full checkpoint again. + +With 32 workers, the process can hold approximately 33 full checkpoint copies at once, in addition to converted arrays. For an LTX2-sized checkpoint, this is a likely host OOM and causes extreme repeated I/O. + +Recommended fix: + +- Load a `.bin` checkpoint once and process its items sequentially. +- Retain parallel chunk loading only for independently readable safetensors. +- Release the source state dictionary as soon as conversion completes. +- Add a regression test that asserts `torch.load` is called once. + +### P1: Tile Auto-Tuning Fails for Standard Flash Attention + +Locations: + +- [`src/maxdiffusion/utils/ltx2_block_benchmark.py`](src/maxdiffusion/utils/ltx2_block_benchmark.py#L206) +- [`src/maxdiffusion/models/attention_flax.py`](src/maxdiffusion/models/attention_flax.py#L565) + +The benchmark always constructs `CustomFlashBlockSizes`. Standard Splash attention expects a full `splash_attention_kernel.BlockSizes` object and accesses fields such as: + +- `use_fused_bwd_kernel` +- `block_q_dkv` +- `block_kv_dkv` +- `block_q_dq` +- `block_kv_dq` + +Those fields do not exist on `CustomFlashBlockSizes`. + +The shipped LTX2 configuration uses `attention: flash`, and its default latent sequence length is 6144, which exceeds the configured flash threshold of 4096. Enabling tile search therefore causes every candidate to error. The broad exception handler converts the failures into benchmark results, after which generation silently retains the old configuration. + +Recommended fix: + +- Construct standard Splash `BlockSizes` for `flash` and plain Ulysses. +- Use `CustomFlashBlockSizes` only for kernels that actually consume its custom fields. +- Build benchmark block sizes through the same factory used by production inference. +- Treat an enabled search with zero successful candidates as an explicit failure. + +### P1: Tile Search Is Not Synchronized Across JAX Processes + +Locations: + +- [`src/maxdiffusion/utils/tile_size_grid_search.py`](src/maxdiffusion/utils/tile_size_grid_search.py#L238) +- [`src/maxdiffusion/utils/tile_size_grid_search.py`](src/maxdiffusion/utils/tile_size_grid_search.py#L386) +- [`src/maxdiffusion/generate_ltx2.py`](src/maxdiffusion/generate_ltx2.py#L142) + +Each process measures local wall-clock time, independently selects its minimum, and immediately writes that result into its local configuration. + +On a multi-controller TPU pod, small timing differences can cause different processes to select different static block sizes. The following distributed model construction or compilation may then use inconsistent static parameters across hosts, risking compilation mismatches or hangs. + +Recommended fix: + +- Execute candidates in identical order on every process. +- Synchronize before and after timed regions. +- Gather per-process measurements. +- Rank candidates using an aggregate or slowest-host measurement. +- Select on process 0 and broadcast the exact winning tuple to all processes. + +### P1: Ring Cross-Attention Can See Only Local KV + +Locations: + +- [`src/maxdiffusion/models/ltx2/attention_ltx2.py`](src/maxdiffusion/models/ltx2/attention_ltx2.py#L457) +- [`src/maxdiffusion/common_types.py`](src/maxdiffusion/common_types.py#L71) +- [`src/maxdiffusion/models/attention_flax.py`](src/maxdiffusion/models/attention_flax.py#L530) + +For cross-attention, ring variants such as `tokamax_ring`, `tokamax_ring_custom`, and `ulysses_ring` are changed to `tokamax_flash`. The logical axis rules still shard cross-attention KV length over the context axis. + +`tokamax_flash` runs a local `shard_map` kernel and performs no ring rotation or KV all-gather. When flash is selected, each query shard therefore attends only to its local KV slice. + +Trigger: + +- Context parallelism greater than one, and +- `flash_min_seq_length=0`, or cross-modal sequences long enough to select flash. + +Impact: + +- Silent numerical corruption rather than a performance-only regression. + +Recommended fix: + +- Keep cross-attention Q sequence-sharded. +- Replicate or all-gather cross-attention KV before routing to local flash. +- Alternatively, route to a cross-attention kernel that globalizes KV. +- Test against a replicated dot-product reference on multiple devices. + +### P2: The Tuned BQ Is Ignored by Tokamax and `ulysses_ring` + +Locations: + +- [`src/maxdiffusion/generate_ltx2.py`](src/maxdiffusion/generate_ltx2.py#L142) +- [`src/maxdiffusion/max_utils.py`](src/maxdiffusion/max_utils.py#L664) + +Tile search updates `block_q`, but production Tokamax and `ulysses_ring` construct their effective forward `block_q` from `block_q_dkv`. The configured `block_q_dkv` remains unchanged, so the model may not use the BQ that the benchmark selected. + +Recommended fix: + +- Apply the winner to every field consumed by the selected kernel. +- Validate the final object returned by `get_flash_block_sizes`. +- Use the same construction path in the benchmark and production model. + +### P2: TPU Text-Encoder Batch Sharding Usually Replicates the Batch + +Location: [`src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py`](src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py#L1111) + +The TPU text-encoder path now shards the batch over `self.mesh.axis_names[0]`. Both supplied LTX2 configurations put `data` first and configure the data-axis size as one. + +As a result, `P("data")` replicates the text batch and defeats the stated Softmax OOM protection. If a user configures a larger data axis that does not divide the prompt batch, `jax.device_put` can fail. + +Recommended fix: + +- Select a combination of mesh axes whose product divides the prompt batch. +- Prefer data-parallel axes, but permit compatible replicated-weight axes. +- Pad and trim non-divisible prompt batches when sharding is required. +- Use a safe replicated fallback. + +### P2: a2v and v2a Masks Are Converted for the Wrong Kernel + +Locations: + +- [`src/maxdiffusion/models/ltx2/transformer_ltx2.py`](src/maxdiffusion/models/ltx2/transformer_ltx2.py#L1320) +- [`src/maxdiffusion/models/attention_flax.py`](src/maxdiffusion/models/attention_flax.py#L1375) + +The transformer converts all masks according to the main attention kernel, even though a2v and v2a use independently configured kernels and may independently fall back from flash to dot product. + +Consequences: + +- Dot-product attention receives a mask but currently ignores it. +- Flash attention can receive an additive mask where it expects binary validity or segment IDs. +- A mask can be ignored or misinterpreted depending on the main kernel rather than the effective per-operation kernel. + +Recommended fix: + +- Keep a canonical boolean mask at the transformer boundary. +- Convert the mask inside the selected backend. +- Make dot-product attention consume masks. +- Preserve per-example masks rather than using the first batch element for every sample. + +### P2: Scanned Weight Stacks Can Contain Uninitialized Layers + +Location: [`src/maxdiffusion/models/ltx2/ltx2_utils.py`](src/maxdiffusion/models/ltx2/ltx2_utils.py#L290) + +Scanned weights are allocated with `np.empty`. If a checkpoint is missing a per-layer tensor, has fewer layers than expected, or is loaded using the wrong layer count, the corresponding slice remains uninitialized. + +The consolidated key and shape still exist, so the current validation cannot detect the missing layer slice. + +Recommended fix: + +- Derive the expected layer count from the initialized state. +- Track the exact layer indices populated for every stacked parameter. +- Raise on missing, duplicate, or out-of-range indices. +- Do not silently accept zero-filled or uninitialized model layers. + +### P2: Dynamic VAE Tuning Permanently Disables Slicing + +Location: [`src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py`](src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py#L2193) + +For batches larger than two, the pipeline sets `self.vae.use_slicing = False` and never restores its prior value. Reusing the pipeline can therefore lose a user-enabled memory-safety setting. The mutation also survives dynamic-sharding failures. + +Recommended fix: + +- Save the original slicing state. +- Apply temporary changes using `try/finally` or a context manager. +- Restore the state after decoding, including exception paths. +- Keep slicing enabled if no useful dynamic batch sharding can be applied. + +### Build and Configuration Issues + +`git diff --check origin/main...HEAD` currently reports trailing whitespace in: + +- `src/maxdiffusion/configs/ltx2_video.yml` +- `src/maxdiffusion/configs/ltx2_3_video.yml` +- `src/maxdiffusion/models/ltx2/ltx2_utils.py` + +Ruff selects `W` rules and CI runs `ruff check .`, so these should be treated as a merge gate. + +The configuration surface should also be audited: + +- `ring` is advertised as supported, but there is no `ring` entry in the attention registry. +- `use_base2_exp` and `use_experimental_scheduler` are present in LTX2 configs but are not wired through the LTX2 model. +- `converted_weights_dir` is present but unused by the LTX2 pipeline. + +Unsupported or unused knobs should be implemented or removed. + +## Implementation Plan + +The fixes should be split into focused, reviewable commits. Correctness and safety should land before performance work. + +### Phase 1: Checkpoint Loading and Scanned-Weight Validation + +Files: + +- `src/maxdiffusion/models/ltx2/ltx2_utils.py` +- LTX2 loader tests + +Work: + +1. Split safetensors and `.bin` conversion paths. +2. Load `.bin` files once. +3. Retain chunked parallelism only for safetensors. +4. Derive the scanned layer count from `eval_shapes`. +5. Track populated layer indices per parameter. +6. Raise on missing, duplicate, or out-of-range layers. +7. Release source tensors as early as possible. + +Exit criteria: + +- `.bin` tests observe one `torch.load`. +- Safetensors tests retain parallel loading. +- Missing-layer and layer-count mismatch tests fail clearly. +- Converted LTX2 and LTX2.3 states match initialized shapes. + +### Phase 2: AOT Cache Safety + +Files: + +- `src/maxdiffusion/generate_ltx2.py` +- `src/maxdiffusion/aot_cache.py` +- `src/maxdiffusion/tests/aot_cache_test.py` + +Work: + +1. Add `LTX2_AOT_ABI_VERSION`. +2. Pass a source or build revision into `ltx2_aot_metadata`. +3. Distinguish dirty source trees. +4. Include the ABI and revision in the cache fingerprint. +5. Improve fingerprint logging. + +Exit criteria: + +- Identical code and configuration produce the same key. +- Source or ABI changes produce a different key. +- Old executables are not loaded after a source change. +- Cold-cache and warm-cache outputs match. + +### Phase 3: Attention Routing and Mask Correctness + +Files: + +- `src/maxdiffusion/models/ltx2/attention_ltx2.py` +- `src/maxdiffusion/models/ltx2/transformer_ltx2.py` +- `src/maxdiffusion/models/attention_flax.py` +- attention tests + +Work: + +1. Define a canonical boolean mask representation. +2. Move backend-specific mask conversion into the backend. +3. Add mask support to dot-product attention. +4. Preserve batch-specific masks. +5. Replicate or gather KV when ring cross-attention is routed to local flash. +6. Audit all ring-to-non-ring cross-attention mappings. + +Exit criteria: + +- Dot, flash, a2v, and v2a masking match a reference implementation. +- Different masks within one batch remain distinct. +- Ring cross-attention matches replicated dot product on a multi-device TPU. +- Normal and forced-flash thresholds are covered. + +### Phase 4: Kernel-Aware, Multi-Host Tile Tuning + +Files: + +- `src/maxdiffusion/utils/ltx2_block_benchmark.py` +- `src/maxdiffusion/utils/tile_size_grid_search.py` +- `src/maxdiffusion/generate_ltx2.py` +- `src/maxdiffusion/max_utils.py` + +Work: + +1. Use the production block-size constructor in the benchmark. +2. Construct standard and custom block-size types according to the actual kernel. +3. Apply winners to every effective field, including `block_q_dkv`. +4. Centralize local sequence-length calculations. +5. Cover every supported ring variant. +6. Add cross-process timing collection. +7. Select on process 0 and broadcast the result. +8. Fail clearly when an explicitly enabled search has no valid candidate. + +Exit criteria: + +- The shipped `attention: flash` configuration produces a valid winner. +- Benchmark and production models use identical effective block sizes. +- All processes receive the same winner despite different local timings. +- Flash, Tokamax, Ulysses, and supported ring variants are covered. + +### Phase 5: TPU Text Sharding and VAE State Isolation + +File: + +- `src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py` + +Work: + +1. Add a pure helper for selecting batch-sharding axes. +2. Handle non-divisible batches using padding and trimming or safe replication. +3. Verify that sharding actually uses multiple devices when possible. +4. Make VAE slicing changes temporary. +5. Restore VAE state after successful and failed decodes. + +Exit criteria: + +- The default context-parallel mesh shards a batch of 16. +- Batch sizes 1, 3, and 16 do not fail placement. +- Repeated pipeline calls do not leak VAE state. +- Failed dynamic sharding does not disable slicing. + +### Phase 6: Public Configuration and CI Cleanup + +Work: + +1. Remove all trailing whitespace. +2. Run Ruff and formatting checks. +3. Implement or remove unused configuration keys. +4. Remove unsupported attention names from configuration comments. +5. Document supported attention kernels and their sharding requirements. + +Exit criteria: + +- `git diff --check origin/main...HEAD` passes. +- `ruff check .` passes. +- Every advertised LTX2 configuration option has a runtime consumer and a test. + +## TPU Qualification Matrix + +Before merging, validate: + +| Dimension | Required coverage | +| --- | --- | +| Model | LTX2 and LTX2.3 | +| Layer execution | `scan_layers=True` and `False` | +| Diffusion loop | `scan_diffusion_loop=True` and `False` | +| KV cache | Enabled and disabled | +| Attention | Flash, Tokamax, and supported ring variants | +| Batch | 1 and 16, plus a non-divisible batch | +| AOT | Cold compile, save, warm load, and source invalidation | +| Checkpoint | Safetensors and `.bin` fallback | +| Masks | Text, a2v, and v2a with different per-example padding | + +Numerical comparisons should use dot-product attention as the reference and BF16-appropriate tolerances. + +## Merge Gates + +The branch should not merge until all of the following are true: + +- `git diff --check` passes. +- `ruff check .` passes. +- The full unit-test suite passes. +- `.bin` loading holds only one full checkpoint copy. +- AOT cache keys change with the source or cache ABI. +- Tile search succeeds with the shipped flash configuration. +- Every JAX process receives the same tile-search winner. +- Ring cross-attention matches a replicated reference. +- Text-encoder sharding is effective and handles non-divisible batches. +- Repeated generation calls leave VAE and pipeline state unchanged. +- No explicitly enabled optimization silently falls back after complete failure. + +## Existing Positive Changes Worth Preserving + +- Full-step-count warmup for `scan_diffusion_loop=True`. +- Rejection of KV caching with `cross_attn_mod=True`. +- Guidance-path deduplication. +- KV-cache parity coverage for scanned and non-scanned layers. +- Named-scope additions for profiling. + +These should remain intact while the correctness and reliability issues above are addressed. diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py index ba722160a..30e69a113 100644 --- a/src/maxdiffusion/aot_cache.py +++ b/src/maxdiffusion/aot_cache.py @@ -72,6 +72,11 @@ def transformer_forward_pass(...): _FORMAT_VERSION = 1 +def _metadata_fingerprint(meta: dict[str, Any]) -> str: + """Returns the stable filename fingerprint for install-time metadata.""" + return hashlib.sha256(repr(sorted(meta.items())).encode()).hexdigest()[:12] + + def _dynamic_signature(args: tuple, kwargs: dict) -> str: """Deterministic digest of everything that selects an executable. @@ -372,21 +377,30 @@ def install(cache_dir: str, meta: dict[str, Any], mesh: Any) -> None: mesh: The pipeline mesh; pins device order for deserialization and provides the context for re-lowering at save time. """ - if not cache_dir: - return - os.makedirs(cache_dir, exist_ok=True) - _STATE.cache_dir = cache_dir - _STATE.fingerprint = hashlib.sha256(repr(sorted(meta.items())).encode()).hexdigest()[:12] - _STATE.mesh = mesh - _STATE.enabled = True + # A process may construct multiple pipelines with different cache settings. + # Finish any previous loads, then reset all install-scoped state even when + # the new cache directory is empty. + wait_for_loads() + _STATE.cache_dir = "" + _STATE.fingerprint = "" + _STATE.mesh = None + _STATE.enabled = False for entry in _REGISTRY: with entry._lock: - # Cached state belongs to the previous install's dir/fingerprint. entry._compiled.clear() entry._out_specs.clear() entry._pending.clear() entry._adapters.clear() entry._on_disk.clear() + + if not cache_dir: + return + os.makedirs(cache_dir, exist_ok=True) + _STATE.cache_dir = cache_dir + _STATE.fingerprint = _metadata_fingerprint(meta) + _STATE.mesh = mesh + _STATE.enabled = True + for entry in _REGISTRY: thread = threading.Thread(target=entry.load_from_disk, name=f"aot-load-{entry.name}", daemon=True) thread.start() _LOAD_THREADS.append(thread) diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml index f9271ef12..e3360350a 100644 --- a/src/maxdiffusion/configs/ltx2_3_video.yml +++ b/src/maxdiffusion/configs/ltx2_3_video.yml @@ -1,7 +1,12 @@ #hardware hardware: 'tpu' skip_jax_distributed_system: False -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +# Supported attention kernels: +# dot_product, flash, tokamax_flash, tokamax_ring, tokamax_ring_custom, +# ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_ring, +# ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_bidir, +# and cudnn_flash_te (GPU only). +attention: 'flash' use_base2_exp: False use_experimental_scheduler: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. @@ -100,7 +105,7 @@ flash_min_seq_length: 4096 dcn_context_parallelism: 1 dcn_tensor_parallelism: 1 ici_data_parallelism: 1 -ici_fsdp_parallelism: 1 +ici_fsdp_parallelism: 1 ici_context_parallelism: -1 # recommended ICI axis to be auto-sharded ici_tensor_parallelism: 1 enable_profiler: False @@ -186,14 +191,14 @@ upsampler_rational_spatial_scale: 2.0 upsampler_output_type: "pil" aot_cache_dir: '' -converted_weights_dir: '' - - +# Immutable package/build revision used when Git metadata is unavailable. +aot_build_revision: '' # Tile-size auto-tuning. When enable_tile_search: True, generate_ltx2 runs a fast one-DiT-block -# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites -# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op). +# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and updates every +# effective flash_block_sizes field with the winner. Default off (no-op). enable_tile_search: False tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep) tile_search_iters: 10 tile_search_out: '' # dir for the results CSV; '' -> print only +tile_search_vmem_limit_bytes: 67108864 # 64 MiB; shared by benchmark and production custom kernels use_kv_cache: False diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml index 6be535010..58877d4a6 100644 --- a/src/maxdiffusion/configs/ltx2_video.yml +++ b/src/maxdiffusion/configs/ltx2_video.yml @@ -1,7 +1,12 @@ #hardware hardware: 'tpu' skip_jax_distributed_system: False -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +# Supported attention kernels: +# dot_product, flash, tokamax_flash, tokamax_ring, tokamax_ring_custom, +# ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_ring, +# ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_bidir, +# and cudnn_flash_te (GPU only). +attention: 'flash' use_base2_exp: False use_experimental_scheduler: False # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. @@ -105,7 +110,7 @@ flash_min_seq_length: 4096 dcn_context_parallelism: 1 dcn_tensor_parallelism: 1 ici_data_parallelism: 1 -ici_fsdp_parallelism: 1 +ici_fsdp_parallelism: 1 ici_context_parallelism: -1 # recommended ICI axis to be auto-sharded ici_tensor_parallelism: 1 enable_profiler: False @@ -183,14 +188,14 @@ upsampler_rational_spatial_scale: 2.0 upsampler_output_type: "pil" aot_cache_dir: '' -converted_weights_dir: '' - - +# Immutable package/build revision used when Git metadata is unavailable. +aot_build_revision: '' # Tile-size auto-tuning. When enable_tile_search: True, generate_ltx2 runs a fast one-DiT-block -# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites -# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op). +# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and updates every +# effective flash_block_sizes field with the winner. Default off (no-op). enable_tile_search: False tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep) tile_search_iters: 10 tile_search_out: '' # dir for the results CSV; '' -> print only +tile_search_vmem_limit_bytes: 67108864 # 64 MiB; shared by benchmark and production custom kernels use_kv_cache: False diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index 637a6ee4d..d2fa72111 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -19,6 +19,7 @@ import time import os import subprocess +import uuid from maxdiffusion.checkpointing.ltx2_checkpointer import LTX2Checkpointer from maxdiffusion import aot_cache, pyconfig, max_logging, max_utils from absl import app @@ -29,6 +30,9 @@ from maxdiffusion.loaders.ltx2_lora_nnx_loader import LTX2NNXLoraLoader +LTX2_AOT_CACHE_ABI = 1 + + def upload_video_to_gcs(output_dir: str, video_path: str): """ Uploads a local video file to a specified Google Cloud Storage bucket. @@ -67,12 +71,22 @@ def delete_file(file_path: str): def get_git_commit_hash(): - """Tries to get the current Git commit hash.""" + """Returns HEAD only when the tracked source tree is clean.""" + source_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) try: - commit_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("utf-8") + commit_hash = subprocess.check_output(["git", "-C", source_root, "rev-parse", "HEAD"]).strip().decode("utf-8") + tracked_changes = subprocess.check_output( + ["git", "-C", source_root, "status", "--porcelain", "--untracked-files=no"], + ).strip() + untracked_source = subprocess.check_output( + ["git", "-C", source_root, "ls-files", "--others", "--exclude-standard", "--", "src/maxdiffusion"], + ).strip() + if tracked_changes or untracked_source: + max_logging.log("Warning: Git source changes detected; persistent LTX2 AOT cache reuse is disabled.") + return f"dirty:{commit_hash}" return commit_hash except subprocess.CalledProcessError: - max_logging.log("Warning: 'git rev-parse HEAD' failed. Not running in a git repo?") + max_logging.log("Warning: unable to determine a clean Git revision. Not running in a Git repo?") return None except FileNotFoundError: max_logging.log("Warning: 'git' command not found.") @@ -115,18 +129,20 @@ def call_pipeline(config, pipeline, prompt, negative_prompt): def maybe_tune_block_sizes(config): - """If enable_tile_search, run a fast one-DiT-block tile-size grid search and overwrite - flash_block_sizes' block_q/block_kv/block_kv_compute with the winner IN PLACE. - """ - keys = config.get_keys() - val = keys.get("enable_tile_search", False) - if str(val).lower() not in ("true", "1", "yes"): + """Tunes and applies the exact block-size fields consumed by production inference.""" + if not _tile_search_enabled(config): return + keys = config.get_keys() from maxdiffusion.utils.tile_size_grid_search import grid_search from maxdiffusion.utils.ltx2_block_benchmark import LTX2BlockBenchmark mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes) - bench = LTX2BlockBenchmark.from_config(config, mesh) + vmem_limit_bytes = int( + keys.get("tile_search_vmem_limit_bytes") + or config.flash_block_sizes.get("vmem_limit_bytes") + or 64 * 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( bench, @@ -136,34 +152,71 @@ def maybe_tune_block_sizes(config): log=max_logging.log, ) if result.best is None: - max_logging.log("[tile-search] no config succeeded; keeping configured flash_block_sizes") - return - fbs = dict(config.flash_block_sizes) - fbs.update({ - "block_q": result.best.bq, - "block_kv": result.best.bkv, - "block_kv_compute": result.best.bkv_compute, - "block_kv_compute_in": result.best.bkv_compute, - }) + raise RuntimeError( + "[tile-search] tuning was explicitly enabled, but no candidate succeeded. " + "Inspect the per-candidate errors instead of running with an untuned configuration." + ) + fbs = max_utils.flash_block_sizes_for_candidate( + config.flash_block_sizes, + config.attention, + result.best.bq, + result.best.bkv, + result.best.bkv_compute, + vmem_limit_bytes=vmem_limit_bytes, + ) config.get_keys()["flash_block_sizes"] = fbs + effective_block_sizes = max_utils.get_flash_block_sizes(config) + if effective_block_sizes is None or effective_block_sizes.block_q != result.best.bq: + effective_bq = None if effective_block_sizes is None else effective_block_sizes.block_q + raise RuntimeError( + f"[tile-search] selected block_q={result.best.bq}, but production resolved block_q={effective_bq}." + ) max_logging.log( f"[tile-search] using block_q={result.best.bq} block_kv={result.best.bkv} " f"(block-bench {result.best.mean_ms:.2f} ms)" ) +def _tile_search_enabled(config) -> bool: + return str(config.get_keys().get("enable_tile_search", False)).lower() in ("true", "1", "yes") + + def _canonical_aot_value(value): return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str) -def ltx2_aot_metadata(config, pipeline): +def _non_reusable_aot_revision(): + """Returns a unique identity so unversioned source can never hit old HLO.""" + return f"unversioned:{uuid.uuid4().hex}" + + +def _resolve_ltx2_aot_source_revision(config, commit_hash=None): + """Prefers an explicit Git revision, then a packaged-build revision.""" + for revision in (commit_hash, getattr(config, "aot_build_revision", None)): + if revision is not None and str(revision).strip(): + return str(revision).strip() + return None + + +def _is_reusable_aot_revision(source_revision) -> bool: + if source_revision is None or not str(source_revision).strip(): + return False + return not str(source_revision).startswith(("dirty:", "unversioned:")) + + +def ltx2_aot_metadata(config, pipeline, source_revision=None): """Returns every graph- and topology-shaping input to LTX2 AOT caching. We deliberately exclude values such as model weights and RNG state: they are executable inputs, not compilation inputs. `use_kv_cache` is also excluded because it is a static argument of `run_diffusion_loop` and therefore part - of that executable's per-call signature. + of that executable's per-call signature. If no clean source/build revision + is available, a per-run identity prevents stale executable reuse. """ + source_revision = str(source_revision).strip() if source_revision is not None else "" + if not source_revision: + source_revision = _non_reusable_aot_revision() + transformer_config = dict(getattr(pipeline.transformer, "config", {})) for key in ( "rngs", @@ -189,6 +242,8 @@ def ltx2_aot_metadata(config, pipeline): device = jax.devices()[0] return { + "ltx2_aot_cache_abi": str(LTX2_AOT_CACHE_ABI), + "source_revision": source_revision, "model": config.pretrained_model_name_or_path, "transformer_architecture": _canonical_aot_value(transformer_config), "attention": getattr(config, "attention", ""), @@ -196,11 +251,16 @@ def ltx2_aot_metadata(config, pipeline): "v2a_attention_kernel": getattr(config, "v2a_attention_kernel", "dot_product"), "flash_block_sizes": _canonical_aot_value(getattr(config, "flash_block_sizes", {})), "flash_min_seq_length": str(getattr(config, "flash_min_seq_length", 4096)), + "use_base2_exp": str(getattr(config, "use_base2_exp", False)), + "use_experimental_scheduler": str(getattr(config, "use_experimental_scheduler", False)), "ulysses_shards": str(getattr(config, "ulysses_shards", -1)), "ulysses_attention_chunks": str(getattr(config, "ulysses_attention_chunks", 1)), "scan_layers": str(getattr(config, "scan_layers", True)), "scan_diffusion_loop": str(getattr(config, "scan_diffusion_loop", True)), + "precision": str(getattr(config, "precision", None)), "remat_policy": str(getattr(config, "remat_policy", "NONE")), + "names_which_can_be_saved": _canonical_aot_value(getattr(config, "names_which_can_be_saved", ())), + "names_which_can_be_offloaded": _canonical_aot_value(getattr(config, "names_which_can_be_offloaded", ())), "spatio_temporal_guidance_blocks": _canonical_aot_value( getattr(config, "spatio_temporal_guidance_blocks", ()) ), @@ -220,6 +280,11 @@ def ltx2_aot_metadata(config, pipeline): def run(config, pipeline=None, filename_prefix="", commit_hash=None): + if pipeline is not None and _tile_search_enabled(config): + raise ValueError( + "enable_tile_search cannot be used with a prebuilt pipeline because block sizes are static in its graph. " + "Tune before constructing the pipeline, or call run without the pipeline argument." + ) if pipeline is None: maybe_tune_block_sizes(config) @@ -299,9 +364,20 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): original_num_steps = config.get_keys().get("num_inference_steps", 40) # Per-shape AOT executable cache + detected_revision = commit_hash if commit_hash is not None else get_git_commit_hash() + source_revision = _resolve_ltx2_aot_source_revision(config, detected_revision) + aot_cache_dir = getattr(config, "aot_cache_dir", "") + if aot_cache_dir and not _is_reusable_aot_revision(source_revision): + max_logging.log( + "[aot] No clean Git commit or aot_build_revision was supplied; " + "persistent LTX2 AOT caching is disabled for this run." + ) + aot_cache_dir = "" + aot_metadata = ltx2_aot_metadata(config, pipeline, source_revision=source_revision) + max_logging.log(f"[aot] LTX2 cache metadata: {_canonical_aot_value(aot_metadata)}") aot_cache.install( - getattr(config, "aot_cache_dir", ""), - meta=ltx2_aot_metadata(config, pipeline), + aot_cache_dir, + meta=aot_metadata, mesh=pipeline.mesh, ) aot_cache.wait_for_loads() diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 1ab5407ac..2c157406d 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -633,6 +633,46 @@ class CustomFlashBlockSizes: vmem_limit_bytes: int | None = None +def flash_block_sizes_for_candidate( + flash_block_sizes: Dict[str, int], + attention: str, + block_q: int, + block_kv: int, + block_kv_compute: int | None = None, + *, + vmem_limit_bytes: int | None = None, +) -> Dict[str, int]: + """Returns a production-compatible block-size mapping for a tuned candidate. + + Splash uses forward and backward block-size fields even for inference-time + padding decisions. Tokamax also derives its effective forward ``block_q`` + from ``block_q_dkv``. Updating only ``block_q`` therefore does not faithfully + apply or benchmark a candidate. + """ + block_kv_compute = block_kv if block_kv_compute is None else block_kv_compute + candidate = dict(flash_block_sizes) + candidate.update({ + "block_q": block_q, + "block_kv": block_kv, + "block_kv_compute": block_kv_compute, + "block_kv_compute_in": block_kv_compute, + "block_q_dkv": block_q, + "block_kv_dkv": block_kv, + "block_kv_dkv_compute": block_kv_compute, + }) + + if not candidate.get("use_fused_bwd_kernel", False): + candidate["block_q_dq"] = block_q + candidate["block_kv_dq"] = block_kv + + if "custom" in attention: + candidate.setdefault("heads_per_tile", 1) + if vmem_limit_bytes is not None: + candidate["vmem_limit_bytes"] = vmem_limit_bytes + + return candidate + + def get_flash_block_sizes(config): """Create custom flash attention BlockSizes.""" flash_block_sizes = None diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index a76da7dea..afcddb22c 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -287,6 +287,13 @@ def _select_flash_block_sizes( else: kv_max_block_size = q_max_block_size + # Custom kernels use a lightweight carrier that omits the standard Splash + # backward fields. A remapped/local standard kernel still needs a complete + # BlockSizes object, including when cross-attention happens to have q_len == + # kv_len. + if flash_block_sizes is not None and not hasattr(flash_block_sizes, "use_fused_bwd_kernel"): + flash_block_sizes = _coerce_tokamax_block_sizes(flash_block_sizes) + # Keep configured block sizes for self-attention, but let # cross-attention derive safe KV-aware sizes when q_len != kv_len. if flash_block_sizes and key_seq_len == query_seq_len: @@ -395,9 +402,9 @@ def _build_padding_segment_ids( """Build splash segment ids that mask q/kv padding and the attention mask. Padding tokens get segment id 0, valid tokens 1. An optional attention_mask - (batch, kv_len) is folded into the kv segment ids; positions beyond the mask - but within key_seq_len default to valid, and positions beyond key_seq_len are - padding. Shared by flash, ulysses, and ulysses+ring kernels. + (batch, kv_len) is folded into the kv segment ids. Positions beyond an + explicit mask are invalid, including sequence padding introduced before the + local flash kernel. Shared by flash, Ulysses, and Ulysses+ring kernels. """ q_indices = jax.lax.broadcasted_iota(jnp.int32, (q_padded_len,), 0) q_segment_ids = (q_indices < query_seq_len).astype(jnp.int32) @@ -406,19 +413,72 @@ def _build_padding_segment_ids( kv_segment_ids = (kv_indices < key_seq_len).astype(jnp.int32) if attention_mask is not None: + if attention_mask.ndim != 2: + raise ValueError(f"attention_mask must have shape [batch, kv_length], got {attention_mask.shape}.") mask_len = min(key_seq_len, attention_mask.shape[1]) - kv_mask_for_batch = attention_mask[0, :mask_len] - # Tokens past the mask but within key_seq_len are assumed valid. + kv_mask_for_batch = attention_mask[:, :mask_len].astype(jnp.int32) + # An explicit mask is authoritative. This also masks sequence padding + # introduced before shard_map to make the KV length context-divisible. if key_seq_len > mask_len: - kv_mask_for_batch = jnp.concatenate([kv_mask_for_batch, jnp.ones((key_seq_len - mask_len,), jnp.int32)], axis=0) + kv_mask_for_batch = jnp.concatenate( + [ + kv_mask_for_batch, + jnp.zeros((attention_mask.shape[0], key_seq_len - mask_len), jnp.int32), + ], + axis=1, + ) # Tokens past key_seq_len are padding. if kv_padded_len > key_seq_len: - kv_mask_for_batch = jnp.concatenate([kv_mask_for_batch, jnp.zeros((kv_padded_len - key_seq_len,), jnp.int32)], axis=0) - kv_segment_ids = (kv_segment_ids * kv_mask_for_batch).astype(jnp.int32) + kv_mask_for_batch = jnp.concatenate( + [ + kv_mask_for_batch, + jnp.zeros((attention_mask.shape[0], kv_padded_len - key_seq_len), jnp.int32), + ], + axis=1, + ) + kv_segment_ids = (kv_segment_ids[None, :] * kv_mask_for_batch).astype(jnp.int32) + q_segment_ids = jnp.broadcast_to(q_segment_ids[None, :], (attention_mask.shape[0], q_padded_len)) return segment_ids_cls(q=q_segment_ids, kv=kv_segment_ids) +def _prepare_attention_mask_for_shard_map( + attention_mask: jax.Array | None, + batch_size: int, + padded_kv_len: int, +) -> jax.Array | None: + """Broadcasts and pads a canonical keep mask before entering shard_map.""" + if attention_mask is None: + return None + if attention_mask.ndim != 2: + raise ValueError(f"attention_mask must have shape [batch, kv_length], got {attention_mask.shape}.") + if attention_mask.shape[0] == 1 and batch_size != 1: + attention_mask = jnp.broadcast_to(attention_mask, (batch_size, attention_mask.shape[1])) + elif attention_mask.shape[0] != batch_size: + raise ValueError( + f"attention_mask batch dimension must be 1 or match the attention batch ({batch_size}), " + f"got {attention_mask.shape[0]}." + ) + + attention_mask = attention_mask.astype(jnp.bool_) + if attention_mask.shape[1] < padded_kv_len: + attention_mask = jnp.pad( + attention_mask, + ((0, 0), (0, padded_kv_len - attention_mask.shape[1])), + constant_values=False, + ) + elif attention_mask.shape[1] > padded_kv_len: + attention_mask = attention_mask[:, :padded_kv_len] + return attention_mask + + +def _mesh_axis_in_spec(axis_spec, mesh_axis: str) -> bool: + """Returns whether a logical-to-mesh axis entry contains mesh_axis.""" + if isinstance(axis_spec, tuple): + return mesh_axis in axis_spec + return axis_spec == mesh_axis + + def _ulysses_head_chunk_ranges(num_heads: int, ulysses_shards: int, num_chunks: int): """Build head-axis ranges for chunked Ulysses all-to-all. @@ -522,19 +582,16 @@ def _tpu_flash_attention( query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards) key, _ = _reshape_data_for_flash(key, heads, num_context_shards) value, _ = _reshape_data_for_flash(value, heads, num_context_shards) + attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) + if attention_mask is not None and attention_kernel == "tokamax_ring_custom": + raise NotImplementedError("tokamax_ring_custom does not support attention_mask.") block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, attention_kernel) q_axis_names = nn.logical_to_mesh_axes(axis_names_q) kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv) + mask_axis_names = nn.logical_to_mesh_axes((axis_names_kv[0], axis_names_kv[2])) - @functools.partial( - shard_map.shard_map, - mesh=mesh, - in_specs=(q_axis_names, kv_axis_names, kv_axis_names), - out_specs=q_axis_names, - check_rep=False, - ) - def wrap_flash_attention(query, key, value): + def wrap_flash_attention(query, key, value, attention_mask): if attention_kernel == "tokamax_ring_custom": # Ring attention backed by the custom dense splash kernel. q stays local, # k/v rotate over the "context" axis (handled inside the ring kernel). @@ -627,9 +684,9 @@ def wrap_flash_attention(query, key, value): ), save_residuals=False, ring_axis=CONTEXT, - # We don't rotate segment ids in tokamax ring attention because our - # segment ids is for padding each kv shard has same segment ids - rotate_segment_ids=False, + # Padding-only IDs are identical on each shard. Explicit masks differ + # by KV shard and must rotate together with K/V. + rotate_segment_ids=attention_mask is not None, ) else: splash_kernel = splash_attention_kernel.make_splash_mha( @@ -641,9 +698,10 @@ def wrap_flash_attention(query, key, value): residual_checkpoint_name=residual_checkpoint_name, ) - vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None)) + segment_ids_in_axes = 0 if attention_mask is not None else None + vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, segment_ids_in_axes)) - if not mask_padding_tokens: + if not mask_padding_tokens and attention_mask is None: segment_ids = None if attention_kernel in ["flash", "tokamax_flash", "tokamax_ring"]: attention_output = vmapped_splash(query, key, value, segment_ids) @@ -700,7 +758,24 @@ def ring_scan_body(carry, _): "Warning, batch dimension should be shardable among the devices in data and fsdp" f" axis, batch dimension: {query.shape[0]}, devices_in_batch_sharding: {devices_in_batch_sharding}" ) - x = wrap_flash_attention(query, key, value) + if attention_mask is None: + sharded_flash_attention = shard_map.shard_map( + lambda q, k, v: wrap_flash_attention(q, k, v, None), + mesh=mesh, + in_specs=(q_axis_names, kv_axis_names, kv_axis_names), + out_specs=q_axis_names, + check_rep=False, + ) + x = sharded_flash_attention(query, key, value) + else: + sharded_flash_attention = shard_map.shard_map( + wrap_flash_attention, + mesh=mesh, + in_specs=(q_axis_names, kv_axis_names, kv_axis_names, mask_axis_names), + out_specs=q_axis_names, + check_rep=False, + ) + x = sharded_flash_attention(query, key, value, attention_mask) # Trim back to original sequence length after context-axis padding. x = x[:, :, :orig_q_seq_len, :] x = _reshape_heads_to_head_dim(x) @@ -745,6 +820,12 @@ def _ulysses_attention( query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_shards) key, _ = _reshape_data_for_flash(key, heads, num_shards) value, _ = _reshape_data_for_flash(value, heads, num_shards) + attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) + if attention_mask is not None and use_custom_kernel: + raise NotImplementedError( + "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " + "(it only handles padding via orig_seq_len); got a non-None attention_mask." + ) num_heads = query.shape[1] # Ulysses only redistributes existing heads across the context mesh; unlike # the earlier draft, we fail fast instead of padding synthetic heads. @@ -759,27 +840,19 @@ def _ulysses_attention( q_axis_names = nn.logical_to_mesh_axes(axis_names_q) kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv) + mask_axis_names = nn.logical_to_mesh_axes((axis_names_kv[0], axis_names_kv[2])) + mask_needs_ulysses_gather = _mesh_axis_in_spec(kv_axis_names[2], axis_name) - @functools.partial( - jax.shard_map, - mesh=mesh, - in_specs=(q_axis_names, kv_axis_names, kv_axis_names), - out_specs=q_axis_names, - check_vma=False, - ) - def wrap_ulysses_attention(query, key, value): + def wrap_ulysses_attention(query, key, value, attention_mask): # Swap sharding: each device gives up a slice of heads and gathers # a slice of sequence, so the local kernel sees the full sequence. query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) key = jax.lax.all_to_all(key, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) value = jax.lax.all_to_all(value, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) + if attention_mask is not None and mask_needs_ulysses_gather: + attention_mask = jax.lax.all_gather(attention_mask, axis_name=axis_name, axis=1, tiled=True) if use_custom_kernel: - if attention_mask is not None: - raise NotImplementedError( - "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " - "(it only handles padding via orig_seq_len); got a non-None attention_mask." - ) bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) if use_base2_exp: @@ -853,7 +926,7 @@ def wrap_ulysses_attention(query, key, value): multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1]) segment_ids = _build_padding_segment_ids(query_seq_len, query.shape[2], key_seq_len, key.shape[2], attention_mask) - if not mask_padding_tokens: + if not mask_padding_tokens and attention_mask is None: segment_ids = None splash_kernel = splash_attention_kernel.make_splash_mha( @@ -864,7 +937,8 @@ def wrap_ulysses_attention(query, key, value): save_residuals=False, residual_checkpoint_name=residual_checkpoint_name, ) - vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None)) + segment_ids_in_axes = 0 if attention_mask is not None else None + vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, segment_ids_in_axes)) attention_output = vmapped_splash(query, key, value, segment_ids) attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) @@ -887,7 +961,9 @@ def wrap_ulysses_attention(query, key, value): # a2a tensors inside the scanned layers (measured 7.0 -> expected ~3.5 # s/step at 720p 81f cp8 CFG). batch = query.shape[0] - fold_batch = batch > 1 and (batch * num_heads) % num_shards == 0 + # Folding batch into heads destroys the one-mask-per-example association. + # Keep the optimization for the common unmasked path only. + fold_batch = attention_mask is None and batch > 1 and (batch * num_heads) % num_shards == 0 if fold_batch: query = query.reshape(1, batch * num_heads, *query.shape[2:]) key = key.reshape(1, batch * num_heads, *key.shape[2:]) @@ -896,6 +972,30 @@ def wrap_ulysses_attention(query, key, value): else: effective_num_heads = num_heads + if attention_mask is None: + sharded_ulysses_attention = jax.shard_map( + lambda q, k, v: wrap_ulysses_attention(q, k, v, None), + mesh=mesh, + in_specs=(q_axis_names, kv_axis_names, kv_axis_names), + out_specs=q_axis_names, + check_vma=False, + ) + + def run_ulysses_attention(q, k, v): + return sharded_ulysses_attention(q, k, v) + + else: + sharded_ulysses_attention = jax.shard_map( + wrap_ulysses_attention, + mesh=mesh, + in_specs=(q_axis_names, kv_axis_names, kv_axis_names, mask_axis_names), + out_specs=q_axis_names, + check_vma=False, + ) + + def run_ulysses_attention(q, k, v): + return sharded_ulysses_attention(q, k, v, attention_mask) + x = _run_chunked_ulysses_attention( query, key, @@ -903,7 +1003,7 @@ def wrap_ulysses_attention(query, key, value): effective_num_heads, num_shards, ulysses_attention_chunks, - wrap_ulysses_attention, + run_ulysses_attention, ) if fold_batch: @@ -974,6 +1074,7 @@ def _ulysses_ring_attention( query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_sequence_shards) key, _ = _reshape_data_for_flash(key, heads, num_sequence_shards) value, _ = _reshape_data_for_flash(value, heads, num_sequence_shards) + attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) num_heads = query.shape[1] block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, "tokamax_ring") @@ -982,20 +1083,18 @@ def _ulysses_ring_attention( kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv) internal_q_axis_names = _replace_mesh_axis_names(q_axis_names, context_axis, internal_sequence_axes) internal_kv_axis_names = _replace_mesh_axis_names(kv_axis_names, context_axis, internal_sequence_axes) + mask_axis_names = nn.logical_to_mesh_axes((axis_names_kv[0], axis_names_kv[2])) + internal_mask_axis_names = _replace_mesh_axis_names(mask_axis_names, context_axis, internal_sequence_axes) + mask_needs_ulysses_gather = _mesh_axis_in_spec(internal_mask_axis_names[1], ulysses_axis) - @functools.partial( - jax.shard_map, - mesh=internal_mesh, - in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names), - out_specs=internal_q_axis_names, - check_vma=False, - ) - def wrap_ulysses_ring_attention(query, key, value): + def wrap_ulysses_ring_attention(query, key, value, attention_mask): # Swap sharding: each device gives up a slice of heads and gathers # a slice of sequence, so the local kernel sees the full sequence. query = jax.lax.all_to_all(query, axis_name=ulysses_axis, split_axis=1, concat_axis=2, tiled=True) key = jax.lax.all_to_all(key, axis_name=ulysses_axis, split_axis=1, concat_axis=2, tiled=True) value = jax.lax.all_to_all(value, axis_name=ulysses_axis, split_axis=1, concat_axis=2, tiled=True) + if attention_mask is not None and mask_needs_ulysses_gather: + attention_mask = jax.lax.all_gather(attention_mask, axis_name=ulysses_axis, axis=1, tiled=True) uses_fused_kernel = block_sizes.use_fused_bwd_kernel block_q_sizes = (block_sizes.block_q, block_sizes.block_q_dkv) @@ -1017,14 +1116,13 @@ def wrap_ulysses_ring_attention(query, key, value): kv_padded_len = key.shape[2] total_kv_len = kv_padded_len * num_ring_shards - # Mask q/kv padding via segment ids, same as the tokamax_ring kernel. Each - # ring shard pads identically so every shard shares the same per-shard ids - # and rotation is unneeded. + # Mask q/kv padding via segment ids, same as the tokamax_ring kernel. + # Padding-only IDs are identical per shard; explicit KV masks rotate. segment_ids = _build_padding_segment_ids( query_seq_len, q_padded_len, key_seq_len, kv_padded_len, attention_mask, tokamax_splash_base.SegmentIds ) - if not mask_padding_tokens: + if not mask_padding_tokens and attention_mask is None: segment_ids = None mask = tokamax_splash_attention_mask.FullMask(_shape=(q_padded_len, total_kv_len)) @@ -1041,9 +1139,10 @@ def wrap_ulysses_ring_attention(query, key, value): save_residuals=False, ring_axis=ring_axis, kv_seq_shards=num_ring_shards, - rotate_segment_ids=False, + rotate_segment_ids=attention_mask is not None, ) - vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None)) + segment_ids_in_axes = 0 if attention_mask is not None else None + vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, segment_ids_in_axes)) attention_output = vmapped_splash(query, key, value, segment_ids) attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) @@ -1063,6 +1162,30 @@ def wrap_ulysses_ring_attention(query, key, value): "Warning, batch dimension should be shardable among the devices in data and fsdp" f" axis, batch dimension: {query.shape[0]}, devices_in_batch_sharding: {devices_in_batch_sharding}" ) + if attention_mask is None: + sharded_ulysses_ring_attention = jax.shard_map( + lambda q, k, v: wrap_ulysses_ring_attention(q, k, v, None), + mesh=internal_mesh, + in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names), + out_specs=internal_q_axis_names, + check_vma=False, + ) + + def run_ulysses_ring_attention(q, k, v): + return sharded_ulysses_ring_attention(q, k, v) + + else: + sharded_ulysses_ring_attention = jax.shard_map( + wrap_ulysses_ring_attention, + mesh=internal_mesh, + in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names, internal_mask_axis_names), + out_specs=internal_q_axis_names, + check_vma=False, + ) + + def run_ulysses_ring_attention(q, k, v): + return sharded_ulysses_ring_attention(q, k, v, attention_mask) + x = _run_chunked_ulysses_attention( query, key, @@ -1070,7 +1193,7 @@ def wrap_ulysses_ring_attention(query, key, value): num_heads, num_ulysses_shards, ulysses_attention_chunks, - wrap_ulysses_ring_attention, + run_ulysses_ring_attention, ) x = jax.lax.with_sharding_constraint(x, q_axis_names) x = x[:, :, :orig_q_seq_len, :] @@ -1274,6 +1397,7 @@ def _apply_attention_dot( split_head_dim: bool, float32_qk_product: bool, use_memory_efficient_attention: bool, + attention_mask: Array = None, ): """Apply Attention.""" if split_head_dim: @@ -1290,7 +1414,7 @@ def _apply_attention_dot( query_states = query_states.astype(jnp.float32) key_states = key_states.astype(jnp.float32) - if use_memory_efficient_attention: + if use_memory_efficient_attention and attention_mask is None: query_states = query_states.transpose(1, 0, 2) key_states = key_states.transpose(1, 0, 2) value_states = value_states.transpose(1, 0, 2) @@ -1324,7 +1448,12 @@ def _apply_attention_dot( attention_scores = jnp.einsum("b i d, b j d->b i j", query_states, key_states) attention_scores = attention_scores * scale + if attention_mask is not None: + attention_scores = attention_scores + attention_mask.astype(attention_scores.dtype) attention_probs = nn.softmax(attention_scores, axis=-1 if split_head_dim else 2) + if attention_mask is not None: + has_valid_key = jnp.any(attention_mask == 0, axis=-1, keepdims=True) + attention_probs = jnp.where(has_valid_key, attention_probs, 0) attention_probs = attention_probs.astype(dtype) @@ -1385,6 +1514,7 @@ def dot_product_kernel(q, k, v, context): context["split_head_dim"], context["float32_qk_product"], context["use_memory_efficient_attention"], + context["attention_mask"], ) @@ -1604,7 +1734,10 @@ def tokamax_ring_kernel(q, k, v, context): context["dtype"], attention_kernel="tokamax_ring", mask_padding_tokens=context["mask_padding_tokens"], + residual_checkpoint_name=context["residual_checkpoint_name"], attention_mask=context["attention_mask"], + use_base2_exp=context["use_base2_exp"], + use_experimental_scheduler=context["use_experimental_scheduler"], ) @@ -1674,7 +1807,28 @@ def _apply_attention( and value.shape[seq_len_idx] >= flash_min_seq_length ) - # Fallback logic + effective_attention_kernel = attention_kernel + if attention_kernel == "dot_product" or use_memory_efficient_attention or not can_use_flash_attention: + effective_attention_kernel = "dot_product" + + # Masks enter the dispatcher as canonical [B, K] keep masks. Adapt them + # only after fallback selection because a configured flash kernel may use + # dot-product attention for short sequences. + if attention_mask is not None: + if attention_mask.ndim != 2: + raise ValueError(f"attention_mask must have shape [batch, kv_length], got {attention_mask.shape}.") + attention_mask = attention_mask.astype(jnp.bool_) + if effective_attention_kernel == "dot_product": + attention_bias = jnp.where( + attention_mask, + jnp.asarray(0.0, dtype=dtype), + jnp.asarray(-10000.0, dtype=dtype), + ) + if split_head_dim: + attention_mask = attention_bias[:, None, None, :] + else: + attention_mask = jnp.repeat(attention_bias, heads, axis=0)[:, None, :] + context = { "heads": heads, "mesh": mesh, @@ -1697,14 +1851,11 @@ def _apply_attention( "dpa_layer": dpa_layer, } - if attention_kernel == "dot_product" or use_memory_efficient_attention or not can_use_flash_attention: - return KERNEL_REGISTRY["dot_product"](query, key, value, context) - # Module-level Registry lookup - if attention_kernel in KERNEL_REGISTRY: - return KERNEL_REGISTRY[attention_kernel](query, key, value, context) + if effective_attention_kernel in KERNEL_REGISTRY: + return KERNEL_REGISTRY[effective_attention_kernel](query, key, value, context) - raise ValueError(f"Unexpected attention kernel {attention_kernel=}.") + raise ValueError(f"Unexpected attention kernel {effective_attention_kernel=}.") def _query_chunk_attention(query, key, value, precision, key_chunk_size: int = 4096): @@ -2133,16 +2284,36 @@ def __init__( self.out_axis_names = out_axis_names self.enable_jax_named_scopes = enable_jax_named_scopes + cross_attention_remapped_to_flash = ( + not is_self_attention + and attention_kernel + in ( + "tokamax_ring", + "tokamax_ring_custom", + "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_bidir", + "ulysses_custom", + "ulysses_custom_fixed_m", + ) + ) + cross_attention_uses_local_kv = not is_self_attention and ( + cross_attention_remapped_to_flash or attention_kernel in ("flash", "tokamax_flash", "cudnn_flash_te") + ) if is_self_attention: axis_names_q = (BATCH, SELF_ATTN_HEAD, SELF_ATTN_Q_LENGTH, D_KV) axis_names_kv = (BATCH, SELF_ATTN_HEAD, SELF_ATTN_KV_LENGTH, D_KV) else: axis_names_q = (BATCH, CROSS_ATTN_HEAD, CROSS_ATTN_Q_LENGTH, D_KV) - axis_names_kv = (BATCH, CROSS_ATTN_HEAD, CROSS_ATTN_KV_LENGTH, D_KV) - if attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring") and not is_self_attention: - attention_kernel = "tokamax_flash" # do not use ring attention for cross attention - if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir") and not is_self_attention: - attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention + axis_names_kv = ( + BATCH, + CROSS_ATTN_HEAD, + None if cross_attention_uses_local_kv else CROSS_ATTN_KV_LENGTH, + D_KV, + ) + if cross_attention_remapped_to_flash: + attention_kernel = "tokamax_flash" self.added_kv_proj_dim = added_kv_proj_dim # New for I2V self.image_seq_len = image_seq_len # New for I2V tpu_type = get_tpu_type() diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py index 601999cf4..71ed4325d 100644 --- a/src/maxdiffusion/models/ltx2/attention_ltx2.py +++ b/src/maxdiffusion/models/ltx2/attention_ltx2.py @@ -347,6 +347,8 @@ def __init__( gated_attn: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, + use_experimental_scheduler: bool = False, enable_jax_named_scopes: bool = False, ): self.heads = heads @@ -455,15 +457,22 @@ def __init__( ) is_self_attention = context_dim is None + cross_attention_remapped_to_flash = not is_self_attention and attention_kernel in ( + "tokamax_ring", + "tokamax_ring_custom", + "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_bidir", + "ulysses_custom", + "ulysses_custom_fixed_m", + ) + cross_attention_uses_local_kv = not is_self_attention and ( + cross_attention_remapped_to_flash or attention_kernel in ("flash", "tokamax_flash", "cudnn_flash_te") + ) if not is_self_attention: - if attention_kernel in ( - "tokamax_ring", - "tokamax_ring_custom", - "ulysses_ring", - ): - attention_kernel = "tokamax_flash" # do not use ring attention for cross attention - if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir"): - attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention + if cross_attention_remapped_to_flash: + attention_kernel = "tokamax_flash" axis_names_q = ( common_types.BATCH, @@ -474,7 +483,7 @@ def __init__( axis_names_kv = ( common_types.BATCH, common_types.CROSS_ATTN_HEAD, - common_types.CROSS_ATTN_KV_LENGTH, + None if cross_attention_uses_local_kv else common_types.CROSS_ATTN_KV_LENGTH, common_types.D_KV, ) else: @@ -504,6 +513,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, ) def compute_kv(self, context: Array) -> Tuple[Array, Array]: diff --git a/src/maxdiffusion/models/ltx2/ltx2_utils.py b/src/maxdiffusion/models/ltx2/ltx2_utils.py index 31819cecc..21e96a006 100644 --- a/src/maxdiffusion/models/ltx2/ltx2_utils.py +++ b/src/maxdiffusion/models/ltx2/ltx2_utils.py @@ -193,12 +193,45 @@ def _torch_tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray: return tensor.numpy() +def _get_eval_shape(value) -> tuple[int, ...]: + if hasattr(value, "shape"): + return tuple(value.shape) + if hasattr(value, "value") and hasattr(value.value, "shape"): + return tuple(value.value.shape) + raise ValueError(f"Unable to determine the initialized shape for {type(value).__name__}") + + +def _get_scanned_layer_shapes(flattened_eval_shapes): + scanned_layer_shapes = {} + for key, value in flattened_eval_shapes.items(): + normalized_key = _tuple_str_to_int(tuple(str(item) for item in key)) + if not normalized_key or normalized_key[0] != "transformer_blocks": + continue + + shape = _get_eval_shape(value) + if not shape: + raise ValueError(f"Scanned parameter {'.'.join(map(str, normalized_key))} has no layer axis") + scanned_layer_shapes[normalized_key] = shape + + if not scanned_layer_shapes: + raise ValueError("scan_layers=True, but eval_shapes contains no transformer_blocks parameters") + + layer_counts = {shape[0] for shape in scanned_layer_shapes.values()} + if len(layer_counts) != 1: + raise ValueError(f"Inconsistent scanned layer counts in eval_shapes: {sorted(layer_counts)}") + + scanned_num_layers = next(iter(layer_counts)) + if scanned_num_layers <= 0: + raise ValueError(f"Invalid scanned layer count derived from eval_shapes: {scanned_num_layers}") + return scanned_num_layers, scanned_layer_shapes + + def load_transformer_weights( pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True, - num_layers: int = 48, + num_layers: Optional[int] = None, scan_layers: bool = True, subfolder: str = "transformer", ): @@ -209,6 +242,20 @@ def load_transformer_weights( device = jax.local_devices(backend=device)[0] max_logging.log(f"Load and port {pretrained_model_name_or_path} {subfolder} on {device}") + flattened_dict = flatten_dict(eval_shapes) + random_flax_state_dict = {} + for key in flattened_dict: + random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key] + + scanned_num_layers = None + scanned_layer_shapes = {} + if scan_layers: + scanned_num_layers, scanned_layer_shapes = _get_scanned_layer_shapes(flattened_dict) + if num_layers is not None and num_layers != scanned_num_layers: + raise ValueError( + f"num_layers={num_layers} does not match the {scanned_num_layers} layers derived from eval_shapes" + ) + index_file = "diffusion_pytorch_model.safetensors.index.json" try: index_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=index_file) @@ -231,26 +278,22 @@ def resolve_shard_path(model_file): t_start = time.perf_counter() - flattened_dict = flatten_dict(eval_shapes) - random_flax_state_dict = {} - for key in flattened_dict: - random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key] - flax_state_dict = {} + populated_layer_indices = {} dict_lock = threading.Lock() - def convert_chunk(ckpt_shard_path, chunk_keys): - if ckpt_shard_path.endswith(".safetensors"): - with safe_open(ckpt_shard_path, framework="pt") as f: - for pt_key in chunk_keys: - tensor = _torch_tensor_to_numpy(f.get_tensor(pt_key)) - process_tensor(pt_key, tensor) - else: - loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu") + def convert_safetensors_chunk(ckpt_shard_path, chunk_keys): + with safe_open(ckpt_shard_path, framework="pt") as f: for pt_key in chunk_keys: - tensor = _torch_tensor_to_numpy(loaded_state_dict[pt_key]) + tensor = _torch_tensor_to_numpy(f.get_tensor(pt_key)) process_tensor(pt_key, tensor) + def convert_bin(ckpt_shard_path): + loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu") + for pt_key, pt_tensor in loaded_state_dict.items(): + tensor = _torch_tensor_to_numpy(pt_tensor) + process_tensor(pt_key, tensor) + def process_tensor(pt_key, tensor): renamed_pt_key = rename_key(pt_key) renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key) @@ -274,48 +317,78 @@ def process_tensor(pt_key, tensor): if "scale_shift_table" in flax_key_str and flax_key_str[-1] in ["kernel", "weight"]: flax_key_str.pop() flax_key = tuple(flax_key_str) - # Convert string indices back to integers only if they were natively numeric keys - # But wait, nnx.List expects strings. _tuple_str_to_int converts '0' to 0. + # Convert string indices back to integers only if they were natively numeric keys + # But wait, nnx.List expects strings. _tuple_str_to_int converts '0' to 0. # If the logical state sharding has '0', _tuple_str_to_int will break the match! # Let's conditionally skip _tuple_str_to_int for nnx.List indices. - - # Actually, Flax traverse_util flattens nnx.List keys to strings in nnx, but _tuple_str_to_int + + # Actually, Flax traverse_util flattens nnx.List keys to strings in nnx, but _tuple_str_to_int # makes them integers. Let's just use _tuple_str_to_int but format it back to string if needed. flax_key = _tuple_str_to_int(flax_key) - + # We must ensure that if scan_layers=False, the index is a string to match NNX List! if not scan_layers and len(flax_key) > 1 and flax_key[0] == "transformer_blocks" and isinstance(flax_key[1], int): flax_key = ("transformer_blocks", str(flax_key[1])) + flax_key[2:] if block_index is not None: + key_name = ".".join(map(str, flax_key)) + if block_index < 0 or block_index >= scanned_num_layers: + raise ValueError( + f"Scanned layer index {block_index} for {key_name} is out of range [0, {scanned_num_layers})" + ) + if flax_key not in scanned_layer_shapes: + raise ValueError(f"Checkpoint tensor {pt_key!r} maps to unexpected scanned parameter {key_name}") + + expected_shape = scanned_layer_shapes[flax_key] + if tuple(flax_tensor.shape) != expected_shape[1:]: + raise ValueError( + f"Shape mismatch for scanned parameter {key_name}: " + f"expected per-layer shape {expected_shape[1:]}, got {tuple(flax_tensor.shape)}" + ) + with dict_lock: + populated = populated_layer_indices.setdefault(flax_key, set()) + if block_index in populated: + raise ValueError(f"Duplicate scanned layer index {block_index} for {key_name}") + stacked = flax_state_dict.get(flax_key) if stacked is None: - stacked = np.empty((num_layers,) + flax_tensor.shape, dtype=flax_tensor.dtype) + stacked = np.zeros(expected_shape, dtype=flax_tensor.dtype) flax_state_dict[flax_key] = stacked - stacked[block_index] = flax_tensor + stacked[block_index] = flax_tensor + populated.add(block_index) else: value = np.array(flax_tensor, dtype=flax_tensor.dtype, copy=True, order="C") with dict_lock: flax_state_dict[flax_key] = value chunk_size = 32 - tasks = [] + safetensors_tasks = [] for model_file in shards: ckpt_shard_path = resolve_shard_path(model_file) if ckpt_shard_path.endswith(".safetensors"): with safe_open(ckpt_shard_path, framework="pt") as f: shard_keys = list(f.keys()) + for i in range(0, len(shard_keys), chunk_size): + safetensors_tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) else: - loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu") - shard_keys = list(loaded_state_dict.keys()) - for i in range(0, len(shard_keys), chunk_size): - tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) - - with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: - futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks] - for future in concurrent.futures.as_completed(futures): - future.result() + convert_bin(ckpt_shard_path) + + if safetensors_tasks: + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: + futures = [executor.submit(convert_safetensors_chunk, path, keys) for path, keys in safetensors_tasks] + for future in concurrent.futures.as_completed(futures): + future.result() + + if scan_layers: + expected_indices = set(range(scanned_num_layers)) + missing_layers = [] + for flax_key in scanned_layer_shapes: + missing = sorted(expected_indices - populated_layer_indices.get(flax_key, set())) + if missing: + missing_layers.append(f"{'.'.join(map(str, flax_key))}: {missing}") + if missing_layers: + raise ValueError(f"Missing scanned layer indices: {'; '.join(missing_layers)}") validate_flax_state_dict(eval_shapes, flax_state_dict) flax_state_dict = unflatten_dict(flax_state_dict) diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py index 98f21d278..6d8fbce56 100644 --- a/src/maxdiffusion/models/ltx2/transformer_ltx2.py +++ b/src/maxdiffusion/models/ltx2/transformer_ltx2.py @@ -57,6 +57,34 @@ class LTX2BlockContext: perturbation_mask: Optional[jax.Array] = None +def _canonicalize_attention_mask(mask: Optional[jax.Array], batch_size: int, name: str) -> Optional[jax.Array]: + """Normalizes an LTX2 keep mask to boolean [batch, kv_length]. + + Binary masks use one for valid tokens. Legacy additive masks use zero for + valid tokens and a negative value for masked tokens. The legacy singleton + rank-3 form disambiguates an all-zero additive mask from an all-masked + rank-2 binary mask. + """ + if mask is None: + return None + was_singleton_rank3 = mask.ndim == 3 and mask.shape[1] == 1 + if was_singleton_rank3: + mask = mask[:, 0, :] + if mask.ndim != 2: + raise ValueError(f"{name} must have shape [batch, kv_length], got {mask.shape}.") + if mask.shape[0] != batch_size: + raise ValueError(f"{name} batch dimension must be {batch_size}, got {mask.shape[0]}.") + if mask.dtype == jnp.bool_: + return mask + has_negative = jnp.any(mask < 0) + if was_singleton_rank3 and jnp.issubdtype(mask.dtype, jnp.floating): + has_positive = jnp.any(mask > 0) + is_legacy_additive = has_negative | ~has_positive + else: + is_legacy_additive = has_negative + return jnp.where(is_legacy_additive, mask == 0, mask > 0) + + class LTX2AdaLayerNormSingle(nnx.Module): """ Adaptive Layer Normalization (AdaLN) single modulation module for LTX-Video/LTX-2. @@ -163,6 +191,8 @@ def __init__( perturbed_attn: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, + use_experimental_scheduler: bool = False, enable_jax_named_scopes: bool = False, ): self.dim = dim @@ -206,6 +236,8 @@ def __init__( gated_attn=gated_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, ) self.audio_norm1 = nnx.RMSNorm( @@ -237,6 +269,8 @@ def __init__( gated_attn=gated_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, ) # 2. Prompt Cross-Attention @@ -265,10 +299,13 @@ def __init__( attention_kernel=self.attention_kernel, rope_type=rope_type, flash_block_sizes=flash_block_sizes, + flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, ) self.audio_norm2 = nnx.RMSNorm( @@ -301,6 +338,8 @@ def __init__( gated_attn=gated_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention @@ -334,6 +373,8 @@ def __init__( gated_attn=gated_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, ) self.video_to_audio_norm = nnx.RMSNorm( @@ -366,6 +407,8 @@ def __init__( gated_attn=gated_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, ) # 4. Feed Forward @@ -789,6 +832,8 @@ def __init__( spatio_temporal_guidance_blocks: Tuple[int, ...] = (), ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, + use_experimental_scheduler: bool = False, enable_jax_named_scopes: bool = False, **kwargs, ): @@ -846,6 +891,8 @@ def __init__( self.a2v_attention_kernel = a2v_attention_kernel self.v2a_attention_kernel = v2a_attention_kernel self.flash_min_seq_length = flash_min_seq_length + self.use_base2_exp = use_base2_exp + self.use_experimental_scheduler = use_experimental_scheduler if sharding_specs is None: sharding_specs = get_sharding_specs("default", "ltx2_dit") @@ -1098,6 +1145,8 @@ def init_block(rngs): perturbed_attn=self.perturbed_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=self.use_base2_exp, + use_experimental_scheduler=self.use_experimental_scheduler, enable_jax_named_scopes=self.enable_jax_named_scopes, ) @@ -1140,6 +1189,8 @@ def init_block(rngs): perturbed_attn=self.perturbed_attn, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=self.use_base2_exp, + use_experimental_scheduler=self.use_experimental_scheduler, enable_jax_named_scopes=self.enable_jax_named_scopes, ) blocks.append(block) @@ -1316,6 +1367,7 @@ def __call__( """ # Determine timestep for audio. audio_timestep = audio_timestep if audio_timestep is not None else timestep + batch_size = hidden_states.shape[0] a2v_cross_attention_mask = None v2a_cross_attention_mask = None @@ -1323,24 +1375,18 @@ def __call__( a2v_cross_attention_mask = attention_kwargs.get("a2v_cross_attention_mask", None) v2a_cross_attention_mask = attention_kwargs.get("v2a_cross_attention_mask", None) - if self.attention_kernel == "dot_product": - if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: - encoder_attention_mask = (1 - encoder_attention_mask.astype(self.dtype)) * -10000.0 - encoder_attention_mask = jnp.expand_dims(encoder_attention_mask, axis=1) - - if audio_encoder_attention_mask is not None and audio_encoder_attention_mask.ndim == 2: - audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.astype(self.dtype)) * -10000.0 - audio_encoder_attention_mask = jnp.expand_dims(audio_encoder_attention_mask, axis=1) - - if a2v_cross_attention_mask is not None and a2v_cross_attention_mask.ndim == 2: - a2v_cross_attention_mask = (1 - a2v_cross_attention_mask.astype(self.dtype)) * -10000.0 - a2v_cross_attention_mask = jnp.expand_dims(a2v_cross_attention_mask, axis=1) - - if v2a_cross_attention_mask is not None and v2a_cross_attention_mask.ndim == 2: - v2a_cross_attention_mask = (1 - v2a_cross_attention_mask.astype(self.dtype)) * -10000.0 - v2a_cross_attention_mask = jnp.expand_dims(v2a_cross_attention_mask, axis=1) - - batch_size = hidden_states.shape[0] + encoder_attention_mask = _canonicalize_attention_mask( + encoder_attention_mask, batch_size, "encoder_attention_mask" + ) + audio_encoder_attention_mask = _canonicalize_attention_mask( + audio_encoder_attention_mask, batch_size, "audio_encoder_attention_mask" + ) + a2v_cross_attention_mask = _canonicalize_attention_mask( + a2v_cross_attention_mask, batch_size, "a2v_cross_attention_mask" + ) + v2a_cross_attention_mask = _canonicalize_attention_mask( + v2a_cross_attention_mask, batch_size, "v2a_cross_attention_mask" + ) # 1. Prepare RoPE positional embeddings with self.named_scope("RoPE Preparation"): diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index fbf820d75..932cd6187 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -16,6 +16,7 @@ from typing import Optional, Any, List, Union from functools import partial +import itertools import time import numpy as np @@ -72,6 +73,38 @@ def _enforce_layout(x, axes): } +def _select_batch_sharding_axes(mesh_shape, batch_size: int) -> tuple[str, ...]: + """Selects the largest existing mesh-axis product that divides ``batch_size``.""" + if batch_size < 1: + raise ValueError(f"batch_size must be positive, got {batch_size}.") + + axes = [(name, int(size)) for name, size in mesh_shape.items() if int(size) > 1] + best_axes: tuple[str, ...] = () + best_shards = 1 + for count in range(1, len(axes) + 1): + for combination in itertools.combinations(axes, count): + shards = int(np.prod([size for _, size in combination])) + if batch_size % shards == 0 and shards > best_shards: + best_axes = tuple(name for name, _ in combination) + best_shards = shards + return best_axes + + +@contextlib.contextmanager +def _temporary_vae_slicing(vae, enabled: Optional[bool]): + """Temporarily changes VAE slicing without leaking state across requests.""" + if enabled is None: + yield + return + + original = vae.use_slicing + vae.use_slicing = enabled + try: + yield + finally: + vae.use_slicing = original + + @flax.struct.dataclass class LTX2PipelineOutput: frames: jax.Array @@ -265,6 +298,8 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): ltx2_config["flash_min_seq_length"] = getattr(config, "flash_min_seq_length", 4096) ltx2_config["ulysses_shards"] = getattr(config, "ulysses_shards", -1) ltx2_config["ulysses_attention_chunks"] = getattr(config, "ulysses_attention_chunks", 1) + ltx2_config["use_base2_exp"] = getattr(config, "use_base2_exp", False) + ltx2_config["use_experimental_scheduler"] = getattr(config, "use_experimental_scheduler", False) ltx2_config["enable_jax_named_scopes"] = getattr(config, "enable_jax_named_scopes", False) ltx2_config["remat_policy"] = config.remat_policy ltx2_config["names_which_can_be_saved"] = config.names_which_can_be_saved @@ -353,11 +388,10 @@ def stages_via_ici(val, sharding) -> bool: try: state[path].set_value(device_put_replicated(val, sharding)) except Exception as e: - max_logging.log(f"Failed to device_put_replicated {path}: {e}") - max_logging.log(f"Trying to use process_allgather for {path}") - val_on_host = jax.experimental.multihost_utils.process_allgather(val, tiled=True) - state[path].set_value(device_put_replicated(val_on_host, sharding)) - del val_on_host + # Every process already loaded the complete host array. A conditional + # process_allgather here can deadlock when only one process fails, and + # tiled=True would concatenate complete parameters along axis 0. + raise RuntimeError(f"Failed to place transformer parameter {path} with sharding {sharding}") from e jax.block_until_ready([state[path].value for path, _, _ in put_specs]) max_logging.log(f"Transformer {subfolder or 'transformer'} weights on device in {time.perf_counter() - t_put_start:.1f}s") @@ -1109,8 +1143,18 @@ def _get_gemma_prompt_embeds( # Distribute the batch dimension across available TPUs to prevent Softmax OOM # (reduces 512MB allocation down to 64MB per TPU for batch size 16) if hasattr(self, "mesh") and self.mesh is not None: - data_axis = self.mesh.axis_names[0] - sharding = NamedSharding(self.mesh, P(data_axis)) + batch_axes = _select_batch_sharding_axes(self.mesh.shape, text_input_ids.shape[0]) + sharding = NamedSharding(self.mesh, P(batch_axes) if batch_axes else P()) + if batch_axes: + max_logging.log( + f"Sharding TPU text-encoder batch across mesh axes {batch_axes} " + f"({int(np.prod([self.mesh.shape[axis] for axis in batch_axes]))} shards)." + ) + else: + max_logging.log( + f"No mesh-axis product divides text batch size {text_input_ids.shape[0]}; " + "using replicated text-encoder inputs." + ) text_input_ids = jax.device_put(text_input_ids, sharding) prompt_attention_mask = jax.device_put(prompt_attention_mask, sharding) else: @@ -2190,27 +2234,26 @@ def __call__( enable_dynamic_vae_sharding = ( getattr(self.config, "enable_dynamic_vae_sharding", True) if hasattr(self, "config") else True ) + temporary_vae_slicing = None if enable_dynamic_vae_sharding and batch_size > 2: - max_logging.log( - f"[Tuning] Disabling VAE slicing and applying dynamic batch sharding to prevent HBM OOM for batch_size {batch_size} > 2" - ) try: - # Disable sequential slicing to avoid JAX concatenating 17GB arrays on the TPU - self.vae.use_slicing = False - # Distribute the batch dimension across the existing mesh to ensure topological compatibility mesh = latents.sharding.mesh - active_axes = [] - current_shards = 1 - - for axis_name, size in mesh.shape.items(): - if size > 1 and batch_size % (current_shards * size) == 0: - active_axes.append(axis_name) - current_shards *= size - + active_axes = _select_batch_sharding_axes(mesh.shape, batch_size) if active_axes: + current_shards = int(np.prod([mesh.shape[axis] for axis in active_axes])) + max_logging.log( + f"[Tuning] Temporarily disabling VAE slicing and sharding batch size {batch_size} " + f"over axes {active_axes} ({current_shards} shards)." + ) batch_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(tuple(active_axes))) latents = jax.lax.with_sharding_constraint(latents, batch_sharding) + temporary_vae_slicing = False + else: + max_logging.log( + f"[Tuning] No mesh-axis product divides VAE batch size {batch_size}; " + "preserving the existing VAE slicing setting." + ) except Exception as e: # pylint: disable=broad-exception-caught max_logging.log(f"[Tuning] Failed to apply batch sharding constraint to VAE: {e}") elif replicate_vae: @@ -2224,15 +2267,16 @@ def __call__( latent_processing_time += time.perf_counter() - t0_latent_processing timings["Latent Processing"] = latent_processing_time - video, video_vae_time, video_post_time = self._decode_latents_to_video( - latents=latents, - generator=generator, - batch_size=batch_size, - decode_timestep=decode_timestep, - decode_noise_scale=decode_noise_scale, - output_type=output_type, - replicate_vae=replicate_vae, - ) + with _temporary_vae_slicing(self.vae, temporary_vae_slicing): + video, video_vae_time, video_post_time = self._decode_latents_to_video( + latents=latents, + generator=generator, + batch_size=batch_size, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + output_type=output_type, + replicate_vae=replicate_vae, + ) # Decode Audio t0_audio_vae = time.perf_counter() diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index 793b8979b..1bb79afbc 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -222,12 +222,15 @@ def user_init(raw_keys): raw_keys["vae_logical_axis_rules"] = _lists_to_tuples(raw_keys["vae_logical_axis_rules"]) # Verify qkv is sharded across sequence. attention = raw_keys["attention"] - if ( - attention in ["ulysses_ring", "ulysses_ring_custom", "ulysses_ring_custom_bidir"] - and raw_keys.get("ulysses_shards", -1) <= 0 - ): + ulysses_ring_attentions = { + "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_bidir", + } + if attention in ulysses_ring_attentions and raw_keys.get("ulysses_shards", -1) <= 0: raise ValueError(f"{attention} requires ulysses_shards to be set from config or command line.") - uses_ulysses_ring_attention = attention == "ulysses_ring" + uses_ulysses_ring_attention = attention in ulysses_ring_attentions uses_ring_attention = "ring" in attention and not uses_ulysses_ring_attention uses_ulysses_attention = "ulysses" in attention and not uses_ulysses_ring_attention uses_uniform_sequence_sharding = raw_keys["attention_sharding_uniform"] diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py index 2154b26ac..916e80ef3 100644 --- a/src/maxdiffusion/tests/aot_cache_test.py +++ b/src/maxdiffusion/tests/aot_cache_test.py @@ -60,6 +60,19 @@ def test_disabled_is_plain_jit(self): np.testing.assert_allclose(result, self._a @ self._b + 1.0) self.assertFalse(aot_cache._STATE.enabled) + def test_empty_install_disables_and_clears_a_previous_cache(self): + self._install() + _toy_fn(self._a, self._b, True) + self.assertTrue(aot_cache._STATE.enabled) + + aot_cache.install("", meta={}, mesh=None) + + self.assertFalse(aot_cache._STATE.enabled) + self.assertEqual(aot_cache._STATE.cache_dir, "") + for entry in aot_cache._REGISTRY: + self.assertFalse(entry._compiled) + self.assertFalse(entry._pending) + def test_record_save_hit_roundtrip(self): self._install() first = _toy_fn(self._a, self._b, True) # miss -> jit + record diff --git a/src/maxdiffusion/tests/attention_test.py b/src/maxdiffusion/tests/attention_test.py index 4421b2006..0fbc3f3b5 100644 --- a/src/maxdiffusion/tests/attention_test.py +++ b/src/maxdiffusion/tests/attention_test.py @@ -235,6 +235,29 @@ def test_select_flash_block_sizes_derives_cross_attn_defaults_for_tokamax(self): self.assertIsNone(cross_attention_block_sizes.block_kv_dq) self.assertTrue(cross_attention_block_sizes.use_fused_bwd_kernel) + def test_select_flash_block_sizes_coerces_custom_carrier_for_local_flash(self): + custom_sizes = max_utils.CustomFlashBlockSizes( + block_q=16, + block_kv=16, + block_kv_compute=16, + block_kv_compute_in=16, + heads_per_tile=1, + ) + query = jnp.zeros((1, 32, 64), dtype=jnp.bfloat16) + + block_sizes = _select_flash_block_sizes( + query=query, + key=query, + flash_block_sizes=custom_sizes, + dtype=jnp.bfloat16, + attention_kernel="tokamax_flash", + ) + + self.assertTrue(block_sizes.use_fused_bwd_kernel) + self.assertEqual(block_sizes.block_q, 16) + self.assertEqual(block_sizes.block_kv, 16) + self.assertEqual(block_sizes.block_kv_compute, 16) + def test_ulysses_head_chunk_ranges_preserve_head_layout_with_remainder(self): ranges = attention_flax._ulysses_head_chunk_ranges(num_heads=40, ulysses_shards=8, num_chunks=2) @@ -486,7 +509,7 @@ def test_ulysses_attention_uses_attention_mask_for_segment_ids(self): query = jnp.arange(batch * length * heads * head_depth, dtype=jnp.float32).reshape(batch, length, heads * head_depth) key = query + 100.0 value = query + 200.0 - attention_mask = jnp.array([[1, 0, 1, 0, 1]], dtype=jnp.int32) + attention_mask = jnp.array([[1, 0, 1, 0, 1], [0, 1, 0, 1, 0]], dtype=jnp.int32) mesh = self._ulysses_mesh() def fake_make_splash_mha(**unused_kwargs): @@ -533,6 +556,88 @@ def fake_kernel(q, k, v, segment_ids): self.assertEqual(output.shape, query.shape) self.assertTrue(jnp.array_equal(output, expected)) + @unittest.skipIf(len(jax.devices()) < 4, "Mask shard_map test requires at least 4 devices.") + def test_flash_and_ulysses_shard_distinct_masks_over_batch_and_sequence(self): + """Masks must follow data sharding and Ulysses must gather their KV axis.""" + batch = 2 + length = 6 + heads = 4 + head_depth = 3 + query = jnp.arange(batch * length * heads * head_depth, dtype=jnp.float32).reshape( + batch, length, heads * head_depth + ) + attention_mask = jnp.array( + [[1, 0, 1, 0, 1, 0], [0, 1, 1, 1, 0, 1]], + dtype=jnp.int32, + ) + devices = np.array(jax.devices()[:4]).reshape(2, 1, 2, 1) + mesh = Mesh(devices, ("data", "fsdp", "context", "tensor")) + + def fake_make_splash_mha(**unused_kwargs): + def fake_kernel(q, k, v, segment_ids): + del k, v + return q + jnp.sum(segment_ids.kv).astype(q.dtype) + + return fake_kernel + + with mock.patch.object( + attention_flax.splash_attention_kernel, + "make_splash_mha", + side_effect=fake_make_splash_mha, + ): + with mesh, nn_partitioning.axis_rules(self._flash_axis_rules()): + flash_output = attention_flax._tpu_flash_attention( + query, + query, + query, + heads=heads, + mesh=mesh, + axis_names_q=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ), + axis_names_kv=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ), + flash_block_sizes=self._ulysses_block_sizes(), + dtype=jnp.float32, + attention_kernel="flash", + attention_mask=attention_mask, + ) + + with mesh, nn_partitioning.axis_rules(self._ulysses_axis_rules()): + ulysses_output = attention_flax._ulysses_attention( + query, + query, + query, + heads=heads, + mesh=mesh, + axis_names_q=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ), + axis_names_kv=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ), + flash_block_sizes=self._ulysses_block_sizes(), + dtype=jnp.float32, + attention_mask=attention_mask, + ) + + expected = query + jnp.sum(attention_mask, axis=1)[:, None, None] + self.assertTrue(jnp.array_equal(flash_output, expected)) + self.assertTrue(jnp.array_equal(ulysses_output, expected)) + @unittest.skipIf(len(jax.devices()) < 4, "Ulysses ring attention layout test requires at least 4 devices.") def test_ulysses_ring_attention_round_trips_query_when_heads_are_divisible(self): """Hybrid Ulysses+ring should preserve layout while only exposing context.""" @@ -767,8 +872,8 @@ def fake_kernel(q, k, v, segment_ids): self.assertEqual(output.shape, query.shape) self.assertTrue(jnp.array_equal(output, query)) - # Same convention as the tokamax_ring kernel: shards pad identically, no rotation. - self.assertEqual(seen, [False]) + # Explicit per-token IDs must rotate with their corresponding K/V shard. + self.assertEqual(seen, [True]) def test_ulysses_ring_attention_raises_when_heads_are_not_divisible_by_ulysses_shards(self): """The hidden all-to-all head split still requires divisible heads.""" diff --git a/src/maxdiffusion/tests/ltx2/test_attention_ltx2.py b/src/maxdiffusion/tests/ltx2/test_attention_ltx2.py index d29fee9a4..3ee55828f 100644 --- a/src/maxdiffusion/tests/ltx2/test_attention_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_attention_ltx2.py @@ -15,6 +15,7 @@ """ import unittest +from unittest import mock import torch import numpy as np import jax @@ -22,7 +23,250 @@ from flax import nnx import pandas as pd from jax.sharding import Mesh +from maxdiffusion.models.attention_flax import ( + KERNEL_REGISTRY, + FlaxWanAttention, + _apply_attention, + _build_padding_segment_ids, +) from maxdiffusion.models.ltx2.attention_ltx2 import LTX2Attention, LTX2RotaryPosEmbed +from maxdiffusion.models.ltx2.transformer_ltx2 import _canonicalize_attention_mask + + +class LTX2AttentionMaskContractTest(unittest.TestCase): + + def test_dot_product_and_flash_fallback_apply_per_example_masks(self): + query = jnp.zeros((2, 1, 2), dtype=jnp.float32) + key = jnp.zeros((2, 2, 2), dtype=jnp.float32) + value = jnp.array( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[1.0, 0.0], [0.0, 1.0]], + ], + dtype=jnp.float32, + ) + attention_mask = jnp.array([[True, False], [False, True]]) + + for attention_kernel in ("dot_product", "flash"): + with self.subTest(attention_kernel=attention_kernel): + output = _apply_attention( + query=query, + key=key, + value=value, + heads=1, + dim_head=2, + split_head_dim=True, + float32_qk_product=True, + attention_kernel=attention_kernel, + flash_min_seq_length=8, + use_memory_efficient_attention=False, + scale=1.0, + dtype=jnp.float32, + mesh=None, + axis_names_q=(None, None, None, None), + axis_names_kv=(None, None, None, None), + flash_block_sizes=None, + dpa_layer=None, + attention_mask=attention_mask, + ) + np.testing.assert_allclose( + np.asarray(output), + np.array([[[1.0, 0.0]], [[0.0, 1.0]]], dtype=np.float32), + atol=1e-6, + ) + + def test_padding_segment_ids_preserve_each_batch_mask(self): + attention_mask = jnp.array([[True, False, True], [False, True, True]]) + segment_ids = _build_padding_segment_ids( + query_seq_len=2, + q_padded_len=4, + key_seq_len=4, + kv_padded_len=5, + attention_mask=attention_mask, + ) + + np.testing.assert_array_equal(np.asarray(segment_ids.q), np.array([[1, 1, 0, 0], [1, 1, 0, 0]])) + np.testing.assert_array_equal(np.asarray(segment_ids.kv), np.array([[1, 0, 1, 0, 0], [0, 1, 1, 0, 0]])) + + def test_dot_product_all_masked_row_returns_zero(self): + query = jnp.zeros((1, 1, 2), dtype=jnp.float32) + key = jnp.zeros((1, 2, 2), dtype=jnp.float32) + value = jnp.ones((1, 2, 2), dtype=jnp.float32) + output = _apply_attention( + query=query, + key=key, + value=value, + heads=1, + dim_head=2, + split_head_dim=True, + float32_qk_product=True, + attention_kernel="dot_product", + flash_min_seq_length=0, + use_memory_efficient_attention=False, + scale=1.0, + dtype=jnp.float32, + mesh=None, + axis_names_q=(None, None, None, None), + axis_names_kv=(None, None, None, None), + flash_block_sizes=None, + dpa_layer=None, + attention_mask=jnp.array([[False, False]]), + ) + + np.testing.assert_array_equal(np.asarray(output), np.zeros((1, 1, 2), dtype=np.float32)) + + def test_forced_flash_receives_boolean_keep_mask(self): + query = jnp.zeros((2, 2, 2), dtype=jnp.float32) + attention_mask = jnp.array([[1, 0], [0, 1]], dtype=jnp.int32) + + def capture_flash_mask(q, _key, _value, context): + self.assertEqual(context["attention_mask"].dtype, jnp.bool_) + np.testing.assert_array_equal(np.asarray(context["attention_mask"]), np.asarray(attention_mask, dtype=bool)) + return q + + with mock.patch.dict(KERNEL_REGISTRY, {"flash": capture_flash_mask}): + output = _apply_attention( + query=query, + key=query, + value=query, + heads=1, + dim_head=2, + split_head_dim=True, + float32_qk_product=True, + attention_kernel="flash", + flash_min_seq_length=0, + use_memory_efficient_attention=False, + scale=1.0, + dtype=jnp.float32, + mesh=None, + axis_names_q=(None, None, None, None), + axis_names_kv=(None, None, None, None), + flash_block_sizes=None, + dpa_layer=None, + attention_mask=attention_mask, + ) + + np.testing.assert_array_equal(np.asarray(output), np.asarray(query)) + + def test_transformer_boundary_canonicalizes_keep_masks(self): + canonical = _canonicalize_attention_mask( + jnp.array([[[1, 0, 1]], [[0, 1, 1]]], dtype=jnp.int32), + batch_size=2, + name="test_mask", + ) + self.assertEqual(canonical.shape, (2, 3)) + self.assertEqual(canonical.dtype, jnp.bool_) + np.testing.assert_array_equal(np.asarray(canonical), np.array([[True, False, True], [False, True, True]])) + + def test_transformer_boundary_accepts_legacy_additive_masks(self): + canonical = _canonicalize_attention_mask( + jnp.array([[[0.0, -10000.0, 0.0]], [[0.0, 0.0, 0.0]]], dtype=jnp.float32), + batch_size=2, + name="test_mask", + ) + np.testing.assert_array_equal(np.asarray(canonical), np.array([[True, False, True], [True, True, True]])) + + def test_rank2_all_zero_float_mask_remains_all_masked(self): + canonical = _canonicalize_attention_mask( + jnp.zeros((2, 3), dtype=jnp.float32), + batch_size=2, + name="test_mask", + ) + np.testing.assert_array_equal(np.asarray(canonical), np.zeros((2, 3), dtype=bool)) + + def test_ring_cross_attention_uses_global_kv_and_wires_kernel_flags(self): + remapped_kernels = ( + "tokamax_ring", + "tokamax_ring_custom", + "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_bidir", + "ulysses_custom", + "ulysses_custom_fixed_m", + ) + + for attention_kernel in remapped_kernels: + with self.subTest(module="ltx2", attention_kernel=attention_kernel): + attention = LTX2Attention( + query_dim=4, + context_dim=4, + heads=1, + dim_head=4, + rngs=nnx.Rngs(0), + attention_kernel=attention_kernel, + flash_min_seq_length=0, + use_base2_exp=True, + use_experimental_scheduler=True, + ) + self.assertEqual(attention.attention_op.attention_kernel, "tokamax_flash") + self.assertIsNone(attention.attention_op.axis_names_kv[2]) + self.assertTrue(attention.attention_op.use_base2_exp) + self.assertTrue(attention.attention_op.use_experimental_scheduler) + + with self.subTest(module="wan", attention_kernel=attention_kernel): + wan_attention = FlaxWanAttention( + rngs=nnx.Rngs(0), + query_dim=4, + cross_attention_dim=4, + heads=1, + dim_head=4, + attention_kernel=attention_kernel, + is_self_attention=False, + ) + self.assertEqual(wan_attention.attention_op.attention_kernel, "tokamax_flash") + self.assertIsNone(wan_attention.attention_op.axis_names_kv[2]) + + local_flash_attention = LTX2Attention( + query_dim=4, + context_dim=4, + heads=1, + dim_head=4, + rngs=nnx.Rngs(0), + attention_kernel="flash", + ) + self.assertEqual(local_flash_attention.attention_op.attention_kernel, "flash") + self.assertIsNone(local_flash_attention.attention_op.axis_names_kv[2]) + + distributed_ulysses_attention = LTX2Attention( + query_dim=4, + context_dim=4, + heads=1, + dim_head=4, + rngs=nnx.Rngs(0), + attention_kernel="ulysses", + ) + self.assertEqual(distributed_ulysses_attention.attention_op.attention_kernel, "ulysses") + self.assertIsNotNone(distributed_ulysses_attention.attention_op.axis_names_kv[2]) + + def test_tokamax_ring_registry_forwards_kernel_flags(self): + query = jnp.zeros((1, 2, 2), dtype=jnp.float32) + context = { + "scale": 1.0, + "heads": 1, + "mesh": None, + "axis_names_q": (None, None, None, None), + "axis_names_kv": (None, None, None, None), + "flash_block_sizes": None, + "dtype": jnp.float32, + "mask_padding_tokens": True, + "residual_checkpoint_name": "ring_residual", + "attention_mask": None, + "use_base2_exp": True, + "use_experimental_scheduler": True, + } + + with mock.patch( + "maxdiffusion.models.attention_flax._tpu_flash_attention", + return_value=query, + ) as mocked_flash: + KERNEL_REGISTRY["tokamax_ring"](query, query, query, context) + + call_kwargs = mocked_flash.call_args.kwargs + self.assertEqual(call_kwargs["residual_checkpoint_name"], "ring_residual") + self.assertTrue(call_kwargs["use_base2_exp"]) + self.assertTrue(call_kwargs["use_experimental_scheduler"]) + # ========================================== # 1. PyTorch Reference Implementations diff --git a/src/maxdiffusion/tests/ltx2/test_generate_ltx2_aot.py b/src/maxdiffusion/tests/ltx2/test_generate_ltx2_aot.py new file mode 100644 index 000000000..c113d9d1a --- /dev/null +++ b/src/maxdiffusion/tests/ltx2/test_generate_ltx2_aot.py @@ -0,0 +1,183 @@ +# 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 +# +# http://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 LTX2 AOT cache identity and invalidation.""" + +from types import SimpleNamespace +import unittest +from unittest import mock + +from maxdiffusion import aot_cache +from maxdiffusion import generate_ltx2 + + +class LTX2AotMetadataTest(unittest.TestCase): + + def setUp(self): + self.config = SimpleNamespace( + pretrained_model_name_or_path="test-model", + use_base2_exp=False, + use_experimental_scheduler=False, + ) + self.pipeline = SimpleNamespace( + transformer=SimpleNamespace(config={}), + mesh=SimpleNamespace(shape={"data": 1}, axis_names=("data",)), + ) + + def _fingerprint(self, revision): + metadata = generate_ltx2.ltx2_aot_metadata(self.config, self.pipeline, source_revision=revision) + return metadata, aot_cache._metadata_fingerprint(metadata) + + def test_source_revision_changes_cache_fingerprint(self): + metadata_a, fingerprint_a = self._fingerprint("revision-a") + metadata_b, fingerprint_b = self._fingerprint("revision-b") + + self.assertEqual(metadata_a["source_revision"], "revision-a") + self.assertEqual(metadata_b["source_revision"], "revision-b") + self.assertNotEqual(fingerprint_a, fingerprint_b) + + def test_ltx2_abi_changes_cache_fingerprint(self): + metadata_a, fingerprint_a = self._fingerprint("same-revision") + with mock.patch.object(generate_ltx2, "LTX2_AOT_CACHE_ABI", generate_ltx2.LTX2_AOT_CACHE_ABI + 1): + metadata_b, fingerprint_b = self._fingerprint("same-revision") + + self.assertNotEqual(metadata_a["ltx2_aot_cache_abi"], metadata_b["ltx2_aot_cache_abi"]) + self.assertNotEqual(fingerprint_a, fingerprint_b) + + def test_unavailable_revision_never_reuses_a_cache_identity(self): + metadata_a, fingerprint_a = self._fingerprint(None) + metadata_b, fingerprint_b = self._fingerprint(None) + + self.assertTrue(metadata_a["source_revision"].startswith("unversioned:")) + self.assertTrue(metadata_b["source_revision"].startswith("unversioned:")) + self.assertNotEqual(fingerprint_a, fingerprint_b) + + def test_build_revision_is_used_for_programmatic_run(self): + self.config.aot_build_revision = "build-123" + self.assertEqual(generate_ltx2._resolve_ltx2_aot_source_revision(self.config), "build-123") + self.assertEqual( + generate_ltx2._resolve_ltx2_aot_source_revision(self.config, commit_hash="commit-456"), + "commit-456", + ) + + def test_attention_math_flags_change_cache_fingerprint(self): + _, fingerprint_a = self._fingerprint("same-revision") + self.config.use_base2_exp = True + _, fingerprint_b = self._fingerprint("same-revision") + self.config.use_experimental_scheduler = True + _, fingerprint_c = self._fingerprint("same-revision") + + self.assertNotEqual(fingerprint_a, fingerprint_b) + self.assertNotEqual(fingerprint_b, fingerprint_c) + + def test_codegen_and_remat_inputs_change_cache_fingerprint(self): + _, fingerprint_a = self._fingerprint("same-revision") + self.config.precision = "highest" + _, fingerprint_b = self._fingerprint("same-revision") + self.config.names_which_can_be_saved = ("attn",) + _, fingerprint_c = self._fingerprint("same-revision") + self.config.names_which_can_be_offloaded = ("mlp",) + _, fingerprint_d = self._fingerprint("same-revision") + + self.assertNotEqual(fingerprint_a, fingerprint_b) + self.assertNotEqual(fingerprint_b, fingerprint_c) + self.assertNotEqual(fingerprint_c, fingerprint_d) + + def test_tile_search_rejects_a_prebuilt_pipeline(self): + config = SimpleNamespace(get_keys=lambda: {"enable_tile_search": True}) + with self.assertRaisesRegex(ValueError, "cannot be used with a prebuilt pipeline"): + generate_ltx2.run(config, pipeline=object()) + + def test_tile_search_forwards_config_and_applies_winner(self): + keys = { + "enable_tile_search": True, + "tile_search_mode": "full", + "tile_search_iters": 3, + "tile_search_out": "results", + "tile_search_vmem_limit_bytes": 123456, + } + config = SimpleNamespace( + get_keys=lambda: keys, + mesh_axes=("data",), + flash_block_sizes={"block_q": 64}, + attention="flash", + ) + best = SimpleNamespace(bq=128, bkv=256, bkv_compute=256, mean_ms=1.5) + benchmark = SimpleNamespace(label="test-benchmark") + + with ( + mock.patch.object(generate_ltx2.max_utils, "create_device_mesh", return_value=["device"]), + mock.patch.object(generate_ltx2.jax.sharding, "Mesh", return_value="mesh"), + mock.patch( + "maxdiffusion.utils.ltx2_block_benchmark.LTX2BlockBenchmark.from_config", + return_value=benchmark, + ) as benchmark_from_config, + mock.patch( + "maxdiffusion.utils.tile_size_grid_search.grid_search", + return_value=SimpleNamespace(best=best), + ) as grid_search, + mock.patch.object( + generate_ltx2.max_utils, + "flash_block_sizes_for_candidate", + return_value={"block_q": 128}, + ) as apply_candidate, + mock.patch.object( + generate_ltx2.max_utils, + "get_flash_block_sizes", + return_value=SimpleNamespace(block_q=128), + ), + ): + generate_ltx2.maybe_tune_block_sizes(config) + + grid_search.assert_called_once_with( + benchmark, + mode="full", + iters=3, + out_dir="results", + log=generate_ltx2.max_logging.log, + ) + benchmark_from_config.assert_called_once_with(config, "mesh", vmem_limit_bytes=123456) + apply_candidate.assert_called_once_with( + config.flash_block_sizes, + "flash", + 128, + 256, + 256, + vmem_limit_bytes=123456, + ) + self.assertEqual(keys["flash_block_sizes"], {"block_q": 128}) + + @mock.patch.object(generate_ltx2.max_logging, "log") + @mock.patch.object( + generate_ltx2.subprocess, + "check_output", + side_effect=[ + b"0123456789abcdef\n", + b" M src/maxdiffusion/generate_ltx2.py\n", + b"", + b"0123456789abcdef\n", + b" M src/maxdiffusion/generate_ltx2.py\n", + b"", + ], + ) + def test_dirty_tree_does_not_reuse_cache_identity(self, _check_output, _log): + revision_a = generate_ltx2.get_git_commit_hash() + revision_b = generate_ltx2.get_git_commit_hash() + self.assertEqual(revision_a, "dirty:0123456789abcdef") + self.assertEqual(revision_b, revision_a) + self.assertFalse(generate_ltx2._is_reusable_aot_revision(revision_a)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py b/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py new file mode 100644 index 000000000..876b4cb0b --- /dev/null +++ b/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py @@ -0,0 +1,167 @@ +""" +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. +""" + +import concurrent.futures +import json +import unittest +from unittest.mock import MagicMock, mock_open, patch + +import numpy as np +import torch +from huggingface_hub.utils import EntryNotFoundError + +from maxdiffusion.models.ltx2 import ltx2_utils + + +class _FakeSafetensorsFile: + + def __init__(self, tensors): + self._tensors = tensors + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def keys(self): + return self._tensors.keys() + + def get_tensor(self, key): + return self._tensors[key] + + +class LTX2TransformerLoaderTest(unittest.TestCase): + + @staticmethod + def _eval_shapes(num_layers=2, parameter_names=("bias",)): + return { + "transformer_blocks": { + parameter_name: np.zeros((num_layers, 1), dtype=np.float32) for parameter_name in parameter_names + } + } + + @staticmethod + def _bin_download(_model_name, *, filename, **_kwargs): + if filename.endswith(".json") or filename.endswith(".safetensors"): + raise EntryNotFoundError("not found") + return "checkpoint.bin" + + def _load_bin(self, checkpoint, eval_shapes=None, **loader_kwargs): + torch_load = MagicMock(return_value=checkpoint) + with ( + patch.object(ltx2_utils.jax, "local_devices", return_value=["cpu"]), + patch.object(ltx2_utils, "hf_hub_download", side_effect=self._bin_download), + patch.object(ltx2_utils.torch, "load", torch_load), + ): + result = ltx2_utils.load_transformer_weights( + "test/model", + eval_shapes or self._eval_shapes(), + "cpu", + **loader_kwargs, + ) + return result, torch_load + + def test_bin_checkpoint_is_loaded_once(self): + checkpoint = { + "transformer_blocks.0.bias": torch.tensor([1.0]), + "transformer_blocks.1.bias": torch.tensor([2.0]), + } + + result, torch_load = self._load_bin(checkpoint) + + torch_load.assert_called_once_with("checkpoint.bin", map_location="cpu") + np.testing.assert_array_equal(result["transformer_blocks"]["bias"], np.array([[1.0], [2.0]])) + + def test_scanned_layer_count_is_derived_from_eval_shapes(self): + checkpoint = { + "transformer_blocks.0.bias": torch.tensor([1.0]), + "transformer_blocks.1.bias": torch.tensor([2.0]), + } + + with self.assertRaisesRegex(ValueError, "num_layers=3 does not match the 2 layers derived from eval_shapes"): + self._load_bin(checkpoint, num_layers=3) + + def test_missing_scanned_layer_raises(self): + checkpoint = {"transformer_blocks.0.bias": torch.tensor([1.0])} + + with self.assertRaisesRegex(ValueError, r"Missing scanned layer indices: transformer_blocks\.bias: \[1\]"): + self._load_bin(checkpoint) + + def test_duplicate_scanned_layer_raises(self): + checkpoint = { + "transformer_blocks.0.bias": torch.tensor([1.0]), + "transformer_blocks.0.alias": torch.tensor([2.0]), + } + original_rename = ltx2_utils.rename_for_ltx2_transformer + + def rename_alias(key): + return original_rename(key).replace(".alias", ".bias") + + with patch.object(ltx2_utils, "rename_for_ltx2_transformer", side_effect=rename_alias): + with self.assertRaisesRegex(ValueError, r"Duplicate scanned layer index 0 for transformer_blocks\.bias"): + self._load_bin(checkpoint) + + def test_out_of_range_scanned_layer_raises(self): + checkpoint = {"transformer_blocks.2.bias": torch.tensor([1.0])} + + with self.assertRaisesRegex( + ValueError, r"Scanned layer index 2 for transformer_blocks\.bias is out of range \[0, 2\)" + ): + self._load_bin(checkpoint) + + def test_safetensors_loading_remains_chunked_and_parallel(self): + parameter_names = tuple(f"param_{index}" for index in range(17)) + checkpoint = { + f"transformer_blocks.{layer_index}.{parameter_name}": torch.tensor([float(layer_index)]) + for parameter_name in parameter_names + for layer_index in range(2) + } + index_data = { + "weight_map": {key: "checkpoint.safetensors" for key in checkpoint}, + } + + def download(_model_name, *, filename, **_kwargs): + if filename.endswith(".json"): + return "index.json" + return filename + + safe_open_mock = MagicMock(side_effect=lambda *_args, **_kwargs: _FakeSafetensorsFile(checkpoint)) + executor_type = concurrent.futures.ThreadPoolExecutor + with ( + patch.object(ltx2_utils.jax, "local_devices", return_value=["cpu"]), + patch.object(ltx2_utils, "hf_hub_download", side_effect=download), + patch.object(ltx2_utils, "safe_open", safe_open_mock), + patch("builtins.open", mock_open(read_data=json.dumps(index_data))), + patch.object(concurrent.futures, "ThreadPoolExecutor", wraps=executor_type) as executor, + ): + result = ltx2_utils.load_transformer_weights( + "test/model", + self._eval_shapes(parameter_names=parameter_names), + "cpu", + ) + + executor.assert_called_once_with(max_workers=32) + self.assertEqual(safe_open_mock.call_count, 3) + for parameter_name in parameter_names: + np.testing.assert_array_equal( + result["transformer_blocks"][parameter_name], + np.array([[0.0], [1.0]]), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py index 348e85438..4eba4d268 100644 --- a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py @@ -20,11 +20,14 @@ import jax.numpy as jnp import numpy as np +from maxdiffusion.max_utils import flash_block_sizes_for_candidate from maxdiffusion.pipelines.ltx2.ltx2_pipeline import ( LTX2Pipeline, _apply_ltx2_guidance, _restore_guidance_latents, + _select_batch_sharding_axes, _select_guidance_latents, + _temporary_vae_slicing, calculate_shift, rescale_noise_cfg, ) @@ -132,6 +135,42 @@ def test_guidance_helpers(self): np.testing.assert_array_equal(restored_latents, stacked_latents) np.testing.assert_array_equal(restored_audio_latents, stacked_audio_latents) + def test_batch_sharding_axes_use_largest_divisor(self): + mesh_shape = {"data": 1, "fsdp": 1, "context": 8, "tensor": 2} + self.assertEqual(_select_batch_sharding_axes(mesh_shape, 16), ("context", "tensor")) + self.assertEqual(_select_batch_sharding_axes(mesh_shape, 8), ("context",)) + self.assertEqual(_select_batch_sharding_axes(mesh_shape, 3), ()) + + def test_tuned_block_sizes_update_effective_tokamax_fields(self): + candidate = flash_block_sizes_for_candidate( + { + "block_q": 2048, + "block_kv": 2048, + "block_kv_compute": 1024, + "block_q_dkv": 2048, + "block_kv_dkv": 2048, + "block_kv_dkv_compute": 2048, + "use_fused_bwd_kernel": True, + }, + "tokamax_flash", + block_q=1024, + block_kv=768, + block_kv_compute=512, + ) + self.assertEqual(candidate["block_q"], 1024) + self.assertEqual(candidate["block_q_dkv"], 1024) + self.assertEqual(candidate["block_kv_dkv"], 768) + self.assertEqual(candidate["block_kv_dkv_compute"], 512) + + def test_temporary_vae_slicing_restores_state_after_error(self): + vae = MagicMock() + vae.use_slicing = True + with self.assertRaisesRegex(RuntimeError, "decode failed"): + with _temporary_vae_slicing(vae, False): + self.assertFalse(vae.use_slicing) + raise RuntimeError("decode failed") + self.assertTrue(vae.use_slicing) + def test_pipeline_init(self): """Test LTX2Pipeline initialization and property extraction.""" mock_vae = MagicMock() diff --git a/src/maxdiffusion/tests/tile_size_grid_search_test.py b/src/maxdiffusion/tests/tile_size_grid_search_test.py index 3a985f072..3419c142f 100644 --- a/src/maxdiffusion/tests/tile_size_grid_search_test.py +++ b/src/maxdiffusion/tests/tile_size_grid_search_test.py @@ -16,15 +16,18 @@ import time import unittest +import numpy as np from maxdiffusion.utils.tile_size_grid_search import ( MXU_TILE, VPU_LANE, + _aggregate_process_measurements, BenchResult, BlockBenchmark, bkv_candidates, bq_candidates, grid_search, + local_tiled_seq_len, padding_of, smart_grid, time_callable, @@ -97,6 +100,19 @@ def test_candidates_not_strictly_256(self): bkvs = bkv_candidates(RING_SEQ, k=3, max_block=vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB)) self.assertTrue(any(b % MXU_TILE != 0 for b in bkvs)) + def test_ring_sequence_length_uses_actual_ring_shards(self): + self.assertEqual(local_tiled_seq_len(6144, "tokamax_ring", context_shards=8, ulysses_shards=-1), 768) + self.assertEqual(local_tiled_seq_len(6144, "ulysses_ring", context_shards=8, ulysses_shards=2), 1536) + self.assertEqual(local_tiled_seq_len(6144, "flash", context_shards=8, ulysses_shards=-1), 6144) + + def test_ring_sequence_length_matches_production_padding(self): + self.assertEqual(local_tiled_seq_len(10, "tokamax_ring", context_shards=4, ulysses_shards=-1), 3) + self.assertEqual(local_tiled_seq_len(10, "ulysses_ring", context_shards=4, ulysses_shards=2), 6) + + def test_hybrid_ring_requires_compatible_shards(self): + with self.assertRaisesRegex(ValueError, "must be divisible"): + local_tiled_seq_len(6144, "ulysses_ring_custom", context_shards=8, ulysses_shards=3) + class TimingTest(unittest.TestCase): @@ -116,6 +132,26 @@ def fake_fn(): class OrchestratorTest(unittest.TestCase): + def test_process_measurements_use_slowest_successful_host(self): + status, mean_ms, std_ms, compile_ms = _aggregate_process_measurements( + np.asarray([ + [0, 10.0, 0.5, 100.0], + [0, 12.0, 0.7, 120.0], + ]) + ) + self.assertEqual(status, "ok") + self.assertEqual((mean_ms, std_ms, compile_ms), (12.0, 0.7, 120.0)) + + def test_process_measurements_reject_candidate_if_any_host_fails(self): + status, mean_ms, _, _ = _aggregate_process_measurements( + np.asarray([ + [0, 10.0, 0.5, 100.0], + [1, np.inf, np.inf, np.inf], + ]) + ) + self.assertEqual(status, "oom") + self.assertIsNone(mean_ms) + def test_smart_search_picks_measured_winner(self): res = grid_search(_MockRingBench(), mode="smart", iters=10, log=lambda *a, **k: None) self.assertIsNotNone(res.best) diff --git a/src/maxdiffusion/utils/ltx2_block_benchmark.py b/src/maxdiffusion/utils/ltx2_block_benchmark.py index 3aefd7ee4..7cf53af32 100644 --- a/src/maxdiffusion/utils/ltx2_block_benchmark.py +++ b/src/maxdiffusion/utils/ltx2_block_benchmark.py @@ -3,6 +3,7 @@ """ import os import functools +from types import SimpleNamespace os.environ.setdefault( "LIBTPU_INIT_ARGS", @@ -38,6 +39,7 @@ BenchResult, BlockBenchmark, grid_search, + local_tiled_seq_len, time_callable, ) @@ -55,11 +57,7 @@ def latent_seq_len(num_frames: int, height: int, width: int) -> int: def tiled_seq_len(full_seq: int, attention: str, context_shards: int, ulysses_shards: int) -> int: - _RING_VARIANTS = {"ulysses_ring_custom", "ulysses_ring_custom_bidir", "tokamax_ring", "ring"} - if attention in _RING_VARIANTS and ulysses_shards >= 1 and context_shards >= 1: - ring_shards = max(1, context_shards // max(1, ulysses_shards)) - return full_seq // ring_shards - return full_seq + return local_tiled_seq_len(full_seq, attention, context_shards, ulysses_shards) @functools.partial(nnx.jit, static_argnames=("num_frames", "height", "width")) @@ -204,16 +202,16 @@ def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2): return BenchResult(bq, bkv, cmp, "oom" if oom else "error", detail=msg[:200]) def _flash_block_sizes(self, bq, bkv, cmp): - from maxdiffusion.max_utils import CustomFlashBlockSizes - - return CustomFlashBlockSizes( - block_q=bq, - block_kv=bkv, - block_kv_compute=cmp, - block_kv_compute_in=cmp, - heads_per_tile=1, + candidate = max_utils.flash_block_sizes_for_candidate( + self._config.flash_block_sizes, + self._attention, + bq, + bkv, + cmp, vmem_limit_bytes=self._vmem, ) + candidate_config = SimpleNamespace(attention=self._attention, flash_block_sizes=candidate) + return max_utils.get_flash_block_sizes(candidate_config) def _build_model(self, bq, bkv, cmp): c = self._config @@ -232,6 +230,8 @@ def _build_model(self, bq, bkv, cmp): flash_min_seq_length=getattr(c, "flash_min_seq_length", 4096), ulysses_shards=getattr(c, "ulysses_shards", -1), ulysses_attention_chunks=getattr(c, "ulysses_attention_chunks", 1), + use_base2_exp=getattr(c, "use_base2_exp", False), + use_experimental_scheduler=getattr(c, "use_experimental_scheduler", False), remat_policy=getattr(c, "remat_policy", "NONE"), scan_layers=False, num_layers=1, diff --git a/src/maxdiffusion/utils/tile_size_grid_search.py b/src/maxdiffusion/utils/tile_size_grid_search.py index 38571a30b..a95b01b62 100644 --- a/src/maxdiffusion/utils/tile_size_grid_search.py +++ b/src/maxdiffusion/utils/tile_size_grid_search.py @@ -31,11 +31,24 @@ import os import sys import jax +import numpy as np +from jax.experimental import multihost_utils # --- the two granularities (see module docstring) -------------------------------- VPU_LANE = 128 # kernel hard floor: block sizes must be multiples of this MXU_TILE = 256 # 256x256 MXU: multiples of this fully pack the systolic array +PURE_RING_ATTENTION_KERNELS = frozenset({ + "tokamax_ring", + "tokamax_ring_custom", +}) +ULYSSES_RING_ATTENTION_KERNELS = frozenset({ + "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_bidir", +}) + def _ceil_div(a: int, b: int) -> int: return (a + b - 1) // b @@ -49,6 +62,26 @@ def _floor_to(x: int, m: int) -> int: return (x // m) * m +def local_tiled_seq_len(full_seq: int, attention: str, context_shards: int, ulysses_shards: int) -> int: + """Returns the sequence length tiled by one local ring kernel invocation.""" + context_shards = max(1, context_shards) + context_local_seq = _ceil_div(full_seq, context_shards) + if attention in PURE_RING_ATTENTION_KERNELS: + return context_local_seq + elif attention in ULYSSES_RING_ATTENTION_KERNELS: + if ulysses_shards < 1: + raise ValueError(f"{attention} requires ulysses_shards >= 1, got {ulysses_shards}.") + if context_shards % ulysses_shards: + raise ValueError( + f"context_shards={context_shards} must be divisible by ulysses_shards={ulysses_shards} for {attention}." + ) + # Production pads the global sequence to a context-shard multiple, then + # Ulysses gathers `ulysses_shards` local sequence chunks per head shard. + return context_local_seq * ulysses_shards + else: + return full_seq + + @dataclass(frozen=True) class Padding: seq_len: int @@ -289,6 +322,74 @@ class SearchResult: mode: str +_STATUS_TO_CODE = {"ok": 0, "oom": 1, "error": 2} + + +def _aggregate_process_measurements(measurements: np.ndarray) -> tuple[str, Optional[float], Optional[float], Optional[float]]: + """Aggregates candidate measurements, rejecting a candidate that fails on any host.""" + measurements = np.asarray(measurements).reshape((-1, 4)) + status_codes = measurements[:, 0].astype(np.int32) + if np.any(status_codes == _STATUS_TO_CODE["error"]): + return "error", None, None, None + if np.any(status_codes == _STATUS_TO_CODE["oom"]): + return "oom", None, None, None + return ( + "ok", + float(np.max(measurements[:, 1])), + float(np.max(measurements[:, 2])), + float(np.max(measurements[:, 3])), + ) + + +def _aggregate_process_result(result: BenchResult) -> BenchResult: + if jax.process_count() == 1: + return result + + local = np.asarray( + [ + _STATUS_TO_CODE.get(result.status, _STATUS_TO_CODE["error"]), + result.mean_ms if result.mean_ms is not None else np.inf, + result.std_ms if result.std_ms is not None else np.inf, + result.compile_ms if result.compile_ms is not None else np.inf, + ], + dtype=np.float32, + ) + gathered = multihost_utils.process_allgather(local, tiled=False) + status, mean_ms, std_ms, compile_ms = _aggregate_process_measurements(gathered) + failed_hosts = int(np.count_nonzero(np.asarray(gathered).reshape((-1, 4))[:, 0])) + detail = result.detail if status == "ok" else f"{status} on {failed_hosts}/{jax.process_count()} process(es)" + return BenchResult( + bq=result.bq, + bkv=result.bkv, + bkv_compute=result.bkv_compute, + status=status, + mean_ms=mean_ms, + std_ms=std_ms, + times_ms=result.times_ms, + compile_ms=compile_ms, + detail=detail, + ) + + +def _broadcast_winner(best: Optional[BenchResult], results: list[BenchResult]) -> Optional[BenchResult]: + if jax.process_count() == 1: + return best + + is_source = jax.process_index() == 0 + payload = np.zeros((4,), dtype=np.int64) + if is_source and best is not None: + payload[:] = (1, best.bq, best.bkv, best.bkv_compute) + payload = np.asarray(multihost_utils.broadcast_one_to_all(payload, is_source=is_source)) + if payload[0] == 0: + return None + + winner_key = tuple(int(value) for value in payload[1:]) + for result in results: + if (result.bq, result.bkv, result.bkv_compute) == winner_key and result.status == "ok": + return result + raise RuntimeError(f"Process 0 selected tile candidate {winner_key}, but it is unavailable on this process.") + + def smart_grid( q_seq: int, kv_seq: int, @@ -373,7 +474,12 @@ def grid_search( log(f"[tile-search] {bench.label}: q_seq={q_seq} kv_seq={kv_seq} mode={mode} " f"-> {len(pairs)} configs (iters={iters})") results: list[BenchResult] = [] for i, (bq, bkv) in enumerate(pairs, 1): + if jax.process_count() > 1: + multihost_utils.sync_global_devices(f"tile_search_candidate_{i}_start") r = bench.run(bq, bkv, bkv_compute=bkv, iters=iters, warmup=warmup) + if jax.process_count() > 1: + multihost_utils.sync_global_devices(f"tile_search_candidate_{i}_complete") + r = _aggregate_process_result(r) results.append(r) tag = "" if bq % MXU_TILE == 0 and bkv % MXU_TILE == 0 else " [½MXU]" compile_note = f" (compile {r.compile_ms/1e3:.0f}s, excluded)" if r.compile_ms else "" @@ -383,7 +489,8 @@ def grid_search( ) ok = [r for r in results if r.status == "ok" and r.mean_ms is not None] - best = min(ok, key=lambda r: r.mean_ms) if ok else None + best = min(ok, key=lambda r: r.mean_ms) if ok and jax.process_index() == 0 else None + best = _broadcast_winner(best, results) _emit(results, best, q_seq, kv_seq, mode, out_dir, log) return SearchResult(best, results, q_seq, kv_seq, mode) From 1657329c86af3f9f594c4b00e3aaf1e2e11411e2 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Mon, 27 Jul 2026 13:23:59 +0530 Subject: [PATCH 08/10] Implement persistent converted-weight caching for LTX2 --- src/maxdiffusion/configs/ltx2_3_video.yml | 2 + src/maxdiffusion/configs/ltx2_video.yml | 2 + .../models/converted_weights_cache.py | 629 ++++++++++++++++++ src/maxdiffusion/models/ltx2/ltx2_utils.py | 84 ++- src/maxdiffusion/models/wan/wan_utils.py | 5 +- .../pipelines/ltx2/ltx2_pipeline.py | 34 +- .../tests/converted_weights_cache_v2_test.py | 203 ++++++ .../tests/ltx2/test_ltx2_utils_loader.py | 49 ++ .../tests/ltx2/test_pipeline_ltx2.py | 18 + 9 files changed, 1010 insertions(+), 16 deletions(-) create mode 100644 src/maxdiffusion/models/converted_weights_cache.py create mode 100644 src/maxdiffusion/tests/converted_weights_cache_v2_test.py diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml index e3360350a..2ff4aecb0 100644 --- a/src/maxdiffusion/configs/ltx2_3_video.yml +++ b/src/maxdiffusion/configs/ltx2_3_video.yml @@ -23,6 +23,8 @@ names_which_can_be_offloaded: [] remat_policy: "NONE" jax_cache_dir: '' +# Local/shared filesystem directory for content-addressed PyTorch->Flax transformer weights ('' = disabled). +converted_weights_dir: '' weights_dtype: 'bfloat16' activations_dtype: 'bfloat16' text_encoder_dtype: 'bfloat16' diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml index 58877d4a6..8a2bbb32f 100644 --- a/src/maxdiffusion/configs/ltx2_video.yml +++ b/src/maxdiffusion/configs/ltx2_video.yml @@ -29,6 +29,8 @@ names_which_can_be_offloaded: [] remat_policy: "NONE" jax_cache_dir: '' +# Local/shared filesystem directory for content-addressed PyTorch->Flax transformer weights ('' = disabled). +converted_weights_dir: '' weights_dtype: 'bfloat16' activations_dtype: 'bfloat16' text_encoder_dtype: 'bfloat16' diff --git a/src/maxdiffusion/models/converted_weights_cache.py b/src/maxdiffusion/models/converted_weights_cache.py new file mode 100644 index 000000000..dfa6c5f34 --- /dev/null +++ b/src/maxdiffusion/models/converted_weights_cache.py @@ -0,0 +1,629 @@ +# 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. + +"""Content-addressed, multi-writer-safe converted-weights cache. + +This module intentionally depends on NumPy, but not JAX. A cache entry is an +immutable UUID-named generation plus an atomically replaced ``current.json`` +pointer. Consequently, readers either observe the old complete generation or +the new complete generation; they never consume a half-written cache. +""" + +from __future__ import annotations + +import concurrent.futures +import hashlib +import json +import numbers +import os +from pathlib import Path +import re +import shutil +import tempfile +import time +from typing import Any, Mapping, Optional +import uuid + +import numpy as np + +try: + import ml_dtypes +except ImportError: # Loading an ml_dtypes cache will fail cleanly without it. + ml_dtypes = None + + +_FORMAT_VERSION = 1 +_MAX_JSON_BYTES = 16 * 1024 * 1024 +_MAX_READ_WORKERS = 16 +_CONCURRENT_PUBLISH_WAIT_SECONDS = 2.0 +_URI_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*://") +_GENERATION_RE = re.compile(r"^[0-9a-f]{32}$") +_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9_.-]+") + + +def _canonical_json(value: Any) -> bytes: + return json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + +def _descriptor_digest(descriptor: Mapping[str, Any]) -> str: + return hashlib.sha256(_canonical_json(descriptor)).hexdigest() + + +def _require_local_path(path: os.PathLike[str] | str, name: str) -> str: + raw_path = os.fspath(path) + if not isinstance(raw_path, str) or not raw_path: + raise ValueError(f"{name} must be a non-empty local filesystem path") + if _URI_RE.match(raw_path): + raise ValueError(f"{name} must be a local filesystem path, not URI {raw_path!r}") + return os.path.abspath(os.path.expanduser(raw_path)) + + +def _key_component(component: Any) -> dict[str, Any]: + if isinstance(component, str): + return {"type": "str", "value": component} + if isinstance(component, numbers.Integral) and not isinstance(component, (bool, np.bool_)): + return {"type": "int", "value": str(int(component))} + raise TypeError(f"cache key components must be str or int, got {type(component).__name__}") + + +def _encode_key(key: tuple[Any, ...]) -> list[dict[str, Any]]: + if not key: + raise ValueError("cache keys must contain at least one component") + return [_key_component(component) for component in key] + + +def _decode_key(encoded: Any) -> tuple[Any, ...]: + if not isinstance(encoded, list) or not encoded: + raise ValueError("encoded cache key must be a non-empty JSON list") + decoded = [] + for component in encoded: + if not isinstance(component, dict) or set(component) != {"type", "value"}: + raise ValueError("malformed typed cache key component") + component_type = component["type"] + value = component["value"] + if component_type == "str" and isinstance(value, str): + decoded.append(value) + elif component_type == "int" and isinstance(value, str) and re.fullmatch(r"-?(0|[1-9][0-9]*)", value): + decoded.append(int(value)) + else: + raise ValueError("malformed typed cache key component") + return tuple(decoded) + + +def _as_key_tuple(key: Any) -> tuple[Any, ...]: + if isinstance(key, tuple): + result = key + elif isinstance(key, (str, numbers.Integral)) and not isinstance(key, (bool, np.bool_)): + result = (key,) + else: + raise TypeError(f"cache mapping keys must be tuples, strings, or integers, got {type(key).__name__}") + _encode_key(result) + return result + + +def _flatten_mapping(mapping: Mapping[Any, Any], prefix: tuple[Any, ...] = ()) -> dict[tuple[Any, ...], Any]: + if not isinstance(mapping, Mapping): + raise TypeError("expected a mapping") + flattened = {} + for raw_key, value in mapping.items(): + key = prefix + _as_key_tuple(raw_key) + if isinstance(value, Mapping): + children = _flatten_mapping(value, key) + overlap = flattened.keys() & children.keys() + if overlap: + raise ValueError(f"duplicate flattened cache key {next(iter(overlap))!r}") + flattened.update(children) + else: + if key in flattened: + raise ValueError(f"duplicate flattened cache key {key!r}") + flattened[key] = value + return flattened + + +def _shape_tuple(shape_or_value: Any) -> tuple[int, ...]: + shape = shape_or_value.shape if hasattr(shape_or_value, "shape") else shape_or_value + if shape is None: + raise TypeError("expected shape cannot be None") + if isinstance(shape, numbers.Integral) and not isinstance(shape, (bool, np.bool_)): + shape = (shape,) + try: + dimensions = tuple(shape) + except TypeError as error: + raise TypeError(f"expected shape must be an iterable of dimensions, got {shape!r}") from error + normalized = [] + for dimension in dimensions: + if not isinstance(dimension, numbers.Integral) or isinstance(dimension, (bool, np.bool_)): + raise TypeError(f"shape dimensions must be integers, got {dimension!r}") + dimension = int(dimension) + if dimension < 0: + raise ValueError(f"shape dimensions must be non-negative, got {dimension}") + normalized.append(dimension) + return tuple(normalized) + + +def _dtype_descriptor(dtype_or_value: Any) -> dict[str, Any]: + if hasattr(dtype_or_value, "dtype") and not isinstance(dtype_or_value, (str, type, np.dtype)): + dtype_or_value = dtype_or_value.dtype + dtype = np.dtype(dtype_or_value) + if dtype.hasobject or dtype.fields is not None or dtype.subdtype is not None: + raise TypeError(f"unsupported converted-weight dtype {dtype!r}") + module = getattr(dtype.type, "__module__", "") + if module == "ml_dtypes" or module.startswith("ml_dtypes."): + return { + "family": "ml_dtypes", + "name": dtype.name, + "itemsize": dtype.itemsize, + } + return { + "family": "numpy", + "name": dtype.name, + "str": dtype.str, + "itemsize": dtype.itemsize, + } + + +def _dtype_from_descriptor(descriptor: Any) -> np.dtype: + if not isinstance(descriptor, dict): + raise ValueError("malformed dtype descriptor") + family = descriptor.get("family") + name = descriptor.get("name") + itemsize = descriptor.get("itemsize") + if not isinstance(name, str) or not isinstance(itemsize, int) or itemsize <= 0: + raise ValueError("malformed dtype descriptor") + if family == "numpy" and set(descriptor) == {"family", "name", "str", "itemsize"}: + dtype_string = descriptor["str"] + if not isinstance(dtype_string, str): + raise ValueError("malformed NumPy dtype descriptor") + dtype = np.dtype(dtype_string) + elif family == "ml_dtypes" and set(descriptor) == {"family", "name", "itemsize"}: + if ml_dtypes is None: + raise ValueError(f"ml_dtypes is required to load logical dtype {name!r}") + dtype_type = getattr(ml_dtypes, name, None) + if dtype_type is None: + raise ValueError(f"unknown ml_dtypes dtype {name!r}") + dtype = np.dtype(dtype_type) + else: + raise ValueError("malformed dtype descriptor") + if dtype.name != name or dtype.itemsize != itemsize or _dtype_descriptor(dtype) != descriptor: + raise ValueError("dtype descriptor does not resolve exactly") + return dtype + + +def _schema( + expected_shapes: Mapping[Any, Any], + expected_dtypes: Mapping[Any, Any], +) -> list[dict[str, Any]]: + shapes = _flatten_mapping(expected_shapes) + dtypes = _flatten_mapping(expected_dtypes) + if not shapes: + raise ValueError("converted-weights schema cannot be empty") + if shapes.keys() != dtypes.keys(): + missing_dtypes = shapes.keys() - dtypes.keys() + missing_shapes = dtypes.keys() - shapes.keys() + raise ValueError( + "shape/dtype key sets differ" + f" (missing dtypes: {sorted(map(repr, missing_dtypes))}," + f" missing shapes: {sorted(map(repr, missing_shapes))})" + ) + entries = [ + { + "key": _encode_key(key), + "shape": list(_shape_tuple(shapes[key])), + "dtype": _dtype_descriptor(dtypes[key]), + } + for key in shapes + ] + entries.sort(key=lambda entry: _canonical_json(entry["key"])) + return entries + + +def _source_identity(source_file: os.PathLike[str] | str) -> dict[str, Any]: + source_path = _require_local_path(source_file, "source file") + if not os.path.isfile(source_path): + raise FileNotFoundError(f"converted-weight source file does not exist: {source_path}") + resolved_path = os.path.realpath(source_path) + stat = os.stat(resolved_path) + + original_parts = Path(source_path).parts + lowered_parts = [part.casefold() for part in original_parts] + snapshot_indices = [index for index, part in enumerate(lowered_parts) if part == "snapshots"] + if snapshot_indices: + snapshot_index = snapshot_indices[-1] + if snapshot_index + 2 < len(original_parts): + revision = original_parts[snapshot_index + 1] + relative_path = Path(*original_parts[snapshot_index + 2 :]).as_posix() + resolved_parts = Path(resolved_path).parts + resolved_lowered = [part.casefold() for part in resolved_parts] + blob_indices = [index for index, part in enumerate(resolved_lowered) if part == "blobs"] + content_id = None + if blob_indices and blob_indices[-1] + 1 < len(resolved_parts): + content_id = "/".join(resolved_parts[blob_indices[-1] + 1 :]) + identity = { + "kind": "huggingface_snapshot", + "revision": revision, + "relative_path": relative_path, + "size": stat.st_size, + } + if content_id is not None: + identity["content_id"] = content_id + else: + # Snapshot revisions are immutable in a healthy HF cache. These + # fields additionally invalidate locally modified/copy-mode entries. + identity["mtime_ns"] = stat.st_mtime_ns + return identity + + return { + "kind": "local_file", + "path": resolved_path, + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + + +def _json_metadata(value: Any) -> Any: + try: + # Round-tripping detaches caller-owned containers and normalizes tuples to + # JSON lists. NaN and Infinity are deliberately rejected. + return json.loads(_canonical_json(value)) + except (TypeError, ValueError) as error: + raise TypeError("extra_metadata must contain only finite JSON values") from error + + +def build_cache_descriptor( + namespace: str, + converter_abi: str | int, + model_id: str, + subfolder: str, + source_files: Any, + scan_layers: bool, + expected_shapes: Mapping[Any, Any], + expected_dtypes: Mapping[Any, Any], + extra_metadata: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: + """Builds the complete content identity and output schema for a conversion.""" + if not isinstance(namespace, str) or not namespace: + raise ValueError("namespace must be a non-empty string") + if isinstance(converter_abi, numbers.Integral) and not isinstance(converter_abi, (bool, np.bool_)): + converter_abi = int(converter_abi) + if converter_abi <= 0: + raise ValueError("integer converter_abi must be positive") + elif not isinstance(converter_abi, str) or not converter_abi: + raise ValueError("converter_abi must be a non-empty string or positive integer") + if not isinstance(model_id, str) or not model_id: + raise ValueError("model_id must be a non-empty string") + if _URI_RE.match(model_id): + raise ValueError(f"model_id must be a Hugging Face id or local path, not URI {model_id!r}") + if not isinstance(subfolder, str): + raise TypeError("subfolder must be a string") + if not isinstance(scan_layers, (bool, np.bool_)): + raise TypeError("scan_layers must be a boolean") + if isinstance(source_files, (str, os.PathLike)): + source_files = [source_files] + else: + source_files = list(source_files) + if not source_files: + raise ValueError("source_files must contain at least one checkpoint file") + + sources = [_source_identity(source_file) for source_file in source_files] + sources.sort(key=_canonical_json) + metadata = {} if extra_metadata is None else _json_metadata(extra_metadata) + if not isinstance(metadata, dict): + raise TypeError("extra_metadata must be a mapping") + descriptor = { + "format_version": _FORMAT_VERSION, + "namespace": namespace, + "converter_abi": converter_abi, + "model": { + "id": model_id, + "subfolder": subfolder, + }, + "sources": sources, + "conversion": { + "scan_layers": bool(scan_layers), + "extra_metadata": metadata, + }, + "schema": _schema(expected_shapes, expected_dtypes), + } + # Ensure the descriptor has one deterministic JSON representation now, not + # when a cache lookup is already in progress. + _canonical_json(descriptor) + return descriptor + + +def _validated_schema(descriptor: Mapping[str, Any]) -> list[tuple[tuple[Any, ...], tuple[int, ...], dict[str, Any]]]: + if not isinstance(descriptor, Mapping) or descriptor.get("format_version") != _FORMAT_VERSION: + raise ValueError("unsupported or malformed cache descriptor") + if not isinstance(descriptor.get("namespace"), str) or not descriptor["namespace"]: + raise ValueError("malformed cache descriptor namespace") + converter_abi = descriptor.get("converter_abi") + if not ( + (isinstance(converter_abi, str) and converter_abi) + or ( + isinstance(converter_abi, numbers.Integral) + and not isinstance(converter_abi, (bool, np.bool_)) + and converter_abi > 0 + ) + ): + raise ValueError("malformed cache descriptor converter ABI") + entries = descriptor.get("schema") + if not isinstance(entries, list) or not entries: + raise ValueError("malformed cache descriptor schema") + validated = [] + seen_keys = set() + for entry in entries: + if not isinstance(entry, dict) or set(entry) != {"key", "shape", "dtype"}: + raise ValueError("malformed cache descriptor schema entry") + key = _decode_key(entry["key"]) + if key in seen_keys: + raise ValueError(f"duplicate cache descriptor key {key!r}") + seen_keys.add(key) + shape = _shape_tuple(entry["shape"]) + dtype_descriptor = entry["dtype"] + _dtype_from_descriptor(dtype_descriptor) + validated.append((key, shape, dtype_descriptor)) + if entries != sorted(entries, key=lambda entry: _canonical_json(entry["key"])): + raise ValueError("cache descriptor schema is not canonically ordered") + _canonical_json(descriptor) + return validated + + +def cache_entry_dir(base_dir: os.PathLike[str] | str, descriptor: Mapping[str, Any]) -> str: + """Returns a path-safe, descriptor-addressed cache entry directory.""" + base_path = _require_local_path(base_dir, "converted weights cache root") + _validated_schema(descriptor) + digest = _descriptor_digest(descriptor) + namespace = _SAFE_NAME_RE.sub("-", descriptor["namespace"]).strip("._-")[:48] or "weights" + return os.path.join(base_path, f"{namespace}-{digest}") + + +def _read_json(path: str) -> Any: + size = os.path.getsize(path) + if size <= 0 or size > _MAX_JSON_BYTES: + raise ValueError(f"invalid JSON metadata size {size}") + with open(path, "rb") as file: + return json.loads(file.read().decode("utf-8")) + + +def _unsigned_dtype(itemsize: int) -> np.dtype: + try: + return np.dtype({1: np.uint8, 2: np.uint16, 4: np.uint32, 8: np.uint64}[itemsize]) + except KeyError as error: + raise ValueError(f"cannot bit-store custom dtype with itemsize {itemsize}") from error + + +def _load_generation( + generation_dir: str, + descriptor: Mapping[str, Any], + schema: list[tuple[tuple[Any, ...], tuple[int, ...], dict[str, Any]]], +) -> dict[tuple[Any, ...], np.ndarray]: + digest = _descriptor_digest(descriptor) + manifest = _read_json(os.path.join(generation_dir, "manifest.json")) + if not isinstance(manifest, dict) or set(manifest) != {"format_version", "descriptor_digest", "tensors"}: + raise ValueError("malformed converted-weights manifest") + if manifest["format_version"] != _FORMAT_VERSION or manifest["descriptor_digest"] != digest: + raise ValueError("converted-weights manifest descriptor mismatch") + tensors = manifest["tensors"] + if not isinstance(tensors, list) or len(tensors) != len(schema): + raise ValueError("converted-weights manifest tensor count mismatch") + + def load_one(index_and_expected): + index, (expected_key, expected_shape, expected_dtype_descriptor) = index_and_expected + metadata = tensors[index] + required_fields = {"key", "shape", "dtype", "file", "bitview", "storage_dtype"} + if not isinstance(metadata, dict) or set(metadata) != required_fields: + raise ValueError("malformed converted-weights tensor metadata") + expected_file = f"tensor-{index:06d}.npy" + if ( + metadata["key"] != descriptor["schema"][index]["key"] + or metadata["shape"] != list(expected_shape) + or metadata["dtype"] != expected_dtype_descriptor + or metadata["file"] != expected_file + or not isinstance(metadata["bitview"], bool) + or not isinstance(metadata["storage_dtype"], str) + ): + raise ValueError(f"converted-weights metadata mismatch for tensor {index}") + tensor_path = os.path.join(generation_dir, expected_file) + with open(tensor_path, "rb") as file: + value = np.load(file, allow_pickle=False) + if not isinstance(value, np.ndarray) or tuple(value.shape) != expected_shape: + raise ValueError(f"converted-weights shape mismatch for {expected_key!r}") + if value.dtype.str != metadata["storage_dtype"]: + raise ValueError(f"converted-weights storage dtype mismatch for {expected_key!r}") + logical_dtype = _dtype_from_descriptor(expected_dtype_descriptor) + if metadata["bitview"]: + storage_dtype = _unsigned_dtype(logical_dtype.itemsize) + if value.dtype != storage_dtype: + raise ValueError(f"converted-weights bit storage mismatch for {expected_key!r}") + value = value.view(logical_dtype) + if _dtype_descriptor(value.dtype) != expected_dtype_descriptor: + raise ValueError(f"converted-weights logical dtype mismatch for {expected_key!r}") + return expected_key, value + + workers = min(_MAX_READ_WORKERS, len(schema)) + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor: + loaded_items = list(executor.map(load_one, enumerate(schema))) + return dict(loaded_items) + + +def try_load_converted_weights( + entry_dir: os.PathLike[str] | str, + descriptor: Mapping[str, Any], +) -> Optional[dict[tuple[Any, ...], np.ndarray]]: + """Eagerly loads and exactly validates the current immutable generation.""" + local_entry_dir = _require_local_path(entry_dir, "converted weights cache entry") + try: + schema = _validated_schema(descriptor) + digest = _descriptor_digest(descriptor) + pointer = _read_json(os.path.join(local_entry_dir, "current.json")) + if not isinstance(pointer, dict) or set(pointer) != {"format_version", "descriptor_digest", "generation"}: + raise ValueError("malformed converted-weights cache pointer") + generation = pointer["generation"] + if ( + pointer["format_version"] != _FORMAT_VERSION + or pointer["descriptor_digest"] != digest + or not isinstance(generation, str) + or not _GENERATION_RE.fullmatch(generation) + ): + raise ValueError("converted-weights cache pointer mismatch") + generation_dir = os.path.join(local_entry_dir, "generations", generation) + return _load_generation(generation_dir, descriptor, schema) + except (EOFError, json.JSONDecodeError, OSError, TypeError, ValueError): + return None + + +def _write_json(path: str, value: Any) -> None: + payload = _canonical_json(value) + with open(path, "xb") as file: + file.write(payload) + file.flush() + os.fsync(file.fileno()) + + +def _fsync_directory(path: str) -> None: + flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) + try: + descriptor = os.open(path, flags) + except OSError: + return + try: + os.fsync(descriptor) + except OSError: + pass + finally: + os.close(descriptor) + + +def save_converted_weights( + entry_dir: os.PathLike[str] | str, + descriptor: Mapping[str, Any], + flat_state_dict: Mapping[Any, Any], +) -> bool: + """Validates and publishes a complete new cache generation. + + False means that validation or publication failed. No pointer is changed + until every tensor and the generation manifest have been written and read + back successfully. + """ + local_entry_dir = _require_local_path(entry_dir, "converted weights cache entry") + temporary_generation = None + generation_dir = None + pointer_replaced = False + pointer_temp = None + try: + schema = _validated_schema(descriptor) + if not isinstance(flat_state_dict, Mapping): + return False + values = {_as_key_tuple(key): np.asarray(value) for key, value in flat_state_dict.items()} + expected_keys = {key for key, _, _ in schema} + if values.keys() != expected_keys: + return False + for key, shape, dtype_descriptor in schema: + value = values[key] + if tuple(value.shape) != shape or _dtype_descriptor(value.dtype) != dtype_descriptor: + return False + + generations_dir = os.path.join(local_entry_dir, "generations") + os.makedirs(generations_dir, exist_ok=True) + generation = uuid.uuid4().hex + temporary_generation = tempfile.mkdtemp(prefix=f".{generation}.", suffix=".tmp", dir=generations_dir) + manifest_entries = [] + for index, (key, shape, dtype_descriptor) in enumerate(schema): + value = np.ascontiguousarray(values[key]) + logical_dtype = _dtype_from_descriptor(dtype_descriptor) + bitview = dtype_descriptor["family"] == "ml_dtypes" + stored = value.view(_unsigned_dtype(logical_dtype.itemsize)) if bitview else value + filename = f"tensor-{index:06d}.npy" + tensor_path = os.path.join(temporary_generation, filename) + with open(tensor_path, "xb") as file: + np.save(file, stored, allow_pickle=False) + file.flush() + os.fsync(file.fileno()) + manifest_entries.append( + { + "key": descriptor["schema"][index]["key"], + "shape": list(shape), + "dtype": dtype_descriptor, + "file": filename, + "bitview": bitview, + "storage_dtype": stored.dtype.str, + } + ) + digest = _descriptor_digest(descriptor) + manifest = { + "format_version": _FORMAT_VERSION, + "descriptor_digest": digest, + "tensors": manifest_entries, + } + _write_json(os.path.join(temporary_generation, "manifest.json"), manifest) + _fsync_directory(temporary_generation) + + generation_dir = os.path.join(generations_dir, generation) + os.replace(temporary_generation, generation_dir) + temporary_generation = None + _fsync_directory(generations_dir) + + pointer = { + "format_version": _FORMAT_VERSION, + "descriptor_digest": digest, + "generation": generation, + } + pointer_temp = os.path.join(local_entry_dir, f".current.{uuid.uuid4().hex}.tmp") + _write_json(pointer_temp, pointer) + os.replace(pointer_temp, os.path.join(local_entry_dir, "current.json")) + pointer_replaced = True + pointer_temp = None + _fsync_directory(local_entry_dir) + return True + except (EOFError, json.JSONDecodeError, OSError, TypeError, ValueError): + # On Windows and some shared filesystems, two valid writers can race while + # replacing the pointer. Treat a concurrently published, fully validated + # generation as success; never infer success from the filesystem error + # alone. + deadline = time.monotonic() + _CONCURRENT_PUBLISH_WAIT_SECONDS + while True: + try: + if try_load_converted_weights(local_entry_dir, descriptor) is not None: + return True + except (OSError, TypeError, ValueError): + pass + if time.monotonic() >= deadline: + return False + time.sleep(0.01) + finally: + if temporary_generation is not None: + shutil.rmtree(temporary_generation, ignore_errors=True) + if pointer_temp is not None: + try: + os.unlink(pointer_temp) + except OSError: + pass + if generation_dir is not None and not pointer_replaced: + # If we failed to publish, clean up our generation folder. + # To be absolutely safe against rare NFS success-but-raise anomalies, + # we verify we didn't accidentally become the current pointer first. + try: + pointer_path = os.path.join(local_entry_dir, "current.json") + if os.path.exists(pointer_path): + current = _read_json(pointer_path) + if isinstance(current, dict) and current.get("generation") == generation: + return + except (OSError, ValueError, TypeError, json.JSONDecodeError): + pass + shutil.rmtree(generation_dir, ignore_errors=True) diff --git a/src/maxdiffusion/models/ltx2/ltx2_utils.py b/src/maxdiffusion/models/ltx2/ltx2_utils.py index 21e96a006..30d464669 100644 --- a/src/maxdiffusion/models/ltx2/ltx2_utils.py +++ b/src/maxdiffusion/models/ltx2/ltx2_utils.py @@ -26,6 +26,14 @@ from safetensors import safe_open from flax.traverse_util import unflatten_dict, flatten_dict from ..modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor, torch2jax, validate_flax_state_dict) +from ..converted_weights_cache import ( + build_cache_descriptor, + cache_entry_dir, + save_converted_weights, + try_load_converted_weights, +) + +LTX2_CONVERTED_WEIGHTS_CACHE_ABI = 1 KNOWN_UPSAMPLER_CONFIGS = { "ltx-2.3-spatial-upscaler-x2-1.0.safetensors": { @@ -201,6 +209,14 @@ def _get_eval_shape(value) -> tuple[int, ...]: raise ValueError(f"Unable to determine the initialized shape for {type(value).__name__}") +def _get_eval_dtype(value) -> np.dtype: + if hasattr(value, "dtype"): + return np.dtype(value.dtype) + if hasattr(value, "value") and hasattr(value.value, "dtype"): + return np.dtype(value.value.dtype) + raise ValueError(f"Unable to determine the initialized dtype for {type(value).__name__}") + + def _get_scanned_layer_shapes(flattened_eval_shapes): scanned_layer_shapes = {} for key, value in flattened_eval_shapes.items(): @@ -234,7 +250,15 @@ def load_transformer_weights( num_layers: Optional[int] = None, scan_layers: bool = True, subfolder: str = "transformer", + cast_dtype_fn=None, + converted_cache_dir: str = "", ): + """Loads and converts an LTX2 transformer checkpoint into host arrays. + + When ``converted_cache_dir`` is set, the final-dtype flat tree is cached + under a content-addressed identity derived from the resolved checkpoint, + converter ABI, scan mode, and exact initialized parameter schema. + """ import threading import concurrent.futures import time @@ -276,6 +300,36 @@ def resolve_shard_path(model_file): except EntryNotFoundError: return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename="diffusion_pytorch_model.bin") + checkpoint_files = [resolve_shard_path(model_file) for model_file in shards] + + expected_shapes = {tuple(key): _get_eval_shape(value) for key, value in flattened_dict.items()} + expected_dtypes = { + tuple(key): np.dtype(cast_dtype_fn(tuple(key))) if cast_dtype_fn is not None else _get_eval_dtype(value) + for key, value in flattened_dict.items() + } + converted_cache_descriptor = None + converted_cache_entry = "" + if converted_cache_dir: + converted_cache_descriptor = build_cache_descriptor( + namespace="ltx2-transformer", + converter_abi=LTX2_CONVERTED_WEIGHTS_CACHE_ABI, + model_id=pretrained_model_name_or_path, + subfolder=subfolder, + source_files=checkpoint_files, + scan_layers=scan_layers, + expected_shapes=expected_shapes, + expected_dtypes=expected_dtypes, + ) + converted_cache_entry = cache_entry_dir(converted_cache_dir, converted_cache_descriptor) + t_cache_load = time.perf_counter() + cached = try_load_converted_weights(converted_cache_entry, converted_cache_descriptor) + if cached is not None: + max_logging.log( + f"Loaded converted {subfolder or 'transformer'} weights from {converted_cache_entry} " + f"in {time.perf_counter() - t_cache_load:.1f}s" + ) + return unflatten_dict(cached) + t_start = time.perf_counter() flax_state_dict = {} @@ -353,19 +407,28 @@ def process_tensor(pt_key, tensor): stacked = flax_state_dict.get(flax_key) if stacked is None: - stacked = np.zeros(expected_shape, dtype=flax_tensor.dtype) + target_dtype = ( + np.dtype(cast_dtype_fn(flax_key)) + if cast_dtype_fn is not None + else expected_dtypes.get(flax_key, np.dtype(flax_tensor.dtype)) + ) + stacked = np.zeros(expected_shape, dtype=target_dtype) flax_state_dict[flax_key] = stacked stacked[block_index] = flax_tensor populated.add(block_index) else: - value = np.array(flax_tensor, dtype=flax_tensor.dtype, copy=True, order="C") + target_dtype = ( + np.dtype(cast_dtype_fn(flax_key)) + if cast_dtype_fn is not None + else expected_dtypes.get(flax_key, np.dtype(flax_tensor.dtype)) + ) + value = np.array(flax_tensor, dtype=target_dtype, copy=True, order="C") with dict_lock: flax_state_dict[flax_key] = value chunk_size = 32 safetensors_tasks = [] - for model_file in shards: - ckpt_shard_path = resolve_shard_path(model_file) + for ckpt_shard_path in checkpoint_files: if ckpt_shard_path.endswith(".safetensors"): with safe_open(ckpt_shard_path, framework="pt") as f: shard_keys = list(f.keys()) @@ -391,6 +454,19 @@ def process_tensor(pt_key, tensor): raise ValueError(f"Missing scanned layer indices: {'; '.join(missing_layers)}") validate_flax_state_dict(eval_shapes, flax_state_dict) + if converted_cache_entry: + t_cache_save = time.perf_counter() + if jax.process_index() == 0: + if save_converted_weights(converted_cache_entry, converted_cache_descriptor, flax_state_dict): + max_logging.log( + f"Saved converted {subfolder or 'transformer'} weights to {converted_cache_entry} " + f"in {time.perf_counter() - t_cache_save:.1f}s" + ) + else: + max_logging.log( + f"Unable to publish converted {subfolder or 'transformer'} weights to {converted_cache_entry}; " + "continuing without a reusable cache entry" + ) flax_state_dict = unflatten_dict(flax_state_dict) max_logging.log(f"Converted weights in {time.perf_counter() - t_start:.1f}s") return flax_state_dict diff --git a/src/maxdiffusion/models/wan/wan_utils.py b/src/maxdiffusion/models/wan/wan_utils.py index e60ffd31c..4c9053219 100644 --- a/src/maxdiffusion/models/wan/wan_utils.py +++ b/src/maxdiffusion/models/wan/wan_utils.py @@ -516,8 +516,9 @@ def convert_chunk(ckpt_shard_path, chunk_keys): validate_flax_state_dict(eval_shapes, flax_state_dict) if converted_cache_dir and not os.path.isdir(converted_cache_dir): t_save = time.perf_counter() - save_converted_weights(converted_cache_dir, flax_state_dict) - max_logging.log(f"Saved converted-weights cache to {converted_cache_dir} in {time.perf_counter() - t_save:.1f}s") + if jax.process_index() == 0: + save_converted_weights(converted_cache_dir, flax_state_dict) + max_logging.log(f"Saved converted-weights cache to {converted_cache_dir} in {time.perf_counter() - t_save:.1f}s") flax_state_dict = unflatten_dict(flax_state_dict) max_logging.log(f"Converted {subfolder or 'transformer'} weights to host arrays in {time.perf_counter() - t_start:.1f}s") return flax_state_dict diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index 932cd6187..a6b66bb55 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -238,23 +238,35 @@ def maybe_rescale(noise, text_noise, scale): logger = logging.get_logger(__name__) +_CAST_EXCLUSION_KEYWORDS = ( + "norm", # All LayerNorm/GroupNorm parameters. + "condition_embedder", # The entire time/text conditioning module. + "scale_shift_table", # The final and AdaLN scale/shift tables. +) + + +def _is_cast_excluded(path_str: str) -> bool: + return any(keyword in path_str.lower() for keyword in _CAST_EXCLUSION_KEYWORDS) + + +def _final_param_dtype(flax_key: tuple, dtype_to_cast) -> np.dtype: + """Returns the final host dtype used for a converted transformer parameter.""" + path_str = ".".join(str(key) for key in flax_key) + if _is_cast_excluded(path_str): + return np.dtype(jnp.float32) + return np.dtype(dtype_to_cast) + def cast_with_exclusion(path, x, dtype_to_cast): """ Casts arrays to dtype_to_cast, but keeps params from any 'norm' layer in float32. """ - exclusion_keywords = [ - "norm", # For all LayerNorm/GroupNorm layers - "condition_embedder", # The entire time/text conditioning module - "scale_shift_table", # Catches both the final and the AdaLN tables - ] - path_str = ".".join(str(k.key) if isinstance(k, jax.tree_util.DictKey) else str(k) for k in path) - if any(keyword in path_str.lower() for keyword in exclusion_keywords): - return x.astype(jnp.float32) - else: - return x.astype(dtype_to_cast) + target_dtype = jnp.float32 if _is_cast_excluded(path_str) else dtype_to_cast + if x.dtype == np.dtype(target_dtype): + return x + return x.astype(target_dtype) def _add_sharding_rule(vs: nnx.Variable, logical_axis_rules) -> nnx.Variable: @@ -336,6 +348,8 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): "cpu", scan_layers=getattr(config, "scan_layers", True), subfolder=subfolder, + cast_dtype_fn=partial(_final_param_dtype, dtype_to_cast=config.weights_dtype), + converted_cache_dir=getattr(config, "converted_weights_dir", ""), ) params = jax.tree_util.tree_map_with_path( diff --git a/src/maxdiffusion/tests/converted_weights_cache_v2_test.py b/src/maxdiffusion/tests/converted_weights_cache_v2_test.py new file mode 100644 index 000000000..57913c875 --- /dev/null +++ b/src/maxdiffusion/tests/converted_weights_cache_v2_test.py @@ -0,0 +1,203 @@ +# 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 the content-addressed converted-weights cache.""" + +import concurrent.futures +import json +import os +import tempfile +import unittest + +import ml_dtypes +import numpy as np + +from maxdiffusion.models.converted_weights_cache import ( + build_cache_descriptor, + cache_entry_dir, + save_converted_weights, + try_load_converted_weights, +) + + +class ConvertedWeightsCacheV2Test(unittest.TestCase): + + def setUp(self): + self._temporary_directory = tempfile.TemporaryDirectory() + self.cache_root = os.path.join(self._temporary_directory.name, "cache") + self.source_file = os.path.join(self._temporary_directory.name, "model.safetensors") + with open(self.source_file, "wb") as file: + file.write(b"checkpoint-v1") + self.flat_state = { + ("blocks", "0", "kernel"): np.arange(12, dtype=np.float32).reshape(3, 4), + ("blocks", 0, "kernel"): np.arange(6, dtype=np.float16).reshape(2, 3), + ("blocks", "bf16", "kernel"): np.arange(8, dtype=np.float32).astype(ml_dtypes.bfloat16).reshape(2, 4), + } + + def tearDown(self): + self._temporary_directory.cleanup() + + def _descriptor( + self, + *, + model_id="org/ltx-2", + converter_abi=1, + source_files=None, + shape_override=None, + dtype_override=None, + ): + shapes = {key: value.shape for key, value in self.flat_state.items()} + dtypes = {key: value.dtype for key, value in self.flat_state.items()} + if shape_override is not None: + shapes[("blocks", "0", "kernel")] = shape_override + if dtype_override is not None: + dtypes[("blocks", "0", "kernel")] = dtype_override + return build_cache_descriptor( + namespace="ltx2-transformer", + converter_abi=converter_abi, + model_id=model_id, + subfolder="transformer", + source_files=[self.source_file] if source_files is None else source_files, + scan_layers=True, + expected_shapes=shapes, + expected_dtypes=dtypes, + extra_metadata={"variant": "ltx2.3", "num_layers": 48}, + ) + + def _entry(self, descriptor): + return cache_entry_dir(self.cache_root, descriptor) + + def test_round_trip_preserves_bfloat16_bits_and_typed_keys(self): + descriptor = self._descriptor() + entry_dir = self._entry(descriptor) + self.assertTrue(save_converted_weights(entry_dir, descriptor, self.flat_state)) + + loaded = try_load_converted_weights(entry_dir, descriptor) + + self.assertIsNotNone(loaded) + self.assertEqual(set(loaded), set(self.flat_state)) + self.assertIn(("blocks", "0", "kernel"), loaded) + self.assertIn(("blocks", 0, "kernel"), loaded) + for key, expected in self.flat_state.items(): + self.assertEqual(loaded[key].dtype, expected.dtype) + if expected.dtype == np.dtype(ml_dtypes.bfloat16): + np.testing.assert_array_equal(loaded[key].view(np.uint16), expected.view(np.uint16)) + else: + np.testing.assert_array_equal(loaded[key], expected) + + def test_descriptor_changes_get_distinct_content_addresses(self): + descriptors = [ + self._descriptor(), + self._descriptor(model_id="org/ltx-2.3"), + self._descriptor(converter_abi=2), + self._descriptor(shape_override=(4, 3)), + self._descriptor(dtype_override=np.float64), + ] + entry_dirs = {self._entry(descriptor) for descriptor in descriptors} + self.assertEqual(len(entry_dirs), len(descriptors)) + + second_source = os.path.join(self._temporary_directory.name, "model-v2.safetensors") + with open(second_source, "wb") as file: + file.write(b"checkpoint-v2-and-a-different-size") + source_changed = self._descriptor(source_files=[second_source]) + self.assertNotIn(self._entry(source_changed), entry_dirs) + + old_descriptor = self._descriptor() + with open(self.source_file, "ab") as file: + file.write(b"-modified") + modified_descriptor = self._descriptor() + self.assertNotEqual(self._entry(old_descriptor), self._entry(modified_descriptor)) + + def test_malformed_missing_and_truncated_cache_return_none(self): + descriptor = self._descriptor() + entry_dir = self._entry(descriptor) + self.assertIsNone(try_load_converted_weights(entry_dir, descriptor)) + + os.makedirs(entry_dir) + with open(os.path.join(entry_dir, "current.json"), "wb") as file: + file.write(b"{") + self.assertIsNone(try_load_converted_weights(entry_dir, descriptor)) + + self.assertTrue(save_converted_weights(entry_dir, descriptor, self.flat_state)) + with open(os.path.join(entry_dir, "current.json"), encoding="utf-8") as file: + pointer = json.load(file) + tensor_path = os.path.join(entry_dir, "generations", pointer["generation"], "tensor-000000.npy") + with open(tensor_path, "wb") as file: + file.write(b"\x93NUMPY") + self.assertIsNone(try_load_converted_weights(entry_dir, descriptor)) + + def test_corruption_is_repaired_by_a_new_generation(self): + descriptor = self._descriptor() + entry_dir = self._entry(descriptor) + self.assertTrue(save_converted_weights(entry_dir, descriptor, self.flat_state)) + pointer_path = os.path.join(entry_dir, "current.json") + with open(pointer_path, encoding="utf-8") as file: + old_generation = json.load(file)["generation"] + old_tensor = os.path.join(entry_dir, "generations", old_generation, "tensor-000001.npy") + with open(old_tensor, "wb") as file: + file.write(b"corrupt") + self.assertIsNone(try_load_converted_weights(entry_dir, descriptor)) + + self.assertTrue(save_converted_weights(entry_dir, descriptor, self.flat_state)) + with open(pointer_path, encoding="utf-8") as file: + new_generation = json.load(file)["generation"] + self.assertNotEqual(new_generation, old_generation) + loaded = try_load_converted_weights(entry_dir, descriptor) + self.assertIsNotNone(loaded) + for key, expected in self.flat_state.items(): + np.testing.assert_array_equal(loaded[key], expected) + + def test_save_rejects_key_shape_and_actual_dtype_mismatches(self): + descriptor = self._descriptor() + entry_dir = self._entry(descriptor) + + missing_key = dict(self.flat_state) + missing_key.pop(("blocks", 0, "kernel")) + self.assertFalse(save_converted_weights(entry_dir, descriptor, missing_key)) + + wrong_shape = dict(self.flat_state) + wrong_shape[("blocks", "0", "kernel")] = np.zeros((4, 3), dtype=np.float32) + self.assertFalse(save_converted_weights(entry_dir, descriptor, wrong_shape)) + + wrong_dtype = dict(self.flat_state) + wrong_dtype[("blocks", "0", "kernel")] = self.flat_state[("blocks", "0", "kernel")].astype(np.float64) + self.assertFalse(save_converted_weights(entry_dir, descriptor, wrong_dtype)) + self.assertIsNone(try_load_converted_weights(entry_dir, descriptor)) + + def test_concurrent_writers_publish_only_valid_generations(self): + descriptor = self._descriptor() + entry_dir = self._entry(descriptor) + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + results = list( + executor.map( + lambda unused: save_converted_weights(entry_dir, descriptor, self.flat_state), + range(16), + ) + ) + self.assertTrue(all(results)) + loaded = try_load_converted_weights(entry_dir, descriptor) + self.assertIsNotNone(loaded) + for key, expected in self.flat_state.items(): + np.testing.assert_array_equal(loaded[key], expected) + + def test_non_local_roots_are_rejected(self): + descriptor = self._descriptor() + with self.assertRaisesRegex(ValueError, "local filesystem path"): + cache_entry_dir("gs://bucket/converted", descriptor) + with self.assertRaisesRegex(ValueError, "local filesystem path"): + self._descriptor(source_files=["https://example.invalid/model.safetensors"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py b/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py index 876b4cb0b..7e12382b9 100644 --- a/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py +++ b/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py @@ -16,6 +16,8 @@ import concurrent.futures import json +import os +import tempfile import unittest from unittest.mock import MagicMock, mock_open, patch @@ -162,6 +164,53 @@ def download(_model_name, *, filename, **_kwargs): np.array([[0.0], [1.0]]), ) + def test_converted_cache_hit_skips_bin_reload(self): + checkpoint = { + "transformer_blocks.0.bias": torch.tensor([1.0]), + "transformer_blocks.1.bias": torch.tensor([2.0]), + } + + with tempfile.TemporaryDirectory() as tmp_dir: + checkpoint_path = os.path.join(tmp_dir, "diffusion_pytorch_model.bin") + with open(checkpoint_path, "wb") as checkpoint_file: + checkpoint_file.write(b"source identity") + + def download(_model_name, *, filename, **_kwargs): + if filename.endswith(".json") or filename.endswith(".safetensors"): + raise EntryNotFoundError("not found") + return checkpoint_path + + torch_load = MagicMock(return_value=checkpoint) + loader_kwargs = { + "cast_dtype_fn": lambda _key: np.dtype(np.float16), + "converted_cache_dir": os.path.join(tmp_dir, "converted"), + } + with ( + patch.object(ltx2_utils.jax, "local_devices", return_value=["cpu"]), + patch.object(ltx2_utils, "hf_hub_download", side_effect=download), + patch.object(ltx2_utils.torch, "load", torch_load), + ): + first = ltx2_utils.load_transformer_weights( + "test/ltx2", + self._eval_shapes(), + "cpu", + **loader_kwargs, + ) + torch_load.assert_called_once_with(checkpoint_path, map_location="cpu") + torch_load.reset_mock() + torch_load.side_effect = AssertionError("cache hit must not reload the PyTorch checkpoint") + + second = ltx2_utils.load_transformer_weights( + "test/ltx2", + self._eval_shapes(), + "cpu", + **loader_kwargs, + ) + + torch_load.assert_not_called() + self.assertEqual(first["transformer_blocks"]["bias"].dtype, np.dtype(np.float16)) + np.testing.assert_array_equal(second["transformer_blocks"]["bias"], first["transformer_blocks"]["bias"]) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py index 4eba4d268..f6730f1c3 100644 --- a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py @@ -24,11 +24,13 @@ from maxdiffusion.pipelines.ltx2.ltx2_pipeline import ( LTX2Pipeline, _apply_ltx2_guidance, + _final_param_dtype, _restore_guidance_latents, _select_batch_sharding_axes, _select_guidance_latents, _temporary_vae_slicing, calculate_shift, + cast_with_exclusion, rescale_noise_cfg, ) @@ -141,6 +143,22 @@ def test_batch_sharding_axes_use_largest_divisor(self): self.assertEqual(_select_batch_sharding_axes(mesh_shape, 8), ("context",)) self.assertEqual(_select_batch_sharding_axes(mesh_shape, 3), ()) + def test_converted_weight_dtype_policy_matches_runtime_cast(self): + self.assertEqual(_final_param_dtype(("proj_in", "kernel"), jnp.bfloat16), np.dtype(jnp.bfloat16)) + self.assertEqual( + _final_param_dtype(("transformer_blocks", "norm1", "scale"), jnp.bfloat16), + np.dtype(jnp.float32), + ) + self.assertEqual( + _final_param_dtype(("condition_embedder", "linear", "kernel"), jnp.bfloat16), + np.dtype(jnp.float32), + ) + + def test_runtime_cast_does_not_copy_cached_final_dtype(self): + value = np.ones((2, 2), dtype=np.float16) + path = (jax.tree_util.DictKey("proj_in"), jax.tree_util.DictKey("kernel")) + self.assertIs(cast_with_exclusion(path, value, np.float16), value) + def test_tuned_block_sizes_update_effective_tokamax_fields(self): candidate = flash_block_sizes_for_candidate( { From aa1c697c30517c7f5681828c8e87c0168f476175 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Mon, 27 Jul 2026 09:45:17 +0000 Subject: [PATCH 09/10] perf: optimize video postprocessing natively on TPU via JAX --- .../pipelines/ltx2/ltx2_pipeline.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index a6b66bb55..608e5c0f7 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -1515,13 +1515,24 @@ def _decode_latents_to_video( video_vae_time = time.perf_counter() - t0_video_vae max_logging.log(f"Video VAE decode time: {video_vae_time:.2f}s") - # VAE outputs (B, T, H, W, C), but video processor expects (B, C, T, H, W) + # Perform rescaling, clipping, and layout transpose natively on TPU via JAX + # VAE outputs (B, T, H, W, C), but downstream processor expects (B, C, T, H, W) t0_video_post = time.perf_counter() + video = jnp.clip((video / 2.0 + 0.5), 0.0, 1.0) + video = jnp.transpose(video, (0, 4, 1, 2, 3)) video = jax.experimental.multihost_utils.process_allgather(video, tiled=True) - video_np = np.array(video).transpose(0, 4, 1, 2, 3) - video = self.video_processor.postprocess_video(torch.from_numpy(video_np), output_type=output_type) + + if output_type == "np": + video = np.array(video) + elif output_type == "pt": + video = torch.from_numpy(np.array(video)) + else: + # Fallback for PIL or custom output types + video_np = np.array(video) + video = self.video_processor.postprocess_video(torch.from_numpy((video_np * 2.0) - 1.0), output_type=output_type) + video_post_time = time.perf_counter() - t0_video_post - max_logging.log(f"Video Post-processing time (numpy+PIL): {video_post_time:.2f}s") + max_logging.log(f"Fast JAX Video Post-processing time: {video_post_time:.2f}s") return video, video_vae_time, video_post_time From 00e01c0cb27cf481f2d3a5abdc1cc7abb987b280 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Mon, 27 Jul 2026 09:50:42 +0000 Subject: [PATCH 10/10] Revert "perf: optimize video postprocessing natively on TPU via JAX" This reverts commit aa1c697c30517c7f5681828c8e87c0168f476175. --- .../pipelines/ltx2/ltx2_pipeline.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index 608e5c0f7..a6b66bb55 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -1515,24 +1515,13 @@ def _decode_latents_to_video( video_vae_time = time.perf_counter() - t0_video_vae max_logging.log(f"Video VAE decode time: {video_vae_time:.2f}s") - # Perform rescaling, clipping, and layout transpose natively on TPU via JAX - # VAE outputs (B, T, H, W, C), but downstream processor expects (B, C, T, H, W) + # VAE outputs (B, T, H, W, C), but video processor expects (B, C, T, H, W) t0_video_post = time.perf_counter() - video = jnp.clip((video / 2.0 + 0.5), 0.0, 1.0) - video = jnp.transpose(video, (0, 4, 1, 2, 3)) video = jax.experimental.multihost_utils.process_allgather(video, tiled=True) - - if output_type == "np": - video = np.array(video) - elif output_type == "pt": - video = torch.from_numpy(np.array(video)) - else: - # Fallback for PIL or custom output types - video_np = np.array(video) - video = self.video_processor.postprocess_video(torch.from_numpy((video_np * 2.0) - 1.0), output_type=output_type) - + video_np = np.array(video).transpose(0, 4, 1, 2, 3) + video = self.video_processor.postprocess_video(torch.from_numpy(video_np), output_type=output_type) video_post_time = time.perf_counter() - t0_video_post - max_logging.log(f"Fast JAX Video Post-processing time: {video_post_time:.2f}s") + max_logging.log(f"Video Post-processing time (numpy+PIL): {video_post_time:.2f}s") return video, video_vae_time, video_post_time