fix: strengthen LTX2 AOT cache key

Author: Perseus14

src/maxdiffusion/generate_ltx2.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from typing import Sequence
+import json
 import jax
 import jax.numpy as jnp
 import time
@@ -151,6 +152,73 @@ def maybe_tune_block_sizes(config):
   )
 
 
+def _canonical_aot_value(value):
+  return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
+
+
+def ltx2_aot_metadata(config, pipeline):
+  """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.
+  """
+  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 {
+      "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)),
+      "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)),
+      "remat_policy": str(getattr(config, "remat_policy", "NONE")),
+      "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 None:
     maybe_tune_block_sizes(config)
@@ -233,17 +301,8 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   # Per-shape AOT executable cache
   aot_cache.install(
       getattr(config, "aot_cache_dir", ""),
-      meta={
-          "model": config.pretrained_model_name_or_path,
-          "attention": getattr(config, "attention", ""),
-          "flash_block_sizes": str(getattr(config, "flash_block_sizes", "")),
-          "mesh_shape": str(pipeline.mesh.shape) if pipeline and hasattr(pipeline, "mesh") and pipeline.mesh else "",
-          "weights_dtype": str(getattr(config, "weights_dtype", "bfloat16")),
-          "activations_dtype": str(getattr(config, "activations_dtype", "bfloat16")),
-          "scan_layers": str(getattr(config, "scan_layers", True)),
-          "jax": jax.__version__,
-      },
-      mesh=pipeline.mesh if pipeline else None,
+      meta=ltx2_aot_metadata(config, pipeline),
+      mesh=pipeline.mesh,
   )
   aot_cache.wait_for_loads()