diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py
index ba722160a..81491b71e 100644
--- a/src/maxdiffusion/aot_cache.py
+++ b/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)
diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml
index a6275cfae..b47ee927f 100644
--- a/src/maxdiffusion/configs/ltx2_3_video.yml
+++ b/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"
\ No newline at end of file
+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
diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml
index 8ddfaba1b..271b07deb 100644
--- a/src/maxdiffusion/configs/ltx2_video.yml
+++ b/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
diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py
index baeededd2..85fb3af41 100644
--- a/src/maxdiffusion/generate_ltx2.py
+++ b/src/maxdiffusion/generate_ltx2.py
@@ -13,18 +13,92 @@
 # limitations under the License.
 
 from typing import Sequence
+import json
 import jax
 import jax.numpy as jnp
 import time
 import os
+import subprocess
+import uuid
 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
 import flax
+import urllib.parse
+import posixpath
+
 from maxdiffusion.utils.export_utils import export_to_video_with_audio
 from maxdiffusion.loaders.ltx2_lora_nnx_loader import LTX2NNXLoraLoader
 
 
+def upload_video_to_gcs(output_dir: str, video_path: str):
+  """
+  Uploads a local video file to a specified Google Cloud Storage bucket.
+  """
+  from google.cloud import storage
+  from google.api_core.exceptions import GoogleAPIError
+
+  try:
+    parsed = urllib.parse.urlparse(output_dir)
+    bucket_name = parsed.netloc
+    folder_name = parsed.path.lstrip("/")
+
+    storage_client = storage.Client()
+    bucket = storage_client.bucket(bucket_name)
+
+    source_file_path = f"./{video_path}"
+    destination_blob_name = posixpath.join(folder_name, "videos", video_path)
+
+    blob = bucket.blob(destination_blob_name)
+
+    max_logging.log(f"Uploading {source_file_path} to {bucket_name}/{destination_blob_name}...")
+    blob.upload_from_filename(source_file_path)
+    max_logging.log(f"Upload complete {source_file_path}.")
+
+  except GoogleAPIError as e:
+    max_logging.log(f"A storage error occurred during upload: {e}")
+
+
+def delete_file(file_path: str):
+  if os.path.exists(file_path):
+    try:
+      os.remove(file_path)
+      max_logging.log(f"Successfully deleted file: {file_path}")
+    except OSError as e:
+      max_logging.log(f"Error deleting file '{file_path}': {e}")
+  else:
+    max_logging.log(f"The file '{file_path}' does not exist.")
+
+
+def get_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", "-C", source_root, "rev-parse", "HEAD"], stderr=subprocess.DEVNULL)
+        .strip()
+        .decode("utf-8")
+    )
+    tracked_changes = subprocess.check_output(
+        ["git", "-C", source_root, "status", "--porcelain", "--untracked-files=no"],
+        stderr=subprocess.DEVNULL,
+    ).strip()
+    untracked_source = subprocess.check_output(
+        ["git", "-C", source_root, "ls-files", "--others", "--exclude-standard", "--", "src/maxdiffusion"],
+        stderr=subprocess.DEVNULL,
+    ).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: 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.")
+    return None
+
+
 jax.config.update("jax_use_shardy_partitioner", True)
 
 
@@ -55,12 +129,164 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
       use_cross_timestep=getattr(config, "use_cross_timestep", None),
       noise_scale=getattr(config, "noise_scale", 1.0),
       dtype=jnp.bfloat16 if getattr(config, "activations_dtype", "bfloat16") == "bfloat16" else jnp.float32,
-      output_type=getattr(config, "upsampler_output_type", "pil"),
+      output_type=getattr(config, "upsampler_output_type", "np_uint8"),
   )
   return out
 
 
