ltx2/svg: address #498 review (profiling note, consistent SVG check, AOT key)

Add svg_attention_enabled() so model construction, the CFG cache/MagCache check and the AOT cache key agree that SVG is off at svg_spatial_density >= 1.0. Key the LTX-2 AOT cache on svg_* settings, including those recorded in the transformer attention_config, only when SVG is active. Note in docs/svg.md that the default 5-step profile never reaches an SVG window that starts at step 10.

Author: jitendra-jalwaniya

docs/svg.md

@@ -163,3 +163,5 @@ The tests cover routing and layout, configuration propagation, schedules, shardi
 For CPU semantics checks, set `JAX_PLATFORMS=cpu` and `XLA_FLAGS=--xla_force_host_platform_device_count=8`. Run the suite separately on an eight-device TPU host to exercise the compiled kernels. Tests restricted to one platform are skipped on the other.
 
 For profiling, enable `enable_jax_named_scopes=True` and capture a short warm denoising interval. Routing, placement, main attention, padding cleanup, merging, and restoration have named scopes, including `svg_route_profile`, `svg_layout_place`, `svg_union_main`, `svg_tail_cleanup`, `svg_lse_merge`, and `svg_layout_restore`. The `svg_kernel_c_tiles…` scope reports the fraction of physical tiles executed, which differs from attention-pair density and total transformer FLOP savings.
+
+The LTX-2 profiling pass runs a separate `profiler_steps`-step denoise (steps `0` through `profiler_steps - 1`), and `profiler_steps` defaults to `5`. With the LTX-2 example above (`svg_active_start_step: 10`), that profile contains only dense steps. Set `profiler_steps` above `svg_active_start_step`, or lower `svg_active_start_step` while profiling, so the trace includes sparse SVG steps.

src/maxdiffusion/generate_ltx2.py

@@ -27,6 +27,7 @@
 
 from maxdiffusion.utils.export_utils import export_to_video_with_audio
 from maxdiffusion.loaders.ltx2_lora_nnx_loader import LTX2NNXLoraLoader
+from maxdiffusion.pipelines.ltx2.ltx2_pipeline import svg_attention_enabled
 
 
 def get_git_commit_hash():
@@ -201,6 +202,13 @@ def ltx2_aot_metadata(config, pipeline, source_revision=None):
   ):
     transformer_config.pop(key, None)
 
+  # SVG settings only shape the compiled graph when SVG is active. When it is
+  # off, leave them out so editing an unused svg_* value keeps the cache.
+  svg_enabled = svg_attention_enabled(config)
+  attention_config = transformer_config.get("attention_config")
+  if not svg_enabled and isinstance(attention_config, dict):
+    transformer_config["attention_config"] = {k: v for k, v in attention_config.items() if not str(k).startswith("svg_")}
+
   config_keys = [
       "attention",
       "a2v_attention_kernel",
@@ -223,26 +231,29 @@ def ltx2_aot_metadata(config, pipeline, source_revision=None):
       "sharding",
       "weights_dtype",
       "activations_dtype",
-      "use_svg_attention",
-      "svg_implementation",
-      "svg_spatial_density",
-      "svg_sample_max_row",
-      "svg_profile_query_count",
-      "svg_profile_seed",
-      "svg_dense_layer_fraction",
-      "svg_dense_timestep_fraction",
-      "svg_active_start_step",
-      "svg_active_end_step",
-      "svg_active_start_layer",
-      "svg_active_end_layer",
-      "svg_num_train_timesteps",
-      "svg_num_layers",
-      "svg_include_first_frame",
-      "svg_global_stride",
-      "svg_global_offset",
-      "svg_flash_block_sizes",
   ]
