diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f05667986f11..577729a1e34f 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -395,6 +395,8 @@ title: SanaTransformer2DModel - local: api/models/sana_video_transformer3d title: SanaVideoTransformer3DModel + - local: api/models/sana_wm_transformer3d + title: SanaWMTransformer3DModel - local: api/models/sd3_transformer2d title: SD3Transformer2DModel - local: api/models/skyreels_v2_transformer_3d @@ -629,6 +631,8 @@ title: Sana Sprint - local: api/pipelines/sana_video title: Sana Video + - local: api/pipelines/sana_wm + title: SANA-WM - local: api/pipelines/shap_e title: Shap-E - local: api/pipelines/stable_cascade diff --git a/docs/source/en/api/models/sana_wm_transformer3d.md b/docs/source/en/api/models/sana_wm_transformer3d.md new file mode 100644 index 000000000000..12392aba1739 --- /dev/null +++ b/docs/source/en/api/models/sana_wm_transformer3d.md @@ -0,0 +1,46 @@ + + +# SanaWMTransformer3DModel + +A 3D Diffusion Transformer (1.6B parameters) for camera-controlled image-to-video generation, used as the stage-1 +sampler of [`SanaWMPipeline`]. The transformer combines: + +* a bidirectional GDN-Triton linear-attention main branch (depth 20, hidden 2240, 20 heads), +* a UCPE (Unified Camera Pose Embedding) camera-control branch that consumes a raymap + Plücker representation of + the requested trajectory, and +* a Wan-style 3D rotary position embedding plus periodic softmax-attention blocks injected every `softmax_every_n` + layers. + +The state-dict layout matches the public SANA-WM release one-to-one — the diffusers wrapper places the inner DiT +under a `_inner.` prefix. See [`SanaWMTransformer3DModel.add_inner_prefix`] for the helper used by the conversion +script. + +The model can be loaded with: + +```python +import torch +from diffusers import SanaWMTransformer3DModel + +transformer = SanaWMTransformer3DModel.from_pretrained( + "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", + subfolder="transformer", + torch_dtype=torch.bfloat16, +) +``` + +## SanaWMTransformer3DModel + +[[autodoc]] SanaWMTransformer3DModel + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput diff --git a/docs/source/en/api/pipelines/sana_wm.md b/docs/source/en/api/pipelines/sana_wm.md new file mode 100644 index 000000000000..19098f4e3646 --- /dev/null +++ b/docs/source/en/api/pipelines/sana_wm.md @@ -0,0 +1,119 @@ + + +# SANA-WM + +SANA-WM is a camera-controlled image-to-video world model built on top of SANA. Given a first-frame image, a text +prompt, and a camera trajectory (either explicit `c2w` poses or a WASD/IJKL action string), it generates a video +whose motion follows the requested camera path. + +Inference runs in two stages: + +1. **Stage 1 — SANA-WM DiT.** A 1.6B-parameter bidirectional DiT with GDN-Triton linear attention and a UCPE + camera-control branch. Sampling uses an LTX-style flow-matching Euler scheduler with per-token timesteps; the + first latent frame is the conditioning anchor. +2. **Stage 2 — LTX-2 refiner (optional).** A sink-bidirectional Euler refiner ([`SanaWMLTX2Refiner`]) that wraps + diffusers' own `LTX2VideoTransformer3DModel` + `LTX2TextConnectors` and a Gemma-3 text encoder, run for 3 + distilled sigma steps. + +Both stages decode through the [`AutoencoderKLLTX2Video`] VAE. + +Available models: + +| Model | Recommended dtype | +|:-----:|:-----------------:| +| [`Efficient-Large-Model/SANA-WM_bidirectional-diffusers`](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional-diffusers) | `torch.bfloat16` | + +> [!TIP] +> SANA-WM is trained at a fixed 704×1280 resolution. The recommended dtype is for the transformer weights — keep +> the text encoder in `torch.bfloat16` and the VAE in `torch.float32` for best numerics. The pipeline expects +> camera intrinsics `[fx, fy, cx, cy]` in *original-image* pixel coordinates; the resize-and-center-crop transform +> is applied internally. + +## Inference + +```python +import torch +from PIL import Image + +from diffusers import SanaWMPipeline +from diffusers.utils import export_to_video + +pipe = SanaWMPipeline.from_pretrained( + "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", + torch_dtype=torch.bfloat16, +) +pipe.enable_model_cpu_offload() # ~45 GB of weights — offload between stages + +output = pipe( + image=Image.open("input.png").convert("RGB"), + prompt="A car driving across a vast desert plain at golden hour.", + action="w-80,jw-40,w-40", # WASD-style action DSL: forward 80f, jump+forward 40f, forward 40f + intrinsics=[800.0, 800.0, 845.0, 464.0], # fx, fy, cx, cy in original-image pixels + num_frames=161, + num_inference_steps=60, + guidance_scale=5.0, + seed=42, +) +export_to_video(list(output.frames), "sana_wm.mp4", fps=16) +``` + +Pass `action=None` and supply your own `c2w` poses (`(F, 4, 4)` numpy array) to drive the camera trajectory +explicitly. Set `use_refiner=False` to skip stage 2. + +If you don't have camera intrinsics, [`pi3-vision`](https://github.com/OliverSFAC/pi3-vision) can estimate them +from a single frame: + +```python +from diffusers.pipelines.sana_wm.cam_utils import estimate_intrinsics_with_pi3x +intrinsics = estimate_intrinsics_with_pi3x(image) # `pip install pi3-vision` +``` + +## Converting the released checkpoint + +If you have the source SANA-WM release (not the pre-converted diffusers snapshot), run the conversion script once: + +```bash +python scripts/sana_wm/convert_sana_wm_to_diffusers.py \ + --src Efficient-Large-Model/SANA-WM_bidirectional \ + --dst ./SANA-WM_bidirectional-diffusers +``` + +Then load from the local path as usual. + +## Components + +- `tokenizer` — [`GemmaTokenizerFast`] +- `text_encoder` — Gemma-2 (returns decoder hidden states) +- `vae` — [`AutoencoderKLLTX2Video`] (LTX-2, spatial ×32 / temporal ×8) +- `transformer` — [`SanaWMTransformer3DModel`], 1.6B-parameter bidirectional DiT +- `scheduler` — [`FlowMatchEulerDiscreteScheduler`] +- `refiner` (optional) — [`SanaWMLTX2Refiner`], wraps `LTX2VideoTransformer3DModel`, `LTX2TextConnectors`, and a + Gemma-3 text encoder + +## SanaWMPipeline + +[[autodoc]] SanaWMPipeline + - all + - __call__ + +## SanaWMLTX2Refiner + +The optional LTX-2 stage-2 refiner is itself a [`DiffusionPipeline`]. [`SanaWMPipeline`] runs it automatically when +`use_refiner=True`, but it can also be used standalone on stage-1 latents. + +[[autodoc]] SanaWMLTX2Refiner + - all + - __call__ + +## SanaWMPipelineOutput + +[[autodoc]] pipelines.sana_wm.pipeline_output.SanaWMPipelineOutput diff --git a/scripts/sana_wm/convert_sana_wm_to_diffusers.py b/scripts/sana_wm/convert_sana_wm_to_diffusers.py new file mode 100644 index 000000000000..8714625e083d --- /dev/null +++ b/scripts/sana_wm/convert_sana_wm_to_diffusers.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. +"""Convert the public SANA-WM release into a diffusers-loadable directory. + +Reads the ``Efficient-Large-Model/SANA-WM_bidirectional`` HF repo (or a local +mirror) and writes a directory ready for ``SanaWMPipeline.from_pretrained(path)``: + + / + ├── model_index.json + ├── tokenizer/ + ├── text_encoder/ + ├── vae/ + ├── transformer/ + ├── scheduler/ + └── refiner/ + ├── transformer/ + ├── connectors/ + ├── text_encoder/ + └── tokenizer/ + +Usage: + python scripts/sana_wm/convert_sana_wm_to_diffusers.py \\ + --src Efficient-Large-Model/SANA-WM_bidirectional \\ + --dst /path/to/SANA-WM_bidirectional-diffusers \\ + [--no-refiner] + +The output is local-only; no upload to the Hub. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from pathlib import Path + +import torch +from huggingface_hub import snapshot_download + + +def _copy_subdir(src: Path, dst: Path) -> None: + if dst.exists(): + shutil.rmtree(dst) + shutil.copytree(src, dst, symlinks=False) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--src", default="Efficient-Large-Model/SANA-WM_bidirectional", help="HF repo or local dir") + parser.add_argument("--dst", required=True, type=Path, help="Output directory") + parser.add_argument("--no-refiner", action="store_true", help="Skip refiner export") + parser.add_argument( + "--torch-dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"], help="Weight dtype" + ) + args = parser.parse_args() + + torch_dtype = getattr(torch, args.torch_dtype) + dst: Path = args.dst.absolute() + dst.mkdir(parents=True, exist_ok=True) + + # Resolve the source on disk (snapshot_download for HF repos, otherwise use as-is). + src_path = Path(args.src) + if not src_path.is_dir(): + print(f"[convert] snapshot_download({args.src}) …") + src_path = Path(snapshot_download(args.src)) + print(f"[convert] source: {src_path}") + + # 1. VAE (already diffusers format under /vae). + print("[convert] vae …") + _copy_subdir(src_path / "vae", dst / "vae") + + # 2. Tokenizer + text encoder — fetch via the configured Gemma-2 repo. + # We save the full ``Gemma2ForCausalLM``; the pipeline grabs the decoder + # at runtime via ``self.text_encoder.model(...)``. This matches the sana + # inference recipe of ``AutoModelForCausalLM.from_pretrained(...).get_decoder()`` + # and avoids subtle state-dict prefix differences when saving just the + # decoder submodule. + print("[convert] tokenizer + text_encoder (gemma-2-2b-it) …") + from transformers import AutoModelForCausalLM, AutoTokenizer + + gemma_repo = "Efficient-Large-Model/gemma-2-2b-it" + tokenizer = AutoTokenizer.from_pretrained(gemma_repo) + tokenizer.padding_side = "right" + tokenizer.save_pretrained(dst / "tokenizer") + text_encoder = AutoModelForCausalLM.from_pretrained(gemma_repo, torch_dtype=torch_dtype) + text_encoder.save_pretrained(dst / "text_encoder") + del text_encoder + + # 3. Transformer (SanaWMTransformer3DModel) — load the public DiT, save in diffusers format. + print("[convert] transformer (SanaWMTransformer3DModel) …") + from diffusers import SanaWMTransformer3DModel + + transformer = SanaWMTransformer3DModel().to(torch_dtype).eval() + dit_ckpt = src_path / "dit" / "sana_wm_1600m_720p.safetensors" + if not dit_ckpt.is_file(): + raise FileNotFoundError(f"DiT checkpoint not found at {dit_ckpt}") + from safetensors.torch import load_file + + sd = load_file(str(dit_ckpt)) + sd.pop("pos_embed", None) # unused at inference (wan_rope is computed on-the-fly) + # The public release keys (``blocks.0...``) load directly into the merged + # SanaWMTransformer3DModel — no ``_inner.`` prefix anymore. + missing, unexpected = transformer.load_state_dict(sd, strict=False) + if missing: + missing_nontrivial = [k for k in missing if not k.endswith(".pos_embed")] + if missing_nontrivial: + print(f" missing keys: {missing_nontrivial[:10]}{' …' if len(missing_nontrivial) > 10 else ''}") + if unexpected: + print(f" unexpected keys: {unexpected[:10]}{' …' if len(unexpected) > 10 else ''}") + transformer.save_pretrained(dst / "transformer") + del transformer, sd + + # 4. Scheduler — FlowMatchEulerDiscreteScheduler config. + print("[convert] scheduler …") + from diffusers import FlowMatchEulerDiscreteScheduler + + FlowMatchEulerDiscreteScheduler(shift=9.8).save_pretrained(dst / "scheduler") + + # 5. Refiner (LTX-2): now a standalone DiffusionPipeline saved in the + # ``refiner/`` subfolder with its own ``model_index.json``. Copy the + # LTX-2 sub-model folders as-is, split out a ``tokenizer/`` folder, add a + # ``scheduler/`` (FlowMatchEulerDiscreteScheduler), and write the manifest. + if not args.no_refiner: + print("[convert] refiner …") + from transformers import AutoTokenizer + + refiner_src = src_path / "refiner" + refiner_dst = dst / "refiner" + refiner_dst.mkdir(exist_ok=True) + for sub in ("transformer", "connectors", "text_encoder"): + if (refiner_src / sub).is_dir(): + _copy_subdir(refiner_src / sub, refiner_dst / sub) + + # Tokenizer lives co-located with the Gemma-3 text encoder in the release; + # re-save it into its own subfolder so it registers as a pipeline component. + refiner_tokenizer = AutoTokenizer.from_pretrained(refiner_src / "text_encoder") + refiner_tokenizer.save_pretrained(refiner_dst / "tokenizer") + + # Scheduler carries the distilled sigma schedule; shift=1.0 leaves the + # explicit sigmas passed at inference time unmodified. + FlowMatchEulerDiscreteScheduler(shift=1.0).save_pretrained(refiner_dst / "scheduler") + + refiner_index = { + "_class_name": "SanaWMLTX2Refiner", + "_diffusers_version": "0.38.0", + "transformer": ["diffusers", "LTX2VideoTransformer3DModel"], + # LTX2TextConnectors lives in diffusers.pipelines.ltx2 (not top-level), + # so the loader resolves it via the pipeline-module path ("ltx2", ...). + "connectors": ["ltx2", "LTX2TextConnectors"], + "tokenizer": ["transformers", type(refiner_tokenizer).__name__], + "text_encoder": ["transformers", "Gemma3ForConditionalGeneration"], + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "text_max_sequence_length": 1024, + } + (refiner_dst / "model_index.json").write_text(json.dumps(refiner_index, indent=2)) + + # 6. model_index.json — the top-level diffusers manifest. + print("[convert] model_index.json …") + index = { + "_class_name": "SanaWMPipeline", + "_diffusers_version": "0.38.0", + "tokenizer": ["transformers", "GemmaTokenizerFast"], + "text_encoder": ["transformers", "Gemma2ForCausalLM"], + "vae": ["diffusers", "AutoencoderKLLTX2Video"], + "transformer": ["diffusers", "SanaWMTransformer3DModel"], + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + } + if not args.no_refiner: + index["refiner"] = ["diffusers", "SanaWMLTX2Refiner"] + (dst / "model_index.json").write_text(json.dumps(index, indent=2)) + + print(f"[convert] done — wrote {dst}") + + +if __name__ == "__main__": + main() diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b7d79b8ee97d..1bc0febb04ff 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -337,6 +337,7 @@ "SanaControlNetModel", "SanaTransformer2DModel", "SanaVideoTransformer3DModel", + "SanaWMTransformer3DModel", "SD3ControlNetModel", "SD3MultiControlNetModel", "SD3Transformer2DModel", @@ -776,6 +777,9 @@ "SanaSprintPipeline", "SanaVideoPipeline", "SanaVideoPipeline", + "SanaWMLTX2Refiner", + "SanaWMPipeline", + "SanaWMPipelineOutput", "SemanticStableDiffusionPipeline", "ShapEImg2ImgPipeline", "ShapEPipeline", @@ -1208,6 +1212,7 @@ SanaControlNetModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, + SanaWMTransformer3DModel, SD3ControlNetModel, SD3MultiControlNetModel, SD3Transformer2DModel, @@ -1621,6 +1626,9 @@ SanaSprintImg2ImgPipeline, SanaSprintPipeline, SanaVideoPipeline, + SanaWMLTX2Refiner, + SanaWMPipeline, + SanaWMPipelineOutput, SemanticStableDiffusionPipeline, ShapEImg2ImgPipeline, ShapEPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 8ba17d896434..d46d35f64e24 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -145,6 +145,7 @@ _import_structure["transformers.transformer_prx"] = ["PRXTransformer2DModel"] _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] _import_structure["transformers.transformer_sana_video"] = ["SanaVideoTransformer3DModel"] + _import_structure["transformers.transformer_sana_wm"] = ["SanaWMTransformer3DModel"] _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] _import_structure["transformers.transformer_stable_audio3"] = ["StableAudio3DiTModel"] @@ -290,6 +291,7 @@ QwenImageTransformer2DModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, + SanaWMTransformer3DModel, SD3Transformer2DModel, SkyReelsV2Transformer3DModel, StableAudio3DiTModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 0e167812ad88..89efe4afc2b0 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -61,6 +61,7 @@ from .transformer_prx import PRXTransformer2DModel from .transformer_qwenimage import QwenImageTransformer2DModel from .transformer_sana_video import SanaVideoTransformer3DModel + from .transformer_sana_wm import SanaWMTransformer3DModel from .transformer_sd3 import SD3Transformer2DModel from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel from .transformer_stable_audio3 import StableAudio3DiTModel diff --git a/src/diffusers/models/transformers/transformer_sana_wm.py b/src/diffusers/models/transformers/transformer_sana_wm.py new file mode 100644 index 000000000000..801b3f81ea3a --- /dev/null +++ b/src/diffusers/models/transformers/transformer_sana_wm.py @@ -0,0 +1,4598 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. +# +# This file is modified from https://github.com/PixArt-alpha/PixArt-sigma + +from __future__ import annotations + +import copy +import math +from collections.abc import Iterable +from copy import deepcopy +from functools import lru_cache, partial +from itertools import repeat as _itertools_repeat +from typing import Any, Callable, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import is_timm_available, logging +from ..embeddings import get_1d_rotary_pos_embed +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from .transformer_sana_wm_kernels import ( + _prepare_ucpe_rope_tables, + _process_camera_conditions_raymats_only, + cam_prep_func, + cam_scan_bidi_chunkwise, + compute_fov_from_fx_xi, + compute_up_lat_map, + fused_bigdn_func, + fused_qk_inv_rms, + prepare_rope_tables, + ucm_unproject_grid_fov, + world_to_ray_mats, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +_CAN_USE_TIMM = is_timm_available() + +if _CAN_USE_TIMM: + from timm.models.layers import DropPath + from timm.models.vision_transformer import Attention as Attention_ + from timm.models.vision_transformer import Mlp +else: + # Several layers below subclass these, so they must exist as classes at module + # import time — this module is imported eagerly by `diffusers.models`. The + # placeholder defers the error to construction time, keeping `import diffusers` + # working without `timm` installed. + class _TimmPlaceholder(nn.Module): + def __init__(self, *args, **kwargs): + raise ImportError("`timm` is required to run SANA-WM. Install it with `pip install timm`.") + + DropPath = Attention_ = Mlp = _TimmPlaceholder + + +class ShortConvolution(nn.Module): + """Depthwise causal 1D convolution over the temporal axis. + + SANA-WM's GDN attention applies a short causal depthwise conv to Q/K/V before the linear-attention kernel. This is + a self-contained PyTorch implementation of the `fla.modules.ShortConvolution` layer the reference implementation + used (with `activation=None`), so the model needs no `fla-core` dependency and can be built on any device. + + Args: + hidden_size (`int`): Number of channels (the conv is depthwise, one group per channel). + kernel_size (`int`): Temporal kernel width. + bias (`bool`, defaults to `False`): Whether to add a per-channel bias. + """ + + def __init__(self, hidden_size: int, kernel_size: int, bias: bool = False, activation: str | None = None) -> None: + super().__init__() + if activation is not None: + raise ValueError(f"SANA-WM only uses `activation=None` short convolutions, got {activation!r}.") + self.hidden_size = hidden_size + self.kernel_size = kernel_size + # Same parameter layout as the reference implementation: (C, 1, K). + self.weight = nn.Parameter(torch.zeros(hidden_size, 1, kernel_size)) + self.bias = nn.Parameter(torch.zeros(hidden_size)) if bias else None + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: + """Apply the causal conv. + + Args: + x (`torch.Tensor`): Input of shape `(batch, seq_len, hidden_size)`. + + Returns: + `tuple[torch.Tensor, None]`: `(output, cache)`; the cache slot is kept for signature compatibility with the + reference implementation but is unused for the bidirectional (non-streaming) forward SANA-WM runs. + """ + seq_len = x.shape[1] + # Left-pad by (K - 1) and drop the tail so output[t] only sees inputs <= t. + y = F.conv1d( + x.transpose(1, 2), + self.weight.to(x.dtype), + None if self.bias is None else self.bias.to(x.dtype), + groups=self.hidden_size, + padding=self.kernel_size - 1, + )[..., :seq_len] + return y.transpose(1, 2), None + + +# ============================================================================ +# Helpers (norms / acts / chunk / weight utilities) + + +# ============================================================================ + +# register activation function here +# name: module, kwargs with default values +REGISTERED_ACT_DICT: dict[str, tuple[type, dict[str, Any]]] = { + "relu": (nn.ReLU, {"inplace": True}), + "relu6": (nn.ReLU6, {"inplace": True}), + "hswish": (nn.Hardswish, {"inplace": True}), + "hsigmoid": (nn.Hardsigmoid, {"inplace": True}), + "swish": (nn.SiLU, {"inplace": True}), + "silu": (nn.SiLU, {"inplace": True}), + "tanh": (nn.Tanh, {}), + "sigmoid": (nn.Sigmoid, {}), + "gelu": (nn.GELU, {"approximate": "tanh"}), + "mish": (nn.Mish, {"inplace": True}), + "identity": (nn.Identity, {}), +} + + +def build_act(name: Optional[str], **kwargs) -> Optional[nn.Module]: + if name in REGISTERED_ACT_DICT: + act_cls, default_args = copy.deepcopy(REGISTERED_ACT_DICT[name]) + for key in default_args: + if key in kwargs: + default_args[key] = kwargs[key] + return act_cls(**default_args) + elif name is None or name.lower() == "none": + return None + else: + raise ValueError(f"do not support: {name}") + + +# register normalization function here +# name: module, kwargs with default values +REGISTERED_NORMALIZATION_DICT: dict[str, tuple[type, dict[str, Any]]] = { + "bn2d": (nn.BatchNorm2d, {"num_features": None, "eps": 1e-5, "momentum": 0.1, "affine": True}), + "syncbn": (nn.SyncBatchNorm, {"num_features": None, "eps": 1e-5, "momentum": 0.1, "affine": True}), + "ln": (nn.LayerNorm, {"normalized_shape": None, "eps": 1e-5, "elementwise_affine": True}), +} + + +def build_norm(name="bn2d", num_features=None, affine=True, **kwargs) -> Optional[nn.Module]: + if name == "ln": + kwargs["normalized_shape"] = num_features + kwargs["elementwise_affine"] = affine + else: + kwargs["num_features"] = num_features + kwargs["affine"] = affine + if name in REGISTERED_NORMALIZATION_DICT: + norm_cls, default_args = copy.deepcopy(REGISTERED_NORMALIZATION_DICT[name]) + for key in default_args: + if key in kwargs: + default_args[key] = kwargs[key] + return norm_cls(**default_args) + elif name is None or name.lower() == "none": + return None + else: + raise ValueError("do not support: %s" % name) + + +class RMSNorm(torch.nn.Module): + def __init__(self, dim: int, scale_factor=1.0, eps: float = 1e-6, norm_dim: int = -1): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + norm_dim (int, optional): The dimension to normalize over. Default is -1 (last dimension). + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + norm_dim (int): The dimension to normalize over. + + """ + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim) * scale_factor) + self.norm_dim = norm_dim + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(self.norm_dim, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + ndim = x.dim() + weight_shape = [1] * ndim + weight_shape[self.norm_dim] = -1 + weight = self.weight.view(*weight_shape) + return (weight * self._norm(x.float())).type_as(x) + + +def _ntuple(n): + def parse(x): + if isinstance(x, Iterable) and not isinstance(x, str): + return x + return tuple(_itertools_repeat(x, n)) + + return parse + + +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) + + +def val2list(x: list or tuple or any, repeat_time=1) -> list: # type: ignore + """Repeat `val` for `repeat_time` times and return the list or val if list/tuple.""" + if isinstance(x, (list, tuple)): + return list(x) + return [x for _ in range(repeat_time)] + + +def val2tuple(x: list or tuple or any, min_len: int = 1, idx_repeat: int = -1) -> tuple: # type: ignore + """Return tuple with min_len by repeating element at idx_repeat.""" + # convert to list first + x = val2list(x) + + # repeat elements if necessary + if len(x) > 0: + x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))] + + return tuple(x) + + +def get_same_padding(kernel_size: int or tuple[int, ...]) -> int or tuple[int, ...]: + if isinstance(kernel_size, tuple): + return tuple([get_same_padding(ks) for ks in kernel_size]) + else: + assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number" + return kernel_size // 2 + + +def chunk_index_from_chunk_size( + T: int, + chunk_size: int, + strategy: str = "uniform", +) -> List[int]: + """Convert chunk_size to chunk_index list with a split strategy. + + Args: + T: Number of latent frames. + chunk_size: Base chunk size for the temporal dimension. + strategy: Chunk split strategy. Supported values: + - "uniform" (default): uniform chunks with optional remainder Example: T=21, chunk_size=4 → + [0,4,8,12,16,20] → sizes [4,4,4,4,4,1] + - "first_frame": first chunk is 1 frame, then uniform chunk_size Example: T=21, chunk_size=4 → + [0,1,5,9,13,17] → sizes [1,4,4,4,4,4] + - "first_plus_one": first chunk is chunk_size + 1, then uniform chunk_size Example: T=21, chunk_size=4 → + [0,5,9,13,17] → sizes [5,4,4,4,4] + + Returns: + List of chunk start indices (not including the final T). + + Raises: + ValueError: If chunk_size or T are invalid, or strategy is unknown. + """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}.") + if T <= 0: + raise ValueError(f"T must be > 0, got {T}.") + + if strategy is None: + strategy = "uniform" + strategy = str(strategy).lower() + + if strategy in ("uniform", "default"): + indices = list(range(0, T, chunk_size)) + # Absorb small remainder into last chunk to avoid degenerate chunks + # (e.g., causal_conv1d crashes on length=1 sequences). + if len(indices) > 1 and (T - indices[-1]) < chunk_size: + indices.pop() + return indices + + if strategy in ("first_frame", "first_frame_alone", "first_frame_only"): + if T <= 1: + return [0] + indices = [0] + list(range(1, T, chunk_size)) + if len(indices) > 2 and (T - indices[-1]) < chunk_size: + indices.pop() + return indices + + if strategy in ("first_plus_one", "first_chunk_plus_one"): + if T <= chunk_size + 1: + return [0] + indices = [0] + list(range(chunk_size + 1, T, chunk_size)) + # Absorb small remainder into last chunk to avoid degenerate chunks + # (e.g., T_latent=41 with chunk_size=3 → last chunk would be 1 frame, + # which crashes causal_conv1d). Merge it into the previous chunk instead. + if len(indices) > 1 and (T - indices[-1]) < chunk_size: + indices.pop() + return indices + + raise ValueError(f"Unknown chunk_split_strategy '{strategy}'. Supported: uniform, first_frame, first_plus_one.") + + +def compute_chunk_sizes(chunk_index: List[int], T: int) -> List[int]: + """Compute actual chunk sizes from chunk_index. + + Args: + chunk_index: List of chunk start indices (e.g., [0, 4, 8, 12]). + T: Total number of frames. + + Returns: + List of chunk sizes (e.g., [4, 4, 4, 1] if T=13). + + Example: + >>> compute_chunk_sizes([0, 4, 8, 12], T=13) [4, 4, 4, 1] >>> compute_chunk_sizes([0, 1, 5, 9], T=13) [1, 4, 4, + 4] + """ + if not chunk_index: + return [] + + # Ensure chunk_index is clean + chunk_index = [idx for idx in chunk_index if 0 <= idx < T] + if not chunk_index: + return [] + + # Add T as the final boundary if not present + if chunk_index[-1] != T: + chunk_index = chunk_index + [T] + + # Compute sizes + sizes = [chunk_index[i + 1] - chunk_index[i] for i in range(len(chunk_index) - 1)] + return sizes + + +def is_uniform_chunking( + chunk_index: List[int], + T: int, + chunk_size: int, +) -> bool: + """Check if chunk_index represents uniform chunking. + + Returns True if all chunks are equal to chunk_size except possibly the last chunk which may be smaller (the + remainder). This is the pattern that allows safe vectorized padding with: pad_t = chunk_size - (T % chunk_size). + + Uniform patterns (return True): + - [0,4,8,12,16,20] with T=21, chunk_size=4 → sizes [4,4,4,4,4,1] ✓ + - [0,4,8,12,16] with T=20, chunk_size=4 → sizes [4,4,4,4,4] ✓ + - [0,4,8] with T=10, chunk_size=4 → sizes [4,4,2] ✓ + + Non-uniform patterns (return False): + - [0,1,5,9,13,17] with T=21, chunk_size=4 → sizes [1,4,4,4,4,4] ✗ + - [0,5,9,13,17] with T=21, chunk_size=4 → sizes [5,4,4,4,4] ✗ + + Args: + chunk_index: List of chunk start indices. + T: Total number of frames. + chunk_size: Expected uniform chunk size. + + Returns: + True if chunking is uniform, False otherwise. + """ + if chunk_size <= 0: + return False + + # Compute actual chunk sizes + sizes = compute_chunk_sizes(chunk_index, T) + + if not sizes: + return True # Empty is trivially uniform + + # Check that all chunks except possibly the last are equal to chunk_size + for i, size in enumerate(sizes): + is_last = i == len(sizes) - 1 + if is_last: + # Last chunk can be <= chunk_size (remainder) + if size > chunk_size: + return False + else: + # All other chunks must be exactly chunk_size + if size != chunk_size: + return False + + return True + + +def normalize_chunk_index( + chunk_index: Optional[List[int]], + T: int, + chunk_size: Optional[int] = None, + chunk_split_strategy: str = "uniform", +) -> Tuple[List[int], bool]: + """Normalize chunk_index and detect if uniform. + + This function handles all the complex logic for: + 1. Converting chunk_size + strategy → chunk_index (if needed) + 2. Cleaning and validating chunk_index + 3. Detecting if the result is uniform (safe for vectorized padding) + + Args: + chunk_index: Optional pre-computed chunk indices. + T: Total number of frames. + chunk_size: Chunk size (required if chunk_index is None or for uniformity check). + chunk_split_strategy: Strategy to use if generating chunk_index from chunk_size. + + Returns: + (normalized_chunk_index, is_uniform): + - normalized_chunk_index: Clean list of chunk start indices + - is_uniform: True if safe to use vectorized path with padding + + Raises: + ValueError: If required parameters are missing or invalid. + """ + # Case 1: chunk_index provided explicitly + if chunk_index is not None: + normalized_chunk_index = list(chunk_index) + + # Clean up: ensure starts with 0 and ends with T + if not normalized_chunk_index or normalized_chunk_index[0] != 0: + normalized_chunk_index = [0] + [idx for idx in normalized_chunk_index if idx > 0] + normalized_chunk_index = [idx for idx in normalized_chunk_index if idx < T] + if not normalized_chunk_index: + normalized_chunk_index = [0] + if normalized_chunk_index[-1] != T: + normalized_chunk_index = normalized_chunk_index + [T] + + # Check if uniform (requires chunk_size for comparison) + if chunk_size is None: + # Can't verify uniformity without chunk_size, assume non-uniform (safe) + is_uniform = False + else: + is_uniform = is_uniform_chunking(normalized_chunk_index, T, chunk_size) + + return normalized_chunk_index, is_uniform + + # Case 2: Generate chunk_index from chunk_size + strategy + if chunk_size is None: + raise ValueError("Either chunk_index or chunk_size must be provided.") + + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}.") + + # Normalize strategy + strategy = "uniform" if chunk_split_strategy is None else str(chunk_split_strategy).lower() + + # Generate chunk_index + chunk_index_gen = chunk_index_from_chunk_size(T, chunk_size, strategy=strategy) + + # Add T as final boundary + if not chunk_index_gen: + chunk_index_gen = [0] + if chunk_index_gen[-1] != T: + chunk_index_gen = chunk_index_gen + [T] + + # Check if uniform + is_uniform = is_uniform_chunking(chunk_index_gen, T, chunk_size) + + return chunk_index_gen, is_uniform + + +# ============================================================================ +# Attention blocks (sana / sana-camctrl / GDN / GDN-camctrl / softmax variants) +# ============================================================================ + +# String-keyed registry for the GDN/softmax attention block variants used by the +# SANA-WM DiT. `SanaWMTransformer3DModel` looks classes up here by its `attn_type` +# / `camctrl_type` config strings. +ATTENTION_BLOCKS: dict[str, type] = {} + + +def _register_block(name: str | None = None): + def deco(cls): + ATTENTION_BLOCKS[name or cls.__name__] = cls + return cls + + return deco + + +def _resolve_attention_block(name: str, *, role: str) -> type: + """Look up an attention class with automatic Triton -> pure-PyTorch fallback. + + The ``*Triton`` attention classes (``BidirectionalGDNTriton``, ``BidirectionalGDNUCPESinglePathLiteLATriton``, + ``BidirectionalGDNUCPESinglePathLiteLABothTriton``) wrap pure-PyTorch ancestor classes and only differ in the + fused-kernel fast path. When Triton isn't usable (CPU-only systems, ROCm without Triton, etc.), we walk the MRO to + find the closest registered non-``Triton`` ancestor and use that instead, with a one-shot log line. + """ + cls = ATTENTION_BLOCKS.get(name) + if cls is None: + raise ValueError(f"Unknown {role}: {name!r}. Available: {sorted(ATTENTION_BLOCKS)}") + if not name.endswith("Triton") or _is_triton_kernels_usable(): + return cls + + for ancestor in cls.__mro__[1:]: + anc_name = ancestor.__name__ + if anc_name.endswith("Triton"): + continue + if ATTENTION_BLOCKS.get(anc_name) is ancestor: + _warn_triton_fallback_once(name, anc_name, role) + return ancestor + # No registered non-Triton ancestor — return the original. The Triton entry + # points each call ``_require_triton`` and will raise a clear error if + # actually invoked. + return cls + + +@lru_cache(maxsize=1) +def _is_triton_kernels_usable() -> bool: + """``triton`` is importable AND the current device can launch its kernels.""" + from .transformer_sana_wm_kernels import is_triton_available # noqa: PLC0415 + + return bool(is_triton_available() and torch.cuda.is_available()) + + +@lru_cache(maxsize=None) +def _warn_triton_fallback_once(requested: str, fallback: str, role: str) -> None: + logger.warning( + f"Triton isn't usable on this device — falling back from {role}={requested!r} " + f"to its pure-PyTorch parent {role}={fallback!r}. Install Triton and run on " + f"CUDA to use the fused-kernel fast path." + ) + + +class ConvLayer(nn.Module): + def __init__( + self, + in_dim: int, + out_dim: int, + kernel_size=3, + stride=1, + dilation=1, + groups=1, + padding: Optional[int] = None, + use_bias=False, + dropout=0.0, + conv_type="2d", + norm="bn2d", + act="relu", + ): + super().__init__() + if padding is None: + padding = get_same_padding(kernel_size) + padding *= dilation + + self.in_dim = in_dim + self.out_dim = out_dim + self.kernel_size = kernel_size + self.stride = stride + self.dilation = dilation + self.groups = groups + self.padding = padding + self.use_bias = use_bias + + self.dropout = nn.Dropout2d(dropout, inplace=False) if dropout > 0 else None + if conv_type == "2d": + self.conv = nn.Conv2d( + in_dim, + out_dim, + kernel_size=(kernel_size, kernel_size), + stride=(stride, stride), + padding=padding, + dilation=(dilation, dilation), + groups=groups, + bias=use_bias, + ) + elif conv_type == "3d": + self.conv = nn.Conv3d( + in_dim, + out_dim, + kernel_size=(kernel_size, kernel_size, kernel_size), + stride=(stride, stride, stride), + padding=padding, + dilation=(dilation, dilation, dilation), + groups=groups, + bias=use_bias, + ) + else: + self.conv = None + + self.norm = build_norm(norm, num_features=out_dim) + self.act = build_act(act) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.dropout is not None: + x = self.dropout(x) + x = self.conv(x) + if self.norm: + x = self.norm(x) + if self.act: + x = self.act(x) + return x + + +# Safe element-count threshold for a single conv call: PyTorch's 2D conv kernels +# (both cuDNN and the ATEN fallback) use 32-bit indexing internally, so very +# large ``(BT, C, H, W)`` inputs (e.g. minute-scale video at default CFG) can +# overflow. Empirically a single call up to ~1 B elements is safe; above that +# we chunk along the leading dim. Set so short videos stay on the original +# fused path (no chunking, no overhead) and long videos transparently split. +_INT32_SAFE_CONV_ELEMENTS = 1 << 30 # 1,073,741,824 + + +class GLUMBConv(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: int, + out_feature=None, + kernel_size=3, + stride=1, + padding: Optional[int] = None, + use_bias=False, + norm=(None, None, None), + act=("silu", "silu", None), + dilation=1, + ): + out_feature = out_feature or in_features + super().__init__() + use_bias = val2tuple(use_bias, 3) + norm = val2tuple(norm, 3) + act = val2tuple(act, 3) + + self.glu_act = build_act(act[1], inplace=False) + self.inverted_conv = ConvLayer( + in_features, + hidden_features * 2, + 1, + use_bias=use_bias[0], + norm=norm[0], + act=act[0], + ) + self.depth_conv = ConvLayer( + hidden_features * 2, + hidden_features * 2, + kernel_size, + stride=stride, + groups=hidden_features * 2, + padding=padding, + use_bias=use_bias[1], + norm=norm[1], + act=None, + dilation=dilation, + ) + self.point_conv = ConvLayer( + hidden_features, + out_feature, + 1, + use_bias=use_bias[2], + norm=norm[2], + act=act[2], + ) + + def _apply_spatial(self, x: torch.Tensor) -> torch.Tensor: + """Fused spatial pipeline: inverted_conv -> depth_conv -> GLU -> point_conv.""" + x = self.inverted_conv(x) + x = self.depth_conv(x) + a, g = torch.chunk(x, 2, dim=1) + g = self.glu_act(g) + return self.point_conv(a * g) + + def _apply_spatial_autochunked(self, x: torch.Tensor) -> torch.Tensor: + """Run :meth:`_apply_spatial`, chunking dim 0 to keep each call under + PyTorch's 32-bit conv indexing limit. No-op for short inputs.""" + BT, _, H, W = x.shape + # Conservative estimate of the largest intermediate (after inverted_conv). + elements_per_bt = self.inverted_conv.conv.out_channels * H * W + max_bt = max(1, _INT32_SAFE_CONV_ELEMENTS // elements_per_bt) + if BT <= max_bt: + return self._apply_spatial(x) + return torch.cat([self._apply_spatial(x[s : s + max_bt]) for s in range(0, BT, max_bt)], dim=0) + + def forward(self, x: torch.Tensor, HW=None) -> torch.Tensor: + B, N, C = x.shape + if HW is None: + H = W = int(N**0.5) + elif len(HW) == 2: + H, W = HW + x = x.reshape(B, H, W, C).permute(0, 3, 1, 2) + elif len(HW) == 3: + T, H, W = HW + x = x.reshape(B * T, H, W, C).permute(0, 3, 1, 2) + + x = self._apply_spatial_autochunked(x) + + if len(HW) == 3: + x = x.reshape(B * T, C, H * W).permute(0, 2, 1) + x = x.reshape(B, N, C) + else: + x = x.reshape(B, C, N).permute(0, 2, 1) + + return x + + +class GLUMBConvTemp(GLUMBConv): + def __init__( + self, + in_features: int, + hidden_features: int, + out_feature=None, + kernel_size=3, + stride=1, + padding: Optional[int] = None, + use_bias=False, + norm=(None, None, None), + act=("silu", "silu", None), + t_kernel_size=3, + ): + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_feature=out_feature, + kernel_size=kernel_size, + stride=stride, + padding=padding, + use_bias=use_bias, + norm=norm, + act=act, + ) + + out_feature = out_feature or in_features + t_padding = t_kernel_size // 2 + self.t_conv = nn.Conv2d( + out_feature, + out_feature, + kernel_size=(t_kernel_size, 1), + stride=1, + padding=(t_padding, 0), + bias=False, + ) + + nn.init.zeros_(self.t_conv.weight) + + def forward(self, x: torch.Tensor, HW=None, **kwargs) -> torch.Tensor: + B, N, C = x.shape + + assert len(HW) == 3, "HW must be a tuple of (T, H, W)" + T, H, W = HW + x = x.reshape(B * T, H, W, C).permute(0, 3, 1, 2) + + x = self._apply_spatial_autochunked(x) + + # Temporal aggregation + x_reshaped = x.view(B, T, C, H * W).permute(0, 2, 1, 3) + x_out = x_reshaped + self.t_conv(x_reshaped) + + x_out = x_out.permute(0, 2, 3, 1).reshape(B, N, C) + + return x_out + + +class DWMlp(Mlp): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + bias=True, + drop=0.0, + kernel_size=3, + stride=1, + dilation=1, + padding=None, + ): + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + act_layer=act_layer, + bias=bias, + drop=drop, + ) + hidden_features = hidden_features or in_features + self.hidden_features = hidden_features + if padding is None: + padding = get_same_padding(kernel_size) + padding *= dilation + + self.conv = nn.Conv2d( + hidden_features, + hidden_features, + kernel_size=(kernel_size, kernel_size), + stride=(stride, stride), + padding=padding, + dilation=(dilation, dilation), + groups=hidden_features, + bias=bias, + ) + + def forward(self, x, HW=None): + B, N, C = x.shape + if HW is None: + H = W = int(N**0.5) + else: + H, W = HW + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = x.reshape(B, H, W, self.hidden_features).permute(0, 3, 1, 2) + x = self.conv(x) + x = x.reshape(B, self.hidden_features, N).permute(0, 2, 1) + x = self.fc2(x) + x = self.drop2(x) + return x + + +class Mlp(Mlp): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.0): + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + act_layer=act_layer, + bias=bias, + drop=drop, + ) + + def forward(self, x, HW=None): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def t2i_modulate(x, shift, scale): + return x * (1 + scale) + shift + + +class MultiHeadCrossAttention(nn.Module): + def __init__(self, d_model, num_heads, attn_drop=0.0, proj_drop=0.0, qk_norm=False, **block_kwargs): + super().__init__() + assert d_model % num_heads == 0, "d_model must be divisible by num_heads" + + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + + self.q_linear = nn.Linear(d_model, d_model) + self.kv_linear = nn.Linear(d_model, d_model * 2) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(d_model, d_model) + self.proj_drop = nn.Dropout(proj_drop) + if qk_norm: + self.q_norm = RMSNorm(d_model, scale_factor=1.0, eps=1e-6) + self.k_norm = RMSNorm(d_model, scale_factor=1.0, eps=1e-6) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + def forward(self, x, cond, mask=None): + # query: img tokens; key/value: condition; mask: if padding tokens + B, N, C = x.shape + q = self.q_linear(x) + kv = self.kv_linear(cond).view(B, -1, 2, C) + k, v = kv.unbind(2) + q = self.q_norm(q).view(B, -1, self.num_heads, self.head_dim) + k = self.k_norm(k).view(B, -1, self.num_heads, self.head_dim) + v = v.view(B, -1, self.num_heads, self.head_dim) + + q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + if mask is not None and mask.ndim == 2: + mask = (1 - mask.to(q.dtype)) * -10000.0 + mask = mask[:, None, None].repeat(1, self.num_heads, 1, 1) + x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) + x = x.transpose(1, 2) + + x = x.view(B, -1, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +################################################################################# +# AMP attention with fp32 softmax to fix loss NaN problem during training # +################################################################################# + + +class T2IFinalLayer(nn.Module): + """ + The final layer of Sana. + """ + + def __init__(self, hidden_size, patch_size, out_channels): + super().__init__() + if isinstance(patch_size, int): + patch_size = [patch_size, patch_size] + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, math.prod(patch_size) * out_channels, bias=True) + self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5) + self.out_channels = out_channels + + def forward_frame_aware(self, x, t): + # t: B,1,F,D + B, N, C = x.shape + num_frames = t.shape[2] + # shift, scale: 2, hidden_size -> 1,1,2,hidden_size -> B,F,2,hidden_size + shift, scale = (self.scale_shift_table[None, None, :, :] + t.transpose(1, 2)).chunk( + 2, dim=-2 + ) # each chunk: B,F,1,D + x = t2i_modulate(self.norm_final(x).reshape(B, num_frames, -1, C), shift, scale).reshape(B, N, C) + x = self.linear(x) + return x + + def forward(self, x, t): + if len(t.shape) > 2: + return self.forward_frame_aware(x, t) + shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1) + x = t2i_modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +################################################################################# +# Embedding Layers for Timesteps and Class Labels # +################################################################################# +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(self.dtype) + t_emb = self.mlp(t_freq) + return t_emb + + @property + def dtype(self): + try: + return next(self.parameters()).dtype + except StopIteration: + return torch.float32 + + +class CaptionEmbedder(nn.Module): + """ + Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. + """ + + def __init__( + self, + in_channels, + hidden_size, + act_layer=nn.GELU(approximate="tanh"), + token_num=120, + ): + super().__init__() + self.y_proj = Mlp( + in_features=in_channels, hidden_features=hidden_size, out_features=hidden_size, act_layer=act_layer, drop=0 + ) + self.register_buffer("y_embedding", nn.Parameter(torch.randn(token_num, in_channels) / in_channels**0.5)) + + def forward(self, caption): + return self.y_proj(caption) + + +class PatchEmbedMS3D(nn.Module): + """3D Image to Patch Embedding""" + + def __init__( + self, + patch_size=(1, 2, 2), + in_chans=3, + embed_dim=768, + kernel_size=None, + padding=0, + norm_layer=None, + flatten=True, + bias=True, + ): + super().__init__() + kernel_size = kernel_size or patch_size + patch_size = to_3tuple(patch_size) + self.kernel_size = kernel_size + self.patch_size = patch_size + self.flatten = flatten + assert patch_size[0] == 1, "Patch size for 3D embedding must be (1, *, *)" + if not padding and kernel_size[-1] % 2 > 0: + padding = get_same_padding(kernel_size) + self.proj = nn.Conv3d( + in_chans, embed_dim, kernel_size=kernel_size, stride=patch_size, padding=padding, bias=bias + ) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCTHW -> BNC + x = self.norm(x) + return x + + +class WanRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size: Tuple[int, int, int], + max_seq_len: int, + theta: float = 10000.0, + fhw_dim: Optional[Tuple[int, int, int]] = None, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + if fhw_dim is not None: + assert attention_head_dim == sum(fhw_dim), ( + f"attention_head_dim {attention_head_dim} must match sum(fhw_dim) {sum(fhw_dim)}" + ) + t_dim, h_dim, w_dim = fhw_dim + else: + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + + freqs = [] + for dim in [t_dim, h_dim, w_dim]: + freq = get_1d_rotary_pos_embed( + dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=torch.float32 + ) + freqs.append(freq) + self.register_buffer("freqs", torch.cat(freqs, dim=1), persistent=False) + + def forward(self, fhw: Tuple[int, int, int]) -> torch.Tensor: + ppf, pph, ppw = fhw + + freqs = self.freqs.split_with_sizes( + [ + self.attention_head_dim // 2 - 2 * (self.attention_head_dim // 6), + self.attention_head_dim // 6, + self.attention_head_dim // 6, + ], + dim=1, + ) + + freqs_f = freqs[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_h = freqs[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_w = freqs[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + freqs = torch.cat([freqs_f, freqs_h, freqs_w], dim=-1).reshape(1, 1, ppf * pph * ppw, -1) + return freqs + + +def apply_rotary_emb( + x: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + use_real: bool = True, + use_real_unbind_dim: int = -1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings + to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are + reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting + tensors contain rotary embeddings and are returned as real tensors. + + Args: + x (`torch.Tensor`): + Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply + freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + """ + if use_real: + cos, sin = freqs_cis # [S, D] + cos = cos[None, None] + sin = sin[None, None] + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Sana + cos = cos.transpose(-1, -2) + sin = sin.transpose(-1, -2) + x_real, x_imag = x.reshape(*x.shape[:-2], -1, 2, x.shape[-1]).unbind(-2) # [B, H, D//2, S] + x_rotated = torch.stack([-x_imag, x_real], dim=-2).flatten(2, 3) + else: + raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + return out + else: + # used for lumina + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + + return x_out.type_as(x) + + +# --------------------------------------------------------------------------- +# Camera-branch dropout +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# UCM (Unified Camera Model) projection / unprojection +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Per-pixel ray transformation (world <-> ray) used by UCPE +# --------------------------------------------------------------------------- + + +def _process_camera_conditions_ucpe(camera_conditions, B, HW, patch_size): + """Convert ``(B, F, 20)`` camera conditions (C2W flat + fx,fy,cx,cy) into + ``(raymats, absmap)``. + + ``raymats`` is ``(B, F, H, W, 4, 4)`` ``ray<-world`` transforms; ``absmap`` is ``(B, F, H, W, 3)`` (up_map 2-ch + + lat_map 1-ch). + """ + F_dim = camera_conditions.shape[1] + c2w_flat = camera_conditions[..., :16] + C_to_W = c2w_flat.view(B, F_dim, 4, 4) + + fx = camera_conditions[..., 16] + fy = camera_conditions[..., 17] + cx = camera_conditions[..., 18] + cy = camera_conditions[..., 19] + H_dim, W_dim = HW[1], HW[2] + image_width = W_dim * patch_size[2] + image_height = H_dim * patch_size[1] + + # xi is fixed at 0 (pinhole) in this stack. + xi = torch.zeros((B, F_dim), device=camera_conditions.device, dtype=camera_conditions.dtype) + x_fov = compute_fov_from_fx_xi( + fx, xi, image_width, device=camera_conditions.device, dtype=camera_conditions.dtype + ).view(B, F_dim) + y_fov = compute_fov_from_fx_xi( + fy, xi, image_height, device=camera_conditions.device, dtype=camera_conditions.dtype + ).view(B, F_dim) + + d_cam = ucm_unproject_grid_fov( + x_fov, + y_fov, + xi, + H_dim, + W_dim, + cx / patch_size[2], + cy / patch_size[1], + device=camera_conditions.device, + dtype=camera_conditions.dtype, + ) + if d_cam.ndim == 4 and d_cam.shape[0] == B * F_dim: + d_cam = d_cam.view(B, F_dim, H_dim, W_dim, 3) + + raymats = world_to_ray_mats(d_cam, C_to_W) # [B, F, H, W, 4, 4] + + up_map, lat_map = compute_up_lat_map( + R=C_to_W[..., :3, :3], + x_fov=x_fov, + y_fov=y_fov, + xi=xi, + height=image_height, + width=image_width, + cx=cx, + cy=cy, + device=camera_conditions.device, + ) + absmap = torch.cat([up_map, lat_map], dim=-1) # (B, F, H, W, 3) + + return raymats, absmap + + +# --------------------------------------------------------------------------- +# Block-diagonal apply primitives shared by camera and main branches +# --------------------------------------------------------------------------- + + +@torch.compile +def _apply_ray_projmat( + feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) + matrix: torch.Tensor, # (batch, seqlen, 4, 4) +) -> torch.Tensor: + """Apply a per-token 4x4 projection matrix to feature channels grouped by 4.""" + (batch, num_heads, seqlen, feat_dim) = feats.shape + D = matrix.shape[-1] + return torch.einsum( + "bnij,bhnkj->bhnki", + matrix, + feats.reshape(batch, num_heads, seqlen, -1, D), + ).reshape(feats.shape) + + +@torch.compile +def _apply_complex_rope( + hidden_states: torch.Tensor, + freqs: torch.Tensor, + inverse: bool = False, +) -> torch.Tensor: + """Apply complex RoPE (compiled: fuses fp64 cast + view_as_complex + multiply chain).""" + x_real = hidden_states.to(torch.float32) + if x_real.stride(-1) != 1: + x_real = x_real.contiguous() + x_complex = torch.view_as_complex(x_real.unflatten(-1, (-1, 2))) + if inverse: + freqs = freqs.conj() + x_out = torch.view_as_real(x_complex * freqs).flatten(-2, -1) + return x_out.type_as(hidden_states) + + +def _apply_block_diagonal( + feats: torch.Tensor, # (..., dim) + func_size_pairs: List[Tuple[Callable[[torch.Tensor], torch.Tensor], int]], +) -> torch.Tensor: + """Apply a block-diagonal function: split features by sizes, transform each, concat.""" + funcs, block_sizes = zip(*func_size_pairs) + assert feats.shape[-1] == sum(block_sizes) + x_blocks = torch.split(feats, block_sizes, dim=-1) + out = torch.cat( + [f(x_block) for f, x_block in zip(funcs, x_blocks)], + dim=-1, + ) + assert out.shape == feats.shape, "Input/output shapes should match." + return out + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Closed-form inverse of a 4x4 SE(3) batch.""" + assert transforms.shape[-2:] == (4, 4) + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + return out + + +# --------------------------------------------------------------------------- +# UCPE apply-fn preparation +# --------------------------------------------------------------------------- + + +def _prepare_ray_apply_fns( + head_dim: int, + P: torch.Tensor, # (batch, seqlen, 4, 4) P = ray<-world + P_T: torch.Tensor, # (batch, seqlen, 4, 4) P_T = world<-ray + P_inv: torch.Tensor, # (batch, seqlen, 4, 4) P_inv = world<-ray + rotary_emb: Optional[torch.Tensor] = None, + apply_vo: bool = True, +) -> Tuple[Callable, Callable, Callable]: + """Build ``(apply_q, apply_kv, apply_o)`` block-diagonal callables for UCPE.""" + if rotary_emb is not None: + rope_fn = partial(_apply_complex_rope, freqs=rotary_emb, inverse=False) + rope_fn_inv = partial(_apply_complex_rope, freqs=rotary_emb, inverse=True) + else: + + def rope_fn(x): + return x + + def rope_fn_inv(x): + return x + + transforms_q = [ + (partial(_apply_ray_projmat, matrix=P_T), head_dim // 2), + (rope_fn, head_dim // 2), + ] + transforms_kv = [ + (partial(_apply_ray_projmat, matrix=P_inv), head_dim // 2), + (rope_fn, head_dim // 2), + ] + if apply_vo: + transforms_o = [ + (partial(_apply_ray_projmat, matrix=P), head_dim // 2), + (rope_fn_inv, head_dim // 2), + ] + else: + + def transforms_o(x): + return x + + apply_fn_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q) + apply_fn_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv) + apply_fn_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o) if apply_vo else transforms_o + + return apply_fn_q, apply_fn_kv, apply_fn_o + + +def _slice_rope_for_cam( + rotary_emb: Optional[torch.Tensor], + head_dim: int, + rope_dim: int, +) -> Optional[torch.Tensor]: + """Re-slice WAN RoPE frequencies for a smaller rope_dim using the same (T, H, W) split.""" + if rotary_emb is None: + return None + orig_t_size = head_dim // 2 - 2 * (head_dim // 6) + orig_h_size = head_dim // 6 + new_t_size = rope_dim // 2 - 2 * (rope_dim // 6) + new_h_size = rope_dim // 6 + new_w_size = rope_dim // 6 + t_part = rotary_emb[..., :new_t_size] + h_part = rotary_emb[..., orig_t_size : orig_t_size + new_h_size] + w_part = rotary_emb[..., orig_t_size + orig_h_size : orig_t_size + orig_h_size + new_w_size] + return torch.cat([t_part, h_part, w_part], dim=-1) + + +def prepare_prope_fns( + camctrl_type: str, + head_dim: int, + camera_conditions: torch.Tensor, + HW: Tuple[int, int, int], + patch_size: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor] = None, + **kwargs, +) -> Tuple[Callable, Callable, Callable]: + """Precompute UCPE apply functions once for a batch (shared across all blocks). + + Only ``camctrl_type == "UCPE"`` is supported. Accepts either precomputed matrices (``cam_pos_embeds`` dict with + ``P``, ``P_inv``, ``pos_embeds_cam``) or raw camera conditions + optional raymats. + """ + if camctrl_type != "UCPE": + raise ValueError(f"Unsupported camctrl_type for prepare_prope_fns: {camctrl_type}") + + B = camera_conditions.shape[0] + + # Priority 1: use precomputed matrices. + if "cam_pos_embeds" in kwargs and kwargs["cam_pos_embeds"] is not None: + cam_pos_embeds = kwargs["cam_pos_embeds"] + P = cam_pos_embeds.get("P") + P_inv = cam_pos_embeds.get("P_inv") + rotary_emb_cam = cam_pos_embeds.get("pos_embeds_cam") + + if P is not None and P_inv is not None: + if P.ndim == 3: + P = P.unsqueeze(0).repeat(B, 1, 1, 1) + if P_inv.ndim == 3: + P_inv = P_inv.unsqueeze(0).repeat(B, 1, 1, 1) + + P_T = P.transpose(-1, -2) + + if rotary_emb_cam is not None and rotary_emb_cam.ndim == 3: + rotary_emb_cam = rotary_emb_cam.unsqueeze(0).repeat(B, 1, 1, 1) + elif rotary_emb_cam is None and rotary_emb is not None: + rotary_emb_cam = _slice_rope_for_cam(rotary_emb, head_dim, head_dim // 2) + elif rotary_emb_cam is None: + rotary_emb_cam = rotary_emb + + return _prepare_ray_apply_fns(head_dim, P, P_T, P_inv, rotary_emb=rotary_emb_cam) + + # Priority 2: online path. + if "raymats" in kwargs and kwargs["raymats"] is not None: + raymats = kwargs["raymats"] + else: + raymats, _ = _process_camera_conditions_ucpe(camera_conditions, B, HW, patch_size) + raymats = raymats.reshape(B, -1, 4, 4) + + P = raymats + P_T = P.transpose(-1, -2) + P_inv = _invert_SE3(P) + + rotary_emb_cam = _slice_rope_for_cam(rotary_emb, head_dim, head_dim // 2) if rotary_emb is not None else None + + return _prepare_ray_apply_fns(head_dim=head_dim, P=P, P_T=P_T, P_inv=P_inv, rotary_emb=rotary_emb_cam) + + +OUTPUT_GATE_INIT_BIAS = 1.278464542761074 # silu(x)=1.0 + + +def flip_and_shift(x, dim=2, shift_val=0.0): + """Flip a sequence and shift it right by one step. + + The operation reverses the sequence, drops the last element, and pads the front with ``shift_val``. + + Example: + [x0, x1, x2, x3] -> flip [x3, x2, x1, x0] -> shift [v, x3, x2, x1] + + Args: + x: Input tensor with a time dimension at ``dim``. + dim: Dimension to flip and shift. + shift_val: Value used for the padded step. + + Returns: + Tensor with the same shape as ``x``. + """ + x_flip = torch.flip(x, dims=[dim]) + x_shifted = x_flip.narrow(dim, 0, x.shape[dim] - 1) + pad_shape = list(x.shape) + pad_shape[dim] = 1 + padding = torch.full(pad_shape, shift_val, device=x.device, dtype=x.dtype) + return torch.cat([padding, x_shifted], dim=dim) + + +class _IdentityForwardContiguousBackward(torch.autograd.Function): + """Identity in forward; force contiguous grad tensor in backward.""" + + @staticmethod + def forward(ctx, x: torch.Tensor) -> torch.Tensor: + return x + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]: + return (grad_output.contiguous(),) + + +def _contiguous_backward(x: torch.Tensor) -> torch.Tensor: + """Ensure downstream backward receives a contiguous gradient buffer.""" + return _IdentityForwardContiguousBackward.apply(x) + + +@torch.compile +def torch_chunk_sana_gdn( + q, + k, + v, + q_rot, + k_rot, + beta, + decay, + recall_gate=None, + chunk_size: int | None = 21, + eps: float = 1e-6, + return_components: bool = False, +): + del recall_gate # Accepted so the chunk and fused scan share one signature; unused by this rule. + + B, H, D, N = q.shape + if beta.ndim not in (3, 4): + raise ValueError(f"Expected beta.ndim in (3, 4), got {beta.ndim}.") + T = beta.shape[2] + if T <= 0: + raise ValueError(f"Expected T > 0, got T={T}.") + if N % T != 0: + raise ValueError(f"Expected N divisible by T, got N={N}, T={T}.") + S = N // T + + target_z = 1.0 + scale = 1.0 + + def to_frame_seq(x): + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q, k, v = to_frame_seq(q), to_frame_seq(k), to_frame_seq(v) + q_rot, k_rot = to_frame_seq(q_rot), to_frame_seq(k_rot) + + if beta.ndim == 4: + beta = beta.unsqueeze(3) + else: + beta = beta.view(B, H, T, 1, 1) + + decay = decay.view(B, H, T, 1, 1) + + # ========================================================================= + # 1. PARALLEL PRE-PROCESSING + # ========================================================================= + + I = torch.eye(D, device=q.device, dtype=q.dtype).view(1, 1, 1, D, D) + + # KV State Matrices: W = g * (I - c * K @ K^T) + k_rot_beta = k_rot * beta + W_kv = decay * (I - scale * torch.matmul(k_rot_beta, k_rot.transpose(-1, -2))) + U_kv = torch.matmul(v * beta, k_rot.transpose(-1, -2)) + + # Z State Matrices: W = g * (I - c * K @ K^T) + k_beta = k * beta + W_z = decay * (I - scale * torch.matmul(k_beta, k.transpose(-1, -2))) + U_z = target_z * k_beta.sum(dim=-1, keepdim=True) # Equivalent to Kt @ bt^T over spatial dim + + # ========================================================================= + # 2. CHUNKING LOGIC + # ========================================================================= + + valid_chunk_index, _ = normalize_chunk_index(None, T, chunk_size) + split_sizes = [valid_chunk_index[i + 1] - valid_chunk_index[i] for i in range(len(valid_chunk_index) - 1)] + + W_kv_c = W_kv.split(split_sizes, dim=2) + U_kv_c = U_kv.split(split_sizes, dim=2) + W_z_c = W_z.split(split_sizes, dim=2) + U_z_c = U_z.split(split_sizes, dim=2) + + # ========================================================================= + # 3. FAST INTRA-CHUNK SCAN OVER DxD SPACE + # ========================================================================= + + S_kv = torch.zeros(B, H, D, D, device=q.device, dtype=q.dtype) + S_z = torch.zeros(B, H, D, 1, device=q.device, dtype=q.dtype) + + out_S_kv = [] + out_S_z = [] + + def _chunk_scan(w_kv, u_kv, w_z, u_z, s_kv, s_z): + c_len = w_kv.shape[2] + s_kv_list, s_z_list = [], [] + for t in range(c_len): + s_kv = torch.matmul(s_kv, w_kv[:, :, t]) + u_kv[:, :, t] + s_z = torch.matmul(w_z[:, :, t], s_z) + u_z[:, :, t] + s_kv_list.append(s_kv) + s_z_list.append(s_z) + return torch.stack(s_kv_list, dim=2), s_kv, torch.stack(s_z_list, dim=2), s_z + + for i in range(len(split_sizes)): + s_kv_all, S_kv, s_z_all, S_z = _chunk_scan(W_kv_c[i], U_kv_c[i], W_z_c[i], U_z_c[i], S_kv, S_z) + out_S_kv.append(s_kv_all) + out_S_z.append(s_z_all) + + S_kv_all = torch.cat(out_S_kv, dim=2) + S_z_all = torch.cat(out_S_z, dim=2) + + # ========================================================================= + # 4. PARALLEL OUTPUT PROJECTION + # ========================================================================= + + out_num = torch.matmul(S_kv_all, q_rot) + out_den = torch.matmul(S_z_all.transpose(-1, -2), q) + + def restore_shape(tensor, target_d): + return tensor.permute(0, 1, 3, 2, 4).reshape(B, H, target_d, N) + + final_num = restore_shape(out_num, D) + final_den = restore_shape(out_den, 1) + + if return_components: + return final_num, final_den + + return final_num / (final_den + eps) + + +# --------------------------------------------------------------------------- +# Compiled helpers for hot-path operations (fuses elementwise chains) +# --------------------------------------------------------------------------- + + +@torch.compile +def _compute_frame_gates( + x: torch.Tensor, + T: int, + S: int, + heads: int, + beta_weight: torch.Tensor, + beta_bias: torch.Tensor, + gate_weight: torch.Tensor, + gate_bias: torch.Tensor, + dt_bias: torch.Tensor, + A_log: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compiled frame gate computation (fuses sigmoid + softplus + exp chain).""" + B, N, C = x.shape + beta = F.linear(x, beta_weight, beta_bias).sigmoid().reshape(B, T, S, heads).permute(0, 3, 1, 2) + x_frame = x.reshape(B, T, S, C).mean(dim=2) + a_out = F.linear(x_frame, gate_weight, gate_bias).float() + dt = dt_bias.float().view(1, 1, -1) + A_val = A_log.float().exp().view(1, 1, -1) + decay = (-A_val * F.softplus(a_out + dt)).exp().transpose(1, 2) + return beta, decay + + +@torch.compile +def _apply_rotary_emb( + hidden_states: torch.Tensor, + freqs: torch.Tensor, +) -> torch.Tensor: + """Compiled rotary embedding application (fuses view_as_complex + multiply chain).""" + x_rotated = torch.view_as_complex( + hidden_states.permute(0, 1, 3, 2).to(torch.float32).unflatten(3, (-1, 2)), + ) + x_out = torch.view_as_real(x_rotated * freqs).flatten(3, 4).permute(0, 1, 3, 2) + return x_out.type_as(hidden_states) + + +@torch.compile +def _apply_output_gate( + out: torch.Tensor, + gate_x: torch.Tensor, + gate_weight: torch.Tensor, + gate_bias: torch.Tensor, +) -> torch.Tensor: + """Compiled output gate (fuses linear + silu + multiply).""" + gate = F.silu(F.linear(gate_x, gate_weight, gate_bias).to(torch.float32)) + return out * gate + + +@_register_block() +class GDN(Attention_): + """Frame-wise Gated Delta Net attention for Sana video. + + This block follows Sana's vanilla linear attention strategy but upgrades it with a Gated Delta Network mechanism: + - Apply ReLU kernel to q/k. + - Apply RoPE only on the numerator (q_rot, k_rot). + - Denominator (Z stream) uses unrotated q/k to maintain mass conservation. + - Gated delta rule is applied across time (T). Gates are computed per-frame (shared spatially), but states are + maintained per-pixel. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + heads: int | None = None, + heads_ratio: float = 1.0, + dim: int = 32, + eps: float = 1e-15, + use_bias: bool = False, + qk_norm: bool = False, + norm_eps: float = 1e-5, + use_output_gate: bool = True, + update_rule_func: str = "torch_chunk_sana_gdn", + chunk_gdn_chunk_size: int = 21, + conv_kernel_size: int = 4, + k_conv_only: bool = True, + **kwargs: object, + ) -> None: + heads = heads or int(out_dim // dim * heads_ratio) + super().__init__(in_dim, num_heads=heads, qkv_bias=use_bias) + + self.in_dim = in_dim + self.out_dim = out_dim + self.heads = heads + self.dim = out_dim // heads + self.eps = eps + self.k_conv_only = k_conv_only + self.key_scale_mode = str(kwargs.pop("key_scale_mode", "dim_spatial")) + + self.kernel_func = nn.ReLU(inplace=False) + + if qk_norm: + self.q_norm = RMSNorm(self.in_dim, scale_factor=1.0, eps=norm_eps) + self.k_norm = RMSNorm(self.in_dim, scale_factor=1.0, eps=norm_eps) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + # Gate projections operate on pooled frame features (B, T, D) -> (B, T, H). + self.beta_proj = nn.Linear(in_dim, heads, bias=True) + self.gate_proj = nn.Linear(in_dim, heads, bias=True) + + A = torch.zeros(self.heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.heads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + + # `recall_gate` is unused by the forward; kept as a buffer for checkpoint compatibility. + self.register_buffer("recall_gate", torch.zeros(1)) + + self.use_output_gate = use_output_gate + if use_output_gate: + self.output_gate = nn.Linear(in_dim, out_dim, bias=True) + else: + self.output_gate = None + + if update_rule_func != "torch_chunk_sana_gdn": + raise ValueError(f"Unsupported update rule function: {update_rule_func}") + self.update_rule_func = partial(torch_chunk_sana_gdn, chunk_size=chunk_gdn_chunk_size) + + # Short Convolutions (FLA causal depthwise Conv1d along T) + self.conv_kernel_size = conv_kernel_size + if conv_kernel_size > 0: + self.conv_k = ShortConvolution( + hidden_size=out_dim, + kernel_size=conv_kernel_size, + activation=None, + ) + if k_conv_only: + self.conv_q = None + self.conv_v = None + else: + self.conv_q = ShortConvolution( + hidden_size=out_dim, + kernel_size=conv_kernel_size, + activation=None, + ) + self.conv_v = ShortConvolution( + hidden_size=out_dim, + kernel_size=conv_kernel_size, + activation=None, + ) + else: + self.conv_q = None + self.conv_k = None + self.conv_v = None + + def _key_scale(self, spatial_tokens: int) -> float: + """Return the post-ReLU key scale used by frame-wise GDN.""" + if self.key_scale_mode == "dim_spatial": + return (self.dim**-0.5) * (spatial_tokens**-0.5) + if self.key_scale_mode == "dim": + return self.dim**-0.5 + if self.key_scale_mode == "none": + return 1.0 + raise ValueError(f"Unsupported GDN key_scale_mode: {self.key_scale_mode}") + + def _apply_output_gate(self, out: torch.Tensor, gate_x: torch.Tensor) -> torch.Tensor: + if not (self.use_output_gate and self.output_gate is not None): + return out + return _apply_output_gate(out, gate_x, self.output_gate.weight, self.output_gate.bias) + + @staticmethod + def _reshape_to_temporal(x: torch.Tensor, HW: tuple[int, int, int]) -> tuple[torch.Tensor, int, int, int]: + """Reshape (B, T*S, C) to (B*S, T, C) for temporal conv. + + Returns: + Reshaped tensor and (B, S, T) for later restoration. + """ + B, N, C = x.shape + T, H, W = HW + S = H * W + # FLA ShortConvolution backward is not reliable on non-contiguous + # strided layouts produced by this permutation path. + x = x.reshape(B, T, S, C).permute(0, 2, 1, 3).contiguous().reshape(B * S, T, C) + return x, B, S, T + + @staticmethod + def _reshape_from_temporal(x: torch.Tensor, B: int, S: int, T: int) -> torch.Tensor: + """Reshape (B*S, T, C) back to (B, T*S, C).""" + x = _contiguous_backward(x) + C = x.shape[-1] + return x.reshape(B, S, T, C).permute(0, 2, 1, 3).reshape(B, T * S, C) + + @staticmethod + def _causal_conv_1d( + x: torch.Tensor, + conv: ShortConvolution, + ) -> torch.Tensor: + """Run causal conv and preserve input dtype. + + Args: + x: Tensor of shape (batch, seq_len, channels). + conv: FLA ``ShortConvolution`` module. + + Returns: + Tensor of same shape and dtype as ``x``. + """ + dtype_in = x.dtype + y, _ = conv(x) + if y.dtype != dtype_in: + y = y.to(dtype_in) + return y + + @staticmethod + def _bidirectional_causal_conv_1d( + x: torch.Tensor, + conv: ShortConvolution, + ) -> torch.Tensor: + """Simulate non-causal conv by combining forward + backward causal passes. + + A causal depthwise Conv1d with kernel ``[w_0, w_1, ..., w_{k-1}]`` computes at time *t*: + + ``y_fwd[t] = w_0 * x[t-k+1] + ... + w_{k-1} * x[t]`` + + Running the same kernel on the time-flipped input and flipping back gives: + + ``y_bwd[t] = w_{k-1} * x[t] + ... + w_0 * x[t+k-1]`` + + Both passes include the current timestep ``x[t]`` with the center weight ``w_{k-1}``. To avoid double-counting + we subtract one copy of the center contribution: + + ``y = y_fwd + y_bwd - w_{k-1} * x`` + + The result is a symmetric temporal filter where every position in the window ``[t-k+1, t+k-1]`` is counted + exactly once. + + Args: + x: Tensor of shape ``(batch, seq_len, channels)``. + conv: FLA ``ShortConvolution`` module (depthwise causal Conv1d). + + Returns: + Tensor of same shape and dtype as ``x``. + """ + dtype_in = x.dtype + + y_fwd, _ = conv(x) + y_bwd, _ = conv(x.flip(1)) + y_bwd = y_bwd.flip(1) + + # Subtract the shared center tap (last weight of the causal kernel). + # ShortConvolution weight shape: (channels, 1, kernel_size). + # The last element along dim=-1 is the weight applied to x[t]. + w_center = conv.weight[:, 0, -1] # (channels,) + center_term = x * w_center.unsqueeze(0).unsqueeze(0) # broadcast over (B, T) + + y = y_fwd + y_bwd - center_term + if y.dtype != dtype_in: + y = y.to(dtype_in) + return y + + def _apply_temporal_short_conv( + self, + x: torch.Tensor, + conv: ShortConvolution, + HW: tuple[int, int, int], + **kwargs: object, + ) -> torch.Tensor: + """Apply causal ShortConvolution along T, with S merged into batch. + + Under CP, a causal conv of kernel size K needs K-1 left-context frames from the previous rank at each boundary. + We use a halo exchange (O(K) communication) instead of a full gather (O(T)). + + Args: + x: Input tensor of shape (B, N, C) where N = T * S. + conv: FLA ``ShortConvolution`` module. + HW: Tuple of (T, H, W) describing the token layout. + **kwargs: Extra keyword arguments (unused in base; subclasses + may consume ``chunk_size``, ``chunk_index``, etc.). + + Returns: + Tensor of shape (B, N, C) after temporal convolution. + """ + del kwargs # unused in base class + + x, B, S, T = self._reshape_to_temporal(x, HW) + x = self._causal_conv_1d(x, conv) + return self._reshape_from_temporal(x, B, S, T) + + @staticmethod + def _apply_rotary_emb( + hidden_states: torch.Tensor, + freqs: torch.Tensor, + ) -> torch.Tensor: + """Apply rotary embeddings (delegates to compiled ``_apply_rotary_emb``).""" + return _apply_rotary_emb(hidden_states, freqs) + + def _compute_frame_gates( + self, + x: torch.Tensor, + hw: tuple[int, int, int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute per-frame gates shared across spatial positions. + + Delegates to the module-level compiled ``_compute_frame_gates``. + """ + T, H, W = hw + S = H * W + return _compute_frame_gates( + x, + T, + S, + self.heads, + self.beta_proj.weight, + self.beta_proj.bias, + self.gate_proj.weight, + self.gate_proj.bias, + self.dt_bias, + self.A_log, + ) + + @staticmethod + def _prepare_frame_valid_masks( + frame_valid_mask: torch.Tensor | None, + *, + B: int, + T: int, + S: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Convert frame-valid mask to token/beta/decay masks used by GDN blocks.""" + if frame_valid_mask is None: + return None, None, None + + m = frame_valid_mask + if m.ndim == 5: + # (B, 1, T, 1, 1) + m = m[:, 0, :, 0, 0] + elif m.ndim == 3 and m.shape[1] == 1: + # (B, 1, T) + m = m[:, 0, :] + elif m.ndim != 2: + raise ValueError( + "frame_valid_mask must be shaped (B, 1, T, 1, 1), (B, 1, T), or (B, T); " + f"got shape={list(frame_valid_mask.shape)}" + ) + + if m.shape[0] != B or m.shape[1] != T: + raise ValueError(f"frame_valid_mask shape mismatch: expected (B={B}, T={T}), got {list(m.shape)}") + + m = m.to(device=device, dtype=dtype) + token_valid_mask = m[:, :, None].expand(B, T, S).reshape(B, T * S) + beta_valid_mask = m.view(B, 1, T, 1) + decay_valid_mask = m.view(B, 1, T) + return token_valid_mask, beta_valid_mask, decay_valid_mask + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + apply_output_gate: bool = True, + **kwargs: object, + ) -> torch.Tensor: + """Apply GDN attention to a token sequence. + + Args: + x: Input tensor of shape (B, N, C). + mask: Unused attention mask (kept for API compatibility). + HW: Tuple of (T, H, W) describing the token layout. + rotary_emb: Optional rotary embeddings for q/k. + block_mask: Unused block mask (kept for API compatibility). + apply_output_gate: When False, return raw attention output + before output gate and projection. + **kwargs: Unused extra arguments. + + Returns: + Tensor of shape (B, N, C) after attention and projection. + """ + del mask, block_mask + frame_valid_mask = kwargs.get("frame_valid_mask", None) + + if HW is None: + raise ValueError("HW (T, H, W) must be provided for GDN attention.") + + B, N, C = x.shape + T, H, W = HW + S = H * W + token_valid_mask, beta_valid_mask, decay_valid_mask = self._prepare_frame_valid_masks( + frame_valid_mask, + B=B, + T=T, + S=S, + device=x.device, + dtype=x.dtype, + ) + if token_valid_mask is not None: + x = x * token_valid_mask.view(B, N, 1) + + # Projections. + qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dim) + q, k, v = qkv.unbind(2) + if token_valid_mask is not None: + token_mask_bnhd = token_valid_mask.view(B, N, 1, 1) + q = q * token_mask_bnhd + k = k * token_mask_bnhd + v = v * token_mask_bnhd + + # Short convolution along T (before norm / kernel activation). + if self.conv_k is not None: + if self.conv_q is not None: + q = self._apply_temporal_short_conv(q.reshape(B, N, C), self.conv_q, HW).reshape( + B, N, self.heads, self.dim + ) + k = self._apply_temporal_short_conv(k.reshape(B, N, C), self.conv_k, HW).reshape( + B, N, self.heads, self.dim + ) + if self.conv_v is not None: + v = self._apply_temporal_short_conv(v.reshape(B, N, C), self.conv_v, HW).reshape( + B, N, self.heads, self.dim + ) + + # Apply Q/K norm on flattened channels (B, N, C) then reshape to heads. + q = self.q_norm(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + + # ReLU kernel. + q = self.kernel_func(q) + k = self.kernel_func(k) + + k_scale = self._key_scale(S) + k = k * k_scale + + # Permute to (B, H, D, N) for processing. + q = q.permute(0, 2, 3, 1) + k = k.permute(0, 2, 3, 1) + v = v.permute(0, 2, 3, 1) + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(B, 1, 1, N) + q = q * token_mask_qkv + k = k * token_mask_qkv + v = v * token_mask_qkv + + # RoPE preparation (numerator only). + if rotary_emb is not None: + q_rot = self._apply_rotary_emb(q, rotary_emb) + k_rot = self._apply_rotary_emb(k, rotary_emb) + else: + q_rot = q + k_rot = k + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(B, 1, 1, N) + q_rot = q_rot * token_mask_qkv + k_rot = k_rot * token_mask_qkv + + # Gate computation (use pre-computed gates when available to avoid + # redundant work in dual-branch CamCtrl models). + precomputed_gates = kwargs.get("precomputed_gates", None) + if precomputed_gates is not None: + beta, decay = precomputed_gates + else: + beta, decay = self._compute_frame_gates(x, HW) + if beta_valid_mask is not None: + beta = beta * beta_valid_mask.to(beta.dtype) + if decay_valid_mask is not None: + decay_m = decay_valid_mask.to(decay.dtype) + decay = decay * decay_m + (1.0 - decay_m) + + # Run the frame-wise GDN update. + # Force FP32 to preserve recurrent stability. + dtype_orig = x.dtype + recall_gate = self.recall_gate + q = q.float() + k = k.float() + v = v.float() + q_rot = q_rot.float() + k_rot = k_rot.float() + beta = beta.float() + decay = decay.float() + recall_gate = recall_gate.float() + + out = self.update_rule_func(q, k, v, q_rot, k_rot, beta, decay, recall_gate=recall_gate, eps=self.eps) + + # Reshape and project output. + if dtype_orig != torch.float32: + out = out.to(dtype_orig) + + out = out.permute(0, 3, 1, 2) + N_out = out.shape[1] + out = out.reshape(B, N_out, C) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N_out, 1).to(out.dtype) + + if apply_output_gate: + out = self._apply_output_gate(out, x) + out = self.proj(out.to(x.dtype)) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N_out, 1).to(out.dtype) + return out + return out + + +@_register_block() +class BidirectionalGDN(GDN): + """Bidirectional GDN attention with forward/backward fusion.""" + + def _apply_temporal_short_conv( + self, + x: torch.Tensor, + conv: ShortConvolution, + HW: tuple[int, int, int], + **kwargs: object, + ) -> torch.Tensor: + """Apply bidirectional (non-causal) ShortConvolution along T. + + Uses the forward+backward causal trick: run the causal conv in both directions and average, yielding a + symmetric temporal filter with a single set of weights. + + Args: + x: Input tensor of shape (B, N, C) where N = T * S. + conv: FLA ``ShortConvolution`` module. + HW: Tuple of (T, H, W) describing the token layout. + **kwargs: Unused. + + Returns: + Tensor of shape (B, N, C) after bidirectional temporal conv. + """ + del kwargs + + x, B, S, T = self._reshape_to_temporal(x, HW) + x = self._bidirectional_causal_conv_1d(x, conv) + return self._reshape_from_temporal(x, B, S, T) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + apply_output_gate: bool = True, + **kwargs: object, + ) -> torch.Tensor: + """Apply bidirectional GDN attention to a token sequence. + + Args: + x: Input tensor of shape (B, N, C). + mask: Unused attention mask (kept for API compatibility). + HW: Tuple of (T, H, W) describing the token layout. + rotary_emb: Optional rotary embeddings for q/k. + block_mask: Unused block mask (kept for API compatibility). + **kwargs: Unused extra arguments. + + Returns: + Tensor of shape (B, N, C) after attention and projection. + """ + del mask, block_mask + frame_valid_mask = kwargs.get("frame_valid_mask", None) + + if HW is None: + raise ValueError("HW (T, H, W) must be provided for GDN attention.") + + B, N, C = x.shape + T, H, W = HW + S = H * W + token_valid_mask, beta_valid_mask, decay_valid_mask = self._prepare_frame_valid_masks( + frame_valid_mask, + B=B, + T=T, + S=S, + device=x.device, + dtype=x.dtype, + ) + if token_valid_mask is not None: + x = x * token_valid_mask.view(B, N, 1) + + # Projections. + qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dim) + q, k, v = qkv.unbind(2) + if token_valid_mask is not None: + token_mask_bnhd = token_valid_mask.view(B, N, 1, 1) + q = q * token_mask_bnhd + k = k * token_mask_bnhd + v = v * token_mask_bnhd + + # Short convolution along T (before norm / kernel activation). + if self.conv_k is not None: + if self.conv_q is not None: + q = self._apply_temporal_short_conv(q.reshape(B, N, C), self.conv_q, HW).reshape( + B, N, self.heads, self.dim + ) + k = self._apply_temporal_short_conv(k.reshape(B, N, C), self.conv_k, HW).reshape( + B, N, self.heads, self.dim + ) + if self.conv_v is not None: + v = self._apply_temporal_short_conv(v.reshape(B, N, C), self.conv_v, HW).reshape( + B, N, self.heads, self.dim + ) + + # Apply Q/K norm on flattened channels (B, N, C) then reshape to heads. + q = self.q_norm(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + + # ReLU kernel. + q = self.kernel_func(q) + k = self.kernel_func(k) + + k_scale = self._key_scale(S) + k = k * k_scale + + # Permute to (B, H, D, N) for processing. + q = q.permute(0, 2, 3, 1) + k = k.permute(0, 2, 3, 1) + v = v.permute(0, 2, 3, 1) + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(B, 1, 1, N) + q = q * token_mask_qkv + k = k * token_mask_qkv + v = v * token_mask_qkv + + # RoPE preparation (numerator only). + if rotary_emb is not None: + q_rot = self._apply_rotary_emb(q, rotary_emb) + k_rot = self._apply_rotary_emb(k, rotary_emb) + else: + q_rot = q + k_rot = k + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(B, 1, 1, N) + q_rot = q_rot * token_mask_qkv + k_rot = k_rot * token_mask_qkv + + # Gate computation (use pre-computed gates when available). + precomputed_gates = kwargs.get("precomputed_gates", None) + if precomputed_gates is not None: + beta, decay = precomputed_gates + else: + beta, decay = self._compute_frame_gates(x, HW) + if beta_valid_mask is not None: + beta = beta * beta_valid_mask.to(beta.dtype) + if decay_valid_mask is not None: + decay_m = decay_valid_mask.to(decay.dtype) + decay = decay * decay_m + (1.0 - decay_m) + + H_eff = q.shape[1] + N_eff = q.shape[3] + T_eff = N_eff // S + + # Run the frame-wise GDN update. + # Force FP32 to preserve recurrent stability. + dtype_orig = x.dtype + recall_gate = self.recall_gate + q = q.float() + k = k.float() + v = v.float() + q_rot = q_rot.float() + k_rot = k_rot.float() + beta = beta.float() + decay = decay.float() + recall_gate = recall_gate.float() + + # Forward pass (inclusive: 1..t). + num_fwd, den_fwd = self.update_rule_func( + q, k, v, q_rot, k_rot, beta, decay, recall_gate=recall_gate, eps=self.eps, return_components=True + ) + + # Backward pass (exclusive: t+1..T). + def to_time_structure(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(B, H_eff, self.dim, T_eff, S).permute(0, 1, 3, 2, 4) + + def from_time_structure(tensor: torch.Tensor) -> torch.Tensor: + return tensor.permute(0, 1, 3, 2, 4).reshape(B, H_eff, self.dim, N_eff) + + q_T = to_time_structure(q) + k_T = to_time_structure(k) + v_T = to_time_structure(v) + q_rot_T = to_time_structure(q_rot) + k_rot_T = to_time_structure(k_rot) + + q_bwd = torch.flip(q_T, dims=[2]) + q_rot_bwd = torch.flip(q_rot_T, dims=[2]) + + k_bwd = flip_and_shift(k_T, dim=2, shift_val=0.0) + v_bwd = flip_and_shift(v_T, dim=2, shift_val=0.0) + k_rot_bwd = flip_and_shift(k_rot_T, dim=2, shift_val=0.0) + beta_bwd = flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = flip_and_shift(decay, dim=2, shift_val=1.0) + + k_bwd_flat = from_time_structure(k_bwd) + v_bwd_flat = from_time_structure(v_bwd) + q_bwd_flat = from_time_structure(q_bwd) + q_rot_bwd_flat = from_time_structure(q_rot_bwd) + k_rot_bwd_flat = from_time_structure(k_rot_bwd) + + num_bwd_flipped, den_bwd_flipped = self.update_rule_func( + q_bwd_flat, + k_bwd_flat, + v_bwd_flat, + q_rot_bwd_flat, + k_rot_bwd_flat, + beta_bwd, + decay_bwd, + recall_gate=recall_gate, + eps=self.eps, + return_components=True, + ) + + def flip_back(tensor: torch.Tensor) -> torch.Tensor: + d_actual = tensor.shape[2] + t_struct = tensor.view(B, H_eff, d_actual, T_eff, S) + return torch.flip(t_struct, dims=[3]).reshape(B, H_eff, d_actual, N_eff) + + num_bwd = flip_back(num_bwd_flipped) + den_bwd = flip_back(den_bwd_flipped) + + total_num = num_fwd + num_bwd + total_den = den_fwd + den_bwd + + out = total_num / (total_den + self.eps) + + # Reshape and project output. + if dtype_orig != torch.float32: + out = out.to(dtype_orig) + + out = out.permute(0, 3, 1, 2) + N_out = out.shape[1] + out = out.reshape(B, N_out, C) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N_out, 1).to(out.dtype) + + if apply_output_gate: + out = self._apply_output_gate(out, x) + out = self.proj(out.to(x.dtype)) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N_out, 1).to(out.dtype) + return out + return out + + +_frame_causal_mask_cache: dict[tuple[int, int, torch.device], torch.Tensor] = {} + + +def _get_frame_causal_mask(T: int, S: int, device: torch.device) -> torch.Tensor: + """Frame-wise block-causal mask: full attention within each frame, + causal across frames. + + Returns a boolean tensor of shape ``(1, 1, T*S, T*S)`` where ``True`` indicates positions that may attend. + """ + key = (T, S, device) + if key not in _frame_causal_mask_cache: + frame_idx = torch.arange(T, device=device).repeat_interleave(S) + mask = frame_idx.unsqueeze(1) >= frame_idx.unsqueeze(0) + _frame_causal_mask_cache[key] = mask.unsqueeze(0).unsqueeze(0) + return _frame_causal_mask_cache[key] + + +def _forward_softmax_attn( + self, + x: torch.Tensor, + HW: tuple[int, int, int], + rotary_emb: torch.Tensor | None, + frame_causal: bool, + apply_output_gate: bool = True, + **kwargs, +) -> torch.Tensor: + """Softmax attention (SDPA) reusing GDN parameters. + + Used by the hybrid GDN+Softmax architecture: every Nth block runs softmax attention instead of the gated-delta + recurrence. Reuses the parent block's QKV/q_norm/k_norm/proj for parameter compatibility. + """ + import torch.nn.functional as F + + B, N, C = x.shape + T, H, W = HW + S = H * W + + frame_valid_mask = kwargs.get("frame_valid_mask", None) + token_valid_mask, _, _ = GDN._prepare_frame_valid_masks( + frame_valid_mask, + B=B, + T=T, + S=S, + device=x.device, + dtype=x.dtype, + ) + if token_valid_mask is not None: + x = x * token_valid_mask.view(B, N, 1) + + qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dim) + q, k, v = qkv.unbind(2) + if token_valid_mask is not None: + m = token_valid_mask.view(B, N, 1, 1) + q, k, v = q * m, k * m, v * m + + q = self.q_norm(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + + if rotary_emb is not None: + q_perm = q.permute(0, 2, 3, 1) + k_perm = k.permute(0, 2, 3, 1) + q_perm = GDN._apply_rotary_emb(q_perm, rotary_emb) + k_perm = GDN._apply_rotary_emb(k_perm, rotary_emb) + q = q_perm.permute(0, 3, 1, 2) + k = k_perm.permute(0, 3, 1, 2) + + if token_valid_mask is not None: + m = token_valid_mask.view(B, N, 1, 1) + q, k, v = q * m, k * m, v * m + + q = q.transpose(1, 2) # (B, H, N, D) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + dtype_orig = x.dtype + if q.dtype == torch.float32: + q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16() + + attn_mask = _get_frame_causal_mask(T, S, x.device) if frame_causal else None + + out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + out = out.transpose(1, 2).reshape(B, N, C).to(dtype_orig) + + if apply_output_gate: + # Re-apply the parent's output projection w/ silu gate; some GDN + # variants split projection into proj_o + proj_gate; match those. + if hasattr(self, "proj_gate"): + out = out * F.silu(self.proj_gate(x)) + out = self.proj(out) + return out + + +# --------------------------------------------------------------------------- +# Base class +# --------------------------------------------------------------------------- + + +@torch.compile(dynamic=True) +def torch_chunk_cam_single_path_delta_rule( + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + chunk_size: int | None = 21, +) -> torch.Tensor: + """Parallel chunk-scan version of the single-path delta-rule recurrence. + + Restructured as a linear recurrence in D x D state space so that Phases 1 (transition-matrix construction) and 3 + (output projection) are fully parallel over T, while Phase 2 (the D x D state scan) is chunked and benefits from + ``@torch.compile``. + + The recurrence: + state[t] = state[t-1] * g[t] + delta_v[t] @ k_rot[t]^T + where delta_v[t] = (v[t] - state[t-1]*g[t] @ k_rot[t]) * beta[t] + + is equivalent to: + state[t] = state[t-1] @ W[t] + U[t] + with: + W[t] = g[t] * (I - beta[t] * k_rot[t] @ k_rot[t]^T) U[t] = beta[t] * v[t] @ k_rot[t]^T + """ + B, H, D, N = q_rot.shape + if beta.ndim not in (3, 4): + raise ValueError(f"Expected beta.ndim in (3, 4), got {beta.ndim}.") + T = beta.shape[2] + if T <= 0: + raise ValueError(f"Expected T > 0, got T={T}.") + if N % T != 0: + raise ValueError(f"Expected N divisible by T, got N={N}, T={T}.") + S = N // T + + def to_frame_seq(x: torch.Tensor) -> torch.Tensor: + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q_rot = to_frame_seq(q_rot) + k_rot = to_frame_seq(k_rot) + v = to_frame_seq(v) + + if beta.ndim == 4: + beta = beta.unsqueeze(3) + else: + beta = beta.view(B, H, T, 1, 1) + decay = decay.view(B, H, T, 1, 1) + + # ========================================================================= + # Phase 1: PARALLEL PRE-PROCESSING (fully parallel over T) + # ========================================================================= + I = torch.eye(D, device=q_rot.device, dtype=q_rot.dtype).view(1, 1, 1, D, D) + + k_rot_beta = k_rot * beta + W_kv = decay * (I - torch.matmul(k_rot_beta, k_rot.transpose(-1, -2))) + U_kv = torch.matmul(v * beta, k_rot.transpose(-1, -2)) + + # ========================================================================= + # Phase 2: CHUNKED SCAN over D x D state space + # ========================================================================= + valid_chunk_index, _ = normalize_chunk_index(None, T, chunk_size) + split_sizes = [valid_chunk_index[i + 1] - valid_chunk_index[i] for i in range(len(valid_chunk_index) - 1)] + + W_kv_c = W_kv.split(split_sizes, dim=2) + U_kv_c = U_kv.split(split_sizes, dim=2) + + S_kv = torch.zeros(B, H, D, D, device=q_rot.device, dtype=q_rot.dtype) + out_S_kv: list[torch.Tensor] = [] + + def _chunk_scan_kv( + w_kv: torch.Tensor, u_kv: torch.Tensor, s_kv: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + c_len = w_kv.shape[2] + s_kv_list: list[torch.Tensor] = [] + for t in range(c_len): + s_kv = torch.matmul(s_kv, w_kv[:, :, t]) + u_kv[:, :, t] + s_kv_list.append(s_kv) + return torch.stack(s_kv_list, dim=2), s_kv + + for i in range(len(split_sizes)): + s_kv_all, S_kv = _chunk_scan_kv(W_kv_c[i], U_kv_c[i], S_kv) + out_S_kv.append(s_kv_all) + + S_kv_all = torch.cat(out_S_kv, dim=2) + + # ========================================================================= + # Phase 3: PARALLEL OUTPUT PROJECTION (no denominator) + # ========================================================================= + out = torch.matmul(S_kv_all, q_rot) # (B, H, T, D, S) + + return out.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + + +class _GDNUCPEBase(GDN): + """Shared camera-branch logic for all GDN + UCPE variants. + + Adds a second attention branch whose positional encoding comes from UCPE per-ray camera transforms instead of the + standard RoPE used by the main branch. + + **Camera-specific parameters** (4 Linear layers per block): + ``q_proj_cam``, ``k_proj_cam``, ``v_proj_cam``, ``out_proj_cam`` + + **Shared with main branch** (no duplication): + QK norms, GDN gates (beta/gate/dt_bias/A_log/recall_gate), output gate, output projection. + + Requires ``cam_dim == in_dim`` and ``cam_heads == heads`` so that all shared parameters have matching dimensions. + + Subclasses only need to override ``_forward_cam_branch`` when the camera branch requires a different recurrence + pattern (e.g. bidirectional or chunk-causal). + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + *, + cam_dim: int, + cam_heads: int, + patch_size: tuple[int, int, int] = (1, 2, 2), + **kwargs: object, + ) -> None: + cam_update_rule_func: str = str(kwargs.pop("cam_update_rule_func", "torch_chunk")) + super().__init__(in_dim, out_dim, **kwargs) + + self.patch_size = patch_size + self.cam_dim = cam_dim + self.cam_heads = cam_heads + self.cam_head_dim = cam_dim // cam_heads + + chunk_gdn_chunk_size = kwargs.get("chunk_gdn_chunk_size", 21) + if cam_update_rule_func != "torch_chunk": + raise ValueError(f"Unsupported cam_update_rule_func: {cam_update_rule_func}") + self._cam_single_path_fn = partial( + torch_chunk_cam_single_path_delta_rule, + chunk_size=chunk_gdn_chunk_size, + ) + + if cam_dim != in_dim: + raise ValueError(f"Parameter sharing requires cam_dim == in_dim, got cam_dim={cam_dim}, in_dim={in_dim}.") + if cam_heads != self.heads: + raise ValueError( + f"Parameter sharing requires cam_heads == heads, got cam_heads={cam_heads}, heads={self.heads}." + ) + if self.cam_head_dim % 4 != 0: + raise ValueError( + "UCPE camera branch requires cam_head_dim divisible by 4, " + f"got {self.cam_head_dim} (cam_dim={cam_dim}, cam_heads={cam_heads})." + ) + + # ---- Camera-specific: QKV + output projections only ---- + self.q_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.k_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.v_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.out_proj_cam = nn.Linear(cam_dim, out_dim, bias=True) + + # Keep branch-specific Q/K norms so camera statistics do not disturb the + # main branch (and vice versa). Start from identical weights. + self.q_norm_cam = deepcopy(self.q_norm) + self.k_norm_cam = deepcopy(self.k_norm) + + nn.init.constant_(self.out_proj_cam.weight, 0) + nn.init.constant_(self.out_proj_cam.bias, 0) + + # Short convolutions for camera branch (matching base GDN variant). + if self.conv_kernel_size > 0: + self.conv_k_cam = ShortConvolution( + hidden_size=cam_dim, + kernel_size=self.conv_kernel_size, + activation=None, + ) + if self.k_conv_only: + self.conv_q_cam = None + self.conv_v_cam = None + else: + self.conv_q_cam = ShortConvolution( + hidden_size=cam_dim, + kernel_size=self.conv_kernel_size, + activation=None, + ) + self.conv_v_cam = ShortConvolution( + hidden_size=cam_dim, + kernel_size=self.conv_kernel_size, + activation=None, + ) + else: + self.conv_q_cam = None + self.conv_k_cam = None + self.conv_v_cam = None + + @staticmethod + def _downscale_to_reference_rms( + ref: torch.Tensor, + transformed: torch.Tensor, + eps: float = 1e-6, + ) -> torch.Tensor: + """Downscale transformed tensor if its channel RMS exceeds reference. + + Args: + ref: Reference tensor with target magnitude, shape (B, H, D, N). + transformed: Tensor to stabilize, shape (B, H, D, N). + eps: Numerical epsilon for RMS. + + Returns: + Stabilized tensor with per-(B,H,N) channel RMS not larger than ref. + """ + ref_rms = ref.square().mean(dim=2, keepdim=True).add(eps).sqrt() + tr_rms = transformed.square().mean(dim=2, keepdim=True).add(eps).sqrt() + scale = (ref_rms / tr_rms.clamp_min(eps)).clamp(max=1.0) + return transformed * scale + + def _stabilize_cam_transforms( + self, + q_cam: torch.Tensor, + k_cam: torch.Tensor, + v_cam: torch.Tensor, + q_cam_trans: torch.Tensor, + k_cam_trans: torch.Tensor, + v_cam_trans: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Optional post-UCPE stabilization hook for experimental variants.""" + del q_cam, k_cam, v_cam + return q_cam_trans, k_cam_trans, v_cam_trans + + # ------------------------------------------------------------------ + # Camera-branch building blocks + # ------------------------------------------------------------------ + + def _prepare_cam_qkv( + self, + x: torch.Tensor, + HW: tuple[int, int, int], + camera_conditions: torch.Tensor, + rotary_emb: torch.Tensor | None, + *, + token_valid_mask: torch.Tensor | None = None, + **kwargs: object, + ) -> tuple: + """Project camera QKV, apply short conv + QK norm + kernel + scaling + UCPE. + + The processing order mirrors the base GDN branch: + project -> mask -> short_conv -> QK_norm -> kernel -> scale -> permute -> UCPE + + Args: + token_valid_mask: Pre-computed mask of shape ``(B, N)`` from the + caller. Avoids redundant ``_prepare_frame_valid_masks`` calls. + + Returns: + (q_cam, k_cam, v_cam_trans, q_cam_trans, k_cam_trans, apply_fn_o, inflation_sq) + + All tensors are shaped ``(B, cam_heads, cam_head_dim, N)``. ``apply_fn_o`` is the UCPE inverse-output transform + closure. ``inflation_sq`` is the energy inflation factor of shape ``(B, cam_heads, 1, N)``. + """ + B, N, C = x.shape + T, H, W = HW + S = H * W + + # Pre-projection token masking (matching base branch). + if token_valid_mask is not None: + x = x * token_valid_mask.view(B, N, 1) + + # Fused camera QKV projection (1 GEMM instead of 3 kernel launches). + qkv_w = torch.cat([self.q_proj_cam.weight, self.k_proj_cam.weight, self.v_proj_cam.weight]) + qkv_b = torch.cat([self.q_proj_cam.bias, self.k_proj_cam.bias, self.v_proj_cam.bias]) + qkv_cam = F.linear(x, qkv_w, qkv_b) + q_cam, k_cam, v_cam = qkv_cam.chunk(3, dim=-1) + + # Post-projection token masking (before conv, matching base branch). + if token_valid_mask is not None: + token_mask = token_valid_mask.view(B, N, 1) + q_cam = q_cam * token_mask + k_cam = k_cam * token_mask + v_cam = v_cam * token_mask + + # Short convolution along T (before norm / kernel activation). + if self.conv_q_cam is not None: + q_cam = self._apply_temporal_short_conv(q_cam, self.conv_q_cam, HW, **kwargs) + if self.conv_k_cam is not None: + k_cam = self._apply_temporal_short_conv(k_cam, self.conv_k_cam, HW, **kwargs) + if self.conv_v_cam is not None: + v_cam = self._apply_temporal_short_conv(v_cam, self.conv_v_cam, HW, **kwargs) + + # Camera-specific QK normalization. + q_cam = self.q_norm_cam(q_cam).reshape(B, N, self.cam_heads, self.cam_head_dim) + k_cam = self.k_norm_cam(k_cam).reshape(B, N, self.cam_heads, self.cam_head_dim) + v_cam = v_cam.reshape(B, N, self.cam_heads, self.cam_head_dim) + + # ReLU kernel (shared). + q_cam = self.kernel_func(q_cam) + k_cam = self.kernel_func(k_cam) + + # FIXED: K scaling -- explicitly use ** for exponentiation! + k_scale = (self.cam_head_dim**-0.5) * (S**-0.5) + k_cam = k_cam * k_scale + + # Permute to (B, H, D, N) for GDN processing. + q_cam = q_cam.permute(0, 2, 3, 1).contiguous() + k_cam = k_cam.permute(0, 2, 3, 1).contiguous() + v_cam = v_cam.permute(0, 2, 3, 1).contiguous() + + # Measure safe geometric norm before UCPE applies translations + pre_ucpe_k_norm = torch.linalg.vector_norm(k_cam, dim=2, keepdim=True).clamp_min(1e-6) + + # UCPE per-ray transforms — reuse model-level cache when available + # to avoid recomputing _process_camera_conditions_ucpe per block. + cached_fns = kwargs.get("prope_fns", None) + if cached_fns is not None: + apply_fn_q, apply_fn_kv, apply_fn_o = cached_fns + else: + apply_fn_q, apply_fn_kv, apply_fn_o = prepare_prope_fns( + camctrl_type="UCPE", + head_dim=self.cam_head_dim, + camera_conditions=camera_conditions, + HW=HW, + patch_size=self.patch_size, + rotary_emb=rotary_emb, + ) + + # UCPE expects (B, h, N, d); our tensors are (B, h, d, N). + # Avoid eager contiguous copies before transforms, and fuse K/V transform + # into one call (same apply_fn_kv), then split back. + q_cam_trans = apply_fn_q(q_cam.transpose(-1, -2)).transpose(-1, -2).contiguous() + kv_cam = torch.cat([k_cam, v_cam], dim=1) + kv_cam_trans = apply_fn_kv(kv_cam.transpose(-1, -2)).transpose(-1, -2).contiguous() + k_cam_trans, v_cam_trans = torch.chunk(kv_cam_trans, chunks=2, dim=1) + + q_cam_trans, k_cam_trans, v_cam_trans = self._stabilize_cam_transforms( + q_cam=q_cam, + k_cam=k_cam, + v_cam=v_cam, + q_cam_trans=q_cam_trans, + k_cam_trans=k_cam_trans, + v_cam_trans=v_cam_trans, + ) + + # Measure inflated geometric norm after UCPE + post_ucpe_k_norm = torch.linalg.vector_norm(k_cam_trans, dim=2, keepdim=True).clamp_min(1e-6) + + # Calculate the squared inflation factor for beta discounting + inflation_sq = (post_ucpe_k_norm / pre_ucpe_k_norm) ** 2 + + return q_cam, k_cam, v_cam_trans, q_cam_trans, k_cam_trans, apply_fn_o, inflation_sq + + def _run_cam_gdn( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + ) -> torch.Tensor: + """Run the shared GDN kernel on camera-branch tensors. + + Uses shared ``self.recall_gate``. Handles FP32 casting. Returns ``num / (den + eps)`` shaped ``(B, H, D, N)``. + """ + recall_gate = self.recall_gate + q = q.float() + k = k.float() + v = v.float() + q_rot = q_rot.float() + k_rot = k_rot.float() + beta = beta.float() + decay = decay.float() + recall_gate = recall_gate.float() + + return self.update_rule_func( + q, + k, + v, + q_rot, + k_rot, + beta, + decay, + recall_gate=recall_gate, + eps=self.eps, + ) + + def _run_cam_gdn_components( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Like ``_run_cam_gdn`` but returns ``(num, den)`` components.""" + recall_gate = self.recall_gate + q = q.float() + k = k.float() + v = v.float() + q_rot = q_rot.float() + k_rot = k_rot.float() + beta = beta.float() + decay = decay.float() + recall_gate = recall_gate.float() + + return self.update_rule_func( + q, + k, + v, + q_rot, + k_rot, + beta, + decay, + recall_gate=recall_gate, + eps=self.eps, + return_components=True, + ) + + def _run_cam_single_path( + self, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + ) -> torch.Tensor: + """Run the numerator-only camera delta-rule recurrence (parallel chunk scan).""" + q_rot = q_rot.float() + k_rot = k_rot.float() + v = v.float() + beta = beta.float() + decay = decay.float() + return self._cam_single_path_fn(q_rot, k_rot, v, beta, decay) + + # ------------------------------------------------------------------ + # Camera-branch forward (forward-only causal -- default) + # ------------------------------------------------------------------ + + def _forward_cam_branch( + self, + x: torch.Tensor, + HW: tuple[int, int, int], + camera_conditions: torch.Tensor, + rotary_emb: torch.Tensor | None, + **kwargs: object, + ) -> torch.Tensor: + """Forward-only causal GDN camera branch with UCPE transforms. + + Subclasses override this for bidirectional / chunk-causal variants. + + Returns raw attention output ``(B, N, C)`` -- no output gate or projection applied (those are shared and + applied in ``forward()``). + """ + B, N, _ = x.shape + T, H, W = HW + S = H * W + dtype_orig = x.dtype + + # Compute masks once; pass token_valid_mask to _prepare_cam_qkv for + # pre-conv masking and reuse here for post-UCPE masking + gate masking. + token_valid_mask, beta_valid_mask, decay_valid_mask = self._prepare_frame_valid_masks( + kwargs.get("frame_valid_mask", None), + B=B, + T=T, + S=S, + device=x.device, + dtype=x.dtype, + ) + + q_cam, k_cam, v_cam_trans, q_cam_trans, k_cam_trans, apply_fn_o, inflation_sq = self._prepare_cam_qkv( + x, + HW, + camera_conditions, + rotary_emb, + token_valid_mask=token_valid_mask, + **kwargs, + ) + + # Re-mask after UCPE transforms (which can reintroduce non-zero values). + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(B, 1, 1, N) + q_cam = q_cam * token_mask_qkv + k_cam = k_cam * token_mask_qkv + v_cam_trans = v_cam_trans * token_mask_qkv + q_cam_trans = q_cam_trans * token_mask_qkv + k_cam_trans = k_cam_trans * token_mask_qkv + + # Shared GDN gates (use pre-computed when available). + precomputed_gates = kwargs.get("precomputed_gates", None) + if precomputed_gates is not None: + beta, decay = precomputed_gates + else: + beta, decay = self._compute_frame_gates(x, HW) + + # Dynamic Beta Discounting: scale beta by UCPE inflation factor. + inflation_sq_spatial = inflation_sq.view(B, self.cam_heads, T, S) + frame_inflation_sq = inflation_sq_spatial.mean(dim=-1) + if beta.ndim == 3: + beta = beta / frame_inflation_sq.clamp_min(1.0) + elif beta.ndim == 4: + beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + + if beta_valid_mask is not None: + beta = beta * beta_valid_mask.to(beta.dtype) + if decay_valid_mask is not None: + decay_m = decay_valid_mask.to(decay.dtype) + decay = decay * decay_m + (1.0 - decay_m) + + out = self._run_cam_gdn( + q_cam, + k_cam, + v_cam_trans, + q_cam_trans, + k_cam_trans, + beta, + decay, + ) + + if dtype_orig != torch.float32: + out = out.to(dtype_orig) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, 1, 1, N).to(out.dtype) + + # Inverse UCPE transform on output. + out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous() + out = out.reshape(B, self.cam_dim, N).permute(0, 2, 1) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N, 1).to(out.dtype) + return out + + # ------------------------------------------------------------------ + # Full forward + # ------------------------------------------------------------------ + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + camera_conditions: torch.Tensor | None = None, + chunk_size: int | None = None, + **kwargs: object, + ) -> torch.Tensor: + """Dual-branch forward: GDN main + UCPE camera. + + Flow: + 1. main_raw = GDN attention (no gate/proj) + 2. cam_raw = GDN+UCPE attention (no gate/proj) + 3. combined = main_raw + out_proj_cam(cam_raw) [zero at init] + 4. output = proj(output_gate(combined)) [shared, once] + """ + # Pre-compute shared gates once for both branches. + if HW is not None: + precomputed_gates = self._compute_frame_gates(x, HW) + else: + precomputed_gates = None + + # Main branch -- raw attention without gate/proj. + main_raw = super().forward( + x, + mask=mask, + HW=HW, + rotary_emb=rotary_emb, + block_mask=block_mask, + apply_output_gate=False, + chunk_size=chunk_size, + precomputed_gates=precomputed_gates, + **kwargs, + ) + + # Camera branch. + cam_contrib: torch.Tensor | int = 0 + if camera_conditions is not None: + if HW is None: + raise ValueError("HW (T, H, W) must be provided for UCPE camera branch.") + cam_raw = self._forward_cam_branch( + x, + HW, + camera_conditions, + rotary_emb, + chunk_size=chunk_size, + precomputed_gates=precomputed_gates, + **kwargs, + ) + cam_contrib = self.out_proj_cam(cam_raw) + + # Combine, then shared gate + projection (applied once). + combined = main_raw + cam_contrib + combined = self._apply_output_gate(combined, x) + return self.proj(combined.to(x.dtype)) + + +# --------------------------------------------------------------------------- +# Concrete variants +# --------------------------------------------------------------------------- + + +class BidirectionalGDNUCPELiteLA(_GDNUCPEBase, BidirectionalGDN): + """Bidirectional GDN with UCPE camera conditioning. + + Main branch: bidirectional GDN (inherited from ``BidirectionalGDN``). Camera branch: bidirectional GDN with UCPE + transforms. + """ + + def _forward_cam_branch( + self, + x: torch.Tensor, + HW: tuple[int, int, int], + camera_conditions: torch.Tensor, + rotary_emb: torch.Tensor | None, + **kwargs: object, + ) -> torch.Tensor: + B, N, C = x.shape + T, H, W = HW + S = H * W + dtype_orig = x.dtype + + token_valid_mask, beta_valid_mask, decay_valid_mask = self._prepare_frame_valid_masks( + kwargs.get("frame_valid_mask", None), + B=B, + T=T, + S=S, + device=x.device, + dtype=x.dtype, + ) + + q_cam, k_cam, v_cam_trans, q_cam_trans, k_cam_trans, apply_fn_o, inflation_sq = self._prepare_cam_qkv( + x, + HW, + camera_conditions, + rotary_emb, + token_valid_mask=token_valid_mask, + **kwargs, + ) + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(B, 1, 1, N) + q_cam = q_cam * token_mask_qkv + k_cam = k_cam * token_mask_qkv + v_cam_trans = v_cam_trans * token_mask_qkv + q_cam_trans = q_cam_trans * token_mask_qkv + k_cam_trans = k_cam_trans * token_mask_qkv + + # Shared GDN gates (use pre-computed when available). + precomputed_gates = kwargs.get("precomputed_gates", None) + if precomputed_gates is not None: + beta, decay = precomputed_gates + else: + beta, decay = self._compute_frame_gates(x, HW) + + # Dynamic Beta Discounting: scale beta by UCPE inflation factor. + inflation_sq_spatial = inflation_sq.view(B, self.cam_heads, T, S) + frame_inflation_sq = inflation_sq_spatial.mean(dim=-1) + if beta.ndim == 3: + beta = beta / frame_inflation_sq.clamp_min(1.0) + elif beta.ndim == 4: + beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + + if beta_valid_mask is not None: + beta = beta * beta_valid_mask.to(beta.dtype) + if decay_valid_mask is not None: + decay_m = decay_valid_mask.to(decay.dtype) + decay = decay * decay_m + (1.0 - decay_m) + + H_heads = self.cam_heads + D_head = self.cam_head_dim + + # -- Forward pass (inclusive 1..t) -- + num_fwd, den_fwd = self._run_cam_gdn_components( + q_cam, + k_cam, + v_cam_trans, + q_cam_trans, + k_cam_trans, + beta, + decay, + ) + + # -- Backward pass (exclusive t+1..T) -- + def to_time(t: torch.Tensor) -> torch.Tensor: + return t.view(B, H_heads, D_head, T, S).permute(0, 1, 3, 2, 4) + + def from_time(t: torch.Tensor) -> torch.Tensor: + return t.permute(0, 1, 3, 2, 4).reshape(B, H_heads, D_head, N) + + q_T = to_time(q_cam) + k_T = to_time(k_cam) + v_T = to_time(v_cam_trans) + q_rot_T = to_time(q_cam_trans) + k_rot_T = to_time(k_cam_trans) + + q_bwd = torch.flip(q_T, dims=[2]) + q_rot_bwd = torch.flip(q_rot_T, dims=[2]) + k_bwd = flip_and_shift(k_T, dim=2, shift_val=0.0) + v_bwd = flip_and_shift(v_T, dim=2, shift_val=0.0) + k_rot_bwd = flip_and_shift(k_rot_T, dim=2, shift_val=0.0) + beta_bwd = flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = flip_and_shift(decay, dim=2, shift_val=1.0) + + num_bwd_f, den_bwd_f = self._run_cam_gdn_components( + from_time(q_bwd), + from_time(k_bwd), + from_time(v_bwd), + from_time(q_rot_bwd), + from_time(k_rot_bwd), + beta_bwd, + decay_bwd, + ) + + def flip_back(tensor: torch.Tensor) -> torch.Tensor: + d = tensor.shape[2] + return torch.flip( + tensor.view(B, H_heads, d, T, S), + dims=[3], + ).reshape(B, H_heads, d, N) + + num_bwd = flip_back(num_bwd_f) + den_bwd = flip_back(den_bwd_f) + out = (num_fwd + num_bwd) / (den_fwd + den_bwd + self.eps) + + if dtype_orig != torch.float32: + out = out.to(dtype_orig) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, 1, 1, N).to(out.dtype) + + out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous() + out = out.reshape(B, self.cam_dim, N).permute(0, 2, 1) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N, 1).to(out.dtype) + return out + + +class BidirectionalGDNUCPELiteLAPostUCPERenorm(BidirectionalGDNUCPELiteLA): + """Bidirectional GDNUCPE with post-UCPE RMS downscaling. + + The raw UCPE transforms are still measured for debug logging, but the transformed camera tensors are downscaled + back to their pre-UCPE RMS envelope before they enter the recurrence. + """ + + def _stabilize_cam_transforms( + self, + q_cam: torch.Tensor, + k_cam: torch.Tensor, + v_cam: torch.Tensor, + q_cam_trans: torch.Tensor, + k_cam_trans: torch.Tensor, + v_cam_trans: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_cam_trans = self._downscale_to_reference_rms(q_cam, q_cam_trans) + k_cam_trans = self._downscale_to_reference_rms(k_cam, k_cam_trans) + v_cam_trans = self._downscale_to_reference_rms(v_cam, v_cam_trans) + return q_cam_trans, k_cam_trans, v_cam_trans + + +@_register_block() +class BidirectionalGDNUCPESinglePathLiteLA(BidirectionalGDNUCPELiteLAPostUCPERenorm): + """Bidirectional UCPE camera branch with numerator-only delta-rule updates. + + This is an experimental ablation that keeps the main branch unchanged, applies UCPE plus post-UCPE RMS downscaling + on the camera tensors, and replaces the camera branch's ``num / den`` recurrence with a single-path delta rule over + the transformed camera stream only. + """ + + def _forward_cam_branch( + self, + x: torch.Tensor, + HW: tuple[int, int, int], + camera_conditions: torch.Tensor, + rotary_emb: torch.Tensor | None, + **kwargs: object, + ) -> torch.Tensor: + B, N, _ = x.shape + T, H, W = HW + S = H * W + dtype_orig = x.dtype + + token_valid_mask, beta_valid_mask, decay_valid_mask = self._prepare_frame_valid_masks( + kwargs.get("frame_valid_mask", None), + B=B, + T=T, + S=S, + device=x.device, + dtype=x.dtype, + ) + + q_cam, _, v_cam_trans, q_cam_trans, k_cam_trans, apply_fn_o, inflation_sq = self._prepare_cam_qkv( + x, + HW, + camera_conditions, + rotary_emb, + token_valid_mask=token_valid_mask, + **kwargs, + ) + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(B, 1, 1, N) + q_cam = q_cam * token_mask_qkv + v_cam_trans = v_cam_trans * token_mask_qkv + q_cam_trans = q_cam_trans * token_mask_qkv + k_cam_trans = k_cam_trans * token_mask_qkv + + precomputed_gates = kwargs.get("precomputed_gates", None) + if precomputed_gates is not None: + beta, decay = precomputed_gates + else: + beta, decay = self._compute_frame_gates(x, HW) + + inflation_sq_spatial = inflation_sq.view(B, self.cam_heads, T, S) + frame_inflation_sq = inflation_sq_spatial.mean(dim=-1) + if beta.ndim == 3: + beta = beta / frame_inflation_sq.clamp_min(1.0) + elif beta.ndim == 4: + beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + + if beta_valid_mask is not None: + beta = beta * beta_valid_mask.to(beta.dtype) + if decay_valid_mask is not None: + decay_m = decay_valid_mask.to(decay.dtype) + decay = decay * decay_m + (1.0 - decay_m) + + H_heads = self.cam_heads + D_head = self.cam_head_dim + out_fwd = self._run_cam_single_path( + q_cam_trans, + k_cam_trans, + v_cam_trans, + beta, + decay, + ) + + def to_time(t: torch.Tensor) -> torch.Tensor: + return t.view(B, H_heads, D_head, T, S).permute(0, 1, 3, 2, 4) + + def from_time(t: torch.Tensor) -> torch.Tensor: + return t.permute(0, 1, 3, 2, 4).reshape(B, H_heads, D_head, N) + + q_rot_T = to_time(q_cam_trans) + k_rot_T = to_time(k_cam_trans) + v_T = to_time(v_cam_trans) + + q_rot_bwd = torch.flip(q_rot_T, dims=[2]) + k_rot_bwd = flip_and_shift(k_rot_T, dim=2, shift_val=0.0) + v_bwd = flip_and_shift(v_T, dim=2, shift_val=0.0) + beta_bwd = flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = flip_and_shift(decay, dim=2, shift_val=1.0) + + out_bwd_f = self._run_cam_single_path( + from_time(q_rot_bwd), + from_time(k_rot_bwd), + from_time(v_bwd), + beta_bwd, + decay_bwd, + ) + + out_bwd = torch.flip( + out_bwd_f.view(B, H_heads, D_head, T, S), + dims=[3], + ).reshape(B, H_heads, D_head, N) + out = out_fwd + out_bwd + + if dtype_orig != torch.float32: + out = out.to(dtype_orig) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, 1, 1, N).to(out.dtype) + + out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous() + out = out.reshape(B, self.cam_dim, N).permute(0, 2, 1) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N, 1).to(out.dtype) + return out + + +def _prepare_cam_qkv_softmax( + self, + x: torch.Tensor, + HW: tuple, + camera_conditions: torch.Tensor, + rotary_emb: torch.Tensor | None, + *, + token_valid_mask: torch.Tensor | None = None, + **kwargs, +) -> tuple: + """Camera branch Q/K/V for softmax attention. + + Mirrors ``_GDNUCPEBase._prepare_cam_qkv`` but skips the ReLU kernel and GDN key scaling — standard softmax SDPA + provides its own 1/sqrt(d_k). Returns ``(q, k, v, apply_fn_o)`` shaped ``(B, cam_heads, cam_head_dim, N)``. + """ + B, N, C = x.shape + + if token_valid_mask is not None: + x = x * token_valid_mask.view(B, N, 1) + + qkv_w = torch.cat([self.q_proj_cam.weight, self.k_proj_cam.weight, self.v_proj_cam.weight]) + qkv_b = torch.cat([self.q_proj_cam.bias, self.k_proj_cam.bias, self.v_proj_cam.bias]) + qkv_cam = F.linear(x, qkv_w, qkv_b) + q_cam, k_cam, v_cam = qkv_cam.chunk(3, dim=-1) + + if token_valid_mask is not None: + m = token_valid_mask.view(B, N, 1) + q_cam, k_cam, v_cam = q_cam * m, k_cam * m, v_cam * m + + if self.conv_q_cam is not None: + q_cam = self._apply_temporal_short_conv(q_cam, self.conv_q_cam, HW, **kwargs) + if self.conv_k_cam is not None: + k_cam = self._apply_temporal_short_conv(k_cam, self.conv_k_cam, HW, **kwargs) + if self.conv_v_cam is not None: + v_cam = self._apply_temporal_short_conv(v_cam, self.conv_v_cam, HW, **kwargs) + + q_cam = self.q_norm_cam(q_cam).reshape(B, N, self.cam_heads, self.cam_head_dim) + k_cam = self.k_norm_cam(k_cam).reshape(B, N, self.cam_heads, self.cam_head_dim) + v_cam = v_cam.reshape(B, N, self.cam_heads, self.cam_head_dim) + + q_cam = q_cam.permute(0, 2, 3, 1).contiguous() + k_cam = k_cam.permute(0, 2, 3, 1).contiguous() + v_cam = v_cam.permute(0, 2, 3, 1).contiguous() + + cached_fns = kwargs.get("prope_fns", None) + if cached_fns is not None: + apply_fn_q, apply_fn_kv, apply_fn_o = cached_fns + else: + apply_fn_q, apply_fn_kv, apply_fn_o = prepare_prope_fns( + camctrl_type="UCPE", + head_dim=self.cam_head_dim, + camera_conditions=camera_conditions, + HW=HW, + patch_size=self.patch_size, + rotary_emb=rotary_emb, + ) + + q_cam_trans = apply_fn_q(q_cam.transpose(-1, -2)).transpose(-1, -2).contiguous() + kv_cam = torch.cat([k_cam, v_cam], dim=1) + kv_cam_trans = apply_fn_kv(kv_cam.transpose(-1, -2)).transpose(-1, -2).contiguous() + k_cam_trans, v_cam_trans = torch.chunk(kv_cam_trans, chunks=2, dim=1) + + q_cam_trans, k_cam_trans, v_cam_trans = self._stabilize_cam_transforms( + q_cam=q_cam, + k_cam=k_cam, + v_cam=v_cam, + q_cam_trans=q_cam_trans, + k_cam_trans=k_cam_trans, + v_cam_trans=v_cam_trans, + ) + return q_cam_trans, k_cam_trans, v_cam_trans, apply_fn_o + + +def _forward_cam_branch_softmax( + self, + x: torch.Tensor, + HW: tuple, + camera_conditions: torch.Tensor, + rotary_emb: torch.Tensor | None, + frame_causal: bool, + **kwargs, +) -> torch.Tensor: + """Bidirectional softmax camera branch (with UCPE transforms). + + Uses ``F.scaled_dot_product_attention`` with optional invalid-key masking. + """ + B, N, _ = x.shape + T, H, W = HW + S = H * W + + token_valid_mask, _, _ = self._prepare_frame_valid_masks( + kwargs.get("frame_valid_mask", None), + B=B, + T=T, + S=S, + device=x.device, + dtype=x.dtype, + ) + + q_cam_trans, k_cam_trans, v_cam_trans, apply_fn_o = _prepare_cam_qkv_softmax( + self, + x, + HW, + camera_conditions, + rotary_emb, + token_valid_mask=token_valid_mask, + **kwargs, + ) + + if token_valid_mask is not None: + m = token_valid_mask.view(B, 1, 1, N) + q_cam_trans, v_cam_trans = q_cam_trans * m, v_cam_trans * m + + q_sdpa = q_cam_trans.transpose(-1, -2) + k_sdpa = k_cam_trans.transpose(-1, -2) + v_sdpa = v_cam_trans.transpose(-1, -2) + + dtype_orig = x.dtype + q_sdpa, k_sdpa, v_sdpa = q_sdpa.float(), k_sdpa.float(), v_sdpa.float() + # SDPA / FlashAttention only supports bf16/fp16; fp32 falls back to math backend. + if q_sdpa.dtype == torch.float32: + q_sdpa, k_sdpa, v_sdpa = q_sdpa.bfloat16(), k_sdpa.bfloat16(), v_sdpa.bfloat16() + + invalid_kv_logit_bias = None + if token_valid_mask is not None and not bool(token_valid_mask.all()): + invalid_kv_logit_bias = torch.where( + token_valid_mask.bool().view(B, 1, 1, -1), + torch.zeros((), dtype=q_sdpa.dtype, device=q_sdpa.device), + torch.full((), -1e9, dtype=q_sdpa.dtype, device=q_sdpa.device), + ) + + # FlashAttention-2 only supports head_dim in {32, 64, 128, 256}. + D = q_sdpa.shape[-1] + _need_pad = D not in (32, 64, 128, 256) and D < 256 + if _need_pad: + _pad_to = 128 if D <= 128 else 256 + _pad_size = _pad_to - D + q_sdpa = F.pad(q_sdpa, (0, _pad_size)) + k_sdpa = F.pad(k_sdpa, (0, _pad_size)) + v_sdpa = F.pad(v_sdpa, (0, _pad_size)) + out = F.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=invalid_kv_logit_bias) + if _need_pad: + out = out[..., :D] + + out = out.transpose(-1, -2) + if out.dtype != dtype_orig: + out = out.to(dtype_orig) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, 1, 1, N).to(out.dtype) + out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous() + out = out.reshape(B, self.cam_dim, N).permute(0, 2, 1) + if token_valid_mask is not None: + out = out * token_valid_mask.view(B, N, 1).to(out.dtype) + return out + + +class _SoftmaxUCPESinglePathLiteLA( + BidirectionalGDNUCPESinglePathLiteLA, +): + """Softmax attention with UCPE camera conditioning (single-path). + + Replaces GDN recurrence with ``F.scaled_dot_product_attention``. Automatically selects the correct masking mode + based on ``chunk_size``: + + - ``chunk_size is None`` or ``chunk_size >= T``: full bidirectional (no mask) + - ``chunk_size < T``: chunk-causal (full within chunks, causal across) + + All parameters match the GDN variants for checkpoint compatibility. GDN-specific parameters are present but unused + in forward. + """ + + def __init__(self, *args, conv_kernel_size: int = 0, **kwargs): + super().__init__(*args, conv_kernel_size=0, **kwargs) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + camera_conditions: torch.Tensor | None = None, + chunk_size: int | None = None, + **kwargs: object, + ) -> torch.Tensor: + main_raw = _forward_softmax_attn( + self, + x, + HW, + rotary_emb, + frame_causal=False, + apply_output_gate=False, + chunk_size=chunk_size, + **kwargs, + ) + + cam_contrib: torch.Tensor | int = 0 + if camera_conditions is not None: + if HW is None: + raise ValueError("HW must be provided for UCPE camera branch.") + cam_raw = _forward_cam_branch_softmax( + self, + x, + HW, + camera_conditions, + rotary_emb, + frame_causal=False, + chunk_size=chunk_size, + **kwargs, + ) + cam_contrib = self.out_proj_cam(cam_raw) + + combined = main_raw + cam_contrib + combined = self._apply_output_gate(combined, x) + return self.proj(combined.to(x.dtype)) + + +# Name used by the `camctrl_type` config string and the block-name mappings below. +BidirectionalSoftmaxUCPESinglePathLiteLA = _SoftmaxUCPESinglePathLiteLA + + +@_register_block() +class BidirectionalGDNTriton(BidirectionalGDN): + """Bidirectional GDN with a fused Triton scan. + + Subclasses :class:`BidirectionalGDN` and only overrides :meth:`forward`. Every learned sub-module (``qkv``, + ``proj``, ``q_norm``, ``k_norm``, ``conv_k``, ``beta_proj``, ``gate_proj``, ``A_log``, ``dt_bias``, + ``output_gate``) and helper (``_apply_temporal_short_conv``, ``_compute_frame_gates``, ``_apply_output_gate``) is + inherited unchanged so existing checkpoints load with zero conversion. + """ + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + apply_output_gate: bool = True, + **kwargs: object, + ) -> torch.Tensor: + # ---- Guards: this path supports inference only. ------------------- + if HW is None: + raise ValueError("BidirectionalGDNTriton requires HW=(T, H, W).") + del mask, block_mask # unused in the bidirectional Triton path + if kwargs.get("frame_valid_mask", None) is not None: + raise NotImplementedError( + "BidirectionalGDNTriton does not support frame_valid_mask (training-only feature)." + ) + if self.conv_q is not None or self.conv_v is not None: + raise NotImplementedError("BidirectionalGDNTriton requires k_conv_only=True; got conv_q or conv_v.") + + B, N, C = x.shape + T, H_s, W_s = HW + S = H_s * W_s + H, D = self.heads, self.dim + if N != T * S: + raise ValueError(f"N={N} != T*S={T * S} for HW={HW}.") + if C != H * D: + raise ValueError(f"C={C} != heads*dim={H * D}.") + + # ---- 1. QKV projection -> (B, N, 3, H, D), kept contiguous. ------- + qkv = self.qkv(x).reshape(B, N, 3, H, D) + + # ---- 2. Bidirectional short conv on K (parent method). ---------- + # ``BidirectionalGDN._apply_temporal_short_conv`` runs the causal + # conv forward + backward then averages, giving a symmetric filter + # with one set of weights. Inherited unchanged. + if self.conv_k is not None: + k_raw = qkv[:, :, 1].contiguous().reshape(B, N, C) + k_conv = self._apply_temporal_short_conv(k_raw, self.conv_k, HW) + qkv[:, :, 1].copy_(k_conv.reshape(B, N, H, D)) + + # ---- 3. Frame gates (precomputed when shared with cam branch). ---- + precomputed_gates = kwargs.get("precomputed_gates", None) + if precomputed_gates is not None: + beta, decay = precomputed_gates + else: + beta, decay = self._compute_frame_gates(x, HW) + beta = beta.contiguous() + decay = decay.contiguous() + + # ---- 4. Full-channel RMSNorm weights. ----------------------------- + if not isinstance(self.q_norm, nn.Identity): + q_nw = self.q_norm.weight.float().contiguous() + k_nw = self.k_norm.weight.float().contiguous() + norm_eps = float(getattr(self.q_norm, "eps", 1e-5)) + else: + q_nw = torch.ones(C, device=x.device, dtype=torch.float32) + k_nw = torch.ones(C, device=x.device, dtype=torch.float32) + norm_eps = 1e-5 + + # ---- 5. Fused Q+K inverse-RMS (single Triton launch). ------------- + q_inv_rms, k_inv_rms = fused_qk_inv_rms(qkv, eps=norm_eps) + + # ---- 6. Expanded RoPE cos/sin tables (N, D). --------------------- + rope_cos, rope_sin = prepare_rope_tables(rotary_emb, N, D, x.device) + + # ---- 7. K scale absorbs Q/K^T variance + spatial mean-pool. ----- + k_scale = (D**-0.5) * (S**-0.5) + + # ---- 8. Fused bidirectional Triton scan over the full sequence. -- + # No ``*_bwd`` overrides: the kernel's ``reverse=True`` path already + # implements the exclusive (t+1..T) reverse recurrence, matching the + # torch ``flip_and_shift`` semantics used in ``BidirectionalGDN``. + out = fused_bigdn_func( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight=q_nw, + k_norm_weight=k_nw, + rope_cos=rope_cos, + rope_sin=rope_sin, + beta=beta, + decay=decay, + F=T, + S=S, + k_scale=k_scale, + eps=self.eps, + ) # (B, N, H, D) + + # ---- 9. Output gate + projection. -------------------------------- + out = out.reshape(B, N, C) + if apply_output_gate: + out = self._apply_output_gate(out, x) + out = self.proj(out.to(x.dtype)) + return out + + +@_register_block() +class BidirectionalGDNUCPESinglePathLiteLATriton(BidirectionalGDNUCPESinglePathLiteLA): + """Bidirectional UCPE camera-controlled GDN with a Triton main branch. + + Inherits the entire camera branch (``_forward_cam_branch``), ``_prepare_cam_qkv``, every sub-module and every + checkpoint key from :class:`BidirectionalGDNUCPESinglePathLiteLA`. The **only** behavioural delta is that the + main-branch GDN scan dispatches through :class:`BidirectionalGDNTriton.forward` instead of the inherited + :class:`BidirectionalGDN.forward`. + + Because ``_GDNUCPEBase.forward`` routes the main branch via ``super().forward(...)`` — which MRO-resolves to + :class:`BidirectionalGDN`, not our Triton variant — we re-implement the dual-branch forward here to explicitly call + ``BidirectionalGDNTriton.forward(self, ...)``. The body is otherwise bit-identical to the parent's ``forward``. + + The cam branch is the inherited torch path; use :class:`BidirectionalGDNUCPESinglePathLiteLABothTriton` for a fully + Triton cam branch. + """ + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + camera_conditions: torch.Tensor | None = None, + chunk_size: int | None = None, + **kwargs: object, + ) -> torch.Tensor: + # Pre-compute shared gates once for both branches. + if HW is not None: + precomputed_gates = self._compute_frame_gates(x, HW) + else: + precomputed_gates = None + + # Main branch — Triton-fused bidirectional scan. + main_raw = BidirectionalGDNTriton.forward( + self, + x, + mask=mask, + HW=HW, + rotary_emb=rotary_emb, + block_mask=block_mask, + apply_output_gate=False, + chunk_size=chunk_size, + precomputed_gates=precomputed_gates, + **kwargs, + ) + + # Camera branch (inherited torch implementation). + cam_contrib: torch.Tensor | int = 0 + if camera_conditions is not None: + if HW is None: + raise ValueError("HW (T, H, W) must be provided for UCPE camera branch.") + cam_raw = self._forward_cam_branch( + x, + HW, + camera_conditions, + rotary_emb, + chunk_size=chunk_size, + precomputed_gates=precomputed_gates, + **kwargs, + ) + cam_contrib = self.out_proj_cam(cam_raw) + + combined = main_raw + cam_contrib + combined = self._apply_output_gate(combined, x) + return self.proj(combined.to(x.dtype)) + + +@_register_block() +class BidirectionalGDNUCPESinglePathLiteLABothTriton(BidirectionalGDNUCPESinglePathLiteLATriton): + """Bidirectional UCPE camera-controlled GDN with **both** branches on Triton. + + Subclasses :class:`BidirectionalGDNUCPESinglePathLiteLATriton` (which already rewires the main GDN scan) and + replaces :meth:`_forward_cam_branch` with a fused Triton camera pipeline: + + 1. Torch QKV linear + bidirectional short conv on K. + 2. UCPE ``P / P_T / P_inv`` from ``camera_conditions``. + 3. Sliced cam-branch RoPE → interleaved ``(N, D/2)`` cos/sin tables. + 4. Fused prep kernel (RMSNorm + ReLU + K-scale + UCPE 4x4 + RoPE), emitting ``inflation_sq`` for Dynamic Beta + Discounting. + 5. Beta discounting via ``inflation_sq`` (mirrors torch path). + 6. Fused forward scan (``reverse=False``) over the full sequence. + 7. Fused reverse scan (``reverse=True``) over the full sequence — the kernel applies flip-and-shift internally, + so no per-chunk loop is needed. + 8. Inverse UCPE (``apply_fn_o``) in torch. + + State-dict keys are identical to :class:`BidirectionalGDNUCPESinglePathLiteLA`. + """ + + def _forward_cam_branch( + self, + x: torch.Tensor, + HW: tuple[int, int, int], + camera_conditions: torch.Tensor, + rotary_emb: torch.Tensor | None, + **kwargs: object, + ) -> torch.Tensor: + # ---- Guards: k_conv_only=True. ---- + if kwargs.get("frame_valid_mask", None) is not None: + raise NotImplementedError( + "BidirectionalGDNUCPESinglePathLiteLABothTriton does not " + "support frame_valid_mask (training-only feature)." + ) + if self.conv_q_cam is not None or self.conv_v_cam is not None: + raise NotImplementedError( + "BidirectionalGDNUCPESinglePathLiteLABothTriton requires " + "k_conv_only=True (conv_q_cam / conv_v_cam must be None)." + ) + + B, N, _ = x.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + dtype_orig = x.dtype + H_heads = self.cam_heads + D_head = self.cam_head_dim + + # ---- 1. QKV linear + bidirectional short conv on K --------------- + qkv_w = torch.cat([self.q_proj_cam.weight, self.k_proj_cam.weight, self.v_proj_cam.weight]) + qkv_b = torch.cat([self.q_proj_cam.bias, self.k_proj_cam.bias, self.v_proj_cam.bias]) + qkv_cam = torch.nn.functional.linear(x, qkv_w, qkv_b) + q_raw, k_raw, v_raw = qkv_cam.chunk(3, dim=-1) + + if self.conv_k_cam is not None: + # Parent routing (BidirectionalGDN) gives the bidirectional + # forward+backward causal conv + average. + k_raw = self._apply_temporal_short_conv(k_raw, self.conv_k_cam, HW) + + q_raw = q_raw.contiguous().view(B, N, H_heads, D_head).contiguous() + k_raw = k_raw.contiguous().view(B, N, H_heads, D_head).contiguous() + v_raw = v_raw.contiguous().view(B, N, H_heads, D_head).contiguous() + + # ---- 2. UCPE P, P_T, P_inv (inline; skip cached prope_fns). ----- + raymats = _process_camera_conditions_raymats_only(camera_conditions, B, HW, self.patch_size) + raymats = raymats.reshape(B, -1, 4, 4) + P = raymats + P_T = P.transpose(-1, -2).contiguous() + P_inv = _invert_SE3(P).contiguous() + + # ---- 3. Sliced cam-branch RoPE + interleaved tables. ------------ + if rotary_emb is not None: + head_dim = D_head + orig_t_size = head_dim // 2 - 2 * (head_dim // 6) + orig_h_size = head_dim // 6 + new_head_dim = head_dim // 2 + new_t_size = new_head_dim // 2 - 2 * (new_head_dim // 6) + new_h_size = new_head_dim // 6 + new_w_size = new_head_dim // 6 + t_part = rotary_emb[..., :new_t_size] + h_part = rotary_emb[..., orig_t_size : orig_t_size + new_h_size] + w_part = rotary_emb[..., orig_t_size + orig_h_size : orig_t_size + orig_h_size + new_w_size] + rotary_emb_cam = torch.cat([t_part, h_part, w_part], dim=-1) + rope_cos, rope_sin = _prepare_ucpe_rope_tables(rotary_emb_cam, N, D_head // 2, x.device) + else: + rotary_emb_cam = None + rope_cos = torch.ones(N, D_head // 2, device=x.device, dtype=torch.float32) + rope_sin = torch.zeros(N, D_head // 2, device=x.device, dtype=torch.float32) + + # ---- 4. Fused Triton prep kernel -------------------------------- + q_norm_w = self.q_norm_cam.weight.float().contiguous() + k_norm_w = self.k_norm_cam.weight.float().contiguous() + k_scale = (D_head**-0.5) * (S**-0.5) + norm_eps_val = float( + getattr( + self.q_norm_cam, + "eps", + getattr(self.q_norm_cam, "variance_epsilon", 1e-6), + ) + ) + q_cam_trans, k_cam_trans, v_cam_trans, inflation_sq = cam_prep_func( + q_raw, + k_raw, + v_raw, + q_norm_weight=q_norm_w, + k_norm_weight=k_norm_w, + proj_q=P_T, + proj_kv=P_inv, + rope_cos=rope_cos, + rope_sin=rope_sin, + k_scale=k_scale, + norm_eps=norm_eps_val, + ) + inflation_sq = inflation_sq.view(B, H_heads, 1, N) + + # ---- 5. Gates + beta discounting ------------------------------- + precomputed_gates = kwargs.get("precomputed_gates", None) + if precomputed_gates is not None: + beta, decay = precomputed_gates + else: + beta, decay = self._compute_frame_gates(x, HW) + + inflation_sq_spatial = inflation_sq.view(B, H_heads, T, S) + frame_inflation_sq = inflation_sq_spatial.mean(dim=-1) + if beta.ndim == 3: + beta = beta / frame_inflation_sq.clamp_min(1.0) + elif beta.ndim == 4: + beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + + # ---- 6. fp32 cast + broadcast beta to (B, H, F, S) ------------- + q_cam_trans = q_cam_trans.float() + k_cam_trans = k_cam_trans.float() + v_cam_trans = v_cam_trans.float() + beta = beta.float() + decay = decay.float() + if beta.ndim == 3: + beta = beta.unsqueeze(-1).expand(B, H_heads, T, S).contiguous() + else: + assert beta.shape == (B, H_heads, T, S), f"beta shape {beta.shape}" + beta = beta.contiguous() + decay = decay.contiguous() + + q_cam_trans = q_cam_trans.contiguous() + k_cam_trans = k_cam_trans.contiguous() + v_cam_trans = v_cam_trans.contiguous() + + # ---- 7. Fused bidirectional chunkwise scan. -------------------- + out = cam_scan_bidi_chunkwise(q_cam_trans, k_cam_trans, v_cam_trans, beta, decay) + + # ---- 9. Cast back to input dtype, then inverse UCPE. ----------- + if dtype_orig != torch.float32: + out = out.to(dtype_orig) + + _, _, apply_fn_o = _prepare_ray_apply_fns( + head_dim=D_head, + P=P, + P_T=P_T, + P_inv=P_inv, + rotary_emb=rotary_emb_cam, + ) + out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous() + out = out.reshape(B, self.cam_dim, -1).permute(0, 2, 1) + return out + + +# ============================================================================ +# DiT base + SANA-WM camera-controlled transformer + public wrapper +# ============================================================================ + + +class SanaVideoMSCamCtrlBlock(nn.Module): + """ + A Sana block with global shared adaptive layer norm zero (adaLN-Zero) conditioning. + """ + + def __init__( + self, + hidden_size, + num_heads, + mlp_ratio=4.0, + drop_path=0.0, + qk_norm=False, + attn_type="flash", + ffn_type="mlp", + mlp_acts=("silu", "silu", None), + linear_head_dim=32, + cross_norm=False, + t_kernel_size=3, + camctrl_type=None, + patch_size=(1, 2, 2), + cam_attn_compress=2, + chunk_size=10, + chunk_split_strategy="uniform", + use_chunk_plucker_post_attn=False, + **block_kwargs, + ): + super().__init__() + self.hidden_size = hidden_size + self.chunk_size = chunk_size + self.chunk_split_strategy = chunk_split_strategy + + if use_chunk_plucker_post_attn: + self.plucker_proj = nn.Linear(hidden_size, hidden_size, bias=True) + nn.init.zeros_(self.plucker_proj.weight) + nn.init.zeros_(self.plucker_proj.bias) + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + # Camera-branch attention. The ``*Triton`` variants share the constructor + # signature with their pure-PyTorch parents (``BidirectionalGDNUCPESinglePathLiteLA``) + # so we can route them through ``_resolve_attention_block`` and get an + # automatic fallback to the parent class when Triton isn't usable. + if camctrl_type in ( + "BidirectionalGDNUCPESinglePathLiteLABothTriton", + "BidirectionalGDNUCPESinglePathLiteLATriton", + "BidirectionalGDNUCPESinglePathLiteLA", + ): + self_num_heads = hidden_size // linear_head_dim + cam_cls = _resolve_attention_block(camctrl_type, role="camctrl_type") + self.attn = cam_cls( + hidden_size, + hidden_size, + heads=self_num_heads, + cam_dim=hidden_size // cam_attn_compress, + cam_heads=max(1, self_num_heads // cam_attn_compress), + eps=1e-8, + qk_norm=qk_norm, + patch_size=patch_size, + **block_kwargs, + ) + elif camctrl_type == "BidirectionalSoftmaxUCPESinglePathLiteLA": + self_num_heads = hidden_size // linear_head_dim + self.attn = BidirectionalSoftmaxUCPESinglePathLiteLA( + hidden_size, + hidden_size, + heads=self_num_heads, + cam_dim=hidden_size // cam_attn_compress, + cam_heads=max(1, self_num_heads // cam_attn_compress), + eps=1e-8, + qk_norm=qk_norm, + patch_size=patch_size, + **block_kwargs, + ) + else: + # Main attention (no camera branch). Auto-falls-back ``*Triton`` to + # the non-Triton parent when Triton isn't usable. + attn_cls = _resolve_attention_block(attn_type, role="attn_type") + self.attn = attn_cls( + hidden_size, + hidden_size, + heads=hidden_size // linear_head_dim, + eps=1e-8, + qk_norm=qk_norm, + ) + + self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, qk_norm=cross_norm, **block_kwargs) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + # MLP + if ffn_type == "glumbconv": + self.mlp = GLUMBConv( + in_features=hidden_size, + hidden_features=int(hidden_size * mlp_ratio), + use_bias=(True, True, False), + norm=(None, None, None), + act=mlp_acts, + ) + elif ffn_type == "GLUMBConvTemp": + self.mlp = GLUMBConvTemp( + in_features=hidden_size, + hidden_features=int(hidden_size * mlp_ratio), + use_bias=(True, True, False), + norm=(None, None, None), + act=mlp_acts, + t_kernel_size=t_kernel_size, + ) + elif ffn_type == "mlp": + + def approx_gelu(): + return nn.GELU(approximate="tanh") + + self.mlp = Mlp( + in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0 + ) + else: + self.mlp = None + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size**0.5) + + @staticmethod + def _build_frame_token_mask( + frame_valid_mask: Optional[torch.Tensor], + *, + B: int, + T: int, + N: int, + device: torch.device, + dtype: torch.dtype, + ) -> Optional[torch.Tensor]: + """Convert frame-valid mask to token mask shaped ``(B, N, 1)``.""" + if frame_valid_mask is None: + return None + + m = frame_valid_mask + if m.ndim == 5: + m = m[:, 0, :, 0, 0] + elif m.ndim == 3 and m.shape[1] == 1: + m = m[:, 0, :] + elif m.ndim != 2: + raise ValueError( + "frame_valid_mask must be shaped (B, 1, T, 1, 1), (B, 1, T), or (B, T); " + f"got shape={list(frame_valid_mask.shape)}" + ) + + if m.shape[0] != B or m.shape[1] != T: + raise ValueError(f"frame_valid_mask shape mismatch: expected (B={B}, T={T}), got {list(m.shape)}") + if T <= 0 or N % T != 0: + raise ValueError(f"Invalid token/frame layout: N={N}, T={T}") + + S = N // T + return m.to(device=device, dtype=dtype).view(B, T, 1).expand(B, T, S).reshape(B, N, 1) + + def forward(self, x, y, t, mask=None, THW=None, rotary_emb=None, block_mask=None, chunk_index=None, **kwargs): + B, N, C = x.shape + num_frames = t.shape[2] + frame_valid_mask = kwargs.get("frame_valid_mask", None) + frame_token_mask = self._build_frame_token_mask( + frame_valid_mask, + B=B, + T=num_frames, + N=N, + device=x.device, + dtype=x.dtype, + ) + if frame_token_mask is not None: + x = x * frame_token_mask + + t = t.reshape(B, num_frames, 6, -1) # B,F,6,D + # scale_shift_table: 6, hidden_size -> 1,1,6,hidden_size + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None, :, :] + t + ).chunk(6, dim=-2) # each chunk: B,F,1,D + self_attn_kwargs = { + "HW": THW, + "rotary_emb": rotary_emb, + "block_mask": block_mask, + "camera_conditions": kwargs.get("camera_conditions", None), + "prope_fns": kwargs.get("prope_fns", None), + "camera_embedding": kwargs.get("camera_embedding", None), + "frame_valid_mask": frame_valid_mask, + } + if chunk_index is not None: + self_attn_kwargs["chunk_index"] = chunk_index[:] # NOTE: important, copy the list + if kwargs.get("chunk_index_global", None) is not None: + self_attn_kwargs["chunk_index_global"] = kwargs.get("chunk_index_global") + chunk_split_strategy = kwargs.get("chunk_split_strategy", self.chunk_split_strategy) + if chunk_split_strategy is not None: + self_attn_kwargs["chunk_split_strategy"] = chunk_split_strategy + + chunk_size = kwargs.get("chunk_size", self.chunk_size) + if chunk_size is not None: + self_attn_kwargs["chunk_size"] = chunk_size + + x_norm1 = self.norm1(x).reshape(B, num_frames, -1, C) + x_msa_in = t2i_modulate(x_norm1, shift_msa, scale_msa).reshape(B, N, C) + if frame_token_mask is not None: + x_msa_in = x_msa_in * frame_token_mask + attn_out = self.attn(x_msa_in, **self_attn_kwargs).reshape(B, num_frames, -1, C) + attn_out = (gate_msa * attn_out).reshape(B, N, C) + if frame_token_mask is not None: + attn_out = attn_out * frame_token_mask + x = x + self.drop_path(attn_out) + if frame_token_mask is not None: + x = x * frame_token_mask + + plucker_emb = kwargs.get("plucker_emb", None) + if plucker_emb is not None and hasattr(self, "plucker_proj"): + x = x + self.plucker_proj(plucker_emb) + + x = x + self.cross_attn(x, y, mask=mask) + if frame_token_mask is not None: + x = x * frame_token_mask + + mlp_kwargs = { + "HW": THW, + "frame_valid_mask": frame_valid_mask, + } + if chunk_index is not None: + mlp_kwargs["chunk_index"] = chunk_index[:] # NOTE: important, copy the list + if kwargs.get("chunk_index_global", None) is not None: + mlp_kwargs["chunk_index_global"] = kwargs.get("chunk_index_global") + if chunk_split_strategy is not None: + mlp_kwargs["chunk_split_strategy"] = chunk_split_strategy + + chunk_size = kwargs.get("chunk_size", self.chunk_size) + if chunk_size is not None: + mlp_kwargs["chunk_size"] = chunk_size + + x_norm2 = self.norm2(x).reshape(B, num_frames, -1, C) + x_mlp_in = t2i_modulate(x_norm2, shift_mlp, scale_mlp).reshape(B, N, C) + if frame_token_mask is not None: + x_mlp_in = x_mlp_in * frame_token_mask + mlp_out = self.mlp(x_mlp_in, **mlp_kwargs).reshape(B, num_frames, -1, C) + mlp_out = (gate_mlp * mlp_out).reshape(B, N, C) + if frame_token_mask is not None: + mlp_out = mlp_out * frame_token_mask + x = x + self.drop_path(mlp_out) + if frame_token_mask is not None: + x = x * frame_token_mask + + return x + + +_GDN_TO_SOFTMAX_CAMCTRL: dict[str, str] = { + "BidirectionalGDNUCPESinglePathLiteLABothTriton": "BidirectionalSoftmaxUCPESinglePathLiteLA", +} + + +def _inject_softmax_layers( + attn_type_list: list, + camctrl_type_list: list, + softmax_every_n: int, +) -> tuple: + """Replace every ``softmax_every_n``-th block's camctrl variant with its softmax counterpart. + + Pattern: for ``softmax_every_n=4``, blocks 3, 7, 11, ... (0-indexed at n-1) use softmax attention; the remaining + blocks keep GDN. Blocks whose camctrl_type has no softmax mapping are left as-is. + """ + attn_out = list(attn_type_list) + camctrl_out = list(camctrl_type_list) + for i in range(len(attn_out)): + if (i + 1) % softmax_every_n != 0: + continue + if camctrl_out[i] in _GDN_TO_SOFTMAX_CAMCTRL: + camctrl_out[i] = _GDN_TO_SOFTMAX_CAMCTRL[camctrl_out[i]] + return attn_out, camctrl_out + + +class SanaWMTransformer3DModel(ModelMixin, ConfigMixin): + r""" + SANA-WM 1600M bidirectional camera-controlled DiT. + + A single-class DiT (depth=20, hidden_size=2240, patch_size=(1,1,1), num_heads=20 — i.e. the public + ``Efficient-Large-Model/SANA-WM_bidirectional`` release). ``save_pretrained`` / ``from_pretrained`` work out of the + box via :class:`~diffusers.configuration_utils.ConfigMixin`. + + Args: + in_channels (`int`, defaults to 128): VAE latent channels (LTX-2). + attn_type (`str`): Main-branch attention, e.g. ``"BidirectionalGDNTriton"``. + camctrl_type (`str`): Camera-branch attention, e.g. + ``"BidirectionalGDNUCPESinglePathLiteLABothTriton"``. + softmax_every_n (`int`, defaults to 4): Inject a softmax block every N blocks. + linear_head_dim (`int`, defaults to 112): GDN head dimension. + ffn_type (`str`, defaults to ``"GLUMBConvTemp"``): FFN. + t_kernel_size (`int`, defaults to 3): Temporal conv kernel. + conv_kernel_size (`int`, defaults to 4): Spatial conv kernel inside attention. + k_conv_only (`bool`, defaults to True): Apply conv only on K. + pos_embed_type (`str`, defaults to ``"wan_rope"``): Position embedding. + qk_norm (`bool`, defaults to True): RMSNorm on Q/K. + cross_norm (`bool`, defaults to True): RMSNorm on cross-attention K. + y_norm (`bool`, defaults to True): Apply ``attention_y_norm`` to text embeddings. + y_norm_scale_factor (`float`, defaults to 0.01): Scale factor for ``attention_y_norm``. + init_cam_from_base (`bool`, defaults to True): Unused; the camera branch is loaded from the checkpoint. + Kept so released `config.json` files load. + chunk_split_strategy (`str`, defaults to ``"first_chunk_plus_one"``). + use_chunk_plucker_post_attn (`bool`, defaults to True). + chunk_plucker_channels (`int`, defaults to 48): ``6 dims * temporal_stride 8``. + chunk_plucker_post_attn_blocks (`int`, defaults to 20): All blocks. + fp32_attention (`bool`, defaults to True): Unused; attention always runs in fp32. Kept so released + `config.json` files load. + image_size (`int`, defaults to 720): Nominal image size. + caption_channels (`int`, defaults to 2304): Gemma-2 hidden size. + model_max_length (`int`, defaults to 300): Max prompt tokens. + + The state-dict is identical to the public sana checkpoint apart from the intentionally-removed ``pos_embed`` + buffer. + """ + + _supports_gradient_checkpointing = False + _no_split_modules = ["SanaVideoMSCamCtrlBlock"] + _repeated_blocks = ["SanaVideoMSCamCtrlBlock"] + _skip_layerwise_casting_patterns = ["x_embedder", "plucker_embedder", "norm"] + # NOTE: `_keep_in_fp32_modules` is intentionally unset. SANA-WM's blocks apply the + # timestep modulation inline (`t2i_modulate`), so holding `t_embedder` / `t_block` / + # `scale_shift_table` in fp32 would upcast the hidden states and feed fp32 activations + # to bf16 weights. Supporting it needs explicit casts in the block forward first. + + @register_to_config + def __init__( + self, + in_channels: int = 128, + num_layers: int = 20, + hidden_size: int = 2240, + num_attention_heads: int = 20, + patch_size: tuple[int, int, int] = (1, 1, 1), + attn_type: str = "BidirectionalGDNTriton", + camctrl_type: str = "BidirectionalGDNUCPESinglePathLiteLABothTriton", + softmax_every_n: int = 4, + linear_head_dim: int = 112, + ffn_type: str = "GLUMBConvTemp", + t_kernel_size: int = 3, + conv_kernel_size: int = 4, + k_conv_only: bool = True, + pos_embed_type: str = "wan_rope", + qk_norm: bool = True, + cross_norm: bool = True, + y_norm: bool = True, + y_norm_scale_factor: float = 0.01, + cam_attn_compress: int = 1, + init_cam_from_base: bool = True, + chunk_split_strategy: str = "first_chunk_plus_one", + use_chunk_plucker_post_attn: bool = True, + chunk_plucker_channels: int = 48, + chunk_plucker_post_attn_blocks: int = 20, + fp32_attention: bool = True, + image_size: int = 720, + caption_channels: int = 2304, + model_max_length: int = 300, + mlp_ratio: float = 3.0, + mlp_acts: tuple = ("silu", "silu", None), + use_pe: bool = True, + learn_sigma: bool = False, + pred_sigma: bool = False, + mixed_precision: str = "bf16", + ) -> None: + super().__init__() + + # The defaults describe the public SANA-WM_bidirectional release; they are + # configurable so a small variant can be built (e.g. for tests). + depth = num_layers + num_heads = num_attention_heads + patch_size = tuple(patch_size) + + # Remaining SanaMSVideoCamCtrl.__init__ defaults not exposed by the config signature. + mlp_acts = list(mlp_acts) + drop_path = 0.0 + pe_interpolation = 1.0 + norm_eps = 1e-5 + patch_embed_kernel = None + cfg_embed = False + timestep_norm_scale_factor = 1.0 + rope_fhw_dim = None + pack_latents = False + camctrl_layers_num = None + chunk_size = 10 + use_chunk_plucker_input = False + + # --- Base DiT config attributes (from Sana.__init__) --- + self.pred_sigma = pred_sigma + self.in_channels = in_channels + self.out_channels = in_channels * 2 if pred_sigma else in_channels + self.hidden_size = hidden_size + self.num_heads = num_heads + self.linear_head_dim = linear_head_dim + self.pe_interpolation = pe_interpolation + self.depth = depth + self.use_pe = use_pe + self.pos_embed_type = pos_embed_type + self.y_norm = y_norm + # NOTE: ``self.config`` is provided (read-only) by ConfigMixin via @register_to_config. + self.timestep_norm_scale_factor = timestep_norm_scale_factor + + self.t_embedder = TimestepEmbedder(hidden_size) + self.cfg_embedder = None + if cfg_embed: + self.cfg_embedder = TimestepEmbedder(hidden_size) + + if self.y_norm: + self.attention_y_norm = RMSNorm(hidden_size, scale_factor=y_norm_scale_factor, eps=norm_eps) + + # --- Video camera-controlled DiT modules (from SanaMSVideoCamCtrl.__init__) --- + self.chunk_size = chunk_size + self.chunk_split_strategy = chunk_split_strategy + self.patch_size = patch_size + self.h = self.w = 0 + + def approx_gelu(): + return nn.GELU(approximate="tanh") + + self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True)) + self.pos_embed_ms = None + self.pack_latents = pack_latents + self.attn_type = attn_type + + self.camctrl_type = camctrl_type + assert self.camctrl_type in [ + "BidirectionalGDNUCPESinglePathLiteLABothTriton", + "BidirectionalSoftmaxUCPESinglePathLiteLA", + ], f"Not supported camera control type: {self.camctrl_type}" + + self.camctrl_layers_num = camctrl_layers_num if camctrl_layers_num is not None else depth + self.cam_attn_compress = cam_attn_compress + + kernel_size = patch_embed_kernel or patch_size + x_embedder_in_channels = in_channels + if self.pack_latents: + x_embedder_in_channels = x_embedder_in_channels * 2 * 2 + self.out_channels = in_channels * 2 * 2 + + self.x_embedder = PatchEmbedMS3D( + patch_size, x_embedder_in_channels, hidden_size, kernel_size=kernel_size, bias=True + ) + + self.y_embedder = CaptionEmbedder( + in_channels=caption_channels, + hidden_size=hidden_size, + act_layer=approx_gelu, + token_num=model_max_length, + ) + + self.use_chunk_plucker_input = use_chunk_plucker_input + self.use_chunk_plucker_post_attn = use_chunk_plucker_post_attn + if self.use_chunk_plucker_input or self.use_chunk_plucker_post_attn: + self.plucker_embedder = PatchEmbedMS3D( + patch_size, chunk_plucker_channels, hidden_size, kernel_size=kernel_size, bias=True + ) + nn.init.zeros_(self.plucker_embedder.proj.weight) + nn.init.zeros_(self.plucker_embedder.proj.bias) + + # UCPE-style camera branch uses a 3-channel absmap (up_map + lat_map). + self.raymap_embedder = PatchEmbedMS3D(patch_size, 3, hidden_size, kernel_size=kernel_size, bias=True) + + if attn_type in ["flash", "FlexLinearAttention", "flex"]: + attention_head_dim = hidden_size // num_heads + else: + attention_head_dim = linear_head_dim + + if use_pe: + if pos_embed_type != "wan_rope": + raise ValueError(f'`pos_embed_type` must be "wan_rope", got {pos_embed_type!r}.') + self.rope = WanRotaryPosEmbed( + attention_head_dim=attention_head_dim, patch_size=patch_size, max_seq_len=1024, fhw_dim=rope_fhw_dim + ) + # stochastic depth decay rule (build on CPU so meta-device construction works) + drop_path = [x.item() for x in torch.linspace(0, drop_path, depth, device="cpu")] + + self.softmax_every_n = softmax_every_n + attn_type_list = [attn_type] * depth + camctrl_type_list = [camctrl_type if i < self.camctrl_layers_num else None for i in range(depth)] + if attn_type in ["flex", "FlexLinearAttention"]: + attn_type_list[0] = "flash" + attn_type_list[1] = "flash" + + if softmax_every_n > 0: + attn_type_list, camctrl_type_list = _inject_softmax_layers( + attn_type_list, + camctrl_type_list, + softmax_every_n, + ) + logger.info( + f"Hybrid attention (softmax_every_n={softmax_every_n}):\n" + f" attn_type_list = {attn_type_list}\n" + f" camctrl_type_list = {camctrl_type_list}" + ) + + self.blocks = nn.ModuleList( + [ + SanaVideoMSCamCtrlBlock( + hidden_size, + num_heads, + mlp_ratio=mlp_ratio, + drop_path=drop_path[i], + qk_norm=qk_norm, + attn_type=attn_type_list[i], + ffn_type=ffn_type, + mlp_acts=mlp_acts, + linear_head_dim=linear_head_dim, + cross_norm=cross_norm, + t_kernel_size=t_kernel_size, + camctrl_type=camctrl_type_list[i], + patch_size=patch_size, + cam_attn_compress=self.cam_attn_compress, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + conv_kernel_size=conv_kernel_size, + k_conv_only=k_conv_only, + use_chunk_plucker_post_attn=( + use_chunk_plucker_post_attn + and (chunk_plucker_post_attn_blocks < 0 or i < chunk_plucker_post_attn_blocks) + ), + ) + for i in range(depth) + ] + ) + self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels) + + if ffn_type == "GLUMBConvTemp": + logger.info(f"{ffn_type} Temporal kernal: {t_kernel_size}") + + self.in_channels = self.out_channels = in_channels + + @staticmethod + def _pack_latents(latents, batch_size, num_channels_latents, height, width, frame): + latents = latents.view(batch_size, num_channels_latents, frame, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 4, 6, 2, 3, 5) + latents = latents.reshape(batch_size, num_channels_latents * 4, frame, height // 2, width // 2) + + return latents + + @staticmethod + def _unpack_latents(latents, height, width, frame): + batch_size, channels, frame, H, W = latents.shape + + assert height % 2 == 0 and width % 2 == 0 + # latent height and width to be divisible by 2. + latents = latents.view(batch_size, channels // 4, 2, 2, frame, height // 2, width // 2) + latents = latents.permute(0, 1, 4, 5, 2, 6, 3) + latents = latents.reshape(batch_size, channels // (2 * 2), frame, height, width) + + return latents + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + return_dict: bool = True, + **kwargs: Any, + ): + """Run the SANA-WM DiT. + + Args: + hidden_states: ``(B, C, T, H, W)`` latents. + timestep: ``(B, 1, T)`` per-frame diffusion timesteps (LTX style). + encoder_hidden_states: ``(B, 1, L, D_caption)`` text embeddings. + encoder_attention_mask: ``(B, L)`` text attention mask (diffusers convention). + mask: Alias for ``encoder_attention_mask`` matching the sana DiT's + kwarg name. If both are passed, ``mask`` takes precedence. + return_dict: If ``True`` (default), returns a :class:`Transformer2DModelOutput`; + otherwise returns a one-tuple ``(sample,)``. + **kwargs: SANA-WM-specific conditioning — at minimum + ``data_info``, ``camera_conditions``, ``chunk_plucker``. + + Returns: + :class:`Transformer2DModelOutput` with ``sample`` of shape ``(B, C, T, H, W)``. + """ + # The sana DiT names its text mask kwarg ``mask``. + # Accept both ``mask=`` (sana convention) and ``encoder_attention_mask=`` + # (diffusers convention); the former wins if both are provided. + if mask is None: + mask = encoder_attention_mask + x = hidden_states + y = encoder_hidden_states + + bs = x.shape[0] + x = x.to(self.dtype) + if self.timestep_norm_scale_factor != 1.0: + timestep = (timestep.float() / self.timestep_norm_scale_factor).to(torch.float32) + else: + timestep = timestep.long().to(torch.float32) + y = y.to(self.dtype) + self.f, self.h, self.w = ( + x.shape[-3] // self.patch_size[0], + x.shape[-2] // self.patch_size[1], + x.shape[-1] // self.patch_size[2], + ) + + data_info = kwargs.get("data_info", {}) + if data_info.get("image_vae_embeds", None) is not None: + x = torch.cat([x, data_info["image_vae_embeds"].to(self.dtype)], dim=1) + cam_embeds = kwargs.get("camera_conditions", None) + if self.pack_latents: + x = self._pack_latents(x, bs, self.in_channels, self.h, self.w, self.f) + if cam_embeds is not None: + cam_embeds = cam_embeds.to(self.dtype) + + self.h = self.h // 2 + self.w = self.w // 2 + + if self.x_embedder.patch_size != self.x_embedder.kernel_size and self.x_embedder.kernel_size == (1, 2, 2): + x = F.pad(x, (0, 1, 0, 1, 0, 0)) + if cam_embeds is not None: + cam_embeds = F.pad(cam_embeds, (0, 1, 0, 1, 0, 0)) + + x = self.x_embedder(x) + if cam_embeds is not None: + # Both surviving camctrl variants are UCPE-style: build raymats + 3-channel + # absmap (up_map + lat_map) from the raw (B,F,20) camera conditions. + raw_cam_conditions = cam_embeds + cam_pos_embeds = kwargs.get("cam_pos_embeds", None) + if cam_pos_embeds is not None and "absmap" in cam_pos_embeds: + cam_embeds = cam_pos_embeds["absmap"] + if "P" in cam_pos_embeds: + kwargs["raymats"] = cam_pos_embeds["P"] + else: + raymats, cam_embeds = _process_camera_conditions_ucpe( + raw_cam_conditions, bs, (self.f, self.h, self.w), self.patch_size + ) + cam_embeds = cam_embeds.permute(0, 4, 1, 2, 3).to(self.dtype) + kwargs["raymats"] = raymats + if not (self.use_chunk_plucker_input or self.use_chunk_plucker_post_attn): + cam_embeds = self.raymap_embedder(cam_embeds) + x = x + cam_embeds + kwargs["camera_embedding"] = cam_embeds + kwargs["camera_conditions"] = raw_cam_conditions + + if self.use_chunk_plucker_input and "chunk_plucker" in kwargs: + plucker_input = kwargs["chunk_plucker"].to(self.dtype) + plucker_emb = self.plucker_embedder(plucker_input) + x = x + plucker_emb + + if self.use_chunk_plucker_post_attn and "chunk_plucker" in kwargs: + plucker_input = kwargs["chunk_plucker"].to(self.dtype) + kwargs["plucker_emb"] = self.plucker_embedder(plucker_input) + + image_pos_embed = kwargs.get("pos_embeds", None) + if self.use_pe and image_pos_embed is None: + image_pos_embed = self.rope((self.f, self.h, self.w)) + elif image_pos_embed is not None: + image_pos_embed = image_pos_embed.to(x.device) + while image_pos_embed.ndim > 4: + image_pos_embed = image_pos_embed.squeeze(1) + + t = self.t_embedder(timestep.flatten()) # (N, D) + t0 = self.t_block(t) + t = t.unflatten(dim=0, sizes=timestep.shape) + t0 = t0.unflatten(dim=0, sizes=timestep.shape) + + y = self.y_embedder(y) # (N, D) + if self.y_norm: + y = self.attention_y_norm(y) + + if mask is None: + raise ValueError( + "`mask` is required: SANA-WM's cross-attention needs the text padding mask to build its attention " + "bias. Pass the prompt attention mask returned by the pipeline's `encode_prompt`." + ) + mask = mask.to(torch.int16) + mask = mask.repeat(y.shape[0] // mask.shape[0], 1) if mask.shape[0] != y.shape[0] else mask + mask = mask.squeeze(1).squeeze(1) + y_lens = mask + + block_mask = None + + if kwargs.get("camera_conditions") is not None: + # Pre-compute UCPE projection functions to share across blocks + # (both surviving camctrl variants are UCPE-style). + if self.attn_type in ["flash", "FlexLinearAttention", "flex"]: + head_dim = self.hidden_size // self.num_heads + else: + head_dim = self.linear_head_dim + + cam_pos_embeds = kwargs.get("cam_pos_embeds", None) + if cam_pos_embeds is not None: + for k, v in cam_pos_embeds.items(): + if isinstance(v, torch.Tensor): + v = v.to(x.device) + if k == "absmap": + while v.ndim > 5: + v = v.squeeze(1) + else: + while v.ndim > 4: + v = v.squeeze(1) + cam_pos_embeds[k] = v + + kwargs["prope_fns"] = prepare_prope_fns( + camctrl_type="UCPE", + head_dim=head_dim, + camera_conditions=kwargs["camera_conditions"], + HW=(self.f, self.h, self.w), + patch_size=self.patch_size, + rotary_emb=image_pos_embed, + raymats=kwargs.get("raymats"), + cam_pos_embeds=cam_pos_embeds, + ) + + for i, block in enumerate(self.blocks): + x = block( + x, + y, + t0, + y_lens, + (self.f, self.h, self.w), + image_pos_embed, + block_mask=block_mask if i > 1 else None, + **kwargs, + ) # (N, T, D) + + x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels) + x = self.unpatchify(x) # (N, out_channels, H, W) + if self.pack_latents: + x = self._unpack_latents(x, self.h * 2, self.w * 2, self.f) + + return Transformer2DModelOutput(sample=x) if return_dict else (x,) + + def unpatchify(self, x): + """ + x: (N, T, patch_size**2 * C) imgs: (N, H, W, C) + """ + c = self.out_channels + p_f, p_h, p_w = self.x_embedder.patch_size + h, w = self.h, self.w + assert self.f * self.h * self.w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], self.f, h, w, p_f, p_h, p_w, c)) + x = torch.einsum("nfhwopqc->ncfohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, self.f * p_f, h * p_h, w * p_w)) + + return imgs diff --git a/src/diffusers/models/transformers/transformer_sana_wm_kernels.py b/src/diffusers/models/transformers/transformer_sana_wm_kernels.py new file mode 100644 index 000000000000..de3a293edf1b --- /dev/null +++ b/src/diffusers/models/transformers/transformer_sana_wm_kernels.py @@ -0,0 +1,3234 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. + +# ruff: noqa: E501 + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Optional, Union + +import torch +import torch.nn.functional as F + + +# Optional Triton import. The kernels below are the fast path on CUDA + Triton +# >= 3.x, but they are not correctness-essential: SanaWMTransformer3DModel has +# pure-PyTorch attention variants for every ``*Triton`` class (the dispatcher +# in ``transformer_sana_wm.py`` auto-falls-back when Triton isn't usable). On +# a Triton-less system, ``@triton.jit`` becomes a no-op so the kernel function +# *definitions* still load (so the module can be imported anywhere), but +# calling any of the Triton-backed entry points raises a clear error. +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + _TRITON_AVAILABLE = False + + class _TritonShim: + """No-op stand-in for ``triton`` / ``triton.language`` on systems without Triton. + + ``@triton.jit`` becomes a pass-through so the @-decorated kernel functions are still defined as plain Python + (and never called on the torch fallback path). Any attribute access returns the same shim so ``tl.constexpr``, + ``tl.load`` etc. evaluate to a harmless sentinel — which is fine as long as no kernel body actually executes. + """ + + def __getattr__(self, name): + return self + + def __call__(self, *args, **kwargs): + if args and callable(args[0]) and not kwargs: + return args[0] + return self + + def jit(self, fn=None, **kwargs): + if fn is None: + return lambda f: f + return fn + + triton = _TritonShim() + tl = _TritonShim() + + +def is_triton_available() -> bool: + """Whether ``triton`` was importable and the kernels in this module can be launched.""" + return _TRITON_AVAILABLE + + +def _require_triton(entry_point: str) -> None: + if not _TRITON_AVAILABLE: + raise RuntimeError( + f"{entry_point} requires the `triton` package to run. Install Triton " + f"or switch to the pure-PyTorch attention variant (e.g. drop the " + f"`Triton` suffix from `attn_type` / `camctrl_type` on " + f"SanaWMTransformer3DModel — the dispatcher does this automatically " + f"when Triton isn't usable)." + ) + + +# ===================================================================== +# GPU-adaptive kernel config +# ===================================================================== + + +def _get_kernel_config() -> dict: + """Return optimal kernel parameters for the current GPU. + + STATE_FP32: use fp32 state_prev when SRAM is large enough. + - bf16 state_prev: ~96KB total SRAM (fits GB10's 101KB). + - fp32 state_prev: ~128KB total SRAM (needs H100's 228KB+). + """ + if not torch.cuda.is_available(): + return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 4, "STATE_FP32": False} + smem = torch.cuda.get_device_properties(0).shared_memory_per_multiprocessor + state_fp32 = smem >= 150 * 1024 # H100 (228KB) yes, GB10 (101KB) no + return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 8, "STATE_FP32": state_fp32} + + +_KCFG = None + + +def _kcfg(): + global _KCFG + if _KCFG is None: + _KCFG = _get_kernel_config() + return _KCFG + + +# precision=0 → IEEE fp32 dots + fp32 state (DOT_PRECISION=2, STATE_FP32=1) +# precision=1 → TF32 dots + fp32 state (DOT_PRECISION=1, STATE_FP32=1) +# precision=2 → bf16 dots + fp32 state (DOT_PRECISION=0, STATE_FP32=1) [default] +# precision=3 → bf16 dots + bf16 state (DOT_PRECISION=0, STATE_FP32=0) +def _precision_params(precision: int) -> tuple: + if precision == 0: + return 2, True + elif precision == 1: + return 1, True + elif precision == 3: + return 0, False + else: # default + return 0, True + + +_env_prec = os.environ.get("FUSED_GDN_PRECISION", None) +PRECISION_OVERRIDE: int | None = int(_env_prec) if _env_prec is not None else None + + +def _resolve_launch_config() -> tuple: + """Returns (prec, dot_prec, state_fp32, num_warps). + + Uses ``PRECISION_OVERRIDE`` when set; otherwise falls back to ``_kcfg()`` (which picks ``STATE_FP32`` based on + per-GPU SRAM). ``num_warps`` is clamped to 4 when dots run on fp32 operands (more registers needed). + """ + cfg = _kcfg() + prec = PRECISION_OVERRIDE if PRECISION_OVERRIDE is not None else 2 + dot_prec, state_fp32 = _precision_params(prec) + if PRECISION_OVERRIDE is None: + state_fp32 = cfg["STATE_FP32"] + nw = cfg["num_warps"] + if dot_prec >= 1: + nw = min(nw, 4) + return prec, dot_prec, state_fp32, nw + + +def prepare_rope_tables(rotary_emb, N: int, D: int, device) -> tuple[torch.Tensor, torch.Tensor]: + """Complex rotary_emb `(1, 1, N, D//2)` → expanded (N, D) cos/sin tables. + + Encodes the interleaved-pair rotation + y[2i] = x[2i]*cos[i] - x[2i+1]*sin[i] y[2i+1] = x[2i]*sin[i] + x[2i+1]*cos[i] + as y[d] = x[d]*cos_exp[d] + x[d^1]*sin_exp[d] where sin_exp[2i] = -sin[i], sin_exp[2i+1] = +sin[i]. + + Returns (cos_exp, sin_exp) both (N, D) float32, contiguous. + """ + if rotary_emb is None: + return ( + torch.ones(N, D, device=device, dtype=torch.float32), + torch.zeros(N, D, device=device, dtype=torch.float32), + ) + freqs = rotary_emb.squeeze(0).squeeze(0) # (N, D//2) complex + cos_half = freqs.real.float() + sin_half = freqs.imag.float() + rope_cos = cos_half.repeat_interleave(2, dim=-1) + rope_sin = torch.stack([-sin_half, sin_half], dim=-1).reshape(N, D) + return rope_cos.contiguous(), rope_sin.contiguous() + + +def _precompute_inv_rms(qkv: torch.Tensor, idx: int, C: int, eps: float = 1e-5) -> torch.Tensor: + """Compute 1/RMS for one component of QKV over the full C = H*D channel dim. + + Args: + qkv: (B, N, 3, H, D) + idx: 0 for Q, 1 for K, 2 for V + C: H*D (channel count) + eps: RMSNorm epsilon + + Returns: + inv_rms: (B, N) float32 + """ + raw = qkv[:, :, idx].float() # (B, N, H, D) + sq_sum = (raw * raw).sum(dim=(-2, -1)) # (B, N) + return torch.rsqrt(sq_sum / C + eps) + + +# ===================================================================== +# Fused single-pass Q+K inverse-RMS Triton kernel +# ===================================================================== +# Single Triton launch that reads each `(b, n)` row of `qkv` once and emits +# both `q_inv_rms[b, n]` and `k_inv_rms[b, n]`. Replaces two separate PyTorch +# scans (cast→square→sum→rsqrt) over `qkv[:, :, 0]` and `qkv[:, :, 1]`. +# +# Layout assumed: `qkv` is (B, N, 3, H, D) contiguous, so the C = H*D channels +# for a given (b, n, qkv_idx) live in a contiguous memory span. + + +@triton.jit +def _fused_qk_inv_rms_kernel( + qkv_ptr, # *T_in (B, N, 3, H, D), contiguous + q_inv_rms_ptr, # *float32 (B, N) + k_inv_rms_ptr, # *float32 (B, N) + N: tl.constexpr, + C: tl.constexpr, # H * D + eps, + BLOCK_C: tl.constexpr, +): + bn_id = tl.program_id(0) + qkv_row_stride = 3 * C + row_base = bn_id * qkv_row_stride + q_base = row_base + k_base = row_base + C + + offs = tl.arange(0, BLOCK_C) + mask = offs < C + + q_vals = tl.load(qkv_ptr + q_base + offs, mask=mask, other=0.0).to(tl.float32) + k_vals = tl.load(qkv_ptr + k_base + offs, mask=mask, other=0.0).to(tl.float32) + + q_sq = tl.sum(q_vals * q_vals, axis=0) + k_sq = tl.sum(k_vals * k_vals, axis=0) + + inv_c = 1.0 / C + q_inv = tl.rsqrt(q_sq * inv_c + eps) + k_inv = tl.rsqrt(k_sq * inv_c + eps) + + tl.store(q_inv_rms_ptr + bn_id, q_inv) + tl.store(k_inv_rms_ptr + bn_id, k_inv) + + +def fused_qk_inv_rms( + qkv: torch.Tensor, + eps: float = 1e-5, +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-pass Triton fused Q+K inverse-RMS. + + Replaces ``(_precompute_inv_rms(qkv, 0, C, eps), _precompute_inv_rms(qkv, 1, C, eps))`` with one launch that reads + each ``(b, n)`` row of ``qkv`` exactly once. + + Args: + qkv: (B, N, 3, H, D) contiguous tensor, any fp dtype. + eps: RMSNorm epsilon. + + Returns: + (q_inv_rms, k_inv_rms), each (B, N) float32 contiguous. + """ + _require_triton("fused_qk_inv_rms") + assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)" + assert qkv.dim() == 5 and qkv.shape[2] == 3, f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}" + B, N, _, H, D = qkv.shape + C = H * D + q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device) + k_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device) + BLOCK_C = triton.next_power_of_2(C) + _fused_qk_inv_rms_kernel[(B * N,)]( + qkv, + q_inv_rms, + k_inv_rms, + N=N, + C=C, + eps=eps, + BLOCK_C=BLOCK_C, + ) + return q_inv_rms, k_inv_rms + + +# ===================================================================== +# Bidirectional GDN entry point (delegates to chunkwise) +# ===================================================================== + + +def fused_bigdn_func( + qkv: torch.Tensor, # (B, N, 3, H, D) + q_inv_rms: torch.Tensor, # (B, N) float32 + k_inv_rms: torch.Tensor, # (B, N) float32 + q_norm_weight: torch.Tensor, # (C,) float32 + k_norm_weight: torch.Tensor, # (C,) float32 + rope_cos: torch.Tensor, # (N, D) float32 + rope_sin: torch.Tensor, # (N, D) float32 + beta: torch.Tensor, # (B, H, F, S) + decay: torch.Tensor, # (B, H, F) + F: int, + S: int, + k_scale: float, + eps: float = 1e-6, +) -> torch.Tensor: + """Bidirectional fused GDN. Returns ``(B, N, H, D)``. + + Thin entry point kept for call-site stability; delegates to :func:`fused_bigdn_bidi_chunkwise` from + ``fused_gdn_chunkwise``. + """ + _require_triton("fused_bigdn_func") + return fused_bigdn_bidi_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F=F, + S=S, + k_scale=k_scale, + eps=eps, + ) + + +# ============================================================================= +# Scalar helpers +# ============================================================================= + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Invert a 4x4 SE(3) matrix batch (closed-form). + + Mirrors the production ``_invert_SE3`` in ``sana_camctrl_blocks.py``; inlined to keep this module dependency-light. + """ + assert transforms.shape[-2:] == (4, 4) + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + return out + + +def _process_camera_conditions_raymats_only( + camera_conditions: torch.Tensor, + B: int, + HW: tuple[int, int, int], + patch_size: tuple[int, int, int], +) -> torch.Tensor: + """Lightweight variant of ``_process_camera_conditions_ucpe`` — raymats only. + + Computes *only* the per-ray ``world -> ray_local`` SE(3) transforms used by UCPE single-path. Skips the + ``compute_up_lat_map`` path (absmap) that the cam branch never consumes — that saves ~1 ms per block on H100. + + Args: + camera_conditions: ``(B, F, 20)`` — ``[c2w_16 | fx | fy | cx | cy]``. + B: Batch size (redundant with ``camera_conditions.shape[0]``; kept + for parity with the production signature). + HW: ``(T_latent, H_latent, W_latent)`` from the caller. + patch_size: ``(pt, ph, pw)`` patch embedding stride. + + Returns: + ``raymats`` of shape ``(B, F, H_latent, W_latent, 4, 4)``. + """ + F_dim = camera_conditions.shape[1] + c2w_flat = camera_conditions[..., :16] + C_to_W = c2w_flat.view(B, F_dim, 4, 4) + + fx = camera_conditions[..., 16] + fy = camera_conditions[..., 17] + cx = camera_conditions[..., 18] + cy = camera_conditions[..., 19] + H_dim, W_dim = HW[1], HW[2] + image_width = W_dim * patch_size[2] + image_height = H_dim * patch_size[1] + + xi = torch.zeros( + (B, F_dim), + device=camera_conditions.device, + dtype=camera_conditions.dtype, + ) + x_fov = compute_fov_from_fx_xi( + fx, + xi, + image_width, + device=camera_conditions.device, + dtype=camera_conditions.dtype, + ).view(B, F_dim) + y_fov = compute_fov_from_fx_xi( + fy, + xi, + image_height, + device=camera_conditions.device, + dtype=camera_conditions.dtype, + ).view(B, F_dim) + + d_cam = ucm_unproject_grid_fov( + x_fov, + y_fov, + xi, + H_dim, + W_dim, + cx / patch_size[2], + cy / patch_size[1], + device=camera_conditions.device, + dtype=camera_conditions.dtype, + ) + if d_cam.ndim == 4 and d_cam.shape[0] == B * F_dim: + d_cam = d_cam.view(B, F_dim, H_dim, W_dim, 3) + + return world_to_ray_mats(d_cam, C_to_W) # (B, F, H, W, 4, 4) + + +def _precompute_cam_inv_rms(raw: torch.Tensor, eps: float) -> torch.Tensor: + """Compute ``1/RMS`` per ``(b, n)`` over full-``C`` channels. + + Args: + raw: ``(B, N, H, D)`` raw QKV projection output (typically fp32). + eps: RMSNorm epsilon. + + Returns: + ``inv_rms`` of shape ``(B, N)`` in fp32, contiguous. + """ + B, N, H, D = raw.shape + C = H * D + sq_sum = (raw.float() * raw.float()).sum(dim=(-1, -2)) # (B, N) + return torch.rsqrt(sq_sum / C + eps).contiguous() + + +def _prepare_ucpe_rope_tables( + rotary_emb_cam: torch.Tensor, + N: int, + D_half: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert complex RoPE ``(1, 1, N, D_half//2)`` to interleaved ``(N, D_half)`` cos/sin. + + Uses the interleaved-pair convention: + y[2i] = x[2i]*cos[i] - x[2i+1]*sin[i] y[2i+1] = x[2i]*sin[i] + x[2i+1]*cos[i] + encoded as ``y[d] = x[d]*cos_exp[d] + x[d^1]*sin_exp[d]`` with + sin_exp[2i] = -sin[i], sin_exp[2i+1] = +sin[i]. + """ + del device # all outputs inherit device from freqs + freqs = rotary_emb_cam.squeeze(0).squeeze(0) # (N, D_half//2) complex + cos_half = freqs.real.float() + sin_half = freqs.imag.float() + rope_cos = cos_half.repeat_interleave(2, dim=-1).contiguous() + rope_sin = torch.stack([-sin_half, sin_half], dim=-1).reshape(N, D_half).contiguous() + return rope_cos, rope_sin + + +# ============================================================================= +# Triton kernels — lifted verbatim from cam_gdn_playground.py::TritonCamBranch +# ============================================================================= + + +_DEFAULT_BLOCK_S = 64 + + +@triton.jit +def _cam_prep_kernel( + q_raw_ptr, # (B, N, H, D) contiguous, any fp dtype + k_raw_ptr, # (B, N, H, D) contiguous (post short-conv on K) + v_raw_ptr, # (B, N, H, D) contiguous + q_inv_rms_ptr, # (B, N) float32 — precomputed over full C channels + k_inv_rms_ptr, # (B, N) float32 + q_norm_w_ptr, # (C,) = (H*D,) float32 + k_norm_w_ptr, # (C,) float32 + proj_q_ptr, # (B, N, 4, 4) — applied to Q first D/2 dims (P_T) + proj_kv_ptr, # (B, N, 4, 4) — applied to K,V first D/2 dims (P_inv) + rope_cos_ptr, # (N, D_rope) float32, D_rope = D//2 + rope_sin_ptr, # (N, D_rope) float32 + # --- outputs in (B, H, D, N) layout, same strides pattern --- + q_out_ptr, + k_out_ptr, + v_out_ptr, + k_pre_norm_sq_ptr, # (B, H, N) float32 — ||k_pre_ucpe||^2 + k_post_norm_sq_ptr, # (B, H, N) float32 — ||k_post_ucpe||^2 + # --- dims --- + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, # head dim + D_HALF: tl.constexpr, # D // 2 + N_GROUPS: tl.constexpr, # D_HALF // 4 + K_SCALE, + # --- tile sizes --- + BLOCK_D_ROPE: tl.constexpr, # next pow2 of D_HALF (rope block) + BLOCK_GROUPS: tl.constexpr, # next pow2 of N_GROUPS +): + """One program per (b, n, h) — processes a single (Q, K, V) head slice. + + Loads the first D_HALF dims as a (N_GROUPS, 4) tile (for the UCPE block-diagonal 4x4 projmat), and the second + D_HALF dims as a (D_HALF,) vector (for RoPE). No redundant loads. + """ + pid = tl.program_id(0) + h_idx = pid % H + bn_idx = pid // H + b_idx = bn_idx // N + n_idx = bn_idx % N + + # layout (B, N, H, D) contiguous + row_base = b_idx * (N * H * D) + n_idx * (H * D) + h_idx * D + nw_off = h_idx * D + + # ---- load inv-RMS (scalar, shared across heads for this token) ---- + q_inv_rms = tl.load(q_inv_rms_ptr + bn_idx).to(tl.float32) + k_inv_rms = tl.load(k_inv_rms_ptr + bn_idx).to(tl.float32) + + # ---- load per-token P matrices (4,4) shared across heads ---- + proj_base = (b_idx * N + n_idx) * 16 + offs_i = tl.arange(0, 4) + offs_j = tl.arange(0, 4) + P_q = tl.load(proj_q_ptr + proj_base + offs_i[:, None] * 4 + offs_j[None, :]).to(tl.float32) + P_kv = tl.load(proj_kv_ptr + proj_base + offs_i[:, None] * 4 + offs_j[None, :]).to(tl.float32) + + # ================================================================== + # Pass 1 — UCPE block-diagonal projmat on first D_HALF dims + # ================================================================== + offs_g = tl.arange(0, BLOCK_GROUPS) + mask_g = offs_g < N_GROUPS + offs_gj = offs_g[:, None] * 4 + offs_j[None, :] # (BLOCK_GROUPS, 4) + mask_gj = mask_g[:, None] + + q_half = tl.load(q_raw_ptr + row_base + offs_gj, mask=mask_gj, other=0.0).to(tl.float32) + k_half = tl.load(k_raw_ptr + row_base + offs_gj, mask=mask_gj, other=0.0).to(tl.float32) + v_half = tl.load(v_raw_ptr + row_base + offs_gj, mask=mask_gj, other=0.0).to(tl.float32) + + q_nw_half = tl.load(q_norm_w_ptr + nw_off + offs_gj, mask=mask_gj, other=0.0).to(tl.float32) + k_nw_half = tl.load(k_norm_w_ptr + nw_off + offs_gj, mask=mask_gj, other=0.0).to(tl.float32) + + q_half = q_half * q_inv_rms * q_nw_half + q_half = tl.where(q_half > 0, q_half, 0.0) + + k_half = k_half * k_inv_rms * k_nw_half + k_half = tl.where(k_half > 0, k_half, 0.0) * K_SCALE + + # Pre-UCPE ||k||^2 contribution from first half + k_half_masked = tl.where(mask_gj, k_half, 0.0) + k_pre_half_sq = tl.sum(k_half_masked * k_half_masked) + + # Apply 4x4 projmat: out[g, i] = sum_j P[i, j] * in[g, j] + # (BLOCK_GROUPS, 1, 4) * (1, 4, 4) -> (BLOCK_GROUPS, 4, 4), sum axis=-1 + q_half_out = tl.sum(q_half[:, None, :] * P_q[None, :, :], axis=-1) + k_half_out = tl.sum(k_half[:, None, :] * P_kv[None, :, :], axis=-1) + v_half_out = tl.sum(v_half[:, None, :] * P_kv[None, :, :], axis=-1) + + # Post-UCPE ||k||^2 contribution from first half + k_half_out_masked = tl.where(mask_gj, k_half_out, 0.0) + k_post_half_sq = tl.sum(k_half_out_masked * k_half_out_masked) + + # ================================================================== + # Pass 2 — RoPE on second D_HALF dims + # ================================================================== + offs_r = tl.arange(0, BLOCK_D_ROPE) + mask_r = offs_r < D_HALF + offs_r_pair = offs_r ^ 1 + mask_r_pair = offs_r_pair < D_HALF + + rope_row = n_idx * D_HALF + cos_v = tl.load(rope_cos_ptr + rope_row + offs_r, mask=mask_r, other=1.0).to(tl.float32) + sin_v = tl.load(rope_sin_ptr + rope_row + offs_r, mask=mask_r, other=0.0).to(tl.float32) + + # Load second-half raw values and their pair partners + rope_base = row_base + D_HALF + q_r = tl.load(q_raw_ptr + rope_base + offs_r, mask=mask_r, other=0.0).to(tl.float32) + k_r = tl.load(k_raw_ptr + rope_base + offs_r, mask=mask_r, other=0.0).to(tl.float32) + v_r = tl.load(v_raw_ptr + rope_base + offs_r, mask=mask_r, other=0.0).to(tl.float32) + q_r_pair = tl.load(q_raw_ptr + rope_base + offs_r_pair, mask=mask_r_pair, other=0.0).to(tl.float32) + k_r_pair = tl.load(k_raw_ptr + rope_base + offs_r_pair, mask=mask_r_pair, other=0.0).to(tl.float32) + v_r_pair = tl.load(v_raw_ptr + rope_base + offs_r_pair, mask=mask_r_pair, other=0.0).to(tl.float32) + + q_nw_r = tl.load(q_norm_w_ptr + nw_off + D_HALF + offs_r, mask=mask_r, other=0.0).to(tl.float32) + k_nw_r = tl.load(k_norm_w_ptr + nw_off + D_HALF + offs_r, mask=mask_r, other=0.0).to(tl.float32) + q_nw_r_pair = tl.load(q_norm_w_ptr + nw_off + D_HALF + offs_r_pair, mask=mask_r_pair, other=0.0).to(tl.float32) + k_nw_r_pair = tl.load(k_norm_w_ptr + nw_off + D_HALF + offs_r_pair, mask=mask_r_pair, other=0.0).to(tl.float32) + + q_r_n = q_r * q_inv_rms * q_nw_r + q_r_n = tl.where(q_r_n > 0, q_r_n, 0.0) + q_r_pair_n = q_r_pair * q_inv_rms * q_nw_r_pair + q_r_pair_n = tl.where(q_r_pair_n > 0, q_r_pair_n, 0.0) + + k_r_n = k_r * k_inv_rms * k_nw_r + k_r_n = tl.where(k_r_n > 0, k_r_n, 0.0) * K_SCALE + k_r_pair_n = k_r_pair * k_inv_rms * k_nw_r_pair + k_r_pair_n = tl.where(k_r_pair_n > 0, k_r_pair_n, 0.0) * K_SCALE + + # Pre-UCPE ||k||^2 contribution from second half (using post-ReLU/scale k_r_n) + k_r_n_masked = tl.where(mask_r, k_r_n, 0.0) + k_pre_rope_sq = tl.sum(k_r_n_masked * k_r_n_masked) + + q_rope_out = q_r_n * cos_v + q_r_pair_n * sin_v + k_rope_out = k_r_n * cos_v + k_r_pair_n * sin_v + v_rope_out = v_r * cos_v + v_r_pair * sin_v + + # Post-UCPE ||k||^2 contribution from second half + k_rope_masked = tl.where(mask_r, k_rope_out, 0.0) + k_post_rope_sq = tl.sum(k_rope_masked * k_rope_masked) + + # Store scalar per-token norm squares + norm_out_idx = (b_idx * H + h_idx) * N + n_idx + tl.store(k_pre_norm_sq_ptr + norm_out_idx, k_pre_half_sq + k_pre_rope_sq) + tl.store(k_post_norm_sq_ptr + norm_out_idx, k_post_half_sq + k_post_rope_sq) + + # ================================================================== + # Store outputs in (B, H, D, N) layout: ptr[b, h, d, n] = base_bh + d*N + n + # ================================================================== + out_base = b_idx * (H * D * N) + h_idx * (D * N) + n_idx + + # First half: d = g*4 + i, write at out_base + d*N (strided by N). + offs_d_half = offs_g[:, None] * 4 + offs_i[None, :] # (BLOCK_GROUPS, 4) + mask_d_half = mask_g[:, None] + tl.store(q_out_ptr + out_base + offs_d_half * N, q_half_out, mask=mask_d_half) + tl.store(k_out_ptr + out_base + offs_d_half * N, k_half_out, mask=mask_d_half) + tl.store(v_out_ptr + out_base + offs_d_half * N, v_half_out, mask=mask_d_half) + + # Second half (RoPE region): d = D_HALF + r + offs_d_r = D_HALF + offs_r # (BLOCK_D_ROPE,) + tl.store(q_out_ptr + out_base + offs_d_r * N, q_rope_out, mask=mask_r) + tl.store(k_out_ptr + out_base + offs_d_r * N, k_rope_out, mask=mask_r) + tl.store(v_out_ptr + out_base + offs_d_r * N, v_rope_out, mask=mask_r) + + +def cam_prep_func( + q_raw: torch.Tensor, + k_raw: torch.Tensor, + v_raw: torch.Tensor, + *, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + proj_q: torch.Tensor, # (B, N, 4, 4) + proj_kv: torch.Tensor, # (B, N, 4, 4) + rope_cos: torch.Tensor, # (N, D//2) + rope_sin: torch.Tensor, # (N, D//2) + k_scale: float, + norm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused RMSNorm + ReLU + (K-scale on K) + UCPE 4x4 + RoPE for the cam branch. + + Args: + q_raw, k_raw, v_raw: ``(B, N, H, D)`` contiguous (any fp dtype). + ``K`` must already have the short convolution applied. + q_norm_weight, k_norm_weight: ``(C,) = (H*D,)`` fp32. + proj_q, proj_kv: ``(B, N, 4, 4)`` fp32 (``P_T`` and ``P_inv`` in UCPE). + rope_cos, rope_sin: ``(N, D//2)`` fp32 interleaved-pair tables. + k_scale: ``(D^-0.5) * (S^-0.5)``. + norm_eps: RMSNorm epsilon. + + Returns: + q_trans, k_trans, v_trans: ``(B, H, D, N)`` same dtype as ``q_raw``. inflation_sq: ``(B, H, N)`` fp32, ratio + ``(||k_post_ucpe|| / ||k_pre_ucpe||)^2`` per token/head. + """ + _require_triton("cam_prep_func") + B, N, H, D = q_raw.shape + assert k_raw.shape == q_raw.shape and v_raw.shape == q_raw.shape + assert D % 2 == 0 and (D // 2) % 4 == 0, f"D={D} must be 2x and (D/2) % 4 == 0" + D_half = D // 2 + N_groups = D_half // 4 + + assert q_raw.is_contiguous() and k_raw.is_contiguous() and v_raw.is_contiguous() + assert proj_q.shape == (B, N, 4, 4) and proj_q.is_contiguous() + assert proj_kv.shape == (B, N, 4, 4) and proj_kv.is_contiguous() + assert rope_cos.shape == (N, D_half) and rope_cos.is_contiguous() + assert rope_sin.shape == (N, D_half) and rope_sin.is_contiguous() + assert q_norm_weight.numel() == H * D and q_norm_weight.dtype == torch.float32 + assert k_norm_weight.numel() == H * D and k_norm_weight.dtype == torch.float32 + + # Precompute inv-RMS over full C channels (shared across heads per token). + q_inv_rms = _precompute_cam_inv_rms(q_raw, norm_eps) + k_inv_rms = _precompute_cam_inv_rms(k_raw, norm_eps) + + out_dtype = q_raw.dtype + q_out = torch.empty(B, H, D, N, dtype=out_dtype, device=q_raw.device) + k_out = torch.empty(B, H, D, N, dtype=out_dtype, device=q_raw.device) + v_out = torch.empty(B, H, D, N, dtype=out_dtype, device=q_raw.device) + k_pre_sq = torch.empty(B, H, N, dtype=torch.float32, device=q_raw.device) + k_post_sq = torch.empty(B, H, N, dtype=torch.float32, device=q_raw.device) + + BLOCK_D_ROPE = triton.next_power_of_2(D_half) + BLOCK_GROUPS = triton.next_power_of_2(N_groups) + + grid = (B * N * H,) + _cam_prep_kernel[grid]( + q_raw, + k_raw, + v_raw, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + q_out, + k_out, + v_out, + k_pre_sq, + k_post_sq, + H=H, + N=N, + D=D, + D_HALF=D_half, + N_GROUPS=N_groups, + K_SCALE=k_scale, + BLOCK_D_ROPE=BLOCK_D_ROPE, + BLOCK_GROUPS=BLOCK_GROUPS, + num_warps=1, + ) + # inflation_sq = (clamp(sqrt(post), 1e-6) / clamp(sqrt(pre), 1e-6))^2 + # = clamp(post, 1e-12) / clamp(pre, 1e-12) (equivalent). + inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12) + return q_out, k_out, v_out, inflation_sq + + +_CAM_IDENTITY_CACHE: dict[ + tuple[str, int | None, int, int, int], tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +] = {} + +# ════════════════════════════════════════════════════════════════ +# Per-architecture launch config (auto-selected via compute capability) +# ════════════════════════════════════════════════════════════════ +# +# Empirically tuned at production config (B=1..8, T=11, S=920, H=20, D=112) on +# A100 / H100 / GB200. Two effects matter: +# +# 1. **Precision sets BLOCK_S**: fp32 operand fragments are 2× the size of +# bf16. BLOCK_S=64 + fp32 → register spills (catastrophic, 40-100× slower). +# BLOCK_S=32 + fp32 → no spills. So fp32 mode forces BLOCK_S=32 everywhere. +# +# 2. **Arch sets BLOCK_S for bf16**: A100 (192 KB SRAM, fewer registers per +# block) prefers BLOCK_S=32 even at bf16. H100/GB200 (228 KB SRAM) tolerate +# BLOCK_S=64 cleanly at bf16. +# +# Each entry: (phase_a_warps, phase_a_BLOCK_S, +# phase_b_warps, phase_b_stages, +# phase_c_warps, phase_c_BLOCK_S, phase_c_stages) + +# ── Launch-config tuning table ───────────────────────────────────── +# +# We tune 8 knobs across 3 phases: +# Phase A : (nw, BS) streaming accumulator in registers +# Phase B : (nw, use_acc, ns) serial-F scan with persistent M in regs +# Phase C : (nw, BS, ns) streams Pass-2 output; loads fp32 M[128,128] +# +# Each arch × precision combination gets a named entry below. Values come from +# empirical sweeps (see commit log: T6 A100/H100 sweep 2026-04-19; Blackwell-DC +# 2026-04-20; Spark GB10 tuning notes in commits 5da52db6 / 3ad104d0) and from +# kernel-structure analysis (Phase B's persistent M[128,128] fp32 is 64 KB → nw +# controls register spread; Phase C's loaded M[128,128] is 64 KB → BS controls +# transient SMEM footprint). +# +# Adding a new arch: pick the closest existing bucket, then override individual +# fields in _CHUNKWISE_SHAPE_OVERRIDES once a targeted sweep lands. + + +@dataclass(frozen=True) +class _PhaseCfg: + nw: int # num_warps + BS: int = 0 # BLOCK_S (Phase A/C only; 0 = N/A for Phase B) + ns: int = 1 # num_stages + use_acc: bool = False # Phase B only: fold A_f via MMA accumulator + + +@dataclass(frozen=True) +class _ChunkwiseCfg: + A: _PhaseCfg + B: _PhaseCfg + C: _PhaseCfg + + def as_tuple(self) -> tuple: + """Flatten to the 8-tuple the legacy API returns.""" + return ( + self.A.nw, + self.A.BS, + self.B.nw, + self.B.ns, + self.B.use_acc, + self.C.nw, + self.C.BS, + self.C.ns, + ) + + +# ────────────────────────────────────────────────────────────────── +# Primary tuning table: (arch_key, prec_key) → _ChunkwiseCfg. +# Arch keys: +# "ampere" sm_80 A100 (164 KB SRAM, no WGMMA) +# "hopper" sm_90 H100 (228 KB SRAM, WGMMA) +# "blackwell_dc" sm_100 B200 / GB200 (228 KB SRAM, WGMMA v2) +# "blackwell_spark" sm_120+ with < 150 KB SRAM 5090 / GB10 (~102 KB SRAM) +# Prec keys: +# "bf16" dot_prec == 0 (bf16 TC, half-size operand fragments) +# "fp32" dot_prec >= 1 (TF32 TC or IEEE Markidis 3-pass; same launch shape) +# ────────────────────────────────────────────────────────────────── +_CHUNKWISE_TUNING: dict[tuple[str, str], _ChunkwiseCfg] = { + # A100: smaller SRAM than Hopper, no WGMMA → bigger CTAs hide MMA latency. + # Phase B fp32 needs nw=32 to spread persistent M across warps (no acc-fusion + # available pre-Hopper, so ns=2 fills the MMA pipeline slot instead). + ("ampere", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=32), + B=_PhaseCfg(nw=8, use_acc=False, ns=1), + C=_PhaseCfg(nw=4, BS=32, ns=1), # nw=4 bf16 C: 27% faster than nw=8 per T6 + ), + ("ampere", "fp32"): _ChunkwiseCfg( + # 2026-04-30 PM retune: Phase A nw=8 → 16 BS=32 yields 8-13× speedup + # across F ∈ {3, 5, 11, 14, 17, 20} (cos=1.0 verified). Old nw=8 was a + # legacy default never re-swept; sweep showed nw=16 dominates every F. + # Closes A100 sink/rolling chunkwise regression where Phase B was + # already optimal (sub-percent tuning gap) — Phase A was the bottleneck. + A=_PhaseCfg(nw=16, BS=32), + B=_PhaseCfg(nw=32, use_acc=False, ns=2), # ns=2 fills pipe (no acc-fusion) + C=_PhaseCfg(nw=16, BS=32, ns=1), # 2026-04-30 retune: nw=16 BS=32 is 2.8x faster (was nw=8 BS=16) + ), + # Hopper (H100): WGMMA + 228 KB SRAM → big tiles win at bf16. + # Phase B fp32 uses acc-fusion (MMA accumulator folds A_f in one op, +12%). + ("hopper", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=64), + B=_PhaseCfg(nw=4, use_acc=False, ns=1), # small CTAs pack better on WGMMA + C=_PhaseCfg(nw=8, BS=32, ns=1), + ), + ("hopper", "fp32"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=32), # fp32 operand 2× bigger → half BS + B=_PhaseCfg( + nw=32, use_acc=False, ns=1 + ), # 2026-04-29 retune: acc_fusion=False is 3x faster post precision-gate fix + C=_PhaseCfg(nw=16, BS=32, ns=1), # 2026-04-30 retune: nw=16 BS=32 is 1.7x faster (was nw=8 BS=16) + ), + # Blackwell-DC (B200 / GB200): 228 KB SRAM + improved WGMMA codegen. + # bf16 likes small CTAs (nw=4); fp32 stays at nw=8 (nw=4 + BS=64 fp32 = 92× regression). + ("blackwell_dc", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=4, BS=64), + B=_PhaseCfg(nw=4, use_acc=False, ns=1), + C=_PhaseCfg(nw=8, BS=64, ns=1), # 228 KB SRAM leaves room for BS=64 bf16 + ), + ("blackwell_dc", "fp32"): _ChunkwiseCfg( + A=_PhaseCfg( + nw=8, BS=128 + ), # 2026-04-30 retune: nw=8 BS=128 ~5% faster at production F=3-6 (sweep across F=3,5,6,11) + B=_PhaseCfg( + nw=32, use_acc=False, ns=3 + ), # 2026-04-29 retune: 14x faster (was nw=8 acc=True 17ms; now nw=32 ns=3 acc=False 1.23ms) + C=_PhaseCfg( + nw=4, BS=64, ns=1 + ), # 2026-04-30 retune: nw=4 BS=64 is 3-5x faster than old nw=8 BS=16 (sweep 2026-04-30) + ), + # Blackwell-Spark (5090 / GB10, ~102 KB SRAM): shares SRAM penalty of small + # chips but not Blackwell-DC's WGMMA-v2 register-spread benefit. Empirically + # behaves like Hopper at fp32 (Phase B wants nw=32 to spread persistent M + # across warps, not nw=8 like DC). BS shrunk one step vs DC; Phase A bf16 + # wants nw=8 (nw=4 tested 22× slower per 2026-04-20 sweep). + # Sweep 2026-04-24 (prod dim F=11 S=920): Phase B nw=32 gives 1.84×/2.65× + # (GB10/5090) at fp32 over prior nw=8 setting. + ("blackwell_spark", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=32), + B=_PhaseCfg(nw=8, use_acc=False, ns=1), # nw=8 (not 4) at bf16: ~5% across F=3,6,11 + # 2026-05-06 P1/P2 retune (5090, F=11 S=920): C.nw=4 BS=32 is ~3.5% + # faster than nw=8 (Phase C is bandwidth-bound, fewer warps schedules + # better on the small SRAM). BS=64 bf16 on Spark OOMs SRAM. + C=_PhaseCfg(nw=4, BS=32, ns=1), + ), + ("blackwell_spark", "fp32"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=16), # fp32 operand 2× bigger → BS=16 (half of DC's 32) + # 2026-05-06 retune: nw=16 OOMs the 102 KB SRAM cap at TF32 on 5090 + # (131 KB needed). nw=8 fits and is within noise of the prior nw=16 + # benchmark. The Phase B D-tile path (auto-enabled on spark, see + # `_pick_phase_b_d_splits`) is ~2.6× faster than this baseline at TF32 + # and ~13% faster at IEEE — these baseline params only apply when + # PHASE_B_D_SPLITS=1 is forced. + B=_PhaseCfg(nw=8, use_acc=False, ns=1), + C=_PhaseCfg(nw=8, BS=16, ns=1), # binding constraint: M.fp32 64 KB + Q stage + ), +} + + +# ────────────────────────────────────────────────────────────────── +# Shape-aware override table: empty by default. Keyed by +# (arch_key, prec_key, shape_hint) +# where shape_hint is a free-form string (e.g. "small_BH", "large_F", +# "B>=8") chosen when populating. Lookup is exact-match; values are +# full `_ChunkwiseCfg` instances (no partial overrides — copy-paste +# from `_CHUNKWISE_TUNING` and edit the one phase you want to change). +# +# Leave empty unless a targeted sweep shows a particular shape regresses +# with the broad arch config. Adding here is strictly additive — base +# table remains the fallback. +# ────────────────────────────────────────────────────────────────── +_CHUNKWISE_SHAPE_OVERRIDES: dict[tuple[str, str, str], _ChunkwiseCfg] = {} + + +# Per-(cap, dot_prec) exact overrides (pins a specific GPU model if the arch +# bucket is wrong for it). Also empty by default. +_ARCH_OVERRIDES: dict = {} + + +def _arch_key(cap: tuple) -> str: + """Map compute capability → named arch bucket in `_CHUNKWISE_TUNING`. + + Blackwell (cap[0] >= 10) is split into "blackwell_dc" and "blackwell_spark" by SRAM size (≥150 KB vs less). Without + CUDA or for unknown archs we default to the conservative "ampere" bucket. + """ + if cap[0] == 8: + return "ampere" + if cap[0] == 9: + return "hopper" + if cap[0] >= 10: + has_big_sram = True + if torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + smem = getattr(props, "shared_memory_per_multiprocessor", 228 * 1024) + has_big_sram = smem >= 150 * 1024 + return "blackwell_dc" if has_big_sram else "blackwell_spark" + return "ampere" + + +def _prec_key(dot_prec: int) -> str: + return "fp32" if dot_prec >= 1 else "bf16" + + +def _auto_config(dot_prec: int, cap: tuple, shape_hint: str | None = None) -> tuple: + """Look up chunkwise kernel launch params from the tuning table. + + Resolution order: + 1. `_ARCH_OVERRIDES[(cap, dot_prec)]` — exact-capability pin, highest priority. + 2. `_CHUNKWISE_SHAPE_OVERRIDES[(arch, prec, shape_hint)]` — sweep-driven overrides. + 3. `_CHUNKWISE_TUNING[(arch, prec)]` — primary per-(arch, prec) table. + 4. Fallback to ("ampere", prec) if the arch is unrecognised. + + Returns the legacy 8-tuple `(a_nw, a_BS, b_nw, b_ns, b_use_acc, c_nw, c_BS, c_ns)` for backward compatibility with + `_get_arch_config` callers. + """ + arch = _arch_key(cap) + prec = _prec_key(dot_prec) + + if shape_hint is not None: + cfg = _CHUNKWISE_SHAPE_OVERRIDES.get((arch, prec, shape_hint)) + if cfg is not None: + return cfg.as_tuple() + + cfg = _CHUNKWISE_TUNING.get((arch, prec)) or _CHUNKWISE_TUNING[("ampere", prec)] + return cfg.as_tuple() + + +def _get_arch_config( + dot_precision: int = 0, + shape_hint: str | None = None, + device: torch.device | int | None = None, +): + """Returns (a_warps, a_BLOCK_S, b_warps, b_stages, b_use_acc_fusion, + c_warps, c_BLOCK_S, c_stages). + + dot_precision: 0=bf16 TC, 1=TF32 TC, 2=IEEE fp32. shape_hint: optional string key for `_CHUNKWISE_SHAPE_OVERRIDES`. + device: device whose capability drives the lookup. Defaults to the + current CUDA device — pass ``qkv.device`` (or any input tensor's device) when launching kernels in + heterogeneous or multi-GPU single-process setups so the right tuning bucket is chosen. + """ + if not torch.cuda.is_available(): + cap = (9, 0) # assume modern when querying from CPU + else: + if device is None: + dev_idx = torch.cuda.current_device() + elif isinstance(device, int): + dev_idx = device + else: + dev_idx = device.index if device.index is not None else torch.cuda.current_device() + cap = torch.cuda.get_device_capability(dev_idx) + key = (cap, dot_precision) + if key in _ARCH_OVERRIDES: + return _ARCH_OVERRIDES[key] + return _auto_config(dot_precision, cap, shape_hint) + + +# ════════════════════════════════════════════════════════════════ +# Phase A — split into KV and Z kernels +# ════════════════════════════════════════════════════════════════ + + +@triton.jit +def _phase_a_kv_kernel( + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + beta_ptr, + k_inv_rms_ptr, + k_norm_w_ptr, + rope_cos_ptr, + rope_sin_ptr, + I_minus_P_kv_ptr, # output: (I - K_rot^T diag(β) K_rot) + A_ptr, # output: K_rot^T diag(β) V + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + K_SCALE, + NORM_EPS: tl.constexpr, + DOT_PRECISION: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, + SKIP_RELU: tl.constexpr = False, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + pid_b = pid // (H * F) + pid_hf = pid % (H * F) + pid_h = pid_hf // F + pid_f = pid_hf % F + bh = pid_b * H + pid_h + N: tl.constexpr = F * S + + qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h + beta_bhf = beta_ptr + bh * (F * S) + pid_f * S + I_P_kv_bhf = I_minus_P_kv_ptr + bh * F * BLOCK_D * BLOCK_D + pid_f * BLOCK_D * BLOCK_D + A_bhf = A_ptr + bh * F * BLOCK_D * BLOCK_D + pid_f * BLOCK_D * BLOCK_D + + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < D + offs_d_pair = offs_d ^ 1 + mask_d_pair = offs_d_pair < D + + nw_offset = pid_h * D + k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) + k_nw_pair = tl.load(k_norm_w_ptr + nw_offset + offs_d_pair, mask=mask_d_pair, other=0.0).to(tl.float32) + + # KV stream accumulators (in-loop fp32 to avoid bf16 round-off compounding) + P_kv_acc = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + A_acc = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + + k_scale = K_SCALE + n_base = pid_f * S + + for s0 in range(0, S, BLOCK_S): + offs_s = s0 + tl.arange(0, BLOCK_S) + mask_s = offs_s < S + mask_sd = mask_s[:, None] & mask_d[None, :] + n_idx = n_base + offs_s + + k_ptrs = qkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d[None, :] * stride_d + v_ptrs = qkv_bh + n_idx[:, None] * stride_n + 2 * stride_3 + offs_d[None, :] * stride_d + K_raw = tl.load(k_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + V_raw = tl.load(v_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + beta_t = tl.load(beta_bhf + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + k_inv_rms = tl.load(k_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0).to(tl.float32) + K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] + if SKIP_RELU: + K = K_normed * k_scale + else: + K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale + + K_pair_raw = tl.reshape( + tl.flip(tl.reshape(K_raw, (BLOCK_S, BLOCK_D // 2, 2)), dim=2), + (BLOCK_S, BLOCK_D), + ) + K_pair_normed = K_pair_raw * k_inv_rms[:, None] * k_nw_pair[None, :] + if SKIP_RELU: + K_pair = K_pair_normed * k_scale + else: + K_pair = tl.where(K_pair_normed > 0, K_pair_normed, 0.0) * k_scale + + rope_ptrs = n_idx[:, None] * D + offs_d[None, :] + Cos = tl.load(rope_cos_ptr + rope_ptrs, mask=mask_sd, other=1.0).to(tl.float32) + Sin = tl.load(rope_sin_ptr + rope_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + K_rot = K * Cos + K_pair * Sin + + beta_Krot = beta_t[:, None] * K_rot + beta_V = beta_t[:, None] * V_raw + + K_rot_T = tl.trans(K_rot) + P_kv_acc += tl.dot( + K_rot_T.to(dot_dtype), beta_Krot.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip + ) + A_acc += tl.dot(K_rot_T.to(dot_dtype), beta_V.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + + # Store bf16 outputs. Padded positions are 0 by construction (K_rot is 0 outside D). + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + diag_in_range = (offs_d[:, None] == offs_d[None, :]) & mask_d[:, None] & mask_d[None, :] + I_minus_P_kv = tl.where(diag_in_range, 1.0 - P_kv_acc, -P_kv_acc) + if DOT_PRECISION >= 1: + tl.store(I_P_kv_bhf + offs_dd, I_minus_P_kv) + tl.store(A_bhf + offs_dd, A_acc) + else: + tl.store(I_P_kv_bhf + offs_dd, I_minus_P_kv.to(tl.bfloat16)) + tl.store(A_bhf + offs_dd, A_acc.to(tl.bfloat16)) + + +@triton.jit +def _phase_a_z_kernel( + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + beta_ptr, + k_inv_rms_ptr, + k_norm_w_ptr, + I_minus_P_z_ptr, # output: (I - K^T diag(β) K) + B_ptr, # output: K^T β + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + K_SCALE, + NORM_EPS: tl.constexpr, + DOT_PRECISION: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Z stream: uses K (no RoPE). Cheaper than KV — no V load, no RoPE compute, + no K_pair derivation.""" + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + pid_b = pid // (H * F) + pid_hf = pid % (H * F) + pid_h = pid_hf // F + pid_f = pid_hf % F + bh = pid_b * H + pid_h + N: tl.constexpr = F * S + + qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h + beta_bhf = beta_ptr + bh * (F * S) + pid_f * S + I_P_z_bhf = I_minus_P_z_ptr + bh * F * BLOCK_D * BLOCK_D + pid_f * BLOCK_D * BLOCK_D + B_bhf = B_ptr + bh * F * BLOCK_D + pid_f * BLOCK_D + + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < D + + nw_offset = pid_h * D + k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) + + P_z_acc = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + B_acc = tl.zeros([BLOCK_D], dtype=tl.float32) + + k_scale = K_SCALE + n_base = pid_f * S + + for s0 in range(0, S, BLOCK_S): + offs_s = s0 + tl.arange(0, BLOCK_S) + mask_s = offs_s < S + mask_sd = mask_s[:, None] & mask_d[None, :] + n_idx = n_base + offs_s + + # Only K_raw needed (no V, no Cos/Sin) + k_ptrs = qkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d[None, :] * stride_d + K_raw = tl.load(k_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + beta_t = tl.load(beta_bhf + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + k_inv_rms = tl.load(k_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0).to(tl.float32) + K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] + K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale + + beta_K = beta_t[:, None] * K + + K_T = tl.trans(K) + P_z_acc += tl.dot(K_T.to(dot_dtype), beta_K.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + B_acc += tl.sum(beta_K, axis=0) + + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + diag_in_range = (offs_d[:, None] == offs_d[None, :]) & mask_d[:, None] & mask_d[None, :] + I_minus_P_z = tl.where(diag_in_range, 1.0 - P_z_acc, -P_z_acc) + + if DOT_PRECISION >= 1: + tl.store(I_P_z_bhf + offs_dd, I_minus_P_z) + else: + tl.store(I_P_z_bhf + offs_dd, I_minus_P_z.to(tl.bfloat16)) + # B stays fp32 (vector, only 0.5 KB, negligible HBM cost) + tl.store(B_bhf + offs_d, B_acc) + + +def phase_a( + qkv: torch.Tensor, + beta: torch.Tensor, + q_inv_rms: torch.Tensor, + k_inv_rms: torch.Tensor, + q_norm_w: torch.Tensor, + k_norm_w: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + F: int, + S: int, + k_scale: float = 1.0, + norm_eps: float = 1e-5, + num_warps: int | None = None, + num_stages: int = 1, + BLOCK_S: int | None = None, + dot_precision: int = 0, + skip_relu: bool = False, + skip_z: bool = False, +): + """Compute (I-P_kv), A, (I-P_z), B for all (B, H, F) via 2 kernels (KV + Z). + + `skip_relu=True` makes the K-stream prep a pure linear chain (no ReLU on K_normed * k_scale). Used by the + camera-branch chunkwise wrapper, where K has already been ReLU'd by the cam_prep kernel and subsequently rotated by + UCPE+RoPE — re-applying ReLU on the rotated values would clobber legitimate negatives. + + `skip_z=True` skips the Phase A Z kernel entirely and returns placeholder tensors for I_P_z and B_z. Used by + NUM_ONLY callers (camera branch) to avoid wasted Z-stream prep when the denominator scan won't be used. + """ + # Auto-pick (num_warps, BLOCK_S) per arch+precision unless overridden + if num_warps is None or BLOCK_S is None: + a_w, a_bs, *_ = _get_arch_config(dot_precision, device=qkv.device) + if num_warps is None: + num_warps = a_w + if BLOCK_S is None: + BLOCK_S = a_bs + B, N, three, H, D = qkv.shape + assert three == 3 and N == F * S + BLOCK_D = triton.next_power_of_2(D) + BH = B * H + + # FAIR-COMPARE PATCH: keep fp32 inter-phase bridge at P0/P1 to match pytorch/fused + bridge_dtype = torch.float32 if dot_precision >= 1 else torch.bfloat16 + I_P_kv = torch.empty(BH, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=bridge_dtype) + A = torch.empty(BH, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=bridge_dtype) + + beta_c = beta.contiguous() + grid = (BH * F,) + + _phase_a_kv_kernel[grid]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta_c, + k_inv_rms, + k_norm_w, + rope_cos, + rope_sin, + I_P_kv, + A, + H=H, + F=F, + S=S, + D=D, + K_SCALE=k_scale, + NORM_EPS=norm_eps, + DOT_PRECISION=dot_precision, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + SKIP_RELU=skip_relu, + num_warps=num_warps, + num_stages=num_stages, + ) + + if skip_z: + # NUM_ONLY callers (camera branch) do not consume the Z scan. Return + # placeholders and let Phase B skip all Z loads/stores as well. + I_P_z = torch.empty(1, device=qkv.device, dtype=bridge_dtype) + B_z = torch.empty(1, device=qkv.device, dtype=torch.float32) + return I_P_kv, A, I_P_z, B_z + + I_P_z = torch.empty(BH, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=bridge_dtype) + # B stays fp32 — small vector (0.5 KB/frame), no benefit to downcast + B_z = torch.empty(BH, F, BLOCK_D, device=qkv.device, dtype=torch.float32) + + _phase_a_z_kernel[grid]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta_c, + k_inv_rms, + k_norm_w, + I_P_z, + B_z, + H=H, + F=F, + S=S, + D=D, + K_SCALE=k_scale, + NORM_EPS=norm_eps, + DOT_PRECISION=dot_precision, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_warps=num_warps, + num_stages=num_stages, + ) + return I_P_kv, A, I_P_z, B_z + + +# ════════════════════════════════════════════════════════════════ +# Phase B — serial scan, uses pre-stored (I - P) so MMA folds in M +# ════════════════════════════════════════════════════════════════ + + +@triton.jit +def _phase_b_kernel( + I_P_kv_ptr, + A_ptr, + I_P_z_ptr, + B_ptr, + decay_ptr, + M_fwd_ptr, + z_fwd_ptr, + M_rev_ptr, + z_rev_ptr, + init_state_kv_ptr, # (BH, BLOCK_D, BLOCK_D) — read when LOAD_INIT_STATE=1 + init_state_z_ptr, # (BH, BLOCK_D) + final_state_kv_ptr, # (BH, BLOCK_D, BLOCK_D) — written when SAVE_FINAL_STATE=1 + final_state_z_ptr, # (BH, BLOCK_D) + BH: tl.constexpr, + F: tl.constexpr, + BLOCK_D: tl.constexpr, + DOT_PRECISION: tl.constexpr, + USE_ACC_FUSION: tl.constexpr, + LOAD_INIT_STATE: tl.constexpr, # forward scan seeded with init state (vs zeros) + SAVE_FINAL_STATE: tl.constexpr, # write M_{F-1} of forward scan to final_state_* + DIRECTION: tl.constexpr, # 0=both, 1=fwd-only, 2=rev-only + COMBINED_HISTORY: tl.constexpr, # 1 → rev branch read-add-stores into M_fwd_ptr + # (M_hist[f] = M_fwd[f] + M_rev[f]); skips the F-1 zero-write so the fwd + # value at F-1 is preserved (rev contribution there is exactly zero anyway). + # Only meaningful when DIRECTION=0. Saves one Phase C launch + one M-shaped + # buffer downstream (Phase C runs once on M_hist instead of twice). + SKIP_Z: tl.constexpr, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + bh = pid + + offs_d = tl.arange(0, BLOCK_D) + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + + # ── Forward scan (skip when DIRECTION=2 i.e. rev-only) ── + if DIRECTION != 2: + if LOAD_INIT_STATE: + M = tl.load(init_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd).to(tl.float32) + if not SKIP_Z: + z = tl.load(init_state_z_ptr + bh * BLOCK_D + offs_d).to(tl.float32) + else: + M = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + for f in range(F): + I_P_kv_f = tl.load(I_P_kv_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd) + A_f = tl.load(A_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd) + g_f = tl.load(decay_ptr + bh * F + f).to(tl.float32) + + # M = g · (I - P_kv) M + A_f + if USE_ACC_FUSION: + # Pre-scale (I-P) by g, accumulate A_f directly via the MMA accumulator. + # Result: A_f + g·(I-P)·M in one MMA — no separate M_temp tensor. + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot(I_P_kv_f.to(dot_dtype), M.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + M = g_f * M_temp + A_f + + tl.store(M_fwd_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd, M) + if not SKIP_Z: + I_P_z_f = tl.load(I_P_z_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d) + # z = g · (I - P_z) z + B_f + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + tl.store(z_fwd_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d, z) + + # Save terminal forward state for state-cached inference (autoregressive sampling). + if SAVE_FINAL_STATE: + tl.store(final_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd, M) + if not SKIP_Z: + tl.store(final_state_z_ptr + bh * BLOCK_D + offs_d, z) + + # ── Reverse scan (skip when DIRECTION=1 i.e. fwd-only) ── + if DIRECTION != 1: + M = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + # COMBINED_HISTORY mode: rev contributions get read-add-stored into the + # fwd buffer (which thereby becomes M_hist = M_fwd + M_rev). The F-1 + # zero-write is skipped so M_hist[F-1] keeps the fwd value (rev value + # there is zero by construction, so no add needed). + if not COMBINED_HISTORY: + tl.store(M_rev_ptr + bh * F * BLOCK_D * BLOCK_D + (F - 1) * BLOCK_D * BLOCK_D + offs_dd, M) + if not SKIP_Z: + tl.store(z_rev_ptr + bh * F * BLOCK_D + (F - 1) * BLOCK_D + offs_d, z) + for f_iter in range(F - 1): + f_src = F - 1 - f_iter + f_dst = f_src - 1 + I_P_kv_f = tl.load(I_P_kv_ptr + bh * F * BLOCK_D * BLOCK_D + f_src * BLOCK_D * BLOCK_D + offs_dd) + A_f = tl.load(A_ptr + bh * F * BLOCK_D * BLOCK_D + f_src * BLOCK_D * BLOCK_D + offs_dd) + g_f = tl.load(decay_ptr + bh * F + f_src).to(tl.float32) + + if USE_ACC_FUSION: + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot(I_P_kv_f.to(dot_dtype), M.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + M = g_f * M_temp + A_f + + if not SKIP_Z: + I_P_z_f = tl.load(I_P_z_ptr + bh * F * BLOCK_D * BLOCK_D + f_src * BLOCK_D * BLOCK_D + offs_dd) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f_src * BLOCK_D + offs_d) + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + + if COMBINED_HISTORY: + # Read-add-store into the fwd buffer. The fwd loop has already + # written M_fwd[f_dst] to this slot; we add the rev contribution + # in place. Stays in L1/L2 since fwd just touched it. + M_addr = M_fwd_ptr + bh * F * BLOCK_D * BLOCK_D + f_dst * BLOCK_D * BLOCK_D + offs_dd + tl.store(M_addr, tl.load(M_addr) + M) + if not SKIP_Z: + z_addr = z_fwd_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d + tl.store(z_addr, tl.load(z_addr) + z) + else: + tl.store(M_rev_ptr + bh * F * BLOCK_D * BLOCK_D + f_dst * BLOCK_D * BLOCK_D + offs_dd, M) + if not SKIP_Z: + tl.store(z_rev_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d, z) + + +def phase_b_triton( + I_P_kv, + A, + I_P_z, + B, + decay, + F, + num_warps=None, + num_stages=None, + use_acc_fusion=None, + dot_precision=0, + init_state_kv=None, + init_state_z=None, + return_final_state=False, + direction=0, + combined_history=False, + skip_z=False, +): + """Phase B serial-F scan over (B*H,). + + Forward scan can be seeded with `init_state_kv`/`init_state_z` (autoregressive sampling chunk > 0) and can write + the terminal `M_{F-1}`/`z_{F-1}` to caller- provided buffers when `return_final_state=True`. + + `direction`: 0=both (default), 1=forward-only, 2=reverse-only. Forward-only skips reverse scan + reverse output + buffers; reverse-only skips forward scan + state load/save. Used by single-direction state-cached entry points. + + `combined_history` (only meaningful with direction=0): the rev branch read-add-stores into the fwd buffer so its + contents become M_hist[f] = M_fwd[f] + M_rev[f] (and same for z). Lets the caller run Phase C exactly once on the + combined history, since Phase C is linear in M and z (`Q @ (M_fwd + M_rev) = Q @ M_fwd + Q @ M_rev`). When set, + M_rev/z_rev outputs are placeholder dummies; only M_fwd/z_fwd carry data. + + `skip_z`: skip the denominator/Z recurrence entirely. Used by camera numerator-only scans where Phase C runs with + `num_only=True`. + + Returns (M_fwd, z_fwd, M_rev, z_rev) — and additionally (final_kv, final_z) when return_final_state=True. + Skipped-direction outputs are returned as a 1-element placeholder tensor (kernel never touches them when DIRECTION + gates them off); callers should always discard the slot they didn't ask for. Reverse scan is always seeded with + zeros (per upstream's bidi state-cache convention — only forward state is cached). + """ + BH = I_P_kv.shape[0] + _, _, BLOCK_D, _ = A.shape # A is always full [BH, F, BLOCK_D, BLOCK_D] + device, fdtype = I_P_kv.device, torch.float32 + + if num_warps is None or num_stages is None or use_acc_fusion is None: + _, _, b_w, b_s, b_acc, *_ = _get_arch_config(dot_precision, device=device) + if num_warps is None: + num_warps = b_w + if num_stages is None: + num_stages = b_s + if use_acc_fusion is None: + use_acc_fusion = b_acc + + if combined_history and direction != 0: + raise ValueError("combined_history=True requires direction=0 (bidi)") + + # Phase B kernel is DIRECTION-gated (constexpr); skipped-direction writes + # never happen, so we can hand it a 1-element placeholder for the inactive + # buffers and free ~4× M_fwd-shaped allocations per single-direction call. + decay_flat = decay.reshape(BH, F).contiguous().float() + + load_init = init_state_kv is not None + dummy = torch.empty(1, device=device, dtype=fdtype) + + def full_M(): + return torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype) + + def full_z(): + return torch.empty(BH, F, BLOCK_D, device=device, dtype=fdtype) + + M_fwd = dummy if direction == 2 else full_M() + z_fwd = dummy if (direction == 2 or skip_z) else full_z() + # Combined-history mode reuses M_fwd/z_fwd as M_hist/z_hist; rev outputs + # become placeholders even though DIRECTION!=1. + M_rev = dummy if (direction == 1 or combined_history) else full_M() + z_rev = dummy if (direction == 1 or combined_history or skip_z) else full_z() + if load_init: + init_kv = init_state_kv.contiguous().view(BH, BLOCK_D, BLOCK_D) + init_z = dummy if skip_z else init_state_z.contiguous().view(BH, BLOCK_D) + else: + init_kv = dummy + init_z = dummy + + if return_final_state: + final_kv = torch.empty(BH, BLOCK_D, BLOCK_D, device=device, dtype=fdtype) + final_z = dummy if skip_z else torch.empty(BH, BLOCK_D, device=device, dtype=fdtype) + else: + final_kv = dummy + final_z = dummy + + d_splits, nw_override, ns_override, acc_override = _pick_phase_b_d_splits(BLOCK_D, dot_precision=dot_precision) + if d_splits > 1: + D_TILE = BLOCK_D // d_splits + # Use D-tile-specific tuning if available, else fall back to baseline tuning + nw_use = nw_override if nw_override is not None else num_warps + ns_use = ns_override if ns_override is not None else num_stages + acc_use = acc_override if acc_override is not None else use_acc_fusion + _phase_b_dtile_kernel[(BH, d_splits)]( + I_P_kv, + A, + I_P_z, + B, + decay_flat, + M_fwd, + z_fwd, + M_rev, + z_rev, + init_kv, + init_z, + final_kv, + final_z, + BH=BH, + F=F, + BLOCK_D=BLOCK_D, + D_TILE=D_TILE, + DOT_PRECISION=dot_precision, + USE_ACC_FUSION=acc_use, + LOAD_INIT_STATE=1 if load_init else 0, + SAVE_FINAL_STATE=1 if return_final_state else 0, + DIRECTION=direction, + COMBINED_HISTORY=1 if combined_history else 0, + SKIP_Z=1 if skip_z else 0, + num_warps=nw_use, + num_stages=ns_use, + ) + else: + _phase_b_kernel[(BH,)]( + I_P_kv, + A, + I_P_z, + B, + decay_flat, + M_fwd, + z_fwd, + M_rev, + z_rev, + init_kv, + init_z, + final_kv, + final_z, + BH=BH, + F=F, + BLOCK_D=BLOCK_D, + DOT_PRECISION=dot_precision, + USE_ACC_FUSION=use_acc_fusion, + LOAD_INIT_STATE=1 if load_init else 0, + SAVE_FINAL_STATE=1 if return_final_state else 0, + DIRECTION=direction, + COMBINED_HISTORY=1 if combined_history else 0, + SKIP_Z=1 if skip_z else 0, + num_warps=num_warps, + num_stages=num_stages, + ) + if return_final_state: + return M_fwd, z_fwd, M_rev, z_rev, final_kv, final_z + return M_fwd, z_fwd, M_rev, z_rev + + +# ════════════════════════════════════════════════════════════════ +# Phase B D-tile — j-axis split for grid parallelism (#118) +# ════════════════════════════════════════════════════════════════ +# Same recurrence as _phase_b_kernel but each program owns a D_TILE-wide +# slice of M's output column dim. Grid: (BH, d_splits). M_new[*, j_tile] +# only depends on M_prev[*, j_tile] and full (I-P_kv) — independent across +# j-tiles. z is unsplittable; only `pid_d == 0` updates/writes z. +@triton.jit +def _phase_b_dtile_kernel( + I_P_kv_ptr, + A_ptr, + I_P_z_ptr, + B_ptr, + decay_ptr, + M_fwd_ptr, + z_fwd_ptr, + M_rev_ptr, + z_rev_ptr, + init_state_kv_ptr, + init_state_z_ptr, + final_state_kv_ptr, + final_state_z_ptr, + BH: tl.constexpr, + F: tl.constexpr, + BLOCK_D: tl.constexpr, + D_TILE: tl.constexpr, + DOT_PRECISION: tl.constexpr, + USE_ACC_FUSION: tl.constexpr, + LOAD_INIT_STATE: tl.constexpr, + SAVE_FINAL_STATE: tl.constexpr, + DIRECTION: tl.constexpr, + COMBINED_HISTORY: tl.constexpr, + SKIP_Z: tl.constexpr, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid_bh = tl.program_id(0) + pid_d = tl.program_id(1) + bh = pid_bh + + offs_d_full = tl.arange(0, BLOCK_D) + offs_d_tile = pid_d * D_TILE + tl.arange(0, D_TILE) + offs_dd_full = offs_d_full[:, None] * BLOCK_D + offs_d_full[None, :] + offs_dd_tile = offs_d_full[:, None] * BLOCK_D + offs_d_tile[None, :] + + is_lead = pid_d == 0 + + if DIRECTION != 2: + if LOAD_INIT_STATE: + M = tl.load(init_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd_tile).to(tl.float32) + else: + M = tl.zeros([BLOCK_D, D_TILE], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + if is_lead and LOAD_INIT_STATE: + z = tl.load(init_state_z_ptr + bh * BLOCK_D + offs_d_full).to(tl.float32) + + for f in range(F): + I_P_kv_f = tl.load(I_P_kv_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd_full) + A_f = tl.load(A_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd_tile) + g_f = tl.load(decay_ptr + bh * F + f).to(tl.float32) + + if USE_ACC_FUSION: + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot(I_P_kv_f.to(dot_dtype), M.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + M = g_f * M_temp + A_f + + tl.store(M_fwd_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd_tile, M) + + if is_lead and not SKIP_Z: + I_P_z_f = tl.load(I_P_z_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd_full) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d_full) + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + tl.store(z_fwd_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d_full, z) + + if SAVE_FINAL_STATE: + tl.store(final_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd_tile, M) + if is_lead and not SKIP_Z: + tl.store(final_state_z_ptr + bh * BLOCK_D + offs_d_full, z) + + if DIRECTION != 1: + M = tl.zeros([BLOCK_D, D_TILE], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + + if not COMBINED_HISTORY: + tl.store(M_rev_ptr + bh * F * BLOCK_D * BLOCK_D + (F - 1) * BLOCK_D * BLOCK_D + offs_dd_tile, M) + if is_lead and not SKIP_Z: + tl.store(z_rev_ptr + bh * F * BLOCK_D + (F - 1) * BLOCK_D + offs_d_full, z) + + for f_iter in range(F - 1): + f_src = F - 1 - f_iter + f_dst = f_src - 1 + I_P_kv_f = tl.load(I_P_kv_ptr + bh * F * BLOCK_D * BLOCK_D + f_src * BLOCK_D * BLOCK_D + offs_dd_full) + A_f = tl.load(A_ptr + bh * F * BLOCK_D * BLOCK_D + f_src * BLOCK_D * BLOCK_D + offs_dd_tile) + g_f = tl.load(decay_ptr + bh * F + f_src).to(tl.float32) + + if USE_ACC_FUSION: + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot(I_P_kv_f.to(dot_dtype), M.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + M = g_f * M_temp + A_f + + if is_lead and not SKIP_Z: + I_P_z_f = tl.load(I_P_z_ptr + bh * F * BLOCK_D * BLOCK_D + f_src * BLOCK_D * BLOCK_D + offs_dd_full) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f_src * BLOCK_D + offs_d_full) + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + + if COMBINED_HISTORY: + M_addr = M_fwd_ptr + bh * F * BLOCK_D * BLOCK_D + f_dst * BLOCK_D * BLOCK_D + offs_dd_tile + tl.store(M_addr, tl.load(M_addr) + M) + if is_lead and not SKIP_Z: + z_addr = z_fwd_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d_full + tl.store(z_addr, tl.load(z_addr) + z) + else: + tl.store(M_rev_ptr + bh * F * BLOCK_D * BLOCK_D + f_dst * BLOCK_D * BLOCK_D + offs_dd_tile, M) + if is_lead and not SKIP_Z: + tl.store(z_rev_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d_full, z) + + +_PHASE_B_DTILE_ARCH_CACHE: dict = {} # (dev, dot_prec) -> (d_splits, nw, ns, acc) + + +# Per-arch D-tile optimum from 2026-04-29 sweep (T=11 B=1 P0 IEEE): +# WGMMA-server (A100 sm_80, H100 sm_90): (d=4, nw=32, ns=1, acc=True) +# Blackwell-family (GB200 sm_100, 5090 sm_120, GB10 sm_121, Ada sm_89): +# (d=8, nw=4, ns=1, acc=False) +# Both clusters were tested across 96 configs (4 ds × 4 nw × 3 ns × 2 acc). +def _pick_phase_b_d_splits(BLOCK_D: int, dot_precision: int = 0): + """Returns (d_splits, nw_override, ns_override, acc_override). + + `d_splits=1` → use baseline `_phase_b_kernel` with `_CHUNKWISE_TUNING` config. `d_splits>1` → use + `_phase_b_dtile_kernel` with overrides for nw/ns/acc. Override via env: PHASE_B_D_SPLITS, PHASE_B_DTILE_NW, + PHASE_B_DTILE_NS, PHASE_B_DTILE_ACC (1=True / 0=False). + """ + import os + + env_d = os.environ.get("PHASE_B_D_SPLITS", None) + if env_d is not None: + d = int(env_d) + if d < 1 or BLOCK_D % d != 0: + return (1, None, None, None) + nw = int(os.environ.get("PHASE_B_DTILE_NW", "0")) or None + ns = int(os.environ.get("PHASE_B_DTILE_NS", "0")) or None + acc_env = os.environ.get("PHASE_B_DTILE_ACC", None) + acc = bool(int(acc_env)) if acc_env is not None else None + return (d, nw, ns, acc) + try: + import torch + + if not torch.cuda.is_available(): + return (1, None, None, None) + dev = torch.cuda.current_device() + cache_key = (dev, dot_precision) + if cache_key not in _PHASE_B_DTILE_ARCH_CACHE: + cap = torch.cuda.get_device_capability(dev) + major, minor = cap[0], cap[1] + if dot_precision == 2: + # IEEE fp32: D-tile dominates baseline on every arch (96-config sweep). + if major == 8 and minor == 0: + cfg = (4, 32, 1, True) # A100 + elif major == 9: + cfg = (4, 32, 1, True) # H100 (Hopper) + elif major == 8 and minor == 9: + cfg = (8, 4, 1, False) # Ada (assume Blackwell-like) + elif major >= 10: + cfg = (8, 4, 1, False) # GB200/B200, 5090, GB10 + else: + cfg = (1, None, None, None) # unknown — baseline + else: + # bf16/TF32: cap-specific dispatch. Multi-arch sweep 2026-05-06 + # (F=11 S=920) determined per-cap whether D-tile beats the + # baseline _phase_b_kernel: + # sm_80 A100: D-tile WIN 1.09× (P1) / 1.02× (P2) — (4,8,2,F). + # sm_90 H100: D-tile WIN ~10% — P1 (4,8,2,F); P2 (8,8,2,F). + # Use (4,8,2,F) for both (P2 within 0.4%). + # sm_100 GB200: D-tile WIN ~12% — (4,8,2,F) both precisions. + # sm_120 5090: D-tile WIN 2.6× (P1) / 1.13× (P2) — (8,8,1,F). + # TF32 baseline OOMs at 102 KB SRAM cap. + # sm_121 GB10: D-tile LOSS 4% — baseline wins. Despite same + # reported SRAM/SM as sm_120, the baseline + # kernel fits all configs up to nw=16 ns=2 on + # sm_121 (Triton/codegen difference between + # consumer-Blackwell variants), so baseline + # saturates the chip without needing D-tile. + if major == 8 and minor == 0: + cfg = (4, 8, 2, False) # A100 + elif major == 9: + cfg = (4, 8, 2, False) # H100 + elif major == 10: + cfg = (4, 8, 2, False) # GB200 / B200 + elif major == 12 and minor == 0: + cfg = (8, 8, 1, False) # 5090 + elif major == 12 and minor == 1: + cfg = (1, None, None, None) # GB10 — baseline wins + else: + cfg = (1, None, None, None) # Ada, unknown + _PHASE_B_DTILE_ARCH_CACHE[cache_key] = cfg + return _PHASE_B_DTILE_ARCH_CACHE[cache_key] + except Exception: + return (1, None, None, None) + + +# ════════════════════════════════════════════════════════════════ +# Phase C — Pass 2 output (per (B, H, F)). Same as v1. +# ════════════════════════════════════════════════════════════════ + + +@triton.jit +def _phase_c_kernel( + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + q_inv_rms_ptr, + q_norm_w_ptr, + rope_cos_ptr, + rope_sin_ptr, + M_ptr, + z_ptr, + num_ptr, + den_ptr, + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + NORM_EPS: tl.constexpr, + DOT_PRECISION: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, + ACCUMULATE: tl.constexpr = False, + SKIP_LAST_F: tl.constexpr = False, + SKIP_RELU: tl.constexpr = False, + NUM_ONLY: tl.constexpr = False, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + pid_b = pid // (H * F) + pid_hf = pid % (H * F) + pid_h = pid_hf // F + pid_f = pid_hf % F + bh = pid_b * H + pid_h + N: tl.constexpr = F * S + + # Reverse-accumulate callers pass SKIP_LAST_F=True: M_rev[F-1] / z_rev[F-1] + # are exactly zero (Phase B initializes the reverse scan with zeros and the + # write loop only fills f 0, Q_normed, 0.0) + Q_pair = tl.where(Q_pair_normed > 0, Q_pair_normed, 0.0) + + rope_ptrs = n_idx[:, None] * D + offs_d[None, :] + Cos = tl.load(rope_cos_ptr + rope_ptrs, mask=mask_sd, other=1.0).to(tl.float32) + Sin = tl.load(rope_sin_ptr + rope_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + Q_rot = Q * Cos + Q_pair * Sin + + num = tl.dot(Q_rot.to(dot_dtype), M_f.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + if not NUM_ONLY: + den = tl.sum(Q * z_f[None, :], axis=1) + + num_ptrs = num_bh + n_idx[:, None] * (H * D) + offs_d[None, :] + if not NUM_ONLY: + den_ptrs = den_bh + n_idx + if ACCUMULATE: + # Used by reverse-direction Phase C: add this pass onto forward's + # already-written buffer instead of allocating a separate one. + prev_num = tl.load(num_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + num = num + prev_num + if not NUM_ONLY: + prev_den = tl.load(den_ptrs, mask=mask_s, other=0.0).to(tl.float32) + den = den + prev_den + if DOT_PRECISION >= 1: + tl.store(num_ptrs, num, mask=mask_sd) + if not NUM_ONLY: + tl.store(den_ptrs, den, mask=mask_s) + else: + tl.store(num_ptrs, num.to(tl.bfloat16), mask=mask_sd) + if not NUM_ONLY: + tl.store(den_ptrs, den.to(tl.bfloat16), mask=mask_s) + + +def phase_c( + qkv, + q_inv_rms, + q_norm_w, + rope_cos, + rope_sin, + M, + z, + F, + S, + num_warps=None, + num_stages=None, + BLOCK_S=None, + dot_precision=0, + num_out=None, + den_out=None, + accumulate=False, + skip_last_frame=False, + skip_relu: bool = False, + num_only: bool = False, +): + """Phase C Pass-2 output. Optionally accumulates into caller-provided + ``num_out``/``den_out`` buffers (used to fuse reverse-direction output into forward-direction buffer without + allocating a separate one — saves ~45 MB at B=1 bf16, ~180 MB at B=4). + + ``skip_last_frame=True`` early-returns the f=F-1 programs. Valid for the reverse-accumulate call only, where + M[F-1]/z[F-1] are guaranteed zero. + + ``skip_relu=True`` matches Phase A KV's flag — used by the camera-branch chunkwise wrapper where Q has already been + ReLU'd by cam_prep before being rotated by UCPE+RoPE; re-applying ReLU on the rotated Q would clobber legitimate + negatives. + + ``num_only=True`` skips the denominator computation and store entirely (kernel writes only ``num_out``; ``den_out`` + is allowed to be None / unallocated). Used by the camera-branch which has no Z scan. + """ + if num_warps is None or num_stages is None or BLOCK_S is None: + *_, c_w, c_bs, c_s = _get_arch_config(dot_precision, device=qkv.device) + if num_warps is None: + num_warps = c_w + if num_stages is None: + num_stages = c_s + if BLOCK_S is None: + BLOCK_S = c_bs + B, N, three, H, D = qkv.shape + BLOCK_D = triton.next_power_of_2(D) + if num_out is None: + num_out = torch.empty( + B, N, H, D, device=qkv.device, dtype=(torch.float32 if dot_precision >= 1 else torch.bfloat16) + ) + if den_out is None and not num_only: + den_out = torch.empty( + B, H, N, device=qkv.device, dtype=(torch.float32 if dot_precision >= 1 else torch.bfloat16) + ) + elif num_only and den_out is None: + # Pass a 1-element placeholder; kernel guards den loads/stores under NUM_ONLY. + den_out = torch.empty(1, device=qkv.device, dtype=(torch.float32 if dot_precision >= 1 else torch.bfloat16)) + + _phase_c_kernel[(B * H * F,)]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + q_inv_rms, + q_norm_w, + rope_cos, + rope_sin, + M, + z, + num_out, + den_out, + H=H, + F=F, + S=S, + D=D, + NORM_EPS=1e-5, + DOT_PRECISION=dot_precision, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + ACCUMULATE=1 if accumulate else 0, + SKIP_LAST_F=skip_last_frame, + SKIP_RELU=skip_relu, + NUM_ONLY=num_only, + num_warps=num_warps, + num_stages=num_stages, + ) + return num_out, den_out + + +def fused_bigdn_bidi_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_w, + k_norm_w, + rope_cos, + rope_sin, + beta, + decay, + F, + S, + k_scale=1.0, + eps=1e-6, + norm_eps=1e-5, + dot_precision=0, + init_state_kv=None, + init_state_z=None, + return_final_state=False, +): + """Bidi chunkwise GDN forward, optionally with state-cache for autoregressive + sampling (chunk 0 = full bidi with state save; chunks > 0 seed forward scan from saved state). Reverse always seeds + from zero per upstream convention. + + Pipeline (2026-04-25 restructure): Phase A once → Phase B direction=0 with combined_history=True (fwd seeded with + init_state and saves final state; rev zero-seeded; rev output summed into fwd buffer in-kernel via read- add-store + so on exit M_hist[f] = M_fwd[f] + M_rev[f]) → Phase C ONCE on M_hist. Phase C linearity `Q @ (M_fwd + M_rev) = Q @ + M_fwd + Q @ M_rev` makes the in-kernel sum exact. + + Replaces the prior 2× Phase B + 2× Phase C pattern. Saves one Phase C launch + one Q+RoPE HBM pass and one M-shape + buffer per call. + """ + I_P_kv, A, I_P_z, B_z = phase_a( + qkv, + beta, + q_inv_rms, + k_inv_rms, + q_norm_w, + k_norm_w, + rope_cos, + rope_sin, + F=F, + S=S, + k_scale=k_scale, + norm_eps=norm_eps, + dot_precision=dot_precision, + ) + + if return_final_state: + M_hist, z_hist, _, _, final_kv, final_z = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=0, + init_state_kv=init_state_kv, + init_state_z=init_state_z, + return_final_state=True, + combined_history=True, + ) + else: + M_hist, z_hist, _, _ = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=0, + init_state_kv=init_state_kv, + init_state_z=init_state_z, + combined_history=True, + ) + num_out, den_out = phase_c( + qkv, + q_inv_rms, + q_norm_w, + rope_cos, + rope_sin, + M_hist, + z_hist, + F=F, + S=S, + dot_precision=dot_precision, + accumulate=False, + ) + del M_hist, z_hist, I_P_kv, A, I_P_z, B_z + + # ── Final divide ── + total_den = den_out.float().permute(0, 2, 1).unsqueeze(-1) # (B, N, H, 1) + out = (num_out.float() / (total_den + eps)).to(qkv.dtype) + del num_out, den_out, total_den + if return_final_state: + B = qkv.shape[0] + H = qkv.shape[3] + D = qkv.shape[4] + BLOCK_D = final_kv.shape[1] + state_kv = final_kv.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D].transpose(-1, -2).contiguous() + state_z = final_z.view(B, H, BLOCK_D)[:, :, :D].unsqueeze(-1).contiguous() + return out, state_kv, state_z + return out + + +def _default_dot_prec(): + """Pull dot_precision from `_resolve_launch_config` (honors PRECISION_OVERRIDE).""" + + _, dot_prec, _, _ = _resolve_launch_config() + return dot_prec + + +def fused_gdn_func_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F, + S, + k_scale, + eps=1e-6, + reverse=False, + dot_precision=None, +): + """Single-direction chunkwise GDN — drop-in for `fused_gdn.fused_gdn_func`. + + Computes only one scan direction (Phase B + Phase C × 1) and returns `(num, den)` shape-compatible with the + upstream function. dot_precision defaults to whatever `_resolve_launch_config` returns (honors module-level + `PRECISION_OVERRIDE`). + """ + if dot_precision is None: + dot_precision = _default_dot_prec() + direction = 2 if reverse else 1 + I_P_kv, A, I_P_z, B_z = phase_a( + qkv, + beta, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F=F, + S=S, + k_scale=k_scale, + dot_precision=dot_precision, + ) + M_fwd, z_fwd, M_rev, z_rev = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=direction, + ) + M_use = M_rev if reverse else M_fwd + z_use = z_rev if reverse else z_fwd + num, den = phase_c( + qkv, q_inv_rms, q_norm_weight, rope_cos, rope_sin, M_use, z_use, F=F, S=S, dot_precision=dot_precision + ) + return num, den + + +def fused_gdn_stateful_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F, + S, + k_scale, + eps=1e-6, + reverse=False, + init_state_kv=None, + init_state_z=None, + return_final_state=False, + dot_precision=None, +): + """Single-direction chunkwise GDN with optional state cache — drop-in for + `fused_gdn.fused_gdn_stateful`. Forward direction supports state load/save (used for autoregressive sampling); + reverse direction always runs fresh (per upstream's bidi state-cache convention). + """ + if dot_precision is None: + dot_precision = _default_dot_prec() + direction = 2 if reverse else 1 + if reverse and (init_state_kv is not None or return_final_state): + raise ValueError( + "fused_gdn_stateful_chunkwise: state cache is forward-only (matching " + "upstream's bidi convention); pass reverse=False or omit state args." + ) + I_P_kv, A, I_P_z, B_z = phase_a( + qkv, + beta, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F=F, + S=S, + k_scale=k_scale, + dot_precision=dot_precision, + ) + # Pad caller-supplied state from (B,H,D,D)/(B,H,D,1) to (BH, BLOCK_D, BLOCK_D)/(BH, BLOCK_D). + # Needed because the state returned by this function is unpadded (B,H,D,D), + # but phase_b_triton's kernel expects the padded layout. + init_kv_padded, init_z_padded = init_state_kv, init_state_z + if init_state_kv is not None: + B_, H_, D_in, D_out = init_state_kv.shape + BLOCK_D_ = I_P_kv.shape[-1] + if D_in != BLOCK_D_ or D_out != BLOCK_D_: + pad_in = BLOCK_D_ - D_in + pad_out = BLOCK_D_ - D_out + init_kv_padded = torch.nn.functional.pad( + init_state_kv.transpose(-1, -2).reshape(B_ * H_, D_out, D_in), (0, pad_in, 0, pad_out) + ).contiguous() + else: + init_kv_padded = init_state_kv.transpose(-1, -2).reshape(B_ * H_, BLOCK_D_, BLOCK_D_).contiguous() + # z: (B, H, D) or (B, H, D, 1) → (BH, BLOCK_D) + z_ = init_state_z.squeeze(-1) if init_state_z.dim() == 4 else init_state_z + Bz_, Hz_, Dz_ = z_.shape + if Dz_ != BLOCK_D_: + init_z_padded = torch.nn.functional.pad(z_.reshape(Bz_ * Hz_, Dz_), (0, BLOCK_D_ - Dz_)).contiguous() + else: + init_z_padded = z_.reshape(Bz_ * Hz_, Dz_).contiguous() + if return_final_state: + M_fwd, z_fwd, M_rev, z_rev, final_kv, final_z = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=direction, + init_state_kv=init_kv_padded, + init_state_z=init_z_padded, + return_final_state=True, + ) + else: + M_fwd, z_fwd, M_rev, z_rev = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=direction, + init_state_kv=init_kv_padded, + init_state_z=init_z_padded, + ) + M_use = M_rev if reverse else M_fwd + z_use = z_rev if reverse else z_fwd + num, den = phase_c( + qkv, q_inv_rms, q_norm_weight, rope_cos, rope_sin, M_use, z_use, F=F, S=S, dot_precision=dot_precision + ) + if return_final_state: + B = qkv.shape[0] + H = qkv.shape[3] + D = qkv.shape[4] + BLOCK_D = final_kv.shape[1] + state_kv = final_kv.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D].transpose(-1, -2).contiguous() + state_z = final_z.view(B, H, BLOCK_D)[:, :, :D].unsqueeze(-1).contiguous() + return num, den, state_kv, state_z + return num, den + + +def fused_bidi_stateful_chunkwise_shared_phase_a( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F, + S, + k_scale, + eps=1e-6, + init_state_kv=None, + init_state_z=None, + dot_precision=None, +): + """Bidi state-cached chunkwise GDN with shared Phase A and combined-history + Phase B. Default chunkwise path for ``_fused_statecached_forward``. + + Pipeline (per layer per step): + 1. Phase A once over qkv — K/V/RoPE pre-norm; was previously duplicated across two streams. + 2. Phase B with direction=0 + combined_history=True — single program does fwd then rev; fwd writes M_hist; rev + read-add-stores into the same buffer so on exit M_hist[f] = M_fwd[f] + M_rev[f] (same for z). Forward branch + loads init_state and saves final state. + 3. Phase C ONCE on M_hist/z_hist — Phase C is linear in M/z so `Q @ (M_fwd + M_rev) = Q @ M_fwd + Q @ M_rev`. + + Returns ``(num_combined, den_combined, state_kv, state_z)`` — caller hands the num/den pair to + ``fused_bidi_merge(num, None, den, None, eps, gate)`` in PRE_SUMMED mode. + + HBM-traffic delta vs the prior 2× Phase C version (per call, B=1 prod): + saved : 1× Phase C Q+RoPE pass (~90 MB) saved : one (B,N,H,D) num and (B,H,N) den allocation cost : Phase B rev + does read-add of M_hist (~14 MB extra per layer) net : ~76 MB saved + 1 fewer kernel launch + + Measured speed on GB10 (sm_121) at H=20, S=920, D=112, vs the prior shared-Phase-A-with-2×-Phase-C path, across + production F values: + P0 IEEE fp32 : 1.26-1.42× (F=3,6,11; B=1,2) P2 bf16+fp32-st : 1.57-1.80× P3 bf16+bf16-st : 1.63-1.96× + Correctness cos ≥ 0.999997 across all cells, state_kv exact. + """ + if dot_precision is None: + dot_precision = _default_dot_prec() + + I_P_kv, A, I_P_z, B_z = phase_a( + qkv, + beta, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F=F, + S=S, + k_scale=k_scale, + dot_precision=dot_precision, + ) + + init_kv_padded, init_z_padded = init_state_kv, init_state_z + if init_state_kv is not None: + B_, H_, D_in, D_out = init_state_kv.shape + BLOCK_D_ = I_P_kv.shape[-1] + if D_in != BLOCK_D_ or D_out != BLOCK_D_: + pad_in = BLOCK_D_ - D_in + pad_out = BLOCK_D_ - D_out + init_kv_padded = torch.nn.functional.pad( + init_state_kv.transpose(-1, -2).reshape(B_ * H_, D_out, D_in), (0, pad_in, 0, pad_out) + ).contiguous() + else: + init_kv_padded = init_state_kv.transpose(-1, -2).reshape(B_ * H_, BLOCK_D_, BLOCK_D_).contiguous() + z_ = init_state_z.squeeze(-1) if init_state_z.dim() == 4 else init_state_z + Bz_, Hz_, Dz_ = z_.shape + if Dz_ != BLOCK_D_: + init_z_padded = torch.nn.functional.pad(z_.reshape(Bz_ * Hz_, Dz_), (0, BLOCK_D_ - Dz_)).contiguous() + else: + init_z_padded = z_.reshape(Bz_ * Hz_, Dz_).contiguous() + + # combined_history=True routes the rev contribution into the fwd buffer → + # M_hist[f] = M_fwd[f] + M_rev[f]. M_rev/z_rev outputs are placeholders. + M_hist, z_hist, _, _, final_kv, final_z = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=0, + init_state_kv=init_kv_padded, + init_state_z=init_z_padded, + return_final_state=True, + combined_history=True, + ) + + num, den = phase_c( + qkv, q_inv_rms, q_norm_weight, rope_cos, rope_sin, M_hist, z_hist, F=F, S=S, dot_precision=dot_precision + ) + + B = qkv.shape[0] + H = qkv.shape[3] + D = qkv.shape[4] + BLOCK_D = final_kv.shape[1] + state_kv = final_kv.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D].transpose(-1, -2).contiguous() + state_z = final_z.view(B, H, BLOCK_D)[:, :, :D].unsqueeze(-1).contiguous() + return num, den, state_kv, state_z + + +def fused_bigdn_stateful_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F, + S, + k_scale, + eps=1e-6, + return_final_state=False, + dot_precision=None, +): + """Drop-in replacement for `fused_gdn.fused_bigdn_stateful` using the + chunkwise pipeline. Same signature, same return shape: + output (B, N, H, D), and if return_final_state: + (state_kv, state_z). + dot_precision defaults to whatever `_resolve_launch_config` returns. + """ + if dot_precision is None: + dot_precision = _default_dot_prec() + if return_final_state: + out, state_kv, state_z = fused_bigdn_bidi_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F=F, + S=S, + k_scale=k_scale, + eps=eps, + dot_precision=dot_precision, + return_final_state=True, + ) + return out, state_kv, state_z + out = fused_bigdn_bidi_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F=F, + S=S, + k_scale=k_scale, + eps=eps, + dot_precision=dot_precision, + ) + return out + + +# ───────────────────────────────────────────────────────────────────────────── +# Camera-branch wrapper — numerator-only single-path delta-rule scan via +# chunkwise. Drop-in for `diffusion.model.ops.fused_cam_gdn.cam_scan_func`. +# +# Cam math expanded: +# state = state * g # apply decay +# state += K^T @ ((V - K @ state) * β) # delta-rule +# Equivalently: +# state_new = g (I - K^T β K) state_old + K^T β V +# = g (I - P_kv) state_old + A +# This is bit-identical to chunkwise's Phase B M update, so the scan kernel +# is reusable. The only differences from main GDN: +# 1. Q/K/V come pre-prepped (cam_prep_kernel did RMSNorm+ReLU+UCPE+RoPE). +# We disable chunkwise's prep with identity tables (k_inv_rms=1, k_nw=1, +# k_scale=1, rope_cos=1, rope_sin=0) AND skip_relu=True (because cam +# applied ReLU BEFORE UCPE; the post-UCPE values can have legitimate +# negatives that re-applying ReLU would clobber). +# 2. No Z denominator scan; output is num-only (out = Q @ M, no /Z). +# skip_z=True elides Phase A Z; num_only=True elides Phase C den compute. +# ───────────────────────────────────────────────────────────────────────────── +def _cam_identity_tables( + *, + B: int, + N: int, + H: int, + D: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Cached identity RMS/RoPE tables used by ``cam_scan_chunkwise``.""" + device_index = device.index if device.type == "cuda" else None + key = (device.type, device_index, B, N, H * D, D) + cached = _CAM_IDENTITY_CACHE.get(key) + if cached is not None: + return cached + + ones_inv_rms = torch.ones(B, N, device=device, dtype=torch.float32) + ones_nw = torch.ones(H * D, device=device, dtype=torch.float32) + ones_cos = torch.ones(N, D, device=device, dtype=torch.float32) + zeros_sin = torch.zeros(N, D, device=device, dtype=torch.float32) + cached = (ones_inv_rms, ones_nw, ones_cos, zeros_sin) + _CAM_IDENTITY_CACHE[key] = cached + return cached + + +def cam_scan_chunkwise( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + reverse: bool = False, + init_state: torch.Tensor | None = None, + save_final_state: bool = False, + dot_precision: int | None = None, +): + """Drop-in chunkwise replacement for `cam_scan_func`. + + Args mirror `cam_scan_func` exactly: + q, k, v: ``(B, H, D, N)`` fp32 contiguous (cam-prep'd: RMSNorm+ReLU+UCPE+RoPE) beta: ``(B, H, F, S)`` fp32 + contiguous decay: ``(B, H, F)`` fp32 contiguous reverse: bwd flip-and-shift semantics (autograd path); not yet + supported. init_state: optional ``(B*H, BLOCK_D, BLOCK_D)`` fp32 — cross-chunk AR state. save_final_state: when + True, also returns ``(out, final_state)``. + + Returns ``out`` of shape ``(B, H, D, N)`` fp32, or ``(out, final_state: (B*H, BLOCK_D, BLOCK_D))`` if + save_final_state=True. + """ + assert q.shape == k.shape == v.shape, f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert beta.is_contiguous() and decay.is_contiguous() + assert q.dtype == torch.float32, f"cam_scan_chunkwise requires fp32 q/k/v (got {q.dtype})" + + if reverse and (init_state is not None or save_final_state): + raise NotImplementedError( + "cam_scan_chunkwise: state passing (init_state / save_final_state) is " + "only supported for the forward direction (reverse=False). The cam " + "branch's anti-causal pass resets per chunk; there is no global " + "cross-prefix state to cache for the reverse direction." + ) + + B, H, D, N = q.shape + F = beta.shape[2] + assert N % F == 0 + S = N // F + assert beta.shape == (B, H, F, S) + assert decay.shape == (B, H, F) + + BLOCK_D = triton.next_power_of_2(D) + + if dot_precision is None: + dot_precision = _default_dot_prec() + + # Repack (B, H, D, N) → (B, N, 3, H, D) for chunkwise's qkv layout. + # Avoid ``stack(...).permute(...).contiguous()`` because that materializes + # two large tensors. Direct packing allocates the destination once. + qkv = torch.empty(B, N, 3, H, D, device=q.device, dtype=q.dtype) + qkv[:, :, 0].copy_(q.permute(0, 3, 1, 2)) + qkv[:, :, 1].copy_(k.permute(0, 3, 1, 2)) + qkv[:, :, 2].copy_(v.permute(0, 3, 1, 2)) + + # Identity prep tables — make chunkwise's RMSNorm + RoPE no-ops. + ones_inv_rms, ones_nw, ones_cos, zeros_sin = _cam_identity_tables(B=B, N=N, H=H, D=D, device=q.device) + + # Phase A (skip_relu=True for cam-prep'd K; skip_z=True since cam has no Z scan). + # k_scale=1.0 because cam_prep already applied K-scale. + I_P_kv, A_, I_P_z, B_z = phase_a( + qkv, + beta, + ones_inv_rms, + ones_inv_rms, + ones_nw, + ones_nw, + ones_cos, + zeros_sin, + F=F, + S=S, + k_scale=1.0, + norm_eps=1e-5, + dot_precision=dot_precision, + skip_relu=True, + skip_z=True, + ) + + # Phase B (forward direction only; cam supports init_state on fwd, save_final + # on fwd; no rev). Pads (B*H, D, D) ↔ (B*H, BLOCK_D, BLOCK_D) inline. + init_kv_padded = None + init_z_padded = None + if init_state is not None: + if init_state.shape != (B * H, BLOCK_D, BLOCK_D): + raise ValueError( + f"cam_scan_chunkwise: init_state shape {tuple(init_state.shape)} " + f"!= expected (B*H, BLOCK_D, BLOCK_D) = {(B * H, BLOCK_D, BLOCK_D)}" + ) + if init_state.dtype != torch.float32: + raise ValueError(f"cam_scan_chunkwise: init_state must be fp32 (got {init_state.dtype}).") + if not init_state.is_contiguous(): + raise ValueError("cam_scan_chunkwise: init_state must be contiguous.") + # Cam stores state as M[K_feat, V_feat]. Chunkwise's Phase B kernel reads + # state with offs_dd = i*BLOCK_D + j where i is the fwd loop's M row. + # Storage layout matches cam's (row-major (D_K, D_V)), so a direct cast + # to fp32 contiguous is enough — no transpose needed. + init_kv_padded = init_state.to(torch.float32).contiguous() + # No Z state in cam — pass zeros to satisfy phase_b_triton. + init_z_padded = torch.zeros(B * H, BLOCK_D, device=q.device, dtype=torch.float32) + + direction = 2 if reverse else 1 + if save_final_state: + M_fwd, z_fwd_out, M_rev, z_rev_out, final_kv, _final_z = phase_b_triton( + I_P_kv, + A_, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=direction, + init_state_kv=init_kv_padded, + init_state_z=init_z_padded, + return_final_state=True, + skip_z=True, + ) + else: + M_fwd, z_fwd_out, M_rev, z_rev_out = phase_b_triton( + I_P_kv, + A_, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=direction, + init_state_kv=init_kv_padded, + init_state_z=init_z_padded, + skip_z=True, + ) + + # For reverse (flip-and-shift bwd), Phase B's reverse mode produces M_rev + # such that M_rev[F-1] = 0 and M_rev[t] = state computed from K/V at frames + # {F-1, F-2, ..., t+1} — exactly cam's REVERSE=1 semantics. + M_use = M_rev if reverse else M_fwd + z_use = z_rev_out if reverse else z_fwd_out + + # Phase C — num-only (NUM_ONLY=True skips den compute + store). + # z is unused with NUM_ONLY but still required by the kernel signature. + num_out, _ = phase_c( + qkv, + ones_inv_rms, + ones_nw, + ones_cos, + zeros_sin, + M_use, + z_use, + F=F, + S=S, + dot_precision=dot_precision, + skip_relu=True, + num_only=True, + ) + + # Convert chunkwise output (B, N, H, D) → cam's (B, H, D, N) layout, fp32. + out = num_out.permute(0, 2, 3, 1).contiguous().to(torch.float32) + + if save_final_state: + return out, final_kv # final_kv already (B*H, BLOCK_D, BLOCK_D) fp32 + return out + + +def cam_scan_bidi_chunkwise( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + dot_precision: int | None = None, +) -> torch.Tensor: + """Bidirectional camera scan using shared chunkwise phases. + + This is equivalent to ``cam_scan_chunkwise(..., reverse=False) + cam_scan_chunkwise(..., reverse=True)`` for full + bidirectional attention, but it packs QKV once, runs Phase A once, combines forward/reverse histories inside Phase + B, and runs Phase C once on the summed state. + """ + _require_triton("cam_scan_bidi_chunkwise") + assert q.shape == k.shape == v.shape, f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert beta.is_contiguous() and decay.is_contiguous() + assert q.dtype == torch.float32, f"cam_scan_bidi_chunkwise requires fp32 q/k/v (got {q.dtype})" + + B, H, D, N = q.shape + F = beta.shape[2] + assert N % F == 0 + S = N // F + assert beta.shape == (B, H, F, S) + assert decay.shape == (B, H, F) + + if dot_precision is None: + dot_precision = _default_dot_prec() + + qkv = torch.empty(B, N, 3, H, D, device=q.device, dtype=q.dtype) + qkv[:, :, 0].copy_(q.permute(0, 3, 1, 2)) + qkv[:, :, 1].copy_(k.permute(0, 3, 1, 2)) + qkv[:, :, 2].copy_(v.permute(0, 3, 1, 2)) + + ones_inv_rms, ones_nw, ones_cos, zeros_sin = _cam_identity_tables(B=B, N=N, H=H, D=D, device=q.device) + I_P_kv, A_, I_P_z, B_z = phase_a( + qkv, + beta, + ones_inv_rms, + ones_inv_rms, + ones_nw, + ones_nw, + ones_cos, + zeros_sin, + F=F, + S=S, + k_scale=1.0, + norm_eps=1e-5, + dot_precision=dot_precision, + skip_relu=True, + skip_z=True, + ) + M_hist, z_hist, _, _ = phase_b_triton( + I_P_kv, + A_, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=0, + combined_history=True, + skip_z=True, + ) + num_out, _ = phase_c( + qkv, + ones_inv_rms, + ones_nw, + ones_cos, + zeros_sin, + M_hist, + z_hist, + F=F, + S=S, + dot_precision=dot_precision, + skip_relu=True, + num_only=True, + ) + return num_out.permute(0, 2, 3, 1).contiguous().to(torch.float32) + + +def cam_scan_pair_chunkwise( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta_fwd: torch.Tensor, + decay_fwd: torch.Tensor, + beta_rev: torch.Tensor, + decay_rev: torch.Tensor, + *, + dot_precision: int | None = None, +) -> torch.Tensor: + """Sum a forward camera scan and a separately-gated reverse scan. + + Chunk-causal camera attention needs the reverse branch to use boundary-masked gates while the forward branch uses + the original gates. This wrapper keeps that exact behavior but shares QKV packing, identity tables, and the final + output layout conversion across the two scans. + """ + assert q.shape == k.shape == v.shape, f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert beta_fwd.is_contiguous() and decay_fwd.is_contiguous() + assert beta_rev.is_contiguous() and decay_rev.is_contiguous() + assert q.dtype == torch.float32, f"cam_scan_pair_chunkwise requires fp32 q/k/v (got {q.dtype})" + + B, H, D, N = q.shape + F = beta_fwd.shape[2] + assert N % F == 0 + S = N // F + assert beta_fwd.shape == beta_rev.shape == (B, H, F, S) + assert decay_fwd.shape == decay_rev.shape == (B, H, F) + + if dot_precision is None: + dot_precision = _default_dot_prec() + + qkv = torch.empty(B, N, 3, H, D, device=q.device, dtype=q.dtype) + qkv[:, :, 0].copy_(q.permute(0, 3, 1, 2)) + qkv[:, :, 1].copy_(k.permute(0, 3, 1, 2)) + qkv[:, :, 2].copy_(v.permute(0, 3, 1, 2)) + + ones_inv_rms, ones_nw, ones_cos, zeros_sin = _cam_identity_tables(B=B, N=N, H=H, D=D, device=q.device) + + I_P_kv, A_, I_P_z, B_z = phase_a( + qkv, + beta_fwd, + ones_inv_rms, + ones_inv_rms, + ones_nw, + ones_nw, + ones_cos, + zeros_sin, + F=F, + S=S, + k_scale=1.0, + norm_eps=1e-5, + dot_precision=dot_precision, + skip_relu=True, + skip_z=True, + ) + M_fwd, z_fwd, _, _ = phase_b_triton( + I_P_kv, + A_, + I_P_z, + B_z, + decay_fwd, + F=F, + dot_precision=dot_precision, + direction=1, + skip_z=True, + ) + num_out, _ = phase_c( + qkv, + ones_inv_rms, + ones_nw, + ones_cos, + zeros_sin, + M_fwd, + z_fwd, + F=F, + S=S, + dot_precision=dot_precision, + skip_relu=True, + num_only=True, + ) + del I_P_kv, A_, I_P_z, B_z, M_fwd, z_fwd + + I_P_kv, A_, I_P_z, B_z = phase_a( + qkv, + beta_rev, + ones_inv_rms, + ones_inv_rms, + ones_nw, + ones_nw, + ones_cos, + zeros_sin, + F=F, + S=S, + k_scale=1.0, + norm_eps=1e-5, + dot_precision=dot_precision, + skip_relu=True, + skip_z=True, + ) + _, _, M_rev, z_rev = phase_b_triton( + I_P_kv, + A_, + I_P_z, + B_z, + decay_rev, + F=F, + dot_precision=dot_precision, + direction=2, + skip_z=True, + ) + phase_c( + qkv, + ones_inv_rms, + ones_nw, + ones_cos, + zeros_sin, + M_rev, + z_rev, + F=F, + S=S, + dot_precision=dot_precision, + num_out=num_out, + accumulate=True, + skip_relu=True, + num_only=True, + ) + return num_out.permute(0, 2, 3, 1).contiguous().to(torch.float32) + + +# ===== camera utility helpers (used by both kernels and the transformer) ===== + + +def compute_fov_from_fx_xi( + fx: Union[torch.Tensor, float], + xi: Union[torch.Tensor, float], + width: int, + device="cpu", + dtype=torch.float32, +): + """Inverse of :func:`compute_fx_from_fov_xi`.""" + + def to_tensor_1d(x): + if torch.is_tensor(x): + return x.to(device=device, dtype=dtype) + return torch.tensor([x], dtype=dtype, device=device) + + fx = to_tensor_1d(fx).reshape(-1) + xi = to_tensor_1d(xi).reshape(-1) + B = max(fx.shape[0], xi.shape[0]) + fx = fx.expand(B) + xi = xi.expand(B) + A = 2.0 * fx / width + phi = torch.atan(1.0 / A) + denom = torch.sqrt(A * A + 1.0) + ratio = (xi / denom).clamp(-1.0, 1.0) + theta = torch.asin(ratio) + phi + x_fov = torch.rad2deg(2.0 * theta) + return x_fov + + +def ucm_unproject_grid_fov( + x_fov: Union[float, torch.Tensor], + y_fov: Union[float, torch.Tensor], + xi: Union[float, torch.Tensor], + height: int, + width: int, + cx: Union[float, torch.Tensor], + cy: Union[float, torch.Tensor], + device: Union[torch.device, str] = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Unproject grid with intrinsics expressed as FoV (degrees) + xi.""" + is_batched = any(torch.is_tensor(p) and p.numel() > 1 for p in [x_fov, y_fov, xi, cx, cy]) + fx = compute_fx_from_fov_xi(x_fov, xi, width, device, dtype) + fy = compute_fx_from_fov_xi(y_fov, xi, height, device, dtype) + d_cam = ucm_unproject_grid( + height=height, + width=width, + fx=fx, + fy=fy, + cx=cx, + cy=cy, + xi=xi if torch.is_tensor(xi) else torch.tensor([xi], dtype=dtype, device=device), + dtype=dtype, + device=device, + y_down=True, + ) + if not is_batched: + d_cam = d_cam[0] + return d_cam + + +def world_to_ray_mats( + d_cam: torch.Tensor, # [H, W, 3], [B, H, W, 3], or [B, T, H, W, 3] + c2w: torch.Tensor, # [B, T, 4, 4] +) -> torch.Tensor: + """Build per-pixel ``ray<-world`` transforms from camera unit rays + C2W poses.""" + if d_cam.ndim == 3: + d_cam = d_cam.unsqueeze(0) + if d_cam.ndim == 4: + B, H, W, _ = d_cam.shape + T = c2w.shape[1] + d_cam = d_cam.unsqueeze(1).expand(-1, T, -1, -1, -1) + elif d_cam.ndim == 5: + B, T, H, W, _ = d_cam.shape + else: + raise ValueError(f"Unsupported d_cam shape: {d_cam.shape}") + + device = d_cam.device + dtype = d_cam.dtype + R_cam = c2w[..., :3, :3] + t_cam = c2w[..., :3, 3] + d_world = torch.einsum("btij,bthwj->bthwi", R_cam, d_cam) + cam_y = R_cam[..., :, 1] + # (B, T, 3) -> (B, T, H, W, 3) + cam_y = cam_y[:, :, None, None, :].expand(-1, -1, H, W, -1) + z_ray = F.normalize(d_world, dim=-1, eps=1e-6) + x_ray = torch.cross(cam_y, z_ray, dim=-1) + x_ray = F.normalize(x_ray, dim=-1, eps=1e-6) + y_ray = torch.cross(z_ray, x_ray, dim=-1) + y_ray = F.normalize(y_ray, dim=-1, eps=1e-6) + R_l2w = torch.stack([x_ray, y_ray, z_ray], dim=-1) + # (B, T, H, W, 3, 3) — transpose last two dims for the world->local rotation. + R_w2l = R_l2w.transpose(-1, -2) + # (B, T, 3) -> (B, T, H, W, 3) + t_world = t_cam[:, :, None, None, :].expand(-1, -1, H, W, -1) + t_w2l = -torch.einsum("bthwij,bthwj->bthwi", R_w2l, t_world) + raymats = torch.zeros(B, T, H, W, 4, 4, device=device, dtype=dtype) + raymats[..., :3, :3] = R_w2l + raymats[..., :3, 3] = t_w2l + raymats[..., 3, 3] = 1.0 + mask = torch.isnan(d_world).any(-1) + raymats[mask] = torch.eye(4, device=device, dtype=dtype) + return raymats + + +def create_grid( + height: int, + width: int, + batch: Optional[int] = None, + dtype: torch.dtype = torch.float32, + device: torch.device = torch.device("cpu"), +) -> torch.Tensor: + """Create a pixel coordinate grid of shape ``(H, W, 3)`` or ``(B, H, W, 3)``.""" + if device.type == "cpu": + assert dtype in (torch.float32, torch.float64), ( + f"ERR: {dtype} is not supported by {device.type}\nIf device is `cpu`, use float32 or float64" + ) + _xs = torch.linspace(0, width - 1, width, dtype=dtype, device=device) + _ys = torch.linspace(0, height - 1, height, dtype=dtype, device=device) + ys, xs = torch.meshgrid([_ys, _xs], indexing="ij") + zs = torch.ones_like(xs, dtype=dtype, device=device) + grid = torch.stack((xs, ys, zs), dim=2) + if batch is not None: + # Prepend a batch dim and broadcast. + grid = grid.unsqueeze(0).expand(batch, *grid.shape) + return grid + + +def ucm_unproject_grid( + height: int, + width: int, + fx: Union[float, torch.Tensor], + fy: Union[float, torch.Tensor], + cx: Union[float, torch.Tensor], + cy: Union[float, torch.Tensor], + xi: Union[float, torch.Tensor], + dtype: torch.dtype = torch.float32, + device: torch.device = torch.device("cpu"), + y_down: bool = True, +) -> torch.Tensor: + """Unproject pixel grid into a camera-frame direction vector using the UCM.""" + fx_, fy_, cx_, cy_, xi_ = fx, fy, cx, cy, xi + + def to_tensor_flatten(x): + if torch.is_tensor(x): + return x.to(device=device, dtype=dtype).reshape(-1) + return torch.tensor([x], dtype=dtype, device=device) + + fx, fy, cx, cy, xi = map(to_tensor_flatten, (fx, fy, cx, cy, xi)) + B = max(fx.shape[0], fy.shape[0], cx.shape[0], cy.shape[0], xi.shape[0]) + fx = fx.expand(B) + fy = fy.expand(B) + cx = cx.expand(B) + cy = cy.expand(B) + xi = xi.expand(B) + + grid = create_grid(height=height, width=width, batch=B, dtype=dtype, device=device) + u = grid[..., 0] + v = grid[..., 1] + fx = fx[:, None, None] + fy = fy[:, None, None] + cx = cx[:, None, None] + cy = cy[:, None, None] + xi = xi[:, None, None] + x = (u - cx) / fx + y = (v - cy) / fy + if not y_down: + y = -y + r2 = x * x + y * y + alpha = xi + torch.sqrt(1 + (1 - xi * xi) * r2) + gamma = alpha / (1 + r2) + X = gamma * x + Y = gamma * y + Z = gamma - xi + d_cam = torch.stack([X, Y, Z], dim=-1) + is_scalar_input = all(not torch.is_tensor(p) for p in (fx_, fy_, cx_, cy_, xi_)) + if is_scalar_input: + return d_cam[0] + else: + return d_cam + + +def compute_fx_from_fov_xi( + x_fov: Union[torch.Tensor, float], + xi: Union[torch.Tensor, float], + width: int, + device: Union[torch.device, str] = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Recover focal length ``fx`` from horizontal FoV (degrees) + UCM xi.""" + + def to_tensor_flatten(x): + if torch.is_tensor(x): + return x.to(device=device, dtype=dtype).view(-1) + return torch.tensor([x], dtype=dtype, device=device) + + x_fov = to_tensor_flatten(x_fov) + xi = to_tensor_flatten(xi) + B = max(x_fov.shape[0], xi.shape[0]) + x_fov = x_fov.expand(B) + xi = xi.expand(B) + theta = torch.deg2rad(0.5 * x_fov) + eps = torch.finfo(dtype).eps + denom = torch.sin(theta).clamp_min(eps) + fx = (width * 0.5) * (torch.cos(theta) + xi) / denom + return fx + + +def project_ucm_points(X, Y, Z, fx, fy, cx, cy, xi): + """Project 3D points in camera frame to UCM image plane.""" + r = torch.sqrt(X * X + Y * Y + Z * Z) + + def reshape_param(p, target): + if torch.is_tensor(p): + if p.numel() == 1: + return p + if p.ndim == 1 and target.ndim == 4: + return p.view(target.shape[0], target.shape[1], 1, 1) + while p.ndim < target.ndim: + p = p.unsqueeze(-1) + return p + + xi = reshape_param(xi, X) + fx = reshape_param(fx, X) + fy = reshape_param(fy, X) + cx = reshape_param(cx, X) + cy = reshape_param(cy, X) + + alpha = Z + xi * r + du = fx * (X / alpha) + cx + dv = fy * (Y / alpha) + cy + return du, dv + + +def project_ucm_points_fov(X, Y, Z, x_fov, y_fov, xi, height, width, cx, cy): + """Project 3D points in camera frame to UCM image plane using FoV-based intrinsics.""" + fx = compute_fx_from_fov_xi(x_fov, xi, width, X.device, X.dtype) + fy = compute_fx_from_fov_xi(y_fov, xi, height, X.device, X.dtype) + return project_ucm_points(X, Y, Z, fx, fy, cx, cy, xi) + + +def compute_up_lat_map( + R: torch.Tensor, + x_fov: torch.Tensor, + y_fov: torch.Tensor, + xi: torch.Tensor, + height: int, + width: int, + cx: torch.Tensor, + cy: torch.Tensor, + device: torch.device = torch.device("cpu"), + delta: float = 0.1, +): + """Compute UCPE absolute embedding maps ``(up_map, lat_map)``. + + ``up_map`` is a 2-channel projected up-direction; ``lat_map`` is a 1-channel latitude. Concatenated they form the + 3-channel absmap consumed by the camera branch. + """ + B, T, _, _ = R.shape + dtype = R.dtype + R = R.float() + d_cam = ucm_unproject_grid_fov( + x_fov=x_fov, + y_fov=y_fov, + xi=xi, + height=height, + width=width, + cx=cx, + cy=cy, + device=device, + dtype=torch.float32, + ) + + if d_cam.ndim == 3: + # (H, W, C) -> (B, T, H, W, C) + d_cam_exp = d_cam[None, None].expand(B, T, -1, -1, -1) + elif d_cam.ndim == 4: + if d_cam.shape[0] == B * T: + d_cam_exp = d_cam.view(B, T, height, width, 3) + else: + # (B, H, W, C) -> (B, T, H, W, C) + d_cam_exp = d_cam.unsqueeze(1).expand(-1, T, -1, -1, -1) + else: + d_cam_exp = d_cam + + mask_exp = d_cam_exp.isnan().any(dim=-1, keepdim=True) + d_world = torch.einsum("btij,bthwj->bthwi", R, d_cam_exp) + d_world = d_world / torch.clamp_min(d_world.norm(dim=-1, keepdim=True), 1e-8) + Xw, Yw, Zw = d_world[..., 0], d_world[..., 1], d_world[..., 2] + lat_map = torch.atan2(-Yw, torch.sqrt(Xw**2 + Zw**2)).unsqueeze(-1) + v = d_world + up_world = torch.tensor([0, -1, 0], device=device, dtype=torch.float32) + k = torch.cross(v, up_world.unsqueeze(0).unsqueeze(0).unsqueeze(0).expand_as(v), dim=-1) + k = k / torch.clamp_min(k.norm(dim=-1, keepdim=True), 1e-8) + delta_t = torch.tensor(delta, device=device, dtype=torch.float32) + cos_eps = torch.cos(delta_t) + sin_eps = torch.sin(delta_t) + v_rot = ( + v * cos_eps + torch.cross(k, v, dim=-1) * sin_eps + k * (k * (v * 1).sum(dim=-1, keepdim=True)) * (1 - cos_eps) + ) + dirs_cam = torch.einsum("btij,bthwj->bthwi", R.transpose(-1, -2), v_rot) + Xs, Ys, Zs = dirs_cam[..., 0], dirs_cam[..., 1], dirs_cam[..., 2] + du, dv = project_ucm_points_fov( + Xs, + Ys, + Zs, + x_fov=x_fov.float(), + y_fov=y_fov.float(), + xi=xi.float(), + height=height, + width=width, + cx=cx.float(), + cy=cy.float(), + ) + grid = create_grid( + height=height, + width=width, + batch=B, + dtype=torch.float32, + device=device, + ) + grid_x = grid[..., 0].unsqueeze(1) + grid_y = grid[..., 1].unsqueeze(1) + up_map = torch.stack((du - grid_x, dv - grid_y), dim=-1) + up_map = up_map / torch.clamp_min(up_map.norm(dim=-1, keepdim=True), 1e-8) + up_map = up_map.to(dtype=dtype) + lat_map = lat_map.to(dtype=dtype) + up_map = up_map.masked_fill(mask_exp, 0.0) + lat_map = lat_map.masked_fill(mask_exp, 0.0) + return up_map, lat_map diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index ef1814bbebcb..8fe2b7689221 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -385,6 +385,11 @@ "SanaVideoPipeline", "SanaImageToVideoPipeline", ] + _import_structure["sana_wm"] = [ + "SanaWMPipeline", + "SanaWMLTX2Refiner", + "SanaWMPipelineOutput", + ] _import_structure["shap_e"] = ["ShapEImg2ImgPipeline", "ShapEPipeline"] _import_structure["stable_audio"] = [ "StableAudioProjectionModel", @@ -865,6 +870,11 @@ SanaSprintPipeline, ) from .sana_video import SanaImageToVideoPipeline, SanaVideoPipeline + from .sana_wm import ( + SanaWMLTX2Refiner, + SanaWMPipeline, + SanaWMPipelineOutput, + ) from .shap_e import ShapEImg2ImgPipeline, ShapEPipeline from .stable_audio import StableAudioPipeline, StableAudioProjectionModel from .stable_audio_3 import ( diff --git a/src/diffusers/pipelines/sana_wm/__init__.py b/src/diffusers/pipelines/sana_wm/__init__.py new file mode 100644 index 000000000000..c33e4f615751 --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/__init__.py @@ -0,0 +1,49 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_output"] = ["SanaWMPipelineOutput"] + _import_structure["pipeline_sana_wm"] = ["SanaWMPipeline"] + _import_structure["refiner"] = ["SanaWMLTX2Refiner"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_output import SanaWMPipelineOutput + from .pipeline_sana_wm import SanaWMPipeline + from .refiner import SanaWMLTX2Refiner +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/sana_wm/cam_utils.py b/src/diffusers/pipelines/sana_wm/cam_utils.py new file mode 100644 index 000000000000..1661b85c536f --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/cam_utils.py @@ -0,0 +1,368 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. + +"""Camera + image utilities for the SANA-WM pipeline. + +* Action-string DSL → camera-to-world trajectory. +* Resize-and-center-crop to (704, 1280) with intrinsics adjustment. +* Plücker / raymap packing for the DiT camera-control branch. +* Optional Pi3X-based intrinsics estimation (only if `pi3` is installed). +""" + +from __future__ import annotations + +import math + +import numpy as np +import torch +from PIL import Image + + +TARGET_HEIGHT = 704 +TARGET_WIDTH = 1280 + +DEFAULT_TRANSLATION_SPEED = 0.05 +DEFAULT_ROTATION_SPEED_DEG = 1.2 +DEFAULT_PITCH_LIMIT_DEG = 85.0 +ALLOWED_ACTION_KEYS: frozenset[str] = frozenset("wasdijkl") + + +# --------------------------------------------------------------------------- +# Action DSL → camera-to-world trajectory +# --------------------------------------------------------------------------- + + +def _rot_x(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64) + + +def _rot_y(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def _parse_action_string(action: str) -> list[list[str]]: + cleaned = "".join(action.replace("，", ",").split()) + if not cleaned: + raise ValueError("action string is empty") + per_frame: list[list[str]] = [] + for segment in cleaned.split(","): + if not segment or "-" not in segment: + raise ValueError(f"Invalid action segment {segment!r}: expected '-'.") + keys_part, dur_str = segment.rsplit("-", 1) + if not dur_str.isdigit() or int(dur_str) <= 0: + raise ValueError(f"Action segment {segment!r} has a non-positive duration {dur_str!r}.") + n = int(dur_str) + keys_lower = keys_part.lower() + if keys_lower == "none": + keys: list[str] = [] + else: + bad = sorted({c for c in keys_lower if c not in ALLOWED_ACTION_KEYS}) + if bad: + raise ValueError( + f"Action segment {segment!r} contains unknown keys {bad}; " + f"allowed: {''.join(sorted(ALLOWED_ACTION_KEYS))}." + ) + keys = sorted(set(keys_lower)) + per_frame.extend([list(keys) for _ in range(n)]) + return per_frame + + +def action_string_to_c2w( + action: str, + *, + translation_speed: float = DEFAULT_TRANSLATION_SPEED, + rotation_speed_deg: float = DEFAULT_ROTATION_SPEED_DEG, + pitch_limit_deg: float = DEFAULT_PITCH_LIMIT_DEG, +) -> np.ndarray: + """Roll out a ``(N+1, 4, 4)`` c2w trajectory from a WASD+IJKL action DSL. + + Coordinate convention: OpenCV (``+X right, +Y down, +Z forward``). WASD translates on the world XZ plane; IJKL + applies pitch / yaw. + """ + per_frame = _parse_action_string(action) + rotate_rad = math.radians(rotation_speed_deg) + pitch_limit_rad = math.radians(pitch_limit_deg) + current = np.eye(4, dtype=np.float64) + poses = [current.copy()] + current_pitch = 0.0 + + for keys in per_frame: + held = set(keys) + R = current[:3, :3] + T_ = current[:3, 3] + + pitch_delta = (rotate_rad if "i" in held else 0.0) - (rotate_rad if "k" in held else 0.0) + new_pitch = current_pitch + pitch_delta + if not (-pitch_limit_rad <= new_pitch <= pitch_limit_rad): + pitch_delta = 0.0 + else: + current_pitch = new_pitch + + yaw_delta = (rotate_rad if "l" in held else 0.0) - (rotate_rad if "j" in held else 0.0) + R_new = _rot_y(yaw_delta) @ R @ _rot_x(pitch_delta) + + forward = R_new[:, 2].copy() + forward[1] = 0.0 + right = R_new[:, 0].copy() + right[1] = 0.0 + if (fn := float(np.linalg.norm(forward))) > 0: + forward /= fn + if (rn := float(np.linalg.norm(right))) > 0: + right /= rn + move = np.zeros(3, dtype=np.float64) + if "w" in held: + move += forward * translation_speed + if "s" in held: + move -= forward * translation_speed + if "d" in held: + move += right * translation_speed + if "a" in held: + move -= right * translation_speed + + current = np.eye(4, dtype=np.float64) + current[:3, :3] = R_new + current[:3, 3] = T_ + move + poses.append(current.copy()) + + return np.stack(poses, axis=0).astype(np.float32) + + +# --------------------------------------------------------------------------- +# Intrinsics handling +# --------------------------------------------------------------------------- + + +def transform_intrinsics_for_crop( + intrinsics_vec4: np.ndarray, + src_size: tuple[int, int], + resized_size: tuple[int, int], + crop_offset: tuple[int, int], +) -> np.ndarray: + """Adjust ``[fx, fy, cx, cy]`` to match a resize-then-center-crop image.""" + src_w, src_h = src_size + rw, rh = resized_size + cl, ct = crop_offset + sx, sy = rw / src_w, rh / src_h + out = intrinsics_vec4.copy() + out[..., 0] *= sx + out[..., 2] = out[..., 2] * sx - cl + out[..., 1] *= sy + out[..., 3] = out[..., 3] * sy - ct + return out + + +def estimate_intrinsics_with_pi3x(image: Image.Image, device: torch.device | str = "cuda") -> np.ndarray: + """Estimate ``[fx, fy, cx, cy]`` for ``image`` using Pi3X. + + Optional helper — requires ``pip install pi3-vision``. The result is in the **original image** pixel grid (not the + cropped one); pass it to [`SanaWMPipeline.__call__`] as ``intrinsics=...``. + """ + try: + from pi3.models.pi3x import Pi3X # type: ignore + from pi3.utils.geometry import recover_intrinsic_from_rays_d # type: ignore + except ImportError as e: # pragma: no cover + raise RuntimeError( + "pi3 is required for intrinsics estimation. Pass `intrinsics` explicitly or `pip install pi3-vision`." + ) from e + + from torchvision import transforms as T # noqa: PLC0415 + + device_t = torch.device(device) + W_orig, H_orig = image.size + pixel_limit = 255_000 + scale = math.sqrt(pixel_limit / (W_orig * H_orig)) if W_orig * H_orig > 0 else 1.0 + W_t, H_t = W_orig * scale, H_orig * scale + k, m = max(1, round(W_t / 14)), max(1, round(H_t / 14)) + while (k * 14) * (m * 14) > pixel_limit: + if k / m > W_t / H_t: + k -= 1 + else: + m -= 1 + W_model, H_model = max(1, k) * 14, max(1, m) * 14 + resized = image.resize((W_model, H_model), Image.Resampling.LANCZOS) + tensor = T.ToTensor()(resized).unsqueeze(0).unsqueeze(0).to(device_t) + + dtype = ( + torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 else torch.float16 + ) + model = Pi3X.from_pretrained("yyfz233/Pi3X").to(device_t).eval() + model.disable_multimodal() + model.requires_grad_(False) + with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype): + out = model(imgs=tensor) + rays_d = torch.nn.functional.normalize(out["local_points"], dim=-1) + K = recover_intrinsic_from_rays_d(rays_d, force_center_principal_point=True)[0, 0] + K = K.detach().cpu().float().numpy() + sx, sy = W_orig / W_model, H_orig / H_model + return np.array([K[0, 0] * sx, K[1, 1] * sy, K[0, 2] * sx, K[1, 2] * sy], dtype=np.float32) + + +# --------------------------------------------------------------------------- +# Image preprocessing +# --------------------------------------------------------------------------- + + +def resize_and_center_crop( + image: Image.Image, + target_h: int = TARGET_HEIGHT, + target_w: int = TARGET_WIDTH, +) -> tuple[Image.Image, tuple[int, int], tuple[int, int], tuple[int, int]]: + """Aspect-preserving resize then center-crop to ``(target_h, target_w)``.""" + src_w, src_h = image.size + scale = max(target_h / src_h, target_w / src_w) + rw = max(target_w, int(round(src_w * scale))) + rh = max(target_h, int(round(src_h * scale))) + resized = image.resize((rw, rh), Image.LANCZOS) + left = (rw - target_w) // 2 + top = (rh - target_h) // 2 + cropped = resized.crop((left, top, left + target_w, top + target_h)) + return cropped, (src_w, src_h), (rw, rh), (left, top) + + +# --------------------------------------------------------------------------- +# Camera condition packing — Plücker + raymap +# --------------------------------------------------------------------------- + + +def compute_raymap( + intrinsics: torch.Tensor, + poses: torch.Tensor, + H: int, + W: int, + *, + use_plucker: bool = True, +) -> torch.Tensor: + """Compute a per-pixel ray geometry map. + + Args: + intrinsics: ``(T, 4)`` ``[fx, fy, cx, cy]`` per frame. + poses: ``(T, 4, 4)`` camera-to-world poses (OpenCV convention). + H: spatial height. + W: spatial width. + use_plucker: if True returns Plücker coordinates ``(d, m)``; otherwise + returns ``(origin, direction)``. + + Returns: + ``(T, H, W, 6)`` tensor. + """ + T = intrinsics.shape[0] + device = intrinsics.device + dtype = intrinsics.dtype + y_grid, x_grid = torch.meshgrid( + torch.arange(H, device=device, dtype=dtype), + torch.arange(W, device=device, dtype=dtype), + indexing="ij", + ) + x_grid = x_grid[None].expand(T, -1, -1) + y_grid = y_grid[None].expand(T, -1, -1) + fx = intrinsics[:, 0].view(T, 1, 1) + fy = intrinsics[:, 1].view(T, 1, 1) + cx = intrinsics[:, 2].view(T, 1, 1) + cy = intrinsics[:, 3].view(T, 1, 1) + dirs_cam = torch.stack( + [(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], + dim=-1, + ) + R = poses[:, :3, :3] + t = poses[:, :3, 3] + dirs_world = torch.einsum("tij,thwj->thwi", R, dirs_cam) + dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True) + origins = t.view(T, 1, 1, 3).expand_as(dirs_world) + if use_plucker: + moments = torch.cross(origins, dirs_world, dim=-1) + return torch.cat([dirs_world, moments], dim=-1) + return torch.cat([origins, dirs_world], dim=-1) + + +def _pose_inverse(T44: torch.Tensor) -> torch.Tensor: + R = T44[..., :3, :3] + t = T44[..., :3, 3:] + Rt = R.transpose(-1, -2) + out = torch.zeros_like(T44) + out[..., :3, :3] = Rt + out[..., :3, 3:] = -Rt @ t + out[..., 3, 3] = 1.0 + return out + + +def prepare_camera( + poses_c2w: np.ndarray, + intrinsics_vec4: np.ndarray, + *, + target_size: tuple[int, int], + vae_stride: tuple[int, int, int], +) -> dict[str, torch.Tensor]: + """Build the DiT-input camera tensors. + + Returns a dict with: + + * ``raymap`` ``(T_lat, 20)`` — flattened (rel-pose, intrinsics) per latent frame + * ``chunk_plucker`` ``(6 * vae_time_stride, T_lat, H_lat, W_lat)`` — Plücker coordinates packed by chunk. + """ + num_frames = poses_c2w.shape[0] + vae_time_stride, vae_spatial_stride = vae_stride[0], vae_stride[-1] + H_pixel, W_pixel = target_size + latent_h = H_pixel // vae_spatial_stride + latent_w = W_pixel // vae_spatial_stride + latent_frames = (num_frames - 1) // vae_time_stride + 1 + + poses = torch.from_numpy(poses_c2w).float() + first_inv = _pose_inverse(poses[0:1]).squeeze(0) + poses_rel = torch.matmul(first_inv, poses[1:]) + poses = torch.cat([torch.eye(4).unsqueeze(0), poses_rel], dim=0) + + intrinsics = torch.from_numpy(intrinsics_vec4).float() + intrinsics_latent = intrinsics.clone() + intrinsics_latent[:, [0, 2]] *= latent_w / float(W_pixel) + intrinsics_latent[:, [1, 3]] *= latent_h / float(H_pixel) + + time_indices = torch.arange(0, num_frames, vae_time_stride) + if len(time_indices) > latent_frames: + time_indices = time_indices[:latent_frames] + + raymap = torch.cat( + [poses[time_indices].reshape(len(time_indices), -1), intrinsics_latent[time_indices]], + dim=-1, + ) + + chunk_starts = time_indices - (vae_time_stride - 1) + chunks = [] + for start in chunk_starts: + s = max(0, int(start)) + e = s + vae_time_stride + chunk_poses, chunk_intrs = poses[s:e], intrinsics_latent[s:e] + if chunk_poses.shape[0] < vae_time_stride: + pad = vae_time_stride - chunk_poses.shape[0] + chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)], dim=0) + chunk_intrs = torch.cat([chunk_intrs, chunk_intrs[-1:].repeat(pad, 1)], dim=0) + plucker = compute_raymap(chunk_intrs, chunk_poses, latent_h, latent_w, use_plucker=True) + chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_h, latent_w)) + chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3) + return {"raymap": raymap, "chunk_plucker": chunk_plucker} + + +def snap_num_frames(n: int, stride: int = 8, *, upper_bound: int | None = None) -> int: + """Snap ``n`` to the nearest ``stride*k + 1`` (LTX-2 VAE constraint).""" + if n < 1: + return 1 + if (n - 1) % stride == 0: + return n + floor_cand = n - ((n - 1) % stride) + ceil_cand = floor_cand + stride + snapped = floor_cand if (n - floor_cand) < (ceil_cand - n) else ceil_cand + if upper_bound is not None and snapped > upper_bound: + snapped = floor_cand + return max(snapped, 1) diff --git a/src/diffusers/pipelines/sana_wm/image_processor.py b/src/diffusers/pipelines/sana_wm/image_processor.py new file mode 100644 index 000000000000..c18326cd33fd --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/image_processor.py @@ -0,0 +1,67 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. + +from __future__ import annotations + +import numpy as np +import PIL.Image +import torch + +from ...configuration_utils import register_to_config +from ...image_processor import VaeImageProcessor +from .cam_utils import TARGET_HEIGHT, TARGET_WIDTH, resize_and_center_crop, transform_intrinsics_for_crop + + +class SanaWMImageProcessor(VaeImageProcessor): + r""" + Image processor for SANA-WM's first-frame input. + + SANA-WM was trained at a fixed 704×1280 resolution with an aspect-preserving *resize + center-crop* transform. The + pipeline also needs to rescale the per-frame camera intrinsics ``[fx, fy, cx, cy]`` to match the crop — + ``preprocess_with_intrinsics`` does both in one call so the two stay in lockstep. + + Args: + vae_scale_factor (`int`, defaults to `32`): + LTX-2 VAE spatial stride. + do_normalize (`bool`, defaults to `True`): + Standard `VaeImageProcessor` [-1, 1] normalization. + """ + + @register_to_config + def __init__(self, vae_scale_factor: int = 32, do_normalize: bool = True) -> None: + super().__init__(vae_scale_factor=vae_scale_factor, do_normalize=do_normalize) + + def preprocess_with_intrinsics( + self, + image: PIL.Image.Image, + intrinsics: np.ndarray, + height: int = TARGET_HEIGHT, + width: int = TARGET_WIDTH, + ) -> tuple[torch.Tensor, np.ndarray]: + """Resize + center-crop the image and rescale ``intrinsics`` to match. + + Args: + image: RGB PIL image (any size). + intrinsics: ``(F, 4)`` ``[fx, fy, cx, cy]`` per frame in original-image pixel coordinates. + height / width: Target crop size (defaults to SANA-WM's training resolution). + + Returns: + ``(pixel_values, intrinsics_cropped)``: + * ``pixel_values`` — ``(1, 3, H, W)`` tensor in `[-1, 1]` (VaeImageProcessor convention). + * ``intrinsics_cropped`` — ``(F, 4)`` array rescaled for the resize + crop. + """ + cropped, src_size, resized_size, crop_offset = resize_and_center_crop(image, height, width) + pixel_values = self.preprocess(cropped, height=height, width=width) + intr = transform_intrinsics_for_crop(intrinsics, src_size, resized_size, crop_offset) + return pixel_values, intr diff --git a/src/diffusers/pipelines/sana_wm/pipeline_output.py b/src/diffusers/pipelines/sana_wm/pipeline_output.py new file mode 100644 index 000000000000..9a007071c32f --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/pipeline_output.py @@ -0,0 +1,43 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. + +from dataclasses import dataclass + +import numpy as np +import PIL.Image +import torch + +from ...utils import BaseOutput + + +@dataclass +class SanaWMPipelineOutput(BaseOutput): + """ + Output class for SANA-WM image-to-video pipeline. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or `list[PIL.Image.Image]`): + Generated video. Shape ``(T, H, W, 3)`` as a float ``np.ndarray`` / ``torch.Tensor`` in ``[0, 1]`` when + ``output_type="np"`` / ``"latent"``, or a list of ``PIL.Image`` of length ``T`` when ``output_type="pil"``. + c2w (`np.ndarray`): + Camera-to-world poses ``(T, 4, 4)`` aligned with ``frames`` (the refiner drops the sink anchor frame; this + array is realigned accordingly when the refiner ran). + latent (`torch.Tensor`, optional): + Latent tensor in LTX-2 VAE space, shape ``(B, C, T_lat, H_lat, W_lat)``. Returned when + ``output_type="latent"``. + """ + + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] + c2w: np.ndarray | None = None + latent: torch.Tensor | None = None diff --git a/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py b/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py new file mode 100644 index 000000000000..f1b53f954169 --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py @@ -0,0 +1,648 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. + +from __future__ import annotations + +import inspect +from pathlib import Path +from typing import Literal + +import numpy as np +import PIL.Image +import torch +from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast + +from ...models import AutoencoderKLLTX2Video, SanaWMTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging, replace_example_docstring +from ...utils.torch_utils import empty_device_cache, randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .cam_utils import ( + TARGET_HEIGHT, + TARGET_WIDTH, + action_string_to_c2w, + prepare_camera, + snap_num_frames, +) +from .image_processor import SanaWMImageProcessor +from .pipeline_output import SanaWMPipelineOutput +from .refiner import SanaWMLTX2Refiner + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from PIL import Image + >>> from diffusers import SanaWMPipeline + + >>> pipe = SanaWMPipeline.from_pretrained( + ... "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", torch_dtype=torch.bfloat16 + ... ).to("cuda") + + >>> output = pipe( + ... image=Image.open("input.png").convert("RGB"), + ... prompt="A car driving across a vast desert plain at golden hour.", + ... action="w-80,jw-40,w-40", + ... intrinsics=[800.0, 800.0, 845.0, 464.0], # fx, fy, cx, cy in original-image pixels + ... num_inference_steps=60, + ... ) + >>> # output.frames is (T, H, W, 3) float np.ndarray in [0, 1] (diffusers convention). + ``` +""" + + +# Default instruction prefix prepended to the user prompt before Gemma-2 encoding. +# SANA-WM was trained with this prefix, so changing it degrades prompt adherence. +DEFAULT_CHI_PROMPT: list[str] = [ + 'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:', + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", +] + + +class SanaWMPipeline(DiffusionPipeline): + r""" + SANA-WM camera-controlled image-to-video pipeline. + + Generates a video from a first-frame image, a text prompt, and a camera trajectory (explicit ``c2w`` poses or a + WASD/IJKL action string). Uses the 1600M bidirectional SANA DiT for stage-1 sampling and the LTX-2 + sink-bidirectional Euler refiner for stage-2 polish; both decode through the LTX-2 VAE. + + Args: + tokenizer ([`GemmaTokenizer`] or [`GemmaTokenizerFast`]): + The Gemma-2 tokenizer. + text_encoder ([`Gemma2PreTrainedModel`]): + The Gemma-2 text encoder. + vae ([`AutoencoderKLLTX2Video`]): + The LTX-2 VAE. + transformer ([`SanaWMTransformer3DModel`]): + The 1600M bidirectional SANA-WM DiT. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + Flow-matching Euler scheduler (LTX-style per-token timesteps). + refiner ([`SanaWMLTX2Refiner`], *optional*): + LTX-2 refiner; if provided, runs 3-step distilled refinement before decoding. If `None`, decode stage-1 + latents directly. + """ + + # ``refiner`` is a nested pipeline (not an nn.Module) so it's excluded from + # the offload sequence; it manages its own sub-module device placement. + model_cpu_offload_seq = "text_encoder->transformer->vae" + _optional_components = ["refiner"] + + def __init__( + self, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + text_encoder: Gemma2PreTrainedModel, + vae: AutoencoderKLLTX2Video, + transformer: SanaWMTransformer3DModel, + scheduler: FlowMatchEulerDiscreteScheduler, + refiner: SanaWMLTX2Refiner | None = None, + ) -> None: + super().__init__() + self.register_modules( + tokenizer=tokenizer, + text_encoder=text_encoder, + vae=vae, + transformer=transformer, + scheduler=scheduler, + refiner=refiner, + ) + # Read VAE strides from the registered component (LTX2Pipeline pattern). + # Fall back to the LTX-2 defaults (32 spatial / 8 temporal) if the VAE + # hasn't been registered yet — matches SANA-WM's training config. + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + # ``image_processor`` handles first-frame input (resize + center-crop + # + [-1, 1] normalization + intrinsics rescale for the crop); + # ``video_processor`` handles the decoded [-1, 1] video -> user-chosen + # ``output_type`` conversion. + self.image_processor = SanaWMImageProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + # The SANA DiT's ``y_embedder`` randomly null-replaces tokens when + # ``self.training=True``. Force eval mode at construction so inference + # is deterministic regardless of how the underlying modules were saved. + if transformer is not None: + transformer.eval() + if vae is not None: + vae.eval() + if text_encoder is not None: + text_encoder.eval() + + # SANA was trained with right-padded prompts; Gemma's default is + # "left", and the saved tokenizer reverts to "left" on load. Pin it. + if tokenizer is not None: + tokenizer.padding_side = "right" + + # SANA-WM trained on LTX-2 VAE in framewise mode with tiling enabled; + # without these flags the VAE encodes the full (B, C, T, H, W) input + # in one shot, which gives subtly different numerics. + if vae is not None: + if hasattr(vae, "enable_tiling"): + vae.enable_tiling() + if hasattr(vae, "use_framewise_encoding"): + vae.use_framewise_encoding = True + vae.use_framewise_decoding = True + vae.tile_sample_stride_num_frames = 64 + vae.tile_sample_min_num_frames = 96 + + def _model_cpu_offload_active(self) -> bool: + """Whether `enable_model_cpu_offload` currently owns module placement. + + Mirrors the check `DiffusionPipeline` uses internally: the hooks list only exists (and is non-empty) while + model CPU offload is installed, and `remove_all_hooks()` empties it again. + """ + return hasattr(self, "_all_hooks") and len(self._all_hooks) > 0 + + # ------------------------------------------------------------------ + # Prompt encoding + # ------------------------------------------------------------------ + + def encode_prompt( + self, + prompt: str, + negative_prompt: str = "", + *, + device: torch.device, + max_sequence_length: int = 300, + chi_prompt: list[str] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode prompt + negative prompt through Gemma-2. + + Mirrors the SANA chi-prompt-prefix trick: the chi prompt is prepended to the user prompt, then a ``select_index + = [0, -L+1, ..., -1]`` slice takes the BOS token plus the last ``max_sequence_length - 1`` tokens. + + Returns: + ``(cond, cond_mask, neg, neg_mask)`` where ``cond`` and ``neg`` are ``(1, 1, L, D)``-shaped Gemma hidden + states and the masks are ``(1, L)``. + """ + chi = "\n".join(chi_prompt) if chi_prompt else "" + if chi: + full_prompt = chi + prompt + max_length_all = len(self.tokenizer.encode(chi)) + max_sequence_length - 2 + else: + full_prompt = prompt + max_length_all = max_sequence_length + + def _encode(text: str, length: int) -> tuple[torch.Tensor, torch.Tensor]: + tok = self.tokenizer( + [text], + max_length=length, + padding="max_length", + truncation=True, + return_tensors="pt", + ).to(device) + # Go through the outer ``Gemma2ForCausalLM`` so the CPU-offload + # hook moves the encoder to GPU; grab the final-layer hidden + # states (== ``Gemma2Model.last_hidden_state``). + out = self.text_encoder( + input_ids=tok.input_ids, + attention_mask=tok.attention_mask, + output_hidden_states=True, + return_dict=True, + ) + return out.hidden_states[-1], tok.attention_mask + + cond, cond_mask = _encode(full_prompt, max_length_all) + select = [0] + list(range(-max_sequence_length + 1, 0)) + cond = cond[:, None][:, :, select] + cond_mask = cond_mask[:, select] + + neg, neg_mask = _encode(negative_prompt, max_sequence_length) + return cond, cond_mask, neg[:, None], neg_mask + + # ------------------------------------------------------------------ + # First-frame VAE encode (deterministic — uses posterior mode) + # ------------------------------------------------------------------ + + def _encode_first_frame( + self, pixel_values: torch.Tensor, device: torch.device, dtype: torch.dtype + ) -> torch.Tensor: + # ``pixel_values`` is ``(1, 3, H, W)`` in [-1, 1] (from ``SanaWMImageProcessor``). + # Add the temporal axis to match the LTX-2 VAE input shape ``(B, C, 1, H, W)``. + img = pixel_values.unsqueeze(2).to(device, dtype=self.vae.dtype) + z = self.vae.encode(img).latent_dist.mode() + latents_mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(z) + latents_std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(z) + z = (z - latents_mean) * self.vae.config.scaling_factor / latents_std + return z.to(dtype) + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Decode latents to a `(B, C, F, H, W)` tensor in `[-1, 1]` (the VAE's native output range). + + Post-processing (e.g. `[-1, 1]` → PIL frames / `np.ndarray` in `[0, 1]`) is handled by + `self.video_processor.postprocess_video` at the call site so callers get the diffusers convention that the + `VideoProcessor` / `export_to_video` helpers assume. + """ + latents = latents.to(self.vae.device, dtype=self.vae.dtype) + latents_mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) + latents_std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) + latents = latents / self.vae.config.scaling_factor * latents_std + latents_mean + return self.vae.decode(latents, return_dict=False)[0] + + # ------------------------------------------------------------------ + # Camera conditioning packing + # ------------------------------------------------------------------ + + def _build_camera_kwargs( + self, + c2w: np.ndarray, + intrinsics_vec4: np.ndarray, + target_size: tuple[int, int], + *, + device: torch.device, + dtype: torch.dtype, + do_cfg: bool, + ) -> dict[str, torch.Tensor]: + cam = prepare_camera( + c2w, + intrinsics_vec4, + target_size=target_size, + vae_stride=( + self.vae_temporal_compression_ratio, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ), + ) + raymap = cam["raymap"].unsqueeze(0).to(device, dtype=dtype) + chunk_plucker = cam["chunk_plucker"].unsqueeze(0).to(device, dtype=dtype) + if do_cfg: + raymap = torch.cat([raymap, raymap], dim=0) + chunk_plucker = torch.cat([chunk_plucker, chunk_plucker], dim=0) + return {"camera_conditions": raymap, "chunk_plucker": chunk_plucker} + + def check_inputs( + self, + image: PIL.Image.Image | str | Path, + c2w: np.ndarray | None, + action: str | None, + intrinsics: np.ndarray | list[float] | None, + num_frames: int, + ) -> tuple[PIL.Image.Image, np.ndarray, np.ndarray]: + """Validate `__call__` inputs and normalize to ``(image_pil, c2w_(F,4,4), intrinsics_(F,4))``. + + Also snaps ``num_frames`` to the VAE-friendly ``8k+1`` and trims the c2w / intrinsics arrays to match. The + cropped image + rescaled intrinsics come later once we know the target resolution. + """ + if isinstance(image, (str, Path)): + image = PIL.Image.open(image).convert("RGB") + + if (c2w is None) == (action is None): + raise ValueError("Provide exactly one of `c2w` or `action`.") + if action is not None: + c2w = action_string_to_c2w(action) + c2w = np.asarray(c2w, dtype=np.float32) + if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): + raise ValueError(f"`c2w` must be `(F, 4, 4)`; got {c2w.shape}.") + + num_frames = min(num_frames, c2w.shape[0]) + num_frames = snap_num_frames(num_frames, stride=self.vae_temporal_compression_ratio, upper_bound=c2w.shape[0]) + c2w = c2w[:num_frames] + + if intrinsics is None: + raise ValueError( + "Pass `intrinsics` as either `[fx, fy, cx, cy]`, a 3x3 K matrix, " + "an `(F, 4)` per-frame [fx,fy,cx,cy], or `(F, 3, 3)` per-frame K — " + "all in original-image pixel coordinates. Use " + "`diffusers.pipelines.sana_wm.cam_utils.estimate_intrinsics_with_pi3x(image)` " + "for an automatic estimate if pi3 is installed." + ) + intr = np.asarray(intrinsics, dtype=np.float32) + # Accept (3, 3), (F, 3, 3), (4,) and (F, 4) — normalize to (F, 4). + if intr.shape == (3, 3): + intr = np.array([intr[0, 0], intr[1, 1], intr[0, 2], intr[1, 2]], dtype=np.float32) + elif intr.ndim == 3 and intr.shape[-2:] == (3, 3): + intr = np.stack([intr[:, 0, 0], intr[:, 1, 1], intr[:, 0, 2], intr[:, 1, 2]], axis=-1) + if intr.shape == (4,): + intr = np.broadcast_to(intr, (num_frames, 4)).copy() + if intr.ndim == 2 and intr.shape[1] == 4 and intr.shape[0] >= num_frames: + # Caller may pass a full-trajectory intrinsics array; trim to match. + intr = intr[:num_frames] + if intr.shape != (num_frames, 4): + raise ValueError( + f"`intrinsics` must be `(4,)`, `(F>={num_frames}, 4)`, `(3, 3)`, or " + f"`(F>={num_frames}, 3, 3)`; got shape {np.asarray(intrinsics).shape}." + ) + return image, c2w, intr + + def prepare_latents( + self, + first_latent: torch.Tensor, + num_frames: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Sample initial latents and pin the first frame as the conditioning anchor. + + Returns ``(latents, condition_mask)`` where ``condition_mask`` has ones on the first frame's tokens (they are + held clean throughout sampling) and zeros elsewhere. + """ + latent_T = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_h = height // self.vae_spatial_compression_ratio + latent_w = width // self.vae_spatial_compression_ratio + latent_channels = first_latent.shape[1] + latents = randn_tensor( + (1, latent_channels, latent_T, latent_h, latent_w), + generator=generator, + device=device, + dtype=dtype, + ) + latents[:, :, :1] = first_latent + condition_mask = torch.zeros_like(latents) + condition_mask[:, :, :1] = 1.0 + return latents, condition_mask + + # ------------------------------------------------------------------ + # __call__ + # ------------------------------------------------------------------ + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + image: PIL.Image.Image | str | Path, + prompt: str, + *, + c2w: np.ndarray | None = None, + action: str | None = None, + intrinsics: np.ndarray | list[float] | None = None, + height: int = TARGET_HEIGHT, + width: int = TARGET_WIDTH, + num_frames: int = 161, + fps: int = 16, + num_inference_steps: int = 60, + guidance_scale: float = 5.0, + negative_prompt: str = "", + generator: torch.Generator | list[torch.Generator] | None = None, + seed: int | None = None, + use_refiner: bool = True, + sink_size: int = 1, + refiner_seed: int = 42, + max_sequence_length: int = 300, + chi_prompt: list[str] | None = None, + output_type: Literal["np", "pil", "latent"] = "np", + return_dict: bool = True, + ) -> SanaWMPipelineOutput | tuple: + r""" + Generate a SANA-WM camera-controlled video. + + Args: + image (`PIL.Image.Image` or `str`): + First-frame image (PIL or path). + prompt (`str`): + Text prompt. + c2w (`np.ndarray`, *optional*): + ``(F, 4, 4)`` camera-to-world poses. Mutually exclusive with `action`. + action (`str`, *optional*): + Action-DSL string e.g. ``"w-80,jw-40,w-40"``. Mutually exclusive with `c2w`. + intrinsics (`np.ndarray` or `list[float]`): + ``[fx, fy, cx, cy]`` in **original-image** pixel coordinates. The pipeline applies the resize+crop + transform internally. + height (`int`, defaults to 704): + Output frame height (fixed for the public model). + width (`int`, defaults to 1280): + Output frame width (fixed for the public model). + num_frames (`int`, defaults to 161): + Target frame count; snapped to ``8k+1`` (LTX-2 VAE constraint). + fps (`int`, defaults to 16): + Output frame rate (also fed to the refiner). + num_inference_steps (`int`, defaults to 60): + Number of stage-1 DiT sampling steps. + guidance_scale (`float`, defaults to 5.0): + Classifier-free guidance scale. + negative_prompt (`str`, defaults to ""): + Optional negative prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or more torch generators to make the noise sampling deterministic. If both `generator` and `seed` + are provided, `generator` takes precedence. + seed (`int`, *optional*): + Convenience shortcut — used only when `generator` is `None`, in which case a fresh + ``torch.Generator(device=execution_device).manual_seed(seed)`` is created. If both are `None`, the + sampling is non-deterministic. + use_refiner (`bool`, defaults to True): + Run the LTX-2 refiner (requires `self.refiner` to be set). + sink_size (`int`, defaults to 1): + Refiner sink-anchor frame count. + refiner_seed (`int`, defaults to 42): + Refiner sampling seed. + max_sequence_length (`int`, defaults to 300): + Max prompt tokens. + chi_prompt (`list[str]`, *optional*): + Override the chi-prompt prefix (default mirrors the public release). + output_type (`"np"`, `"pil"`, or `"latent"`, defaults to `"np"`): + Output format. + return_dict (`bool`, defaults to True): + Return [`SanaWMPipelineOutput`] vs tuple. + + Returns: + [`SanaWMPipelineOutput`] with `.frames` of shape ``(T, H, W, 3)``, float ``np.ndarray`` in ``[0, 1]`` for + `output_type="np"`, a list of ``PIL.Image.Image`` of length ``T`` for `"pil"`, or the raw latent tensor for + `"latent"`. + + Examples: + """ + image, c2w, intr = self.check_inputs(image, c2w, action, intrinsics, num_frames) + num_frames = c2w.shape[0] + pixel_values, intr = self.image_processor.preprocess_with_intrinsics(image, intr, height, width) + + device = self._execution_device + dtype = self.transformer.dtype + + cond, cond_mask, neg, neg_mask = self.encode_prompt( + prompt, + negative_prompt, + device=device, + max_sequence_length=max_sequence_length, + chi_prompt=chi_prompt or DEFAULT_CHI_PROMPT, + ) + + first_latent = self._encode_first_frame(pixel_values, device, dtype) + cam_kwargs = self._build_camera_kwargs( + c2w, intr, (height, width), device=device, dtype=dtype, do_cfg=guidance_scale > 1.0 + ) + + if generator is None and seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + do_cfg = guidance_scale > 1.0 + + # Stage-1 denoising — LTX-style flow-matching Euler with per-token + # timesteps. The first latent frame is the conditioning anchor: its + # per-token timestep is pinned to 0 so it is never denoised away. + latents, condition_mask = self.prepare_latents( + first_latent, num_frames, height, width, dtype, device, generator + ) + timesteps, _ = retrieve_timesteps(self.scheduler, num_inference_steps, device, None) + + prompt_embeds = torch.cat([neg, cond], dim=0) if do_cfg else cond + mask_cfg = torch.cat([neg_mask, cond_mask], dim=0) if do_cfg else cond_mask + model_kwargs = { + "data_info": { + "img_hw": torch.tensor([[height, width]], dtype=torch.float, device=device), + }, + "mask": mask_cfg, + **cam_kwargs, + } + + for t in self.progress_bar(timesteps): + cond_mask_input = torch.cat([condition_mask] * 2) if do_cfg else condition_mask + latent_model_input = torch.cat([latents] * 2) if do_cfg else latents + timestep = t.expand(cond_mask_input.shape).float() + timestep = torch.min(timestep, (1.0 - cond_mask_input) * 1000.0) + + noise_pred = self.transformer( + latent_model_input, + timestep[:, :1, :, 0, 0], # (B, 1, T) + prompt_embeds, + return_dict=False, + **model_kwargs, + )[0] + + if do_cfg: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + timestep = timestep.chunk(2)[0] + + B, C, F, H, W = latents.shape + denoised = self.scheduler.step( + -noise_pred.reshape(B, C, -1).transpose(1, 2), + t, + latents.reshape(B, C, -1).transpose(1, 2), + per_token_timesteps=timestep.reshape(B, C, -1)[:, 0], + return_dict=False, + )[0] + denoised = denoised.transpose(1, 2).reshape(B, C, F, H, W) + keep_clean = t / 1000.0 - 1e-6 < (1.0 - condition_mask) + latents = torch.where(keep_clean, denoised, latents).to(dtype) + + if output_type == "latent": + return SanaWMPipelineOutput(frames=latents, c2w=c2w, latent=latents) if return_dict else (latents,) + + if use_refiner and self.refiner is not None: + # Stage-1 is done; free the parent's GPU-resident weights so the + # refiner (nested pipeline, manages its own placement) has the device + # to itself. Skip when accelerate offload is active — it owns + # placement then. The VAE is moved back for decode below. + if not self._model_cpu_offload_active(): + self.text_encoder.to("cpu") + self.transformer.to("cpu") + self.vae.to("cpu") + empty_device_cache(device.type) + # The refiner is a nested pipeline, so it doesn't follow the parent's + # ``.to(device)`` / offload hooks. Rather than bulk-moving its (~87 GB) + # weights up front, pass the execution device and let it move its own + # sub-modules on/off GPU as it runs (peak VRAM ~= largest sub-model). + refined = self.refiner( + latents, + prompt, + fps=float(fps), + sink_size=sink_size, + seed=refiner_seed, + device=device, + ) + # Bring the VAE back for decode (moved to CPU above to free the GPU + # for the refiner). No-op under accelerate offload. + if not self._model_cpu_offload_active(): + self.vae.to(device) + empty_device_cache(device.type) + decoded = self._decode_latents(refined) # (B=1, C=3, F, H, W) in [-1, 1] + decoded = decoded[:, :, 1:] # refiner drops the sink anchor frame + video_c2w = c2w[1:num_frames] + else: + decoded = self._decode_latents(latents) + video_c2w = c2w[:num_frames] + + # ``VideoProcessor.postprocess_video`` handles the standard [-1, 1] -> + # requested output_type conversion (uint8 PIL frames, float np.ndarray + # in [0, 1], or the raw pt tensor). + frames = self.video_processor.postprocess_video(decoded, output_type=output_type)[0] + + if not return_dict: + return (frames,) + return SanaWMPipelineOutput(frames=frames, c2w=video_c2w, latent=latents) diff --git a/src/diffusers/pipelines/sana_wm/refiner.py b/src/diffusers/pipelines/sana_wm/refiner.py new file mode 100644 index 000000000000..9c3979182145 --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/refiner.py @@ -0,0 +1,1016 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. + +"""LTX-2 chunk-causal AR refiner used as SANA-WM stage 2. + +Wraps diffusers' own ``LTX2VideoTransformer3DModel`` + ``LTX2TextConnectors`` plus a Gemma-3 text encoder. The +transformer's public forward always runs the audio stream and does not expose the streaming sink/current self-attention +mask this refiner was trained with, so we run a video-only forward in-place with a sink/current attention split. + +Refinement is chunk-causal / autoregressive (``block_size=3``, ``kv_max_frames=11``): ``block_size`` latent frames are +processed at a time over a sliding window of ``[source_sink + recent_history + active_block]`` K/V. The model was +trained with this contract; per-block compute is bounded by the window size, so total cost scales linearly with video +length. +""" + +from __future__ import annotations + +import torch +from torch import nn +from tqdm.auto import tqdm + +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils.torch_utils import empty_device_cache, randn_tensor +from ..pipeline_utils import DiffusionPipeline + + +# Sigma schedule for the 3-step distilled refiner (matches the public release). +STAGE_2_DISTILLED_SIGMA_VALUES: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0) + + +class SanaWMLTX2Refiner(DiffusionPipeline): + r""" + LTX-2 sink-bidirectional Euler refiner — SANA-WM stage 2, as a standalone pipeline. + + Wraps the diffusers LTX-2 components (transformer + text connectors + Gemma-3 text encoder + tokenizer) plus a + [`FlowMatchEulerDiscreteScheduler`] that carries the distilled sigma schedule and performs the Euler steps. It is + registered as an optional component of [`SanaWMPipeline`] and can also be used on its own to refine stage-1 + latents. + + Args: + transformer ([`LTX2VideoTransformer3DModel`]): + The LTX-2 video DiT. + connectors ([`LTX2TextConnectors`]): + LTX-2 text connectors. + tokenizer: + Gemma-3 tokenizer. + text_encoder: + Gemma-3 text encoder. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + Flow-matching Euler scheduler. Constructed with ``shift=1.0`` so the distilled sigmas pass through + unmodified. + text_max_sequence_length (`int`, defaults to 1024): + Maximum tokens passed to the Gemma-3 tokenizer. + """ + + model_cpu_offload_seq = "text_encoder->connectors->transformer" + + def __init__( + self, + transformer, + connectors, + tokenizer, + text_encoder, + scheduler: FlowMatchEulerDiscreteScheduler, + text_max_sequence_length: int = 1024, + ) -> None: + super().__init__() + self.register_modules( + transformer=transformer, + connectors=connectors, + tokenizer=tokenizer, + text_encoder=text_encoder, + scheduler=scheduler, + ) + self.register_to_config(text_max_sequence_length=int(text_max_sequence_length)) + self.text_max_sequence_length = int(text_max_sequence_length) + + # ------------------------------------------------------------------ + # forward + # ------------------------------------------------------------------ + + @torch.no_grad() + def __call__( + self, + sana_latent: torch.Tensor, + prompt: str, + *, + fps: float, + sink_size: int = 1, + seed: int = 42, + progress: bool = True, + block_size: int = 3, + kv_max_frames: int = 11, + sigmas: tuple[float, ...] = STAGE_2_DISTILLED_SIGMA_VALUES, + device: str | torch.device | None = None, + ) -> torch.Tensor: + """Run the LTX-2 refiner and return refined VAE latents. + + Uses the chunk-causal AR recipe the model was trained on (``block_size=3``, ``kv_max_frames=11``): a sliding + window of ``[source_sink + recent_history + active_block]`` K/V is fed to the transformer one block at a time, + so per-block compute is bounded and total refinement cost scales linearly with video length. + + Args: + sana_latent: ``(B, C, F, H, W)`` stage-1 latent. + prompt: text prompt. + fps: video frame rate (drives LTX-2 RoPE temporal scaling). + sink_size: how many leading raw ``z_sana`` frames to anchor as the + attention sink (canonical: 1). + seed: noise seed for the FM endpoint. + progress: show a tqdm bar. + block_size: latent frames per AR block (canonical: 3). + kv_max_frames: maximum context+active frames retained in the + sliding window (canonical: 11 = 1 sink + 10 recent). + sigmas: descending Euler schedule terminating at 0.0 (canonical + 3-step distilled: ``(0.909375, 0.725, 0.421875, 0.0)``). Fed to ``self.scheduler`` (minus the trailing + 0.0, which the scheduler appends itself). + device: execution device for the refiner's sub-modules. If ``None``, falls back to where the transformer + currently lives. The refiner moves each sub-module on/off this device as it runs. + + Returns: + `torch.Tensor`: Refined VAE latents of shape ``(B, C, F, H, W)`` — the first ``sink_size`` frames carry the + raw stage-1 sink latents unchanged, the rest carry the refined output. + """ + if sana_latent.shape[2] <= sink_size: + raise ValueError(f"Stage-1 latent has {sana_latent.shape[2]} frames but sink_size={sink_size}.") + + dtype = next(self.transformer.parameters()).dtype + # The refiner moves its own sub-modules on/off ``device`` as it runs (so + # peak VRAM ~= the largest single sub-model, not the sum). Callers pass + # the execution device explicitly; otherwise fall back to where the + # transformer currently lives. + if device is None: + device = next(self.transformer.parameters()).device + device = torch.device(device) + + # Load the distilled sigma schedule into the scheduler. Drop the trailing + # 0.0 — ``FlowMatchEulerDiscreteScheduler.set_timesteps`` appends the + # terminal 0.0 itself, so ``self.scheduler.sigmas`` reproduces ``sigmas``. + self.scheduler.set_timesteps(sigmas=list(sigmas[:-1]), device=device) + sigmas_t = self.scheduler.sigmas.to(device=device, dtype=torch.float32) + + # Free transformer GPU memory while we run the text encoder. + self.transformer.to("cpu") + empty_device_cache(device.type) + prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt, device=device, dtype=dtype) + + self.transformer.to(device) + z = sana_latent.to(device=device, dtype=dtype) + + return self._refine_latents_ar( + z=z, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + fps=fps, + sigmas=sigmas_t, + source_sink_frames=int(sink_size), + block_size=int(block_size), + kv_max_frames=int(kv_max_frames), + seed=int(seed), + progress=bool(progress), + dtype=dtype, + device=device, + ) + + def _refine_latents_ar( + self, + *, + z: torch.Tensor, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + fps: float, + sigmas: torch.Tensor, + source_sink_frames: int, + block_size: int, + kv_max_frames: int, + seed: int, + progress: bool, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Chunk-causal AR refinement — thin wrapper around ``_RefinerChunkRunner``. + + Implements the canonical ``rf_shifted_sink`` KV-cache contract end-to-end: + + 1. Pre-capture **pre-RoPE** sink K/V from raw ``z_sana[:source_sink_frames]`` at σ=0. The sink frames + themselves are **never refined** — they sit unchanged in the output volume. + 2. AR blocks cover frames ``[source_sink_frames, T_full)`` in ``block_size``-frame chunks. For each block: + - Initialize ``x_t = (1-σ₀)·z_sana_block + σ₀·ε`` (single eps per block). + - 3-step deterministic Euler. Each step injects the per-layer prefix ``{sink_k_pre, sink_v, sink_pe, + history_k, history_v}`` where ``sink_pe`` is rebuilt at ``sink_rope_offset = active_start - history_frames + - source_sink_frames`` so the sink slides to sit immediately before the bounded working cache. + - Capture **post-RoPE** K/V from the refined block under the same prefix; append to ``history_kv_post`` and + trim to ``kv_max_frames - source_sink_frames``. + + The returned tensor has the same shape ``(B, C, T_full, H, W)`` as ``z``; the first ``source_sink_frames`` + slots carry the raw sink latents unchanged, the rest carry the refined output. + """ + runner = _RefinerChunkRunner( + self, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + fps=fps, + sigmas=sigmas, + source_sink_frames=int(source_sink_frames), + block_size=int(block_size), + kv_max_frames=int(kv_max_frames), + seed=int(seed), + spatial_shape=(int(z.shape[3]), int(z.shape[4])), + dtype=dtype, + device=device, + ) + + T_full = z.shape[2] + sink_size = int(source_sink_frames) + # Output keeps the raw sink prefix verbatim; AR blocks fill frames + # [sink_size, T_full). + output = z.clone() + n_active = max(T_full - sink_size, 0) + n_blocks = (n_active + block_size - 1) // block_size if n_active > 0 else 0 + + iterator = range(n_blocks) + if progress: + iterator = tqdm(iterator, desc="refiner-ar", unit="block", total=n_blocks) + + for block_idx in iterator: + block_start = sink_size + block_idx * block_size + block_end = min(block_start + block_size, T_full) + clean_block = z[:, :, block_start:block_end] + refined = runner.refine_block( + block_idx=block_idx, + clean_block=clean_block, + block_start=block_start, + block_end=block_end, + sink_seed_frames=(z[:, :, :sink_size] if block_idx == 0 else None), + ) + output[:, :, block_start:block_end] = refined + + return output + + def _predict_x0_active_block( + self, + *, + active: torch.Tensor, + active_positions: list[int], + sigma_cur: float, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + fps: float, + kv_prefix_per_layer: list[dict[str, object]] | None, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Forward through the transformer on the active block only and return x0. + + The active block's Q attends to ``[prefix, current]`` K/V via the ``_tf_kv_prefix`` hook on every + self-attention block. All active tokens carry the same ``sigma_cur``. + """ + latent_tokens = _pack_latents( + active, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + batch_size, seq_len, _ = latent_tokens.shape + timestep_scalar = float(sigma_cur) * float(self.transformer.config.timestep_scale_multiplier) + model_timestep = torch.full((batch_size, seq_len), timestep_scalar, dtype=torch.float32, device=device) + + video_rotary_emb = _build_rotary_emb_for_absolute_positions( + transformer=self.transformer, + batch_size=batch_size, + frame_positions=active_positions, + height=int(active.shape[3]), + width=int(active.shape[4]), + device=device, + fps=float(fps), + ) + + _set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer) + try: + velocity = self._forward_video_only_with_rope( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=model_timestep, + encoder_attention_mask=prompt_attention_mask, + video_rotary_emb=video_rotary_emb, + n_context_tokens=0, + ) + finally: + _clear_kv_prefix_on_blocks(self.transformer) + + # FM x0 prediction: x_t - σ_cur · v. + raw_sigma = torch.full((batch_size, seq_len, 1), float(sigma_cur), dtype=torch.float32, device=device) + denoised_tokens = latent_tokens.float() - velocity.float() * raw_sigma + return _unpack_latents( + denoised_tokens.to(dtype), + num_frames=int(active.shape[2]), + height=int(active.shape[3]), + width=int(active.shape[4]), + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + + def _capture_block_kv( + self, + *, + clean_block: torch.Tensor, + frame_positions: list[int], + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + fps: float, + capture_mode: str, + kv_prefix_per_layer: list[dict[str, object]] | None, + device: torch.device, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Run one forward at σ=0 with capture hooks; return per-layer (K, V). + + ``capture_mode='pre_rope'`` saves PRE-RoPE K/V (so a future window can re-RoPE the sink to its shifted offset). + ``capture_mode='post_rope'`` saves POST-RoPE K/V (ready to concatenate directly into the next window's prefix). + """ + latent_tokens = _pack_latents( + clean_block, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + batch_size, seq_len, _ = latent_tokens.shape + model_timestep = torch.zeros(batch_size, seq_len, dtype=torch.float32, device=device) + + video_rotary_emb = _build_rotary_emb_for_absolute_positions( + transformer=self.transformer, + batch_size=batch_size, + frame_positions=frame_positions, + height=int(clean_block.shape[3]), + width=int(clean_block.shape[4]), + device=device, + fps=float(fps), + ) + + _set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer) + _set_capture_flag_on_blocks(self.transformer, capture_mode, enable=True) + try: + _ = self._forward_video_only_with_rope( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=model_timestep, + encoder_attention_mask=prompt_attention_mask, + video_rotary_emb=video_rotary_emb, + n_context_tokens=0, + ) + finally: + _set_capture_flag_on_blocks(self.transformer, capture_mode, enable=False) + _clear_kv_prefix_on_blocks(self.transformer) + + return _collect_captured_kv_from_blocks(self.transformer, capture_mode) + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + + def _encode_prompt( + self, prompt: str, *, device: torch.device, dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + tokenizer = self.tokenizer + tokenizer.padding_side = "left" + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + text_inputs = tokenizer( + [prompt.strip()], + padding="max_length", + max_length=self.text_max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + input_ids = text_inputs.input_ids.to(device) + attention_mask = text_inputs.attention_mask.to(device) + + self.text_encoder.to(device) + text_backbone = getattr(self.text_encoder, "model", self.text_encoder) + outputs = text_backbone(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) + hidden_states = torch.stack(outputs.hidden_states, dim=-1) + sequence_lengths = attention_mask.sum(dim=-1) + prompt_embeds = _pack_text_embeds( + hidden_states, + sequence_lengths, + device=device, + padding_side=tokenizer.padding_side, + ).to(dtype=dtype) + + # Release the text encoder once we have the prompt embeds — otherwise it + # stays resident on GPU through the entire (much longer) AR refinement. + self.text_encoder.to("cpu") + del outputs, hidden_states + empty_device_cache(device.type) + + self.connectors.to(device) + connector_prompt_embeds, _, connector_attention_mask = self.connectors(prompt_embeds, attention_mask) + self.connectors.to("cpu") + del prompt_embeds, attention_mask + empty_device_cache(device.type) + + return ( + connector_prompt_embeds.to(device=device, dtype=dtype), + connector_attention_mask.to(device=device), + ) + + def _forward_video_only_with_rope( + self, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_attention_mask: torch.Tensor | None, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor], + n_context_tokens: int, + ) -> torch.Tensor: + """Shared body of ``_forward_video_only`` that takes a pre-built RoPE. + + Used by the AR refinement path where each block forward needs custom per-frame absolute positions in the source + video. + """ + transformer = self.transformer + batch_size = hidden_states.size(0) + + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: + encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + hidden_states = transformer.proj_in(hidden_states) + temb, embedded_timestep = transformer.time_embed( + timestep.flatten(), + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(batch_size, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) + + encoder_hidden_states = transformer.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) + + for block in transformer.transformer_blocks: + hidden_states = _forward_video_block( + block=block, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + n_context_tokens=n_context_tokens, + ) + + scale_shift_values = transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None] + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + hidden_states = transformer.norm_out(hidden_states) + hidden_states = hidden_states * (1 + scale) + shift + return transformer.proj_out(hidden_states) + + +class _RefinerChunkRunner: + """Stateful per-AR-block driver for :class:`SanaWMLTX2Refiner`. + + Owns the rolling KV state that the chunk-causal AR recipe accumulates as refiner blocks complete: + + * ``_sink_kv_pre``: per-layer pre-RoPE K/V captured from the first ``source_sink_frames`` raw stage-1 latents at + σ=0. Lazily filled on the first call to :meth:`refine_block`. + * ``_history_kv_post``: per-layer post-RoPE K/V of every refined block already produced, trimmed to ``kv_max_frames + - source_sink_frames`` frames so the sliding window stays bounded. + * ``_history_frames``: number of frames currently in ``_history_kv_post``. + """ + + def __init__( + self, + refiner: SanaWMLTX2Refiner, + *, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + fps: float, + sigmas: torch.Tensor, + source_sink_frames: int, + block_size: int, + kv_max_frames: int, + seed: int, + spatial_shape: tuple[int, int], + dtype: torch.dtype, + device: torch.device, + ) -> None: + self._refiner = refiner + self._prompt_embeds = prompt_embeds + self._prompt_attention_mask = prompt_attention_mask + self._fps = float(fps) + self._sigmas = sigmas + self._sigma_max = float(sigmas[0]) + self._n_steps = int(sigmas.numel() - 1) + self._source_sink_frames = int(source_sink_frames) + self._block_size = int(block_size) + self._max_history_frames = int(kv_max_frames) - int(source_sink_frames) + self._device = device + self._dtype = dtype + self._generator = torch.Generator(device=self._device).manual_seed(int(seed)) + + transformer = refiner.transformer + self._n_layers = len(transformer.transformer_blocks) + H, W = spatial_shape + self._H, self._W = int(H), int(W) + # ``_pack_latents`` emits ``(T // patch_size_t) * (H // p) * (W // p)`` tokens, + # so a single latent frame contributes ``(H // p) * (W // p) / patch_size_t`` + # tokens. (No-op for LTX-2, which uses ``patch_size_t=1``.) + self._tokens_per_frame = ( + int(H // transformer.config.patch_size) + * int(W // transformer.config.patch_size) + // int(transformer.config.patch_size_t) + ) + + self._sink_kv_pre: list[tuple[torch.Tensor, torch.Tensor]] | None = None + self._history_kv_post: list[tuple[torch.Tensor, torch.Tensor] | None] = [None] * self._n_layers + self._history_frames: int = 0 + + def refine_block( + self, + *, + block_idx: int, + clean_block: torch.Tensor, + block_start: int, + block_end: int, + sink_seed_frames: torch.Tensor | None = None, + ) -> torch.Tensor: + """Refine one AR block; advance internal KV state. + + Args: + block_idx: 0-based block index in the AR schedule. + clean_block: ``(B, C, active_len, H, W)`` clean stage-1 latents + covering frames ``[block_start, block_end)``. + block_start: absolute latent-frame index of the active block's + first frame (drives the ``rf_shifted_sink`` RoPE offset). Must be >= ``source_sink_frames``. + block_end: absolute latent-frame index just past the active block. + sink_seed_frames: ``(B, C, source_sink_frames, H, W)`` raw sink + latents used once on the first call to pre-capture the pre-RoPE sink K/V at ``sigma=0`` with frame + positions ``[0, source_sink_frames)``. + """ + refiner = self._refiner + device = self._device + B = int(clean_block.shape[0]) + active_len = block_end - block_start + if block_start < self._source_sink_frames: + raise ValueError( + f"block_start={block_start} overlaps the source sink (source_sink_frames={self._source_sink_frames})." + ) + + # 1) On the first call: pre-capture PRE-RoPE sink K/V from the supplied + # raw sink latents at sigma=0 with absolute positions [0, sink_size). + if self._sink_kv_pre is None: + if sink_seed_frames is None: + raise ValueError("First refine_block call requires sink_seed_frames (raw stage-1 sink latents).") + if sink_seed_frames.shape[2] != self._source_sink_frames: + raise ValueError( + f"sink_seed_frames has {sink_seed_frames.shape[2]} frames " + f"but source_sink_frames={self._source_sink_frames}." + ) + source_sink = sink_seed_frames.contiguous() + self._sink_kv_pre = refiner._capture_block_kv( + clean_block=source_sink, + frame_positions=list(range(self._source_sink_frames)), + prompt_embeds=self._prompt_embeds, + prompt_attention_mask=self._prompt_attention_mask, + fps=self._fps, + capture_mode="pre_rope", + kv_prefix_per_layer=None, + device=device, + ) + + # 2) Build per-window kv_prefix dict per layer. + sink_rope_offset = block_start - self._history_frames - self._source_sink_frames + sink_pe = _build_rotary_emb_for_absolute_positions( + transformer=refiner.transformer, + batch_size=B, + frame_positions=list(range(sink_rope_offset, sink_rope_offset + self._source_sink_frames)), + height=self._H, + width=self._W, + device=device, + fps=self._fps, + ) + kv_prefix_per_layer: list[dict[str, object]] = [] + for layer_idx in range(self._n_layers): + hk = self._history_kv_post[layer_idx] + kv_prefix_per_layer.append( + { + "mode": "rf_shifted_sink", + "sink_k_pre": self._sink_kv_pre[layer_idx][0], + "sink_v": self._sink_kv_pre[layer_idx][1], + "sink_pe": sink_pe, + "history_k": (hk[0] if hk is not None else None), + "history_v": (hk[1] if hk is not None else None), + } + ) + + # 3) FM endpoint at sigma=sigma0: single epsilon per block. + eps = randn_tensor(clean_block.shape, generator=self._generator, device=device, dtype=self._dtype) + x_t = ((1.0 - self._sigma_max) * clean_block.float() + self._sigma_max * eps.float()).to(self._dtype) + + # Reset the shared scheduler to step 0 for this block's Euler run (blocks + # are processed sequentially, so re-seeding the schedule per block is safe). + scheduler = refiner.scheduler + scheduler.set_timesteps(sigmas=[float(s) for s in self._sigmas[:-1]], device=device) + timesteps = scheduler.timesteps + + active_positions = list(range(int(block_start), int(block_end))) + for level, t in enumerate(timesteps): + sigma_cur = float(self._sigmas[level].item()) + pred_x0 = refiner._predict_x0_active_block( + active=x_t, + active_positions=active_positions, + sigma_cur=sigma_cur, + prompt_embeds=self._prompt_embeds, + prompt_attention_mask=self._prompt_attention_mask, + fps=self._fps, + kv_prefix_per_layer=kv_prefix_per_layer, + dtype=self._dtype, + device=device, + ) + if sigma_cur <= 1.0e-6: + x_t = pred_x0.to(self._dtype) + else: + # FM velocity from x0; the scheduler applies the Euler update. + velocity = (x_t.float() - pred_x0.float()) / sigma_cur + x_t = scheduler.step(velocity, t, x_t.float(), return_dict=False)[0].to(self._dtype) + + # 4) Capture POST-RoPE K/V for this refined block under the same prefix. + block_kv_post = refiner._capture_block_kv( + clean_block=x_t, + frame_positions=active_positions, + prompt_embeds=self._prompt_embeds, + prompt_attention_mask=self._prompt_attention_mask, + fps=self._fps, + capture_mode="post_rope", + kv_prefix_per_layer=kv_prefix_per_layer, + device=device, + ) + for layer_idx in range(self._n_layers): + new_k, new_v = block_kv_post[layer_idx] + old = self._history_kv_post[layer_idx] + if old is None: + self._history_kv_post[layer_idx] = (new_k, new_v) + else: + self._history_kv_post[layer_idx] = ( + torch.cat([old[0], new_k], dim=1), + torch.cat([old[1], new_v], dim=1), + ) + self._history_frames += active_len + + if self._max_history_frames > 0 and self._history_frames > self._max_history_frames: + keep_tokens = self._max_history_frames * self._tokens_per_frame + for layer_idx in range(self._n_layers): + hk = self._history_kv_post[layer_idx] + if hk is not None: + self._history_kv_post[layer_idx] = (hk[0][:, -keep_tokens:], hk[1][:, -keep_tokens:]) + self._history_frames = self._max_history_frames + + return x_t + + +# ------------------------------------------------------------------------- +# private helpers (block + attention + packing) +# ------------------------------------------------------------------------- + + +def _build_rotary_emb_for_absolute_positions( + *, + transformer: nn.Module, + batch_size: int, + frame_positions: list[int], + height: int, + width: int, + device: torch.device, + fps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reimplement ``LTX2VideoRotaryPosEmbed.prepare_video_coords`` with explicit per-frame positions. + + The default helper assumes contiguous ``torch.arange(num_frames)`` which is fine for bidirectional inference; the + sliding-window AR refiner needs to keep each frame's absolute index in the source video so RoPE captures the + correct temporal phase across the sink + recent + active window. + """ + rope = transformer.rope + patch_size_t = int(rope.patch_size_t) + patch_size = int(rope.patch_size) + f_positions = torch.tensor(frame_positions, dtype=torch.float32, device=device) + if patch_size_t > 1: + # Each patch covers ``patch_size_t`` latent frames; pick the start of each patch. + f_positions = f_positions[::patch_size_t] + grid_h = torch.arange(start=0, end=height, step=patch_size, dtype=torch.float32, device=device) + grid_w = torch.arange(start=0, end=width, step=patch_size, dtype=torch.float32, device=device) + grid = torch.meshgrid(f_positions, grid_h, grid_w, indexing="ij") + grid = torch.stack(grid, dim=0) + + patch_size_delta = torch.tensor((patch_size_t, patch_size, patch_size), dtype=grid.dtype, device=device) + patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) + latent_coords = torch.stack([grid, patch_ends], dim=-1) + latent_coords = latent_coords.flatten(1, 3).unsqueeze(0).repeat(batch_size, 1, 1, 1) + + scale_tensor = torch.tensor(rope.scale_factors, device=device) + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 + pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) + pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + rope.causal_offset - rope.scale_factors[0]).clamp(min=0) + pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / float(fps) + return rope(pixel_coords, device=device) + + +def _forward_video_block( + *, + block: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor], + encoder_attention_mask: torch.Tensor | None, + n_context_tokens: int, +) -> torch.Tensor: + batch_size = hidden_states.size(0) + + norm_hidden_states = block.norm1(hidden_states) + num_ada_params = block.scale_shift_table.shape[0] + ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch_size, temb.size(1), num_ada_params, -1 + ) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + + attn_hidden_states = _streaming_self_attention( + attn=block.attn1, + hidden_states=norm_hidden_states, + query_rotary_emb=video_rotary_emb, + n_context_tokens=n_context_tokens, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa + + norm_hidden_states = block.norm2(hidden_states) + attn_hidden_states = block.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + query_rotary_emb=None, + attention_mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states + + norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp + hidden_states = hidden_states + block.ff(norm_hidden_states) * gate_mlp + return hidden_states + + +def _streaming_self_attention( + *, + attn: nn.Module, + hidden_states: torch.Tensor, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor], + n_context_tokens: int, +) -> torch.Tensor: + """LTX-2 self-attention with sink/current streaming mask + AR KV-cache hooks. + + Two modes layered on top of vanilla diffusers self-attention, selected by ``n_context_tokens`` and per-block hook + attributes (set by the AR refiner): + + * ``n_context_tokens > 0`` (legacy single-shot path): sink queries attend sink only, current queries attend ``[sink + + current]`` via two SDPA calls. + + * ``n_context_tokens == 0`` (AR mode): Q comes from the active block only; the per-block ``_tf_kv_prefix`` dict + (``rf_shifted_sink``) supplies the pre-RoPE sink K/V (re-RoPE'd here with its sliding offset PE) and the + post-RoPE recent-history K/V, concatenated before SDPA. The ``_kv_cache_capture`` and ``_tf_capture_kv`` hooks + record K/V into the module for the AR orchestrator to read back. + """ + from ...models.attention_dispatch import dispatch_attention_fn # noqa: PLC0415 + from ...models.transformers.transformer_ltx2 import ( # noqa: PLC0415 + apply_interleaved_rotary_emb, + apply_split_rotary_emb, + ) + + gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + # KV-cache capture / inject hooks for ``rf_shifted_sink`` AR refinement: + # - ``_kv_cache_capture`` saves PRE-RoPE (post-norm) K/V so a future window + # can re-apply RoPE at its shifted sink offset. + # - ``_tf_capture_kv`` saves POST-RoPE K/V so the next window can directly + # concatenate the recent history. + # - ``_tf_kv_prefix`` (a dict with ``mode='rf_shifted_sink'``) prepends a + # re-RoPE'd sink + already-post-RoPE recent history before SDPA. + if getattr(attn, "_kv_cache_capture", False): + attn._cached_kv_pre = (key.detach().clone(), value.detach().clone()) + + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb(key, query_rotary_emb) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb(key, query_rotary_emb) + else: + raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") + + if getattr(attn, "_tf_capture_kv", False): + attn._cached_kv_post = (key.detach().clone(), value.detach().clone()) + + tf_prefix = getattr(attn, "_tf_kv_prefix", None) + if isinstance(tf_prefix, dict) and tf_prefix.get("mode") == "rf_shifted_sink": + prefix_k_parts: list[torch.Tensor] = [] + prefix_v_parts: list[torch.Tensor] = [] + sink_k_pre = tf_prefix.get("sink_k_pre") + sink_v = tf_prefix.get("sink_v") + if sink_k_pre is not None and sink_v is not None and sink_k_pre.shape[1] > 0: + sink_pe = tf_prefix.get("sink_pe") + if sink_pe is None: + raise RuntimeError("rf_shifted_sink prefix requires a sink_pe RoPE tuple.") + sink_k_pre_dt = sink_k_pre.to(key.dtype) + if attn.rope_type == "interleaved": + sink_k = apply_interleaved_rotary_emb(sink_k_pre_dt, sink_pe) + else: + sink_k = apply_split_rotary_emb(sink_k_pre_dt, sink_pe) + prefix_k_parts.append(sink_k) + prefix_v_parts.append(sink_v.to(value.dtype)) + history_k = tf_prefix.get("history_k") + history_v = tf_prefix.get("history_v") + if history_k is not None and history_v is not None and history_k.shape[1] > 0: + prefix_k_parts.append(history_k.to(key.dtype)) + prefix_v_parts.append(history_v.to(value.dtype)) + if prefix_k_parts: + key = torch.cat([*prefix_k_parts, key], dim=1) + value = torch.cat([*prefix_v_parts, value], dim=1) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + processor = attn.processor + backend = getattr(processor, "_attention_backend", None) + parallel_config = getattr(processor, "_parallel_config", None) + + # AR mode (n_context_tokens == 0): Q from active block attends to the + # injected prefix + current K/V in one SDPA call. Legacy single-shot + # mode keeps the sink-self / current-cross split. + if n_context_tokens <= 0 or n_context_tokens >= query.shape[1]: + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + else: + context_hidden_states = dispatch_attention_fn( + query[:, :n_context_tokens], + key[:, :n_context_tokens], + value[:, :n_context_tokens], + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + current_hidden_states = dispatch_attention_fn( + query[:, n_context_tokens:], + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + hidden_states = torch.cat([context_hidden_states, current_hidden_states], dim=1) + + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + if gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) + gates = 2.0 * torch.sigmoid(gate_logits) + hidden_states = hidden_states * gates.unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +def _set_kv_prefix_on_blocks( + transformer: nn.Module, + kv_prefix_per_layer: list[dict[str, object]] | None, +) -> None: + """Attach a per-layer KV-prefix dict to each ``attn1`` for the AR refiner.""" + blocks = transformer.transformer_blocks + if kv_prefix_per_layer is None: + _clear_kv_prefix_on_blocks(transformer) + return + if len(kv_prefix_per_layer) != len(blocks): + raise RuntimeError( + f"kv_prefix_per_layer has {len(kv_prefix_per_layer)} entries but transformer has {len(blocks)} blocks." + ) + for block, prefix in zip(blocks, kv_prefix_per_layer): + block.attn1._tf_kv_prefix = prefix + + +def _clear_kv_prefix_on_blocks(transformer: nn.Module) -> None: + for block in transformer.transformer_blocks: + block.attn1._tf_kv_prefix = None + + +def _set_capture_flag_on_blocks(transformer: nn.Module, mode: str, *, enable: bool) -> None: + """Toggle ``_kv_cache_capture`` (pre-RoPE) or ``_tf_capture_kv`` (post-RoPE) per block.""" + if mode == "pre_rope": + attr = "_kv_cache_capture" + clear_attr = "_cached_kv_pre" + elif mode == "post_rope": + attr = "_tf_capture_kv" + clear_attr = "_cached_kv_post" + else: + raise ValueError(f"capture_mode must be 'pre_rope' or 'post_rope', got {mode!r}") + for block in transformer.transformer_blocks: + setattr(block.attn1, attr, bool(enable)) + if enable and hasattr(block.attn1, clear_attr): + setattr(block.attn1, clear_attr, None) + + +def _collect_captured_kv_from_blocks( + transformer: nn.Module, + mode: str, +) -> list[tuple[torch.Tensor, torch.Tensor]]: + attr = "_cached_kv_pre" if mode == "pre_rope" else "_cached_kv_post" + out: list[tuple[torch.Tensor, torch.Tensor]] = [] + for block in transformer.transformer_blocks: + cached = getattr(block.attn1, attr, None) + if cached is None: + raise RuntimeError(f"Expected {attr!r} on attn1 after capture forward, but found None.") + out.append(cached) + # Release the reference so the orchestrator owns the only handle. + setattr(block.attn1, attr, None) + return out + + +def _pack_text_embeds( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + device: str | torch.device, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, +) -> torch.Tensor: + batch_size, seq_len, hidden_dim, _ = text_hidden_states.shape + original_dtype = text_hidden_states.dtype + + token_indices = torch.arange(seq_len, device=device).unsqueeze(0) + if padding_side == "right": + mask = token_indices < sequence_lengths[:, None] + elif padding_side == "left": + start_indices = seq_len - sequence_lengths[:, None] + mask = token_indices >= start_indices + else: + raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") + mask = mask[:, :, None, None] + + masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) + num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) + + x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) + x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) + + normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized_hidden_states = normalized_hidden_states * scale_factor + normalized_hidden_states = normalized_hidden_states.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, normalized_hidden_states.shape[-1]) + normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) + return normalized_hidden_states.to(dtype=original_dtype) + + +def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + batch_size, _, num_frames, height, width = latents.shape + latents = latents.reshape( + batch_size, + -1, + num_frames // patch_size_t, + patch_size_t, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + + +def _unpack_latents( + latents: torch.Tensor, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, +) -> torch.Tensor: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index f34008252ab6..76e51ce26f22 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2100,6 +2100,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class SanaWMTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class SD3ControlNetModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 376596d632ea..c61fc58418fd 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3887,6 +3887,51 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class SanaWMLTX2Refiner(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class SanaWMPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class SanaWMPipelineOutput(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class SemanticStableDiffusionPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/pipelines/sana_wm/__init__.py b/tests/pipelines/sana_wm/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/sana_wm/test_sana_wm.py b/tests/pipelines/sana_wm/test_sana_wm.py new file mode 100644 index 000000000000..9465730d6282 --- /dev/null +++ b/tests/pipelines/sana_wm/test_sana_wm.py @@ -0,0 +1,253 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# 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 +# +# http://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. + +"""SANA-WM CPU unit tests. + +Covers the standalone helpers (action DSL, intrinsics math, resize-and-crop), +the public-surface registration, and the Triton -> pure-PyTorch attention +fallback. +""" + +import unittest + +import numpy as np +from PIL import Image + +from diffusers import SanaWMPipeline, SanaWMPipelineOutput +from diffusers.pipelines.sana_wm import SanaWMLTX2Refiner +from diffusers.pipelines.sana_wm.cam_utils import ( + TARGET_HEIGHT, + TARGET_WIDTH, + action_string_to_c2w, + resize_and_center_crop, + snap_num_frames, + transform_intrinsics_for_crop, +) + + +class SanaWMCamUtilsTests(unittest.TestCase): + """Pure-numpy/PIL helpers — no torch.cuda required.""" + + def test_action_dsl_forward_only(self): + c2w = action_string_to_c2w("w-5", translation_speed=0.1) + # 5 action frames + leading identity = 6 total + self.assertEqual(c2w.shape, (6, 4, 4)) + self.assertEqual(c2w.dtype, np.float32) + # First frame is identity (the anchor). + np.testing.assert_allclose(c2w[0], np.eye(4, dtype=np.float32), atol=1e-6) + # 'w' moves forward (+Z in OpenCV convention). + self.assertAlmostEqual(float(c2w[-1, 2, 3]), 0.5, places=5) + # No yaw / pitch -> rotation is identity throughout. + for i in range(c2w.shape[0]): + np.testing.assert_allclose(c2w[i, :3, :3], np.eye(3), atol=1e-6) + + def test_action_dsl_concat_segments(self): + c2w = action_string_to_c2w("w-3,a-2", translation_speed=0.1) + self.assertEqual(c2w.shape, (6, 4, 4)) # 3 + 2 + identity anchor + + def test_action_dsl_rejects_bad_input(self): + with self.assertRaises(ValueError): + action_string_to_c2w("") + with self.assertRaises(ValueError): + action_string_to_c2w("x-5") # 'x' is not in WASD/IJKL + with self.assertRaises(ValueError): + action_string_to_c2w("w-0") # zero-length segment + + def test_action_dsl_none_segment_is_idle(self): + c2w = action_string_to_c2w("none-3", translation_speed=0.1) + self.assertEqual(c2w.shape, (4, 4, 4)) + # No motion -> all frames are identity. + for i in range(c2w.shape[0]): + np.testing.assert_allclose(c2w[i], np.eye(4), atol=1e-6) + + def test_transform_intrinsics_for_crop_scalar(self): + # (fx, fy, cx, cy) for a 1000x500 source, resized to 1280x704, then + # center-cropped to 1280x704 (no extra crop offset). + intr = np.array([800.0, 800.0, 500.0, 250.0], dtype=np.float32) + out = transform_intrinsics_for_crop(intr, src_size=(1000, 500), resized_size=(1280, 704), crop_offset=(0, 0)) + self.assertAlmostEqual(float(out[0]), 800.0 * 1280 / 1000, places=4) # fx scales with x + self.assertAlmostEqual(float(out[1]), 800.0 * 704 / 500, places=4) + self.assertAlmostEqual(float(out[2]), 500.0 * 1280 / 1000, places=4) + self.assertAlmostEqual(float(out[3]), 250.0 * 704 / 500, places=4) + + def test_transform_intrinsics_for_crop_with_offset(self): + intr = np.array([800.0, 800.0, 500.0, 250.0], dtype=np.float32) + # After resize, an extra crop offset shifts the principal point. + out = transform_intrinsics_for_crop( + intr, src_size=(1000, 500), resized_size=(2000, 1000), crop_offset=(360, 148) + ) + self.assertAlmostEqual(float(out[2]), 500.0 * 2.0 - 360.0, places=4) + self.assertAlmostEqual(float(out[3]), 250.0 * 2.0 - 148.0, places=4) + + def test_resize_and_center_crop_default_target(self): + src = Image.new("RGB", (1691, 930)) + cropped, src_size, resized_size, crop_offset = resize_and_center_crop(src) + self.assertEqual(cropped.size, (TARGET_WIDTH, TARGET_HEIGHT)) + self.assertEqual(src_size, (1691, 930)) + # Resize preserves aspect; one of the resized dimensions equals the target. + rw, rh = resized_size + self.assertTrue(rw >= TARGET_WIDTH and rh >= TARGET_HEIGHT) + cl, ct = crop_offset + self.assertGreaterEqual(cl, 0) + self.assertGreaterEqual(ct, 0) + # Center crop produces 0 offset on the dimension that hit the target exactly. + self.assertTrue(cl == 0 or ct == 0) + + def test_snap_num_frames_to_8k_plus_1(self): + # The LTX-2 VAE requires (8k + 1)-shaped temporal dim. ``snap_num_frames`` + # rounds to the nearest such value (ties break to the ceil). + for n in [1, 9, 17, 81, 161, 321, 801]: + self.assertEqual(snap_num_frames(n), n) + self.assertEqual(snap_num_frames(2), 1) + self.assertEqual(snap_num_frames(10), 9) # 10 is closer to 9 than 17 + self.assertEqual(snap_num_frames(80), 81) # 80 is closer to 81 than 73 + self.assertEqual(snap_num_frames(100), 97) # 100 is closer to 97 than 105 + # ``upper_bound`` caps the result (the snap falls back to the floor). + self.assertLessEqual(snap_num_frames(100, upper_bound=100), 100) + self.assertEqual(snap_num_frames(100, upper_bound=100), 97) + + +class SanaWMRegistrationTests(unittest.TestCase): + """Verify the SANA-WM symbols are reachable through the public diffusers surface.""" + + def test_top_level_symbols(self): + import diffusers + + for name in ("SanaWMPipeline", "SanaWMTransformer3DModel", "SanaWMLTX2Refiner", "SanaWMPipelineOutput"): + self.assertTrue(hasattr(diffusers, name), msg=f"{name!r} not exported from diffusers top-level") + + def test_pipeline_output_dataclass(self): + import torch + + frames = np.zeros((3, 8, 8, 3), dtype=np.float32) + c2w = np.broadcast_to(np.eye(4, dtype=np.float32), (3, 4, 4)).copy() + latent = torch.zeros(1, 16, 1, 4, 4) + out = SanaWMPipelineOutput(frames=frames, c2w=c2w, latent=latent) + self.assertEqual(tuple(out.frames.shape), (3, 8, 8, 3)) + self.assertEqual(tuple(out.c2w.shape), (3, 4, 4)) + self.assertEqual(tuple(out.latent.shape), (1, 16, 1, 4, 4)) + + def test_refiner_is_pipeline_with_ar_call_defaults(self): + import inspect + + from diffusers import DiffusionPipeline + + # The refiner is a standalone DiffusionPipeline. + self.assertTrue(issubclass(SanaWMLTX2Refiner, DiffusionPipeline)) + + # Its denoising entry point is ``__call__`` with the canonical AR defaults. + params = inspect.signature(SanaWMLTX2Refiner.__call__).parameters + self.assertIn("block_size", params) + self.assertIn("kv_max_frames", params) + # AR mode is on by default. + self.assertEqual(params["block_size"].default, 3) + self.assertEqual(params["kv_max_frames"].default, 11) + + def test_pipeline_call_intrinsics_signature(self): + import inspect + + params = inspect.signature(SanaWMPipeline.__call__).parameters + self.assertIn("intrinsics", params) + self.assertIn("c2w", params) + self.assertIn("action", params) + self.assertIn("use_refiner", params) + + +class SanaWMTritonFallbackTests(unittest.TestCase): + """When Triton isn't usable, ``*Triton`` attention classes should auto-fall-back + to their non-Triton parents at dispatch time so the model works on CPU / + ROCm-without-Triton without users having to know the variant names. + """ + + def test_kernels_module_imports_with_triton_hidden(self): + # Simulate a Triton-less environment and reload the kernels module from + # scratch — it must still import (definitions of @triton.jit kernels + # become no-op shims) and the pure-torch helpers must still work. + import importlib + import sys + + # Make sure diffusers is loaded first (its loaders module hard-imports triton). + import diffusers # noqa: F401 + + orig_triton = sys.modules.get("triton") + sys.modules["triton"] = None + sys.modules.pop("diffusers.models.transformers.transformer_sana_wm_kernels", None) + try: + kernels = importlib.import_module("diffusers.models.transformers.transformer_sana_wm_kernels") + self.assertFalse(kernels.is_triton_available()) + # Pure-torch helpers must still be callable. + self.assertTrue(callable(kernels.prepare_rope_tables)) + self.assertTrue(callable(kernels.compute_fov_from_fx_xi)) + finally: + sys.modules.pop("diffusers.models.transformers.transformer_sana_wm_kernels", None) + if orig_triton is not None: + sys.modules["triton"] = orig_triton + else: + sys.modules.pop("triton", None) + # Restore the real kernels module for downstream tests. + importlib.import_module("diffusers.models.transformers.transformer_sana_wm_kernels") + + def test_resolve_attention_block_cpu_fallback(self): + # On a CPU-only test host, _is_triton_kernels_usable() returns False and + # ``*Triton`` attn types should resolve to their non-Triton ancestors. + import torch + + from diffusers.models.transformers.transformer_sana_wm import ( + _is_triton_kernels_usable, + _resolve_attention_block, + ) + + if torch.cuda.is_available() and _is_triton_kernels_usable(): + self.skipTest("Triton is usable on this host; fallback path not exercised.") + + expected = { + "BidirectionalGDNTriton": "BidirectionalGDN", + "BidirectionalGDNUCPESinglePathLiteLATriton": "BidirectionalGDNUCPESinglePathLiteLA", + "BidirectionalGDNUCPESinglePathLiteLABothTriton": "BidirectionalGDNUCPESinglePathLiteLA", + # Already non-Triton: should resolve to itself. + "BidirectionalGDN": "BidirectionalGDN", + "BidirectionalGDNUCPESinglePathLiteLA": "BidirectionalGDNUCPESinglePathLiteLA", + } + for requested, expected_name in expected.items(): + cls = _resolve_attention_block(requested, role="attn_type") + self.assertEqual( + cls.__name__, + expected_name, + msg=f"_resolve_attention_block({requested!r}) -> {cls.__name__}, expected {expected_name}", + ) + + def test_triton_entry_point_raises_clean_error_without_triton(self): + # ``_require_triton`` should raise a clear RuntimeError when invoked on + # a Triton-less host (regardless of CUDA availability — the kernels + # need both). + import importlib + import sys + + import diffusers # noqa: F401 + + orig_triton = sys.modules.get("triton") + sys.modules["triton"] = None + sys.modules.pop("diffusers.models.transformers.transformer_sana_wm_kernels", None) + try: + kernels = importlib.import_module("diffusers.models.transformers.transformer_sana_wm_kernels") + with self.assertRaises(RuntimeError) as ctx: + kernels._require_triton("test_entry_point") + self.assertIn("triton", str(ctx.exception).lower()) + finally: + sys.modules.pop("diffusers.models.transformers.transformer_sana_wm_kernels", None) + if orig_triton is not None: + sys.modules["triton"] = orig_triton + else: + sys.modules.pop("triton", None) + importlib.import_module("diffusers.models.transformers.transformer_sana_wm_kernels")