+def maybe_tune_block_sizes(config):
+  """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)
+  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,
+      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:
+    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 _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. 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",
+      "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 {
+      "source_revision": source_revision,
+      "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)),
+      "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", ())),
+      "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 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)
+
   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}")
@@ -136,14 +362,50 @@ 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
+  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(
+      aot_cache_dir,
+      meta=aot_metadata,
+      mesh=pipeline.mesh,
+  )
+  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
 
-  max_logging.log(f"🚀 Starting warmup compilation pass ({original_num_steps} steps)...")
-  _ = call_pipeline(config, pipeline, prompt, negative_prompt)
+  # 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)...")
+  with aot_cache.warmup_mode():
+    _ = 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}")
@@ -184,7 +446,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/max_utils.py b/src/maxdiffusion/max_utils.py
index 28b747894..37027c27d 100644
--- a/src/maxdiffusion/max_utils.py
+++ b/src/maxdiffusion/max_utils.py
@@ -697,6 +697,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 c56d8c64c..51d6f7b6f 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):
@@ -1914,7 +2065,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 +2160,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
@@ -2133,16 +2284,32 @@ 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 214690c98..71ed4325d 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
@@ -88,13 +89,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 +172,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 +247,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,12 +345,18 @@ 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,
+      use_base2_exp: bool = False,
+      use_experimental_scheduler: bool = False,
+      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")
@@ -392,10 +391,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 +456,50 @@ def __init__(
           dtype=dtype,
       )
 
+    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 cross_attention_remapped_to_flash:
+        attention_kernel = "tokamax_flash"
+
+      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,
+          None if cross_attention_uses_local_kv else 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 +507,25 @@ 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,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
 
+  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 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,
@@ -466,26 +534,31 @@ 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
 
     # 1. Project and Norm
-    with jax.named_scope("QKV Projection"):
+    with self.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"):
+    with self.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"):
+    with self.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
 
@@ -504,7 +577,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/ltx2_utils.py b/src/maxdiffusion/models/ltx2/ltx2_utils.py
index 7cf7fa761..45d71e27c 100644
--- a/src/maxdiffusion/models/ltx2/ltx2_utils.py
+++ b/src/maxdiffusion/models/ltx2/ltx2_utils.py
@@ -15,8 +15,14 @@
 """
 
 import json
-from typing import Optional
+import os
+import shutil
+import time
+import threading
+import concurrent.futures
+from typing import Optional, Callable
 import torch
+import numpy as np
 import jax
 import jax.numpy as jnp
 from maxdiffusion import max_logging
@@ -26,6 +32,74 @@
 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)
 
+
+def _tuple_str_to_int(in_tuple):
+  out_list = []
+  for item in in_tuple:
+    try:
+      out_list.append(int(item))
+    except ValueError:
+      out_list.append(item)
+  return tuple(out_list)
+
+
+def _converted_key_to_filename(flax_key: tuple) -> str:
+  return ".".join(str(k) for k in flax_key) + ".npy"
+
+
+def try_load_converted_weights(cache_dir: str, eval_shapes: dict, cast_dtype_fn: Optional[Callable]) -> Optional[dict]:
+  manifest_path = os.path.join(cache_dir, "manifest.json")
+  if not os.path.isfile(manifest_path):
+    return None
+  try:
+    with open(manifest_path, "r") as f:
+      manifest = json.load(f)
+    expected_keys = set(flatten_dict(eval_shapes).keys())
+
+    def load_one(key_str, meta):
+      flax_key = _tuple_str_to_int(tuple(key_str.split(".")))
+      logical_dtype = np.dtype(meta["dtype"])
+      if cast_dtype_fn is not None and logical_dtype != np.dtype(cast_dtype_fn(flax_key)):
+        raise ValueError(f"dtype policy changed for {key_str}")
+      value = np.load(os.path.join(cache_dir, meta["file"]))
+      if meta.get("bitview"):
+        value = value.view(logical_dtype)
+      if tuple(value.shape) != tuple(meta["shape"]):
+        raise ValueError(f"shape changed for {key_str}")
+      return flax_key, value
+
+    flax_state_dict = {}
+    with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
+      for flax_key, value in executor.map(lambda kv: load_one(*kv), manifest.items()):
+        flax_state_dict[flax_key] = value
+    if set(flax_state_dict.keys()) != expected_keys:
+      return None
+    return unflatten_dict(flax_state_dict)
+  except (OSError, ValueError, KeyError, TypeError) as e:
+    max_logging.log(f"Converted-weights cache unusable ({e}); reconverting")
+    return None
+
+
+def save_converted_weights(cache_dir: str, flat_state_dict: dict) -> None:
+  tmp_dir = f"{cache_dir}.tmp.{os.getpid()}"
+  os.makedirs(tmp_dir, exist_ok=True)
+  manifest = {}
+  uint_by_width = {1: np.uint8, 2: np.uint16, 4: np.uint32}
+  for flax_key, value in flat_state_dict.items():
+    filename = _converted_key_to_filename(flax_key)
+    bitview = value.dtype.kind not in "fiub"
+    stored = value.view(uint_by_width[value.dtype.itemsize]) if bitview else value
+    np.save(os.path.join(tmp_dir, filename), stored)
+    key_str = ".".join(str(k) for k in flax_key)
+    manifest[key_str] = {"file": filename, "shape": list(value.shape), "dtype": str(value.dtype), "bitview": bitview}
+  with open(os.path.join(tmp_dir, "manifest.json"), "w") as f:
+    json.dump(manifest, f)
+  try:
+    os.rename(tmp_dir, cache_dir)
+  except OSError:
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+
+
 KNOWN_UPSAMPLER_CONFIGS = {
     "ltx-2.3-spatial-upscaler-x2-1.0.safetensors": {
         "spatial_upsample": True,
@@ -184,47 +258,255 @@ 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 _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_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():
+    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",
+    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 concurrent.futures
+
   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)
+  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 = {}
-    cpu = jax.local_devices(backend="cpu")[0]
-    flattened_dict = flatten_dict(eval_shapes)
+  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")
 
-    random_flax_state_dict = {}
-    for key in flattened_dict:
-      random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key]
+  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()))
 