+  if svg_enabled:
+    config_keys += [
+        "svg_implementation",
+        "svg_spatial_density",
+        "svg_sample_max_row",
+        "svg_profile_query_count",
+        "svg_profile_seed",
+        "svg_dense_layer_fraction",
+        "svg_dense_timestep_fraction",
+        "svg_active_start_step",
+        "svg_active_end_step",
+        "svg_active_start_layer",
+        "svg_active_end_layer",
+        "svg_num_train_timesteps",
+        "svg_num_layers",
+        "svg_include_first_frame",
+        "svg_global_stride",
+        "svg_global_offset",
+        "svg_flash_block_sizes",
+    ]
   config_dict = {k: getattr(config, k, None) for k in config_keys}
+  config_dict["use_svg_attention"] = svg_enabled
 
   device = jax.devices()[0]
   return {

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -287,6 +287,18 @@ def _add_sharding_rule(vs: nnx.Variable, logical_axis_rules) -> nnx.Variable:
   return vs
 
 
+def svg_attention_enabled(config) -> bool:
+  """Returns whether SVG sparse attention is actually active for `config`.
+
+  A density of 1.0 or more is dense attention, so SVG stays off even when
+  `use_svg_attention` is set.
+  """
+  if not bool(getattr(config, "use_svg_attention", False)):
+    return False
+  density = getattr(config, "svg_spatial_density", None)
+  return float(0.25 if density is None else density) < 1.0
+
+
 def create_sharded_logical_transformer(
     devices_array: np.array,
     mesh: Mesh,
@@ -341,18 +353,14 @@ def _cfg(name, default):
     value = getattr(config, name, None)
     return default if value is None else value
 
-  expert_density = float(_cfg("svg_spatial_density", 0.25))
-
-  use_svg = bool(getattr(config, "use_svg_attention", False)) and (expert_density < 1.0)
-
   ltx2_config["attention_config"] = {
       "use_base2_exp": getattr(config, "use_base2_exp", False),
       "use_experimental_scheduler": getattr(config, "use_experimental_scheduler", False),
       "ulysses_shards": getattr(config, "ulysses_shards", -1),
       "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1),
-      "use_svg_attention": use_svg,
+      "use_svg_attention": svg_attention_enabled(config),
       "svg_implementation": _cfg("svg_implementation", "official_svg"),
-      "svg_spatial_density": expert_density,
+      "svg_spatial_density": float(_cfg("svg_spatial_density", 0.25)),
       "svg_sample_max_row": _cfg("svg_sample_max_row", 10000),
       "svg_profile_query_count": _cfg("svg_profile_query_count", 64),
       "svg_profile_seed": _cfg("svg_profile_seed", 0),
@@ -1866,7 +1874,7 @@ def __call__(
     if stg_scale > 0.0 and guidance_scale <= 1.0:
       raise ValueError("Spatio-temporal guidance requires guidance_scale > 1.0.")
 
-    if self.config and bool(getattr(self.config, "use_svg_attention", False)):
+    if self.config and svg_attention_enabled(self.config):
       if getattr(self.config, "use_cfg_cache", False) or getattr(self.config, "use_magcache", False):
         raise ValueError("SVG sparse attention cannot be combined with CFG cache or MagCache.")
 

src/maxdiffusion/tests/ltx2/test_generate_ltx2_aot.py

@@ -88,6 +88,45 @@ def test_codegen_and_remat_inputs_change_cache_fingerprint(self):
     self.assertNotEqual(fingerprint_b, fingerprint_c)
     self.assertNotEqual(fingerprint_c, fingerprint_d)
 
+  def _set_svg(self, **svg):
+    for name, value in svg.items():
+      setattr(self.config, name, value)
+    attention_config = {"use_base2_exp": False}
+    attention_config.update({k: v for k, v in svg.items() if k.startswith("svg_")})
+    # Like create_sharded_logical_transformer, record the effective SVG flag.
+    attention_config["use_svg_attention"] = generate_ltx2.svg_attention_enabled(self.config)
+    self.pipeline.transformer.config = {"attention_config": attention_config}
+
+  def test_unused_svg_settings_do_not_change_cache_fingerprint(self):
+    self._set_svg(use_svg_attention=False, svg_spatial_density=0.25, svg_active_start_step=10)
+    _, fingerprint_a = self._fingerprint("same-revision")
+    self._set_svg(use_svg_attention=False, svg_spatial_density=0.5, svg_active_start_step=20)
+    _, fingerprint_b = self._fingerprint("same-revision")
+
+    self.assertEqual(fingerprint_a, fingerprint_b)
+
+  def test_svg_settings_change_cache_fingerprint_when_svg_is_active(self):
+    self._set_svg(use_svg_attention=True, svg_spatial_density=0.25, svg_active_start_step=10)
+    _, fingerprint_a = self._fingerprint("same-revision")
+    self._set_svg(use_svg_attention=True, svg_spatial_density=0.5, svg_active_start_step=10)
+    _, fingerprint_b = self._fingerprint("same-revision")
+    self._set_svg(use_svg_attention=True, svg_spatial_density=0.5, svg_active_start_step=20)
+    _, fingerprint_c = self._fingerprint("same-revision")
+
+    self.assertNotEqual(fingerprint_a, fingerprint_b)
+    self.assertNotEqual(fingerprint_b, fingerprint_c)
+
+  def test_svg_at_density_one_shares_the_dense_cache_fingerprint(self):
+    self._set_svg(use_svg_attention=False, svg_spatial_density=0.25)
+    _, dense = self._fingerprint("same-revision")
+    self._set_svg(use_svg_attention=True, svg_spatial_density=1.0)
+    _, density_one = self._fingerprint("same-revision")
+    self._set_svg(use_svg_attention=True, svg_spatial_density=0.25)
+    _, sparse = self._fingerprint("same-revision")
+
+    self.assertEqual(dense, density_one)
+    self.assertNotEqual(dense, sparse)
+
   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"):

src/maxdiffusion/tests/ltx2/test_svg_config_propagation_ltx2.py

@@ -277,6 +277,33 @@ def test_cfg_and_magcache_incompatibility_validation(self):
     with self.assertRaisesRegex(ValueError, "SVG sparse attention cannot be combined with CFG cache or MagCache"):
       pipeline(prompt="test prompt")
 
+  def test_cfg_and_magcache_allowed_when_svg_density_is_one(self):
+    class ReachedEncodeError(Exception):
+      pass
+
+    pipeline = LTX2Pipeline.__new__(LTX2Pipeline)
+    pipeline.check_inputs = MagicMock()
+    pipeline.encode_prompt = MagicMock(side_effect=ReachedEncodeError)
+    for cache_flags in ({"use_cfg_cache": True, "use_magcache": False}, {"use_cfg_cache": False, "use_magcache": True}):
+      pipeline.config = SimpleNamespace(use_svg_attention=True, svg_spatial_density=1.0, **cache_flags)
+      # SVG is dense at density 1.0, so the call must get past the SVG check.
+      with self.assertRaises(ReachedEncodeError):
+        pipeline(prompt="test prompt")
+
+  def test_svg_attention_enabled(self):
+    tests = [
+        ("flag off", SimpleNamespace(use_svg_attention=False, svg_spatial_density=0.25), False),
+        ("flag missing", SimpleNamespace(), False),
+        ("sparse density", SimpleNamespace(use_svg_attention=True, svg_spatial_density=0.25), True),
+        ("density one", SimpleNamespace(use_svg_attention=True, svg_spatial_density=1.0), False),
+        ("density above one", SimpleNamespace(use_svg_attention=True, svg_spatial_density=1.5), False),
+        ("density missing uses default", SimpleNamespace(use_svg_attention=True), True),
+        ("density none uses default", SimpleNamespace(use_svg_attention=True, svg_spatial_density=None), True),
+    ]
+    for name, config, want in tests:
+      with self.subTest(name):
+        self.assertEqual(ltx2_pipeline.svg_attention_enabled(config), want)
+
 
 if __name__ == "__main__":
   unittest.main()