Merge pull request #480 from Ravi-VK:svg-head-local-pr-clean

PiperOrigin-RevId: 987128543

Author: Google-ML-Automation

src/maxdiffusion/aot_cache.py

@@ -73,6 +73,77 @@ def transformer_forward_pass(...):
 _FORMAT_VERSION = 1
 
 
+def _is_graphdef(x: Any) -> bool:
+  return hasattr(x, "attributes") and hasattr(x, "nodes")
+
+
+def _graphdef_desc(gd: Any) -> str:
+  """Extracts a process-deterministic digest of static attributes in an nnx.GraphDef."""
+  items = []
+  for k, v in getattr(gd, "attributes", ()):
+    if str(k).startswith("_pytree__"):
+      continue
+    if hasattr(v, "value"):
+      try:
+        s = json.dumps(
+            v.value,
+            sort_keys=True,
+            default=lambda o: re.sub(r"0x[0-9a-fA-F]+", "@", repr(o)),
+        )
+      except Exception:  # noqa: BLE001
+        s = re.sub(r"0x[0-9a-fA-F]+", "@", repr(v.value))
+      items.append(f"{k}:{s}")
+  return "GraphDef(" + ",".join(items) + ")"
+
+
+def extract_svg_meta(config: Any, pipeline: Any = None) -> dict[str, Any]:
+  """Extracts all graph-affecting SVG configuration settings for AOT cache identity."""
+  svg_keys = (
+      "use_svg_attention",
+      "svg_implementation",
+      "svg_spatial_density",
+      "svg_high_noise_density",
+      "svg_low_noise_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",
+  )
+  meta: dict[str, Any] = {}
+  if config is not None:
+    for k in svg_keys:
+      val = getattr(config, k, None) if not isinstance(config, dict) else config.get(k, None)
+      if val is not None:
+        meta[k] = str(val)
+  if pipeline is not None:
+    for attr in ("transformer", "high_noise_transformer", "low_noise_transformer"):
+      t = getattr(pipeline, attr, None)
+      if t is not None:
+        t_cfg = getattr(t, "config", None)
+        attn_cfg = (
+            getattr(t_cfg, "attention_config", None)
+            or (t_cfg.get("attention_config") if isinstance(t_cfg, dict) else None)
+        )
+        if isinstance(attn_cfg, dict):
+          meta[f"{attr}_svg_config"] = json.dumps(
+              {k: v for k, v in attn_cfg.items() if "svg" in str(k)},
+              sort_keys=True,
+              default=str,
+          )
+  return meta
+
+
 def _metadata_fingerprint(meta: dict[str, Any]) -> str:
   """Returns the stable filename fingerprint for install-time metadata."""
   serialized = json.dumps(meta, sort_keys=True, default=str)
@@ -87,15 +158,18 @@ def _dynamic_signature(args: tuple, kwargs: dict) -> str:
   addresses and hash-order-dependent content that differ per process and
   made signatures never match across restarts (measured: every array
   part stable, only the treedef part unstable). Static graph metadata