-    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):
+      return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=model_file)
 
-      pt_tuple_key = tuple(renamed_pt_key.split("."))
+  except EntryNotFoundError:
+    shards = ["diffusion_pytorch_model.safetensors"]
+
+    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")
 
-      flax_key, flax_tensor = get_key_and_value(
-          pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers, num_layers
+  checkpoint_files = [resolve_shard_path(model_file) for model_file in shards]
+
+  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()
+  }
+
+  if converted_cache_dir:
+    t_cache_load = time.perf_counter()
+    cached = try_load_converted_weights(converted_cache_dir, unflatten_dict(flattened_dict), cast_dtype_fn)
+    if cached is not None:
+      max_logging.log(
+          f"Loaded converted {subfolder or 'transformer'} weights from {converted_cache_dir} "
+          f"in {time.perf_counter() - t_cache_load:.1f}s"
       )
+      return cached
 
-      flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu)
+  t_start = time.perf_counter()
 
-    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_state_dict = {}
+  populated_layer_indices = {}
+  dict_lock = threading.Lock()
+
+  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(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)
+    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:]
+
+    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)
+
+    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:
+          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:
+      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 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())
+      for i in range(0, len(shard_keys), chunk_size):
+        safetensors_tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size]))
+    else:
+      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)
+  if converted_cache_dir and not os.path.isdir(converted_cache_dir):
+    t_cache_save = time.perf_counter()
+    if jax.process_index() == 0:
+      save_converted_weights(converted_cache_dir, flax_state_dict)
+      max_logging.log(
+          f"Saved converted {subfolder or 'transformer'} weights to {converted_cache_dir} "
+          f"in {time.perf_counter() - t_cache_save:.1f}s"
+      )
+  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..0718f1c85 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
@@ -27,6 +28,61 @@
 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
+
+
+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):
@@ -89,7 +145,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,12 +189,18 @@ 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,
+      use_base2_exp: bool = False,
+      use_experimental_scheduler: bool = False,
+      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")
@@ -151,6 +218,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,
@@ -166,6 +234,10 @@ 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,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
 
     self.audio_norm1 = nnx.RMSNorm(
@@ -179,6 +251,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,
@@ -194,6 +267,10 @@ 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,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
 
     # 2. Prompt Cross-Attention
@@ -208,6 +285,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,
@@ -221,8 +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(
@@ -236,6 +319,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,
@@ -252,6 +336,10 @@ 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,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
 
     # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
@@ -266,6 +354,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,
@@ -279,9 +368,13 @@ 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,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
 
     self.video_to_audio_norm = nnx.RMSNorm(
@@ -295,6 +388,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,
@@ -311,6 +405,10 @@ 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,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
     )
 
     # 4. Feed Forward
@@ -361,69 +459,102 @@ 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 named_scope(self, name: str):
+    return jax.named_scope(name) if self.enable_jax_named_scopes else contextlib.nullcontext()
+
   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",
+      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.
+
+    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
     batch_size = hidden_states.shape[0]
 
     axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
@@ -455,14 +586,14 @@ 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, :]
 
     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,
@@ -486,14 +617,14 @@ 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, :]
 
     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,
@@ -504,10 +635,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
@@ -515,21 +646,23 @@ 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):
+    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
@@ -537,13 +670,15 @@ 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):
+    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
 
@@ -590,7 +725,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,
@@ -606,7 +741,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,
@@ -695,6 +830,11 @@ 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,
+      use_base2_exp: bool = False,
+      use_experimental_scheduler: bool = False,
+      enable_jax_named_scopes: bool = False,
       **kwargs,
   ):
     self.spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks
@@ -747,9 +887,12 @@ 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
+    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")
@@ -760,6 +903,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 +1025,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 +1104,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 +1143,11 @@ 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,
+          use_base2_exp=self.use_base2_exp,
+          use_experimental_scheduler=self.use_experimental_scheduler,
+          enable_jax_named_scopes=self.enable_jax_named_scopes,
       )
 
     if self.scan_layers:
