ltx2: fix custom block size fallback, framewise VAE tiling, v6 tile-search VMEM cap

Author: jitendra-jalwaniya

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -125,6 +125,8 @@ profiler_steps: 5
 enable_jax_named_scopes: False
 
 replicate_vae: False
+enable_vae_tiling: False
+enable_vae_slicing: False
 
 run_text_encoder_on_tpu: False
 # Dynamically disables VAE slicing and distributes the batch dimension to avoid HBM OOM for larger batch sizes.

src/maxdiffusion/configs/ltx2_video.yml

@@ -131,6 +131,8 @@ enable_jax_named_scopes: False
 
 replicate_vae: False
 use_bwe: False
+enable_vae_tiling: False
+enable_vae_slicing: False
 
 run_text_encoder_on_tpu: False
 # Dynamically disables VAE slicing and distributes the batch dimension to avoid HBM OOM for larger batch sizes.

src/maxdiffusion/generate_ltx2.py

@@ -105,6 +105,9 @@ def maybe_tune_block_sizes(config):
   vmem_limit_bytes = int(
       keys.get("tile_search_vmem_limit_bytes") or config.flash_block_sizes.get("vmem_limit_bytes") or 64 * 1024 * 1024
   )
+  if "v6" in jax.devices()[0].device_kind.lower():
+    vmem_limit_bytes = min(vmem_limit_bytes, 32 * 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(

src/maxdiffusion/models/attention_flax.py

@@ -381,7 +381,7 @@ def _extract_custom_block_sizes(flash_block_sizes):
   bq = 4864
   bkv = 1024
   bkv_compute = 1024
-  bkv_compute_in = 1024
+  bkv_compute_in = None
   heads_per_tile = 1
   vmem_limit_bytes = None
   if flash_block_sizes is not None:
@@ -390,22 +390,25 @@ def _extract_custom_block_sizes(flash_block_sizes):
       bq = get("block_q", None) or bq
       bkv = get("block_kv", None) or bkv
       bkv_compute = get("block_kv_compute", None) or bkv_compute
-      bkv_compute_in = get("block_kv_compute_in", None) or bkv_compute_in
+      bkv_compute_in = get("block_kv_compute_in", None)
       heads_per_tile = get("heads_per_tile", None) or heads_per_tile
       vmem_limit_bytes = get("vmem_limit_bytes", None) or vmem_limit_bytes
     else:
       bq = getattr(flash_block_sizes, "block_q", None) or bq
       bkv = getattr(flash_block_sizes, "block_kv", None) or bkv
       bkv_compute = getattr(flash_block_sizes, "block_kv_compute", None) or bkv_compute
-      bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", None) or bkv_compute_in
+      bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", None)
       heads_per_tile = getattr(flash_block_sizes, "heads_per_tile", None) or heads_per_tile
       vmem_limit_bytes = getattr(flash_block_sizes, "vmem_limit_bytes", None) or vmem_limit_bytes
   # A BlockSizes object carries heads_per_tile=None when the config dict omitted
   # it; getattr then returns that None instead of the default, so coerce it back
   # to 1 (the custom-kernel default) to keep the `heads_per_tile > 1` guards safe.
   if heads_per_tile is None:
     heads_per_tile = 1
-  bkv_compute_in = min(bkv_compute, bkv_compute_in)
+  if bkv_compute_in is None:
+    bkv_compute_in = bkv_compute
+  else:
+    bkv_compute_in = min(bkv_compute, bkv_compute_in)
   return bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes
 
 

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -558,9 +558,11 @@ def enable_vae_tiling(self):
     if hasattr(self.vae, "enable_tiling"):
       self.vae.enable_tiling()
     self.vae.use_tiling = True
+    self.vae.use_framewise_decoding = True
 
   def disable_vae_tiling(self):
     self.vae.use_tiling = False
+    self.vae.use_framewise_decoding = False
 
   @classmethod
   def load_tokenizer(cls, config: HyperParameters):