-  not visible in key paths (attention kernel, dtypes, model path) is
-  covered by the install-time config fingerprint in the filename.
+  in GraphDef leaves is extracted deterministically via ``_graphdef_desc``.
   Array leaves contribute shape/dtype; non-array leaves (python scalars,
   None flags) contribute an address-stripped repr.
   """
-  leaves_with_paths = jax.tree_util.tree_flatten_with_path((args, kwargs))[0]
+  leaves_with_paths = jax.tree_util.tree_flatten_with_path(
+      (args, kwargs), is_leaf=_is_graphdef
+  )[0]
   parts = []
   for path, leaf in leaves_with_paths:
-    if hasattr(leaf, "shape") and hasattr(leaf, "dtype"):
+    if _is_graphdef(leaf):
+      desc = _graphdef_desc(leaf)
+    elif hasattr(leaf, "shape") and hasattr(leaf, "dtype"):
       desc = f"{tuple(leaf.shape)}:{leaf.dtype}"
     else:
       desc = re.sub(r"0x[0-9a-fA-F]+", "@", repr(leaf))

src/maxdiffusion/configs/base_wan_14b.yml

@@ -83,7 +83,26 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
+# Sparse VideoGen (SVG) Attention configuration
+use_svg_attention: False
+svg_spatial_density: 0.25
+svg_sample_max_row: 10000
+svg_profile_query_count: 64
+svg_profile_seed: 0
+svg_active_start_step: -1
+svg_active_end_step: -1
+svg_active_start_layer: -1
+svg_active_end_layer: -1
+svg_include_first_frame: True
+svg_high_noise_density: -1.0
+svg_low_noise_density: -1.0
+# Tiling for the sparse SVG kernel only. Empty means "reuse flash_block_sizes",
+# which is rarely what you want: the dense ring kernel is tuned for large kv
+# compute blocks and the sparse kernel for small ones. Must stay a dict so the
+# command line can override it with JSON.
+svg_flash_block_sizes: {}
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
+
 use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -80,7 +80,26 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
+# Sparse VideoGen (SVG) Attention configuration
+use_svg_attention: False
+svg_spatial_density: 0.25
+svg_sample_max_row: 10000
+svg_profile_query_count: 64
+svg_profile_seed: 0
+svg_active_start_step: -1
+svg_active_end_step: -1
+svg_active_start_layer: -1
+svg_active_end_layer: -1
+svg_include_first_frame: True
+svg_high_noise_density: -1.0
+svg_low_noise_density: -1.0
+# Tiling for the sparse SVG kernel only. Empty means "reuse flash_block_sizes",
+# which is rarely what you want: the dense ring kernel is tuned for large kv
+# compute blocks and the sparse kernel for small ones. Must stay a dict so the
+# command line can override it with JSON.
+svg_flash_block_sizes: {}
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
+
 use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.

src/maxdiffusion/configs/base_wan_27b.yml

@@ -83,7 +83,29 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
+# Sparse VideoGen (SVG) Attention configuration
+use_svg_attention: False
+svg_spatial_density: 0.35
+svg_sample_max_row: 10000
+svg_profile_query_count: 64
+svg_profile_seed: 0
+svg_active_start_step: -1
+svg_active_end_step: -1
+svg_active_start_layer: -1
+svg_active_end_layer: -1
+svg_include_first_frame: True
+svg_high_noise_density: -1.0
+svg_low_noise_density: -1.0
+# Tiling for the sparse SVG kernel only. Empty means "reuse flash_block_sizes",
+# which is rarely what you want: the dense ring kernel is tuned for large kv
+# compute blocks and the sparse kernel for small ones. Must stay a dict so the
+# command line can override it with JSON. The tuned v6e-8 720p values are
+# {"block_q":3328,"block_kv":2816,"block_kv_compute":256,
+#  "block_kv_compute_in":256,"heads_per_tile":1,"vmem_limit_bytes":67108864}.
+svg_flash_block_sizes: {}
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, tokamax_ring_custom, ulysses, ulysses_custom, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_bidir
+
+
 #
 # Best 2D-ring / USP (Ulysses x ring) configs for WAN2.2-T2V-A14B (720x1280, 81 frames)
 # Set attention=ulysses_ring_custom and ulysses_shards=U (ring degree R=CP/U):
@@ -314,6 +336,7 @@ max_train_steps: 1500
 num_train_epochs: 1
 seed: 0
 output_dir: 'sdxl-model-finetuned'
+tensorboard_dir: ''
 per_device_batch_size: 1.0
 # If global_batch_size % jax.device_count is not 0, use FSDP sharding.
 global_batch_size: 0

src/maxdiffusion/generate_wan.py

@@ -314,6 +314,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
           "activations_dtype": str(config.activations_dtype),
           "scan_layers": str(config.scan_layers),
           "jax": jax.__version__,
+          **aot_cache.extract_svg_meta(config, pipeline),
       },
       mesh=pipeline.mesh,
   )

src/maxdiffusion/kernels/custom_svg_attention_dispatch.py

@@ -0,0 +1,116 @@
+"""
+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.
+
+Production entry points for SVG attention.
+
+This module is the top layer of the SVG kernel stack and exists to keep that
+stack acyclic:
+
+  custom_splash_attention                       (dense primitives)
+    -> custom_svg_static_range_attention        (tiling, partial kernels,
+                                                 exact reference builder)
+      -> custom_svg_balanced_rounding_attention (boundary selection, production
+                                                 builder)
+        -> custom_svg_attention_dispatch        (this module)
+
+Each arrow points from a dependency to its dependent, and no arrow runs the
+other way. Callers that want "the SVG kernel we ship" use this module; callers
+that want a specific implementation import the corresponding layer directly.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxdiffusion.kernels import custom_svg_balanced_rounding_attention as balanced
+from maxdiffusion.kernels import custom_svg_static_range_attention as static_range
+
+SVGBlockSizes = static_range.SVGBlockSizes
+
+
+def make_svg_static_range_mha(
+    *,
+    block_sizes,
+    orig_q_seq_len,
+    orig_kv_seq_len,
+    band_width,
+    frame_size,
+    include_first_frame=True,
+    bkv_compute_in=None,
+    use_base2_exp=True,
+    use_experimental_scheduler=False,
+    vmem_limit_bytes=None,
+):
+  """Build the production balanced-rounding SVG implementation.
+
+  The upstream path intentionally has no environment-variable policy switch.
+  Production uses the validated `global_balanced`, scale=1.0 policy. The exact
+  reference remains available explicitly via
+  `custom_svg_static_range_attention.make_svg_exact_static_range_mha`.
+  """
+  return balanced.make_svg_balanced_rounding_mha(
+      policy="global_balanced",
+      budget_scale=1.0,
+      block_sizes=block_sizes,
+      orig_q_seq_len=orig_q_seq_len,
+      orig_kv_seq_len=orig_kv_seq_len,
+      band_width=band_width,
+      frame_size=frame_size,
+      include_first_frame=include_first_frame,
+      bkv_compute_in=bkv_compute_in,
+      use_base2_exp=use_base2_exp,
+      use_experimental_scheduler=use_experimental_scheduler,
+      vmem_limit_bytes=vmem_limit_bytes,
+  )
+
+
+def custom_svg_static_range_attention(
+    query,
+    key,
+    value,
+    band_width,
+    anchor_width=0,
+    global_stride=0,
+    global_offset=0,
+    mesh: Any = None,
+    axis_names_q=None,
+    axis_names_kv=None,
+    dtype=jnp.bfloat16,
+    block_sizes=None,
+    use_base2_exp=True,
+    use_experimental_scheduler=False,
+):
+  del global_stride, global_offset, mesh, axis_names_q, axis_names_kv, dtype
+  if block_sizes is None:
+    block_sizes = SVGBlockSizes()
+  qlen = query.shape[2]
+  klen = key.shape[2]
+  bw = int(np.asarray(band_width).max()) if isinstance(band_width, (jax.Array, np.ndarray)) else int(band_width)
+  kernel = make_svg_static_range_mha(
+      block_sizes=block_sizes,
+      orig_q_seq_len=qlen,
+      orig_kv_seq_len=klen,
+      band_width=bw,
+      frame_size=int(anchor_width) if anchor_width > 0 else 0,
+      include_first_frame=anchor_width > 0,
+      use_base2_exp=use_base2_exp,
+      use_experimental_scheduler=use_experimental_scheduler,
+  )
+  out = jax.vmap(kernel, in_axes=(0, 0, 0))(query, key, value)
+  return jnp.swapaxes(out, 2, 3)