@@ -999,6 +1157,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 +1187,11 @@ 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,
+            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)
       self.transformer_blocks = nnx.List(blocks)
@@ -1035,7 +1199,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 +1218,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 +1236,140 @@ 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,
+      audio_encoder_hidden_states: jax.Array,
+      num_frames: int,
+      height: int,
+      width: int,
+      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)
+      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)
+
+    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 precompute_time_embeds(
+      self,
+      timesteps: jax.Array,
+      sigmas: jax.Array,
+      hidden_dtype: jnp.dtype,
+      use_cross_timestep: bool = False,
+  ) -> Dict[str, jax.Array]:
+    """
+    Precomputes the timestep embeddings and cross-attention scale/shifts
+    over an array of timesteps AOT outside the diffusion loop.
+
+    Args:
+      timesteps: 1D array of video timesteps.
+      sigmas: 1D array of sigmas.
+      hidden_dtype: Data type for the embeddings.
+      use_cross_timestep: Whether to use cross timestep modulation.
+
+    Returns:
+      A dictionary of precomputed embedding arrays.
+    """
+    temb, embedded_timestep = self.time_embed(timesteps, hidden_dtype=hidden_dtype)
+    temb_audio, audio_embedded_timestep = self.audio_time_embed(timesteps, hidden_dtype=hidden_dtype)
+
+    res = {
+        "temb": jnp.expand_dims(temb, 1),
+        "embedded_timestep": jnp.expand_dims(embedded_timestep, 1),
+        "temb_audio": jnp.expand_dims(temb_audio, 1),
+        "audio_embedded_timestep": jnp.expand_dims(audio_embedded_timestep, 1),
+    }
+
+    if self.cross_attn_mod:
+      temb_prompt, _ = self.prompt_adaln(sigmas, hidden_dtype=hidden_dtype)
+      temb_prompt_audio, _ = self.audio_prompt_adaln(sigmas, hidden_dtype=hidden_dtype)
+      res["temb_prompt"] = jnp.expand_dims(temb_prompt, 1)
+      res["temb_prompt_audio"] = jnp.expand_dims(temb_prompt_audio, 1)
+
+    if use_cross_timestep:
+      video_ca_timestep = sigmas
+      audio_ca_timestep = sigmas
+    else:
+      video_ca_timestep = timesteps
+      audio_ca_timestep = timesteps
+
+    timestep_cross_attn_gate_scale_factor = self.cross_attn_timestep_scale_multiplier / self.timestep_scale_multiplier
+
+    video_cross_attn_scale_shift, _ = self.av_cross_attn_video_scale_shift(video_ca_timestep, hidden_dtype=hidden_dtype)
+    video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate(
+        video_ca_timestep * timestep_cross_attn_gate_scale_factor, hidden_dtype=hidden_dtype
+    )
+
+    audio_cross_attn_scale_shift, _ = self.av_cross_attn_audio_scale_shift(audio_ca_timestep, hidden_dtype=hidden_dtype)
+    audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate(
+        audio_ca_timestep * timestep_cross_attn_gate_scale_factor, hidden_dtype=hidden_dtype
+    )
+
+    res["video_cross_attn_scale_shift"] = jnp.expand_dims(video_cross_attn_scale_shift, 1)
+    res["video_cross_attn_a2v_gate"] = jnp.expand_dims(video_cross_attn_a2v_gate, 1)
+    res["audio_cross_attn_scale_shift"] = jnp.expand_dims(audio_cross_attn_scale_shift, 1)
+    res["audio_cross_attn_v2a_gate"] = jnp.expand_dims(audio_cross_attn_v2a_gate, 1)
+
+    return res
+
   def __call__(
       self,
       hidden_states: jax.Array,
@@ -1084,149 +1394,261 @@ 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,
+      time_embed_cache: Optional[Dict[str, 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
+    batch_size = hidden_states.shape[0]
 
-    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)
+    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)
 
-    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 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)
+    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"]
+        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"):
+    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(
-          timestep.flatten(),
-          hidden_dtype=hidden_states.dtype,
-      )
-      temb = temb.reshape(batch_size, -1, temb.shape[-1])
-      embedded_timestep = embedded_timestep.reshape(batch_size, -1, embedded_timestep.shape[-1])
-
-      temb_audio, audio_embedded_timestep = self.audio_time_embed(
-          audio_timestep.flatten(),
-          hidden_dtype=audio_hidden_states.dtype,
-      )
-      temb_audio = temb_audio.reshape(batch_size, -1, temb_audio.shape[-1])
-      audio_embedded_timestep = audio_embedded_timestep.reshape(batch_size, -1, audio_embedded_timestep.shape[-1])
+      if time_embed_cache is not None:
+        temb = time_embed_cache["temb"]
+        embedded_timestep = time_embed_cache["embedded_timestep"]
+        temb_audio = time_embed_cache["temb_audio"]
+        audio_embedded_timestep = time_embed_cache["audio_embedded_timestep"]
+
+        # Broadcast batch size
+        temb = jnp.repeat(temb, batch_size, axis=0)
+        embedded_timestep = jnp.repeat(embedded_timestep, batch_size, axis=0)
+        temb_audio = jnp.repeat(temb_audio, batch_size, axis=0)
+        audio_embedded_timestep = jnp.repeat(audio_embedded_timestep, batch_size, axis=0)
+
+        # Reshape to expected shapes
+        temb = temb.reshape(batch_size, -1, temb.shape[-1])
+        embedded_timestep = embedded_timestep.reshape(batch_size, -1, embedded_timestep.shape[-1])
+        temb_audio = temb_audio.reshape(batch_size, -1, temb_audio.shape[-1])
+        audio_embedded_timestep = audio_embedded_timestep.reshape(batch_size, -1, audio_embedded_timestep.shape[-1])
+
+        if self.cross_attn_mod and sigma is not None:
+          temb_prompt = time_embed_cache["temb_prompt"]
+          temb_prompt_audio = time_embed_cache["temb_prompt_audio"]
+
+          temb_prompt = jnp.repeat(temb_prompt, batch_size, axis=0)
+          temb_prompt_audio = jnp.repeat(temb_prompt_audio, batch_size, axis=0)
+
+          temb_prompt = temb_prompt.reshape(batch_size, -1, temb_prompt.shape[-1])
+          temb_prompt_audio = temb_prompt_audio.reshape(batch_size, -1, temb_prompt_audio.shape[-1])
+        else:
+          temb_prompt = None
+          temb_prompt_audio = None
 
