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

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

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

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")

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