Port PR 447 (Tile-size auto-tuning) to LTX2 and LTX2.3

Author: Perseus14

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -2,7 +2,8 @@
 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 
-
+use_base2_exp: True
+use_experimental_scheduler: True
 # 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.
@@ -111,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
@@ -183,3 +188,11 @@ 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

@@ -2,7 +2,8 @@
 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 
-
+use_base2_exp: True
+use_experimental_scheduler: True
 # 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.
@@ -116,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
 
@@ -180,3 +185,11 @@ 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

@@ -113,7 +113,55 @@ 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, 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}")