-      if self.cross_attn_mod and sigma is not None:
-        audio_sigma = audio_sigma if audio_sigma is not None else sigma
-        temb_prompt, _ = self.prompt_adaln(
-            sigma.flatten(),
+      else:
+        temb, embedded_timestep = self.time_embed(
+            timestep.flatten(),
             hidden_dtype=hidden_states.dtype,
         )
-        temb_prompt_audio, _ = self.audio_prompt_adaln(
-            audio_sigma.flatten(),
+        temb = temb.reshape(batch_size, -1, temb.shape[-1])
+        embedded_timestep = embedded_timestep.reshape(batch_size, -1, embedded_timestep.shape[-1])
+
+        temb_audio, audio_embedded_timestep = self.audio_time_embed(
+            audio_timestep.flatten(),
             hidden_dtype=audio_hidden_states.dtype,
         )
-        temb_prompt = temb_prompt.reshape(batch_size, -1, temb_prompt.shape[-1])
-        temb_prompt_audio = temb_prompt_audio.reshape(batch_size, -1, temb_prompt_audio.shape[-1])
-      else:
-        temb_prompt = None
-        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"
-        video_ca_timestep = audio_sigma.flatten()
-        audio_ca_timestep = sigma.flatten()
+        temb_audio = temb_audio.reshape(batch_size, -1, temb_audio.shape[-1])
+        audio_embedded_timestep = audio_embedded_timestep.reshape(batch_size, -1, audio_embedded_timestep.shape[-1])
+
+        if self.cross_attn_mod and sigma is not None:
+          audio_sigma = audio_sigma if audio_sigma is not None else sigma
+          temb_prompt, _ = self.prompt_adaln(
+              sigma.flatten(),
+              hidden_dtype=hidden_states.dtype,
+          )
+          temb_prompt_audio, _ = self.audio_prompt_adaln(
+              audio_sigma.flatten(),
+              hidden_dtype=audio_hidden_states.dtype,
+          )
+          temb_prompt = temb_prompt.reshape(batch_size, -1, temb_prompt.shape[-1])
+          temb_prompt_audio = temb_prompt_audio.reshape(batch_size, -1, temb_prompt_audio.shape[-1])
+        else:
+          temb_prompt = None
+          temb_prompt_audio = None
+
+      if time_embed_cache is not None:
+        video_cross_attn_scale_shift = time_embed_cache["video_cross_attn_scale_shift"]
+        video_cross_attn_a2v_gate = time_embed_cache["video_cross_attn_a2v_gate"]
+        audio_cross_attn_scale_shift = time_embed_cache["audio_cross_attn_scale_shift"]
+        audio_cross_attn_v2a_gate = time_embed_cache["audio_cross_attn_v2a_gate"]
+
+        video_cross_attn_scale_shift = jnp.repeat(video_cross_attn_scale_shift, batch_size, axis=0)
+        video_cross_attn_a2v_gate = jnp.repeat(video_cross_attn_a2v_gate, batch_size, axis=0)
+        audio_cross_attn_scale_shift = jnp.repeat(audio_cross_attn_scale_shift, batch_size, axis=0)
+        audio_cross_attn_v2a_gate = jnp.repeat(audio_cross_attn_v2a_gate, batch_size, axis=0)
+
+        video_cross_attn_scale_shift = video_cross_attn_scale_shift.reshape(
+            batch_size, -1, video_cross_attn_scale_shift.shape[-1]
+        )
+        video_cross_attn_a2v_gate = video_cross_attn_a2v_gate.reshape(batch_size, -1, video_cross_attn_a2v_gate.shape[-1])
+        audio_cross_attn_scale_shift = audio_cross_attn_scale_shift.reshape(
+            batch_size, -1, audio_cross_attn_scale_shift.shape[-1]
+        )
+        audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.reshape(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
       else:
-        video_ca_timestep = timestep.flatten()
-        audio_ca_timestep = audio_timestep.flatten() if audio_timestep is not None else timestep.flatten()
+        if use_cross_timestep:
+          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:
+          video_ca_timestep = timestep.flatten()
+          audio_ca_timestep = audio_timestep.flatten() if audio_timestep is not None else timestep.flatten()
 
-      video_cross_attn_scale_shift, _ = self.av_cross_attn_video_scale_shift(
-          video_ca_timestep,
-          hidden_dtype=hidden_states.dtype,
-      )
-      video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate(
-          video_ca_timestep * timestep_cross_attn_gate_scale_factor,
-          hidden_dtype=hidden_states.dtype,
-      )
-      video_cross_attn_scale_shift = video_cross_attn_scale_shift.reshape(
-          batch_size, -1, video_cross_attn_scale_shift.shape[-1]
-      )
-      video_cross_attn_a2v_gate = video_cross_attn_a2v_gate.reshape(batch_size, -1, video_cross_attn_a2v_gate.shape[-1])
+        video_cross_attn_scale_shift, _ = self.av_cross_attn_video_scale_shift(
+            video_ca_timestep,
+            hidden_dtype=hidden_states.dtype,
+        )
+        video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate(
+            video_ca_timestep * timestep_cross_attn_gate_scale_factor,
+            hidden_dtype=hidden_states.dtype,
+        )
+        video_cross_attn_scale_shift = video_cross_attn_scale_shift.reshape(
+            batch_size, -1, video_cross_attn_scale_shift.shape[-1]
+        )
+        video_cross_attn_a2v_gate = video_cross_attn_a2v_gate.reshape(batch_size, -1, video_cross_attn_a2v_gate.shape[-1])
 
-      audio_cross_attn_scale_shift, _ = self.av_cross_attn_audio_scale_shift(
-          audio_ca_timestep,
-          hidden_dtype=audio_hidden_states.dtype,
-      )
-      audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate(
-          audio_ca_timestep * timestep_cross_attn_gate_scale_factor,
-          hidden_dtype=audio_hidden_states.dtype,
-      )
-      audio_cross_attn_scale_shift = audio_cross_attn_scale_shift.reshape(
-          batch_size, -1, audio_cross_attn_scale_shift.shape[-1]
-      )
-      audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.reshape(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
+        audio_cross_attn_scale_shift, _ = self.av_cross_attn_audio_scale_shift(
+            audio_ca_timestep,
+            hidden_dtype=audio_hidden_states.dtype,
+        )
+        audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate(
+            audio_ca_timestep * timestep_cross_attn_gate_scale_factor,
+            hidden_dtype=audio_hidden_states.dtype,
+        )
+        audio_cross_attn_scale_shift = audio_cross_attn_scale_shift.reshape(
+            batch_size, -1, audio_cross_attn_scale_shift.shape[-1]
+        )
+        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)
 
         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,
+          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,
+      )
+
+      def apply_block(block, context: LTX2BlockContext, mask, block_cached_kv=None):
+        orig_perturbation_mask = context.perturbation_mask
+        context = context.replace(perturbation_mask=mask)
+        with self.named_scope("Transformer Layer"):
+          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,
+        )
+        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 +1657,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,42 +1712,42 @@ 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"):
+    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/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 19069b6f9..4b215b43d 100644
--- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py
+++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py
@@ -16,10 +16,27 @@
 
 from typing import Optional, Any, List, Union
 from functools import partial
+import itertools
+import os
+import numpy as np
+
+
+def converted_weights_cache_dir(config, subfolder: str) -> str:
+  """Per-(model, subfolder, dtype, scan) dir for memoized converted weights."""
+  base = getattr(config, "converted_weights_dir", "")
+  if not base:
+    return ""
+  model_tag = (
+      getattr(config, "ltx2_transformer_pretrained_model_name_or_path", "")
+      or getattr(config, "pretrained_model_name_or_path", "")
+  ).replace("/", "--")
+  scan_layers = getattr(config, "scan_layers", True)
+  return os.path.join(base, f"{model_tag}--{subfolder or 'transformer'}--{config.weights_dtype}--scan{scan_layers}")
+
 
 import time
-import numpy as np
 import torch
+from PIL import Image
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
@@ -56,9 +73,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,
@@ -66,6 +89,38 @@
 }
 
 
