Perf(LTX2): Comprehensive XLA, Memory, and Transformer Code Quality Optimizations

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,16 @@ 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"
+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

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,15 @@ 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

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,57 @@ 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
+
+  vmem = config.flash_block_sizes.get("vmem_limit_bytes", None) if config.flash_block_sizes else None
+  if vmem is None:
+    import os
+    import re
+
+    m = re.search(r"--xla_tpu_scoped_vmem_limit_kib=(\d+)", os.environ.get("LIBTPU_INIT_ARGS", ""))
+    vmem = int(m.group(1)) * 1024 if m else 32 * 1024 * 1024
+
+  mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes)
+  bench = LTX2BlockBenchmark.from_config(config, mesh, vmem_limit_bytes=vmem)
+  max_logging.log(f"[tile-search] tuning block sizes for {bench.label} (vmem={vmem/1024/1024:.1f}MB) 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) if config.flash_block_sizes else {}
+  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,
+      "vmem_limit_bytes": vmem,
+  })
+  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 +239,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 +310,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")

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

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]
@@ -352,6 +341,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
@@ -445,17 +436,37 @@ 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,
         scale=dim_head**-0.5,
         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 __call__(
@@ -485,7 +496,7 @@ def __call__(
     # 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