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

src/maxdiffusion/aot_cache.py

@@ -58,6 +58,7 @@ def transformer_forward_pass(...):
 import hashlib
 import inspect
 import os
+import json
 import pickle
 import re
 import threading
@@ -72,6 +73,12 @@ def transformer_forward_pass(...):
 _FORMAT_VERSION = 1
 
 
+def _metadata_fingerprint(meta: dict[str, Any]) -> str:
+  """Returns the stable filename fingerprint for install-time metadata."""
+  serialized = json.dumps(meta, sort_keys=True)
+  return hashlib.sha256(serialized.encode("utf-8")).hexdigest()[:12]
+
+
 def _dynamic_signature(args: tuple, kwargs: dict) -> str:
   """Deterministic digest of everything that selects an executable.
 
@@ -372,21 +379,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)

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -1,7 +1,18 @@
 #hardware
 hardware: 'tpu'
 skip_jax_distributed_system: False
+# 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.
+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 
@@ -12,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'
@@ -94,7 +107,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
@@ -106,6 +119,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 +190,17 @@ 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: "np_uint8"
+
+aot_cache_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 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

src/maxdiffusion/configs/ltx2_video.yml

@@ -1,7 +1,18 @@
 #hardware
 hardware: 'tpu'
 skip_jax_distributed_system: False
+# 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.
+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 
@@ -18,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'
@@ -99,7 +112,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
@@ -111,6 +124,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
 
@@ -170,4 +187,17 @@ 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: "np_uint8"
+
+aot_cache_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 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