+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
@@ -81,6 +136,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
@@ -88,25 +145,144 @@ 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__)
 
+_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:
@@ -148,6 +324,11 @@ 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["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
   ltx2_config["names_which_can_be_offloaded"] = config.names_which_can_be_offloaded
@@ -183,16 +364,68 @@ 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=converted_weights_cache_dir(config, subfolder),
     )
 
   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:
+        # 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")
+
   state = nnx.from_flat_state(state)
 
   transformer = nnx.merge(graphdef, state, rest_of_state)
@@ -242,6 +475,14 @@ def retrieve_timesteps(
   return scheduler_state
 
 
+def _tpu_format_pil(video):
+  return jnp.clip((video / 2.0 + 0.5) * 255.0, 0, 255).astype(jnp.uint8)
+
+
+def _tpu_format_pt(video):
+  return jnp.transpose(jnp.clip((video / 2.0 + 0.5), 0.0, 1.0), (0, 4, 1, 2, 3))
+
+
 class LTX2Pipeline:
   """
   Pipeline for LTX-2.
@@ -364,7 +605,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 +658,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 +714,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 +786,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 +849,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 +886,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 +930,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 +991,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 +1050,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 +1161,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 +1180,43 @@ 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:
+          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:
+          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 +1238,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 +1347,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 +1358,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 +1442,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 +1466,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 +1478,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 +1505,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
@@ -1202,13 +1539,35 @@ 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)
+    # Fast post-processing: Format on TPU to save bandwidth, then gather and wrap
     t0_video_post = time.perf_counter()
+
+    # 1. Native TPU Formatting
+    if output_type in ["pil", "np_uint8"]:
+      # PIL and raw numpy both need (B, T, H, W, C) in uint8.
+      video = _tpu_format_pil(video)
+    elif output_type in ["np", "pt"]:
+      # PyTorch/NumPy expect (B, C, T, H, W) in float [0.0, 1.0].
+      video = _tpu_format_pt(video)
+
+    # 2. Host Transfer
     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)
+    video_np = np.array(video)
+
+    # 3. Host Assembly
+    if output_type == "pil":
+      video = [[Image.fromarray(frame) for frame in batch] for batch in video_np]
+    elif output_type in ["np", "np_uint8"]:
+      video = video_np
+    elif output_type == "pt":
+      video = torch.from_numpy(video_np)
+    else:
+      # Fallback for custom output_type
+      video_np = video_np.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"Video Post-processing time (numpy+PIL): {video_post_time:.2f}s")
+    max_logging.log(f"Fast Video Post-processing time: {video_post_time:.2f}s")
 
     return video, video_vae_time, video_post_time
 
@@ -1235,6 +1594,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 +1604,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 +1651,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 +1702,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 +1722,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 +1745,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 +1767,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 +1824,25 @@ 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,
     )
+    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()
-    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 +1954,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 +2019,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 +2032,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,9 +2074,31 @@ 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
+        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]
@@ -1668,8 +2107,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)
 
@@ -1693,89 +2137,51 @@ 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
-          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(
-              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:
-            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
@@ -1804,7 +2210,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 +2257,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)
@@ -1880,27 +2294,26 @@ def convert_to_vel(lat, x0, sig):
     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:
@@ -1914,15 +2327,16 @@ def convert_to_vel(lat, x0, sig):
     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()
@@ -1956,7 +2370,41 @@ def convert_to_vel(lat, x0, sig):
 
 
 @partial(
-    jax.jit,
+    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=(
         "latent_num_frames",
         "latent_height",
@@ -1988,6 +2436,9 @@ def transformer_forward_pass(
     audio_sigma=None,
     use_cross_timestep=False,
     is_cfg_stg_mode: bool = False,
+    kv_cache=None,
+    rope_cache=None,
+    time_embed_cache=None,
 ):
   """Forward pass for the transformer."""
   # pylint: disable=too-many-positional-arguments,unused-argument
@@ -2039,22 +2490,17 @@ 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,
+      time_embed_cache=time_embed_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 +2511,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,153 +2546,146 @@ 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
-
-  # 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
+  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,
+    )
 
-  def scan_body(carry, inputs):
-    t, sigma_t = inputs
-    latents, audio_latents, s_state = carry
+  if use_kv_cache:
+    # Precompute timestep embeddings AOT outside the diffusion loop
+    time_embed_cache_full = transformer.precompute_time_embeds(
+        timesteps_jax,
+        sigmas,
+        hidden_dtype=latents_jax.dtype,
+        use_cross_timestep=use_cross_timestep,
+    )
+  else:
+    time_embed_cache_full = None
 
+  if not scan_layers:
     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)
-
-        x0_combined = x0_text + cfg_delta + stg_delta + video_modality_delta
+      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)
 
-        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 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
-
-        if audio_guidance_rescale > 0:
-          x0_audio_combined = rescale_noise_cfg(x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale)
+  def scan_body(carry, inputs):
+    if use_kv_cache:
+      t, sigma_t, time_embed_cache_step = inputs
+    else:
+      t, sigma_t = inputs
+      time_embed_cache_step = None
 
-        noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t)
+    latents, audio_latents, s_state = carry
 
-      # ... (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)
+    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,
+        time_embed_cache=time_embed_cache_step,
+    )
 
-        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)
-      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
+    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
+    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)
+  if use_kv_cache:
+    scan_inputs = (timesteps_jax, sigmas, time_embed_cache_full)
+  else:
+    scan_inputs = (timesteps_jax, sigmas)
+
+  final_carry, _ = jax.lax.scan(scan_body, initial_carry, scan_inputs)
 
   return final_carry[0], final_carry[1]
diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py
index e84571165..d1121ca3f 100644
--- a/src/maxdiffusion/pyconfig.py
+++ b/src/maxdiffusion/pyconfig.py
@@ -225,12 +225,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..07a85512e 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,86 @@ 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 +870,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..3249fceaa
--- /dev/null
+++ b/src/maxdiffusion/tests/ltx2/test_generate_ltx2_aot.py
@@ -0,0 +1,175 @@
+# 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_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..f5fe81667
--- /dev/null
+++ b/src/maxdiffusion/tests/ltx2/test_ltx2_utils_loader.py
@@ -0,0 +1,214 @@
+"""
+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 os
+import tempfile
+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": dict.fromkeys(checkpoint, "checkpoint.safetensors"),
+    }
+
+    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]]),
+      )
+
+  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 831699474..a88aee01c 100644
--- a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py
+++ b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py
@@ -16,10 +16,23 @@
 
 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.max_utils import flash_block_sizes_for_candidate
