diff --git a/docs/images/svg/attention-tiles.png b/docs/images/svg/attention-tiles.png new file mode 100644 index 000000000..3af737eae Binary files /dev/null and b/docs/images/svg/attention-tiles.png differ diff --git a/docs/images/svg/dense-vs-svg.png b/docs/images/svg/dense-vs-svg.png new file mode 100644 index 000000000..ebf99af31 Binary files /dev/null and b/docs/images/svg/dense-vs-svg.png differ diff --git a/docs/images/svg/head-patterns.png b/docs/images/svg/head-patterns.png new file mode 100644 index 000000000..b26f26c70 Binary files /dev/null and b/docs/images/svg/head-patterns.png differ diff --git a/docs/svg.md b/docs/svg.md new file mode 100644 index 000000000..3ed29f315 --- /dev/null +++ b/docs/svg.md @@ -0,0 +1,147 @@ +# Sparse VideoGen attention on TPUs + +Video diffusion models generate a video through a sequence of denoising steps. At each step, attention lets each video token gather information from other tokens across space and time. This becomes expensive as the resolution and number of frames grow: dense attention considers every query–key pair, even though many interactions contribute very little to the output. + +Sparse spatiotemporal attention takes advantage of this structure. Instead of attending everywhere, a query attends to a smaller set of positions chosen to capture the spatial and temporal information it needs. + +## How SVG chooses where to attend + +[Sparse VideoGen (SVG)](https://arxiv.org/abs/2502.01776) observes that attention heads often favor different patterns. Spatial heads concentrate attention within a frame or nearby frames. Temporal heads concentrate attention around corresponding spatial positions across frames. These patterns let us approximate dense attention while computing fewer interactions. + +![Attention from the same query in a spatial head and a temporal head, shown across six latent frames.](images/svg/head-patterns.png) + +*Observed attention weights for the same query in two heads. The spatial head concentrates 94.8% of its attention on the query frame, while the temporal head places substantial attention near corresponding positions in other frames. Cyan squares mark the query's spatial position; `m` is the attention mass in each displayed frame. Colors use a shared logarithmic scale. These are examples of observed attention, rather than the masks themselves.* + +SVG makes the choice separately for each head: + +1. **Profile a few queries.** Compute dense attention outputs for a small sample of query tokens, using all keys. +2. **Compare two patterns.** Compute the sampled outputs under spatial and temporal masks, then measure each one's error relative to dense attention. +3. **Use the better approximation.** Select the lower-error pattern and apply sparse attention to all queries in that head. + +The choice is recomputed at each active layer and denoising step. A head does not need to keep the same assignment throughout generation. Sparsity is configurable, so the same method can trade a smaller approximation error for a larger reduction in computation. + +This implementation reimplements routing, token placement, and kernel execution for Wan in MaxDiffusion. The [original SVG implementation](https://github.com/svg-project/Sparse-VideoGen) provides the reference method. + +## Making sparse attention efficient on TPU + +Skipping query–key interactions only helps if the hardware can skip the corresponding work efficiently. Our implementation arranges tokens so that both spatial and temporal patterns can use the same local-band attention kernel. Spatial heads keep frame-major order; temporal heads group corresponding spatial positions across frames. Outputs are restored to their original order afterward. + +TPUs compute attention in tiles. A tile can lie entirely inside the sparse pattern, entirely outside it, or cross its boundary. + +A local attention band over a query–key tile grid, highlighting full, boundary, and skipped tiles. + +*Sparse pattern before tile rounding. Query and key indices refer to the selected token layout. Blue indicates retained interactions and gray indicates skipped interactions. The prefix anchor is omitted for clarity.* + +We round boundary tiles to either keep or skip them, approximately preserving the attention-pair budget of the original pattern. This slightly changes which interactions are retained, but lets all selected interior tiles run through one kernel without per-token sparse masking. Only tiles touching sequence padding need an additional validity mask; their outputs are combined with the main result using a numerically stable merge. + +Token placement and restoration run after the Ulysses exchange, on each device's local heads. This limits the layout work to the heads that device will actually process. + +The masks also include an optional prefix anchor. The `svg_include_first_frame` option retains the first `H × W` keys in the selected layout. For temporal heads, that prefix spans spatial positions across frames, so it does not correspond to the original first video frame. + +## Configuration and usage + +SVG is disabled by default. To enable it, set `use_svg_attention=True` and choose the densities and the steps and layers where sparsity should be active. Calls outside that interval continue to use dense attention. + +For example, these overrides select the moderate Wan2.2 policy: + +```yaml +use_svg_attention: True +svg_high_noise_density: 0.50 +svg_low_noise_density: 0.20 +svg_active_start_step: 11 +svg_active_end_step: 40 +svg_active_start_layer: 1 +svg_active_end_layer: 40 +svg_profile_query_count: 64 +svg_sample_max_row: 10000 +svg_profile_seed: 0 +svg_include_first_frame: True +``` + +Step and layer intervals are zero-based and half-open: `[11, 40)` includes steps 11 through 39. Steps refer to denoising iterations, not noise-timestep values. + +| Option | What it controls | +|---|---| +| `svg_spatial_density` | Density for single-expert Wan models; applies to either selected head pattern. | +| `svg_high_noise_density`, `svg_low_noise_density` | Separate densities for Wan2.2's two experts. | +| `svg_active_start_step`, `svg_active_end_step` | Denoising steps where SVG is enabled. | +| `svg_active_start_layer`, `svg_active_end_layer` | Transformer layers where SVG is enabled. | +| `svg_profile_query_count` | Number of queries sampled when choosing each head's pattern. | +| `svg_sample_max_row` | Limits query sampling to this prefix of the sequence. | +| `svg_profile_seed` | Seed for reproducible query sampling. | +| `svg_include_first_frame` | Enables the prefix anchor in the selected layout. | +| `svg_flash_block_sizes` | Sparse kernel tiling; an empty mapping uses `flash_block_sizes`. | + +Density controls the local-band width. The anchor and tile rounding affect the actual number of retained interactions, so density is not itself the fraction of total transformer FLOPs retained. + +The following is an example attention configuration for eight devices. Add it and the policy above to a Wan2.2 configuration used by `src/maxdiffusion/generate_wan.py`: + +```yaml +attention: ulysses_ring_custom_fixed_m +ici_data_parallelism: 2 +ici_context_parallelism: 4 +ulysses_shards: 2 +flash_block_sizes: + block_q: 6400 + block_kv: 2048 + block_kv_compute: 2048 + block_kv_compute_in: 1024 + heads_per_tile: 1 + vmem_limit_bytes: 67108864 +svg_flash_block_sizes: + block_q: 3328 + block_kv: 2816 + block_kv_compute: 256 + block_kv_compute_in: 256 + heads_per_tile: 1 + vmem_limit_bytes: 67108864 +``` + +Dense calls retain their configured ring split. SVG exchanges over the whole context axis, giving Ring2/Ulysses2 for dense calls and Ring1/Ulysses4 for sparse calls in this example. Tile sizes may need tuning for other shapes and devices. + +### Supported configurations + +SVG supports inference through the four custom Ulysses/ring attention backends. It requires matching self-attention QKV shapes, a matching video-token grid, and `heads_per_tile=1`. Heads and sequence length must divide evenly across the context shards, including any heads created by folding an unsharded batch. + +Training, Animate, external attention masks, periodic support, and chunked Ulysses are unsupported. SVG also cannot be combined with CFG cache or MagCache. + +## Performance and quality + +At 720p, SVG offers a configurable tradeoff between denoising latency and similarity to dense generation. The following Wan2.2 results use TPU v6e-8, 81 frames, and 40 denoising steps. Times are medians of three warm runs against same-node optimized fixed-M dense controls, using one prompt and seed. + +| Policy | Estimated total transformer FLOPs saved | Dense denoising | SVG denoising | Denoising speedup | PSNR (dB) | +|---|---:|---:|---:|---:|---:| +| Conservative | ≈27.2% | 153.46 s | 136.11 s | **1.13×** | **26.47** | +| Moderate | ≈32.2% | 153.37 s | 127.94 s | **1.20×** | **26.14** | +| Aggressive | ≈37.3% | 153.50 s | 119.86 s | **1.28×** | **24.80** | + +*PSNR is measured against dense outputs using FFmpeg's aggregate YUV metric.* + +In an earlier evaluation using the moderate SVG policy, Wan2.2 at 720p retained **98.1% of the dense baseline’s mean VBench dimension score** across a 31-prompt, 16-dimension screening subset. + +## Qualitative example + +![Three rows of giraffe video frames, with dense attention on the left and SVG on the right.](images/svg/dense-vs-svg.png) + +*Dense attention (left) and SVG (right), shown at three video frames. The overall scene and subject arrangement remain similar, with visible differences in ground texture, background detail, and coat patterns.* + +This illustration is separate from the benchmark measurements above. + +## Tests and profiling + +Run the SVG tests from the repository root: + +```bash +python -m pytest -q \ + src/maxdiffusion/tests/wan/svg_attention_test.py \ + src/maxdiffusion/tests/wan/svg_balanced_rounding_test.py \ + src/maxdiffusion/tests/wan/svg_config_propagation_test.py \ + src/maxdiffusion/tests/wan/svg_head_local_test.py \ + src/maxdiffusion/tests/wan/wan_pipeline_signature_test.py +``` + +The tests cover routing and layout, configuration propagation, schedules, sharding, unsupported configurations, and numerical agreement with attention references. Production-kernel tests cover sparse and density-one support, aligned and padded sequences, and natural and base-2 exponentials. + +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. diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py index c07af3e24..0e7cc344e 100644 --- a/src/maxdiffusion/aot_cache.py +++ b/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)) diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml index 837bbe98b..4c2b2b30c 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/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. diff --git a/src/maxdiffusion/configs/base_wan_1_3b.yml b/src/maxdiffusion/configs/base_wan_1_3b.yml index 6f5ea10b6..89bcfb885 100644 --- a/src/maxdiffusion/configs/base_wan_1_3b.yml +++ b/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. diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index bf8e1c740..f281b3110 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/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 diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index 662069ffd..a8a184f48 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/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, ) diff --git a/src/maxdiffusion/kernels/custom_svg_attention_dispatch.py b/src/maxdiffusion/kernels/custom_svg_attention_dispatch.py new file mode 100644 index 000000000..3bdd3ed6c --- /dev/null +++ b/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) diff --git a/src/maxdiffusion/kernels/custom_svg_balanced_rounding_attention.py b/src/maxdiffusion/kernels/custom_svg_balanced_rounding_attention.py new file mode 100644 index 000000000..4ec9b0b21 --- /dev/null +++ b/src/maxdiffusion/kernels/custom_svg_balanced_rounding_attention.py @@ -0,0 +1,370 @@ +""" +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. + +Budget-preserving tile-native boundary rounding for SVG attention. + +FULL tiles retain exact SVG semantics. Each exact SVG BOUNDARY tile is rounded +UP (execute the full real rectangle with padding-only masking) or DOWN (omit +it). The host policy chooses UP tiles so rounded pair work matches the exact +boundary-pair budget as closely as possible. If a query tile would have no +retained keys, one boundary tile is restored; its cost is included in the +reported budget error. Nonempty support takes priority over budget matching. + +Production execution then unions FULL tiles with selected rounded BOUNDARY +tiles. Tiles fully inside the physical sequence run together through one +mask-free Pallas partial; only tiles touching the Q/KV sequence tail use the +padding-only cleanup partial. This keeps the hot path uniform and avoids a +FULL-vs-BOUNDARY decision inside the kernel. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax +import jax.numpy as jnp +import numpy as np + +from maxdiffusion.kernels import custom_svg_balanced_rounding_partial as partial_impl +from maxdiffusion.kernels import custom_svg_static_range_attention as exact_svg + + +@dataclass(frozen=True) +class BoundaryTileStat: + qi: int + kj: int + real_pairs: int + exact_pairs: int + + @property + def alpha(self): + return float(self.exact_pairs) / max(float(self.real_pairs), 1.0) + + +def _pack_rows(rows, qtiles): + width = max(1, max((len(r) for r in rows), default=0)) + table = np.zeros((qtiles, width), dtype=np.int32) + active = np.zeros((qtiles,), dtype=np.int32) + for qi, row in enumerate(rows): + if row: + table[qi, : len(row)] = np.asarray(row, np.int32) + active[qi] = len(row) + return table, active + + +def _exact_pairs_for_tile( + *, + qi, + kj, + q_seq_len, + kv_seq_len, + bq, + bkv, + band_width, + frame_size, + include_first_frame, +): + q0, k0 = qi * bq, kj * bkv + q1, k1 = min(q_seq_len, q0 + bq), min(kv_seq_len, k0 + bkv) + if q1 <= q0 or k1 <= k0: + return 0, 0 + + qs = np.arange(q0, q1, dtype=np.int64) + lo = np.maximum(k0, qs - int(band_width)) + hi = np.minimum(k1 - 1, qs + int(band_width)) + local = np.maximum(0, hi - lo + 1) + if include_first_frame: + sink_hi = min(k1 - 1, int(frame_size) - 1) + sink_count = max(0, sink_hi - k0 + 1) + if sink_count: + overlap = np.maximum( + 0, + np.minimum(hi, sink_hi) - np.maximum(lo, k0) + 1, + ) + exact = local + sink_count - overlap + else: + exact = local + else: + exact = local + return int((q1 - q0) * (k1 - k0)), int(exact.sum()) + + +def build_boundary_stats( + *, + orig_q_seq_len, + orig_kv_seq_len, + block_sizes, + band_width, + frame_size, + include_first_frame=True, +): + bq, bkv = int(block_sizes.block_q), int(block_sizes.block_kv) + fm, fa, bm, ba = exact_svg._classify_tiles( + int(orig_q_seq_len), + int(orig_kv_seq_len), + bq, + bkv, + int(band_width), + int(frame_size), + bool(include_first_frame), + ) + stats = [] + for qi in range(bm.shape[0]): + for slot in range(int(ba[qi])): + kj = int(bm[qi, slot]) + rp, ep = _exact_pairs_for_tile( + qi=qi, + kj=kj, + q_seq_len=int(orig_q_seq_len), + kv_seq_len=int(orig_kv_seq_len), + bq=bq, + bkv=bkv, + band_width=int(band_width), + frame_size=int(frame_size), + include_first_frame=bool(include_first_frame), + ) + stats.append(BoundaryTileStat(qi, kj, rp, ep)) + return fm, fa, bm, ba, stats + + +def _closest_prefix(items, target): + ordered = sorted( + items, + key=lambda x: (-x.alpha, -x.exact_pairs, x.qi, x.kj), + ) + best_k, best_err, running = 0, abs(float(target)), 0 + for k, tile in enumerate(ordered, 1): + running += tile.real_pairs + err = abs(float(running) - float(target)) + if err < best_err: + best_k, best_err = k, err + return {(x.qi, x.kj) for x in ordered[:best_k]} + + +def select_boundary_tiles(stats, *, policy="global_balanced", budget_scale=1.0): + policy = str(policy).lower().strip() + if policy == "up": + return {(x.qi, x.kj) for x in stats} + if policy == "down": + return set() + if policy == "nearest": + return {(x.qi, x.kj) for x in stats if x.alpha >= 0.5} + + target = float(budget_scale) * sum(x.exact_pairs for x in stats) + if policy == "global_balanced": + return _closest_prefix(stats, target) + if policy == "row_balanced": + out = set() + for qi in sorted({x.qi for x in stats}): + row = [x for x in stats if x.qi == qi] + out |= _closest_prefix( + row, + float(budget_scale) * sum(x.exact_pairs for x in row), + ) + return out + raise ValueError(f"unknown SVG boundary rounding policy {policy!r}") + + +def build_selected_boundary_table( + *, + stats, + qtiles, + full_active=None, + policy="global_balanced", + budget_scale=1.0, +): + selected = select_boundary_tiles( + stats, + policy=policy, + budget_scale=budget_scale, + ) + # A global budget alone can leave real queries with no keys. Preserve the + # original selection except for empty rows, where validity takes priority + # over the approximate pair budget. A rounded tile covers every real query + # in its query block, even if its exact support only covered some of them. + coverage_added = set() + if full_active is not None: + covered = {qi for qi in range(qtiles) if full_active[qi] > 0} + covered.update(qi for qi, _ in selected) + for qi in range(qtiles): + if qi in covered: + continue + candidates = [x for x in stats if x.qi == qi and x.exact_pairs > 0] + if not candidates: + raise ValueError(f"SVG query tile {qi} has no valid attention support") + tile = min(candidates, key=lambda x: (-x.alpha, -x.exact_pairs, x.kj)) + coverage_added.add((tile.qi, tile.kj)) + selected |= coverage_added + rows = [[] for _ in range(qtiles)] + for tile in stats: + if (tile.qi, tile.kj) in selected: + rows[tile.qi].append(tile.kj) + for row in rows: + row.sort() + + table, active = _pack_rows(rows, qtiles) + exact = int(sum(x.exact_pairs for x in stats)) + rounded = int(sum(x.real_pairs for x in stats if (x.qi, x.kj) in selected)) + retained = int(sum(x.exact_pairs for x in stats if (x.qi, x.kj) in selected)) + target = float(budget_scale) * exact + report = { + "coverage_tiles_added": len(coverage_added), + "coverage_pairs_added": int(sum(x.real_pairs for x in stats if (x.qi, x.kj) in coverage_added)), + "policy": policy, + "budget_scale": float(budget_scale), + "boundary_tiles_total": len(stats), + "boundary_tiles_selected": len(selected), + "exact_boundary_pairs": exact, + "rounded_boundary_pairs": rounded, + "target_boundary_pairs": target, + "budget_error_pairs": rounded - target, + "budget_error_fraction": (rounded - target) / max(target, 1.0), + "retained_exact_pairs": retained, + "dropped_exact_pairs": exact - retained, + "added_outside_pairs": rounded - retained, + "boundary_exact_recall": retained / max(exact, 1), + "rounded_precision": retained / max(rounded, 1), + } + return table, active, report + + +def build_union_tail_tables( + *, + full_table, + full_active, + selected_boundary_table, + selected_boundary_active, + orig_q_seq_len, + orig_kv_seq_len, + block_sizes, +): + """Split the rounded physical tile set into mask-free main and tail tables. + + FULL and selected rounded BOUNDARY tiles are semantically identical once the + boundary decision has been made: both execute the whole real rectangle. They + therefore share one mask-free table whenever the whole hardware tile lies + inside the physical Q/KV sequence. Only tiles touching sequence padding are + sent to the padding-aware cleanup table. + """ + qtiles = full_table.shape[0] + bq = int(block_sizes.block_q) + bkv = int(block_sizes.block_kv) + main_rows = [[] for _ in range(qtiles)] + tail_rows = [[] for _ in range(qtiles)] + + for qi in range(qtiles): + row = { + *(int(x) for x in full_table[qi, : int(full_active[qi])]), + *(int(x) for x in selected_boundary_table[qi, : int(selected_boundary_active[qi])]), + } + q_full = (qi + 1) * bq <= int(orig_q_seq_len) + for kj in sorted(row): + k_full = (kj + 1) * bkv <= int(orig_kv_seq_len) + (main_rows if q_full and k_full else tail_rows)[qi].append(kj) + + main_table, main_active = _pack_rows(main_rows, qtiles) + tail_table, tail_active = _pack_rows(tail_rows, qtiles) + return main_table, main_active, tail_table, tail_active + + +def make_svg_balanced_rounding_mha( + *, + policy="global_balanced", + budget_scale=1.0, + 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, +): + fm, fa, bm, ba, stats = build_boundary_stats( + orig_q_seq_len=orig_q_seq_len, + orig_kv_seq_len=orig_kv_seq_len, + block_sizes=block_sizes, + band_width=band_width, + frame_size=frame_size, + include_first_frame=include_first_frame, + ) + sm, sa, budget = build_selected_boundary_table( + stats=stats, + qtiles=bm.shape[0], + full_active=fa, + policy=policy, + budget_scale=budget_scale, + ) + mm, ma, tm, ta = build_union_tail_tables( + full_table=fm, + full_active=fa, + selected_boundary_table=sm, + selected_boundary_active=sa, + orig_q_seq_len=orig_q_seq_len, + orig_kv_seq_len=orig_kv_seq_len, + block_sizes=block_sizes, + ) + common = { + "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, + } + main = exact_svg._make_partial( + table_np=mm, + active_np=ma, + mask_mode="none", + **common, + ) + tail = partial_impl.make_padding_partial_from_table( + table_np=tm, + active_np=ta, + **common, + ) + tail_tiles = int(ta.sum()) + + def attention(q, k, v): + with jax.named_scope("svg_union_main"): + om, lm = main(q, k, v) + if tail_tiles == 0: + return om + with jax.named_scope("svg_tail_cleanup"): + ot, lt = tail(q, k, v) + with jax.named_scope("svg_lse_merge"): + m = jnp.maximum(lm, lt) + exp = jnp.exp2 if use_base2_exp else jnp.exp + wm, wt = exp(lm - m), exp(lt - m) + den = wm + wt + return (om.astype(jnp.float32) * (wm / den)[:, None, :] + ot.astype(jnp.float32) * (wt / den)[:, None, :]).astype( + q.dtype + ) + + attention.full_tiles = int(fa.sum()) + attention.original_boundary_tiles = int(ba.sum()) + attention.selected_boundary_tiles = int(sa.sum()) + attention.union_main_tiles = int(ma.sum()) + attention.tail_cleanup_tiles = tail_tiles + attention.rounding_budget = budget + attention.rounding_policy = policy + attention.rounding_budget_scale = float(budget_scale) + return attention diff --git a/src/maxdiffusion/kernels/custom_svg_balanced_rounding_partial.py b/src/maxdiffusion/kernels/custom_svg_balanced_rounding_partial.py new file mode 100644 index 000000000..cc0c8a102 --- /dev/null +++ b/src/maxdiffusion/kernels/custom_svg_balanced_rounding_partial.py @@ -0,0 +1,260 @@ +""" +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. + +Pallas partial attention with padding-only masks for sequence-edge tiles. +""" + +from __future__ import annotations + +import functools + +import jax +import jax.numpy as jnp +import numpy as np +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu + +from maxdiffusion.kernels import custom_splash_attention as dense_custom + +NUM_SUBLANES = dense_custom.NUM_SUBLANES +NT_DIM_NUMBERS = dense_custom.NT_DIM_NUMBERS +MASK_VALUE = dense_custom.DEFAULT_MASK_VALUE + + +def _build_padding_metadata(table_np, active_np, *, q_seq_len, kv_seq_len, bq, bkv): + q_valid = np.zeros(table_np.shape, dtype=np.int32) + k_valid = np.zeros(table_np.shape, dtype=np.int32) + for qi in range(table_np.shape[0]): + q0 = qi * bq + qv = max(0, min(bq, q_seq_len - q0)) + for slot in range(int(active_np[qi])): + kj = int(table_np[qi, slot]) + k0 = kj * bkv + q_valid[qi, slot] = qv + k_valid[qi, slot] = max(0, min(bkv, kv_seq_len - k0)) + return q_valid, k_valid + + +def _padding_kernel( + kv_map_ref, + active_ref, + qvalid_ref, + kvalid_ref, + q_ref, + k_ref, + v_ref, + out_ref, + lse_ref, + m_ref, + l_ref, + oacc_ref, + *, + grid_width, + bq, + bkv, + bc, + bci, + dv, + use_base2_exp, +): + qi = pl.program_id(1) + slot = pl.program_id(2) + active = slot < active_ref[qi] + exp = jnp.exp2 if use_base2_exp else jnp.exp + log = jnp.log2 if use_base2_exp else jnp.log + + @pl.when(slot == 0) + def init(): + m_ref[...] = jnp.full_like(m_ref, MASK_VALUE) + l_ref[...] = jnp.zeros_like(l_ref) + oacc_ref[...] = jnp.zeros_like(oacc_ref) + + @pl.when(active) + def run(): + q = q_ref[...] + m_prev = m_ref[...] + l_prev = l_ref[...] + o_prev = oacc_ref[...] + + def body(ci, carry): + m_c, l_c, o_c = carry + off = ci * bc + sl = pl.ds(off, bc) + kc = k_ref[sl, :] + vc = v_ref[sl, :] + scores = lax.dot_general( + kc, + q, + NT_DIM_NUMBERS, + preferred_element_type=jnp.float32, + ) + + qc = jnp.arange(bq, dtype=jnp.int32)[None, :] + kr = off + jnp.arange(bc, dtype=jnp.int32)[:, None] + valid = (qc < qvalid_ref[qi, slot]) & (kr < kvalid_ref[qi, slot]) + scores = jnp.where(valid, scores, jnp.asarray(MASK_VALUE, scores.dtype)) + + for inner in range(0, bc, bci): + z = scores[inner : inner + bci] + vv = vc[inner : inner + bci] + m_chunk = z.max(axis=0)[None, :] + m_new = jnp.maximum(m_c, m_chunk) + p = exp(z - m_new[0:1]) + l_chunk = p.sum(axis=0, keepdims=True) + alpha = exp(m_c - m_new) + l_new = l_chunk + alpha * l_c + o_chunk = lax.dot_general( + vv, + p.astype(q_ref.dtype), + (((0,), (0,)), ((), ())), + preferred_element_type=jnp.float32, + ) + o_c = alpha[0:1] * o_c + o_chunk + m_c, l_c = m_new, l_new + return m_c, l_c, o_c + + m_prev, l_prev, o_prev = lax.fori_loop( + 0, + bkv // bc, + body, + (m_prev, l_prev, o_prev), + unroll=True, + ) + m_ref[...] = m_prev + l_ref[...] = l_prev + oacc_ref[...] = o_prev + + @pl.when(slot == grid_width - 1) + def finish(): + l = l_ref[...] + m = m_ref[...] + has_value = l > 0 + safe_l = jnp.where(has_value, l, jnp.ones_like(l)) + inv_l = jnp.tile(1.0 / safe_l, (dv // NUM_SUBLANES, 1)) + out = oacc_ref[...] * inv_l + out_ref[...] = jnp.where( + jnp.tile(has_value, (dv // NUM_SUBLANES, 1)), + out, + jnp.zeros_like(out), + ).astype(out_ref.dtype) + lse_ref[...] = jnp.where( + has_value, + m + log(safe_l), + jnp.asarray(MASK_VALUE, m.dtype), + ).astype(lse_ref.dtype) + + +def make_padding_partial_from_table( + *, + table_np, + active_np, + 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, +): + del band_width, frame_size, include_first_frame + bq = int(block_sizes.block_q) + bkv = int(block_sizes.block_kv) + bc = int(block_sizes.block_kv_compute) + bci = int(bkv_compute_in if bkv_compute_in is not None else block_sizes.block_kv_compute_in) + if bkv % bc or bc % bci: + raise ValueError(f"invalid blocks bkv={bkv}, bc={bc}, bci={bci}") + + q_valid, k_valid = _build_padding_metadata( + table_np, + active_np, + q_seq_len=int(orig_q_seq_len), + kv_seq_len=int(orig_kv_seq_len), + bq=bq, + bkv=bkv, + ) + scalars = tuple(jnp.asarray(x) for x in (table_np, active_np, q_valid, k_valid)) + qtiles, width = table_np.shape + + def partial(q, k, v): + heads, _, dq = q.shape + dv = v.shape[-1] + if dv % NUM_SUBLANES: + raise NotImplementedError(f"head_dim_v={dv} must be divisible by {NUM_SUBLANES}") + + def qmap(h, qi, slot, *refs): + del slot, refs + return (h, qi, 0) + + def kvmap(h, qi, slot, kvmap_ref, *refs): + del refs + return (h, kvmap_ref[qi, slot], 0) + + def outmap(h, qi, slot, *refs): + del slot, refs + return (h, 0, qi) + + outs = pl.pallas_call( + functools.partial( + _padding_kernel, + grid_width=width, + bq=bq, + bkv=bkv, + bc=bc, + bci=bci, + dv=dv, + use_base2_exp=bool(use_base2_exp), + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=4, + in_specs=[ + pl.BlockSpec((None, bq, dq), qmap), + pl.BlockSpec((None, bkv, dq), kvmap), + pl.BlockSpec((None, bkv, dv), kvmap), + ], + out_specs=[ + pl.BlockSpec((None, dv, bq), outmap), + pl.BlockSpec((None, NUM_SUBLANES, bq), outmap), + ], + scratch_shapes=[ + pltpu.VMEM((NUM_SUBLANES, bq), jnp.float32), + pltpu.VMEM((NUM_SUBLANES, bq), jnp.float32), + pltpu.VMEM((dv, bq), jnp.float32), + ], + grid=(heads, qtiles, width), + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": use_experimental_scheduler}, + disable_bounds_checks=True, + skip_device_barrier=True, + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=[ + jax.ShapeDtypeStruct((heads, dv, qtiles * bq), q.dtype), + jax.ShapeDtypeStruct((heads, NUM_SUBLANES, qtiles * bq), jnp.float32), + ], + )(*scalars, q, k, v) + return ( + outs[-2][:, :, :orig_q_seq_len], + outs[-1][:, 0, :orig_q_seq_len], + ) + + partial.boundary_tiles = int(active_np.sum()) + partial.boundary_table_shape = tuple(table_np.shape) + return partial diff --git a/src/maxdiffusion/kernels/custom_svg_static_range_attention.py b/src/maxdiffusion/kernels/custom_svg_static_range_attention.py new file mode 100644 index 000000000..9c1963991 --- /dev/null +++ b/src/maxdiffusion/kernels/custom_svg_static_range_attention.py @@ -0,0 +1,378 @@ +""" +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. + +Static-range SVG attention with exact and tile-rounded support. + +Mask behavior is selected at trace time to keep full tiles branch-free. + +This module is a leaf of the SVG kernel stack: it depends only on the dense +splash-attention primitives and must not import the balanced-rounding layer or +the dispatch layer. The production entry points live in +`custom_svg_attention_dispatch`. +""" + +from __future__ import annotations + +import dataclasses +import functools +import math + +import jax +import jax.numpy as jnp +import numpy as np +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu + +from maxdiffusion.kernels import custom_splash_attention as dense_custom + +NUM_SUBLANES = dense_custom.NUM_SUBLANES +NT_DIM_NUMBERS = dense_custom.NT_DIM_NUMBERS +DEFAULT_MASK_VALUE = dense_custom.DEFAULT_MASK_VALUE + + +@dataclasses.dataclass(frozen=True) +class SVGBlockSizes: + block_q: int = 3328 + block_kv: int = 2048 + block_kv_compute: int = 256 + block_kv_compute_in: int = 256 + + +def _classify_tiles(q_seq_len, kv_seq_len, bq, bkv, band_width, frame_size, include_first_frame): + q_tiles = math.ceil(q_seq_len / bq) + kv_tiles = math.ceil(kv_seq_len / bkv) + sink_last = (frame_size - 1) // bkv if include_first_frame else -1 + full_rows, boundary_rows = [], [] + + for qi in range(q_tiles): + q0 = qi * bq + q1 = min(q_seq_len, q0 + bq) - 1 + band0 = max(0, q0 - band_width) + band1 = min(kv_seq_len - 1, q1 + band_width) + first, last = band0 // bkv, band1 // bkv + live = set(range(first, last + 1)) + if include_first_frame: + live.update(range(0, sink_last + 1)) + + full, boundary = [], [] + for kj in sorted(x for x in live if 0 <= x < kv_tiles): + k0 = kj * bkv + k1 = min(kv_seq_len, k0 + bkv) - 1 + q_real_full = q0 + bq <= q_seq_len + k_real_full = k0 + bkv <= kv_seq_len + sink_full = include_first_frame and k1 < frame_size + local_full = max(abs(q0 - k1), abs(q1 - k0)) <= band_width + (full if q_real_full and k_real_full and (sink_full or local_full) else boundary).append(kj) + full_rows.append(full) + boundary_rows.append(boundary) + + def pack(rows): + width = max(1, max((len(r) for r in rows), default=0)) + table = np.zeros((q_tiles, width), np.int32) + active = np.zeros((q_tiles,), np.int32) + for qi, row in enumerate(rows): + if row: + table[qi, : len(row)] = row + active[qi] = len(row) + return table, active + + fm, fa = pack(full_rows) + bm, ba = pack(boundary_rows) + return fm, fa, bm, ba + + +def _partial_kernel( + kv_map_ref, + active_counts_ref, + q_ref, + k_ref, + v_ref, + o_ref, + lse_ref, + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + *, + mask_mode, + mask_value, + grid_width, + bq, + bkv, + bkv_compute, + bkv_compute_in, + head_dim_v, + q_seq_len, + kv_seq_len, + band_width, + frame_size, + include_first_frame, + use_base2_exp, +): + """Run one static tile-table partial. + + `mask_mode` is a Python/static argument, so Pallas specializes three distinct + kernels rather than carrying a runtime tile-kind branch through FULL work: + `none` for exact FULL tiles, `padding` for rounded boundary rectangles, and + `exact` for the exact SVG boundary predicate. + """ + if mask_mode not in ("none", "padding", "exact"): + raise ValueError(f"unknown SVG mask_mode {mask_mode!r}") + + float32 = jnp.float32 + repeats, rem = divmod(head_dim_v, NUM_SUBLANES) + if rem: + raise NotImplementedError(f"head_dim_v={head_dim_v} must divide {NUM_SUBLANES}") + + exp = jnp.exp2 if use_base2_exp else jnp.exp + log = jnp.log2 if use_base2_exp else jnp.log + qi = pl.program_id(1) + slot = pl.program_id(2) + active = slot < active_counts_ref[qi] + + @pl.when(slot == 0) + def init(): + m_scratch_ref[...] = jnp.full_like(m_scratch_ref, mask_value) + l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref) + o_scratch_ref[...] = jnp.zeros_like(o_scratch_ref) + + def body(ci, _): + m_prev = m_scratch_ref[...] + l_prev = l_scratch_ref[...] + o_prev = o_scratch_ref[:] + q = q_ref[...] + off = ci * bkv_compute + sl = pl.ds(off, bkv_compute) + kc = k_ref[sl, :] + vc = v_ref[sl, :] + qk = lax.dot_general(kc, q, NT_DIM_NUMBERS, preferred_element_type=float32) + + if mask_mode != "none": + kj = kv_map_ref[qi, slot] + qids = qi * bq + jnp.arange(bq, dtype=jnp.int32)[None, :] + kids = kj * bkv + off + jnp.arange(bkv_compute, dtype=jnp.int32)[:, None] + valid = (qids < q_seq_len) & (kids < kv_seq_len) + if mask_mode == "exact": + local = jnp.abs(qids - kids) <= band_width + anchor = (kids < frame_size) if include_first_frame else jnp.zeros_like(local) + valid = valid & (local | anchor) + qk = jnp.where(valid, qk, jnp.asarray(mask_value, qk.dtype)) + + for i in range(0, bkv_compute, bkv_compute_in): + z = qk[i : i + bkv_compute_in] + vv = vc[i : i + bkv_compute_in] + mc = z.max(axis=0)[None, :] + mn = jnp.maximum(m_prev, mc) + p = exp(z - mn[0:1]) + lc = p.sum(axis=0, keepdims=True) + alpha = exp(m_prev - mn) + ln = lc + alpha * l_prev + oc = lax.dot_general( + vv, + p.astype(q_ref.dtype), + (((0,), (0,)), ((), ())), + preferred_element_type=float32, + ) + o_prev = alpha[0:1] * o_prev + oc + m_prev, l_prev = mn, ln + + m_scratch_ref[...] = m_prev + l_scratch_ref[...] = l_prev + o_scratch_ref[:] = o_prev + + @pl.when(active) + def run(): + lax.fori_loop(0, bkv // bkv_compute, body, None, unroll=True) + + @pl.when(slot == grid_width - 1) + def finish(): + l = l_scratch_ref[...] + m = m_scratch_ref[...] + has = l > 0 + safe = jnp.where(has, l, jnp.ones_like(l)) + inv = jnp.tile(1.0 / safe, (repeats, 1)) + out = o_scratch_ref[...] * inv + o_ref[...] = jnp.where(jnp.tile(has, (repeats, 1)), out, jnp.zeros_like(out)).astype(o_ref.dtype) + lse_ref[...] = jnp.where( + has, + m + log(safe), + jnp.asarray(mask_value, m.dtype), + ).astype(lse_ref.dtype) + + +def _make_partial( + *, + table_np, + active_np, + mask_mode="none", + exact_boundary_mask=None, + block_sizes, + orig_q_seq_len, + orig_kv_seq_len, + band_width, + frame_size, + include_first_frame, + bkv_compute_in=None, + use_base2_exp=True, + use_experimental_scheduler=False, + vmem_limit_bytes=None, +): + """Build a specialized partial for one static tile table. + + `exact_boundary_mask` remains as a compatibility alias for older callers. + """ + if exact_boundary_mask is not None: + mask_mode = "exact" if exact_boundary_mask else "none" + + bq = int(block_sizes.block_q) + bkv = int(block_sizes.block_kv) + bc = int(block_sizes.block_kv_compute) + bci = int(bkv_compute_in if bkv_compute_in is not None else block_sizes.block_kv_compute_in) + if bkv % bc or bc % bci: + raise ValueError(f"invalid blocks bkv={bkv}, bc={bc}, bci={bci}") + + kv_map = jnp.asarray(table_np) + active = jnp.asarray(active_np) + height, width = table_np.shape + + def partial(q, k, v): + heads, _, dq = q.shape + dv = v.shape[-1] + if heads != k.shape[0]: + raise NotImplementedError("static-range SVG kernel supports MHA only") + + def qmap(h, i, j, *refs): + del j, refs + return (h, i, 0) + + def kvmap(h, i, j, kvref, *refs): + del refs + return (h, kvref[i, j], 0) + + def outmap(h, i, j, *refs): + del j, refs + return (h, 0, i) + + outs = pl.pallas_call( + functools.partial( + _partial_kernel, + mask_mode=str(mask_mode), + mask_value=DEFAULT_MASK_VALUE, + grid_width=width, + bq=bq, + bkv=bkv, + bkv_compute=bc, + bkv_compute_in=bci, + head_dim_v=dv, + q_seq_len=int(orig_q_seq_len), + kv_seq_len=int(orig_kv_seq_len), + band_width=int(band_width), + frame_size=int(frame_size), + include_first_frame=bool(include_first_frame), + use_base2_exp=bool(use_base2_exp), + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=2, + in_specs=[ + pl.BlockSpec((None, bq, dq), qmap), + pl.BlockSpec((None, bkv, dq), kvmap), + pl.BlockSpec((None, bkv, dv), kvmap), + ], + out_specs=[ + pl.BlockSpec((None, dv, bq), outmap), + pl.BlockSpec((None, NUM_SUBLANES, bq), outmap), + ], + scratch_shapes=[ + pltpu.VMEM((NUM_SUBLANES, bq), jnp.float32), + pltpu.VMEM((NUM_SUBLANES, bq), jnp.float32), + pltpu.VMEM((dv, bq), jnp.float32), + ], + grid=(heads, height, width), + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": use_experimental_scheduler}, + disable_bounds_checks=True, + skip_device_barrier=True, + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=[ + jax.ShapeDtypeStruct((heads, dv, orig_q_seq_len), q.dtype), + jax.ShapeDtypeStruct((heads, NUM_SUBLANES, orig_q_seq_len), jnp.float32), + ], + )(kv_map, active, q, k, v) + return outs[-2], outs[-1][:, 0, :] + + partial.boundary_tiles = int(active_np.sum()) + partial.boundary_table_shape = tuple(table_np.shape) + return partial + + +def make_svg_exact_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 exact two-pass SVG reference implementation.""" + bq, bkv = int(block_sizes.block_q), int(block_sizes.block_kv) + fm, fa, bm, ba = _classify_tiles( + orig_q_seq_len, + orig_kv_seq_len, + bq, + bkv, + int(band_width), + int(frame_size), + bool(include_first_frame), + ) + common = { + "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, + } + full = _make_partial(table_np=fm, active_np=fa, mask_mode="none", **common) + boundary = _make_partial(table_np=bm, active_np=ba, mask_mode="exact", **common) + + def attention(q, k, v): + of, lf = full(q, k, v) + ob, lb = boundary(q, k, v) + m = jnp.maximum(lf, lb) + exp = jnp.exp2 if use_base2_exp else jnp.exp + wf, wb = exp(lf - m), exp(lb - m) + den = wf + wb + return (of.astype(jnp.float32) * (wf / den)[:, None, :] + ob.astype(jnp.float32) * (wb / den)[:, None, :]).astype( + q.dtype + ) + + attention.full_tiles = int(fa.sum()) + attention.boundary_tiles = int(ba.sum()) + attention.full_table_shape = tuple(fm.shape) + attention.boundary_table_shape = tuple(bm.shape) + attention.physical_tiles = attention.full_tiles + attention.boundary_tiles + return attention diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 7b2ba0df7..3db542586 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -35,12 +35,15 @@ from ..kernels import custom_splash_attention as custom_splash +from ..kernels import custom_svg_attention_dispatch +from ..kernels import custom_svg_static_range_attention from . import quantizations from .modeling_flax_utils import get_activation LOG2E = math.log2(math.e) Array = common_types.Array + Mesh = common_types.Mesh DType = common_types.DType BlockSizes = common_types.BlockSizes @@ -402,6 +405,7 @@ def _extract_custom_block_sizes(flash_block_sizes): # 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) return bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes @@ -596,6 +600,8 @@ def _tpu_flash_attention( use_experimental_scheduler: bool = False, is_causal: bool = False, preserve_asymmetric_block_sizes: bool = False, + spatiotemporal_config: Optional[dict] = None, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, ) -> jax.Array: """TPU Flash Attention""" @@ -859,6 +865,8 @@ def _ulysses_attention( use_fixed_m: bool = False, ulysses_attention_chunks: int = 1, preserve_asymmetric_block_sizes: bool = False, + spatiotemporal_config: Optional[dict] = None, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -1676,6 +1684,8 @@ def ulysses_custom_kernel(q, k, v, context): use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), ulysses_attention_chunks=context["ulysses_attention_chunks"], + spatiotemporal_config=context.get("spatiotemporal_config"), + spatiotemporal_shape=context.get("spatiotemporal_shape"), ) @@ -1863,6 +1873,8 @@ def tokamax_flash_kernel(q, k, v, context): use_experimental_scheduler=context["use_experimental_scheduler"], is_causal=context.get("is_causal", False), preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), + spatiotemporal_config=context.get("spatiotemporal_config"), + spatiotemporal_shape=context.get("spatiotemporal_shape"), ) @@ -1942,6 +1954,8 @@ def _apply_attention( ulysses_attention_chunks: int = 1, is_causal: bool = False, preserve_asymmetric_block_sizes: bool = False, + spatiotemporal_config: Optional[dict] = None, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, ): """Routes to different attention kernels using a module-level registry.""" @@ -2009,8 +2023,21 @@ def _apply_attention( "dpa_layer": dpa_layer, "is_causal": is_causal, "preserve_asymmetric_block_sizes": preserve_asymmetric_block_sizes, + "spatiotemporal_config": spatiotemporal_config, + "spatiotemporal_shape": spatiotemporal_shape, } + if spatiotemporal_config and spatiotemporal_config.get("use_svg_attention"): + if effective_attention_kernel not in ( + "ulysses_custom", + "ulysses_custom_fixed_m", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + ): + raise ValueError("Head-local SVG requires a custom Ulysses attention backend.") + # Dense uses its configured ring split; SVG exchanges over the full context axis. + return _head_local_svg_attention(query, key, value, context) + # Module-level Registry lookup if effective_attention_kernel in KERNEL_REGISTRY: with jax.named_scope(f"kernel_{effective_attention_kernel}"): @@ -2019,6 +2046,109 @@ def _apply_attention( raise ValueError(f"Unexpected attention kernel {effective_attention_kernel=}.") +def _head_local_svg_attention(query, key, value, context): + from .wan.transformers import svg_attention, svg_head_local + + cfg = context["spatiotemporal_config"] + grid = context["spatiotemporal_shape"] + mesh = context["mesh"] + cp = mesh.shape[CONTEXT] + if context["attention_mask"] is not None or cfg.get("global_stride", 0): + raise ValueError("Head-local SVG does not support external or periodic masks.") + if context["ulysses_attention_chunks"] != 1: + raise ValueError("Head-local SVG does not implement chunked Ulysses attention.") + q, k, v = (_unflatten_heads(x, context["heads"]) if x.ndim == 3 else x for x in (query, key, value)) + if grid is None or q.shape != k.shape or q.shape != v.shape or q.shape[2] != math.prod(grid): + raise ValueError("Head-local SVG requires matched self-attention QKV and token grid.") + batch, heads = q.shape[0], q.shape[1] + # Fold unsharded batches into heads to match the dense Ulysses layout. + devices_in_batch_sharding = mesh.shape["data"] * (mesh.shape["fsdp"] if "fsdp" in mesh.shape else 1) + fold_batch = batch > 1 and devices_in_batch_sharding == 1 and (batch * heads) % cp == 0 + local_heads = batch * heads if fold_batch else heads + if local_heads % cp or q.shape[2] % cp: + raise ValueError("SVG heads and sequence length must divide evenly across Ulysses shards.") + q, k, v = (svg_head_local.inference_only(x) for x in (q, k, v)) + qspec = nn.logical_to_mesh_axes(context["axis_names_q"]) + kvspec = nn.logical_to_mesh_axes(context["axis_names_kv"]) + if qspec != kvspec or qspec[1:] != (None, CONTEXT, None): + raise ValueError("Head-local SVG requires sequence sharding and unsharded heads.") + profile_key = jax.random.PRNGKey(int(cfg["profile_seed"])) + for index in (cfg.get("svg_layer_index"), cfg.get("svg_step_index")): + if index is not None: + profile_key = jax.random.fold_in(profile_key, jnp.asarray(index, jnp.uint32)) + if cfg.get("svg_step_index") is None and cfg.get("svg_timestep") is not None: + profile_key = jax.random.fold_in(profile_key, jnp.max(jnp.asarray(cfg["svg_timestep"])).astype(jnp.uint32)) + with jax.named_scope("svg_routing"): + route = svg_attention.svg_profile_temporal_heads( + q, + k, + v, + grid, + int(cfg["profile_query_count"]), + profile_key, + context["scale"], + sample_max_row=int(cfg.get("sample_max_row", 10000)), + ) + if fold_batch: + q, k, v = (x.reshape(1, local_heads, *x.shape[2:]) for x in (q, k, v)) + route = route.reshape(1, local_heads) + # Sparse and dense attention use independently configured tiles. + bq, bkv, bc, bci, hpt, vmem = _extract_custom_block_sizes( + cfg.get("custom_flash_block_sizes") or context["flash_block_sizes"] + ) + if hpt != 1: + raise ValueError("SVG requires heads_per_tile=1.") + blocks = custom_svg_static_range_attention.SVGBlockSizes( + block_q=bq, + block_kv=bkv, + block_kv_compute=bc, + block_kv_compute_in=bci, + ) + + def core(q, k, v): + k = k * context["scale"] + if context["use_base2_exp"]: + q = q * LOG2E + q, dim, n = _pad_data_for_flash(q, q.shape[1], bq) + k, _, nk = _pad_data_for_flash(k, k.shape[1], bkv) + v, _, _ = _pad_data_for_flash(v, v.shape[1], bkv) + kernel = custom_svg_attention_dispatch.make_svg_static_range_mha( + block_sizes=blocks, + orig_q_seq_len=n, + orig_kv_seq_len=nk, + band_width=int(cfg["band_width"]), + frame_size=int(grid[1] * grid[2]), + include_first_frame=bool(cfg.get("include_first_frame", True)), + bkv_compute_in=bci, + use_base2_exp=context["use_base2_exp"], + use_experimental_scheduler=context["use_experimental_scheduler"], + vmem_limit_bytes=vmem, + ) + # Report executed tile fraction, which differs from real attention-pair density. + executed = kernel.union_main_tiles + kernel.tail_cleanup_tiles + total = -(-n // bq) * -(-nk // bkv) + with jax.named_scope(f"svg_kernel_c_tiles{executed}of{total}_d{executed / total:.3f}"): + out = jax.vmap(kernel)(q, k, v) + return jnp.swapaxes(out, 2, 3)[:, :, :n, :dim].astype(q.dtype) + + out = svg_head_local.exchange_local( + q, + k, + v, + route, + mesh=mesh, + qspec=qspec, + kvspec=kvspec, + ulysses_axis=CONTEXT, + place=lambda q, k, v, r: svg_attention.svg_placement_permute(q, k, v, r, grid), + restore=lambda o, r: svg_attention.svg_placement_unpermute(o, r, grid), + core=core, + ) + if fold_batch: + out = out.reshape(batch, heads, *out.shape[2:]) + return _reshape_heads_to_head_dim(out) + + def _query_chunk_attention(query, key, value, precision, key_chunk_size: int = 4096): """Multi-head dot product attention with a limited number of queries.""" num_kv, num_heads, k_features = key.shape[-3:] @@ -2296,6 +2426,8 @@ def apply_attention( value: Array, attention_mask: Array = None, preserve_asymmetric_block_sizes: bool = False, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, + sparse_config_override: Optional[dict] = None, ): return _apply_attention( query=query, @@ -2323,6 +2455,8 @@ def apply_attention( ulysses_shards=(self.ulysses_shards if hasattr(self, "ulysses_shards") else -1), ulysses_attention_chunks=(self.ulysses_attention_chunks if hasattr(self, "ulysses_attention_chunks") else 1), preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes, + spatiotemporal_config=sparse_config_override, + spatiotemporal_shape=spatiotemporal_shape, ) @@ -2378,6 +2512,8 @@ def apply_attention( value: Array, attention_mask: Array = None, preserve_asymmetric_block_sizes: bool = False, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, + sparse_config_override: Optional[dict] = None, ): return _apply_attention( query=query, @@ -2404,6 +2540,8 @@ def apply_attention( ulysses_attention_chunks=self.ulysses_attention_chunks, is_causal=self.is_causal, preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes, + spatiotemporal_config=sparse_config_override, + spatiotemporal_shape=spatiotemporal_shape, ) @@ -2447,9 +2585,51 @@ def __init__( "use_experimental_scheduler": False, "ulysses_shards": -1, "ulysses_attention_chunks": 1, + "use_svg_attention": False, + "svg_implementation": "official_svg", + "svg_spatial_density": 0.25, + "svg_sample_max_row": 10000, + "svg_profile_query_count": 64, + "svg_profile_seed": 0, + "svg_dense_layer_fraction": 0.0, + "svg_dense_timestep_fraction": 0.0, + "svg_active_start_step": -1, + "svg_active_end_step": -1, + "svg_active_start_layer": -1, + "svg_active_end_layer": -1, + "svg_num_train_timesteps": 1000, + "svg_num_layers": 40, + "svg_include_first_frame": True, + "svg_global_stride": 0, + "svg_global_offset": 0, + "svg_high_noise_density": -1.0, + "svg_low_noise_density": -1.0, + "svg_flash_block_sizes": None, **(attention_config or {}), } + self.use_svg_attention = attention_config["use_svg_attention"] + self.svg_implementation = attention_config["svg_implementation"] + self.svg_spatial_density = attention_config["svg_spatial_density"] + self.svg_sample_max_row = attention_config["svg_sample_max_row"] + self.svg_profile_query_count = attention_config["svg_profile_query_count"] + self.svg_profile_seed = attention_config["svg_profile_seed"] + self.svg_dense_layer_fraction = attention_config["svg_dense_layer_fraction"] + self.svg_dense_timestep_fraction = attention_config["svg_dense_timestep_fraction"] + self.svg_active_start_step = attention_config["svg_active_start_step"] + self.svg_active_end_step = attention_config["svg_active_end_step"] + self.svg_active_start_layer = attention_config["svg_active_start_layer"] + self.svg_active_end_layer = attention_config["svg_active_end_layer"] + self.svg_num_train_timesteps = attention_config["svg_num_train_timesteps"] + self.svg_num_layers = attention_config["svg_num_layers"] + self.svg_include_first_frame = attention_config["svg_include_first_frame"] + self.svg_global_stride = attention_config["svg_global_stride"] + self.svg_global_offset = attention_config["svg_global_offset"] + self.svg_high_noise_density = attention_config["svg_high_noise_density"] + self.svg_low_noise_density = attention_config["svg_low_noise_density"] + self.svg_flash_block_sizes = attention_config["svg_flash_block_sizes"] + self.is_self_attention = is_self_attention + if attention_kernel in {"flash", "cudnn_flash_te"} and mesh is None: raise ValueError(f"The flash attention kernel requires a value for mesh, but mesh is {self.mesh}") self.dim_head = dim_head @@ -2693,12 +2873,25 @@ def __call__( deterministic: bool = True, rngs: nnx.Rngs = None, cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, + svg_layer_index: Optional[int | jax.Array] = None, + svg_timestep: Optional[int | float | jax.Array] = None, + svg_step_index: Optional[int | jax.Array] = None, ) -> jax.Array: hidden_states = nn.with_logical_constraint(hidden_states, (BATCH, LENGTH, HEAD)) if encoder_hidden_states is not None: encoder_hidden_states = nn.with_logical_constraint(encoder_hidden_states, (BATCH, LENGTH, HEAD)) dtype = hidden_states.dtype - is_self_attention = encoder_hidden_states is None + is_self_attention = getattr( + self, + "is_self_attention", + encoder_hidden_states is None or encoder_hidden_states is hidden_states, + ) + if self.use_svg_attention and is_self_attention: + if not deterministic: + raise ValueError("SVG attention supports deterministic inference only.") + if spatiotemporal_shape is None: + raise ValueError("SVG attention requires spatiotemporal_shape.") if encoder_hidden_states is None: encoder_hidden_states = hidden_states @@ -2741,14 +2934,74 @@ def __call__( key_proj = checkpoint_name(key_proj, "key_proj") value_proj = checkpoint_name(value_proj, "value_proj") - with jax.named_scope("apply_attention"): - attn_output = self.attention_op.apply_attention( - query_proj, - key_proj, - value_proj, - attention_mask=encoder_attention_mask, + if self.use_svg_attention and is_self_attention and spatiotemporal_shape is not None: + from .wan.transformers import svg_attention + + is_active = svg_attention.is_svg_active( + step_index=svg_step_index, + layer_index=svg_layer_index, + timestep=svg_timestep, + start_step=self.svg_active_start_step, + end_step=self.svg_active_end_step, + start_layer=self.svg_active_start_layer, + end_layer=self.svg_active_end_layer, + dense_layer_fraction=self.svg_dense_layer_fraction, + dense_timestep_fraction=self.svg_dense_timestep_fraction, + num_train_timesteps=self.svg_num_train_timesteps, + num_layers=self.svg_num_layers, ) + def run_dense(_): + return self.attention_op.apply_attention( + query_proj, + key_proj, + value_proj, + attention_mask=encoder_attention_mask, + ) + + def run_sparse_svg(_): + execution_band_width = svg_attention.svg_execution_band_width( + spatiotemporal_shape, + self.svg_spatial_density, + ) + sparse_config = { + "use_svg_attention": True, + "mask_type": "svg_spatial", + "band_width": execution_band_width, + "include_first_frame": self.svg_include_first_frame, + "global_stride": self.svg_global_stride, + "global_offset": self.svg_global_offset, + "profile_query_count": self.svg_profile_query_count, + "profile_seed": self.svg_profile_seed, + "sample_max_row": self.svg_sample_max_row, + "custom_flash_block_sizes": self.svg_flash_block_sizes, + "svg_step_index": svg_step_index, + "svg_layer_index": svg_layer_index, + "svg_timestep": svg_timestep, + } + return self.attention_op.apply_attention( + query_proj, + key_proj, + value_proj, + attention_mask=encoder_attention_mask, + spatiotemporal_shape=spatiotemporal_shape, + sparse_config_override=sparse_config, + ) + + with jax.named_scope("apply_attention"): + if isinstance(is_active, bool): + attn_output = run_sparse_svg(None) if is_active else run_dense(None) + else: + attn_output = jax.lax.cond(is_active, run_sparse_svg, run_dense, operand=None) + else: + with jax.named_scope("apply_attention"): + attn_output = self.attention_op.apply_attention( + query_proj, + key_proj, + value_proj, + attention_mask=encoder_attention_mask, + ) + else: # NEW PATH for I2V CROSS-ATTENTION with self.conditional_named_scope("proj_query"): diff --git a/src/maxdiffusion/models/wan/transformers/svg_attention.py b/src/maxdiffusion/models/wan/transformers/svg_attention.py new file mode 100644 index 000000000..7ffd1e11a --- /dev/null +++ b/src/maxdiffusion/models/wan/transformers/svg_attention.py @@ -0,0 +1,355 @@ +""" +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. + +Sparse VideoGen (SVG) per-head routing, layout placement, and scheduling helpers. +""" + +from __future__ import annotations + +import math +from numbers import Integral +from typing import Tuple + +import jax +import jax.numpy as jnp + + +def svg_execution_band_width(token_grid: Tuple[int, int, int], density: float) -> int: + """Compute SVG symmetric band width from retained density with 128-token block ceiling. + + For sequence length N and retained pair density s: + w = N * (1.0 - sqrt(1.0 - s)) + rounded up to the nearest multiple of 128. + + Edge cases: + density == 1.0 -> N - 1 (dense band covering all token pairs) + density <= 0.0 or density > 1.0 -> raises ValueError + """ + sequence_length = math.prod(token_grid) + if sequence_length <= 0: + raise ValueError(f"SVG sequence length must be positive, got {sequence_length}") + density = float(density) + if density <= 0.0 or density > 1.0: + raise ValueError(f"svg_retained_density must be in (0.0, 1.0], got {density}") + if density >= 1.0: + return sequence_length - 1 + + width = sequence_length * (1.0 - math.sqrt(1.0 - density)) + band_width = int(math.ceil(width / 128.0)) * 128 + return min(max(band_width, 0), sequence_length - 1) + + +def svg_probe_masks( + sampled_rows: jax.Array, + token_grid: Tuple[int, int, int], + block_size: int = 128, +) -> Tuple[jax.Array, jax.Array]: + """Construct fixed spatial and temporal profiler masks for sampled query rows. + + Spatial mask: + (k < frame_size) | (abs(q // 128 - k // 128) < (2 * frame_size) // 128) + Temporal mask: + (k_tm < frame_size) | (abs(q_tm // 128 - k_tm // 128) < (2 * frame_size) // 128) + where q_tm = (q % frame_size) * frames + (q // frame_size) + k_tm = (k % frame_size) * frames + (k // frame_size) + """ + frames, height, width = token_grid + frame_size = height * width + sequence_length = frames * frame_size + + key_indices = jnp.arange(sequence_length, dtype=jnp.int32) + query_indices = sampled_rows.astype(jnp.int32) + + block_thres_blocks = (2 * frame_size) // block_size + + q_blk = query_indices[:, None] // block_size + k_blk = key_indices[None, :] // block_size + + q_tm = (query_indices[:, None] % frame_size) * frames + (query_indices[:, None] // frame_size) + k_tm = (key_indices[None, :] % frame_size) * frames + (key_indices[None, :] // frame_size) + + q_tm_blk = q_tm // block_size + k_tm_blk = k_tm // block_size + + spatial_mask = (key_indices[None, :] < frame_size) | (jnp.abs(q_blk - k_blk) < block_thres_blocks) + temporal_mask = (k_tm < frame_size) | (jnp.abs(q_tm_blk - k_tm_blk) < block_thres_blocks) + + return spatial_mask, temporal_mask + + +def svg_profile_temporal_heads( + query: jax.Array, + key: jax.Array, + value: jax.Array, + token_grid: Tuple[int, int, int], + query_count: int, + profile_key: jax.Array, + scale: float, + sample_max_row: int = 10000, +) -> jax.Array: + """Profile each head independently using sampled dense reconstruction MSE. + + `scale` is the ordinary attention logit scale. The profiler intentionally + uses natural-exp softmax. The production base-2 kernel multiplies Q by + log2(e) before exp2, which is mathematically equivalent to this natural-exp + formulation, so no additional LOG2E factor belongs in the profiler. + """ + with jax.named_scope("svg_route_profile"): + sequence_length = query.shape[2] + sample_pool_size = min(max(int(sample_max_row), 1), sequence_length) + sample_count = min(max(int(query_count), 1), sample_pool_size) + + if sample_count >= sample_pool_size: + sampled_rows = jnp.arange(sample_count, dtype=jnp.int32) + else: + sampled_rows = jax.random.randint( + profile_key, + (sample_count,), + minval=0, + maxval=sample_pool_size, + dtype=jnp.int32, + ) + + sampled_query = jnp.take(query, sampled_rows, axis=2) + spatial_mask, temporal_mask = svg_probe_masks(sampled_rows, token_grid) + + sampled_qk = ( + jnp.einsum( + "bhqd,bhkd->bhqk", + sampled_query.astype(jnp.float32), + key.astype(jnp.float32), + ) + * scale + ) + + dense_weights = jax.nn.softmax(sampled_qk, axis=-1) + dense_output = jnp.einsum("bhqk,bhkd->bhqd", dense_weights, value.astype(jnp.float32)) + + spatial_logits = jnp.where(spatial_mask[None, None, :, :], sampled_qk, -1e9) + spatial_weights = jax.nn.softmax(spatial_logits, axis=-1) + spatial_output = jnp.einsum("bhqk,bhkd->bhqd", spatial_weights, value.astype(jnp.float32)) + + temporal_logits = jnp.where(temporal_mask[None, None, :, :], sampled_qk, -1e9) + temporal_weights = jax.nn.softmax(temporal_logits, axis=-1) + temporal_output = jnp.einsum("bhqk,bhkd->bhqd", temporal_weights, value.astype(jnp.float32)) + + spatial_error = jnp.mean(jnp.square(spatial_output - dense_output), axis=(-2, -1)) + temporal_error = jnp.mean(jnp.square(temporal_output - dense_output), axis=(-2, -1)) + return temporal_error < spatial_error + + +def svg_token_major_indices(token_grid: Tuple[int, int, int]) -> Tuple[jax.Array, jax.Array]: + """Return forward (frame-major -> token-major) and inverse permutation index arrays.""" + frames, height, width = token_grid + frame_size = height * width + sequence_length = frames * frame_size + + j = jnp.arange(sequence_length, dtype=jnp.int32) + forward_gather_idx = (j % frames) * frame_size + (j // frames) + inverse_gather_idx = (j % frame_size) * frames + (j // frame_size) + return forward_gather_idx, inverse_gather_idx + + +def svg_placement_permute( + query: jax.Array, + key: jax.Array, + value: jax.Array, + is_temporal: jax.Array, + token_grid: Tuple[int, int, int], +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Place temporal heads into token-major order, leaving spatial heads unchanged.""" + with jax.named_scope("svg_layout_place"): + frames, height, width = (int(v) for v in token_grid) + frame_size = height * width + b, heads, n, d = query.shape + + q_tm = query.reshape(b, heads, frames, frame_size, d).transpose(0, 1, 3, 2, 4).reshape(b, heads, n, d) + k_tm = key.reshape(b, heads, frames, frame_size, d).transpose(0, 1, 3, 2, 4).reshape(b, heads, n, d) + v_tm = value.reshape(b, heads, frames, frame_size, d).transpose(0, 1, 3, 2, 4).reshape(b, heads, n, d) + + cond = is_temporal[:, :, None, None] + q_out = jnp.where(cond, q_tm, query) + k_out = jnp.where(cond, k_tm, key) + v_out = jnp.where(cond, v_tm, value) + return q_out, k_out, v_out + + +def svg_placement_unpermute( + output: jax.Array, + is_temporal: jax.Array, + token_grid: Tuple[int, int, int], +) -> jax.Array: + """Restore temporal-head output from token-major to frame-major order.""" + with jax.named_scope("svg_layout_restore"): + frames, height, width = (int(v) for v in token_grid) + frame_size = height * width + b, heads, n, d = output.shape + out_fm = output.reshape(b, heads, frame_size, frames, d).transpose(0, 1, 3, 2, 4).reshape(b, heads, n, d) + return jnp.where(is_temporal[:, :, None, None], out_fm, output) + + +def is_svg_active( + step_index: int | jax.Array | None = None, + layer_index: int | jax.Array | None = None, + timestep: int | float | jax.Array | None = None, + start_step: int = -1, + end_step: int = -1, + start_layer: int = -1, + end_layer: int = -1, + dense_layer_fraction: float = 0.0, + dense_timestep_fraction: float = 0.0, + num_train_timesteps: int = 1000, + num_layers: int = 40, +) -> bool | jax.Array: + """Evaluate whether SVG attention is active for the current step and layer.""" + has_start_step = start_step >= 0 + has_end_step = end_step >= 0 + if has_start_step != has_end_step: + raise ValueError( + f"Incomplete explicit SVG step schedule: start_step={start_step}, end_step={end_step}. " + "Both bounds must be non-negative or both unset (< 0)." + ) + has_explicit_step = has_start_step and has_end_step + + has_start_layer = start_layer >= 0 + has_end_layer = end_layer >= 0 + if has_start_layer != has_end_layer: + raise ValueError( + f"Incomplete explicit SVG layer schedule: start_layer={start_layer}, end_layer={end_layer}. " + "Both bounds must be non-negative or both unset (< 0)." + ) + has_explicit_layer = has_start_layer and has_end_layer + + if has_explicit_step and step_index is None: + raise ValueError("Explicit SVG step schedule requires step_index.") + if has_explicit_layer and layer_index is None: + raise ValueError("Explicit SVG layer schedule requires layer_index.") + + # 1. Evaluate layer interval / fraction + if has_explicit_layer: + if isinstance(layer_index, Integral): + layer_active: bool | jax.Array = bool(start_layer <= layer_index < end_layer) + else: + layer_arr = jnp.asarray(layer_index) + layer_active = jnp.logical_and(layer_arr >= start_layer, layer_arr < end_layer) + else: + dense_layer_count = math.ceil(dense_layer_fraction * num_layers) + if dense_layer_count > 0 and layer_index is not None: + if isinstance(layer_index, Integral): + layer_active = bool(layer_index >= dense_layer_count) + else: + layer_active = jnp.asarray(layer_index) >= dense_layer_count + else: + layer_active = True + + if isinstance(layer_active, bool) and not layer_active: + return False + + # 2. Evaluate step interval / fraction + if has_explicit_step: + if isinstance(step_index, Integral): + step_active: bool | jax.Array = bool(start_step <= step_index < end_step) + else: + step_arr = jnp.asarray(step_index) + step_active = jnp.logical_and(step_arr >= start_step, step_arr < end_step) + else: + if dense_timestep_fraction > 0.0 and timestep is not None: + first_sparse_timestep = (1.0 - dense_timestep_fraction) * num_train_timesteps + if isinstance(timestep, (Integral, float)): + step_active = bool(timestep < first_sparse_timestep) + else: + step_active = jnp.max(jnp.asarray(timestep)) < first_sparse_timestep + else: + step_active = True + + if isinstance(step_active, bool) and not step_active: + return False + + if isinstance(layer_active, bool) and isinstance(step_active, bool): + return bool(layer_active and step_active) + if isinstance(layer_active, bool): + return step_active + if isinstance(step_active, bool): + return layer_active + return jnp.logical_and(step_active, layer_active) + + +def place_sequence_for_mask( + tensor: jax.Array, + is_temporal: jax.Array, + token_grid: Tuple[int, int, int], +) -> jax.Array: + """Place a single tensor (query, key, or value) into per-head layout.""" + frames, height, width = (int(v) for v in token_grid) + frame_size = height * width + b, heads, n, d = tensor.shape + t_tm = tensor.reshape(b, heads, frames, frame_size, d).transpose(0, 1, 3, 2, 4).reshape(b, heads, n, d) + return jnp.where(is_temporal[:, :, None, None], t_tm, tensor) + + +def unplace_sequence_for_mask( + tensor: jax.Array, + is_temporal: jax.Array, + token_grid: Tuple[int, int, int], +) -> jax.Array: + """Unplace a single tensor from per-head layout back to frame-major.""" + return svg_placement_unpermute(tensor, is_temporal, token_grid) + + +class SVGHeadRouting: + + def __init__(self, is_temporal: jax.Array): + self.is_temporal = is_temporal + + +def route_heads( + query: jax.Array, + key: jax.Array, + value: jax.Array, + token_grid: Tuple[int, int, int], + spatial_density: float = 0.25, + temporal_density: float = 0.25, + sample_query_count: int = 64, + sample_max_row: int = 10000, + profile_seed: int = 0, + scale: float = 1.0, + include_first_frame: bool = True, + global_stride: int = 0, + global_offset: int = 0, + layer_index: int | jax.Array | None = None, + step_index: int | jax.Array | None = None, + timestep: int | float | jax.Array | None = None, +) -> SVGHeadRouting: + """Profile and route heads between spatial and temporal attention layouts.""" + del spatial_density, temporal_density, include_first_frame, global_stride, global_offset + profile_key = jax.random.PRNGKey(profile_seed) + if layer_index is not None: + profile_key = jax.random.fold_in(profile_key, jnp.asarray(layer_index, dtype=jnp.uint32)) + if step_index is not None: + profile_key = jax.random.fold_in(profile_key, jnp.asarray(step_index, dtype=jnp.uint32)) + elif timestep is not None: + profile_key = jax.random.fold_in(profile_key, jnp.max(jnp.asarray(timestep)).astype(jnp.uint32)) + + is_temporal = svg_profile_temporal_heads( + query, + key, + value, + token_grid, + sample_query_count, + profile_key, + scale, + sample_max_row=sample_max_row, + ) + return SVGHeadRouting(is_temporal=is_temporal) diff --git a/src/maxdiffusion/models/wan/transformers/svg_head_local.py b/src/maxdiffusion/models/wan/transformers/svg_head_local.py new file mode 100644 index 000000000..9da987fd7 --- /dev/null +++ b/src/maxdiffusion/models/wan/transformers/svg_head_local.py @@ -0,0 +1,53 @@ +""" +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. + +Head-local placement for pure-Ulysses SVG inference. +""" + +from functools import partial + +import jax +from jax.sharding import PartitionSpec as P + + +@jax.custom_jvp +def inference_only(x): + return x + + +@inference_only.defjvp +def _reject_derivative(primals, tangents): + raise NotImplementedError("SVG attention supports inference only.") + + +def exchange_local(q, k, v, route, *, mesh, qspec, kvspec, ulysses_axis, place, restore, core): + """Route sharding follows the exact head blocks selected by all_to_all.""" + if qspec != kvspec or qspec[1] is not None: + raise ValueError("Require identical sequence-sharded QKV with unsharded heads") + route_spec = P(qspec[0], ulysses_axis) + + @partial(jax.shard_map, mesh=mesh, in_specs=(qspec, kvspec, kvspec, route_spec), out_specs=qspec, check_vma=False) + def local(q, k, v, route): + a2a = partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True) + q, k, v = (a2a(x, split_axis=1, concat_axis=2) for x in (q, k, v)) + with jax.named_scope("head_local_placement"): + q, k, v = place(q, k, v, route) + out = core(q, k, v) + with jax.named_scope("head_local_restore"): + out = restore(out, route) + out = a2a(out, split_axis=2, concat_axis=1) + return out + + return local(q, k, v, route) diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index 4cdfd0ca1..f147e4414 100644 --- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py +++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py @@ -448,6 +448,10 @@ def __call__( rngs: nnx.Rngs = None, encoder_attention_mask: Optional[jax.Array] = None, cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, + spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, + svg_layer_index: Optional[int | jax.Array] = None, + svg_timestep: Optional[int | float | jax.Array] = None, + svg_step_index: Optional[int | jax.Array] = None, ): with self.conditional_named_scope("transformer_block"): # Support both global [B, 6, dim] and per-token [B, seq_len, 6, dim] temb. @@ -494,6 +498,10 @@ def __call__( rotary_emb=rotary_emb, deterministic=deterministic, rngs=rngs, + spatiotemporal_shape=spatiotemporal_shape, + svg_layer_index=svg_layer_index, + svg_timestep=svg_timestep, + svg_step_index=svg_step_index, ) with self.conditional_named_scope("self_attn_residual"): hidden_states = (hidden_states.astype(jnp.float32) + attn_output * gate_msa).astype(hidden_states.dtype) @@ -773,6 +781,7 @@ def __call__( kv_cache: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None, rotary_emb: Optional[jax.Array] = None, encoder_attention_mask: Optional[jax.Array] = None, + svg_step_index: Optional[int | jax.Array] = None, ) -> Union[jax.Array, Tuple[jax.Array, jax.Array], Dict[str, jax.Array]]: hidden_states = nn.with_logical_constraint(hidden_states, ("batch", None, None, None, None)) batch_size, _, num_frames, height, width = hidden_states.shape @@ -848,9 +857,9 @@ def _run_all_blocks(h): def scan_fn(carry, block_input): hidden_states_carry, rngs_carry = carry if kv_cache is not None: - block, layer_kv_cache = block_input + block, layer_kv_cache, layer_index = block_input else: - block = block_input + block, layer_index = block_input layer_kv_cache = None hidden_states = block( @@ -862,6 +871,10 @@ def scan_fn(carry, block_input): rngs_carry, encoder_attention_mask, cached_kv=layer_kv_cache, + spatiotemporal_shape=(post_patch_num_frames, post_patch_height, post_patch_width), + svg_layer_index=layer_index, + svg_timestep=timestep, + svg_step_index=svg_step_index, ) new_carry = (hidden_states, rngs_carry) return new_carry, None @@ -874,10 +887,11 @@ def scan_fn(carry, block_input): ) initial_carry = (h, rngs) + layer_indices = jnp.arange(self.num_layers, dtype=jnp.int32) if kv_cache is not None: - scan_input = (self.blocks, kv_cache) + scan_input = (self.blocks, kv_cache, layer_indices) else: - scan_input = self.blocks + scan_input = (self.blocks, layer_indices) final_carry, _ = nnx.scan( rematted_block_forward, @@ -904,6 +918,12 @@ def layer_forward(hidden_states, l_kv): rngs, encoder_attention_mask=encoder_attention_mask, cached_kv=l_kv, + spatiotemporal_shape=(post_patch_num_frames, post_patch_height, post_patch_width), + # Python int: lets is_svg_active resolve statically inactive + # layers at trace time instead of emitting a lax.cond. + svg_layer_index=i, + svg_timestep=timestep, + svg_step_index=svg_step_index, ) rematted_layer_forward = self.gradient_checkpoint.apply( diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index 6464ce6a6..cf866e5be 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -13,7 +13,7 @@ # limitations under the License. from abc import abstractmethod -from typing import List, Union, Optional, Tuple +from typing import Any, List, Union, Optional, Tuple from functools import partial from maxdiffusion.image_processor import PipelineImageInput import numpy as np @@ -56,7 +56,6 @@ FlaxCLIPVisionModel = None import PIL - TORCH_DTYPE_MAP = { "bfloat16": torch.bfloat16, "float16": torch.float16, @@ -342,11 +341,45 @@ def create_model(rngs: nnx.Rngs, wan_config: dict): wan_config["mask_padding_tokens"] = config.mask_padding_tokens wan_config["scan_layers"] = config.scan_layers wan_config["enable_jax_named_scopes"] = config.enable_jax_named_scopes + high_density = getattr(config, "svg_high_noise_density", None) + low_density = getattr(config, "svg_low_noise_density", None) + high_density = -1.0 if high_density is None else float(high_density) + low_density = -1.0 if low_density is None else float(low_density) + + if subfolder == "transformer" and high_density >= 0: + expert_density = high_density + elif subfolder == "transformer_2" and low_density >= 0: + expert_density = low_density + else: + expert_density = float(getattr(config, "svg_spatial_density", 0.25)) + + use_svg_for_expert = bool(getattr(config, "use_svg_attention", False)) and (expert_density < 1.0) + wan_config["attention_config"] = { "use_base2_exp": config.use_base2_exp, "use_experimental_scheduler": config.use_experimental_scheduler, "ulysses_shards": getattr(config, "ulysses_shards", -1), "ulysses_attention_chunks": getattr(config, "ulysses_attention_chunks", 1), + "use_svg_attention": use_svg_for_expert, + "svg_implementation": getattr(config, "svg_implementation", "official_svg"), + "svg_spatial_density": expert_density, + "svg_sample_max_row": getattr(config, "svg_sample_max_row", 10000), + "svg_profile_query_count": getattr(config, "svg_profile_query_count", 64), + "svg_profile_seed": getattr(config, "svg_profile_seed", 0), + "svg_dense_layer_fraction": getattr(config, "svg_dense_layer_fraction", 0.0), + "svg_dense_timestep_fraction": getattr(config, "svg_dense_timestep_fraction", 0.0), + "svg_active_start_step": getattr(config, "svg_active_start_step", -1), + "svg_active_end_step": getattr(config, "svg_active_end_step", -1), + "svg_active_start_layer": getattr(config, "svg_active_start_layer", -1), + "svg_active_end_layer": getattr(config, "svg_active_end_layer", -1), + "svg_num_train_timesteps": getattr(config, "svg_num_train_timesteps", 1000), + "svg_num_layers": getattr(config, "svg_num_layers", wan_config.get("num_layers", 40)), + "svg_include_first_frame": getattr(config, "svg_include_first_frame", True), + "svg_global_stride": getattr(config, "svg_global_stride", 0), + "svg_global_offset": getattr(config, "svg_global_offset", 0), + "svg_high_noise_density": high_density, + "svg_low_noise_density": low_density, + "svg_flash_block_sizes": getattr(config, "svg_flash_block_sizes", None) or None, } # 2. eval_shape - will not use flops or create weights on device @@ -1350,12 +1383,87 @@ def _prepare_model_inputs( num_frames, ) + def _validate_svg_cache_compatibility( + self, + use_cfg_cache: bool = False, + use_magcache: bool = False, + ) -> None: + """Validates that SVG sparse attention is not combined with incompatible caches.""" + validate_svg_cache_compatibility( + self, + use_cfg_cache=use_cfg_cache, + use_magcache=use_magcache, + ) + @abstractmethod def __call__(self, **kwargs): """Runs the inference pipeline.""" pass +def _has_svg_enabled(obj: Any) -> bool: + """Returns True if obj (transformer, GraphDef, or config) has effective SVG enabled.""" + if obj is None: + return False + if bool(getattr(obj, "use_svg_attention", False)): + return True + cfg = getattr(obj, "config", None) + if cfg is not None: + attn_cfg = ( + getattr(cfg, "attention_config", None) + or (cfg.get("attention_config") if isinstance(cfg, dict) else None) + ) + if isinstance(attn_cfg, dict) and bool(attn_cfg.get("use_svg_attention", False)): + return True + if hasattr(obj, "attributes") and hasattr(obj, "nodes"): + for k, v in getattr(obj, "attributes", ()): + if k == "use_svg_attention" and getattr(v, "value", False) is True: + return True + if k == "config": + val = getattr(v, "value", None) + ac = ( + getattr(val, "attention_config", None) + or (val.get("attention_config") if isinstance(val, dict) else None) + ) + if isinstance(ac, dict) and bool(ac.get("use_svg_attention", False)): + return True + return False + + +def validate_svg_cache_compatibility( + target: Any, + use_cfg_cache: bool = False, + use_magcache: bool = False, + *extra_targets: Any, +) -> None: + """Raises ValueError if SVG sparse attention is active and CFG cache or MagCache is enabled.""" + if not (use_cfg_cache or use_magcache): + return + svg_active = _has_svg_enabled(target) + transformers_found = False + for attr in ("transformer", "high_noise_transformer", "low_noise_transformer"): + t = getattr(target, attr, None) + if t is not None: + transformers_found = True + if _has_svg_enabled(t): + svg_active = True + for extra in extra_targets: + if _has_svg_enabled(extra): + svg_active = True + if not svg_active and not transformers_found: + cfg = getattr(target, "config", target) + if cfg is not None and bool(getattr(cfg, "use_svg_attention", False)): + h_d = getattr(cfg, "svg_high_noise_density", None) + l_d = getattr(cfg, "svg_low_noise_density", None) + s_d = float(getattr(cfg, "svg_spatial_density", 0.25)) + eff_high = float(h_d) if (h_d is not None and float(h_d) >= 0) else s_d + eff_low = float(l_d) if (l_d is not None and float(l_d) >= 0) else s_d + if eff_high < 1.0 or eff_low < 1.0: + svg_active = True + if svg_active: + raise ValueError("SVG sparse attention cannot be combined with CFG cache or MagCache.") + + @partial( aot_cache.cached_jit, static_argnames=( @@ -1380,6 +1488,7 @@ def transformer_forward_pass( kv_cache=None, rotary_emb=None, encoder_attention_mask=None, + svg_step_index=None, ): if do_classifier_free_guidance and latents.shape[0] != prompt_embeds.shape[0]: latents = jnp.concatenate([latents, latents], axis=0) @@ -1395,6 +1504,7 @@ def transformer_forward_pass( kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=svg_step_index, ) if return_residual: @@ -1453,6 +1563,7 @@ def transformer_forward_pass_full_cfg( kv_cache=None, rotary_emb=None, encoder_attention_mask=None, + svg_step_index=None, ): """Full CFG forward pass. @@ -1474,6 +1585,7 @@ def transformer_forward_pass_full_cfg( kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=svg_step_index, ) noise_cond = noise_pred[:bsz] noise_uncond = noise_pred[bsz:] @@ -1498,6 +1610,7 @@ def transformer_forward_pass_cfg_cache( kv_cache=None, rotary_emb=None, encoder_attention_mask=None, + svg_step_index=None, ): """CFG-Cache forward pass with FFT frequency-domain compensation. @@ -1526,6 +1639,7 @@ def transformer_forward_pass_cfg_cache( kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=svg_step_index, ) # FFT over spatial dims (H, W) — last 2 dims of [B, C, F, H, W] diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py index a6b97cc10..0cd0d7f7b 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py @@ -19,6 +19,7 @@ transformer_forward_pass_cfg_cache, init_magcache, magcache_step, + validate_svg_cache_compatibility, ) from ...models.wan.transformers.transformer_wan import WanModel from typing import List, Union, Optional @@ -146,6 +147,7 @@ def __call__( f"use_cfg_cache=True requires guidance_scale > 1.0 (got {guidance_scale}). " "CFG cache accelerates classifier-free guidance, which is disabled when guidance_scale <= 1.0." ) + self._validate_svg_cache_compatibility(use_cfg_cache=use_cfg_cache, use_magcache=use_magcache) trace = {} t_cond_start = time.perf_counter() @@ -261,6 +263,8 @@ def run_inference_2_1( except Exception: pass + validate_svg_cache_compatibility(config, use_cfg_cache, use_magcache, graphdef) + if use_cfg_cache and do_cfg and bsz % data_shards != 0: max_logging.log( f"Warning: Disabling CFG cache because batch size {bsz} is not divisible by data shards {data_shards}. This often happens with data_parallelism > 1 and per_device_batch_size = 1." @@ -372,6 +376,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=current_scheduler_state.step_index, ) else: timestep = jnp.broadcast_to(t, bsz) @@ -387,6 +392,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=current_scheduler_state.step_index, ) new_latents, new_scheduler_state = scheduler.step( @@ -437,6 +443,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) if not skip_blocks: @@ -465,6 +472,7 @@ def scan_body(carry, t): kv_cache=kv_cache_cond, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask_cond, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) elif do_cfg: @@ -485,6 +493,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) else: @@ -501,6 +510,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py index 29546ff9a..f55bd79e5 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py @@ -19,6 +19,7 @@ transformer_forward_pass_cfg_cache, init_magcache, magcache_step, + validate_svg_cache_compatibility, ) from ...models.wan.transformers.transformer_wan import WanModel from typing import List, Union, Optional @@ -193,6 +194,8 @@ def __call__( "SenCache requires classifier-free guidance to be enabled for both transformer phases." ) + self._validate_svg_cache_compatibility(use_cfg_cache=use_cfg_cache, use_magcache=use_magcache) + trace = {} t_cond_start = time.perf_counter() @@ -323,6 +326,8 @@ def run_inference_2_2( except Exception: pass + validate_svg_cache_compatibility(config, use_cfg_cache, use_magcache, low_noise_graphdef, high_noise_graphdef) + if use_cfg_cache and do_classifier_free_guidance and bsz % data_shards != 0: max_logging.log( f"Warning: Disabling CFG cache because batch size {bsz} is not divisible by data shards {data_shards}. This often happens with data_parallelism > 1 and per_device_batch_size = 1." @@ -451,6 +456,7 @@ def run_inference_2_2( kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) if skip_blocks: @@ -554,6 +560,7 @@ def run_inference_2_2( kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) ref_noise_pred = noise_pred ref_latent = latents @@ -592,6 +599,7 @@ def compute_fn(): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) return out @@ -716,6 +724,7 @@ def compute_fn(): kv_cache=kv_cache_cond, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask_cond, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) else: # ── Full CFG step: doubled batch, store raw cond/uncond for cache ── @@ -736,6 +745,7 @@ def compute_fn(): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() @@ -768,6 +778,7 @@ def high_noise_branch(operands): r_emb, mask_high, _, + step_idx, ) = operands return transformer_forward_pass( high_noise_graphdef, @@ -781,6 +792,7 @@ def high_noise_branch(operands): kv_cache=kv_cache_high, rotary_emb=r_emb, encoder_attention_mask=mask_high, + svg_step_index=step_idx, ) def low_noise_branch(operands): @@ -793,6 +805,7 @@ def low_noise_branch(operands): r_emb, _, mask_low, + step_idx, ) = operands return transformer_forward_pass( low_noise_graphdef, @@ -806,6 +819,7 @@ def low_noise_branch(operands): kv_cache=kv_cache_low, rotary_emb=r_emb, encoder_attention_mask=mask_low, + svg_step_index=step_idx, ) if scan_diffusion_loop: @@ -814,8 +828,9 @@ def low_noise_branch(operands): step_index=jnp.array(0, dtype=jnp.int32), ) - def scan_body(carry, t): + def scan_body(carry, scan_elem): current_latents, current_scheduler_state = carry + t, step_idx = scan_elem timestep = jnp.broadcast_to(t, (bsz * 2 if do_classifier_free_guidance else bsz,)) use_high_noise = jnp.greater_equal(t, boundary) @@ -833,6 +848,7 @@ def scan_body(carry, t): rotary_emb, encoder_attention_mask_high, encoder_attention_mask_low, + step_idx, ), ) @@ -843,8 +859,9 @@ def scan_body(carry, t): return (new_latents, new_scheduler_state), None initial_carry = (latents, scheduler_state) + scan_input = (timesteps, jnp.arange(num_inference_steps, dtype=jnp.int32)) - final_carry, _ = jax.lax.scan(scan_body, initial_carry, timesteps) + final_carry, _ = jax.lax.scan(scan_body, initial_carry, scan_input) final_latents, _ = final_carry return final_latents @@ -888,6 +905,7 @@ def scan_body(carry, t): kv_cache=_kv, rotary_emb=rotary_emb, encoder_attention_mask=_mask, + svg_step_index=jnp.asarray(0, dtype=jnp.int32), ) ) @@ -929,6 +947,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) else: timestep = jnp.broadcast_to(t, bsz) @@ -944,6 +963,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py index 52ad66c46..3ee802e9d 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py @@ -14,7 +14,13 @@ from maxdiffusion import max_logging from maxdiffusion.image_processor import PipelineImageInput -from .wan_pipeline import WanPipeline, transformer_forward_pass, init_magcache, magcache_step +from .wan_pipeline import ( + WanPipeline, + transformer_forward_pass, + init_magcache, + magcache_step, + validate_svg_cache_compatibility, +) from ...models.wan.transformers.transformer_wan import WanModel from typing import List, Union, Optional, Tuple from ...pyconfig import HyperParameters @@ -185,6 +191,7 @@ def __call__( last_image: Optional[PipelineImageInput] = None, output_type: Optional[str] = "np", rng: Optional[jax.Array] = None, + use_cfg_cache: bool = False, use_magcache: bool = False, magcache_thresh: Optional[float] = None, magcache_K: Optional[int] = None, @@ -192,6 +199,7 @@ def __call__( use_kv_cache: bool = False, ): config = getattr(self, "config", None) + self._validate_svg_cache_compatibility(use_cfg_cache=use_cfg_cache, use_magcache=use_magcache) if max_sequence_length is None: max_sequence_length = getattr(config, "max_sequence_length", 512) @@ -350,6 +358,7 @@ def run_inference_2_1_i2v( use_kv_cache: bool = False, ): do_cfg = guidance_scale > 1.0 + validate_svg_cache_compatibility(config, False, use_magcache, graphdef) if use_magcache and do_cfg: magcache_init = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base) @@ -422,6 +431,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=current_scheduler_state.step_index, ) noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1)) @@ -481,6 +491,7 @@ def scan_body(carry, t): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) if use_magcache and do_cfg: noise_pred, residual_x_cur = outputs diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py index 91e77f1b8..fd249a7de 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py @@ -21,6 +21,7 @@ transformer_forward_pass_cfg_cache, init_magcache, magcache_step, + validate_svg_cache_compatibility, ) from ...models.wan.transformers.transformer_wan import WanModel from typing import List, Union, Optional, Tuple @@ -276,6 +277,8 @@ def __call__( "SenCache requires classifier-free guidance to be enabled for both transformer phases." ) + self._validate_svg_cache_compatibility(use_cfg_cache=use_cfg_cache, use_magcache=use_magcache) + height = height or self.config.height width = width or self.config.width num_frames = num_frames or self.config.num_frames @@ -452,6 +455,8 @@ def run_inference_2_2_i2v( except Exception: pass + validate_svg_cache_compatibility(config, use_cfg_cache, use_magcache, low_noise_graphdef, high_noise_graphdef) + if use_cfg_cache and do_classifier_free_guidance and bsz % data_shards != 0: max_logging.log( f"Warning: Disabling CFG cache because batch size {bsz} is not divisible by data shards {data_shards}. This often happens with data_parallelism > 1 and per_device_batch_size = 1." @@ -585,6 +590,7 @@ def run_inference_2_2_i2v( kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1)) @@ -681,6 +687,7 @@ def run_inference_2_2_i2v( kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1)) ref_noise_pred = noise_pred @@ -722,6 +729,7 @@ def compute_fn(): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) return jnp.transpose(out, (0, 2, 3, 4, 1)) @@ -853,6 +861,7 @@ def compute_fn(): kv_cache=kv_cache_cond, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask_cond, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) else: # ── Full CFG step: doubled batch, store raw cond/uncond for cache ── @@ -876,6 +885,7 @@ def compute_fn(): kv_cache=kv_cache, rotary_emb=rotary_emb, encoder_attention_mask=encoder_attention_mask, + svg_step_index=jnp.asarray(step, dtype=jnp.int32), ) noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1)) # BCFHW -> BFHWC @@ -894,6 +904,7 @@ def high_noise_branch(operands): r_emb, mask_high, _, + svg_step_index, ) = operands return transformer_forward_pass( high_noise_graphdef, @@ -908,6 +919,7 @@ def high_noise_branch(operands): kv_cache=kv_cache_high, rotary_emb=r_emb, encoder_attention_mask=mask_high, + svg_step_index=svg_step_index, ) def low_noise_branch(operands): @@ -921,6 +933,7 @@ def low_noise_branch(operands): r_emb, _, mask_low, + svg_step_index, ) = operands return transformer_forward_pass( low_noise_graphdef, @@ -935,6 +948,7 @@ def low_noise_branch(operands): kv_cache=kv_cache_low, rotary_emb=r_emb, encoder_attention_mask=mask_low, + svg_step_index=svg_step_index, ) if do_classifier_free_guidance: @@ -979,6 +993,7 @@ def scan_body(carry, t): rotary_emb, encoder_attention_mask_high, encoder_attention_mask_low, + current_scheduler_state.step_index, ), ) noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1)) @@ -1023,6 +1038,7 @@ def scan_body(carry, t): rotary_emb, encoder_attention_mask_high, encoder_attention_mask_low, + jnp.asarray(step, dtype=jnp.int32), )) noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1)) latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py index 916e80ef3..25ddd1ca4 100644 --- a/src/maxdiffusion/tests/aot_cache_test.py +++ b/src/maxdiffusion/tests/aot_cache_test.py @@ -179,6 +179,71 @@ def test_signature_deterministic_across_processes(self): ] self.assertEqual(outs[0], outs[1]) + def test_graphdef_svg_config_changes_dynamic_signature_and_prevents_reuse(self): + from flax import nnx + + class ToyExpert(nnx.Module): + + def __init__(self, use_svg: bool, density: float): + self.use_svg_attention = use_svg + self.svg_spatial_density = density + self.config = {"attention_config": {"use_svg_attention": use_svg, "svg_spatial_density": density}} + self.w = nnx.Param(jnp.ones((8, 8))) + + def __call__(self, x): + scale = self.svg_spatial_density if self.use_svg_attention else 1.0 + return (x @ self.w[...]) * scale + + @aot_cache.cached_jit + def forward(graphdef, state, x): + model = nnx.merge(graphdef, state) + return model(x) + + self._install() + gd_dense, state_dense = nnx.split(ToyExpert(False, 1.0)) + gd_svg50, state_svg50 = nnx.split(ToyExpert(True, 0.5)) + gd_svg25, state_svg25 = nnx.split(ToyExpert(True, 0.25)) + + sig_dense = aot_cache._dynamic_signature((gd_dense, state_dense, self._a), {}) + sig_svg50 = aot_cache._dynamic_signature((gd_svg50, state_svg50, self._a), {}) + sig_svg25 = aot_cache._dynamic_signature((gd_svg25, state_svg25, self._a), {}) + self.assertEqual(len({sig_dense, sig_svg50, sig_svg25}), 3) + + # Cache dense executable + out_dense = forward(gd_dense, state_dense, self._a) + self.assertEqual(aot_cache.save_pending(), 1) + np.testing.assert_allclose(np.asarray(out_dense), np.full((8, 8), 8.0)) + + # Call SVG expert with identical input shapes; must NOT reuse dense executable + out_svg50 = forward(gd_svg50, state_svg50, self._a) + np.testing.assert_allclose(np.asarray(out_svg50), np.full((8, 8), 4.0)) + self.assertEqual(aot_cache.save_pending(), 1) + + def test_extract_svg_meta_changes_fingerprint_on_svg_changes(self): + from types import SimpleNamespace + + cfg_dense = SimpleNamespace(use_svg_attention=False, svg_spatial_density=1.0) + cfg_svg50 = SimpleNamespace(use_svg_attention=True, svg_spatial_density=0.5) + cfg_svg25 = SimpleNamespace(use_svg_attention=True, svg_spatial_density=0.25) + + meta_dense = aot_cache.extract_svg_meta(cfg_dense) + meta_svg50 = aot_cache.extract_svg_meta(cfg_svg50) + meta_svg25 = aot_cache.extract_svg_meta(cfg_svg25) + + fp_dense = aot_cache._metadata_fingerprint(meta_dense) + fp_svg50 = aot_cache._metadata_fingerprint(meta_svg50) + fp_svg25 = aot_cache._metadata_fingerprint(meta_svg25) + self.assertEqual(len({fp_dense, fp_svg50, fp_svg25}), 3) + + def test_step_array_preserves_dynamic_signature_across_steps(self): + sigs = { + aot_cache._dynamic_signature( + (self._a,), {"svg_step_index": jnp.asarray(step, dtype=jnp.int32)} + ) + for step in range(40) + } + self.assertEqual(len(sigs), 1) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/tests/wan/svg_attention_test.py b/src/maxdiffusion/tests/wan/svg_attention_test.py new file mode 100644 index 000000000..638d03ac0 --- /dev/null +++ b/src/maxdiffusion/tests/wan/svg_attention_test.py @@ -0,0 +1,702 @@ +""" +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. + +Comprehensive CPU and numerical reference tests for Sparse VideoGen (SVG). +""" + +import math +import unittest + +import jax +import jax.numpy as jnp +import numpy as np + +from maxdiffusion.models.wan.transformers import svg_attention +from maxdiffusion.kernels import custom_svg_static_range_attention as static_kernel + + +def _numpy_upstream_get_attention_mask(mask_name: str, sample_mse_max_row: int, num_frame: int, frame_size: int): + """Exact upstream mask reference from Sparse VideoGen.""" + seq_len = num_frame * frame_size + block_size = 128 + block_thres = frame_size * 2 + num_block = math.ceil(seq_len / block_size) + pixel_attn_mask = np.zeros((seq_len, seq_len), dtype=bool) + pixel_attn_mask[:, :frame_size] = 1 + + for i in range(num_block): + for j in range(num_block): + if abs(i - j) < block_thres // block_size: + pixel_attn_mask[i * block_size : (i + 1) * block_size, j * block_size : (j + 1) * block_size] = 1 + + if mask_name == "spatial": + attention_mask = pixel_attn_mask + else: + pixel_attn_mask = ( + pixel_attn_mask.reshape(frame_size, num_frame, frame_size, num_frame) + .transpose(1, 0, 3, 2) + .reshape(frame_size * num_frame, frame_size * num_frame) + ) + attention_mask = pixel_attn_mask + + return attention_mask[:sample_mse_max_row] + + +def _require_tpu(test): + """These tests call Pallas directly; CPU only supports interpret mode.""" + if jax.devices()[0].platform != "tpu": + test.skipTest("Requires TPU: Pallas runs only in interpret mode on CPU") + + +class SVGAttentionUnitTest(unittest.TestCase): + + def test_01_sample_row_domain(self): + F, H, W = 21, 45, 80 + L = F * H * W + key = jax.random.PRNGKey(42) + + sampled_rows = jax.random.randint(key, (64,), minval=0, maxval=min(10000, L), dtype=jnp.int32) + self.assertEqual(sampled_rows.shape, (64,)) + self.assertTrue(jnp.all(sampled_rows >= 0)) + self.assertTrue(jnp.all(sampled_rows < 10000)) + self.assertTrue(jnp.all(sampled_rows < L)) + + L_small = 1280 + sampled_small = jax.random.randint(key, (64,), minval=0, maxval=min(10000, L_small), dtype=jnp.int32) + self.assertTrue(jnp.all(sampled_small >= 0)) + self.assertTrue(jnp.all(sampled_small < 1280)) + + def test_02_fixed_profiler_spatial_mask(self): + F, H, W = 5, 16, 16 + P = H * W + L = F * P + token_grid = (F, H, W) + sampled_rows = jnp.array([0, 10, 127, 128, 255, 256, 500, 1000, 1279], dtype=jnp.int32) + + jax_spatial, _ = svg_attention.svg_probe_masks(sampled_rows, token_grid, block_size=128) + np_spatial = _numpy_upstream_get_attention_mask("spatial", L, F, P)[np.asarray(sampled_rows), :] + np.testing.assert_array_equal(np.asarray(jax_spatial), np_spatial) + + def test_03_fixed_profiler_temporal_mask(self): + F, H, W = 5, 16, 16 + P = H * W + L = F * P + token_grid = (F, H, W) + sampled_rows = jnp.array([0, 10, 127, 128, 255, 256, 500, 1000, 1279], dtype=jnp.int32) + + _, jax_temporal = svg_attention.svg_probe_masks(sampled_rows, token_grid, block_size=128) + np_temporal = _numpy_upstream_get_attention_mask("temporal", L, F, P)[np.asarray(sampled_rows), :] + np.testing.assert_array_equal(np.asarray(jax_temporal), np_temporal) + + def test_04_first_frame_sink_semantics(self): + F, H, W = 5, 16, 16 + P = H * W + L = F * P + token_grid = (F, H, W) + sampled_rows = jnp.array([600, 1000], dtype=jnp.int32) + + jax_spatial, jax_temporal = svg_attention.svg_probe_masks(sampled_rows, token_grid, block_size=128) + self.assertTrue(np.all(np.asarray(jax_spatial[:, :P]))) + np_temporal = _numpy_upstream_get_attention_mask("temporal", L, F, P)[np.asarray(sampled_rows), :] + np.testing.assert_array_equal(np.asarray(jax_temporal), np_temporal) + + def test_05_band_width_calculation(self): + F, H, W = 21, 45, 80 + L = F * H * W + token_grid = (F, H, W) + + w_25 = svg_attention.svg_execution_band_width(token_grid, 0.25) + self.assertEqual(w_25, 10240) + self.assertEqual(w_25 % 128, 0) + self.assertEqual(svg_attention.svg_execution_band_width(token_grid, 1.0), L - 1) + + def test_06_density_boundary_conditions(self): + token_grid = (5, 4, 4) + self.assertEqual(svg_attention.svg_execution_band_width(token_grid, 1.0), 79) + self.assertGreater(svg_attention.svg_execution_band_width(token_grid, 0.5), 0) + with self.assertRaises(ValueError): + svg_attention.svg_execution_band_width(token_grid, 0.0) + with self.assertRaises(ValueError): + svg_attention.svg_execution_band_width(token_grid, -0.1) + with self.assertRaises(ValueError): + svg_attention.svg_execution_band_width(token_grid, 1.05) + + def test_07_placement_permutation_and_inverse(self): + F, H, W = 4, 3, 2 + L = F * H * W + token_grid = (F, H, W) + B, num_heads, D = 2, 4, 8 + + rng = jax.random.PRNGKey(123) + x = jax.random.normal(rng, (B, num_heads, L, D)) + + is_temporal_spatial = jnp.zeros((B, num_heads), dtype=bool) + placed_sp = svg_attention.svg_placement_permute(x, x, x, is_temporal_spatial, token_grid)[0] + restored_sp = svg_attention.svg_placement_unpermute(placed_sp, is_temporal_spatial, token_grid) + np.testing.assert_allclose(np.asarray(restored_sp), np.asarray(x), atol=1e-6) + np.testing.assert_allclose(np.asarray(placed_sp), np.asarray(x), atol=1e-6) + + is_temporal_all = jnp.ones((B, num_heads), dtype=bool) + placed_tp = svg_attention.svg_placement_permute(x, x, x, is_temporal_all, token_grid)[0] + restored_tp = svg_attention.svg_placement_unpermute(placed_tp, is_temporal_all, token_grid) + np.testing.assert_allclose(np.asarray(restored_tp), np.asarray(x), atol=1e-6) + + def test_08_mixed_per_head_permutation(self): + F, H, W = 3, 2, 2 + L = F * H * W + token_grid = (F, H, W) + B, num_heads, D = 1, 4, 8 + + q = jax.random.normal(jax.random.PRNGKey(1), (B, num_heads, L, D)) + k = jax.random.normal(jax.random.PRNGKey(2), (B, num_heads, L, D)) + v = jax.random.normal(jax.random.PRNGKey(3), (B, num_heads, L, D)) + is_temporal = jnp.array([[False, True, True, False]]) + + q_placed, _, _ = svg_attention.svg_placement_permute(q, k, v, is_temporal, token_grid) + forward_idx, _ = svg_attention.svg_token_major_indices(token_grid) + np.testing.assert_allclose(np.asarray(q_placed[:, 0]), np.asarray(q[:, 0])) + np.testing.assert_allclose(np.asarray(q_placed[:, 3]), np.asarray(q[:, 3])) + np.testing.assert_allclose(np.asarray(q_placed[:, 1]), np.asarray(jnp.take(q[:, 1], forward_idx, axis=1))) + np.testing.assert_allclose(np.asarray(q_placed[:, 2]), np.asarray(jnp.take(q[:, 2], forward_idx, axis=1))) + q_restored = svg_attention.svg_placement_unpermute(q_placed, is_temporal, token_grid) + np.testing.assert_allclose(np.asarray(q_restored), np.asarray(q), atol=1e-6) + + def test_09_routing_argmin_oracle(self): + F, H, W = 3, 2, 2 + P = H * W + L = F * P + token_grid = (F, H, W) + B, num_heads, D = 1, 2, 16 + + key = jax.random.PRNGKey(10) + k1, k2, k3, k4 = jax.random.split(key, 4) + q = jax.random.normal(k1, (B, num_heads, L, D)) + k = jax.random.normal(k2, (B, num_heads, L, D)) + v = jax.random.normal(k3, (B, num_heads, L, D)) + + scale = 1.0 / math.sqrt(D) + jax_is_temporal = svg_attention.svg_profile_temporal_heads( + q, k, v, token_grid, query_count=12, profile_key=k4, scale=scale, sample_max_row=12 + ) + + sampled_rows = np.arange(12, dtype=np.int32) + sampled_q = np.asarray(q)[:, :, sampled_rows, :] + key_np = np.asarray(k) + val_np = np.asarray(v) + scores = np.einsum("bhqd,bhkd->bhqk", sampled_q, key_np) * scale + dense_w = np.exp(scores - np.max(scores, axis=-1, keepdims=True)) + dense_w /= np.sum(dense_w, axis=-1, keepdims=True) + golden = np.einsum("bhqk,bhkd->bhqd", dense_w, val_np) + + sp_mask = _numpy_upstream_get_attention_mask("spatial", 12, F, P) + tp_mask = _numpy_upstream_get_attention_mask("temporal", 12, F, P) + + def mask_out(mask): + masked_scores = np.where(mask[None, None, :, :], scores, -1e9) + weights = np.exp(masked_scores - np.max(masked_scores, axis=-1, keepdims=True)) + weights /= np.sum(weights, axis=-1, keepdims=True) + return np.einsum("bhqk,bhkd->bhqd", weights, val_np) + + sp_out = mask_out(sp_mask) + tp_out = mask_out(tp_mask) + sp_err = np.mean((sp_out - golden) ** 2, axis=(-2, -1)) + tp_err = np.mean((tp_out - golden) ** 2, axis=(-2, -1)) + np.testing.assert_array_equal(np.asarray(jax_is_temporal), tp_err < sp_err) + + def test_09b_base2_execution_scale_is_softmax_equivalent(self): + logits = jax.random.normal(jax.random.PRNGKey(91), (4, 17), dtype=jnp.float32) + natural = jax.nn.softmax(logits, axis=-1) + log2e = math.log2(math.e) + exp2_weights = jnp.exp2(logits * log2e - jnp.max(logits * log2e, axis=-1, keepdims=True)) + exp2_weights /= jnp.sum(exp2_weights, axis=-1, keepdims=True) + np.testing.assert_allclose(np.asarray(natural), np.asarray(exp2_weights), rtol=2e-6, atol=2e-6) + + def test_inactive_static_step_with_traced_layer(self): + for step in (0, 40): + + def active(layer): + result = svg_attention.is_svg_active( + step_index=step, layer_index=layer, start_step=11, end_step=40, start_layer=1, end_layer=40 + ) + self.assertIs(result, False) + return result + + self.assertFalse(bool(jax.jit(active)(jnp.asarray(2)))) + + active = jax.jit( + lambda step, layer: svg_attention.is_svg_active( + step_index=step, layer_index=layer, start_step=11, end_step=40, start_layer=1, end_layer=40 + ) + ) + self.assertTrue(bool(active(jnp.asarray(11), jnp.asarray(1)))) + self.assertFalse(bool(active(jnp.asarray(10), jnp.asarray(1)))) + + def test_10_step_scheduling_exact_counts(self): + num_inference_steps = 40 + active_start_step = 12 + active_end_step = 40 + step_is_sparse = [(active_start_step <= s < active_end_step) for s in range(num_inference_steps)] + self.assertEqual(sum(not s for s in step_is_sparse), 12) + self.assertEqual(sum(step_is_sparse), 28) + self.assertFalse(step_is_sparse[11]) + self.assertTrue(step_is_sparse[12]) + + def test_11_layer_scheduling_exact_counts(self): + num_layers = 40 + active_start_layer = 1 + active_end_layer = 40 + layer_is_sparse = [(active_start_layer <= l < active_end_layer) for l in range(num_layers)] + self.assertEqual(sum(not l for l in layer_is_sparse), 1) + self.assertEqual(sum(layer_is_sparse), 39) + self.assertFalse(layer_is_sparse[0]) + self.assertTrue(layer_is_sparse[39]) + + def test_12_static_tile_classifier_properties(self): + q_seq_len = 1000 + kv_seq_len = 1000 + bq = 128 + bkv = 128 + band_width = 256 + frame_size = 128 + include_first_frame = True + + full_map, full_act, boundary_map, bnd_act = static_kernel._classify_tiles( + q_seq_len, kv_seq_len, bq, bkv, band_width, frame_size, include_first_frame + ) + q_tiles = math.ceil(q_seq_len / bq) + kv_tiles = math.ceil(kv_seq_len / bkv) + for qi in range(q_tiles): + full_set = set(full_map[qi, : full_act[qi]]) + bnd_set = set(boundary_map[qi, : bnd_act[qi]]) + self.assertEqual(len(full_set.intersection(bnd_set)), 0) + q0 = qi * bq + for kj in full_set: + k0 = kj * bkv + self.assertTrue(q0 + bq <= q_seq_len) + self.assertTrue(k0 + bkv <= kv_seq_len) + sink_full = include_first_frame and (k0 + bkv - 1 < frame_size) + local_full = max(abs(q0 - (k0 + bkv - 1)), abs((q0 + bq - 1) - k0)) <= band_width + self.assertTrue(sink_full or local_full) + q1 = min(q_seq_len, q0 + bq) - 1 + band0 = max(0, q0 - band_width) + band1 = min(kv_seq_len - 1, q1 + band_width) + expected_live = set(range(band0 // bkv, min(kv_tiles, (band1 // bkv) + 1))) + if include_first_frame: + expected_live.update(range(0, min(kv_tiles, ((frame_size - 1) // bkv) + 1))) + self.assertEqual(full_set.union(bnd_set), expected_live) + + def test_13_single_execution_matches_dual_kernel_reference(self): + F, H, W = 3, 2, 2 + P = H * W + L = F * P + token_grid = (F, H, W) + B, num_heads, D = 1, 4, 8 + + key = jax.random.PRNGKey(99) + k1, k2, k3 = jax.random.split(key, 3) + q = jax.random.normal(k1, (B, num_heads, L, D)) + k = jax.random.normal(k2, (B, num_heads, L, D)) + v = jax.random.normal(k3, (B, num_heads, L, D)) + scale = 1.0 / math.sqrt(D) + is_temporal = jnp.array([[False, True, True, False]]) + band_width = 3 + + common_mask = (jnp.arange(L)[None, :] < P) | (jnp.abs(jnp.arange(L)[:, None] - jnp.arange(L)[None, :]) <= band_width) + q_placed, k_placed, v_placed = svg_attention.svg_placement_permute(q, k, v, is_temporal, token_grid) + scores_placed = jnp.einsum("bhqd,bhkd->bhqk", q_placed, k_placed) * scale + logits_placed = jnp.where(common_mask[None, None, :, :], scores_placed, -1e9) + weights_placed = jax.nn.softmax(logits_placed, axis=-1) + out_placed = jnp.einsum("bhqk,bhkd->bhqd", weights_placed, v_placed) + out_single_exec = svg_attention.svg_placement_unpermute(out_placed, is_temporal, token_grid) + + scores_sp = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale + logits_sp = jnp.where(common_mask[None, None, :, :], scores_sp, -1e9) + weights_sp = jax.nn.softmax(logits_sp, axis=-1) + out_spatial = jnp.einsum("bhqk,bhkd->bhqd", weights_sp, v) + tm_idx = (jnp.arange(L) % P) * F + (jnp.arange(L) // P) + temporal_mask = (tm_idx[None, :] < P) | (jnp.abs(tm_idx[:, None] - tm_idx[None, :]) <= band_width) + logits_tp = jnp.where(temporal_mask[None, None, :, :], scores_sp, -1e9) + weights_tp = jax.nn.softmax(logits_tp, axis=-1) + out_temporal = jnp.einsum("bhqk,bhkd->bhqd", weights_tp, v) + out_bruteforce = jnp.where(is_temporal[:, :, None, None], out_temporal, out_spatial) + np.testing.assert_allclose(np.asarray(out_single_exec), np.asarray(out_bruteforce), rtol=1e-5, atol=1e-5) + + def test_14_exact_sparse_placed_matches_dense_masked_reference(self): + _require_tpu(self) + F, H, W = 5, 16, 16 + P = H * W + L = F * P + token_grid = (F, H, W) + B, num_heads, D = 1, 4, 128 + + key = jax.random.PRNGKey(42) + k1, k2, k3 = jax.random.split(key, 3) + q = jax.random.normal(k1, (B, num_heads, L, D), dtype=jnp.bfloat16) + k = jax.random.normal(k2, (B, num_heads, L, D), dtype=jnp.bfloat16) + v = jax.random.normal(k3, (B, num_heads, L, D), dtype=jnp.bfloat16) + scale = 1.0 / math.sqrt(D) + test_cases = [ + ("all_spatial", jnp.zeros((B, num_heads), dtype=bool)), + ("all_temporal", jnp.ones((B, num_heads), dtype=bool)), + ("mixed", jnp.array([[False, True, True, False]])), + ] + band_width = 256 + block_sizes = static_kernel.SVGBlockSizes(block_q=128, block_kv=128, block_kv_compute=128, block_kv_compute_in=128) + scores_fp32 = jnp.einsum("bhqd,bhkd->bhqk", q.astype(jnp.float32), k.astype(jnp.float32)) * scale + sp_mask = (jnp.arange(L)[None, :] < P) | (jnp.abs(jnp.arange(L)[:, None] - jnp.arange(L)[None, :]) <= band_width) + tm_idx = (jnp.arange(L) % P) * F + (jnp.arange(L) // P) + tp_mask = (tm_idx[None, :] < P) | (jnp.abs(tm_idx[:, None] - tm_idx[None, :]) <= band_width) + + for case_name, is_temporal in test_cases: + with self.subTest(case=case_name): + head_masks = jnp.where(is_temporal[:, :, None, None], tp_mask[None, None, :, :], sp_mask[None, None, :, :]) + logits = jnp.where(head_masks, scores_fp32, -1e9) + weights = jax.nn.softmax(logits, axis=-1) + out_ref = jnp.einsum("bhqk,bhkd->bhqd", weights, v.astype(jnp.float32)).astype(jnp.bfloat16) + + q_placed, k_placed, v_placed = svg_attention.svg_placement_permute(q, k, v, is_temporal, token_grid) + q_local = q_placed * math.log2(math.e) + k_scaled = k_placed * scale + svg_kernel = static_kernel.make_svg_exact_static_range_mha( + block_sizes=block_sizes, + orig_q_seq_len=L, + orig_kv_seq_len=L, + band_width=band_width, + frame_size=P, + include_first_frame=True, + use_base2_exp=True, + ) + vmapped_svg = jax.vmap(svg_kernel, in_axes=(0, 0, 0)) + out_placed = vmapped_svg(q_local, k_scaled, v_placed) + out_placed = jnp.swapaxes(out_placed, 2, 3) + out_actual = svg_attention.svg_placement_unpermute(out_placed, is_temporal, token_grid) + rel_l2 = jnp.linalg.norm(out_actual.astype(jnp.float32) - out_ref.astype(jnp.float32)) / jnp.linalg.norm( + out_ref.astype(jnp.float32) + ) + max_abs = jnp.max(jnp.abs(out_actual.astype(jnp.float32) - out_ref.astype(jnp.float32))) + self.assertLess(float(rel_l2), 0.01, f"{case_name} rel_l2 {float(rel_l2):.6f} exceeds tolerance") + self.assertLess(float(max_abs), 0.05, f"{case_name} max_abs {float(max_abs):.6f} exceeds tolerance") + + def test_15_density_one_matches_dense_reference(self): + _require_tpu(self) + F, H, W = 4, 16, 16 + P = H * W + L = F * P + token_grid = (F, H, W) + B, num_heads, D = 1, 2, 128 + + key = jax.random.PRNGKey(777) + k1, k2, k3 = jax.random.split(key, 3) + q = jax.random.normal(k1, (B, num_heads, L, D), dtype=jnp.bfloat16) + k = jax.random.normal(k2, (B, num_heads, L, D), dtype=jnp.bfloat16) + v = jax.random.normal(k3, (B, num_heads, L, D), dtype=jnp.bfloat16) + scale = 1.0 / math.sqrt(D) + is_temporal = jnp.array([[False, True]]) + band_width = svg_attention.svg_execution_band_width(token_grid, 1.0) + self.assertEqual(band_width, L - 1) + + scores_fp32 = jnp.einsum("bhqd,bhkd->bhqk", q.astype(jnp.float32), k.astype(jnp.float32)) * scale + weights = jax.nn.softmax(scores_fp32, axis=-1) + out_ref = jnp.einsum("bhqk,bhkd->bhqd", weights, v.astype(jnp.float32)).astype(jnp.bfloat16) + + block_sizes = static_kernel.SVGBlockSizes(block_q=128, block_kv=128, block_kv_compute=128, block_kv_compute_in=128) + q_placed, k_placed, v_placed = svg_attention.svg_placement_permute(q, k, v, is_temporal, token_grid) + q_local = q_placed * math.log2(math.e) + k_scaled = k_placed * scale + svg_kernel = static_kernel.make_svg_exact_static_range_mha( + block_sizes=block_sizes, + orig_q_seq_len=L, + orig_kv_seq_len=L, + band_width=band_width, + frame_size=P, + include_first_frame=True, + use_base2_exp=True, + ) + vmapped_svg = jax.vmap(svg_kernel, in_axes=(0, 0, 0)) + out_placed = vmapped_svg(q_local, k_scaled, v_placed) + out_placed = jnp.swapaxes(out_placed, 2, 3) + out_actual = svg_attention.svg_placement_unpermute(out_placed, is_temporal, token_grid) + rel_l2 = jnp.linalg.norm(out_actual.astype(jnp.float32) - out_ref.astype(jnp.float32)) / jnp.linalg.norm( + out_ref.astype(jnp.float32) + ) + self.assertLess(float(rel_l2), 0.01) + + def test_profile_respects_sample_pool_limit(self): + from unittest.mock import patch + + q = jnp.zeros((1, 1, 128, 8), dtype=jnp.float32) + for pool_size, query_count in ((32, 64), (32, 32), (32, 8), (256, 256), (1, 64)): + with patch.object(svg_attention, "svg_probe_masks", wraps=svg_attention.svg_probe_masks) as masks: + svg_attention.svg_profile_temporal_heads( + q, q, q, (2, 8, 8), query_count, jax.random.PRNGKey(0), 1.0, sample_max_row=pool_size + ) + rows = np.asarray(masks.call_args.args[0]) + self.assertEqual(len(rows), min(query_count, pool_size, 128)) + self.assertTrue(np.all((rows >= 0) & (rows < min(pool_size, 128)))) + + def test_low_noise_svg_nested_config_rejects_caches(self): + from types import SimpleNamespace + from maxdiffusion.pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2 + + pipeline = WanPipeline2_2.__new__(WanPipeline2_2) + pipeline.use_svg_attention = False + pipeline.low_noise_transformer = SimpleNamespace(config=SimpleNamespace(attention_config={"use_svg_attention": True})) + for flag in ("use_cfg_cache", "use_magcache"): + with self.assertRaisesRegex(ValueError, "SVG sparse attention cannot be combined"): + pipeline(prompt="test", guidance_scale_low=5.0, guidance_scale_high=5.0, **{flag: True}) + + def test_oversized_profile_request_matches_capped_request_under_jit(self): + from unittest.mock import patch + + token_grid = (4, 8, 8) + rng = np.random.default_rng(42) + q, k, v = [jnp.asarray(rng.normal(size=(1, 4, 256, 8)), dtype=jnp.float32) for _ in range(3)] + for pool_size, query_count in ((32, 64), (32, 512), (512, 512), (0, 64)): + with self.subTest(pool_size=pool_size, query_count=query_count): + count = min(max(pool_size, 1), 256) + + def route(q, requested): + return svg_attention.svg_profile_temporal_heads( + q, k, v, token_grid, requested, jax.random.PRNGKey(0), 8**-0.5, sample_max_row=pool_size + ) + + with patch.object(svg_attention, "svg_probe_masks", wraps=svg_attention.svg_probe_masks) as masks: + expected = route(q, count) + actual = route(q, query_count) + np.testing.assert_array_equal(np.asarray(masks.call_args.args[0]), np.arange(count)) + np.testing.assert_array_equal(actual, expected) + compiled = jax.jit(lambda query: route(query, query_count)) + np.testing.assert_array_equal(compiled(q), expected) + # Queries outside the eligible pool cannot influence the route. + changed = q.at[:, :, count:, :].set(1e3) + np.testing.assert_array_equal(compiled(changed), expected) + + def test_disabled_low_noise_svg_allows_caches(self): + from types import SimpleNamespace + from unittest.mock import Mock + from maxdiffusion.pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2 + + for config in ( + None, + SimpleNamespace(), + SimpleNamespace(attention_config=None), + SimpleNamespace(attention_config={}), + SimpleNamespace(attention_config={"use_svg_attention": False}), + ): + for flag in ("use_cfg_cache", "use_magcache"): + with self.subTest(config=config, flag=flag): + pipeline = WanPipeline2_2.__new__(WanPipeline2_2) + pipeline.use_svg_attention = False + pipeline.low_noise_transformer = SimpleNamespace(config=config) + pipeline._prepare_model_inputs = Mock(side_effect=RuntimeError("reached input preparation")) + with self.assertRaisesRegex(RuntimeError, "reached input preparation"): + pipeline(prompt="test", **{flag: True}) + pipeline._prepare_model_inputs.assert_called_once() + + def test_svg_cache_check_without_loaded_low_noise_transformer(self): + from unittest.mock import Mock + from maxdiffusion.pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2 + + pipeline = WanPipeline2_2.__new__(WanPipeline2_2) + pipeline.use_svg_attention = False + pipeline.low_noise_transformer = None + pipeline._prepare_model_inputs = Mock(side_effect=RuntimeError("reached input preparation")) + with self.assertRaisesRegex(RuntimeError, "reached input preparation"): + pipeline(prompt="test") + pipeline._prepare_model_inputs.assert_called_once() + + def test_svg_cache_incompatibility_fail_closed(self): + from types import SimpleNamespace + from maxdiffusion.pipelines.wan.wan_pipeline_2_1 import run_inference_2_1 + from maxdiffusion.pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2 + + for cache_flag in ["use_cfg_cache", "use_magcache"]: + cfg = SimpleNamespace(use_svg_attention=True) + with self.assertRaises(ValueError): + run_inference_2_1( + graphdef=None, + sharded_state=None, + rest_of_state=None, + latents=jnp.zeros((1, 16, 21, 45, 80), dtype=jnp.bfloat16), + prompt_embeds=None, + negative_prompt_embeds=None, + guidance_scale=5.0, + num_inference_steps=40, + scheduler=None, + scheduler_state=None, + config=cfg, + use_cfg_cache=(cache_flag == "use_cfg_cache"), + use_magcache=(cache_flag == "use_magcache"), + ) + + pipeline_22 = WanPipeline2_2.__new__(WanPipeline2_2) + pipeline_22.use_svg_attention = True + pipeline_22.low_noise_transformer = type("T", (), {"config": type("C", (), {"use_svg_attention": True})()})() + for cache_flag in ["use_cfg_cache", "use_magcache"]: + with self.assertRaises(ValueError): + pipeline_22( + prompt="test prompt", + use_cfg_cache=(cache_flag == "use_cfg_cache"), + use_magcache=(cache_flag == "use_magcache"), + ) + + def test_partially_specified_schedules(self): + # 1. Layer-only schedule: [1, 40), steps unrestricted + self.assertIs( + svg_attention.is_svg_active(step_index=0, layer_index=0, start_layer=1, end_layer=40), + False, + ) + self.assertIs( + svg_attention.is_svg_active(step_index=0, layer_index=1, start_layer=1, end_layer=40), + True, + ) + self.assertIs( + svg_attention.is_svg_active(step_index=None, layer_index=0, start_layer=1, end_layer=40), + False, + ) + self.assertIs( + svg_attention.is_svg_active(step_index=None, layer_index=1, start_layer=1, end_layer=40), + True, + ) + # Statically inactive layer short-circuits to Python False even with JAX array step_index + self.assertIs( + svg_attention.is_svg_active(step_index=jnp.asarray(15), layer_index=0, start_layer=1, end_layer=40), + False, + ) + self.assertTrue( + bool(svg_attention.is_svg_active(step_index=jnp.asarray(15), layer_index=1, start_layer=1, end_layer=40)) + ) + + # 2. Step-only schedule: [11, 40), layers unrestricted + for layer in (0, 1, 20, 39): + self.assertIs( + svg_attention.is_svg_active(step_index=10, layer_index=layer, start_step=11, end_step=40), + False, + ) + self.assertIs( + svg_attention.is_svg_active(step_index=11, layer_index=layer, start_step=11, end_step=40), + True, + ) + self.assertTrue( + bool(svg_attention.is_svg_active(step_index=jnp.asarray(11), layer_index=layer, start_step=11, end_step=40)) + ) + self.assertFalse( + bool(svg_attention.is_svg_active(step_index=jnp.asarray(10), layer_index=layer, start_step=11, end_step=40)) + ) + + # 3. Fully specified schedule: steps [11, 40), layers [1, 40) + self.assertIs( + svg_attention.is_svg_active( + step_index=11, layer_index=0, start_step=11, end_step=40, start_layer=1, end_layer=40 + ), + False, + ) + self.assertIs( + svg_attention.is_svg_active( + step_index=10, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40 + ), + False, + ) + self.assertIs( + svg_attention.is_svg_active( + step_index=11, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40 + ), + True, + ) + + # 4. Incomplete schedules must raise ValueError + with self.assertRaisesRegex(ValueError, "Incomplete explicit SVG step schedule"): + svg_attention.is_svg_active(step_index=0, layer_index=0, start_step=11, end_step=-1) + with self.assertRaisesRegex(ValueError, "Incomplete explicit SVG layer schedule"): + svg_attention.is_svg_active(step_index=0, layer_index=0, start_layer=1, end_layer=-1) + + def test_high_only_low_only_and_i2v_svg_cache_incompatibility(self): + from types import SimpleNamespace + from unittest.mock import Mock + from maxdiffusion.pipelines.wan.wan_pipeline_2_1 import WanPipeline2_1 + from maxdiffusion.pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2 + from maxdiffusion.pipelines.wan.wan_pipeline_i2v_2p1 import WanPipelineI2V_2_1 + from maxdiffusion.pipelines.wan.wan_pipeline_i2v_2p2 import WanPipelineI2V_2_2 + + def make_transformer(use_svg: bool): + return SimpleNamespace(config=SimpleNamespace(attention_config={"use_svg_attention": use_svg})) + + # Dual-expert pipelines (T2V 2.2 and I2V 2.2) without pipeline.use_svg_attention set + for pipe_cls in (WanPipeline2_2, WanPipelineI2V_2_2): + extra_kwargs = {"image": None} if "I2V" in pipe_cls.__name__ else {} + for high_svg, low_svg in ((True, False), (False, True), (True, True), (False, False)): + for flag in ("use_cfg_cache", "use_magcache"): + with self.subTest(pipe=pipe_cls.__name__, high=high_svg, low=low_svg, flag=flag): + pipe = pipe_cls.__new__(pipe_cls) + pipe.config = SimpleNamespace(height=480, width=832, num_frames=81, max_sequence_length=512) + pipe.vae_scale_factor_temporal = 4 + pipe.high_noise_transformer = make_transformer(high_svg) + pipe.low_noise_transformer = make_transformer(low_svg) + pipe._prepare_model_inputs = Mock(side_effect=RuntimeError("reached input preparation")) + pipe._prepare_model_inputs_i2v = pipe._prepare_model_inputs + if high_svg or low_svg: + with self.assertRaisesRegex(ValueError, "SVG sparse attention cannot be combined"): + pipe(prompt="test", guidance_scale_low=5.0, guidance_scale_high=5.0, **{flag: True}, **extra_kwargs) + else: + with self.assertRaisesRegex(RuntimeError, "reached input preparation"): + pipe(prompt="test", guidance_scale_low=5.0, guidance_scale_high=5.0, **{flag: True}, **extra_kwargs) + + # Single-expert pipelines (T2V 2.1 and I2V 2.1) + for pipe_cls in (WanPipeline2_1, WanPipelineI2V_2_1): + extra_kwargs = {"image": None} if "I2V" in pipe_cls.__name__ else {} + for use_svg in (True, False): + for flag in ("use_cfg_cache", "use_magcache"): + with self.subTest(pipe=pipe_cls.__name__, use_svg=use_svg, flag=flag): + pipe = pipe_cls.__new__(pipe_cls) + pipe.config = SimpleNamespace(height=480, width=832, num_frames=81, max_sequence_length=512) + pipe.vae_scale_factor_temporal = 4 + pipe.transformer = make_transformer(use_svg) + pipe._prepare_model_inputs = Mock(side_effect=RuntimeError("reached input preparation")) + pipe._prepare_model_inputs_i2v = pipe._prepare_model_inputs + if use_svg: + with self.assertRaisesRegex(ValueError, "SVG sparse attention cannot be combined"): + pipe(prompt="test", guidance_scale=5.0, **{flag: True}, **extra_kwargs) + else: + with self.assertRaisesRegex(RuntimeError, "reached input preparation"): + pipe(prompt="test", guidance_scale=5.0, **{flag: True}, **extra_kwargs) + + def test_i2v_step_argument_preserves_aot_signature(self): + from maxdiffusion import aot_cache + import inspect + from maxdiffusion.pipelines.wan import wan_pipeline_i2v_2p1, wan_pipeline_i2v_2p2 + + # 1. Verify source code passes jnp.asarray(step, dtype=jnp.int32) rather than raw Python int step + src_2p1 = inspect.getsource(wan_pipeline_i2v_2p1.run_inference_2_1_i2v) + src_2p2 = inspect.getsource(wan_pipeline_i2v_2p2.run_inference_2_2_i2v) + self.assertNotIn("svg_step_index=step,", src_2p1) + self.assertNotIn("svg_step_index=step,", src_2p2) + self.assertIn("svg_step_index=jnp.asarray(step, dtype=jnp.int32)", src_2p1) + self.assertIn("svg_step_index=jnp.asarray(step, dtype=jnp.int32)", src_2p2) + + # 2. Verify dynamic signature is identical across all 40 denoising steps + dummy_latents = jnp.zeros((2, 16, 21, 45, 80), dtype=jnp.bfloat16) + dummy_ts = jnp.zeros((2,), dtype=jnp.int32) + sigs = { + aot_cache._dynamic_signature( + (dummy_latents, dummy_ts), + {"svg_step_index": jnp.asarray(step, dtype=jnp.int32)}, + ) + for step in range(40) + } + self.assertEqual(len(sigs), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/wan/svg_balanced_rounding_test.py b/src/maxdiffusion/tests/wan/svg_balanced_rounding_test.py new file mode 100644 index 000000000..a184d927a --- /dev/null +++ b/src/maxdiffusion/tests/wan/svg_balanced_rounding_test.py @@ -0,0 +1,367 @@ +""" +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. + +Tests for SVG balanced boundary rounding. +""" + +from __future__ import annotations + +import math + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from maxdiffusion.kernels import custom_svg_attention_dispatch as dispatch +from maxdiffusion.kernels import custom_svg_balanced_rounding_attention as balanced +from maxdiffusion.kernels import custom_svg_static_range_attention as static_range +from maxdiffusion.kernels import custom_svg_balanced_rounding_partial as padding + + +def _bs(): + return static_range.SVGBlockSizes( + block_q=3328, + block_kv=2816, + block_kv_compute=256, + block_kv_compute_in=256, + ) + + +def _band_width(density: float, n: int = 75600): + if density >= 1.0: + return n - 1 + return min( + n - 1, + int(math.ceil((n * (1.0 - math.sqrt(1.0 - density))) / 128.0)) * 128, + ) + + +def _budget(density: float, n: int = 75600, frame_size: int = 3600): + _, _, bm, _, stats = balanced.build_boundary_stats( + orig_q_seq_len=n, + orig_kv_seq_len=n, + block_sizes=_bs(), + band_width=_band_width(density, n), + frame_size=frame_size, + include_first_frame=True, + ) + _, active, report = balanced.build_selected_boundary_table( + stats=stats, + qtiles=bm.shape[0], + policy="global_balanced", + budget_scale=1.0, + ) + return active, report + + +def test_global_balanced_matches_exact_pair_budget_720p_operating_points(): + for density in (0.65, 0.50, 0.35, 0.20, 0.15): + _, report = _budget(density) + assert abs(report["budget_error_fraction"]) <= 0.02, (density, report) + assert report["boundary_tiles_selected"] > 0 + assert report["rounded_boundary_pairs"] > 0 + assert report["exact_boundary_pairs"] > 0 + + +def test_round_up_and_down_bracket_exact_boundary_budget(): + _, _, bm, _, stats = balanced.build_boundary_stats( + orig_q_seq_len=75600, + orig_kv_seq_len=75600, + block_sizes=_bs(), + band_width=22144, + frame_size=3600, + include_first_frame=True, + ) + _, _, up = balanced.build_selected_boundary_table( + stats=stats, + qtiles=bm.shape[0], + policy="up", + ) + _, _, down = balanced.build_selected_boundary_table( + stats=stats, + qtiles=bm.shape[0], + policy="down", + ) + assert up["rounded_boundary_pairs"] >= up["exact_boundary_pairs"] + assert down["rounded_boundary_pairs"] == 0 + + +def test_selected_table_is_deterministic(): + _, _, bm, _, stats = balanced.build_boundary_stats( + orig_q_seq_len=75600, + orig_kv_seq_len=75600, + block_sizes=_bs(), + band_width=22144, + frame_size=3600, + include_first_frame=True, + ) + a_table, a_active, a_report = balanced.build_selected_boundary_table( + stats=stats, + qtiles=bm.shape[0], + policy="global_balanced", + budget_scale=1.0, + ) + b_table, b_active, b_report = balanced.build_selected_boundary_table( + stats=stats, + qtiles=bm.shape[0], + policy="global_balanced", + budget_scale=1.0, + ) + np.testing.assert_array_equal(a_table, b_table) + np.testing.assert_array_equal(a_active, b_active) + assert a_report == b_report + + +def test_union_tail_preserves_physical_tiles_and_isolates_sequence_tail(): + n = 75600 + bs = _bs() + fm, fa, bm, _, stats = balanced.build_boundary_stats( + orig_q_seq_len=n, + orig_kv_seq_len=n, + block_sizes=bs, + band_width=_band_width(0.15, n), + frame_size=3600, + include_first_frame=True, + ) + sm, sa, _ = balanced.build_selected_boundary_table( + stats=stats, + qtiles=bm.shape[0], + policy="global_balanced", + budget_scale=1.0, + ) + mm, ma, tm, ta = balanced.build_union_tail_tables( + full_table=fm, + full_active=fa, + selected_boundary_table=sm, + selected_boundary_active=sa, + orig_q_seq_len=n, + orig_kv_seq_len=n, + block_sizes=bs, + ) + + original = { + (qi, int(kj)) + for qi in range(fm.shape[0]) + for table, active in ((fm, fa), (sm, sa)) + for kj in table[qi, : int(active[qi])] + } + main = {(qi, int(kj)) for qi in range(mm.shape[0]) for kj in mm[qi, : int(ma[qi])]} + tail = {(qi, int(kj)) for qi in range(tm.shape[0]) for kj in tm[qi, : int(ta[qi])]} + + assert main.isdisjoint(tail) + assert main | tail == original + assert len(main) == 115 + assert len(tail) == 6 + assert all((qi + 1) * bs.block_q > n or (kj + 1) * bs.block_kv > n for qi, kj in tail) + assert all((qi + 1) * bs.block_q <= n and (kj + 1) * bs.block_kv <= n for qi, kj in main) + + +def test_production_builder_is_fixed_global_balanced(): + kernel = dispatch.make_svg_static_range_mha( + block_sizes=_bs(), + orig_q_seq_len=75600, + orig_kv_seq_len=75600, + band_width=14720, + frame_size=3600, + include_first_frame=True, + ) + assert kernel.rounding_policy == "global_balanced" + assert kernel.rounding_budget_scale == 1.0 + assert abs(kernel.rounding_budget["budget_error_fraction"]) < 0.02 + assert kernel.union_main_tiles > 0 + assert kernel.tail_cleanup_tiles > 0 + assert kernel.union_main_tiles + kernel.tail_cleanup_tiles == (kernel.full_tiles + kernel.selected_boundary_tiles) + + +def test_exact_reference_builder_remains_available(): + kernel = static_range.make_svg_exact_static_range_mha( + block_sizes=_bs(), + orig_q_seq_len=75600, + orig_kv_seq_len=75600, + band_width=14720, + frame_size=3600, + include_first_frame=True, + ) + assert kernel.full_tiles > 0 + assert kernel.boundary_tiles > 0 + + +@pytest.mark.parametrize("n", [384, 401], ids=["aligned", "padded"]) +@pytest.mark.parametrize("dense_support", [False, True], ids=["sparse", "density_one"]) +@pytest.mark.parametrize("use_base2_exp", [False, True], ids=["exp", "exp2"]) +def test_production_kernel_matches_rounded_support(n, dense_support, use_base2_exp): + """Check compiled main/tail execution against its real-token support.""" + if jax.default_backend() != "tpu": + pytest.skip("Requires TPU compilation and execution") + block = 128 + padded = math.ceil(n / block) * block + band = n - 1 if dense_support else block + bs = static_range.SVGBlockSizes(block_q=block, block_kv=block, block_kv_compute=block, block_kv_compute_in=block) + support_args = { + "block_sizes": bs, + "orig_q_seq_len": n, + "orig_kv_seq_len": n, + "band_width": band, + "frame_size": 128, + "include_first_frame": True, + } + full, full_active, boundary, _, stats = balanced.build_boundary_stats(**support_args) + selected, selected_active, _ = balanced.build_selected_boundary_table( + stats=stats, qtiles=boundary.shape[0], full_active=full_active, policy="global_balanced", budget_scale=1.0 + ) + mask = np.zeros((n, n), dtype=bool) + for table, active in ((full, full_active), (selected, selected_active)): + for qi in range(table.shape[0]): + for ki in table[qi, : int(active[qi])]: + mask[qi * block : (qi + 1) * block, int(ki) * block : (int(ki) + 1) * block] = True + assert mask.any(axis=1).all() + assert mask.all() if dense_support else not mask.all() + + rng = np.random.default_rng(41) + q, k, v = (jnp.asarray(rng.normal(size=(1, n, 128)), dtype=jnp.bfloat16) for _ in range(3)) + # Match the BF16 input scaling used by dispatch before forming the FP32 oracle. + k = k * (128**-0.5) + q = q * math.log2(math.e) if use_base2_exp else q + logits = jnp.einsum("hqd,hkd->hqk", q.astype(jnp.float32), k.astype(jnp.float32)) + if use_base2_exp: + logits = logits * math.log(2) + weights = jax.nn.softmax(jnp.where(jnp.asarray(mask)[None], logits, -jnp.inf), axis=-1) + expected = jnp.einsum("hqk,hkd->hqd", weights, v.astype(jnp.float32)) + # Large finite padding makes accidental inclusion visible without NaN propagation. + inputs = [jnp.pad(x, ((0, 0), (0, padded - n), (0, 0)), constant_values=64) for x in (q, k, v)] + kernel = dispatch.make_svg_static_range_mha(**support_args, use_base2_exp=use_base2_exp) + assert kernel.union_main_tiles > 0 + assert (kernel.tail_cleanup_tiles > 0) == (n != padded) + actual = jnp.swapaxes(jax.jit(kernel)(*inputs), 1, 2)[:, :n, :].astype(jnp.float32) + actual, expected = np.asarray(actual), np.asarray(expected) + assert np.isfinite(actual).all() + relative_error = np.linalg.norm(actual - expected) / np.linalg.norm(expected) + assert relative_error < 0.01, relative_error + np.testing.assert_allclose(actual, expected, rtol=0, atol=0.05) + + +def test_padding_partial_rejects_unaligned_value_dimension(): + block = 128 + bs = static_range.SVGBlockSizes(block_q=block, block_kv=block, block_kv_compute=block, block_kv_compute_in=block) + kernel = padding.make_padding_partial_from_table( + table_np=np.zeros((1, 1), dtype=np.int32), + active_np=np.ones((1,), dtype=np.int32), + block_sizes=bs, + orig_q_seq_len=127, + orig_kv_seq_len=127, + band_width=128, + frame_size=128, + ) + q = jnp.zeros((1, block, 128), dtype=jnp.bfloat16) + v = jnp.zeros((1, block, padding.NUM_SUBLANES + 1), dtype=jnp.bfloat16) + with pytest.raises(NotImplementedError, match="must be divisible"): + kernel(q, q, v) + + +@pytest.mark.parametrize("n", [128, 401, 1024]) +@pytest.mark.parametrize("include_first_frame", [False, True]) +def test_rounding_keeps_support_for_every_query(n, include_first_frame): + block = 128 + bs = static_range.SVGBlockSizes(block_q=block, block_kv=block, block_kv_compute=block, block_kv_compute_in=block) + args = { + "orig_q_seq_len": n, + "orig_kv_seq_len": n, + "block_sizes": bs, + "band_width": 1, + "frame_size": 1, + "include_first_frame": include_first_frame, + } + fm, fa, bm, _, stats = balanced.build_boundary_stats(**args) + sm, sa, report = balanced.build_selected_boundary_table(stats=stats, qtiles=bm.shape[0], full_active=fa) + assert np.all(fa + sa > 0) + assert report["coverage_tiles_added"] > 0 + actual_pairs = sum(x.real_pairs for x in stats if x.kj in sm[x.qi, : sa[x.qi]]) + assert report["rounded_boundary_pairs"] == actual_pairs + assert report["budget_error_pairs"] == actual_pairs - report["target_boundary_pairs"] + # Constant values must yield ones for every real query, including when the + # global budget is smaller than the cost of one physical tile per row. + if jax.default_backend() != "tpu": + return # Host support and accounting assertions above also run on CPU. + padded = math.ceil(n / block) * block + q = jnp.zeros((1, padded, 128), dtype=jnp.bfloat16) + v = jnp.ones_like(q).at[:, n:, :].set(64) + kernel = dispatch.make_svg_static_range_mha(**args) + actual = np.asarray(jax.jit(kernel)(q, q, v)) + assert np.isfinite(actual).all() + np.testing.assert_allclose(actual, 1.0, rtol=0, atol=0.01) + + +def test_coverage_guard_preserves_already_supported_selection(): + bs = _bs() + for density in (0.65, 0.50, 0.35, 0.20, 0.15): + _, fa, bm, _, stats = balanced.build_boundary_stats( + orig_q_seq_len=75600, + orig_kv_seq_len=75600, + block_sizes=bs, + band_width=_band_width(density), + frame_size=3600, + include_first_frame=True, + ) + old_table, old_active, _ = balanced.build_selected_boundary_table(stats=stats, qtiles=bm.shape[0]) + table, active, report = balanced.build_selected_boundary_table(stats=stats, qtiles=bm.shape[0], full_active=fa) + assert np.all(fa + active > 0) + if np.all(fa + old_active > 0): + np.testing.assert_array_equal(table, old_table) + np.testing.assert_array_equal(active, old_active) + assert report["coverage_tiles_added"] == 0 + + +def test_coverage_guard_rejects_rows_without_candidates(): + with pytest.raises(ValueError, match="no valid attention support"): + balanced.build_selected_boundary_table(stats=[], qtiles=1, full_active=np.zeros(1, dtype=np.int32)) + + +def test_coverage_repair_only_adds_support_to_empty_rows(): + # Row 0 has a full tile, row 1 wins the global budget, and row 2 would + # otherwise be dropped. Repair must not spend additional work on row 0. + stats = [ + balanced.BoundaryTileStat(0, 0, 100, 1), + balanced.BoundaryTileStat(1, 1, 100, 90), + balanced.BoundaryTileStat(2, 2, 100, 1), + ] + table, active, report = balanced.build_selected_boundary_table( + stats=stats, qtiles=3, full_active=np.array([1, 0, 0], dtype=np.int32) + ) + np.testing.assert_array_equal(active, [0, 1, 1]) + assert table[1, 0] == 1 + assert table[2, 0] == 2 + assert report["coverage_tiles_added"] == 1 + assert report["coverage_pairs_added"] == 100 + assert report["rounded_boundary_pairs"] == 200 + assert report["retained_exact_pairs"] == 91 + assert report["dropped_exact_pairs"] == 1 + assert report["budget_error_pairs"] == 108 + # Candidate order must not change the repaired support or accounting. + other_table, other_active, other_report = balanced.build_selected_boundary_table( + stats=list(reversed(stats)), qtiles=3, full_active=np.array([1, 0, 0], dtype=np.int32) + ) + np.testing.assert_array_equal(other_table, table) + np.testing.assert_array_equal(other_active, active) + assert other_report == report + + +def test_coverage_repair_rejects_zero_overlap_candidates(): + with pytest.raises(ValueError, match="query tile 0 has no valid attention support"): + balanced.build_selected_boundary_table( + stats=[balanced.BoundaryTileStat(0, 0, 128 * 128, 0)], + qtiles=1, + full_active=np.zeros(1, dtype=np.int32), + ) diff --git a/src/maxdiffusion/tests/wan/svg_config_propagation_test.py b/src/maxdiffusion/tests/wan/svg_config_propagation_test.py new file mode 100644 index 000000000..6a3cdb087 --- /dev/null +++ b/src/maxdiffusion/tests/wan/svg_config_propagation_test.py @@ -0,0 +1,216 @@ +""" +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. + +Configuration propagation tests for Sparse VideoGen attention. +""" + +from types import SimpleNamespace +import unittest + +from flax import nnx +import jax +import numpy as np +from jax.sharding import Mesh + +from maxdiffusion.pipelines.wan.wan_pipeline import create_sharded_logical_transformer + + +class SVGConfigPropagationTest(unittest.TestCase): + + def test_svg_config_propagation_through_transformer_construction(self): + test_wan_config = { + "num_layers": 2, + "num_attention_heads": 2, + "attention_head_dim": 32, + "in_channels": 16, + "patch_size": (1, 2, 2), + "text_dim": 64, + "freq_dim": 64, + "ffn_dim": 64, + "eps": 1e-6, + } + + devices = np.array(jax.devices()[:1]).reshape((1, 1)) + mesh = Mesh(devices, ("data", "fsdp")) + rngs = nnx.Rngs(0) + + cfg = SimpleNamespace( + use_svg_attention=True, + svg_high_noise_density=0.50, + svg_low_noise_density=0.20, + svg_spatial_density=0.25, + svg_active_start_step=8, + svg_active_end_step=30, + svg_active_start_layer=1, + svg_active_end_layer=40, + precision="DEFAULT", + flash_block_sizes={}, + activations_dtype="bfloat16", + weights_dtype="bfloat16", + attention="dot_product", + remat_policy="none", + names_which_can_be_saved=[], + names_which_can_be_offloaded=[], + flash_min_seq_length=0, + dropout=0.0, + mask_padding_tokens=False, + scan_layers=False, + enable_jax_named_scopes=False, + use_base2_exp=False, + use_experimental_scheduler=False, + logical_axis_rules=(), + model_type="T2V", + model_name="wan2.2", + ) + + m_high = create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=rngs, + config=cfg, + restored_checkpoint={"wan_config": dict(test_wan_config), "wan_state": {}}, + subfolder="transformer", + ) + self.assertTrue(m_high.blocks[0].attn1.use_svg_attention) + self.assertEqual(m_high.blocks[0].attn1.svg_spatial_density, 0.50) + self.assertEqual(m_high.blocks[0].attn1.svg_active_start_step, 8) + self.assertEqual(m_high.blocks[0].attn1.svg_active_end_step, 30) + self.assertEqual(m_high.blocks[0].attn1.svg_active_start_layer, 1) + self.assertEqual(m_high.blocks[0].attn1.svg_active_end_layer, 40) + + m_low = create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=rngs, + config=cfg, + restored_checkpoint={"wan_config": dict(test_wan_config), "wan_state": {}}, + subfolder="transformer_2", + ) + self.assertTrue(m_low.blocks[0].attn1.use_svg_attention) + self.assertEqual(m_low.blocks[0].attn1.svg_spatial_density, 0.20) + + cfg.svg_high_noise_density = None + cfg.svg_low_noise_density = None + for subfolder in ("transformer", "transformer_2"): + fallback = create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=nnx.Rngs(0), + config=cfg, + restored_checkpoint={"wan_config": dict(test_wan_config), "wan_state": {}}, + subfolder=subfolder, + ) + self.assertEqual(fallback.blocks[0].attn1.svg_spatial_density, 0.25) + self.assertTrue(fallback.blocks[0].attn1.use_svg_attention) + + cfg_dense = SimpleNamespace( + use_svg_attention=False, + precision="DEFAULT", + flash_block_sizes={}, + activations_dtype="bfloat16", + weights_dtype="bfloat16", + attention="dot_product", + remat_policy="none", + names_which_can_be_saved=[], + names_which_can_be_offloaded=[], + flash_min_seq_length=0, + dropout=0.0, + mask_padding_tokens=False, + scan_layers=False, + enable_jax_named_scopes=False, + use_base2_exp=False, + use_experimental_scheduler=False, + logical_axis_rules=(), + model_type="T2V", + model_name="wan2.2", + ) + m_dense = create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=rngs, + config=cfg_dense, + restored_checkpoint={"wan_config": dict(test_wan_config), "wan_state": {}}, + subfolder="transformer_2", + ) + self.assertFalse(m_dense.blocks[0].attn1.use_svg_attention) + + def test_svg_block_sizes_are_independent_of_the_dense_block_sizes(self): + """The sparse kernel must be tunable separately from the dense one. + + The two kernels want opposite tilings: the dense ring kernel is tuned for + large kv compute blocks, the sparse kernel for small ones. Before this key + existed the sparse path silently reused the dense sizes, so a tuned sparse + configuration was unreachable from config. + """ + test_wan_config = { + "num_layers": 1, + "num_attention_heads": 2, + "attention_head_dim": 32, + "in_channels": 16, + "patch_size": (1, 2, 2), + "text_dim": 64, + "freq_dim": 64, + "ffn_dim": 64, + "eps": 1e-6, + } + devices = np.array(jax.devices()[:1]).reshape((1, 1)) + mesh = Mesh(devices, ("data", "fsdp")) + sparse_blocks = {"block_q": 3328, "block_kv": 2816, "block_kv_compute": 256, "block_kv_compute_in": 256} + common = { + "use_svg_attention": True, + "svg_spatial_density": 0.25, + "precision": "DEFAULT", + "activations_dtype": "bfloat16", + "weights_dtype": "bfloat16", + "attention": "dot_product", + "remat_policy": "none", + "names_which_can_be_saved": [], + "names_which_can_be_offloaded": [], + "flash_min_seq_length": 0, + "dropout": 0.0, + "mask_padding_tokens": False, + "scan_layers": False, + "enable_jax_named_scopes": False, + "use_base2_exp": False, + "use_experimental_scheduler": False, + "logical_axis_rules": (), + "model_type": "T2V", + "model_name": "wan2.2", + } + + def build(**extra): + return create_sharded_logical_transformer( + devices_array=devices, + mesh=mesh, + rngs=nnx.Rngs(0), + config=SimpleNamespace(**common, **extra), + restored_checkpoint={"wan_config": dict(test_wan_config), "wan_state": {}}, + subfolder="transformer", + ) + + tuned = build(flash_block_sizes={}, svg_flash_block_sizes=sparse_blocks) + self.assertEqual(tuned.blocks[0].attn1.svg_flash_block_sizes, sparse_blocks) + + # An empty override must mean "reuse the dense sizes", not "use {}", which + # would silently fall back to the kernel's own unrelated defaults. + default = build(flash_block_sizes={}, svg_flash_block_sizes={}) + self.assertIsNone(default.blocks[0].attn1.svg_flash_block_sizes) + + absent = build(flash_block_sizes={}) + self.assertIsNone(absent.blocks[0].attn1.svg_flash_block_sizes) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/wan/svg_head_local_test.py b/src/maxdiffusion/tests/wan/svg_head_local_test.py new file mode 100644 index 000000000..630879311 --- /dev/null +++ b/src/maxdiffusion/tests/wan/svg_head_local_test.py @@ -0,0 +1,460 @@ +""" +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. + +Synthetic wiring tests on eight CPU devices; these are not TPU benchmarks. +""" + +import re +from functools import partial + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from maxdiffusion.models.wan.transformers.svg_head_local import exchange_local, inference_only + +from maxdiffusion.models.wan.transformers import svg_attention as svg + + +@pytest.mark.parametrize("routing", ["mixed", "spatial", "temporal"]) +def test_exchange_preserves_head_routes_and_attention(routing): + if len(jax.devices()) != 8 or any(d.platform != "cpu" for d in jax.devices()): + pytest.skip("Requires eight CPU devices") + mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ("data", "context")) + ps = P("data", None, "context", None) + grid = (3, 2, 4) + rng = np.random.default_rng(41) + inputs = tuple( + jax.device_put(rng.normal(size=(4, 8, 24, 4)).astype(np.float32), NamedSharding(mesh, ps)) for _ in range(3) + ) + # Different patterns within and between head blocks and data replicas. + route = np.arange(32).reshape(4, 8) % 3 == 1 + if routing != "mixed": + route[:] = routing == "temporal" + route = jax.device_put(route, NamedSharding(mesh, P("data", None))) + place = partial(svg.svg_placement_permute, token_grid=grid) + restore = partial(svg.svg_placement_unpermute, token_grid=grid) + + def core(q, k, v): + # Position-dependent attention makes missing/wrong placement observable. + score = jnp.einsum("bhid,bhjd->bhij", q, k) * 0.5 + pos = jnp.arange(q.shape[2]) + mask = jnp.abs(pos[:, None] - pos[None, :]) <= 3 + return jnp.einsum("bhij,bhjd->bhid", jax.nn.softmax(jnp.where(mask, score, -jnp.inf)), v) + + def reference(q, k, v, r): + return restore(core(*place(q, k, v, r)), r) + + def candidate(q, k, v, r): + return exchange_local( + q, k, v, r, mesh=mesh, qspec=ps, kvspec=ps, ulysses_axis="context", place=place, restore=restore, core=core + ) + + expected = jax.jit(reference)(*inputs, route) + executable = jax.jit(candidate).lower(*inputs, route).compile() + actual = executable(*inputs, route) + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + assert actual.sharding.is_equivalent_to(NamedSharding(mesh, ps), 4) + assert "all-to-all" in executable.as_text().lower() + if routing == "mixed": + # Sensitivity control: wrong routing must not accidentally pass this test. + wrong = jax.jit(reference)(*inputs, jnp.logical_not(route)) + assert np.max(np.abs(np.asarray(actual - wrong))) > 0.1 + + +def test_reject_head_sharded_inputs(): + with pytest.raises(ValueError, match="unsharded heads"): + exchange_local( + None, + None, + None, + None, + mesh=None, + qspec=P("data", "ulysses", None, None), + kvspec=P("data", "ulysses", None, None), + ulysses_axis="context", + place=None, + restore=None, + core=None, + ) + + +@pytest.mark.parametrize("transform", [jax.grad, lambda f: lambda x: jax.jvp(f, (x,), (jnp.ones_like(x),))]) +def test_inference_rejects_derivatives(transform): + with pytest.raises(NotImplementedError, match="inference only"): + transform(lambda x: inference_only(x).sum())(jnp.ones((2,))) + + +def test_static_inactive_layer_skips_dynamic_step(): + def call(step): + assert svg.is_svg_active(step, 0, start_step=0, end_step=40, start_layer=1, end_layer=40) is False + return step + + assert jax.jit(call)(3) == 3 + + +@pytest.mark.parametrize("base2", [False, True]) +def test_pallas_scratch_matches_dense(monkeypatch, base2): + from jax.experimental import pallas as pl + from maxdiffusion.kernels import custom_svg_attention_dispatch as dispatch + from maxdiffusion.kernels import custom_svg_static_range_attention as kernel + + original = pl.pallas_call + monkeypatch.setattr(pl, "pallas_call", lambda *a, **kw: original(*a, **dict(kw, interpret=True))) + rng = np.random.default_rng(12) + n = 257 + q, k, v = [jnp.asarray(rng.normal(size=(2, 384, 128)).astype(np.float32) * 0.1) for _ in range(3)] + # Padding must not contribute, even when its values dominate real tokens. + k = k.at[:, n:].set(100) + v = v.at[:, n:].set(100) + blocks = kernel.SVGBlockSizes(block_q=128, block_kv=128, block_kv_compute=128, block_kv_compute_in=128) + call = dispatch.make_svg_static_range_mha( + block_sizes=blocks, + orig_q_seq_len=n, + orig_kv_seq_len=n, + band_width=n, + frame_size=1, + use_base2_exp=base2, + ) + # TPU lowers float32 matmuls to bf16 passes by default, which perturbs the + # reference below by ~1e-2 and would make this exactness check test the + # hardware default rather than the kernel. Pin precision instead of loosening + # the tolerance; on CPU this is already the effective behaviour. + with jax.default_matmul_precision("highest"): + actual = call(q * (np.log2(np.e) if base2 else 1), k, v).transpose(0, 2, 1) + scores = jnp.einsum("hid,hjd->hij", q[:, :n], k[:, :n]) + expected = jnp.einsum("hij,hjd->hid", jax.nn.softmax(scores), v[:, :n]) + np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5) + + +# --- End-to-end dispatcher tests ------------------------------------------ +# +# These drive `_apply_attention` itself, so they cover the parts the unit +# tests above cannot: the SVG gate, logical->mesh axis resolution, the +# all_to_all, route sharding, flash padding and the final unpad. Pallas runs +# in interpret mode, so they check semantics only, never performance. + +_AXIS_RULES = ( + ("activation_batch", "data"), + ("activation_length", "context"), + ("activation_heads", None), + ("activation_kv", None), + ("activation_self_attn_heads", None), + ("activation_self_attn_q_length", "context"), + ("activation_self_attn_kv_length", "context"), +) +_GRID = (5, 2, 4) # 40 tokens: not a multiple of the 16-wide q block. +_HEADS, _DIM = 8, 4 # dim < 128 forces head-dim padding inside the kernel. + + +def _axis_names(): + from maxdiffusion.common_types import BATCH, D_KV, SELF_ATTN_HEAD, SELF_ATTN_KV_LENGTH, SELF_ATTN_Q_LENGTH + + return (BATCH, SELF_ATTN_HEAD, SELF_ATTN_Q_LENGTH, D_KV), (BATCH, SELF_ATTN_HEAD, SELF_ATTN_KV_LENGTH, D_KV) + + +def _interpret_pallas(monkeypatch): + from jax.experimental import pallas as pl + + original = pl.pallas_call + monkeypatch.setattr(pl, "pallas_call", lambda *a, **kw: original(*a, **dict(kw, interpret=True))) + + +def _apply_svg( + monkeypatch, + mesh, + qkv, + route, + *, + band_width, + scale, + base2=False, + block=16, + lower=False, + kernel="ulysses_custom", + ulysses_shards=None, + **overrides, +): + """Run the production dispatcher on the head-local SVG path.""" + import flax.linen as nn + from maxdiffusion.models import attention_flax + + monkeypatch.setattr(svg, "svg_profile_temporal_heads", lambda *a, **kw: jnp.asarray(route)) + axis_names_q, axis_names_kv = _axis_names() + cfg = { + "use_svg_attention": True, + "profile_seed": 0, + "profile_query_count": 4, + "band_width": band_width, + "custom_flash_block_sizes": { + "block_q": block, + "block_kv": block, + "block_kv_compute": block, + "block_kv_compute_in": block, + "heads_per_tile": 1, + }, + } + cfg.update(overrides) + + def run(q, k, v): + return attention_flax._apply_attention( + query=q, + key=k, + value=v, + heads=_HEADS, + dim_head=_DIM, + split_head_dim=True, + float32_qk_product=False, + attention_kernel=kernel, + flash_min_seq_length=0, + use_memory_efficient_attention=False, + scale=scale, + dtype=jnp.float32, + mesh=mesh, + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + flash_block_sizes=None, + dpa_layer=None, + use_base2_exp=base2, + ulysses_shards=mesh.shape["context"] if ulysses_shards is None else ulysses_shards, + spatiotemporal_config=cfg, + spatiotemporal_shape=_GRID, + ) + + with mesh, nn.logical_axis_rules(_AXIS_RULES): + if lower: + return jax.jit(run).lower(*qkv).as_text(debug_info=True) + return jax.jit(run)(*qkv) + + +def _dense_reference(q, k, v, scale): + scores = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale + out = jnp.einsum("bhqk,bhkd->bhqd", jax.nn.softmax(scores, axis=-1), v) + return jnp.transpose(out, (0, 2, 1, 3)).reshape(out.shape[0], out.shape[2], -1) + + +def _random_qkv(batch, seed=7): + rng = np.random.default_rng(seed) + n = int(np.prod(_GRID)) + return tuple(jnp.asarray(rng.normal(size=(batch, _HEADS, n, _DIM)).astype(np.float32)) for _ in range(3)) + + +@pytest.mark.parametrize("base2", [False, True]) +def test_dispatcher_full_support_matches_dense(monkeypatch, base2): + """Density 1.0 must reproduce plain softmax attention exactly. + + Full support makes place/restore an exact inverse pair, so any surviving + difference is a real defect: a missing or doubled logit scale, padding + tokens leaking into the softmax, or a mis-stitched all_to_all. + """ + if len(jax.devices()) != 8: + pytest.skip("Requires eight devices") + _interpret_pallas(monkeypatch) + mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ("data", "context")) + q, k, v = _random_qkv(2) + route = np.arange(16).reshape(2, 8) % 3 == 1 + # See test_pallas_scratch_matches_dense: TPU's default bf16 matmul passes + # perturb the dense reference well past this tolerance, so pin the precision + # rather than weaken the exactness requirement. + with jax.default_matmul_precision("highest"): + actual = _apply_svg(monkeypatch, mesh, (q, k, v), route, band_width=int(np.prod(_GRID)), scale=0.37, base2=base2) + expected = _dense_reference(q, k, v, 0.37) + np.testing.assert_allclose(actual, expected, rtol=2e-5, atol=2e-5) + + +@pytest.mark.parametrize("routing", ["mixed", "spatial", "temporal"]) +def test_dispatcher_sparse_is_sharding_invariant(monkeypatch, routing): + """Head-local execution must not depend on how heads are distributed. + + The band is narrow and the sequence is ragged, so this exercises partial + tiles, the union tail and the log-sum-exp merge under a real all_to_all. + """ + if len(jax.devices()) != 8: + pytest.skip("Requires eight devices") + _interpret_pallas(monkeypatch) + route = np.arange(8).reshape(1, 8) % 3 == 1 + if routing != "mixed": + route[:] = routing == "temporal" + qkv = _random_qkv(1, seed=19) + sharded = Mesh(np.array(jax.devices()[:4]).reshape(1, 4), ("data", "context")) + single = Mesh(np.array(jax.devices()[:1]).reshape(1, 1), ("data", "context")) + kwargs = {"band_width": 8, "scale": 0.25} + actual = _apply_svg(monkeypatch, sharded, qkv, route, **kwargs) + expected = _apply_svg(monkeypatch, single, qkv, route, **kwargs) + # A narrow band must actually drop mass; otherwise this test is vacuous. + assert np.max(np.abs(np.asarray(expected - _dense_reference(*qkv, 0.25)))) > 1e-3 + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_dispatcher_batch_fold_matches_unfolded(monkeypatch): + """Folding batch into heads is a pure layout change, not a semantic one.""" + if len(jax.devices()) != 8: + pytest.skip("Requires eight devices") + _interpret_pallas(monkeypatch) + qkv = _random_qkv(2, seed=23) + route = np.arange(16).reshape(2, 8) % 4 < 2 + kwargs = {"band_width": 8, "scale": 0.25} + # data=1 folds (batch>1, nothing shards batch); data=2 cannot. + folded = _apply_svg(monkeypatch, Mesh(np.array(jax.devices()).reshape(1, 8), ("data", "context")), qkv, route, **kwargs) + unfolded = _apply_svg(monkeypatch, Mesh(np.array(jax.devices()).reshape(2, 4), ("data", "context")), qkv, route, **kwargs) + np.testing.assert_allclose(folded, unfolded, rtol=1e-5, atol=1e-5) + # The two batch entries have different routes, so folding must not mix them. + assert np.max(np.abs(np.asarray(folded[0] - folded[1]))) > 1e-3 + + +@pytest.mark.parametrize( + "overrides,error", + [ + ({"attention_kernel": "flash"}, "custom Ulysses attention backend"), + ({"ulysses_attention_chunks": 2}, "chunked Ulysses"), + ({"global_stride": 4}, "external or periodic masks"), + ({"spatiotemporal_shape": None}, "matched self-attention QKV"), + ({"heads_per_tile": 2}, "heads_per_tile=1"), + ], +) +def test_dispatcher_rejects_unsupported_configurations(monkeypatch, overrides, error): + """Unsupported combinations must fail loudly rather than silently run dense.""" + import flax.linen as nn + from maxdiffusion.models import attention_flax + + if len(jax.devices()) != 8: + pytest.skip("Requires eight devices") + monkeypatch.setattr(svg, "svg_profile_temporal_heads", lambda *a, **kw: jnp.zeros((1, _HEADS), bool)) + mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ("data", "context")) + axis_names_q, axis_names_kv = _axis_names() + block = {"block_q": 16, "block_kv": 16, "block_kv_compute": 16, "block_kv_compute_in": 16, "heads_per_tile": 1} + if "heads_per_tile" in overrides: + block["heads_per_tile"] = overrides.pop("heads_per_tile") + cfg = { + "use_svg_attention": True, + "profile_seed": 0, + "profile_query_count": 4, + "band_width": 8, + "custom_flash_block_sizes": block, + } + cfg.update({key: overrides.pop(key) for key in ("global_stride",) if key in overrides}) + kwargs = { + "heads": _HEADS, + "dim_head": _DIM, + "split_head_dim": True, + "float32_qk_product": False, + "attention_kernel": "ulysses_custom", + "flash_min_seq_length": 0, + "use_memory_efficient_attention": False, + "scale": 0.25, + "dtype": jnp.float32, + "mesh": mesh, + "axis_names_q": axis_names_q, + "axis_names_kv": axis_names_kv, + "flash_block_sizes": None, + "dpa_layer": None, + "ulysses_shards": 4, + "spatiotemporal_config": cfg, + "spatiotemporal_shape": _GRID, + } + kwargs.update(overrides) + q, k, v = _random_qkv(1) + with mesh, nn.logical_axis_rules(_AXIS_RULES): + with pytest.raises(ValueError, match=error): + jax.eval_shape(lambda a, b, c: attention_flax._apply_attention(query=a, key=b, value=c, **kwargs), q, k, v) + + +def test_realized_density_is_recorded_in_profile_scope(monkeypatch): + """A profile must show the realized density, not just that SVG was entered. + + Density is not the requested fraction: tile rounding, the anchor and ragged + edges all move it. Recording the executed tile count is what makes an XProf + trace sufficient to prove the sparse path really ran sparsely. + """ + if len(jax.devices()) != 8: + pytest.skip("Requires eight devices") + _interpret_pallas(monkeypatch) + mesh = Mesh(np.array(jax.devices()[:4]).reshape(1, 4), ("data", "context")) + qkv = _random_qkv(1, seed=31) + route = np.zeros((1, _HEADS), bool) + sparse = _apply_svg(monkeypatch, mesh, qkv, route, band_width=8, scale=0.25, lower=True) + dense = _apply_svg(monkeypatch, mesh, qkv, route, band_width=int(np.prod(_GRID)), scale=0.25, lower=True) + pattern = r"svg_kernel_c_tiles(\d+)of(\d+)_d[0-9.]+" + (sparse_exec, total), *_ = re.findall(pattern, sparse) or [(None, None)] + (dense_exec, _), *_ = re.findall(pattern, dense) or [(None, None)] + assert sparse_exec and dense_exec, "realized density missing from the profile scope" + assert int(dense_exec) == int(total), "full support must execute every tile" + assert int(sparse_exec) < int(dense_exec), "a narrow band must execute fewer tiles" + + +def test_dynamic_step_index_stays_traced(): + """A traced step index must yield a traced predicate, not a Python bool. + + The step index cannot be static: it is a non-static argument of the jitted + transformer pass, so only the layer dimension can short-circuit statically. + """ + captured = {} + + def call(step): + captured["active"] = svg.is_svg_active(step, 3, start_step=0, end_step=40, start_layer=1, end_layer=40) + return step + + jax.jit(call)(3) + assert isinstance(captured["active"], jax.Array) and not isinstance(captured["active"], bool) + + +@pytest.mark.parametrize("kernel", ["ulysses_ring_custom", "ulysses_ring_custom_fixed_m"]) +def test_svg_runs_under_a_ring_configured_dense_backend(monkeypatch, kernel): + """SVG must not force the dense arm to give up its ring split. + + The validated recipe pairs a Ring2/Ulysses2 dense arm with pure-Ulysses + sparse attention. `ulysses_shards` configures only the dense arm, which + re-meshes the context axis inside its own shard_map; head-local SVG always + exchanges over the whole context axis. A gate demanding + `ulysses_shards == context` rejected that recipe outright, and because + lax.cond traces both branches it did so even on dense-scheduled steps. + """ + if len(jax.devices()) != 8: + pytest.skip("Requires eight devices") + _interpret_pallas(monkeypatch) + mesh = Mesh(np.array(jax.devices()[:4]).reshape(1, 4), ("data", "context")) + qkv = _random_qkv(1, seed=17) + route = np.arange(_HEADS).reshape(1, _HEADS) % 2 == 0 + common = {"band_width": 8, "scale": 0.31, "kernel": kernel} + ring2 = _apply_svg(monkeypatch, mesh, qkv, route, ulysses_shards=2, **common) + ring1 = _apply_svg(monkeypatch, mesh, qkv, route, ulysses_shards=4, **common) + # Same sparse result either way: proof that the sparse arm ignores the dense + # arm's ring/ulysses split rather than merely tolerating it. + np.testing.assert_array_equal(np.asarray(ring2), np.asarray(ring1)) + reference = _apply_svg(monkeypatch, mesh, qkv, route, band_width=8, scale=0.31, kernel="ulysses_custom") + np.testing.assert_allclose(np.asarray(ring2), np.asarray(reference), rtol=2e-5, atol=2e-5) + + +def test_svg_block_sizes_override_the_dense_tiling(monkeypatch): + """A configured SVG tiling must reach the kernel and change its geometry.""" + if len(jax.devices()) != 8: + pytest.skip("Requires eight devices") + _interpret_pallas(monkeypatch) + mesh = Mesh(np.array(jax.devices()[:4]).reshape(1, 4), ("data", "context")) + qkv = _random_qkv(1, seed=23) + route = np.zeros((1, _HEADS), bool) + pattern = r"svg_kernel_c_tiles(\d+)of(\d+)_d[0-9.]+" + + def tiles(block): + text = _apply_svg(monkeypatch, mesh, qkv, route, band_width=8, scale=0.25, block=block, lower=True) + found = re.findall(pattern, text) + assert found, "sparse kernel did not run" + return tuple(int(x) for x in found[0]) + + coarse_exec, coarse_total = tiles(16) + fine_exec, fine_total = tiles(8) + assert fine_total > coarse_total, "a smaller tile must produce more tiles overall" + assert fine_exec != coarse_exec, "the configured tiling must change what the kernel executes" diff --git a/src/maxdiffusion/tests/wan/wan_pipeline_signature_test.py b/src/maxdiffusion/tests/wan/wan_pipeline_signature_test.py new file mode 100644 index 000000000..71267a250 --- /dev/null +++ b/src/maxdiffusion/tests/wan/wan_pipeline_signature_test.py @@ -0,0 +1,58 @@ +""" +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. + +Check that Wan pipeline calls use keywords accepted by local helpers. +""" + +import ast +import os + +import pytest + +PIPELINE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "pipelines", "wan") + + +def _module_paths(): + return sorted( + os.path.join(PIPELINE_DIR, name) for name in os.listdir(PIPELINE_DIR) if name.endswith(".py") and name != "__init__.py" + ) + + +def _accepted_keywords(func: ast.FunctionDef): + spec = func.args + if spec.kwarg is not None: + return None # **kwargs accepts anything + names = {a.arg for a in spec.args} | {a.arg for a in spec.kwonlyargs} | {a.arg for a in spec.posonlyargs} + return names + + +@pytest.mark.parametrize("path", _module_paths(), ids=os.path.basename) +def test_pipeline_keyword_arguments_match_local_definitions(path): + tree = ast.parse(open(path, encoding="utf-8").read(), filename=path) + defs = {node.name: node for node in tree.body if isinstance(node, ast.FunctionDef)} + problems = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): + continue + target = defs.get(node.func.id) + if target is None: + continue + accepted = _accepted_keywords(target) + if accepted is None: + continue + for kw in node.keywords: + if kw.arg is not None and kw.arg not in accepted: + problems.append(f"{os.path.basename(path)}:{node.lineno}: {node.func.id}() does not accept '{kw.arg}'") + assert not problems, "\n".join(problems)