+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,
+)
 
 
 class LTX2PipelineTest(unittest.TestCase):
@@ -55,6 +68,127 @@ 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_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_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(
+        {
+            "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()
@@ -188,13 +322,17 @@ 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_config.seed = 0
     mock_mesh = MagicMock()
+    mock_Mesh.return_value = mock_mesh
 
     mock_common = {
         "vae": MagicMock(),
@@ -217,11 +355,13 @@ def test_load_and_init(self, mock_quantize, mock_create_common, mock_load_transf
 
     pipeline, transformer = LTX2Pipeline._load_and_init(mock_config, None, vae_only=False, load_transformer=True)
 
+    from unittest.mock import ANY
+
     # Assert load_transformer was called with the components
     mock_load_transformer.assert_called_once_with(
-        devices_array=mock_common["devices_array"],
-        mesh=mock_common["mesh"],
-        rngs=mock_common["rngs"],
+        devices_array=mock_create_device_mesh.return_value,
+        mesh=mock_mesh,
+        rngs=ANY,
         config=mock_config,
         restored_checkpoint=None,
     )
diff --git a/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py b/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py
index f8c5b8d62..1e0b22c16 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,
@@ -137,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 = {
+            "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)
@@ -203,7 +308,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 +321,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/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
new file mode 100644
index 000000000..7cf53af32
--- /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
+from types import SimpleNamespace
+
+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,
+    local_tiled_seq_len,
+    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:
+  return local_tiled_seq_len(full_seq, attention, context_shards, ulysses_shards)
+
+
+@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):
+    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
+    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),
+        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,
+    )
+    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()
diff --git a/src/maxdiffusion/utils/tile_size_grid_search.py b/src/maxdiffusion/utils/tile_size_grid_search.py
index 38571a30b..3b7afb90b 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,76 @@ 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 +476,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 +491,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)