From 63378616a6390a00ac6eff4bc327522df553997b Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 15 Jul 2026 15:15:39 -0700 Subject: [PATCH 01/36] feat(sana): add SANA-WM integration shim Add a workspace integration package for the upstream NVlabs/Sana SANA-WM bidirectional world model. The package registers the `sana-wm-bidirectional` runner, resolves upstream Sana assets and Hugging Face checkpoints, loads images, prompts, camera trajectories, and intrinsics, then writes generated videos through the upstream pipeline. Port CPU-safe camera trajectory and intrinsics helpers so the integration can be checked without running GPU inference. Add smoke and camera tests covering runner registration, config derivation, action parsing, and geometry transforms. Signed-off-by: Aidan Foster --- integrations/sana/README.md | 77 +++ integrations/sana/pyproject.toml | 50 ++ integrations/sana/sana_wm/__init__.py | 23 + integrations/sana/sana_wm/camera.py | 513 ++++++++++++++++++ integrations/sana/sana_wm/config.py | 56 ++ integrations/sana/sana_wm/constants.py | 49 ++ integrations/sana/sana_wm/placeholder.py | 77 +++ integrations/sana/sana_wm/runner.py | 435 +++++++++++++++ .../sana/tests/parity_check/README.md | 21 + integrations/sana/tests/test_camera.py | 142 +++++ integrations/sana/tests/test_smoke.py | 105 ++++ pyproject.toml | 2 + uv.lock | 21 + 13 files changed, 1571 insertions(+) create mode 100644 integrations/sana/README.md create mode 100644 integrations/sana/pyproject.toml create mode 100644 integrations/sana/sana_wm/__init__.py create mode 100644 integrations/sana/sana_wm/camera.py create mode 100644 integrations/sana/sana_wm/config.py create mode 100644 integrations/sana/sana_wm/constants.py create mode 100644 integrations/sana/sana_wm/placeholder.py create mode 100644 integrations/sana/sana_wm/runner.py create mode 100644 integrations/sana/tests/parity_check/README.md create mode 100644 integrations/sana/tests/test_camera.py create mode 100644 integrations/sana/tests/test_smoke.py diff --git a/integrations/sana/README.md b/integrations/sana/README.md new file mode 100644 index 000000000..88d2c2911 --- /dev/null +++ b/integrations/sana/README.md @@ -0,0 +1,77 @@ + + +# `sana_wm` + +FlashDreams runner shim for +[SANA-WM](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional), +the 2.6B bidirectional camera-controlled world model from the +[NVlabs/Sana](https://github.com/NVlabs/Sana) repository. + +This first slice follows the dependency-isolated integration pattern: +FlashDreams owns the CLI registration, config surface, and CPU-tested +camera/action helpers, while Stage-1 DiT, LTX2 VAE/refiner, Pi3X, and +custom attention kernels are imported from an installed or local upstream +Sana checkout at runtime. + +## Runner + +| slug | description | +| --- | --- | +| `sana-wm-bidirectional` | SANA-WM bidirectional I2V world model using upstream NVlabs/Sana Stage-1 + LTX-2 refiner. | + +The FlashDreams package is named `sana_wm` rather than `sana` so it does +not shadow upstream's own top-level `sana` package. + +## Setup + +Clone Sana next to FlashDreams or pass its path explicitly: + +```bash +git clone https://github.com/NVlabs/Sana ../Sana +bash ../Sana/environment_setup.sh sana +``` + +The runner auto-detects `../Sana`. You can also set `SANA_ROOT` or pass +`--upstream-sana-root /path/to/Sana`. + +## Run + +```bash +uv run flashdreams-run sana-wm-bidirectional \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --action "w-100,dw-60,w-100,aw-60" \ + --num-frames 321 \ + --output-dir outputs/sana_wm +``` + +To use an explicit trajectory: + +```bash +uv run flashdreams-run sana-wm-bidirectional \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 321 \ + --output-dir outputs/sana_wm +``` + +`--no-refiner` skips the LTX-2 refiner and decodes Stage-1 latents +directly, matching upstream's fast debugging path. Full generation still +downloads the public Hugging Face artefacts on first use. + +## Tests + +CPU-safe tests cover import, runner config, action parsing, intrinsics +handling, frame snapping, and SANA-WM camera-conditioning tensor shapes: + +```bash +uv run --extra dev pytest integrations/sana/tests +``` + +Full upstream parity and generated-video checks are heavyweight GPU +workflows and should live under `tests/parity_check/`. diff --git a/integrations/sana/pyproject.toml b/integrations/sana/pyproject.toml new file mode 100644 index 000000000..d0baed354 --- /dev/null +++ b/integrations/sana/pyproject.toml @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-sana-wm" +version = "0.1.0" +description = "SANA-WM bidirectional world-model runner shim for flashdreams." +readme = "README.md" +requires-python = ">=3.10" +# The upstream SANA repository owns the Stage-1 DiT, LTX2 VAE/refiner, +# Pi3X intrinsics estimator, and custom attention kernels. Those heavy +# dependencies are intentionally installed from the upstream SANA checkout, +# not pulled into every FlashDreams workspace sync. +dependencies = [ + "flashdreams", +] + +[tool.uv.sources] +flashdreams = { workspace = true } + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", +] + +[project.entry-points."flashdreams.runner_configs"] +"sana-wm-bidirectional" = "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL" + +[tool.setuptools.packages.find] +include = ["sana_wm*"] +exclude = ["tests"] + +[tool.uv] +managed = true diff --git a/integrations/sana/sana_wm/__init__.py b/integrations/sana/sana_wm/__init__.py new file mode 100644 index 000000000..f4b7dbd28 --- /dev/null +++ b/integrations/sana/sana_wm/__init__.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 bidirectional runner shim for FlashDreams.""" + +from sana_wm.config import RUNNER_CONFIGS, RUNNER_SANA_WM_BIDIRECTIONAL + +__all__ = [ + "RUNNER_CONFIGS", + "RUNNER_SANA_WM_BIDIRECTIONAL", +] diff --git a/integrations/sana/sana_wm/camera.py b/integrations/sana/sana_wm/camera.py new file mode 100644 index 000000000..97daf41ba --- /dev/null +++ b/integrations/sana/sana_wm/camera.py @@ -0,0 +1,513 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 camera-pose DSL, intrinsics, and Plucker conditioning helpers.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_TEMPORAL_COMPRESSION, +) + +FPS = 16 +"""Camera-control integration frame rate.""" + +DEFAULT_TRANSLATION_SPEED = 0.025 +"""Default per-frame translation magnitude used by upstream SANA-WM.""" + +DEFAULT_ROTATION_SPEED_DEG = 0.6 +"""Default per-frame rotation magnitude in degrees.""" + +DEFAULT_PITCH_LIMIT_DEG = 60.0 +"""Maximum absolute camera pitch in degrees.""" + +TAU_PRESS = 0.45 +"""Velocity smoothing time constant when a key is pressed.""" + +TAU_COAST = 1.0 +"""Velocity smoothing time constant when controls are released.""" + +DSL_KEY_TO_CONTROL: dict[str, str] = { + "w": "forward", + "s": "back", + "a": "yaw_left", + "d": "yaw_right", + "i": "pitch_up", + "k": "pitch_down", + "j": "strafe_left", + "l": "strafe_right", +} +"""Mapping from SANA-WM action DSL letters to canonical controls.""" + +ALLOWED_ACTION_KEYS = frozenset(DSL_KEY_TO_CONTROL) +"""Action DSL keys accepted in ``-`` segments.""" + + +@dataclass +class VelocityState: + """Per-frame camera velocity in OpenCV camera coordinates.""" + + tx: float = 0.0 + """Forward translation velocity.""" + + sx: float = 0.0 + """Rightward strafe velocity.""" + + yaw: float = 0.0 + """Positive yaw-right angular velocity in radians.""" + + pitch: float = 0.0 + """Positive pitch-up angular velocity in radians.""" + + def snap_to(self, target: "VelocityState") -> None: + """Copy all velocity components from ``target``.""" + self.tx, self.sx = target.tx, target.sx + self.yaw, self.pitch = target.yaw, target.pitch + + def step_toward(self, target: "VelocityState", dt: float) -> None: + """Ease velocity components toward ``target`` over one timestep.""" + for attr in ("tx", "sx", "yaw", "pitch"): + cur = getattr(self, attr) + tgt = getattr(target, attr) + tau = TAU_PRESS if abs(tgt) > 1e-12 else TAU_COAST + alpha = 1.0 - math.exp(-dt / tau) + setattr(self, attr, cur + alpha * (tgt - cur)) + + +class CameraPoseIntegrator: + """Integrate camera velocity into camera-to-world poses.""" + + def __init__( + self, pitch_limit_rad: float = math.radians(DEFAULT_PITCH_LIMIT_DEG) + ) -> None: + self.pose = np.eye(4, dtype=np.float64) + self.pitch = 0.0 + self.pitch_limit = float(pitch_limit_rad) + + def step(self, velocity: VelocityState) -> np.ndarray: + """Advance one frame and return the current camera-to-world pose.""" + new_pitch = max( + -self.pitch_limit, min(self.pitch_limit, self.pitch + velocity.pitch) + ) + pitch_step = new_pitch - self.pitch + self.pitch = new_pitch + + rotation = self.pose[:3, :3] + rotation_new = _rot_y(velocity.yaw) @ rotation @ _rot_x(pitch_step) + + forward = rotation_new[:, 2].copy() + forward[1] = 0.0 + right = rotation_new[:, 0].copy() + right[1] = 0.0 + forward_norm = float(np.linalg.norm(forward)) + right_norm = float(np.linalg.norm(right)) + if forward_norm > 0: + forward /= forward_norm + 1e-6 + if right_norm > 0: + right /= right_norm + 1e-6 + + translation = self.pose[:3, 3] + forward * velocity.tx + right * velocity.sx + self.pose = np.eye(4, dtype=np.float64) + self.pose[:3, :3] = rotation_new + self.pose[:3, 3] = translation + return self.pose.copy() + + +def _rot_x(angle_rad: float) -> np.ndarray: + c, s = math.cos(angle_rad), math.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 = math.cos(angle_rad), math.sin(angle_rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def controls_to_target_velocity( + controls: set[str], + *, + translation_speed: float = DEFAULT_TRANSLATION_SPEED, + rotation_speed_rad: float | None = None, +) -> VelocityState: + """Map canonical control tokens to a target velocity.""" + if rotation_speed_rad is None: + rotation_speed_rad = math.radians(DEFAULT_ROTATION_SPEED_DEG) + forward = (1.0 if "forward" in controls else 0.0) - ( + 1.0 if "back" in controls else 0.0 + ) + strafe = (1.0 if "strafe_right" in controls else 0.0) - ( + 1.0 if "strafe_left" in controls else 0.0 + ) + yaw = (1.0 if "yaw_right" in controls else 0.0) - ( + 1.0 if "yaw_left" in controls else 0.0 + ) + pitch = (1.0 if "pitch_up" in controls else 0.0) - ( + 1.0 if "pitch_down" in controls else 0.0 + ) + return VelocityState( + tx=forward * translation_speed, + sx=strafe * translation_speed, + yaw=yaw * rotation_speed_rad, + pitch=pitch * rotation_speed_rad, + ) + + +def parse_action_string(action: str) -> list[list[str]]: + """Expand a SANA-WM action string into per-frame held keys. + + Args: + action: Comma-separated ``-`` segments. ``none-N`` + holds the pose for ``N`` frames. + + Returns: + Per-frame lists of DSL keys. + + Raises: + ValueError: The action string is empty or contains an invalid segment. + """ + 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, duration_part = segment.rsplit("-", 1) + if not duration_part.isdigit() or int(duration_part) <= 0: + raise ValueError(f"Action segment {segment!r} has a non-positive duration.") + keys_lower = keys_part.lower() + if keys_lower == "none": + keys: list[str] = [] + else: + bad = sorted({key for key in keys_lower if key 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([keys] * int(duration_part)) + 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, + smooth: bool = True, +) -> np.ndarray: + """Roll out a camera-to-world trajectory from an action string. + + Args: + action: Comma-separated ``-`` segments. + translation_speed: Per-frame translation magnitude. + rotation_speed_deg: Per-frame angular magnitude in degrees. + pitch_limit_deg: Maximum absolute pitch in degrees. + smooth: Whether to use upstream's velocity smoothing model. + + Returns: + ``[F + 1, 4, 4]`` camera-to-world matrices in OpenCV convention. + """ + per_frame = parse_action_string(action) + integrator = CameraPoseIntegrator(math.radians(pitch_limit_deg)) + velocity = VelocityState() + poses = [integrator.pose.copy()] + last_controls: set[str] = set() + dt = 1.0 / FPS + rotation_speed_rad = math.radians(rotation_speed_deg) + + for keys in per_frame: + controls = {DSL_KEY_TO_CONTROL[key] for key in keys} + target = controls_to_target_velocity( + controls, + translation_speed=translation_speed, + rotation_speed_rad=rotation_speed_rad, + ) + if smooth: + if controls - last_controls: + velocity.snap_to(target) + else: + velocity.step_toward(target, dt) + last_controls = controls + else: + velocity = target + poses.append(integrator.step(velocity)) + + return np.stack(poses, axis=0).astype(np.float32) + + +def fit_intrinsics_sequence(arr: np.ndarray, num_frames: int) -> np.ndarray: + """Return ``arr`` fitted to ``num_frames`` along axis 0.""" + arr = np.asarray(arr, dtype=np.float32) + if arr.shape[0] == num_frames: + return arr.copy() + if arr.shape[0] > num_frames: + return arr[:num_frames].copy() + if arr.shape[0] == 1: + return np.broadcast_to(arr[:1], (num_frames, *arr.shape[1:])).copy() + + old_t = np.linspace(0.0, 1.0, arr.shape[0], dtype=np.float32) + new_t = np.linspace(0.0, 1.0, num_frames, dtype=np.float32) + flat = arr.reshape(arr.shape[0], -1) + fitted = np.empty((num_frames, flat.shape[1]), dtype=np.float32) + for idx in range(flat.shape[1]): + fitted[:, idx] = np.interp(new_t, old_t, flat[:, idx]).astype(np.float32) + return fitted.reshape((num_frames, *arr.shape[1:])) + + +def load_intrinsics(path: Path, num_frames: int) -> np.ndarray: + """Return ``[num_frames, 4]`` intrinsics as ``[fx, fy, cx, cy]``. + + Args: + path: ``.npy`` file shaped ``[3, 3]``, ``[F, 3, 3]``, ``[4]``, or + ``[F, 4]``. + num_frames: Number of frames required by the rollout. + + Returns: + Per-frame vector intrinsics. + + Raises: + ValueError: The file shape is unsupported. + """ + arr = np.load(path).astype(np.float32) + if arr.shape == (4,): + return np.broadcast_to(arr, (num_frames, 4)).copy() + if arr.shape == (3, 3): + vector = np.array( + [arr[0, 0], arr[1, 1], arr[0, 2], arr[1, 2]], dtype=np.float32 + ) + return np.broadcast_to(vector, (num_frames, 4)).copy() + if arr.ndim == 2 and arr.shape[1] == 4: + return fit_intrinsics_sequence(arr, num_frames) + if arr.ndim == 3 and arr.shape[1:] == (3, 3): + matrix = fit_intrinsics_sequence(arr, num_frames) + return np.stack( + [matrix[:, 0, 0], matrix[:, 1, 1], matrix[:, 0, 2], matrix[:, 1, 2]], + axis=1, + ) + raise ValueError( + f"Unsupported intrinsics shape {arr.shape} for num_frames={num_frames}; " + "expected (3,3), (F,3,3), (4,), or (F,4)." + ) + + +def snap_num_frames( + num_frames: int, + *, + stride: int = SANA_WM_VAE_TEMPORAL_COMPRESSION, + upper_bound: int | None = None, +) -> int: + """Snap frame count to the nearest ``stride * k + 1`` value.""" + if num_frames < 1: + return 1 + if (num_frames - 1) % stride == 0: + return num_frames + floor_cand = num_frames - ((num_frames - 1) % stride) + ceil_cand = floor_cand + stride + snapped = ( + floor_cand + if (num_frames - floor_cand) < (ceil_cand - num_frames) + else ceil_cand + ) + if upper_bound is not None and snapped > upper_bound: + snapped = floor_cand + return max(snapped, 1) + + +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 resize-then-center-crop pixels.""" + src_w, src_h = src_size + resized_w, resized_h = resized_size + crop_left, crop_top = crop_offset + sx, sy = resized_w / src_w, resized_h / src_h + out = intrinsics_vec4.copy() + out[..., 0] *= sx + out[..., 2] = out[..., 2] * sx - crop_left + out[..., 1] *= sy + out[..., 3] = out[..., 3] * sy - crop_top + return out + + +def resize_center_crop_geometry( + src_size: tuple[int, int], + *, + target_h: int = DEFAULT_VIDEO_HEIGHT, + target_w: int = DEFAULT_VIDEO_WIDTH, +) -> tuple[tuple[int, int], tuple[int, int]]: + """Return resize dimensions and center-crop offset for an input image size. + + Args: + src_size: Source image size as ``(width, height)``. + target_h: Target crop height. + target_w: Target crop width. + + Returns: + ``(resized_size, crop_offset)`` where sizes are ``(width, height)`` + and offsets are ``(left, top)``. + """ + src_w, src_h = src_size + scale = max(target_h / src_h, target_w / src_w) + resized_w = max(target_w, int(round(src_w * scale))) + resized_h = max(target_h, int(round(src_h * scale))) + left = (resized_w - target_w) // 2 + top = (resized_h - target_h) // 2 + return (resized_w, resized_h), (left, top) + + +def get_pose_inverse(transform: torch.Tensor) -> torch.Tensor: + """Return inverse homogeneous transforms using rotation transpose.""" + rotation = transform[..., :3, :3] + translation = transform[..., :3, 3] + rotation_inv = rotation.transpose(-1, -2) + translation_inv = -torch.matmul(rotation_inv, translation.unsqueeze(-1)).squeeze(-1) + out = torch.eye(4, dtype=transform.dtype, device=transform.device).repeat( + transform.shape[:-2] + (1, 1) + ) + out[..., :3, :3] = rotation_inv + out[..., :3, 3] = translation_inv + return out + + +def compute_raymap( + intrinsics: torch.Tensor, + poses: torch.Tensor, + height: int, + width: int, + *, + use_plucker: bool = True, +) -> torch.Tensor: + """Compute world-space ray directions and moments for camera poses. + + Args: + intrinsics: ``[T, 4]`` tensor of ``[fx, fy, cx, cy]``. + poses: ``[T, 4, 4]`` camera-to-world matrices. + height: Raymap height. + width: Raymap width. + use_plucker: ``True`` returns ``[direction, moment]``; otherwise + returns ``[origin, direction]``. + + Returns: + ``[T, height, width, 6]`` raymap tensor. + """ + num_frames = intrinsics.shape[0] + device, dtype = intrinsics.device, intrinsics.dtype + y_grid, x_grid = torch.meshgrid( + torch.arange(height, device=device, dtype=dtype), + torch.arange(width, device=device, dtype=dtype), + indexing="ij", + ) + x_grid = x_grid[None].expand(num_frames, -1, -1) + y_grid = y_grid[None].expand(num_frames, -1, -1) + + fx = intrinsics[:, 0].view(num_frames, 1, 1) + fy = intrinsics[:, 1].view(num_frames, 1, 1) + cx = intrinsics[:, 2].view(num_frames, 1, 1) + cy = intrinsics[:, 3].view(num_frames, 1, 1) + + dirs_cam = torch.stack( + [(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], + dim=-1, + ) + rotation = poses[:, :3, :3] + translation = poses[:, :3, 3] + dirs_world = torch.einsum("tij,thwj->thwi", rotation, dirs_cam) + dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True) + origins = translation.view(num_frames, 1, 1, 3).expand_as(dirs_world) + if use_plucker: + return torch.cat([dirs_world, torch.cross(origins, dirs_world, dim=-1)], dim=-1) + return torch.cat([origins, dirs_world], dim=-1) + + +def prepare_camera( + poses_c2w: np.ndarray, + intrinsics_vec4: np.ndarray, + *, + target_size: tuple[int, int] = (DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH), + vae_stride: tuple[int, int, int] = ( + SANA_WM_VAE_TEMPORAL_COMPRESSION, + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_SPATIAL_COMPRESSION, + ), +) -> dict[str, torch.Tensor]: + """Build SANA-WM raymap and chunk-Plucker conditioning tensors.""" + num_frames = poses_c2w.shape[0] + vae_time_stride, vae_spatial_stride = vae_stride[0], vae_stride[-1] + pixel_h, pixel_w = target_size + latent_h = pixel_h // vae_spatial_stride + latent_w = pixel_w // vae_spatial_stride + latent_frames = (num_frames - 1) // vae_time_stride + 1 + + poses = torch.from_numpy(poses_c2w).float() + first_inv = get_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(pixel_w) + intrinsics_latent[:, [1, 3]] *= latent_h / float(pixel_h) + + 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, + ) + + chunks = [] + for start in time_indices - (vae_time_stride - 1): + chunk_start = max(0, int(start)) + chunk_end = chunk_start + vae_time_stride + chunk_poses = poses[chunk_start:chunk_end] + chunk_intrinsics = intrinsics_latent[chunk_start:chunk_end] + 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)]) + chunk_intrinsics = torch.cat( + [chunk_intrinsics, chunk_intrinsics[-1:].repeat(pad, 1)] + ) + plucker = compute_raymap( + chunk_intrinsics, + 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} diff --git a/integrations/sana/sana_wm/config.py b/integrations/sana/sana_wm/config.py new file mode 100644 index 000000000..ffa72922e --- /dev/null +++ b/integrations/sana/sana_wm/config.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Static configs for the SANA-WM bidirectional runner.""" + +from __future__ import annotations + +from flashdreams.infra.diffusion.model import DiffusionModelConfig +from flashdreams.infra.diffusion.scheduler import FlowMatchSchedulerConfig +from flashdreams.infra.pipeline import StreamInferencePipelineConfig +from flashdreams.infra.runner import RunnerConfig +from sana_wm.placeholder import SanaWMPlaceholderTransformerConfig +from sana_wm.runner import SanaWMRunnerConfig + +PIPELINE_SANA_WM_BIDIRECTIONAL = StreamInferencePipelineConfig( + name="sana-wm-bidirectional", + diffusion_model=DiffusionModelConfig( + transformer=SanaWMPlaceholderTransformerConfig(), + scheduler=FlowMatchSchedulerConfig(), + seed=42, + ), +) +"""Schema placeholder for the upstream-delegating SANA-WM runner.""" + +RUNNER_SANA_WM_BIDIRECTIONAL = SanaWMRunnerConfig( + runner_name=PIPELINE_SANA_WM_BIDIRECTIONAL.name, + description=( + "SANA-WM bidirectional I2V world model (upstream NVlabs/Sana " + "Stage-1 + LTX-2 refiner runner)." + ), + pipeline=PIPELINE_SANA_WM_BIDIRECTIONAL, +) +"""Default SANA-WM bidirectional runner config.""" + +RUNNER_CONFIGS: dict[str, RunnerConfig] = { + cfg.runner_name: cfg for cfg in (RUNNER_SANA_WM_BIDIRECTIONAL,) +} +"""SANA-WM runner configs keyed by ``runner_name``.""" + +__all__ = [ + "PIPELINE_SANA_WM_BIDIRECTIONAL", + "RUNNER_CONFIGS", + "RUNNER_SANA_WM_BIDIRECTIONAL", +] diff --git a/integrations/sana/sana_wm/constants.py b/integrations/sana/sana_wm/constants.py new file mode 100644 index 000000000..5251099ca --- /dev/null +++ b/integrations/sana/sana_wm/constants.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Constants for the public SANA-WM bidirectional release.""" + +SANA_WM_HF_REPO = "Efficient-Large-Model/SANA-WM_bidirectional" +"""Hugging Face repository containing SANA-WM bidirectional artefacts.""" + +SANA_WM_MODEL_PATH = f"hf://{SANA_WM_HF_REPO}/dit/sana_wm_1600m_720p.safetensors" +"""Default Stage-1 SANA-WM DiT checkpoint.""" + +SANA_WM_CONFIG_PATH = f"hf://{SANA_WM_HF_REPO}/config.yaml" +"""Default upstream inference YAML.""" + +SANA_WM_REFINER_ROOT = f"hf://{SANA_WM_HF_REPO}/refiner" +"""Default LTX-2 refiner root.""" + +SANA_WM_REFINER_GEMMA_ROOT = f"hf://{SANA_WM_HF_REPO}/refiner/text_encoder" +"""Default Gemma text-encoder root used by the refiner.""" + +DEFAULT_VIDEO_HEIGHT = 704 +"""SANA-WM bidirectional output height in pixels.""" + +DEFAULT_VIDEO_WIDTH = 1280 +"""SANA-WM bidirectional output width in pixels.""" + +SANA_WM_VAE_TEMPORAL_COMPRESSION = 8 +"""Temporal compression ratio of the LTX2 VAE used by SANA-WM.""" + +SANA_WM_VAE_SPATIAL_COMPRESSION = 32 +"""Spatial compression ratio of the LTX2 VAE used by SANA-WM.""" + +DEFAULT_FPS = 16 +"""Frame rate used by the public SANA-WM examples.""" + +DEFAULT_ACTION = "w-100,dw-60,w-100,aw-60" +"""Default SANA-WM demo action string from upstream ``demo_0``.""" diff --git a/integrations/sana/sana_wm/placeholder.py b/integrations/sana/sana_wm/placeholder.py new file mode 100644 index 000000000..ef46f1690 --- /dev/null +++ b/integrations/sana/sana_wm/placeholder.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Placeholder FlashDreams pipeline objects for the upstream SANA-WM runner.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.nn as nn +from torch import Tensor + +from flashdreams.infra.diffusion.transformer import ( + Transformer, + TransformerAutoregressiveCache, + TransformerConfig, +) + + +@dataclass(kw_only=True) +class SanaWMPlaceholderTransformerConfig(TransformerConfig): + """Transformer config used only to satisfy the runner config schema.""" + + _target: type["SanaWMPlaceholderTransformer"] = field( + default_factory=lambda: SanaWMPlaceholderTransformer + ) + + +class SanaWMPlaceholderTransformer(Transformer[TransformerAutoregressiveCache]): + """Non-executable transformer for the upstream-delegating SANA-WM runner.""" + + def __init__(self, config: SanaWMPlaceholderTransformerConfig) -> None: + super().__init__(config) + self._dummy = nn.Parameter(torch.empty(0)) + + @property + def latent_shape(self) -> tuple[int, ...]: + """Raise because SANA-WM generation is delegated to upstream Sana.""" + raise RuntimeError( + "SANA-WM bidirectional uses the custom SanaWMRunner path; " + "the placeholder FlashDreams pipeline is not executable." + ) + + def predict_flow( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: TransformerAutoregressiveCache, + input: Any = None, + ) -> Tensor: + """Raise because SANA-WM generation is delegated to upstream Sana.""" + raise RuntimeError( + "SANA-WM bidirectional uses the custom SanaWMRunner path; " + "the placeholder FlashDreams pipeline is not executable." + ) + + def patchify_and_maybe_split_cp(self, x: Any) -> Any: + """Return ``x`` unchanged for schema-only construction.""" + return x + + def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: + """Return ``x`` unchanged for schema-only construction.""" + return x diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py new file mode 100644 index 000000000..1a999673c --- /dev/null +++ b/integrations/sana/sana_wm/runner.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 bidirectional runner that delegates execution to upstream Sana.""" + +from __future__ import annotations + +import importlib +import os +import sys +from dataclasses import dataclass, field +from pathlib import Path +from types import ModuleType +from typing import Literal + +import numpy as np +import torch +from loguru import logger + +from flashdreams.core.io.disk import preflight_runtime_write_paths +from flashdreams.infra.runner import Runner, RunnerConfig +from sana_wm.constants import ( + DEFAULT_ACTION, + DEFAULT_FPS, + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, + SANA_WM_MODEL_PATH, + SANA_WM_REFINER_GEMMA_ROOT, + SANA_WM_REFINER_ROOT, + SANA_WM_VAE_TEMPORAL_COMPRESSION, +) + +SamplingAlgo = Literal[ + "auto", + "flow_euler_ltx", + "flow_euler", + "flow_dpm-solver", + "chunk_flow_euler", + "self_forcing", +] +"""Sampling algorithms exposed by upstream SANA-WM inference.""" + + +@dataclass(kw_only=True) +class SanaWMRunnerConfig(RunnerConfig): + """Runner config for the SANA-WM bidirectional release.""" + + _target: type["SanaWMRunner"] = field(default_factory=lambda: SanaWMRunner) + + upstream_sana_root: Path | None = None + """Path to a cloned NVlabs/Sana checkout. ``None`` auto-detects + ``$SANA_ROOT`` / ``../Sana`` or imports an installed Sana package.""" + + device: str = "auto" + """Torch device for upstream SANA-WM. ``"auto"`` picks CUDA when available.""" + + image_path: Path | None = None + """Path to the first-frame RGB image. Required at ``run()`` time.""" + + prompt: str = "" + """Inline text prompt. A non-empty value wins over ``prompt_path``.""" + + prompt_path: Path | None = None + """Fallback prompt file read when ``prompt`` is empty.""" + + camera_path: Path | None = None + """Optional ``.npy`` camera-to-world trajectory shaped ``[F, 4, 4]``.""" + + intrinsics_path: Path | None = None + """Optional ``.npy`` intrinsics shaped ``[3, 3]``, ``[F, 3, 3]``, + ``[4]``, or ``[F, 4]``. When absent, upstream estimates intrinsics + with Pi3X.""" + + action: str | None = DEFAULT_ACTION + """Action DSL used when ``camera_path`` is not provided.""" + + translation_speed: float = 0.025 + """Per-frame action translation speed.""" + + rotation_speed_deg: float = 0.6 + """Per-frame action rotation speed in degrees.""" + + name: str = "sana_wm" + """Output filename stem passed to upstream ``write_video``.""" + + num_frames: int = 161 + """Requested output frames before LTX2-VAE stride snapping.""" + + fps: int = DEFAULT_FPS + """Output video frame rate.""" + + step: int = 60 + """Stage-1 DiT sampling steps.""" + + cfg_scale: float = 5.0 + """Classifier-free guidance scale for Stage 1.""" + + flow_shift: float | None = None + """Optional scheduler flow-shift override.""" + + sampling_algo: SamplingAlgo = "auto" + """Stage-1 sampler. ``"auto"`` follows upstream config defaults.""" + + chunk_interval_k: float | None = None + """ChunkFlowEuler interval ratio override.""" + + num_cached_blocks: int = 2 + """Self-forcing sampler cached-block count.""" + + sink_token: bool = False + """Whether self-forcing keeps chunk 0 as a permanent sink.""" + + num_frame_per_block: int = 3 + """Self-forcing latent frames per AR block.""" + + denoising_step_list: str = "" + """Comma-separated self-forcing timesteps ending in ``0``.""" + + save_stage1: bool = False + """Also decode the unrefined Stage-1 latent when the refiner is enabled.""" + + negative_prompt: str = "" + """Negative prompt used when ``cfg_scale > 1``.""" + + seed: int = 42 + """Stage-1 random seed.""" + + no_action_overlay: bool = False + """Skip upstream action-overlay compositing on generated videos.""" + + config_path: str = SANA_WM_CONFIG_PATH + """Upstream inference YAML path or ``hf://`` URI.""" + + model_path: str = SANA_WM_MODEL_PATH + """Stage-1 checkpoint path or ``hf://`` URI.""" + + no_refiner: bool = False + """Skip the LTX-2 refiner and decode Stage-1 latents directly.""" + + refiner_root: str = SANA_WM_REFINER_ROOT + """LTX-2 refiner root path or ``hf://`` URI.""" + + refiner_gemma_root: str = SANA_WM_REFINER_GEMMA_ROOT + """Gemma text-encoder root for the LTX-2 refiner.""" + + refiner_seed: int = 42 + """Refiner random seed.""" + + sink_size: int = 1 + """Number of sink latent frames used by the refiner.""" + + refiner_block_size: int | None = None + """Optional refiner AR block size; ``None`` uses sink-bidirectional mode.""" + + refiner_kv_max_frames: int = 11 + """Maximum refiner KV context frames in AR mode.""" + + offload_vae: bool = False + """Move the VAE to CPU between encode/decode phases.""" + + offload_refiner: bool = False + """Build and release the refiner only around refinement.""" + + offload_text_encoder: bool = False + """Move the Stage-1 text encoder to CPU between prompt encodes.""" + + +class SanaWMRunner(Runner[SanaWMRunnerConfig, object]): + """CLI driver for upstream SANA-WM bidirectional inference.""" + + config: SanaWMRunnerConfig + + def __init__(self, config: SanaWMRunnerConfig) -> None: + self.config = config + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if torch.distributed.is_initialized(): + self.world_size = torch.distributed.get_world_size() + self.global_rank = torch.distributed.get_rank() + else: + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + self.global_rank = int(os.environ.get("RANK", "0")) + self.is_rank_zero = self.global_rank == 0 + preflight_runtime_write_paths(output_dir=config.output_dir) + + def _resolve_prompt(self) -> str: + """Resolve the prompt from inline text or ``prompt_path``.""" + if self.config.prompt: + return self.config.prompt + if self.config.prompt_path is None: + raise ValueError("SanaWMRunner requires --prompt or --prompt-path.") + prompt = self.config.prompt_path.read_text( + encoding="utf-8", errors="replace" + ).strip() + if not prompt: + raise ValueError(f"Prompt file is empty: {self.config.prompt_path}") + return prompt + + def _resolve_device(self) -> torch.device: + """Return the device used by upstream SANA-WM.""" + if self.config.device == "auto": + if torch.cuda.is_available(): + return torch.device(f"cuda:{self.local_rank}") + return torch.device("cpu") + if self.config.device == "cuda" and torch.cuda.is_available(): + return torch.device(f"cuda:{self.local_rank}") + return torch.device(self.config.device) + + def _resolve_trajectory(self, upstream: ModuleType) -> np.ndarray: + """Load or roll out the camera-to-world trajectory.""" + if self.config.camera_path is not None: + c2w = np.load(self.config.camera_path).astype(np.float32) + if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): + raise ValueError( + f"--camera-path must be a [F, 4, 4] .npy; got {c2w.shape}." + ) + return c2w + if not self.config.action: + raise ValueError("SanaWMRunner requires --camera-path or --action.") + return upstream.action_string_to_c2w( + self.config.action, + translation_speed=self.config.translation_speed, + rotation_speed_deg=self.config.rotation_speed_deg, + ) + + def _denoising_steps(self) -> list[int] | None: + """Parse optional self-forcing denoising timesteps.""" + if not self.config.denoising_step_list: + return None + steps = [ + int(item.strip()) + for item in self.config.denoising_step_list.split(",") + if item.strip() + ] + if not steps or steps[-1] != 0: + raise ValueError("--denoising-step-list must end with 0.") + return steps + + def run(self) -> None: + """Run upstream SANA-WM bidirectional inference and write outputs.""" + cfg = self.config + if cfg.image_path is None: + raise ValueError("SanaWMRunner requires --image-path.") + + upstream = _import_upstream_sana(cfg.upstream_sana_root) + device = self._resolve_device() + prompt = self._resolve_prompt() + + image = upstream.Image.open(cfg.image_path).convert("RGB") + c2w_full = self._resolve_trajectory(upstream) + num_frames = min(cfg.num_frames, c2w_full.shape[0]) + snapped = upstream._snap_num_frames( + num_frames, + stride=SANA_WM_VAE_TEMPORAL_COMPRESSION, + upper_bound=c2w_full.shape[0], + ) + if snapped != cfg.num_frames and self.is_rank_zero: + logger.warning( + "SANA-WM requires num_frames = 8k+1; requested {} snapped to {} " + "(trajectory has {} frames).", + cfg.num_frames, + snapped, + c2w_full.shape[0], + ) + num_frames = snapped + c2w = c2w_full[:num_frames] + + cropped, src_size, resized_size, crop_offset = upstream.resize_and_center_crop( + image, + target_h=DEFAULT_VIDEO_HEIGHT, + target_w=DEFAULT_VIDEO_WIDTH, + ) + if cfg.intrinsics_path is None: + intrinsics_src = np.broadcast_to( + upstream.estimate_intrinsics_with_pi3x( + image, device, upstream.get_root_logger() + ), + (num_frames, 4), + ).copy() + else: + intrinsics_src = upstream.load_intrinsics(cfg.intrinsics_path, num_frames) + intrinsics_vec4 = upstream.transform_intrinsics_for_crop( + intrinsics_src, src_size, resized_size, crop_offset + ) + + inference_config = upstream.pyrallis.parse( + config_class=upstream.InferenceConfig, + config_path=upstream.resolve_hf_path(cfg.config_path), + args=[], + ) + refiner = ( + None + if cfg.no_refiner + else upstream.RefinerSettings( + root=cfg.refiner_root, + gemma_root=cfg.refiner_gemma_root, + sink_size=cfg.sink_size, + seed=cfg.refiner_seed, + block_size=cfg.refiner_block_size, + kv_max_frames=cfg.refiner_kv_max_frames, + ) + ) + pipeline = upstream.SanaWMPipeline( + config=inference_config, + model_path=upstream.resolve_hf_path(cfg.model_path), + device=device, + refiner=refiner, + offload_vae=cfg.offload_vae, + offload_refiner=cfg.offload_refiner, + offload_text_encoder=cfg.offload_text_encoder, + logger=upstream.get_root_logger(), + ) + + sampling_algo = cfg.sampling_algo + if sampling_algo == "auto": + sampling_algo = ( + inference_config.scheduler.vis_sampler + if inference_config.scheduler.vis_sampler + in {"chunk_flow_euler", "self_forcing"} + else "flow_euler_ltx" + ) + params = upstream.GenerationParams( + num_frames=num_frames, + fps=cfg.fps, + step=cfg.step, + cfg_scale=cfg.cfg_scale, + flow_shift=cfg.flow_shift, + seed=cfg.seed, + negative_prompt=cfg.negative_prompt, + sampling_algo=sampling_algo, + chunk_interval_k=cfg.chunk_interval_k, + num_cached_blocks=cfg.num_cached_blocks, + sink_token=cfg.sink_token, + num_frame_per_block=cfg.num_frame_per_block, + denoising_step_list=self._denoising_steps(), + save_stage1=cfg.save_stage1, + ) + + result = pipeline.generate(cropped, prompt, c2w, intrinsics_vec4, params) + video_hwc = result["video"] + if not cfg.no_action_overlay: + video_hwc = upstream.apply_overlay(video_hwc, result["c2w"]) + + if not self.is_rank_zero: + return + upstream.write_video( + cfg.output_dir, + cfg.name, + video_hwc, + params.fps, + upstream.get_root_logger(), + ) + + stage1_video = result.get("stage1_video") + if stage1_video is not None: + stage1_hwc = stage1_video + if not cfg.no_action_overlay: + stage1_hwc = upstream.apply_overlay(stage1_hwc, result["stage1_c2w"]) + upstream.write_video( + cfg.output_dir, + f"{cfg.name}_stage1", + stage1_hwc, + params.fps, + upstream.get_root_logger(), + ) + + +def _candidate_upstream_roots(explicit_root: Path | None) -> list[Path]: + """Return candidate NVlabs/Sana checkout roots in priority order.""" + roots: list[Path] = [] + if explicit_root is not None: + roots.append(explicit_root) + for env_name in ("SANA_ROOT", "SANA_REPO", "SANA_WM_UPSTREAM_ROOT"): + value = os.environ.get(env_name) + if value: + roots.append(Path(value)) + repo_root = Path(__file__).resolve().parents[3] + roots.append(repo_root.parent / "Sana") + roots.append(Path.cwd().parent / "Sana") + return roots + + +def _is_upstream_sana_root(path: Path) -> bool: + """Return ``True`` when ``path`` looks like an NVlabs/Sana checkout.""" + return ( + path.expanduser() / "inference_video_scripts" / "wm" / "inference_sana_wm.py" + ).is_file() + + +def _import_upstream_sana(explicit_root: Path | None) -> ModuleType: + """Import upstream SANA-WM inference after adding a checkout to ``sys.path``.""" + for root in _candidate_upstream_roots(explicit_root): + candidate = root.expanduser().resolve() + if _is_upstream_sana_root(candidate): + candidate_str = str(candidate) + if candidate_str not in sys.path: + sys.path.insert(0, candidate_str) + break + + try: + return importlib.import_module("inference_video_scripts.wm.inference_sana_wm") + except ModuleNotFoundError as exc: + roots = ", ".join( + str(path) for path in _candidate_upstream_roots(explicit_root) + ) + raise ModuleNotFoundError( + "Could not import upstream SANA-WM inference. Clone NVlabs/Sana " + "next to this repo, pass --upstream-sana-root, or set SANA_ROOT; " + "then install the upstream Sana inference dependencies. " + f"Checked roots: {roots}. Original import error: {exc}" + ) from exc + except Exception as exc: + raise RuntimeError( + "Importing upstream SANA-WM inference failed. Ensure the upstream " + "Sana checkout dependencies are installed in the active environment." + ) from exc + + +__all__ = [ + "SamplingAlgo", + "SanaWMRunner", + "SanaWMRunnerConfig", +] diff --git a/integrations/sana/tests/parity_check/README.md b/integrations/sana/tests/parity_check/README.md new file mode 100644 index 000000000..08d5d5abf --- /dev/null +++ b/integrations/sana/tests/parity_check/README.md @@ -0,0 +1,21 @@ + + +# SANA-WM parity check + +This directory is reserved for the opt-in SANA-WM upstream parity +harness. It should follow the existing integration pattern: + +- clone or reuse a pinned NVlabs/Sana checkout; +- install upstream-heavy dependencies in an isolated environment; +- run upstream `inference_video_scripts/wm/inference_sana_wm.py` and + `flashdreams-run sana-wm-bidirectional` on the same image, prompt, + camera, intrinsics, seed, and sampler settings; +- compare decoded MP4 frames and report mean / max `|Delta|`. + +Do not put the full checkpoint download, refiner execution, or MP4 +generation path in `ci_cpu`. Add a bounded `ci_gpu` gate only if CI +pre-provisions the required artefacts and the test skips cleanly when +they are absent. diff --git a/integrations/sana/tests/test_camera.py b/integrations/sana/tests/test_camera.py new file mode 100644 index 000000000..43d0b7275 --- /dev/null +++ b/integrations/sana/tests/test_camera.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""CPU-safe tests for SANA-WM camera and intrinsics helpers.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from sana_wm.camera import ( + action_string_to_c2w, + fit_intrinsics_sequence, + load_intrinsics, + prepare_camera, + resize_center_crop_geometry, + snap_num_frames, + transform_intrinsics_for_crop, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_action_string_rolls_out_identity_plus_motion() -> None: + """Expand ``w-3`` to an identity frame plus three motion frames.""" + c2w = action_string_to_c2w("w-3", smooth=False) + + assert c2w.shape == (4, 4, 4) + np.testing.assert_allclose(c2w[0], np.eye(4), atol=1e-6) + assert np.all(np.diff(c2w[:, 2, 3]) > 0) + + +def test_action_string_rejects_unknown_keys() -> None: + """Reject invalid action DSL tokens.""" + with pytest.raises(ValueError, match="unknown keys"): + action_string_to_c2w("q-3") + + +def test_snap_num_frames_matches_ltx2_stride() -> None: + """Snap requested frames to nearest ``8k + 1`` value.""" + assert snap_num_frames(321) == 321 + assert snap_num_frames(322) == 321 + assert snap_num_frames(325) == 329 + assert snap_num_frames(325, upper_bound=326) == 321 + + +def test_fit_intrinsics_sequence_interpolates_short_sequences() -> None: + """Interpolate per-frame intrinsics when the source sequence is shorter.""" + source = np.array([[10.0, 20.0, 1.0, 2.0], [30.0, 40.0, 3.0, 4.0]]) + + fitted = fit_intrinsics_sequence(source, 3) + + np.testing.assert_allclose( + fitted, + np.array( + [ + [10.0, 20.0, 1.0, 2.0], + [20.0, 30.0, 2.0, 3.0], + [30.0, 40.0, 3.0, 4.0], + ], + dtype=np.float32, + ), + ) + + +@pytest.mark.parametrize( + ("array", "expected"), + [ + ( + np.array([100.0, 110.0, 50.0, 55.0], dtype=np.float32), + np.array([[100.0, 110.0, 50.0, 55.0]] * 3, dtype=np.float32), + ), + ( + np.array( + [[100.0, 0.0, 50.0], [0.0, 110.0, 55.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ), + np.array([[100.0, 110.0, 50.0, 55.0]] * 3, dtype=np.float32), + ), + ], +) +def test_load_intrinsics_accepts_static_shapes( + tmp_path: Path, array: np.ndarray, expected: np.ndarray +) -> None: + """Load static vector and matrix intrinsics as per-frame vectors.""" + path = tmp_path / "intrinsics.npy" + np.save(path, array) + + loaded = load_intrinsics(path, num_frames=3) + + np.testing.assert_allclose(loaded, expected) + + +def test_resize_crop_geometry_and_intrinsics_transform() -> None: + """Map source intrinsics through SANA-WM resize and center-crop geometry.""" + src_size = (640, 480) + resized_size, crop_offset = resize_center_crop_geometry(src_size) + intrinsics = np.array([[400.0, 420.0, 320.0, 240.0]], dtype=np.float32) + + transformed = transform_intrinsics_for_crop( + intrinsics, src_size, resized_size, crop_offset + ) + + assert resized_size == (1280, 960) + assert crop_offset == (0, 128) + np.testing.assert_allclose( + transformed, + np.array([[800.0, 840.0, 640.0, 352.0]], dtype=np.float32), + ) + + +def test_prepare_camera_shapes_for_sana_wm_resolution() -> None: + """Build raymap and chunk-Plucker tensors at 704x1280 SANA-WM shape.""" + num_frames = 17 + poses = np.broadcast_to(np.eye(4, dtype=np.float32), (num_frames, 4, 4)).copy() + poses[:, 2, 3] = np.linspace(0.0, 1.0, num_frames) + intrinsics = np.broadcast_to( + np.array([900.0, 900.0, 640.0, 352.0], dtype=np.float32), + (num_frames, 4), + ).copy() + + camera = prepare_camera(poses, intrinsics) + + assert camera["raymap"].shape == (3, 20) + assert camera["chunk_plucker"].shape == (48, 3, 22, 40) + assert camera["raymap"].dtype == torch.float32 + assert camera["chunk_plucker"].dtype == torch.float32 diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py new file mode 100644 index 000000000..c361941bc --- /dev/null +++ b/integrations/sana/tests/test_smoke.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""CPU-safe smoke tests for the SANA-WM runner shim.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover - Python < 3.11 fallback + import tomli as tomllib + +from sana_wm.config import ( + PIPELINE_SANA_WM_BIDIRECTIONAL, + RUNNER_CONFIGS, + RUNNER_SANA_WM_BIDIRECTIONAL, +) +from sana_wm.constants import ( + SANA_WM_CONFIG_PATH, + SANA_WM_HF_REPO, + SANA_WM_MODEL_PATH, +) +from sana_wm.runner import SanaWMRunner, SanaWMRunnerConfig + +from flashdreams.infra.config import derive_config + +pytestmark = pytest.mark.ci_cpu + +ENTRY_POINT_GROUP = "flashdreams.runner_configs" + + +def test_runner_config_is_registered() -> None: + """Expose one SANA-WM runner keyed by its public slug.""" + assert RUNNER_CONFIGS == {"sana-wm-bidirectional": RUNNER_SANA_WM_BIDIRECTIONAL} + + +def test_runner_name_mirrors_pipeline_name() -> None: + """Keep ``flashdreams-run `` aligned with the wrapped config slug.""" + assert RUNNER_SANA_WM_BIDIRECTIONAL.runner_name == ( + PIPELINE_SANA_WM_BIDIRECTIONAL.name + ) + + +def test_runner_has_description() -> None: + """Provide non-empty CLI help text for the runner registry.""" + assert RUNNER_SANA_WM_BIDIRECTIONAL.description.strip() + + +def test_hf_defaults_point_at_bidirectional_release() -> None: + """Pin every default SANA-WM artefact to the bidirectional HF repo.""" + assert SANA_WM_HF_REPO == "Efficient-Large-Model/SANA-WM_bidirectional" + assert SANA_WM_MODEL_PATH == ( + "hf://Efficient-Large-Model/SANA-WM_bidirectional/" + "dit/sana_wm_1600m_720p.safetensors" + ) + assert SANA_WM_CONFIG_PATH == ( + "hf://Efficient-Large-Model/SANA-WM_bidirectional/config.yaml" + ) + + +def test_runner_setup_does_not_import_upstream_sana() -> None: + """Construct the runner without touching upstream Sana dependencies.""" + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + image_path=Path("missing.png"), + prompt="demo", + ) + + runner = cfg.setup() + + assert isinstance(runner, SanaWMRunner) + assert runner.config.image_path == Path("missing.png") + + +def test_runner_config_type() -> None: + """Keep the exported literal on the SANA-WM runner config subclass.""" + assert isinstance(RUNNER_SANA_WM_BIDIRECTIONAL, SanaWMRunnerConfig) + + +def test_pyproject_entry_point_matches_runner_literal() -> None: + """Keep the package entry point aligned with ``RUNNER_CONFIGS``.""" + pyproject = tomllib.loads( + Path("integrations/sana/pyproject.toml").read_text(encoding="utf-8") + ) + entry_points = pyproject["project"]["entry-points"][ENTRY_POINT_GROUP] + + assert entry_points == { + "sana-wm-bidirectional": "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL" + } diff --git a/pyproject.toml b/pyproject.toml index 250048b84..49b04c628 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ extraPaths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/sana", "integrations/self_forcing", "integrations/wan21", "integrations/wan22", @@ -74,6 +75,7 @@ extra-paths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/sana", "integrations/self_forcing", "integrations/wan21", "integrations/wan22", diff --git a/uv.lock b/uv.lock index 2b0d03215..55413b69e 100644 --- a/uv.lock +++ b/uv.lock @@ -27,6 +27,7 @@ members = [ "flashdreams-hy-worldplay", "flashdreams-lingbot", "flashdreams-omnidreams", + "flashdreams-sana-wm", "flashdreams-self-forcing", "flashdreams-wan21", "flashdreams-wan22", @@ -1567,6 +1568,26 @@ requires-dist = [ ] provides-extras = ["interactive-drive", "dev"] +[[package]] +name = "flashdreams-sana-wm" +version = "0.1.0" +source = { editable = "integrations/sana" } +dependencies = [ + { name = "flashdreams" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashdreams", editable = "flashdreams" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, +] +provides-extras = ["dev"] + [[package]] name = "flashdreams-self-forcing" version = "0.1.0" From f504eac0f041b6de962b62a3c7efbe1f7b02baba Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 15 Jul 2026 16:23:10 -0700 Subject: [PATCH 02/36] docs(sana): document upstream env install flow Document that the SANA-WM integration is expected to run inside the upstream Sana environment created by `environment_setup.sh sana`. The setup notes keep Sana's pinned PyTorch and model dependencies intact, then install FlashDreams and the local integration in editable mode without dependency resolution. Include the expected checkout layout and call out Transformer Engine as an optional dependency for quantized inference while BF16 remains the default validated path. Signed-off-by: Aidan Foster --- integrations/sana/README.md | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index 88d2c2911..6ae421109 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -31,16 +31,38 @@ Clone Sana next to FlashDreams or pass its path explicitly: ```bash git clone https://github.com/NVlabs/Sana ../Sana -bash ../Sana/environment_setup.sh sana +cd ../Sana +bash environment_setup.sh sana +conda activate sana ``` The runner auto-detects `../Sana`. You can also set `SANA_ROOT` or pass `--upstream-sana-root /path/to/Sana`. +Install the FlashDreams packages into the same environment without +letting pip resolve or upgrade dependencies. The upstream Sana environment +owns the model runtime stack and pins packages such as `transformers` and +`huggingface-hub`. + +```bash +cd ../flashdreams +python -m pip install --no-deps -e flashdreams -e integrations/sana +``` + +If FlashDreams was installed without `--no-deps` and pip upgraded Sana's +pinned packages, repair the env before running generation: + +```bash +python -m pip install --no-deps \ + "huggingface-hub==0.36.0" \ + "transformers==4.57.3" +python -m pip install --no-deps -e flashdreams -e integrations/sana +``` + ## Run ```bash -uv run flashdreams-run sana-wm-bidirectional \ +flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --action "w-100,dw-60,w-100,aw-60" \ @@ -51,7 +73,7 @@ uv run flashdreams-run sana-wm-bidirectional \ To use an explicit trajectory: ```bash -uv run flashdreams-run sana-wm-bidirectional \ +flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ From 41ee2bb57b0c892c4a9fffb47f152451b0b2fe8b Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 15 Jul 2026 16:26:49 -0700 Subject: [PATCH 03/36] docs(sana): clarify boolean CLI values Update the SANA-WM README examples to pass explicit values for boolean Tyro options such as `--no-refiner True`. This matches the generated CLI contract and avoids the confusing "Missing value" error when users treat those fields as flag-style booleans. Signed-off-by: Aidan Foster --- integrations/sana/README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index 6ae421109..d22d06231 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -82,9 +82,11 @@ flashdreams-run sana-wm-bidirectional \ --output-dir outputs/sana_wm ``` -`--no-refiner` skips the LTX-2 refiner and decodes Stage-1 latents -directly, matching upstream's fast debugging path. Full generation still -downloads the public Hugging Face artefacts on first use. +Set `--no-refiner True` to skip the LTX-2 refiner and decode Stage-1 +latents directly, matching upstream's fast debugging path. FlashDreams' +current CLI expects explicit boolean values, so pass `True` or `False` +for boolean fields. Full generation still downloads the public Hugging +Face artefacts on first use. ## Tests From 7e6af884d939b399650e479fa70e5787c6920915 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 15 Jul 2026 17:07:19 -0700 Subject: [PATCH 04/36] docs(sana): record validated BF16 smoke tests Record the BF16 commands used to validate both the Stage-1-only and refined SANA-WM paths, including the expected output locations under `outputs/sana_wm*/sana_wm_generated.mp4`. Note the upstream warnings that appeared during successful runs so they are not mistaken for failures: the null-embed load warning, missing `pos_embed`, and the optional Apex RMSNorm fallback. These commands establish BF16 as the known-good baseline for the integration. Signed-off-by: Aidan Foster --- integrations/sana/README.md | 45 +++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index d22d06231..d18cb7a2a 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -88,6 +88,51 @@ current CLI expects explicit boolean values, so pass `True` or `False` for boolean fields. Full generation still downloads the public Hugging Face artefacts on first use. +## Validated BF16 smoke tests + +The following commands were validated on an NVIDIA RTX PRO 6000 Blackwell +with upstream Sana installed in the `sana` conda environment. They use +the default BF16 path and do not require Transformer Engine. + +Stage-1 plus VAE decode: + +```bash +flashdreams-run sana-wm-bidirectional \ + --upstream-sana-root ../Sana \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm \ + --no-refiner True +``` + +Expected output: + +```text +outputs/sana_wm/sana_wm_generated.mp4 +``` + +Stage-1 plus LTX-2 refiner: + +```bash +flashdreams-run sana-wm-bidirectional \ + --upstream-sana-root ../Sana \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_refiner +``` + +Expected output: + +```text +outputs/sana_wm_refiner/sana_wm_generated.mp4 +``` + ## Tests CPU-safe tests cover import, runner config, action parsing, intrinsics From de39905ac8aee4d0093d819267f482ea322a425d Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 16 Jul 2026 08:21:19 -0700 Subject: [PATCH 05/36] feat(sana): add FP8 and FP4 precision options Add first-class `stage1_precision` and `refiner_precision` runner options for BF16, FP8, and FP4. The runner maps those options onto upstream SANA-WM Transformer Engine environment selectors inside a scoped environment so external NVFP4 variables cannot accidentally change the default BF16 path. Validate quantized precision requests before loading checkpoints. FP8 now requires a CUDA device, Hopper-or-newer GPU capability, CUDA 12.9 or newer as reported by PyTorch, and Transformer Engine's Float8BlockScaling recipe. FP4 requires a CUDA device, Blackwell GPU capability, and Transformer Engine's NVFP4 recipe. Extend the SANA README with precision requirements and an FP4 example command, and add CPU-safe tests for the env mapping, cleanup behavior, and early capability errors. Signed-off-by: Aidan Foster --- integrations/sana/README.md | 33 +++ integrations/sana/sana_wm/runner.py | 394 +++++++++++++++++++------- integrations/sana/tests/test_smoke.py | 117 +++++++- 3 files changed, 437 insertions(+), 107 deletions(-) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index d18cb7a2a..705d721a3 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -133,6 +133,39 @@ Expected output: outputs/sana_wm_refiner/sana_wm_generated.mp4 ``` +## FP8 and FP4 precision + +The runner defaults to BF16. Quantized inference is opt-in: + +| option | hardware | dependency | +| --- | --- | --- | +| `--stage1-precision fp8` / `--refiner-precision fp8` | Hopper or newer (`sm_90+`) with CUDA 12.9+ | Transformer Engine | +| `--stage1-precision fp4` / `--refiner-precision fp4` | Blackwell (`sm_100+`) | Transformer Engine with NVFP4 recipes | + +The runner validates these requirements before loading checkpoints. If the +GPU or Transformer Engine install cannot support the requested precision, +it fails early with a targeted error instead of falling through to a deep +upstream import or model-build failure. + +Example Blackwell FP4 run: + +```bash +flashdreams-run sana-wm-bidirectional \ + --upstream-sana-root ../Sana \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_fp4 \ + --stage1-precision fp4 \ + --refiner-precision fp4 +``` + +Use `fp8` for Hopper-class GPUs when the PyTorch environment reports CUDA +12.9 or newer. Keep both precision flags at `bf16` if Transformer Engine is +unavailable or if maximum compatibility is preferred. + ## Tests CPU-safe tests cover import, runner config, action parsing, intrinsics diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index 1a999673c..56c5981b2 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -20,6 +20,8 @@ import importlib import os import sys +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from types import ModuleType @@ -53,6 +55,9 @@ ] """Sampling algorithms exposed by upstream SANA-WM inference.""" +Precision = Literal["bf16", "fp8", "fp4"] +"""SANA-WM Stage-1/refiner precision modes.""" + @dataclass(kw_only=True) class SanaWMRunnerConfig(RunnerConfig): @@ -147,9 +152,20 @@ class SanaWMRunnerConfig(RunnerConfig): model_path: str = SANA_WM_MODEL_PATH """Stage-1 checkpoint path or ``hf://`` URI.""" + stage1_precision: Precision = "bf16" + """Stage-1 DiT compute precision. ``"bf16"`` is the default; ``"fp8"`` + requires Hopper or newer plus Transformer Engine; ``"fp4"`` requires + Blackwell plus Transformer Engine.""" + no_refiner: bool = False """Skip the LTX-2 refiner and decode Stage-1 latents directly.""" + refiner_precision: Precision = "bf16" + """LTX-2 refiner compute precision. ``"bf16"`` is the default; + ``"fp8"`` requires Hopper or newer plus Transformer Engine; ``"fp4"`` + requires Blackwell plus Transformer Engine. Ignored when + ``no_refiner`` is ``True``.""" + refiner_root: str = SANA_WM_REFINER_ROOT """LTX-2 refiner root path or ``hf://`` URI.""" @@ -254,128 +270,293 @@ def run(self) -> None: if cfg.image_path is None: raise ValueError("SanaWMRunner requires --image-path.") - upstream = _import_upstream_sana(cfg.upstream_sana_root) device = self._resolve_device() - prompt = self._resolve_prompt() - - image = upstream.Image.open(cfg.image_path).convert("RGB") - c2w_full = self._resolve_trajectory(upstream) - num_frames = min(cfg.num_frames, c2w_full.shape[0]) - snapped = upstream._snap_num_frames( - num_frames, - stride=SANA_WM_VAE_TEMPORAL_COMPRESSION, - upper_bound=c2w_full.shape[0], - ) - if snapped != cfg.num_frames and self.is_rank_zero: - logger.warning( - "SANA-WM requires num_frames = 8k+1; requested {} snapped to {} " - "(trajectory has {} frames).", - cfg.num_frames, - snapped, - c2w_full.shape[0], - ) - num_frames = snapped - c2w = c2w_full[:num_frames] - - cropped, src_size, resized_size, crop_offset = upstream.resize_and_center_crop( - image, - target_h=DEFAULT_VIDEO_HEIGHT, - target_w=DEFAULT_VIDEO_WIDTH, - ) - if cfg.intrinsics_path is None: - intrinsics_src = np.broadcast_to( - upstream.estimate_intrinsics_with_pi3x( - image, device, upstream.get_root_logger() - ), - (num_frames, 4), - ).copy() - else: - intrinsics_src = upstream.load_intrinsics(cfg.intrinsics_path, num_frames) - intrinsics_vec4 = upstream.transform_intrinsics_for_crop( - intrinsics_src, src_size, resized_size, crop_offset + _validate_precision_request( + device=device, + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, ) - inference_config = upstream.pyrallis.parse( - config_class=upstream.InferenceConfig, - config_path=upstream.resolve_hf_path(cfg.config_path), - args=[], + precision_env = _precision_env_updates( + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, ) - refiner = ( - None - if cfg.no_refiner - else upstream.RefinerSettings( - root=cfg.refiner_root, - gemma_root=cfg.refiner_gemma_root, - sink_size=cfg.sink_size, - seed=cfg.refiner_seed, - block_size=cfg.refiner_block_size, - kv_max_frames=cfg.refiner_kv_max_frames, + with _temporary_environment(precision_env): + upstream = _import_upstream_sana(cfg.upstream_sana_root) + prompt = self._resolve_prompt() + + image = upstream.Image.open(cfg.image_path).convert("RGB") + c2w_full = self._resolve_trajectory(upstream) + num_frames = min(cfg.num_frames, c2w_full.shape[0]) + snapped = upstream._snap_num_frames( + num_frames, + stride=SANA_WM_VAE_TEMPORAL_COMPRESSION, + upper_bound=c2w_full.shape[0], + ) + if snapped != cfg.num_frames and self.is_rank_zero: + logger.warning( + "SANA-WM requires num_frames = 8k+1; requested {} snapped to {} " + "(trajectory has {} frames).", + cfg.num_frames, + snapped, + c2w_full.shape[0], + ) + num_frames = snapped + c2w = c2w_full[:num_frames] + + cropped, src_size, resized_size, crop_offset = ( + upstream.resize_and_center_crop( + image, + target_h=DEFAULT_VIDEO_HEIGHT, + target_w=DEFAULT_VIDEO_WIDTH, + ) + ) + if cfg.intrinsics_path is None: + intrinsics_src = np.broadcast_to( + upstream.estimate_intrinsics_with_pi3x( + image, device, upstream.get_root_logger() + ), + (num_frames, 4), + ).copy() + else: + intrinsics_src = upstream.load_intrinsics( + cfg.intrinsics_path, num_frames + ) + intrinsics_vec4 = upstream.transform_intrinsics_for_crop( + intrinsics_src, src_size, resized_size, crop_offset ) - ) - pipeline = upstream.SanaWMPipeline( - config=inference_config, - model_path=upstream.resolve_hf_path(cfg.model_path), - device=device, - refiner=refiner, - offload_vae=cfg.offload_vae, - offload_refiner=cfg.offload_refiner, - offload_text_encoder=cfg.offload_text_encoder, - logger=upstream.get_root_logger(), - ) - sampling_algo = cfg.sampling_algo - if sampling_algo == "auto": - sampling_algo = ( - inference_config.scheduler.vis_sampler - if inference_config.scheduler.vis_sampler - in {"chunk_flow_euler", "self_forcing"} - else "flow_euler_ltx" + inference_config = upstream.pyrallis.parse( + config_class=upstream.InferenceConfig, + config_path=upstream.resolve_hf_path(cfg.config_path), + args=[], + ) + refiner = ( + None + if cfg.no_refiner + else upstream.RefinerSettings( + root=cfg.refiner_root, + gemma_root=cfg.refiner_gemma_root, + sink_size=cfg.sink_size, + seed=cfg.refiner_seed, + block_size=cfg.refiner_block_size, + kv_max_frames=cfg.refiner_kv_max_frames, + ) + ) + pipeline = upstream.SanaWMPipeline( + config=inference_config, + model_path=upstream.resolve_hf_path(cfg.model_path), + device=device, + refiner=refiner, + offload_vae=cfg.offload_vae, + offload_refiner=cfg.offload_refiner, + offload_text_encoder=cfg.offload_text_encoder, + logger=upstream.get_root_logger(), ) - params = upstream.GenerationParams( - num_frames=num_frames, - fps=cfg.fps, - step=cfg.step, - cfg_scale=cfg.cfg_scale, - flow_shift=cfg.flow_shift, - seed=cfg.seed, - negative_prompt=cfg.negative_prompt, - sampling_algo=sampling_algo, - chunk_interval_k=cfg.chunk_interval_k, - num_cached_blocks=cfg.num_cached_blocks, - sink_token=cfg.sink_token, - num_frame_per_block=cfg.num_frame_per_block, - denoising_step_list=self._denoising_steps(), - save_stage1=cfg.save_stage1, - ) - result = pipeline.generate(cropped, prompt, c2w, intrinsics_vec4, params) - video_hwc = result["video"] - if not cfg.no_action_overlay: - video_hwc = upstream.apply_overlay(video_hwc, result["c2w"]) - - if not self.is_rank_zero: - return - upstream.write_video( - cfg.output_dir, - cfg.name, - video_hwc, - params.fps, - upstream.get_root_logger(), - ) + sampling_algo = cfg.sampling_algo + if sampling_algo == "auto": + sampling_algo = ( + inference_config.scheduler.vis_sampler + if inference_config.scheduler.vis_sampler + in {"chunk_flow_euler", "self_forcing"} + else "flow_euler_ltx" + ) + params = upstream.GenerationParams( + num_frames=num_frames, + fps=cfg.fps, + step=cfg.step, + cfg_scale=cfg.cfg_scale, + flow_shift=cfg.flow_shift, + seed=cfg.seed, + negative_prompt=cfg.negative_prompt, + sampling_algo=sampling_algo, + chunk_interval_k=cfg.chunk_interval_k, + num_cached_blocks=cfg.num_cached_blocks, + sink_token=cfg.sink_token, + num_frame_per_block=cfg.num_frame_per_block, + denoising_step_list=self._denoising_steps(), + save_stage1=cfg.save_stage1, + ) - stage1_video = result.get("stage1_video") - if stage1_video is not None: - stage1_hwc = stage1_video + result = pipeline.generate(cropped, prompt, c2w, intrinsics_vec4, params) + video_hwc = result["video"] if not cfg.no_action_overlay: - stage1_hwc = upstream.apply_overlay(stage1_hwc, result["stage1_c2w"]) + video_hwc = upstream.apply_overlay(video_hwc, result["c2w"]) + + if not self.is_rank_zero: + return upstream.write_video( cfg.output_dir, - f"{cfg.name}_stage1", - stage1_hwc, + cfg.name, + video_hwc, params.fps, upstream.get_root_logger(), ) + stage1_video = result.get("stage1_video") + if stage1_video is not None: + stage1_hwc = stage1_video + if not cfg.no_action_overlay: + stage1_hwc = upstream.apply_overlay( + stage1_hwc, result["stage1_c2w"] + ) + upstream.write_video( + cfg.output_dir, + f"{cfg.name}_stage1", + stage1_hwc, + params.fps, + upstream.get_root_logger(), + ) + + +def _precision_env_updates( + *, + stage1_precision: Precision, + refiner_precision: Precision, + refiner_enabled: bool, +) -> dict[str, str | None]: + """Return upstream SANA-WM precision env overrides.""" + updates: dict[str, str | None] = {"SANA_WM_STAGE1_NVFP4": None} + if stage1_precision != "bf16": + updates.update( + { + "SANA_WM_STAGE1_NVFP4": "self_attn+cross+ffn", + "SANA_WM_STAGE1_NVFP4_SCOPE": "block", + "SANA_WM_STAGE1_LINEARIZE_FFN": None, + "SANA_WM_STAGE1_QUANT": _quant_env_value(stage1_precision), + } + ) + + updates["SANA_WM_REFINER_NVFP4"] = None + if refiner_enabled and refiner_precision != "bf16": + updates.update( + { + "SANA_WM_REFINER_NVFP4": "1", + "SANA_WM_REFINER_QUANT": _quant_env_value(refiner_precision), + } + ) + return updates + + +def _quant_env_value(precision: Precision) -> str: + """Return upstream Transformer Engine recipe selector for a precision.""" + if precision == "fp8": + return "fp8block" + if precision == "fp4": + return "nvfp4" + raise ValueError(f"{precision!r} does not use Transformer Engine quantization.") + + +@contextmanager +def _temporary_environment( + updates: Mapping[str, str | None], +) -> Iterator[None]: + """Temporarily set or unset environment variables.""" + previous = {key: os.environ.get(key) for key in updates} + try: + for key, value in updates.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + yield + finally: + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _validate_precision_request( + *, + device: torch.device, + stage1_precision: Precision, + refiner_precision: Precision, + refiner_enabled: bool, +) -> None: + """Fail early when requested quantized precision cannot run.""" + active_precisions = [stage1_precision] + if refiner_enabled: + active_precisions.append(refiner_precision) + quantized = [precision for precision in active_precisions if precision != "bf16"] + if not quantized: + return + + if device.type != "cuda" or not torch.cuda.is_available(): + raise ValueError( + "SANA-WM fp8/fp4 precision requires a CUDA device; " + f"resolved device is {device}." + ) + + capability = torch.cuda.get_device_capability(device) + major, minor = capability + sm_name = f"sm_{major}{minor}" + if "fp8" in quantized and major < 9: + raise ValueError( + "SANA-WM fp8 precision requires a Hopper or newer GPU " + f"(sm_90+); detected {sm_name}." + ) + if "fp8" in quantized and not _torch_cuda_version_at_least(12, 9): + cuda_version = torch.version.cuda or "unknown" + raise ValueError( + "SANA-WM fp8 precision uses upstream Sana's Transformer Engine " + "Float8BlockScaling path, which requires CUDA 12.9 or newer. " + f"This PyTorch environment reports CUDA {cuda_version}." + ) + if "fp4" in quantized and major < 10: + raise ValueError( + "SANA-WM fp4/NVFP4 precision requires a Blackwell GPU " + f"(sm_100+); detected {sm_name}. Use bf16 or fp8 on this GPU." + ) + + _validate_transformer_engine(quantized) + + +def _torch_cuda_version_at_least(major: int, minor: int) -> bool: + """Return whether ``torch.version.cuda`` is at least ``major.minor``.""" + version = torch.version.cuda + if version is None: + return False + parts = version.split(".") + try: + current_major = int(parts[0]) + current_minor = int(parts[1]) if len(parts) > 1 else 0 + except (ValueError, IndexError): + return False + return (current_major, current_minor) >= (major, minor) + + +def _validate_transformer_engine(precisions: list[Precision]) -> None: + """Ensure Transformer Engine has the recipes required by precision modes.""" + required_recipes: set[str] = set() + if "fp8" in precisions: + required_recipes.add("Float8BlockScaling") + if "fp4" in precisions: + required_recipes.add("NVFP4BlockScaling") + try: + import transformer_engine.common.recipe as te_recipe + import transformer_engine.pytorch # noqa: F401 + except Exception as exc: + raise RuntimeError( + "SANA-WM fp8/fp4 precision requires NVIDIA Transformer Engine. " + "Install it in the upstream Sana environment, for example with " + "`pip install --no-build-isolation 'transformer_engine[pytorch]'`, " + "or run with the default bf16 precision." + ) from exc + + missing = sorted( + recipe for recipe in required_recipes if not hasattr(te_recipe, recipe) + ) + if missing: + raise RuntimeError( + "Installed Transformer Engine does not provide the recipe(s) " + f"required by the requested SANA-WM precision: {', '.join(missing)}." + ) + def _candidate_upstream_roots(explicit_root: Path | None) -> list[Path]: """Return candidate NVlabs/Sana checkout roots in priority order.""" @@ -429,6 +610,7 @@ def _import_upstream_sana(explicit_root: Path | None) -> ModuleType: __all__ = [ + "Precision", "SamplingAlgo", "SanaWMRunner", "SanaWMRunnerConfig", diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index c361941bc..6c7c032ed 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -17,9 +17,11 @@ from __future__ import annotations +import os from pathlib import Path import pytest +import torch try: import tomllib @@ -36,7 +38,13 @@ SANA_WM_HF_REPO, SANA_WM_MODEL_PATH, ) -from sana_wm.runner import SanaWMRunner, SanaWMRunnerConfig +from sana_wm.runner import ( + SanaWMRunner, + SanaWMRunnerConfig, + _precision_env_updates, + _temporary_environment, + _validate_precision_request, +) from flashdreams.infra.config import derive_config @@ -93,6 +101,113 @@ def test_runner_config_type() -> None: assert isinstance(RUNNER_SANA_WM_BIDIRECTIONAL, SanaWMRunnerConfig) +def test_runner_defaults_to_bf16_precision() -> None: + """Keep the default path on upstream SANA-WM's BF16 config.""" + assert RUNNER_SANA_WM_BIDIRECTIONAL.stage1_precision == "bf16" + assert RUNNER_SANA_WM_BIDIRECTIONAL.refiner_precision == "bf16" + + +def test_bf16_precision_clears_upstream_quant_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Prevent external NVFP4 env vars from changing the default runner path.""" + monkeypatch.setenv("SANA_WM_STAGE1_NVFP4", "1") + monkeypatch.setenv("SANA_WM_REFINER_NVFP4", "1") + updates = _precision_env_updates( + stage1_precision="bf16", + refiner_precision="bf16", + refiner_enabled=True, + ) + + with _temporary_environment(updates): + assert "SANA_WM_STAGE1_NVFP4" not in os.environ + assert "SANA_WM_REFINER_NVFP4" not in os.environ + + assert os.environ["SANA_WM_STAGE1_NVFP4"] == "1" + assert os.environ["SANA_WM_REFINER_NVFP4"] == "1" + + +def test_fp8_precision_sets_upstream_quant_env() -> None: + """Map FlashDreams fp8 fields to upstream SANA-WM env selectors.""" + updates = _precision_env_updates( + stage1_precision="fp8", + refiner_precision="fp8", + refiner_enabled=True, + ) + + assert updates["SANA_WM_STAGE1_NVFP4"] == "self_attn+cross+ffn" + assert updates["SANA_WM_STAGE1_QUANT"] == "fp8block" + assert updates["SANA_WM_STAGE1_LINEARIZE_FFN"] is None + assert updates["SANA_WM_REFINER_NVFP4"] == "1" + assert updates["SANA_WM_REFINER_QUANT"] == "fp8block" + + +def test_fp4_precision_sets_upstream_quant_env() -> None: + """Map FlashDreams fp4 fields to upstream SANA-WM NVFP4 env selectors.""" + updates = _precision_env_updates( + stage1_precision="fp4", + refiner_precision="fp4", + refiner_enabled=True, + ) + + assert updates["SANA_WM_STAGE1_QUANT"] == "nvfp4" + assert updates["SANA_WM_REFINER_QUANT"] == "nvfp4" + + +def test_quantized_precision_requires_cuda() -> None: + """Reject fp8/fp4 before importing upstream Sana on CPU-only devices.""" + with pytest.raises(ValueError, match="requires a CUDA device"): + _validate_precision_request( + device=torch.device("cpu"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + ) + + +def test_fp8_precision_requires_hopper(monkeypatch: pytest.MonkeyPatch) -> None: + """Reject FP8 on pre-Hopper GPUs before checking Transformer Engine.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (8, 9)) + + with pytest.raises(ValueError, match="requires a Hopper or newer GPU"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + ) + + +def test_fp8_precision_requires_cuda_129(monkeypatch: pytest.MonkeyPatch) -> None: + """Reject upstream FP8 block scaling when torch reports CUDA 12.8.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch.version, "cuda", "12.8") + + with pytest.raises(ValueError, match="requires CUDA 12.9 or newer"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + ) + + +def test_fp4_precision_requires_blackwell(monkeypatch: pytest.MonkeyPatch) -> None: + """Reject NVFP4 on Hopper-class GPUs before checking Transformer Engine.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0)) + + with pytest.raises(ValueError, match="requires a Blackwell GPU"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp4", + refiner_precision="bf16", + refiner_enabled=True, + ) + + def test_pyproject_entry_point_matches_runner_literal() -> None: """Keep the package entry point aligned with ``RUNNER_CONFIGS``.""" pyproject = tomllib.loads( From a710bd0043ee1035bbef44ec4af77d5f5e685518 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 16 Jul 2026 16:47:40 -0700 Subject: [PATCH 06/36] Add native SANA-WM Stage-1 FlashDreams integration Route the public SANA-WM runner through FlashDreams pipeline, diffusion-model, transformer, and scheduler boundaries instead of the upstream inference harness. Keep the upstream harness only as an explicit reference runner with separate documentation. Replace the vendored upstream diffusion package with a FlashDreams-owned Stage-1 DiT module, direct checkpoint loading, local YAML config parsing, direct diffusers LTX2 VAE encode/decode, and native runner wiring for Stage-1 BF16 execution. Add Transformer Engine-free native FP8 and FP4 Stage-1 Linear replacements using PyTorch scaled-mm and Triton NVFP4 activation quantization, with precision/backend validation before checkpoint loading. Keep the native scope honest: require no-refiner for native runs, keep action overlays and upstream TE quantization on the separate upstream-reference harness, and document the remaining refiner and GDN parity gaps separately from the native runner docs. Expand CPU-safe SANA tests to pin native/reference runner separation, config loading without upstream types, native Stage-1 checkpoint schema, Stage-1 forward shape, runtime release, VAE tiling safeguards, and native quant replacement behavior. Signed-off-by: Aidan Foster --- integrations/sana/README.md | 176 ++- integrations/sana/docs/upstream_reference.md | 103 ++ integrations/sana/pyproject.toml | 15 +- integrations/sana/sana_wm/__init__.py | 2 +- integrations/sana/sana_wm/_tools.py | 86 ++ integrations/sana/sana_wm/config.py | 73 +- integrations/sana/sana_wm/native_diffusion.py | 228 ++++ integrations/sana/sana_wm/native_quant.py | 564 +++++++++ .../sana/sana_wm/native_transformer.py | 1085 +++++++++++++++++ integrations/sana/sana_wm/runner.py | 403 +++++- integrations/sana/sana_wm/stage1_model.py | 514 ++++++++ .../{placeholder.py => upstream_reference.py} | 41 +- .../sana/tests/parity_check/README.md | 12 +- .../sana/tests/test_native_quant_cuda.py | 79 ++ integrations/sana/tests/test_smoke.py | 581 ++++++++- 15 files changed, 3785 insertions(+), 177 deletions(-) create mode 100644 integrations/sana/docs/upstream_reference.md create mode 100644 integrations/sana/sana_wm/_tools.py create mode 100644 integrations/sana/sana_wm/native_diffusion.py create mode 100644 integrations/sana/sana_wm/native_quant.py create mode 100644 integrations/sana/sana_wm/native_transformer.py create mode 100644 integrations/sana/sana_wm/stage1_model.py rename integrations/sana/sana_wm/{placeholder.py => upstream_reference.py} (55%) create mode 100644 integrations/sana/tests/test_native_quant_cuda.py diff --git a/integrations/sana/README.md b/integrations/sana/README.md index 705d721a3..aa0500e82 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -5,175 +5,135 @@ SPDX-License-Identifier: Apache-2.0 # `sana_wm` -FlashDreams runner shim for +FlashDreams SANA-WM integration for [SANA-WM](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional), -the 2.6B bidirectional camera-controlled world model from the -[NVlabs/Sana](https://github.com/NVlabs/Sana) repository. +the 2.6B bidirectional camera-controlled world model released from NVlabs/Sana. -This first slice follows the dependency-isolated integration pattern: -FlashDreams owns the CLI registration, config surface, and CPU-tested -camera/action helpers, while Stage-1 DiT, LTX2 VAE/refiner, Pi3X, and -custom attention kernels are imported from an installed or local upstream -Sana checkout at runtime. +The `sana-wm-bidirectional` runner is the native FlashDreams path. It uses +FlashDreams config, runner, pipeline, diffusion-model, scheduler, transformer, +camera-conditioning, VAE decode, and output-writing boundaries. The Stage-1 +DiT module is implemented in this package and loads the public SANA-WM +checkpoint directly; it does not import or install a vendored upstream +`diffusion` package. + +Current native scope: + +| component | status | +| --- | --- | +| Stage-1 BF16 | Native FlashDreams path, GPU smoke-tested. | +| Stage-1 FP8 | Native PyTorch `_scaled_mm` backend, GPU smoke-tested. | +| Stage-1 FP4 | Native Triton quantization plus PyTorch `_scaled_mm`, GPU smoke-tested. | +| VAE decode | Native direct `diffusers` LTX2 VAE use with low-memory tiling. | +| LTX-2 refiner | Not native yet; the native runner requires `--no-refiner True`. | +| GDN attention parity | Not complete; smoke-tested, not quality-parity validated. | + +The separate `sana-wm-upstream-reference` runner is documented only in +`docs/upstream_reference.md`. Keep reference-harness details out of this +native integration README. ## Runner | slug | description | | --- | --- | -| `sana-wm-bidirectional` | SANA-WM bidirectional I2V world model using upstream NVlabs/Sana Stage-1 + LTX-2 refiner. | +| `sana-wm-bidirectional` | Native FlashDreams SANA-WM Stage-1 runner. | +| `sana-wm-native-bidirectional` | Alias for the native FlashDreams runner. | -The FlashDreams package is named `sana_wm` rather than `sana` so it does -not shadow upstream's own top-level `sana` package. +The FlashDreams package is named `sana_wm` rather than `sana` so it does not +shadow the upstream project name. ## Setup -Clone Sana next to FlashDreams or pass its path explicitly: +Install FlashDreams and the SANA integration into an environment with the +project's GPU runtime dependencies: ```bash -git clone https://github.com/NVlabs/Sana ../Sana -cd ../Sana -bash environment_setup.sh sana -conda activate sana +uv sync --package flashdreams-sana-wm --extra dev ``` -The runner auto-detects `../Sana`. You can also set `SANA_ROOT` or pass -`--upstream-sana-root /path/to/Sana`. - -Install the FlashDreams packages into the same environment without -letting pip resolve or upgrade dependencies. The upstream Sana environment -owns the model runtime stack and pins packages such as `transformers` and -`huggingface-hub`. +For an existing environment, install the local packages directly: ```bash -cd ../flashdreams -python -m pip install --no-deps -e flashdreams -e integrations/sana +python -m pip install -e flashdreams -e integrations/sana ``` -If FlashDreams was installed without `--no-deps` and pip upgraded Sana's -pinned packages, repair the env before running generation: - -```bash -python -m pip install --no-deps \ - "huggingface-hub==0.36.0" \ - "transformers==4.57.3" -python -m pip install --no-deps -e flashdreams -e integrations/sana -``` +The native runner does not require a local NVlabs/Sana checkout. The checkout +is only useful as a source of demo assets such as +`asset/sana_wm/demo_0.png`, `demo_0.txt`, `demo_0_pose.npy`, and +`demo_0_intrinsics.npy`. ## Run -```bash -flashdreams-run sana-wm-bidirectional \ - --image-path ../Sana/asset/sana_wm/demo_0.png \ - --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ - --action "w-100,dw-60,w-100,aw-60" \ - --num-frames 321 \ - --output-dir outputs/sana_wm -``` - -To use an explicit trajectory: +The native runner currently requires explicit intrinsics and no refiner: ```bash +PYTORCH_ALLOC_CONF=expandable_segments:True \ flashdreams-run sana-wm-bidirectional \ - --image-path ../Sana/asset/sana_wm/demo_0.png \ - --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ - --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ - --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ - --num-frames 321 \ - --output-dir outputs/sana_wm -``` - -Set `--no-refiner True` to skip the LTX-2 refiner and decode Stage-1 -latents directly, matching upstream's fast debugging path. FlashDreams' -current CLI expects explicit boolean values, so pass `True` or `False` -for boolean fields. Full generation still downloads the public Hugging -Face artefacts on first use. - -## Validated BF16 smoke tests - -The following commands were validated on an NVIDIA RTX PRO 6000 Blackwell -with upstream Sana installed in the `sana` conda environment. They use -the default BF16 path and do not require Transformer Engine. - -Stage-1 plus VAE decode: - -```bash -flashdreams-run sana-wm-bidirectional \ - --upstream-sana-root ../Sana \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ - --output-dir outputs/sana_wm \ + --output-dir outputs/sana_wm_native_bf16 \ --no-refiner True ``` Expected output: ```text -outputs/sana_wm/sana_wm_generated.mp4 +outputs/sana_wm_native_bf16/sana_wm_generated.mp4 ``` -Stage-1 plus LTX-2 refiner: +## FP8 and FP4 + +The runner defaults to BF16. Quantized Stage-1 inference is opt-in and uses +native FlashDreams/PyTorch backends by default, not Transformer Engine. + +| option | hardware | backend | +| --- | --- | --- | +| `--stage1-precision fp8` | Hopper or newer (`sm_90+`) | native E4M3 `_scaled_mm` | +| `--stage1-precision fp4` | Blackwell (`sm_100+`) | native Triton NVFP4 plus `_scaled_mm` | +| `--quant-backend torch-fp8` | Hopper or newer (`sm_90+`) | force the native FP8 backend | +| `--quant-backend torch-fp4` | Blackwell (`sm_100+`) | force the native FP4 backend | + +FP8 smoke: ```bash +PYTORCH_ALLOC_CONF=expandable_segments:True \ flashdreams-run sana-wm-bidirectional \ - --upstream-sana-root ../Sana \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ - --output-dir outputs/sana_wm_refiner -``` - -Expected output: - -```text -outputs/sana_wm_refiner/sana_wm_generated.mp4 + --output-dir outputs/sana_wm_native_fp8 \ + --no-refiner True \ + --stage1-precision fp8 ``` -## FP8 and FP4 precision - -The runner defaults to BF16. Quantized inference is opt-in: - -| option | hardware | dependency | -| --- | --- | --- | -| `--stage1-precision fp8` / `--refiner-precision fp8` | Hopper or newer (`sm_90+`) with CUDA 12.9+ | Transformer Engine | -| `--stage1-precision fp4` / `--refiner-precision fp4` | Blackwell (`sm_100+`) | Transformer Engine with NVFP4 recipes | - -The runner validates these requirements before loading checkpoints. If the -GPU or Transformer Engine install cannot support the requested precision, -it fails early with a targeted error instead of falling through to a deep -upstream import or model-build failure. - -Example Blackwell FP4 run: +FP4 smoke: ```bash +PYTORCH_ALLOC_CONF=expandable_segments:True \ flashdreams-run sana-wm-bidirectional \ - --upstream-sana-root ../Sana \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ - --output-dir outputs/sana_wm_fp4 \ - --stage1-precision fp4 \ - --refiner-precision fp4 + --output-dir outputs/sana_wm_native_fp4 \ + --no-refiner True \ + --stage1-precision fp4 ``` -Use `fp8` for Hopper-class GPUs when the PyTorch environment reports CUDA -12.9 or newer. Keep both precision flags at `bf16` if Transformer Engine is -unavailable or if maximum compatibility is preferred. - ## Tests -CPU-safe tests cover import, runner config, action parsing, intrinsics -handling, frame snapping, and SANA-WM camera-conditioning tensor shapes: +CPU-safe tests cover import, config boundaries, action parsing, intrinsics, +camera conditioning, native Stage-1 checkpoint schema, native Stage-1 CPU +forward shape, VAE tiling, and native low-precision backend selection: ```bash -uv run --extra dev pytest integrations/sana/tests +uv run --extra dev pytest integrations/sana/tests/test_smoke.py ``` -Full upstream parity and generated-video checks are heavyweight GPU -workflows and should live under `tests/parity_check/`. +GPU generation and native-vs-reference quality parity checks are heavyweight +manual workflows and should stay out of `ci_cpu`. diff --git a/integrations/sana/docs/upstream_reference.md b/integrations/sana/docs/upstream_reference.md new file mode 100644 index 000000000..21a339fc0 --- /dev/null +++ b/integrations/sana/docs/upstream_reference.md @@ -0,0 +1,103 @@ + + +# SANA-WM Upstream Reference Harness + +This document covers the temporary upstream-reference runner. It exists only +for parity/debug work while the native FlashDreams integration is being +validated. Do not use this as user-facing documentation for the real SANA-WM +FlashDreams integration. + +## Runner + +| slug | description | +| --- | --- | +| `sana-wm-upstream-reference` | FlashDreams-packaged NVlabs/Sana reference harness. | + +This runner delegates generation to the SANA reference pipeline object. It is +expected to be removed once native generation is validated. + +## Setup + +The reference harness needs a local or importable NVlabs/Sana checkout: + +```bash +git clone https://github.com/NVlabs/Sana ../Sana +cd ../Sana +bash environment_setup.sh sana +conda activate sana +``` + +The runner auto-detects `../Sana`. You can also set `SANA_ROOT` or pass +`--upstream-sana-root /path/to/Sana`. + +## Run + +```bash +flashdreams-run sana-wm-upstream-reference \ + --upstream-sana-root ../Sana \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 321 \ + --output-dir outputs/sana_wm_reference +``` + +The upstream reference harness can estimate intrinsics with Pi3X when +`--intrinsics-path` is omitted. The native runner currently requires explicit +intrinsics instead. + +## Transformer Engine Reference Notes + +Upstream SANA does not use Transformer Engine here for positional embeddings. +Where the reference path uses Transformer Engine, it is as a quantized +`Linear`/GEMM backend: + +- Stage 1 replaces selected `torch.nn.Linear` layers with + `transformer_engine.pytorch.Linear`, then wraps transformer block or linear + forwards in `te.fp8_autocast(...)`. +- `SANA_WM_STAGE1_QUANT=nvfp4` selects `NVFP4BlockScaling`; `fp8block` selects + `Float8BlockScaling`. +- The default FlashDreams FP4 Stage-1 selector maps to upstream + `self_attn+cross+ffn`, covering self-attention, cross-attention, and MLP + linear layers. In the validated reference demo run, upstream reported 120 + converted linear layers and 20 wrapped transformer blocks. +- The refiner applies the same pattern to eligible LTX-2 transformer linear + layers after skipping input/output projection, audio-only, caption, and time + embedding modules. + +Use `--quant-backend upstream-te` only on this reference harness when comparing +native low-precision behavior against upstream Transformer Engine behavior. + +```bash +flashdreams-run sana-wm-upstream-reference \ + --upstream-sana-root ../Sana \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_reference_fp4 \ + --stage1-precision fp4 \ + --refiner-precision fp4 \ + --quant-backend upstream-te +``` + +FlashDreams already removed a separate Transformer Engine dependency for RoPE +via `flashdreams.core.attention.rope_kernel`, but that replacement is not a +drop-in for SANA FP4/FP8. Native SANA replaces eligible Stage-1 and refiner +`Linear` modules directly, including activation quantization/scales, weight +conversion, GEMM, and output rescaling. + +Relevant FlashDreams components for future hardening: + +- `flashdreams.core.attention.rope_kernel`: the prior narrow TE-replacement + pattern and test strategy. +- `integrations/omnidreams/omnidreams_singleview/python/cosmos_fp8_utils.py`: + FP8 weight quantization and prepared weight/scale aliases for Cosmos blocks. +- `integrations/omnidreams/omnidreams_singleview/src/dit_streaming/kernels/`: + CUDA/CUTLASS FP8 GEMM, BF16 GEMM, INT8 block quantization, and SageAttention3 + FP4 attention/cache kernels. diff --git a/integrations/sana/pyproject.toml b/integrations/sana/pyproject.toml index d0baed354..df7ebf34b 100644 --- a/integrations/sana/pyproject.toml +++ b/integrations/sana/pyproject.toml @@ -20,15 +20,18 @@ build-backend = "setuptools.build_meta" [project] name = "flashdreams-sana-wm" version = "0.1.0" -description = "SANA-WM bidirectional world-model runner shim for flashdreams." +description = "SANA-WM bidirectional world-model integration for flashdreams." readme = "README.md" requires-python = ">=3.10" -# The upstream SANA repository owns the Stage-1 DiT, LTX2 VAE/refiner, -# Pi3X intrinsics estimator, and custom attention kernels. Those heavy -# dependencies are intentionally installed from the upstream SANA checkout, -# not pulled into every FlashDreams workspace sync. dependencies = [ + "diffusers>=0.36", "flashdreams", + "imageio[ffmpeg]>=2.31", + "Pillow>=10", + "PyYAML>=6.0", + "safetensors>=0.5", + "torchvision>=0.26", + "transformers>=5.0,<6", ] [tool.uv.sources] @@ -41,6 +44,8 @@ dev = [ [project.entry-points."flashdreams.runner_configs"] "sana-wm-bidirectional" = "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL" +"sana-wm-upstream-reference" = "sana_wm.config:RUNNER_SANA_WM_UPSTREAM_REFERENCE" +"sana-wm-native-bidirectional" = "sana_wm.config:RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL" [tool.setuptools.packages.find] include = ["sana_wm*"] diff --git a/integrations/sana/sana_wm/__init__.py b/integrations/sana/sana_wm/__init__.py index f4b7dbd28..77dfd6e32 100644 --- a/integrations/sana/sana_wm/__init__.py +++ b/integrations/sana/sana_wm/__init__.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SANA-WM bidirectional runner shim for FlashDreams.""" +"""SANA-WM native and upstream-reference runners for FlashDreams.""" from sana_wm.config import RUNNER_CONFIGS, RUNNER_SANA_WM_BIDIRECTIONAL diff --git a/integrations/sana/sana_wm/_tools.py b/integrations/sana/sana_wm/_tools.py new file mode 100644 index 000000000..7562bfcba --- /dev/null +++ b/integrations/sana/sana_wm/_tools.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Local file and checkpoint helpers for the SANA-WM integration.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import torch + +HF_URI_SCHEME = "hf://" + + +def resolve_hf_path(path: str | Path) -> str: + """Resolve a local path or ``hf://owner/repo/subpath`` URI to a local path.""" + path_str = str(path) + if not path_str or Path(path_str).exists(): + return path_str + if not path_str.startswith(HF_URI_SCHEME): + return path_str + + from huggingface_hub import snapshot_download + + parts = path_str[len(HF_URI_SCHEME) :].split("/", 2) + if len(parts) < 2 or not parts[0] or not parts[1]: + raise ValueError( + f"Invalid HF path {path_str!r}; expected hf:///[/]." + ) + repo_id = f"{parts[0]}/{parts[1]}" + subpath = parts[2] if len(parts) > 2 else "" + allow_patterns = None + if subpath: + allow_patterns = [subpath, f"{subpath}/*", f"{subpath}/**"] + local_root = snapshot_download(repo_id=repo_id, allow_patterns=allow_patterns) + return os.path.join(local_root, subpath) if subpath else local_root + + +def find_model(model_name: str) -> dict[str, Any]: + """Load a SANA checkpoint from a local path or resolved HF artefact.""" + resolved = resolve_hf_path(model_name) + if not os.path.isfile(resolved): + raise FileNotFoundError(f"Could not find SANA checkpoint at {resolved}") + + if resolved.endswith(".safetensors"): + import safetensors.torch + + return {"state_dict": safetensors.torch.load_file(resolved, device="cpu")} + if resolved.endswith(".safetensors.index.json"): + import safetensors.torch + + with open(resolved, encoding="utf-8") as handle: + index = json.load(handle)["weight_map"] + state_dict = {} + for shard in sorted(set(index.values())): + shard_path = os.path.join(os.path.dirname(resolved), shard) + state_dict.update(safetensors.torch.load_file(shard_path, device="cpu")) + return {"state_dict": state_dict} + return torch.load(resolved, map_location=lambda storage, _loc: storage) + + +def hf_download_or_fpath(path: str | Path) -> str: + """Compatibility alias retained for local helper call sites.""" + return resolve_hf_path(path) + + +__all__ = [ + "find_model", + "hf_download_or_fpath", + "resolve_hf_path", +] diff --git a/integrations/sana/sana_wm/config.py b/integrations/sana/sana_wm/config.py index ffa72922e..faa9d8aa2 100644 --- a/integrations/sana/sana_wm/config.py +++ b/integrations/sana/sana_wm/config.py @@ -13,44 +13,99 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Static configs for the SANA-WM bidirectional runner.""" +"""Static configs for SANA-WM native and upstream-reference runners.""" from __future__ import annotations +from typing import cast + +from flashdreams.infra.config import derive_config from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler import FlowMatchSchedulerConfig from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.infra.runner import RunnerConfig -from sana_wm.placeholder import SanaWMPlaceholderTransformerConfig +from sana_wm.native_diffusion import SanaWMDiffusionModelConfig +from sana_wm.native_transformer import SanaWMNativeTransformerConfig from sana_wm.runner import SanaWMRunnerConfig +from sana_wm.upstream_reference import SanaWMUpstreamReferenceTransformerConfig + + +def _reference_pipeline(name: str) -> StreamInferencePipelineConfig: + return StreamInferencePipelineConfig( + name=name, + diffusion_model=DiffusionModelConfig( + transformer=SanaWMUpstreamReferenceTransformerConfig(), + scheduler=FlowMatchSchedulerConfig(), + seed=42, + ), + ) + + +PIPELINE_SANA_WM_UPSTREAM_REFERENCE = _reference_pipeline("sana-wm-upstream-reference") +"""Schema marker for the upstream-delegating SANA-WM reference runner.""" PIPELINE_SANA_WM_BIDIRECTIONAL = StreamInferencePipelineConfig( name="sana-wm-bidirectional", - diffusion_model=DiffusionModelConfig( - transformer=SanaWMPlaceholderTransformerConfig(), + diffusion_model=SanaWMDiffusionModelConfig( + transformer=SanaWMNativeTransformerConfig(), scheduler=FlowMatchSchedulerConfig(), seed=42, ), ) -"""Schema placeholder for the upstream-delegating SANA-WM runner.""" +"""Native FlashDreams SANA-WM pipeline.""" + +PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL = cast( + StreamInferencePipelineConfig, + derive_config(PIPELINE_SANA_WM_BIDIRECTIONAL, name="sana-wm-native-bidirectional"), +) +"""Backward-compatible native SANA-WM pipeline alias.""" + +RUNNER_SANA_WM_UPSTREAM_REFERENCE = SanaWMRunnerConfig( + runner_name=PIPELINE_SANA_WM_UPSTREAM_REFERENCE.name, + description=( + "SANA-WM upstream reference harness (NVlabs/Sana Stage-1 + LTX-2 refiner)." + ), + pipeline=PIPELINE_SANA_WM_UPSTREAM_REFERENCE, +) +"""Explicit SANA-WM upstream reference runner config.""" RUNNER_SANA_WM_BIDIRECTIONAL = SanaWMRunnerConfig( runner_name=PIPELINE_SANA_WM_BIDIRECTIONAL.name, description=( - "SANA-WM bidirectional I2V world model (upstream NVlabs/Sana " - "Stage-1 + LTX-2 refiner runner)." + "SANA-WM bidirectional I2V native FlashDreams runner " + "(Stage-1 DiT)." ), pipeline=PIPELINE_SANA_WM_BIDIRECTIONAL, + execution_backend="native-flashdreams", + no_refiner=True, +) +"""SANA-WM native FlashDreams runner config.""" + +RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL = SanaWMRunnerConfig( + runner_name=PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL.name, + description="SANA-WM native FlashDreams pipeline alias.", + pipeline=PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL, + execution_backend="native-flashdreams", + no_refiner=True, ) -"""Default SANA-WM bidirectional runner config.""" +"""Backward-compatible native SANA-WM FlashDreams runner alias.""" RUNNER_CONFIGS: dict[str, RunnerConfig] = { - cfg.runner_name: cfg for cfg in (RUNNER_SANA_WM_BIDIRECTIONAL,) + cfg.runner_name: cfg + for cfg in ( + RUNNER_SANA_WM_BIDIRECTIONAL, + RUNNER_SANA_WM_UPSTREAM_REFERENCE, + RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, + ) } """SANA-WM runner configs keyed by ``runner_name``.""" __all__ = [ "PIPELINE_SANA_WM_BIDIRECTIONAL", + "PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL", + "PIPELINE_SANA_WM_UPSTREAM_REFERENCE", "RUNNER_CONFIGS", "RUNNER_SANA_WM_BIDIRECTIONAL", + "RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL", + "RUNNER_SANA_WM_UPSTREAM_REFERENCE", ] diff --git a/integrations/sana/sana_wm/native_diffusion.py b/integrations/sana/sana_wm/native_diffusion.py new file mode 100644 index 000000000..be719ffab --- /dev/null +++ b/integrations/sana/sana_wm/native_diffusion.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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-specific diffusion model for first-frame-pinned LTX Euler sampling.""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass, field +from typing import Any, cast + +import torch +from loguru import logger +from torch import Tensor + +from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig +from sana_wm.native_transformer import ( + SanaWMNativeTransformer, + SanaWMNativeTransformerCache, +) + + +@dataclass(kw_only=True) +class SanaWMDiffusionModelConfig(DiffusionModelConfig): + """Diffusion model config for SANA-WM's LTX-style Stage-1 sampler.""" + + _target: type["SanaWMDiffusionModel"] = field( + default_factory=lambda: SanaWMDiffusionModel + ) + + +class SanaWMDiffusionModel(DiffusionModel[SanaWMNativeTransformerCache]): + """Run SANA-WM Stage-1 denoising behind the FlashDreams diffusion boundary.""" + + transformer: SanaWMNativeTransformer + + def __init__(self, config: SanaWMDiffusionModelConfig) -> None: + super().__init__(config) + if not isinstance(self.transformer, SanaWMNativeTransformer): + raise TypeError( + "SanaWMDiffusionModel requires SanaWMNativeTransformerConfig." + ) + + def generate( + self, + autoregressive_index: int, + cache: SanaWMNativeTransformerCache, + input: Any = None, + ) -> tuple[Tensor, "DiffusionModel.FinalState[SanaWMNativeTransformerCache]"]: + """Run SANA-WM's first-frame-pinned LTX Euler denoising loop.""" + del input + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional native inference has one AR step.") + conditioning = cache.conditioning + if conditioning is None: + raise RuntimeError("SANA-WM native diffusion cache has no conditioning.") + + cache.start(autoregressive_index) + latents = self.transformer.initial_latents(conditioning) + if torch.cuda.is_available(): + torch.cuda.synchronize() + t0 = time.perf_counter() + clean_latent = self._sample_ltx_euler(latents, cache) + if torch.cuda.is_available(): + torch.cuda.synchronize() + logger.info( + "[timing] stage1 sample: {:.3f}s (latent shape {})", + time.perf_counter() - t0, + tuple(clean_latent.shape), + ) + + final_state = DiffusionModel.FinalState( + clean_latent=clean_latent, + autoregressive_index=autoregressive_index, + cache=cache, + input=None, + ) + return clean_latent, final_state + + def finalize( + self, + final_state: "DiffusionModel.FinalState[SanaWMNativeTransformerCache]", + ) -> None: + """Finalize the one-shot cache without an extra model forward.""" + final_state.cache.finalize(final_state.autoregressive_index) + + def _sample_ltx_euler( + self, + latents: Tensor, + cache: SanaWMNativeTransformerCache, + ) -> Tensor: + conditioning = cache.conditioning + assert conditioning is not None + + from diffusers import FlowMatchEulerDiscreteScheduler + from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import ( + retrieve_timesteps, + ) + from tqdm import tqdm + + scheduler = FlowMatchEulerDiscreteScheduler(shift=conditioning.flow_shift) + timesteps, _ = retrieve_timesteps( + scheduler, + conditioning.steps, + self.device, + None, + ) + do_cfg = conditioning.cfg_scale > 1.0 + condition_frame_info = dict( + cast( + dict[int, float], + cast(dict[str, object], conditioning.model_kwargs["data_info"]).get( + "condition_frame_info", + {}, + ), + ) + ) + condition_mask = torch.zeros_like(latents) + image_cond_noise_scale = 0.0 + for frame_idx, frame_weight in condition_frame_info.items(): + condition_mask[:, :, int(frame_idx)] = 1 + image_cond_noise_scale = max(image_cond_noise_scale, float(frame_weight)) + + prompt_embeds = conditioning.condition + if do_cfg: + if conditioning.uncondition is None: + raise RuntimeError("CFG was requested without negative prompt embeds.") + prompt_embeds = torch.cat( + [conditioning.uncondition, conditioning.condition], + dim=0, + ) + + init_latents = latents.clone() + generator = torch.Generator(device=self.device).manual_seed(conditioning.seed) + iterator = enumerate(timesteps) + if os.getenv("DPM_TQDM", "False") != "True": + iterator = tqdm(list(iterator)) + + for _, timestep_scalar in iterator: + if image_cond_noise_scale > 0: + latents = _add_noise_to_conditioning_latents( + t=timestep_scalar / 1000.0, + init_latents=init_latents, + latents=latents, + noise_scale=image_cond_noise_scale, + conditioning_mask=condition_mask, + generator=generator, + ) + + condition_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 = timestep_scalar.expand(condition_mask_input.shape).float() + timestep = torch.min(timestep, (1 - condition_mask_input) * 1000.0) + + noise_pred = self.transformer.predict_flow( + noisy_latent=latent_model_input, + timestep=timestep[:, :1, :, 0, 0], + cache=cache, + input=prompt_embeds, + ) + + if do_cfg: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + conditioning.cfg_scale * ( + noise_pred_text - noise_pred_uncond + ) + timestep = timestep.chunk(2)[0] + + latents_dtype = latents.dtype + latents_shape = latents.shape + batch_size, channels, _frames, _height, _width = latents_shape + denoised_latents = scheduler.step( + -noise_pred.reshape(batch_size, channels, -1).transpose(1, 2), + timestep_scalar, + latents.reshape(batch_size, channels, -1).transpose(1, 2), + per_token_timesteps=timestep.reshape(batch_size, channels, -1)[:, 0], + return_dict=False, + )[0] + denoised_latents = denoised_latents.transpose(1, 2).reshape(latents_shape) + tokens_to_denoise_mask = timestep_scalar / 1000 - 1e-6 < ( + 1.0 - condition_mask + ) + latents = torch.where(tokens_to_denoise_mask, denoised_latents, latents) + if latents.dtype != latents_dtype: + latents = latents.to(latents_dtype) + + return latents.detach() + + +def _add_noise_to_conditioning_latents( + *, + t: Tensor, + init_latents: Tensor, + latents: Tensor, + noise_scale: float, + conditioning_mask: Tensor, + generator: torch.Generator, + eps: float = 1e-6, +) -> Tensor: + from diffusers.utils.torch_utils import randn_tensor + + noise = randn_tensor( + latents.shape, + generator=generator, + device=latents.device, + dtype=latents.dtype, + ) + need_to_noise = conditioning_mask > (1.0 - eps) + noised_latents = init_latents + noise_scale * noise * (t**2) + return torch.where(need_to_noise, noised_latents, latents) + + +__all__ = ["SanaWMDiffusionModel", "SanaWMDiffusionModelConfig"] diff --git a/integrations/sana/sana_wm/native_quant.py b/integrations/sana/sana_wm/native_quant.py new file mode 100644 index 000000000..50b878223 --- /dev/null +++ b/integrations/sana/sana_wm/native_quant.py @@ -0,0 +1,564 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Native low-precision linear helpers for SANA-WM. + +Upstream SANA's FP8/FP4 path uses Transformer Engine by replacing selected +``torch.nn.Linear`` modules and entering ``te.fp8_autocast`` around the +corresponding transformer blocks. This module provides TE-free replacements for +the same replacement point: inference-only FP8 and NVFP4 Linear modules backed +by PyTorch's native ``torch._scaled_mm`` kernel, with Triton used for NVFP4 +activation quantization. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any, Literal + +import torch +import triton +import triton.language as tl +from torch import nn + +FP8_MAX_E4M3 = 448.0 +FP8_SCALE_EPS = 1.0e-12 +FP4_MAX_E2M1 = 6.0 +NVFP4_BLOCK_SIZE = 16 +NVFP4_SCALE_EPS = 1.5258789e-05 + + +@dataclass(frozen=True) +class TorchScaledMMFP8Recipe: + """Marker recipe used while patching upstream SANA's TE recipe hooks.""" + + name: str = "torch_scaled_mm_fp8" + precision: Literal["fp8"] = "fp8" + + +@dataclass(frozen=True) +class TorchScaledMMFP4Recipe: + """Marker recipe used while patching upstream SANA's NVFP4 hooks.""" + + name: str = "torch_scaled_mm_fp4" + precision: Literal["fp4"] = "fp4" + + +NativeQuantRecipe = TorchScaledMMFP8Recipe | TorchScaledMMFP4Recipe + + +@triton.jit +def _pack_fp32_to_fp4_pairs(values): + packed = tl.inline_asm_elementwise( + asm=""" + { + .reg .b8 byte0, byte1, byte2, byte3; + cvt.rn.satfinite.e2m1x2.f32 byte0, $5, $1; + cvt.rn.satfinite.e2m1x2.f32 byte1, $6, $2; + cvt.rn.satfinite.e2m1x2.f32 byte2, $7, $3; + cvt.rn.satfinite.e2m1x2.f32 byte3, $8, $4; + mov.b32 $0, {byte0, byte1, byte2, byte3}; + } + """, + constraints=("=r,r,r,r,r,r,r,r,r"), + args=values, + dtype=tl.uint8, + is_pure=True, + pack=4, + ) + return packed + + +@triton.jit +def _quantize_nvfp4_kernel( + input_ptr, + qdata_ptr, + scale_ptr, + stride_m, + stride_k, + rows: tl.constexpr, + cols: tl.constexpr, + mask_scales: tl.constexpr, +): + fp4_max = 6.0 + fp8_max = 448.0 + scale_eps = 1.5258789e-05 + + pid_k = tl.program_id(0) + pid_m = tl.program_id(1) + + offsets_m = pid_m * 128 + tl.arange(0, 128)[:, None] + offsets_k = pid_k * 64 + tl.arange(0, 64)[None, :] + mask = (offsets_m < rows) & (offsets_k < cols) + values = tl.load( + input_ptr + offsets_m * stride_m + offsets_k * stride_k, + mask=mask, + other=0.0, + ) + values = values.to(tl.float32).reshape(128, 4, 16) + block_amax = tl.max(tl.abs(values), axis=2) + scales_f32 = block_amax / fp4_max + scales_f32 = tl.clamp(scales_f32, scale_eps, fp8_max) + scales = scales_f32.to(tl.float8e4nv) + quantized = tl.div_rn(values, scales.to(tl.float32)[:, :, None]) + + if mask_scales: + scale_offsets_k = pid_k * 4 + tl.arange(0, 4)[None, :] + scale_mask = (offsets_m < rows) & (scale_offsets_k < tl.cdiv(cols, 16)) + scales = tl.where(scale_mask, scales, 0.0) + + packed_scales = scales.reshape(4, 32, 4).permute(1, 0, 2).reshape(32, 16) + scale_m = tl.arange(0, 32)[:, None] + scale_k = tl.arange(0, 16)[None, :] + tl.store( + scale_ptr + (pid_m * tl.num_programs(0) + pid_k) * (32 * 16) + scale_m * 16 + scale_k, + packed_scales, + ) + + packed = _pack_fp32_to_fp4_pairs(quantized.reshape(128, 32, 2).split()) + q_offsets_m = pid_m * 128 + tl.arange(0, 128)[:, None] + q_offsets_k = pid_k * 32 + tl.arange(0, 32)[None, :] + q_mask = (q_offsets_m < rows) & (q_offsets_k < cols // 2) + tl.store(qdata_ptr + q_offsets_m * (cols // 2) + q_offsets_k, packed, mask=q_mask) + + +def _ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def _swizzled_nvfp4_scale_shape(rows: int, cols: int) -> tuple[int, int]: + return _ceil_div(rows, 128) * 32, _ceil_div(cols, 64) * 16 + + +def quantize_nvfp4_swizzled(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D CUDA tensor to NVFP4 qdata plus swizzled E4M3 scales.""" + _require_fp4_dtype() + if input.dim() != 2: + raise ValueError(f"NVFP4 quantization requires a 2D tensor, got {input.dim()}D.") + if not input.is_cuda: + raise RuntimeError("NVFP4 quantization requires CUDA input.") + rows, cols = input.shape + if cols % NVFP4_BLOCK_SIZE != 0: + raise ValueError( + f"NVFP4 quantization requires last dim divisible by {NVFP4_BLOCK_SIZE}, got {cols}." + ) + input_2d = input.contiguous() + scale_rows, scale_cols = _swizzled_nvfp4_scale_shape(rows, cols) + qdata = torch.empty((rows, cols // 2), device=input.device, dtype=torch.uint8) + scales = torch.empty( + (scale_rows, scale_cols), + device=input.device, + dtype=torch.float8_e4m3fn, + ) + grid = (_ceil_div(cols, 64), _ceil_div(rows, 128)) + _quantize_nvfp4_kernel[grid]( + input_2d, + qdata, + scales, + input_2d.stride(0), + input_2d.stride(1), + rows, + cols, + mask_scales=(rows % 128 != 0 or cols % 64 != 0), + ) + return qdata, scales + + +class TorchScaledMMFP8Linear(nn.Module): + """Inference-only FP8 Linear using ``torch._scaled_mm``. + + The module stores the source weight as E4M3 FP8 with per-output-channel + scales. At runtime it quantizes the flattened activation rows to E4M3 with + per-row scales, then calls ``torch._scaled_mm`` and returns BF16 output. + This intentionally mirrors the shape contract of ``nn.Linear`` so it can + replace eligible upstream SANA linear layers without changing call sites. + """ + + in_features: int + out_features: int + + def __init__( + self, + *, + weight: torch.Tensor, + weight_fp8: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + ) -> None: + super().__init__() + if weight.shape != weight_fp8.shape: + raise ValueError( + "weight and weight_fp8 must have matching shape, got " + f"{tuple(weight.shape)} and {tuple(weight_fp8.shape)}." + ) + if weight_fp8.dim() != 2: + raise ValueError(f"weight_fp8 must be 2D, got {tuple(weight_fp8.shape)}.") + if weight_scale.shape != (weight_fp8.shape[0], 1): + raise ValueError( + "weight_scale must have shape " + f"({weight_fp8.shape[0]}, 1), got {tuple(weight_scale.shape)}." + ) + if bias is not None and bias.shape != (weight_fp8.shape[0],): + raise ValueError( + f"bias must have shape ({weight_fp8.shape[0]},), got {tuple(bias.shape)}." + ) + self.out_features = int(weight_fp8.shape[0]) + self.in_features = int(weight_fp8.shape[1]) + self.out_dtype = out_dtype + self.register_buffer("weight", weight.detach().contiguous()) + self.register_buffer("weight_fp8", weight_fp8.contiguous()) + self.register_buffer("weight_scale", weight_scale.to(torch.float32).contiguous()) + if bias is None: + self.register_buffer("bias", None) + else: + self.register_buffer("bias", bias.detach().to(out_dtype).contiguous()) + + @classmethod + def from_linear( + cls, + source: nn.Linear, + *, + out_dtype: torch.dtype, + ) -> "TorchScaledMMFP8Linear": + """Create an FP8 replacement from a source ``nn.Linear``.""" + _require_fp8_dtype() + weight_f32 = source.weight.detach().to(torch.float32) + weight_scale = ( + weight_f32.abs().amax(dim=1, keepdim=True).clamp_min(FP8_SCALE_EPS) + / FP8_MAX_E4M3 + ) + weight_fp8 = torch.clamp( + weight_f32 / weight_scale, + -FP8_MAX_E4M3, + FP8_MAX_E4M3, + ).to(torch.float8_e4m3fn) + bias = source.bias.detach() if source.bias is not None else None + replacement = cls( + weight=source.weight.detach().to(device=source.weight.device), + weight_fp8=weight_fp8.to(device=source.weight.device), + weight_scale=weight_scale.to(device=source.weight.device), + bias=bias.to(device=source.weight.device) if bias is not None else None, + out_dtype=out_dtype, + ) + replacement.train(source.training) + return replacement + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Apply the quantized linear projection.""" + if input.shape[-1] != self.in_features: + raise ValueError( + f"expected input last dim {self.in_features}, got {input.shape[-1]}." + ) + if not input.is_cuda: + raise RuntimeError("TorchScaledMMFP8Linear requires CUDA input.") + _require_scaled_mm() + _require_fp8_dtype() + + leading_shape = input.shape[:-1] + input_2d = input.reshape(-1, self.in_features) + if input_2d.numel() == 0: + return input.new_empty( + (*leading_shape, self.out_features), + dtype=self.out_dtype, + ) + + input_f32 = input_2d.to(torch.float32) + input_scale = ( + input_f32.abs().amax(dim=1, keepdim=True).clamp_min(FP8_SCALE_EPS) + / FP8_MAX_E4M3 + ) + input_fp8 = torch.clamp( + input_f32 / input_scale, + -FP8_MAX_E4M3, + FP8_MAX_E4M3, + ).to(torch.float8_e4m3fn) + input_fp8 = input_fp8.contiguous() + + # ``_scaled_mm`` expects the B operand as a column-major transpose view, + # so keep the transpose stride instead of forcing it contiguous. + output = torch._scaled_mm( + input_fp8, + self.weight_fp8.t(), + scale_a=input_scale.contiguous(), + scale_b=self.weight_scale.t().contiguous(), + out_dtype=self.out_dtype, + ) + if self.bias is not None: + output = output + self.bias.to(device=output.device, dtype=output.dtype) + return output.reshape(*leading_shape, self.out_features) + + +class TorchScaledMMFP4Linear(nn.Module): + """Inference-only NVFP4 Linear using Triton quantization and ``_scaled_mm``.""" + + in_features: int + out_features: int + + def __init__( + self, + *, + weight: torch.Tensor, + weight_qdata: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + ) -> None: + super().__init__() + expected_weight_shape = (weight_qdata.shape[0], weight_qdata.shape[1] * 2) + if weight.shape != expected_weight_shape: + raise ValueError( + "weight must match unpacked weight_qdata shape " + f"{expected_weight_shape}, got {tuple(weight.shape)}." + ) + if weight_qdata.dim() != 2: + raise ValueError(f"weight_qdata must be 2D, got {tuple(weight_qdata.shape)}.") + if weight_qdata.shape[1] * 2 % NVFP4_BLOCK_SIZE != 0: + raise ValueError( + "weight_qdata must represent a K dimension divisible by " + f"{NVFP4_BLOCK_SIZE}, got packed shape {tuple(weight_qdata.shape)}." + ) + if bias is not None and bias.shape != (weight_qdata.shape[0],): + raise ValueError( + f"bias must have shape ({weight_qdata.shape[0]},), got {tuple(bias.shape)}." + ) + self.out_features = int(weight_qdata.shape[0]) + self.in_features = int(weight_qdata.shape[1] * 2) + self.out_dtype = out_dtype + self.register_buffer("weight", weight.detach().contiguous()) + self.register_buffer("weight_qdata", weight_qdata.contiguous()) + self.register_buffer("weight_scale", weight_scale.contiguous()) + if bias is None: + self.register_buffer("bias", None) + else: + self.register_buffer("bias", bias.detach().to(out_dtype).contiguous()) + + @classmethod + def from_linear( + cls, + source: nn.Linear, + *, + out_dtype: torch.dtype, + ) -> "TorchScaledMMFP4Linear": + """Create an NVFP4 replacement from a source ``nn.Linear``.""" + _require_fp4_dtype() + weight_qdata, weight_scale = quantize_nvfp4_swizzled(source.weight.detach()) + bias = source.bias.detach() if source.bias is not None else None + replacement = cls( + weight=source.weight.detach().to(device=source.weight.device), + weight_qdata=weight_qdata, + weight_scale=weight_scale, + bias=bias.to(device=source.weight.device) if bias is not None else None, + out_dtype=out_dtype, + ) + replacement.train(source.training) + return replacement + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Apply the quantized linear projection.""" + if input.shape[-1] != self.in_features: + raise ValueError( + f"expected input last dim {self.in_features}, got {input.shape[-1]}." + ) + if not input.is_cuda: + raise RuntimeError("TorchScaledMMFP4Linear requires CUDA input.") + _require_scaled_mm() + _require_fp4_dtype() + + leading_shape = input.shape[:-1] + input_2d = input.reshape(-1, self.in_features) + if input_2d.numel() == 0: + return input.new_empty( + (*leading_shape, self.out_features), + dtype=self.out_dtype, + ) + + input_qdata, input_scale = quantize_nvfp4_swizzled(input_2d) + output = torch._scaled_mm( + input_qdata.view(torch.float4_e2m1fn_x2), + self.weight_qdata.t().view(torch.float4_e2m1fn_x2), + input_scale.view(torch.float8_e4m3fn), + self.weight_scale.view(torch.float8_e4m3fn), + bias=self.bias.to(device=input.device, dtype=self.out_dtype) + if self.bias is not None + else None, + out_dtype=self.out_dtype, + ) + return output.reshape(*leading_shape, self.out_features) + + +def replace_linear_with_torch_fp8( + module: nn.Module, + *, + recipe: Any, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None = None, + prefix: str = "", +) -> tuple[int, int]: + """Replace eligible ``nn.Linear`` modules with ``TorchScaledMMFP8Linear``.""" + return _replace_linear_with_native_quant( + module, + recipe=TorchScaledMMFP8Recipe(), + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=prefix, + ) + + +def replace_linear_with_torch_fp4( + module: nn.Module, + *, + recipe: Any, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None = None, + prefix: str = "", +) -> tuple[int, int]: + """Replace eligible ``nn.Linear`` modules with ``TorchScaledMMFP4Linear``.""" + return _replace_linear_with_native_quant( + module, + recipe=TorchScaledMMFP4Recipe(), + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=prefix, + ) + + +def replace_linear_with_native_quant( + module: nn.Module, + *, + recipe: NativeQuantRecipe, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None = None, + prefix: str = "", +) -> tuple[int, int]: + """Replace eligible ``nn.Linear`` modules with the requested native backend. + + The signature intentionally matches upstream SANA's + ``_replace_linear_with_te_nvfp4`` helper so the runner can patch the backend + without modifying upstream source files. + """ + return _replace_linear_with_native_quant( + module, + recipe=recipe, + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=prefix, + ) + + +def _replace_linear_with_native_quant( + module: nn.Module, + *, + recipe: NativeQuantRecipe, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None, + prefix: str, +) -> tuple[int, int]: + converted = 0 + skipped = 0 + for name, child in list(module.named_children()): + child_prefix = f"{prefix}.{name}" if prefix else name + if _name_matches(skip_patterns, child_prefix): + skipped += 1 + continue + if isinstance(child, nn.Linear): + if include_patterns is not None and not _name_matches( + include_patterns, + child_prefix, + ): + skipped += 1 + continue + if child.in_features % 16 != 0 or child.out_features % 16 != 0: + skipped += 1 + continue + if recipe.precision == "fp4" and child.in_features % 32 != 0: + skipped += 1 + continue + out_dtype = ( + params_dtype + if params_dtype in {torch.bfloat16, torch.float16} + else torch.bfloat16 + ) + if recipe.precision == "fp8": + replacement = TorchScaledMMFP8Linear.from_linear( + child, + out_dtype=out_dtype, + ) + elif recipe.precision == "fp4": + replacement = TorchScaledMMFP4Linear.from_linear( + child, + out_dtype=out_dtype, + ) + else: + raise ValueError(f"Unsupported native quant recipe: {recipe!r}.") + setattr(module, name, replacement) + converted += 1 + continue + child_converted, child_skipped = _replace_linear_with_native_quant( + child, + recipe=recipe, + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=child_prefix, + ) + converted += child_converted + skipped += child_skipped + return converted, skipped + + +def _name_matches(patterns: tuple[str, ...], name: str) -> bool: + return any(re.search(pattern, name) for pattern in patterns) + + +def _require_scaled_mm() -> None: + if not hasattr(torch, "_scaled_mm"): + raise RuntimeError("torch._scaled_mm is required for native SANA quantization.") + + +def _require_fp8_dtype() -> None: + if not hasattr(torch, "float8_e4m3fn"): + raise RuntimeError("torch.float8_e4m3fn is required for the native SANA FP8 backend.") + + +def _require_fp4_dtype() -> None: + if not hasattr(torch, "float4_e2m1fn_x2"): + raise RuntimeError( + "torch.float4_e2m1fn_x2 is required for the native SANA FP4 backend." + ) + if not hasattr(torch, "float8_e4m3fn"): + raise RuntimeError( + "torch.float8_e4m3fn scales are required for the native SANA FP4 backend." + ) + + +__all__ = [ + "TorchScaledMMFP4Linear", + "TorchScaledMMFP4Recipe", + "TorchScaledMMFP8Linear", + "TorchScaledMMFP8Recipe", + "quantize_nvfp4_swizzled", + "replace_linear_with_native_quant", + "replace_linear_with_torch_fp4", + "replace_linear_with_torch_fp8", +] diff --git a/integrations/sana/sana_wm/native_transformer.py b/integrations/sana/sana_wm/native_transformer.py new file mode 100644 index 000000000..d8be6efb1 --- /dev/null +++ b/integrations/sana/sana_wm/native_transformer.py @@ -0,0 +1,1085 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Native FlashDreams adapter for the SANA-WM Stage-1 DiT.""" + +from __future__ import annotations + +import gc +import os +import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any, Literal, cast + +import numpy as np +import torch +import torch.nn as nn +import yaml +from loguru import logger +from torch import Tensor + +from flashdreams.infra.diffusion.transformer import ( + Transformer, + TransformerAutoregressiveCache, + TransformerConfig, +) +from sana_wm._tools import find_model, resolve_hf_path +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, + SANA_WM_MODEL_PATH, + SANA_WM_REFINER_GEMMA_ROOT, + SANA_WM_REFINER_ROOT, +) +from sana_wm.native_quant import ( + TorchScaledMMFP4Recipe, + TorchScaledMMFP8Recipe, + replace_linear_with_native_quant, +) +from sana_wm.stage1_model import SanaWMStage1Model + +Precision = Literal["bf16", "fp8", "fp4"] +QuantBackend = Literal["auto", "upstream-te", "torch-fp8", "torch-fp4", "native"] + +_STAGE1_QUANT_SKIP_DEFAULTS = ( + "^x_embedder", + "^raymap_embedder", + "^plucker_embedder", + "^t_embedder", + "^t_block", + "^y_embedder", + "^final_layer", + # SANA attention code accesses output_gate.weight directly; replacing it + # with a module that only has quantized buffers breaks that call site. + r"\.output_gate$", +) +_STAGE1_QUANT_INCLUDE_DEFAULTS = ( + r"^blocks\.\d+\.attn\.qkv$", + r"^blocks\.\d+\.attn\.proj$", + r"^blocks\.\d+\.attn\.beta_proj$", + r"^blocks\.\d+\.attn\.gate_proj$", + r"^blocks\.\d+\.cross_attn\.", + r"^blocks\.\d+\.mlp\.inverted_conv\.linear$", + r"^blocks\.\d+\.mlp\.point_conv\.linear$", +) + + +@dataclass(kw_only=True) +class SanaWMStage1Conditioning: + """Per-rollout inputs needed by the SANA-WM Stage-1 sampler.""" + + condition: Tensor + uncondition: Tensor | None + model_kwargs: dict[str, object] + first_latent: Tensor + latent_shape: tuple[int, int, int, int, int] + cfg_scale: float + flow_shift: float + steps: int + seed: int + + +@dataclass(kw_only=True) +class SanaWMNativeTransformerCache(TransformerAutoregressiveCache): + """AR cache for the one-shot SANA-WM Stage-1 rollout.""" + + conditioning: SanaWMStage1Conditioning | None = None + + +@dataclass(kw_only=True) +class SanaWMNativeTransformerConfig(TransformerConfig): + """Config for the native SANA-WM Stage-1 transformer adapter.""" + + _target: type["SanaWMNativeTransformer"] = field( + default_factory=lambda: SanaWMNativeTransformer + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + checkpoint_path: str = SANA_WM_MODEL_PATH + """SANA-WM Stage-1 checkpoint path or ``hf://`` URI.""" + + stage1_precision: Precision = "bf16" + """Stage-1 precision requested by the runner.""" + + quant_backend: QuantBackend = "auto" + """Low-precision linear replacement backend for FP8/FP4 modes.""" + + refiner_root: str = SANA_WM_REFINER_ROOT + """LTX-2 refiner root path or ``hf://`` URI.""" + + refiner_gemma_root: str = SANA_WM_REFINER_GEMMA_ROOT + """Gemma text-encoder root for the refiner.""" + + refiner_precision: Precision = "bf16" + """Refiner precision requested by the runner.""" + + offload_vae: bool = False + """Move the VAE to CPU between encode/decode phases.""" + + offload_text_encoder: bool = False + """Move the Stage-1 text encoder to CPU between prompt encodes.""" + + offload_refiner: bool = False + """Release the LTX-2 refiner after each refinement.""" + + height: int = DEFAULT_VIDEO_HEIGHT + """Pixel height used by the public SANA-WM bidirectional release.""" + + width: int = DEFAULT_VIDEO_WIDTH + """Pixel width used by the public SANA-WM bidirectional release.""" + + vae_tile_sample_min_width: int = 256 + """Pixel width tile size used for native LTX-2 VAE decode.""" + + vae_tile_sample_stride_width: int = 224 + """Pixel width stride used for native LTX-2 VAE decode.""" + + vae_tile_sample_min_height: int = 256 + """Pixel height tile size used for native LTX-2 VAE decode.""" + + vae_tile_sample_stride_height: int = 192 + """Pixel height stride used for native LTX-2 VAE decode.""" + + vae_tile_sample_min_num_frames: int = 24 + """Pixel-frame temporal tile size used for native LTX-2 VAE decode.""" + + vae_tile_sample_stride_num_frames: int = 8 + """Pixel-frame temporal tile stride used for native LTX-2 VAE decode.""" + + vae_oom_retry_tile_sample_min_width: int = 128 + """Smaller pixel width tile size for one VAE decode OOM retry.""" + + vae_oom_retry_tile_sample_stride_width: int = 64 + """Smaller pixel width tile stride for one VAE decode OOM retry.""" + + vae_oom_retry_tile_sample_min_height: int = 128 + """Smaller pixel height tile size for one VAE decode OOM retry.""" + + vae_oom_retry_tile_sample_stride_height: int = 64 + """Smaller pixel height tile stride for one VAE decode OOM retry.""" + + +class SanaWMNativeTransformer(Transformer[SanaWMNativeTransformerCache]): + """Native FlashDreams adapter for the SANA-WM Stage-1 model call.""" + + def __init__(self, config: SanaWMNativeTransformerConfig) -> None: + super().__init__(config) + self.config: SanaWMNativeTransformerConfig = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._model_path: str | None = None + self.weight_dtype: torch.dtype | None = None + self._model_built = False + self._vae_built = False + self._text_encoder_built = False + self._refiner_built = False + self._stage1_quantized = False + self._streaming_prompt_cache: dict[ + tuple[object, ...], tuple[Tensor, Tensor, Tensor, Tensor] + ] = {} + + @property + def latent_shape(self) -> tuple[int, ...]: + """Return the current rollout latent shape or the public default.""" + if self._runtime_config is None: + return (1, 128, 21, 22, 40) + cfg = self._ensure_runtime_config() + latent_t = (161 - 1) // int(cfg.vae.vae_stride[0]) + 1 + return ( + 1, + int(cfg.vae.vae_latent_dim), + latent_t, + self.config.height // int(cfg.vae.vae_stride[-1]), + self.config.width // int(cfg.vae.vae_stride[-1]), + ) + + def initialize_autoregressive_cache( + self, + *, + conditioning: SanaWMStage1Conditioning | None = None, + **_: Any, + ) -> SanaWMNativeTransformerCache: + """Build a cache containing per-rollout SANA-WM conditioning.""" + return SanaWMNativeTransformerCache(conditioning=conditioning) + + def predict_flow( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: SanaWMNativeTransformerCache, + input: Any = None, + ) -> Tensor: + """Execute one SANA-WM DiT flow prediction.""" + conditioning = _require_conditioning(cache) + prompt_embeds = cast(Tensor, input) if input is not None else conditioning.condition + self._ensure_model() + return self.model( + noisy_latent, + timestep, + prompt_embeds, + **conditioning.model_kwargs, + ) + + def finalize_kv_cache( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: SanaWMNativeTransformerCache, + input: Any = None, + ) -> None: + """SANA-WM bidirectional inference has no streaming KV cache to advance.""" + del noisy_latent, timestep, cache, input + + def patchify_and_maybe_split_cp(self, x: Any) -> Any: + """SANA-WM Stage-1 latents are already in model layout.""" + return x + + def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: + """SANA-WM Stage-1 latents are already in model layout.""" + return x + + def prepare_conditioning( + self, + *, + image: Any, + prompt: str, + camera: Mapping[str, Tensor], + num_frames: int, + fps: int, + steps: int, + cfg_scale: float, + flow_shift: float | None, + seed: int, + negative_prompt: str, + ) -> SanaWMStage1Conditioning: + """Encode first-frame, text, and camera tensors for one rollout.""" + del fps + cfg = self._ensure_runtime_config() + weight_dtype = self._ensure_weight_dtype() + self._ensure_vae() + if self.config.offload_vae: + self.vae.to(self.device) + + from torchvision import transforms as T + + img = (T.ToTensor()(image) * 2.0 - 1.0).unsqueeze(0).unsqueeze(2) + first_latent = self._vae_encode( + img.to(self.device, dtype=self.vae_dtype), + ).to(weight_dtype) + if self.config.offload_vae: + self.vae.to("cpu") + torch.cuda.empty_cache() + + cond, cond_mask, neg, neg_mask = self._encode_prompts( + prompt, + negative_prompt, + ) + cond, cond_mask, neg, neg_mask = self._pad_text_for_quant( + cond, + cond_mask, + neg, + neg_mask, + ) + + raymap = camera["raymap"].unsqueeze(0).to(self.device, dtype=weight_dtype) + chunk_plucker = camera["chunk_plucker"].unsqueeze(0).to( + self.device, + dtype=weight_dtype, + ) + if cfg_scale > 1.0: + mask_cfg = torch.cat([neg_mask, cond_mask], dim=0) + raymap_cfg = torch.cat([raymap, raymap], dim=0) + chunk_plucker_cfg = torch.cat([chunk_plucker, chunk_plucker], dim=0) + uncondition = neg + else: + mask_cfg = cond_mask + raymap_cfg = raymap + chunk_plucker_cfg = chunk_plucker + uncondition = None + + latent_t = (num_frames - 1) // int(cfg.vae.vae_stride[0]) + 1 + latent_h = self.config.height // int(cfg.vae.vae_stride[-1]) + latent_w = self.config.width // int(cfg.vae.vae_stride[-1]) + chunk_index = self._chunk_index(latent_t) + model_kwargs: dict[str, object] = { + "data_info": { + "img_hw": torch.tensor( + [[self.config.height, self.config.width]], + dtype=torch.float, + device=self.device, + ), + "condition_frame_info": {0: 0.0}, + }, + "mask": mask_cfg, + "camera_conditions": raymap_cfg, + "chunk_plucker": chunk_plucker_cfg, + } + if chunk_index is not None: + model_kwargs["chunk_index"] = chunk_index + + return SanaWMStage1Conditioning( + condition=cond, + uncondition=uncondition, + model_kwargs=model_kwargs, + first_latent=first_latent, + latent_shape=( + 1, + int(first_latent.shape[1]), + latent_t, + latent_h, + latent_w, + ), + cfg_scale=float(cfg_scale), + flow_shift=self._resolve_flow_shift(flow_shift), + steps=int(steps), + seed=int(seed), + ) + + def initial_latents(self, conditioning: SanaWMStage1Conditioning) -> Tensor: + """Draw Stage-1 initial noise and pin the encoded first frame.""" + generator = torch.Generator(device=self.device).manual_seed(conditioning.seed) + latents = torch.randn( + conditioning.latent_shape, + dtype=self._ensure_weight_dtype(), + device=self.device, + generator=generator, + ) + latents[:, :, :1] = conditioning.first_latent + return latents + + def decode_latents(self, latents: Tensor) -> np.ndarray: + """Decode SANA-WM VAE latents to ``uint8`` HWC video.""" + self._ensure_vae() + if self.config.offload_vae: + self.vae.to(self.device) + samples = latents.to(device=self.device, dtype=self.vae_dtype) + if torch.cuda.is_available(): + torch.cuda.synchronize() + t0 = time.perf_counter() + retry_decode = False + try: + decoded = self._vae_decode(samples) + except torch.OutOfMemoryError: + retry_decode = True + + if retry_decode: + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + gc.collect() + logger.warning( + "[sana-vae] decode OOM; retrying with smaller tiles " + "width={} stride_width={} height={} stride_height={}", + self.config.vae_oom_retry_tile_sample_min_width, + self.config.vae_oom_retry_tile_sample_stride_width, + self.config.vae_oom_retry_tile_sample_min_height, + self.config.vae_oom_retry_tile_sample_stride_height, + ) + self._configure_vae_tiling( + tile_sample_min_width=( + self.config.vae_oom_retry_tile_sample_min_width + ), + tile_sample_stride_width=( + self.config.vae_oom_retry_tile_sample_stride_width + ), + tile_sample_min_height=( + self.config.vae_oom_retry_tile_sample_min_height + ), + tile_sample_stride_height=( + self.config.vae_oom_retry_tile_sample_stride_height + ), + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + decoded = self._vae_decode(samples) + if torch.cuda.is_available(): + torch.cuda.synchronize() + logger.info( + "[timing] vae decode: {:.3f}s (latent T={} -> pixels {})", + time.perf_counter() - t0, + latents.shape[2], + tuple(decoded.shape) if isinstance(decoded, Tensor) else "list", + ) + if isinstance(decoded, list): + decoded = torch.stack(decoded, dim=0) + video = ( + torch.clamp(127.5 * decoded + 127.5, 0, 255) + .permute(0, 2, 3, 4, 1) + .to("cpu", dtype=torch.uint8) + .numpy()[0] + ) + if self.config.offload_vae: + self.vae.to("cpu") + del samples, decoded + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return video + + def release_stage1_runtime( + self, + cache: SanaWMNativeTransformerCache | None = None, + ) -> None: + """Release Stage-1-only tensors before VAE/refiner work.""" + free_before_gib: float | None = None + if torch.cuda.is_available(): + try: + free_before, _total = torch.cuda.mem_get_info() + free_before_gib = free_before / (1024**3) + except RuntimeError: + free_before_gib = None + if cache is not None: + cache.conditioning = None + self._streaming_prompt_cache.clear() + + for attr in ("model", "text_encoder"): + module = getattr(self, attr, None) + if module is None: + continue + try: + module.to("meta") + except Exception: + try: + module.to("cpu") + except Exception: + pass + setattr(self, attr, None) + + if hasattr(self, "tokenizer"): + self.tokenizer = None + self._model_built = False + self._text_encoder_built = False + self._stage1_quantized = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + try: + free_after, _total = torch.cuda.mem_get_info() + if free_before_gib is not None: + logger.info( + "[stage1] released Stage-1 runtime before decode/refine " + "(free CUDA memory: {:.2f} -> {:.2f} GiB)", + free_before_gib, + free_after / (1024**3), + ) + else: + logger.info("[stage1] released Stage-1 runtime before decode/refine") + except RuntimeError: + logger.info("[stage1] released Stage-1 runtime before decode/refine") + gc.collect() + + def refine_latents( + self, + *, + latents: Tensor, + prompt: str, + fps: int, + sink_size: int, + seed: int, + block_size: int | None, + kv_max_frames: int, + ) -> Tensor: + """Run the native LTX-2 refiner once that path is implemented.""" + self._ensure_refiner() + sigmas = torch.tensor( + self._stage2_sigmas(), + dtype=torch.float32, + device=self.device, + ) + logger.info( + "[refiner] {}-step Euler, start_sigma={:.4f}", + len(sigmas) - 1, + float(sigmas[0]), + ) + refined = self.refiner.refine_latents( + latents, + prompt, + fps=float(fps), + sink_size=int(sink_size), + seed=int(seed), + progress=True, + block_size=block_size, + kv_max_frames=int(kv_max_frames), + ) + if self.config.offload_refiner: + self._release_refiner() + return refined + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is not None: + return self._runtime_config + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_weight_dtype(self) -> torch.dtype: + if self.weight_dtype is None: + self.weight_dtype = _get_weight_dtype( + self._ensure_runtime_config().model.mixed_precision + ) + return self.weight_dtype + + def _ensure_vae(self) -> None: + if self._vae_built: + return + cfg = self._ensure_runtime_config() + self.vae_dtype = _get_weight_dtype(cfg.vae.weight_dtype) + cfg.vae.vae_pretrained = resolve_hf_path(cfg.vae.vae_pretrained) + self.vae = _get_vae( + cfg.vae.vae_type, + cfg.vae.vae_pretrained, + device=self.device, + dtype=self.vae_dtype, + config=cfg.vae, + ) + if hasattr(self.vae, "enable_tiling"): + self._configure_vae_tiling() + self._vae_built = True + + def _ensure_text_encoder(self) -> None: + if self._text_encoder_built: + return + cfg = self._ensure_runtime_config() + self.tokenizer, self.text_encoder = _get_tokenizer_and_text_encoder( + name=cfg.text_encoder.text_encoder_name, + device=self.device, + ) + if self.config.offload_text_encoder: + self.text_encoder.to("cpu") + self._text_encoder_built = True + + def _ensure_model(self) -> None: + if self._model_built: + self._prepare_stage1_quant() + return + cfg = self._ensure_runtime_config() + weight_dtype = self._ensure_weight_dtype() + model = SanaWMStage1Model().to(self.device) + logger.info( + "[Sana] Loaded native {} ({:,} params)", + cfg.model.model, + sum(p.numel() for p in model.parameters()), + ) + self._model_path = resolve_hf_path(self.config.checkpoint_path) + state = find_model(self._model_path) + if "generator" in state: + state = state["generator"] + if "state_dict" not in state: + state = { + "state_dict": { + (k[len("model.") :] if k.startswith("model.") else k): v + for k, v in state.items() + } + } + missing, unexpected = model.load_state_dict(state["state_dict"], strict=True) + if missing: + logger.warning("[Sana] Missing keys: {}", missing) + if unexpected: + logger.warning("[Sana] Unexpected keys: {}", unexpected) + self.model = model.eval().to(weight_dtype) + self._model_built = True + self._prepare_stage1_quant() + + def _ensure_refiner(self) -> None: + raise RuntimeError( + "Native SANA-WM refiner execution is not implemented yet. " + "Use --no-refiner True for native Stage-1 GPU tests, or " + "sana-wm-upstream-reference for the temporary upstream refiner harness." + ) + + def _release_refiner(self) -> None: + if not self._refiner_built: + return + del self.refiner + self._refiner_built = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() + + def _configure_vae_tiling( + self, + *, + tile_sample_min_width: int | None = None, + tile_sample_stride_width: int | None = None, + tile_sample_min_height: int | None = None, + tile_sample_stride_height: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_num_frames: int | None = None, + ) -> None: + vae = self.vae + min_width = int(tile_sample_min_width or self.config.vae_tile_sample_min_width) + stride_width = int( + tile_sample_stride_width or self.config.vae_tile_sample_stride_width + ) + min_height = int( + tile_sample_min_height or self.config.vae_tile_sample_min_height + ) + stride_height = int( + tile_sample_stride_height or self.config.vae_tile_sample_stride_height + ) + spatial_ratio = int(getattr(vae, "spatial_compression_ratio", 1)) + stride_width = _avoid_degenerate_tile_tail( + sample_extent=self.config.width, + sample_tile_min=min_width, + sample_stride=stride_width, + compression_ratio=spatial_ratio, + ) + stride_height = _avoid_degenerate_tile_tail( + sample_extent=self.config.height, + sample_tile_min=min_height, + sample_stride=stride_height, + compression_ratio=spatial_ratio, + ) + min_frames = int( + tile_sample_min_num_frames + or self.config.vae_tile_sample_min_num_frames + ) + stride_frames = int( + tile_sample_stride_num_frames + or self.config.vae_tile_sample_stride_num_frames + ) + kwargs = { + "tile_sample_min_height": min_height, + "tile_sample_stride_height": stride_height, + "tile_sample_min_width": min_width, + "tile_sample_stride_width": stride_width, + "tile_sample_min_num_frames": min_frames, + "tile_sample_stride_num_frames": stride_frames, + } + if hasattr(vae, "enable_tiling"): + try: + vae.enable_tiling(**kwargs) + except TypeError: + vae.enable_tiling() + for name, value in kwargs.items(): + if hasattr(vae, name): + setattr(vae, name, value) + if hasattr(vae, "use_framewise_encoding"): + vae.use_framewise_encoding = True + if hasattr(vae, "use_framewise_decoding"): + vae.use_framewise_decoding = True + logger.info( + "[sana-vae] tiling width={} stride_width={} height={} " + "stride_height={} frames={} stride_frames={}", + min_width, + stride_width, + min_height, + stride_height, + min_frames, + stride_frames, + ) + + def _encode_prompts( + self, + prompt: str, + negative_prompt: str, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + cfg = self._ensure_runtime_config() + self._ensure_text_encoder() + max_length = cfg.text_encoder.model_max_length + chi_prompt = "\n".join(cfg.text_encoder.chi_prompt or []) + if chi_prompt: + prompt = chi_prompt + prompt + max_length_all = len(self.tokenizer.encode(chi_prompt)) + max_length - 2 + else: + max_length_all = max_length + + key = ( + prompt, + negative_prompt, + str(self.device), + str(self._ensure_weight_dtype()), + self.config.stage1_precision, + self.config.quant_backend, + ) + if key in self._streaming_prompt_cache: + return self._streaming_prompt_cache[key] + + move_text_encoder = self.config.offload_text_encoder or ( + next(self.text_encoder.parameters()).device != self.device + ) + if move_text_encoder: + self.text_encoder.to(self.device) + + def encode(text: str, length: int) -> tuple[Tensor, Tensor]: + tokens = self.tokenizer( + [text], + max_length=length, + padding="max_length", + truncation=True, + return_tensors="pt", + ).to(self.device) + return self.text_encoder(tokens.input_ids, tokens.attention_mask)[0], ( + tokens.attention_mask + ) + + try: + cond, cond_mask = encode(prompt, max_length_all) + select = [0] + list(range(-max_length + 1, 0)) + cond = cond[:, None][:, :, select] + cond_mask = cond_mask[:, select] + neg, neg_mask = encode(negative_prompt, max_length) + result = (cond, cond_mask, neg[:, None], neg_mask) + self._streaming_prompt_cache.clear() + self._streaming_prompt_cache[key] = result + return result + finally: + if move_text_encoder: + self.text_encoder.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _pad_text_for_quant( + self, + cond: Tensor, + cond_mask: Tensor, + neg: Tensor, + neg_mask: Tensor, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + if self.config.stage1_precision == "bf16": + return cond, cond_mask, neg, neg_mask + multiple = int(os.environ.get("SANA_WM_STAGE1_NVFP4_TEXT_PAD_MULTIPLE", "8")) + if multiple <= 1: + return cond, cond_mask, neg, neg_mask + + def pad_pair(text: Tensor, mask: Tensor) -> tuple[Tensor, Tensor]: + pad = (-text.shape[-2]) % multiple + if pad == 0: + return text, mask + text_shape = list(text.shape) + text_shape[-2] = pad + mask_shape = list(mask.shape) + mask_shape[-1] = pad + return ( + torch.cat([text, text.new_zeros(text_shape)], dim=-2), + torch.cat([mask, mask.new_zeros(mask_shape)], dim=-1), + ) + + cond, cond_mask = pad_pair(cond, cond_mask) + neg, neg_mask = pad_pair(neg, neg_mask) + return cond, cond_mask, neg, neg_mask + + def _prepare_stage1_quant(self) -> None: + if self._stage1_quantized or self.config.stage1_precision == "bf16": + return + if self.config.quant_backend == "upstream-te": + raise ValueError( + "Native FlashDreams SANA execution does not use Transformer Engine; " + "select --quant-backend native, torch-fp8, or torch-fp4." + ) + if self.config.stage1_precision == "fp8": + recipe = TorchScaledMMFP8Recipe() + else: + recipe = TorchScaledMMFP4Recipe() + converted, skipped = replace_linear_with_native_quant( + self.model, + recipe=recipe, + params_dtype=self._ensure_weight_dtype(), + skip_patterns=_STAGE1_QUANT_SKIP_DEFAULTS, + include_patterns=_STAGE1_QUANT_INCLUDE_DEFAULTS, + ) + if converted <= 0: + raise RuntimeError( + f"SANA-WM native {self.config.stage1_precision} converted no " + f"Stage-1 Linear layers; skipped={skipped}." + ) + self._stage1_quantized = True + logger.info( + "[stage1-native-quant] precision={} converted {} Linear layers (skipped {})", + self.config.stage1_precision, + converted, + skipped, + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _resolve_flow_shift(self, override: float | None) -> float: + cfg = self._ensure_runtime_config() + if override is not None: + return float(override) + if cfg.scheduler.inference_flow_shift is not None: + return float(cfg.scheduler.inference_flow_shift) + return float(cfg.scheduler.flow_shift) + + def _chunk_index(self, latent_frames: int) -> list[int] | None: + return _chunk_index_from_config( + self._ensure_runtime_config(), + num_frames=latent_frames, + ) + + def _vae_encode(self, images: Tensor) -> Tensor: + return _vae_encode_ltx2( + self._ensure_runtime_config().vae.vae_type, + self.vae, + images, + device=self.device, + ) + + def _vae_decode(self, latents: Tensor) -> Tensor: + return _vae_decode_ltx2( + self._ensure_runtime_config().vae.vae_type, + self.vae, + latents, + ) + + @staticmethod + def _stage2_sigmas() -> tuple[float, ...]: + return (0.909375, 0.725, 0.421875, 0.0) + + +def _require_conditioning( + cache: SanaWMNativeTransformerCache, +) -> SanaWMStage1Conditioning: + if cache.conditioning is None: + raise RuntimeError("SANA-WM native cache was initialized without conditioning.") + return cache.conditioning + + +def _avoid_degenerate_tile_tail( + *, + sample_extent: int, + sample_tile_min: int, + sample_stride: int, + compression_ratio: int, +) -> int: + """Choose a tile stride that avoids final latent tiles of size one.""" + if compression_ratio <= 1: + return sample_stride + latent_extent = max(1, sample_extent // compression_ratio) + latent_tile_min = max(1, sample_tile_min // compression_ratio) + requested_latent_stride = max(1, sample_stride // compression_ratio) + + def tail_size(latent_stride: int) -> int: + last_start = ((latent_extent - 1) // latent_stride) * latent_stride + return latent_extent - last_start + + candidates = range( + min(requested_latent_stride, latent_tile_min), + 1, + -1, + ) + for latent_stride in candidates: + if tail_size(latent_stride) > 1: + return latent_stride * compression_ratio + + for latent_stride in range( + requested_latent_stride + 1, + latent_tile_min + 1, + ): + if tail_size(latent_stride) > 1: + return latent_stride * compression_ratio + + return sample_stride + + +def _load_inference_config(config_path: str) -> Any: + with open(resolve_hf_path(config_path), encoding="utf-8") as handle: + raw = yaml.safe_load(handle) or {} + if not isinstance(raw, dict): + raise TypeError(f"SANA-WM config must be a mapping, got {type(raw).__name__}.") + raw.setdefault("work_dir", "") + return _to_namespace(raw) + + +def _get_vae(*args: Any, **kwargs: Any) -> nn.Module: + name, model_path = args[:2] + device = kwargs["device"] + dtype = kwargs["dtype"] + if "LTX2VAE_diffusers" not in str(name): + raise ValueError(f"Unsupported native SANA-WM VAE type: {name!r}") + from diffusers import AutoencoderKLLTX2Video + + return ( + AutoencoderKLLTX2Video.from_pretrained( + model_path, + subfolder="vae", + torch_dtype=dtype, + ) + .to(device) + .eval() + ) + + +def _get_tokenizer_and_text_encoder(*args: Any, **kwargs: Any) -> tuple[Any, nn.Module]: + name = kwargs.get("name", args[0] if args else "T5") + device = kwargs.get("device", "cuda") + model_id = _TEXT_ENCODER_MODEL_IDS.get(str(name)) + if model_id is None: + raise ValueError(f"Unsupported native SANA-WM text encoder: {name!r}") + from transformers import AutoModelForCausalLM, AutoTokenizer, T5EncoderModel, T5Tokenizer + + if "T5" in str(name): + tokenizer = T5Tokenizer.from_pretrained(model_id) + text_encoder = T5EncoderModel.from_pretrained( + model_id, + torch_dtype=torch.float16, + ).to(device) + return tokenizer, text_encoder.eval() + + tokenizer = AutoTokenizer.from_pretrained(model_id) + tokenizer.padding_side = "right" + text_encoder = ( + AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16) + .get_decoder() + .to(device) + .eval() + ) + return tokenizer, text_encoder + + +def _get_weight_dtype(value: str) -> torch.dtype: + normalized = str(value).lower() + if normalized in {"bf16", "bfloat16"}: + return torch.bfloat16 + if normalized in {"fp16", "float16", "half"}: + return torch.float16 + if normalized in {"fp32", "float32", "float"}: + return torch.float32 + raise ValueError(f"Unsupported dtype value: {value!r}") + + +class _ConfigNamespace(SimpleNamespace): + """Simple YAML-backed config object with dict-like ``get`` support.""" + + def get(self, name: str, default: Any = None) -> Any: + return getattr(self, name, default) + + +_TEXT_ENCODER_MODEL_IDS = { + "T5": "DeepFloyd/t5-v1_1-xxl", + "T5-small": "google/t5-v1_1-small", + "T5-base": "google/t5-v1_1-base", + "T5-large": "google/t5-v1_1-large", + "T5-xl": "google/t5-v1_1-xl", + "T5-xxl": "google/t5-v1_1-xxl", + "gemma-2b": "google/gemma-2b", + "gemma-2b-it": "google/gemma-2b-it", + "gemma-2-2b": "google/gemma-2-2b", + "gemma-2-2b-it": "Efficient-Large-Model/gemma-2-2b-it", + "gemma-2-9b": "google/gemma-2-9b", + "gemma-2-9b-it": "google/gemma-2-9b-it", +} + + +def _vae_encode_ltx2(name: str, vae: nn.Module, images: Tensor, *, device: torch.device) -> Tensor: + if "LTX2VAE_diffusers" not in name: + raise ValueError(f"Unsupported native SANA-WM VAE encode type: {name!r}") + dtype = images.dtype + vae_device = next(vae.parameters()).device + vae_dtype = next(vae.parameters()).dtype + posterior = vae.encode(images.to(device=vae_device, dtype=vae_dtype)).latent_dist + z = posterior.mode() + latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(z.device, z.dtype) + latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(z.device, z.dtype) + z = (z - latents_mean) * vae.config.scaling_factor / latents_std + return z.to(device=device, dtype=dtype) + + +def _vae_decode_ltx2(name: str, vae: nn.Module, latents: Tensor) -> Tensor: + if "LTX2VAE_diffusers" not in name: + raise ValueError(f"Unsupported native SANA-WM VAE decode type: {name!r}") + vae_device = next(vae.parameters()).device + vae_dtype = next(vae.parameters()).dtype + latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to( + latents.device, + latents.dtype, + ) + latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to( + latents.device, + latents.dtype, + ) + scaled = latents * latents_std / vae.config.scaling_factor + latents_mean + return vae.decode( + scaled.to(device=vae_device, dtype=vae_dtype), + temb=None, + return_dict=False, + )[0] + + +def _to_namespace(value: Any) -> Any: + if isinstance(value, dict): + return _ConfigNamespace( + **{str(key): _to_namespace(child) for key, child in value.items()} + ) + if isinstance(value, list): + return [_to_namespace(child) for child in value] + return value + + +def _chunk_index_from_config(config: Any, *, num_frames: int) -> list[int] | None: + model = getattr(config, "model", None) + if model is None: + return None + chunk_index = model.get("chunk_index", None) + chunk_size = model.get("chunk_size", None) + strategy = model.get("chunk_split_strategy", "uniform") + if chunk_index is not None: + if not isinstance(chunk_index, (list, tuple)): + raise TypeError( + f"chunk_index must be a list, got {type(chunk_index).__name__}" + ) + if len(chunk_index) == 0: + raise ValueError("chunk_index cannot be empty.") + return [int(index) for index in chunk_index] + if chunk_size is None: + return None + return _chunk_index_from_chunk_size( + num_frames, + int(chunk_size), + strategy=str(strategy), + ) + + +def _chunk_index_from_chunk_size( + num_frames: int, + chunk_size: int, + *, + strategy: str = "uniform", +) -> list[int]: + if num_frames <= 0: + raise ValueError(f"num_frames must be > 0, got {num_frames}.") + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}.") + normalized = strategy.lower() + if normalized in {"uniform", "default"}: + indices = list(range(0, num_frames, chunk_size)) + if len(indices) > 1 and (num_frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if normalized in {"first_frame", "first_frame_alone", "first_frame_only"}: + if num_frames <= 1: + return [0] + indices = [0] + list(range(1, num_frames, chunk_size)) + if len(indices) > 2 and (num_frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if normalized in {"first_plus_one", "first_chunk_plus_one"}: + if num_frames <= chunk_size + 1: + return [0] + indices = [0] + list(range(chunk_size + 1, num_frames, chunk_size)) + if len(indices) > 1 and (num_frames - indices[-1]) < chunk_size: + indices.pop() + return indices + raise ValueError(f"Unknown chunk_split_strategy {strategy!r}.") + + +__all__ = [ + "SanaWMNativeTransformer", + "SanaWMNativeTransformerCache", + "SanaWMNativeTransformerConfig", + "SanaWMStage1Conditioning", +] diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index 56c5981b2..09795c905 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SANA-WM bidirectional runner that delegates execution to upstream Sana.""" +"""SANA-WM bidirectional runner and upstream-reference harness.""" from __future__ import annotations @@ -32,7 +32,17 @@ from loguru import logger from flashdreams.core.io.disk import preflight_runtime_write_paths +from flashdreams.infra.config import derive_config +from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.infra.runner import Runner, RunnerConfig +from sana_wm.camera import ( + action_string_to_c2w, + load_intrinsics, + prepare_camera, + resize_center_crop_geometry, + snap_num_frames, + transform_intrinsics_for_crop, +) from sana_wm.constants import ( DEFAULT_ACTION, DEFAULT_FPS, @@ -44,6 +54,7 @@ SANA_WM_REFINER_ROOT, SANA_WM_VAE_TEMPORAL_COMPRESSION, ) +from sana_wm.native_transformer import SanaWMNativeTransformer SamplingAlgo = Literal[ "auto", @@ -58,6 +69,15 @@ Precision = Literal["bf16", "fp8", "fp4"] """SANA-WM Stage-1/refiner precision modes.""" +QuantBackend = Literal["auto", "upstream-te", "torch-fp8", "torch-fp4", "native"] +"""Low-precision linear backend for FP8/FP4 SANA-WM paths.""" + +ResolvedQuantBackend = Literal["upstream-te", "torch-fp8", "torch-fp4", "native"] +"""Concrete low-precision linear backend selected after resolving ``auto``.""" + +ExecutionBackend = Literal["upstream-reference", "native-flashdreams"] +"""Execution path selected by a SANA-WM runner config.""" + @dataclass(kw_only=True) class SanaWMRunnerConfig(RunnerConfig): @@ -143,7 +163,7 @@ class SanaWMRunnerConfig(RunnerConfig): seed: int = 42 """Stage-1 random seed.""" - no_action_overlay: bool = False + no_action_overlay: bool = True """Skip upstream action-overlay compositing on generated videos.""" config_path: str = SANA_WM_CONFIG_PATH @@ -154,17 +174,26 @@ class SanaWMRunnerConfig(RunnerConfig): stage1_precision: Precision = "bf16" """Stage-1 DiT compute precision. ``"bf16"`` is the default; ``"fp8"`` - requires Hopper or newer plus Transformer Engine; ``"fp4"`` requires - Blackwell plus Transformer Engine.""" + requires Hopper or newer; ``"fp4"`` requires Blackwell.""" no_refiner: bool = False """Skip the LTX-2 refiner and decode Stage-1 latents directly.""" refiner_precision: Precision = "bf16" """LTX-2 refiner compute precision. ``"bf16"`` is the default; - ``"fp8"`` requires Hopper or newer plus Transformer Engine; ``"fp4"`` - requires Blackwell plus Transformer Engine. Ignored when - ``no_refiner`` is ``True``.""" + ``"fp8"`` requires Hopper or newer; ``"fp4"`` requires Blackwell. + Ignored when ``no_refiner`` is ``True``.""" + + quant_backend: QuantBackend = "auto" + """Backend for quantized linear layers. ``"auto"`` uses the native + FlashDreams/PyTorch low-precision path. ``"upstream-te"`` is accepted only + by the upstream reference harness. ``"torch-fp8"`` and ``"torch-fp4"`` + select one native ``torch._scaled_mm`` replacement explicitly; ``"native"`` + allows both native FP8 and native FP4.""" + + execution_backend: ExecutionBackend = "upstream-reference" + """Whether this config runs the upstream reference harness or the native + FlashDreams path.""" refiner_root: str = SANA_WM_REFINER_ROOT """LTX-2 refiner root path or ``hf://`` URI.""" @@ -195,7 +224,7 @@ class SanaWMRunnerConfig(RunnerConfig): class SanaWMRunner(Runner[SanaWMRunnerConfig, object]): - """CLI driver for upstream SANA-WM bidirectional inference.""" + """CLI driver for SANA-WM native and upstream-reference configs.""" config: SanaWMRunnerConfig @@ -234,7 +263,7 @@ def _resolve_device(self) -> torch.device: return torch.device(f"cuda:{self.local_rank}") return torch.device(self.config.device) - def _resolve_trajectory(self, upstream: ModuleType) -> np.ndarray: + def _resolve_trajectory(self) -> np.ndarray: """Load or roll out the camera-to-world trajectory.""" if self.config.camera_path is not None: c2w = np.load(self.config.camera_path).astype(np.float32) @@ -245,7 +274,7 @@ def _resolve_trajectory(self, upstream: ModuleType) -> np.ndarray: return c2w if not self.config.action: raise ValueError("SanaWMRunner requires --camera-path or --action.") - return upstream.action_string_to_c2w( + return action_string_to_c2w( self.config.action, translation_speed=self.config.translation_speed, rotation_speed_deg=self.config.rotation_speed_deg, @@ -265,18 +294,142 @@ def _denoising_steps(self) -> list[int] | None: return steps def run(self) -> None: - """Run upstream SANA-WM bidirectional inference and write outputs.""" + """Run SANA-WM bidirectional inference and write outputs.""" + if self.config.execution_backend == "native-flashdreams": + self._run_native() + return + self._run_upstream_reference() + + def _run_native(self) -> None: + """Run the native FlashDreams SANA-WM pipeline.""" cfg = self.config if cfg.image_path is None: raise ValueError("SanaWMRunner requires --image-path.") + if not cfg.no_refiner: + raise ValueError( + "Native SANA-WM currently supports Stage-1 execution only. " + "Use --no-refiner True, or use sana-wm-upstream-reference for " + "the temporary upstream refiner harness." + ) + if not cfg.no_action_overlay: + raise ValueError( + "Native SANA-WM does not include the upstream action overlay. " + "Use --no-action-overlay True." + ) device = self._resolve_device() + quant_backend: QuantBackend = ( + "native" if cfg.quant_backend == "auto" else cfg.quant_backend + ) _validate_precision_request( device=device, stage1_precision=cfg.stage1_precision, refiner_precision=cfg.refiner_precision, refiner_enabled=not cfg.no_refiner, + quant_backend=quant_backend, + ) + resolved_quant_backend = _resolve_quant_backend( + quant_backend, + _active_quantized_precisions( + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, + ), + ) + precision_env = _precision_env_updates( + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, ) + with _temporary_environment(precision_env): + prompt = self._resolve_prompt() + image, c2w, intrinsics_vec4, num_frames = self._prepare_native_inputs( + device + ) + pipeline_cfg = _native_pipeline_config( + cfg, + quant_backend=quant_backend, + ) + pipeline = pipeline_cfg.setup().to(device).eval() + transformer = pipeline.diffusion_model.transformer + if not isinstance(transformer, SanaWMNativeTransformer): + raise TypeError( + "Native SANA-WM runner resolved a non-SANA transformer: " + f"{type(transformer).__name__}." + ) + sampling_algo = self._native_sampling_algo() + if sampling_algo != "flow_euler_ltx": + raise ValueError( + "Native SANA-WM currently supports flow_euler_ltx only; " + f"got {sampling_algo!r}. Use sana-wm-upstream-reference " + "for chunk/self-forcing samplers." + ) + camera = prepare_camera( + c2w, + intrinsics_vec4, + target_size=(DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH), + ) + conditioning = transformer.prepare_conditioning( + image=image, + prompt=prompt, + camera=camera, + num_frames=num_frames, + fps=cfg.fps, + steps=cfg.step, + cfg_scale=cfg.cfg_scale, + flow_shift=cfg.flow_shift, + seed=cfg.seed, + negative_prompt=cfg.negative_prompt, + ) + cache = pipeline.initialize_cache( + transformer_context={"conditioning": conditioning} + ) + sana_latent = pipeline.generate(0, cache) + pipeline.finalize(0, cache) + transformer.release_stage1_runtime(cache) + del conditioning, cache + + video_hwc = transformer.decode_latents(sana_latent) + if not self.is_rank_zero: + return + _write_video(cfg.output_dir, cfg.name, video_hwc, cfg.fps) + + def _run_upstream_reference(self) -> None: + """Run the explicit upstream SANA-WM reference harness.""" + cfg = self.config + if cfg.image_path is None: + raise ValueError("SanaWMRunner requires --image-path.") + + device = self._resolve_device() + _validate_precision_request( + device=device, + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, + quant_backend=cfg.quant_backend, + ) + resolved_quant_backend = _resolve_quant_backend( + cfg.quant_backend, + _active_quantized_precisions( + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, + ), + ) + if ( + _active_quantized_precisions( + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, + ) + and resolved_quant_backend != "upstream-te" + ): + raise ValueError( + "The upstream reference runner uses upstream SANA's own " + "quantized path. Pass --quant-backend upstream-te for " + "FP8/FP4 reference comparisons, or use sana-wm-bidirectional " + "for the native FlashDreams low-precision backend." + ) precision_env = _precision_env_updates( stage1_precision=cfg.stage1_precision, @@ -288,7 +441,7 @@ def run(self) -> None: prompt = self._resolve_prompt() image = upstream.Image.open(cfg.image_path).convert("RGB") - c2w_full = self._resolve_trajectory(upstream) + c2w_full = self._resolve_trajectory() num_frames = min(cfg.num_frames, c2w_full.shape[0]) snapped = upstream._snap_num_frames( num_frames, @@ -411,6 +564,107 @@ def run(self) -> None: upstream.get_root_logger(), ) + def _prepare_native_inputs( + self, + device: torch.device, + ) -> tuple[object, np.ndarray, np.ndarray, int]: + """Load/crop the input image and prepare c2w/intrinsics for native run.""" + del device + from PIL import Image + + cfg = self.config + assert cfg.image_path is not None + image = Image.open(cfg.image_path).convert("RGB") + c2w_full = self._resolve_trajectory() + num_frames = min(cfg.num_frames, c2w_full.shape[0]) + snapped = snap_num_frames( + num_frames, + stride=SANA_WM_VAE_TEMPORAL_COMPRESSION, + upper_bound=c2w_full.shape[0], + ) + if snapped != cfg.num_frames and self.is_rank_zero: + logger.warning( + "SANA-WM requires num_frames = 8k+1; requested {} snapped to {} " + "(trajectory has {} frames).", + cfg.num_frames, + snapped, + c2w_full.shape[0], + ) + num_frames = snapped + c2w = c2w_full[:num_frames] + + resized_size, crop_offset = resize_center_crop_geometry( + image.size, + target_h=DEFAULT_VIDEO_HEIGHT, + target_w=DEFAULT_VIDEO_WIDTH, + ) + resized = image.resize(resized_size, Image.LANCZOS) + left, top = crop_offset + cropped = resized.crop( + ( + left, + top, + left + DEFAULT_VIDEO_WIDTH, + top + DEFAULT_VIDEO_HEIGHT, + ) + ) + if cfg.intrinsics_path is None: + raise ValueError( + "The native SANA-WM runner currently requires --intrinsics-path. " + "Use sana-wm-upstream-reference if you need Pi3X intrinsics estimation." + ) + intrinsics_src = load_intrinsics(cfg.intrinsics_path, num_frames) + intrinsics_vec4 = transform_intrinsics_for_crop( + intrinsics_src, + image.size, + resized_size, + crop_offset, + ) + return cropped, c2w, intrinsics_vec4, num_frames + + def _native_sampling_algo(self) -> SamplingAlgo: + """Resolve the native sampler selection.""" + if self.config.sampling_algo == "auto": + return "flow_euler_ltx" + return self.config.sampling_algo + + +def _native_pipeline_config( + cfg: SanaWMRunnerConfig, + *, + quant_backend: QuantBackend, +) -> StreamInferencePipelineConfig: + """Apply CLI runtime fields to the native SANA-WM pipeline literal.""" + return derive_config( + cfg.pipeline, + diffusion_model=dict( + seed=cfg.seed, + transformer=dict( + config_path=cfg.config_path, + checkpoint_path=cfg.model_path, + stage1_precision=cfg.stage1_precision, + quant_backend=quant_backend, + refiner_root=cfg.refiner_root, + refiner_gemma_root=cfg.refiner_gemma_root, + refiner_precision=cfg.refiner_precision, + offload_vae=cfg.offload_vae, + offload_refiner=cfg.offload_refiner, + offload_text_encoder=cfg.offload_text_encoder, + ), + ), + ) + + +def _write_video(output_dir: Path, name: str, video_hwc: np.ndarray, fps: int) -> Path: + """Write an HWC uint8 video to the runner output directory.""" + import imageio.v3 as iio + + output_dir.mkdir(parents=True, exist_ok=True) + video_path = output_dir / f"{name}_generated.mp4" + iio.imwrite(video_path, video_hwc, fps=fps) + logger.info("Saved {}", video_path) + return video_path + def _precision_env_updates( *, @@ -418,7 +672,7 @@ def _precision_env_updates( refiner_precision: Precision, refiner_enabled: bool, ) -> dict[str, str | None]: - """Return upstream SANA-WM precision env overrides.""" + """Return SANA-WM precision environment overrides.""" updates: dict[str, str | None] = {"SANA_WM_STAGE1_NVFP4": None} if stage1_precision != "bf16": updates.update( @@ -442,12 +696,12 @@ def _precision_env_updates( def _quant_env_value(precision: Precision) -> str: - """Return upstream Transformer Engine recipe selector for a precision.""" + """Return the SANA quantization recipe selector for a precision.""" if precision == "fp8": return "fp8block" if precision == "fp4": return "nvfp4" - raise ValueError(f"{precision!r} does not use Transformer Engine quantization.") + raise ValueError(f"{precision!r} does not use quantized SANA execution.") @contextmanager @@ -477,12 +731,14 @@ def _validate_precision_request( stage1_precision: Precision, refiner_precision: Precision, refiner_enabled: bool, + quant_backend: QuantBackend, ) -> None: """Fail early when requested quantized precision cannot run.""" - active_precisions = [stage1_precision] - if refiner_enabled: - active_precisions.append(refiner_precision) - quantized = [precision for precision in active_precisions if precision != "bf16"] + quantized = _active_quantized_precisions( + stage1_precision=stage1_precision, + refiner_precision=refiner_precision, + refiner_enabled=refiner_enabled, + ) if not quantized: return @@ -500,20 +756,58 @@ def _validate_precision_request( "SANA-WM fp8 precision requires a Hopper or newer GPU " f"(sm_90+); detected {sm_name}." ) - if "fp8" in quantized and not _torch_cuda_version_at_least(12, 9): - cuda_version = torch.version.cuda or "unknown" - raise ValueError( - "SANA-WM fp8 precision uses upstream Sana's Transformer Engine " - "Float8BlockScaling path, which requires CUDA 12.9 or newer. " - f"This PyTorch environment reports CUDA {cuda_version}." - ) if "fp4" in quantized and major < 10: raise ValueError( "SANA-WM fp4/NVFP4 precision requires a Blackwell GPU " f"(sm_100+); detected {sm_name}. Use bf16 or fp8 on this GPU." ) - _validate_transformer_engine(quantized) + resolved_backend = _resolve_quant_backend(quant_backend, quantized) + if resolved_backend == "torch-fp8": + _validate_torch_fp8_backend(quantized) + elif resolved_backend == "torch-fp4": + _validate_torch_fp4_backend(quantized) + elif resolved_backend == "native": + _validate_native_quant_backend(quantized) + elif resolved_backend == "upstream-te": + if "fp8" in quantized and not _torch_cuda_version_at_least(12, 9): + cuda_version = torch.version.cuda or "unknown" + raise ValueError( + "SANA-WM fp8 precision uses upstream Sana's Transformer Engine " + "Float8BlockScaling path, which requires CUDA 12.9 or newer. " + f"This PyTorch environment reports CUDA {cuda_version}. Use " + "--quant-backend torch-fp8 or --quant-backend native to try " + "the native PyTorch FP8 backend." + ) + _validate_transformer_engine(quantized) + else: + raise ValueError(f"Unsupported SANA-WM quant backend: {quant_backend!r}.") + + +def _active_quantized_precisions( + *, + stage1_precision: Precision, + refiner_precision: Precision, + refiner_enabled: bool, +) -> list[Precision]: + """Return active non-BF16 precision requests in execution order.""" + active_precisions = [stage1_precision] + if refiner_enabled: + active_precisions.append(refiner_precision) + return [precision for precision in active_precisions if precision != "bf16"] + + +def _resolve_quant_backend( + quant_backend: QuantBackend, + quantized_precisions: list[Precision], +) -> ResolvedQuantBackend: + """Resolve ``auto`` to the native low-precision backend.""" + del quantized_precisions + if quant_backend == "auto": + return "native" + if quant_backend in {"upstream-te", "torch-fp8", "torch-fp4", "native"}: + return quant_backend + raise ValueError(f"Unsupported SANA-WM quant backend: {quant_backend!r}.") def _torch_cuda_version_at_least(major: int, minor: int) -> bool: @@ -530,6 +824,59 @@ def _torch_cuda_version_at_least(major: int, minor: int) -> bool: return (current_major, current_minor) >= (major, minor) +def _validate_torch_fp8_backend(precisions: list[Precision]) -> None: + """Validate the native PyTorch scaled-MM backend request.""" + unsupported = sorted({precision for precision in precisions if precision != "fp8"}) + if unsupported: + raise ValueError( + "SANA-WM --quant-backend torch-fp8 currently supports fp8 only; " + f"unsupported requested precision(s): {', '.join(unsupported)}. " + "Use --quant-backend torch-fp4 or --quant-backend native for fp4." + ) + _validate_native_fp8_primitives() + + +def _validate_torch_fp4_backend(precisions: list[Precision]) -> None: + """Validate the native PyTorch NVFP4 backend request.""" + unsupported = sorted({precision for precision in precisions if precision != "fp4"}) + if unsupported: + raise ValueError( + "SANA-WM --quant-backend torch-fp4 currently supports fp4 only; " + f"unsupported requested precision(s): {', '.join(unsupported)}. " + "Use --quant-backend torch-fp8 or --quant-backend native for fp8." + ) + _validate_native_fp4_primitives() + + +def _validate_native_quant_backend(precisions: list[Precision]) -> None: + """Validate native PyTorch low-precision primitives.""" + if "fp8" in precisions: + _validate_native_fp8_primitives() + if "fp4" in precisions: + _validate_native_fp4_primitives() + + +def _validate_native_fp8_primitives() -> None: + if not hasattr(torch, "_scaled_mm") or not hasattr(torch, "float8_e4m3fn"): + raise RuntimeError( + "SANA-WM --quant-backend torch-fp8 requires PyTorch with " + "torch._scaled_mm and torch.float8_e4m3fn support." + ) + + +def _validate_native_fp4_primitives() -> None: + missing = [ + name + for name in ("_scaled_mm", "float4_e2m1fn_x2", "float8_e4m3fn") + if not hasattr(torch, name) + ] + if missing: + raise RuntimeError( + "SANA-WM native fp4 requires PyTorch with " + f"{', '.join(f'torch.{name}' for name in missing)} support." + ) + + def _validate_transformer_engine(precisions: list[Precision]) -> None: """Ensure Transformer Engine has the recipes required by precision modes.""" required_recipes: set[str] = set() @@ -610,7 +957,9 @@ def _import_upstream_sana(explicit_root: Path | None) -> ModuleType: __all__ = [ + "ExecutionBackend", "Precision", + "QuantBackend", "SamplingAlgo", "SanaWMRunner", "SanaWMRunnerConfig", diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py new file mode 100644 index 000000000..74f21cff2 --- /dev/null +++ b/integrations/sana/sana_wm/stage1_model.py @@ -0,0 +1,514 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""FlashDreams-owned SANA-WM Stage-1 DiT module definitions. + +This file intentionally models the public SANA-WM bidirectional Stage-1 +checkpoint schema without importing the upstream SANA ``diffusion`` package. +The module names and tensor shapes are checkpoint-facing API: keep them stable +unless the public checkpoint contract changes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + + +@dataclass(frozen=True) +class SanaWMStage1Spec: + """Static architecture values for the public SANA-WM bidirectional DiT.""" + + latent_channels: int = 128 + hidden_size: int = 2240 + text_dim: int = 2304 + timestep_dim: int = 256 + depth: int = 20 + num_heads: int = 20 + head_dim: int = 112 + max_text_length: int = 300 + latent_grid_size: tuple[int, int] = (22, 40) + mlp_ratio: int = 3 + conv_kernel_size: int = 4 + temporal_kernel_size: int = 3 + plucker_channels: int = 48 + raymap_channels: int = 3 + softmax_every_n: int = 4 + + @property + def mlp_inner_size(self) -> int: + """Return the GLUMBConv hidden expansion width.""" + return self.hidden_size * self.mlp_ratio * 2 + + @property + def gated_mlp_size(self) -> int: + """Return the width consumed by the pointwise output projection.""" + return self.hidden_size * self.mlp_ratio + + def block_uses_gdn(self, index: int) -> bool: + """Return whether a block has GDN convolution checkpoint tensors.""" + return (index + 1) % self.softmax_every_n != 0 + + +SANA_WM_STAGE1_SPEC = SanaWMStage1Spec() +"""Architecture spec for the public SANA-WM bidirectional Stage-1 checkpoint.""" + + +class RMSNorm(nn.Module): + """RMSNorm parameter container matching the SANA checkpoint schema.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(hidden_size)) + self.eps = eps + + def forward(self, x: Tensor) -> Tensor: + """Apply RMS normalization using the stored scale parameter.""" + dtype = x.dtype + normed = x.float() * torch.rsqrt( + x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps + ) + return (normed * self.weight).to(dtype=dtype) + + +class Conv3dProjector(nn.Module): + """Named 1x1x1 projection used by latent, ray, and Plucker embedders.""" + + def __init__(self, in_channels: int, hidden_size: int) -> None: + super().__init__() + self.proj = nn.Conv3d(in_channels, hidden_size, kernel_size=1) + + def forward(self, x: Tensor) -> Tensor: + """Project a 5D tensor to hidden channels.""" + return self.proj(x) + + +class TimestepEmbedder(nn.Module): + """Two-layer timestep embedder with checkpoint-compatible names.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(spec.timestep_dim, spec.hidden_size), + nn.SiLU(), + nn.Linear(spec.hidden_size, spec.hidden_size), + ) + + def forward(self, t: Tensor) -> Tensor: + """Project timestep features into the DiT hidden width.""" + t_freq = _timestep_embedding(t.flatten(), self.mlp[0].in_features) + return self.mlp(t_freq.to(device=t.device, dtype=self.mlp[0].weight.dtype)) + + +class TextProjection(nn.Module): + """Two-layer text projection matching ``y_embedder.y_proj`` keys.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.fc1 = nn.Linear(spec.text_dim, spec.hidden_size) + self.act = nn.SiLU() + self.fc2 = nn.Linear(spec.hidden_size, spec.hidden_size) + + def forward(self, y: Tensor) -> Tensor: + """Project text encoder activations into the DiT hidden width.""" + return self.fc2(self.act(self.fc1(y))) + + +class TextEmbedder(nn.Module): + """Text embedding/projection container matching SANA checkpoint keys.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.y_embedding = nn.Parameter(torch.empty(spec.max_text_length, spec.text_dim)) + self.y_proj = TextProjection(spec) + + def forward(self, y: Tensor | None) -> Tensor: + """Project text embeddings, using the learned null embedding when absent.""" + if y is None: + y = self.y_embedding.unsqueeze(0) + if y.ndim == 4 and y.shape[1] == 1: + y = y.squeeze(1) + return self.y_proj(y) + + +class FinalLayer(nn.Module): + """Final AdaLN and latent-channel projection container.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.norm_final = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.scale_shift_table = nn.Parameter(torch.empty(2, spec.hidden_size)) + self.linear = nn.Linear(spec.hidden_size, spec.latent_channels) + + def forward(self, x: Tensor, t: Tensor, *, frames: int) -> Tensor: + """Project hidden tokens back to latent channels.""" + if t.ndim > 2: + batch, tokens, channels = x.shape + shift, scale = ( + self.scale_shift_table[None, None, :, :] + t.transpose(1, 2) + ).chunk(2, dim=-2) + x = _modulate( + self.norm_final(x).reshape(batch, frames, -1, channels), + shift, + scale, + ).reshape(batch, tokens, channels) + else: + shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk( + 2, + dim=1, + ) + x = _modulate(self.norm_final(x), shift, scale) + return self.linear(x) + + +class Conv2dContainer(nn.Module): + """Expose a convolution under a stable ``.conv`` attribute.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + *, + groups: int = 1, + bias: bool = True, + padding: int | tuple[int, int] = 0, + ) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size=kernel_size, + groups=groups, + bias=bias, + padding=padding, + ) + + def forward(self, x: Tensor) -> Tensor: + """Apply the contained convolution.""" + return self.conv(x) + + +class GLUMBConvTemp(nn.Module): + """Checkpoint-compatible GLUMBConvTemp feed-forward container.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.inverted_conv = Conv2dContainer(spec.hidden_size, spec.mlp_inner_size, 1) + self.depth_conv = Conv2dContainer( + spec.mlp_inner_size, + spec.mlp_inner_size, + 3, + groups=spec.mlp_inner_size, + padding=1, + ) + self.point_conv = Conv2dContainer( + spec.gated_mlp_size, + spec.hidden_size, + 1, + bias=False, + ) + self.t_conv = nn.Conv2d( + spec.hidden_size, + spec.hidden_size, + kernel_size=(spec.temporal_kernel_size, 1), + padding=(spec.temporal_kernel_size // 2, 0), + bias=False, + ) + + def forward(self, x: Tensor, *, frames: int, height: int, width: int) -> Tensor: + """Run spatial GLU plus temporal aggregation.""" + batch, tokens, channels = x.shape + x_2d = x.reshape(batch * frames, height, width, channels).permute(0, 3, 1, 2) + x_2d = self.inverted_conv(x_2d) + x_2d = self.depth_conv(x_2d) + value, gate = x_2d.chunk(2, dim=1) + x_2d = self.point_conv(value * F.silu(gate)) + + x_time = x_2d.view(batch, frames, channels, height * width).permute(0, 2, 1, 3) + x_time = x_time + self.t_conv(x_time) + return x_time.permute(0, 2, 3, 1).reshape(batch, tokens, channels) + + +class Stage1SelfAttention(nn.Module): + """Self/camera attention parameter container for one Stage-1 block.""" + + def __init__(self, spec: SanaWMStage1Spec, *, use_gdn_convs: bool) -> None: + super().__init__() + self.A_log = nn.Parameter(torch.empty(spec.num_heads)) + self.beta_proj = nn.Linear(spec.hidden_size, spec.num_heads) + self.dt_bias = nn.Parameter(torch.empty(spec.num_heads)) + self.gate_proj = nn.Linear(spec.hidden_size, spec.num_heads) + self.k_norm = RMSNorm(spec.hidden_size) + self.k_norm_cam = RMSNorm(spec.hidden_size) + self.k_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + self.out_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + self.output_gate = nn.Linear(spec.hidden_size, spec.hidden_size) + self.proj = nn.Linear(spec.hidden_size, spec.hidden_size) + self.q_norm = RMSNorm(spec.hidden_size) + self.q_norm_cam = RMSNorm(spec.hidden_size) + self.q_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + self.qkv = nn.Linear(spec.hidden_size, 3 * spec.hidden_size, bias=False) + self.recall_gate = nn.Parameter(torch.empty(1)) + self.v_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + if use_gdn_convs: + self.conv_k = nn.Conv1d( + spec.hidden_size, + spec.hidden_size, + kernel_size=spec.conv_kernel_size, + groups=spec.hidden_size, + bias=False, + ) + self.conv_k_cam = nn.Conv1d( + spec.hidden_size, + spec.hidden_size, + kernel_size=spec.conv_kernel_size, + groups=spec.hidden_size, + bias=False, + ) + + def forward(self, x: Tensor, *_args: object, **_kwargs: object) -> Tensor: + """Run a memory-bounded native self-attention approximation. + + The public checkpoint's GDN blocks are linear-recurrent attention + blocks. Until the native GDN scan is implemented, this path uses the + checkpoint QKV/value/projection tensors without materializing an + all-token softmax attention matrix. + """ + batch, tokens, channels = x.shape + qkv = self.qkv(x).view(batch, tokens, 3, channels) + q, _k, v = qkv.unbind(dim=2) + q = self.q_norm(q) + v = v * torch.sigmoid(q) + out = F.silu(self.output_gate(x).float()).to(dtype=v.dtype) * v + return self.proj(out.to(dtype=self.proj.weight.dtype)) + + +class Stage1CrossAttention(nn.Module): + """Cross-attention parameter container for one Stage-1 block.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.num_heads = spec.num_heads + self.head_dim = spec.hidden_size // spec.num_heads + self.k_norm = RMSNorm(spec.hidden_size) + self.kv_linear = nn.Linear(spec.hidden_size, 2 * spec.hidden_size) + self.proj = nn.Linear(spec.hidden_size, spec.hidden_size) + self.q_linear = nn.Linear(spec.hidden_size, spec.hidden_size) + self.q_norm = RMSNorm(spec.hidden_size) + + def forward( + self, + x: Tensor, + y: Tensor, + *, + mask: Tensor | None = None, + **_kwargs: object, + ) -> Tensor: + """Cross-attention execution is implemented in a later native pass.""" + batch, tokens, channels = x.shape + q = self.q_norm(self.q_linear(x)).view( + batch, + tokens, + self.num_heads, + self.head_dim, + ) + kv = self.kv_linear(y).view(batch, -1, 2, channels) + k, v = kv.unbind(dim=2) + k = self.k_norm(k).view(batch, -1, self.num_heads, self.head_dim) + v = v.view(batch, -1, self.num_heads, self.head_dim) + + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + attn_mask = None + if mask is not None: + attn_mask = (1 - mask.to(dtype=q.dtype)) * -10000.0 + attn_mask = attn_mask[:, None, None, :] + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + ) + out = out.transpose(1, 2).reshape(batch, tokens, channels) + return self.proj(out.to(dtype=self.proj.weight.dtype)) + + +class SanaWMStage1Block(nn.Module): + """One checkpoint-compatible SANA-WM Stage-1 transformer block.""" + + def __init__(self, spec: SanaWMStage1Spec, *, index: int) -> None: + super().__init__() + self.norm1 = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.norm2 = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.scale_shift_table = nn.Parameter(torch.empty(6, spec.hidden_size)) + self.attn = Stage1SelfAttention( + spec, + use_gdn_convs=spec.block_uses_gdn(index), + ) + self.cross_attn = Stage1CrossAttention(spec) + self.mlp = GLUMBConvTemp(spec) + self.plucker_proj = nn.Linear(spec.hidden_size, spec.hidden_size) + + def forward( + self, + x: Tensor, + y: Tensor, + t: Tensor, + *, + frames: int, + height: int, + width: int, + mask: Tensor | None = None, + plucker_emb: Tensor | None = None, + **kwargs: object, + ) -> Tensor: + """Run one native Stage-1 transformer block.""" + batch, tokens, channels = x.shape + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None, :, :] + t.reshape(batch, frames, 6, -1) + ).chunk(6, dim=-2) + + x_norm = self.norm1(x).reshape(batch, frames, -1, channels) + attn_in = _modulate(x_norm, shift_msa, scale_msa).reshape(batch, tokens, channels) + attn_out = self.attn(attn_in, **kwargs).reshape(batch, frames, -1, channels) + x = x + (gate_msa * attn_out).reshape(batch, tokens, channels) + + if plucker_emb is not None: + x = x + self.plucker_proj(plucker_emb) + + x = x + self.cross_attn(x, y, mask=mask) + + x_norm = self.norm2(x).reshape(batch, frames, -1, channels) + mlp_in = _modulate(x_norm, shift_mlp, scale_mlp).reshape(batch, tokens, channels) + mlp_out = self.mlp(mlp_in, frames=frames, height=height, width=width) + mlp_out = mlp_out.reshape(batch, frames, -1, channels) + return x + (gate_mlp * mlp_out).reshape(batch, tokens, channels) + + +class SanaWMStage1Model(nn.Module): + """Checkpoint-compatible native Stage-1 SANA-WM DiT shell.""" + + def __init__(self, spec: SanaWMStage1Spec = SANA_WM_STAGE1_SPEC) -> None: + super().__init__() + self.spec = spec + self.x_embedder = Conv3dProjector(spec.latent_channels, spec.hidden_size) + self.t_embedder = TimestepEmbedder(spec) + self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(spec.hidden_size, 6 * spec.hidden_size)) + self.y_embedder = TextEmbedder(spec) + self.attention_y_norm = RMSNorm(spec.hidden_size) + self.raymap_embedder = Conv3dProjector(spec.raymap_channels, spec.hidden_size) + self.plucker_embedder = Conv3dProjector(spec.plucker_channels, spec.hidden_size) + self.blocks = nn.ModuleList( + SanaWMStage1Block(spec, index=index) for index in range(spec.depth) + ) + num_pos_tokens = spec.latent_grid_size[0] * spec.latent_grid_size[1] // 2 + 44 + self.pos_embed = nn.Parameter(torch.empty(1, num_pos_tokens, spec.hidden_size)) + self.final_layer = FinalLayer(spec) + + def forward( + self, + x: Tensor, + timestep: Tensor, + y: Tensor, + *, + mask: Tensor | None = None, + chunk_plucker: Tensor | None = None, + **kwargs: object, + ) -> Tensor: + """Run the native SANA-WM Stage-1 DiT.""" + batch, _channels, frames, height, width = x.shape + x = self.x_embedder(x) + x = x.permute(0, 2, 3, 4, 1).reshape(batch, frames * height * width, -1) + + plucker_emb = None + if chunk_plucker is not None: + plucker_emb = self.plucker_embedder(chunk_plucker) + plucker_emb = plucker_emb.permute(0, 2, 3, 4, 1).reshape_as(x) + + y = self.y_embedder(y) + y = self.attention_y_norm(y) + if mask is not None and mask.ndim > 2: + mask = mask.squeeze(1).squeeze(1) + + timestep_embed = self.t_embedder(timestep.flatten()) + timestep_embed = timestep_embed.unflatten(0, timestep.shape) + block_t = self.t_block(timestep_embed).reshape(batch, 1, frames, 6 * self.spec.hidden_size) + + for block in self.blocks: + x = block( + x, + y, + block_t, + frames=frames, + height=height, + width=width, + mask=mask, + plucker_emb=plucker_emb, + **kwargs, + ) + + x = self.final_layer(x, timestep_embed, frames=frames) + return x.reshape(batch, frames, height, width, -1).permute(0, 4, 1, 2, 3) + + +def _modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor: + return x * (1 + scale) + shift + + +def _timestep_embedding(t: Tensor, dim: int, max_period: int = 10000) -> Tensor: + 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 + + +__all__ = [ + "SANA_WM_STAGE1_SPEC", + "GLUMBConvTemp", + "RMSNorm", + "SanaWMStage1Block", + "SanaWMStage1Model", + "SanaWMStage1Spec", + "Stage1CrossAttention", + "Stage1SelfAttention", +] diff --git a/integrations/sana/sana_wm/placeholder.py b/integrations/sana/sana_wm/upstream_reference.py similarity index 55% rename from integrations/sana/sana_wm/placeholder.py rename to integrations/sana/sana_wm/upstream_reference.py index ef46f1690..6b87504e2 100644 --- a/integrations/sana/sana_wm/placeholder.py +++ b/integrations/sana/sana_wm/upstream_reference.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Placeholder FlashDreams pipeline objects for the upstream SANA-WM runner.""" +"""Schema-only FlashDreams objects for the upstream SANA-WM reference runner.""" from __future__ import annotations @@ -32,27 +32,34 @@ @dataclass(kw_only=True) -class SanaWMPlaceholderTransformerConfig(TransformerConfig): - """Transformer config used only to satisfy the runner config schema.""" +class SanaWMUpstreamReferenceTransformerConfig(TransformerConfig): + """Transformer config used by the upstream-reference runner. - _target: type["SanaWMPlaceholderTransformer"] = field( - default_factory=lambda: SanaWMPlaceholderTransformer + This is intentionally not an executable SANA DiT adapter. It exists so the + CLI can expose the upstream reference path through the normal runner config + registry without pretending that upstream execution is the native path. + """ + + _target: type["SanaWMUpstreamReferenceTransformer"] = field( + default_factory=lambda: SanaWMUpstreamReferenceTransformer ) -class SanaWMPlaceholderTransformer(Transformer[TransformerAutoregressiveCache]): - """Non-executable transformer for the upstream-delegating SANA-WM runner.""" +class SanaWMUpstreamReferenceTransformer( + Transformer[TransformerAutoregressiveCache] +): + """Non-executable marker for the upstream-delegating reference runner.""" - def __init__(self, config: SanaWMPlaceholderTransformerConfig) -> None: + def __init__(self, config: SanaWMUpstreamReferenceTransformerConfig) -> None: super().__init__(config) self._dummy = nn.Parameter(torch.empty(0)) @property def latent_shape(self) -> tuple[int, ...]: - """Raise because SANA-WM generation is delegated to upstream Sana.""" + """Raise because this config marks the upstream reference path.""" raise RuntimeError( - "SANA-WM bidirectional uses the custom SanaWMRunner path; " - "the placeholder FlashDreams pipeline is not executable." + "This SANA-WM runner is an upstream reference harness. Use " + "sana-wm-bidirectional for native FlashDreams execution." ) def predict_flow( @@ -62,10 +69,10 @@ def predict_flow( cache: TransformerAutoregressiveCache, input: Any = None, ) -> Tensor: - """Raise because SANA-WM generation is delegated to upstream Sana.""" + """Raise because this config marks the upstream reference path.""" raise RuntimeError( - "SANA-WM bidirectional uses the custom SanaWMRunner path; " - "the placeholder FlashDreams pipeline is not executable." + "This SANA-WM runner delegates generation to upstream NVlabs/Sana " + "instead of executing through Transformer.predict_flow." ) def patchify_and_maybe_split_cp(self, x: Any) -> Any: @@ -75,3 +82,9 @@ def patchify_and_maybe_split_cp(self, x: Any) -> Any: def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: """Return ``x`` unchanged for schema-only construction.""" return x + + +__all__ = [ + "SanaWMUpstreamReferenceTransformer", + "SanaWMUpstreamReferenceTransformerConfig", +] diff --git a/integrations/sana/tests/parity_check/README.md b/integrations/sana/tests/parity_check/README.md index 08d5d5abf..62b13519b 100644 --- a/integrations/sana/tests/parity_check/README.md +++ b/integrations/sana/tests/parity_check/README.md @@ -5,16 +5,20 @@ SPDX-License-Identifier: Apache-2.0 # SANA-WM parity check -This directory is reserved for the opt-in SANA-WM upstream parity +This directory is reserved for the opt-in SANA-WM native-vs-upstream parity harness. It should follow the existing integration pattern: - clone or reuse a pinned NVlabs/Sana checkout; - install upstream-heavy dependencies in an isolated environment; -- run upstream `inference_video_scripts/wm/inference_sana_wm.py` and - `flashdreams-run sana-wm-bidirectional` on the same image, prompt, - camera, intrinsics, seed, and sampler settings; +- run upstream `inference_video_scripts/wm/inference_sana_wm.py` and the + native `flashdreams-run sana-wm-bidirectional` path on the same image, + prompt, camera, intrinsics, seed, and sampler settings; - compare decoded MP4 frames and report mean / max `|Delta|`. +The temporary FlashDreams-packaged upstream-reference runner is documented in +`../../docs/upstream_reference.md`. Keep those wrapper details out of the native +integration README. + Do not put the full checkpoint download, refiner execution, or MP4 generation path in `ci_cpu`. Add a bounded `ci_gpu` gate only if CI pre-provisions the required artefacts and the test skips cleanly when diff --git a/integrations/sana/tests/test_native_quant_cuda.py b/integrations/sana/tests/test_native_quant_cuda.py new file mode 100644 index 000000000..c32e2e65a --- /dev/null +++ b/integrations/sana/tests/test_native_quant_cuda.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""CUDA smoke tests for SANA-WM native low-precision linear replacements.""" + +from __future__ import annotations + +import pytest +import torch + +from sana_wm.native_quant import TorchScaledMMFP4Linear, TorchScaledMMFP8Linear + +pytestmark = pytest.mark.ci_gpu + + +def _require_native_quant_cuda() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for native quantization smokes") + missing = [ + name + for name in ("_scaled_mm", "float8_e4m3fn", "float4_e2m1fn_x2") + if not hasattr(torch, name) + ] + if missing: + pytest.skip(f"PyTorch lacks native quantization primitive(s): {missing}") + + +def test_native_fp8_linear_runs_scaled_mm_on_cuda() -> None: + """Exercise the TE-free E4M3 FP8 Linear replacement on CUDA.""" + _require_native_quant_cuda() + source = torch.nn.Linear(32, 64, bias=True, device="cuda", dtype=torch.bfloat16) + quantized = TorchScaledMMFP8Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + inputs = torch.randn(4, 5, 32, device="cuda", dtype=torch.bfloat16) + + output = quantized(inputs) + + assert quantized.weight.shape == source.weight.shape + assert quantized.bias is not None + assert output.shape == (4, 5, 64) + assert output.dtype == torch.bfloat16 + assert torch.isfinite(output.float()).all() + + +def test_native_fp4_linear_runs_scaled_mm_on_blackwell() -> None: + """Exercise the TE-free E2M1 NVFP4 Linear replacement on Blackwell CUDA.""" + _require_native_quant_cuda() + major, minor = torch.cuda.get_device_capability() + if major < 10: + pytest.skip(f"NVFP4 requires Blackwell-class CUDA, got sm_{major}{minor}") + + source = torch.nn.Linear(32, 64, bias=True, device="cuda", dtype=torch.bfloat16) + quantized = TorchScaledMMFP4Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + inputs = torch.randn(4, 5, 32, device="cuda", dtype=torch.bfloat16) + + output = quantized(inputs) + + assert quantized.weight.shape == source.weight.shape + assert quantized.bias is not None + assert output.shape == (4, 5, 64) + assert output.dtype == torch.bfloat16 + assert torch.isfinite(output.float()).all() diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 6c7c032ed..692795bad 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CPU-safe smoke tests for the SANA-WM runner shim.""" +"""CPU-safe smoke tests for the SANA-WM native and reference configs.""" from __future__ import annotations @@ -30,8 +30,12 @@ from sana_wm.config import ( PIPELINE_SANA_WM_BIDIRECTIONAL, + PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL, + PIPELINE_SANA_WM_UPSTREAM_REFERENCE, RUNNER_CONFIGS, RUNNER_SANA_WM_BIDIRECTIONAL, + RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, + RUNNER_SANA_WM_UPSTREAM_REFERENCE, ) from sana_wm.constants import ( SANA_WM_CONFIG_PATH, @@ -42,9 +46,30 @@ SanaWMRunner, SanaWMRunnerConfig, _precision_env_updates, + _resolve_quant_backend, _temporary_environment, _validate_precision_request, ) +from sana_wm.native_diffusion import SanaWMDiffusionModelConfig +from sana_wm.native_transformer import ( + SanaWMNativeTransformerCache, + SanaWMNativeTransformerConfig, + SanaWMStage1Conditioning, + _avoid_degenerate_tile_tail, + _load_inference_config, +) +from sana_wm.native_quant import ( + TorchScaledMMFP4Linear, + TorchScaledMMFP8Linear, + replace_linear_with_torch_fp4, + replace_linear_with_torch_fp8, +) +from sana_wm.stage1_model import ( + SANA_WM_STAGE1_SPEC, + SanaWMStage1Model, + SanaWMStage1Spec, +) +from sana_wm.upstream_reference import SanaWMUpstreamReferenceTransformerConfig from flashdreams.infra.config import derive_config @@ -54,8 +79,12 @@ def test_runner_config_is_registered() -> None: - """Expose one SANA-WM runner keyed by its public slug.""" - assert RUNNER_CONFIGS == {"sana-wm-bidirectional": RUNNER_SANA_WM_BIDIRECTIONAL} + """Expose SANA-WM native and upstream-reference runner slugs.""" + assert RUNNER_CONFIGS == { + "sana-wm-bidirectional": RUNNER_SANA_WM_BIDIRECTIONAL, + "sana-wm-upstream-reference": RUNNER_SANA_WM_UPSTREAM_REFERENCE, + "sana-wm-native-bidirectional": RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, + } def test_runner_name_mirrors_pipeline_name() -> None: @@ -63,11 +92,313 @@ def test_runner_name_mirrors_pipeline_name() -> None: assert RUNNER_SANA_WM_BIDIRECTIONAL.runner_name == ( PIPELINE_SANA_WM_BIDIRECTIONAL.name ) + assert RUNNER_SANA_WM_UPSTREAM_REFERENCE.runner_name == ( + PIPELINE_SANA_WM_UPSTREAM_REFERENCE.name + ) + assert RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL.runner_name == ( + PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL.name + ) def test_runner_has_description() -> None: """Provide non-empty CLI help text for the runner registry.""" assert RUNNER_SANA_WM_BIDIRECTIONAL.description.strip() + assert RUNNER_SANA_WM_UPSTREAM_REFERENCE.description.strip() + assert RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL.description.strip() + + +def test_reference_and_native_pipelines_are_distinct() -> None: + """Keep the upstream harness separate from the native FlashDreams target.""" + reference_transformer = ( + PIPELINE_SANA_WM_UPSTREAM_REFERENCE.diffusion_model.transformer + ) + main_transformer = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.transformer + native_transformer = ( + PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL.diffusion_model.transformer + ) + + assert isinstance(reference_transformer, SanaWMUpstreamReferenceTransformerConfig) + assert isinstance( + PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model, + SanaWMDiffusionModelConfig, + ) + assert isinstance(main_transformer, SanaWMNativeTransformerConfig) + assert isinstance(native_transformer, SanaWMNativeTransformerConfig) + assert RUNNER_SANA_WM_BIDIRECTIONAL.execution_backend == "native-flashdreams" + assert RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL.execution_backend == ( + "native-flashdreams" + ) + + +def test_native_transformer_contract_shape_and_conditioning_guard() -> None: + """Pin the native SANA-WM Stage-1 boundary to the public model layout.""" + transformer_cfg = SanaWMNativeTransformerConfig() + transformer = transformer_cfg.setup() + + assert transformer.latent_shape == (1, 128, 21, 22, 40) + cache = transformer.initialize_autoregressive_cache() + + with pytest.raises(RuntimeError, match="without conditioning"): + transformer.predict_flow( + noisy_latent=torch.empty(transformer.latent_shape), + timestep=torch.tensor(1000.0), + cache=cache, + ) + + +def test_native_inference_config_loads_yaml_without_upstream_types( + tmp_path: Path, +) -> None: + """Parse SANA-WM YAML with local config objects, not SANA pyrallis types.""" + config_path = tmp_path / "config.yaml" + config_path.write_text( + """ +model: + model: SanaMSVideoCamCtrl_1600M_P1_D20 + mixed_precision: bf16 + chunk_split_strategy: first_chunk_plus_one +vae: + vae_type: LTX2VAE_diffusers + vae_pretrained: hf://example/model + weight_dtype: bfloat16 + vae_latent_dim: 128 + vae_stride: [8, 32, 32] +text_encoder: + text_encoder_name: gemma-2-2b-it + model_max_length: 300 + chi_prompt: ["prefix"] +scheduler: + flow_shift: 9.95 + inference_flow_shift: 9.8 +""", + encoding="utf-8", + ) + + cfg = _load_inference_config(str(config_path)) + + assert cfg.model.model == "SanaMSVideoCamCtrl_1600M_P1_D20" + assert cfg.model.get("missing", "fallback") == "fallback" + assert cfg.vae.vae_stride == [8, 32, 32] + assert cfg.text_encoder.chi_prompt == ["prefix"] + assert cfg.scheduler.inference_flow_shift == 9.8 + assert cfg.work_dir == "" + + +def test_native_stage1_model_matches_checkpoint_schema() -> None: + """Pin the FlashDreams-owned Stage-1 module to the public checkpoint schema.""" + state = SanaWMStage1Model().state_dict() + + assert len(state) == 872 + assert tuple(state["x_embedder.proj.weight"].shape) == (2240, 128, 1, 1, 1) + assert tuple(state["raymap_embedder.proj.weight"].shape) == (2240, 3, 1, 1, 1) + assert tuple(state["plucker_embedder.proj.weight"].shape) == ( + 2240, + 48, + 1, + 1, + 1, + ) + assert tuple(state["pos_embed"].shape) == (1, 484, 2240) + assert tuple(state["y_embedder.y_embedding"].shape) == (300, 2304) + assert tuple(state["blocks.0.attn.qkv.weight"].shape) == (6720, 2240) + assert tuple(state["blocks.0.attn.conv_k.weight"].shape) == (2240, 1, 4) + assert tuple(state["blocks.0.cross_attn.kv_linear.weight"].shape) == ( + 4480, + 2240, + ) + assert tuple(state["blocks.19.mlp.t_conv.weight"].shape) == (2240, 2240, 3, 1) + assert tuple(state["final_layer.linear.weight"].shape) == (128, 2240) + + for block_index in range(SANA_WM_STAGE1_SPEC.depth): + has_gdn_conv = f"blocks.{block_index}.attn.conv_k.weight" in state + assert has_gdn_conv is SANA_WM_STAGE1_SPEC.block_uses_gdn(block_index) + + +def test_native_stage1_forward_preserves_latent_shape() -> None: + """Exercise the native Stage-1 forward path on a small CPU-safe spec.""" + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=2, + num_heads=4, + head_dim=4, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + model = SanaWMStage1Model(spec) + latents = torch.randn(1, 4, 3, 2, 2) + timesteps = torch.ones(1, 1, 3) + text = torch.randn(1, 1, 5, 12) + mask = torch.ones(1, 5) + plucker = torch.randn(1, 6, 3, 2, 2) + + out = model(latents, timesteps, text, mask=mask, chunk_plucker=plucker) + + assert out.shape == latents.shape + + +def test_native_transformer_releases_stage1_runtime() -> None: + """Free Stage-1-only modules and conditioning before decode/refine.""" + transformer = SanaWMNativeTransformerConfig().setup() + transformer.model = torch.nn.Linear(1, 1) + transformer.text_encoder = torch.nn.Linear(1, 1) + transformer.tokenizer = object() + transformer._model_built = True + transformer._text_encoder_built = True + transformer._stage1_quantized = True + transformer._streaming_prompt_cache[("prompt",)] = ( + torch.empty(1), + torch.empty(1), + torch.empty(1), + torch.empty(1), + ) + cache = SanaWMNativeTransformerCache( + conditioning=SanaWMStage1Conditioning( + condition=torch.empty(1), + uncondition=None, + model_kwargs={}, + first_latent=torch.empty(1), + latent_shape=(1, 1, 1, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + ) + + transformer.release_stage1_runtime(cache) + + assert cache.conditioning is None + assert transformer.model is None + assert transformer.text_encoder is None + assert transformer.tokenizer is None + assert transformer._streaming_prompt_cache == {} + assert transformer._model_built is False + assert transformer._text_encoder_built is False + assert transformer._stage1_quantized is False + + +def test_native_vae_tiling_uses_low_memory_tiles() -> None: + """Keep native VAE decode on smaller tiles than upstream defaults.""" + + class DummyVAE: + def __init__(self) -> None: + self.calls: list[dict[str, int]] = [] + self.tile_sample_min_height = 512 + self.tile_sample_stride_height = 448 + self.tile_sample_min_width = 512 + self.tile_sample_stride_width = 448 + self.tile_sample_min_num_frames = 96 + self.tile_sample_stride_num_frames = 64 + self.use_framewise_encoding = False + self.use_framewise_decoding = False + self.spatial_compression_ratio = 32 + + def enable_tiling(self, **kwargs: int) -> None: + self.calls.append(kwargs) + + transformer = SanaWMNativeTransformerConfig().setup() + transformer.vae = DummyVAE() + + transformer._configure_vae_tiling() + + assert transformer.vae.calls == [ + { + "tile_sample_min_height": 256, + "tile_sample_stride_height": 192, + "tile_sample_min_width": 256, + "tile_sample_stride_width": 224, + "tile_sample_min_num_frames": 24, + "tile_sample_stride_num_frames": 8, + } + ] + assert transformer.vae.tile_sample_min_height == 256 + assert transformer.vae.tile_sample_stride_height == 192 + assert transformer.vae.tile_sample_min_width == 256 + assert transformer.vae.tile_sample_stride_width == 224 + assert transformer.vae.tile_sample_min_num_frames == 24 + assert transformer.vae.tile_sample_stride_num_frames == 8 + assert transformer.vae.use_framewise_encoding is True + assert transformer.vae.use_framewise_decoding is True + + +def test_native_decode_retries_vae_oom_with_smaller_tiles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Retry native VAE decode with smaller tiles after a CUDA OOM.""" + + class DummyVAE: + def __init__(self) -> None: + self.tile_sample_min_height = 256 + self.tile_sample_stride_height = 224 + self.tile_sample_min_width = 256 + self.tile_sample_stride_width = 224 + self.tile_sample_min_num_frames = 24 + self.tile_sample_stride_num_frames = 8 + self.use_framewise_encoding = False + self.use_framewise_decoding = False + self.spatial_compression_ratio = 32 + + def to(self, *_args: object, **_kwargs: object) -> "DummyVAE": + return self + + def enable_tiling(self, **kwargs: int) -> None: + for name, value in kwargs.items(): + setattr(self, name, value) + + transformer = SanaWMNativeTransformerConfig().setup() + transformer.vae = DummyVAE() + transformer.vae_dtype = torch.float32 + monkeypatch.setattr(transformer, "_ensure_vae", lambda: None) + calls = 0 + + def fake_decode(_samples: torch.Tensor) -> torch.Tensor: + nonlocal calls + calls += 1 + if calls == 1: + raise torch.OutOfMemoryError("test OOM") + return torch.zeros((1, 3, 1, 2, 2), dtype=torch.float32) + + monkeypatch.setattr(transformer, "_vae_decode", fake_decode) + + video = transformer.decode_latents(torch.zeros((1, 4, 1, 1, 1))) + + assert calls == 2 + assert video.shape == (1, 2, 2, 3) + assert transformer.vae.tile_sample_min_height == 128 + assert transformer.vae.tile_sample_stride_height == 64 + assert transformer.vae.tile_sample_min_width == 128 + assert transformer.vae.tile_sample_stride_width == 64 + + +def test_native_vae_tiling_avoids_degenerate_latent_tails() -> None: + """Avoid last spatial VAE tiles with size one after compression.""" + assert ( + _avoid_degenerate_tile_tail( + sample_extent=704, + sample_tile_min=256, + sample_stride=224, + compression_ratio=32, + ) + == 192 + ) + assert ( + _avoid_degenerate_tile_tail( + sample_extent=1280, + sample_tile_min=128, + sample_stride=112, + compression_ratio=32, + ) + == 64 + ) def test_hf_defaults_point_at_bidirectional_release() -> None: @@ -99,12 +430,22 @@ def test_runner_setup_does_not_import_upstream_sana() -> None: def test_runner_config_type() -> None: """Keep the exported literal on the SANA-WM runner config subclass.""" assert isinstance(RUNNER_SANA_WM_BIDIRECTIONAL, SanaWMRunnerConfig) + assert isinstance(RUNNER_SANA_WM_UPSTREAM_REFERENCE, SanaWMRunnerConfig) + assert isinstance(RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, SanaWMRunnerConfig) def test_runner_defaults_to_bf16_precision() -> None: """Keep the default path on upstream SANA-WM's BF16 config.""" assert RUNNER_SANA_WM_BIDIRECTIONAL.stage1_precision == "bf16" assert RUNNER_SANA_WM_BIDIRECTIONAL.refiner_precision == "bf16" + assert RUNNER_SANA_WM_BIDIRECTIONAL.quant_backend == "auto" + + +def test_auto_quant_backend_resolves_to_native_backend() -> None: + """Keep default quantization on the TE-free native backend.""" + assert _resolve_quant_backend("auto", ["fp4"]) == "native" + assert _resolve_quant_backend("auto", ["fp8", "fp4"]) == "native" + assert _resolve_quant_backend("auto", []) == "native" def test_bf16_precision_clears_upstream_quant_env( @@ -142,8 +483,8 @@ def test_fp8_precision_sets_upstream_quant_env() -> None: assert updates["SANA_WM_REFINER_QUANT"] == "fp8block" -def test_fp4_precision_sets_upstream_quant_env() -> None: - """Map FlashDreams fp4 fields to upstream SANA-WM NVFP4 env selectors.""" +def test_fp4_precision_sets_sana_quant_env() -> None: + """Map FlashDreams fp4 fields to SANA-WM NVFP4 env selectors.""" updates = _precision_env_updates( stage1_precision="fp4", refiner_precision="fp4", @@ -155,18 +496,30 @@ def test_fp4_precision_sets_upstream_quant_env() -> None: def test_quantized_precision_requires_cuda() -> None: - """Reject fp8/fp4 before importing upstream Sana on CPU-only devices.""" + """Reject fp8/fp4 before loading checkpoints on CPU-only devices.""" with pytest.raises(ValueError, match="requires a CUDA device"): _validate_precision_request( device=torch.device("cpu"), stage1_precision="fp8", refiner_precision="bf16", refiner_enabled=True, + quant_backend="upstream-te", ) +def test_bf16_precision_does_not_require_cuda_or_quant_backend() -> None: + """Keep BF16 as the TE-free compatibility path.""" + _validate_precision_request( + device=torch.device("cpu"), + stage1_precision="bf16", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="auto", + ) + + def test_fp8_precision_requires_hopper(monkeypatch: pytest.MonkeyPatch) -> None: - """Reject FP8 on pre-Hopper GPUs before checking Transformer Engine.""" + """Reject FP8 on pre-Hopper GPUs before checking quantization backend.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (8, 9)) @@ -176,6 +529,7 @@ def test_fp8_precision_requires_hopper(monkeypatch: pytest.MonkeyPatch) -> None: stage1_precision="fp8", refiner_precision="bf16", refiner_enabled=True, + quant_backend="upstream-te", ) @@ -191,11 +545,12 @@ def test_fp8_precision_requires_cuda_129(monkeypatch: pytest.MonkeyPatch) -> Non stage1_precision="fp8", refiner_precision="bf16", refiner_enabled=True, + quant_backend="upstream-te", ) def test_fp4_precision_requires_blackwell(monkeypatch: pytest.MonkeyPatch) -> None: - """Reject NVFP4 on Hopper-class GPUs before checking Transformer Engine.""" + """Reject NVFP4 on Hopper-class GPUs before checking quantization backend.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0)) @@ -205,9 +560,197 @@ def test_fp4_precision_requires_blackwell(monkeypatch: pytest.MonkeyPatch) -> No stage1_precision="fp4", refiner_precision="bf16", refiner_enabled=True, + quant_backend="upstream-te", ) +def test_torch_fp8_backend_skips_transformer_engine_cuda_129_gate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Allow native FP8 validation in the CUDA 12.8 env that blocks TE FP8.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch.version, "cuda", "12.8") + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp8", + ) + + +def test_auto_fp8_backend_uses_native_primitives( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Route default FP8 to the TE-free native backend.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch.version, "cuda", "12.8") + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="auto", + ) + + +def test_native_runner_reaches_normal_input_validation() -> None: + """Do not silently delegate the native runner to the upstream harness.""" + cfg = derive_config( + RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, + image_path=Path("missing.png"), + prompt="demo", + ) + runner = cfg.setup() + + with pytest.raises(FileNotFoundError, match="missing.png"): + runner.run() + + +def test_torch_fp8_backend_rejects_fp4(monkeypatch: pytest.MonkeyPatch) -> None: + """Do not silently route FP4 to a backend without an NVFP4 kernel.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + + with pytest.raises(ValueError, match="supports fp8 only"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp4", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp8", + ) + + +def test_torch_fp4_backend_skips_transformer_engine(monkeypatch: pytest.MonkeyPatch) -> None: + """Allow native FP4 validation without importing Transformer Engine.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float4_e2m1fn_x2", torch.uint8, raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp4", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp4", + ) + + +def test_torch_fp4_backend_rejects_fp8(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep precision-specific native backends explicit.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + + with pytest.raises(ValueError, match="supports fp4 only"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp4", + ) + + +def test_native_backend_allows_mixed_fp8_fp4_without_cuda_129( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Allow mixed native FP8/FP4 requests in the CUDA 12.8 env that blocks TE FP8.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch.version, "cuda", "12.8") + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float4_e2m1fn_x2", torch.uint8, raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="fp4", + refiner_enabled=True, + quant_backend="native", + ) + + +def test_torch_fp8_linear_replaces_matching_modules() -> None: + """Provide a TE-free replacement for eligible upstream Linear modules.""" + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("torch.float8_e4m3fn is required for native FP8 replacement") + + module = torch.nn.Sequential( + torch.nn.Linear(16, 32, bias=True), + torch.nn.Sequential(torch.nn.Linear(16, 32, bias=False)), + torch.nn.Linear(15, 32, bias=False), + ) + + converted, skipped = replace_linear_with_torch_fp8( + module, + recipe=None, + params_dtype=torch.bfloat16, + skip_patterns=(), + include_patterns=("^0$", "^1\\.0$", "^2$"), + ) + + assert converted == 2 + assert skipped == 1 + assert isinstance(module[0], TorchScaledMMFP8Linear) + assert isinstance(module[1][0], TorchScaledMMFP8Linear) + assert isinstance(module[2], torch.nn.Linear) + assert module[0].weight.shape == (32, 16) + assert module[0].bias is not None + + +def test_torch_fp4_linear_replaces_matching_modules( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Route eligible FP4 modules through the native replacement helper.""" + module = torch.nn.Sequential( + torch.nn.Linear(32, 64, bias=True), + torch.nn.Sequential(torch.nn.Linear(32, 64, bias=False)), + torch.nn.Linear(16, 64, bias=False), + ) + + @classmethod + def from_linear( + cls: type[TorchScaledMMFP4Linear], + source: torch.nn.Linear, + *, + out_dtype: torch.dtype, + ) -> TorchScaledMMFP4Linear: + del out_dtype + instance = cls.__new__(cls) + torch.nn.Module.__init__(instance) + instance.in_features = source.in_features + instance.out_features = source.out_features + return instance + + monkeypatch.setattr(TorchScaledMMFP4Linear, "from_linear", from_linear) + + converted, skipped = replace_linear_with_torch_fp4( + module, + recipe=None, + params_dtype=torch.bfloat16, + skip_patterns=(), + include_patterns=("^0$", "^1\\.0$", "^2$"), + ) + + assert converted == 2 + assert skipped == 1 + assert isinstance(module[0], TorchScaledMMFP4Linear) + assert isinstance(module[1][0], TorchScaledMMFP4Linear) + assert isinstance(module[2], torch.nn.Linear) + + def test_pyproject_entry_point_matches_runner_literal() -> None: """Keep the package entry point aligned with ``RUNNER_CONFIGS``.""" pyproject = tomllib.loads( @@ -216,5 +759,25 @@ def test_pyproject_entry_point_matches_runner_literal() -> None: entry_points = pyproject["project"]["entry-points"][ENTRY_POINT_GROUP] assert entry_points == { - "sana-wm-bidirectional": "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL" + "sana-wm-bidirectional": "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL", + "sana-wm-upstream-reference": ( + "sana_wm.config:RUNNER_SANA_WM_UPSTREAM_REFERENCE" + ), + "sana-wm-native-bidirectional": ( + "sana_wm.config:RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL" + ), } + + +def test_pyproject_excludes_upstream_diffusion_package() -> None: + """Do not install the vendored upstream SANA ``diffusion`` package.""" + pyproject = tomllib.loads( + Path("integrations/sana/pyproject.toml").read_text(encoding="utf-8") + ) + + packages = pyproject["tool"]["setuptools"]["packages"]["find"] + dependencies = set(pyproject["project"]["dependencies"]) + + assert packages["include"] == ["sana_wm*"] + assert "diffusers>=0.36" in dependencies + assert "transformers>=5.0,<6" in dependencies From 5c87bfddbce0b9c6b05d4e4121befcb53c4ab0e9 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 16 Jul 2026 18:36:20 -0700 Subject: [PATCH 07/36] Reduce native SANA-WM Stage-1 memory use Run classifier-free guidance as sequential negative and positive Stage-1 forwards instead of doubling the full model batch. This reduces peak activation memory for long native SANA-WM rollouts. Release large Stage-1 block temporaries earlier and make the VAE decode OOM retry reduce temporal tile size as well as spatial tile size. Add CPU coverage for sequential CFG kwargs and the lower-memory VAE retry configuration. Signed-off-by: Aidan Foster --- integrations/sana/sana_wm/native_diffusion.py | 72 ++++++++++++------- .../sana/sana_wm/native_transformer.py | 36 +++++++--- integrations/sana/sana_wm/stage1_model.py | 8 ++- integrations/sana/tests/test_smoke.py | 29 ++++++++ 4 files changed, 109 insertions(+), 36 deletions(-) diff --git a/integrations/sana/sana_wm/native_diffusion.py b/integrations/sana/sana_wm/native_diffusion.py index be719ffab..35270d4a0 100644 --- a/integrations/sana/sana_wm/native_diffusion.py +++ b/integrations/sana/sana_wm/native_diffusion.py @@ -134,14 +134,10 @@ def _sample_ltx_euler( condition_mask[:, :, int(frame_idx)] = 1 image_cond_noise_scale = max(image_cond_noise_scale, float(frame_weight)) - prompt_embeds = conditioning.condition - if do_cfg: - if conditioning.uncondition is None: - raise RuntimeError("CFG was requested without negative prompt embeds.") - prompt_embeds = torch.cat( - [conditioning.uncondition, conditioning.condition], - dim=0, - ) + condition_kwargs = _condition_model_kwargs(conditioning.model_kwargs) + uncondition_kwargs = _uncondition_model_kwargs(conditioning.model_kwargs) + if do_cfg and conditioning.uncondition is None: + raise RuntimeError("CFG was requested without negative prompt embeds.") init_latents = latents.clone() generator = torch.Generator(device=self.device).manual_seed(conditioning.seed) @@ -160,26 +156,36 @@ def _sample_ltx_euler( generator=generator, ) - condition_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 = timestep_scalar.expand(condition_mask_input.shape).float() - timestep = torch.min(timestep, (1 - condition_mask_input) * 1000.0) - - noise_pred = self.transformer.predict_flow( - noisy_latent=latent_model_input, - timestep=timestep[:, :1, :, 0, 0], - cache=cache, - input=prompt_embeds, - ) - + timestep = timestep_scalar.expand(condition_mask.shape).float() + timestep = torch.min(timestep, (1 - condition_mask) * 1000.0) if do_cfg: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + assert conditioning.uncondition is not None + noise_pred_uncond = self.transformer.predict_flow( + noisy_latent=latents, + timestep=timestep[:, :1, :, 0, 0], + cache=cache, + input=conditioning.uncondition, + model_kwargs=uncondition_kwargs, + ) + noise_pred_text = self.transformer.predict_flow( + noisy_latent=latents, + timestep=timestep[:, :1, :, 0, 0], + cache=cache, + input=conditioning.condition, + model_kwargs=condition_kwargs, + ) noise_pred = noise_pred_uncond + conditioning.cfg_scale * ( noise_pred_text - noise_pred_uncond ) - timestep = timestep.chunk(2)[0] + del noise_pred_uncond, noise_pred_text + else: + noise_pred = self.transformer.predict_flow( + noisy_latent=latents, + timestep=timestep[:, :1, :, 0, 0], + cache=cache, + input=conditioning.condition, + model_kwargs=condition_kwargs, + ) latents_dtype = latents.dtype latents_shape = latents.shape @@ -202,6 +208,24 @@ def _sample_ltx_euler( return latents.detach() +def _condition_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: + """Return model kwargs for the positive prompt branch.""" + return { + key: value + for key, value in model_kwargs.items() + if key != "negative_mask" + } + + +def _uncondition_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: + """Return model kwargs for the negative prompt branch.""" + kwargs = _condition_model_kwargs(model_kwargs) + negative_mask = model_kwargs.get("negative_mask") + if negative_mask is not None: + kwargs["mask"] = negative_mask + return kwargs + + def _add_noise_to_conditioning_latents( *, t: Tensor, diff --git a/integrations/sana/sana_wm/native_transformer.py b/integrations/sana/sana_wm/native_transformer.py index d8be6efb1..59ebf538f 100644 --- a/integrations/sana/sana_wm/native_transformer.py +++ b/integrations/sana/sana_wm/native_transformer.py @@ -175,6 +175,12 @@ class SanaWMNativeTransformerConfig(TransformerConfig): vae_oom_retry_tile_sample_stride_height: int = 64 """Smaller pixel height tile stride for one VAE decode OOM retry.""" + vae_oom_retry_tile_sample_min_num_frames: int = 16 + """Smaller temporal tile size for one VAE decode OOM retry.""" + + vae_oom_retry_tile_sample_stride_num_frames: int = 4 + """Smaller temporal tile stride for one VAE decode OOM retry.""" + class SanaWMNativeTransformer(Transformer[SanaWMNativeTransformerCache]): """Native FlashDreams adapter for the SANA-WM Stage-1 model call.""" @@ -225,16 +231,18 @@ def predict_flow( timestep: Tensor, cache: SanaWMNativeTransformerCache, input: Any = None, + model_kwargs: dict[str, object] | None = None, ) -> Tensor: """Execute one SANA-WM DiT flow prediction.""" conditioning = _require_conditioning(cache) prompt_embeds = cast(Tensor, input) if input is not None else conditioning.condition + kwargs = conditioning.model_kwargs if model_kwargs is None else model_kwargs self._ensure_model() return self.model( noisy_latent, timestep, prompt_embeds, - **conditioning.model_kwargs, + **kwargs, ) def finalize_kv_cache( @@ -303,15 +311,11 @@ def prepare_conditioning( self.device, dtype=weight_dtype, ) + model_kwargs_extra: dict[str, object] = {} if cfg_scale > 1.0: - mask_cfg = torch.cat([neg_mask, cond_mask], dim=0) - raymap_cfg = torch.cat([raymap, raymap], dim=0) - chunk_plucker_cfg = torch.cat([chunk_plucker, chunk_plucker], dim=0) + model_kwargs_extra["negative_mask"] = neg_mask uncondition = neg else: - mask_cfg = cond_mask - raymap_cfg = raymap - chunk_plucker_cfg = chunk_plucker uncondition = None latent_t = (num_frames - 1) // int(cfg.vae.vae_stride[0]) + 1 @@ -327,9 +331,10 @@ def prepare_conditioning( ), "condition_frame_info": {0: 0.0}, }, - "mask": mask_cfg, - "camera_conditions": raymap_cfg, - "chunk_plucker": chunk_plucker_cfg, + "mask": cond_mask, + "camera_conditions": raymap, + "chunk_plucker": chunk_plucker, + **model_kwargs_extra, } if chunk_index is not None: model_kwargs["chunk_index"] = chunk_index @@ -386,11 +391,14 @@ def decode_latents(self, latents: Tensor) -> np.ndarray: gc.collect() logger.warning( "[sana-vae] decode OOM; retrying with smaller tiles " - "width={} stride_width={} height={} stride_height={}", + "width={} stride_width={} height={} stride_height={} " + "frames={} stride_frames={}", self.config.vae_oom_retry_tile_sample_min_width, self.config.vae_oom_retry_tile_sample_stride_width, self.config.vae_oom_retry_tile_sample_min_height, self.config.vae_oom_retry_tile_sample_stride_height, + self.config.vae_oom_retry_tile_sample_min_num_frames, + self.config.vae_oom_retry_tile_sample_stride_num_frames, ) self._configure_vae_tiling( tile_sample_min_width=( @@ -405,6 +413,12 @@ def decode_latents(self, latents: Tensor) -> np.ndarray: tile_sample_stride_height=( self.config.vae_oom_retry_tile_sample_stride_height ), + tile_sample_min_num_frames=( + self.config.vae_oom_retry_tile_sample_min_num_frames + ), + tile_sample_stride_num_frames=( + self.config.vae_oom_retry_tile_sample_stride_num_frames + ), ) if torch.cuda.is_available(): torch.cuda.empty_cache() diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index 74f21cff2..10d0733ec 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -403,8 +403,10 @@ def forward( x_norm = self.norm1(x).reshape(batch, frames, -1, channels) attn_in = _modulate(x_norm, shift_msa, scale_msa).reshape(batch, tokens, channels) + del x_norm, shift_msa, scale_msa attn_out = self.attn(attn_in, **kwargs).reshape(batch, frames, -1, channels) x = x + (gate_msa * attn_out).reshape(batch, tokens, channels) + del attn_in, attn_out, gate_msa if plucker_emb is not None: x = x + self.plucker_proj(plucker_emb) @@ -413,9 +415,13 @@ def forward( x_norm = self.norm2(x).reshape(batch, frames, -1, channels) mlp_in = _modulate(x_norm, shift_mlp, scale_mlp).reshape(batch, tokens, channels) + del x_norm, shift_mlp, scale_mlp mlp_out = self.mlp(mlp_in, frames=frames, height=height, width=width) + del mlp_in mlp_out = mlp_out.reshape(batch, frames, -1, channels) - return x + (gate_mlp * mlp_out).reshape(batch, tokens, channels) + out = x + (gate_mlp * mlp_out).reshape(batch, tokens, channels) + del gate_mlp, mlp_out + return out class SanaWMStage1Model(nn.Module): diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 692795bad..282f12229 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -50,6 +50,10 @@ _temporary_environment, _validate_precision_request, ) +from sana_wm.native_diffusion import ( + _condition_model_kwargs, + _uncondition_model_kwargs, +) from sana_wm.native_diffusion import SanaWMDiffusionModelConfig from sana_wm.native_transformer import ( SanaWMNativeTransformerCache, @@ -377,6 +381,31 @@ def fake_decode(_samples: torch.Tensor) -> torch.Tensor: assert transformer.vae.tile_sample_stride_height == 64 assert transformer.vae.tile_sample_min_width == 128 assert transformer.vae.tile_sample_stride_width == 64 + assert transformer.vae.tile_sample_min_num_frames == 16 + assert transformer.vae.tile_sample_stride_num_frames == 4 + + +def test_cfg_model_kwargs_do_not_duplicate_camera_tensors() -> None: + """Run CFG as sequential branches instead of doubling Stage-1 activations.""" + camera = torch.zeros((1, 3, 3, 2, 2)) + cond_mask = torch.ones((1, 8)) + neg_mask = torch.zeros((1, 8)) + kwargs = { + "mask": cond_mask, + "negative_mask": neg_mask, + "camera_conditions": camera, + "chunk_plucker": torch.zeros((1, 6, 3, 2, 2)), + } + + cond_kwargs = _condition_model_kwargs(kwargs) + neg_kwargs = _uncondition_model_kwargs(kwargs) + + assert cond_kwargs["mask"] is cond_mask + assert neg_kwargs["mask"] is neg_mask + assert cond_kwargs["camera_conditions"] is camera + assert neg_kwargs["camera_conditions"] is camera + assert "negative_mask" not in cond_kwargs + assert "negative_mask" not in neg_kwargs def test_native_vae_tiling_avoids_degenerate_latent_tails() -> None: From c00b3c58f357172a3d3f240f7da2e0ce1e19bbf9 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 16 Jul 2026 18:43:19 -0700 Subject: [PATCH 08/36] Fix SANA-WM VAE retry temporal stride Keep retry temporal tiling compatible with LTX2's temporal compression ratio so diffusers does not derive a zero latent stride during OOM retry. Update the SANA smoke coverage for the clamped retry stride and sync the uv lock metadata for the integration dependencies declared by the native SANA package. Signed-off-by: Aidan Foster --- .../sana/sana_wm/native_transformer.py | 6 ++- integrations/sana/tests/test_smoke.py | 3 +- uv.lock | 40 +++++++++---------- 3 files changed, 27 insertions(+), 22 deletions(-) diff --git a/integrations/sana/sana_wm/native_transformer.py b/integrations/sana/sana_wm/native_transformer.py index 59ebf538f..2b9775f77 100644 --- a/integrations/sana/sana_wm/native_transformer.py +++ b/integrations/sana/sana_wm/native_transformer.py @@ -178,7 +178,7 @@ class SanaWMNativeTransformerConfig(TransformerConfig): vae_oom_retry_tile_sample_min_num_frames: int = 16 """Smaller temporal tile size for one VAE decode OOM retry.""" - vae_oom_retry_tile_sample_stride_num_frames: int = 4 + vae_oom_retry_tile_sample_stride_num_frames: int = 8 """Smaller temporal tile stride for one VAE decode OOM retry.""" @@ -666,6 +666,10 @@ def _configure_vae_tiling( tile_sample_stride_num_frames or self.config.vae_tile_sample_stride_num_frames ) + temporal_ratio = int(getattr(vae, "temporal_compression_ratio", 1)) + if temporal_ratio > 1: + min_frames = max(min_frames, temporal_ratio) + stride_frames = max(stride_frames, temporal_ratio) kwargs = { "tile_sample_min_height": min_height, "tile_sample_stride_height": stride_height, diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 282f12229..ae960fd71 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -350,6 +350,7 @@ def __init__(self) -> None: self.use_framewise_encoding = False self.use_framewise_decoding = False self.spatial_compression_ratio = 32 + self.temporal_compression_ratio = 8 def to(self, *_args: object, **_kwargs: object) -> "DummyVAE": return self @@ -382,7 +383,7 @@ def fake_decode(_samples: torch.Tensor) -> torch.Tensor: assert transformer.vae.tile_sample_min_width == 128 assert transformer.vae.tile_sample_stride_width == 64 assert transformer.vae.tile_sample_min_num_frames == 16 - assert transformer.vae.tile_sample_stride_num_frames == 4 + assert transformer.vae.tile_sample_stride_num_frames == 8 def test_cfg_model_kwargs_do_not_duplicate_camera_tensors() -> None: diff --git a/uv.lock b/uv.lock index 55413b69e..407fb184b 100644 --- a/uv.lock +++ b/uv.lock @@ -1573,7 +1573,16 @@ name = "flashdreams-sana-wm" version = "0.1.0" source = { editable = "integrations/sana" } dependencies = [ + { name = "diffusers" }, { name = "flashdreams" }, + { name = "imageio", extra = ["ffmpeg"] }, + { name = "pillow" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torchvision", version = "0.27.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torchvision", version = "0.27.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "transformers" }, ] [package.optional-dependencies] @@ -1583,8 +1592,15 @@ dev = [ [package.metadata] requires-dist = [ + { name = "diffusers", specifier = ">=0.36" }, { name = "flashdreams", editable = "flashdreams" }, + { name = "imageio", extras = ["ffmpeg"], specifier = ">=2.31" }, + { name = "pillow", specifier = ">=10" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "safetensors", specifier = ">=0.5" }, + { name = "torchvision", specifier = ">=0.26" }, + { name = "transformers", specifier = ">=5.0,<6" }, ] provides-extras = ["dev"] @@ -3221,26 +3237,10 @@ name = "numpy" version = "2.4.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.13.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.12.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.13.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.12.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", + "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version >= '3.12' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } wheels = [ From 7de18d69548dd7f98651c62fd49fb1cae6a23563 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 16 Jul 2026 18:48:36 -0700 Subject: [PATCH 09/36] Run native SANA VAE decode in inference mode Wrap native SANA conditioning and VAE decode in torch.inference_mode so full-length runs do not retain autograd activations during first-frame encode, prompt encode, or LTX2 VAE decode. Extend the VAE OOM retry smoke test to assert both decode attempts run under inference mode. Signed-off-by: Aidan Foster --- integrations/sana/sana_wm/native_transformer.py | 2 ++ integrations/sana/tests/test_smoke.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/integrations/sana/sana_wm/native_transformer.py b/integrations/sana/sana_wm/native_transformer.py index 2b9775f77..b847761dc 100644 --- a/integrations/sana/sana_wm/native_transformer.py +++ b/integrations/sana/sana_wm/native_transformer.py @@ -263,6 +263,7 @@ def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: """SANA-WM Stage-1 latents are already in model layout.""" return x + @torch.inference_mode() def prepare_conditioning( self, *, @@ -369,6 +370,7 @@ def initial_latents(self, conditioning: SanaWMStage1Conditioning) -> Tensor: latents[:, :, :1] = conditioning.first_latent return latents + @torch.inference_mode() def decode_latents(self, latents: Tensor) -> np.ndarray: """Decode SANA-WM VAE latents to ``uint8`` HWC video.""" self._ensure_vae() diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index ae960fd71..2b839d7d7 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -364,10 +364,12 @@ def enable_tiling(self, **kwargs: int) -> None: transformer.vae_dtype = torch.float32 monkeypatch.setattr(transformer, "_ensure_vae", lambda: None) calls = 0 + inference_modes: list[bool] = [] def fake_decode(_samples: torch.Tensor) -> torch.Tensor: nonlocal calls calls += 1 + inference_modes.append(torch.is_inference_mode_enabled()) if calls == 1: raise torch.OutOfMemoryError("test OOM") return torch.zeros((1, 3, 1, 2, 2), dtype=torch.float32) @@ -377,6 +379,7 @@ def fake_decode(_samples: torch.Tensor) -> torch.Tensor: video = transformer.decode_latents(torch.zeros((1, 4, 1, 1, 1))) assert calls == 2 + assert inference_modes == [True, True] assert video.shape == (1, 2, 2, 3) assert transformer.vae.tile_sample_min_height == 128 assert transformer.vae.tile_sample_stride_height == 64 From 46b6e660bac13ae862d4c4aa722ad4b614a59543 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 16 Jul 2026 19:52:32 -0700 Subject: [PATCH 10/36] Implement native SANA-WM refinement Replace the Stage-1 placeholder attention path with checkpoint-compatible GDN, softmax, camera UCPE, output-gate, and FFN behavior. Add a native LTX-2 refiner that loads diffusers transformer/connectors and Gemma text conditioning directly from the SANA-WM checkpoint, then wire the native runner to refine by default before VAE decode. Keep BF16/FP8/FP4 on FlashDreams/PyTorch native paths, add accelerate for low-memory refiner loading, update docs, and cover the new refiner construction plus latent packing in CPU tests. Validation: uv run --package flashdreams-sana-wm pytest integrations/sana/tests/test_smoke.py integrations/sana/tests/test_camera.py; uv run --package flashdreams-sana-wm flashdreams-run --no-instantiate sana-wm-bidirectional; no-instantiate checks for fp8, fp4, and --no-refiner True. --- integrations/sana/README.md | 40 +- integrations/sana/pyproject.toml | 1 + integrations/sana/sana_wm/config.py | 4 +- integrations/sana/sana_wm/native_refiner.py | 601 +++++++++++++ .../sana/sana_wm/native_transformer.py | 30 +- integrations/sana/sana_wm/runner.py | 45 +- integrations/sana/sana_wm/stage1_model.py | 791 +++++++++++++++++- integrations/sana/tests/test_smoke.py | 157 +++- uv.lock | 23 + 9 files changed, 1628 insertions(+), 64 deletions(-) create mode 100644 integrations/sana/sana_wm/native_refiner.py diff --git a/integrations/sana/README.md b/integrations/sana/README.md index aa0500e82..db4830004 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -13,19 +13,18 @@ The `sana-wm-bidirectional` runner is the native FlashDreams path. It uses FlashDreams config, runner, pipeline, diffusion-model, scheduler, transformer, camera-conditioning, VAE decode, and output-writing boundaries. The Stage-1 DiT module is implemented in this package and loads the public SANA-WM -checkpoint directly; it does not import or install a vendored upstream +checkpoint directly; it does not import or install the NVlabs/Sana `diffusion` package. Current native scope: | component | status | | --- | --- | -| Stage-1 BF16 | Native FlashDreams path, GPU smoke-tested. | -| Stage-1 FP8 | Native PyTorch `_scaled_mm` backend, GPU smoke-tested. | -| Stage-1 FP4 | Native Triton quantization plus PyTorch `_scaled_mm`, GPU smoke-tested. | +| Stage-1 BF16 | Native FlashDreams DiT execution. | +| Stage-1 FP8 | Native PyTorch `_scaled_mm` backend. | +| Stage-1 FP4 | Native Triton quantization plus PyTorch `_scaled_mm`. | | VAE decode | Native direct `diffusers` LTX2 VAE use with low-memory tiling. | -| LTX-2 refiner | Not native yet; the native runner requires `--no-refiner True`. | -| GDN attention parity | Not complete; smoke-tested, not quality-parity validated. | +| LTX-2 refiner | Native direct `diffusers` LTX-2 transformer and Gemma connector use. | The separate `sana-wm-upstream-reference` runner is documented only in `docs/upstream_reference.md`. Keep reference-harness details out of this @@ -35,7 +34,7 @@ native integration README. | slug | description | | --- | --- | -| `sana-wm-bidirectional` | Native FlashDreams SANA-WM Stage-1 runner. | +| `sana-wm-bidirectional` | Native FlashDreams SANA-WM Stage-1 + LTX-2 refiner runner. | | `sana-wm-native-bidirectional` | Alias for the native FlashDreams runner. | The FlashDreams package is named `sana_wm` rather than `sana` so it does not @@ -50,12 +49,6 @@ project's GPU runtime dependencies: uv sync --package flashdreams-sana-wm --extra dev ``` -For an existing environment, install the local packages directly: - -```bash -python -m pip install -e flashdreams -e integrations/sana -``` - The native runner does not require a local NVlabs/Sana checkout. The checkout is only useful as a source of demo assets such as `asset/sana_wm/demo_0.png`, `demo_0.txt`, `demo_0_pose.npy`, and @@ -63,18 +56,17 @@ is only useful as a source of demo assets such as ## Run -The native runner currently requires explicit intrinsics and no refiner: +The native runner requires explicit intrinsics: ```bash PYTORCH_ALLOC_CONF=expandable_segments:True \ -flashdreams-run sana-wm-bidirectional \ +uv run flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ - --output-dir outputs/sana_wm_native_bf16 \ - --no-refiner True + --output-dir outputs/sana_wm_native_bf16 ``` Expected output: @@ -83,6 +75,8 @@ Expected output: outputs/sana_wm_native_bf16/sana_wm_generated.mp4 ``` +For Stage-1-only diagnostics, add `--no-refiner True`. + ## FP8 and FP4 The runner defaults to BF16. Quantized Stage-1 inference is opt-in and uses @@ -99,30 +93,30 @@ FP8 smoke: ```bash PYTORCH_ALLOC_CONF=expandable_segments:True \ -flashdreams-run sana-wm-bidirectional \ +uv run flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ --output-dir outputs/sana_wm_native_fp8 \ - --no-refiner True \ - --stage1-precision fp8 + --stage1-precision fp8 \ + --refiner-precision fp8 ``` FP4 smoke: ```bash PYTORCH_ALLOC_CONF=expandable_segments:True \ -flashdreams-run sana-wm-bidirectional \ +uv run flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ --output-dir outputs/sana_wm_native_fp4 \ - --no-refiner True \ - --stage1-precision fp4 + --stage1-precision fp4 \ + --refiner-precision fp4 ``` ## Tests diff --git a/integrations/sana/pyproject.toml b/integrations/sana/pyproject.toml index df7ebf34b..780bf446b 100644 --- a/integrations/sana/pyproject.toml +++ b/integrations/sana/pyproject.toml @@ -24,6 +24,7 @@ description = "SANA-WM bidirectional world-model integration for flashdreams." readme = "README.md" requires-python = ">=3.10" dependencies = [ + "accelerate>=1.0", "diffusers>=0.36", "flashdreams", "imageio[ffmpeg]>=2.31", diff --git a/integrations/sana/sana_wm/config.py b/integrations/sana/sana_wm/config.py index faa9d8aa2..73734b634 100644 --- a/integrations/sana/sana_wm/config.py +++ b/integrations/sana/sana_wm/config.py @@ -73,11 +73,10 @@ def _reference_pipeline(name: str) -> StreamInferencePipelineConfig: runner_name=PIPELINE_SANA_WM_BIDIRECTIONAL.name, description=( "SANA-WM bidirectional I2V native FlashDreams runner " - "(Stage-1 DiT)." + "(Stage-1 DiT + LTX-2 refiner)." ), pipeline=PIPELINE_SANA_WM_BIDIRECTIONAL, execution_backend="native-flashdreams", - no_refiner=True, ) """SANA-WM native FlashDreams runner config.""" @@ -86,7 +85,6 @@ def _reference_pipeline(name: str) -> StreamInferencePipelineConfig: description="SANA-WM native FlashDreams pipeline alias.", pipeline=PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL, execution_backend="native-flashdreams", - no_refiner=True, ) """Backward-compatible native SANA-WM FlashDreams runner alias.""" diff --git a/integrations/sana/sana_wm/native_refiner.py b/integrations/sana/sana_wm/native_refiner.py new file mode 100644 index 000000000..084940608 --- /dev/null +++ b/integrations/sana/sana_wm/native_refiner.py @@ -0,0 +1,601 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Native LTX-2 refiner used by the SANA-WM integration.""" + +from __future__ import annotations + +import gc +from pathlib import Path +from typing import Literal + +import torch +import torch.nn as nn +import torch.nn.functional as F +from loguru import logger +from torch import Tensor + +from sana_wm.native_quant import ( + TorchScaledMMFP4Recipe, + TorchScaledMMFP8Recipe, + replace_linear_with_native_quant, +) + +Precision = Literal["bf16", "fp8", "fp4"] +QuantBackend = Literal["auto", "upstream-te", "torch-fp8", "torch-fp4", "native"] + +_REFINER_QUANT_SKIP_DEFAULTS = ( + r"^proj_in$", + r"^proj_out$", + r"(^|\.)audio_", + r"audio_to_video", + r"video_to_audio", + r"av_cross_attn", + r"caption_projection", + r"time_embed", +) + + +class SanaWMLTX2Refiner(nn.Module): + """Run the SANA-WM LTX-2 latent refiner without importing a Sana checkout.""" + + def __init__( + self, + *, + refiner_root: str | Path, + gemma_root: str | Path, + dtype: torch.dtype, + device: torch.device | str, + precision: Precision = "bf16", + quant_backend: QuantBackend = "native", + text_max_sequence_length: int = 1024, + ) -> None: + super().__init__() + self.refiner_root = Path(refiner_root) + self.gemma_root = Path(gemma_root) + self.dtype = dtype + self.device = torch.device(device) + self.precision = precision + self.quant_backend = quant_backend + self.text_max_sequence_length = int(text_max_sequence_length) + self._quantized = False + self.transformer, self.connectors = self._load_diffusers_components() + + def refine_latents( + self, + sana_latent: Tensor, + prompt: str, + *, + fps: float, + sink_size: int = 1, + seed: int = 42, + progress: bool = True, + block_size: int | None = None, + kv_max_frames: int = 11, + sigmas: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0), + ) -> Tensor: + """Refine Stage-1 VAE latents with the sink-bidirectional LTX-2 path.""" + del kv_max_frames + if block_size is not None: + logger.warning( + "[refiner] --refiner-block-size={} requested; running the " + "sink-bidirectional LTX-2 path used by the bidirectional model.", + block_size, + ) + if sana_latent.shape[2] <= sink_size: + raise ValueError( + f"Stage-1 latent has {sana_latent.shape[2]} frames but " + f"sink_size={sink_size}." + ) + + prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt) + self.transformer.to(self.device) + self.transformer.eval() + self._prepare_quantization() + + z = sana_latent.to(device=self.device, dtype=self.dtype) + sigmas_t = torch.tensor(sigmas, dtype=torch.float32, device=self.device) + start_sigma = float(sigmas_t[0]) + sink = z[:, :, :sink_size].contiguous() + current = z[:, :, sink_size:].contiguous() + generator = torch.Generator(device=self.device).manual_seed(int(seed)) + eps = torch.randn( + current.shape, + generator=generator, + device=self.device, + dtype=self.dtype, + ) + noisy = (1.0 - start_sigma) * current + start_sigma * eps + + iterator = range(len(sigmas_t) - 1) + if progress: + from tqdm.auto import tqdm + + iterator = tqdm(iterator, desc="refiner", unit="step") + + for step_index in iterator: + sigma = sigmas_t[step_index] + denoised = self._predict_current_x0( + sink=sink, + noisy_current=noisy, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + sigma=sigma, + fps=fps, + ) + noisy_tokens = _pack_latents( + noisy, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + velocity = (noisy_tokens.float() - denoised.float()) / sigma.float() + next_tokens = ( + noisy_tokens.float() + + velocity * (sigmas_t[step_index + 1] - sigma).float() + ) + noisy = _unpack_latents( + next_tokens.to(self.dtype), + num_frames=noisy.shape[2], + height=noisy.shape[3], + width=noisy.shape[4], + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + + return torch.cat([sink, noisy], dim=2) + + def _load_diffusers_components(self) -> tuple[nn.Module, nn.Module]: + from diffusers.models.transformers.transformer_ltx2 import ( + LTX2VideoTransformer3DModel, + ) + from diffusers.pipelines.ltx2 import LTX2TextConnectors + + transformer = LTX2VideoTransformer3DModel.from_pretrained( + self.refiner_root, + subfolder="transformer", + torch_dtype=self.dtype, + ).eval() + connectors = LTX2TextConnectors.from_pretrained( + self.refiner_root, + subfolder="connectors", + torch_dtype=self.dtype, + ).eval() + return transformer, connectors + + def _prepare_quantization(self) -> None: + if self._quantized or self.precision == "bf16": + return + if self.quant_backend == "upstream-te": + raise ValueError( + "Native SANA-WM refiner execution does not use Transformer Engine; " + "select --quant-backend native, torch-fp8, or torch-fp4." + ) + recipe = ( + TorchScaledMMFP8Recipe() + if self.precision == "fp8" + else TorchScaledMMFP4Recipe() + ) + converted, skipped = replace_linear_with_native_quant( + self.transformer, + recipe=recipe, + params_dtype=self.dtype, + skip_patterns=_REFINER_QUANT_SKIP_DEFAULTS, + ) + if converted <= 0: + raise RuntimeError( + f"SANA-WM native refiner {self.precision} converted no Linear " + f"layers; skipped={skipped}." + ) + self._quantized = True + + @torch.inference_mode() + def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: + from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + + tokenizer = AutoTokenizer.from_pretrained(self.gemma_root) + 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(self.device) + attention_mask = text_inputs.attention_mask.to(self.device) + + text_encoder = Gemma3ForConditionalGeneration.from_pretrained( + self.gemma_root, + torch_dtype=self.dtype, + low_cpu_mem_usage=True, + ).eval() + text_encoder.to(self.device) + text_backbone = getattr(text_encoder, "model", 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) + del text_encoder, text_backbone, outputs + _empty_cuda_cache() + + prompt_embeds = _pack_text_embeds( + hidden_states, + sequence_lengths, + device=self.device, + padding_side=tokenizer.padding_side, + ).to(dtype=self.dtype) + del hidden_states + _empty_cuda_cache() + + self.connectors.to(self.device) + connector_prompt_embeds, _, connector_attention_mask = self.connectors( + prompt_embeds, + attention_mask, + ) + self.connectors.to("cpu") + del prompt_embeds, attention_mask + _empty_cuda_cache() + return ( + connector_prompt_embeds.to(device=self.device, dtype=self.dtype), + connector_attention_mask.to(device=self.device), + ) + + def _predict_current_x0( + self, + *, + sink: Tensor, + noisy_current: Tensor, + prompt_embeds: Tensor, + prompt_attention_mask: Tensor, + sigma: Tensor, + fps: float, + ) -> Tensor: + full_latent = torch.cat([sink, noisy_current], dim=2) + batch_size, _, num_frames, height, width = full_latent.shape + latent_tokens = _pack_latents( + full_latent, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + n_context_tokens = _pack_latents( + sink, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ).shape[1] + + raw_timestep = torch.zeros( + batch_size, + latent_tokens.shape[1], + 1, + dtype=torch.float32, + device=self.device, + ) + raw_timestep[:, n_context_tokens:, 0] = sigma.float() + model_timestep = ( + raw_timestep.squeeze(-1) + * float(self.transformer.config.timestep_scale_multiplier) + ) + + velocity = self._forward_video_only( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=model_timestep, + encoder_attention_mask=prompt_attention_mask, + num_frames=num_frames, + height=height, + width=width, + fps=fps, + n_context_tokens=n_context_tokens, + ) + denoised = latent_tokens.float() - velocity.float() * raw_timestep + return denoised[:, n_context_tokens:, :].to(self.dtype) + + def _forward_video_only( + self, + *, + hidden_states: Tensor, + encoder_hidden_states: Tensor, + timestep: Tensor, + encoder_attention_mask: Tensor | None, + num_frames: int, + height: int, + width: int, + fps: float, + n_context_tokens: int, + ) -> Tensor: + transformer = self.transformer + batch_size = hidden_states.size(0) + encoder_attention_mask = _prepare_encoder_attention_mask( + encoder_attention_mask, + hidden_states.dtype, + ) + video_coords = transformer.rope.prepare_video_coords( + batch_size, + num_frames, + height, + width, + hidden_states.device, + fps=fps, + ) + video_rotary_emb = transformer.rope(video_coords, device=hidden_states.device) + + 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) + + +def _forward_video_block( + *, + block: nn.Module, + hidden_states: Tensor, + encoder_hidden_states: Tensor, + temb: Tensor, + video_rotary_emb: tuple[Tensor, Tensor], + encoder_attention_mask: Tensor | None, + n_context_tokens: int, +) -> 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 + return hidden_states + block.ff(norm_hidden_states) * gate_mlp + + +def _streaming_self_attention( + *, + attn: nn.Module, + hidden_states: Tensor, + query_rotary_emb: tuple[Tensor, Tensor], + n_context_tokens: int, +) -> Tensor: + from diffusers.models.transformers.transformer_ltx2 import ( + 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) + + 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}") + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if n_context_tokens <= 0 or n_context_tokens >= query.shape[1]: + hidden_states = _refiner_attention(query, key, value) + else: + context_hidden_states = _refiner_attention( + query[:, :n_context_tokens], + key[:, :n_context_tokens], + value[:, :n_context_tokens], + ) + current_hidden_states = _refiner_attention( + query[:, n_context_tokens:], + key, + value, + ) + 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)).unsqueeze(-1) + hidden_states = hidden_states * gates + hidden_states = hidden_states.flatten(2, 3) + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +def _refiner_attention(query: Tensor, key: Tensor, value: Tensor) -> Tensor: + hidden_states = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + return hidden_states.transpose(1, 2) + + +def _pack_text_embeds( + text_hidden_states: Tensor, + sequence_lengths: Tensor, + *, + device: torch.device | str, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, +) -> 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!r}." + ) + mask = mask[:, :, None, None] + + masked = text_hidden_states.masked_fill(~mask, 0.0) + denom = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + 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 = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized = normalized * scale_factor + normalized = normalized.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, normalized.shape[-1]) + normalized = normalized.masked_fill(~mask_flat, 0.0) + return normalized.to(dtype=original_dtype) + + +def _pack_latents(latents: Tensor, patch_size: int = 1, patch_size_t: int = 1) -> Tensor: + batch_size, _, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7) + return latents.flatten(4, 7).flatten(1, 3) + + +def _unpack_latents( + latents: Tensor, + *, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, +) -> Tensor: + batch_size = latents.size(0) + latents = latents.reshape( + batch_size, + num_frames, + height, + width, + -1, + patch_size_t, + patch_size, + patch_size, + ) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7) + return latents.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +def _prepare_encoder_attention_mask(mask: Tensor | None, dtype: torch.dtype) -> Tensor | None: + if mask is None: + return None + if mask.ndim != 2: + return mask + if bool(torch.all(mask)): + return None + return ((1 - mask.to(dtype)) * -10000.0).unsqueeze(1) + + +def _empty_cuda_cache() -> None: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() diff --git a/integrations/sana/sana_wm/native_transformer.py b/integrations/sana/sana_wm/native_transformer.py index b847761dc..bebf59e7f 100644 --- a/integrations/sana/sana_wm/native_transformer.py +++ b/integrations/sana/sana_wm/native_transformer.py @@ -499,6 +499,10 @@ def release_stage1_runtime( logger.info("[stage1] released Stage-1 runtime before decode/refine") gc.collect() + def release_refiner_runtime(self) -> None: + """Release native refiner tensors before VAE decode.""" + self._release_refiner() + def refine_latents( self, *, @@ -510,7 +514,7 @@ def refine_latents( block_size: int | None, kv_max_frames: int, ) -> Tensor: - """Run the native LTX-2 refiner once that path is implemented.""" + """Run the native LTX-2 refiner.""" self._ensure_refiner() sigmas = torch.tensor( self._stage2_sigmas(), @@ -611,11 +615,27 @@ def _ensure_model(self) -> None: self._prepare_stage1_quant() def _ensure_refiner(self) -> None: - raise RuntimeError( - "Native SANA-WM refiner execution is not implemented yet. " - "Use --no-refiner True for native Stage-1 GPU tests, or " - "sana-wm-upstream-reference for the temporary upstream refiner harness." + if self._refiner_built: + return + from sana_wm.native_refiner import SanaWMLTX2Refiner + + compute_dtype = ( + torch.bfloat16 + if self.config.refiner_precision in {"fp8", "fp4"} + else _get_weight_dtype(self.config.refiner_precision) + ) + quant_backend: QuantBackend = ( + "native" if self.config.quant_backend == "auto" else self.config.quant_backend + ) + self.refiner = SanaWMLTX2Refiner( + refiner_root=resolve_hf_path(self.config.refiner_root), + gemma_root=resolve_hf_path(self.config.refiner_gemma_root), + dtype=compute_dtype, + device=self.device, + precision=self.config.refiner_precision, + quant_backend=quant_backend, ) + self._refiner_built = True def _release_refiner(self) -> None: if not self._refiner_built: diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index 09795c905..f0266c667 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -305,12 +305,6 @@ def _run_native(self) -> None: cfg = self.config if cfg.image_path is None: raise ValueError("SanaWMRunner requires --image-path.") - if not cfg.no_refiner: - raise ValueError( - "Native SANA-WM currently supports Stage-1 execution only. " - "Use --no-refiner True, or use sana-wm-upstream-reference for " - "the temporary upstream refiner harness." - ) if not cfg.no_action_overlay: raise ValueError( "Native SANA-WM does not include the upstream action overlay. " @@ -360,9 +354,8 @@ def _run_native(self) -> None: sampling_algo = self._native_sampling_algo() if sampling_algo != "flow_euler_ltx": raise ValueError( - "Native SANA-WM currently supports flow_euler_ltx only; " - f"got {sampling_algo!r}. Use sana-wm-upstream-reference " - "for chunk/self-forcing samplers." + "Native SANA-WM requires flow_euler_ltx for the " + f"bidirectional runner; got {sampling_algo!r}." ) camera = prepare_camera( c2w, @@ -389,10 +382,38 @@ def _run_native(self) -> None: transformer.release_stage1_runtime(cache) del conditioning, cache - video_hwc = transformer.decode_latents(sana_latent) + stage1_video_hwc: np.ndarray | None = None + output_latent = sana_latent + if not cfg.no_refiner: + output_latent = transformer.refine_latents( + latents=sana_latent, + prompt=prompt, + fps=cfg.fps, + sink_size=cfg.sink_size, + seed=cfg.refiner_seed, + block_size=cfg.refiner_block_size, + kv_max_frames=cfg.refiner_kv_max_frames, + ) + transformer.release_refiner_runtime() + elif cfg.save_stage1: + logger.info( + "Native SANA-WM is already running without the refiner; " + "--save-stage1 does not create an extra output." + ) + + video_hwc = transformer.decode_latents(output_latent) + if self.is_rank_zero and cfg.save_stage1 and not cfg.no_refiner: + stage1_video_hwc = transformer.decode_latents(sana_latent) if not self.is_rank_zero: return _write_video(cfg.output_dir, cfg.name, video_hwc, cfg.fps) + if stage1_video_hwc is not None: + _write_video( + cfg.output_dir, + f"{cfg.name}_stage1", + stage1_video_hwc, + cfg.fps, + ) def _run_upstream_reference(self) -> None: """Run the explicit upstream SANA-WM reference harness.""" @@ -829,7 +850,7 @@ def _validate_torch_fp8_backend(precisions: list[Precision]) -> None: unsupported = sorted({precision for precision in precisions if precision != "fp8"}) if unsupported: raise ValueError( - "SANA-WM --quant-backend torch-fp8 currently supports fp8 only; " + "SANA-WM --quant-backend torch-fp8 accepts fp8 only; " f"unsupported requested precision(s): {', '.join(unsupported)}. " "Use --quant-backend torch-fp4 or --quant-backend native for fp4." ) @@ -841,7 +862,7 @@ def _validate_torch_fp4_backend(precisions: list[Precision]) -> None: unsupported = sorted({precision for precision in precisions if precision != "fp4"}) if unsupported: raise ValueError( - "SANA-WM --quant-backend torch-fp4 currently supports fp4 only; " + "SANA-WM --quant-backend torch-fp4 accepts fp4 only; " f"unsupported requested precision(s): {', '.join(unsupported)}. " "Use --quant-backend torch-fp8 or --quant-backend native for fp8." ) diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index 10d0733ec..f657c0b0e 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -23,6 +23,7 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass import math @@ -78,6 +79,7 @@ def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.empty(hidden_size)) self.eps = eps + nn.init.ones_(self.weight) def forward(self, x: Tensor) -> Tensor: """Apply RMS normalization using the stored scale parameter.""" @@ -236,12 +238,13 @@ def __init__(self, spec: SanaWMStage1Spec) -> None: padding=(spec.temporal_kernel_size // 2, 0), bias=False, ) + nn.init.zeros_(self.t_conv.weight) def forward(self, x: Tensor, *, frames: int, height: int, width: int) -> Tensor: """Run spatial GLU plus temporal aggregation.""" batch, tokens, channels = x.shape x_2d = x.reshape(batch * frames, height, width, channels).permute(0, 3, 1, 2) - x_2d = self.inverted_conv(x_2d) + x_2d = F.silu(self.inverted_conv(x_2d)) x_2d = self.depth_conv(x_2d) value, gate = x_2d.chunk(2, dim=1) x_2d = self.point_conv(value * F.silu(gate)) @@ -256,6 +259,11 @@ class Stage1SelfAttention(nn.Module): def __init__(self, spec: SanaWMStage1Spec, *, use_gdn_convs: bool) -> None: super().__init__() + self.heads = spec.num_heads + self.dim = spec.head_dim + self.eps = 1e-6 + self.use_gdn_convs = use_gdn_convs + self.patch_size = (1, 1, 1) self.A_log = nn.Parameter(torch.empty(spec.num_heads)) self.beta_proj = nn.Linear(spec.hidden_size, spec.num_heads) self.dt_bias = nn.Parameter(torch.empty(spec.num_heads)) @@ -287,22 +295,242 @@ def __init__(self, spec: SanaWMStage1Spec, *, use_gdn_convs: bool) -> None: groups=spec.hidden_size, bias=False, ) + nn.init.zeros_(self.A_log) + nn.init.constant_(self.dt_bias, -5.0) + nn.init.zeros_(self.recall_gate) - def forward(self, x: Tensor, *_args: object, **_kwargs: object) -> Tensor: - """Run a memory-bounded native self-attention approximation. + def forward( + self, + x: Tensor, + *_args: object, + HW: tuple[int, int, int] | None = None, + rotary_emb: Tensor | None = None, + camera_conditions: Tensor | None = None, + apply_output_gate: bool = True, + **kwargs: object, + ) -> Tensor: + """Run native SANA-WM self/camera attention.""" + if HW is None: + raise ValueError("SANA-WM Stage-1 attention requires HW=(T, H, W).") + batch, tokens, channels = x.shape + if channels != self.heads * self.dim: + raise ValueError( + f"channels={channels} != heads*dim={self.heads * self.dim}" + ) - The public checkpoint's GDN blocks are linear-recurrent attention - blocks. Until the native GDN scan is implemented, this path uses the - checkpoint QKV/value/projection tensors without materializing an - all-token softmax attention matrix. - """ + precomputed_gates = self._compute_frame_gates(x, HW) + if self.use_gdn_convs: + main_raw = self._forward_gdn_main( + x, + HW=HW, + rotary_emb=rotary_emb, + precomputed_gates=precomputed_gates, + ) + else: + main_raw = self._forward_softmax_main( + x, + HW=HW, + rotary_emb=rotary_emb, + ) + + cam_contrib: Tensor | int = 0 + if camera_conditions is not None: + if self.use_gdn_convs: + cam_raw = self._forward_gdn_camera( + x, + HW=HW, + rotary_emb=rotary_emb, + camera_conditions=camera_conditions, + precomputed_gates=precomputed_gates, + ) + else: + cam_raw = self._forward_softmax_camera( + x, + HW=HW, + rotary_emb=rotary_emb, + camera_conditions=camera_conditions, + ) + cam_contrib = self.out_proj_cam(cam_raw) + + combined = main_raw + cam_contrib + if apply_output_gate: + combined = _apply_output_gate(combined, x, self.output_gate) + combined = self.proj(combined.to(dtype=self.proj.weight.dtype)) + del kwargs + return combined + + def _forward_gdn_main( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + precomputed_gates: tuple[Tensor, Tensor], + ) -> Tensor: batch, tokens, channels = x.shape - qkv = self.qkv(x).view(batch, tokens, 3, channels) - q, _k, v = qkv.unbind(dim=2) - q = self.q_norm(q) - v = v * torch.sigmoid(q) - out = F.silu(self.output_gate(x).float()).to(dtype=v.dtype) * v - return self.proj(out.to(dtype=self.proj.weight.dtype)) + qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim) + if hasattr(self, "conv_k"): + k_raw = qkv[:, :, 1].reshape(batch, tokens, channels) + k_conv = _apply_bidirectional_temporal_conv(k_raw, self.conv_k, HW) + qkv[:, :, 1] = k_conv.reshape(batch, tokens, self.heads, self.dim) + + q = self.q_norm(qkv[:, :, 0].reshape(batch, tokens, channels)).reshape( + batch, + tokens, + self.heads, + self.dim, + ) + k = self.k_norm(qkv[:, :, 1].reshape(batch, tokens, channels)).reshape( + batch, + tokens, + self.heads, + self.dim, + ) + v = qkv[:, :, 2] + q = F.relu(q.float()) + k = F.relu(k.float()) * _gdn_key_scale(self.dim, HW) + v = v.float() + q_rot = _apply_complex_rope(q, rotary_emb) + k_rot = _apply_complex_rope(k, rotary_emb) + beta, decay = precomputed_gates + out = _bidirectional_gdn_scan( + q=q, + k=k, + q_rot=q_rot, + k_rot=k_rot, + v=v, + beta=beta, + decay=decay, + HW=HW, + eps=self.eps, + ) + return out.reshape(batch, tokens, channels).to(dtype=x.dtype) + + def _forward_softmax_main( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + ) -> Tensor: + del HW + batch, tokens, channels = x.shape + qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim) + q, k, v = qkv.unbind(dim=2) + q = self.q_norm(q.reshape(batch, tokens, channels)).reshape( + batch, + tokens, + self.heads, + self.dim, + ) + k = self.k_norm(k.reshape(batch, tokens, channels)).reshape( + batch, + tokens, + self.heads, + self.dim, + ) + q = _apply_complex_rope(q, rotary_emb) + k = _apply_complex_rope(k, rotary_emb) + out = F.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + dropout_p=0.0, + is_causal=False, + ) + return out.transpose(1, 2).reshape(batch, tokens, channels).to(dtype=x.dtype) + + def _forward_gdn_camera( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + camera_conditions: Tensor, + precomputed_gates: tuple[Tensor, Tensor], + ) -> Tensor: + q_cam, k_cam, v_cam = _camera_qkv(self, x, HW) + q_trans, k_trans, v_trans, inflation_sq, output_projector = _prepare_ucpe_qkv( + q_cam, + k_cam, + v_cam, + camera_conditions=camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + q_norm_weight=self.q_norm_cam.weight, + k_norm_weight=self.k_norm_cam.weight, + norm_eps=self.q_norm_cam.eps, + ) + beta, decay = precomputed_gates + frame_inflation = inflation_sq.reshape( + x.shape[0], + self.heads, + HW[0], + HW[1] * HW[2], + ).mean(dim=-1) + beta = beta / frame_inflation.unsqueeze(-1).clamp_min(1.0) + out = _bidirectional_numerator_scan( + q=q_trans, + k=k_trans, + v=v_trans, + beta=beta, + decay=decay, + HW=HW, + ) + out = output_projector(out) + return out.reshape(x.shape[0], x.shape[1], -1).to(dtype=x.dtype) + + def _forward_softmax_camera( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + camera_conditions: Tensor, + ) -> Tensor: + q_cam, k_cam, v_cam = _camera_qkv(self, x, HW) + q_trans, k_trans, v_trans, output_projector = _prepare_ucpe_qkv_softmax( + q_cam, + k_cam, + v_cam, + camera_conditions=camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + q_norm_weight=self.q_norm_cam.weight, + k_norm_weight=self.k_norm_cam.weight, + norm_eps=self.q_norm_cam.eps, + ) + out = F.scaled_dot_product_attention( + q_trans.transpose(1, 2), + k_trans.transpose(1, 2), + v_trans.transpose(1, 2), + dropout_p=0.0, + is_causal=False, + ) + return output_projector(out.transpose(1, 2)).reshape( + x.shape[0], + x.shape[1], + -1, + ).to(dtype=x.dtype) + + def _compute_frame_gates( + self, + x: Tensor, + HW: tuple[int, int, int], + ) -> tuple[Tensor, Tensor]: + batch, tokens, channels = x.shape + frames, height, width = HW + spatial = height * width + if tokens != frames * spatial: + raise ValueError(f"tokens={tokens} != T*H*W={frames * spatial}") + beta = self.beta_proj(x).sigmoid() + beta = beta.reshape(batch, frames, spatial, self.heads).permute(0, 3, 1, 2) + x_frame = x.reshape(batch, frames, spatial, channels).mean(dim=2) + gate = self.gate_proj(x_frame).float() + dt = self.dt_bias.float().view(1, 1, -1) + a_val = self.A_log.float().exp().view(1, 1, -1) + decay = (-a_val * F.softplus(gate + dt)).exp().transpose(1, 2) + return beta.float(), decay.float() class Stage1CrossAttention(nn.Module): @@ -326,7 +554,7 @@ def forward( mask: Tensor | None = None, **_kwargs: object, ) -> Tensor: - """Cross-attention execution is implemented in a later native pass.""" + """Run text cross-attention with checkpoint-compatible Q/K norms.""" batch, tokens, channels = x.shape q = self.q_norm(self.q_linear(x)).view( batch, @@ -464,6 +692,18 @@ def forward( plucker_emb = self.plucker_embedder(chunk_plucker) plucker_emb = plucker_emb.permute(0, 2, 3, 4, 1).reshape_as(x) + rotary_emb = _wan_rope_complex( + self.spec.head_dim, + frames, + height, + width, + x.device, + ) + camera_conditions = kwargs.get("camera_conditions") + block_kwargs = { + key: value for key, value in kwargs.items() if key != "camera_conditions" + } + y = self.y_embedder(y) y = self.attention_y_norm(y) if mask is not None and mask.ndim > 2: @@ -483,7 +723,14 @@ def forward( width=width, mask=mask, plucker_emb=plucker_emb, - **kwargs, + HW=(frames, height, width), + rotary_emb=rotary_emb, + camera_conditions=( + camera_conditions.to(device=x.device, dtype=x.dtype) + if isinstance(camera_conditions, Tensor) + else None + ), + **block_kwargs, ) x = self.final_layer(x, timestep_embed, frames=frames) @@ -508,6 +755,518 @@ def _timestep_embedding(t: Tensor, dim: int, max_period: int = 10000) -> Tensor: return embedding +def _gdn_key_scale(head_dim: int, HW: tuple[int, int, int]) -> float: + return (head_dim**-0.5) * ((HW[1] * HW[2]) ** -0.5) + + +def _apply_output_gate(out: Tensor, gate_x: Tensor, gate: nn.Linear) -> Tensor: + gate_values = F.silu(F.linear(gate_x, gate.weight, gate.bias).float()) + return out * gate_values.to(dtype=out.dtype) + + +def _apply_bidirectional_temporal_conv( + x: Tensor, + conv: nn.Conv1d, + HW: tuple[int, int, int], +) -> Tensor: + batch, tokens, channels = x.shape + frames, height, width = HW + spatial = height * width + if tokens != frames * spatial: + raise ValueError(f"tokens={tokens} != T*H*W={frames * spatial}") + temporal = ( + x.reshape(batch, frames, spatial, channels) + .permute(0, 2, 3, 1) + .reshape(batch * spatial, channels, frames) + ) + y_fwd = _causal_depthwise_conv(temporal, conv) + y_bwd = _causal_depthwise_conv(temporal.flip(-1), conv).flip(-1) + center = conv.weight[:, 0, -1].view(1, channels, 1) + out = y_fwd + y_bwd - temporal * center + return ( + out.reshape(batch, spatial, channels, frames) + .permute(0, 3, 1, 2) + .reshape(batch, tokens, channels) + .to(dtype=x.dtype) + ) + + +def _causal_depthwise_conv(x: Tensor, conv: nn.Conv1d) -> Tensor: + kernel = int(conv.kernel_size[0]) + padded = F.pad(x, (kernel - 1, 0)) + return F.conv1d( + padded, + conv.weight.to(dtype=x.dtype), + bias=None, + stride=1, + padding=0, + dilation=1, + groups=conv.groups, + ) + + +def _wan_rope_complex( + head_dim: int, + frames: int, + height: int, + width: int, + device: torch.device, +) -> Tensor: + t_size = head_dim // 2 - 2 * (head_dim // 6) + h_size = head_dim // 6 + w_size = head_dim // 6 + freqs_t = _axis_rope_complex(frames, t_size, device) + freqs_h = _axis_rope_complex(height, h_size, device) + freqs_w = _axis_rope_complex(width, w_size, device) + expanded_t = freqs_t[:, None, None, :].expand(frames, height, width, t_size) + expanded_h = freqs_h[None, :, None, :].expand(frames, height, width, h_size) + expanded_w = freqs_w[None, None, :, :].expand(frames, height, width, w_size) + freqs = torch.cat([expanded_t, expanded_h, expanded_w], dim=-1) + return freqs.reshape(1, 1, frames * height * width, -1) + + +def _axis_rope_complex(length: int, complex_dims: int, device: torch.device) -> Tensor: + if complex_dims == 0: + return torch.empty(length, 0, dtype=torch.complex64, device=device) + dim = complex_dims * 2 + exponent = torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim + freqs = 1.0 / (10000.0**exponent) + positions = torch.arange(length, dtype=torch.float32, device=device) + angles = positions[:, None] * freqs[None] + return torch.polar(torch.ones_like(angles), angles) + + +def _apply_complex_rope(x: Tensor, rotary_emb: Tensor | None) -> Tensor: + if rotary_emb is None: + return x + batch, tokens, heads, dim = x.shape + freqs = rotary_emb.squeeze(0).squeeze(0) + if freqs.shape[0] != tokens or freqs.shape[1] != dim // 2: + raise ValueError( + f"RoPE shape {tuple(freqs.shape)} is incompatible with {(tokens, dim)}." + ) + x_float = x.float() + x_complex = torch.view_as_complex( + x_float.reshape(batch, tokens, heads, dim // 2, 2) + ) + rotated = torch.view_as_real(x_complex * freqs[None, :, None, :]).flatten(-2) + return rotated.to(dtype=x.dtype) + + +def _slice_rope_for_camera(rotary_emb: Tensor | None, head_dim: int) -> Tensor | None: + if rotary_emb is None: + return None + 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, + ] + return torch.cat([t_part, h_part, w_part], dim=-1) + + +def _bidirectional_gdn_scan( + *, + q: Tensor, + k: Tensor, + q_rot: Tensor, + k_rot: Tensor, + v: Tensor, + beta: Tensor, + decay: Tensor, + HW: tuple[int, int, int], + eps: float, +) -> Tensor: + m_hist, z_hist = _gdn_histories( + k=k, + k_rot=k_rot, + v=v, + beta=beta, + decay=decay, + HW=HW, + include_denominator=True, + ) + batch, tokens, heads, dim = q.shape + frames, height, width = HW + spatial = height * width + q = q.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4) + q_rot = q_rot.reshape(batch, frames, spatial, heads, dim).permute( + 0, + 3, + 1, + 2, + 4, + ) + num = torch.einsum("bhfsd,bhfde->bhfse", q_rot.float(), m_hist) + den = torch.einsum("bhfsd,bhfd->bhfs", q.float(), z_hist) + out = num / (den[..., None] + eps) + return out.permute(0, 2, 3, 1, 4).reshape(batch, tokens, heads, dim) + + +def _bidirectional_numerator_scan( + *, + q: Tensor, + k: Tensor, + v: Tensor, + beta: Tensor, + decay: Tensor, + HW: tuple[int, int, int], +) -> Tensor: + m_hist, _z_hist = _gdn_histories( + k=k, + k_rot=k, + v=v, + beta=beta, + decay=decay, + HW=HW, + include_denominator=False, + ) + batch, tokens, heads, dim = q.shape + frames, height, width = HW + spatial = height * width + q = q.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4) + out = torch.einsum("bhfsd,bhfde->bhfse", q.float(), m_hist) + return out.permute(0, 2, 3, 1, 4).reshape(batch, tokens, heads, dim) + + +def _gdn_histories( + *, + k: Tensor, + k_rot: Tensor, + v: Tensor, + beta: Tensor, + decay: Tensor, + HW: tuple[int, int, int], + include_denominator: bool, +) -> tuple[Tensor, Tensor | None]: + batch, tokens, heads, dim = k.shape + frames, height, width = HW + spatial = height * width + k = k.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4).float() + k_rot = ( + k_rot.reshape(batch, frames, spatial, heads, dim) + .permute(0, 3, 1, 2, 4) + .float() + ) + v = v.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4).float() + eye = torch.eye(dim, dtype=torch.float32, device=k.device).view(1, 1, dim, dim) + m = torch.zeros(batch, heads, dim, dim, dtype=torch.float32, device=k.device) + z = torch.zeros(batch, heads, dim, dtype=torch.float32, device=k.device) + m_hist = torch.empty( + batch, + heads, + frames, + dim, + dim, + dtype=torch.float32, + device=k.device, + ) + z_hist = ( + torch.empty(batch, heads, frames, dim, dtype=torch.float32, device=k.device) + if include_denominator + else None + ) + for frame in range(frames): + beta_f = beta[:, :, frame] + decay_f = decay[:, :, frame] + p_kv = torch.einsum( + "bhsd,bhse->bhde", + k_rot[:, :, frame], + beta_f[..., None] * k_rot[:, :, frame], + ) + a_val = torch.einsum( + "bhsd,bhse->bhde", + k_rot[:, :, frame], + beta_f[..., None] * v[:, :, frame], + ) + m = decay_f[..., None, None] * torch.einsum("bhde,bhef->bhdf", eye - p_kv, m) + m = m + a_val + m_hist[:, :, frame] = m + if include_denominator and z_hist is not None: + p_z = torch.einsum( + "bhsd,bhse->bhde", + k[:, :, frame], + beta_f[..., None] * k[:, :, frame], + ) + b_z = (beta_f[..., None] * k[:, :, frame]).sum(dim=2) + z = decay_f[..., None] * torch.einsum("bhde,bhe->bhd", eye - p_z, z) + z = z + b_z + z_hist[:, :, frame] = z + + m = torch.zeros_like(m) + z = torch.zeros_like(z) + for src in range(frames - 1, 0, -1): + dst = src - 1 + beta_f = beta[:, :, src] + decay_f = decay[:, :, src] + p_kv = torch.einsum( + "bhsd,bhse->bhde", + k_rot[:, :, src], + beta_f[..., None] * k_rot[:, :, src], + ) + a_val = torch.einsum( + "bhsd,bhse->bhde", + k_rot[:, :, src], + beta_f[..., None] * v[:, :, src], + ) + m = decay_f[..., None, None] * torch.einsum("bhde,bhef->bhdf", eye - p_kv, m) + m = m + a_val + m_hist[:, :, dst] += m + if include_denominator and z_hist is not None: + p_z = torch.einsum( + "bhsd,bhse->bhde", + k[:, :, src], + beta_f[..., None] * k[:, :, src], + ) + b_z = (beta_f[..., None] * k[:, :, src]).sum(dim=2) + z = decay_f[..., None] * torch.einsum("bhde,bhe->bhd", eye - p_z, z) + z = z + b_z + z_hist[:, :, dst] += z + return m_hist, z_hist + + +def _camera_qkv( + module: Stage1SelfAttention, + x: Tensor, + HW: tuple[int, int, int], +) -> tuple[Tensor, Tensor, Tensor]: + batch, tokens, _channels = x.shape + q_raw = module.q_proj_cam(x).reshape(batch, tokens, module.heads, module.dim) + k_raw_flat = module.k_proj_cam(x) + if hasattr(module, "conv_k_cam"): + k_raw_flat = _apply_bidirectional_temporal_conv(k_raw_flat, module.conv_k_cam, HW) + k_raw = k_raw_flat.reshape(batch, tokens, module.heads, module.dim) + v_raw = module.v_proj_cam(x).reshape(batch, tokens, module.heads, module.dim) + return q_raw, k_raw, v_raw + + +def _prepare_ucpe_qkv( + q_raw: Tensor, + k_raw: Tensor, + v_raw: Tensor, + *, + camera_conditions: Tensor, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + q_norm_weight: Tensor, + k_norm_weight: Tensor, + norm_eps: float, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Callable[[Tensor], Tensor]]: + batch, tokens, heads, dim = q_raw.shape + q_inv = _inv_rms(q_raw, norm_eps) + k_inv = _inv_rms(k_raw, norm_eps) + q_weight = q_norm_weight.float().view(heads, dim) + k_weight = k_norm_weight.float().view(heads, dim) + q_norm = F.relu(q_raw.float() * q_inv[:, :, None, None] * q_weight[None, None]) + k_norm = F.relu(k_raw.float() * k_inv[:, :, None, None] * k_weight[None, None]) + k_norm = k_norm * _gdn_key_scale(dim, HW) + v_float = v_raw.float() + raymats = _camera_ray_mats(camera_conditions.float(), HW) + matrix_q = raymats.reshape(batch, tokens, 4, 4).transpose(-1, -2).contiguous() + matrix_kv = _invert_se3(raymats.reshape(batch, tokens, 4, 4)).contiguous() + rope_cam = _slice_rope_for_camera(rotary_emb, dim) + + q_trans, _q_pre, _q_post = _ucpe_transform( + q_norm, + matrix_q, + rope_cam, + inverse_rope=False, + ) + k_trans, k_pre_sq, k_post_sq = _ucpe_transform( + k_norm, + matrix_kv, + rope_cam, + inverse_rope=False, + ) + v_trans, _v_pre, _v_post = _ucpe_transform( + v_float, + matrix_kv, + rope_cam, + inverse_rope=False, + ) + inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12) + + def output_projector(out: Tensor) -> Tensor: + projected, _pre, _post = _ucpe_transform( + out.float(), + raymats.reshape(batch, tokens, 4, 4), + rope_cam, + inverse_rope=True, + ) + return projected.to(dtype=out.dtype) + + return q_trans, k_trans, v_trans, inflation_sq, output_projector + + +def _prepare_ucpe_qkv_softmax( + q_raw: Tensor, + k_raw: Tensor, + v_raw: Tensor, + *, + camera_conditions: Tensor, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + q_norm_weight: Tensor, + k_norm_weight: Tensor, + norm_eps: float, +) -> tuple[Tensor, Tensor, Tensor, Callable[[Tensor], Tensor]]: + batch, tokens, heads, dim = q_raw.shape + q_inv = _inv_rms(q_raw, norm_eps) + k_inv = _inv_rms(k_raw, norm_eps) + q_weight = q_norm_weight.float().view(heads, dim) + k_weight = k_norm_weight.float().view(heads, dim) + q_norm = q_raw.float() * q_inv[:, :, None, None] * q_weight[None, None] + k_norm = k_raw.float() * k_inv[:, :, None, None] * k_weight[None, None] + v_float = v_raw.float() + raymats = _camera_ray_mats(camera_conditions.float(), HW) + matrix_q = raymats.reshape(batch, tokens, 4, 4).transpose(-1, -2).contiguous() + matrix_kv = _invert_se3(raymats.reshape(batch, tokens, 4, 4)).contiguous() + rope_cam = _slice_rope_for_camera(rotary_emb, dim) + q_trans, _q_pre, _q_post = _ucpe_transform( + q_norm, + matrix_q, + rope_cam, + inverse_rope=False, + ) + k_trans, _k_pre, _k_post = _ucpe_transform( + k_norm, + matrix_kv, + rope_cam, + inverse_rope=False, + ) + v_trans, _v_pre, _v_post = _ucpe_transform( + v_float, + matrix_kv, + rope_cam, + inverse_rope=False, + ) + + def output_projector(out: Tensor) -> Tensor: + projected, _pre, _post = _ucpe_transform( + out.float(), + raymats.reshape(batch, tokens, 4, 4), + rope_cam, + inverse_rope=True, + ) + return projected.to(dtype=out.dtype) + + return q_trans, k_trans, v_trans, output_projector + + +def _inv_rms(x: Tensor, eps: float) -> Tensor: + batch, tokens, heads, dim = x.shape + channels = heads * dim + return torch.rsqrt(x.float().pow(2).sum(dim=(-1, -2)) / channels + eps) + + +def _ucpe_transform( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> tuple[Tensor, Tensor, Tensor]: + batch, tokens, heads, dim = x.shape + half = dim // 2 + if half % 4 != 0: + raise ValueError(f"UCPE requires head_dim/2 divisible by 4, got {half}.") + first = x[..., :half].reshape(batch, tokens, heads, half // 4, 4) + first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), first.float()) + first_out = first_out.reshape(batch, tokens, heads, half) + second = x[..., half:] + if inverse_rope and rotary_emb is not None: + second_out = _apply_complex_rope(second, rotary_emb.conj()) + else: + second_out = _apply_complex_rope(second, rotary_emb) + out = torch.cat([first_out, second_out], dim=-1) + pre_sq = x.float().pow(2).sum(dim=-1).transpose(1, 2).contiguous() + post_sq = out.float().pow(2).sum(dim=-1).transpose(1, 2).contiguous() + return out, pre_sq, post_sq + + +def _camera_ray_mats( + camera_conditions: Tensor, + HW: tuple[int, int, int], +) -> Tensor: + batch, frames = camera_conditions.shape[:2] + latent_frames, height, width = HW + if frames != latent_frames: + raise ValueError( + f"camera_conditions has {frames} frames but latent grid has {latent_frames}." + ) + c2w = camera_conditions[..., :16].reshape(batch, frames, 4, 4) + fx = camera_conditions[..., 16] + fy = camera_conditions[..., 17] + cx = camera_conditions[..., 18] + cy = camera_conditions[..., 19] + y_grid, x_grid = torch.meshgrid( + torch.arange(height, dtype=torch.float32, device=camera_conditions.device), + torch.arange(width, dtype=torch.float32, device=camera_conditions.device), + indexing="ij", + ) + x = (x_grid.view(1, 1, height, width) - cx[:, :, None, None]) / fx[ + :, + :, + None, + None, + ].clamp_min(1e-6) + y = (y_grid.view(1, 1, height, width) - cy[:, :, None, None]) / fy[ + :, + :, + None, + None, + ].clamp_min(1e-6) + dirs_cam = torch.stack([x, y, torch.ones_like(x)], dim=-1) + dirs_cam = F.normalize(dirs_cam, dim=-1, eps=1e-6) + rotation = c2w[..., :3, :3] + translation = c2w[..., :3, 3] + dirs_world = torch.einsum("btij,bthwj->bthwi", rotation, dirs_cam) + cam_y = rotation[..., :, 1].view(batch, frames, 1, 1, 3).expand_as(dirs_world) + z_ray = F.normalize(dirs_world, dim=-1, eps=1e-6) + x_ray = F.normalize(torch.cross(cam_y, z_ray, dim=-1), dim=-1, eps=1e-6) + y_ray = F.normalize(torch.cross(z_ray, x_ray, dim=-1), dim=-1, eps=1e-6) + ray_to_world = torch.stack([x_ray, y_ray, z_ray], dim=-1) + world_to_ray = ray_to_world.transpose(-1, -2) + origin = translation.view(batch, frames, 1, 1, 3).expand_as(dirs_world) + trans = -torch.einsum("bthwij,bthwj->bthwi", world_to_ray, origin) + mats = torch.zeros( + batch, + frames, + height, + width, + 4, + 4, + dtype=torch.float32, + device=camera_conditions.device, + ) + mats[..., :3, :3] = world_to_ray + mats[..., :3, 3] = trans + mats[..., 3, 3] = 1.0 + return mats + + +def _invert_se3(transforms: Tensor) -> Tensor: + rotation_inv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = rotation_inv + out[..., :3, 3] = -torch.einsum( + "...ij,...j->...i", + rotation_inv, + transforms[..., :3, 3], + ) + out[..., 3, 3] = 1.0 + return out + + __all__ = [ "SANA_WM_STAGE1_SPEC", "GLUMBConvTemp", diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 2b839d7d7..eae39bf05 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -23,6 +23,9 @@ import pytest import torch +import sana_wm.native_refiner as native_refiner +import sana_wm.native_transformer as native_transformer_module + try: import tomllib except ModuleNotFoundError: # pragma: no cover - Python < 3.11 fallback @@ -62,6 +65,7 @@ _avoid_degenerate_tile_tail, _load_inference_config, ) +from sana_wm.native_refiner import SanaWMLTX2Refiner, _pack_latents, _unpack_latents from sana_wm.native_quant import ( TorchScaledMMFP4Linear, TorchScaledMMFP8Linear, @@ -72,6 +76,7 @@ SANA_WM_STAGE1_SPEC, SanaWMStage1Model, SanaWMStage1Spec, + Stage1SelfAttention, ) from sana_wm.upstream_reference import SanaWMUpstreamReferenceTransformerConfig @@ -249,6 +254,41 @@ def test_native_stage1_forward_preserves_latent_shape() -> None: assert out.shape == latents.shape +def test_stage1_self_attention_uses_camera_conditions() -> None: + """Changing camera conditions must affect native camera attention.""" + torch.manual_seed(0) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + attn = Stage1SelfAttention(spec, use_gdn_convs=True).eval() + hidden = torch.randn(1, 8, 16) + camera = torch.zeros(1, 2, 20) + camera[..., :16] = torch.eye(4).flatten() + camera[..., 16:] = torch.tensor([1.0, 1.0, 0.5, 0.5]) + shifted_camera = camera.clone() + shifted_camera[..., 3] = 0.25 + + base = attn(hidden, HW=(2, 2, 2), camera_conditions=camera) + shifted = attn(hidden, HW=(2, 2, 2), camera_conditions=shifted_camera) + + assert base.shape == hidden.shape + assert not torch.allclose(base, shifted) + + def test_native_transformer_releases_stage1_runtime() -> None: """Free Stage-1-only modules and conditioning before decode/refine.""" transformer = SanaWMNativeTransformerConfig().setup() @@ -366,7 +406,7 @@ def enable_tiling(self, **kwargs: int) -> None: calls = 0 inference_modes: list[bool] = [] - def fake_decode(_samples: torch.Tensor) -> torch.Tensor: + def decode_once(_samples: torch.Tensor) -> torch.Tensor: nonlocal calls calls += 1 inference_modes.append(torch.is_inference_mode_enabled()) @@ -374,7 +414,7 @@ def fake_decode(_samples: torch.Tensor) -> torch.Tensor: raise torch.OutOfMemoryError("test OOM") return torch.zeros((1, 3, 1, 2, 2), dtype=torch.float32) - monkeypatch.setattr(transformer, "_vae_decode", fake_decode) + monkeypatch.setattr(transformer, "_vae_decode", decode_once) video = transformer.decode_latents(torch.zeros((1, 4, 1, 1, 1))) @@ -472,6 +512,112 @@ def test_runner_defaults_to_bf16_precision() -> None: assert RUNNER_SANA_WM_BIDIRECTIONAL.stage1_precision == "bf16" assert RUNNER_SANA_WM_BIDIRECTIONAL.refiner_precision == "bf16" assert RUNNER_SANA_WM_BIDIRECTIONAL.quant_backend == "auto" + assert RUNNER_SANA_WM_BIDIRECTIONAL.no_refiner is False + + +def test_native_refiner_is_flashdreams_owned( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Build the native refiner adapter instead of raising or importing Sana.""" + calls: dict[str, object] = {} + + class DummyRefiner: + def __init__(self, **kwargs: object) -> None: + calls.update(kwargs) + + monkeypatch.setattr(native_refiner, "SanaWMLTX2Refiner", DummyRefiner) + monkeypatch.setattr( + native_transformer_module, + "resolve_hf_path", + lambda value: f"/resolved/{value}", + ) + transformer = SanaWMNativeTransformerConfig().setup() + + transformer._ensure_refiner() + + assert transformer._refiner_built is True + assert isinstance(transformer.refiner, DummyRefiner) + assert calls["refiner_root"] == ( + "/resolved/hf://Efficient-Large-Model/SANA-WM_bidirectional/refiner" + ) + assert calls["gemma_root"] == ( + "/resolved/hf://Efficient-Large-Model/SANA-WM_bidirectional/refiner/text_encoder" + ) + assert calls["dtype"] is torch.bfloat16 + assert calls["precision"] == "bf16" + assert calls["quant_backend"] == "native" + + +def test_refiner_latent_pack_round_trips() -> None: + """Preserve LTX-2 latent layout across token packing and unpacking.""" + latents = torch.arange(1 * 4 * 3 * 2 * 2, dtype=torch.float32).reshape( + 1, + 4, + 3, + 2, + 2, + ) + + packed = _pack_latents(latents, patch_size=1, patch_size_t=1) + unpacked = _unpack_latents( + packed, + num_frames=3, + height=2, + width=2, + patch_size=1, + patch_size_t=1, + ) + + torch.testing.assert_close(unpacked, latents) + + +def test_refiner_block_size_request_still_executes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Do not turn native refiner block-size requests into a hard failure.""" + + class DummyTransformer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = type("Config", (), {"patch_size": 1, "patch_size_t": 1})() + + refiner = SanaWMLTX2Refiner.__new__(SanaWMLTX2Refiner) + torch.nn.Module.__init__(refiner) + refiner.device = torch.device("cpu") + refiner.dtype = torch.float32 + refiner.transformer = DummyTransformer() + monkeypatch.setattr(refiner, "_prepare_quantization", lambda: None) + monkeypatch.setattr( + refiner, + "_encode_prompt", + lambda _prompt: (torch.zeros((1, 1, 1)), torch.ones((1, 1))), + ) + + def predict_current_x0( + *, + sink: torch.Tensor, + noisy_current: torch.Tensor, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + sigma: torch.Tensor, + fps: float, + ) -> torch.Tensor: + del sink, prompt_embeds, prompt_attention_mask, sigma, fps + return torch.zeros_like(_pack_latents(noisy_current)) + + monkeypatch.setattr(refiner, "_predict_current_x0", predict_current_x0) + latents = torch.zeros((1, 2, 3, 1, 1)) + + refined = refiner.refine_latents( + latents, + "demo", + fps=16.0, + block_size=1, + progress=False, + sigmas=(0.5, 0.0), + ) + + assert refined.shape == latents.shape def test_auto_quant_backend_resolves_to_native_backend() -> None: @@ -653,7 +799,7 @@ def test_torch_fp8_backend_rejects_fp4(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) - with pytest.raises(ValueError, match="supports fp8 only"): + with pytest.raises(ValueError, match="accepts fp8 only"): _validate_precision_request( device=torch.device("cuda:0"), stage1_precision="fp4", @@ -685,7 +831,7 @@ def test_torch_fp4_backend_rejects_fp8(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) - with pytest.raises(ValueError, match="supports fp4 only"): + with pytest.raises(ValueError, match="accepts fp4 only"): _validate_precision_request( device=torch.device("cuda:0"), stage1_precision="fp8", @@ -803,7 +949,7 @@ def test_pyproject_entry_point_matches_runner_literal() -> None: def test_pyproject_excludes_upstream_diffusion_package() -> None: - """Do not install the vendored upstream SANA ``diffusion`` package.""" + """Do not install the NVlabs/Sana ``diffusion`` package.""" pyproject = tomllib.loads( Path("integrations/sana/pyproject.toml").read_text(encoding="utf-8") ) @@ -811,6 +957,7 @@ def test_pyproject_excludes_upstream_diffusion_package() -> None: packages = pyproject["tool"]["setuptools"]["packages"]["find"] dependencies = set(pyproject["project"]["dependencies"]) + assert "accelerate>=1.0" in dependencies assert packages["include"] == ["sana_wm*"] assert "diffusers>=0.36" in dependencies assert "transformers>=5.0,<6" in dependencies diff --git a/uv.lock b/uv.lock index 407fb184b..4a76b98c6 100644 --- a/uv.lock +++ b/uv.lock @@ -77,6 +77,27 @@ test = [ { name = "tomli", specifier = ">=2.0" }, ] +[[package]] +name = "accelerate" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "packaging" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8d/75/94cd5d389649578aca399e5aa822637eec18319a1dadc400ffe2f9a7493f/accelerate-1.14.0.tar.gz", hash = "sha256:41b9c4377a54e0b460a959b0defa1b736e4ca0a2373252d9a539964c2afe3c8d", size = 412167, upload-time = "2026-06-11T13:45:52.326Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl", hash = "sha256:e94390c2863b873be18f623f9df48a0d8fe5eff13ea7f1a00092b0a7904888c6", size = 389246, upload-time = "2026-06-11T13:45:50.477Z" }, +] + [[package]] name = "accessible-pygments" version = "0.0.5" @@ -1573,6 +1594,7 @@ name = "flashdreams-sana-wm" version = "0.1.0" source = { editable = "integrations/sana" } dependencies = [ + { name = "accelerate" }, { name = "diffusers" }, { name = "flashdreams" }, { name = "imageio", extra = ["ffmpeg"] }, @@ -1592,6 +1614,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "accelerate", specifier = ">=1.0" }, { name = "diffusers", specifier = ">=0.36" }, { name = "flashdreams", editable = "flashdreams" }, { name = "imageio", extras = ["ffmpeg"], specifier = ">=2.31" }, From 0f77ccfb33933bda2c5fd001e0fc9091d73d353a Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 16 Jul 2026 19:58:18 -0700 Subject: [PATCH 11/36] Run native SANA refiner under inference mode The native LTX-2 refiner prompt encoder returns inference tensors. Keep the whole refiner denoise path inside torch.inference_mode so diffusers Linear calls do not try to save those tensors for autograd. Add a CPU regression assertion that the refiner prediction callback executes with inference mode enabled. Validation: uv run --package flashdreams-sana-wm pytest integrations/sana/tests/test_smoke.py integrations/sana/tests/test_camera.py. --- integrations/sana/sana_wm/native_refiner.py | 1 + integrations/sana/tests/test_smoke.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/integrations/sana/sana_wm/native_refiner.py b/integrations/sana/sana_wm/native_refiner.py index 084940608..4ea247c9a 100644 --- a/integrations/sana/sana_wm/native_refiner.py +++ b/integrations/sana/sana_wm/native_refiner.py @@ -73,6 +73,7 @@ def __init__( self._quantized = False self.transformer, self.connectors = self._load_diffusers_components() + @torch.inference_mode() def refine_latents( self, sana_latent: Tensor, diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index eae39bf05..ddb05d8f6 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -586,6 +586,7 @@ def __init__(self) -> None: refiner.device = torch.device("cpu") refiner.dtype = torch.float32 refiner.transformer = DummyTransformer() + inference_modes: list[bool] = [] monkeypatch.setattr(refiner, "_prepare_quantization", lambda: None) monkeypatch.setattr( refiner, @@ -603,6 +604,7 @@ def predict_current_x0( fps: float, ) -> torch.Tensor: del sink, prompt_embeds, prompt_attention_mask, sigma, fps + inference_modes.append(torch.is_inference_mode_enabled()) return torch.zeros_like(_pack_latents(noisy_current)) monkeypatch.setattr(refiner, "_predict_current_x0", predict_current_x0) @@ -618,6 +620,7 @@ def predict_current_x0( ) assert refined.shape == latents.shape + assert inference_modes == [True] def test_auto_quant_backend_resolves_to_native_backend() -> None: From 52617e89673ecd3e8eab3e49f7917696db356d95 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 17 Jul 2026 09:28:38 -0700 Subject: [PATCH 12/36] Clean up SANA-WM integration surface Signed-off-by: Aidan Foster --- integrations/sana/README.md | 67 ++- integrations/sana/docs/upstream_reference.md | 103 ---- integrations/sana/pyproject.toml | 2 - integrations/sana/sana_wm/__init__.py | 2 +- integrations/sana/sana_wm/camera.py | 4 +- integrations/sana/sana_wm/config.py | 71 +-- integrations/sana/sana_wm/constants.py | 4 +- .../{native_diffusion.py => diffusion.py} | 26 +- .../sana_wm/{native_quant.py => quant.py} | 53 +- .../sana_wm/{native_refiner.py => refiner.py} | 19 +- integrations/sana/sana_wm/runner.py | 470 +++--------------- integrations/sana/sana_wm/stage1_model.py | 12 +- .../{native_transformer.py => transformer.py} | 87 ++-- .../sana/sana_wm/upstream_reference.py | 90 ---- .../sana/tests/parity_check/README.md | 25 - ...ative_quant_cuda.py => test_quant_cuda.py} | 18 +- integrations/sana/tests/test_smoke.py | 199 +++----- 17 files changed, 269 insertions(+), 983 deletions(-) delete mode 100644 integrations/sana/docs/upstream_reference.md rename integrations/sana/sana_wm/{native_diffusion.py => diffusion.py} (91%) rename integrations/sana/sana_wm/{native_quant.py => quant.py} (91%) rename integrations/sana/sana_wm/{native_refiner.py => refiner.py} (96%) rename integrations/sana/sana_wm/{native_transformer.py => transformer.py} (92%) delete mode 100644 integrations/sana/sana_wm/upstream_reference.py delete mode 100644 integrations/sana/tests/parity_check/README.md rename integrations/sana/tests/{test_native_quant_cuda.py => test_quant_cuda.py} (81%) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index db4830004..a8a9f18d5 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -9,36 +9,28 @@ FlashDreams SANA-WM integration for [SANA-WM](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional), the 2.6B bidirectional camera-controlled world model released from NVlabs/Sana. -The `sana-wm-bidirectional` runner is the native FlashDreams path. It uses -FlashDreams config, runner, pipeline, diffusion-model, scheduler, transformer, -camera-conditioning, VAE decode, and output-writing boundaries. The Stage-1 -DiT module is implemented in this package and loads the public SANA-WM -checkpoint directly; it does not import or install the NVlabs/Sana -`diffusion` package. +The `sana-wm-bidirectional` runner uses FlashDreams config, runner, pipeline, +diffusion-model, scheduler, transformer, camera-conditioning, VAE decode, and +output-writing boundaries. The Stage-1 DiT module in this package loads the +public SANA-WM checkpoint directly. -Current native scope: +Current scope: | component | status | | --- | --- | -| Stage-1 BF16 | Native FlashDreams DiT execution. | -| Stage-1 FP8 | Native PyTorch `_scaled_mm` backend. | -| Stage-1 FP4 | Native Triton quantization plus PyTorch `_scaled_mm`. | -| VAE decode | Native direct `diffusers` LTX2 VAE use with low-memory tiling. | -| LTX-2 refiner | Native direct `diffusers` LTX-2 transformer and Gemma connector use. | - -The separate `sana-wm-upstream-reference` runner is documented only in -`docs/upstream_reference.md`. Keep reference-harness details out of this -native integration README. +| Stage-1 BF16 | FlashDreams DiT execution. | +| Stage-1 FP8 | PyTorch `_scaled_mm` backend. | +| Stage-1 FP4 | Triton quantization plus PyTorch `_scaled_mm`. | +| VAE decode | Direct `diffusers` LTX2 VAE use with low-memory tiling. | +| LTX-2 refiner | Direct `diffusers` LTX-2 transformer and Gemma connector use. | ## Runner | slug | description | | --- | --- | -| `sana-wm-bidirectional` | Native FlashDreams SANA-WM Stage-1 + LTX-2 refiner runner. | -| `sana-wm-native-bidirectional` | Alias for the native FlashDreams runner. | +| `sana-wm-bidirectional` | SANA-WM Stage-1 + LTX-2 refiner runner. | -The FlashDreams package is named `sana_wm` rather than `sana` so it does not -shadow the upstream project name. +The FlashDreams package is named `sana_wm`. ## Setup @@ -49,14 +41,13 @@ project's GPU runtime dependencies: uv sync --package flashdreams-sana-wm --extra dev ``` -The native runner does not require a local NVlabs/Sana checkout. The checkout -is only useful as a source of demo assets such as -`asset/sana_wm/demo_0.png`, `demo_0.txt`, `demo_0_pose.npy`, and -`demo_0_intrinsics.npy`. +The examples below use the public demo image, prompt, camera, and intrinsics +files from the SANA-WM release. Equivalent inputs with the same shapes work as +well. ## Run -The native runner requires explicit intrinsics: +The runner requires explicit intrinsics: ```bash PYTORCH_ALLOC_CONF=expandable_segments:True \ @@ -66,13 +57,13 @@ uv run flashdreams-run sana-wm-bidirectional \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ - --output-dir outputs/sana_wm_native_bf16 + --output-dir outputs/sana_wm_bf16 ``` Expected output: ```text -outputs/sana_wm_native_bf16/sana_wm_generated.mp4 +outputs/sana_wm_bf16/sana_wm_generated.mp4 ``` For Stage-1-only diagnostics, add `--no-refiner True`. @@ -80,14 +71,14 @@ For Stage-1-only diagnostics, add `--no-refiner True`. ## FP8 and FP4 The runner defaults to BF16. Quantized Stage-1 inference is opt-in and uses -native FlashDreams/PyTorch backends by default, not Transformer Engine. +Torch-based backends by default. | option | hardware | backend | | --- | --- | --- | -| `--stage1-precision fp8` | Hopper or newer (`sm_90+`) | native E4M3 `_scaled_mm` | -| `--stage1-precision fp4` | Blackwell (`sm_100+`) | native Triton NVFP4 plus `_scaled_mm` | -| `--quant-backend torch-fp8` | Hopper or newer (`sm_90+`) | force the native FP8 backend | -| `--quant-backend torch-fp4` | Blackwell (`sm_100+`) | force the native FP4 backend | +| `--stage1-precision fp8` | Hopper or newer (`sm_90+`) | E4M3 `_scaled_mm` | +| `--stage1-precision fp4` | Blackwell (`sm_100+`) | Triton NVFP4 plus `_scaled_mm` | +| `--quant-backend torch-fp8` | Hopper or newer (`sm_90+`) | force FP8 | +| `--quant-backend torch-fp4` | Blackwell (`sm_100+`) | force FP4 | FP8 smoke: @@ -99,7 +90,7 @@ uv run flashdreams-run sana-wm-bidirectional \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ - --output-dir outputs/sana_wm_native_fp8 \ + --output-dir outputs/sana_wm_fp8 \ --stage1-precision fp8 \ --refiner-precision fp8 ``` @@ -114,7 +105,7 @@ uv run flashdreams-run sana-wm-bidirectional \ --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ --num-frames 161 \ - --output-dir outputs/sana_wm_native_fp4 \ + --output-dir outputs/sana_wm_fp4 \ --stage1-precision fp4 \ --refiner-precision fp4 ``` @@ -122,12 +113,12 @@ uv run flashdreams-run sana-wm-bidirectional \ ## Tests CPU-safe tests cover import, config boundaries, action parsing, intrinsics, -camera conditioning, native Stage-1 checkpoint schema, native Stage-1 CPU -forward shape, VAE tiling, and native low-precision backend selection: +camera conditioning, Stage-1 checkpoint schema, Stage-1 CPU forward shape, VAE +tiling, and low-precision backend selection: ```bash uv run --extra dev pytest integrations/sana/tests/test_smoke.py ``` -GPU generation and native-vs-reference quality parity checks are heavyweight -manual workflows and should stay out of `ci_cpu`. +GPU generation checks are heavyweight manual workflows and should stay out of +`ci_cpu`. diff --git a/integrations/sana/docs/upstream_reference.md b/integrations/sana/docs/upstream_reference.md deleted file mode 100644 index 21a339fc0..000000000 --- a/integrations/sana/docs/upstream_reference.md +++ /dev/null @@ -1,103 +0,0 @@ - - -# SANA-WM Upstream Reference Harness - -This document covers the temporary upstream-reference runner. It exists only -for parity/debug work while the native FlashDreams integration is being -validated. Do not use this as user-facing documentation for the real SANA-WM -FlashDreams integration. - -## Runner - -| slug | description | -| --- | --- | -| `sana-wm-upstream-reference` | FlashDreams-packaged NVlabs/Sana reference harness. | - -This runner delegates generation to the SANA reference pipeline object. It is -expected to be removed once native generation is validated. - -## Setup - -The reference harness needs a local or importable NVlabs/Sana checkout: - -```bash -git clone https://github.com/NVlabs/Sana ../Sana -cd ../Sana -bash environment_setup.sh sana -conda activate sana -``` - -The runner auto-detects `../Sana`. You can also set `SANA_ROOT` or pass -`--upstream-sana-root /path/to/Sana`. - -## Run - -```bash -flashdreams-run sana-wm-upstream-reference \ - --upstream-sana-root ../Sana \ - --image-path ../Sana/asset/sana_wm/demo_0.png \ - --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ - --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ - --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ - --num-frames 321 \ - --output-dir outputs/sana_wm_reference -``` - -The upstream reference harness can estimate intrinsics with Pi3X when -`--intrinsics-path` is omitted. The native runner currently requires explicit -intrinsics instead. - -## Transformer Engine Reference Notes - -Upstream SANA does not use Transformer Engine here for positional embeddings. -Where the reference path uses Transformer Engine, it is as a quantized -`Linear`/GEMM backend: - -- Stage 1 replaces selected `torch.nn.Linear` layers with - `transformer_engine.pytorch.Linear`, then wraps transformer block or linear - forwards in `te.fp8_autocast(...)`. -- `SANA_WM_STAGE1_QUANT=nvfp4` selects `NVFP4BlockScaling`; `fp8block` selects - `Float8BlockScaling`. -- The default FlashDreams FP4 Stage-1 selector maps to upstream - `self_attn+cross+ffn`, covering self-attention, cross-attention, and MLP - linear layers. In the validated reference demo run, upstream reported 120 - converted linear layers and 20 wrapped transformer blocks. -- The refiner applies the same pattern to eligible LTX-2 transformer linear - layers after skipping input/output projection, audio-only, caption, and time - embedding modules. - -Use `--quant-backend upstream-te` only on this reference harness when comparing -native low-precision behavior against upstream Transformer Engine behavior. - -```bash -flashdreams-run sana-wm-upstream-reference \ - --upstream-sana-root ../Sana \ - --image-path ../Sana/asset/sana_wm/demo_0.png \ - --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ - --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ - --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ - --num-frames 161 \ - --output-dir outputs/sana_wm_reference_fp4 \ - --stage1-precision fp4 \ - --refiner-precision fp4 \ - --quant-backend upstream-te -``` - -FlashDreams already removed a separate Transformer Engine dependency for RoPE -via `flashdreams.core.attention.rope_kernel`, but that replacement is not a -drop-in for SANA FP4/FP8. Native SANA replaces eligible Stage-1 and refiner -`Linear` modules directly, including activation quantization/scales, weight -conversion, GEMM, and output rescaling. - -Relevant FlashDreams components for future hardening: - -- `flashdreams.core.attention.rope_kernel`: the prior narrow TE-replacement - pattern and test strategy. -- `integrations/omnidreams/omnidreams_singleview/python/cosmos_fp8_utils.py`: - FP8 weight quantization and prepared weight/scale aliases for Cosmos blocks. -- `integrations/omnidreams/omnidreams_singleview/src/dit_streaming/kernels/`: - CUDA/CUTLASS FP8 GEMM, BF16 GEMM, INT8 block quantization, and SageAttention3 - FP4 attention/cache kernels. diff --git a/integrations/sana/pyproject.toml b/integrations/sana/pyproject.toml index 780bf446b..2d63b6f77 100644 --- a/integrations/sana/pyproject.toml +++ b/integrations/sana/pyproject.toml @@ -45,8 +45,6 @@ dev = [ [project.entry-points."flashdreams.runner_configs"] "sana-wm-bidirectional" = "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL" -"sana-wm-upstream-reference" = "sana_wm.config:RUNNER_SANA_WM_UPSTREAM_REFERENCE" -"sana-wm-native-bidirectional" = "sana_wm.config:RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL" [tool.setuptools.packages.find] include = ["sana_wm*"] diff --git a/integrations/sana/sana_wm/__init__.py b/integrations/sana/sana_wm/__init__.py index 77dfd6e32..c59615ca6 100644 --- a/integrations/sana/sana_wm/__init__.py +++ b/integrations/sana/sana_wm/__init__.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SANA-WM native and upstream-reference runners for FlashDreams.""" +"""SANA-WM runner for FlashDreams.""" from sana_wm.config import RUNNER_CONFIGS, RUNNER_SANA_WM_BIDIRECTIONAL diff --git a/integrations/sana/sana_wm/camera.py b/integrations/sana/sana_wm/camera.py index 97daf41ba..af8aa406e 100644 --- a/integrations/sana/sana_wm/camera.py +++ b/integrations/sana/sana_wm/camera.py @@ -35,7 +35,7 @@ """Camera-control integration frame rate.""" DEFAULT_TRANSLATION_SPEED = 0.025 -"""Default per-frame translation magnitude used by upstream SANA-WM.""" +"""Default per-frame translation magnitude used by SANA-WM.""" DEFAULT_ROTATION_SPEED_DEG = 0.6 """Default per-frame rotation magnitude in degrees.""" @@ -229,7 +229,7 @@ def action_string_to_c2w( translation_speed: Per-frame translation magnitude. rotation_speed_deg: Per-frame angular magnitude in degrees. pitch_limit_deg: Maximum absolute pitch in degrees. - smooth: Whether to use upstream's velocity smoothing model. + smooth: Whether to use the SANA-WM velocity smoothing model. Returns: ``[F + 1, 4, 4]`` camera-to-world matrices in OpenCV convention. diff --git a/integrations/sana/sana_wm/config.py b/integrations/sana/sana_wm/config.py index 73734b634..5d9a176c8 100644 --- a/integrations/sana/sana_wm/config.py +++ b/integrations/sana/sana_wm/config.py @@ -13,97 +13,42 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Static configs for SANA-WM native and upstream-reference runners.""" +"""Static configs for the SANA-WM runner.""" from __future__ import annotations -from typing import cast - -from flashdreams.infra.config import derive_config -from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler import FlowMatchSchedulerConfig from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.infra.runner import RunnerConfig -from sana_wm.native_diffusion import SanaWMDiffusionModelConfig -from sana_wm.native_transformer import SanaWMNativeTransformerConfig +from sana_wm.diffusion import SanaWMDiffusionModelConfig from sana_wm.runner import SanaWMRunnerConfig -from sana_wm.upstream_reference import SanaWMUpstreamReferenceTransformerConfig - - -def _reference_pipeline(name: str) -> StreamInferencePipelineConfig: - return StreamInferencePipelineConfig( - name=name, - diffusion_model=DiffusionModelConfig( - transformer=SanaWMUpstreamReferenceTransformerConfig(), - scheduler=FlowMatchSchedulerConfig(), - seed=42, - ), - ) +from sana_wm.transformer import SanaWMTransformerConfig -PIPELINE_SANA_WM_UPSTREAM_REFERENCE = _reference_pipeline("sana-wm-upstream-reference") -"""Schema marker for the upstream-delegating SANA-WM reference runner.""" - PIPELINE_SANA_WM_BIDIRECTIONAL = StreamInferencePipelineConfig( name="sana-wm-bidirectional", diffusion_model=SanaWMDiffusionModelConfig( - transformer=SanaWMNativeTransformerConfig(), + transformer=SanaWMTransformerConfig(), scheduler=FlowMatchSchedulerConfig(), seed=42, ), ) -"""Native FlashDreams SANA-WM pipeline.""" - -PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL = cast( - StreamInferencePipelineConfig, - derive_config(PIPELINE_SANA_WM_BIDIRECTIONAL, name="sana-wm-native-bidirectional"), -) -"""Backward-compatible native SANA-WM pipeline alias.""" - -RUNNER_SANA_WM_UPSTREAM_REFERENCE = SanaWMRunnerConfig( - runner_name=PIPELINE_SANA_WM_UPSTREAM_REFERENCE.name, - description=( - "SANA-WM upstream reference harness (NVlabs/Sana Stage-1 + LTX-2 refiner)." - ), - pipeline=PIPELINE_SANA_WM_UPSTREAM_REFERENCE, -) -"""Explicit SANA-WM upstream reference runner config.""" +"""FlashDreams SANA-WM pipeline.""" RUNNER_SANA_WM_BIDIRECTIONAL = SanaWMRunnerConfig( runner_name=PIPELINE_SANA_WM_BIDIRECTIONAL.name, - description=( - "SANA-WM bidirectional I2V native FlashDreams runner " - "(Stage-1 DiT + LTX-2 refiner)." - ), + description="SANA-WM bidirectional I2V runner (Stage-1 DiT + LTX-2 refiner).", pipeline=PIPELINE_SANA_WM_BIDIRECTIONAL, - execution_backend="native-flashdreams", -) -"""SANA-WM native FlashDreams runner config.""" - -RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL = SanaWMRunnerConfig( - runner_name=PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL.name, - description="SANA-WM native FlashDreams pipeline alias.", - pipeline=PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL, - execution_backend="native-flashdreams", ) -"""Backward-compatible native SANA-WM FlashDreams runner alias.""" +"""SANA-WM runner config.""" RUNNER_CONFIGS: dict[str, RunnerConfig] = { - cfg.runner_name: cfg - for cfg in ( - RUNNER_SANA_WM_BIDIRECTIONAL, - RUNNER_SANA_WM_UPSTREAM_REFERENCE, - RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, - ) + cfg.runner_name: cfg for cfg in (RUNNER_SANA_WM_BIDIRECTIONAL,) } """SANA-WM runner configs keyed by ``runner_name``.""" __all__ = [ "PIPELINE_SANA_WM_BIDIRECTIONAL", - "PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL", - "PIPELINE_SANA_WM_UPSTREAM_REFERENCE", "RUNNER_CONFIGS", "RUNNER_SANA_WM_BIDIRECTIONAL", - "RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL", - "RUNNER_SANA_WM_UPSTREAM_REFERENCE", ] diff --git a/integrations/sana/sana_wm/constants.py b/integrations/sana/sana_wm/constants.py index 5251099ca..d45f0c974 100644 --- a/integrations/sana/sana_wm/constants.py +++ b/integrations/sana/sana_wm/constants.py @@ -22,7 +22,7 @@ """Default Stage-1 SANA-WM DiT checkpoint.""" SANA_WM_CONFIG_PATH = f"hf://{SANA_WM_HF_REPO}/config.yaml" -"""Default upstream inference YAML.""" +"""Default inference YAML.""" SANA_WM_REFINER_ROOT = f"hf://{SANA_WM_HF_REPO}/refiner" """Default LTX-2 refiner root.""" @@ -46,4 +46,4 @@ """Frame rate used by the public SANA-WM examples.""" DEFAULT_ACTION = "w-100,dw-60,w-100,aw-60" -"""Default SANA-WM demo action string from upstream ``demo_0``.""" +"""Default SANA-WM demo action string.""" diff --git a/integrations/sana/sana_wm/native_diffusion.py b/integrations/sana/sana_wm/diffusion.py similarity index 91% rename from integrations/sana/sana_wm/native_diffusion.py rename to integrations/sana/sana_wm/diffusion.py index 35270d4a0..2c26149b3 100644 --- a/integrations/sana/sana_wm/native_diffusion.py +++ b/integrations/sana/sana_wm/diffusion.py @@ -27,9 +27,9 @@ from torch import Tensor from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig -from sana_wm.native_transformer import ( - SanaWMNativeTransformer, - SanaWMNativeTransformerCache, +from sana_wm.transformer import ( + SanaWMTransformer, + SanaWMTransformerCache, ) @@ -42,31 +42,31 @@ class SanaWMDiffusionModelConfig(DiffusionModelConfig): ) -class SanaWMDiffusionModel(DiffusionModel[SanaWMNativeTransformerCache]): +class SanaWMDiffusionModel(DiffusionModel[SanaWMTransformerCache]): """Run SANA-WM Stage-1 denoising behind the FlashDreams diffusion boundary.""" - transformer: SanaWMNativeTransformer + transformer: SanaWMTransformer def __init__(self, config: SanaWMDiffusionModelConfig) -> None: super().__init__(config) - if not isinstance(self.transformer, SanaWMNativeTransformer): + if not isinstance(self.transformer, SanaWMTransformer): raise TypeError( - "SanaWMDiffusionModel requires SanaWMNativeTransformerConfig." + "SanaWMDiffusionModel requires SanaWMTransformerConfig." ) def generate( self, autoregressive_index: int, - cache: SanaWMNativeTransformerCache, + cache: SanaWMTransformerCache, input: Any = None, - ) -> tuple[Tensor, "DiffusionModel.FinalState[SanaWMNativeTransformerCache]"]: + ) -> tuple[Tensor, "DiffusionModel.FinalState[SanaWMTransformerCache]"]: """Run SANA-WM's first-frame-pinned LTX Euler denoising loop.""" del input if autoregressive_index != 0: - raise ValueError("SANA-WM bidirectional native inference has one AR step.") + raise ValueError("SANA-WM bidirectional inference has one AR step.") conditioning = cache.conditioning if conditioning is None: - raise RuntimeError("SANA-WM native diffusion cache has no conditioning.") + raise RuntimeError("SANA-WM diffusion cache has no conditioning.") cache.start(autoregressive_index) latents = self.transformer.initial_latents(conditioning) @@ -92,7 +92,7 @@ def generate( def finalize( self, - final_state: "DiffusionModel.FinalState[SanaWMNativeTransformerCache]", + final_state: "DiffusionModel.FinalState[SanaWMTransformerCache]", ) -> None: """Finalize the one-shot cache without an extra model forward.""" final_state.cache.finalize(final_state.autoregressive_index) @@ -100,7 +100,7 @@ def finalize( def _sample_ltx_euler( self, latents: Tensor, - cache: SanaWMNativeTransformerCache, + cache: SanaWMTransformerCache, ) -> Tensor: conditioning = cache.conditioning assert conditioning is not None diff --git a/integrations/sana/sana_wm/native_quant.py b/integrations/sana/sana_wm/quant.py similarity index 91% rename from integrations/sana/sana_wm/native_quant.py rename to integrations/sana/sana_wm/quant.py index 50b878223..3ff3ee0e1 100644 --- a/integrations/sana/sana_wm/native_quant.py +++ b/integrations/sana/sana_wm/quant.py @@ -13,15 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Native low-precision linear helpers for SANA-WM. - -Upstream SANA's FP8/FP4 path uses Transformer Engine by replacing selected -``torch.nn.Linear`` modules and entering ``te.fp8_autocast`` around the -corresponding transformer blocks. This module provides TE-free replacements for -the same replacement point: inference-only FP8 and NVFP4 Linear modules backed -by PyTorch's native ``torch._scaled_mm`` kernel, with Triton used for NVFP4 -activation quantization. -""" +"""Low-precision linear helpers for SANA-WM.""" from __future__ import annotations @@ -43,7 +35,7 @@ @dataclass(frozen=True) class TorchScaledMMFP8Recipe: - """Marker recipe used while patching upstream SANA's TE recipe hooks.""" + """FP8 recipe marker for SANA-WM Linear replacement.""" name: str = "torch_scaled_mm_fp8" precision: Literal["fp8"] = "fp8" @@ -51,13 +43,13 @@ class TorchScaledMMFP8Recipe: @dataclass(frozen=True) class TorchScaledMMFP4Recipe: - """Marker recipe used while patching upstream SANA's NVFP4 hooks.""" + """FP4 recipe marker for SANA-WM Linear replacement.""" name: str = "torch_scaled_mm_fp4" precision: Literal["fp4"] = "fp4" -NativeQuantRecipe = TorchScaledMMFP8Recipe | TorchScaledMMFP4Recipe +QuantRecipe = TorchScaledMMFP8Recipe | TorchScaledMMFP4Recipe @triton.jit @@ -184,7 +176,7 @@ class TorchScaledMMFP8Linear(nn.Module): scales. At runtime it quantizes the flattened activation rows to E4M3 with per-row scales, then calls ``torch._scaled_mm`` and returns BF16 output. This intentionally mirrors the shape contract of ``nn.Linear`` so it can - replace eligible upstream SANA linear layers without changing call sites. + replace eligible SANA-WM linear layers without changing call sites. """ in_features: int @@ -410,7 +402,7 @@ def replace_linear_with_torch_fp8( prefix: str = "", ) -> tuple[int, int]: """Replace eligible ``nn.Linear`` modules with ``TorchScaledMMFP8Linear``.""" - return _replace_linear_with_native_quant( + return _replace_linear_with_quant( module, recipe=TorchScaledMMFP8Recipe(), params_dtype=params_dtype, @@ -430,7 +422,7 @@ def replace_linear_with_torch_fp4( prefix: str = "", ) -> tuple[int, int]: """Replace eligible ``nn.Linear`` modules with ``TorchScaledMMFP4Linear``.""" - return _replace_linear_with_native_quant( + return _replace_linear_with_quant( module, recipe=TorchScaledMMFP4Recipe(), params_dtype=params_dtype, @@ -440,22 +432,17 @@ def replace_linear_with_torch_fp4( ) -def replace_linear_with_native_quant( +def replace_linear_with_quant( module: nn.Module, *, - recipe: NativeQuantRecipe, + recipe: QuantRecipe, params_dtype: torch.dtype, skip_patterns: tuple[str, ...], include_patterns: tuple[str, ...] | None = None, prefix: str = "", ) -> tuple[int, int]: - """Replace eligible ``nn.Linear`` modules with the requested native backend. - - The signature intentionally matches upstream SANA's - ``_replace_linear_with_te_nvfp4`` helper so the runner can patch the backend - without modifying upstream source files. - """ - return _replace_linear_with_native_quant( + """Replace eligible ``nn.Linear`` modules with the requested backend.""" + return _replace_linear_with_quant( module, recipe=recipe, params_dtype=params_dtype, @@ -465,10 +452,10 @@ def replace_linear_with_native_quant( ) -def _replace_linear_with_native_quant( +def _replace_linear_with_quant( module: nn.Module, *, - recipe: NativeQuantRecipe, + recipe: QuantRecipe, params_dtype: torch.dtype, skip_patterns: tuple[str, ...], include_patterns: tuple[str, ...] | None, @@ -510,11 +497,11 @@ def _replace_linear_with_native_quant( out_dtype=out_dtype, ) else: - raise ValueError(f"Unsupported native quant recipe: {recipe!r}.") + raise ValueError(f"Unsupported SANA-WM quant recipe: {recipe!r}.") setattr(module, name, replacement) converted += 1 continue - child_converted, child_skipped = _replace_linear_with_native_quant( + child_converted, child_skipped = _replace_linear_with_quant( child, recipe=recipe, params_dtype=params_dtype, @@ -533,22 +520,22 @@ def _name_matches(patterns: tuple[str, ...], name: str) -> bool: def _require_scaled_mm() -> None: if not hasattr(torch, "_scaled_mm"): - raise RuntimeError("torch._scaled_mm is required for native SANA quantization.") + raise RuntimeError("torch._scaled_mm is required for SANA-WM quantization.") def _require_fp8_dtype() -> None: if not hasattr(torch, "float8_e4m3fn"): - raise RuntimeError("torch.float8_e4m3fn is required for the native SANA FP8 backend.") + raise RuntimeError("torch.float8_e4m3fn is required for the SANA-WM FP8 backend.") def _require_fp4_dtype() -> None: if not hasattr(torch, "float4_e2m1fn_x2"): raise RuntimeError( - "torch.float4_e2m1fn_x2 is required for the native SANA FP4 backend." + "torch.float4_e2m1fn_x2 is required for the SANA-WM FP4 backend." ) if not hasattr(torch, "float8_e4m3fn"): raise RuntimeError( - "torch.float8_e4m3fn scales are required for the native SANA FP4 backend." + "torch.float8_e4m3fn scales are required for the SANA-WM FP4 backend." ) @@ -558,7 +545,7 @@ def _require_fp4_dtype() -> None: "TorchScaledMMFP8Linear", "TorchScaledMMFP8Recipe", "quantize_nvfp4_swizzled", - "replace_linear_with_native_quant", + "replace_linear_with_quant", "replace_linear_with_torch_fp4", "replace_linear_with_torch_fp8", ] diff --git a/integrations/sana/sana_wm/native_refiner.py b/integrations/sana/sana_wm/refiner.py similarity index 96% rename from integrations/sana/sana_wm/native_refiner.py rename to integrations/sana/sana_wm/refiner.py index 4ea247c9a..c29844d93 100644 --- a/integrations/sana/sana_wm/native_refiner.py +++ b/integrations/sana/sana_wm/refiner.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Native LTX-2 refiner used by the SANA-WM integration.""" +"""LTX-2 refiner used by the SANA-WM integration.""" from __future__ import annotations @@ -27,14 +27,14 @@ from loguru import logger from torch import Tensor -from sana_wm.native_quant import ( +from sana_wm.quant import ( TorchScaledMMFP4Recipe, TorchScaledMMFP8Recipe, - replace_linear_with_native_quant, + replace_linear_with_quant, ) Precision = Literal["bf16", "fp8", "fp4"] -QuantBackend = Literal["auto", "upstream-te", "torch-fp8", "torch-fp4", "native"] +QuantBackend = Literal["auto", "torch", "torch-fp8", "torch-fp4"] _REFINER_QUANT_SKIP_DEFAULTS = ( r"^proj_in$", @@ -59,7 +59,7 @@ def __init__( dtype: torch.dtype, device: torch.device | str, precision: Precision = "bf16", - quant_backend: QuantBackend = "native", + quant_backend: QuantBackend = "torch", text_max_sequence_length: int = 1024, ) -> None: super().__init__() @@ -178,17 +178,12 @@ def _load_diffusers_components(self) -> tuple[nn.Module, nn.Module]: def _prepare_quantization(self) -> None: if self._quantized or self.precision == "bf16": return - if self.quant_backend == "upstream-te": - raise ValueError( - "Native SANA-WM refiner execution does not use Transformer Engine; " - "select --quant-backend native, torch-fp8, or torch-fp4." - ) recipe = ( TorchScaledMMFP8Recipe() if self.precision == "fp8" else TorchScaledMMFP4Recipe() ) - converted, skipped = replace_linear_with_native_quant( + converted, skipped = replace_linear_with_quant( self.transformer, recipe=recipe, params_dtype=self.dtype, @@ -196,7 +191,7 @@ def _prepare_quantization(self) -> None: ) if converted <= 0: raise RuntimeError( - f"SANA-WM native refiner {self.precision} converted no Linear " + f"SANA-WM refiner {self.precision} converted no Linear " f"layers; skipped={skipped}." ) self._quantized = True diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index f0266c667..02d1163d0 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -13,18 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SANA-WM bidirectional runner and upstream-reference harness.""" +"""SANA-WM bidirectional runner.""" from __future__ import annotations -import importlib import os -import sys from collections.abc import Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path -from types import ModuleType from typing import Literal import numpy as np @@ -54,30 +51,20 @@ SANA_WM_REFINER_ROOT, SANA_WM_VAE_TEMPORAL_COMPRESSION, ) -from sana_wm.native_transformer import SanaWMNativeTransformer - -SamplingAlgo = Literal[ - "auto", - "flow_euler_ltx", - "flow_euler", - "flow_dpm-solver", - "chunk_flow_euler", - "self_forcing", -] -"""Sampling algorithms exposed by upstream SANA-WM inference.""" +from sana_wm.transformer import SanaWMTransformer + +SamplingAlgo = Literal["auto", "flow_euler_ltx"] +"""Sampling algorithms exposed by the SANA-WM runner.""" Precision = Literal["bf16", "fp8", "fp4"] """SANA-WM Stage-1/refiner precision modes.""" -QuantBackend = Literal["auto", "upstream-te", "torch-fp8", "torch-fp4", "native"] +QuantBackend = Literal["auto", "torch", "torch-fp8", "torch-fp4"] """Low-precision linear backend for FP8/FP4 SANA-WM paths.""" -ResolvedQuantBackend = Literal["upstream-te", "torch-fp8", "torch-fp4", "native"] +ResolvedQuantBackend = Literal["torch", "torch-fp8", "torch-fp4"] """Concrete low-precision linear backend selected after resolving ``auto``.""" -ExecutionBackend = Literal["upstream-reference", "native-flashdreams"] -"""Execution path selected by a SANA-WM runner config.""" - @dataclass(kw_only=True) class SanaWMRunnerConfig(RunnerConfig): @@ -85,12 +72,8 @@ class SanaWMRunnerConfig(RunnerConfig): _target: type["SanaWMRunner"] = field(default_factory=lambda: SanaWMRunner) - upstream_sana_root: Path | None = None - """Path to a cloned NVlabs/Sana checkout. ``None`` auto-detects - ``$SANA_ROOT`` / ``../Sana`` or imports an installed Sana package.""" - device: str = "auto" - """Torch device for upstream SANA-WM. ``"auto"`` picks CUDA when available.""" + """Torch device for SANA-WM. ``"auto"`` picks CUDA when available.""" image_path: Path | None = None """Path to the first-frame RGB image. Required at ``run()`` time.""" @@ -106,8 +89,7 @@ class SanaWMRunnerConfig(RunnerConfig): intrinsics_path: Path | None = None """Optional ``.npy`` intrinsics shaped ``[3, 3]``, ``[F, 3, 3]``, - ``[4]``, or ``[F, 4]``. When absent, upstream estimates intrinsics - with Pi3X.""" + ``[4]``, or ``[F, 4]``.""" action: str | None = DEFAULT_ACTION """Action DSL used when ``camera_path`` is not provided.""" @@ -119,7 +101,7 @@ class SanaWMRunnerConfig(RunnerConfig): """Per-frame action rotation speed in degrees.""" name: str = "sana_wm" - """Output filename stem passed to upstream ``write_video``.""" + """Output filename stem.""" num_frames: int = 161 """Requested output frames before LTX2-VAE stride snapping.""" @@ -137,22 +119,7 @@ class SanaWMRunnerConfig(RunnerConfig): """Optional scheduler flow-shift override.""" sampling_algo: SamplingAlgo = "auto" - """Stage-1 sampler. ``"auto"`` follows upstream config defaults.""" - - chunk_interval_k: float | None = None - """ChunkFlowEuler interval ratio override.""" - - num_cached_blocks: int = 2 - """Self-forcing sampler cached-block count.""" - - sink_token: bool = False - """Whether self-forcing keeps chunk 0 as a permanent sink.""" - - num_frame_per_block: int = 3 - """Self-forcing latent frames per AR block.""" - - denoising_step_list: str = "" - """Comma-separated self-forcing timesteps ending in ``0``.""" + """Stage-1 sampler. ``"auto"`` uses ``"flow_euler_ltx"``.""" save_stage1: bool = False """Also decode the unrefined Stage-1 latent when the refiner is enabled.""" @@ -163,11 +130,8 @@ class SanaWMRunnerConfig(RunnerConfig): seed: int = 42 """Stage-1 random seed.""" - no_action_overlay: bool = True - """Skip upstream action-overlay compositing on generated videos.""" - config_path: str = SANA_WM_CONFIG_PATH - """Upstream inference YAML path or ``hf://`` URI.""" + """SANA-WM inference YAML path or ``hf://`` URI.""" model_path: str = SANA_WM_MODEL_PATH """Stage-1 checkpoint path or ``hf://`` URI.""" @@ -185,15 +149,9 @@ class SanaWMRunnerConfig(RunnerConfig): Ignored when ``no_refiner`` is ``True``.""" quant_backend: QuantBackend = "auto" - """Backend for quantized linear layers. ``"auto"`` uses the native - FlashDreams/PyTorch low-precision path. ``"upstream-te"`` is accepted only - by the upstream reference harness. ``"torch-fp8"`` and ``"torch-fp4"`` - select one native ``torch._scaled_mm`` replacement explicitly; ``"native"`` - allows both native FP8 and native FP4.""" - - execution_backend: ExecutionBackend = "upstream-reference" - """Whether this config runs the upstream reference harness or the native - FlashDreams path.""" + """Backend for quantized linear layers. ``"auto"`` and ``"torch"`` allow + both FP8 and FP4 Torch replacements. ``"torch-fp8"`` and ``"torch-fp4"`` + select one ``torch._scaled_mm`` replacement explicitly.""" refiner_root: str = SANA_WM_REFINER_ROOT """LTX-2 refiner root path or ``hf://`` URI.""" @@ -224,7 +182,7 @@ class SanaWMRunnerConfig(RunnerConfig): class SanaWMRunner(Runner[SanaWMRunnerConfig, object]): - """CLI driver for SANA-WM native and upstream-reference configs.""" + """CLI driver for SANA-WM configs.""" config: SanaWMRunnerConfig @@ -254,7 +212,7 @@ def _resolve_prompt(self) -> str: return prompt def _resolve_device(self) -> torch.device: - """Return the device used by upstream SANA-WM.""" + """Return the device used by SANA-WM.""" if self.config.device == "auto": if torch.cuda.is_available(): return torch.device(f"cuda:{self.local_rank}") @@ -280,55 +238,27 @@ def _resolve_trajectory(self) -> np.ndarray: rotation_speed_deg=self.config.rotation_speed_deg, ) - def _denoising_steps(self) -> list[int] | None: - """Parse optional self-forcing denoising timesteps.""" - if not self.config.denoising_step_list: - return None - steps = [ - int(item.strip()) - for item in self.config.denoising_step_list.split(",") - if item.strip() - ] - if not steps or steps[-1] != 0: - raise ValueError("--denoising-step-list must end with 0.") - return steps - def run(self) -> None: """Run SANA-WM bidirectional inference and write outputs.""" - if self.config.execution_backend == "native-flashdreams": - self._run_native() - return - self._run_upstream_reference() - - def _run_native(self) -> None: - """Run the native FlashDreams SANA-WM pipeline.""" cfg = self.config if cfg.image_path is None: raise ValueError("SanaWMRunner requires --image-path.") - if not cfg.no_action_overlay: - raise ValueError( - "Native SANA-WM does not include the upstream action overlay. " - "Use --no-action-overlay True." - ) device = self._resolve_device() - quant_backend: QuantBackend = ( - "native" if cfg.quant_backend == "auto" else cfg.quant_backend + quant_backend = _resolve_quant_backend( + cfg.quant_backend, + _active_quantized_precisions( + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, + ), ) _validate_precision_request( device=device, stage1_precision=cfg.stage1_precision, refiner_precision=cfg.refiner_precision, refiner_enabled=not cfg.no_refiner, - quant_backend=quant_backend, - ) - resolved_quant_backend = _resolve_quant_backend( - quant_backend, - _active_quantized_precisions( - stage1_precision=cfg.stage1_precision, - refiner_precision=cfg.refiner_precision, - refiner_enabled=not cfg.no_refiner, - ), + quant_backend=cfg.quant_backend, ) precision_env = _precision_env_updates( stage1_precision=cfg.stage1_precision, @@ -337,24 +267,22 @@ def _run_native(self) -> None: ) with _temporary_environment(precision_env): prompt = self._resolve_prompt() - image, c2w, intrinsics_vec4, num_frames = self._prepare_native_inputs( - device - ) - pipeline_cfg = _native_pipeline_config( + image, c2w, intrinsics_vec4, num_frames = self._prepare_inputs() + pipeline_cfg = _pipeline_config( cfg, quant_backend=quant_backend, ) pipeline = pipeline_cfg.setup().to(device).eval() transformer = pipeline.diffusion_model.transformer - if not isinstance(transformer, SanaWMNativeTransformer): + if not isinstance(transformer, SanaWMTransformer): raise TypeError( - "Native SANA-WM runner resolved a non-SANA transformer: " + "SANA-WM runner resolved a non-SANA transformer: " f"{type(transformer).__name__}." ) - sampling_algo = self._native_sampling_algo() + sampling_algo = self._sampling_algo() if sampling_algo != "flow_euler_ltx": raise ValueError( - "Native SANA-WM requires flow_euler_ltx for the " + "SANA-WM requires flow_euler_ltx for the " f"bidirectional runner; got {sampling_algo!r}." ) camera = prepare_camera( @@ -397,7 +325,7 @@ def _run_native(self) -> None: transformer.release_refiner_runtime() elif cfg.save_stage1: logger.info( - "Native SANA-WM is already running without the refiner; " + "SANA-WM is already running without the refiner; " "--save-stage1 does not create an extra output." ) @@ -415,182 +343,8 @@ def _run_native(self) -> None: cfg.fps, ) - def _run_upstream_reference(self) -> None: - """Run the explicit upstream SANA-WM reference harness.""" - cfg = self.config - if cfg.image_path is None: - raise ValueError("SanaWMRunner requires --image-path.") - - device = self._resolve_device() - _validate_precision_request( - device=device, - stage1_precision=cfg.stage1_precision, - refiner_precision=cfg.refiner_precision, - refiner_enabled=not cfg.no_refiner, - quant_backend=cfg.quant_backend, - ) - resolved_quant_backend = _resolve_quant_backend( - cfg.quant_backend, - _active_quantized_precisions( - stage1_precision=cfg.stage1_precision, - refiner_precision=cfg.refiner_precision, - refiner_enabled=not cfg.no_refiner, - ), - ) - if ( - _active_quantized_precisions( - stage1_precision=cfg.stage1_precision, - refiner_precision=cfg.refiner_precision, - refiner_enabled=not cfg.no_refiner, - ) - and resolved_quant_backend != "upstream-te" - ): - raise ValueError( - "The upstream reference runner uses upstream SANA's own " - "quantized path. Pass --quant-backend upstream-te for " - "FP8/FP4 reference comparisons, or use sana-wm-bidirectional " - "for the native FlashDreams low-precision backend." - ) - - precision_env = _precision_env_updates( - stage1_precision=cfg.stage1_precision, - refiner_precision=cfg.refiner_precision, - refiner_enabled=not cfg.no_refiner, - ) - with _temporary_environment(precision_env): - upstream = _import_upstream_sana(cfg.upstream_sana_root) - prompt = self._resolve_prompt() - - image = upstream.Image.open(cfg.image_path).convert("RGB") - c2w_full = self._resolve_trajectory() - num_frames = min(cfg.num_frames, c2w_full.shape[0]) - snapped = upstream._snap_num_frames( - num_frames, - stride=SANA_WM_VAE_TEMPORAL_COMPRESSION, - upper_bound=c2w_full.shape[0], - ) - if snapped != cfg.num_frames and self.is_rank_zero: - logger.warning( - "SANA-WM requires num_frames = 8k+1; requested {} snapped to {} " - "(trajectory has {} frames).", - cfg.num_frames, - snapped, - c2w_full.shape[0], - ) - num_frames = snapped - c2w = c2w_full[:num_frames] - - cropped, src_size, resized_size, crop_offset = ( - upstream.resize_and_center_crop( - image, - target_h=DEFAULT_VIDEO_HEIGHT, - target_w=DEFAULT_VIDEO_WIDTH, - ) - ) - if cfg.intrinsics_path is None: - intrinsics_src = np.broadcast_to( - upstream.estimate_intrinsics_with_pi3x( - image, device, upstream.get_root_logger() - ), - (num_frames, 4), - ).copy() - else: - intrinsics_src = upstream.load_intrinsics( - cfg.intrinsics_path, num_frames - ) - intrinsics_vec4 = upstream.transform_intrinsics_for_crop( - intrinsics_src, src_size, resized_size, crop_offset - ) - - inference_config = upstream.pyrallis.parse( - config_class=upstream.InferenceConfig, - config_path=upstream.resolve_hf_path(cfg.config_path), - args=[], - ) - refiner = ( - None - if cfg.no_refiner - else upstream.RefinerSettings( - root=cfg.refiner_root, - gemma_root=cfg.refiner_gemma_root, - sink_size=cfg.sink_size, - seed=cfg.refiner_seed, - block_size=cfg.refiner_block_size, - kv_max_frames=cfg.refiner_kv_max_frames, - ) - ) - pipeline = upstream.SanaWMPipeline( - config=inference_config, - model_path=upstream.resolve_hf_path(cfg.model_path), - device=device, - refiner=refiner, - offload_vae=cfg.offload_vae, - offload_refiner=cfg.offload_refiner, - offload_text_encoder=cfg.offload_text_encoder, - logger=upstream.get_root_logger(), - ) - - sampling_algo = cfg.sampling_algo - if sampling_algo == "auto": - sampling_algo = ( - inference_config.scheduler.vis_sampler - if inference_config.scheduler.vis_sampler - in {"chunk_flow_euler", "self_forcing"} - else "flow_euler_ltx" - ) - params = upstream.GenerationParams( - num_frames=num_frames, - fps=cfg.fps, - step=cfg.step, - cfg_scale=cfg.cfg_scale, - flow_shift=cfg.flow_shift, - seed=cfg.seed, - negative_prompt=cfg.negative_prompt, - sampling_algo=sampling_algo, - chunk_interval_k=cfg.chunk_interval_k, - num_cached_blocks=cfg.num_cached_blocks, - sink_token=cfg.sink_token, - num_frame_per_block=cfg.num_frame_per_block, - denoising_step_list=self._denoising_steps(), - save_stage1=cfg.save_stage1, - ) - - result = pipeline.generate(cropped, prompt, c2w, intrinsics_vec4, params) - video_hwc = result["video"] - if not cfg.no_action_overlay: - video_hwc = upstream.apply_overlay(video_hwc, result["c2w"]) - - if not self.is_rank_zero: - return - upstream.write_video( - cfg.output_dir, - cfg.name, - video_hwc, - params.fps, - upstream.get_root_logger(), - ) - - stage1_video = result.get("stage1_video") - if stage1_video is not None: - stage1_hwc = stage1_video - if not cfg.no_action_overlay: - stage1_hwc = upstream.apply_overlay( - stage1_hwc, result["stage1_c2w"] - ) - upstream.write_video( - cfg.output_dir, - f"{cfg.name}_stage1", - stage1_hwc, - params.fps, - upstream.get_root_logger(), - ) - - def _prepare_native_inputs( - self, - device: torch.device, - ) -> tuple[object, np.ndarray, np.ndarray, int]: - """Load/crop the input image and prepare c2w/intrinsics for native run.""" - del device + def _prepare_inputs(self) -> tuple[object, np.ndarray, np.ndarray, int]: + """Load/crop the input image and prepare c2w/intrinsics.""" from PIL import Image cfg = self.config @@ -630,10 +384,7 @@ def _prepare_native_inputs( ) ) if cfg.intrinsics_path is None: - raise ValueError( - "The native SANA-WM runner currently requires --intrinsics-path. " - "Use sana-wm-upstream-reference if you need Pi3X intrinsics estimation." - ) + raise ValueError("SanaWMRunner requires --intrinsics-path.") intrinsics_src = load_intrinsics(cfg.intrinsics_path, num_frames) intrinsics_vec4 = transform_intrinsics_for_crop( intrinsics_src, @@ -643,19 +394,19 @@ def _prepare_native_inputs( ) return cropped, c2w, intrinsics_vec4, num_frames - def _native_sampling_algo(self) -> SamplingAlgo: - """Resolve the native sampler selection.""" + def _sampling_algo(self) -> SamplingAlgo: + """Resolve the sampler selection.""" if self.config.sampling_algo == "auto": return "flow_euler_ltx" return self.config.sampling_algo -def _native_pipeline_config( +def _pipeline_config( cfg: SanaWMRunnerConfig, *, - quant_backend: QuantBackend, + quant_backend: ResolvedQuantBackend, ) -> StreamInferencePipelineConfig: - """Apply CLI runtime fields to the native SANA-WM pipeline literal.""" + """Apply CLI runtime fields to the SANA-WM pipeline literal.""" return derive_config( cfg.pipeline, diffusion_model=dict( @@ -788,19 +539,8 @@ def _validate_precision_request( _validate_torch_fp8_backend(quantized) elif resolved_backend == "torch-fp4": _validate_torch_fp4_backend(quantized) - elif resolved_backend == "native": - _validate_native_quant_backend(quantized) - elif resolved_backend == "upstream-te": - if "fp8" in quantized and not _torch_cuda_version_at_least(12, 9): - cuda_version = torch.version.cuda or "unknown" - raise ValueError( - "SANA-WM fp8 precision uses upstream Sana's Transformer Engine " - "Float8BlockScaling path, which requires CUDA 12.9 or newer. " - f"This PyTorch environment reports CUDA {cuda_version}. Use " - "--quant-backend torch-fp8 or --quant-backend native to try " - "the native PyTorch FP8 backend." - ) - _validate_transformer_engine(quantized) + elif resolved_backend == "torch": + _validate_quant_backend(quantized) else: raise ValueError(f"Unsupported SANA-WM quant backend: {quant_backend!r}.") @@ -822,62 +562,48 @@ def _resolve_quant_backend( quant_backend: QuantBackend, quantized_precisions: list[Precision], ) -> ResolvedQuantBackend: - """Resolve ``auto`` to the native low-precision backend.""" + """Resolve ``auto`` to the Torch low-precision backend.""" del quantized_precisions if quant_backend == "auto": - return "native" - if quant_backend in {"upstream-te", "torch-fp8", "torch-fp4", "native"}: + return "torch" + if quant_backend in {"torch", "torch-fp8", "torch-fp4"}: return quant_backend raise ValueError(f"Unsupported SANA-WM quant backend: {quant_backend!r}.") -def _torch_cuda_version_at_least(major: int, minor: int) -> bool: - """Return whether ``torch.version.cuda`` is at least ``major.minor``.""" - version = torch.version.cuda - if version is None: - return False - parts = version.split(".") - try: - current_major = int(parts[0]) - current_minor = int(parts[1]) if len(parts) > 1 else 0 - except (ValueError, IndexError): - return False - return (current_major, current_minor) >= (major, minor) - - def _validate_torch_fp8_backend(precisions: list[Precision]) -> None: - """Validate the native PyTorch scaled-MM backend request.""" + """Validate the PyTorch scaled-MM backend request.""" unsupported = sorted({precision for precision in precisions if precision != "fp8"}) if unsupported: raise ValueError( "SANA-WM --quant-backend torch-fp8 accepts fp8 only; " f"unsupported requested precision(s): {', '.join(unsupported)}. " - "Use --quant-backend torch-fp4 or --quant-backend native for fp4." + "Use --quant-backend torch-fp4 or --quant-backend torch for fp4." ) - _validate_native_fp8_primitives() + _validate_torch_fp8_primitives() def _validate_torch_fp4_backend(precisions: list[Precision]) -> None: - """Validate the native PyTorch NVFP4 backend request.""" + """Validate the PyTorch NVFP4 backend request.""" unsupported = sorted({precision for precision in precisions if precision != "fp4"}) if unsupported: raise ValueError( "SANA-WM --quant-backend torch-fp4 accepts fp4 only; " f"unsupported requested precision(s): {', '.join(unsupported)}. " - "Use --quant-backend torch-fp8 or --quant-backend native for fp8." + "Use --quant-backend torch-fp8 or --quant-backend torch for fp8." ) - _validate_native_fp4_primitives() + _validate_torch_fp4_primitives() -def _validate_native_quant_backend(precisions: list[Precision]) -> None: - """Validate native PyTorch low-precision primitives.""" +def _validate_quant_backend(precisions: list[Precision]) -> None: + """Validate PyTorch low-precision primitives.""" if "fp8" in precisions: - _validate_native_fp8_primitives() + _validate_torch_fp8_primitives() if "fp4" in precisions: - _validate_native_fp4_primitives() + _validate_torch_fp4_primitives() -def _validate_native_fp8_primitives() -> None: +def _validate_torch_fp8_primitives() -> None: if not hasattr(torch, "_scaled_mm") or not hasattr(torch, "float8_e4m3fn"): raise RuntimeError( "SANA-WM --quant-backend torch-fp8 requires PyTorch with " @@ -885,7 +611,7 @@ def _validate_native_fp8_primitives() -> None: ) -def _validate_native_fp4_primitives() -> None: +def _validate_torch_fp4_primitives() -> None: missing = [ name for name in ("_scaled_mm", "float4_e2m1fn_x2", "float8_e4m3fn") @@ -893,92 +619,12 @@ def _validate_native_fp4_primitives() -> None: ] if missing: raise RuntimeError( - "SANA-WM native fp4 requires PyTorch with " + "SANA-WM fp4 requires PyTorch with " f"{', '.join(f'torch.{name}' for name in missing)} support." ) -def _validate_transformer_engine(precisions: list[Precision]) -> None: - """Ensure Transformer Engine has the recipes required by precision modes.""" - required_recipes: set[str] = set() - if "fp8" in precisions: - required_recipes.add("Float8BlockScaling") - if "fp4" in precisions: - required_recipes.add("NVFP4BlockScaling") - try: - import transformer_engine.common.recipe as te_recipe - import transformer_engine.pytorch # noqa: F401 - except Exception as exc: - raise RuntimeError( - "SANA-WM fp8/fp4 precision requires NVIDIA Transformer Engine. " - "Install it in the upstream Sana environment, for example with " - "`pip install --no-build-isolation 'transformer_engine[pytorch]'`, " - "or run with the default bf16 precision." - ) from exc - - missing = sorted( - recipe for recipe in required_recipes if not hasattr(te_recipe, recipe) - ) - if missing: - raise RuntimeError( - "Installed Transformer Engine does not provide the recipe(s) " - f"required by the requested SANA-WM precision: {', '.join(missing)}." - ) - - -def _candidate_upstream_roots(explicit_root: Path | None) -> list[Path]: - """Return candidate NVlabs/Sana checkout roots in priority order.""" - roots: list[Path] = [] - if explicit_root is not None: - roots.append(explicit_root) - for env_name in ("SANA_ROOT", "SANA_REPO", "SANA_WM_UPSTREAM_ROOT"): - value = os.environ.get(env_name) - if value: - roots.append(Path(value)) - repo_root = Path(__file__).resolve().parents[3] - roots.append(repo_root.parent / "Sana") - roots.append(Path.cwd().parent / "Sana") - return roots - - -def _is_upstream_sana_root(path: Path) -> bool: - """Return ``True`` when ``path`` looks like an NVlabs/Sana checkout.""" - return ( - path.expanduser() / "inference_video_scripts" / "wm" / "inference_sana_wm.py" - ).is_file() - - -def _import_upstream_sana(explicit_root: Path | None) -> ModuleType: - """Import upstream SANA-WM inference after adding a checkout to ``sys.path``.""" - for root in _candidate_upstream_roots(explicit_root): - candidate = root.expanduser().resolve() - if _is_upstream_sana_root(candidate): - candidate_str = str(candidate) - if candidate_str not in sys.path: - sys.path.insert(0, candidate_str) - break - - try: - return importlib.import_module("inference_video_scripts.wm.inference_sana_wm") - except ModuleNotFoundError as exc: - roots = ", ".join( - str(path) for path in _candidate_upstream_roots(explicit_root) - ) - raise ModuleNotFoundError( - "Could not import upstream SANA-WM inference. Clone NVlabs/Sana " - "next to this repo, pass --upstream-sana-root, or set SANA_ROOT; " - "then install the upstream Sana inference dependencies. " - f"Checked roots: {roots}. Original import error: {exc}" - ) from exc - except Exception as exc: - raise RuntimeError( - "Importing upstream SANA-WM inference failed. Ensure the upstream " - "Sana checkout dependencies are installed in the active environment." - ) from exc - - __all__ = [ - "ExecutionBackend", "Precision", "QuantBackend", "SamplingAlgo", diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index f657c0b0e..87cbfe21f 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -13,10 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""FlashDreams-owned SANA-WM Stage-1 DiT module definitions. +"""SANA-WM Stage-1 DiT module definitions. -This file intentionally models the public SANA-WM bidirectional Stage-1 -checkpoint schema without importing the upstream SANA ``diffusion`` package. The module names and tensor shapes are checkpoint-facing API: keep them stable unless the public checkpoint contract changes. """ @@ -309,7 +307,7 @@ def forward( apply_output_gate: bool = True, **kwargs: object, ) -> Tensor: - """Run native SANA-WM self/camera attention.""" + """Run SANA-WM self/camera attention.""" if HW is None: raise ValueError("SANA-WM Stage-1 attention requires HW=(T, H, W).") batch, tokens, channels = x.shape @@ -623,7 +621,7 @@ def forward( plucker_emb: Tensor | None = None, **kwargs: object, ) -> Tensor: - """Run one native Stage-1 transformer block.""" + """Run one Stage-1 transformer block.""" batch, tokens, channels = x.shape shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.scale_shift_table[None, None, :, :] + t.reshape(batch, frames, 6, -1) @@ -653,7 +651,7 @@ def forward( class SanaWMStage1Model(nn.Module): - """Checkpoint-compatible native Stage-1 SANA-WM DiT shell.""" + """Checkpoint-compatible Stage-1 SANA-WM DiT shell.""" def __init__(self, spec: SanaWMStage1Spec = SANA_WM_STAGE1_SPEC) -> None: super().__init__() @@ -682,7 +680,7 @@ def forward( chunk_plucker: Tensor | None = None, **kwargs: object, ) -> Tensor: - """Run the native SANA-WM Stage-1 DiT.""" + """Run the SANA-WM Stage-1 DiT.""" batch, _channels, frames, height, width = x.shape x = self.x_embedder(x) x = x.permute(0, 2, 3, 4, 1).reshape(batch, frames * height * width, -1) diff --git a/integrations/sana/sana_wm/native_transformer.py b/integrations/sana/sana_wm/transformer.py similarity index 92% rename from integrations/sana/sana_wm/native_transformer.py rename to integrations/sana/sana_wm/transformer.py index bebf59e7f..56b8482fc 100644 --- a/integrations/sana/sana_wm/native_transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Native FlashDreams adapter for the SANA-WM Stage-1 DiT.""" +"""FlashDreams adapter for the SANA-WM Stage-1 DiT.""" from __future__ import annotations @@ -46,15 +46,15 @@ SANA_WM_REFINER_GEMMA_ROOT, SANA_WM_REFINER_ROOT, ) -from sana_wm.native_quant import ( +from sana_wm.quant import ( TorchScaledMMFP4Recipe, TorchScaledMMFP8Recipe, - replace_linear_with_native_quant, + replace_linear_with_quant, ) from sana_wm.stage1_model import SanaWMStage1Model Precision = Literal["bf16", "fp8", "fp4"] -QuantBackend = Literal["auto", "upstream-te", "torch-fp8", "torch-fp4", "native"] +QuantBackend = Literal["auto", "torch", "torch-fp8", "torch-fp4"] _STAGE1_QUANT_SKIP_DEFAULTS = ( "^x_embedder", @@ -95,18 +95,18 @@ class SanaWMStage1Conditioning: @dataclass(kw_only=True) -class SanaWMNativeTransformerCache(TransformerAutoregressiveCache): +class SanaWMTransformerCache(TransformerAutoregressiveCache): """AR cache for the one-shot SANA-WM Stage-1 rollout.""" conditioning: SanaWMStage1Conditioning | None = None @dataclass(kw_only=True) -class SanaWMNativeTransformerConfig(TransformerConfig): - """Config for the native SANA-WM Stage-1 transformer adapter.""" +class SanaWMTransformerConfig(TransformerConfig): + """Config for the SANA-WM Stage-1 transformer adapter.""" - _target: type["SanaWMNativeTransformer"] = field( - default_factory=lambda: SanaWMNativeTransformer + _target: type["SanaWMTransformer"] = field( + default_factory=lambda: SanaWMTransformer ) config_path: str = SANA_WM_CONFIG_PATH @@ -146,22 +146,22 @@ class SanaWMNativeTransformerConfig(TransformerConfig): """Pixel width used by the public SANA-WM bidirectional release.""" vae_tile_sample_min_width: int = 256 - """Pixel width tile size used for native LTX-2 VAE decode.""" + """Pixel width tile size used for LTX-2 VAE decode.""" vae_tile_sample_stride_width: int = 224 - """Pixel width stride used for native LTX-2 VAE decode.""" + """Pixel width stride used for LTX-2 VAE decode.""" vae_tile_sample_min_height: int = 256 - """Pixel height tile size used for native LTX-2 VAE decode.""" + """Pixel height tile size used for LTX-2 VAE decode.""" vae_tile_sample_stride_height: int = 192 - """Pixel height stride used for native LTX-2 VAE decode.""" + """Pixel height stride used for LTX-2 VAE decode.""" vae_tile_sample_min_num_frames: int = 24 - """Pixel-frame temporal tile size used for native LTX-2 VAE decode.""" + """Pixel-frame temporal tile size used for LTX-2 VAE decode.""" vae_tile_sample_stride_num_frames: int = 8 - """Pixel-frame temporal tile stride used for native LTX-2 VAE decode.""" + """Pixel-frame temporal tile stride used for LTX-2 VAE decode.""" vae_oom_retry_tile_sample_min_width: int = 128 """Smaller pixel width tile size for one VAE decode OOM retry.""" @@ -182,12 +182,12 @@ class SanaWMNativeTransformerConfig(TransformerConfig): """Smaller temporal tile stride for one VAE decode OOM retry.""" -class SanaWMNativeTransformer(Transformer[SanaWMNativeTransformerCache]): - """Native FlashDreams adapter for the SANA-WM Stage-1 model call.""" +class SanaWMTransformer(Transformer[SanaWMTransformerCache]): + """FlashDreams adapter for the SANA-WM Stage-1 model call.""" - def __init__(self, config: SanaWMNativeTransformerConfig) -> None: + def __init__(self, config: SanaWMTransformerConfig) -> None: super().__init__(config) - self.config: SanaWMNativeTransformerConfig = config + self.config: SanaWMTransformerConfig = config self._dummy = nn.Parameter(torch.empty(0)) self._runtime_config: Any | None = None self._model_path: str | None = None @@ -221,15 +221,15 @@ def initialize_autoregressive_cache( *, conditioning: SanaWMStage1Conditioning | None = None, **_: Any, - ) -> SanaWMNativeTransformerCache: + ) -> SanaWMTransformerCache: """Build a cache containing per-rollout SANA-WM conditioning.""" - return SanaWMNativeTransformerCache(conditioning=conditioning) + return SanaWMTransformerCache(conditioning=conditioning) def predict_flow( self, noisy_latent: Tensor, timestep: Tensor, - cache: SanaWMNativeTransformerCache, + cache: SanaWMTransformerCache, input: Any = None, model_kwargs: dict[str, object] | None = None, ) -> Tensor: @@ -249,7 +249,7 @@ def finalize_kv_cache( self, noisy_latent: Tensor, timestep: Tensor, - cache: SanaWMNativeTransformerCache, + cache: SanaWMTransformerCache, input: Any = None, ) -> None: """SANA-WM bidirectional inference has no streaming KV cache to advance.""" @@ -450,7 +450,7 @@ def decode_latents(self, latents: Tensor) -> np.ndarray: def release_stage1_runtime( self, - cache: SanaWMNativeTransformerCache | None = None, + cache: SanaWMTransformerCache | None = None, ) -> None: """Release Stage-1-only tensors before VAE/refiner work.""" free_before_gib: float | None = None @@ -500,7 +500,7 @@ def release_stage1_runtime( gc.collect() def release_refiner_runtime(self) -> None: - """Release native refiner tensors before VAE decode.""" + """Release refiner tensors before VAE decode.""" self._release_refiner() def refine_latents( @@ -514,7 +514,7 @@ def refine_latents( block_size: int | None, kv_max_frames: int, ) -> Tensor: - """Run the native LTX-2 refiner.""" + """Run the LTX-2 refiner.""" self._ensure_refiner() sigmas = torch.tensor( self._stage2_sigmas(), @@ -590,7 +590,7 @@ def _ensure_model(self) -> None: weight_dtype = self._ensure_weight_dtype() model = SanaWMStage1Model().to(self.device) logger.info( - "[Sana] Loaded native {} ({:,} params)", + "[Sana] Loaded {} ({:,} params)", cfg.model.model, sum(p.numel() for p in model.parameters()), ) @@ -617,7 +617,7 @@ def _ensure_model(self) -> None: def _ensure_refiner(self) -> None: if self._refiner_built: return - from sana_wm.native_refiner import SanaWMLTX2Refiner + from sana_wm.refiner import SanaWMLTX2Refiner compute_dtype = ( torch.bfloat16 @@ -625,7 +625,7 @@ def _ensure_refiner(self) -> None: else _get_weight_dtype(self.config.refiner_precision) ) quant_backend: QuantBackend = ( - "native" if self.config.quant_backend == "auto" else self.config.quant_backend + "torch" if self.config.quant_backend == "auto" else self.config.quant_backend ) self.refiner = SanaWMLTX2Refiner( refiner_root=resolve_hf_path(self.config.refiner_root), @@ -816,16 +816,11 @@ def pad_pair(text: Tensor, mask: Tensor) -> tuple[Tensor, Tensor]: def _prepare_stage1_quant(self) -> None: if self._stage1_quantized or self.config.stage1_precision == "bf16": return - if self.config.quant_backend == "upstream-te": - raise ValueError( - "Native FlashDreams SANA execution does not use Transformer Engine; " - "select --quant-backend native, torch-fp8, or torch-fp4." - ) if self.config.stage1_precision == "fp8": recipe = TorchScaledMMFP8Recipe() else: recipe = TorchScaledMMFP4Recipe() - converted, skipped = replace_linear_with_native_quant( + converted, skipped = replace_linear_with_quant( self.model, recipe=recipe, params_dtype=self._ensure_weight_dtype(), @@ -834,12 +829,12 @@ def _prepare_stage1_quant(self) -> None: ) if converted <= 0: raise RuntimeError( - f"SANA-WM native {self.config.stage1_precision} converted no " + f"SANA-WM {self.config.stage1_precision} converted no " f"Stage-1 Linear layers; skipped={skipped}." ) self._stage1_quantized = True logger.info( - "[stage1-native-quant] precision={} converted {} Linear layers (skipped {})", + "[stage1-quant] precision={} converted {} Linear layers (skipped {})", self.config.stage1_precision, converted, skipped, @@ -882,10 +877,10 @@ def _stage2_sigmas() -> tuple[float, ...]: def _require_conditioning( - cache: SanaWMNativeTransformerCache, + cache: SanaWMTransformerCache, ) -> SanaWMStage1Conditioning: if cache.conditioning is None: - raise RuntimeError("SANA-WM native cache was initialized without conditioning.") + raise RuntimeError("SANA-WM cache was initialized without conditioning.") return cache.conditioning @@ -940,7 +935,7 @@ def _get_vae(*args: Any, **kwargs: Any) -> nn.Module: device = kwargs["device"] dtype = kwargs["dtype"] if "LTX2VAE_diffusers" not in str(name): - raise ValueError(f"Unsupported native SANA-WM VAE type: {name!r}") + raise ValueError(f"Unsupported SANA-WM VAE type: {name!r}") from diffusers import AutoencoderKLLTX2Video return ( @@ -959,7 +954,7 @@ def _get_tokenizer_and_text_encoder(*args: Any, **kwargs: Any) -> tuple[Any, nn. device = kwargs.get("device", "cuda") model_id = _TEXT_ENCODER_MODEL_IDS.get(str(name)) if model_id is None: - raise ValueError(f"Unsupported native SANA-WM text encoder: {name!r}") + raise ValueError(f"Unsupported SANA-WM text encoder: {name!r}") from transformers import AutoModelForCausalLM, AutoTokenizer, T5EncoderModel, T5Tokenizer if "T5" in str(name): @@ -1017,7 +1012,7 @@ def get(self, name: str, default: Any = None) -> Any: def _vae_encode_ltx2(name: str, vae: nn.Module, images: Tensor, *, device: torch.device) -> Tensor: if "LTX2VAE_diffusers" not in name: - raise ValueError(f"Unsupported native SANA-WM VAE encode type: {name!r}") + raise ValueError(f"Unsupported SANA-WM VAE encode type: {name!r}") dtype = images.dtype vae_device = next(vae.parameters()).device vae_dtype = next(vae.parameters()).dtype @@ -1031,7 +1026,7 @@ def _vae_encode_ltx2(name: str, vae: nn.Module, images: Tensor, *, device: torch def _vae_decode_ltx2(name: str, vae: nn.Module, latents: Tensor) -> Tensor: if "LTX2VAE_diffusers" not in name: - raise ValueError(f"Unsupported native SANA-WM VAE decode type: {name!r}") + raise ValueError(f"Unsupported SANA-WM VAE decode type: {name!r}") vae_device = next(vae.parameters()).device vae_dtype = next(vae.parameters()).dtype latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to( @@ -1118,8 +1113,8 @@ def _chunk_index_from_chunk_size( __all__ = [ - "SanaWMNativeTransformer", - "SanaWMNativeTransformerCache", - "SanaWMNativeTransformerConfig", + "SanaWMTransformer", + "SanaWMTransformerCache", + "SanaWMTransformerConfig", "SanaWMStage1Conditioning", ] diff --git a/integrations/sana/sana_wm/upstream_reference.py b/integrations/sana/sana_wm/upstream_reference.py deleted file mode 100644 index 6b87504e2..000000000 --- a/integrations/sana/sana_wm/upstream_reference.py +++ /dev/null @@ -1,90 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -"""Schema-only FlashDreams objects for the upstream SANA-WM reference runner.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import torch -import torch.nn as nn -from torch import Tensor - -from flashdreams.infra.diffusion.transformer import ( - Transformer, - TransformerAutoregressiveCache, - TransformerConfig, -) - - -@dataclass(kw_only=True) -class SanaWMUpstreamReferenceTransformerConfig(TransformerConfig): - """Transformer config used by the upstream-reference runner. - - This is intentionally not an executable SANA DiT adapter. It exists so the - CLI can expose the upstream reference path through the normal runner config - registry without pretending that upstream execution is the native path. - """ - - _target: type["SanaWMUpstreamReferenceTransformer"] = field( - default_factory=lambda: SanaWMUpstreamReferenceTransformer - ) - - -class SanaWMUpstreamReferenceTransformer( - Transformer[TransformerAutoregressiveCache] -): - """Non-executable marker for the upstream-delegating reference runner.""" - - def __init__(self, config: SanaWMUpstreamReferenceTransformerConfig) -> None: - super().__init__(config) - self._dummy = nn.Parameter(torch.empty(0)) - - @property - def latent_shape(self) -> tuple[int, ...]: - """Raise because this config marks the upstream reference path.""" - raise RuntimeError( - "This SANA-WM runner is an upstream reference harness. Use " - "sana-wm-bidirectional for native FlashDreams execution." - ) - - def predict_flow( - self, - noisy_latent: Tensor, - timestep: Tensor, - cache: TransformerAutoregressiveCache, - input: Any = None, - ) -> Tensor: - """Raise because this config marks the upstream reference path.""" - raise RuntimeError( - "This SANA-WM runner delegates generation to upstream NVlabs/Sana " - "instead of executing through Transformer.predict_flow." - ) - - def patchify_and_maybe_split_cp(self, x: Any) -> Any: - """Return ``x`` unchanged for schema-only construction.""" - return x - - def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: - """Return ``x`` unchanged for schema-only construction.""" - return x - - -__all__ = [ - "SanaWMUpstreamReferenceTransformer", - "SanaWMUpstreamReferenceTransformerConfig", -] diff --git a/integrations/sana/tests/parity_check/README.md b/integrations/sana/tests/parity_check/README.md deleted file mode 100644 index 62b13519b..000000000 --- a/integrations/sana/tests/parity_check/README.md +++ /dev/null @@ -1,25 +0,0 @@ - - -# SANA-WM parity check - -This directory is reserved for the opt-in SANA-WM native-vs-upstream parity -harness. It should follow the existing integration pattern: - -- clone or reuse a pinned NVlabs/Sana checkout; -- install upstream-heavy dependencies in an isolated environment; -- run upstream `inference_video_scripts/wm/inference_sana_wm.py` and the - native `flashdreams-run sana-wm-bidirectional` path on the same image, - prompt, camera, intrinsics, seed, and sampler settings; -- compare decoded MP4 frames and report mean / max `|Delta|`. - -The temporary FlashDreams-packaged upstream-reference runner is documented in -`../../docs/upstream_reference.md`. Keep those wrapper details out of the native -integration README. - -Do not put the full checkpoint download, refiner execution, or MP4 -generation path in `ci_cpu`. Add a bounded `ci_gpu` gate only if CI -pre-provisions the required artefacts and the test skips cleanly when -they are absent. diff --git a/integrations/sana/tests/test_native_quant_cuda.py b/integrations/sana/tests/test_quant_cuda.py similarity index 81% rename from integrations/sana/tests/test_native_quant_cuda.py rename to integrations/sana/tests/test_quant_cuda.py index c32e2e65a..9034c0c7f 100644 --- a/integrations/sana/tests/test_native_quant_cuda.py +++ b/integrations/sana/tests/test_quant_cuda.py @@ -13,33 +13,33 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CUDA smoke tests for SANA-WM native low-precision linear replacements.""" +"""CUDA smoke tests for SANA-WM low-precision linear replacements.""" from __future__ import annotations import pytest import torch -from sana_wm.native_quant import TorchScaledMMFP4Linear, TorchScaledMMFP8Linear +from sana_wm.quant import TorchScaledMMFP4Linear, TorchScaledMMFP8Linear pytestmark = pytest.mark.ci_gpu -def _require_native_quant_cuda() -> None: +def _require_quant_cuda() -> None: if not torch.cuda.is_available(): - pytest.skip("CUDA is required for native quantization smokes") + pytest.skip("CUDA is required for quantization smokes") missing = [ name for name in ("_scaled_mm", "float8_e4m3fn", "float4_e2m1fn_x2") if not hasattr(torch, name) ] if missing: - pytest.skip(f"PyTorch lacks native quantization primitive(s): {missing}") + pytest.skip(f"PyTorch lacks quantization primitive(s): {missing}") -def test_native_fp8_linear_runs_scaled_mm_on_cuda() -> None: +def test_fp8_linear_runs_scaled_mm_on_cuda() -> None: """Exercise the TE-free E4M3 FP8 Linear replacement on CUDA.""" - _require_native_quant_cuda() + _require_quant_cuda() source = torch.nn.Linear(32, 64, bias=True, device="cuda", dtype=torch.bfloat16) quantized = TorchScaledMMFP8Linear.from_linear( source, @@ -56,9 +56,9 @@ def test_native_fp8_linear_runs_scaled_mm_on_cuda() -> None: assert torch.isfinite(output.float()).all() -def test_native_fp4_linear_runs_scaled_mm_on_blackwell() -> None: +def test_fp4_linear_runs_scaled_mm_on_blackwell() -> None: """Exercise the TE-free E2M1 NVFP4 Linear replacement on Blackwell CUDA.""" - _require_native_quant_cuda() + _require_quant_cuda() major, minor = torch.cuda.get_device_capability() if major < 10: pytest.skip(f"NVFP4 requires Blackwell-class CUDA, got sm_{major}{minor}") diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index ddb05d8f6..5cb0a8dfd 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CPU-safe smoke tests for the SANA-WM native and reference configs.""" +"""CPU-safe smoke tests for the SANA-WM configs.""" from __future__ import annotations @@ -23,8 +23,8 @@ import pytest import torch -import sana_wm.native_refiner as native_refiner -import sana_wm.native_transformer as native_transformer_module +import sana_wm.refiner as refiner_module +import sana_wm.transformer as transformer_module try: import tomllib @@ -33,12 +33,8 @@ from sana_wm.config import ( PIPELINE_SANA_WM_BIDIRECTIONAL, - PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL, - PIPELINE_SANA_WM_UPSTREAM_REFERENCE, RUNNER_CONFIGS, RUNNER_SANA_WM_BIDIRECTIONAL, - RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, - RUNNER_SANA_WM_UPSTREAM_REFERENCE, ) from sana_wm.constants import ( SANA_WM_CONFIG_PATH, @@ -53,20 +49,20 @@ _temporary_environment, _validate_precision_request, ) -from sana_wm.native_diffusion import ( +from sana_wm.diffusion import ( _condition_model_kwargs, _uncondition_model_kwargs, ) -from sana_wm.native_diffusion import SanaWMDiffusionModelConfig -from sana_wm.native_transformer import ( - SanaWMNativeTransformerCache, - SanaWMNativeTransformerConfig, +from sana_wm.diffusion import SanaWMDiffusionModelConfig +from sana_wm.transformer import ( + SanaWMTransformerCache, + SanaWMTransformerConfig, SanaWMStage1Conditioning, _avoid_degenerate_tile_tail, _load_inference_config, ) -from sana_wm.native_refiner import SanaWMLTX2Refiner, _pack_latents, _unpack_latents -from sana_wm.native_quant import ( +from sana_wm.refiner import SanaWMLTX2Refiner, _pack_latents, _unpack_latents +from sana_wm.quant import ( TorchScaledMMFP4Linear, TorchScaledMMFP8Linear, replace_linear_with_torch_fp4, @@ -78,7 +74,6 @@ SanaWMStage1Spec, Stage1SelfAttention, ) -from sana_wm.upstream_reference import SanaWMUpstreamReferenceTransformerConfig from flashdreams.infra.config import derive_config @@ -88,11 +83,9 @@ def test_runner_config_is_registered() -> None: - """Expose SANA-WM native and upstream-reference runner slugs.""" + """Expose the SANA-WM runner slug.""" assert RUNNER_CONFIGS == { "sana-wm-bidirectional": RUNNER_SANA_WM_BIDIRECTIONAL, - "sana-wm-upstream-reference": RUNNER_SANA_WM_UPSTREAM_REFERENCE, - "sana-wm-native-bidirectional": RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, } @@ -101,47 +94,27 @@ def test_runner_name_mirrors_pipeline_name() -> None: assert RUNNER_SANA_WM_BIDIRECTIONAL.runner_name == ( PIPELINE_SANA_WM_BIDIRECTIONAL.name ) - assert RUNNER_SANA_WM_UPSTREAM_REFERENCE.runner_name == ( - PIPELINE_SANA_WM_UPSTREAM_REFERENCE.name - ) - assert RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL.runner_name == ( - PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL.name - ) def test_runner_has_description() -> None: """Provide non-empty CLI help text for the runner registry.""" assert RUNNER_SANA_WM_BIDIRECTIONAL.description.strip() - assert RUNNER_SANA_WM_UPSTREAM_REFERENCE.description.strip() - assert RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL.description.strip() -def test_reference_and_native_pipelines_are_distinct() -> None: - """Keep the upstream harness separate from the native FlashDreams target.""" - reference_transformer = ( - PIPELINE_SANA_WM_UPSTREAM_REFERENCE.diffusion_model.transformer - ) - main_transformer = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.transformer - native_transformer = ( - PIPELINE_SANA_WM_NATIVE_BIDIRECTIONAL.diffusion_model.transformer - ) +def test_pipeline_uses_sana_diffusion_model() -> None: + """Keep the public runner wired to the SANA-WM diffusion model.""" + transformer = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.transformer - assert isinstance(reference_transformer, SanaWMUpstreamReferenceTransformerConfig) assert isinstance( PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model, SanaWMDiffusionModelConfig, ) - assert isinstance(main_transformer, SanaWMNativeTransformerConfig) - assert isinstance(native_transformer, SanaWMNativeTransformerConfig) - assert RUNNER_SANA_WM_BIDIRECTIONAL.execution_backend == "native-flashdreams" - assert RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL.execution_backend == ( - "native-flashdreams" - ) + assert isinstance(transformer, SanaWMTransformerConfig) -def test_native_transformer_contract_shape_and_conditioning_guard() -> None: - """Pin the native SANA-WM Stage-1 boundary to the public model layout.""" - transformer_cfg = SanaWMNativeTransformerConfig() +def test_transformer_contract_shape_and_conditioning_guard() -> None: + """Pin the SANA-WM Stage-1 boundary to the public model layout.""" + transformer_cfg = SanaWMTransformerConfig() transformer = transformer_cfg.setup() assert transformer.latent_shape == (1, 128, 21, 22, 40) @@ -155,10 +128,10 @@ def test_native_transformer_contract_shape_and_conditioning_guard() -> None: ) -def test_native_inference_config_loads_yaml_without_upstream_types( +def test_inference_config_loads_yaml( tmp_path: Path, ) -> None: - """Parse SANA-WM YAML with local config objects, not SANA pyrallis types.""" + """Parse SANA-WM YAML with local config objects.""" config_path = tmp_path / "config.yaml" config_path.write_text( """ @@ -193,8 +166,8 @@ def test_native_inference_config_loads_yaml_without_upstream_types( assert cfg.work_dir == "" -def test_native_stage1_model_matches_checkpoint_schema() -> None: - """Pin the FlashDreams-owned Stage-1 module to the public checkpoint schema.""" +def test_stage1_model_matches_checkpoint_schema() -> None: + """Pin the Stage-1 module to the public checkpoint schema.""" state = SanaWMStage1Model().state_dict() assert len(state) == 872 @@ -223,8 +196,8 @@ def test_native_stage1_model_matches_checkpoint_schema() -> None: assert has_gdn_conv is SANA_WM_STAGE1_SPEC.block_uses_gdn(block_index) -def test_native_stage1_forward_preserves_latent_shape() -> None: - """Exercise the native Stage-1 forward path on a small CPU-safe spec.""" +def test_stage1_forward_preserves_latent_shape() -> None: + """Exercise the Stage-1 forward path on a small CPU-safe spec.""" spec = SanaWMStage1Spec( latent_channels=4, hidden_size=16, @@ -255,7 +228,7 @@ def test_native_stage1_forward_preserves_latent_shape() -> None: def test_stage1_self_attention_uses_camera_conditions() -> None: - """Changing camera conditions must affect native camera attention.""" + """Changing camera conditions must affect camera attention.""" torch.manual_seed(0) spec = SanaWMStage1Spec( latent_channels=4, @@ -289,9 +262,9 @@ def test_stage1_self_attention_uses_camera_conditions() -> None: assert not torch.allclose(base, shifted) -def test_native_transformer_releases_stage1_runtime() -> None: +def test_transformer_releases_stage1_runtime() -> None: """Free Stage-1-only modules and conditioning before decode/refine.""" - transformer = SanaWMNativeTransformerConfig().setup() + transformer = SanaWMTransformerConfig().setup() transformer.model = torch.nn.Linear(1, 1) transformer.text_encoder = torch.nn.Linear(1, 1) transformer.tokenizer = object() @@ -304,7 +277,7 @@ def test_native_transformer_releases_stage1_runtime() -> None: torch.empty(1), torch.empty(1), ) - cache = SanaWMNativeTransformerCache( + cache = SanaWMTransformerCache( conditioning=SanaWMStage1Conditioning( condition=torch.empty(1), uncondition=None, @@ -330,8 +303,8 @@ def test_native_transformer_releases_stage1_runtime() -> None: assert transformer._stage1_quantized is False -def test_native_vae_tiling_uses_low_memory_tiles() -> None: - """Keep native VAE decode on smaller tiles than upstream defaults.""" +def test_vae_tiling_uses_low_memory_tiles() -> None: + """Keep VAE decode on the configured low-memory tiles.""" class DummyVAE: def __init__(self) -> None: @@ -349,7 +322,7 @@ def __init__(self) -> None: def enable_tiling(self, **kwargs: int) -> None: self.calls.append(kwargs) - transformer = SanaWMNativeTransformerConfig().setup() + transformer = SanaWMTransformerConfig().setup() transformer.vae = DummyVAE() transformer._configure_vae_tiling() @@ -374,10 +347,10 @@ def enable_tiling(self, **kwargs: int) -> None: assert transformer.vae.use_framewise_decoding is True -def test_native_decode_retries_vae_oom_with_smaller_tiles( +def test_decode_retries_vae_oom_with_smaller_tiles( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Retry native VAE decode with smaller tiles after a CUDA OOM.""" + """Retry VAE decode with smaller tiles after a CUDA OOM.""" class DummyVAE: def __init__(self) -> None: @@ -399,7 +372,7 @@ def enable_tiling(self, **kwargs: int) -> None: for name, value in kwargs.items(): setattr(self, name, value) - transformer = SanaWMNativeTransformerConfig().setup() + transformer = SanaWMTransformerConfig().setup() transformer.vae = DummyVAE() transformer.vae_dtype = torch.float32 monkeypatch.setattr(transformer, "_ensure_vae", lambda: None) @@ -452,7 +425,7 @@ def test_cfg_model_kwargs_do_not_duplicate_camera_tensors() -> None: assert "negative_mask" not in neg_kwargs -def test_native_vae_tiling_avoids_degenerate_latent_tails() -> None: +def test_vae_tiling_avoids_degenerate_latent_tails() -> None: """Avoid last spatial VAE tiles with size one after compression.""" assert ( _avoid_degenerate_tile_tail( @@ -486,8 +459,8 @@ def test_hf_defaults_point_at_bidirectional_release() -> None: ) -def test_runner_setup_does_not_import_upstream_sana() -> None: - """Construct the runner without touching upstream Sana dependencies.""" +def test_runner_setup_preserves_cli_fields() -> None: + """Construct the runner and preserve CLI override fields.""" cfg = derive_config( RUNNER_SANA_WM_BIDIRECTIONAL, image_path=Path("missing.png"), @@ -503,35 +476,33 @@ def test_runner_setup_does_not_import_upstream_sana() -> None: def test_runner_config_type() -> None: """Keep the exported literal on the SANA-WM runner config subclass.""" assert isinstance(RUNNER_SANA_WM_BIDIRECTIONAL, SanaWMRunnerConfig) - assert isinstance(RUNNER_SANA_WM_UPSTREAM_REFERENCE, SanaWMRunnerConfig) - assert isinstance(RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, SanaWMRunnerConfig) def test_runner_defaults_to_bf16_precision() -> None: - """Keep the default path on upstream SANA-WM's BF16 config.""" + """Keep the default runner on BF16 precision.""" assert RUNNER_SANA_WM_BIDIRECTIONAL.stage1_precision == "bf16" assert RUNNER_SANA_WM_BIDIRECTIONAL.refiner_precision == "bf16" assert RUNNER_SANA_WM_BIDIRECTIONAL.quant_backend == "auto" assert RUNNER_SANA_WM_BIDIRECTIONAL.no_refiner is False -def test_native_refiner_is_flashdreams_owned( +def test_refiner_is_flashdreams_owned( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Build the native refiner adapter instead of raising or importing Sana.""" + """Build the refiner adapter.""" calls: dict[str, object] = {} class DummyRefiner: def __init__(self, **kwargs: object) -> None: calls.update(kwargs) - monkeypatch.setattr(native_refiner, "SanaWMLTX2Refiner", DummyRefiner) + monkeypatch.setattr(refiner_module, "SanaWMLTX2Refiner", DummyRefiner) monkeypatch.setattr( - native_transformer_module, + transformer_module, "resolve_hf_path", lambda value: f"/resolved/{value}", ) - transformer = SanaWMNativeTransformerConfig().setup() + transformer = SanaWMTransformerConfig().setup() transformer._ensure_refiner() @@ -545,7 +516,7 @@ def __init__(self, **kwargs: object) -> None: ) assert calls["dtype"] is torch.bfloat16 assert calls["precision"] == "bf16" - assert calls["quant_backend"] == "native" + assert calls["quant_backend"] == "torch" def test_refiner_latent_pack_round_trips() -> None: @@ -574,7 +545,7 @@ def test_refiner_latent_pack_round_trips() -> None: def test_refiner_block_size_request_still_executes( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Do not turn native refiner block-size requests into a hard failure.""" + """Do not turn refiner block-size requests into a hard failure.""" class DummyTransformer(torch.nn.Module): def __init__(self) -> None: @@ -623,14 +594,14 @@ def predict_current_x0( assert inference_modes == [True] -def test_auto_quant_backend_resolves_to_native_backend() -> None: - """Keep default quantization on the TE-free native backend.""" - assert _resolve_quant_backend("auto", ["fp4"]) == "native" - assert _resolve_quant_backend("auto", ["fp8", "fp4"]) == "native" - assert _resolve_quant_backend("auto", []) == "native" +def test_auto_quant_backend_resolves_to_torch_backend() -> None: + """Keep default quantization on the Torch backend.""" + assert _resolve_quant_backend("auto", ["fp4"]) == "torch" + assert _resolve_quant_backend("auto", ["fp8", "fp4"]) == "torch" + assert _resolve_quant_backend("auto", []) == "torch" -def test_bf16_precision_clears_upstream_quant_env( +def test_bf16_precision_clears_quant_env( monkeypatch: pytest.MonkeyPatch, ) -> None: """Prevent external NVFP4 env vars from changing the default runner path.""" @@ -650,8 +621,8 @@ def test_bf16_precision_clears_upstream_quant_env( assert os.environ["SANA_WM_REFINER_NVFP4"] == "1" -def test_fp8_precision_sets_upstream_quant_env() -> None: - """Map FlashDreams fp8 fields to upstream SANA-WM env selectors.""" +def test_fp8_precision_sets_quant_env() -> None: + """Map FlashDreams fp8 fields to SANA-WM env selectors.""" updates = _precision_env_updates( stage1_precision="fp8", refiner_precision="fp8", @@ -685,7 +656,7 @@ def test_quantized_precision_requires_cuda() -> None: stage1_precision="fp8", refiner_precision="bf16", refiner_enabled=True, - quant_backend="upstream-te", + quant_backend="torch", ) @@ -711,23 +682,7 @@ def test_fp8_precision_requires_hopper(monkeypatch: pytest.MonkeyPatch) -> None: stage1_precision="fp8", refiner_precision="bf16", refiner_enabled=True, - quant_backend="upstream-te", - ) - - -def test_fp8_precision_requires_cuda_129(monkeypatch: pytest.MonkeyPatch) -> None: - """Reject upstream FP8 block scaling when torch reports CUDA 12.8.""" - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) - monkeypatch.setattr(torch.version, "cuda", "12.8") - - with pytest.raises(ValueError, match="requires CUDA 12.9 or newer"): - _validate_precision_request( - device=torch.device("cuda:0"), - stage1_precision="fp8", - refiner_precision="bf16", - refiner_enabled=True, - quant_backend="upstream-te", + quant_backend="torch", ) @@ -742,14 +697,14 @@ def test_fp4_precision_requires_blackwell(monkeypatch: pytest.MonkeyPatch) -> No stage1_precision="fp4", refiner_precision="bf16", refiner_enabled=True, - quant_backend="upstream-te", + quant_backend="torch", ) -def test_torch_fp8_backend_skips_transformer_engine_cuda_129_gate( +def test_torch_fp8_backend_validates_primitives( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Allow native FP8 validation in the CUDA 12.8 env that blocks TE FP8.""" + """Allow FP8 validation with the required PyTorch primitives.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) monkeypatch.setattr(torch.version, "cuda", "12.8") @@ -765,10 +720,10 @@ def test_torch_fp8_backend_skips_transformer_engine_cuda_129_gate( ) -def test_auto_fp8_backend_uses_native_primitives( +def test_auto_fp8_backend_uses_torch_primitives( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Route default FP8 to the TE-free native backend.""" + """Route default FP8 to the Torch backend.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) monkeypatch.setattr(torch.version, "cuda", "12.8") @@ -784,10 +739,10 @@ def test_auto_fp8_backend_uses_native_primitives( ) -def test_native_runner_reaches_normal_input_validation() -> None: - """Do not silently delegate the native runner to the upstream harness.""" +def test_runner_reaches_normal_input_validation() -> None: + """Reach normal file input validation.""" cfg = derive_config( - RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL, + RUNNER_SANA_WM_BIDIRECTIONAL, image_path=Path("missing.png"), prompt="demo", ) @@ -812,8 +767,8 @@ def test_torch_fp8_backend_rejects_fp4(monkeypatch: pytest.MonkeyPatch) -> None: ) -def test_torch_fp4_backend_skips_transformer_engine(monkeypatch: pytest.MonkeyPatch) -> None: - """Allow native FP4 validation without importing Transformer Engine.""" +def test_torch_fp4_backend_validates_primitives(monkeypatch: pytest.MonkeyPatch) -> None: + """Allow FP4 validation with the required PyTorch primitives.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) @@ -830,7 +785,7 @@ def test_torch_fp4_backend_skips_transformer_engine(monkeypatch: pytest.MonkeyPa def test_torch_fp4_backend_rejects_fp8(monkeypatch: pytest.MonkeyPatch) -> None: - """Keep precision-specific native backends explicit.""" + """Keep precision-specific backends explicit.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) @@ -844,10 +799,10 @@ def test_torch_fp4_backend_rejects_fp8(monkeypatch: pytest.MonkeyPatch) -> None: ) -def test_native_backend_allows_mixed_fp8_fp4_without_cuda_129( +def test_torch_backend_allows_mixed_fp8_fp4( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Allow mixed native FP8/FP4 requests in the CUDA 12.8 env that blocks TE FP8.""" + """Allow mixed FP8/FP4 requests on the Torch backend.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) monkeypatch.setattr(torch.version, "cuda", "12.8") @@ -860,14 +815,14 @@ def test_native_backend_allows_mixed_fp8_fp4_without_cuda_129( stage1_precision="fp8", refiner_precision="fp4", refiner_enabled=True, - quant_backend="native", + quant_backend="torch", ) def test_torch_fp8_linear_replaces_matching_modules() -> None: - """Provide a TE-free replacement for eligible upstream Linear modules.""" + """Provide a Torch replacement for eligible Linear modules.""" if not hasattr(torch, "float8_e4m3fn"): - pytest.skip("torch.float8_e4m3fn is required for native FP8 replacement") + pytest.skip("torch.float8_e4m3fn is required for FP8 replacement") module = torch.nn.Sequential( torch.nn.Linear(16, 32, bias=True), @@ -895,7 +850,7 @@ def test_torch_fp8_linear_replaces_matching_modules() -> None: def test_torch_fp4_linear_replaces_matching_modules( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Route eligible FP4 modules through the native replacement helper.""" + """Route eligible FP4 modules through the replacement helper.""" module = torch.nn.Sequential( torch.nn.Linear(32, 64, bias=True), torch.nn.Sequential(torch.nn.Linear(32, 64, bias=False)), @@ -942,17 +897,11 @@ def test_pyproject_entry_point_matches_runner_literal() -> None: assert entry_points == { "sana-wm-bidirectional": "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL", - "sana-wm-upstream-reference": ( - "sana_wm.config:RUNNER_SANA_WM_UPSTREAM_REFERENCE" - ), - "sana-wm-native-bidirectional": ( - "sana_wm.config:RUNNER_SANA_WM_NATIVE_BIDIRECTIONAL" - ), } -def test_pyproject_excludes_upstream_diffusion_package() -> None: - """Do not install the NVlabs/Sana ``diffusion`` package.""" +def test_pyproject_package_selection() -> None: + """Keep the package metadata aligned with the integration package.""" pyproject = tomllib.loads( Path("integrations/sana/pyproject.toml").read_text(encoding="utf-8") ) From 26a66be62b119daebcec7254472a833abe2773e8 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 17 Jul 2026 11:25:56 -0700 Subject: [PATCH 13/36] Harden SANA-WM FlashDreams integration boundaries Move SANA-WM prompt, first-frame, camera, decode, and refiner orchestration into explicit pipeline components instead of keeping it inside the transformer or runner. The public runner now drives the normal StreamInferencePipeline contract, with a Sana conditioning encoder, Stage-1 transformer, LTX-style scheduler, and video decoder/refiner boundary. Keep the native Stage-1 DiT implementation but make SanaWMDiffusionModelConfig instantiate the shared FlashDreams DiffusionModel. Add generic infra hooks for transformer-provided initial noise and scheduler access to per-step encoder context so Sana can pin the first-frame latent and run per-token LTX Euler sampling without owning a custom diffusion loop. Replace direct diffusers FlowMatch Euler use with a SanaWMLTXEulerScheduler that matches the LTX per-token behavior and keeps the scheduler-specific timestep handling isolated behind the FlashDreams scheduler interface. Restore batched CFG as the default Stage-1 path and remove manual early temporary deletions added during the earlier OOM investigation. Keep inference-mode execution and default-false offload/tiling fallback paths where they remain useful inference hygiene or explicit user controls. Expand CPU coverage for runner registration, component config routing, prompt/first-frame/camera conditioning shapes, scheduler parity, base DiffusionModel instantiation, CFG batching, VAE decode/refiner boundaries, and Stage-1 checkpoint shape coverage. Signed-off-by: Aidan Foster --- .../flashdreams/infra/diffusion/model/base.py | 21 +- .../infra/diffusion/scheduler/base.py | 6 +- .../infra/diffusion/scheduler/fm.py | 3 + .../infra/diffusion/scheduler/fm_euler.py | 3 + .../infra/diffusion/scheduler/fm_unipc.py | 3 + .../infra/diffusion/transformer/base.py | 31 + integrations/sana/sana_wm/conditioning.py | 552 +++++++++++++++ integrations/sana/sana_wm/config.py | 8 +- integrations/sana/sana_wm/decoder.py | 488 +++++++++++++ integrations/sana/sana_wm/diffusion.py | 231 +------ integrations/sana/sana_wm/runner.py | 127 ++-- integrations/sana/sana_wm/scheduler.py | 310 +++++++++ integrations/sana/sana_wm/stage1_model.py | 8 +- integrations/sana/sana_wm/transformer.py | 648 +++--------------- integrations/sana/tests/test_smoke.py | 566 +++++++++++++-- 15 files changed, 2089 insertions(+), 916 deletions(-) create mode 100644 integrations/sana/sana_wm/conditioning.py create mode 100644 integrations/sana/sana_wm/decoder.py create mode 100644 integrations/sana/sana_wm/scheduler.py diff --git a/flashdreams/flashdreams/infra/diffusion/model/base.py b/flashdreams/flashdreams/infra/diffusion/model/base.py index 02dbb909c..67cef7532 100644 --- a/flashdreams/flashdreams/infra/diffusion/model/base.py +++ b/flashdreams/flashdreams/infra/diffusion/model/base.py @@ -173,19 +173,19 @@ def generate( self.latent_shape, device=self.device, dtype=self.dtype ) dummy_latent = self.transformer.unpatchify_and_maybe_gather_cp(dummy_latent) - initial_noise = torch.randn( - dummy_latent.shape, - device=self.device, - dtype=self.dtype, - generator=self.rng, + initial_noise = self.transformer.initial_noise( + latent_shape=tuple(dummy_latent.shape), + rng=self.rng, + cache=cache, + input=input, ) else: - initial_noise = torch.randn( - self.latent_shape, - device=self.device, - dtype=self.dtype, - generator=self.rng, + initial_noise = self.transformer.initial_noise( + latent_shape=self.latent_shape, + rng=self.rng, + cache=cache, + input=input, ) def predict_flow(noisy_latent: Tensor, timestep: Tensor) -> Tensor: @@ -212,6 +212,7 @@ def predict_flow(noisy_latent: Tensor, timestep: Tensor) -> Tensor: initial_noise=initial_noise, predict_flow=predict_flow, rng=self.rng, + context=input, ) if self.config.noise_in_unpatchified_shape: diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/base.py b/flashdreams/flashdreams/infra/diffusion/scheduler/base.py index 46511d93b..4bb0f1dd6 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/base.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/base.py @@ -19,7 +19,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Protocol +from typing import Any, Protocol import torch import torch.nn as nn @@ -69,6 +69,7 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, + context: Any = None, ) -> Tensor: """Run the full denoising loop and return the clean latent. @@ -83,6 +84,9 @@ def sample( ``num_inference_steps`` times. rng: Generator on the same device. Used by self-forcing renoise loops; pure ODE solvers ignore it. + context: Optional per-AR-step encoder output. Most schedulers + ignore it; schedulers with token-local constraints can use it + without taking over the diffusion-model orchestration. Returns: Clean latent with the same shape, device, and dtype as diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py b/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py index baaabf78b..75a6cbc4e 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py @@ -18,6 +18,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Any import torch from torch import Tensor @@ -212,6 +213,7 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, + context: Any = None, ) -> Tensor: """Run the self-forcing flow-match denoising loop. @@ -220,6 +222,7 @@ def sample( sigma before the network forward. Schedule arithmetic auto-promotes to fp32; the result is cast back to ``initial_noise.dtype``. """ + del context input_dtype = initial_noise.dtype sigmas = self.denoising_sigmas timesteps = self.denoising_step_list diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py index 76498d987..2723e0bd0 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py @@ -29,6 +29,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Any import numpy as np import torch @@ -186,6 +187,7 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, + context: Any = None, ) -> Tensor: """Run the explicit Euler denoising loop. @@ -199,6 +201,7 @@ def sample( ``rng`` is unused (deterministic ODE) but accepted for interface conformance. """ + del context input_dtype = initial_noise.dtype N = self.config.num_inference_steps diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py index a9ce940a7..ff5ca162d 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py @@ -18,6 +18,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Any import numpy as np import torch @@ -331,6 +332,7 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, + context: Any = None, ) -> Tensor: """Run the order-2 UniPC predictor-corrector denoising loop. @@ -340,6 +342,7 @@ def sample( fp32; the result is cast back to ``initial_noise.dtype``. ``rng`` is unused (deterministic ODE) but accepted for interface conformance. """ + del context input_dtype = initial_noise.dtype N = self.timesteps.shape[0] diff --git a/flashdreams/flashdreams/infra/diffusion/transformer/base.py b/flashdreams/flashdreams/infra/diffusion/transformer/base.py index c904e3d3a..1ff10749c 100644 --- a/flashdreams/flashdreams/infra/diffusion/transformer/base.py +++ b/flashdreams/flashdreams/infra/diffusion/transformer/base.py @@ -169,6 +169,37 @@ def initialize_autoregressive_cache(self, **context: Any) -> TransformerCacheT: """ return cast("TransformerCacheT", TransformerAutoregressiveCache()) + def initial_noise( + self, + *, + latent_shape: tuple[int, ...], + rng: torch.Generator | None, + cache: TransformerCacheT, + input: Any = None, + ) -> Tensor: + """Draw the starting latent for one denoising loop. + + Default is Gaussian noise in ``latent_shape``. Model integrations can + override this to seed from encoder outputs, pin I2V frames, or attach + per-step state to the cache while still using ``DiffusionModel``. + + Args: + latent_shape: Shape requested by the diffusion model. + rng: Per-model generator on ``self.device``, or ``None``. + cache: Per-rollout AR cache. + input: Same patchified encoder output passed to ``predict_flow``. + + Returns: + Initial latent tensor for scheduler sampling. + """ + del cache, input + return torch.randn( + latent_shape, + device=self.device, + dtype=self.dtype, + generator=rng, + ) + def postprocess_clean_latent( self, clean_latent: Tensor, diff --git a/integrations/sana/sana_wm/conditioning.py b/integrations/sana/sana_wm/conditioning.py new file mode 100644 index 000000000..6be23039b --- /dev/null +++ b/integrations/sana/sana_wm/conditioning.py @@ -0,0 +1,552 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 prompt, first-frame, and camera conditioning components.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from torch import Tensor + +from flashdreams.infra.encoder import ( + Encoder, + EncoderConfig, + StreamingEncoder, + StreamingEncoderCache, +) +from sana_wm.camera import prepare_camera +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, +) +from sana_wm.transformer import ( + QuantBackend, + SanaWMStage1Conditioning, + _chunk_index_from_config, + _get_tokenizer_and_text_encoder, + _get_vae, + _get_weight_dtype, + _load_inference_config, + _vae_encode_ltx2, +) + + +@dataclass(kw_only=True) +class SanaWMTextPromptRequest: + """Raw text prompt inputs for SANA-WM Stage 1.""" + + prompt: str + negative_prompt: str = "" + + +@dataclass(kw_only=True) +class SanaWMTextConditioning: + """Encoded positive and negative prompt tensors.""" + + condition: Tensor + condition_mask: Tensor + negative: Tensor + negative_mask: Tensor + + +@dataclass(kw_only=True) +class SanaWMCameraRequest: + """Raw camera trajectory inputs for SANA-WM conditioning.""" + + poses_c2w: np.ndarray + intrinsics_vec4: np.ndarray + + +@dataclass(kw_only=True) +class SanaWMI2VConditioning: + """Encoded first-frame and camera tensors for Stage-1 diffusion.""" + + first_latent: Tensor + camera: dict[str, Tensor] + num_frames: int + + +@dataclass(kw_only=True) +class SanaWMI2VConditioningRequest: + """Raw one-shot SANA-WM I2V rollout inputs.""" + + image: Any + prompt: str + poses_c2w: np.ndarray + intrinsics_vec4: np.ndarray + num_frames: int + fps: int + steps: int + cfg_scale: float + flow_shift: float | None + seed: int + negative_prompt: str = "" + + +@dataclass(kw_only=True) +class SanaWMTextPromptEncoderConfig(EncoderConfig): + """Config for the Stage-1 prompt encoder component.""" + + _target: type["SanaWMTextPromptEncoder"] = field( + default_factory=lambda: SanaWMTextPromptEncoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + stage1_precision: str = "bf16" + """Stage-1 precision; quantized paths pad text tokens for scaled-MM.""" + + quant_backend: QuantBackend = "auto" + """Low-precision backend selector, included in the prompt cache key.""" + + offload_text_encoder: bool = False + """Move the text encoder back to CPU after prompt encoding.""" + + +class SanaWMTextPromptEncoder(Encoder): + """Encode SANA-WM Stage-1 positive and negative prompts.""" + + config: SanaWMTextPromptEncoderConfig + + def __init__(self, config: SanaWMTextPromptEncoderConfig) -> None: + super().__init__(config) + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._text_encoder_built = False + self._prompt_cache: dict[ + tuple[object, ...], tuple[Tensor, Tensor, Tensor, Tensor] + ] = {} + + @property + def device(self) -> torch.device: + return self._dummy.device + + def forward(self, input: SanaWMTextPromptRequest) -> SanaWMTextConditioning: + """Encode prompt strings into Stage-1 text embeddings and masks.""" + cond, cond_mask, neg, neg_mask = self._encode_prompts( + input.prompt, + input.negative_prompt, + ) + cond, cond_mask, neg, neg_mask = self._pad_text_for_quant( + cond, + cond_mask, + neg, + neg_mask, + ) + return SanaWMTextConditioning( + condition=cond, + condition_mask=cond_mask, + negative=neg, + negative_mask=neg_mask, + ) + + def release_runtime(self) -> None: + """Release prompt encoder tensors.""" + self._prompt_cache.clear() + for attr in ("text_encoder", "tokenizer"): + if hasattr(self, attr): + setattr(self, attr, None) + self._text_encoder_built = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_weight_dtype(self) -> torch.dtype: + return _get_weight_dtype(self._ensure_runtime_config().model.mixed_precision) + + def _ensure_text_encoder(self) -> None: + if self._text_encoder_built: + return + cfg = self._ensure_runtime_config() + self.tokenizer, self.text_encoder = _get_tokenizer_and_text_encoder( + name=cfg.text_encoder.text_encoder_name, + device=self.device, + ) + if self.config.offload_text_encoder: + self.text_encoder.to("cpu") + self._text_encoder_built = True + + def _encode_prompts( + self, + prompt: str, + negative_prompt: str, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + cfg = self._ensure_runtime_config() + self._ensure_text_encoder() + max_length = cfg.text_encoder.model_max_length + chi_prompt = "\n".join(cfg.text_encoder.chi_prompt or []) + if chi_prompt: + prompt = chi_prompt + prompt + max_length_all = len(self.tokenizer.encode(chi_prompt)) + max_length - 2 + else: + max_length_all = max_length + + key = ( + prompt, + negative_prompt, + str(self.device), + str(self._ensure_weight_dtype()), + self.config.stage1_precision, + self.config.quant_backend, + ) + if key in self._prompt_cache: + return self._prompt_cache[key] + + move_text_encoder = self.config.offload_text_encoder or ( + _module_device(self.text_encoder) != self.device + ) + if move_text_encoder: + self.text_encoder.to(self.device) + + def encode(text: str, length: int) -> tuple[Tensor, Tensor]: + tokens = self.tokenizer( + [text], + max_length=length, + padding="max_length", + truncation=True, + return_tensors="pt", + ).to(self.device) + return self.text_encoder(tokens.input_ids, tokens.attention_mask)[0], ( + tokens.attention_mask + ) + + try: + cond, cond_mask = encode(prompt, max_length_all) + select = [0] + list(range(-max_length + 1, 0)) + cond = cond[:, None][:, :, select] + cond_mask = cond_mask[:, select] + neg, neg_mask = encode(negative_prompt, max_length) + result = (cond, cond_mask, neg[:, None], neg_mask) + self._prompt_cache.clear() + self._prompt_cache[key] = result + return result + finally: + if move_text_encoder: + self.text_encoder.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _pad_text_for_quant( + self, + cond: Tensor, + cond_mask: Tensor, + neg: Tensor, + neg_mask: Tensor, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + if self.config.stage1_precision == "bf16": + return cond, cond_mask, neg, neg_mask + multiple = int(os.environ.get("SANA_WM_STAGE1_NVFP4_TEXT_PAD_MULTIPLE", "8")) + if multiple <= 1: + return cond, cond_mask, neg, neg_mask + + def pad_pair(text: Tensor, mask: Tensor) -> tuple[Tensor, Tensor]: + pad = (-text.shape[-2]) % multiple + if pad == 0: + return text, mask + text_shape = list(text.shape) + text_shape[-2] = pad + mask_shape = list(mask.shape) + mask_shape[-1] = pad + return ( + torch.cat([text, text.new_zeros(text_shape)], dim=-2), + torch.cat([mask, mask.new_zeros(mask_shape)], dim=-1), + ) + + cond, cond_mask = pad_pair(cond, cond_mask) + neg, neg_mask = pad_pair(neg, neg_mask) + return cond, cond_mask, neg, neg_mask + + +@dataclass(kw_only=True) +class SanaWMFirstFrameEncoderConfig(EncoderConfig): + """Config for the first-frame VAE encoder component.""" + + _target: type["SanaWMFirstFrameEncoder"] = field( + default_factory=lambda: SanaWMFirstFrameEncoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + offload_vae: bool = False + """Move the VAE to CPU after first-frame encoding.""" + + +class SanaWMFirstFrameEncoder(Encoder): + """Encode the input first frame with the LTX-2 VAE.""" + + config: SanaWMFirstFrameEncoderConfig + + def __init__(self, config: SanaWMFirstFrameEncoderConfig) -> None: + super().__init__(config) + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._vae_built = False + + @property + def device(self) -> torch.device: + return self._dummy.device + + def forward(self, input: Any) -> Tensor: + """Encode a PIL-like RGB image into a single latent frame.""" + from torchvision import transforms as T + + cfg = self._ensure_runtime_config() + weight_dtype = _get_weight_dtype(cfg.model.mixed_precision) + self._ensure_vae() + if self.config.offload_vae: + self.vae.to(self.device) + + image = (T.ToTensor()(input) * 2.0 - 1.0).unsqueeze(0).unsqueeze(2) + latent = _vae_encode_ltx2( + cfg.vae.vae_type, + self.vae, + image.to(self.device, dtype=self.vae_dtype), + device=self.device, + ).to(weight_dtype) + if self.config.offload_vae: + self.vae.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return latent + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_vae(self) -> None: + if self._vae_built: + return + cfg = self._ensure_runtime_config() + self.vae_dtype = _get_weight_dtype(cfg.vae.weight_dtype) + from sana_wm._tools import resolve_hf_path + + cfg.vae.vae_pretrained = resolve_hf_path(cfg.vae.vae_pretrained) + self.vae = _get_vae( + cfg.vae.vae_type, + cfg.vae.vae_pretrained, + device=self.device, + dtype=self.vae_dtype, + config=cfg.vae, + ) + self._vae_built = True + + +@dataclass(kw_only=True) +class SanaWMCameraConditioningEncoderConfig(EncoderConfig): + """Config for the camera/raymap conditioning component.""" + + _target: type["SanaWMCameraConditioningEncoder"] = field( + default_factory=lambda: SanaWMCameraConditioningEncoder + ) + + height: int = DEFAULT_VIDEO_HEIGHT + width: int = DEFAULT_VIDEO_WIDTH + + +class SanaWMCameraConditioningEncoder(Encoder): + """Build SANA-WM raymap and chunk-Plucker camera tensors.""" + + config: SanaWMCameraConditioningEncoderConfig + + def __init__(self, config: SanaWMCameraConditioningEncoderConfig) -> None: + super().__init__(config) + self.config = config + + def forward(self, input: SanaWMCameraRequest) -> dict[str, Tensor]: + """Encode raw poses and intrinsics into Stage-1 camera tensors.""" + return prepare_camera( + input.poses_c2w, + input.intrinsics_vec4, + target_size=(self.config.height, self.config.width), + ) + + +@dataclass(kw_only=True) +class SanaWMConditioningEncoderConfig(EncoderConfig): + """Config for the pipeline-level SANA-WM conditioning encoder.""" + + _target: type["SanaWMConditioningEncoder"] = field( + default_factory=lambda: SanaWMConditioningEncoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + text_encoder: SanaWMTextPromptEncoderConfig = field( + default_factory=SanaWMTextPromptEncoderConfig + ) + first_frame_encoder: SanaWMFirstFrameEncoderConfig = field( + default_factory=SanaWMFirstFrameEncoderConfig + ) + camera_encoder: SanaWMCameraConditioningEncoderConfig = field( + default_factory=SanaWMCameraConditioningEncoderConfig + ) + height: int = DEFAULT_VIDEO_HEIGHT + width: int = DEFAULT_VIDEO_WIDTH + + +class SanaWMConditioningEncoder(StreamingEncoder[StreamingEncoderCache]): + """Prepare all per-rollout inputs for the Sana Stage-1 sampler.""" + + config: SanaWMConditioningEncoderConfig + + def __init__(self, config: SanaWMConditioningEncoderConfig) -> None: + super().__init__(config) + self.config = config + self.text_encoder = config.text_encoder.setup() + self.first_frame_encoder = config.first_frame_encoder.setup() + self.camera_encoder = config.camera_encoder.setup() + self._runtime_config: Any | None = None + + def initialize_autoregressive_cache(self, **_context: Any) -> StreamingEncoderCache: + """Return an empty cache for the one-shot conditioning encoder.""" + return StreamingEncoderCache() + + @torch.inference_mode() + def forward( + self, + input: SanaWMI2VConditioningRequest, + autoregressive_index: int = 0, + cache: StreamingEncoderCache | None = None, + ) -> SanaWMStage1Conditioning: + """Encode raw rollout inputs into Stage-1 conditioning.""" + del cache + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + + cfg = self._ensure_runtime_config() + text = self.text_encoder( + SanaWMTextPromptRequest( + prompt=input.prompt, + negative_prompt=input.negative_prompt, + ) + ) + first_latent = self.first_frame_encoder(input.image) + weight_dtype = _get_weight_dtype(cfg.model.mixed_precision) + camera = self.camera_encoder( + SanaWMCameraRequest( + poses_c2w=input.poses_c2w, + intrinsics_vec4=input.intrinsics_vec4, + ) + ) + raymap = camera["raymap"].unsqueeze(0).to( + first_latent.device, + dtype=weight_dtype, + ) + chunk_plucker = camera["chunk_plucker"].unsqueeze(0).to( + first_latent.device, + dtype=weight_dtype, + ) + + model_kwargs_extra: dict[str, object] = {} + if input.cfg_scale > 1.0: + model_kwargs_extra["negative_mask"] = text.negative_mask + uncondition = text.negative + else: + uncondition = None + + vae_stride = cfg.vae.vae_stride + latent_t = (input.num_frames - 1) // int(vae_stride[0]) + 1 + latent_h = self.config.height // int(vae_stride[-1]) + latent_w = self.config.width // int(vae_stride[-1]) + chunk_index = _chunk_index_from_config(cfg, num_frames=latent_t) + model_kwargs: dict[str, object] = { + "data_info": { + "img_hw": torch.tensor( + [[self.config.height, self.config.width]], + dtype=torch.float, + device=first_latent.device, + ), + "condition_frame_info": {0: 0.0}, + }, + "mask": text.condition_mask, + "camera_conditions": raymap, + "chunk_plucker": chunk_plucker, + **model_kwargs_extra, + } + if chunk_index is not None: + model_kwargs["chunk_index"] = chunk_index + + return SanaWMStage1Conditioning( + condition=text.condition, + uncondition=uncondition, + model_kwargs=model_kwargs, + first_latent=first_latent, + latent_shape=( + 1, + int(first_latent.shape[1]), + latent_t, + latent_h, + latent_w, + ), + cfg_scale=float(input.cfg_scale), + flow_shift=self._resolve_flow_shift(input.flow_shift), + steps=int(input.steps), + seed=int(input.seed), + ) + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _resolve_flow_shift(self, override: float | None) -> float: + cfg = self._ensure_runtime_config() + if override is not None: + return float(override) + if cfg.scheduler.inference_flow_shift is not None: + return float(cfg.scheduler.inference_flow_shift) + return float(cfg.scheduler.flow_shift) + + +def _module_device(module: nn.Module) -> torch.device: + try: + return next(module.parameters()).device + except StopIteration: + return torch.device("cpu") + + +__all__ = [ + "SanaWMCameraConditioningEncoder", + "SanaWMCameraConditioningEncoderConfig", + "SanaWMCameraRequest", + "SanaWMConditioningEncoder", + "SanaWMConditioningEncoderConfig", + "SanaWMFirstFrameEncoder", + "SanaWMFirstFrameEncoderConfig", + "SanaWMI2VConditioning", + "SanaWMI2VConditioningRequest", + "SanaWMTextConditioning", + "SanaWMTextPromptEncoder", + "SanaWMTextPromptEncoderConfig", + "SanaWMTextPromptRequest", +] diff --git a/integrations/sana/sana_wm/config.py b/integrations/sana/sana_wm/config.py index 5d9a176c8..e08b2c679 100644 --- a/integrations/sana/sana_wm/config.py +++ b/integrations/sana/sana_wm/config.py @@ -17,21 +17,25 @@ from __future__ import annotations -from flashdreams.infra.diffusion.scheduler import FlowMatchSchedulerConfig from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.infra.runner import RunnerConfig +from sana_wm.conditioning import SanaWMConditioningEncoderConfig +from sana_wm.decoder import SanaWMVideoDecoderConfig from sana_wm.diffusion import SanaWMDiffusionModelConfig from sana_wm.runner import SanaWMRunnerConfig +from sana_wm.scheduler import SanaWMLTXEulerSchedulerConfig from sana_wm.transformer import SanaWMTransformerConfig PIPELINE_SANA_WM_BIDIRECTIONAL = StreamInferencePipelineConfig( name="sana-wm-bidirectional", + encoder=SanaWMConditioningEncoderConfig(), diffusion_model=SanaWMDiffusionModelConfig( transformer=SanaWMTransformerConfig(), - scheduler=FlowMatchSchedulerConfig(), + scheduler=SanaWMLTXEulerSchedulerConfig(), seed=42, ), + decoder=SanaWMVideoDecoderConfig(), ) """FlashDreams SANA-WM pipeline.""" diff --git a/integrations/sana/sana_wm/decoder.py b/integrations/sana/sana_wm/decoder.py new file mode 100644 index 000000000..698e1d6aa --- /dev/null +++ b/integrations/sana/sana_wm/decoder.py @@ -0,0 +1,488 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 latent refiner and VAE decoder components.""" + +from __future__ import annotations + +import gc +import time +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from loguru import logger +from torch import Tensor + +from flashdreams.infra.config import InstantiateConfig +from flashdreams.infra.decoder import ( + DecoderConfig, + StreamingDecoder, + StreamingDecoderCache, +) +from sana_wm._tools import resolve_hf_path +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, + SANA_WM_REFINER_GEMMA_ROOT, + SANA_WM_REFINER_ROOT, + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_TEMPORAL_COMPRESSION, +) +from sana_wm.transformer import ( + Precision, + QuantBackend, + _avoid_degenerate_tile_tail, + _get_vae, + _get_weight_dtype, + _load_inference_config, + _vae_decode_ltx2, +) + + +@dataclass(kw_only=True) +class SanaWMDecodedVideo: + """Decoded SANA-WM video outputs.""" + + video_hwc: np.ndarray + stage1_video_hwc: np.ndarray | None = None + + +@dataclass(kw_only=True) +class SanaWMVideoDecoderCache(StreamingDecoderCache): + """Per-rollout settings for the SANA-WM decoder/refiner stage.""" + + prompt: str = "" + fps: int = 16 + save_stage1: bool = False + refiner_seed: int = 42 + sink_size: int = 1 + refiner_block_size: int | None = None + refiner_kv_max_frames: int = 11 + + +@dataclass(kw_only=True) +class SanaWMLTX2VAEDecoderConfig(InstantiateConfig): + """Config for the LTX-2 VAE latent-to-video decoder.""" + + _target: type["SanaWMLTX2VAEDecoder"] = field( + default_factory=lambda: SanaWMLTX2VAEDecoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + height: int = DEFAULT_VIDEO_HEIGHT + width: int = DEFAULT_VIDEO_WIDTH + + offload_vae: bool = False + """Move the VAE to CPU between decode calls.""" + + vae_tile_sample_min_width: int = 256 + vae_tile_sample_stride_width: int = 224 + vae_tile_sample_min_height: int = 256 + vae_tile_sample_stride_height: int = 192 + vae_tile_sample_min_num_frames: int = 24 + vae_tile_sample_stride_num_frames: int = 8 + + vae_oom_retry_tile_sample_min_width: int = 128 + vae_oom_retry_tile_sample_stride_width: int = 64 + vae_oom_retry_tile_sample_min_height: int = 128 + vae_oom_retry_tile_sample_stride_height: int = 64 + vae_oom_retry_tile_sample_min_num_frames: int = 16 + vae_oom_retry_tile_sample_stride_num_frames: int = 8 + + +class SanaWMLTX2VAEDecoder(nn.Module): + """Decode SANA-WM LTX-2 VAE latents into HWC uint8 video.""" + + config: SanaWMLTX2VAEDecoderConfig + + def __init__(self, config: SanaWMLTX2VAEDecoderConfig) -> None: + super().__init__() + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._vae_built = False + + @property + def device(self) -> torch.device: + return self._dummy.device + + @torch.inference_mode() + def decode_latents(self, latents: Tensor) -> np.ndarray: + """Decode VAE latents to ``uint8`` HWC video.""" + self._ensure_vae() + if self.config.offload_vae: + self.vae.to(self.device) + samples = latents.to(device=self.device, dtype=self.vae_dtype) + if torch.cuda.is_available(): + torch.cuda.synchronize() + t0 = time.perf_counter() + retry_decode = False + try: + decoded = self._vae_decode(samples) + except torch.OutOfMemoryError: + retry_decode = True + + if retry_decode: + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + gc.collect() + logger.warning( + "[sana-vae] decode OOM; retrying with smaller tiles " + "width={} stride_width={} height={} stride_height={} " + "frames={} stride_frames={}", + self.config.vae_oom_retry_tile_sample_min_width, + self.config.vae_oom_retry_tile_sample_stride_width, + self.config.vae_oom_retry_tile_sample_min_height, + self.config.vae_oom_retry_tile_sample_stride_height, + self.config.vae_oom_retry_tile_sample_min_num_frames, + self.config.vae_oom_retry_tile_sample_stride_num_frames, + ) + self._configure_vae_tiling( + tile_sample_min_width=self.config.vae_oom_retry_tile_sample_min_width, + tile_sample_stride_width=( + self.config.vae_oom_retry_tile_sample_stride_width + ), + tile_sample_min_height=self.config.vae_oom_retry_tile_sample_min_height, + tile_sample_stride_height=( + self.config.vae_oom_retry_tile_sample_stride_height + ), + tile_sample_min_num_frames=( + self.config.vae_oom_retry_tile_sample_min_num_frames + ), + tile_sample_stride_num_frames=( + self.config.vae_oom_retry_tile_sample_stride_num_frames + ), + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + decoded = self._vae_decode(samples) + if torch.cuda.is_available(): + torch.cuda.synchronize() + logger.info( + "[timing] vae decode: {:.3f}s (latent T={} -> pixels {})", + time.perf_counter() - t0, + latents.shape[2], + tuple(decoded.shape) if isinstance(decoded, Tensor) else "list", + ) + if isinstance(decoded, list): + decoded = torch.stack(decoded, dim=0) + video = ( + torch.clamp(127.5 * decoded + 127.5, 0, 255) + .permute(0, 2, 3, 4, 1) + .to("cpu", dtype=torch.uint8) + .numpy()[0] + ) + if self.config.offload_vae: + self.vae.to("cpu") + del samples, decoded + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return video + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_vae(self) -> None: + if self._vae_built: + return + cfg = self._ensure_runtime_config() + self.vae_dtype = _get_weight_dtype(cfg.vae.weight_dtype) + cfg.vae.vae_pretrained = resolve_hf_path(cfg.vae.vae_pretrained) + self.vae = _get_vae( + cfg.vae.vae_type, + cfg.vae.vae_pretrained, + device=self.device, + dtype=self.vae_dtype, + config=cfg.vae, + ) + if hasattr(self.vae, "enable_tiling"): + self._configure_vae_tiling() + self._vae_built = True + + def _configure_vae_tiling( + self, + *, + tile_sample_min_width: int | None = None, + tile_sample_stride_width: int | None = None, + tile_sample_min_height: int | None = None, + tile_sample_stride_height: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_num_frames: int | None = None, + ) -> None: + vae = self.vae + min_width = int(tile_sample_min_width or self.config.vae_tile_sample_min_width) + stride_width = int( + tile_sample_stride_width or self.config.vae_tile_sample_stride_width + ) + min_height = int(tile_sample_min_height or self.config.vae_tile_sample_min_height) + stride_height = int( + tile_sample_stride_height or self.config.vae_tile_sample_stride_height + ) + spatial_ratio = int(getattr(vae, "spatial_compression_ratio", 1)) + stride_width = _avoid_degenerate_tile_tail( + sample_extent=self.config.width, + sample_tile_min=min_width, + sample_stride=stride_width, + compression_ratio=spatial_ratio, + ) + stride_height = _avoid_degenerate_tile_tail( + sample_extent=self.config.height, + sample_tile_min=min_height, + sample_stride=stride_height, + compression_ratio=spatial_ratio, + ) + min_frames = int( + tile_sample_min_num_frames + or self.config.vae_tile_sample_min_num_frames + ) + stride_frames = int( + tile_sample_stride_num_frames + or self.config.vae_tile_sample_stride_num_frames + ) + temporal_ratio = int(getattr(vae, "temporal_compression_ratio", 1)) + if temporal_ratio > 1: + min_frames = max(min_frames, temporal_ratio) + stride_frames = max(stride_frames, temporal_ratio) + kwargs = { + "tile_sample_min_height": min_height, + "tile_sample_stride_height": stride_height, + "tile_sample_min_width": min_width, + "tile_sample_stride_width": stride_width, + "tile_sample_min_num_frames": min_frames, + "tile_sample_stride_num_frames": stride_frames, + } + if hasattr(vae, "enable_tiling"): + try: + vae.enable_tiling(**kwargs) + except TypeError: + vae.enable_tiling() + for name, value in kwargs.items(): + if hasattr(vae, name): + setattr(vae, name, value) + if hasattr(vae, "use_framewise_encoding"): + vae.use_framewise_encoding = True + if hasattr(vae, "use_framewise_decoding"): + vae.use_framewise_decoding = True + logger.info( + "[sana-vae] tiling width={} stride_width={} height={} " + "stride_height={} frames={} stride_frames={}", + min_width, + stride_width, + min_height, + stride_height, + min_frames, + stride_frames, + ) + + def _vae_decode(self, latents: Tensor) -> Tensor: + return _vae_decode_ltx2( + self._ensure_runtime_config().vae.vae_type, + self.vae, + latents, + ) + + +@dataclass(kw_only=True) +class SanaWMLTX2LatentRefinerConfig(InstantiateConfig): + """Config for the optional LTX-2 latent refiner component.""" + + _target: type["SanaWMLTX2LatentRefiner"] = field( + default_factory=lambda: SanaWMLTX2LatentRefiner + ) + + refiner_root: str = SANA_WM_REFINER_ROOT + refiner_gemma_root: str = SANA_WM_REFINER_GEMMA_ROOT + refiner_precision: Precision = "bf16" + quant_backend: QuantBackend = "torch" + offload_refiner: bool = False + + +class SanaWMLTX2LatentRefiner(nn.Module): + """Run the optional LTX-2 refinement stage over Stage-1 latents.""" + + config: SanaWMLTX2LatentRefinerConfig + + def __init__(self, config: SanaWMLTX2LatentRefinerConfig) -> None: + super().__init__() + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._refiner_built = False + + @property + def device(self) -> torch.device: + return self._dummy.device + + @torch.inference_mode() + def refine_latents( + self, + *, + latents: Tensor, + prompt: str, + fps: int, + sink_size: int, + seed: int, + block_size: int | None, + kv_max_frames: int, + ) -> Tensor: + """Run the LTX-2 refiner.""" + self._ensure_refiner() + refined = self.refiner.refine_latents( + latents, + prompt, + fps=float(fps), + sink_size=int(sink_size), + seed=int(seed), + progress=True, + block_size=block_size, + kv_max_frames=int(kv_max_frames), + ) + if self.config.offload_refiner: + self.release_runtime() + return refined + + def release_runtime(self) -> None: + """Release refiner tensors.""" + if not self._refiner_built: + return + del self.refiner + self._refiner_built = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() + + def _ensure_refiner(self) -> None: + if self._refiner_built: + return + from sana_wm.refiner import SanaWMLTX2Refiner + + compute_dtype = ( + torch.bfloat16 + if self.config.refiner_precision in {"fp8", "fp4"} + else _get_weight_dtype(self.config.refiner_precision) + ) + self.refiner = SanaWMLTX2Refiner( + refiner_root=resolve_hf_path(self.config.refiner_root), + gemma_root=resolve_hf_path(self.config.refiner_gemma_root), + dtype=compute_dtype, + device=self.device, + precision=self.config.refiner_precision, + quant_backend=self.config.quant_backend, + ) + self._refiner_built = True + + +@dataclass(kw_only=True) +class SanaWMVideoDecoderConfig(DecoderConfig): + """Config for the SANA-WM latent refiner plus VAE decode boundary.""" + + _target: type["SanaWMVideoDecoder"] = field( + default_factory=lambda: SanaWMVideoDecoder + ) + + vae_decoder: SanaWMLTX2VAEDecoderConfig = field( + default_factory=SanaWMLTX2VAEDecoderConfig + ) + refiner: SanaWMLTX2LatentRefinerConfig | None = field( + default_factory=SanaWMLTX2LatentRefinerConfig + ) + + +class SanaWMVideoDecoder(StreamingDecoder[SanaWMVideoDecoderCache]): + """Decode Stage-1 latents, optionally through the LTX-2 refiner.""" + + config: SanaWMVideoDecoderConfig + + def __init__(self, config: SanaWMVideoDecoderConfig) -> None: + super().__init__(config) + self.config = config + self.vae_decoder = config.vae_decoder.setup() + self.refiner = config.refiner.setup() if config.refiner is not None else None + + def initialize_autoregressive_cache( + self, + **context: Any, + ) -> SanaWMVideoDecoderCache: + """Build per-rollout decode/refiner settings.""" + return SanaWMVideoDecoderCache(**context) + + @torch.inference_mode() + def forward( + self, + input: Tensor, + autoregressive_index: int = 0, + cache: SanaWMVideoDecoderCache | None = None, + ) -> SanaWMDecodedVideo: + """Refine/decode one SANA-WM latent rollout.""" + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + cache = cache or SanaWMVideoDecoderCache() + stage1_latent = input + output_latent = input + if self.refiner is not None: + output_latent = self.refiner.refine_latents( + latents=stage1_latent, + prompt=cache.prompt, + fps=cache.fps, + sink_size=cache.sink_size, + seed=cache.refiner_seed, + block_size=cache.refiner_block_size, + kv_max_frames=cache.refiner_kv_max_frames, + ) + self.refiner.release_runtime() + elif cache.save_stage1: + logger.info( + "SANA-WM is already running without the refiner; " + "--save-stage1 does not create an extra output." + ) + + video_hwc = self.vae_decoder.decode_latents(output_latent) + stage1_video_hwc = None + if cache.save_stage1 and self.refiner is not None: + stage1_video_hwc = self.vae_decoder.decode_latents(stage1_latent) + return SanaWMDecodedVideo( + video_hwc=video_hwc, + stage1_video_hwc=stage1_video_hwc, + ) + + @property + def spatial_compression_ratio(self) -> int: + """Pixel side divided by latent side.""" + return SANA_WM_VAE_SPATIAL_COMPRESSION + + @property + def temporal_compression_ratio(self) -> int: + """Pixel frame compression ratio after the first latent.""" + return SANA_WM_VAE_TEMPORAL_COMPRESSION + + +__all__ = [ + "SanaWMDecodedVideo", + "SanaWMLTX2LatentRefiner", + "SanaWMLTX2LatentRefinerConfig", + "SanaWMLTX2VAEDecoder", + "SanaWMLTX2VAEDecoderConfig", + "SanaWMVideoDecoder", + "SanaWMVideoDecoderCache", + "SanaWMVideoDecoderConfig", +] diff --git a/integrations/sana/sana_wm/diffusion.py b/integrations/sana/sana_wm/diffusion.py index 2c26149b3..6b67f01b5 100644 --- a/integrations/sana/sana_wm/diffusion.py +++ b/integrations/sana/sana_wm/diffusion.py @@ -13,240 +13,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SANA-WM-specific diffusion model for first-frame-pinned LTX Euler sampling.""" +"""SANA-WM diffusion config bound to the shared FlashDreams diffusion model.""" from __future__ import annotations -import os -import time from dataclasses import dataclass, field -from typing import Any, cast - -import torch -from loguru import logger -from torch import Tensor from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig -from sana_wm.transformer import ( - SanaWMTransformer, - SanaWMTransformerCache, -) @dataclass(kw_only=True) class SanaWMDiffusionModelConfig(DiffusionModelConfig): - """Diffusion model config for SANA-WM's LTX-style Stage-1 sampler.""" - - _target: type["SanaWMDiffusionModel"] = field( - default_factory=lambda: SanaWMDiffusionModel - ) - - -class SanaWMDiffusionModel(DiffusionModel[SanaWMTransformerCache]): - """Run SANA-WM Stage-1 denoising behind the FlashDreams diffusion boundary.""" - - transformer: SanaWMTransformer - - def __init__(self, config: SanaWMDiffusionModelConfig) -> None: - super().__init__(config) - if not isinstance(self.transformer, SanaWMTransformer): - raise TypeError( - "SanaWMDiffusionModel requires SanaWMTransformerConfig." - ) - - def generate( - self, - autoregressive_index: int, - cache: SanaWMTransformerCache, - input: Any = None, - ) -> tuple[Tensor, "DiffusionModel.FinalState[SanaWMTransformerCache]"]: - """Run SANA-WM's first-frame-pinned LTX Euler denoising loop.""" - del input - if autoregressive_index != 0: - raise ValueError("SANA-WM bidirectional inference has one AR step.") - conditioning = cache.conditioning - if conditioning is None: - raise RuntimeError("SANA-WM diffusion cache has no conditioning.") - - cache.start(autoregressive_index) - latents = self.transformer.initial_latents(conditioning) - if torch.cuda.is_available(): - torch.cuda.synchronize() - t0 = time.perf_counter() - clean_latent = self._sample_ltx_euler(latents, cache) - if torch.cuda.is_available(): - torch.cuda.synchronize() - logger.info( - "[timing] stage1 sample: {:.3f}s (latent shape {})", - time.perf_counter() - t0, - tuple(clean_latent.shape), - ) - - final_state = DiffusionModel.FinalState( - clean_latent=clean_latent, - autoregressive_index=autoregressive_index, - cache=cache, - input=None, - ) - return clean_latent, final_state - - def finalize( - self, - final_state: "DiffusionModel.FinalState[SanaWMTransformerCache]", - ) -> None: - """Finalize the one-shot cache without an extra model forward.""" - final_state.cache.finalize(final_state.autoregressive_index) - - def _sample_ltx_euler( - self, - latents: Tensor, - cache: SanaWMTransformerCache, - ) -> Tensor: - conditioning = cache.conditioning - assert conditioning is not None - - from diffusers import FlowMatchEulerDiscreteScheduler - from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import ( - retrieve_timesteps, - ) - from tqdm import tqdm - - scheduler = FlowMatchEulerDiscreteScheduler(shift=conditioning.flow_shift) - timesteps, _ = retrieve_timesteps( - scheduler, - conditioning.steps, - self.device, - None, - ) - do_cfg = conditioning.cfg_scale > 1.0 - condition_frame_info = dict( - cast( - dict[int, float], - cast(dict[str, object], conditioning.model_kwargs["data_info"]).get( - "condition_frame_info", - {}, - ), - ) - ) - condition_mask = torch.zeros_like(latents) - image_cond_noise_scale = 0.0 - for frame_idx, frame_weight in condition_frame_info.items(): - condition_mask[:, :, int(frame_idx)] = 1 - image_cond_noise_scale = max(image_cond_noise_scale, float(frame_weight)) - - condition_kwargs = _condition_model_kwargs(conditioning.model_kwargs) - uncondition_kwargs = _uncondition_model_kwargs(conditioning.model_kwargs) - if do_cfg and conditioning.uncondition is None: - raise RuntimeError("CFG was requested without negative prompt embeds.") - - init_latents = latents.clone() - generator = torch.Generator(device=self.device).manual_seed(conditioning.seed) - iterator = enumerate(timesteps) - if os.getenv("DPM_TQDM", "False") != "True": - iterator = tqdm(list(iterator)) - - for _, timestep_scalar in iterator: - if image_cond_noise_scale > 0: - latents = _add_noise_to_conditioning_latents( - t=timestep_scalar / 1000.0, - init_latents=init_latents, - latents=latents, - noise_scale=image_cond_noise_scale, - conditioning_mask=condition_mask, - generator=generator, - ) - - timestep = timestep_scalar.expand(condition_mask.shape).float() - timestep = torch.min(timestep, (1 - condition_mask) * 1000.0) - if do_cfg: - assert conditioning.uncondition is not None - noise_pred_uncond = self.transformer.predict_flow( - noisy_latent=latents, - timestep=timestep[:, :1, :, 0, 0], - cache=cache, - input=conditioning.uncondition, - model_kwargs=uncondition_kwargs, - ) - noise_pred_text = self.transformer.predict_flow( - noisy_latent=latents, - timestep=timestep[:, :1, :, 0, 0], - cache=cache, - input=conditioning.condition, - model_kwargs=condition_kwargs, - ) - noise_pred = noise_pred_uncond + conditioning.cfg_scale * ( - noise_pred_text - noise_pred_uncond - ) - del noise_pred_uncond, noise_pred_text - else: - noise_pred = self.transformer.predict_flow( - noisy_latent=latents, - timestep=timestep[:, :1, :, 0, 0], - cache=cache, - input=conditioning.condition, - model_kwargs=condition_kwargs, - ) - - latents_dtype = latents.dtype - latents_shape = latents.shape - batch_size, channels, _frames, _height, _width = latents_shape - denoised_latents = scheduler.step( - -noise_pred.reshape(batch_size, channels, -1).transpose(1, 2), - timestep_scalar, - latents.reshape(batch_size, channels, -1).transpose(1, 2), - per_token_timesteps=timestep.reshape(batch_size, channels, -1)[:, 0], - return_dict=False, - )[0] - denoised_latents = denoised_latents.transpose(1, 2).reshape(latents_shape) - tokens_to_denoise_mask = timestep_scalar / 1000 - 1e-6 < ( - 1.0 - condition_mask - ) - latents = torch.where(tokens_to_denoise_mask, denoised_latents, latents) - if latents.dtype != latents_dtype: - latents = latents.to(latents_dtype) - - return latents.detach() - - -def _condition_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: - """Return model kwargs for the positive prompt branch.""" - return { - key: value - for key, value in model_kwargs.items() - if key != "negative_mask" - } - - -def _uncondition_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: - """Return model kwargs for the negative prompt branch.""" - kwargs = _condition_model_kwargs(model_kwargs) - negative_mask = model_kwargs.get("negative_mask") - if negative_mask is not None: - kwargs["mask"] = negative_mask - return kwargs - + """Diffusion model config for SANA-WM's Stage-1 sampler. -def _add_noise_to_conditioning_latents( - *, - t: Tensor, - init_latents: Tensor, - latents: Tensor, - noise_scale: float, - conditioning_mask: Tensor, - generator: torch.Generator, - eps: float = 1e-6, -) -> Tensor: - from diffusers.utils.torch_utils import randn_tensor + Sana-specific behavior lives in the configured transformer and scheduler; + this config intentionally instantiates the common FlashDreams + :class:`DiffusionModel` instead of a custom ``generate`` implementation. + """ - noise = randn_tensor( - latents.shape, - generator=generator, - device=latents.device, - dtype=latents.dtype, - ) - need_to_noise = conditioning_mask > (1.0 - eps) - noised_latents = init_latents + noise_scale * noise * (t**2) - return torch.where(need_to_noise, noised_latents, latents) + _target: type[DiffusionModel] = field(default_factory=lambda: DiffusionModel) -__all__ = ["SanaWMDiffusionModel", "SanaWMDiffusionModelConfig"] +__all__ = ["SanaWMDiffusionModelConfig"] diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index 02d1163d0..48cada5f6 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -35,11 +35,11 @@ from sana_wm.camera import ( action_string_to_c2w, load_intrinsics, - prepare_camera, resize_center_crop_geometry, snap_num_frames, transform_intrinsics_for_crop, ) +from sana_wm.conditioning import SanaWMI2VConditioningRequest from sana_wm.constants import ( DEFAULT_ACTION, DEFAULT_FPS, @@ -51,7 +51,7 @@ SANA_WM_REFINER_ROOT, SANA_WM_VAE_TEMPORAL_COMPRESSION, ) -from sana_wm.transformer import SanaWMTransformer +from sana_wm.decoder import SanaWMDecodedVideo SamplingAlgo = Literal["auto", "flow_euler_ltx"] """Sampling algorithms exposed by the SANA-WM runner.""" @@ -273,73 +273,54 @@ def run(self) -> None: quant_backend=quant_backend, ) pipeline = pipeline_cfg.setup().to(device).eval() - transformer = pipeline.diffusion_model.transformer - if not isinstance(transformer, SanaWMTransformer): - raise TypeError( - "SANA-WM runner resolved a non-SANA transformer: " - f"{type(transformer).__name__}." - ) sampling_algo = self._sampling_algo() if sampling_algo != "flow_euler_ltx": raise ValueError( "SANA-WM requires flow_euler_ltx for the " f"bidirectional runner; got {sampling_algo!r}." ) - camera = prepare_camera( - c2w, - intrinsics_vec4, - target_size=(DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH), - ) - conditioning = transformer.prepare_conditioning( - image=image, - prompt=prompt, - camera=camera, - num_frames=num_frames, - fps=cfg.fps, - steps=cfg.step, - cfg_scale=cfg.cfg_scale, - flow_shift=cfg.flow_shift, - seed=cfg.seed, - negative_prompt=cfg.negative_prompt, - ) cache = pipeline.initialize_cache( - transformer_context={"conditioning": conditioning} + decoder_context={ + "prompt": prompt, + "fps": cfg.fps, + "save_stage1": cfg.save_stage1, + "refiner_seed": cfg.refiner_seed, + "sink_size": cfg.sink_size, + "refiner_block_size": cfg.refiner_block_size, + "refiner_kv_max_frames": cfg.refiner_kv_max_frames, + } ) - sana_latent = pipeline.generate(0, cache) - pipeline.finalize(0, cache) - transformer.release_stage1_runtime(cache) - del conditioning, cache - - stage1_video_hwc: np.ndarray | None = None - output_latent = sana_latent - if not cfg.no_refiner: - output_latent = transformer.refine_latents( - latents=sana_latent, + decoded = pipeline.generate( + 0, + cache, + input=SanaWMI2VConditioningRequest( + image=image, prompt=prompt, + poses_c2w=c2w, + intrinsics_vec4=intrinsics_vec4, + num_frames=num_frames, fps=cfg.fps, - sink_size=cfg.sink_size, - seed=cfg.refiner_seed, - block_size=cfg.refiner_block_size, - kv_max_frames=cfg.refiner_kv_max_frames, - ) - transformer.release_refiner_runtime() - elif cfg.save_stage1: - logger.info( - "SANA-WM is already running without the refiner; " - "--save-stage1 does not create an extra output." + steps=cfg.step, + cfg_scale=cfg.cfg_scale, + flow_shift=cfg.flow_shift, + seed=cfg.seed, + negative_prompt=cfg.negative_prompt, + ), + ) + pipeline.finalize(0, cache) + if not isinstance(decoded, SanaWMDecodedVideo): + raise TypeError( + "SANA-WM pipeline decoder returned " + f"{type(decoded).__name__}, expected SanaWMDecodedVideo." ) - - video_hwc = transformer.decode_latents(output_latent) - if self.is_rank_zero and cfg.save_stage1 and not cfg.no_refiner: - stage1_video_hwc = transformer.decode_latents(sana_latent) if not self.is_rank_zero: return - _write_video(cfg.output_dir, cfg.name, video_hwc, cfg.fps) - if stage1_video_hwc is not None: + _write_video(cfg.output_dir, cfg.name, decoded.video_hwc, cfg.fps) + if decoded.stage1_video_hwc is not None: _write_video( cfg.output_dir, f"{cfg.name}_stage1", - stage1_video_hwc, + decoded.stage1_video_hwc, cfg.fps, ) @@ -416,13 +397,43 @@ def _pipeline_config( checkpoint_path=cfg.model_path, stage1_precision=cfg.stage1_precision, quant_backend=quant_backend, - refiner_root=cfg.refiner_root, - refiner_gemma_root=cfg.refiner_gemma_root, - refiner_precision=cfg.refiner_precision, - offload_vae=cfg.offload_vae, - offload_refiner=cfg.offload_refiner, + ), + ), + encoder=dict( + config_path=cfg.config_path, + text_encoder=dict( + config_path=cfg.config_path, + stage1_precision=cfg.stage1_precision, + quant_backend=quant_backend, offload_text_encoder=cfg.offload_text_encoder, ), + first_frame_encoder=dict( + config_path=cfg.config_path, + offload_vae=cfg.offload_vae, + ), + camera_encoder=dict( + height=DEFAULT_VIDEO_HEIGHT, + width=DEFAULT_VIDEO_WIDTH, + ), + height=DEFAULT_VIDEO_HEIGHT, + width=DEFAULT_VIDEO_WIDTH, + ), + decoder=dict( + vae_decoder=dict( + config_path=cfg.config_path, + offload_vae=cfg.offload_vae, + ), + refiner=( + None + if cfg.no_refiner + else dict( + refiner_root=cfg.refiner_root, + refiner_gemma_root=cfg.refiner_gemma_root, + refiner_precision=cfg.refiner_precision, + quant_backend=quant_backend, + offload_refiner=cfg.offload_refiner, + ) + ), ), ) diff --git a/integrations/sana/sana_wm/scheduler.py b/integrations/sana/sana_wm/scheduler.py new file mode 100644 index 000000000..4bf1ed288 --- /dev/null +++ b/integrations/sana/sana_wm/scheduler.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 LTX-style Euler scheduler boundary.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Any, cast + +import numpy as np +import torch +from torch import Tensor + +from flashdreams.infra.diffusion.scheduler import ( + FlowPredictor, + Scheduler, + SchedulerConfig, +) +from sana_wm.transformer import SanaWMStage1Conditioning + + +@dataclass(kw_only=True) +class SanaWMLTXEulerSchedulerConfig(SchedulerConfig): + """Config for SANA-WM's LTX-style flow-matching Euler scheduler.""" + + _target: type["SanaWMLTXEulerScheduler"] = field( + default_factory=lambda: SanaWMLTXEulerScheduler + ) + + num_inference_steps: int = 60 + """Default number of Euler steps when using the generic scheduler API.""" + + shift: float = 9.8 + """Default flow-match schedule shift when using the generic scheduler API.""" + + num_train_timesteps: int = 1000 + """Training timestep scale used by the public SANA-WM release.""" + + +class SanaWMLTXEulerScheduler(Scheduler): + """Euler scheduler with SANA-WM per-token timestep support.""" + + config: SanaWMLTXEulerSchedulerConfig + + def __init__(self, config: SanaWMLTXEulerSchedulerConfig) -> None: + super().__init__(config) + self.config = config + + def timesteps( + self, + *, + num_inference_steps: int, + shift: float, + device: torch.device | str, + ) -> Tensor: + """Return diffusers-compatible FlowMatch Euler timesteps.""" + steps = int(num_inference_steps) + if steps <= 0: + raise ValueError(f"num_inference_steps must be > 0, got {steps}.") + num_train_timesteps = int(self.config.num_train_timesteps) + init_timesteps = np.linspace( + 1, + num_train_timesteps, + num_train_timesteps, + dtype=np.float32, + )[::-1].copy() + init_sigmas = init_timesteps / num_train_timesteps + init_sigmas = shift * init_sigmas / (1.0 + (shift - 1.0) * init_sigmas) + sigma_min = float(init_sigmas[-1]) + sigma_max = float(init_sigmas[0]) + timesteps = np.linspace( + sigma_max * num_train_timesteps, + sigma_min * num_train_timesteps, + steps, + ) + sigmas = timesteps / num_train_timesteps + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) + return torch.from_numpy(sigmas * num_train_timesteps).to(device=device) + + def step_ltx( + self, + *, + model_output: Tensor, + timestep: Tensor, + next_timestep: Tensor, + sample: Tensor, + per_token_timesteps: Tensor | None = None, + schedule_timesteps: Tensor | None = None, + ) -> Tensor: + """Apply one SANA-WM LTX Euler step in token layout. + + Args: + model_output: Flow tensor in ``[B, N, C]`` token layout. + timestep: Current scalar training-scale timestep. + next_timestep: Next scalar training-scale timestep. + sample: Current latent in ``[B, N, C]`` token layout. + per_token_timesteps: Optional ``[B, N]`` timestep table. Tokens + with timestep ``0`` stay at sigma zero, which matches the + first-frame-pinning branch used by SANA-WM. + schedule_timesteps: Optional full scheduler timestep table used + to find the next lower sigma for each per-token timestep. + + Returns: + Updated sample in ``[B, N, C]`` layout. + """ + num_train_timesteps = float(self.config.num_train_timesteps) + sample = sample.to(torch.float32) + if per_token_timesteps is None: + sigma = timestep.to(device=sample.device, dtype=sample.dtype) + sigma_next = next_timestep.to(device=sample.device, dtype=sample.dtype) + dt = (sigma_next - sigma) / num_train_timesteps + return (sample + dt * model_output).to(model_output.dtype) + + per_token_sigmas = ( + per_token_timesteps.to(device=sample.device, dtype=sample.dtype) + / num_train_timesteps + ) + if schedule_timesteps is not None: + sigmas = ( + schedule_timesteps.to(device=sample.device, dtype=sample.dtype) + / num_train_timesteps + ) + sigmas = sigmas[:, None, None] + lower_mask = sigmas < per_token_sigmas[None] - 1e-6 + lower_sigmas = (lower_mask * sigmas).max(dim=0).values + else: + next_sigma_scalar = ( + next_timestep.to(device=sample.device, dtype=sample.dtype) + / num_train_timesteps + ) + lower_sigmas = torch.where( + per_token_sigmas > next_sigma_scalar + 1e-6, + next_sigma_scalar.expand_as(per_token_sigmas), + torch.zeros_like(per_token_sigmas), + ) + dt = per_token_sigmas - lower_sigmas + return sample + dt.unsqueeze(-1) * model_output + + def sample( + self, + initial_noise: Tensor, + predict_flow: FlowPredictor, + rng: torch.Generator | None = None, + context: Any = None, + ) -> Tensor: + """Run the generic scalar-timestep Euler loop.""" + del rng + if isinstance(context, SanaWMStage1Conditioning): + return self._sample_conditioned( + initial_noise=initial_noise, + predict_flow=predict_flow, + conditioning=context, + ) + + timesteps = self.timesteps( + num_inference_steps=self.config.num_inference_steps, + shift=self.config.shift, + device=initial_noise.device, + ) + noisy = initial_noise + for index, timestep in enumerate(timesteps[:-1]): + flow = predict_flow(noisy, timestep.to(dtype=initial_noise.dtype)) + noisy = self.step_ltx( + model_output=flow, + timestep=timestep, + next_timestep=timesteps[index + 1], + sample=noisy, + ) + return noisy.to(initial_noise.dtype) + + def _sample_conditioned( + self, + *, + initial_noise: Tensor, + predict_flow: FlowPredictor, + conditioning: SanaWMStage1Conditioning, + ) -> Tensor: + """Run SANA-WM's first-frame-pinned LTX Euler denoising loop.""" + latents = initial_noise + timesteps = self.timesteps( + num_inference_steps=conditioning.steps, + shift=conditioning.flow_shift, + device=latents.device, + ) + condition_frame_info = dict( + cast( + dict[int, float], + cast(dict[str, object], conditioning.model_kwargs["data_info"]).get( + "condition_frame_info", + {}, + ), + ) + ) + condition_mask = torch.zeros_like(latents) + image_cond_noise_scale = 0.0 + for frame_idx, frame_weight in condition_frame_info.items(): + condition_mask[:, :, int(frame_idx)] = 1 + image_cond_noise_scale = max(image_cond_noise_scale, float(frame_weight)) + + init_latents = latents.clone() + generator = torch.Generator(device=latents.device).manual_seed( + conditioning.seed + ) + iterator = enumerate(timesteps[:-1]) + if os.getenv("DPM_TQDM", "False") == "True": + from tqdm import tqdm + + iterator = tqdm(list(iterator)) + + for step_index, timestep_scalar in iterator: + if image_cond_noise_scale > 0: + latents = _add_noise_to_conditioning_latents( + t=timestep_scalar / self.config.num_train_timesteps, + init_latents=init_latents, + latents=latents, + noise_scale=image_cond_noise_scale, + conditioning_mask=condition_mask, + generator=generator, + ) + + timestep = timestep_scalar.expand(condition_mask.shape).float() + timestep = torch.min( + timestep, + (1 - condition_mask) * float(self.config.num_train_timesteps), + ) + noise_pred = predict_flow(latents, timestep[:, :1, :, 0, 0]) + + latents_dtype = latents.dtype + latents_shape = latents.shape + batch_size, channels, _frames, _height, _width = latents_shape + denoised_latents = self.step_ltx( + model_output=-noise_pred.reshape( + batch_size, + channels, + -1, + ).transpose(1, 2), + timestep=timestep_scalar, + next_timestep=timesteps[step_index + 1], + sample=latents.reshape(batch_size, channels, -1).transpose(1, 2), + per_token_timesteps=timestep.reshape(batch_size, channels, -1)[:, 0], + schedule_timesteps=timesteps, + ) + denoised_latents = denoised_latents.transpose(1, 2).reshape(latents_shape) + tokens_to_denoise_mask = ( + timestep_scalar / self.config.num_train_timesteps - 1e-6 + ) < (1.0 - condition_mask) + latents = torch.where(tokens_to_denoise_mask, denoised_latents, latents) + if latents.dtype != latents_dtype: + latents = latents.to(latents_dtype) + + return latents.detach() + + def add_noise( + self, + clean_input: Tensor, + timestep: Tensor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Apply forward flow-match corruption at ``timestep``.""" + sigma = timestep.to(device=clean_input.device, dtype=clean_input.dtype) + sigma = sigma / float(self.config.num_train_timesteps) + noise = torch.randn( + clean_input.shape, + generator=rng, + device=clean_input.device, + dtype=clean_input.dtype, + ) + return ((1.0 - sigma) * clean_input + sigma * noise).to(clean_input.dtype) + + +def _add_noise_to_conditioning_latents( + *, + t: Tensor, + init_latents: Tensor, + latents: Tensor, + noise_scale: float, + conditioning_mask: Tensor, + generator: torch.Generator, + eps: float = 1e-6, +) -> Tensor: + noise = torch.randn( + latents.shape, + generator=generator, + device=latents.device, + dtype=latents.dtype, + ) + need_to_noise = conditioning_mask > (1.0 - eps) + noised_latents = init_latents + noise_scale * noise * (t**2) + return torch.where(need_to_noise, noised_latents, latents) + + +__all__ = [ + "SanaWMLTXEulerScheduler", + "SanaWMLTXEulerSchedulerConfig", +] diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index 87cbfe21f..e6615e9a9 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -629,10 +629,8 @@ def forward( x_norm = self.norm1(x).reshape(batch, frames, -1, channels) attn_in = _modulate(x_norm, shift_msa, scale_msa).reshape(batch, tokens, channels) - del x_norm, shift_msa, scale_msa attn_out = self.attn(attn_in, **kwargs).reshape(batch, frames, -1, channels) x = x + (gate_msa * attn_out).reshape(batch, tokens, channels) - del attn_in, attn_out, gate_msa if plucker_emb is not None: x = x + self.plucker_proj(plucker_emb) @@ -641,13 +639,9 @@ def forward( x_norm = self.norm2(x).reshape(batch, frames, -1, channels) mlp_in = _modulate(x_norm, shift_mlp, scale_mlp).reshape(batch, tokens, channels) - del x_norm, shift_mlp, scale_mlp mlp_out = self.mlp(mlp_in, frames=frames, height=height, width=width) - del mlp_in mlp_out = mlp_out.reshape(batch, frames, -1, channels) - out = x + (gate_mlp * mlp_out).reshape(batch, tokens, channels) - del gate_mlp, mlp_out - return out + return x + (gate_mlp * mlp_out).reshape(batch, tokens, channels) class SanaWMStage1Model(nn.Module): diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py index 56b8482fc..335db4e9c 100644 --- a/integrations/sana/sana_wm/transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -18,14 +18,10 @@ from __future__ import annotations import gc -import os -import time -from collections.abc import Mapping from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any, Literal, cast -import numpy as np import torch import torch.nn as nn import yaml @@ -43,8 +39,6 @@ DEFAULT_VIDEO_WIDTH, SANA_WM_CONFIG_PATH, SANA_WM_MODEL_PATH, - SANA_WM_REFINER_GEMMA_ROOT, - SANA_WM_REFINER_ROOT, ) from sana_wm.quant import ( TorchScaledMMFP4Recipe, @@ -121,66 +115,12 @@ class SanaWMTransformerConfig(TransformerConfig): quant_backend: QuantBackend = "auto" """Low-precision linear replacement backend for FP8/FP4 modes.""" - refiner_root: str = SANA_WM_REFINER_ROOT - """LTX-2 refiner root path or ``hf://`` URI.""" - - refiner_gemma_root: str = SANA_WM_REFINER_GEMMA_ROOT - """Gemma text-encoder root for the refiner.""" - - refiner_precision: Precision = "bf16" - """Refiner precision requested by the runner.""" - - offload_vae: bool = False - """Move the VAE to CPU between encode/decode phases.""" - - offload_text_encoder: bool = False - """Move the Stage-1 text encoder to CPU between prompt encodes.""" - - offload_refiner: bool = False - """Release the LTX-2 refiner after each refinement.""" - height: int = DEFAULT_VIDEO_HEIGHT """Pixel height used by the public SANA-WM bidirectional release.""" width: int = DEFAULT_VIDEO_WIDTH """Pixel width used by the public SANA-WM bidirectional release.""" - vae_tile_sample_min_width: int = 256 - """Pixel width tile size used for LTX-2 VAE decode.""" - - vae_tile_sample_stride_width: int = 224 - """Pixel width stride used for LTX-2 VAE decode.""" - - vae_tile_sample_min_height: int = 256 - """Pixel height tile size used for LTX-2 VAE decode.""" - - vae_tile_sample_stride_height: int = 192 - """Pixel height stride used for LTX-2 VAE decode.""" - - vae_tile_sample_min_num_frames: int = 24 - """Pixel-frame temporal tile size used for LTX-2 VAE decode.""" - - vae_tile_sample_stride_num_frames: int = 8 - """Pixel-frame temporal tile stride used for LTX-2 VAE decode.""" - - vae_oom_retry_tile_sample_min_width: int = 128 - """Smaller pixel width tile size for one VAE decode OOM retry.""" - - vae_oom_retry_tile_sample_stride_width: int = 64 - """Smaller pixel width tile stride for one VAE decode OOM retry.""" - - vae_oom_retry_tile_sample_min_height: int = 128 - """Smaller pixel height tile size for one VAE decode OOM retry.""" - - vae_oom_retry_tile_sample_stride_height: int = 64 - """Smaller pixel height tile stride for one VAE decode OOM retry.""" - - vae_oom_retry_tile_sample_min_num_frames: int = 16 - """Smaller temporal tile size for one VAE decode OOM retry.""" - - vae_oom_retry_tile_sample_stride_num_frames: int = 8 - """Smaller temporal tile stride for one VAE decode OOM retry.""" - class SanaWMTransformer(Transformer[SanaWMTransformerCache]): """FlashDreams adapter for the SANA-WM Stage-1 model call.""" @@ -193,13 +133,7 @@ def __init__(self, config: SanaWMTransformerConfig) -> None: self._model_path: str | None = None self.weight_dtype: torch.dtype | None = None self._model_built = False - self._vae_built = False - self._text_encoder_built = False - self._refiner_built = False self._stage1_quantized = False - self._streaming_prompt_cache: dict[ - tuple[object, ...], tuple[Tensor, Tensor, Tensor, Tensor] - ] = {} @property def latent_shape(self) -> tuple[int, ...]: @@ -234,15 +168,24 @@ def predict_flow( model_kwargs: dict[str, object] | None = None, ) -> Tensor: """Execute one SANA-WM DiT flow prediction.""" + if isinstance(input, SanaWMStage1Conditioning): + cache.conditioning = input + return self._predict_conditioned( + noisy_latent=noisy_latent, + timestep=timestep, + conditioning=input, + ) + conditioning = _require_conditioning(cache) - prompt_embeds = cast(Tensor, input) if input is not None else conditioning.condition + prompt_embeds = ( + cast(Tensor, input) if input is not None else conditioning.condition + ) kwargs = conditioning.model_kwargs if model_kwargs is None else model_kwargs - self._ensure_model() - return self.model( + return self._predict_with_prompt( noisy_latent, timestep, prompt_embeds, - **kwargs, + kwargs, ) def finalize_kv_cache( @@ -263,100 +206,33 @@ def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: """SANA-WM Stage-1 latents are already in model layout.""" return x - @torch.inference_mode() - def prepare_conditioning( + def initial_noise( self, *, - image: Any, - prompt: str, - camera: Mapping[str, Tensor], - num_frames: int, - fps: int, - steps: int, - cfg_scale: float, - flow_shift: float | None, - seed: int, - negative_prompt: str, - ) -> SanaWMStage1Conditioning: - """Encode first-frame, text, and camera tensors for one rollout.""" - del fps - cfg = self._ensure_runtime_config() - weight_dtype = self._ensure_weight_dtype() - self._ensure_vae() - if self.config.offload_vae: - self.vae.to(self.device) - - from torchvision import transforms as T - - img = (T.ToTensor()(image) * 2.0 - 1.0).unsqueeze(0).unsqueeze(2) - first_latent = self._vae_encode( - img.to(self.device, dtype=self.vae_dtype), - ).to(weight_dtype) - if self.config.offload_vae: - self.vae.to("cpu") - torch.cuda.empty_cache() - - cond, cond_mask, neg, neg_mask = self._encode_prompts( - prompt, - negative_prompt, - ) - cond, cond_mask, neg, neg_mask = self._pad_text_for_quant( - cond, - cond_mask, - neg, - neg_mask, - ) - - raymap = camera["raymap"].unsqueeze(0).to(self.device, dtype=weight_dtype) - chunk_plucker = camera["chunk_plucker"].unsqueeze(0).to( - self.device, - dtype=weight_dtype, - ) - model_kwargs_extra: dict[str, object] = {} - if cfg_scale > 1.0: - model_kwargs_extra["negative_mask"] = neg_mask - uncondition = neg + latent_shape: tuple[int, ...], + rng: torch.Generator | None, + cache: SanaWMTransformerCache, + input: Any = None, + ) -> Tensor: + """Draw SANA-WM's first-frame-pinned Stage-1 initial latent.""" + del latent_shape, rng + if isinstance(input, SanaWMStage1Conditioning): + cache.conditioning = input + conditioning = input else: - uncondition = None - - latent_t = (num_frames - 1) // int(cfg.vae.vae_stride[0]) + 1 - latent_h = self.config.height // int(cfg.vae.vae_stride[-1]) - latent_w = self.config.width // int(cfg.vae.vae_stride[-1]) - chunk_index = self._chunk_index(latent_t) - model_kwargs: dict[str, object] = { - "data_info": { - "img_hw": torch.tensor( - [[self.config.height, self.config.width]], - dtype=torch.float, - device=self.device, - ), - "condition_frame_info": {0: 0.0}, - }, - "mask": cond_mask, - "camera_conditions": raymap, - "chunk_plucker": chunk_plucker, - **model_kwargs_extra, - } - if chunk_index is not None: - model_kwargs["chunk_index"] = chunk_index - - return SanaWMStage1Conditioning( - condition=cond, - uncondition=uncondition, - model_kwargs=model_kwargs, - first_latent=first_latent, - latent_shape=( - 1, - int(first_latent.shape[1]), - latent_t, - latent_h, - latent_w, - ), - cfg_scale=float(cfg_scale), - flow_shift=self._resolve_flow_shift(flow_shift), - steps=int(steps), - seed=int(seed), - ) + conditioning = _require_conditioning(cache) + return self.initial_latents(conditioning) + + def postprocess_clean_latent( + self, + clean_latent: Tensor, + cache: SanaWMTransformerCache, + input: Any = None, + ) -> Tensor: + """Release Stage-1 runtime before the pipeline enters decode/refine.""" + del input + self.release_stage1_runtime(cache) + return clean_latent def initial_latents(self, conditioning: SanaWMStage1Conditioning) -> Tensor: """Draw Stage-1 initial noise and pin the encoded first frame.""" @@ -370,84 +246,6 @@ def initial_latents(self, conditioning: SanaWMStage1Conditioning) -> Tensor: latents[:, :, :1] = conditioning.first_latent return latents - @torch.inference_mode() - def decode_latents(self, latents: Tensor) -> np.ndarray: - """Decode SANA-WM VAE latents to ``uint8`` HWC video.""" - self._ensure_vae() - if self.config.offload_vae: - self.vae.to(self.device) - samples = latents.to(device=self.device, dtype=self.vae_dtype) - if torch.cuda.is_available(): - torch.cuda.synchronize() - t0 = time.perf_counter() - retry_decode = False - try: - decoded = self._vae_decode(samples) - except torch.OutOfMemoryError: - retry_decode = True - - if retry_decode: - if torch.cuda.is_available(): - torch.cuda.synchronize() - torch.cuda.empty_cache() - gc.collect() - logger.warning( - "[sana-vae] decode OOM; retrying with smaller tiles " - "width={} stride_width={} height={} stride_height={} " - "frames={} stride_frames={}", - self.config.vae_oom_retry_tile_sample_min_width, - self.config.vae_oom_retry_tile_sample_stride_width, - self.config.vae_oom_retry_tile_sample_min_height, - self.config.vae_oom_retry_tile_sample_stride_height, - self.config.vae_oom_retry_tile_sample_min_num_frames, - self.config.vae_oom_retry_tile_sample_stride_num_frames, - ) - self._configure_vae_tiling( - tile_sample_min_width=( - self.config.vae_oom_retry_tile_sample_min_width - ), - tile_sample_stride_width=( - self.config.vae_oom_retry_tile_sample_stride_width - ), - tile_sample_min_height=( - self.config.vae_oom_retry_tile_sample_min_height - ), - tile_sample_stride_height=( - self.config.vae_oom_retry_tile_sample_stride_height - ), - tile_sample_min_num_frames=( - self.config.vae_oom_retry_tile_sample_min_num_frames - ), - tile_sample_stride_num_frames=( - self.config.vae_oom_retry_tile_sample_stride_num_frames - ), - ) - if torch.cuda.is_available(): - torch.cuda.empty_cache() - decoded = self._vae_decode(samples) - if torch.cuda.is_available(): - torch.cuda.synchronize() - logger.info( - "[timing] vae decode: {:.3f}s (latent T={} -> pixels {})", - time.perf_counter() - t0, - latents.shape[2], - tuple(decoded.shape) if isinstance(decoded, Tensor) else "list", - ) - if isinstance(decoded, list): - decoded = torch.stack(decoded, dim=0) - video = ( - torch.clamp(127.5 * decoded + 127.5, 0, 255) - .permute(0, 2, 3, 4, 1) - .to("cpu", dtype=torch.uint8) - .numpy()[0] - ) - if self.config.offload_vae: - self.vae.to("cpu") - del samples, decoded - if torch.cuda.is_available(): - torch.cuda.empty_cache() - return video - def release_stage1_runtime( self, cache: SanaWMTransformerCache | None = None, @@ -462,9 +260,8 @@ def release_stage1_runtime( free_before_gib = None if cache is not None: cache.conditioning = None - self._streaming_prompt_cache.clear() - for attr in ("model", "text_encoder"): + for attr in ("model",): module = getattr(self, attr, None) if module is None: continue @@ -477,10 +274,7 @@ def release_stage1_runtime( pass setattr(self, attr, None) - if hasattr(self, "tokenizer"): - self.tokenizer = None self._model_built = False - self._text_encoder_built = False self._stage1_quantized = False if torch.cuda.is_available(): torch.cuda.empty_cache() @@ -494,51 +288,55 @@ def release_stage1_runtime( free_after / (1024**3), ) else: - logger.info("[stage1] released Stage-1 runtime before decode/refine") + logger.info( + "[stage1] released Stage-1 runtime before decode/refine" + ) except RuntimeError: logger.info("[stage1] released Stage-1 runtime before decode/refine") gc.collect() - def release_refiner_runtime(self) -> None: - """Release refiner tensors before VAE decode.""" - self._release_refiner() - - def refine_latents( + def _predict_conditioned( self, *, - latents: Tensor, - prompt: str, - fps: int, - sink_size: int, - seed: int, - block_size: int | None, - kv_max_frames: int, + noisy_latent: Tensor, + timestep: Tensor, + conditioning: SanaWMStage1Conditioning, ) -> Tensor: - """Run the LTX-2 refiner.""" - self._ensure_refiner() - sigmas = torch.tensor( - self._stage2_sigmas(), - dtype=torch.float32, - device=self.device, + if conditioning.cfg_scale <= 1.0: + return self._predict_with_prompt( + noisy_latent, + timestep, + conditioning.condition, + _condition_model_kwargs(conditioning.model_kwargs), + ) + if conditioning.uncondition is None: + raise RuntimeError("CFG was requested without negative prompt embeds.") + + noise_pred = self._predict_with_prompt( + torch.cat([noisy_latent, noisy_latent], dim=0), + torch.cat([timestep, timestep], dim=0), + torch.cat([conditioning.uncondition, conditioning.condition], dim=0), + _batched_cfg_model_kwargs(conditioning.model_kwargs), ) - logger.info( - "[refiner] {}-step Euler, start_sigma={:.4f}", - len(sigmas) - 1, - float(sigmas[0]), + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0) + return noise_pred_uncond + conditioning.cfg_scale * ( + noise_pred_text - noise_pred_uncond ) - refined = self.refiner.refine_latents( - latents, - prompt, - fps=float(fps), - sink_size=int(sink_size), - seed=int(seed), - progress=True, - block_size=block_size, - kv_max_frames=int(kv_max_frames), + + def _predict_with_prompt( + self, + noisy_latent: Tensor, + timestep: Tensor, + prompt_embeds: Tensor, + model_kwargs: dict[str, object], + ) -> Tensor: + self._ensure_model() + return self.model( + noisy_latent, + timestep, + prompt_embeds, + **model_kwargs, ) - if self.config.offload_refiner: - self._release_refiner() - return refined def _ensure_runtime_config(self) -> Any: if self._runtime_config is not None: @@ -553,35 +351,6 @@ def _ensure_weight_dtype(self) -> torch.dtype: ) return self.weight_dtype - def _ensure_vae(self) -> None: - if self._vae_built: - return - cfg = self._ensure_runtime_config() - self.vae_dtype = _get_weight_dtype(cfg.vae.weight_dtype) - cfg.vae.vae_pretrained = resolve_hf_path(cfg.vae.vae_pretrained) - self.vae = _get_vae( - cfg.vae.vae_type, - cfg.vae.vae_pretrained, - device=self.device, - dtype=self.vae_dtype, - config=cfg.vae, - ) - if hasattr(self.vae, "enable_tiling"): - self._configure_vae_tiling() - self._vae_built = True - - def _ensure_text_encoder(self) -> None: - if self._text_encoder_built: - return - cfg = self._ensure_runtime_config() - self.tokenizer, self.text_encoder = _get_tokenizer_and_text_encoder( - name=cfg.text_encoder.text_encoder_name, - device=self.device, - ) - if self.config.offload_text_encoder: - self.text_encoder.to("cpu") - self._text_encoder_built = True - def _ensure_model(self) -> None: if self._model_built: self._prepare_stage1_quant() @@ -614,205 +383,6 @@ def _ensure_model(self) -> None: self._model_built = True self._prepare_stage1_quant() - def _ensure_refiner(self) -> None: - if self._refiner_built: - return - from sana_wm.refiner import SanaWMLTX2Refiner - - compute_dtype = ( - torch.bfloat16 - if self.config.refiner_precision in {"fp8", "fp4"} - else _get_weight_dtype(self.config.refiner_precision) - ) - quant_backend: QuantBackend = ( - "torch" if self.config.quant_backend == "auto" else self.config.quant_backend - ) - self.refiner = SanaWMLTX2Refiner( - refiner_root=resolve_hf_path(self.config.refiner_root), - gemma_root=resolve_hf_path(self.config.refiner_gemma_root), - dtype=compute_dtype, - device=self.device, - precision=self.config.refiner_precision, - quant_backend=quant_backend, - ) - self._refiner_built = True - - def _release_refiner(self) -> None: - if not self._refiner_built: - return - del self.refiner - self._refiner_built = False - if torch.cuda.is_available(): - torch.cuda.empty_cache() - gc.collect() - - def _configure_vae_tiling( - self, - *, - tile_sample_min_width: int | None = None, - tile_sample_stride_width: int | None = None, - tile_sample_min_height: int | None = None, - tile_sample_stride_height: int | None = None, - tile_sample_min_num_frames: int | None = None, - tile_sample_stride_num_frames: int | None = None, - ) -> None: - vae = self.vae - min_width = int(tile_sample_min_width or self.config.vae_tile_sample_min_width) - stride_width = int( - tile_sample_stride_width or self.config.vae_tile_sample_stride_width - ) - min_height = int( - tile_sample_min_height or self.config.vae_tile_sample_min_height - ) - stride_height = int( - tile_sample_stride_height or self.config.vae_tile_sample_stride_height - ) - spatial_ratio = int(getattr(vae, "spatial_compression_ratio", 1)) - stride_width = _avoid_degenerate_tile_tail( - sample_extent=self.config.width, - sample_tile_min=min_width, - sample_stride=stride_width, - compression_ratio=spatial_ratio, - ) - stride_height = _avoid_degenerate_tile_tail( - sample_extent=self.config.height, - sample_tile_min=min_height, - sample_stride=stride_height, - compression_ratio=spatial_ratio, - ) - min_frames = int( - tile_sample_min_num_frames - or self.config.vae_tile_sample_min_num_frames - ) - stride_frames = int( - tile_sample_stride_num_frames - or self.config.vae_tile_sample_stride_num_frames - ) - temporal_ratio = int(getattr(vae, "temporal_compression_ratio", 1)) - if temporal_ratio > 1: - min_frames = max(min_frames, temporal_ratio) - stride_frames = max(stride_frames, temporal_ratio) - kwargs = { - "tile_sample_min_height": min_height, - "tile_sample_stride_height": stride_height, - "tile_sample_min_width": min_width, - "tile_sample_stride_width": stride_width, - "tile_sample_min_num_frames": min_frames, - "tile_sample_stride_num_frames": stride_frames, - } - if hasattr(vae, "enable_tiling"): - try: - vae.enable_tiling(**kwargs) - except TypeError: - vae.enable_tiling() - for name, value in kwargs.items(): - if hasattr(vae, name): - setattr(vae, name, value) - if hasattr(vae, "use_framewise_encoding"): - vae.use_framewise_encoding = True - if hasattr(vae, "use_framewise_decoding"): - vae.use_framewise_decoding = True - logger.info( - "[sana-vae] tiling width={} stride_width={} height={} " - "stride_height={} frames={} stride_frames={}", - min_width, - stride_width, - min_height, - stride_height, - min_frames, - stride_frames, - ) - - def _encode_prompts( - self, - prompt: str, - negative_prompt: str, - ) -> tuple[Tensor, Tensor, Tensor, Tensor]: - cfg = self._ensure_runtime_config() - self._ensure_text_encoder() - max_length = cfg.text_encoder.model_max_length - chi_prompt = "\n".join(cfg.text_encoder.chi_prompt or []) - if chi_prompt: - prompt = chi_prompt + prompt - max_length_all = len(self.tokenizer.encode(chi_prompt)) + max_length - 2 - else: - max_length_all = max_length - - key = ( - prompt, - negative_prompt, - str(self.device), - str(self._ensure_weight_dtype()), - self.config.stage1_precision, - self.config.quant_backend, - ) - if key in self._streaming_prompt_cache: - return self._streaming_prompt_cache[key] - - move_text_encoder = self.config.offload_text_encoder or ( - next(self.text_encoder.parameters()).device != self.device - ) - if move_text_encoder: - self.text_encoder.to(self.device) - - def encode(text: str, length: int) -> tuple[Tensor, Tensor]: - tokens = self.tokenizer( - [text], - max_length=length, - padding="max_length", - truncation=True, - return_tensors="pt", - ).to(self.device) - return self.text_encoder(tokens.input_ids, tokens.attention_mask)[0], ( - tokens.attention_mask - ) - - try: - cond, cond_mask = encode(prompt, max_length_all) - select = [0] + list(range(-max_length + 1, 0)) - cond = cond[:, None][:, :, select] - cond_mask = cond_mask[:, select] - neg, neg_mask = encode(negative_prompt, max_length) - result = (cond, cond_mask, neg[:, None], neg_mask) - self._streaming_prompt_cache.clear() - self._streaming_prompt_cache[key] = result - return result - finally: - if move_text_encoder: - self.text_encoder.to("cpu") - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - def _pad_text_for_quant( - self, - cond: Tensor, - cond_mask: Tensor, - neg: Tensor, - neg_mask: Tensor, - ) -> tuple[Tensor, Tensor, Tensor, Tensor]: - if self.config.stage1_precision == "bf16": - return cond, cond_mask, neg, neg_mask - multiple = int(os.environ.get("SANA_WM_STAGE1_NVFP4_TEXT_PAD_MULTIPLE", "8")) - if multiple <= 1: - return cond, cond_mask, neg, neg_mask - - def pad_pair(text: Tensor, mask: Tensor) -> tuple[Tensor, Tensor]: - pad = (-text.shape[-2]) % multiple - if pad == 0: - return text, mask - text_shape = list(text.shape) - text_shape[-2] = pad - mask_shape = list(mask.shape) - mask_shape[-1] = pad - return ( - torch.cat([text, text.new_zeros(text_shape)], dim=-2), - torch.cat([mask, mask.new_zeros(mask_shape)], dim=-1), - ) - - cond, cond_mask = pad_pair(cond, cond_mask) - neg, neg_mask = pad_pair(neg, neg_mask) - return cond, cond_mask, neg, neg_mask - def _prepare_stage1_quant(self) -> None: if self._stage1_quantized or self.config.stage1_precision == "bf16": return @@ -842,39 +412,6 @@ def _prepare_stage1_quant(self) -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() - def _resolve_flow_shift(self, override: float | None) -> float: - cfg = self._ensure_runtime_config() - if override is not None: - return float(override) - if cfg.scheduler.inference_flow_shift is not None: - return float(cfg.scheduler.inference_flow_shift) - return float(cfg.scheduler.flow_shift) - - def _chunk_index(self, latent_frames: int) -> list[int] | None: - return _chunk_index_from_config( - self._ensure_runtime_config(), - num_frames=latent_frames, - ) - - def _vae_encode(self, images: Tensor) -> Tensor: - return _vae_encode_ltx2( - self._ensure_runtime_config().vae.vae_type, - self.vae, - images, - device=self.device, - ) - - def _vae_decode(self, latents: Tensor) -> Tensor: - return _vae_decode_ltx2( - self._ensure_runtime_config().vae.vae_type, - self.vae, - latents, - ) - - @staticmethod - def _stage2_sigmas() -> tuple[float, ...]: - return (0.909375, 0.725, 0.421875, 0.0) - def _require_conditioning( cache: SanaWMTransformerCache, @@ -884,6 +421,30 @@ def _require_conditioning( return cache.conditioning +def _condition_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: + """Return model kwargs for the positive prompt branch.""" + return { + key: value + for key, value in model_kwargs.items() + if key != "negative_mask" + } + + +def _batched_cfg_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: + """Return model kwargs for a single batched negative/positive CFG forward.""" + kwargs = _condition_model_kwargs(model_kwargs) + mask = kwargs.get("mask") + negative_mask = model_kwargs.get("negative_mask") + if isinstance(mask, Tensor): + mask_uncond = negative_mask if isinstance(negative_mask, Tensor) else mask + kwargs["mask"] = torch.cat([mask_uncond, mask], dim=0) + for key in ("camera_conditions", "chunk_plucker"): + value = kwargs.get(key) + if isinstance(value, Tensor): + kwargs[key] = torch.cat([value, value], dim=0) + return kwargs + + def _avoid_degenerate_tile_tail( *, sample_extent: int, @@ -955,7 +516,12 @@ def _get_tokenizer_and_text_encoder(*args: Any, **kwargs: Any) -> tuple[Any, nn. model_id = _TEXT_ENCODER_MODEL_IDS.get(str(name)) if model_id is None: raise ValueError(f"Unsupported SANA-WM text encoder: {name!r}") - from transformers import AutoModelForCausalLM, AutoTokenizer, T5EncoderModel, T5Tokenizer + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + T5EncoderModel, + T5Tokenizer, + ) if "T5" in str(name): tokenizer = T5Tokenizer.from_pretrained(model_id) diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 5cb0a8dfd..c626a5b26 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -19,12 +19,15 @@ import os from pathlib import Path +from types import SimpleNamespace +import numpy as np import pytest import torch +import sana_wm.conditioning as conditioning_module +import sana_wm.decoder as decoder_module import sana_wm.refiner as refiner_module -import sana_wm.transformer as transformer_module try: import tomllib @@ -44,16 +47,31 @@ from sana_wm.runner import ( SanaWMRunner, SanaWMRunnerConfig, + _pipeline_config, _precision_env_updates, _resolve_quant_backend, _temporary_environment, _validate_precision_request, ) -from sana_wm.diffusion import ( - _condition_model_kwargs, - _uncondition_model_kwargs, +from sana_wm.conditioning import ( + SanaWMCameraConditioningEncoderConfig, + SanaWMCameraRequest, + SanaWMConditioningEncoderConfig, + SanaWMFirstFrameEncoderConfig, + SanaWMTextPromptEncoderConfig, + SanaWMTextPromptRequest, +) +from sana_wm.decoder import ( + SanaWMDecodedVideo, + SanaWMLTX2LatentRefinerConfig, + SanaWMLTX2VAEDecoderConfig, + SanaWMVideoDecoderConfig, ) from sana_wm.diffusion import SanaWMDiffusionModelConfig +from sana_wm.scheduler import ( + SanaWMLTXEulerScheduler, + SanaWMLTXEulerSchedulerConfig, +) from sana_wm.transformer import ( SanaWMTransformerCache, SanaWMTransformerConfig, @@ -76,6 +94,7 @@ ) from flashdreams.infra.config import derive_config +from flashdreams.infra.diffusion.model import DiffusionModel pytestmark = pytest.mark.ci_cpu @@ -102,14 +121,164 @@ def test_runner_has_description() -> None: def test_pipeline_uses_sana_diffusion_model() -> None: - """Keep the public runner wired to the SANA-WM diffusion model.""" + """Keep the public runner wired to explicit SANA-WM boundaries.""" + pipeline = PIPELINE_SANA_WM_BIDIRECTIONAL transformer = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.transformer - assert isinstance( - PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model, - SanaWMDiffusionModelConfig, - ) + assert isinstance(pipeline.encoder, SanaWMConditioningEncoderConfig) + assert isinstance(pipeline.decoder, SanaWMVideoDecoderConfig) + assert isinstance(pipeline.diffusion_model, SanaWMDiffusionModelConfig) + assert pipeline.diffusion_model._target is DiffusionModel assert isinstance(transformer, SanaWMTransformerConfig) + assert isinstance(pipeline.diffusion_model.scheduler, SanaWMLTXEulerSchedulerConfig) + assert isinstance(pipeline.decoder.refiner, SanaWMLTX2LatentRefinerConfig) + + +def test_sana_diffusion_config_instantiates_base_model() -> None: + """Use FlashDreams' shared diffusion model rather than a Sana-only runner.""" + model = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.setup() + + assert type(model) is DiffusionModel + assert isinstance(model.transformer.config, SanaWMTransformerConfig) + assert isinstance(model.scheduler, SanaWMLTXEulerScheduler) + + +def test_runner_pipeline_config_routes_runtime_fields_to_components() -> None: + """Apply CLI overrides to the component that owns each runtime field.""" + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + config_path="local_config.yaml", + model_path="local_model.safetensors", + stage1_precision="fp8", + no_refiner=True, + offload_vae=True, + offload_text_encoder=True, + ) + + pipeline = _pipeline_config(cfg, quant_backend="torch-fp8") + + assert isinstance(pipeline.encoder, SanaWMConditioningEncoderConfig) + assert isinstance(pipeline.decoder, SanaWMVideoDecoderConfig) + assert pipeline.encoder.config_path == "local_config.yaml" + assert pipeline.encoder.text_encoder.config_path == "local_config.yaml" + assert pipeline.encoder.text_encoder.stage1_precision == "fp8" + assert pipeline.encoder.text_encoder.quant_backend == "torch-fp8" + assert pipeline.encoder.text_encoder.offload_text_encoder is True + assert pipeline.encoder.first_frame_encoder.offload_vae is True + assert pipeline.decoder.vae_decoder.config_path == "local_config.yaml" + assert pipeline.decoder.vae_decoder.offload_vae is True + assert pipeline.decoder.refiner is None + assert pipeline.diffusion_model.transformer.checkpoint_path == ( + "local_model.safetensors" + ) + + +def test_sana_ltx_scheduler_step_pins_zero_timestep_tokens() -> None: + """Keep first-frame tokens fixed in the per-token LTX Euler step.""" + scheduler = SanaWMLTXEulerSchedulerConfig(num_inference_steps=4).setup() + + timesteps = scheduler.timesteps( + num_inference_steps=4, + shift=5.0, + device=torch.device("cpu"), + ) + sample = torch.ones((1, 2, 1)) + model_output = torch.ones_like(sample) + stepped = scheduler.step_ltx( + model_output=model_output, + timestep=torch.tensor(1000.0), + next_timestep=torch.tensor(500.0), + sample=sample, + per_token_timesteps=torch.tensor([[1000.0, 0.0]]), + schedule_timesteps=torch.tensor([1000.0, 500.0, 0.0]), + ) + + assert isinstance(scheduler, SanaWMLTXEulerScheduler) + assert timesteps.shape == (5,) + assert float(timesteps[-1]) == 0.0 + torch.testing.assert_close( + stepped, + torch.tensor([[[1.5], [1.0]]]), + ) + + +def test_sana_ltx_scheduler_matches_diffusers_per_token_step() -> None: + """Match diffusers FlowMatch Euler for SANA-WM's per-token branch.""" + from diffusers import FlowMatchEulerDiscreteScheduler + + ours = SanaWMLTXEulerSchedulerConfig(num_inference_steps=4).setup() + schedule_timesteps = ours.timesteps( + num_inference_steps=4, + shift=5.0, + device=torch.device("cpu"), + ) + upstream = FlowMatchEulerDiscreteScheduler(shift=5.0) + upstream.set_timesteps(4, device=torch.device("cpu")) + torch.testing.assert_close(schedule_timesteps[:-1], upstream.timesteps) + torch.testing.assert_close(schedule_timesteps / 1000.0, upstream.sigmas) + + sample = torch.tensor([[[1.0, -0.5], [0.25, 0.5], [-1.0, 0.75]]]) + model_output = torch.tensor([[[0.2, 0.4], [-0.3, 0.1], [0.5, -0.2]]]) + per_token_timesteps = torch.stack( + [ + torch.tensor( + [ + float(upstream.timesteps[0]), + 0.0, + float(upstream.timesteps[0]), + ] + ) + ] + ) + + expected = upstream.step( + model_output, + upstream.timesteps[0], + sample, + per_token_timesteps=per_token_timesteps, + return_dict=False, + )[0] + actual = ours.step_ltx( + model_output=model_output, + timestep=schedule_timesteps[0], + next_timestep=schedule_timesteps[1], + sample=sample, + per_token_timesteps=per_token_timesteps, + schedule_timesteps=schedule_timesteps, + ) + + torch.testing.assert_close(actual, expected) + + +def test_sana_scheduler_context_path_keeps_conditioned_frame_fixed() -> None: + """Let the scheduler own SANA-WM's per-token first-frame constraint.""" + scheduler = SanaWMLTXEulerSchedulerConfig(num_inference_steps=1).setup() + initial_noise = torch.zeros((1, 1, 2, 1, 1), dtype=torch.float32) + initial_noise[:, :, 0] = 5.0 + conditioning = SanaWMStage1Conditioning( + condition=torch.ones((1, 1, 1, 1)), + uncondition=None, + model_kwargs={"data_info": {"condition_frame_info": {0: 0.0}}}, + first_latent=torch.empty((1, 1, 1, 1, 1)), + latent_shape=tuple(initial_noise.shape), + cfg_scale=1.0, + flow_shift=5.0, + steps=1, + seed=0, + ) + + def predict_flow(noisy_latent: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + assert timestep.shape == (1, 1, 2) + return -torch.ones_like(noisy_latent) + + sampled = scheduler.sample( + initial_noise=initial_noise, + predict_flow=predict_flow, + context=conditioning, + ) + + torch.testing.assert_close(sampled[:, :, 0], initial_noise[:, :, 0]) + assert torch.all(sampled[:, :, 1] > initial_noise[:, :, 1]) def test_transformer_contract_shape_and_conditioning_guard() -> None: @@ -128,6 +297,127 @@ def test_transformer_contract_shape_and_conditioning_guard() -> None: ) +def test_transformer_initial_noise_uses_conditioning_payload() -> None: + """Generate first-frame-pinned noise through the shared transformer hook.""" + transformer = SanaWMTransformerConfig().setup() + transformer.weight_dtype = torch.float32 + first_latent = torch.full((1, 1, 1, 1, 1), 7.0) + conditioning = SanaWMStage1Conditioning( + condition=torch.empty((1, 1, 1, 1)), + uncondition=None, + model_kwargs={}, + first_latent=first_latent, + latent_shape=(1, 1, 2, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=123, + ) + cache = transformer.initialize_autoregressive_cache() + + noise = transformer.initial_noise( + latent_shape=(1, 1, 2, 1, 1), + rng=None, + cache=cache, + input=conditioning, + ) + + assert cache.conditioning is conditioning + assert noise.shape == (1, 1, 2, 1, 1) + torch.testing.assert_close(noise[:, :, :1], first_latent) + + +def test_transformer_predict_flow_applies_cfg_from_conditioning_input() -> None: + """Keep CFG inside the transformer boundary used by base diffusion.""" + + class DummyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.calls: list[dict[str, torch.Tensor]] = [] + + def forward( + self, + noisy_latent: torch.Tensor, + timestep: torch.Tensor, + prompt_embeds: torch.Tensor, + *, + mask: torch.Tensor, + camera_conditions: torch.Tensor, + chunk_plucker: torch.Tensor, + **_kwargs: object, + ) -> torch.Tensor: + assert "negative_mask" not in _kwargs + self.calls.append( + { + "noisy_latent": noisy_latent, + "timestep": timestep, + "prompt_embeds": prompt_embeds, + "mask": mask, + "camera_conditions": camera_conditions, + "chunk_plucker": chunk_plucker, + "data_info": _kwargs["data_info"], + } + ) + branch_values = prompt_embeds.flatten(1).mean(dim=1).reshape( + -1, + 1, + 1, + 1, + 1, + ) + return torch.ones_like(noisy_latent) * branch_values + + transformer = SanaWMTransformerConfig().setup() + dummy_model = DummyModel() + transformer.model = dummy_model + transformer._model_built = True + cond_mask = torch.ones((1, 1)) + neg_mask = torch.zeros((1, 1)) + camera = torch.zeros((1, 3, 1, 1, 1)) + chunk_plucker = torch.ones((1, 6, 1, 1, 1)) + conditioning = SanaWMStage1Conditioning( + condition=torch.ones((1, 1, 1, 1)), + uncondition=torch.zeros((1, 1, 1, 1)), + model_kwargs={ + "mask": cond_mask, + "negative_mask": neg_mask, + "camera_conditions": camera, + "chunk_plucker": chunk_plucker, + "data_info": {"condition_frame_info": {0: 0.0}}, + }, + first_latent=torch.empty((1, 1, 1, 1, 1)), + latent_shape=(1, 1, 1, 1, 1), + cfg_scale=2.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + + out = transformer.predict_flow( + noisy_latent=torch.zeros((1, 1, 1, 1, 1)), + timestep=torch.zeros((1, 1, 1)), + cache=SanaWMTransformerCache(), + input=conditioning, + ) + + torch.testing.assert_close(out, torch.full((1, 1, 1, 1, 1), 2.0)) + assert len(dummy_model.calls) == 1 + call = dummy_model.calls[0] + assert call["noisy_latent"].shape == (2, 1, 1, 1, 1) + assert call["timestep"].shape == (2, 1, 1) + assert call["prompt_embeds"].shape == (2, 1, 1, 1) + torch.testing.assert_close(call["mask"], torch.cat([neg_mask, cond_mask], dim=0)) + torch.testing.assert_close( + call["camera_conditions"], + torch.cat([camera, camera], dim=0), + ) + torch.testing.assert_close( + call["chunk_plucker"], + torch.cat([chunk_plucker, chunk_plucker], dim=0), + ) + assert call["data_info"] == {"condition_frame_info": {0: 0.0}} + + def test_inference_config_loads_yaml( tmp_path: Path, ) -> None: @@ -166,6 +456,170 @@ def test_inference_config_loads_yaml( assert cfg.work_dir == "" +def test_text_prompt_encoder_outputs_padded_prompt_schema( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Encode prompts behind the explicit text component boundary.""" + + class DummyTokens: + def __init__(self, max_length: int) -> None: + self.input_ids = torch.arange(max_length).reshape(1, max_length) + self.attention_mask = torch.ones((1, max_length), dtype=torch.long) + + def to(self, _device: torch.device) -> "DummyTokens": + return self + + class DummyTokenizer: + def encode(self, _text: str) -> list[int]: + return [1, 2] + + def __call__( + self, + *_args: object, + max_length: int, + **_kwargs: object, + ) -> DummyTokens: + return DummyTokens(max_length) + + class DummyTextEncoder(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(0)) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor]: + del attention_mask + hidden = torch.ones((1, input_ids.shape[1], 12), dtype=torch.float32) + return (hidden,) + + monkeypatch.setenv("SANA_WM_STAGE1_NVFP4_TEXT_PAD_MULTIPLE", "8") + monkeypatch.setattr( + conditioning_module, + "_load_inference_config", + lambda _path: SimpleNamespace( + model=SimpleNamespace(mixed_precision="bf16"), + text_encoder=SimpleNamespace( + text_encoder_name="T5-small", + model_max_length=5, + chi_prompt=["prefix"], + ), + ), + ) + monkeypatch.setattr( + conditioning_module, + "_get_tokenizer_and_text_encoder", + lambda **_kwargs: (DummyTokenizer(), DummyTextEncoder()), + ) + encoder = SanaWMTextPromptEncoderConfig( + config_path="dummy.yaml", + stage1_precision="fp4", + ).setup() + + encoded = encoder( + SanaWMTextPromptRequest( + prompt="drive forward", + negative_prompt="low quality", + ) + ) + + assert encoded.condition.shape == (1, 1, 8, 12) + assert encoded.condition_mask.shape == (1, 8) + assert encoded.negative.shape == (1, 1, 8, 12) + assert encoded.negative_mask.shape == (1, 8) + + +def test_first_frame_encoder_outputs_latent_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Encode the first frame behind the named I2V component.""" + from PIL import Image + + class DummyLatentDist: + def mode(self) -> torch.Tensor: + return torch.ones((1, 4, 1, 2, 3), dtype=torch.float32) + + class DummyPosterior: + latent_dist = DummyLatentDist() + + class DummyVAE(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(0)) + self.latents_mean = torch.zeros(4) + self.latents_std = torch.ones(4) + self.config = SimpleNamespace(scaling_factor=1.0) + + def encode(self, _images: torch.Tensor) -> DummyPosterior: + return DummyPosterior() + + monkeypatch.setattr( + conditioning_module, + "_load_inference_config", + lambda _path: SimpleNamespace( + model=SimpleNamespace(mixed_precision="bf16"), + vae=SimpleNamespace( + weight_dtype="float32", + vae_type="LTX2VAE_diffusers", + vae_pretrained="local", + ), + ), + ) + monkeypatch.setattr( + conditioning_module, + "_get_vae", + lambda *_args, **_kwargs: DummyVAE(), + ) + encoder = SanaWMFirstFrameEncoderConfig(config_path="dummy.yaml").setup() + + latent = encoder(Image.new("RGB", (4, 4))) + + assert latent.shape == (1, 4, 1, 2, 3) + assert latent.dtype == torch.bfloat16 + + +def test_camera_conditioning_encoder_outputs_sana_schema() -> None: + """Build raymap/chunk-Plucker tensors behind the camera component.""" + poses = np.broadcast_to(np.eye(4, dtype=np.float32), (17, 4, 4)).copy() + intrinsics = np.broadcast_to( + np.array([900.0, 900.0, 640.0, 352.0], dtype=np.float32), + (17, 4), + ).copy() + encoder = SanaWMCameraConditioningEncoderConfig().setup() + + camera = encoder( + SanaWMCameraRequest( + poses_c2w=poses, + intrinsics_vec4=intrinsics, + ) + ) + + assert camera["raymap"].shape == (3, 20) + assert camera["chunk_plucker"].shape == (48, 3, 22, 40) + + +def test_video_decoder_returns_structured_video(monkeypatch: pytest.MonkeyPatch) -> None: + """Decode Stage-1 latents through the explicit decoder component.""" + decoder = SanaWMVideoDecoderConfig(refiner=None).setup() + monkeypatch.setattr( + decoder.vae_decoder, + "decode_latents", + lambda _latents: np.zeros((1, 2, 2, 3), dtype=np.uint8), + ) + + decoded = decoder( + torch.zeros((1, 4, 1, 1, 1)), + autoregressive_index=0, + cache=None, + ) + + assert isinstance(decoded, SanaWMDecodedVideo) + assert decoded.video_hwc.shape == (1, 2, 2, 3) + assert decoded.stage1_video_hwc is None + + def test_stage1_model_matches_checkpoint_schema() -> None: """Pin the Stage-1 module to the public checkpoint schema.""" state = SanaWMStage1Model().state_dict() @@ -266,17 +720,8 @@ def test_transformer_releases_stage1_runtime() -> None: """Free Stage-1-only modules and conditioning before decode/refine.""" transformer = SanaWMTransformerConfig().setup() transformer.model = torch.nn.Linear(1, 1) - transformer.text_encoder = torch.nn.Linear(1, 1) - transformer.tokenizer = object() transformer._model_built = True - transformer._text_encoder_built = True transformer._stage1_quantized = True - transformer._streaming_prompt_cache[("prompt",)] = ( - torch.empty(1), - torch.empty(1), - torch.empty(1), - torch.empty(1), - ) cache = SanaWMTransformerCache( conditioning=SanaWMStage1Conditioning( condition=torch.empty(1), @@ -295,11 +740,7 @@ def test_transformer_releases_stage1_runtime() -> None: assert cache.conditioning is None assert transformer.model is None - assert transformer.text_encoder is None - assert transformer.tokenizer is None - assert transformer._streaming_prompt_cache == {} assert transformer._model_built is False - assert transformer._text_encoder_built is False assert transformer._stage1_quantized is False @@ -322,12 +763,12 @@ def __init__(self) -> None: def enable_tiling(self, **kwargs: int) -> None: self.calls.append(kwargs) - transformer = SanaWMTransformerConfig().setup() - transformer.vae = DummyVAE() + decoder = SanaWMLTX2VAEDecoderConfig().setup() + decoder.vae = DummyVAE() - transformer._configure_vae_tiling() + decoder._configure_vae_tiling() - assert transformer.vae.calls == [ + assert decoder.vae.calls == [ { "tile_sample_min_height": 256, "tile_sample_stride_height": 192, @@ -337,14 +778,14 @@ def enable_tiling(self, **kwargs: int) -> None: "tile_sample_stride_num_frames": 8, } ] - assert transformer.vae.tile_sample_min_height == 256 - assert transformer.vae.tile_sample_stride_height == 192 - assert transformer.vae.tile_sample_min_width == 256 - assert transformer.vae.tile_sample_stride_width == 224 - assert transformer.vae.tile_sample_min_num_frames == 24 - assert transformer.vae.tile_sample_stride_num_frames == 8 - assert transformer.vae.use_framewise_encoding is True - assert transformer.vae.use_framewise_decoding is True + assert decoder.vae.tile_sample_min_height == 256 + assert decoder.vae.tile_sample_stride_height == 192 + assert decoder.vae.tile_sample_min_width == 256 + assert decoder.vae.tile_sample_stride_width == 224 + assert decoder.vae.tile_sample_min_num_frames == 24 + assert decoder.vae.tile_sample_stride_num_frames == 8 + assert decoder.vae.use_framewise_encoding is True + assert decoder.vae.use_framewise_decoding is True def test_decode_retries_vae_oom_with_smaller_tiles( @@ -372,10 +813,10 @@ def enable_tiling(self, **kwargs: int) -> None: for name, value in kwargs.items(): setattr(self, name, value) - transformer = SanaWMTransformerConfig().setup() - transformer.vae = DummyVAE() - transformer.vae_dtype = torch.float32 - monkeypatch.setattr(transformer, "_ensure_vae", lambda: None) + decoder = SanaWMLTX2VAEDecoderConfig().setup() + decoder.vae = DummyVAE() + decoder.vae_dtype = torch.float32 + monkeypatch.setattr(decoder, "_ensure_vae", lambda: None) calls = 0 inference_modes: list[bool] = [] @@ -387,42 +828,19 @@ def decode_once(_samples: torch.Tensor) -> torch.Tensor: raise torch.OutOfMemoryError("test OOM") return torch.zeros((1, 3, 1, 2, 2), dtype=torch.float32) - monkeypatch.setattr(transformer, "_vae_decode", decode_once) + monkeypatch.setattr(decoder, "_vae_decode", decode_once) - video = transformer.decode_latents(torch.zeros((1, 4, 1, 1, 1))) + video = decoder.decode_latents(torch.zeros((1, 4, 1, 1, 1))) assert calls == 2 assert inference_modes == [True, True] assert video.shape == (1, 2, 2, 3) - assert transformer.vae.tile_sample_min_height == 128 - assert transformer.vae.tile_sample_stride_height == 64 - assert transformer.vae.tile_sample_min_width == 128 - assert transformer.vae.tile_sample_stride_width == 64 - assert transformer.vae.tile_sample_min_num_frames == 16 - assert transformer.vae.tile_sample_stride_num_frames == 8 - - -def test_cfg_model_kwargs_do_not_duplicate_camera_tensors() -> None: - """Run CFG as sequential branches instead of doubling Stage-1 activations.""" - camera = torch.zeros((1, 3, 3, 2, 2)) - cond_mask = torch.ones((1, 8)) - neg_mask = torch.zeros((1, 8)) - kwargs = { - "mask": cond_mask, - "negative_mask": neg_mask, - "camera_conditions": camera, - "chunk_plucker": torch.zeros((1, 6, 3, 2, 2)), - } - - cond_kwargs = _condition_model_kwargs(kwargs) - neg_kwargs = _uncondition_model_kwargs(kwargs) - - assert cond_kwargs["mask"] is cond_mask - assert neg_kwargs["mask"] is neg_mask - assert cond_kwargs["camera_conditions"] is camera - assert neg_kwargs["camera_conditions"] is camera - assert "negative_mask" not in cond_kwargs - assert "negative_mask" not in neg_kwargs + assert decoder.vae.tile_sample_min_height == 128 + assert decoder.vae.tile_sample_stride_height == 64 + assert decoder.vae.tile_sample_min_width == 128 + assert decoder.vae.tile_sample_stride_width == 64 + assert decoder.vae.tile_sample_min_num_frames == 16 + assert decoder.vae.tile_sample_stride_num_frames == 8 def test_vae_tiling_avoids_degenerate_latent_tails() -> None: @@ -498,16 +916,16 @@ def __init__(self, **kwargs: object) -> None: monkeypatch.setattr(refiner_module, "SanaWMLTX2Refiner", DummyRefiner) monkeypatch.setattr( - transformer_module, + decoder_module, "resolve_hf_path", lambda value: f"/resolved/{value}", ) - transformer = SanaWMTransformerConfig().setup() + refiner = SanaWMLTX2LatentRefinerConfig().setup() - transformer._ensure_refiner() + refiner._ensure_refiner() - assert transformer._refiner_built is True - assert isinstance(transformer.refiner, DummyRefiner) + assert refiner._refiner_built is True + assert isinstance(refiner.refiner, DummyRefiner) assert calls["refiner_root"] == ( "/resolved/hf://Efficient-Large-Model/SANA-WM_bidirectional/refiner" ) From 80ef8155e41c1bce4b5bef092d613c3e6bc274ab Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 17 Jul 2026 13:32:13 -0700 Subject: [PATCH 14/36] Keep SANA-WM models resident across reuse without inflating peak memory Make the SANA-WM pipeline reuse-friendly: the Stage-1 DiT, LTX-2 refiner, and Gemma text encoder are no longer torn down and rebuilt on every generate(). Previously each rollout unconditionally released Stage-1 before decode, released the refiner after refine, and reloaded the ~20 GB Gemma encoder from disk per prompt, so a reused pipeline re-paid a full checkpoint reload and re-quantization (~10 s bf16, ~17 s fp8) plus an ~8-15 s Gemma reload on every call. - transformer: gate the Stage-1 teardown behind a new offload_stage1 flag (default keeps it resident); still drop per-rollout conditioning. - decoder: drop the unconditional refiner release (refine_latents already releases when offload_refiner is set). - refiner: cache the Gemma tokenizer/encoder in CPU RAM instead of reloading per call, and move it to the GPU only for the encode forward so it never inflates peak memory during denoise/decode. Fix deprecated torch_dtype and drop dead _empty_cuda_cache/gc. - runner: expose --offload-stage1 and route it to the transformer. - Add [timing] instrumentation for the Stage-1 build, refiner build, and Gemma load. Measured on Blackwell: warm-reuse pass drops from ~26 s to ~5 s (fp8); full-res fp8 demo_0 peaks at ~54/98 GB with no OOM. Co-Authored-By: Claude Opus 4.8 (1M context) --- integrations/sana/sana_wm/decoder.py | 10 ++- integrations/sana/sana_wm/refiner.py | 78 ++++++++++++++++-------- integrations/sana/sana_wm/runner.py | 6 ++ integrations/sana/sana_wm/transformer.py | 24 +++++++- integrations/sana/tests/test_smoke.py | 54 ++++++++++++++++ 5 files changed, 142 insertions(+), 30 deletions(-) diff --git a/integrations/sana/sana_wm/decoder.py b/integrations/sana/sana_wm/decoder.py index 698e1d6aa..f3367f6d7 100644 --- a/integrations/sana/sana_wm/decoder.py +++ b/integrations/sana/sana_wm/decoder.py @@ -376,6 +376,7 @@ def _ensure_refiner(self) -> None: return from sana_wm.refiner import SanaWMLTX2Refiner + t0 = time.perf_counter() compute_dtype = ( torch.bfloat16 if self.config.refiner_precision in {"fp8", "fp4"} @@ -390,6 +391,11 @@ def _ensure_refiner(self) -> None: quant_backend=self.config.quant_backend, ) self._refiner_built = True + logger.info( + "[timing] refiner build: {:.3f}s (precision={})", + time.perf_counter() - t0, + self.config.refiner_precision, + ) @dataclass(kw_only=True) @@ -449,7 +455,9 @@ def forward( block_size=cache.refiner_block_size, kv_max_frames=cache.refiner_kv_max_frames, ) - self.refiner.release_runtime() + # refine_latents() already releases the refiner when offload_refiner + # is set; keep it resident otherwise so a reused pipeline avoids a + # full refiner rebuild + re-quantization on the next call. elif cache.save_stage1: logger.info( "SANA-WM is already running without the refiner; " diff --git a/integrations/sana/sana_wm/refiner.py b/integrations/sana/sana_wm/refiner.py index c29844d93..8cfce1f96 100644 --- a/integrations/sana/sana_wm/refiner.py +++ b/integrations/sana/sana_wm/refiner.py @@ -17,7 +17,7 @@ from __future__ import annotations -import gc +import time from pathlib import Path from typing import Literal @@ -71,6 +71,7 @@ def __init__( self.quant_backend = quant_backend self.text_max_sequence_length = int(text_max_sequence_length) self._quantized = False + self._text_encoder_built = False self.transformer, self.connectors = self._load_diffusers_components() @torch.inference_mode() @@ -196,14 +197,44 @@ def _prepare_quantization(self) -> None: ) self._quantized = True - @torch.inference_mode() - def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: + def _ensure_text_encoder(self) -> None: + """Load the Gemma tokenizer + encoder once and cache them in CPU RAM. + + The encoder is ~20 GB; reloading it from disk on every refine call + dominated repeated-use latency, so the built module is kept cached and + a reused refiner pays the load once. It stays on the CPU between calls + and is moved to the GPU only for the encode forward (see + :meth:`_encode_prompt`) so it never inflates peak memory during the + denoise/decode phases. When ``offload_refiner`` is set the whole + refiner (and this encoder) is released between runs by + ``release_runtime``. + """ + if self._text_encoder_built: + return from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + t0 = time.perf_counter() tokenizer = AutoTokenizer.from_pretrained(self.gemma_root) tokenizer.padding_side = "left" if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token + self.tokenizer = tokenizer + # Built on CPU; moved to the GPU on demand for each encode call only. + self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained( + self.gemma_root, + dtype=self.dtype, + low_cpu_mem_usage=True, + ).eval() + self._text_encoder_built = True + logger.info( + "[timing] refiner text-encoder build: {:.3f}s", + time.perf_counter() - t0, + ) + + @torch.inference_mode() + def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: + self._ensure_text_encoder() + tokenizer = self.tokenizer text_inputs = tokenizer( [prompt.strip()], @@ -216,22 +247,23 @@ def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: input_ids = text_inputs.input_ids.to(self.device) attention_mask = text_inputs.attention_mask.to(self.device) - text_encoder = Gemma3ForConditionalGeneration.from_pretrained( - self.gemma_root, - torch_dtype=self.dtype, - low_cpu_mem_usage=True, - ).eval() - text_encoder.to(self.device) - text_backbone = getattr(text_encoder, "model", 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) - del text_encoder, text_backbone, outputs - _empty_cuda_cache() + # Pull the cached encoder onto the GPU only for the forward, then park + # it back on the CPU so it is absent during denoise/decode. + self.text_encoder.to(self.device) + try: + 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) + del outputs + finally: + self.text_encoder.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() prompt_embeds = _pack_text_embeds( hidden_states, @@ -240,7 +272,6 @@ def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: padding_side=tokenizer.padding_side, ).to(dtype=self.dtype) del hidden_states - _empty_cuda_cache() self.connectors.to(self.device) connector_prompt_embeds, _, connector_attention_mask = self.connectors( @@ -249,7 +280,6 @@ def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: ) self.connectors.to("cpu") del prompt_embeds, attention_mask - _empty_cuda_cache() return ( connector_prompt_embeds.to(device=self.device, dtype=self.dtype), connector_attention_mask.to(device=self.device), @@ -589,9 +619,3 @@ def _prepare_encoder_attention_mask(mask: Tensor | None, dtype: torch.dtype) -> if bool(torch.all(mask)): return None return ((1 - mask.to(dtype)) * -10000.0).unsqueeze(1) - - -def _empty_cuda_cache() -> None: - if torch.cuda.is_available(): - torch.cuda.empty_cache() - gc.collect() diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index 48cada5f6..deb66530f 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -174,6 +174,11 @@ class SanaWMRunnerConfig(RunnerConfig): offload_vae: bool = False """Move the VAE to CPU between encode/decode phases.""" + offload_stage1: bool = False + """Tear down the Stage-1 DiT after sampling to free memory for decode/ + refine. Default keeps it resident (fastest); enable on memory-constrained + GPUs.""" + offload_refiner: bool = False """Build and release the refiner only around refinement.""" @@ -397,6 +402,7 @@ def _pipeline_config( checkpoint_path=cfg.model_path, stage1_precision=cfg.stage1_precision, quant_backend=quant_backend, + offload_stage1=cfg.offload_stage1, ), ), encoder=dict( diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py index 335db4e9c..dd5b1ac19 100644 --- a/integrations/sana/sana_wm/transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -18,6 +18,7 @@ from __future__ import annotations import gc +import time from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any, Literal, cast @@ -121,6 +122,11 @@ class SanaWMTransformerConfig(TransformerConfig): width: int = DEFAULT_VIDEO_WIDTH """Pixel width used by the public SANA-WM bidirectional release.""" + offload_stage1: bool = False + """Tear down the Stage-1 DiT after sampling, before decode/refine. Default + keeps it resident (fastest when memory is available); enable only to free + Stage-1 memory for the VAE/refiner on memory-constrained GPUs.""" + class SanaWMTransformer(Transformer[SanaWMTransformerCache]): """FlashDreams adapter for the SANA-WM Stage-1 model call.""" @@ -229,9 +235,17 @@ def postprocess_clean_latent( cache: SanaWMTransformerCache, input: Any = None, ) -> Tensor: - """Release Stage-1 runtime before the pipeline enters decode/refine.""" + """Release Stage-1 runtime before decode/refine only when offloading. + + By default the Stage-1 DiT stays resident so a reused pipeline never + reloads and re-quantizes it; ``offload_stage1`` restores the old + free-before-decode behavior for memory-constrained GPUs. + """ del input - self.release_stage1_runtime(cache) + if self.config.offload_stage1: + self.release_stage1_runtime(cache) + elif cache is not None: + cache.conditioning = None return clean_latent def initial_latents(self, conditioning: SanaWMStage1Conditioning) -> Tensor: @@ -355,6 +369,7 @@ def _ensure_model(self) -> None: if self._model_built: self._prepare_stage1_quant() return + t0 = time.perf_counter() cfg = self._ensure_runtime_config() weight_dtype = self._ensure_weight_dtype() model = SanaWMStage1Model().to(self.device) @@ -382,6 +397,11 @@ def _ensure_model(self) -> None: self.model = model.eval().to(weight_dtype) self._model_built = True self._prepare_stage1_quant() + logger.info( + "[timing] stage1 build+load+quant: {:.3f}s (precision={})", + time.perf_counter() - t0, + self.config.stage1_precision, + ) def _prepare_stage1_quant(self) -> None: if self._stage1_quantized or self.config.stage1_precision == "bf16": diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index c626a5b26..6c9c76fcf 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -153,6 +153,7 @@ def test_runner_pipeline_config_routes_runtime_fields_to_components() -> None: no_refiner=True, offload_vae=True, offload_text_encoder=True, + offload_stage1=True, ) pipeline = _pipeline_config(cfg, quant_backend="torch-fp8") @@ -171,6 +172,7 @@ def test_runner_pipeline_config_routes_runtime_fields_to_components() -> None: assert pipeline.diffusion_model.transformer.checkpoint_path == ( "local_model.safetensors" ) + assert pipeline.diffusion_model.transformer.offload_stage1 is True def test_sana_ltx_scheduler_step_pins_zero_timestep_tokens() -> None: @@ -744,6 +746,58 @@ def test_transformer_releases_stage1_runtime() -> None: assert transformer._stage1_quantized is False +def _stage1_conditioning_cache() -> SanaWMTransformerCache: + return SanaWMTransformerCache( + conditioning=SanaWMStage1Conditioning( + condition=torch.empty(1), + uncondition=None, + model_kwargs={}, + first_latent=torch.empty(1), + latent_shape=(1, 1, 1, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + ) + + +def test_postprocess_keeps_stage1_resident_by_default() -> None: + """Default keeps the Stage-1 DiT resident so a reused pipeline never reloads.""" + transformer = SanaWMTransformerConfig().setup() + assert transformer.config.offload_stage1 is False + model = torch.nn.Linear(1, 1) + transformer.model = model + transformer._model_built = True + transformer._stage1_quantized = True + cache = _stage1_conditioning_cache() + latent = torch.zeros(1) + + result = transformer.postprocess_clean_latent(latent, cache) + + assert result is latent + assert transformer.model is model + assert transformer._model_built is True + assert transformer._stage1_quantized is True + # Per-rollout conditioning is still dropped even when the model stays warm. + assert cache.conditioning is None + + +def test_postprocess_releases_stage1_when_offloading() -> None: + """offload_stage1 restores the free-before-decode behavior.""" + transformer = SanaWMTransformerConfig(offload_stage1=True).setup() + transformer.model = torch.nn.Linear(1, 1) + transformer._model_built = True + transformer._stage1_quantized = True + cache = _stage1_conditioning_cache() + + transformer.postprocess_clean_latent(torch.zeros(1), cache) + + assert transformer.model is None + assert transformer._model_built is False + assert cache.conditioning is None + + def test_vae_tiling_uses_low_memory_tiles() -> None: """Keep VAE decode on the configured low-memory tiles.""" From 4e4a8cd23e5edf1fa2b86171ccc4c21cb1f67fd1 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 17 Jul 2026 14:06:44 -0700 Subject: [PATCH 15/36] Derive SANA-WM intrinsics from the first frame when none provided Make --intrinsics-path optional. When omitted, intrinsics are derived from the first-frame size with a centered principal point and a focal length set by a horizontal field of view (new --intrinsics-hfov-deg, default 90 to match the public demo intrinsics). This lets a run be driven by just an image and a prompt (with --action supplying the camera trajectory). - camera: add default_intrinsics_vec4(), returning [F, 4] intrinsics in source-image pixels so they flow through transform_intrinsics_for_crop identically to a loaded file. - runner: derive intrinsics when intrinsics_path is None, log the derivation, and expose --intrinsics-hfov-deg; explicit files behave as before. - README: document optional intrinsics and a minimal image+prompt run. - tests: cover the helper (centering, FOV math, bad-FOV rejection) and the runner derive path. Co-Authored-By: Claude Opus 4.8 (1M context) --- integrations/sana/README.md | 17 ++++++++++-- integrations/sana/sana_wm/camera.py | 36 ++++++++++++++++++++++++++ integrations/sana/sana_wm/runner.py | 24 ++++++++++++++--- integrations/sana/tests/test_camera.py | 21 +++++++++++++++ integrations/sana/tests/test_smoke.py | 26 +++++++++++++++++++ 5 files changed, 119 insertions(+), 5 deletions(-) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index a8a9f18d5..48e3c0b2a 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -47,8 +47,6 @@ well. ## Run -The runner requires explicit intrinsics: - ```bash PYTORCH_ALLOC_CONF=expandable_segments:True \ uv run flashdreams-run sana-wm-bidirectional \ @@ -66,6 +64,21 @@ Expected output: outputs/sana_wm_bf16/sana_wm_generated.mp4 ``` +`--intrinsics-path` is optional. When omitted, intrinsics are derived from the +first-frame size assuming a centered principal point and a horizontal field of +view of `--intrinsics-hfov-deg` (default `90`, matching the demo intrinsics). +Likewise `--camera-path` may be replaced by an `--action` DSL string. A minimal +run therefore needs only an image and a prompt: + +```bash +uv run flashdreams-run sana-wm-bidirectional \ + --image-path my_frame.png \ + --prompt "a scene description; describe the world's own motion" \ + --action "w-100,dw-60,w-101" \ + --num-frames 161 \ + --output-dir outputs/mine +``` + For Stage-1-only diagnostics, add `--no-refiner True`. ## FP8 and FP4 diff --git a/integrations/sana/sana_wm/camera.py b/integrations/sana/sana_wm/camera.py index af8aa406e..363534100 100644 --- a/integrations/sana/sana_wm/camera.py +++ b/integrations/sana/sana_wm/camera.py @@ -317,6 +317,42 @@ def load_intrinsics(path: Path, num_frames: int) -> np.ndarray: ) +def default_intrinsics_vec4( + src_size: tuple[int, int], + num_frames: int, + *, + hfov_deg: float = 90.0, +) -> np.ndarray: + """Derive ``[fx, fy, cx, cy]`` intrinsics from an image size. + + Used when the caller does not supply an intrinsics file. Assumes square + pixels and a principal point at the image center, with focal length set by + a horizontal field of view. The public SANA-WM demo intrinsics correspond + to roughly ``90`` degrees; lower values are narrower/more zoomed-in. The + returned intrinsics are in the source image's pixel coordinates, matching + :func:`load_intrinsics`, so they flow through + :func:`transform_intrinsics_for_crop` identically. + + Args: + src_size: Source image size as ``(width, height)``. + num_frames: Number of frames required by the rollout. + hfov_deg: Horizontal field of view in degrees. + + Returns: + ``[num_frames, 4]`` intrinsics as ``[fx, fy, cx, cy]``. + """ + src_w, src_h = src_size + if src_w <= 0 or src_h <= 0: + raise ValueError(f"src_size must be positive, got {src_size}.") + if not 0.0 < hfov_deg < 180.0: + raise ValueError(f"hfov_deg must be in (0, 180), got {hfov_deg}.") + focal = 0.5 * src_w / math.tan(math.radians(hfov_deg) / 2.0) + vector = np.array( + [focal, focal, src_w / 2.0, src_h / 2.0], dtype=np.float32 + ) + return np.broadcast_to(vector, (num_frames, 4)).copy() + + def snap_num_frames( num_frames: int, *, diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index deb66530f..ea0e01dd3 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -34,6 +34,7 @@ from flashdreams.infra.runner import Runner, RunnerConfig from sana_wm.camera import ( action_string_to_c2w, + default_intrinsics_vec4, load_intrinsics, resize_center_crop_geometry, snap_num_frames, @@ -89,7 +90,14 @@ class SanaWMRunnerConfig(RunnerConfig): intrinsics_path: Path | None = None """Optional ``.npy`` intrinsics shaped ``[3, 3]``, ``[F, 3, 3]``, - ``[4]``, or ``[F, 4]``.""" + ``[4]``, or ``[F, 4]``. When omitted, intrinsics are derived from the + first-frame size using ``intrinsics_hfov_deg`` with a centered principal + point.""" + + intrinsics_hfov_deg: float = 90.0 + """Horizontal field of view in degrees used to derive intrinsics when + ``intrinsics_path`` is not provided. The public demo intrinsics correspond + to ~90 degrees; lower values are narrower/more zoomed-in.""" action: str | None = DEFAULT_ACTION """Action DSL used when ``camera_path`` is not provided.""" @@ -370,8 +378,18 @@ def _prepare_inputs(self) -> tuple[object, np.ndarray, np.ndarray, int]: ) ) if cfg.intrinsics_path is None: - raise ValueError("SanaWMRunner requires --intrinsics-path.") - intrinsics_src = load_intrinsics(cfg.intrinsics_path, num_frames) + intrinsics_src = default_intrinsics_vec4( + image.size, num_frames, hfov_deg=cfg.intrinsics_hfov_deg + ) + if self.is_rank_zero: + logger.info( + "No --intrinsics-path provided; deriving intrinsics from " + "image size {} at hfov={} deg (principal point centered).", + image.size, + cfg.intrinsics_hfov_deg, + ) + else: + intrinsics_src = load_intrinsics(cfg.intrinsics_path, num_frames) intrinsics_vec4 = transform_intrinsics_for_crop( intrinsics_src, image.size, diff --git a/integrations/sana/tests/test_camera.py b/integrations/sana/tests/test_camera.py index 43d0b7275..ba7a52ed8 100644 --- a/integrations/sana/tests/test_camera.py +++ b/integrations/sana/tests/test_camera.py @@ -25,6 +25,7 @@ from sana_wm.camera import ( action_string_to_c2w, + default_intrinsics_vec4, fit_intrinsics_sequence, load_intrinsics, prepare_camera, @@ -36,6 +37,26 @@ pytestmark = pytest.mark.ci_cpu +def test_default_intrinsics_vec4_centers_and_matches_hfov() -> None: + vec = default_intrinsics_vec4((1000, 500), num_frames=3, hfov_deg=90.0) + assert vec.shape == (3, 4) + fx, fy, cx, cy = vec[0] + # 90 deg hfov with square pixels -> fx = 0.5 * W / tan(45) = 0.5 * W. + assert fx == pytest.approx(500.0) + assert fy == pytest.approx(500.0) + # Principal point at the image center. + assert cx == pytest.approx(500.0) + assert cy == pytest.approx(250.0) + # Same intrinsics for every frame. + assert np.allclose(vec, vec[0]) + + +def test_default_intrinsics_vec4_rejects_bad_fov() -> None: + for bad in (0.0, 180.0, 200.0): + with pytest.raises(ValueError): + default_intrinsics_vec4((640, 360), num_frames=1, hfov_deg=bad) + + def test_action_string_rolls_out_identity_plus_motion() -> None: """Expand ``w-3`` to an identity frame plus three motion frames.""" c2w = action_string_to_c2w("w-3", smooth=False) diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 6c9c76fcf..f39c3eaf3 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -945,6 +945,32 @@ def test_runner_setup_preserves_cli_fields() -> None: assert runner.config.image_path == Path("missing.png") +def test_runner_derives_intrinsics_when_omitted(tmp_path: Path) -> None: + """Omitting --intrinsics-path derives per-frame intrinsics from the frame.""" + from PIL import Image + + image_path = tmp_path / "frame.png" + Image.new("RGB", (640, 360), color=(10, 20, 30)).save(image_path) + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + image_path=image_path, + prompt="demo", + intrinsics_path=None, + action="w-40", + num_frames=25, + ) + runner = cfg.setup() + + _image, c2w, intrinsics_vec4, num_frames = runner._prepare_inputs() + + assert num_frames == 25 + assert c2w.shape[0] == num_frames + assert intrinsics_vec4.shape == (num_frames, 4) + # Derived focal lengths are finite and positive after the crop transform. + assert np.all(np.isfinite(intrinsics_vec4)) + assert np.all(intrinsics_vec4[:, :2] > 0) + + def test_runner_config_type() -> None: """Keep the exported literal on the SANA-WM runner config subclass.""" assert isinstance(RUNNER_SANA_WM_BIDIRECTIONAL, SanaWMRunnerConfig) From 5449238e58025be25a106a1e5cd4d29ce42c3464 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 17 Jul 2026 14:19:19 -0700 Subject: [PATCH 16/36] Make SANA integration obey FlashDreams contracts Move SANA-WM first-frame timestep pinning and zero-flow behavior out of the scheduler context path and into the SANA transformer. The shared diffusion model now calls schedulers through the generic initial_noise/predict_flow/rng interface again, and the SANA LTX Euler scheduler only owns timestep construction and update math. Make the decoder/refiner boundary conform to FlashDreams video decoder contracts by subclassing StreamingVideoDecoder and publishing the SANA 8k+1 temporal sizing. Remove runner, cache, and refiner AR block-size/KV-frame knobs that were accepted but ignored, leaving the supported sink-bidirectional LTX-2 refiner path explicit. Route runner step and flow_shift overrides into the scheduler config instead of through SANA conditioning, and update SANA smoke coverage for scheduler isolation, transformer first-frame pinning, decoder sizing, CFG, and refiner behavior. Tests:\n- uv run --package flashdreams-sana-wm --extra dev pytest integrations/sana/tests Signed-off-by: Aidan Foster --- .../flashdreams/infra/diffusion/model/base.py | 1 - .../infra/diffusion/scheduler/base.py | 6 +- .../infra/diffusion/scheduler/fm.py | 3 - .../infra/diffusion/scheduler/fm_euler.py | 3 - .../infra/diffusion/scheduler/fm_unipc.py | 3 - integrations/sana/sana_wm/decoder.py | 43 +++++-- integrations/sana/sana_wm/refiner.py | 9 -- integrations/sana/sana_wm/runner.py | 12 +- integrations/sana/sana_wm/scheduler.py | 114 ------------------ integrations/sana/sana_wm/transformer.py | 93 +++++++++++++- integrations/sana/tests/test_smoke.py | 74 +++++++++--- 11 files changed, 185 insertions(+), 176 deletions(-) diff --git a/flashdreams/flashdreams/infra/diffusion/model/base.py b/flashdreams/flashdreams/infra/diffusion/model/base.py index 67cef7532..57b4e5171 100644 --- a/flashdreams/flashdreams/infra/diffusion/model/base.py +++ b/flashdreams/flashdreams/infra/diffusion/model/base.py @@ -212,7 +212,6 @@ def predict_flow(noisy_latent: Tensor, timestep: Tensor) -> Tensor: initial_noise=initial_noise, predict_flow=predict_flow, rng=self.rng, - context=input, ) if self.config.noise_in_unpatchified_shape: diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/base.py b/flashdreams/flashdreams/infra/diffusion/scheduler/base.py index 4bb0f1dd6..46511d93b 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/base.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/base.py @@ -19,7 +19,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, Protocol +from typing import Protocol import torch import torch.nn as nn @@ -69,7 +69,6 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, - context: Any = None, ) -> Tensor: """Run the full denoising loop and return the clean latent. @@ -84,9 +83,6 @@ def sample( ``num_inference_steps`` times. rng: Generator on the same device. Used by self-forcing renoise loops; pure ODE solvers ignore it. - context: Optional per-AR-step encoder output. Most schedulers - ignore it; schedulers with token-local constraints can use it - without taking over the diffusion-model orchestration. Returns: Clean latent with the same shape, device, and dtype as diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py b/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py index 75a6cbc4e..baaabf78b 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/fm.py @@ -18,7 +18,6 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any import torch from torch import Tensor @@ -213,7 +212,6 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, - context: Any = None, ) -> Tensor: """Run the self-forcing flow-match denoising loop. @@ -222,7 +220,6 @@ def sample( sigma before the network forward. Schedule arithmetic auto-promotes to fp32; the result is cast back to ``initial_noise.dtype``. """ - del context input_dtype = initial_noise.dtype sigmas = self.denoising_sigmas timesteps = self.denoising_step_list diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py index 2723e0bd0..76498d987 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_euler.py @@ -29,7 +29,6 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any import numpy as np import torch @@ -187,7 +186,6 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, - context: Any = None, ) -> Tensor: """Run the explicit Euler denoising loop. @@ -201,7 +199,6 @@ def sample( ``rng`` is unused (deterministic ODE) but accepted for interface conformance. """ - del context input_dtype = initial_noise.dtype N = self.config.num_inference_steps diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py index ff5ca162d..a9ce940a7 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/fm_unipc.py @@ -18,7 +18,6 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any import numpy as np import torch @@ -332,7 +331,6 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, - context: Any = None, ) -> Tensor: """Run the order-2 UniPC predictor-corrector denoising loop. @@ -342,7 +340,6 @@ def sample( fp32; the result is cast back to ``initial_noise.dtype``. ``rng`` is unused (deterministic ODE) but accepted for interface conformance. """ - del context input_dtype = initial_noise.dtype N = self.timesteps.shape[0] diff --git a/integrations/sana/sana_wm/decoder.py b/integrations/sana/sana_wm/decoder.py index f3367f6d7..34c53b8df 100644 --- a/integrations/sana/sana_wm/decoder.py +++ b/integrations/sana/sana_wm/decoder.py @@ -31,8 +31,8 @@ from flashdreams.infra.config import InstantiateConfig from flashdreams.infra.decoder import ( DecoderConfig, - StreamingDecoder, StreamingDecoderCache, + StreamingVideoDecoder, ) from sana_wm._tools import resolve_hf_path from sana_wm.constants import ( @@ -72,8 +72,6 @@ class SanaWMVideoDecoderCache(StreamingDecoderCache): save_stage1: bool = False refiner_seed: int = 42 sink_size: int = 1 - refiner_block_size: int | None = None - refiner_kv_max_frames: int = 11 @dataclass(kw_only=True) @@ -342,8 +340,6 @@ def refine_latents( fps: int, sink_size: int, seed: int, - block_size: int | None, - kv_max_frames: int, ) -> Tensor: """Run the LTX-2 refiner.""" self._ensure_refiner() @@ -354,8 +350,6 @@ def refine_latents( sink_size=int(sink_size), seed=int(seed), progress=True, - block_size=block_size, - kv_max_frames=int(kv_max_frames), ) if self.config.offload_refiner: self.release_runtime() @@ -414,7 +408,7 @@ class SanaWMVideoDecoderConfig(DecoderConfig): ) -class SanaWMVideoDecoder(StreamingDecoder[SanaWMVideoDecoderCache]): +class SanaWMVideoDecoder(StreamingVideoDecoder[SanaWMVideoDecoderCache]): """Decode Stage-1 latents, optionally through the LTX-2 refiner.""" config: SanaWMVideoDecoderConfig @@ -452,8 +446,6 @@ def forward( fps=cache.fps, sink_size=cache.sink_size, seed=cache.refiner_seed, - block_size=cache.refiner_block_size, - kv_max_frames=cache.refiner_kv_max_frames, ) # refine_latents() already releases the refiner when offload_refiner # is set; keep it resident otherwise so a reused pipeline avoids a @@ -483,6 +475,37 @@ def temporal_compression_ratio(self) -> int: """Pixel frame compression ratio after the first latent.""" return SANA_WM_VAE_TEMPORAL_COMPRESSION + def get_output_temporal_size( + self, + autoregressive_index: int, + input_temporal_size: int, + ) -> int: + """Return decoded pixel frames for a SANA-WM latent sequence.""" + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + if input_temporal_size <= 0: + raise ValueError( + f"input_temporal_size must be positive, got {input_temporal_size}." + ) + return 1 + (input_temporal_size - 1) * self.temporal_compression_ratio + + def get_input_temporal_size( + self, + autoregressive_index: int, + output_temporal_size: int, + ) -> int: + """Return latent frames required to decode ``output_temporal_size`` pixels.""" + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + ratio = self.temporal_compression_ratio + remainder = (output_temporal_size - 1) % ratio + if output_temporal_size <= 0 or remainder != 0: + raise ValueError( + "SANA-WM output frame count must be positive and equal to 8k+1; " + f"got {output_temporal_size}." + ) + return ((output_temporal_size - 1) // ratio) + 1 + __all__ = [ "SanaWMDecodedVideo", diff --git a/integrations/sana/sana_wm/refiner.py b/integrations/sana/sana_wm/refiner.py index 8cfce1f96..86e465c1e 100644 --- a/integrations/sana/sana_wm/refiner.py +++ b/integrations/sana/sana_wm/refiner.py @@ -84,18 +84,9 @@ def refine_latents( sink_size: int = 1, seed: int = 42, progress: bool = True, - block_size: int | None = None, - kv_max_frames: int = 11, sigmas: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0), ) -> Tensor: """Refine Stage-1 VAE latents with the sink-bidirectional LTX-2 path.""" - del kv_max_frames - if block_size is not None: - logger.warning( - "[refiner] --refiner-block-size={} requested; running the " - "sink-bidirectional LTX-2 path used by the bidirectional model.", - block_size, - ) if sana_latent.shape[2] <= sink_size: raise ValueError( f"Stage-1 latent has {sana_latent.shape[2]} frames but " diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index ea0e01dd3..6debabab4 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -173,12 +173,6 @@ class SanaWMRunnerConfig(RunnerConfig): sink_size: int = 1 """Number of sink latent frames used by the refiner.""" - refiner_block_size: int | None = None - """Optional refiner AR block size; ``None`` uses sink-bidirectional mode.""" - - refiner_kv_max_frames: int = 11 - """Maximum refiner KV context frames in AR mode.""" - offload_vae: bool = False """Move the VAE to CPU between encode/decode phases.""" @@ -299,8 +293,6 @@ def run(self) -> None: "save_stage1": cfg.save_stage1, "refiner_seed": cfg.refiner_seed, "sink_size": cfg.sink_size, - "refiner_block_size": cfg.refiner_block_size, - "refiner_kv_max_frames": cfg.refiner_kv_max_frames, } ) decoded = pipeline.generate( @@ -411,10 +403,14 @@ def _pipeline_config( quant_backend: ResolvedQuantBackend, ) -> StreamInferencePipelineConfig: """Apply CLI runtime fields to the SANA-WM pipeline literal.""" + scheduler_updates: dict[str, object] = {"num_inference_steps": cfg.step} + if cfg.flow_shift is not None: + scheduler_updates["shift"] = cfg.flow_shift return derive_config( cfg.pipeline, diffusion_model=dict( seed=cfg.seed, + scheduler=scheduler_updates, transformer=dict( config_path=cfg.config_path, checkpoint_path=cfg.model_path, diff --git a/integrations/sana/sana_wm/scheduler.py b/integrations/sana/sana_wm/scheduler.py index 4bf1ed288..168c377db 100644 --- a/integrations/sana/sana_wm/scheduler.py +++ b/integrations/sana/sana_wm/scheduler.py @@ -17,9 +17,7 @@ from __future__ import annotations -import os from dataclasses import dataclass, field -from typing import Any, cast import numpy as np import torch @@ -30,7 +28,6 @@ Scheduler, SchedulerConfig, ) -from sana_wm.transformer import SanaWMStage1Conditioning @dataclass(kw_only=True) @@ -156,17 +153,9 @@ def sample( initial_noise: Tensor, predict_flow: FlowPredictor, rng: torch.Generator | None = None, - context: Any = None, ) -> Tensor: """Run the generic scalar-timestep Euler loop.""" del rng - if isinstance(context, SanaWMStage1Conditioning): - return self._sample_conditioned( - initial_noise=initial_noise, - predict_flow=predict_flow, - conditioning=context, - ) - timesteps = self.timesteps( num_inference_steps=self.config.num_inference_steps, shift=self.config.shift, @@ -183,88 +172,6 @@ def sample( ) return noisy.to(initial_noise.dtype) - def _sample_conditioned( - self, - *, - initial_noise: Tensor, - predict_flow: FlowPredictor, - conditioning: SanaWMStage1Conditioning, - ) -> Tensor: - """Run SANA-WM's first-frame-pinned LTX Euler denoising loop.""" - latents = initial_noise - timesteps = self.timesteps( - num_inference_steps=conditioning.steps, - shift=conditioning.flow_shift, - device=latents.device, - ) - condition_frame_info = dict( - cast( - dict[int, float], - cast(dict[str, object], conditioning.model_kwargs["data_info"]).get( - "condition_frame_info", - {}, - ), - ) - ) - condition_mask = torch.zeros_like(latents) - image_cond_noise_scale = 0.0 - for frame_idx, frame_weight in condition_frame_info.items(): - condition_mask[:, :, int(frame_idx)] = 1 - image_cond_noise_scale = max(image_cond_noise_scale, float(frame_weight)) - - init_latents = latents.clone() - generator = torch.Generator(device=latents.device).manual_seed( - conditioning.seed - ) - iterator = enumerate(timesteps[:-1]) - if os.getenv("DPM_TQDM", "False") == "True": - from tqdm import tqdm - - iterator = tqdm(list(iterator)) - - for step_index, timestep_scalar in iterator: - if image_cond_noise_scale > 0: - latents = _add_noise_to_conditioning_latents( - t=timestep_scalar / self.config.num_train_timesteps, - init_latents=init_latents, - latents=latents, - noise_scale=image_cond_noise_scale, - conditioning_mask=condition_mask, - generator=generator, - ) - - timestep = timestep_scalar.expand(condition_mask.shape).float() - timestep = torch.min( - timestep, - (1 - condition_mask) * float(self.config.num_train_timesteps), - ) - noise_pred = predict_flow(latents, timestep[:, :1, :, 0, 0]) - - latents_dtype = latents.dtype - latents_shape = latents.shape - batch_size, channels, _frames, _height, _width = latents_shape - denoised_latents = self.step_ltx( - model_output=-noise_pred.reshape( - batch_size, - channels, - -1, - ).transpose(1, 2), - timestep=timestep_scalar, - next_timestep=timesteps[step_index + 1], - sample=latents.reshape(batch_size, channels, -1).transpose(1, 2), - per_token_timesteps=timestep.reshape(batch_size, channels, -1)[:, 0], - schedule_timesteps=timesteps, - ) - denoised_latents = denoised_latents.transpose(1, 2).reshape(latents_shape) - tokens_to_denoise_mask = ( - timestep_scalar / self.config.num_train_timesteps - 1e-6 - ) < (1.0 - condition_mask) - latents = torch.where(tokens_to_denoise_mask, denoised_latents, latents) - if latents.dtype != latents_dtype: - latents = latents.to(latents_dtype) - - return latents.detach() - def add_noise( self, clean_input: Tensor, @@ -283,27 +190,6 @@ def add_noise( return ((1.0 - sigma) * clean_input + sigma * noise).to(clean_input.dtype) -def _add_noise_to_conditioning_latents( - *, - t: Tensor, - init_latents: Tensor, - latents: Tensor, - noise_scale: float, - conditioning_mask: Tensor, - generator: torch.Generator, - eps: float = 1e-6, -) -> Tensor: - noise = torch.randn( - latents.shape, - generator=generator, - device=latents.device, - dtype=latents.dtype, - ) - need_to_noise = conditioning_mask > (1.0 - eps) - noised_latents = init_latents + noise_scale * noise * (t**2) - return torch.where(need_to_noise, noised_latents, latents) - - __all__ = [ "SanaWMLTXEulerScheduler", "SanaWMLTXEulerSchedulerConfig", diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py index dd5b1ac19..e9e0abd72 100644 --- a/integrations/sana/sana_wm/transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -316,26 +316,34 @@ def _predict_conditioned( timestep: Tensor, conditioning: SanaWMStage1Conditioning, ) -> Tensor: + model_timestep = _conditioned_frame_timestep( + noisy_latent=noisy_latent, + timestep=timestep, + conditioning=conditioning, + num_train_timesteps=1000, + ) if conditioning.cfg_scale <= 1.0: - return self._predict_with_prompt( + flow = self._predict_with_prompt( noisy_latent, - timestep, + model_timestep, conditioning.condition, _condition_model_kwargs(conditioning.model_kwargs), ) + return _zero_conditioned_frame_flow(flow, conditioning) if conditioning.uncondition is None: raise RuntimeError("CFG was requested without negative prompt embeds.") noise_pred = self._predict_with_prompt( torch.cat([noisy_latent, noisy_latent], dim=0), - torch.cat([timestep, timestep], dim=0), + torch.cat([model_timestep, model_timestep], dim=0), torch.cat([conditioning.uncondition, conditioning.condition], dim=0), _batched_cfg_model_kwargs(conditioning.model_kwargs), ) noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0) - return noise_pred_uncond + conditioning.cfg_scale * ( + flow = noise_pred_uncond + conditioning.cfg_scale * ( noise_pred_text - noise_pred_uncond ) + return _zero_conditioned_frame_flow(flow, conditioning) def _predict_with_prompt( self, @@ -465,6 +473,83 @@ def _batched_cfg_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, obje return kwargs +def _conditioned_frame_timestep( + *, + noisy_latent: Tensor, + timestep: Tensor, + conditioning: SanaWMStage1Conditioning, + num_train_timesteps: int, +) -> Tensor: + """Expand scheduler timesteps and pin conditioned frames to timestep zero.""" + batch, _channels, frames, _height, _width = noisy_latent.shape + if timestep.ndim == 0: + frame_timestep = timestep.reshape(1, 1, 1).expand(batch, 1, frames) + elif timestep.ndim == 1: + frame_timestep = timestep.reshape(batch, 1, 1).expand(batch, 1, frames) + elif timestep.ndim == 3: + frame_timestep = timestep + elif timestep.ndim == noisy_latent.ndim: + frame_timestep = timestep[:, :1, :, 0, 0] + else: + raise ValueError( + "SANA-WM timestep must be scalar, [B], [B, 1, T], or latent-shaped; " + f"got shape={tuple(timestep.shape)}." + ) + if frame_timestep.shape != (batch, 1, frames): + raise ValueError( + "SANA-WM timestep shape is incompatible with the latent: " + f"got {tuple(frame_timestep.shape)}, expected {(batch, 1, frames)}." + ) + + frame_timestep = frame_timestep.to(device=noisy_latent.device, dtype=torch.float32) + condition_frame_mask = _condition_frame_mask( + conditioning, + batch=batch, + frames=frames, + device=noisy_latent.device, + ) + max_timestep = (1.0 - condition_frame_mask) * float(num_train_timesteps) + return torch.minimum(frame_timestep, max_timestep) + + +def _zero_conditioned_frame_flow( + flow: Tensor, + conditioning: SanaWMStage1Conditioning, +) -> Tensor: + """Prevent scalar scheduler updates from moving pinned condition frames.""" + batch, _channels, frames, _height, _width = flow.shape + condition_frame_mask = _condition_frame_mask( + conditioning, + batch=batch, + frames=frames, + device=flow.device, + ).to(dtype=torch.bool) + return torch.where(condition_frame_mask[:, :, :, None, None], 0, flow) + + +def _condition_frame_mask( + conditioning: SanaWMStage1Conditioning, + *, + batch: int, + frames: int, + device: torch.device, +) -> Tensor: + data_info = conditioning.model_kwargs.get("data_info", {}) + condition_frame_info = ( + data_info.get("condition_frame_info", {}) + if isinstance(data_info, dict) + else {} + ) + mask = torch.zeros(batch, 1, frames, dtype=torch.float32, device=device) + if not isinstance(condition_frame_info, dict): + return mask + for frame_idx in condition_frame_info: + index = int(frame_idx) + if 0 <= index < frames: + mask[:, :, index] = 1.0 + return mask + + def _avoid_degenerate_tile_tail( *, sample_extent: int, diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index f39c3eaf3..6c774ec0e 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -94,6 +94,7 @@ ) from flashdreams.infra.config import derive_config +from flashdreams.infra.decoder import StreamingVideoDecoder from flashdreams.infra.diffusion.model import DiffusionModel pytestmark = pytest.mark.ci_cpu @@ -134,6 +135,21 @@ def test_pipeline_uses_sana_diffusion_model() -> None: assert isinstance(pipeline.decoder.refiner, SanaWMLTX2LatentRefinerConfig) +def test_sana_decoder_uses_video_decoder_contract() -> None: + """Expose SANA-WM VAE temporal sizing through the FlashDreams decoder API.""" + decoder = SanaWMVideoDecoderConfig().setup() + + assert isinstance(decoder, StreamingVideoDecoder) + assert decoder.spatial_compression_ratio == 32 + assert decoder.temporal_compression_ratio == 8 + assert decoder.get_output_temporal_size(0, 21) == 161 + assert decoder.get_input_temporal_size(0, 161) == 21 + with pytest.raises(ValueError, match="8k\\+1"): + decoder.get_input_temporal_size(0, 160) + with pytest.raises(ValueError, match="one AR step"): + decoder.get_output_temporal_size(1, 21) + + def test_sana_diffusion_config_instantiates_base_model() -> None: """Use FlashDreams' shared diffusion model rather than a Sana-only runner.""" model = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.setup() @@ -150,6 +166,8 @@ def test_runner_pipeline_config_routes_runtime_fields_to_components() -> None: config_path="local_config.yaml", model_path="local_model.safetensors", stage1_precision="fp8", + step=7, + flow_shift=6.5, no_refiner=True, offload_vae=True, offload_text_encoder=True, @@ -173,6 +191,8 @@ def test_runner_pipeline_config_routes_runtime_fields_to_components() -> None: "local_model.safetensors" ) assert pipeline.diffusion_model.transformer.offload_stage1 is True + assert pipeline.diffusion_model.scheduler.num_inference_steps == 7 + assert pipeline.diffusion_model.scheduler.shift == 6.5 def test_sana_ltx_scheduler_step_pins_zero_timestep_tokens() -> None: @@ -252,9 +272,10 @@ def test_sana_ltx_scheduler_matches_diffusers_per_token_step() -> None: torch.testing.assert_close(actual, expected) -def test_sana_scheduler_context_path_keeps_conditioned_frame_fixed() -> None: - """Let the scheduler own SANA-WM's per-token first-frame constraint.""" +def test_sana_transformer_keeps_conditioned_frame_fixed_with_generic_scheduler() -> None: + """Keep SANA conditioning out of the scheduler and inside the transformer.""" scheduler = SanaWMLTXEulerSchedulerConfig(num_inference_steps=1).setup() + transformer = SanaWMTransformerConfig().setup() initial_noise = torch.zeros((1, 1, 2, 1, 1), dtype=torch.float32) initial_noise[:, :, 0] = 5.0 conditioning = SanaWMStage1Conditioning( @@ -269,14 +290,33 @@ def test_sana_scheduler_context_path_keeps_conditioned_frame_fixed() -> None: seed=0, ) + class DummyModel(torch.nn.Module): + def forward( + self, + noisy_latent: torch.Tensor, + timestep: torch.Tensor, + _prompt_embeds: torch.Tensor, + **_kwargs: object, + ) -> torch.Tensor: + assert timestep.shape == (1, 1, 2) + torch.testing.assert_close(timestep[:, :, 0], torch.zeros((1, 1))) + assert torch.all(timestep[:, :, 1] > 0) + return -torch.ones_like(noisy_latent) + + transformer.model = DummyModel() + transformer._model_built = True + def predict_flow(noisy_latent: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: - assert timestep.shape == (1, 1, 2) - return -torch.ones_like(noisy_latent) + return transformer.predict_flow( + noisy_latent=noisy_latent, + timestep=timestep, + cache=SanaWMTransformerCache(), + input=conditioning, + ) sampled = scheduler.sample( initial_noise=initial_noise, predict_flow=predict_flow, - context=conditioning, ) torch.testing.assert_close(sampled[:, :, 0], initial_noise[:, :, 0]) @@ -375,8 +415,8 @@ def forward( transformer._model_built = True cond_mask = torch.ones((1, 1)) neg_mask = torch.zeros((1, 1)) - camera = torch.zeros((1, 3, 1, 1, 1)) - chunk_plucker = torch.ones((1, 6, 1, 1, 1)) + camera = torch.zeros((1, 3, 2, 1, 1)) + chunk_plucker = torch.ones((1, 6, 2, 1, 1)) conditioning = SanaWMStage1Conditioning( condition=torch.ones((1, 1, 1, 1)), uncondition=torch.zeros((1, 1, 1, 1)), @@ -388,7 +428,7 @@ def forward( "data_info": {"condition_frame_info": {0: 0.0}}, }, first_latent=torch.empty((1, 1, 1, 1, 1)), - latent_shape=(1, 1, 1, 1, 1), + latent_shape=(1, 1, 2, 1, 1), cfg_scale=2.0, flow_shift=1.0, steps=1, @@ -396,17 +436,20 @@ def forward( ) out = transformer.predict_flow( - noisy_latent=torch.zeros((1, 1, 1, 1, 1)), - timestep=torch.zeros((1, 1, 1)), + noisy_latent=torch.zeros((1, 1, 2, 1, 1)), + timestep=torch.full((1, 1, 2), 1000.0), cache=SanaWMTransformerCache(), input=conditioning, ) - torch.testing.assert_close(out, torch.full((1, 1, 1, 1, 1), 2.0)) + torch.testing.assert_close( + out, + torch.tensor([[[[[0.0]], [[2.0]]]]]), + ) assert len(dummy_model.calls) == 1 call = dummy_model.calls[0] - assert call["noisy_latent"].shape == (2, 1, 1, 1, 1) - assert call["timestep"].shape == (2, 1, 1) + assert call["noisy_latent"].shape == (2, 1, 2, 1, 1) + assert call["timestep"].shape == (2, 1, 2) assert call["prompt_embeds"].shape == (2, 1, 1, 1) torch.testing.assert_close(call["mask"], torch.cat([neg_mask, cond_mask], dim=0)) torch.testing.assert_close( @@ -1040,10 +1083,10 @@ def test_refiner_latent_pack_round_trips() -> None: torch.testing.assert_close(unpacked, latents) -def test_refiner_block_size_request_still_executes( +def test_refiner_sink_bidirectional_path_runs_under_inference_mode( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Do not turn refiner block-size requests into a hard failure.""" + """Keep the supported LTX-2 refiner path explicit and inference-only.""" class DummyTransformer(torch.nn.Module): def __init__(self) -> None: @@ -1083,7 +1126,6 @@ def predict_current_x0( latents, "demo", fps=16.0, - block_size=1, progress=False, sigmas=(0.5, 0.0), ) From 5a82a1e697c06b44bf983321469464d51442ae7f Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 17 Jul 2026 14:34:02 -0700 Subject: [PATCH 17/36] Fit SANA trajectories to requested frame count Make the SANA-WM runner choose the snapped rollout length from --num-frames before preparing camera conditioning. Action-derived trajectories now repeat or truncate to that snapped length, so a short action prompt no longer caps the generated video. Fit explicit --camera-path trajectories to the same frame count instead of treating the file length as an upper bound. This matches the existing intrinsics fitting behavior and prevents a 321-frame pose file from truncating a longer requested rollout. Document the updated behavior and add CPU coverage for action repetition, explicit camera-path fitting, and runner input preparation. Tests: - uv run --package flashdreams-sana-wm --extra dev pytest integrations/sana/tests Signed-off-by: Aidan Foster --- integrations/sana/README.md | 6 ++- integrations/sana/sana_wm/camera.py | 61 ++++++++++++++++++++++++++ integrations/sana/sana_wm/runner.py | 34 +++++++++----- integrations/sana/tests/test_camera.py | 32 ++++++++++++++ integrations/sana/tests/test_smoke.py | 30 ++++++++++++- 5 files changed, 148 insertions(+), 15 deletions(-) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index 48e3c0b2a..142108a5d 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -67,8 +67,10 @@ outputs/sana_wm_bf16/sana_wm_generated.mp4 `--intrinsics-path` is optional. When omitted, intrinsics are derived from the first-frame size assuming a centered principal point and a horizontal field of view of `--intrinsics-hfov-deg` (default `90`, matching the demo intrinsics). -Likewise `--camera-path` may be replaced by an `--action` DSL string. A minimal -run therefore needs only an image and a prompt: +Likewise `--camera-path` may be replaced by an `--action` DSL string, which is +repeated or truncated to the requested snapped frame count. Explicit camera +paths are fitted to the same frame count instead of capping the video length. A +minimal run therefore needs only an image and a prompt: ```bash uv run flashdreams-run sana-wm-bidirectional \ diff --git a/integrations/sana/sana_wm/camera.py b/integrations/sana/sana_wm/camera.py index 363534100..063345910 100644 --- a/integrations/sana/sana_wm/camera.py +++ b/integrations/sana/sana_wm/camera.py @@ -221,6 +221,7 @@ def action_string_to_c2w( rotation_speed_deg: float = DEFAULT_ROTATION_SPEED_DEG, pitch_limit_deg: float = DEFAULT_PITCH_LIMIT_DEG, smooth: bool = True, + num_frames: int | None = None, ) -> np.ndarray: """Roll out a camera-to-world trajectory from an action string. @@ -230,11 +231,25 @@ def action_string_to_c2w( rotation_speed_deg: Per-frame angular magnitude in degrees. pitch_limit_deg: Maximum absolute pitch in degrees. smooth: Whether to use the SANA-WM velocity smoothing model. + num_frames: Optional output frame count. When provided, the action + prompt is repeated or truncated to produce exactly this many + camera poses. Returns: ``[F + 1, 4, 4]`` camera-to-world matrices in OpenCV convention. """ per_frame = parse_action_string(action) + if num_frames is not None: + if num_frames < 1: + raise ValueError(f"num_frames must be positive, got {num_frames}.") + target_steps = num_frames - 1 + if target_steps == 0: + per_frame = [] + elif len(per_frame) < target_steps: + repeats = math.ceil(target_steps / len(per_frame)) + per_frame = (per_frame * repeats)[:target_steps] + else: + per_frame = per_frame[:target_steps] integrator = CameraPoseIntegrator(math.radians(pitch_limit_deg)) velocity = VelocityState() poses = [integrator.pose.copy()] @@ -262,6 +277,52 @@ def action_string_to_c2w( return np.stack(poses, axis=0).astype(np.float32) +def fit_camera_trajectory(c2w: np.ndarray, num_frames: int) -> np.ndarray: + """Return a camera-to-world trajectory fitted to ``num_frames`` poses. + + Translations are interpolated linearly. Rotations are interpolated in + matrix space and projected back to the nearest proper rotation, avoiding a + new SciPy dependency for runner-time camera fitting. + """ + c2w = np.asarray(c2w, dtype=np.float32) + if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): + raise ValueError(f"camera trajectory must be [F, 4, 4], got {c2w.shape}.") + if num_frames < 1: + raise ValueError(f"num_frames must be positive, got {num_frames}.") + if c2w.shape[0] == num_frames: + return c2w.copy() + if c2w.shape[0] == 1: + return np.broadcast_to(c2w[:1], (num_frames, 4, 4)).copy() + + old_t = np.linspace(0.0, 1.0, c2w.shape[0], dtype=np.float32) + new_t = np.linspace(0.0, 1.0, num_frames, dtype=np.float32) + fitted = np.broadcast_to(np.eye(4, dtype=np.float32), (num_frames, 4, 4)).copy() + + for axis in range(3): + fitted[:, axis, 3] = np.interp( + new_t, + old_t, + c2w[:, axis, 3], + ).astype(np.float32) + + rotations = c2w[:, :3, :3].reshape(c2w.shape[0], 9) + interp_rotations = np.empty((num_frames, 9), dtype=np.float32) + for idx in range(9): + interp_rotations[:, idx] = np.interp( + new_t, + old_t, + rotations[:, idx], + ).astype(np.float32) + for frame_idx, rotation in enumerate(interp_rotations.reshape(num_frames, 3, 3)): + u, _, vh = np.linalg.svd(rotation.astype(np.float64), full_matrices=False) + projected = u @ vh + if np.linalg.det(projected) < 0.0: + u[:, -1] *= -1.0 + projected = u @ vh + fitted[frame_idx, :3, :3] = projected.astype(np.float32) + return fitted + + def fit_intrinsics_sequence(arr: np.ndarray, num_frames: int) -> np.ndarray: """Return ``arr`` fitted to ``num_frames`` along axis 0.""" arr = np.asarray(arr, dtype=np.float32) diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index 6debabab4..a3bc1a46d 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -35,6 +35,7 @@ from sana_wm.camera import ( action_string_to_c2w, default_intrinsics_vec4, + fit_camera_trajectory, load_intrinsics, resize_center_crop_geometry, snap_num_frames, @@ -100,7 +101,8 @@ class SanaWMRunnerConfig(RunnerConfig): to ~90 degrees; lower values are narrower/more zoomed-in.""" action: str | None = DEFAULT_ACTION - """Action DSL used when ``camera_path`` is not provided.""" + """Action DSL used to derive camera motion when ``camera_path`` is not + provided. The action is repeated or truncated to match ``num_frames``.""" translation_speed: float = 0.025 """Per-frame action translation speed.""" @@ -228,21 +230,34 @@ def _resolve_device(self) -> torch.device: return torch.device(f"cuda:{self.local_rank}") return torch.device(self.config.device) - def _resolve_trajectory(self) -> np.ndarray: - """Load or roll out the camera-to-world trajectory.""" + def _resolve_trajectory(self, *, num_frames: int) -> np.ndarray: + """Load, fit, or roll out the camera-to-world trajectory.""" if self.config.camera_path is not None: c2w = np.load(self.config.camera_path).astype(np.float32) if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): raise ValueError( f"--camera-path must be a [F, 4, 4] .npy; got {c2w.shape}." ) - return c2w + if c2w.shape[0] != num_frames and self.is_rank_zero: + logger.info( + "Fitting --camera-path trajectory from {} to {} frames.", + c2w.shape[0], + num_frames, + ) + return fit_camera_trajectory(c2w, num_frames) if not self.config.action: raise ValueError("SanaWMRunner requires --camera-path or --action.") + if self.is_rank_zero: + logger.info( + "No --camera-path provided; deriving a {}-frame trajectory " + "from --action.", + num_frames, + ) return action_string_to_c2w( self.config.action, translation_speed=self.config.translation_speed, rotation_speed_deg=self.config.rotation_speed_deg, + num_frames=num_frames, ) def run(self) -> None: @@ -336,23 +351,18 @@ def _prepare_inputs(self) -> tuple[object, np.ndarray, np.ndarray, int]: cfg = self.config assert cfg.image_path is not None image = Image.open(cfg.image_path).convert("RGB") - c2w_full = self._resolve_trajectory() - num_frames = min(cfg.num_frames, c2w_full.shape[0]) snapped = snap_num_frames( - num_frames, + cfg.num_frames, stride=SANA_WM_VAE_TEMPORAL_COMPRESSION, - upper_bound=c2w_full.shape[0], ) if snapped != cfg.num_frames and self.is_rank_zero: logger.warning( - "SANA-WM requires num_frames = 8k+1; requested {} snapped to {} " - "(trajectory has {} frames).", + "SANA-WM requires num_frames = 8k+1; requested {} snapped to {}.", cfg.num_frames, snapped, - c2w_full.shape[0], ) num_frames = snapped - c2w = c2w_full[:num_frames] + c2w = self._resolve_trajectory(num_frames=num_frames) resized_size, crop_offset = resize_center_crop_geometry( image.size, diff --git a/integrations/sana/tests/test_camera.py b/integrations/sana/tests/test_camera.py index ba7a52ed8..e54d14eb9 100644 --- a/integrations/sana/tests/test_camera.py +++ b/integrations/sana/tests/test_camera.py @@ -26,6 +26,7 @@ from sana_wm.camera import ( action_string_to_c2w, default_intrinsics_vec4, + fit_camera_trajectory, fit_intrinsics_sequence, load_intrinsics, prepare_camera, @@ -66,6 +67,19 @@ def test_action_string_rolls_out_identity_plus_motion() -> None: assert np.all(np.diff(c2w[:, 2, 3]) > 0) +def test_action_string_repeats_to_requested_frame_count() -> None: + """Use action prompts to derive the requested number of camera poses.""" + c2w = action_string_to_c2w("w-1", smooth=False, num_frames=5) + + assert c2w.shape == (5, 4, 4) + np.testing.assert_allclose(c2w[0], np.eye(4), atol=1e-6) + np.testing.assert_allclose( + c2w[:, 2, 3], + [0.0, 0.025, 0.05, 0.075, 0.1], + atol=1e-6, + ) + + def test_action_string_rejects_unknown_keys() -> None: """Reject invalid action DSL tokens.""" with pytest.raises(ValueError, match="unknown keys"): @@ -99,6 +113,24 @@ def test_fit_intrinsics_sequence_interpolates_short_sequences() -> None: ) +def test_fit_camera_trajectory_interpolates_short_sequences() -> None: + """Fit explicit camera paths to the requested rollout length.""" + source = np.broadcast_to(np.eye(4, dtype=np.float32), (2, 4, 4)).copy() + source[1, 2, 3] = 2.0 + + fitted = fit_camera_trajectory(source, 5) + + assert fitted.shape == (5, 4, 4) + np.testing.assert_allclose(fitted[:, 2, 3], [0.0, 0.5, 1.0, 1.5, 2.0]) + np.testing.assert_allclose( + fitted[:, 3], + np.broadcast_to(np.array([0.0, 0.0, 0.0, 1.0]), (5, 4)), + ) + for rotation in fitted[:, :3, :3]: + np.testing.assert_allclose(rotation.T @ rotation, np.eye(3), atol=1e-5) + assert np.linalg.det(rotation) == pytest.approx(1.0, abs=1e-5) + + @pytest.mark.parametrize( ("array", "expected"), [ diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 6c774ec0e..ad41e950b 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -999,7 +999,7 @@ def test_runner_derives_intrinsics_when_omitted(tmp_path: Path) -> None: image_path=image_path, prompt="demo", intrinsics_path=None, - action="w-40", + action="w-4", num_frames=25, ) runner = cfg.setup() @@ -1014,6 +1014,34 @@ def test_runner_derives_intrinsics_when_omitted(tmp_path: Path) -> None: assert np.all(intrinsics_vec4[:, :2] > 0) +def test_runner_fits_camera_path_to_requested_frames(tmp_path: Path) -> None: + """Do not let short explicit trajectories cap the output frame count.""" + from PIL import Image + + image_path = tmp_path / "frame.png" + camera_path = tmp_path / "camera.npy" + Image.new("RGB", (640, 360), color=(10, 20, 30)).save(image_path) + c2w = np.broadcast_to(np.eye(4, dtype=np.float32), (3, 4, 4)).copy() + c2w[:, 2, 3] = [0.0, 1.0, 2.0] + np.save(camera_path, c2w) + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + image_path=image_path, + prompt="demo", + intrinsics_path=None, + camera_path=camera_path, + num_frames=25, + ) + runner = cfg.setup() + + _image, fitted_c2w, intrinsics_vec4, num_frames = runner._prepare_inputs() + + assert num_frames == 25 + assert fitted_c2w.shape[0] == num_frames + assert intrinsics_vec4.shape == (num_frames, 4) + np.testing.assert_allclose(fitted_c2w[[0, -1], 2, 3], [0.0, 2.0]) + + def test_runner_config_type() -> None: """Keep the exported literal on the SANA-WM runner config subclass.""" assert isinstance(RUNNER_SANA_WM_BIDIRECTIONAL, SanaWMRunnerConfig) From 3f4bbfdf7df12839ef4105d798ef72ab4e1bdffc Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 20 Jul 2026 16:10:22 -0700 Subject: [PATCH 18/36] Adapt runner for refactor --- integrations/sana/sana_wm/runner.py | 42 +++++++++++++---------------- 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index a3bc1a46d..ff7e23ed3 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -28,10 +28,10 @@ import torch from loguru import logger -from flashdreams.core.io.disk import preflight_runtime_write_paths from flashdreams.infra.config import derive_config from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ensure_output_dir, resolve_prompt_value, runner_artifact_path from sana_wm.camera import ( action_string_to_c2w, default_intrinsics_vec4, @@ -110,9 +110,6 @@ class SanaWMRunnerConfig(RunnerConfig): rotation_speed_deg: float = 0.6 """Per-frame action rotation speed in degrees.""" - name: str = "sana_wm" - """Output filename stem.""" - num_frames: int = 161 """Requested output frames before LTX2-VAE stride snapping.""" @@ -205,20 +202,15 @@ def __init__(self, config: SanaWMRunnerConfig) -> None: self.world_size = int(os.environ.get("WORLD_SIZE", "1")) self.global_rank = int(os.environ.get("RANK", "0")) self.is_rank_zero = self.global_rank == 0 - preflight_runtime_write_paths(output_dir=config.output_dir) def _resolve_prompt(self) -> str: """Resolve the prompt from inline text or ``prompt_path``.""" - if self.config.prompt: - return self.config.prompt - if self.config.prompt_path is None: + cfg = self.config + if cfg.prompt: + return resolve_prompt_value(cfg.prompt) + if cfg.prompt_path is None: raise ValueError("SanaWMRunner requires --prompt or --prompt-path.") - prompt = self.config.prompt_path.read_text( - encoding="utf-8", errors="replace" - ).strip() - if not prompt: - raise ValueError(f"Prompt file is empty: {self.config.prompt_path}") - return prompt + return resolve_prompt_value(cfg.prompt_path) def _resolve_device(self) -> torch.device: """Return the device used by SANA-WM.""" @@ -335,11 +327,15 @@ def run(self) -> None: ) if not self.is_rank_zero: return - _write_video(cfg.output_dir, cfg.name, decoded.video_hwc, cfg.fps) + ensure_output_dir(cfg.output_dir) + _write_video( + runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4"), + decoded.video_hwc, + cfg.fps, + ) if decoded.stage1_video_hwc is not None: _write_video( - cfg.output_dir, - f"{cfg.name}_stage1", + runner_artifact_path(cfg.output_dir, f"{cfg.runner_name}_stage1", "mp4"), decoded.stage1_video_hwc, cfg.fps, ) @@ -468,15 +464,13 @@ def _pipeline_config( ) -def _write_video(output_dir: Path, name: str, video_hwc: np.ndarray, fps: int) -> Path: - """Write an HWC uint8 video to the runner output directory.""" +def _write_video(path: Path, video_hwc: np.ndarray, fps: int) -> Path: + """Write an HWC uint8 video to ``path``.""" import imageio.v3 as iio - output_dir.mkdir(parents=True, exist_ok=True) - video_path = output_dir / f"{name}_generated.mp4" - iio.imwrite(video_path, video_hwc, fps=fps) - logger.info("Saved {}", video_path) - return video_path + iio.imwrite(path, video_hwc, fps=fps) + logger.info("Saved {}", path) + return path def _precision_env_updates( From e26d882657456d9d02ad1acb841e156b4ff06868 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 21 Jul 2026 11:06:38 -0700 Subject: [PATCH 19/36] Prune dead env-var code --- integrations/sana/sana_wm/runner.py | 173 ++++++++------------------ integrations/sana/tests/test_smoke.py | 50 -------- 2 files changed, 53 insertions(+), 170 deletions(-) diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index ff7e23ed3..bce53ea96 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -18,8 +18,6 @@ from __future__ import annotations import os -from collections.abc import Iterator, Mapping -from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Literal @@ -274,71 +272,65 @@ def run(self) -> None: refiner_enabled=not cfg.no_refiner, quant_backend=cfg.quant_backend, ) - precision_env = _precision_env_updates( - stage1_precision=cfg.stage1_precision, - refiner_precision=cfg.refiner_precision, - refiner_enabled=not cfg.no_refiner, + prompt = self._resolve_prompt() + image, c2w, intrinsics_vec4, num_frames = self._prepare_inputs() + pipeline_cfg = _pipeline_config( + cfg, + quant_backend=quant_backend, ) - with _temporary_environment(precision_env): - prompt = self._resolve_prompt() - image, c2w, intrinsics_vec4, num_frames = self._prepare_inputs() - pipeline_cfg = _pipeline_config( - cfg, - quant_backend=quant_backend, - ) - pipeline = pipeline_cfg.setup().to(device).eval() - sampling_algo = self._sampling_algo() - if sampling_algo != "flow_euler_ltx": - raise ValueError( - "SANA-WM requires flow_euler_ltx for the " - f"bidirectional runner; got {sampling_algo!r}." - ) - cache = pipeline.initialize_cache( - decoder_context={ - "prompt": prompt, - "fps": cfg.fps, - "save_stage1": cfg.save_stage1, - "refiner_seed": cfg.refiner_seed, - "sink_size": cfg.sink_size, - } + pipeline = pipeline_cfg.setup().to(device).eval() + sampling_algo = self._sampling_algo() + if sampling_algo != "flow_euler_ltx": + raise ValueError( + "SANA-WM requires flow_euler_ltx for the " + f"bidirectional runner; got {sampling_algo!r}." ) - decoded = pipeline.generate( - 0, - cache, - input=SanaWMI2VConditioningRequest( - image=image, - prompt=prompt, - poses_c2w=c2w, - intrinsics_vec4=intrinsics_vec4, - num_frames=num_frames, - fps=cfg.fps, - steps=cfg.step, - cfg_scale=cfg.cfg_scale, - flow_shift=cfg.flow_shift, - seed=cfg.seed, - negative_prompt=cfg.negative_prompt, - ), + cache = pipeline.initialize_cache( + decoder_context={ + "prompt": prompt, + "fps": cfg.fps, + "save_stage1": cfg.save_stage1, + "refiner_seed": cfg.refiner_seed, + "sink_size": cfg.sink_size, + } + ) + decoded = pipeline.generate( + 0, + cache, + input=SanaWMI2VConditioningRequest( + image=image, + prompt=prompt, + poses_c2w=c2w, + intrinsics_vec4=intrinsics_vec4, + num_frames=num_frames, + fps=cfg.fps, + steps=cfg.step, + cfg_scale=cfg.cfg_scale, + flow_shift=cfg.flow_shift, + seed=cfg.seed, + negative_prompt=cfg.negative_prompt, + ), + ) + pipeline.finalize(0, cache) + if not isinstance(decoded, SanaWMDecodedVideo): + raise TypeError( + "SANA-WM pipeline decoder returned " + f"{type(decoded).__name__}, expected SanaWMDecodedVideo." ) - pipeline.finalize(0, cache) - if not isinstance(decoded, SanaWMDecodedVideo): - raise TypeError( - "SANA-WM pipeline decoder returned " - f"{type(decoded).__name__}, expected SanaWMDecodedVideo." - ) - if not self.is_rank_zero: - return - ensure_output_dir(cfg.output_dir) + if not self.is_rank_zero: + return + ensure_output_dir(cfg.output_dir) + _write_video( + runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4"), + decoded.video_hwc, + cfg.fps, + ) + if decoded.stage1_video_hwc is not None: _write_video( - runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4"), - decoded.video_hwc, + runner_artifact_path(cfg.output_dir, f"{cfg.runner_name}_stage1", "mp4"), + decoded.stage1_video_hwc, cfg.fps, ) - if decoded.stage1_video_hwc is not None: - _write_video( - runner_artifact_path(cfg.output_dir, f"{cfg.runner_name}_stage1", "mp4"), - decoded.stage1_video_hwc, - cfg.fps, - ) def _prepare_inputs(self) -> tuple[object, np.ndarray, np.ndarray, int]: """Load/crop the input image and prepare c2w/intrinsics.""" @@ -473,65 +465,6 @@ def _write_video(path: Path, video_hwc: np.ndarray, fps: int) -> Path: return path -def _precision_env_updates( - *, - stage1_precision: Precision, - refiner_precision: Precision, - refiner_enabled: bool, -) -> dict[str, str | None]: - """Return SANA-WM precision environment overrides.""" - updates: dict[str, str | None] = {"SANA_WM_STAGE1_NVFP4": None} - if stage1_precision != "bf16": - updates.update( - { - "SANA_WM_STAGE1_NVFP4": "self_attn+cross+ffn", - "SANA_WM_STAGE1_NVFP4_SCOPE": "block", - "SANA_WM_STAGE1_LINEARIZE_FFN": None, - "SANA_WM_STAGE1_QUANT": _quant_env_value(stage1_precision), - } - ) - - updates["SANA_WM_REFINER_NVFP4"] = None - if refiner_enabled and refiner_precision != "bf16": - updates.update( - { - "SANA_WM_REFINER_NVFP4": "1", - "SANA_WM_REFINER_QUANT": _quant_env_value(refiner_precision), - } - ) - return updates - - -def _quant_env_value(precision: Precision) -> str: - """Return the SANA quantization recipe selector for a precision.""" - if precision == "fp8": - return "fp8block" - if precision == "fp4": - return "nvfp4" - raise ValueError(f"{precision!r} does not use quantized SANA execution.") - - -@contextmanager -def _temporary_environment( - updates: Mapping[str, str | None], -) -> Iterator[None]: - """Temporarily set or unset environment variables.""" - previous = {key: os.environ.get(key) for key in updates} - try: - for key, value in updates.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - yield - finally: - for key, value in previous.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - - def _validate_precision_request( *, device: torch.device, diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index ad41e950b..7630a625e 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -17,7 +17,6 @@ from __future__ import annotations -import os from pathlib import Path from types import SimpleNamespace @@ -48,9 +47,7 @@ SanaWMRunner, SanaWMRunnerConfig, _pipeline_config, - _precision_env_updates, _resolve_quant_backend, - _temporary_environment, _validate_precision_request, ) from sana_wm.conditioning import ( @@ -1169,53 +1166,6 @@ def test_auto_quant_backend_resolves_to_torch_backend() -> None: assert _resolve_quant_backend("auto", []) == "torch" -def test_bf16_precision_clears_quant_env( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Prevent external NVFP4 env vars from changing the default runner path.""" - monkeypatch.setenv("SANA_WM_STAGE1_NVFP4", "1") - monkeypatch.setenv("SANA_WM_REFINER_NVFP4", "1") - updates = _precision_env_updates( - stage1_precision="bf16", - refiner_precision="bf16", - refiner_enabled=True, - ) - - with _temporary_environment(updates): - assert "SANA_WM_STAGE1_NVFP4" not in os.environ - assert "SANA_WM_REFINER_NVFP4" not in os.environ - - assert os.environ["SANA_WM_STAGE1_NVFP4"] == "1" - assert os.environ["SANA_WM_REFINER_NVFP4"] == "1" - - -def test_fp8_precision_sets_quant_env() -> None: - """Map FlashDreams fp8 fields to SANA-WM env selectors.""" - updates = _precision_env_updates( - stage1_precision="fp8", - refiner_precision="fp8", - refiner_enabled=True, - ) - - assert updates["SANA_WM_STAGE1_NVFP4"] == "self_attn+cross+ffn" - assert updates["SANA_WM_STAGE1_QUANT"] == "fp8block" - assert updates["SANA_WM_STAGE1_LINEARIZE_FFN"] is None - assert updates["SANA_WM_REFINER_NVFP4"] == "1" - assert updates["SANA_WM_REFINER_QUANT"] == "fp8block" - - -def test_fp4_precision_sets_sana_quant_env() -> None: - """Map FlashDreams fp4 fields to SANA-WM NVFP4 env selectors.""" - updates = _precision_env_updates( - stage1_precision="fp4", - refiner_precision="fp4", - refiner_enabled=True, - ) - - assert updates["SANA_WM_STAGE1_QUANT"] == "nvfp4" - assert updates["SANA_WM_REFINER_QUANT"] == "nvfp4" - - def test_quantized_precision_requires_cuda() -> None: """Reject fp8/fp4 before loading checkpoints on CPU-only devices.""" with pytest.raises(ValueError, match="requires a CUDA device"): From 2fe96afe06f63ee85b2bcc81b32fc030e2e49032 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 22 Jul 2026 11:37:48 -0700 Subject: [PATCH 20/36] Add WIP model card --- docs/source/index.rst | 8 ++ docs/source/models/index.rst | 9 +++ docs/source/models/sana_wm.rst | 144 +++++++++++++++++++++++++++++++++ 3 files changed, 161 insertions(+) create mode 100644 docs/source/models/sana_wm.rst diff --git a/docs/source/index.rst b/docs/source/index.rst index 16530b788..49971cc80 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -261,6 +261,14 @@ invocation, the checkpoint source, and the per-implementation knobs. Bidirectional Cosmos-Predict2 reference implementations (T2V / I2V, 2B). + .. grid-item-card:: SANA-WM (bidirectional) + :class-card: fd-feature + :link: models/sana_wm + :link-type: doc + + Bidirectional camera-controlled world model (Stage-1 DiT + LTX-2 + refiner, 2.6B). + .. Master toctree: one flat entry per top-level navbar item. Order here = order in the navbar. diff --git a/docs/source/models/index.rst b/docs/source/models/index.rst index 78d767002..3a68a014a 100644 --- a/docs/source/models/index.rst +++ b/docs/source/models/index.rst @@ -28,6 +28,7 @@ Models flashvsr hy_worldplay lingbot_world + sana_wm wan21 FlashDreams runs a growing family of world and video models (text-to-video, @@ -172,6 +173,14 @@ uses, and the settings you can tune. Bidirectional Cosmos-Predict2 reference implementations (T2V / I2V, 2B). + .. grid-item-card:: SANA-WM + :class-card: fd-feature + :link: /models/sana_wm + :link-type: doc + + Bidirectional camera-controlled world model (Stage-1 DiT + LTX-2 + refiner, 2.6B). + .. container:: fd-eyebrow Super-resolution diff --git a/docs/source/models/sana_wm.rst b/docs/source/models/sana_wm.rst new file mode 100644 index 000000000..f8cf6bb75 --- /dev/null +++ b/docs/source/models/sana_wm.rst @@ -0,0 +1,144 @@ +.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +.. SPDX-License-Identifier: Apache-2.0 +.. +.. 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 +=================================== + +.. container:: fd-cta-row + + .. button-link:: https://nvlabs.github.io/Sana/ + :color: primary + + Project page + + .. button-link:: https://arxiv.org/abs/2410.10629 + :color: primary + + arXiv paper + + .. button-link:: https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional + :color: primary + + Model page + + .. button-link:: https://github.com/NVlabs/Sana + :color: primary + + Official code + +SANA-WM is a 2.6B bidirectional, camera-controlled world model from the +`NVlabs/Sana `_ family. Given a first frame, a +text prompt, and a camera trajectory it renders a video clip in a single +bidirectional pass. FlashDreams runs it through the ``sana-wm-bidirectional`` +runner, which pairs a native Stage-1 DiT (loading the public SANA-WM checkpoint +directly) with an LTX-2 refiner for the final decode. + +Requirements +------------ + +- **PyTorch**: >= 2.9. +- **Precision**: BF16 by default. FP8 Stage-1/refiner inference is available on + Hopper or newer GPUs (``sm_90+``) and FP4 on Blackwell (``sm_100+``); both + lower the memory footprint relative to the BF16 default. + +Installation +------------ + +.. code-block:: bash + + # from the repo root + uv sync --package flashdreams-sana-wm --extra dev + +Running the method +------------------ + +To run SANA-WM, launch the ``sana-wm-bidirectional`` runner with a first-frame +image, a prompt, and a camera trajectory. Set +``PYTORCH_ALLOC_CONF=expandable_segments:True`` to keep the allocator from +fragmenting across the Stage-1 + refiner handoff: + +.. code-block:: bash + + PYTORCH_ALLOC_CONF=expandable_segments:True \ + uv run flashdreams-run sana-wm-bidirectional \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_bf16 + +The demo image, prompt, camera, and intrinsics files ship with the SANA-WM +release; any inputs with matching shapes work as well. + +Optional inputs and knobs +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +- ``--intrinsics-path`` is optional. When omitted, intrinsics are derived from + the first-frame size, assuming a centered principal point and a horizontal + field of view of ``--intrinsics-hfov-deg`` (default ``90``, matching the demo + intrinsics). +- ``--camera-path`` can be replaced by an ``--action`` DSL string, so a minimal + run needs only an image and a prompt: + + .. code-block:: bash + + uv run flashdreams-run sana-wm-bidirectional \ + --image-path my_frame.png \ + --prompt "a scene description; describe the world's own motion" \ + --action "w-100,dw-60,w-101" \ + --num-frames 161 \ + --output-dir outputs/mine + +- ``--no-refiner True`` runs Stage-1 only, for diagnostics. +- Quantized Stage-1/refiner inference is opt-in via ``--stage1-precision`` / + ``--refiner-precision`` (``fp8`` on ``sm_90+``, ``fp4`` on ``sm_100+``), with + ``--quant-backend`` (``torch-fp8`` / ``torch-fp4``) to force a backend. + +To inspect all supported CLI arguments and their default values, run: + +.. code-block:: bash + + uv run flashdreams-run sana-wm-bidirectional --help + +What to expect +-------------- + +- **Model checkpoint**: pulled from + ``huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional`` on first run and + cached under ``$HF_HOME``. The LTX-2 refiner and VAE weights are fetched from + their respective ``diffusers`` repositories on first use. +- **First launch**: a few minutes for the download and warmup; subsequent + launches reuse the caches. +- **Outputs**: ``outputs//sana_wm_generated.mp4``. + +See :doc:`/developer_guides/inference_pipeline_overview` for what one pass does +end-to-end. + +Citation +-------- + +If you use SANA-WM, please cite the original SANA work: + +.. code-block:: bibtex + + @misc{xie2024sana, + title={SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers}, + author={Enze Xie and Junsong Chen and Junyu Chen and Han Cai and Haotian Tang and Yujun Lin and Zhekai Zhang and Muyang Li and Ligeng Zhu and Yao Lu and Song Han}, + year={2024}, + eprint={2410.10629}, + archivePrefix={arXiv}, + primaryClass={cs.CV} + } From ec6031e3d60a6e766a8ae8afb6c66fe0560f2efe Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 22 Jul 2026 16:50:56 -0700 Subject: [PATCH 21/36] Fix FP4 quality using NVFP4 scaling --- integrations/sana/README.md | 16 +- integrations/sana/sana_wm/quant.py | 166 +++++++++++++++++++-- integrations/sana/sana_wm/refiner.py | 14 ++ integrations/sana/sana_wm/stage1_model.py | 61 +++++++- integrations/sana/sana_wm/transformer.py | 30 +++- integrations/sana/tests/test_quant_cuda.py | 32 ++++ integrations/sana/tests/test_smoke.py | 85 ++++++++++- 7 files changed, 377 insertions(+), 27 deletions(-) diff --git a/integrations/sana/README.md b/integrations/sana/README.md index 142108a5d..3da123d38 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -61,7 +61,7 @@ uv run flashdreams-run sana-wm-bidirectional \ Expected output: ```text -outputs/sana_wm_bf16/sana_wm_generated.mp4 +outputs/sana_wm_bf16/sana-wm-bidirectional.mp4 ``` `--intrinsics-path` is optional. When omitted, intrinsics are derived from the @@ -83,10 +83,10 @@ uv run flashdreams-run sana-wm-bidirectional \ For Stage-1-only diagnostics, add `--no-refiner True`. -## FP8 and FP4 +## BF16, FP8, and FP4 -The runner defaults to BF16. Quantized Stage-1 inference is opt-in and uses -Torch-based backends by default. +The runner defaults to BF16. Quantized Stage-1 and refiner inference are opt-in +and use Torch-based backends by default. | option | hardware | backend | | --- | --- | --- | @@ -95,6 +95,14 @@ Torch-based backends by default. | `--quant-backend torch-fp8` | Hopper or newer (`sm_90+`) | force FP8 | | `--quant-backend torch-fp4` | Blackwell (`sm_100+`) | force FP4 | +Stage-1 FP8 and FP4 follow the upstream precision CLI scope: self-attention, +cross-attention, and linearized FFN pointwise projections. This integration uses +Torch/Triton scaled-MM replacements instead of TransformerEngine. FP4 uses +NVFP4 W4A4 GEMMs with tiled 16-wide Hadamard rotation enabled, two-level +per-tensor-plus-block scaling, and 2D 16x16 block scaling for weights, matching +the upstream recipe behavior that keeps Stage-1 camera/action conditioning +coherent. + FP8 smoke: ```bash diff --git a/integrations/sana/sana_wm/quant.py b/integrations/sana/sana_wm/quant.py index 3ff3ee0e1..d285e73f1 100644 --- a/integrations/sana/sana_wm/quant.py +++ b/integrations/sana/sana_wm/quant.py @@ -30,7 +30,27 @@ FP8_SCALE_EPS = 1.0e-12 FP4_MAX_E2M1 = 6.0 NVFP4_BLOCK_SIZE = 16 -NVFP4_SCALE_EPS = 1.5258789e-05 +NVFP4_E4M3_SCALE_EPS = 0.015625 +NVFP4_GLOBAL_SCALE_EPS = 1.0e-12 +_RHT16_SCALE = 0.25 +_RHT16_SIGNS = ( + 1.0, + -1.0, + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + 1.0, + 1.0, + -1.0, + -1.0, + -1.0, + -1.0, + 1.0, + 1.0, +) @dataclass(frozen=True) @@ -47,6 +67,9 @@ class TorchScaledMMFP4Recipe: name: str = "torch_scaled_mm_fp4" precision: Literal["fp4"] = "fp4" + use_rht: bool = True + use_global_scale: bool = True + weight_scale_2d: bool = True QuantRecipe = TorchScaledMMFP8Recipe | TorchScaledMMFP4Recipe @@ -79,15 +102,21 @@ def _quantize_nvfp4_kernel( input_ptr, qdata_ptr, scale_ptr, + global_scale_ptr, stride_m, stride_k, rows: tl.constexpr, cols: tl.constexpr, mask_scales: tl.constexpr, + scale_2d: tl.constexpr, + has_global_scale: tl.constexpr, ): fp4_max = 6.0 fp8_max = 448.0 - scale_eps = 1.5258789e-05 + scale_eps = 0.015625 + global_scale = 1.0 + if has_global_scale: + global_scale = tl.load(global_scale_ptr).to(tl.float32) pid_k = tl.program_id(0) pid_m = tl.program_id(1) @@ -100,12 +129,31 @@ def _quantize_nvfp4_kernel( mask=mask, other=0.0, ) - values = values.to(tl.float32).reshape(128, 4, 16) - block_amax = tl.max(tl.abs(values), axis=2) - scales_f32 = block_amax / fp4_max - scales_f32 = tl.clamp(scales_f32, scale_eps, fp8_max) - scales = scales_f32.to(tl.float8e4nv) - quantized = tl.div_rn(values, scales.to(tl.float32)[:, :, None]) + values = values.to(tl.float32) + if scale_2d: + values_4d = values.reshape(8, 16, 4, 16) + block_amax = tl.max(tl.max(tl.abs(values_4d), axis=3), axis=1) + scales_f32 = block_amax / fp4_max + if has_global_scale: + scales_f32 = scales_f32 / global_scale + scales_f32 = tl.clamp(scales_f32, scale_eps, fp8_max) + scales_tile = scales_f32.to(tl.float8e4nv) + scales = tl.broadcast_to(tl.expand_dims(scales_tile, 1), (8, 16, 4)) + scales = scales.reshape(128, 4) + else: + values_3d = values.reshape(128, 4, 16) + block_amax = tl.max(tl.abs(values_3d), axis=2) + scales_f32 = block_amax / fp4_max + if has_global_scale: + scales_f32 = scales_f32 / global_scale + scales_f32 = tl.clamp(scales_f32, scale_eps, fp8_max) + scales = scales_f32.to(tl.float8e4nv) + + values = values.reshape(128, 4, 16) + scale_product = scales.to(tl.float32) + if has_global_scale: + scale_product = scale_product * global_scale + quantized = tl.div_rn(values, scale_product[:, :, None]) if mask_scales: scale_offsets_k = pid_k * 4 + tl.arange(0, 4)[None, :] @@ -135,7 +183,12 @@ def _swizzled_nvfp4_scale_shape(rows: int, cols: int) -> tuple[int, int]: return _ceil_div(rows, 128) * 32, _ceil_div(cols, 64) * 16 -def quantize_nvfp4_swizzled(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: +def quantize_nvfp4_swizzled( + input: torch.Tensor, + *, + global_scale: torch.Tensor | None = None, + scale_2d: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: """Quantize a 2D CUDA tensor to NVFP4 qdata plus swizzled E4M3 scales.""" _require_fp4_dtype() if input.dim() != 2: @@ -160,15 +213,45 @@ def quantize_nvfp4_swizzled(input: torch.Tensor) -> tuple[torch.Tensor, torch.Te input_2d, qdata, scales, + global_scale if global_scale is not None else input_2d, input_2d.stride(0), input_2d.stride(1), rows, cols, mask_scales=(rows % 128 != 0 or cols % 64 != 0), + scale_2d=scale_2d, + has_global_scale=global_scale is not None, ) return qdata, scales +def nvfp4_global_scale(input: torch.Tensor) -> torch.Tensor: + """Return the per-tensor FP32 global scale for hierarchical NVFP4.""" + amax = input.detach().abs().amax().to(torch.float32) + scale = amax / (FP8_MAX_E4M3 * FP4_MAX_E2M1) + return scale.clamp_min(NVFP4_GLOBAL_SCALE_EPS).reshape(()) + + +def apply_rht16(input: torch.Tensor) -> torch.Tensor: + """Apply the tiled 16-wide random Hadamard transform used by NVFP4.""" + if input.shape[-1] % 16 != 0: + raise ValueError( + f"RHT16 requires the last dimension to be divisible by 16, got {input.shape[-1]}." + ) + original_shape = input.shape + work = input.to(torch.float32).reshape(-1, 16) + signs = torch.tensor(_RHT16_SIGNS, device=input.device, dtype=work.dtype) + work = work * signs + width = 1 + while width < 16: + work = work.reshape(-1, 16 // (2 * width), 2, width) + left = work[:, :, 0, :] + right = work[:, :, 1, :] + work = torch.stack((left + right, left - right), dim=2).reshape(-1, 16) + width *= 2 + return (work * _RHT16_SCALE).reshape(original_shape).contiguous() + + class TorchScaledMMFP8Linear(nn.Module): """Inference-only FP8 Linear using ``torch._scaled_mm``. @@ -306,8 +389,10 @@ def __init__( weight: torch.Tensor, weight_qdata: torch.Tensor, weight_scale: torch.Tensor, + weight_global_scale: torch.Tensor | None, bias: torch.Tensor | None, out_dtype: torch.dtype, + use_rht: bool, ) -> None: super().__init__() expected_weight_shape = (weight_qdata.shape[0], weight_qdata.shape[1] * 2) @@ -330,9 +415,17 @@ def __init__( self.out_features = int(weight_qdata.shape[0]) self.in_features = int(weight_qdata.shape[1] * 2) self.out_dtype = out_dtype + self.use_rht = use_rht self.register_buffer("weight", weight.detach().contiguous()) self.register_buffer("weight_qdata", weight_qdata.contiguous()) self.register_buffer("weight_scale", weight_scale.contiguous()) + if weight_global_scale is None: + self.register_buffer("weight_global_scale", None) + else: + self.register_buffer( + "weight_global_scale", + weight_global_scale.detach().to(torch.float32).reshape(()), + ) if bias is None: self.register_buffer("bias", None) else: @@ -344,17 +437,32 @@ def from_linear( source: nn.Linear, *, out_dtype: torch.dtype, + use_rht: bool = True, + use_global_scale: bool = True, + weight_scale_2d: bool = True, ) -> "TorchScaledMMFP4Linear": """Create an NVFP4 replacement from a source ``nn.Linear``.""" _require_fp4_dtype() - weight_qdata, weight_scale = quantize_nvfp4_swizzled(source.weight.detach()) + weight_for_quant = source.weight.detach() + if use_rht: + weight_for_quant = apply_rht16(weight_for_quant) + weight_global_scale = ( + nvfp4_global_scale(weight_for_quant) if use_global_scale else None + ) + weight_qdata, weight_scale = quantize_nvfp4_swizzled( + weight_for_quant, + global_scale=weight_global_scale, + scale_2d=weight_scale_2d, + ) bias = source.bias.detach() if source.bias is not None else None replacement = cls( weight=source.weight.detach().to(device=source.weight.device), weight_qdata=weight_qdata, weight_scale=weight_scale, + weight_global_scale=weight_global_scale, bias=bias.to(device=source.weight.device) if bias is not None else None, out_dtype=out_dtype, + use_rht=use_rht, ) replacement.train(source.training) return replacement @@ -378,17 +486,42 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: dtype=self.out_dtype, ) - input_qdata, input_scale = quantize_nvfp4_swizzled(input_2d) + input_for_quant = apply_rht16(input_2d) if self.use_rht else input_2d + input_global_scale = ( + nvfp4_global_scale(input_for_quant) + if self.weight_global_scale is not None + else None + ) + input_qdata, input_scale = quantize_nvfp4_swizzled( + input_for_quant, + global_scale=input_global_scale, + ) + add_bias_after_scale = ( + self.bias is not None + and self.weight_global_scale is not None + and input_global_scale is not None + ) output = torch._scaled_mm( input_qdata.view(torch.float4_e2m1fn_x2), self.weight_qdata.t().view(torch.float4_e2m1fn_x2), input_scale.view(torch.float8_e4m3fn), self.weight_scale.view(torch.float8_e4m3fn), - bias=self.bias.to(device=input.device, dtype=self.out_dtype) - if self.bias is not None - else None, + bias=None + if add_bias_after_scale + else ( + self.bias.to(device=input.device, dtype=self.out_dtype) + if self.bias is not None + else None + ), out_dtype=self.out_dtype, ) + if input_global_scale is not None and self.weight_global_scale is not None: + output = output * ( + input_global_scale.to(device=output.device, dtype=output.dtype) + * self.weight_global_scale.to(device=output.device, dtype=output.dtype) + ) + if add_bias_after_scale: + output = output + self.bias.to(device=output.device, dtype=output.dtype) return output.reshape(*leading_shape, self.out_features) @@ -495,6 +628,9 @@ def _replace_linear_with_quant( replacement = TorchScaledMMFP4Linear.from_linear( child, out_dtype=out_dtype, + use_rht=recipe.use_rht, + use_global_scale=recipe.use_global_scale, + weight_scale_2d=recipe.weight_scale_2d, ) else: raise ValueError(f"Unsupported SANA-WM quant recipe: {recipe!r}.") @@ -544,6 +680,8 @@ def _require_fp4_dtype() -> None: "TorchScaledMMFP4Recipe", "TorchScaledMMFP8Linear", "TorchScaledMMFP8Recipe", + "apply_rht16", + "nvfp4_global_scale", "quantize_nvfp4_swizzled", "replace_linear_with_quant", "replace_linear_with_torch_fp4", diff --git a/integrations/sana/sana_wm/refiner.py b/integrations/sana/sana_wm/refiner.py index 86e465c1e..0d7c61107 100644 --- a/integrations/sana/sana_wm/refiner.py +++ b/integrations/sana/sana_wm/refiner.py @@ -187,6 +187,20 @@ def _prepare_quantization(self) -> None: f"layers; skipped={skipped}." ) self._quantized = True + recipe_detail = "" + if isinstance(recipe, TorchScaledMMFP4Recipe): + recipe_detail = ( + f" rht={recipe.use_rht}" + f" global_scale={recipe.use_global_scale}" + f" weight_scale_2d={recipe.weight_scale_2d}" + ) + logger.info( + "[refiner-quant] precision={}{} converted {} Linear layers (skipped {})", + self.precision, + recipe_detail, + converted, + skipped, + ) def _ensure_text_encoder(self) -> None: """Load the Gemma tokenizer + encoder once and cache them in CPU RAM. diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index e6615e9a9..522b2b44c 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -252,6 +252,62 @@ def forward(self, x: Tensor, *, frames: int, height: int, width: int) -> Tensor: return x_time.permute(0, 2, 3, 1).reshape(batch, tokens, channels) +class _LinearizedPointwiseConv2d(nn.Module): + """Run an existing 1x1 ``Conv2dContainer`` through a Linear module.""" + + def __init__(self, conv_layer: nn.Module) -> None: + super().__init__() + conv = getattr(conv_layer, "conv", None) + if not isinstance(conv, nn.Conv2d): + raise ValueError("expected Conv2dContainer.conv to be nn.Conv2d.") + if conv.kernel_size != (1, 1) or conv.stride != (1, 1) or conv.padding != (0, 0): + raise ValueError("only exact 1x1 pointwise Conv2d can be linearized.") + if conv.dilation != (1, 1) or conv.groups != 1: + raise ValueError("grouped or dilated pointwise Conv2d cannot be linearized.") + self.linear = nn.Linear( + conv.in_channels, + conv.out_channels, + bias=conv.bias is not None, + device=conv.weight.device, + dtype=conv.weight.dtype, + ) + with torch.no_grad(): + self.linear.weight.copy_(conv.weight.flatten(1)) + if conv.bias is not None: + self.linear.bias.copy_(conv.bias) + + def forward(self, x: Tensor) -> Tensor: + """Apply the pointwise projection while preserving NCHW layout.""" + if x.dim() != 4: + raise ValueError(f"expected NCHW input, got shape {tuple(x.shape)}.") + batch, _channels, height, width = x.shape + x_nhwc = x.permute(0, 2, 3, 1).reshape(batch * height * width, -1) + y = self.linear(x_nhwc) + return y.reshape(batch, height, width, -1).permute(0, 3, 1, 2).contiguous() + + +def linearize_stage1_ffn_for_quant(module: nn.Module) -> tuple[int, int]: + """Expose Stage-1 pointwise FFN convolutions as Linears for quantization.""" + converted = 0 + skipped = 0 + for child in module.modules(): + if not isinstance(child, GLUMBConvTemp): + continue + for attr in ("inverted_conv", "point_conv"): + pointwise = getattr(child, attr) + if isinstance(pointwise, _LinearizedPointwiseConv2d): + continue + try: + replacement = _LinearizedPointwiseConv2d(pointwise) + except ValueError: + skipped += 1 + continue + replacement.train(pointwise.training) + setattr(child, attr, replacement) + converted += 1 + return converted, skipped + + class Stage1SelfAttention(nn.Module): """Self/camera attention parameter container for one Stage-1 block.""" @@ -751,8 +807,8 @@ def _gdn_key_scale(head_dim: int, HW: tuple[int, int, int]) -> float: return (head_dim**-0.5) * ((HW[1] * HW[2]) ** -0.5) -def _apply_output_gate(out: Tensor, gate_x: Tensor, gate: nn.Linear) -> Tensor: - gate_values = F.silu(F.linear(gate_x, gate.weight, gate.bias).float()) +def _apply_output_gate(out: Tensor, gate_x: Tensor, gate: nn.Module) -> Tensor: + gate_values = F.silu(gate(gate_x).float()) return out * gate_values.to(dtype=out.dtype) @@ -1268,4 +1324,5 @@ def _invert_se3(transforms: Tensor) -> Tensor: "SanaWMStage1Spec", "Stage1CrossAttention", "Stage1SelfAttention", + "linearize_stage1_ffn_for_quant", ] diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py index e9e0abd72..0b07b2810 100644 --- a/integrations/sana/sana_wm/transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -46,7 +46,7 @@ TorchScaledMMFP8Recipe, replace_linear_with_quant, ) -from sana_wm.stage1_model import SanaWMStage1Model +from sana_wm.stage1_model import SanaWMStage1Model, linearize_stage1_ffn_for_quant Precision = Literal["bf16", "fp8", "fp4"] QuantBackend = Literal["auto", "torch", "torch-fp8", "torch-fp4"] @@ -59,15 +59,13 @@ "^t_block", "^y_embedder", "^final_layer", - # SANA attention code accesses output_gate.weight directly; replacing it - # with a module that only has quantized buffers breaks that call site. - r"\.output_gate$", ) _STAGE1_QUANT_INCLUDE_DEFAULTS = ( r"^blocks\.\d+\.attn\.qkv$", r"^blocks\.\d+\.attn\.proj$", r"^blocks\.\d+\.attn\.beta_proj$", r"^blocks\.\d+\.attn\.gate_proj$", + r"^blocks\.\d+\.attn\.output_gate$", r"^blocks\.\d+\.cross_attn\.", r"^blocks\.\d+\.mlp\.inverted_conv\.linear$", r"^blocks\.\d+\.mlp\.point_conv\.linear$", @@ -418,12 +416,19 @@ def _prepare_stage1_quant(self) -> None: recipe = TorchScaledMMFP8Recipe() else: recipe = TorchScaledMMFP4Recipe() + linearized, linearize_skipped = linearize_stage1_ffn_for_quant(self.model) + if linearized > 0 or linearize_skipped > 0: + logger.info( + "[stage1-quant] linearized {} FFN pointwise convs (skipped {})", + linearized, + linearize_skipped, + ) converted, skipped = replace_linear_with_quant( self.model, recipe=recipe, params_dtype=self._ensure_weight_dtype(), skip_patterns=_STAGE1_QUANT_SKIP_DEFAULTS, - include_patterns=_STAGE1_QUANT_INCLUDE_DEFAULTS, + include_patterns=_stage1_quant_include_patterns(), ) if converted <= 0: raise RuntimeError( @@ -431,9 +436,17 @@ def _prepare_stage1_quant(self) -> None: f"Stage-1 Linear layers; skipped={skipped}." ) self._stage1_quantized = True + recipe_detail = "" + if isinstance(recipe, TorchScaledMMFP4Recipe): + recipe_detail = ( + f" rht={recipe.use_rht}" + f" global_scale={recipe.use_global_scale}" + f" weight_scale_2d={recipe.weight_scale_2d}" + ) logger.info( - "[stage1-quant] precision={} converted {} Linear layers (skipped {})", + "[stage1-quant] precision={}{} converted {} Linear layers (skipped {})", self.config.stage1_precision, + recipe_detail, converted, skipped, ) @@ -441,6 +454,11 @@ def _prepare_stage1_quant(self) -> None: torch.cuda.empty_cache() +def _stage1_quant_include_patterns() -> tuple[str, ...]: + """Return Stage-1 Linear names eligible for FP8/FP4 quantization.""" + return _STAGE1_QUANT_INCLUDE_DEFAULTS + + def _require_conditioning( cache: SanaWMTransformerCache, ) -> SanaWMStage1Conditioning: diff --git a/integrations/sana/tests/test_quant_cuda.py b/integrations/sana/tests/test_quant_cuda.py index 9034c0c7f..ff6265760 100644 --- a/integrations/sana/tests/test_quant_cuda.py +++ b/integrations/sana/tests/test_quant_cuda.py @@ -77,3 +77,35 @@ def test_fp4_linear_runs_scaled_mm_on_blackwell() -> None: assert output.shape == (4, 5, 64) assert output.dtype == torch.bfloat16 assert torch.isfinite(output.float()).all() + + +def test_quantized_linears_roughly_track_source_layer() -> None: + """Catch gross scale/layout regressions beyond expected FP4/FP8 noise.""" + _require_quant_cuda() + major, minor = torch.cuda.get_device_capability() + if major < 10: + pytest.skip(f"NVFP4 requires Blackwell-class CUDA, got sm_{major}{minor}") + + torch.manual_seed(123) + source = torch.nn.Linear(32, 64, bias=True, device="cuda", dtype=torch.bfloat16) + inputs = torch.randn(4, 5, 32, device="cuda", dtype=torch.bfloat16) + reference = source(inputs).float() + + fp8 = TorchScaledMMFP8Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + fp4 = TorchScaledMMFP4Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + + fp8_error = (fp8(inputs).float() - reference).abs() + fp4_error = (fp4(inputs).float() - reference).abs() + fp8_rel_mae = fp8_error.mean() / reference.abs().mean().clamp_min(1e-6) + fp4_rel_mae = fp4_error.mean() / reference.abs().mean().clamp_min(1e-6) + + assert fp8_error.max().item() <= 0.10 + assert fp8_rel_mae.item() <= 0.06 + assert fp4_error.max().item() <= 0.35 + assert fp4_rel_mae.item() <= 0.20 diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 7630a625e..e079e56ed 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -75,19 +75,24 @@ SanaWMStage1Conditioning, _avoid_degenerate_tile_tail, _load_inference_config, + _stage1_quant_include_patterns, ) from sana_wm.refiner import SanaWMLTX2Refiner, _pack_latents, _unpack_latents from sana_wm.quant import ( + apply_rht16, + nvfp4_global_scale, TorchScaledMMFP4Linear, TorchScaledMMFP8Linear, replace_linear_with_torch_fp4, replace_linear_with_torch_fp8, ) from sana_wm.stage1_model import ( + GLUMBConvTemp, SANA_WM_STAGE1_SPEC, SanaWMStage1Model, SanaWMStage1Spec, Stage1SelfAttention, + linearize_stage1_ffn_for_quant, ) from flashdreams.infra.config import derive_config @@ -1052,6 +1057,81 @@ def test_runner_defaults_to_bf16_precision() -> None: assert RUNNER_SANA_WM_BIDIRECTIONAL.no_refiner is False +def test_stage1_quant_scope_matches_upstream_precision_cli() -> None: + """Keep FP8 and FP4 on upstream's self-attn + cross-attn + FFN scope.""" + patterns = _stage1_quant_include_patterns() + + assert patterns is not None + assert r"^blocks\.\d+\.attn\.qkv$" in patterns + assert r"^blocks\.\d+\.attn\.output_gate$" in patterns + assert r"^blocks\.\d+\.cross_attn\." in patterns + assert r"^blocks\.\d+\.mlp\.inverted_conv\.linear$" in patterns + assert r"^blocks\.\d+\.mlp\.point_conv\.linear$" in patterns + + +def test_stage1_ffn_linearization_preserves_forward() -> None: + """Expose pointwise FFN convs as Linear modules without changing math.""" + torch.manual_seed(0) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + mlp = GLUMBConvTemp(spec).eval() + inputs = torch.randn(1, 8, 16) + reference = mlp(inputs, frames=2, height=2, width=2) + + converted, skipped = linearize_stage1_ffn_for_quant(mlp) + output = mlp(inputs, frames=2, height=2, width=2) + + assert converted == 2 + assert skipped == 0 + assert hasattr(mlp.inverted_conv, "linear") + assert hasattr(mlp.point_conv, "linear") + torch.testing.assert_close(output, reference) + + +def test_fp4_rht16_is_orthogonal() -> None: + """The tiled Hadamard rotation must preserve dot products before quantization.""" + inputs = torch.randn(3, 4, 32) + weights = torch.randn(5, 32) + + rotated_inputs = apply_rht16(inputs) + rotated_weights = apply_rht16(weights) + + torch.testing.assert_close( + rotated_inputs.reshape(-1, 32) @ rotated_weights.t(), + inputs.reshape(-1, 32) @ weights.t(), + atol=1.0e-5, + rtol=1.0e-5, + ) + + +def test_fp4_rht16_rejects_unaligned_last_dimension() -> None: + """NVFP4 RHT is tiled in groups of sixteen values.""" + with pytest.raises(ValueError, match="divisible by 16"): + apply_rht16(torch.randn(2, 15)) + + +def test_fp4_global_scale_uses_nvfp4_dynamic_range() -> None: + """The hierarchical FP4 global scale should target E4M3 x E2M1 range.""" + inputs = torch.tensor([-2688.0, 0.0, 1344.0]) + + assert nvfp4_global_scale(inputs).item() == pytest.approx(1.0) + + def test_refiner_is_flashdreams_owned( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1381,8 +1461,11 @@ def from_linear( source: torch.nn.Linear, *, out_dtype: torch.dtype, + use_rht: bool = True, + use_global_scale: bool = True, + weight_scale_2d: bool = True, ) -> TorchScaledMMFP4Linear: - del out_dtype + del out_dtype, use_rht, use_global_scale, weight_scale_2d instance = cls.__new__(cls) torch.nn.Module.__init__(instance) instance.in_features = source.in_features From 8c6a6e4c484991d42f705e9e1cbe081eb4776d4c Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 23 Jul 2026 17:06:38 -0700 Subject: [PATCH 22/36] Add benchmarking scripts --- docs/source/models/sana_wm.rst | 5 +- integrations/sana/README.md | 26 +- .../sana/tests/parity_check/.gitignore | 9 + .../parity_check/BENCHMARK_FINALIZATION.md | 61 + .../sana/tests/parity_check/PERF_HANDOFF.md | 68 + .../sana/tests/parity_check/README.md | 97 + integrations/sana/tests/parity_check/bench.sh | 263 +++ .../sana/tests/parity_check/bench_summary.py | 327 ++++ .../tests/parity_check/bench_sweep_summary.py | 128 ++ .../sana/tests/parity_check/changes.patch | 215 +++ .../parity_check/compat/mmcv/__init__.py | 116 ++ .../tests/parity_check/compat/mmcv/runner.py | 35 + .../compat/mmcv/utils/__init__.py | 13 + .../parity_check/compat/mmcv/utils/logging.py | 8 + .../sana/tests/parity_check/diff_parity.py | 79 + .../sana/tests/parity_check/pyproject.toml | 47 + integrations/sana/tests/parity_check/run.sh | 185 ++ .../sana/tests/parity_check/run_native.py | 319 ++++ integrations/sana/tests/parity_check/uv.lock | 1676 +++++++++++++++++ 19 files changed, 3669 insertions(+), 8 deletions(-) create mode 100644 integrations/sana/tests/parity_check/.gitignore create mode 100644 integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md create mode 100644 integrations/sana/tests/parity_check/PERF_HANDOFF.md create mode 100644 integrations/sana/tests/parity_check/README.md create mode 100644 integrations/sana/tests/parity_check/bench.sh create mode 100644 integrations/sana/tests/parity_check/bench_summary.py create mode 100644 integrations/sana/tests/parity_check/bench_sweep_summary.py create mode 100644 integrations/sana/tests/parity_check/changes.patch create mode 100644 integrations/sana/tests/parity_check/compat/mmcv/__init__.py create mode 100644 integrations/sana/tests/parity_check/compat/mmcv/runner.py create mode 100644 integrations/sana/tests/parity_check/compat/mmcv/utils/__init__.py create mode 100644 integrations/sana/tests/parity_check/compat/mmcv/utils/logging.py create mode 100644 integrations/sana/tests/parity_check/diff_parity.py create mode 100644 integrations/sana/tests/parity_check/pyproject.toml create mode 100644 integrations/sana/tests/parity_check/run.sh create mode 100644 integrations/sana/tests/parity_check/run_native.py create mode 100644 integrations/sana/tests/parity_check/uv.lock diff --git a/docs/source/models/sana_wm.rst b/docs/source/models/sana_wm.rst index f8cf6bb75..b58ccb795 100644 --- a/docs/source/models/sana_wm.rst +++ b/docs/source/models/sana_wm.rst @@ -65,13 +65,10 @@ Running the method ------------------ To run SANA-WM, launch the ``sana-wm-bidirectional`` runner with a first-frame -image, a prompt, and a camera trajectory. Set -``PYTORCH_ALLOC_CONF=expandable_segments:True`` to keep the allocator from -fragmenting across the Stage-1 + refiner handoff: +image, a prompt, and a camera trajectory: .. code-block:: bash - PYTORCH_ALLOC_CONF=expandable_segments:True \ uv run flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ diff --git a/integrations/sana/README.md b/integrations/sana/README.md index 3da123d38..16f64236b 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -21,7 +21,7 @@ Current scope: | Stage-1 BF16 | FlashDreams DiT execution. | | Stage-1 FP8 | PyTorch `_scaled_mm` backend. | | Stage-1 FP4 | Triton quantization plus PyTorch `_scaled_mm`. | -| VAE decode | Direct `diffusers` LTX2 VAE use with low-memory tiling. | +| VAE decode | Direct `diffusers` LTX2 VAE use with upstream-matched tiling. | | LTX-2 refiner | Direct `diffusers` LTX-2 transformer and Gemma connector use. | ## Runner @@ -48,7 +48,6 @@ well. ## Run ```bash -PYTORCH_ALLOC_CONF=expandable_segments:True \ uv run flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ @@ -106,7 +105,6 @@ coherent. FP8 smoke: ```bash -PYTORCH_ALLOC_CONF=expandable_segments:True \ uv run flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ @@ -121,7 +119,6 @@ uv run flashdreams-run sana-wm-bidirectional \ FP4 smoke: ```bash -PYTORCH_ALLOC_CONF=expandable_segments:True \ uv run flashdreams-run sana-wm-bidirectional \ --image-path ../Sana/asset/sana_wm/demo_0.png \ --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ @@ -133,6 +130,27 @@ uv run flashdreams-run sana-wm-bidirectional \ --refiner-precision fp4 ``` +## Parity and benchmark + +The upstream comparison harness lives under `tests/parity_check/`, matching the +pattern used by the other benchmarked FlashDreams integrations. + +```bash +cd integrations/sana/tests/parity_check + +# Upstream + FlashDreams parity artifacts and frame diff. +bash run.sh + +# Matched upstream-vs-FlashDreams benchmark report. +DEVICE_LABEL="RTX PRO 6000 Blackwell" bash bench.sh +``` + +`bench.sh` writes `outputs/bench/bench.md` for review and `outputs/bench/perf.md` +for model-card chart data. Both use the same benchmark metric: post-load +generation latency per generated frame. Stage-1 DiT, conditioning/encode, VAE +decode, optional refiner, and memory rows are reported as the breakdown of that +metric. + ## Tests CPU-safe tests cover import, config boundaries, action parsing, intrinsics, diff --git a/integrations/sana/tests/parity_check/.gitignore b/integrations/sana/tests/parity_check/.gitignore new file mode 100644 index 000000000..2df75f55e --- /dev/null +++ b/integrations/sana/tests/parity_check/.gitignore @@ -0,0 +1,9 @@ +/Sana/ +/.venv/ +/outputs/ +__pycache__/ +*.pyc +*.npy +*.mp4 +*.json +*.log diff --git a/integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md b/integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md new file mode 100644 index 000000000..d6c6f2654 --- /dev/null +++ b/integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md @@ -0,0 +1,61 @@ + + +# SANA-WM Benchmark Finalization + +Use this after the perf optimization pass is done and the final BF16, FP8, and +FP4 numbers are ready for the model card. + +1. Run the precision sweep: + + ```bash + cd integrations/sana/tests/parity_check + DEVICE_LABEL="" BENCH_PRECISIONS=bf16,fp8,fp4 bash bench.sh + ``` + +2. Copy `outputs/bench/perf.md` to + `docs/source/_static/performance/sana_wm/perf-.md`. + +3. Add a `Profiling benchmark` section to `docs/source/models/sana_wm.rst` + between `What to expect` and `Citation`. + +4. The section should describe post-load generation latency per generated frame + for FlashDreams SANA-WM versus the official `NVlabs/Sana` implementation + under matched BF16, FP8, and FP4 settings. + +5. Use the standard benchmark chart block: + + ```rst + Profiling benchmark + ------------------- + + Here is the profiling benchmark on post-load generation latency per generated + frame for FlashDreams SANA-WM compared to the + `official SANA-WM implementation `_ under + matched BF16, FP8, and FP4 settings. + + .. raw:: html + +
+
+
+

+ This chart shows post-load generation latency per generated frame in milliseconds on a single GPU. + For the official SANA-WM implementation, see + this instruction. +

+
+
+ + ``` + +Keep the chart metric unchanged: post-load generation latency per generated +frame, in ms/frame. diff --git a/integrations/sana/tests/parity_check/PERF_HANDOFF.md b/integrations/sana/tests/parity_check/PERF_HANDOFF.md new file mode 100644 index 000000000..aaed1d654 --- /dev/null +++ b/integrations/sana/tests/parity_check/PERF_HANDOFF.md @@ -0,0 +1,68 @@ + + +# SANA-WM Perf Handoff + +This note is only for the next SANA-WM performance investigation. It should not +be copied into the model card or public README. + +## Observed Delta + +The best current BF16, no-refiner, 121-frame, 60-step run still has +FlashDreams slower than upstream: + +| artifact | official | FlashDreams | gap | +| --- | ---: | ---: | ---: | +| `outputs/bf16_step60_bidirectional_default/bench.md` | 640.89 ms/frame | 663.60 ms/frame | +22.71 ms/frame | + +Component medians from that run: + +| component | official | FlashDreams | +| --- | ---: | ---: | +| wall | 77.55 s | 80.30 s | +| Stage-1 DiT | 75,565.73 ms | 78,237.04 ms | +| VAE decode row | 1,340.43 ms | 1,584.82 ms | +| peak memory | 14.37 GiB | 17.98 GiB | + +The main target is the 60-step Stage-1 DiT path. The decode row also still has +a smaller gap. + +## Perf Theories Tested + +| theory | result | evidence | +| --- | --- | --- | +| FlashDreams was slower because VAE tiling used low-memory tiles. | Partly right. | Early probes showed FlashDreams VAE decode around 5.2 s. Switching to upstream-style tiles dropped it to about 1.5 s. This fixed a large initial problem, but not the remaining 60-step gap. | +| Remaining VAE gap is mostly raw VAE forward kernels. | Probably wrong. | Later internal VAE timing was close to upstream, while the harness decode row remained slower. The remaining cost is likely after the VAE forward: clamp, permute, CPU transfer, numpy materialization, or `torch.cuda.empty_cache()`. | +| Matching upstream's bidirectional chunk default would close the gap. | Correct for config parity, insufficient for speed. | `SanaWMStage1Spec.chunk_size=None` with `chunk_split_strategy="first_chunk_plus_one"` matches upstream. The 8-step probe became slightly favorable to FlashDreams, but the 60-step run stayed slower. | +| Softmax camera Q/K/V transforms staying in higher precision were slowing blocks. | Right. | After casting transformed softmax-camera Q/K/V back to BF16, targeted one-step block timings dropped from roughly 106-107 ms to roughly 49-52 ms for those blocks. This change is kept. | +| Forcing cuDNN SDPA would make the default comparison faster. | Wrong. | The forced-cuDNN probe was much slower and used about 80 GiB. Keep `FORCE_CUDNN_SDPA=1` only as a backend-isolation probe. | +| `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` helps this path. | Unproven. | No run showed a speed or memory benefit for SANA-WM. The harness and docs do not set allocator overrides. | +| Direct SDPA on head dim 112 is faster than padding to 128. | Not supported by clean evidence. | The direct-D112 probe had suspicious backend and memory behavior. Current code keeps padding to 128. | +| `torch.inference_mode()` could reduce framework overhead versus `no_grad`. | Untested. | Upstream wraps generation with `@torch.inference_mode()`. FlashDreams framework generation uses `@torch.no_grad()`. A harness-only test was started, then reverted before measurement. | + +## Next Probes + +1. Reproduce the 60-step BF16 no-refiner benchmark once before changing code, + using `bench.sh`, to confirm the current gap on the current branch. +2. Profile Stage-1 DiT without using profiler-perturbed timings as headline + numbers. Focus on block-level attention, cross-attention, FFN, camera QKV, + and Python/framework overhead. +3. Test `torch.inference_mode()` around the FlashDreams generation path and + compare against the same command with `no_grad`. +4. Split FlashDreams decode timing into VAE forward, clamp, permute, CPU copy, + numpy conversion, and `empty_cache`. +5. Test disabling or configuring `torch.cuda.empty_cache()` after decode. The + optimization goal allows spending more memory for speed. + +## Local Provenance + +These paths are under ignored `outputs/` directories and may not exist in a +fresh checkout: + +- `outputs/bf16_step60_bidirectional_default/bench.md` +- `outputs/bf16_step8_bidirectional_default/bench.md` +- `outputs/bf16_step60_softmax_bf16/bench.md` +- `outputs/bench_121_noalloc/bench.md` +- `outputs/fd_profile_step1_softmax_bf16/stats.json` diff --git a/integrations/sana/tests/parity_check/README.md b/integrations/sana/tests/parity_check/README.md new file mode 100644 index 000000000..d0119a828 --- /dev/null +++ b/integrations/sana/tests/parity_check/README.md @@ -0,0 +1,97 @@ + + +# SANA-WM parity and benchmark harness + +This harness compares upstream +[`NVlabs/Sana`](https://github.com/NVlabs/Sana) SANA-WM against the in-tree +FlashDreams `sana-wm-bidirectional` integration on the same demo image, prompt, +camera trajectory, intrinsics, seed, resolution, and precision settings. + +All dependencies live in this directory's isolated `./.venv`. `run.sh` reuses or +clones the upstream checkout, pins it to `6298508`, and applies `changes.patch` +idempotently. The patch is instrumentation-only: it adds upstream CLI flags for +frame dumps, stats JSON, precision selection, cuDNN SDPA selection, and Stage-1 +`torch.compile` without changing the generation algorithm. A small +`compat/mmcv` shim covers the registry/logging imports needed by the inference +path. + +## Run parity + +```bash +cd integrations/sana/tests/parity_check +bash run.sh +``` + +Defaults: + +- `SANA_REPO=$HOME/dev/Sana` +- upstream pin `6298508` +- demo input `asset/sana_wm/demo_0.*` +- `NUM_FRAMES=121` +- `SEED=42` +- `NO_REFINER=1` to isolate Stage-1 plus SANA VAE decode +- `STAGE1_PRECISION=bf16` +- `REFINER_PRECISION=$STAGE1_PRECISION` +- `QUANT_BACKEND=auto` +- `FORCE_CUDNN_SDPA=0` + +Outputs are written under `outputs/parity/`: + +- `upstream/frames.npy` +- `upstream/stats.json` +- `flashdreams/frames.npy` +- `flashdreams/stats.json` +- `parity.json` + +Set `NO_REFINER=0` to compare the full Stage-1 + LTX-2 refiner path. Set +`COMPILE_STAGE1=1` to wrap each Stage-1 DiT with `torch.compile`; this is opt-in +because the pinned upstream SANA-WM stack can fail during TorchInductor Triton +compilation on current PyTorch/Triton builds. + +## Run benchmark + +```bash +cd integrations/sana/tests/parity_check +bash bench.sh +``` + +Benchmark defaults discard one warmup run and measure three additional runs: +`WARMUP_RUNS=1 MEASURED_RUNS=3 COMPILE_STAGE1=0 NO_REFINER=1 FORCE_CUDNN_SDPA=0`. +Outputs are under `outputs/bench/`: + +- `bench.json` - machine-readable inputs, medians, p90s, memory, and stage + timings. +- `bench.md` - human-readable report. +- `perf.md` - chart-ready model-card data using the same benchmark metric as + `bench.md`. + +The benchmark metric is post-load generation latency per generated frame. Model +construction, checkpoint loading, video writing, and frame dumps are outside the +timing boundary. With the default `NO_REFINER=1`, the timed work covers +conditioning, Stage-1 DiT, and SANA VAE decode. Set `NO_REFINER=0` to benchmark +the full Stage-1 + LTX-2 refiner path with the same timing boundary. + +`bench.md` also reports Stage-1 DiT, conditioning/encode, VAE decode, optional +refiner, and memory breakdowns. Those rows explain the benchmark result; they +are not a second benchmark metric. + +Set `FORCE_CUDNN_SDPA=1` only for backend-isolation probes. It is not the +default SANA-WM benchmark setting. + +Set `DEVICE_LABEL` when generating chart data for docs: + +```bash +DEVICE_LABEL="RTX PRO 6000 Blackwell" bash bench.sh +``` + +Run the precision sweep used by the SANA-WM model card with: + +```bash +DEVICE_LABEL="RTX PRO 6000 Blackwell" BENCH_PRECISIONS=bf16,fp8,fp4 bash bench.sh +``` + +The scripts do not set allocator overrides or GPU wait loops. If GPU contention +matters in your environment, handle it outside the harness. diff --git a/integrations/sana/tests/parity_check/bench.sh b/integrations/sana/tests/parity_check/bench.sh new file mode 100644 index 000000000..2769d7e59 --- /dev/null +++ b/integrations/sana/tests/parity_check/bench.sh @@ -0,0 +1,263 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# Stack-matched SANA-WM upstream vs FlashDreams benchmark. The chart metric is +# post-load generation latency per generated frame. Defaults to Stage-1-only so +# conditioning, Stage-1 DiT, and SANA VAE decode are measured without the LTX-2 +# refiner. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +_abspath() { + case "$1" in + /*) printf '%s\n' "$1" ;; + *) printf '%s/%s\n' "${PWD}" "$1" ;; + esac +} + +SANA_REPO="$(_abspath "${SANA_REPO:-${HOME}/dev/Sana}")" +PATCH_FILE="${SCRIPT_DIR}/changes.patch" +REPO_URL="https://github.com/NVlabs/Sana.git" +PIN_COMMIT="6298508" + +OUTPUT_DIR="$(_abspath "${OUTPUT_DIR:-${SCRIPT_DIR}/outputs/bench}")" +UPSTREAM_ROOT="${OUTPUT_DIR}/upstream" +NATIVE_ROOT="${OUTPUT_DIR}/flashdreams" +IMAGE_PATH="$(_abspath "${IMAGE_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.png}")" +PROMPT_PATH="$(_abspath "${PROMPT_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.txt}")" +CAMERA_PATH="$(_abspath "${CAMERA_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_pose.npy}")" +INTRINSICS_PATH="$(_abspath "${INTRINSICS_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_intrinsics.npy}")" +NUM_FRAMES="${NUM_FRAMES:-121}" +FPS="${FPS:-16}" +STEP="${STEP:-60}" +CFG_SCALE="${CFG_SCALE:-5.0}" +SEED="${SEED:-42}" +NO_REFINER="${NO_REFINER:-1}" +FORCE_CUDNN_SDPA="${FORCE_CUDNN_SDPA:-0}" +COMPILE_STAGE1="${COMPILE_STAGE1:-0}" +WARMUP_RUNS="${WARMUP_RUNS:-1}" +MEASURED_RUNS="${MEASURED_RUNS:-3}" +DEVICE_LABEL="${DEVICE_LABEL:-GPU}" +STAGE1_PRECISION="${STAGE1_PRECISION:-bf16}" +REFINER_PRECISION="${REFINER_PRECISION:-${STAGE1_PRECISION}}" +QUANT_BACKEND="${QUANT_BACKEND:-auto}" +CHART_LABEL="${CHART_LABEL:-${DEVICE_LABEL}}" +BENCH_PRECISIONS="${BENCH_PRECISIONS:-}" + +_is_true() { + case "${1,,}" in + 1|true|yes|on) return 0 ;; + *) return 1 ;; + esac +} + +if [[ -n "${BENCH_PRECISIONS}" ]]; then + mkdir -p "${OUTPUT_DIR}" + SWEEP_ITEMS=() + IFS=',' read -r -a PRECISION_LIST <<< "${BENCH_PRECISIONS}" + for PRECISION_RAW in "${PRECISION_LIST[@]}"; do + PRECISION="${PRECISION_RAW//[[:space:]]/}" + if [[ -z "${PRECISION}" ]]; then + continue + fi + case "${PRECISION}" in + bf16|fp8|fp4) ;; + *) + echo "[bench] ERROR: unsupported BENCH_PRECISIONS entry: ${PRECISION}" >&2 + exit 1 + ;; + esac + PRECISION_LABEL="${PRECISION^^}" + PRECISION_OUTPUT_DIR="${OUTPUT_DIR}/${PRECISION}" + echo "[bench] precision sweep row ${PRECISION_LABEL} -> ${PRECISION_OUTPUT_DIR}" + BENCH_PRECISIONS="" \ + OUTPUT_DIR="${PRECISION_OUTPUT_DIR}" \ + STAGE1_PRECISION="${PRECISION}" \ + REFINER_PRECISION="${PRECISION}" \ + QUANT_BACKEND="${QUANT_BACKEND}" \ + CHART_LABEL="${PRECISION_LABEL}" \ + bash "${SCRIPT_DIR}/bench.sh" + SWEEP_ITEMS+=(--item "${PRECISION_LABEL}:${PRECISION_OUTPUT_DIR}/bench.json") + done + if [[ "${#SWEEP_ITEMS[@]}" -eq 0 ]]; then + echo "[bench] ERROR: BENCH_PRECISIONS produced no precision rows." >&2 + exit 1 + fi + echo "[bench] aggregating precision sweep -> ${OUTPUT_DIR}/perf.md" + ( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/bench_sweep_summary.py" \ + "${SWEEP_ITEMS[@]}" \ + --output-json "${OUTPUT_DIR}/bench.json" \ + --output-md "${OUTPUT_DIR}/bench.md" \ + --output-chart-md "${OUTPUT_DIR}/perf.md" ) + echo "[bench] done." + echo " summary: ${OUTPUT_DIR}/bench.md" + echo " chart data: ${OUTPUT_DIR}/perf.md" + exit 0 +fi + +if [[ ! -d "${SANA_REPO}/.git" ]]; then + echo "[setup] cloning ${REPO_URL} -> ${SANA_REPO}" + git clone "${REPO_URL}" "${SANA_REPO}" +else + echo "[setup] repo already present at ${SANA_REPO}, skipping clone" +fi + +cd "${SANA_REPO}" +CURRENT_COMMIT="$(git rev-parse --short HEAD)" +if [[ "${CURRENT_COMMIT}" != "${PIN_COMMIT}" ]]; then + echo "[setup] checking out pinned commit ${PIN_COMMIT}" + git checkout "${PIN_COMMIT}" +else + echo "[setup] already at pinned commit ${PIN_COMMIT}, skipping checkout" +fi + +if git apply --reverse --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] patch already applied, skipping" +elif git apply --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] applying ${PATCH_FILE}" + git apply "${PATCH_FILE}" +else + echo "[setup] ERROR: ${PATCH_FILE} neither cleanly applies nor is already applied." >&2 + exit 1 +fi + +echo "[setup] ensuring Python deps via uv sync (isolated venv)" +( cd "${SCRIPT_DIR}" && uv sync ) +UPSTREAM_PYTHONPATH="${SCRIPT_DIR}/compat:${SANA_REPO}" +if [[ -n "${PYTHONPATH:-}" ]]; then + UPSTREAM_PYTHONPATH="${UPSTREAM_PYTHONPATH}:${PYTHONPATH}" +fi + +mkdir -p "${UPSTREAM_ROOT}" "${NATIVE_ROOT}" + +UPSTREAM_REFINER_ARGS=() +NATIVE_REFINER_ARGS=() +if _is_true "${NO_REFINER}"; then + UPSTREAM_REFINER_ARGS+=(--no_refiner) + NATIVE_REFINER_ARGS+=(--no-refiner) +fi +UPSTREAM_BACKEND_ARGS=() +NATIVE_BACKEND_ARGS=() +if _is_true "${FORCE_CUDNN_SDPA}"; then + UPSTREAM_BACKEND_ARGS+=(--force_cudnn_sdpa) + NATIVE_BACKEND_ARGS+=(--force-cudnn-sdpa) +fi +UPSTREAM_COMPILE_ARGS=() +NATIVE_COMPILE_ARGS=() +if _is_true "${COMPILE_STAGE1}"; then + UPSTREAM_COMPILE_ARGS+=(--compile_stage1) + NATIVE_COMPILE_ARGS+=(--compile-stage1) +fi +UPSTREAM_PRECISION_ARGS=( + --stage1_precision "${STAGE1_PRECISION}" + --refiner_precision "${REFINER_PRECISION}" +) +NATIVE_PRECISION_ARGS=( + --stage1-precision "${STAGE1_PRECISION}" + --refiner-precision "${REFINER_PRECISION}" + --quant-backend "${QUANT_BACKEND}" +) + +TOTAL_RUNS=$(( WARMUP_RUNS + MEASURED_RUNS )) +for ((i = 0; i < TOTAL_RUNS; i++)); do + UPSTREAM_OUT="${UPSTREAM_ROOT}/run_${i}" + NATIVE_OUT="${NATIVE_ROOT}/run_${i}" + mkdir -p "${UPSTREAM_OUT}" "${NATIVE_OUT}" + + echo "[bench] upstream run ${i}/${TOTAL_RUNS}" + ( cd "${SCRIPT_DIR}" && \ + PYTHONPATH="${UPSTREAM_PYTHONPATH}" \ + uv run python "${SANA_REPO}/inference_video_scripts/wm/inference_sana_wm.py" \ + --image "${IMAGE_PATH}" \ + --prompt "${PROMPT_PATH}" \ + --camera "${CAMERA_PATH}" \ + --intrinsics "${INTRINSICS_PATH}" \ + --output_dir "${UPSTREAM_OUT}" \ + --name upstream \ + --num_frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg_scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --no_action_overlay \ + --stats_json "${UPSTREAM_OUT}/stats.json" \ + "${UPSTREAM_PRECISION_ARGS[@]}" \ + "${UPSTREAM_BACKEND_ARGS[@]}" \ + "${UPSTREAM_COMPILE_ARGS[@]}" \ + "${UPSTREAM_REFINER_ARGS[@]}" ) + + echo "[bench] FlashDreams run ${i}/${TOTAL_RUNS}" + ( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/run_native.py" \ + --image-path "${IMAGE_PATH}" \ + --prompt-path "${PROMPT_PATH}" \ + --camera-path "${CAMERA_PATH}" \ + --intrinsics-path "${INTRINSICS_PATH}" \ + --output-dir "${NATIVE_OUT}" \ + --name flashdreams \ + --num-frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg-scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --stats-json "${NATIVE_OUT}/stats.json" \ + "${NATIVE_PRECISION_ARGS[@]}" \ + "${NATIVE_BACKEND_ARGS[@]}" \ + "${NATIVE_COMPILE_ARGS[@]}" \ + "${NATIVE_REFINER_ARGS[@]}" ) +done + +SUMMARY_JSON="${OUTPUT_DIR}/bench.json" +SUMMARY_MD="${OUTPUT_DIR}/bench.md" +SUMMARY_CHART_MD="${OUTPUT_DIR}/perf.md" +SUMMARY_FLAGS=() +if _is_true "${NO_REFINER}"; then + SUMMARY_FLAGS+=(--no-refiner) +fi +if _is_true "${COMPILE_STAGE1}"; then + SUMMARY_FLAGS+=(--compile-stage1) +fi +if _is_true "${FORCE_CUDNN_SDPA}"; then + SUMMARY_FLAGS+=(--force-cudnn-sdpa) +fi +echo "[bench] summarising -> ${SUMMARY_MD}" +( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/bench_summary.py" \ + --upstream-dir "${UPSTREAM_ROOT}" \ + --flashdreams-dir "${NATIVE_ROOT}" \ + --warmup-runs "${WARMUP_RUNS}" \ + --image-path "${IMAGE_PATH}" \ + --prompt-path "${PROMPT_PATH}" \ + --camera-path "${CAMERA_PATH}" \ + --intrinsics-path "${INTRINSICS_PATH}" \ + --num-frames "${NUM_FRAMES}" \ + --seed "${SEED}" \ + --device-label "${DEVICE_LABEL}" \ + --chart-label "${CHART_LABEL}" \ + --stage1-precision "${STAGE1_PRECISION}" \ + --refiner-precision "${REFINER_PRECISION}" \ + --quant-backend "${QUANT_BACKEND}" \ + "${SUMMARY_FLAGS[@]}" \ + --output-json "${SUMMARY_JSON}" \ + --output-md "${SUMMARY_MD}" \ + --output-chart-md "${SUMMARY_CHART_MD}" ) + +echo "[bench] done." +echo " summary: ${SUMMARY_MD}" +echo " chart data: ${SUMMARY_CHART_MD}" diff --git a/integrations/sana/tests/parity_check/bench_summary.py b/integrations/sana/tests/parity_check/bench_summary.py new file mode 100644 index 000000000..c0d8bbf22 --- /dev/null +++ b/integrations/sana/tests/parity_check/bench_summary.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Summarize SANA-WM upstream and FlashDreams benchmark stats.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path +from typing import Any + + +def _load_stats(root: Path) -> list[dict[str, Any]]: + paths = sorted(root.glob("run_*/stats.json")) + if not paths and (root / "stats.json").exists(): + paths = [root / "stats.json"] + return [json.loads(path.read_text(encoding="utf-8")) | {"_path": str(path)} for path in paths] + + +def _median(values: list[float]) -> float | None: + return statistics.median(values) if values else None + + +def _p90(values: list[float]) -> float | None: + if not values: + return None + ordered = sorted(values) + index = max(0, min(len(ordered) - 1, round(0.9 * (len(ordered) - 1)))) + return ordered[index] + + +def _upstream_stage_ms(item: dict[str, Any], key: str) -> float | None: + timings = item.get("timings_s") + if not isinstance(timings, dict): + return None + for candidate in key.split("|"): + value = timings.get(candidate) + if isinstance(value, (int, float)): + return float(value) * 1000.0 + return None + + +def _native_stage_ms(item: dict[str, Any], key: str) -> float | None: + stats = item.get("stats_ms") + if not isinstance(stats, dict): + return None + value = stats.get(key) + return float(value) if isinstance(value, (int, float)) else None + + +def _upstream_mem_gib(item: dict[str, Any]) -> float | None: + value = item.get("mem_peak_gib") + return float(value) if isinstance(value, (int, float)) else None + + +def _native_mem_gib(item: dict[str, Any]) -> float | None: + stats = item.get("stats_ms") + if not isinstance(stats, dict): + return None + value = stats.get("mem_peak_gib") + return float(value) if isinstance(value, (int, float)) else None + + +def _collect( + items: list[dict[str, Any]], + warmup_runs: int, + stage_reader, + mem_reader, + stages: dict[str, str], +) -> dict[str, Any]: + kept = items[warmup_runs:] + rows: dict[str, Any] = { + "runs_total": len(items), + "warmup_runs": warmup_runs, + "runs_measured": len(kept), + "paths": [item.get("_path") for item in kept], + } + wall = [ + float(item["wall_s"]) + for item in kept + if isinstance(item.get("wall_s"), (int, float)) + ] + rows["wall_median_s"] = _median(wall) + rows["wall_p90_s"] = _p90(wall) + memory = [ + value + for item in kept + if (value := mem_reader(item)) is not None + ] + rows["mem_peak_median_gib"] = _median(memory) + rows["mem_peak_p90_gib"] = _p90(memory) + for label, key in stages.items(): + values = [ + value + for item in kept + if (value := stage_reader(item, key)) is not None + ] + rows[f"{label}_median_ms"] = _median(values) + rows[f"{label}_p90_ms"] = _p90(values) + return rows + + +def _fmt(value: Any, suffix: str = "") -> str: + if value is None: + return "n/a" + if isinstance(value, float): + return f"{value:.2f}{suffix}" + return f"{value}{suffix}" + + +def _ms_per_frame(wall_s: float | None, num_frames: int) -> float | None: + if wall_s is None or num_frames <= 0: + return None + return wall_s * 1000.0 / num_frames + + +def _metric_value(summary: dict[str, Any], side: str, key: str) -> float | None: + value = summary[side].get(key) + return float(value) if isinstance(value, (int, float)) else None + + +def _generation_ms_per_frame( + summary: dict[str, Any], + side: str, + *, + percentile: str = "median", +) -> float | None: + num_frames = int(summary["inputs"]["num_frames"]) + return _ms_per_frame( + _metric_value(summary, side, f"wall_{percentile}_s"), + num_frames, + ) + + +def _sum_optional(*values: float | None) -> float | None: + if any(value is None for value in values): + return None + return sum(float(value) for value in values) + + +def _render_markdown(summary: dict[str, Any]) -> str: + upstream = summary["upstream"] + native = summary["flashdreams"] + rows = [ + "# SANA-WM parity harness benchmark", + "", + "## Inputs", + "", + f"- image: `{summary['inputs']['image_path']}`", + f"- prompt: `{summary['inputs']['prompt_path']}`", + f"- camera: `{summary['inputs']['camera_path']}`", + f"- intrinsics: `{summary['inputs']['intrinsics_path']}`", + f"- num_frames: `{summary['inputs']['num_frames']}`", + f"- seed: `{summary['inputs']['seed']}`", + f"- no_refiner: `{summary['inputs']['no_refiner']}`", + f"- stage1_precision: `{summary['inputs']['stage1_precision']}`", + f"- refiner_precision: `{summary['inputs']['refiner_precision']}`", + f"- quant_backend: `{summary['inputs']['quant_backend']}`", + f"- compile_stage1: `{summary['inputs']['compile_stage1']}`", + f"- force_cudnn_sdpa: `{summary['inputs']['force_cudnn_sdpa']}`", + f"- warmup runs discarded: `{summary['inputs']['warmup_runs']}`", + "", + "## Benchmark metric", + "", + "The chart metric is post-load generation latency per generated frame.", + "Model construction, checkpoint loading, video writing, and frame dumps are outside this timing boundary.", + "With the default `NO_REFINER=1`, the timed work covers conditioning, Stage-1 DiT, and SANA VAE decode.", + "", + "| metric | upstream | FlashDreams |", + "| --- | ---: | ---: |", + f"| measured runs | {upstream['runs_measured']} | {native['runs_measured']} |", + f"| generation median / frame | {_fmt(_generation_ms_per_frame(summary, 'upstream'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams'), ' ms')} |", + f"| generation p90 / frame | {_fmt(_generation_ms_per_frame(summary, 'upstream', percentile='p90'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams', percentile='p90'), ' ms')} |", + f"| wall median | {_fmt(upstream['wall_median_s'], ' s')} | {_fmt(native['wall_median_s'], ' s')} |", + f"| wall p90 | {_fmt(upstream['wall_p90_s'], ' s')} | {_fmt(native['wall_p90_s'], ' s')} |", + "", + "## Timing breakdown", + "", + "| stage | upstream median | FlashDreams median |", + "| --- | ---: | ---: |", + f"| Stage-1 incl. conditioning median | {_fmt(upstream['stage1_total_median_ms'], ' ms')} | {_fmt(_sum_optional(native['encode_median_ms'], native['dit_median_ms']), ' ms')} |", + f"| Stage-1 DiT median | {_fmt(upstream['dit_median_ms'], ' ms')} | {_fmt(native['dit_median_ms'], ' ms')} |", + f"| conditioning/encode median | n/a | {_fmt(native['encode_median_ms'], ' ms')} |", + f"| VAE decode median | {_fmt(upstream['vae_decode_median_ms'], ' ms')} | {_fmt(native['vae_decode_median_ms'], ' ms')} |", + f"| peak memory median | {_fmt(upstream['mem_peak_median_gib'], ' GiB')} | {_fmt(native['mem_peak_median_gib'], ' GiB')} |", + ] + if upstream.get("refiner_median_ms") is not None: + rows.extend( + [ + f"| refiner median | {_fmt(upstream['refiner_median_ms'], ' ms')} | n/a |", + ] + ) + rows.append("") + return "\n".join(rows) + + +def _render_chart_markdown(summary: dict[str, Any], device_label: str) -> str: + official = _generation_ms_per_frame(summary, "upstream") + flashdreams = _generation_ms_per_frame(summary, "flashdreams") + if official is None or flashdreams is None: + raise ValueError("cannot render chart data without wall-clock stats") + return "\n".join( + [ + "# SANA-WM Benchmark Data (ms/frame)", + "", + "| device | official | flashdreams |", + "| --- | ---: | ---: |", + f"| {device_label} | {official:.2f} | {flashdreams:.2f} |", + "", + ] + ) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream-dir", type=Path, required=True) + parser.add_argument("--flashdreams-dir", type=Path, required=True) + parser.add_argument("--warmup-runs", type=int, default=1) + parser.add_argument("--image-path", type=Path, required=True) + parser.add_argument("--prompt-path", type=Path, required=True) + parser.add_argument("--camera-path", type=Path, required=True) + parser.add_argument("--intrinsics-path", type=Path, required=True) + parser.add_argument("--num-frames", type=int, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--no-refiner", action="store_true") + parser.add_argument("--stage1-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument("--refiner-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument( + "--quant-backend", + choices=["auto", "torch", "torch-fp8", "torch-fp4"], + default="auto", + ) + parser.add_argument("--compile-stage1", action="store_true") + parser.add_argument("--force-cudnn-sdpa", action="store_true") + parser.add_argument("--device-label", default="GPU") + parser.add_argument("--chart-label", default=None) + parser.add_argument("--output-json", type=Path, required=True) + parser.add_argument("--output-md", type=Path, required=True) + parser.add_argument("--output-chart-md", type=Path, default=None) + args = parser.parse_args(argv) + + upstream_items = _load_stats(args.upstream_dir) + native_items = _load_stats(args.flashdreams_dir) + summary = { + "inputs": { + "image_path": str(args.image_path), + "prompt_path": str(args.prompt_path), + "camera_path": str(args.camera_path), + "intrinsics_path": str(args.intrinsics_path), + "num_frames": args.num_frames, + "seed": args.seed, + "no_refiner": args.no_refiner, + "stage1_precision": args.stage1_precision, + "refiner_precision": args.refiner_precision, + "quant_backend": args.quant_backend, + "compile_stage1": args.compile_stage1, + "force_cudnn_sdpa": args.force_cudnn_sdpa, + "warmup_runs": args.warmup_runs, + "device_label": args.device_label, + "chart_label": args.chart_label or args.device_label, + }, + "upstream": _collect( + upstream_items, + args.warmup_runs, + _upstream_stage_ms, + _upstream_mem_gib, + { + "dit": "stage1_dit_s|stage1_sample_s", + "stage1_total": "stage1_sample_s", + "refiner": "refiner_s", + "vae_decode": "vae_decode_s", + }, + ), + "flashdreams": _collect( + native_items, + args.warmup_runs, + _native_stage_ms, + _native_mem_gib, + { + "encode": "encode_ms", + "dit": "diffuse_ms", + "vae_decode": "decode_ms", + }, + ), + } + summary["benchmark"] = { + "metric": "generation_ms_per_frame", + "unit": "ms/frame", + "timing_boundary": ( + "pipeline.generate after model setup; excludes model construction, " + "checkpoint loading, video writing, and frame dumps" + ), + "device_label": args.device_label, + "chart_label": args.chart_label or args.device_label, + "official": _generation_ms_per_frame(summary, "upstream"), + "flashdreams": _generation_ms_per_frame(summary, "flashdreams"), + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + args.output_md.parent.mkdir(parents=True, exist_ok=True) + report = _render_markdown(summary) + args.output_md.write_text(report, encoding="utf-8") + if args.output_chart_md is not None: + args.output_chart_md.parent.mkdir(parents=True, exist_ok=True) + args.output_chart_md.write_text( + _render_chart_markdown(summary, args.chart_label or args.device_label), + encoding="utf-8", + ) + print(report) + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/bench_sweep_summary.py b/integrations/sana/tests/parity_check/bench_sweep_summary.py new file mode 100644 index 000000000..c2343325f --- /dev/null +++ b/integrations/sana/tests/parity_check/bench_sweep_summary.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Aggregate SANA-WM precision benchmark summaries into chart data.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def _load_item(raw: str) -> dict[str, Any]: + if ":" not in raw: + raise argparse.ArgumentTypeError( + f"expected LABEL:/path/to/bench.json, got {raw!r}" + ) + label, path_raw = raw.split(":", 1) + label = label.strip() + path = Path(path_raw) + if not label: + raise argparse.ArgumentTypeError(f"empty label in {raw!r}") + if not path.exists(): + raise argparse.ArgumentTypeError(f"bench summary does not exist: {path}") + summary = json.loads(path.read_text(encoding="utf-8")) + benchmark = summary.get("benchmark") + if not isinstance(benchmark, dict): + raise argparse.ArgumentTypeError(f"{path} has no benchmark object") + official = benchmark.get("official") + flashdreams = benchmark.get("flashdreams") + if not isinstance(official, (int, float)) or not isinstance( + flashdreams, (int, float) + ): + raise argparse.ArgumentTypeError( + f"{path} benchmark must contain numeric official and flashdreams values" + ) + return { + "label": label, + "path": str(path), + "official": float(official), + "flashdreams": float(flashdreams), + "inputs": summary.get("inputs", {}), + } + + +def _render_chart(rows: list[dict[str, Any]]) -> str: + lines = [ + "# SANA-WM Benchmark Data (ms/frame)", + "", + "| precision | official | flashdreams |", + "| --- | ---: | ---: |", + ] + for row in rows: + lines.append( + f"| {row['label']} | {row['official']:.2f} | {row['flashdreams']:.2f} |" + ) + lines.append("") + return "\n".join(lines) + + +def _render_report(rows: list[dict[str, Any]]) -> str: + lines = [ + "# SANA-WM precision benchmark sweep", + "", + "The chart metric is post-load generation latency per generated frame.", + "", + "| precision | official | FlashDreams | source |", + "| --- | ---: | ---: | --- |", + ] + for row in rows: + lines.append( + f"| {row['label']} | {row['official']:.2f} ms/frame | " + f"{row['flashdreams']:.2f} ms/frame | `{row['path']}` |" + ) + lines.append("") + return "\n".join(lines) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--item", + action="append", + type=_load_item, + required=True, + help="Precision row as LABEL:/path/to/bench.json. Repeat once per row.", + ) + parser.add_argument("--output-json", type=Path, required=True) + parser.add_argument("--output-md", type=Path, required=True) + parser.add_argument("--output-chart-md", type=Path, required=True) + args = parser.parse_args(argv) + + payload = { + "benchmark": { + "metric": "generation_ms_per_frame", + "unit": "ms/frame", + "timing_boundary": ( + "pipeline.generate after model setup; excludes model construction, " + "checkpoint loading, video writing, and frame dumps" + ), + }, + "rows": args.item, + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + args.output_md.parent.mkdir(parents=True, exist_ok=True) + report = _render_report(args.item) + args.output_md.write_text(report, encoding="utf-8") + args.output_chart_md.parent.mkdir(parents=True, exist_ok=True) + args.output_chart_md.write_text(_render_chart(args.item), encoding="utf-8") + print(report) + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/changes.patch b/integrations/sana/tests/parity_check/changes.patch new file mode 100644 index 000000000..10f349d4c --- /dev/null +++ b/integrations/sana/tests/parity_check/changes.patch @@ -0,0 +1,215 @@ +diff --git a/inference_video_scripts/wm/inference_sana_wm.py b/inference_video_scripts/wm/inference_sana_wm.py +index 5e07197..17db817 100644 +--- a/inference_video_scripts/wm/inference_sana_wm.py ++++ b/inference_video_scripts/wm/inference_sana_wm.py +@@ -937,6 +937,101 @@ def write_video(output_dir: Path, name: str, video_hwc: np.ndarray, fps: int, lo + return video_path + + ++def write_json(path: Path, payload: dict[str, object]) -> None: ++ path.parent.mkdir(parents=True, exist_ok=True) ++ path.write_text(json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8") ++ ++ ++def apply_backend_defaults() -> None: ++ torch.backends.cudnn.benchmark = True ++ torch.set_float32_matmul_precision("high") ++ if torch.cuda.is_available(): ++ torch.backends.cuda.enable_flash_sdp(False) ++ torch.backends.cuda.enable_math_sdp(True) ++ torch.backends.cuda.enable_mem_efficient_sdp(False) ++ torch.backends.cuda.enable_cudnn_sdp(True) ++ try: ++ import torch._inductor.config as inductor_config ++ ++ inductor_config.coordinate_descent_tuning = True ++ inductor_config.epilogue_fusion = True ++ except Exception: ++ pass ++ ++ ++def apply_precision_args(args: argparse.Namespace, logger: logging.Logger) -> None: ++ stage1_precision = args.stage1_precision ++ refiner_precision = "bf16" if args.no_refiner else args.refiner_precision ++ ++ if stage1_precision == "bf16": ++ os.environ.pop("SANA_WM_STAGE1_NVFP4", None) ++ os.environ.pop("SANA_WM_STAGE1_NVFP4_MODE", None) ++ os.environ.pop("SANA_WM_STAGE1_NVFP4_SCOPE", None) ++ os.environ.pop("SANA_WM_STAGE1_LINEARIZE_FFN", None) ++ os.environ.pop("SANA_WM_STAGE1_QUANT", None) ++ else: ++ os.environ["SANA_WM_STAGE1_NVFP4"] = "1" ++ os.environ["SANA_WM_STAGE1_NVFP4_MODE"] = "self_attn+cross+ffn" ++ os.environ["SANA_WM_STAGE1_NVFP4_SCOPE"] = "block" ++ os.environ["SANA_WM_STAGE1_LINEARIZE_FFN"] = "1" ++ os.environ["SANA_WM_STAGE1_QUANT"] = ( ++ "fp8block" if stage1_precision == "fp8" else "nvfp4" ++ ) ++ ++ if refiner_precision == "bf16": ++ os.environ.pop("SANA_WM_REFINER_NVFP4", None) ++ os.environ.pop("SANA_WM_REFINER_QUANT", None) ++ else: ++ os.environ["SANA_WM_REFINER_NVFP4"] = "1" ++ os.environ["SANA_WM_REFINER_QUANT"] = ( ++ "fp8block" if refiner_precision == "fp8" else "nvfp4" ++ ) ++ ++ if stage1_precision != "bf16" or refiner_precision != "bf16": ++ need = ( ++ "NVFP4BlockScaling" ++ if "fp4" in (stage1_precision, refiner_precision) ++ else "Float8BlockScaling" ++ ) ++ try: ++ import transformer_engine.common.recipe as te_recipe ++ import transformer_engine.pytorch # noqa: F401 ++ ++ if not hasattr(te_recipe, need): ++ raise ImportError(f"installed Transformer Engine lacks {need}") ++ except Exception as exc: ++ raise SystemExit( ++ "--stage1_precision/--refiner_precision fp8/fp4 require NVIDIA " ++ f"Transformer Engine >= 2.x for {need} ({exc})." ++ ) ++ ++ if "fp4" in (stage1_precision, refiner_precision) and torch.cuda.is_available(): ++ major, minor = torch.cuda.get_device_capability() ++ if major < 10: ++ logger.warning( ++ "fp4 (NVFP4) requires a Blackwell GPU (sm_100+); detected sm_%d%d.", ++ major, ++ minor, ++ ) ++ logger.info("[precision] stage1=%s refiner=%s", stage1_precision, refiner_precision) ++ ++ ++def time_bound_method(obj: object, name: str, timing_key: str, timings: dict[str, float]) -> None: ++ original = getattr(obj, name) ++ ++ def wrapped(*args, **kwargs): ++ if torch.cuda.is_available(): ++ torch.cuda.synchronize() ++ t0 = time.perf_counter() ++ result = original(*args, **kwargs) ++ if torch.cuda.is_available(): ++ torch.cuda.synchronize() ++ timings[timing_key] = time.perf_counter() - t0 ++ return result ++ ++ setattr(obj, name, wrapped) ++ ++ + # ============================================================================ + # Pipeline + # ============================================================================ +@@ -2132,6 +2227,30 @@ def _build_parser() -> argparse.ArgumentParser: + action="store_true", + help="Skip rendering the WASD + joystick overlay on the output video.", + ) ++ p.add_argument("--dump_frames", type=Path, default=None, help="Optional .npy path for decoded uint8 frames.") ++ p.add_argument("--stats_json", type=Path, default=None, help="Optional JSON path for timing and memory stats.") ++ p.add_argument( ++ "--force_cudnn_sdpa", ++ action="store_true", ++ help="Force PyTorch scaled_dot_product_attention to use the cuDNN backend.", ++ ) ++ p.add_argument( ++ "--compile_stage1", ++ action="store_true", ++ help="Wrap the Stage-1 DiT with torch.compile(mode='max-autotune-no-cudagraphs').", ++ ) ++ p.add_argument( ++ "--stage1_precision", ++ choices=["bf16", "fp8", "fp4"], ++ default="bf16", ++ help="Stage-1 DiT precision for benchmark parity.", ++ ) ++ p.add_argument( ++ "--refiner_precision", ++ choices=["bf16", "fp8", "fp4"], ++ default="bf16", ++ help="LTX-2 refiner precision for benchmark parity.", ++ ) + + # Weights and config. + p.add_argument( +@@ -2222,7 +2341,11 @@ def _snap_num_frames(n: int, stride: int = 8, *, upper_bound: int | None = None) + def main() -> None: + args = _build_parser().parse_args() + ++ if args.force_cudnn_sdpa: ++ apply_backend_defaults() ++ + logger = get_root_logger() ++ apply_precision_args(args, logger) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + image = Image.open(args.image).convert("RGB") +@@ -2276,6 +2399,15 @@ def main() -> None: + offload_refiner=args.offload_refiner, + logger=logger, + ) ++ if args.compile_stage1: ++ pipeline.model = torch.compile(pipeline.model, mode="max-autotune-no-cudagraphs") ++ ++ timings: dict[str, float] = {} ++ time_bound_method(pipeline, "_sample_stage1", "stage1_sample_s", timings) ++ time_bound_method(pipeline, "_dispatch_solver", "stage1_dit_s", timings) ++ time_bound_method(pipeline, "_decode_with_sana_vae", "vae_decode_s", timings) ++ if refiner is not None: ++ time_bound_method(pipeline, "_refine", "refiner_s", timings) + + denoising_step_list: list[int] | None = None + if args.denoising_step_list: +@@ -2308,14 +2440,48 @@ def main() -> None: + save_stage1=args.save_stage1, + ) + ++ if torch.cuda.is_available(): ++ torch.cuda.reset_peak_memory_stats(device) ++ torch.cuda.synchronize(device) ++ wall_start = time.perf_counter() + out = pipeline.generate(cropped, prompt, c2w, intrinsics_vec4, params) ++ if torch.cuda.is_available(): ++ torch.cuda.synchronize(device) ++ wall_s = time.perf_counter() - wall_start + video_hwc = out["video"] + + if not args.no_action_overlay: + logger.info("Compositing action overlay onto the output video.") + video_hwc = apply_overlay(video_hwc, out["c2w"]) + +- write_video(args.output_dir, args.name, video_hwc, params.fps, logger) ++ video_path = write_video(args.output_dir, args.name, video_hwc, params.fps, logger) ++ ++ if args.dump_frames is not None: ++ args.dump_frames.parent.mkdir(parents=True, exist_ok=True) ++ np.save(args.dump_frames, np.asarray(video_hwc, dtype=np.uint8)) ++ ++ if args.stats_json is not None: ++ stats: dict[str, object] = { ++ "backend": "upstream", ++ "entrypoint": "inference_video_scripts/wm/inference_sana_wm.py", ++ "video_path": str(video_path), ++ "video_shape": list(np.asarray(video_hwc).shape), ++ "timings_s": timings, ++ "num_frames": num_frames, ++ "seed": args.seed, ++ "fps": params.fps, ++ "step": params.step, ++ "cfg_scale": params.cfg_scale, ++ "no_refiner": args.no_refiner, ++ "stage1_precision": args.stage1_precision, ++ "refiner_precision": args.refiner_precision, ++ "compile_stage1": args.compile_stage1, ++ "force_cudnn_sdpa": args.force_cudnn_sdpa, ++ "wall_s": wall_s, ++ } ++ if torch.cuda.is_available(): ++ stats["mem_peak_gib"] = torch.cuda.max_memory_allocated(device) / (1024**3) ++ write_json(args.stats_json, stats) + + stage1_video = out.get("stage1_video") + if stage1_video is not None: diff --git a/integrations/sana/tests/parity_check/compat/mmcv/__init__.py b/integrations/sana/tests/parity_check/compat/mmcv/__init__.py new file mode 100644 index 000000000..9b2aebaae --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/__init__.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal MMCV compatibility used by the SANA-WM parity harness. + +The upstream inference entrypoint only needs registry construction plus a few +utility functions at import time. This shim avoids installing old mmcv builds in +the isolated parity venv. +""" + +from __future__ import annotations + +import os +import pickle +from pathlib import Path +from typing import Any, Callable + + +class Registry: + """Small subset of ``mmcv.Registry`` used by SANA model registration.""" + + def __init__(self, name: str) -> None: + self.name = name + self.module_dict: dict[str, Any] = {} + + def __contains__(self, key: str) -> bool: + return key in self.module_dict + + def get(self, key: str) -> Any: + return self.module_dict.get(key) + + def register_module( + self, + module: Any | None = None, + *, + name: str | None = None, + force: bool = False, + ) -> Callable[[Any], Any] | Any: + """Register ``module`` or return a class decorator.""" + + def _register(obj: Any) -> Any: + key = name or obj.__name__ + if key in self.module_dict and not force: + raise KeyError(f"{key!r} is already registered in {self.name}.") + self.module_dict[key] = obj + return obj + + if module is not None: + return _register(module) + return _register + + def build( + self, + cfg: dict[str, Any], + *, + default_args: dict[str, Any] | None = None, + ) -> Any: + return build_from_cfg(cfg, self, default_args=default_args) + + +class Config(dict): + """Tiny dict-backed stand-in for import-time ``mmcv.Config`` references.""" + + @classmethod + def fromfile(cls, filename: str | os.PathLike[str]) -> "Config": + raise NotImplementedError( + "mmcv.Config.fromfile is not implemented in the parity harness shim." + ) + + +def build_from_cfg( + cfg: dict[str, Any], + registry: Registry, + *, + default_args: dict[str, Any] | None = None, +) -> Any: + """Instantiate a registered object from a config dict.""" + + if not isinstance(cfg, dict): + raise TypeError(f"cfg must be a dict, got {type(cfg).__name__}.") + if "type" not in cfg: + raise KeyError("cfg must contain key 'type'.") + args = dict(default_args or {}) + args.update({k: v for k, v in cfg.items() if k != "type"}) + obj_type = cfg["type"] + if isinstance(obj_type, str): + obj_cls = registry.get(obj_type) + if obj_cls is None: + raise KeyError(f"{obj_type!r} is not registered in {registry.name}.") + else: + obj_cls = obj_type + return obj_cls(**args) + + +def mkdir_or_exist(path: str | os.PathLike[str]) -> None: + Path(path).mkdir(parents=True, exist_ok=True) + + +def dump(obj: Any, file: str | os.PathLike[str]) -> None: + with open(file, "wb") as handle: + pickle.dump(obj, handle) + + +def load(file: str | os.PathLike[str]) -> Any: + with open(file, "rb") as handle: + return pickle.load(handle) + + +__all__ = [ + "Config", + "Registry", + "build_from_cfg", + "dump", + "load", + "mkdir_or_exist", +] diff --git a/integrations/sana/tests/parity_check/compat/mmcv/runner.py b/integrations/sana/tests/parity_check/compat/mmcv/runner.py new file mode 100644 index 000000000..3e82685a2 --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/runner.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small ``mmcv.runner`` compatibility layer for SANA inference imports.""" + +from __future__ import annotations + +import torch.distributed as dist + + +def get_dist_info() -> tuple[int, int]: + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(), dist.get_world_size() + return 0, 1 + + +class DefaultOptimizerConstructor: + pass + + +OPTIMIZERS = {} +OPTIMIZER_BUILDERS = {} + + +def build_optimizer(*_args, **_kwargs): + raise NotImplementedError("Optimizer construction is not available in the parity harness.") + + +__all__ = [ + "DefaultOptimizerConstructor", + "OPTIMIZER_BUILDERS", + "OPTIMIZERS", + "build_optimizer", + "get_dist_info", +] diff --git a/integrations/sana/tests/parity_check/compat/mmcv/utils/__init__.py b/integrations/sana/tests/parity_check/compat/mmcv/utils/__init__.py new file mode 100644 index 000000000..3acb82c06 --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/utils/__init__.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal ``mmcv.utils`` namespace for SANA inference imports.""" + +from __future__ import annotations + +import torch.nn as nn + +_BatchNorm = nn.modules.batchnorm._BatchNorm +_InstanceNorm = nn.modules.instancenorm._InstanceNorm + +__all__ = ["_BatchNorm", "_InstanceNorm"] diff --git a/integrations/sana/tests/parity_check/compat/mmcv/utils/logging.py b/integrations/sana/tests/parity_check/compat/mmcv/utils/logging.py new file mode 100644 index 000000000..bb2396088 --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/utils/logging.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Logging globals expected by upstream SANA.""" + +logger_initialized: dict[str, bool] = {} + +__all__ = ["logger_initialized"] diff --git a/integrations/sana/tests/parity_check/diff_parity.py b/integrations/sana/tests/parity_check/diff_parity.py new file mode 100644 index 000000000..db9d395e2 --- /dev/null +++ b/integrations/sana/tests/parity_check/diff_parity.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Compute SANA-WM frame parity as mean absolute uint8 delta.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np + + +def _load_frames(path: Path) -> np.ndarray: + frames = np.load(path) + if frames.ndim != 4 or frames.shape[-1] != 3: + raise ValueError(f"{path} must contain [T,H,W,3] frames; got {frames.shape}.") + if frames.dtype != np.uint8: + frames = np.clip(frames, 0, 255).astype(np.uint8) + return frames + + +def _summary(upstream: np.ndarray, flashdreams: np.ndarray) -> dict[str, Any]: + if upstream.shape != flashdreams.shape: + raise ValueError( + "shape mismatch: " + f"upstream={tuple(upstream.shape)} flashdreams={tuple(flashdreams.shape)}" + ) + diff = np.abs(upstream.astype(np.int16) - flashdreams.astype(np.int16)) + per_frame = diff.reshape(diff.shape[0], -1).mean(axis=1) + return { + "shape": list(upstream.shape), + "mean_abs_delta": float(diff.mean()), + "mean_abs_delta_over_255": float(diff.mean() / 255.0), + "max_abs_delta": int(diff.max()), + "per_frame_mean_abs_delta": [float(v) for v in per_frame], + "per_frame_max_abs_delta": [ + int(v) for v in diff.reshape(diff.shape[0], -1).max(axis=1) + ], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream", type=Path, required=True) + parser.add_argument("--flashdreams", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + summary = _summary(_load_frames(args.upstream), _load_frames(args.flashdreams)) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + + print( + "mean |Delta|: " + f"{summary['mean_abs_delta']:.4f} / 255 " + f"({summary['mean_abs_delta_over_255']:.6f})" + ) + print(f"max |Delta|: {summary['max_abs_delta']} / 255") + for idx, value in enumerate(summary["per_frame_mean_abs_delta"]): + print(f"frame {idx:04d}: mean |Delta| {value:.4f} / 255") + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/pyproject.toml b/integrations/sana/tests/parity_check/pyproject.toml new file mode 100644 index 000000000..0f2fbdf09 --- /dev/null +++ b/integrations/sana/tests/parity_check/pyproject.toml @@ -0,0 +1,47 @@ +[project] +name = "sana-wm-parity-check" +version = "0.0.0" +description = "Isolated venv for the SANA-WM upstream parity and benchmark harness." +requires-python = ">=3.12,<3.13" +dependencies = [ + "flashdreams", + "flashdreams-sana-wm", + "accelerate>=1.3", + "diffusers>=0.37", + "einops>=0.7", + "flash-linear-attention>=0.4.2", + "ftfy>=6.0", + "huggingface-hub>=0.36", + "imageio[ffmpeg]>=2.31", + "numpy>=1.24,<2.5", + "omegaconf>=2.3", + "opencv-python-headless>=4.8", + "pillow>=10", + "protobuf>=7.35.0,<8", + "pyrallis>=0.3", + "pytz>=2024.0", + "pyyaml>=6.0", + "qwen-vl-utils>=0.0.8", + "safetensors>=0.5", + "scipy>=1.11", + "sentencepiece>=0.2", + "termcolor>=2.4", + "timm>=0.6.13", + "torch>=2.11", + "torchvision>=0.26", + "tqdm>=4.60", + "transformer-engine[pytorch,core-cu13]>=2.12; sys_platform != 'win32'", + "transformers>=5.0,<6", + "triton>=3.6; sys_platform == 'linux'", +] + +[tool.uv.sources] +flashdreams = { path = "../../../../flashdreams", editable = true } +flashdreams-sana-wm = { path = "../..", editable = true } + +[tool.uv] +managed = true +override-dependencies = [ + "nvidia-cublas>=13.4", +] +no-build-isolation-package = ["transformer-engine-torch"] diff --git a/integrations/sana/tests/parity_check/run.sh b/integrations/sana/tests/parity_check/run.sh new file mode 100644 index 000000000..a5460210d --- /dev/null +++ b/integrations/sana/tests/parity_check/run.sh @@ -0,0 +1,185 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# Patch pinned upstream SANA-WM with instrumentation, run it and FlashDreams on +# the same demo input, dump decoded uint8 frames, and compute mean |Delta| / 255. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" + +_abspath() { + case "$1" in + /*) printf '%s\n' "$1" ;; + *) printf '%s/%s\n' "${PWD}" "$1" ;; + esac +} + +SANA_REPO="$(_abspath "${SANA_REPO:-${HOME}/dev/Sana}")" +PATCH_FILE="${SCRIPT_DIR}/changes.patch" +REPO_URL="https://github.com/NVlabs/Sana.git" +PIN_COMMIT="6298508" + +OUTPUT_DIR="$(_abspath "${OUTPUT_DIR:-${SCRIPT_DIR}/outputs/parity}")" +UPSTREAM_OUT="${OUTPUT_DIR}/upstream" +NATIVE_OUT="${OUTPUT_DIR}/flashdreams" +IMAGE_PATH="$(_abspath "${IMAGE_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.png}")" +PROMPT_PATH="$(_abspath "${PROMPT_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.txt}")" +CAMERA_PATH="$(_abspath "${CAMERA_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_pose.npy}")" +INTRINSICS_PATH="$(_abspath "${INTRINSICS_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_intrinsics.npy}")" +NUM_FRAMES="${NUM_FRAMES:-121}" +FPS="${FPS:-16}" +STEP="${STEP:-60}" +CFG_SCALE="${CFG_SCALE:-5.0}" +SEED="${SEED:-42}" +NO_REFINER="${NO_REFINER:-1}" +FORCE_CUDNN_SDPA="${FORCE_CUDNN_SDPA:-0}" +COMPILE_STAGE1="${COMPILE_STAGE1:-0}" +STAGE1_PRECISION="${STAGE1_PRECISION:-bf16}" +REFINER_PRECISION="${REFINER_PRECISION:-${STAGE1_PRECISION}}" +QUANT_BACKEND="${QUANT_BACKEND:-auto}" + +_is_true() { + case "${1,,}" in + 1|true|yes|on) return 0 ;; + *) return 1 ;; + esac +} + +if [[ ! -d "${SANA_REPO}/.git" ]]; then + echo "[setup] cloning ${REPO_URL} -> ${SANA_REPO}" + git clone "${REPO_URL}" "${SANA_REPO}" +else + echo "[setup] repo already present at ${SANA_REPO}, skipping clone" +fi + +cd "${SANA_REPO}" +CURRENT_COMMIT="$(git rev-parse --short HEAD)" +if [[ "${CURRENT_COMMIT}" != "${PIN_COMMIT}" ]]; then + echo "[setup] checking out pinned commit ${PIN_COMMIT}" + git checkout "${PIN_COMMIT}" +else + echo "[setup] already at pinned commit ${PIN_COMMIT}, skipping checkout" +fi + +if git apply --reverse --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] patch already applied, skipping" +elif git apply --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] applying ${PATCH_FILE}" + git apply "${PATCH_FILE}" +else + echo "[setup] ERROR: ${PATCH_FILE} neither cleanly applies nor is already applied." >&2 + exit 1 +fi + +echo "[setup] ensuring Python deps via uv sync (isolated venv)" +( cd "${SCRIPT_DIR}" && uv sync ) +UPSTREAM_PYTHONPATH="${SCRIPT_DIR}/compat:${SANA_REPO}" +if [[ -n "${PYTHONPATH:-}" ]]; then + UPSTREAM_PYTHONPATH="${UPSTREAM_PYTHONPATH}:${PYTHONPATH}" +fi + +mkdir -p "${UPSTREAM_OUT}" "${NATIVE_OUT}" + +UPSTREAM_FRAMES="${UPSTREAM_OUT}/frames.npy" +NATIVE_FRAMES="${NATIVE_OUT}/frames.npy" +UPSTREAM_STATS="${UPSTREAM_OUT}/stats.json" +NATIVE_STATS="${NATIVE_OUT}/stats.json" + +UPSTREAM_REFINER_ARGS=() +NATIVE_REFINER_ARGS=() +if _is_true "${NO_REFINER}"; then + UPSTREAM_REFINER_ARGS+=(--no_refiner) + NATIVE_REFINER_ARGS+=(--no-refiner) +fi +UPSTREAM_BACKEND_ARGS=() +NATIVE_BACKEND_ARGS=() +if _is_true "${FORCE_CUDNN_SDPA}"; then + UPSTREAM_BACKEND_ARGS+=(--force_cudnn_sdpa) + NATIVE_BACKEND_ARGS+=(--force-cudnn-sdpa) +fi +UPSTREAM_COMPILE_ARGS=() +NATIVE_COMPILE_ARGS=() +if _is_true "${COMPILE_STAGE1}"; then + UPSTREAM_COMPILE_ARGS+=(--compile_stage1) + NATIVE_COMPILE_ARGS+=(--compile-stage1) +fi +UPSTREAM_PRECISION_ARGS=( + --stage1_precision "${STAGE1_PRECISION}" + --refiner_precision "${REFINER_PRECISION}" +) +NATIVE_PRECISION_ARGS=( + --stage1-precision "${STAGE1_PRECISION}" + --refiner-precision "${REFINER_PRECISION}" + --quant-backend "${QUANT_BACKEND}" +) + +echo "[run] upstream SANA-WM -> ${UPSTREAM_OUT}" +( cd "${SCRIPT_DIR}" && \ + PYTHONPATH="${UPSTREAM_PYTHONPATH}" \ + uv run python "${SANA_REPO}/inference_video_scripts/wm/inference_sana_wm.py" \ + --image "${IMAGE_PATH}" \ + --prompt "${PROMPT_PATH}" \ + --camera "${CAMERA_PATH}" \ + --intrinsics "${INTRINSICS_PATH}" \ + --output_dir "${UPSTREAM_OUT}" \ + --name upstream \ + --num_frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg_scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --no_action_overlay \ + --dump_frames "${UPSTREAM_FRAMES}" \ + --stats_json "${UPSTREAM_STATS}" \ + "${UPSTREAM_PRECISION_ARGS[@]}" \ + "${UPSTREAM_BACKEND_ARGS[@]}" \ + "${UPSTREAM_COMPILE_ARGS[@]}" \ + "${UPSTREAM_REFINER_ARGS[@]}" ) + +echo "[run] FlashDreams SANA-WM -> ${NATIVE_OUT}" +( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/run_native.py" \ + --image-path "${IMAGE_PATH}" \ + --prompt-path "${PROMPT_PATH}" \ + --camera-path "${CAMERA_PATH}" \ + --intrinsics-path "${INTRINSICS_PATH}" \ + --output-dir "${NATIVE_OUT}" \ + --name flashdreams \ + --num-frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg-scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --dump-frames "${NATIVE_FRAMES}" \ + --stats-json "${NATIVE_STATS}" \ + "${NATIVE_PRECISION_ARGS[@]}" \ + "${NATIVE_BACKEND_ARGS[@]}" \ + "${NATIVE_COMPILE_ARGS[@]}" \ + "${NATIVE_REFINER_ARGS[@]}" ) + +echo "[diff] summarising parity -> ${OUTPUT_DIR}/parity.json" +( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/diff_parity.py" \ + --upstream "${UPSTREAM_FRAMES}" \ + --flashdreams "${NATIVE_FRAMES}" \ + --output "${OUTPUT_DIR}/parity.json" ) + +echo "[run] done." +echo " upstream frames : ${UPSTREAM_FRAMES}" +echo " flashdreams frames: ${NATIVE_FRAMES}" +echo " parity JSON : ${OUTPUT_DIR}/parity.json" diff --git a/integrations/sana/tests/parity_check/run_native.py b/integrations/sana/tests/parity_check/run_native.py new file mode 100644 index 000000000..c93903fb9 --- /dev/null +++ b/integrations/sana/tests/parity_check/run_native.py @@ -0,0 +1,319 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Run the FlashDreams SANA-WM pipeline with frame dumps and profiler JSON.""" + +from __future__ import annotations + +import argparse +from contextlib import nullcontext +import json +import os +import time +from pathlib import Path +from typing import Any + +import imageio.v3 as iio +import numpy as np +import torch +from PIL import Image + +from flashdreams.infra.config import derive_config +from sana_wm.camera import ( + default_intrinsics_vec4, + load_intrinsics, + resize_center_crop_geometry, + snap_num_frames, + transform_intrinsics_for_crop, +) +from sana_wm.conditioning import SanaWMI2VConditioningRequest +from sana_wm.config import RUNNER_SANA_WM_BIDIRECTIONAL +from sana_wm.constants import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH +from sana_wm.decoder import SanaWMDecodedVideo +from sana_wm.runner import ( + _active_quantized_precisions, + _pipeline_config, + _resolve_quant_backend, + _validate_precision_request, +) + + +def _apply_backend_defaults() -> None: + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if torch.cuda.is_available(): + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_cudnn_sdp(True) + try: + import torch._inductor.config as inductor_config + + inductor_config.coordinate_descent_tuning = True + inductor_config.epilogue_fusion = True + except Exception: + pass + + +def _resolve_device(device: str) -> torch.device: + if device == "auto": + if torch.cuda.is_available(): + return torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") + return torch.device("cpu") + if device == "cuda" and torch.cuda.is_available(): + return torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") + return torch.device(device) + + +def _prepare_inputs( + *, + image_path: Path, + camera_path: Path, + intrinsics_path: Path | None, + intrinsics_hfov_deg: float, + num_frames_requested: int, +) -> tuple[Image.Image, np.ndarray, np.ndarray, int]: + image = Image.open(image_path).convert("RGB") + c2w_full = np.load(camera_path).astype(np.float32) + if c2w_full.ndim != 3 or c2w_full.shape[1:] != (4, 4): + raise ValueError(f"--camera-path must be [F,4,4]; got {c2w_full.shape}.") + + num_frames = min(num_frames_requested, c2w_full.shape[0]) + num_frames = snap_num_frames(num_frames, stride=8, upper_bound=c2w_full.shape[0]) + c2w = c2w_full[:num_frames] + + resized_size, crop_offset = resize_center_crop_geometry( + image.size, + target_h=DEFAULT_VIDEO_HEIGHT, + target_w=DEFAULT_VIDEO_WIDTH, + ) + resized = image.resize(resized_size, Image.LANCZOS) + left, top = crop_offset + cropped = resized.crop( + (left, top, left + DEFAULT_VIDEO_WIDTH, top + DEFAULT_VIDEO_HEIGHT) + ) + if intrinsics_path is None: + intrinsics_src = default_intrinsics_vec4( + image.size, + num_frames, + hfov_deg=intrinsics_hfov_deg, + ) + else: + intrinsics_src = load_intrinsics(intrinsics_path, num_frames) + intrinsics_vec4 = transform_intrinsics_for_crop( + intrinsics_src, + image.size, + resized_size, + crop_offset, + ) + return cropped, c2w, intrinsics_vec4, num_frames + + +def _install_stage1_compile_hook(pipeline: Any) -> None: + transformer = pipeline.diffusion_model.transformer + original_ensure_model = transformer._ensure_model + compiled = {"done": False} + + def _ensure_model_with_compile() -> None: + original_ensure_model() + if compiled["done"]: + return + transformer.model = torch.compile( + transformer.model, + mode="max-autotune-no-cudagraphs", + ) + compiled["done"] = True + + transformer._ensure_model = _ensure_model_with_compile + + +def _write_video(output_dir: Path, name: str, frames: np.ndarray, fps: int) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + path = output_dir / f"{name}_generated.mp4" + iio.imwrite(path, frames, fps=fps) + return path + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--image-path", type=Path, required=True) + parser.add_argument("--prompt-path", type=Path, required=True) + parser.add_argument("--camera-path", type=Path, required=True) + parser.add_argument("--intrinsics-path", type=Path, default=None) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--name", default="flashdreams") + parser.add_argument("--dump-frames", type=Path, default=None) + parser.add_argument("--stats-json", type=Path, default=None) + parser.add_argument("--num-frames", type=int, default=161) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--step", type=int, default=60) + parser.add_argument("--cfg-scale", type=float, default=5.0) + parser.add_argument("--flow-shift", type=float, default=None) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--negative-prompt", default="") + parser.add_argument("--no-refiner", action="store_true") + parser.add_argument("--save-stage1", action="store_true") + parser.add_argument("--device", default="auto") + parser.add_argument("--intrinsics-hfov-deg", type=float, default=90.0) + parser.add_argument("--stage1-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument("--refiner-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument( + "--quant-backend", + choices=["auto", "torch", "torch-fp8", "torch-fp4"], + default="auto", + ) + parser.add_argument("--force-cudnn-sdpa", action="store_true") + parser.add_argument("--compile-stage1", action="store_true") + args = parser.parse_args() + + if args.force_cudnn_sdpa: + _apply_backend_defaults() + + prompt = args.prompt_path.read_text(encoding="utf-8", errors="replace").strip() + if not prompt: + raise ValueError(f"Prompt file is empty: {args.prompt_path}") + + device = _resolve_device(args.device) + quantized = _active_quantized_precisions( + stage1_precision=args.stage1_precision, + refiner_precision=args.refiner_precision, + refiner_enabled=not args.no_refiner, + ) + quant_backend = _resolve_quant_backend(args.quant_backend, quantized) + _validate_precision_request( + device=device, + stage1_precision=args.stage1_precision, + refiner_precision=args.refiner_precision, + refiner_enabled=not args.no_refiner, + quant_backend=args.quant_backend, + ) + with nullcontext(): + image, c2w, intrinsics_vec4, num_frames = _prepare_inputs( + image_path=args.image_path, + camera_path=args.camera_path, + intrinsics_path=args.intrinsics_path, + intrinsics_hfov_deg=args.intrinsics_hfov_deg, + num_frames_requested=args.num_frames, + ) + runner_cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + output_dir=args.output_dir, + runner_name=args.name, + num_frames=num_frames, + fps=args.fps, + step=args.step, + cfg_scale=args.cfg_scale, + flow_shift=args.flow_shift, + seed=args.seed, + negative_prompt=args.negative_prompt, + no_refiner=args.no_refiner, + save_stage1=args.save_stage1, + refiner_seed=args.seed, + stage1_precision=args.stage1_precision, + refiner_precision=args.refiner_precision, + quant_backend=args.quant_backend, + ) + pipeline_cfg = _pipeline_config(runner_cfg, quant_backend=quant_backend) + pipeline_cfg = derive_config(pipeline_cfg, enable_sync_and_profile=True) + pipeline = pipeline_cfg.setup().to(device).eval() + if args.compile_stage1: + _install_stage1_compile_hook(pipeline) + encoder = getattr(pipeline, "encoder", None) + if encoder is not None: + text_encoder = getattr(encoder, "text_encoder", None) + if text_encoder is not None: + text_encoder._ensure_text_encoder() + first_frame_encoder = getattr(encoder, "first_frame_encoder", None) + if first_frame_encoder is not None: + first_frame_encoder._ensure_vae() + pipeline.diffusion_model.transformer._ensure_model() + decoder = getattr(pipeline, "decoder", None) + vae_decoder = getattr(decoder, "vae_decoder", None) + if vae_decoder is not None: + vae_decoder._ensure_vae() + + cache = pipeline.initialize_cache( + decoder_context={ + "prompt": prompt, + "fps": args.fps, + "save_stage1": args.save_stage1, + "refiner_seed": args.seed, + "sink_size": runner_cfg.sink_size, + } + ) + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats(device) + torch.cuda.synchronize(device) + generation_start = time.perf_counter() + decoded = pipeline.generate( + 0, + cache, + input=SanaWMI2VConditioningRequest( + image=image, + prompt=prompt, + poses_c2w=c2w, + intrinsics_vec4=intrinsics_vec4, + num_frames=num_frames, + fps=args.fps, + steps=args.step, + cfg_scale=args.cfg_scale, + flow_shift=args.flow_shift, + seed=args.seed, + negative_prompt=args.negative_prompt, + ), + ) + if torch.cuda.is_available(): + torch.cuda.synchronize(device) + wall_s = time.perf_counter() - generation_start + stats = pipeline.finalize(0, cache) or {} + + if not isinstance(decoded, SanaWMDecodedVideo): + raise TypeError(f"expected SanaWMDecodedVideo, got {type(decoded).__name__}") + + frames = np.asarray(decoded.video_hwc, dtype=np.uint8) + video_path = _write_video(args.output_dir, args.name, frames, args.fps) + if args.dump_frames is not None: + args.dump_frames.parent.mkdir(parents=True, exist_ok=True) + np.save(args.dump_frames, frames) + if args.stats_json is not None: + _write_json( + args.stats_json, + { + "backend": "flashdreams", + "runner": "sana-wm-bidirectional", + "video_path": str(video_path), + "video_shape": list(frames.shape), + "num_frames": num_frames, + "seed": args.seed, + "fps": args.fps, + "step": args.step, + "cfg_scale": args.cfg_scale, + "no_refiner": args.no_refiner, + "compile_stage1": args.compile_stage1, + "force_cudnn_sdpa": args.force_cudnn_sdpa, + "wall_s": wall_s, + "stats_ms": stats, + }, + ) + print(f"[native] wrote {video_path}") + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/uv.lock b/integrations/sana/tests/parity_check/uv.lock new file mode 100644 index 000000000..baa519ed3 --- /dev/null +++ b/integrations/sana/tests/parity_check/uv.lock @@ -0,0 +1,1676 @@ +version = 1 +revision = 3 +requires-python = "==3.12.*" +resolution-markers = [ + "sys_platform == 'win32'", + "sys_platform != 'win32'", +] + +[manifest] +overrides = [{ name = "nvidia-cublas", specifier = ">=13.4" }] + +[[package]] +name = "accelerate" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8d/75/94cd5d389649578aca399e5aa822637eec18319a1dadc400ffe2f9a7493f/accelerate-1.14.0.tar.gz", hash = "sha256:41b9c4377a54e0b460a959b0defa1b736e4ca0a2373252d9a539964c2afe3c8d", size = 412167, upload-time = "2026-06-11T13:45:52.326Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl", hash = "sha256:e94390c2863b873be18f623f9df48a0d8fe5eff13ea7f1a00092b0a7904888c6", size = 389246, upload-time = "2026-06-11T13:45:50.477Z" }, +] + +[[package]] +name = "annotated-doc" +version = "0.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/ba/046ceea27344560984e26a590f90bc7f4a75b06701f653222458922b558c/annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4", size = 7288, upload-time = "2025-11-10T22:07:42.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, +] + +[[package]] +name = "annotated-types" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" }, +] + +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + +[[package]] +name = "anyio" +version = "4.14.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl", hash = "sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494", size = 125813, upload-time = "2026-07-12T20:29:05.763Z" }, +] + +[[package]] +name = "av" +version = "18.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/a4/570a5a35c8638aba01e739925846c35fdd6b0756a15526766d0a4dd3b7df/av-18.0.0.tar.gz", hash = "sha256:4ef7e72c3d3a872584a1215173b16e0226811037f40dcdbf75992631098df1ba", size = 4340222, upload-time = "2026-07-02T06:37:58.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/4a/9e3463df030e063d757fa12f0f39be6541b45b06b5bad48c2ce361b924bf/av-18.0.0-cp311-abi3-macosx_11_0_x86_64.whl", hash = "sha256:149289d40e732a6e49c9530bc245b49d9964cfd1c8c9e06778703b7d5bba6b25", size = 22499354, upload-time = "2026-07-02T06:36:58.751Z" }, + { url = "https://files.pythonhosted.org/packages/77/b3/2576a44b4f39c7462ced4c17fec04c756f7b0f3c5cb940d124173e417d6a/av-18.0.0-cp311-abi3-macosx_14_0_arm64.whl", hash = "sha256:35274c20d2ad3b4774fe632bcef2e34af79858ddf899352339cc3babbc13a484", size = 18175248, upload-time = "2026-07-02T06:37:01.741Z" }, + { url = "https://files.pythonhosted.org/packages/84/74/6732f17b96dc23fd23b876b2805435855abdc8a3b397142be4e581165de8/av-18.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:4d683b7747a0ba9222b8a5f81e41db5f796e7f64473454ec4fe2548e083c2fa0", size = 33387843, upload-time = "2026-07-02T06:37:05.097Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b9/7708c43fed7ae28b4a1bad060b4221e3334cd827cec24f7165902a6ac1f4/av-18.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ae56b40b6f8b067a8ad2dac664fbfbabac7f7a55b9a7bb031eb99289252bc017", size = 35536910, upload-time = "2026-07-02T06:37:08.806Z" }, + { url = "https://files.pythonhosted.org/packages/5a/94/eba99691d184f6a395a242d54dc370e2fd2265e95bbc98e2963a0fdbdd6c/av-18.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:ea2e8ebbce521f21b55df9400e00d721623c9020ef158f5a188a96130be0743f", size = 38984619, upload-time = "2026-07-02T06:37:11.861Z" }, + { url = "https://files.pythonhosted.org/packages/c9/cf/0d7aee07fe16aa9ffdf96043c14bed5485a52c0dea4259de87aa306ecab4/av-18.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef96dabb3e50dac249913145dff5424b302b257fd95dcb64be3c7b7a8aef16d1", size = 34451176, upload-time = "2026-07-02T06:37:15.154Z" }, + { url = "https://files.pythonhosted.org/packages/76/92/810da80b12680d4c4fe235bd1b4003289be9213ac7f114b77b8ecf0e3b3e/av-18.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:0f65518a184613e41536f29e8758c8e3d8293e46bf5bef108f04f925bbfa3f44", size = 36619869, upload-time = "2026-07-02T06:37:18.495Z" }, + { url = "https://files.pythonhosted.org/packages/11/85/0f121ff43dc5a70696676c98a8f1674e2fa787614c2abaacb15fa1a9bc99/av-18.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:aaf4d354d2beaa6651e4f92e54409a578bde64f79c0beef9a30b388d06f7c629", size = 27556236, upload-time = "2026-07-02T06:37:21.388Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f6/2509754d4d2356abc6fc0ea3d57c12ade29bac23a1fb7fc215a53ca518fb/av-18.0.0-cp311-abi3-win_arm64.whl", hash = "sha256:adac2b3833b6cb9bd6cb52664a522b94db453615b3675b1dbb26e13fe1c80da6", size = 20221133, upload-time = "2026-07-02T06:37:23.88Z" }, +] + +[[package]] +name = "boto3" +version = "1.43.53" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6c/da/0e90eab875f2eb4b8708fef2c198f2559ed1e451a1016f1cd4fcdcfbfbe3/boto3-1.43.53.tar.gz", hash = "sha256:c80425acab314d7af09609562053f565139e1fe49108eacfcc1601ebfaee235b", size = 112678, upload-time = "2026-07-21T19:28:53.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/27/7e72d25fdde77668b7bd4fa47381192dd2aa64fb77265e4bab786fd9fe2a/boto3-1.43.53-py3-none-any.whl", hash = "sha256:5383e705d8a976a14f23bb8c113c07a396931a019db98fce4cdc68650ec6e4d8", size = 140025, upload-time = "2026-07-21T19:28:50.957Z" }, +] + +[[package]] +name = "botocore" +version = "1.43.53" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/6b/ebcefacc4de3cd4f1c449540d86877f76c2f5e586a620831012decbb2b2c/botocore-1.43.53.tar.gz", hash = "sha256:36d93dd8db68ee75f6b61ca9f775161b8168844e4601698701530e6efdded141", size = 15720336, upload-time = "2026-07-21T19:28:41.547Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/e5/1b60e394f0fff97ee70dd16913382b7dffda85b98e639c0c9e8ff56cfaa7/botocore-1.43.53-py3-none-any.whl", hash = "sha256:b7ee9a70d187e5348883c820990ccd9436ab14e2bd6622741fc96fe561e816b8", size = 15404628, upload-time = "2026-07-21T19:28:37.475Z" }, +] + +[[package]] +name = "certifi" +version = "2026.7.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bd/2a/23f34ec9d04624958e137efdc394888716353190e75f25dd22c7a2c7a8aa/charset_normalizer-3.4.9.tar.gz", hash = "sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b", size = 152439, upload-time = "2026-07-07T14:34:58.454Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/70/4a/ecbd131485c07fcdfad54e28946d513e3da22ef3b4bd854dcafae54ec739/charset_normalizer-3.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:45b0cc4e3556cd875e09102988d1ab8356c998b596c9fced84547c8138b487a0", size = 319300, upload-time = "2026-07-07T14:33:15.666Z" }, + { url = "https://files.pythonhosted.org/packages/ec/96/5d9364e3342d69f3a045e1777bc47c85c383e6e9466d561b33fdb419d1f9/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b2aff1c7b3884512b9512c3eaadd9bab39fb45042ffaaa1dd08ff2b9f8109d9", size = 215802, upload-time = "2026-07-07T14:33:17.031Z" }, + { url = "https://files.pythonhosted.org/packages/4b/4c/5361f9aa7f2cb58d94f2ab831b3d493f69efb1d239654b4744e3c09527cb/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9104ed0bd76a429d46f9ec0dbc9b08ad1d2dcdf2b00a5a0daa1c145329b35b44", size = 237171, upload-time = "2026-07-07T14:33:18.576Z" }, + { url = "https://files.pythonhosted.org/packages/50/78/ce342ca4ff30b2eb49fe6d9578df85974f90c67d294113e94efdd9664cbd/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7b86a2b16095d250c6f58b3d9b2eee6f4147754344f3dab0922f7c9bf7d226c9", size = 233075, upload-time = "2026-07-07T14:33:20.084Z" }, + { url = "https://files.pythonhosted.org/packages/01/c4/4fa4c8b3097a11f3c5f09a35b72ed6855fb1d332469504962ab7bafcc702/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5e226f6218febc71f6c1fc2fafb91c226f75bdc1d8fb12d66823716e891608fd", size = 224256, upload-time = "2026-07-07T14:33:21.747Z" }, + { url = "https://files.pythonhosted.org/packages/87/3a/ad914516df7e358a81aae018caa5e0470ba827fa6d763b1d2e87d920a5f6/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:90c44bc373b7687f6948b693cceaea1348ae0975d7474746559494468e3c1d84", size = 208784, upload-time = "2026-07-07T14:33:23.313Z" }, + { url = "https://files.pythonhosted.org/packages/d7/74/3c12f9755717dfe5c5c87da63f35d765fa0c00382ec26bf23f7fae34f2ba/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9cdef90ae47919cae358d8ab15797a800ed41da7aba5d72419fb510729e2ed4b", size = 219928, upload-time = "2026-07-07T14:33:24.814Z" }, + { url = "https://files.pythonhosted.org/packages/33/9a/895095b83e7907abd6d3d99aad3a38ad0d9686cc186cb0c94c24320fe63e/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:60f44ade2cf573dad7a277e6f8ca9a51a21dda572b13bd7d8539bb3cd5dbedde", size = 218489, upload-time = "2026-07-07T14:33:26.42Z" }, + { url = "https://files.pythonhosted.org/packages/a1/34/ef5c05f412f42520d7709b7d3784d19640839eb7366ded1755511585429f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a1786910334ed46ab1dd73222f2cd1e05c2c3bb39f6dddb4f8b36fc382058a39", size = 210267, upload-time = "2026-07-07T14:33:27.952Z" }, + { url = "https://files.pythonhosted.org/packages/83/dc/9b29fa4412b318bf3bfea985c35d67eb55e04b59a7c3f2237168b0e0be6f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:03d07803992c6c7bbc976327f34b18b6160327fc81cb82c9d504720ac0be3b62", size = 226030, upload-time = "2026-07-07T14:33:29.397Z" }, + { url = "https://files.pythonhosted.org/packages/0e/42/6dbc00b8cd16011691203e33570fa42ed5746599a2e878112d16eab403a3/charset_normalizer-3.4.9-cp312-cp312-win32.whl", hash = "sha256:78841cccf1af7b40f6f716338d50c0902dbe88d9f800b3c973b7a9a0a693a642", size = 151185, upload-time = "2026-07-07T14:33:30.781Z" }, + { url = "https://files.pythonhosted.org/packages/80/cc/f920afd1a23c58ccd53c1d36085a71893a4737ff5e66e0371efab6809850/charset_normalizer-3.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:4b3dac63058cc36820b0dd072f89898604e2d39686fe05321729d00d8ac185a0", size = 162557, upload-time = "2026-07-07T14:33:32.176Z" }, + { url = "https://files.pythonhosted.org/packages/f0/e6/0386d43a261ff4e4b30c5857af7df877254b46bec7b9d1b74b6bf969a90b/charset_normalizer-3.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:78fa18e436a1a0e58dbd7e02fc4473f3f32cceb12df9dfca542d075961c307d2", size = 152665, upload-time = "2026-07-07T14:33:33.711Z" }, + { url = "https://files.pythonhosted.org/packages/98/2b/f97f1c193fb855c345d678f5077d6926034db0722df74c8f057020e05a25/charset_normalizer-3.4.9-py3-none-any.whl", hash = "sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5", size = 64538, upload-time = "2026-07-07T14:34:56.993Z" }, +] + +[[package]] +name = "click" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/d4/81420972a676e8ffea40450d8c8c92943e7218a78fe9b64359836cc9876b/click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6", size = 338000, upload-time = "2026-06-24T17:45:15.148Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "cuda-bindings" +version = "13.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" }, + { url = "https://files.pythonhosted.org/packages/39/2a/6d2e9047d1fb243dbaa364b01e0297534b9ed7fd27dba1c9f361519cf69b/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a", size = 6657965, upload-time = "2026-05-29T23:11:52.227Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/b4/d088047afe39827556df21118cac9ffd20cc3f968c99a7681494d1eb333c/cuda_pathfinder-1.6.0-py3-none-any.whl", hash = "sha256:1503af579d8379c24bdd65528379bc57039b0455be9f5f9686cf8e473a1fce51", size = 54591, upload-time = "2026-07-21T15:03:56.224Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.3.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/c7/a79086a62c98befcdb8349656c6f114e2db3b8b2422f6e25c97a7f2a9a3c/cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f", size = 2512, upload-time = "2026-04-14T00:50:08.173Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cufile = [ + { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +curand = [ + { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cusolver = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +nvtx = [ + { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] + +[[package]] +name = "diffusers" +version = "0.39.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "importlib-metadata" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "regex" }, + { name = "requests" }, + { name = "safetensors" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1a/81/6095237b86a3116c4789f28c4435d5296c00c0fc74ffde99008fd6b3a36c/diffusers-0.39.0.tar.gz", hash = "sha256:14bb1d98c85a0e463d734c99aaa73b480a7bc9bad22af30fbf730ef8f09c1d67", size = 4651240, upload-time = "2026-07-03T08:48:47.904Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/3f/7469c46e9d22307ea686bab687d70e6bf328722952f9d10339f5e913e608/diffusers-0.39.0-py3-none-any.whl", hash = "sha256:912aca51b5787365110806e984d5555735bf8a461073bb8459029d0bca7870ef", size = 5631176, upload-time = "2026-07-03T08:48:45.337Z" }, +] + +[[package]] +name = "docstring-parser" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/4d/f332313098c1de1b2d2ff91cf2674415cc7cddab2ca1b01ae29774bd5fdf/docstring_parser-0.18.0.tar.gz", hash = "sha256:292510982205c12b1248696f44959db3cdd1740237a968ea1e2e7a900eeb2015", size = 29341, upload-time = "2026-04-14T04:09:19.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, +] + +[[package]] +name = "einops" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, +] + +[[package]] +name = "filelock" +version = "3.32.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c0/80/8232b582c4b318b817cf1274ba74976b07b34d35ef439b3eb948f98645a1/filelock-3.32.0.tar.gz", hash = "sha256:7be2ad23a14607ccc71808e68fe30848aeace7058ace17852f68e2a68e310402", size = 213757, upload-time = "2026-07-21T13:17:42.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/79/b4c714bef36bc4ec2beeae1e0c124f0223888cd8c6feb1cdc56038116920/filelock-3.32.0-py3-none-any.whl", hash = "sha256:d396bea984af47333ef05e50eae7eff88c84256de6112aea0ec48a233c064fe3", size = 97732, upload-time = "2026-07-21T13:17:41.55Z" }, +] + +[[package]] +name = "fla-core" +version = "0.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "einops" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4e/62/99e149f19a447ce809d7f4fa64ae61c073da337505b9db7f389502470820/fla_core-0.5.1.tar.gz", hash = "sha256:7f3cf56edfbaa9115f4937d1181372e5c7b11809ad8eb2e411fffc3caf729f48", size = 498800, upload-time = "2026-06-18T18:17:15.377Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/78/a55ee7a62515dcb9220770dd99dfe59ac6599da8af84ad1d20f9f407df4d/fla_core-0.5.1-py3-none-any.whl", hash = "sha256:02150d34aa1e37f6b8ed9b2feec5d29af93680573e5077ed981238179c11fb06", size = 702955, upload-time = "2026-06-18T18:17:12.229Z" }, +] + +[[package]] +name = "flash-linear-attention" +version = "0.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fla-core" }, + { name = "transformers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4c/e8/8f115be585a046795e4a1f7ade727889219bb66b8d96cc668b8c9e437c0e/flash_linear_attention-0.5.1.tar.gz", hash = "sha256:8840fd4c37de8b0612dc8fd493867f3d330672ba2f17c024a2ce37239634e247", size = 221733, upload-time = "2026-06-18T18:17:16.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/3c/5819fb19dc071302ca818616a4e64d4454e1f1193929eb8365c8d38e9052/flash_linear_attention-0.5.1-py3-none-any.whl", hash = "sha256:9022862f0a238752372c81290694b8b2ed1cb2d13fc40d340ffeeb554d6bee5c", size = 403096, upload-time = "2026-06-18T18:17:14.025Z" }, +] + +[[package]] +name = "flashdreams" +source = { editable = "../../../../flashdreams" } +dependencies = [ + { name = "boto3" }, + { name = "botocore" }, + { name = "einops" }, + { name = "filelock" }, + { name = "ftfy" }, + { name = "huggingface-hub" }, + { name = "loguru" }, + { name = "numpy" }, + { name = "nvidia-ml-py" }, + { name = "safetensors" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, + { name = "tqdm" }, + { name = "transformers" }, + { name = "triton-windows", marker = "sys_platform == 'win32'" }, + { name = "tyro" }, + { name = "urllib3" }, +] + +[package.metadata] +requires-dist = [ + { name = "aiohttp", marker = "extra == 'serving'", specifier = ">=3.9" }, + { name = "aiortc", marker = "extra == 'serving'", specifier = ">=1.9" }, + { name = "boto3", specifier = ">=1.35" }, + { name = "botocore", specifier = ">=1.35" }, + { name = "einops", specifier = ">=0.7" }, + { name = "filelock", specifier = ">=3" }, + { name = "ftfy", specifier = ">=6.0" }, + { name = "huggingface-hub", specifier = ">=0.33" }, + { name = "loguru", specifier = ">=0.7" }, + { name = "mediapy", marker = "extra == 'dev'", specifier = ">=1.1" }, + { name = "mediapy", marker = "extra == 'examples'", specifier = ">=1.1" }, + { name = "mediapy", marker = "extra == 'runners'", specifier = ">=1.1" }, + { name = "numpy", specifier = ">=1.24,<2.5" }, + { name = "nvidia-ml-py", specifier = ">=12.0" }, + { name = "opencv-python-headless", marker = "extra == 'examples'", specifier = ">=4.5" }, + { name = "opencv-python-headless", marker = "extra == 'runners'", specifier = ">=4.5" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pytest-manual-marker", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "safetensors", specifier = ">=0.4" }, + { name = "scipy", marker = "extra == 'examples'", specifier = ">=1.11" }, + { name = "scipy", marker = "extra == 'runners'", specifier = ">=1.11" }, + { name = "tomli", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "torch", marker = "sys_platform != 'win32'", specifier = ">=2.9" }, + { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu130" }, + { name = "tqdm", specifier = ">=4.60" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "sys_platform != 'win32' and extra == 'dev'", specifier = ">=2.12" }, + { name = "transformers", specifier = ">=5.0,<6" }, + { name = "triton-windows", marker = "sys_platform == 'win32'", specifier = ">=3.5.0" }, + { name = "tyro", specifier = ">=1.0" }, + { name = "urllib3", specifier = ">=2.7.0" }, +] +provides-extras = ["dev", "examples", "runners", "serving"] + +[package.metadata.requires-dev] +cuda12 = [ + { name = "torch", marker = "sys_platform != 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu128", conflict = { package = "flashdreams", group = "cuda12" } }, + { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu130" }, + { name = "torchvision", marker = "sys_platform != 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu128", conflict = { package = "flashdreams", group = "cuda12" } }, + { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu130" }, +] +cuda13 = [ + { name = "torch", marker = "sys_platform != 'win32'", specifier = ">=2.9" }, + { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu130" }, + { name = "torchvision", marker = "sys_platform != 'win32'", specifier = ">=0.24" }, + { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu130" }, +] + +[[package]] +name = "flashdreams-sana-wm" +version = "0.1.0" +source = { editable = "../../" } +dependencies = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "flashdreams" }, + { name = "imageio", extra = ["ffmpeg"] }, + { name = "pillow" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torchvision" }, + { name = "transformers" }, +] + +[package.metadata] +requires-dist = [ + { name = "accelerate", specifier = ">=1.0" }, + { name = "diffusers", specifier = ">=0.36" }, + { name = "flashdreams", editable = "../../../../flashdreams" }, + { name = "imageio", extras = ["ffmpeg"], specifier = ">=2.31" }, + { name = "pillow", specifier = ">=10" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "safetensors", specifier = ">=0.5" }, + { name = "torchvision", specifier = ">=0.26" }, + { name = "transformers", specifier = ">=5.0,<6" }, +] +provides-extras = ["dev"] + +[[package]] +name = "fsspec" +version = "2026.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/10/a1/ae4e3e5003468d6391d2c77b6fa1cd73bd5d13511d81c642d7b28ac90ed4/fsspec-2026.6.0.tar.gz", hash = "sha256:f5bac145310fe30e16e1471bd6840b2d990d609e872251d7e674241822abf01a", size = 313646, upload-time = "2026-06-16T01:57:28.105Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/22/4222d7ddf3da30f363edaa98e329c2bce6c65497c9cb2810931c8b2c0fbc/fsspec-2026.6.0-py3-none-any.whl", hash = "sha256:02e0b71817df9b2169dc30a16832045764def1191b43dcff5bb85bdee212d2a1", size = 203949, upload-time = "2026-06-16T01:57:26.358Z" }, +] + +[[package]] +name = "ftfy" +version = "6.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a5/d3/8650919bc3c7c6e90ee3fa7fd618bf373cbbe55dff043bd67353dbb20cd8/ftfy-6.3.1.tar.gz", hash = "sha256:9b3c3d90f84fb267fe64d375a07b7f8912d817cf86009ae134aa03e1819506ec", size = 308927, upload-time = "2024-10-26T00:50:35.149Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/6e/81d47999aebc1b155f81eca4477a616a70f238a2549848c38983f3c22a82/ftfy-6.3.1-py3-none-any.whl", hash = "sha256:7c70eb532015cd2f9adb53f101fb6c7945988d023a085d127d1573dc49dd0083", size = 44821, upload-time = "2024-10-26T00:50:33.425Z" }, +] + +[[package]] +name = "h11" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, +] + +[[package]] +name = "hf-xet" +version = "1.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/63/39/67be8d71f900d9a55761b6022821d6679fb56c64f1b6063d5af2c2606727/hf_xet-1.5.2.tar.gz", hash = "sha256:73044bd31bae33c984af832d19c752a0dffb67518fee9ddbd91d616e1101cf47", size = 903674, upload-time = "2026-07-16T17:29:56.833Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/ba/2b70603c7552db82baeb2623e2336898304a17328845151be4fe1f48d420/hf_xet-1.5.2-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:f922b8f5fb84f1dd3d7ab7a1316354a1bca9b1c73ecfc19c76e51a2a49d29799", size = 4033760, upload-time = "2026-07-16T17:29:43.884Z" }, + { url = "https://files.pythonhosted.org/packages/60/ac/b097a86a1e4a6098f3a79382643ab09d5733d87ccc864877ad1e12b49b70/hf_xet-1.5.2-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:045f84440c55cdeb659cf1a1dd48c77bcd0d2e93632e2fea8f2c3bdee79f38ed", size = 3841438, upload-time = "2026-07-16T17:29:45.539Z" }, + { url = "https://files.pythonhosted.org/packages/d3/35/db860aa3a0780660324a506ad4b3d322ddc6ecbba4b9340aed0942cbf21c/hf_xet-1.5.2-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:db78c39c83d6279daddc98e2238f373ab8980685556d42472b4ec51abcf03e8c", size = 4428006, upload-time = "2026-07-16T17:29:46.996Z" }, + { url = "https://files.pythonhosted.org/packages/af/6b/832dd980af4b0c3ae0660e309285f2ffcdff2faa38129390dbb47aa4a3f9/hf_xet-1.5.2-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:7db73c810500c54c6760be8c39d4b2e476974de85424c50063efc22fdda13025", size = 4221099, upload-time = "2026-07-16T17:29:48.525Z" }, + { url = "https://files.pythonhosted.org/packages/9e/05/ae50f0d34e3254e6c3e208beb2519f6b8673016fc4b3643badaf6450d186/hf_xet-1.5.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:6395cfe3c9cbead4f16b31808b0e67eac428b66c656f856e99636adaddea878f", size = 4420766, upload-time = "2026-07-16T17:29:50.092Z" }, + { url = "https://files.pythonhosted.org/packages/07/a9/c050bc2743a2bcd68928bfee157b08681667a164a24ec95fbfcfcd717e08/hf_xet-1.5.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:cde8cd167126bb6109b2ceb19b844433a4988643e8f3e01dd9dd0e4a34535097", size = 4636716, upload-time = "2026-07-16T17:29:51.62Z" }, + { url = "https://files.pythonhosted.org/packages/e9/f8/68b01c5c2edb56ac9a67b3d076ffddcb90867abaee923923eb34e7a14e76/hf_xet-1.5.2-cp38-abi3-win_amd64.whl", hash = "sha256:ecf63d1cb69a9a7319910f8f83fcf9b46e7a32dfcf4b8f8eeddb55f647306e65", size = 3988373, upload-time = "2026-07-16T17:29:53.395Z" }, + { url = "https://files.pythonhosted.org/packages/39/c6/988383e9dc17294d536fcbcd6fd16eed882e411ad16c954984a53e47b09c/hf_xet-1.5.2-cp38-abi3-win_arm64.whl", hash = "sha256:1da28519496eb7c8094c11e4d25509b4a468457a0302d58136099db2fd9a671d", size = 3816957, upload-time = "2026-07-16T17:29:54.991Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, +] + +[[package]] +name = "huggingface-hub" +version = "1.24.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "httpx" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "tqdm" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/9b/d3bb4e7d792835daf34dd7091bbc7d7b4e0437d9388f1ea7239cce49f478/huggingface_hub-1.24.0.tar.gz", hash = "sha256:18431ff4daae0749aa9ba102fc952e314c98e1d30ebdec5319d85ca0a83e1ae5", size = 921848, upload-time = "2026-07-17T09:54:01.022Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/c3/aeaaf3911d2529614be18d1c8b5496afc185560e76568063d517283318af/huggingface_hub-1.24.0-py3-none-any.whl", hash = "sha256:6ed4120a84a6beec900640aa7e346bd766a6b7341e41526fef5dc8bd81fb7d59", size = 771904, upload-time = "2026-07-17T09:53:59.106Z" }, +] + +[[package]] +name = "idna" +version = "3.18" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, +] + +[[package]] +name = "imageio" +version = "2.37.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/48/62/aa770a9307508d2a2a2c62d536a49347bffe9e55322db27838d3c93d0b07/imageio-2.37.4.tar.gz", hash = "sha256:e45cbc5e83502047fb138f7f585f7f105a136a57eea5f4b3cfc6ce1b52720bd3", size = 390173, upload-time = "2026-07-20T05:26:11.369Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3e/2d/ca050652104bab2cf55e569db2a178b1b61cb041fef28307f2db383f6d9f/imageio-2.37.4-py3-none-any.whl", hash = "sha256:1ab2e22c8debf700f24c3ac43e8f95f3b3a8110c83b93411e97b4b0b2cd1c7e6", size = 318000, upload-time = "2026-07-20T05:26:09.874Z" }, +] + +[package.optional-dependencies] +ffmpeg = [ + { name = "imageio-ffmpeg" }, + { name = "psutil" }, +] + +[[package]] +name = "imageio-ffmpeg" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, + { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" }, + { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" }, + { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" }, +] + +[[package]] +name = "importlib-metadata" +version = "9.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "zipp" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + +[[package]] +name = "loguru" +version = "0.7.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "win32-setctime", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3a/05/a1dae3dffd1116099471c643b8924f5aa6524411dc6c63fdae648c4f1aca/loguru-0.7.3.tar.gz", hash = "sha256:19480589e77d47b8d85b2c827ad95d49bf31b0dcde16593892eb51dd18706eb6", size = 63559, upload-time = "2024-12-06T11:20:56.608Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/29/0348de65b8cc732daa3e33e67806420b2ae89bdce2b04af740289c5c6c8c/loguru-0.7.3-py3-none-any.whl", hash = "sha256:31a33c10c8e1e10422bfd431aeb5d351c7cf7fa671e3c4df004162264b28220c", size = 61595, upload-time = "2024-12-06T11:20:54.538Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" }, + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "mypy-extensions" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/2a/3d7b5ac8aac24feaf9ad7ed58f45b0bbc06d37e4338ae84c9f2298b570f9/numpy-2.4.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:001fbb8e08d942dd57599e781f2472269ee7f2755fae407b4f67b2f0b17da3f1", size = 16689119, upload-time = "2026-05-18T23:33:54.065Z" }, + { url = "https://files.pythonhosted.org/packages/ea/12/92c4c131527599e8288d6918e888d88726f84d805d784b771f32408aeaef/numpy-2.4.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ebfb099f8dcf083deef3ac1ca4c1503f387cf76296fcb3816b66f5ecb5f54fdb", size = 14699246, upload-time = "2026-05-18T23:33:57.621Z" }, + { url = "https://files.pythonhosted.org/packages/ad/fe/c0a6b7b2ca128a8fb228575147073b660656734b8ebe4d76c8fd748dcc79/numpy-2.4.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:3213d622a0283a39a93d188f3cf72b26862df52fbb4ca3697f51705016523d41", size = 5204410, upload-time = "2026-05-18T23:34:00.302Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d4/9770d14ba719432bb90a421bfd443872ed0f70f7264b64bec12ea363d5fd/numpy-2.4.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:357cc07a6d7b0b182ff02249616a03742827ebb1277546b5c7cd7f7620a45698", size = 6551240, upload-time = "2026-05-18T23:34:02.852Z" }, + { url = "https://files.pythonhosted.org/packages/c9/c6/50a46a6205feba2343f1d6d17438107c5dc491ed1c736e6ea68689fd906b/numpy-2.4.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f9fb9157b4ce2971008323afe46053787b526ef624fea915b261468a8421a0f", size = 15671012, upload-time = "2026-05-18T23:34:05.485Z" }, + { url = "https://files.pythonhosted.org/packages/99/60/14115e6364fa676c5397c2ad3004e527e9aa487abf5d0706ec81bbd08529/numpy-2.4.6-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90f9849678c75fe7afa2d348ac842c168b0a4d3d61919687216dfc547976d853", size = 16645538, upload-time = "2026-05-18T23:34:09.265Z" }, + { url = "https://files.pythonhosted.org/packages/ae/c5/693cbe59e57db94d2231fa519ca3978dc9e19da5a8f088588f5c6e947ff2/numpy-2.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c1a2af6c6ef86344a6b0db6b97834208bf598db514f2b155042439b62605601a", size = 17020706, upload-time = "2026-05-18T23:34:13.053Z" }, + { url = "https://files.pythonhosted.org/packages/ef/fc/85b7c4eff9b4966ade25c2273cf7e7012e92366c032058653934b37de044/numpy-2.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e5805d5a22fd19c8ccff10a9561f9df94436b0545619ea579db2d3c35294bce2", size = 18368541, upload-time = "2026-05-18T23:34:17.024Z" }, + { url = "https://files.pythonhosted.org/packages/f6/81/e1b27545deedce7f4a0b348618c6b62d74e36a4dc9ccd42f3eb2f85eee32/numpy-2.4.6-cp312-cp312-win32.whl", hash = "sha256:e3eeb0aabd6bd5ce64faae67e9935203a6991b4bc2a485a767fbafb2c5125f45", size = 5962825, upload-time = "2026-05-18T23:34:20.3Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ca/feab00bd44aa5fe1ad2c18f08b4d3bb92e26484b0b1d1443897809ed528c/numpy-2.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:d8e8286dd7cea7895157318d1b91cdacac64c479f3cbc8dce548331728484751", size = 12321687, upload-time = "2026-05-18T23:34:23.095Z" }, + { url = "https://files.pythonhosted.org/packages/63/cf/5a6d34850a39d1093558564f77ee8e8e0bee5061151b8f05a55711001ec7/numpy-2.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:4081eb135ac24158bd51cdfbef16f1c64df7063b1143f24731387137c092bec8", size = 10221482, upload-time = "2026-05-18T23:34:25.876Z" }, +] + +[[package]] +name = "nvdlfw-inspect" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" } }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/86/94188e03e5d4dd7b73c390b0cddcde5618b3799c18e327b2bf15763f6137/nvdlfw_inspect-0.2.2-py3-none-any.whl", hash = "sha256:8a4dc2814c5a4cd19ae304170b9bfa514538ef3c3eb243a45a82404ec3cb279d", size = 30964, upload-time = "2025-12-03T10:52:01.933Z" }, +] + +[[package]] +name = "nvidia-cublas" +version = "13.6.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/4f/22/7c08d8e93f6a2e4879ac83c12696aecaf17a0fc2c9e8d204caceaa3b8426/nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:946f6a252b1cc72d8de912c75975fd6d8ba44f67d4e5044fe764ddb909f4a688", size = 518599300, upload-time = "2026-06-29T16:54:56.859Z" }, + { url = "https://files.pythonhosted.org/packages/fd/6c/173c7a3db77a6592210f73f194f0f8ed5e51b6ec61cfed7b1eee06ac5fd3/nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b82c80c886cea6da6e149a5c3bdba274f12b7e4ec4b00a050b916b0446fb4153", size = 410473152, upload-time = "2026-06-29T16:55:54.297Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, + { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, + { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime" +version = "13.0.96" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, + { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu13" +version = "9.20.0.48" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, + { url = "https://files.pythonhosted.org/packages/6e/5e/edb9c0ae051602c3ccaffe424256463636d639e27d7f302dde9975ef9e7a/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304", size = 366173588, upload-time = "2026-03-09T19:29:34.474Z" }, +] + +[[package]] +name = "nvidia-cufft" +version = "12.0.0.61" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, +] + +[[package]] +name = "nvidia-cufile" +version = "1.15.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, + { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, +] + +[[package]] +name = "nvidia-curand" +version = "10.4.0.35" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, +] + +[[package]] +name = "nvidia-cusolver" +version = "12.0.4.66" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, + { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, +] + +[[package]] +name = "nvidia-cusparse" +version = "12.6.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu13" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" }, + { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" }, +] + +[[package]] +name = "nvidia-ml-py" +version = "13.610.43" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/b5/a8fbc356f768fa5c9cfd646668fd7d34bf55bdd1c6e20754642a64d930d4/nvidia_ml_py-13.610.43.tar.gz", hash = "sha256:65437eb73d68d0c62c931ca4d45038472faff03bd0b8729abba4b899f70d60f2", size = 52109, upload-time = "2026-06-01T18:54:08.829Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl", hash = "sha256:f13c72698edef492f985cc225f14faafe68ae065a2e407f45bdf6f4b9b43fde8", size = 53163, upload-time = "2026-06-01T18:54:07.704Z" }, +] + +[[package]] +name = "nvidia-nccl-cu13" +version = "2.29.7" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/72/0d/daf50d44177ee0cbc7ff0a0c91eb5ff676c82be42f9a970bc7597f440c3a/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5", size = 206014712, upload-time = "2026-03-03T05:34:20.843Z" }, + { url = "https://files.pythonhosted.org/packages/67/f4/58e4e91b6919367c7aafb8e36fce9aad1a3047e536bf7e2fd560927d3a4c/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d", size = 205976000, upload-time = "2026-03-03T05:36:24.472Z" }, +] + +[[package]] +name = "nvidia-nvjitlink" +version = "13.3.33" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f0/ee/580ca6f29dcab0221db8706badca1bbbb084f1975c4d4e83329c3a7e31f0/nvidia_nvjitlink-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5", size = 40742423, upload-time = "2026-05-26T16:54:51.613Z" }, + { url = "https://files.pythonhosted.org/packages/69/30/45414e35ff2eee7db3da037e5707037ccf9d2b5218ffbdb055ea4d5aa98a/nvidia_nvjitlink-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e", size = 39168635, upload-time = "2026-05-26T16:54:13.906Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu13" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, + { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, +] + +[[package]] +name = "nvidia-nvtx" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, + { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, +] + +[[package]] +name = "omegaconf" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/3d/e4b57b8d9008c6ebe0d5eff901f91d5700cf7bdb8c8863df817463a7fd5e/omegaconf-2.3.1.tar.gz", hash = "sha256:e5e7de64aeebeddaf8e6d3f7a783b32ac2a01c0fbd9c878012caecb891a1f42a", size = 3298472, upload-time = "2026-06-11T05:05:12.885Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/0e/152509871bf30df6fc38569f52a2db9b55dd41aae957adae50a053ac7778/omegaconf-2.3.1-py3-none-any.whl", hash = "sha256:3d701d14e9a8828f1edd28bb70b725908b34277cdd72cf7d6a83f94dadc6b6a0", size = 79502, upload-time = "2026-06-11T05:05:09.954Z" }, +] + +[[package]] +name = "onnx" +version = "1.22.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/04/19/8ea73a64b368b75fe339771a20a02bc61ea1f551484c9e3d9d0bfbd0450f/onnx-1.22.0.tar.gz", hash = "sha256:ef40c0aaf0b643857ea9306fc7eddce17eaf9fb0407e4801f1fc5758443a38e0", size = 12024721, upload-time = "2026-06-15T12:50:05.354Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/6a/481561f1093834376ed493e4ca42a73e5be0d50031f2969c86593bdc7c96/onnx-1.22.0-cp312-abi3-macosx_12_0_universal2.whl", hash = "sha256:596fbf0490947533c1c1045ba860851dc9fb77471023dac9a71ba5b42ceab103", size = 20167081, upload-time = "2026-06-15T12:49:32.078Z" }, + { url = "https://files.pythonhosted.org/packages/84/55/b34fc2aa30aa54b4a775402d24c4082242c720283a274fe976ac8eb94480/onnx-1.22.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ae5a563f281cd9d2845622cecf6c092a57e4ee1b138f66fdbbdd4200567a5e16", size = 18889249, upload-time = "2026-06-15T12:49:34.7Z" }, + { url = "https://files.pythonhosted.org/packages/09/a6/bd32357e6cc1ecb473afd78193d7231724f284435d2db25696ecfaaa1503/onnx-1.22.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:955e02e1f6d385b53d52f9cd7b9cdf5caf417c300bcfe3c64c6d542be763845b", size = 19106514, upload-time = "2026-06-15T12:49:37.424Z" }, + { url = "https://files.pythonhosted.org/packages/5a/9d/3af461ac6c714b8b369cb71499659932f4f12cfb066250b62f7567c3d530/onnx-1.22.0-cp312-abi3-pyemscripten_2025_0_wasm32.whl", hash = "sha256:82e9f27fc1223cb06d68a56bed6f9d3caf3d0dad1b61bce45006d529b15bd94c", size = 16966387, upload-time = "2026-06-15T12:49:40.918Z" }, +] + +[[package]] +name = "onnx-ir" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/e6/672fefb2f108d077f58181a7babf4c0f8d1182a30353ffc9c79c63afc5ee/onnx_ir-0.2.1.tar.gz", hash = "sha256:8b8b10a93f43e65962104de6070c43c5dacb0e3cdfefc7c8059dd83c9db64f35", size = 144279, upload-time = "2026-04-20T20:21:47.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/aa/f7a53321c60b9ad9ee184b6018292ed6b5389947592a2c8c09c736bb7f9e/onnx_ir-0.2.1-py3-none-any.whl", hash = "sha256:c7285da889312f91882de2092e298a9eeeefbfc1d1951c49d983992967eb09a7", size = 166792, upload-time = "2026-04-20T20:21:46.357Z" }, +] + +[[package]] +name = "onnxscript" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c4/3a/4d79bce3f460e0df7fed54a92ce80827f25da66511da368bb00783ad8d20/onnxscript-0.7.1.tar.gz", hash = "sha256:309fb86484b11fa4ded90dba580e0d63f1a0827588e521cecaf2eeddb46d6e86", size = 618160, upload-time = "2026-06-29T23:33:21.526Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/bd/a0c8e737b6afda10e42a597787d53d5b66e00268df6f59184701eeae37d9/onnxscript-0.7.1-py3-none-any.whl", hash = "sha256:544763b7fdef49940cdd9412ff5135cbae96d59ac6bc1921457f21280f40f4b7", size = 721970, upload-time = "2026-06-29T23:33:23.298Z" }, +] + +[[package]] +name = "opencv-python-headless" +version = "5.0.0.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/99/76b7c80252aa83c1af16393454aafd125a0287101afe8deb0a6821af0e30/opencv_python_headless-5.0.0.93.tar.gz", hash = "sha256:b82f9831daab90b725c7c1ee1b36cb5732c367096ac76d119e64e14eb70d5f3c", size = 81817738, upload-time = "2026-07-02T07:01:06.039Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/7c/8c8097891c509d98cd128493835c95631c80be6a8f37ed9d25716c2e16f1/opencv_python_headless-5.0.0.93-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:030ca5e0837a2963ab36ef896baa9767eb8d2b83353fb28af5a521e40dd8756f", size = 48322581, upload-time = "2026-07-02T05:50:34.207Z" }, + { url = "https://files.pythonhosted.org/packages/90/8c/eab2ad388c3cbab2a350c10c2ef19ce6bd099240afc31789032c996bab52/opencv_python_headless-5.0.0.93-cp37-abi3-macosx_14_0_x86_64.whl", hash = "sha256:1e55af3abfb462eeeabe5c775f12bdb36216d8a93a3583d69e6bd6e1d6ba7d00", size = 34782894, upload-time = "2026-07-02T05:51:39.856Z" }, + { url = "https://files.pythonhosted.org/packages/ec/78/afca939f40ffe2b2380bfa86f812b2f7d4acc5a27b27dc41b49cad7ce7b4/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:10818d91510e05c04568ae12b5cd120779c70c01bf897b001a6221fe430df80f", size = 36521085, upload-time = "2026-07-02T06:55:24.429Z" }, + { url = "https://files.pythonhosted.org/packages/2b/97/8170e9819764c47e436c130d3ff6cfb73b58f923eae9d3a03d8982b04aec/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:09a872a157c1376ab922a69bbf22f9a95bcc7b658a9d8b436a60212b02b2eeb4", size = 56563598, upload-time = "2026-07-02T06:55:47.355Z" }, + { url = "https://files.pythonhosted.org/packages/3a/98/1a28a7101e31801042b3098871a74b76c61581d328ef40774ff4edb53a56/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:840bd717c21e5c11cadadc022a823315ea417f961213d06b4df010e019eb16f4", size = 39648433, upload-time = "2026-07-02T06:56:04.255Z" }, + { url = "https://files.pythonhosted.org/packages/9b/21/f6ef335f6e65724aa78b8d792b48d40a48c381715f1e62f5a5049e09d07e/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ed709fdf9aa0bd1f2ed8549e71d19449b03a675bb581eb292285f6861953be37", size = 61204038, upload-time = "2026-07-02T06:56:41.823Z" }, + { url = "https://files.pythonhosted.org/packages/d0/8f/b8756467ea991449a293797f6b3fa80fcfdd29598a0a60d1cd5715b96e61/opencv_python_headless-5.0.0.93-cp37-abi3-win32.whl", hash = "sha256:c6bcd96b185975ea240d22cfdb15a1f6d080cc95264cfbe2621f21bb144d89b9", size = 35411237, upload-time = "2026-07-02T05:50:12.901Z" }, + { url = "https://files.pythonhosted.org/packages/b8/88/763b967f7efd7226b82c9fae16d560cba049b1f0c036647e65c610fd636e/opencv_python_headless-5.0.0.93-cp37-abi3-win_amd64.whl", hash = "sha256:829717b6a95554f273e49e357cee3b3a2a26b6f4842fbc1bed2b45bdd8f87e0e", size = 43825962, upload-time = "2026-07-02T05:50:09.627Z" }, +] + +[[package]] +name = "packaging" +version = "26.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, +] + +[[package]] +name = "pillow" +version = "12.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/bf/fb3ebff8ddcb76aac5a01389251bbbb9519922a9b520d8247c1ca864a25d/pillow-12.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ba09209fbe443b4acccebe845d8a138b89a8f4fbaeedd44953490b5315d5e965", size = 5345969, upload-time = "2026-07-01T11:54:06.397Z" }, + { url = "https://files.pythonhosted.org/packages/d8/66/9a386a92561f402389a4fc70c18838bf6d35eb5eb5c6850b4b2dc64f5048/pillow-12.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7", size = 4780323, upload-time = "2026-07-01T11:54:09.351Z" }, + { url = "https://files.pythonhosted.org/packages/25/27/ac8f99618ffd3dde21db0f4d4b1d2ab00c0880595bfd17df103f7f39fd0c/pillow-12.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9", size = 6266838, upload-time = "2026-07-01T11:54:11.71Z" }, + { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" }, + { url = "https://files.pythonhosted.org/packages/eb/51/8b08617af3ad95e33ce6d7dd2c99ed6c8298f7fb131636303956be022e25/pillow-12.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c", size = 6344383, upload-time = "2026-07-01T11:54:15.756Z" }, + { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" }, + { url = "https://files.pythonhosted.org/packages/20/20/25e0f4dc178a6bc0696793720055519a0de89e7661dae886992decbd2f81/pillow-12.3.0-cp312-cp312-win32.whl", hash = "sha256:dbce0b29841537a2fa4a214c2bbf14de3587c9680caa9b4e217568472490b28f", size = 6472684, upload-time = "2026-07-01T11:54:19.839Z" }, + { url = "https://files.pythonhosted.org/packages/45/89/da2f7971a317f83d807fdd4065c0af40208e59e692cc43d315a71a0e96d1/pillow-12.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a2b55dd6b2a4c4b7d87ffa56bdb33fdc5fdb9a462173861a7bc097f17d91cb09", size = 7227137, upload-time = "2026-07-01T11:54:22.025Z" }, + { url = "https://files.pythonhosted.org/packages/de/47/4845a0a6c0dbf1db8456bd9fc791f13c5ced7ced20606d08a0aacfd25b49/pillow-12.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:331b624368d4f1d069149002f25f44bc61c8919ce8ddb3c45bdad8f6e2d89510", size = 2568267, upload-time = "2026-07-01T11:54:24.051Z" }, +] + +[[package]] +name = "protobuf" +version = "7.35.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/03/8aeeb7458d22546bf64b5250ca1daeb5ff757d900e8e4a7476c6f0db843e/protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6", size = 433226, upload-time = "2026-06-11T21:55:31.719Z" }, + { url = "https://files.pythonhosted.org/packages/37/4b/dfb89eb0e652a1ff073c39a59fb5e3a83cfe9b57a2c83fa6d78270101767/protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799", size = 328847, upload-time = "2026-06-11T21:55:34.035Z" }, + { url = "https://files.pythonhosted.org/packages/0f/58/dc12f2cd484951524af6e3382c785869b9b3fb5e52ee95ae23add53ee8f9/protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4", size = 344030, upload-time = "2026-06-11T21:55:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" }, + { url = "https://files.pythonhosted.org/packages/d8/bc/6d6c7ba8709c85f8f2c390b2b118d6fb08a783676a572271851bf45a7d22/protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30", size = 428945, upload-time = "2026-06-11T21:55:37.034Z" }, + { url = "https://files.pythonhosted.org/packages/0a/19/8d0cb6f20a1ef7b18f1c8986ad5783f22f84cce39c6ce9a6e645ea55192e/protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87", size = 439996, upload-time = "2026-06-11T21:55:38.123Z" }, + { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, +] + +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + +[[package]] +name = "pydantic" +version = "2.13.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/18/a5/b60d21ac674192f8ab0ba4e9fd860690f9b4a6e51ca5df118733b487d8d6/pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6", size = 844775, upload-time = "2026-05-06T13:43:05.343Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/7b/122376b1fd3c62c1ed9dc80c931ace4844b3c55407b6fb2d199377c9736f/pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba", size = 472262, upload-time = "2026-05-06T13:43:02.641Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.46.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/56/921726b776ace8d8f5db44c4ef961006580d91dc52b803c489fafd1aa249/pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1", size = 471464, upload-time = "2026-05-06T13:37:06.98Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/8c/af022f0af448d7747c5154288d46b5f2bc5f17366eaa0e23e9aa04d59f3b/pydantic_core-2.46.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2", size = 2106158, upload-time = "2026-05-06T13:38:57.215Z" }, + { url = "https://files.pythonhosted.org/packages/19/95/6195171e385007300f0f5574592e467c568becce2d937a0b6804f218bc49/pydantic_core-2.46.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f", size = 1951724, upload-time = "2026-05-06T13:37:02.697Z" }, + { url = "https://files.pythonhosted.org/packages/8e/bc/f47d1ff9cbb1620e1b5b697eef06010035735f07820180e74178226b27b3/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7", size = 1975742, upload-time = "2026-05-06T13:37:09.448Z" }, + { url = "https://files.pythonhosted.org/packages/5b/11/9b9a5b0306345664a2da6410877af6e8082481b5884b3ddd78d47c6013ce/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7", size = 2052418, upload-time = "2026-05-06T13:37:38.234Z" }, + { url = "https://files.pythonhosted.org/packages/f1/b7/a65fec226f5d78fc39f4a13c4cc0c768c22b113438f60c14adc9d2865038/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712", size = 2232274, upload-time = "2026-05-06T13:38:27.753Z" }, + { url = "https://files.pythonhosted.org/packages/68/f0/92039db98b907ef49269a8271f67db9cb78ae2fc68062ef7e4e77adb5f61/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4", size = 2309940, upload-time = "2026-05-06T13:38:05.353Z" }, + { url = "https://files.pythonhosted.org/packages/5f/97/2aab507d3d00ca626e8e57c1eac6a79e4e5fbcc63eb99733ff55d1717f65/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce", size = 2094516, upload-time = "2026-05-06T13:39:10.577Z" }, + { url = "https://files.pythonhosted.org/packages/22/37/a8aca44d40d737dde2bc05b3c6c07dff0de07ce6f82e9f3167aeaf4d5dea/pydantic_core-2.46.4-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987", size = 2136854, upload-time = "2026-05-06T13:40:22.59Z" }, + { url = "https://files.pythonhosted.org/packages/24/99/fcef1b79238c06a8cbec70819ac722ba76e02bc8ada9b0fd66eba40da01b/pydantic_core-2.46.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b", size = 2180306, upload-time = "2026-05-06T13:40:10.666Z" }, + { url = "https://files.pythonhosted.org/packages/ae/6c/fc44000918855b42779d007ae63b0532794739027b2f417321cddbc44f6a/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458", size = 2190044, upload-time = "2026-05-06T13:40:43.231Z" }, + { url = "https://files.pythonhosted.org/packages/6b/65/d9cadc9f1920d7a127ad2edba16c1db7916e59719285cd6c94600b0080ba/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b", size = 2329133, upload-time = "2026-05-06T13:39:57.365Z" }, + { url = "https://files.pythonhosted.org/packages/d0/cf/c873d91679f3a30bcf5e7ac280ce5573483e72295307685120d0d5ad3416/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c", size = 2374464, upload-time = "2026-05-06T13:38:06.976Z" }, + { url = "https://files.pythonhosted.org/packages/9d/1d/8987ad40f65ae1432753072f214fb5c74fe47ffbd0698bb9cbbb585664f8/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7", size = 2095527, upload-time = "2026-05-06T13:39:52.283Z" }, + { url = "https://files.pythonhosted.org/packages/64/d3/84c282a7eee1d3ac4c0377546ef5a1ea436ce26840d9ac3b7ed54a377507/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df", size = 1936024, upload-time = "2026-05-06T13:40:15.671Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ca/eac61596cdeb4d7e174d3dc0bd8a6238f14f75f97a24e7b7db4c7e7340a0/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526", size = 1990696, upload-time = "2026-05-06T13:38:34.717Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, +] + +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + +[[package]] +name = "pyrallis" +version = "0.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, + { name = "typing-inspect" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/ff/4d865cd0166b3e0b0b09b1f5399114fe27c836b74f45f580d61efecf28dd/pyrallis-0.3.1.tar.gz", hash = "sha256:ab7298f31c633d4858ec3b045a30e6cc9b9f7ab50f21ea93aa3fb28532ac4b6e", size = 33799, upload-time = "2022-03-14T08:23:25.805Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/21/896ce8acd58566f38ffa179b833fd52b4c285fd4e90842ad6b8362372639/pyrallis-0.3.1-py3-none-any.whl", hash = "sha256:632370b563486495f5f9e7caf86cddffab8351214a3bdc60ae0d23b261ebdb59", size = 33419, upload-time = "2022-03-14T08:23:24.193Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "pytz" +version = "2026.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ff/46/dd499ec9038423421951e4fad73051febaa13d2df82b4064f87af8b8c0c3/pytz-2026.2.tar.gz", hash = "sha256:0e60b47b29f21574376f218fe21abc009894a2321ea16c6754f3cad6eb7cdd6a", size = 320861, upload-time = "2026-05-04T01:35:29.667Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/dd/96da98f892250475bdf2328112d7468abdd4acc7b902b6af23f4ed958ea0/pytz-2026.2-py2.py3-none-any.whl", hash = "sha256:04156e608bee23d3792fd45c94ae47fae1036688e75032eea2e3bf0323d1f126", size = 510141, upload-time = "2026-05-04T01:35:27.408Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, +] + +[[package]] +name = "qwen-vl-utils" +version = "0.0.14" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "av" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b6/b1/ad4fc2260a3badd278b38d642f3b987412f1f6682f0ef2b31b0572d5caa8/qwen_vl_utils-0.0.14.tar.gz", hash = "sha256:9c7cad5ae803b3a10f8bb7194deb12aeacdd032f92f4224e880c73587a7346ad", size = 8453, upload-time = "2025-09-23T09:38:57.532Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/43/80f67e0336cb2fc725f8e06f7fe35c1d0fe946f4d2b8b2175e797e07349e/qwen_vl_utils-0.0.14-py3-none-any.whl", hash = "sha256:5e28657bfd031e56bd447c5901b58ddfc3835285ed100f4c56580e0ade054e96", size = 8120, upload-time = "2025-09-23T09:38:56.297Z" }, +] + +[[package]] +name = "regex" +version = "2026.7.19" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/20/98/04b13f1ddfb63158025291c02e03eb42fbb7acb51d091d541050eb4e35e8/regex-2026.7.19.tar.gz", hash = "sha256:7e77b324909c1617cbb4c668677e2c6ae13f44d7c1de0d4f15f2e3c10f3315b5", size = 416440, upload-time = "2026-07-19T00:19:48.923Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/b9/d11d7e501ac8fd7d617684423ebb9561e0b998481c1e4cbc0cb212c5d74a/regex-2026.7.19-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2cc3460cedf7579948486eab03bc9ad7089df4d7281c0f47f4afe03e8d13f02d", size = 496778, upload-time = "2026-07-19T00:17:05.677Z" }, + { url = "https://files.pythonhosted.org/packages/3f/a9/a5ab6f312f24318019170dc485d5421fe4f89e43a98640da50d95a8a7041/regex-2026.7.19-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:0e9554c8785eac5cffe6300f69a91f58ba72bc88a5f8d661235ad7c6aa5b8ccd", size = 297122, upload-time = "2026-07-19T00:17:07.59Z" }, + { url = "https://files.pythonhosted.org/packages/b3/63/4cab4d7f2d384a144d420b763d97674cb70619c878ea6fcd7640d0e62143/regex-2026.7.19-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d7da47a0f248977f08e2cb659ff3c17ddc13a4d39b3a7baa0a81bf5b415430f6", size = 292009, upload-time = "2026-07-19T00:17:09.648Z" }, + { url = "https://files.pythonhosted.org/packages/22/85/102a81b218298957d4ea7d2f084fae537a71add9d6ff93c8e67284c5f45e/regex-2026.7.19-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:93db40c8de0815baab96a06e08a984bac71f989d13bab789e382158c5d426797", size = 796708, upload-time = "2026-07-19T00:17:11.542Z" }, + { url = "https://files.pythonhosted.org/packages/78/b5/dc136af5629938a037cd2b304c12240e132ec92f38be8ff9cc89af2a1f2d/regex-2026.7.19-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:66bd62c59a5427746e8c44becae1d9b99d22fb13f30f492083dfb9ad7c45cc18", size = 865651, upload-time = "2026-07-19T00:17:13.312Z" }, + { url = "https://files.pythonhosted.org/packages/e0/75/67402ae3cd9c8c988a4c805d15ee3eef015e7ca4cb112cf3e640fc1f4153/regex-2026.7.19-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1649eb39fcc9ea80c4d2f110fde2b8ab2aef3877b98f02ab9b14e961f418c511", size = 911756, upload-time = "2026-07-19T00:17:15.015Z" }, + { url = "https://files.pythonhosted.org/packages/2a/8e/096d00c7c480ef2ff4265349b14e2261d4ab787ba1f74e2e80d1c58079c3/regex-2026.7.19-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9dce8ec9695f531a1b8a6f314fd4b393adcccf2ea861db480cdf97a301d01a68", size = 801798, upload-time = "2026-07-19T00:17:17.208Z" }, + { url = "https://files.pythonhosted.org/packages/f0/41/e7ecac6edb5722417f85cc67eaf386322fbe8acf6918ec2fdc37c20dd9d0/regex-2026.7.19-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:3080a7fd38ef049bd489e01c970c97dd84ff446a885b0f1f6b26d9b1ad13ce11", size = 776933, upload-time = "2026-07-19T00:17:19.347Z" }, + { url = "https://files.pythonhosted.org/packages/6f/69/03c9b3f058d66403e0ca2c938696e81d51cd4c6d47ec5265f02f96948d9a/regex-2026.7.19-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1d793a7988e04fcb1e2e135567443d82173225d657419ec09414a9b5a145b986", size = 784338, upload-time = "2026-07-19T00:17:21.057Z" }, + { url = "https://files.pythonhosted.org/packages/f6/f7/b38ab3d43f284afbb618fcd15d0e77eb786ae461ce1f6bc7494619ddc0f2/regex-2026.7.19-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:e8b0abe7d870f53ca5143895fef7d1041a0c831a140d3dc2c760dd7ba25d4a8b", size = 860452, upload-time = "2026-07-19T00:17:23.119Z" }, + { url = "https://files.pythonhosted.org/packages/15/5c/ff60ef0571121714f3cf9920bc183071e384a10b556d042e0fdb06cc07a5/regex-2026.7.19-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:4e5413bd5f13d3a4e3539ca98f70f75e7fca92518dd7f117f030ebedd10b60cb", size = 765958, upload-time = "2026-07-19T00:17:24.81Z" }, + { url = "https://files.pythonhosted.org/packages/aa/0f/bd34021162c0ab47f9a315bd56cd5642e920c8e5668a75ef6c6a6fca590d/regex-2026.7.19-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:73b133a9e6fb512858e7f065e96f1180aa46646bc74a83aea62f1d314f3dd035", size = 851765, upload-time = "2026-07-19T00:17:26.993Z" }, + { url = "https://files.pythonhosted.org/packages/2a/20/a2ca43edade0595cccfdc98636739f536d9e26898e7dbddc2b9e98898953/regex-2026.7.19-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:dbe6493fbd27321b1d1f2dd4f5c7e5bd4d8b1d7cab7f32fd67db3d0b2ed8248a", size = 789714, upload-time = "2026-07-19T00:17:28.699Z" }, + { url = "https://files.pythonhosted.org/packages/5d/47/e02db4015d424fc83c00ea0ac8c5e5ec14397943de9abf909d5ce3a25931/regex-2026.7.19-cp312-cp312-win32.whl", hash = "sha256:ddd67571c10869f65a5d7dde536d1e066e306cc90de57d7de4d5f34802428bb5", size = 267157, upload-time = "2026-07-19T00:17:31.051Z" }, + { url = "https://files.pythonhosted.org/packages/08/8e/c780c131f79b42ed22d1bd7da4096c2c35f813e835acd02ef0f018bd892c/regex-2026.7.19-cp312-cp312-win_amd64.whl", hash = "sha256:e30d40268a28d54ce0437031750497004c22602b8e3ab891f759b795a003b312", size = 277777, upload-time = "2026-07-19T00:17:32.848Z" }, + { url = "https://files.pythonhosted.org/packages/3e/4c/e4d7e086449bdf379d89774bf1f89dc4a41943f3c5a6125a03905b34b5fb/regex-2026.7.19-cp312-cp312-win_arm64.whl", hash = "sha256:de9208bb427130c82a5dbfd104f92c8876fc9559278c880b3002755bbbe9c83d", size = 277136, upload-time = "2026-07-19T00:17:34.803Z" }, +] + +[[package]] +name = "requests" +version = "2.34.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, +] + +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, +] + +[[package]] +name = "s3transfer" +version = "0.19.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/65/da/4bef7ce7bb989b222aa4785a413896dbec53306dfc59c6ce7d16a7ffbd6a/s3transfer-0.19.1.tar.gz", hash = "sha256:d3d6371dc3f1e5c5427b2b457bcf13bcf87bec334c95aed18642eae61f6926f3", size = 165354, upload-time = "2026-07-10T19:32:04.849Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/23/e84c64ad0e8bc59cd1b2ef98def848deff0ef3456c542afe74d51e9e8c85/s3transfer-0.19.1-py3-none-any.whl", hash = "sha256:d5fd7005ee39307455ad5f310b5ea67f4b1960d7fed5b3671ee50c249de675de", size = 90072, upload-time = "2026-07-10T19:32:03.673Z" }, +] + +[[package]] +name = "safetensors" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, +] + +[[package]] +name = "sana-wm-parity-check" +version = "0.0.0" +source = { virtual = "." } +dependencies = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "einops" }, + { name = "flash-linear-attention" }, + { name = "flashdreams" }, + { name = "flashdreams-sana-wm" }, + { name = "ftfy" }, + { name = "huggingface-hub" }, + { name = "imageio", extra = ["ffmpeg"] }, + { name = "numpy" }, + { name = "omegaconf" }, + { name = "opencv-python-headless" }, + { name = "pillow" }, + { name = "protobuf" }, + { name = "pyrallis" }, + { name = "pytz" }, + { name = "pyyaml" }, + { name = "qwen-vl-utils" }, + { name = "safetensors" }, + { name = "scipy" }, + { name = "sentencepiece" }, + { name = "termcolor" }, + { name = "timm" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, + { name = "torchvision" }, + { name = "tqdm" }, + { name = "transformer-engine", extra = ["core-cu13", "pytorch"], marker = "sys_platform != 'win32'" }, + { name = "transformers" }, + { name = "triton", marker = "sys_platform == 'linux'" }, +] + +[package.metadata] +requires-dist = [ + { name = "accelerate", specifier = ">=1.3" }, + { name = "diffusers", specifier = ">=0.37" }, + { name = "einops", specifier = ">=0.7" }, + { name = "flash-linear-attention", specifier = ">=0.4.2" }, + { name = "flashdreams", editable = "../../../../flashdreams" }, + { name = "flashdreams-sana-wm", editable = "../../" }, + { name = "ftfy", specifier = ">=6.0" }, + { name = "huggingface-hub", specifier = ">=0.36" }, + { name = "imageio", extras = ["ffmpeg"], specifier = ">=2.31" }, + { name = "numpy", specifier = ">=1.24,<2.5" }, + { name = "omegaconf", specifier = ">=2.3" }, + { name = "opencv-python-headless", specifier = ">=4.8" }, + { name = "pillow", specifier = ">=10" }, + { name = "protobuf", specifier = ">=7.35.0,<8" }, + { name = "pyrallis", specifier = ">=0.3" }, + { name = "pytz", specifier = ">=2024.0" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "qwen-vl-utils", specifier = ">=0.0.8" }, + { name = "safetensors", specifier = ">=0.5" }, + { name = "scipy", specifier = ">=1.11" }, + { name = "sentencepiece", specifier = ">=0.2" }, + { name = "termcolor", specifier = ">=2.4" }, + { name = "timm", specifier = ">=0.6.13" }, + { name = "torch", specifier = ">=2.11" }, + { name = "torchvision", specifier = ">=0.26" }, + { name = "tqdm", specifier = ">=4.60" }, + { name = "transformer-engine", extras = ["pytorch", "core-cu13"], marker = "sys_platform != 'win32'", specifier = ">=2.12" }, + { name = "transformers", specifier = ">=5.0,<6" }, + { name = "triton", marker = "sys_platform == 'linux'", specifier = ">=3.6" }, +] + +[[package]] +name = "scipy" +version = "1.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/19/ca10ead60b0acc80b2b833c2c4a4f2ff753d0f58b811f70d911c7e94a25c/scipy-1.18.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:7bd21faaf5a1a3b2eff922d02db5f191b99a6518db9078a8fb23169f6d22259a", size = 31056519, upload-time = "2026-06-19T14:59:45.203Z" }, + { url = "https://files.pythonhosted.org/packages/96/72/1e6442a00cd2924d361aa1b642ab6373ec35c6fabf311a760be9f76e0f13/scipy-1.18.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:265915e79107de9f946b855e50d7470d5893ec3f54b342e1aa6201cbdcd8bb6b", size = 28681889, upload-time = "2026-06-19T14:59:48.103Z" }, + { url = "https://files.pythonhosted.org/packages/9b/2d/11dd93d21e147a73ba22bd75c0b9208d3a2e0ec76d53170ce7d9029b1015/scipy-1.18.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:9ab7b758be6940954a713ee466e2043e9f6e2ed965c1fce5c91039f4be3d90a9", size = 20423580, upload-time = "2026-06-19T14:59:50.665Z" }, + { url = "https://files.pythonhosted.org/packages/9c/01/93552f75e0d2a7dd115a45e59209c51e8d514daff02fc887d2623be06fe1/scipy-1.18.0-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:97b6cddaaee0a779ef6b5ca83c9604b27cc16b2b8fc22c142652df8793319fb8", size = 23054441, upload-time = "2026-06-19T14:59:53.564Z" }, + { url = "https://files.pythonhosted.org/packages/3c/23/21f5e703643d66f21faa6b4c73195bfcad70c55efcb4f1ab327cd7c4101a/scipy-1.18.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:52a96e21517c7292375c0e27dd796a811f03fcea5fd4d108fdfea8145dcf17ab", size = 33968720, upload-time = "2026-06-19T14:59:56.415Z" }, + { url = "https://files.pythonhosted.org/packages/dd/aa/1b939f6c67ed68635bb538e6752d3dacc02f66535182e939a89581a44e9c/scipy-1.18.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1f55797419e16e7f30cf88ffb3113ce0467f00cfe3f70d5c281730b21769bfc2", size = 35287115, upload-time = "2026-06-19T14:59:59.411Z" }, + { url = "https://files.pythonhosted.org/packages/b6/ff/eec46be7e9234208f801062b53e1983085eddebd693f6c9bfb03b459830d/scipy-1.18.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ad033410e2e0672ffdc1042110cef20e1c46f8fd0616cee1d44d8d58fad8fc11", size = 35577989, upload-time = "2026-06-19T15:00:02.235Z" }, + { url = "https://files.pythonhosted.org/packages/84/ca/210d4759c7210bb7d269437421959b39a33434e2776b60c5cb8a763bb30a/scipy-1.18.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4a55985d54c769c872e64b7f4c8a81cc30ef700cc04296abbbf3705439c126de", size = 37421717, upload-time = "2026-06-19T15:00:05.102Z" }, + { url = "https://files.pythonhosted.org/packages/2b/54/9a9edb45345bd6744da5ddfb6628e5d5185920494c6a67ec45b6381004cb/scipy-1.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:71ccc8faa2dd16ac310233203474a8b5cb67f10dedd54a3116d34943f4b19132", size = 36597428, upload-time = "2026-06-19T15:00:08.112Z" }, + { url = "https://files.pythonhosted.org/packages/99/0e/33f32a2a58987e26aec0f7df252cbbad1e90ae77bdbc76f40dd4ed0cf0ea/scipy-1.18.0-cp312-cp312-win_arm64.whl", hash = "sha256:d88363fd9d8fbd3511bd273f1a49efb2a540773ddf92a91d57498ce7dd7f3e76", size = 24351481, upload-time = "2026-06-19T15:00:11.103Z" }, +] + +[[package]] +name = "sentencepiece" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/33/ea3cb3839607eb175da835244a798f797f478c5ddf0e8ecdf57ea85a4c70/sentencepiece-0.2.2.tar.gz", hash = "sha256:3d2b5e824b5622038dc7b490897efe05ebbbb9e7350fc142f3ecc8789ef9bdf6", size = 8218435, upload-time = "2026-07-12T08:39:34.701Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/13/7a562289c8d5b49ebdf3f9c1e8ab67cf14a8743b1d90c8f406bfdec36b72/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:1edb10e520e4bddf74d85b0f5ae74cc2d60c2b448885080bfb618bc2b3a49f6b", size = 2188384, upload-time = "2026-07-12T08:38:28.486Z" }, + { url = "https://files.pythonhosted.org/packages/85/d1/912f14fd5eae168aba726ffb6a9a2dc1c71fe7676c53da6f5c442b886d4a/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f7c06c751c19d923435a54bff4f7e66e728fad160e8da28254f133abc9725820", size = 1441553, upload-time = "2026-07-12T08:38:30.552Z" }, + { url = "https://files.pythonhosted.org/packages/bd/44/caa9cab5f261a019e2808bc5046152775dc57352ba9cbae7525e9e7a1ed4/sentencepiece-0.2.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38111ed1f79268f399c505028023d5eaaf0ab4e5eafceb709468b0d3323e7838", size = 1347176, upload-time = "2026-07-12T08:38:32.211Z" }, + { url = "https://files.pythonhosted.org/packages/19/90/cd798935668cff71d309d8ff10385844ecf216b1fe454f1993ed8bf2cb91/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cbce24284f51f71d10a42b7b9c964dcb9048b28f1c8e5db40bcbcb6f428cba6a", size = 1325200, upload-time = "2026-07-12T08:38:33.689Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2d/37e3da037318a70066ded0d51bc2a7f35491ae6338dd993d5eb1503fc3b5/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c8a168b040bc61681293f79a949b5d911c8e25086f4260285b8d97ab5f1195da", size = 1397736, upload-time = "2026-07-12T08:38:35.771Z" }, + { url = "https://files.pythonhosted.org/packages/8d/11/753fca2e6b109be3ab7867abf357dfe48677fe726ae5a5363d0b54ca9450/sentencepiece-0.2.2-cp312-cp312-win_amd64.whl", hash = "sha256:7c6e7bf684dc12145bfa685d3060beaea55139134ba848289bee514ed42e7383", size = 1248030, upload-time = "2026-07-12T08:38:37.604Z" }, + { url = "https://files.pythonhosted.org/packages/e2/0a/70efbe861ca182d7d4b6e1a20f58e043400848fa9f2915229f082e221648/sentencepiece-0.2.2-cp312-cp312-win_arm64.whl", hash = "sha256:76ff5814db72e7462dece042d7593cdf102b8ec82c2b1cc201a2add34ee3050d", size = 1187325, upload-time = "2026-07-12T08:38:39.348Z" }, +] + +[[package]] +name = "setuptools" +version = "83.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, +] + +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "termcolor" +version = "3.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/79/cf31d7a93a8fdc6aa0fbb665be84426a8c5a557d9240b6239e9e11e35fc5/termcolor-3.3.0.tar.gz", hash = "sha256:348871ca648ec6a9a983a13ab626c0acce02f515b9e1983332b17af7979521c5", size = 14434, upload-time = "2025-12-29T12:55:21.882Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl", hash = "sha256:cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5", size = 7734, upload-time = "2025-12-29T12:55:20.718Z" }, +] + +[[package]] +name = "timm" +version = "1.0.28" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, + { name = "torchvision" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/03/e41389ac641747bfec48d016fde8be1eade1901e6f2c1aedcb0c8cb4b5d9/timm-1.0.28.tar.gz", hash = "sha256:3789d313fdd5541a327b60180d70dbb4bdec73db8ff0655e413db3c3d134a9a4", size = 2451413, upload-time = "2026-07-11T17:24:32.615Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/76/de1bfac17d183c49c6d0887903d3064ced51cf1d9ba7a8d611c1a8808c4f/timm-1.0.28-py3-none-any.whl", hash = "sha256:e577b88da96b3a722ea5e2f042455ce6f715d398304d8e63b17d126ed7d89968", size = 2597944, upload-time = "2026-07-11T17:24:30.869Z" }, +] + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/73/6f/f80cfef4a312e1fb34baf7d85c72d4411afde10978d4657f8cdd811d3ccc/tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917", size = 372115, upload-time = "2026-01-05T10:45:15.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/97/5dbfabf04c7e348e655e907ed27913e03db0923abb5dfdd120d7b25630e1/tokenizers-0.22.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c", size = 3100275, upload-time = "2026-01-05T10:41:02.158Z" }, + { url = "https://files.pythonhosted.org/packages/2e/47/174dca0502ef88b28f1c9e06b73ce33500eedfac7a7692108aec220464e7/tokenizers-0.22.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001", size = 2981472, upload-time = "2026-01-05T10:41:00.276Z" }, + { url = "https://files.pythonhosted.org/packages/d6/84/7990e799f1309a8b87af6b948f31edaa12a3ed22d11b352eaf4f4b2e5753/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7", size = 3290736, upload-time = "2026-01-05T10:40:32.165Z" }, + { url = "https://files.pythonhosted.org/packages/78/59/09d0d9ba94dcd5f4f1368d4858d24546b4bdc0231c2354aa31d6199f0399/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd", size = 3168835, upload-time = "2026-01-05T10:40:38.847Z" }, + { url = "https://files.pythonhosted.org/packages/47/50/b3ebb4243e7160bda8d34b731e54dd8ab8b133e50775872e7a434e524c28/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5", size = 3521673, upload-time = "2026-01-05T10:40:56.614Z" }, + { url = "https://files.pythonhosted.org/packages/e0/fa/89f4cb9e08df770b57adb96f8cbb7e22695a4cb6c2bd5f0c4f0ebcf33b66/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e", size = 3724818, upload-time = "2026-01-05T10:40:44.507Z" }, + { url = "https://files.pythonhosted.org/packages/64/04/ca2363f0bfbe3b3d36e95bf67e56a4c88c8e3362b658e616d1ac185d47f2/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b", size = 3379195, upload-time = "2026-01-05T10:40:51.139Z" }, + { url = "https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67", size = 3274982, upload-time = "2026-01-05T10:40:58.331Z" }, + { url = "https://files.pythonhosted.org/packages/1d/28/5f9f5a4cc211b69e89420980e483831bcc29dade307955cc9dc858a40f01/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4", size = 9478245, upload-time = "2026-01-05T10:41:04.053Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fb/66e2da4704d6aadebf8cb39f1d6d1957df667ab24cff2326b77cda0dcb85/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a", size = 9560069, upload-time = "2026-01-05T10:45:10.673Z" }, + { url = "https://files.pythonhosted.org/packages/16/04/fed398b05caa87ce9b1a1bb5166645e38196081b225059a6edaff6440fac/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a", size = 9899263, upload-time = "2026-01-05T10:45:12.559Z" }, + { url = "https://files.pythonhosted.org/packages/05/a1/d62dfe7376beaaf1394917e0f8e93ee5f67fea8fcf4107501db35996586b/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5", size = 10033429, upload-time = "2026-01-05T10:45:14.333Z" }, + { url = "https://files.pythonhosted.org/packages/fd/18/a545c4ea42af3df6effd7d13d250ba77a0a86fb20393143bbb9a92e434d4/tokenizers-0.22.2-cp39-abi3-win32.whl", hash = "sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92", size = 2502363, upload-time = "2026-01-05T10:45:20.593Z" }, + { url = "https://files.pythonhosted.org/packages/65/71/0670843133a43d43070abeb1949abfdef12a86d490bea9cd9e18e37c5ff7/tokenizers-0.22.2-cp39-abi3-win_amd64.whl", hash = "sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48", size = 2747786, upload-time = "2026-01-05T10:45:18.411Z" }, + { url = "https://files.pythonhosted.org/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, +] + +[[package]] +name = "torch" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform != 'win32'", +] +dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/3a/ed0f4d4d1dcde03bced7aac9a28e800abcdc0cbd06b6775044c9fbd877b7/torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027", size = 111213045, upload-time = "2026-07-08T16:05:22.997Z" }, + { url = "https://files.pythonhosted.org/packages/df/a9/f6a2a4d763ff1df02e9a64c477029db614295bc9367f4131223791ccc243/torch-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:572df8be8ffb4599c88cbd6a0726f1f854f4da65d2e3c09f0e2c2283333cd6d4", size = 427210998, upload-time = "2026-07-08T16:04:37.708Z" }, + { url = "https://files.pythonhosted.org/packages/f3/82/fea946351658e6534db52d2cc12bc53087cbf87f9440c5f180f367c1950b/torch-2.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:796633c4cdf0fe2cdced72d8f88f22e73dbcfce83132763162f6d4bff13b820b", size = 526605292, upload-time = "2026-07-08T16:04:22.81Z" }, +] + +[[package]] +name = "torch" +version = "2.13.0+cu130" +source = { registry = "https://download.pytorch.org/whl/cu130" } +resolution-markers = [ + "sys_platform == 'win32'", +] +dependencies = [ + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.13.0%2Bcu130-cp312-cp312-win_amd64.whl", hash = "sha256:2efab1e83604ca628c6d85b9e188c153690980498d1297081a9dad704919303c", upload-time = "2026-07-08T20:26:27Z" }, +] + +[[package]] +name = "torchvision" +version = "0.28.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/49/c1cab1ecbb3ff1a380a3f99283db1dee61b8afe354f6352c643b65937130/torchvision-0.28.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:e9f54c30cd52e3ef7fd034cc69b7bb7e0964e1c8f8743e018ab92e95b40f9eee", size = 1856020, upload-time = "2026-07-08T16:07:52.182Z" }, + { url = "https://files.pythonhosted.org/packages/f0/4c/95233776e2def960e5abb7a07931230a545f43717a56a1e1140162033598/torchvision-0.28.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5cf78ebc401ce64ae19b8c55de866bb836797d559a4de9c25ccbe74cfa642d3a", size = 7842127, upload-time = "2026-07-08T16:07:53.446Z" }, + { url = "https://files.pythonhosted.org/packages/93/e4/e9b2495d0d57b9f60d63c57d0a910410a81b4b073bf70917bef815291119/torchvision-0.28.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:028a3d481b37d785605620d7cdad897064c5a55bae2aa1f2658766333e291940", size = 7675040, upload-time = "2026-07-08T16:07:58.017Z" }, + { url = "https://files.pythonhosted.org/packages/7c/9c/55ed9cb6dfe3ee9c837df5cd0e758372e5829aa38b8dd71343aa632cc4e2/torchvision-0.28.0-cp312-cp312-win_amd64.whl", hash = "sha256:87dc16b2df427c1318ad335f1e2be2b3b15b2cf20f7934c83b0505a48425ee5d", size = 4085785, upload-time = "2026-07-08T16:07:50.928Z" }, +] + +[[package]] +name = "tqdm" +version = "4.69.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8c/69/40407dfc835517f058b603dbf37a6df094d8582b015a51eddc988febbcb7/tqdm-4.69.0.tar.gz", hash = "sha256:700c5e85dcd5f009dd6222588a29180a193a748247a5d855b4d67db93d79a53b", size = 792569, upload-time = "2026-07-17T18:09:06.2Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/21/99a0cdaf54eb35e77623c41b5a2c9472ee4404bba687052791fe2aba6773/tqdm-4.69.0-py3-none-any.whl", hash = "sha256:9979978912be667a6ef21fd5d8abf54e324e63d82f7f43c360792ebc2bc4e622", size = 676680, upload-time = "2026-07-17T18:09:04.172Z" }, +] + +[[package]] +name = "transformer-engine" +version = "2.17.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/26/e501a474a3f36c2d439561fa4f2094a3c0344df439500dda55937a8fccf7/transformer_engine-2.17.0-py3-none-any.whl", hash = "sha256:c274fd74ea2e4caa7132921e1ecfa3847931abbed9f6b8cab61346cdb833bc76", size = 1001710, upload-time = "2026-07-09T00:36:15.453Z" }, +] + +[package.optional-dependencies] +core-cu13 = [ + { name = "transformer-engine-cu13" }, +] +pytorch = [ + { name = "transformer-engine-torch" }, +] + +[[package]] +name = "transformer-engine-cu13" +version = "2.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "importlib-metadata" }, + { name = "packaging" }, + { name = "pydantic" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/21/5e0ec562539798ffd375c218b8bd7612b387455ccc79ae14576f411c9609/transformer_engine_cu13-2.17.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:1c8e6627cf02358201f513ab902c7b4f82f179214bd775edc97ce0f2001d82da", size = 244570183, upload-time = "2026-07-09T00:36:35.507Z" }, +] + +[[package]] +name = "transformer-engine-torch" +version = "2.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "einops" }, + { name = "nvdlfw-inspect" }, + { name = "onnx" }, + { name = "onnxscript" }, + { name = "packaging" }, + { name = "pydantic" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" } }, + { name = "transformer-engine-cu13" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/f1/ad613157261ac71262f5f260acf07fe22ae5ee0bb0675622c4f19a3e8afb/transformer_engine_torch-2.17.0.tar.gz", hash = "sha256:52ca109cdbc987ca02f87735fa375da8a02e117acb87e25dcb7e27dca37aa87f", size = 369712, upload-time = "2026-07-09T00:36:28.065Z" } + +[[package]] +name = "transformers" +version = "5.14.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, + { name = "typer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/fb/2a2ba88f325e68a921d8b69ff63b477830b2e73ade9a3c8c8cab2f06d741/transformers-5.14.1.tar.gz", hash = "sha256:60d196c27781eacf8637e2b533f517582907ad6f9ae142046d6b69431a5b2173", size = 9295927, upload-time = "2026-07-16T09:41:57.773Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/67/8d85ca2323233ae3c0365a659c4e52ee1f587b440e4bc577e7d8e4416d0f/transformers-5.14.1-py3-none-any.whl", hash = "sha256:9db974c4079ede2d1a3ea7ca5a240df33f2cc26fc2b36ba64c5f2a4f43b6e725", size = 11625234, upload-time = "2026-07-16T09:41:54.143Z" }, +] + +[[package]] +name = "triton" +version = "3.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/fa/f856e24deb462d5f18bd4b5a746957862ab9b6ee5834bda60605ec348366/triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1", size = 184692359, upload-time = "2026-06-17T20:03:48.288Z" }, + { url = "https://files.pythonhosted.org/packages/c4/6f/fb96d15db6f36d6eae4cafb998c2e0353bf59d7c4ea1662d7497f269134a/triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728", size = 197719725, upload-time = "2026-06-17T19:53:20.419Z" }, +] + +[[package]] +name = "triton-windows" +version = "3.7.1.post27" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/76/30/325b420efd0047e119679c646a9a410db216069800ec009fae3da26c69a3/triton_windows-3.7.1.post27-cp312-cp312-win_amd64.whl", hash = "sha256:f5406230d7dbf6965bc4051fcad27b81c39ba4a4bfde06f494dd7ff4eb325a9e", size = 49683004, upload-time = "2026-06-21T16:48:14.02Z" }, +] + +[[package]] +name = "typeguard" +version = "4.5.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/67/1c/dfba5c4633cafc4c701f237d2ba63b416805047fd6d96aab4cfc40969f98/typeguard-4.5.2.tar.gz", hash = "sha256:5a16dcac23502039299c97c8941651bc33d7ea8cc4b2f7d6bbb1b528f6eea423", size = 80240, upload-time = "2026-05-14T12:59:40.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/29/74eeb4d3f3ae61ca096b018ad486b3b3c74b17bec09ab4edab721cbefec3/typeguard-4.5.2-py3-none-any.whl", hash = "sha256:fcf9de18bd945cdb4c7b996e12b4c51ce83f92f191314a6d7cf1739586ec98cf", size = 36748, upload-time = "2026-05-14T12:59:39.473Z" }, +] + +[[package]] +name = "typer" +version = "0.27.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "rich" }, + { name = "shellingham" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/37/78/fda3361b56efc27944f24225f6ecd13d96d6fcfe37bd0eb34e2f4c63f9fc/typer-0.27.0.tar.gz", hash = "sha256:629bd12ea5d13a17148125d9a264f949eb171fb3f120f9b04d85873cab054fa5", size = 203430, upload-time = "2026-07-15T19:21:07.007Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/40/03/26a383c9e58c213199d1aad1c3d353cfc22d4444ec6d2c0bf8ad02523843/typer-0.27.0-py3-none-any.whl", hash = "sha256:6f4b27631e47f077871b7dc30e933ec0131c1390fbe0e387ea5574b5bac9ccf1", size = 122716, upload-time = "2026-07-15T19:21:05.553Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "typing-inspect" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mypy-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/74/1789779d91f1961fa9438e9a8710cdae6bd138c80d7303996933d117264a/typing_inspect-0.9.0.tar.gz", hash = "sha256:b23fc42ff6f6ef6954e4852c1fb512cdd18dbea03134f91f856a95ccc9461f78", size = 13825, upload-time = "2023-05-24T20:25:47.612Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827, upload-time = "2023-05-24T20:25:45.287Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "tyro" +version = "1.0.15" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "docstring-parser" }, + { name = "typeguard" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/78/a5749a6c1ee9abc2999e294f339f8f72476d1a60bb95fc0e86156aafed3b/tyro-1.0.15.tar.gz", hash = "sha256:3f1d60887723eecb9c489f195d11f079c4a1f33df74b723552ad31ec57c667bb", size = 593822, upload-time = "2026-06-20T08:48:28.364Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/28/d607636187cf6c18eb72efb5d65c1d1b9e451db03d84676f835dd488fdbd/tyro-1.0.15-py3-none-any.whl", hash = "sha256:982da1d566005f1b2a6b56f6be6c6929c96f0e5fab9d61bc6097573ddfaf8d13", size = 215045, upload-time = "2026-06-20T08:48:27.104Z" }, +] + +[[package]] +name = "urllib3" +version = "2.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, +] + +[[package]] +name = "wcwidth" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/34/74/c6428f875774288bec1396f5bfcbc2d925700a4dad61727fd5f2b12f249d/wcwidth-0.8.2.tar.gz", hash = "sha256:91fbef97204b96a3d4d421609b80340b760cf33e26da123ff243d76b1fda8dda", size = 1466253, upload-time = "2026-06-29T18:11:11.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/42/3e5985a0a7e57de470b320c6d6a1a67c844f6737a587f3d44dd13d1819e7/wcwidth-0.8.2-py3-none-any.whl", hash = "sha256:d63947694a0539a1d51e01eda7caf800c291020e6cdd7e28ad7b14dd33ad4f85", size = 323166, upload-time = "2026-06-29T18:11:09.888Z" }, +] + +[[package]] +name = "win32-setctime" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b3/8f/705086c9d734d3b663af0e9bb3d4de6578d08f46b1b101c2442fd9aecaa2/win32_setctime-1.2.0.tar.gz", hash = "sha256:ae1fdf948f5640aae05c511ade119313fb6a30d7eabe25fef9764dca5873c4c0", size = 4867, upload-time = "2024-12-07T15:28:28.314Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/07/c6fe3ad3e685340704d314d765b7912993bcb8dc198f0e7a89382d37974b/win32_setctime-1.2.0-py3-none-any.whl", hash = "sha256:95d644c4e708aba81dc3704a116d8cbc974d70b3bdb8be1d150e36be6e9d1390", size = 4083, upload-time = "2024-12-07T15:28:26.465Z" }, +] + +[[package]] +name = "zipp" +version = "4.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" }, +] From 69c4daad2342842f305481f27d059e129e324b3d Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 23 Jul 2026 17:07:47 -0700 Subject: [PATCH 23/36] Remove local docs --- .../parity_check/BENCHMARK_FINALIZATION.md | 61 ----------------- .../sana/tests/parity_check/PERF_HANDOFF.md | 68 ------------------- 2 files changed, 129 deletions(-) delete mode 100644 integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md delete mode 100644 integrations/sana/tests/parity_check/PERF_HANDOFF.md diff --git a/integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md b/integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md deleted file mode 100644 index d6c6f2654..000000000 --- a/integrations/sana/tests/parity_check/BENCHMARK_FINALIZATION.md +++ /dev/null @@ -1,61 +0,0 @@ - - -# SANA-WM Benchmark Finalization - -Use this after the perf optimization pass is done and the final BF16, FP8, and -FP4 numbers are ready for the model card. - -1. Run the precision sweep: - - ```bash - cd integrations/sana/tests/parity_check - DEVICE_LABEL="" BENCH_PRECISIONS=bf16,fp8,fp4 bash bench.sh - ``` - -2. Copy `outputs/bench/perf.md` to - `docs/source/_static/performance/sana_wm/perf-.md`. - -3. Add a `Profiling benchmark` section to `docs/source/models/sana_wm.rst` - between `What to expect` and `Citation`. - -4. The section should describe post-load generation latency per generated frame - for FlashDreams SANA-WM versus the official `NVlabs/Sana` implementation - under matched BF16, FP8, and FP4 settings. - -5. Use the standard benchmark chart block: - - ```rst - Profiling benchmark - ------------------- - - Here is the profiling benchmark on post-load generation latency per generated - frame for FlashDreams SANA-WM compared to the - `official SANA-WM implementation `_ under - matched BF16, FP8, and FP4 settings. - - .. raw:: html - -
-
-
-

- This chart shows post-load generation latency per generated frame in milliseconds on a single GPU. - For the official SANA-WM implementation, see - this instruction. -

-
-
- - ``` - -Keep the chart metric unchanged: post-load generation latency per generated -frame, in ms/frame. diff --git a/integrations/sana/tests/parity_check/PERF_HANDOFF.md b/integrations/sana/tests/parity_check/PERF_HANDOFF.md deleted file mode 100644 index aaed1d654..000000000 --- a/integrations/sana/tests/parity_check/PERF_HANDOFF.md +++ /dev/null @@ -1,68 +0,0 @@ - - -# SANA-WM Perf Handoff - -This note is only for the next SANA-WM performance investigation. It should not -be copied into the model card or public README. - -## Observed Delta - -The best current BF16, no-refiner, 121-frame, 60-step run still has -FlashDreams slower than upstream: - -| artifact | official | FlashDreams | gap | -| --- | ---: | ---: | ---: | -| `outputs/bf16_step60_bidirectional_default/bench.md` | 640.89 ms/frame | 663.60 ms/frame | +22.71 ms/frame | - -Component medians from that run: - -| component | official | FlashDreams | -| --- | ---: | ---: | -| wall | 77.55 s | 80.30 s | -| Stage-1 DiT | 75,565.73 ms | 78,237.04 ms | -| VAE decode row | 1,340.43 ms | 1,584.82 ms | -| peak memory | 14.37 GiB | 17.98 GiB | - -The main target is the 60-step Stage-1 DiT path. The decode row also still has -a smaller gap. - -## Perf Theories Tested - -| theory | result | evidence | -| --- | --- | --- | -| FlashDreams was slower because VAE tiling used low-memory tiles. | Partly right. | Early probes showed FlashDreams VAE decode around 5.2 s. Switching to upstream-style tiles dropped it to about 1.5 s. This fixed a large initial problem, but not the remaining 60-step gap. | -| Remaining VAE gap is mostly raw VAE forward kernels. | Probably wrong. | Later internal VAE timing was close to upstream, while the harness decode row remained slower. The remaining cost is likely after the VAE forward: clamp, permute, CPU transfer, numpy materialization, or `torch.cuda.empty_cache()`. | -| Matching upstream's bidirectional chunk default would close the gap. | Correct for config parity, insufficient for speed. | `SanaWMStage1Spec.chunk_size=None` with `chunk_split_strategy="first_chunk_plus_one"` matches upstream. The 8-step probe became slightly favorable to FlashDreams, but the 60-step run stayed slower. | -| Softmax camera Q/K/V transforms staying in higher precision were slowing blocks. | Right. | After casting transformed softmax-camera Q/K/V back to BF16, targeted one-step block timings dropped from roughly 106-107 ms to roughly 49-52 ms for those blocks. This change is kept. | -| Forcing cuDNN SDPA would make the default comparison faster. | Wrong. | The forced-cuDNN probe was much slower and used about 80 GiB. Keep `FORCE_CUDNN_SDPA=1` only as a backend-isolation probe. | -| `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` helps this path. | Unproven. | No run showed a speed or memory benefit for SANA-WM. The harness and docs do not set allocator overrides. | -| Direct SDPA on head dim 112 is faster than padding to 128. | Not supported by clean evidence. | The direct-D112 probe had suspicious backend and memory behavior. Current code keeps padding to 128. | -| `torch.inference_mode()` could reduce framework overhead versus `no_grad`. | Untested. | Upstream wraps generation with `@torch.inference_mode()`. FlashDreams framework generation uses `@torch.no_grad()`. A harness-only test was started, then reverted before measurement. | - -## Next Probes - -1. Reproduce the 60-step BF16 no-refiner benchmark once before changing code, - using `bench.sh`, to confirm the current gap on the current branch. -2. Profile Stage-1 DiT without using profiler-perturbed timings as headline - numbers. Focus on block-level attention, cross-attention, FFN, camera QKV, - and Python/framework overhead. -3. Test `torch.inference_mode()` around the FlashDreams generation path and - compare against the same command with `no_grad`. -4. Split FlashDreams decode timing into VAE forward, clamp, permute, CPU copy, - numpy conversion, and `empty_cache`. -5. Test disabling or configuring `torch.cuda.empty_cache()` after decode. The - optimization goal allows spending more memory for speed. - -## Local Provenance - -These paths are under ignored `outputs/` directories and may not exist in a -fresh checkout: - -- `outputs/bf16_step60_bidirectional_default/bench.md` -- `outputs/bf16_step8_bidirectional_default/bench.md` -- `outputs/bf16_step60_softmax_bf16/bench.md` -- `outputs/bench_121_noalloc/bench.md` -- `outputs/fd_profile_step1_softmax_bf16/stats.json` From 8a3ef240a70f63fb9f18a90cded0f65c52001a6e Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Thu, 23 Jul 2026 17:15:55 -0700 Subject: [PATCH 24/36] More benchmarking scripts, missing from prev commits --- .../tests/test_parity_benchmark_summary.py | 242 ++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 integrations/sana/tests/test_parity_benchmark_summary.py diff --git a/integrations/sana/tests/test_parity_benchmark_summary.py b/integrations/sana/tests/test_parity_benchmark_summary.py new file mode 100644 index 000000000..cebdc2d5f --- /dev/null +++ b/integrations/sana/tests/test_parity_benchmark_summary.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""CPU-safe tests for the SANA-WM parity benchmark summary.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import ModuleType + +import pytest + +pytestmark = pytest.mark.ci_cpu + + +def _load_bench_summary() -> ModuleType: + path = Path("integrations/sana/tests/parity_check/bench_summary.py") + spec = importlib.util.spec_from_file_location("sana_wm_bench_summary", path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _load_bench_sweep_summary() -> ModuleType: + path = Path("integrations/sana/tests/parity_check/bench_sweep_summary.py") + spec = importlib.util.spec_from_file_location("sana_wm_bench_sweep_summary", path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_benchmark_summary_uses_generation_ms_per_frame_for_chart() -> None: + module = _load_bench_summary() + summary = { + "inputs": { + "image_path": "image.png", + "prompt_path": "prompt.txt", + "camera_path": "pose.npy", + "intrinsics_path": "intrinsics.npy", + "num_frames": 20, + "seed": 42, + "no_refiner": True, + "stage1_precision": "fp8", + "refiner_precision": "fp8", + "quant_backend": "torch-fp8", + "compile_stage1": False, + "force_cudnn_sdpa": True, + "warmup_runs": 1, + }, + "upstream": { + "runs_measured": 3, + "wall_median_s": 2.0, + "wall_p90_s": 2.2, + "stage1_total_median_ms": 1600.0, + "dit_median_ms": 1500.0, + "vae_decode_median_ms": 300.0, + "mem_peak_median_gib": 10.0, + }, + "flashdreams": { + "runs_measured": 3, + "wall_median_s": 1.5, + "wall_p90_s": 1.8, + "encode_median_ms": 100.0, + "dit_median_ms": 1100.0, + "vae_decode_median_ms": 250.0, + "mem_peak_median_gib": 12.0, + }, + } + + report = module._render_markdown(summary) + chart = module._render_chart_markdown(summary, "Test GPU") + + assert "## Benchmark metric" in report + assert "- stage1_precision: `fp8`" in report + assert "generation median / frame | 100.00 ms | 75.00 ms" in report + assert "## Timing breakdown" in report + assert "| Stage-1 DiT median | 1500.00 ms | 1100.00 ms |" in report + assert "Stage-1 DiT metric" not in report + assert "Generation after model load" not in report + assert chart == ( + "# SANA-WM Benchmark Data (ms/frame)\n" + "\n" + "| device | official | flashdreams |\n" + "| --- | ---: | ---: |\n" + "| Test GPU | 100.00 | 75.00 |\n" + ) + + +def test_benchmark_summary_writes_chart_data(tmp_path: Path) -> None: + module = _load_bench_summary() + upstream = tmp_path / "upstream" / "run_0" + flashdreams = tmp_path / "flashdreams" / "run_0" + upstream.mkdir(parents=True) + flashdreams.mkdir(parents=True) + (upstream / "stats.json").write_text( + json.dumps( + { + "wall_s": 4.0, + "mem_peak_gib": 10.0, + "timings_s": { + "stage1_sample_s": 3.5, + "stage1_dit_s": 3.0, + "vae_decode_s": 0.5, + }, + } + ), + encoding="utf-8", + ) + (flashdreams / "stats.json").write_text( + json.dumps( + { + "wall_s": 2.0, + "stats_ms": { + "encode_ms": 100.0, + "diffuse_ms": 1500.0, + "decode_ms": 400.0, + "mem_peak_gib": 12.0, + }, + } + ), + encoding="utf-8", + ) + out_json = tmp_path / "bench.json" + out_md = tmp_path / "bench.md" + out_chart = tmp_path / "perf.md" + + module.main( + [ + "--upstream-dir", + str(tmp_path / "upstream"), + "--flashdreams-dir", + str(tmp_path / "flashdreams"), + "--warmup-runs", + "0", + "--image-path", + "image.png", + "--prompt-path", + "prompt.txt", + "--camera-path", + "pose.npy", + "--intrinsics-path", + "intrinsics.npy", + "--num-frames", + "40", + "--seed", + "42", + "--device-label", + "Test GPU", + "--chart-label", + "BF16", + "--stage1-precision", + "bf16", + "--refiner-precision", + "bf16", + "--output-json", + str(out_json), + "--output-md", + str(out_md), + "--output-chart-md", + str(out_chart), + ] + ) + + assert "| BF16 | 100.00 | 50.00 |" in out_chart.read_text(encoding="utf-8") + assert "generation median / frame | 100.00 ms | 50.00 ms" in out_md.read_text( + encoding="utf-8" + ) + summary = json.loads(out_json.read_text(encoding="utf-8")) + assert summary["benchmark"] == { + "metric": "generation_ms_per_frame", + "unit": "ms/frame", + "timing_boundary": ( + "pipeline.generate after model setup; excludes model construction, " + "checkpoint loading, video writing, and frame dumps" + ), + "device_label": "Test GPU", + "chart_label": "BF16", + "official": 100.0, + "flashdreams": 50.0, + } + + +def test_benchmark_sweep_summary_writes_precision_chart_data(tmp_path: Path) -> None: + module = _load_bench_sweep_summary() + bf16_json = tmp_path / "bf16.json" + fp8_json = tmp_path / "fp8.json" + bf16_json.write_text( + json.dumps({"benchmark": {"official": 100.0, "flashdreams": 90.0}}), + encoding="utf-8", + ) + fp8_json.write_text( + json.dumps({"benchmark": {"official": 80.0, "flashdreams": 70.0}}), + encoding="utf-8", + ) + out_json = tmp_path / "bench.json" + out_md = tmp_path / "bench.md" + out_chart = tmp_path / "perf.md" + + module.main( + [ + "--item", + f"BF16:{bf16_json}", + "--item", + f"FP8:{fp8_json}", + "--output-json", + str(out_json), + "--output-md", + str(out_md), + "--output-chart-md", + str(out_chart), + ] + ) + + assert out_chart.read_text(encoding="utf-8") == ( + "# SANA-WM Benchmark Data (ms/frame)\n" + "\n" + "| precision | official | flashdreams |\n" + "| --- | ---: | ---: |\n" + "| BF16 | 100.00 | 90.00 |\n" + "| FP8 | 80.00 | 70.00 |\n" + ) + payload = json.loads(out_json.read_text(encoding="utf-8")) + assert payload["benchmark"]["metric"] == "generation_ms_per_frame" + assert [row["label"] for row in payload["rows"]] == ["BF16", "FP8"] From 9b48f4171893f11f095275655a23366d029b1b3c Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 24 Jul 2026 13:44:24 -0700 Subject: [PATCH 25/36] Optmizations to catch up to upstream --- integrations/sana/sana_wm/decoder.py | 21 +- integrations/sana/sana_wm/ops/__init__.py | 4 + .../sana/sana_wm/ops/fused_cam_gdn.py | 2448 +++++++++++++++++ integrations/sana/sana_wm/ops/fused_gdn.py | 2249 +++++++++++++++ .../sana/sana_wm/ops/fused_gdn_chunkwise.py | 2269 +++++++++++++++ integrations/sana/sana_wm/refiner.py | 98 +- integrations/sana/sana_wm/runner.py | 35 +- integrations/sana/sana_wm/stage1_model.py | 630 ++++- integrations/sana/sana_wm/transformer.py | 10 +- 9 files changed, 7652 insertions(+), 112 deletions(-) create mode 100644 integrations/sana/sana_wm/ops/__init__.py create mode 100644 integrations/sana/sana_wm/ops/fused_cam_gdn.py create mode 100644 integrations/sana/sana_wm/ops/fused_gdn.py create mode 100644 integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py diff --git a/integrations/sana/sana_wm/decoder.py b/integrations/sana/sana_wm/decoder.py index 34c53b8df..38fcf5794 100644 --- a/integrations/sana/sana_wm/decoder.py +++ b/integrations/sana/sana_wm/decoder.py @@ -91,12 +91,12 @@ class SanaWMLTX2VAEDecoderConfig(InstantiateConfig): offload_vae: bool = False """Move the VAE to CPU between decode calls.""" - vae_tile_sample_min_width: int = 256 - vae_tile_sample_stride_width: int = 224 - vae_tile_sample_min_height: int = 256 - vae_tile_sample_stride_height: int = 192 - vae_tile_sample_min_num_frames: int = 24 - vae_tile_sample_stride_num_frames: int = 8 + vae_tile_sample_min_width: int = 512 + vae_tile_sample_stride_width: int = 448 + vae_tile_sample_min_height: int = 512 + vae_tile_sample_stride_height: int = 448 + vae_tile_sample_min_num_frames: int = 96 + vae_tile_sample_stride_num_frames: int = 64 vae_oom_retry_tile_sample_min_width: int = 128 vae_oom_retry_tile_sample_stride_width: int = 64 @@ -185,8 +185,10 @@ def decode_latents(self, latents: Tensor) -> np.ndarray: decoded = torch.stack(decoded, dim=0) video = ( torch.clamp(127.5 * decoded + 127.5, 0, 255) + .to(torch.uint8) .permute(0, 2, 3, 4, 1) - .to("cpu", dtype=torch.uint8) + .contiguous() + .cpu() .numpy()[0] ) if self.config.offload_vae: @@ -314,6 +316,8 @@ class SanaWMLTX2LatentRefinerConfig(InstantiateConfig): refiner_precision: Precision = "bf16" quant_backend: QuantBackend = "torch" offload_refiner: bool = False + cache_text_encoder: bool = False + """Keep Gemma cached on CPU after prompt encoding for repeated pipeline use.""" class SanaWMLTX2LatentRefiner(nn.Module): @@ -383,6 +387,7 @@ def _ensure_refiner(self) -> None: device=self.device, precision=self.config.refiner_precision, quant_backend=self.config.quant_backend, + cache_text_encoder=self.config.cache_text_encoder, ) self._refiner_built = True logger.info( @@ -457,6 +462,8 @@ def forward( ) video_hwc = self.vae_decoder.decode_latents(output_latent) + if self.refiner is not None: + video_hwc = video_hwc[1:] stage1_video_hwc = None if cache.save_stage1 and self.refiner is not None: stage1_video_hwc = self.vae_decoder.decode_latents(stage1_latent) diff --git a/integrations/sana/sana_wm/ops/__init__.py b/integrations/sana/sana_wm/ops/__init__.py new file mode 100644 index 000000000..07bb2847d --- /dev/null +++ b/integrations/sana/sana_wm/ops/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SANA-WM integration-local Triton kernels.""" diff --git a/integrations/sana/sana_wm/ops/fused_cam_gdn.py b/integrations/sana/sana_wm/ops/fused_cam_gdn.py new file mode 100644 index 000000000..6f6229637 --- /dev/null +++ b/integrations/sana/sana_wm/ops/fused_cam_gdn.py @@ -0,0 +1,2448 @@ +"""Triton-fused camera-branch UCPE single-path delta rule. + +Companion to :mod:`diffusion.model.ops.fused_gdn` (main GDN +branch). This module fuses the *camera* branch of +:class:`ChunkCausalGDNUCPESinglePathLiteLA` through two Triton kernels: + +1. ``_cam_prep_kernel`` — fuses, per ``(batch, token, head)``, the RMSNorm + over the full ``C`` channels, ReLU on Q/K, K-scale on K, the UCPE 4x4 + block-diagonal projection matrix on the first ``D/2`` dims, and the + interleaved-pair complex RoPE on the second ``D/2`` dims. Q, K, V are + processed in one pass. The kernel also emits per-token pre-UCPE and + post-UCPE ``||k||^2`` so the caller can compute the inflation-squared + factor used for Dynamic Beta Discounting. + +2. ``_cam_scan_kernel`` — fuses the numerator-only single-path delta-rule + scan per ``(batch, head)``. ``REVERSE=1`` implements the + ``flip_and_shift`` backward pass semantics directly, avoiding the + torch-side flips in the per-chunk backward loop. + +The prep and scan kernels are runtime paths. Torch implementations below are +kept only for fallback backward paths and focused validation. + +Notes: + - V skips RMSNorm / ReLU / K-scale but receives the same UCPE 4x4 + + RoPE transforms as K (apply_fn_kv in reference). + - The short convolution on K and the inverse UCPE output transform + (``apply_fn_o``) stay in PyTorch — they are single lightweight ops + not on the critical path. +""" + +# ruff: noqa: E501 + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +from .fused_gdn_chunkwise import cam_scan_chunkwise + +# ============================================================================= +# Scalar helpers +# ============================================================================= + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Invert a 4x4 SE(3) matrix batch (closed-form). + + Mirrors the reference ``_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 reference 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 +# ============================================================================= + + +_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) + + +@triton.jit +def _cam_prep_bwd_kernel( + # --- forward inputs (replayed for ReLU mask + k_post_kscale recompute) --- + q_raw_ptr, # (B, N, H, D) contiguous, any fp dtype + k_raw_ptr, # (B, N, H, D) contiguous (post short-conv on K) + q_norm_w_ptr, # (C,) = (H*D,) float32 + k_norm_w_ptr, # (C,) float32 + q_inv_rms_ptr, # (B, N) float32 — saved from forward + k_inv_rms_ptr, # (B, N) 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 + # --- upstream gradients (B, H, D, N) layout matching forward outputs --- + d_q_out_ptr, # grad of q_out (any dtype, cast to fp32 on load) + eff_d_k_out_ptr, # grad of k_out + inflation_sq contribution through k_out (fp32) + d_v_out_ptr, # grad of v_out + # --- inflation_sq direct grad to k_post_kscale^2 sum (B, H, N) fp32 --- + d_pre_k_sq_ptr, + # --- outputs --- + d_q_post_norm_ptr, # (B, N, H, D) fp32 — grad after RoPE^T+UCPE^T+Kscale+ReLU; consumed by torch RMSNorm bwd + d_k_post_norm_ptr, # (B, N, H, D) fp32 — same for K + dv_raw_ptr, # (B, N, H, D) fp32 — final dv_raw (V skips norm/ReLU/Kscale) + # --- dims --- + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + D_HALF: tl.constexpr, + N_GROUPS: tl.constexpr, + 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) — matches the forward kernel's parallelism. + + Implements the bwd of the fused fwd: + forward order: RMSNorm -> ReLU -> [K-scale on K] -> UCPE (first D/2) + -> RoPE (second D/2) -> output (B, H, D, N). + backward order: RoPE^T (second D/2) -> UCPE^T (first D/2) + -> K-scale (only K) -> ReLU mask -> emit d_post_norm + intermediates for the cross-head RMSNorm bwd handled + outside (in :func:`_cam_prep_bwd_dispatch`). + + The full-channel RMSNorm bwd's outer-product term and per-channel weight + grad both require a sum over ``H*D`` (the full ``C``) per token, which + couples heads. We deliberately leave that step in PyTorch (a couple of + fused element-wise + reduction ops) — see :func:`_cam_prep_bwd_dispatch`. + + Inflation handling: the kernel takes an *effective* ``dO_k`` that already + includes the contribution from ``grad_inflation_sq`` flowing through + ``k_out`` (i.e. ``eff_dO_k = grad_k + 2 * k_out * d_post_k_sq``), plus a + per-(b, h, n) scalar ``d_pre_k_sq`` that is the chain-rule contribution + of ``grad_inflation_sq`` into the pre-UCPE ``||k_post_kscale||^2`` sum. + Inside the kernel we recompute ``k_post_kscale`` (= post-norm * ReLU * + K_SCALE) and add ``2 * k_post_kscale[d] * d_pre_k_sq`` as a direct + contribution to ``d_k_post_kscale``. + """ + pid = tl.program_id(0) + h_idx = pid % H + bn_idx = pid // H + b_idx = bn_idx // N + n_idx = bn_idx % N + + # ---- load saved scalars: inv-RMS (per b, n) and d_pre_k_sq (per b, h, n) ---- + 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) + bhn_idx = (b_idx * H + h_idx) * N + n_idx + d_pre_k_sq = tl.load(d_pre_k_sq_ptr + bhn_idx).to(tl.float32) + + # ---- load per-token P matrices (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) + + # ---- layout offsets ---- + row_base = b_idx * (N * H * D) + n_idx * (H * D) + h_idx * D # (B, N, H, D) + nw_off = h_idx * D + out_base_BHDN = b_idx * (H * D * N) + h_idx * (D * N) + n_idx # (B, H, D, N) for dO_* + norm_base = row_base # d_*_post_norm and dv_raw share the (B, N, H, D) layout + + # ============================================================ + # First half — UCPE region + # ============================================================ + 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] + + # Recompute post-norm (pre-ReLU) Q/K for the ReLU mask + k_post_kscale. + q_half_raw = tl.load(q_raw_ptr + row_base + offs_gj, mask=mask_gj, other=0.0).to(tl.float32) + k_half_raw = tl.load(k_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_post_norm_half = q_half_raw * q_inv_rms * q_nw_half + k_post_norm_half = k_half_raw * k_inv_rms * k_nw_half + q_relu_mask_half = q_post_norm_half > 0 + k_relu_mask_half = k_post_norm_half > 0 + + # k_post_kscale = relu(k_post_norm) * K_SCALE (used for direct inflation contribution) + k_post_relu_half = tl.where(k_relu_mask_half, k_post_norm_half, 0.0) + k_post_kscale_half = k_post_relu_half * K_SCALE + + # Load upstream gradients for first half. + # In (B, H, D, N): ptr[b, h, d, n] = out_base_BHDN + d * N. d = g*4 + i. + offs_d_half = offs_g[:, None] * 4 + offs_i[None, :] # (BLOCK_GROUPS, 4) + mask_d_half = mask_g[:, None] + dO_q_half = tl.load(d_q_out_ptr + out_base_BHDN + offs_d_half * N, mask=mask_d_half, other=0.0).to(tl.float32) + dO_k_eff_half = tl.load(eff_d_k_out_ptr + out_base_BHDN + offs_d_half * N, mask=mask_d_half, other=0.0).to( + tl.float32 + ) + dO_v_half = tl.load(d_v_out_ptr + out_base_BHDN + offs_d_half * N, mask=mask_d_half, other=0.0).to(tl.float32) + + # UCPE^T: din[g, j] = sum_i P[i, j] * dout[g, i] + # forward: out[g, i] = sum_j q[g, j] * P[i, j] -- so bwd sums over i + d_q_post_relu_half = tl.sum(dO_q_half[:, :, None] * P_q[None, :, :], axis=1) + d_k_post_kscale_via_ucpe_half = tl.sum(dO_k_eff_half[:, :, None] * P_kv[None, :, :], axis=1) + d_v_first_half = tl.sum(dO_v_half[:, :, None] * P_kv[None, :, :], axis=1) + + # K direct inflation contribution: 2 * k_post_kscale * d_pre_k_sq + d_k_post_kscale_half = d_k_post_kscale_via_ucpe_half + 2.0 * k_post_kscale_half * d_pre_k_sq + + # K-scale bwd (multiply by K_SCALE) + d_k_post_relu_half = d_k_post_kscale_half * K_SCALE + + # ReLU mask + d_q_post_norm_half = tl.where(q_relu_mask_half, d_q_post_relu_half, 0.0) + d_k_post_norm_half = tl.where(k_relu_mask_half, d_k_post_relu_half, 0.0) + + # Mask out-of-bounds groups to 0 explicitly (for safety on uneven N_GROUPS). + d_q_post_norm_half = tl.where(mask_d_half, d_q_post_norm_half, 0.0) + d_k_post_norm_half = tl.where(mask_d_half, d_k_post_norm_half, 0.0) + d_v_first_half = tl.where(mask_d_half, d_v_first_half, 0.0) + + # Store at (B, N, H, D), d = g*4+i + tl.store(d_q_post_norm_ptr + norm_base + offs_d_half, d_q_post_norm_half, mask=mask_d_half) + tl.store(d_k_post_norm_ptr + norm_base + offs_d_half, d_k_post_norm_half, mask=mask_d_half) + tl.store(dv_raw_ptr + norm_base + offs_d_half, d_v_first_half, mask=mask_d_half) + + # ============================================================ + # Second half — RoPE region + # ============================================================ + 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_pair = tl.load(rope_sin_ptr + rope_row + offs_r_pair, mask=mask_r_pair, other=0.0).to(tl.float32) + + # Recompute post-norm (pre-ReLU) Q/K for the ReLU mask. + rope_base = row_base + D_HALF + q_r_raw = tl.load(q_raw_ptr + rope_base + offs_r, mask=mask_r, other=0.0).to(tl.float32) + k_r_raw = tl.load(k_raw_ptr + rope_base + offs_r, mask=mask_r, 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_post_norm_r = q_r_raw * q_inv_rms * q_nw_r + k_post_norm_r = k_r_raw * k_inv_rms * k_nw_r + q_relu_mask_r = q_post_norm_r > 0 + k_relu_mask_r = k_post_norm_r > 0 + + k_post_relu_r = tl.where(k_relu_mask_r, k_post_norm_r, 0.0) + k_post_kscale_r = k_post_relu_r * K_SCALE + + # Load upstream gradients (second half) — direct + pair. + offs_d_r = D_HALF + offs_r + offs_d_r_pair = D_HALF + offs_r_pair + dO_q_r = tl.load(d_q_out_ptr + out_base_BHDN + offs_d_r * N, mask=mask_r, other=0.0).to(tl.float32) + dO_k_eff_r = tl.load(eff_d_k_out_ptr + out_base_BHDN + offs_d_r * N, mask=mask_r, other=0.0).to(tl.float32) + dO_v_r = tl.load(d_v_out_ptr + out_base_BHDN + offs_d_r * N, mask=mask_r, other=0.0).to(tl.float32) + dO_q_r_pair = tl.load(d_q_out_ptr + out_base_BHDN + offs_d_r_pair * N, mask=mask_r_pair, other=0.0).to(tl.float32) + dO_k_eff_r_pair = tl.load(eff_d_k_out_ptr + out_base_BHDN + offs_d_r_pair * N, mask=mask_r_pair, other=0.0).to( + tl.float32 + ) + dO_v_r_pair = tl.load(d_v_out_ptr + out_base_BHDN + offs_d_r_pair * N, mask=mask_r_pair, other=0.0).to(tl.float32) + + # RoPE^T: forward y[r] = x[r]*cos[r] + x[r^1]*sin[r] + # bwd dx[r] = dy[r]*cos[r] + dy[r^1]*sin[r^1] + d_q_post_relu_r = dO_q_r * cos_v + dO_q_r_pair * sin_v_pair + d_k_post_kscale_via_rope_r = dO_k_eff_r * cos_v + dO_k_eff_r_pair * sin_v_pair + d_v_second_r = dO_v_r * cos_v + dO_v_r_pair * sin_v_pair + + # K direct inflation contribution + d_k_post_kscale_r = d_k_post_kscale_via_rope_r + 2.0 * k_post_kscale_r * d_pre_k_sq + + # K-scale bwd + d_k_post_relu_r = d_k_post_kscale_r * K_SCALE + + # ReLU mask + d_q_post_norm_r = tl.where(q_relu_mask_r, d_q_post_relu_r, 0.0) + d_k_post_norm_r = tl.where(k_relu_mask_r, d_k_post_relu_r, 0.0) + + # Out-of-bound mask + d_q_post_norm_r = tl.where(mask_r, d_q_post_norm_r, 0.0) + d_k_post_norm_r = tl.where(mask_r, d_k_post_norm_r, 0.0) + d_v_second_r = tl.where(mask_r, d_v_second_r, 0.0) + + norm_offs_r = D_HALF + offs_r + tl.store(d_q_post_norm_ptr + norm_base + norm_offs_r, d_q_post_norm_r, mask=mask_r) + tl.store(d_k_post_norm_ptr + norm_base + norm_offs_r, d_k_post_norm_r, mask=mask_r) + tl.store(dv_raw_ptr + norm_base + norm_offs_r, d_v_second_r, mask=mask_r) + + +@triton.jit +def _cam_scan_kernel( + # --- inputs (B, H, D, N) contiguous, fp32 --- + q_ptr, + k_ptr, + v_ptr, + # --- gates --- + beta_ptr, # (B, H, F, S) contiguous + decay_ptr, # (B, H, F) contiguous + # --- output (B, H, D, N) fp32 --- + out_ptr, + # --- saved state snapshots (used when SAVE_STATES=1) --- + state_pre_ptr, # (B, H, F, BLOCK_D, BLOCK_D) fp32 — state after decay, before update + state_post_ptr, # (B, H, F, BLOCK_D, BLOCK_D) fp32 — state after update + # --- forward-direction cache state (used when LOAD_INIT_STATE / SAVE_FINAL_STATE) --- + init_state_ptr, # (B*H, BLOCK_D, BLOCK_D) fp32 — state at end of prefix + final_state_ptr, # (B*H, BLOCK_D, BLOCK_D) fp32 — state after last frame's update + # --- dims --- + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + N: tl.constexpr, # F * S + REVERSE: tl.constexpr, + SAVE_STATES: tl.constexpr, + LOAD_INIT_STATE: tl.constexpr, + SAVE_FINAL_STATE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """One program per (b, h) — runs the full numerator-only delta-rule scan. + + When ``SAVE_STATES=1`` the kernel additionally writes per-frame snapshots + of ``state_prev`` (state after applying ``decay``, before the K@delta + update) to ``state_pre_ptr`` and ``state_curr`` (state after the update) + to ``state_post_ptr``, both indexed by ``q_frame``. The bwd kernel + (``_cam_scan_bwd_kernel``) consumes these snapshots for both + ``REVERSE=0`` and ``REVERSE=1``. In ``REVERSE=1`` the slot at + ``q_frame=F-1`` always holds the all-zero state (skip-update, decay=1 + on the zero initial state), and the bwd kernel reads exactly that — + no special-case load is needed. + + When ``LOAD_INIT_STATE=1`` (forward direction only — wrapper enforces + ``REVERSE=0``) the per-program ``state_curr`` is initialized from + ``init_state_ptr`` instead of zero. The convention is: the loaded value + is the state AT THE END of a prefix sequence (i.e., AFTER the prefix's + last update, BEFORE any further decay applied here). On the very first + frame, the kernel's own ``state_curr *= g`` then applies ``decay[0]`` + to this loaded state — which is exactly the decay that the global + sequence's f=K-th frame would have applied. This keeps split/resume + state trajectories identical from frame K onwards. + + When ``SAVE_FINAL_STATE=1`` (forward direction only) the final + ``state_curr`` (after the last frame's update) is written to + ``final_state_ptr``. This is the state to be loaded with + ``LOAD_INIT_STATE`` for a downstream segment. + """ + pid = tl.program_id(0) + pid_b = pid // H + pid_h = pid % H + bh = pid_b * H + pid_h + + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < D + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + mask_dd = mask_d[:, None] & mask_d[None, :] + + base_bh = pid_b * (H * D * N) + pid_h * (D * N) + q_bh = q_ptr + base_bh + k_bh = k_ptr + base_bh + v_bh = v_ptr + base_bh + out_bh = out_ptr + base_bh + beta_bh = beta_ptr + bh * F * S + decay_bh = decay_ptr + bh * F + if SAVE_STATES: + spre_bh = state_pre_ptr + bh * F * BLOCK_D * BLOCK_D + spost_bh = state_post_ptr + bh * F * BLOCK_D * BLOCK_D + + # State: (D_k, D_v) in the upstream convention. Here we call rows "k-dim" + # (input dim of state) and cols "v-dim" (output dim of state). + if LOAD_INIT_STATE: + init_bh = init_state_ptr + bh * BLOCK_D * BLOCK_D + state_curr = tl.load(init_bh + offs_dd, mask=mask_dd, other=0.0).to(tl.float32) + else: + state_curr = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + + for f_iter in range(F): + if REVERSE: + q_frame = F - 1 - f_iter + kv_frame = F - f_iter if f_iter > 0 else 0 + skip_update = f_iter == 0 + else: + q_frame = f_iter + kv_frame = f_iter + skip_update = False + + if REVERSE and f_iter == 0: + g = 1.0 + else: + g = tl.load(decay_bh + kv_frame).to(tl.float32) + state_curr = state_curr * g + state_prev = state_curr # fp32 snapshot (same tensor, kept for clarity) + + if SAVE_STATES: + tl.store( + spre_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, + state_prev, + mask=mask_dd, + ) + + if skip_update == 0: + kv_n_base = kv_frame * S + f_beta = beta_bh + kv_frame * 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 = kv_n_base + offs_s + + # Load K, V tiles (BLOCK_S, BLOCK_D) from (B, H, D, N) layout: + # ptr[s, d] = base_bh + offs_d[d] * N + n_idx[s] + k_ptrs = k_bh + offs_d[None, :] * N + n_idx[:, None] + v_ptrs = v_bh + offs_d[None, :] * N + n_idx[:, None] + K = tl.load(k_ptrs, mask=mask_sd, other=0.0) + V = tl.load(v_ptrs, mask=mask_sd, other=0.0) + + bt = tl.load(f_beta + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + # V_pred = K @ state_prev : (BLOCK_S, BLOCK_D) + V_pred = tl.dot( + K, + state_prev, + out_dtype=tl.float32, + input_precision="tf32", + ) + dv = (V - V_pred) * bt[:, None] + state_curr += tl.dot( + tl.trans(K), + dv, + out_dtype=tl.float32, + input_precision="tf32", + ) + + if SAVE_STATES: + tl.store( + spost_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, + state_curr, + mask=mask_dd, + ) + + # --- Pass 2: output --- + state_out = state_curr + q_n_base = q_frame * 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 = q_n_base + offs_s + + q_ptrs = q_bh + offs_d[None, :] * N + n_idx[:, None] + Q = tl.load(q_ptrs, mask=mask_sd, other=0.0) + + # num = Q @ state_out : (BLOCK_S, BLOCK_D), rows=S, cols=D_v + num = tl.dot( + Q, + state_out, + out_dtype=tl.float32, + input_precision="tf32", + ) + + # Store transposed into (B, H, D, N): + # ptr[d, s] = out_bh + offs_d[d] * N + n_idx[s] + out_ptrs = out_bh + offs_d[:, None] * N + n_idx[None, :] + mask_ds = mask_d[:, None] & mask_s[None, :] + tl.store(out_ptrs, tl.trans(num), mask=mask_ds) + + if SAVE_FINAL_STATE: + final_bh = final_state_ptr + bh * BLOCK_D * BLOCK_D + tl.store(final_bh + offs_dd, state_curr, mask=mask_dd) + + +# ============================================================================= +# Backward Triton kernel +# ============================================================================= + + +@triton.jit +def _cam_scan_bwd_kernel( + # --- forward inputs (B, H, D, N) fp32 contiguous --- + q_ptr, + k_ptr, + v_ptr, + # --- gates --- + beta_ptr, # (B, H, F, S) fp32 + decay_ptr, # (B, H, F) fp32 + # --- saved state snapshots (B*H, F, BLOCK_D, BLOCK_D) fp32, indexed by q_frame --- + state_pre_ptr, # state after decay, before update + state_post_ptr, # state after update + # --- upstream gradient (B, H, D, N) fp32 --- + grad_out_ptr, + # --- output gradients --- + dq_ptr, # (B, H, D, N) fp32 + dk_ptr, + dv_ptr, + dbeta_ptr, # (B, H, F, S) fp32 + ddecay_ptr, # (B, H, F) fp32 + # --- dims --- + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + N: tl.constexpr, # F * S + REVERSE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Reverse-time backward of the numerator-only delta-rule scan. + + One program per ``(b, h)``. Walks the same ``f_iter`` index space as the + forward kernel — but in reverse time order — replaying the recurrence + using the per-``q_frame`` state snapshots saved by the forward pass with + ``SAVE_STATES=1``. Accumulates gradients into ``dq``, ``dk``, ``dv``, + ``dbeta``, ``ddecay``. + + Forward indexing (matched here exactly):: + + REVERSE=False: q_frame = kv_frame = f_iter, skip_update = False + REVERSE=True : q_frame = F-1-f_iter, + kv_frame = F-f_iter (skip when f_iter==0) + g = decay[kv_frame] (1.0 when f_iter==0) + + Per-iteration backward derivation (when ``not skip_update``): + + ds_post += Q.T @ d_out # accumulate output grad + dQ[q] = d_out @ s_post.T + + ddelta = K @ ds_post # via K.T @ delta term + dV[kv] = ddelta * beta[kv,:] + dbeta[kv,:] = sum_d (ddelta * (V - K @ s_pre)) + dV_pred = -ddelta * beta[kv,:] + dK[kv] = delta @ ds_post.T + dV_pred @ s_pre.T + ds_pre = ds_post + K.T @ dV_pred # direct + V_pred path + + ddecay[kv] = sum(ds_pre * s_pre[q]) / g (= 0 when state_in is 0) + ds_post[next-bwd-iter] = ds_pre * g # propagate through decay + + For the ``skip_update`` branch (``REVERSE=True`` and ``f_iter==0``) the + update / decay path is bypassed: ``state_post == state_pre == 0`` so + ``dQ == 0``, ``ds_post`` is left unchanged across the iter, and + no writes to ``dK``, ``dV``, ``dbeta`` or ``ddecay`` happen at + ``kv_frame == 0``. Caller MUST pre-zero those output buffers (the + dispatch wrapper uses ``torch.zeros_like``). + + The ``ddecay`` formula uses ``state_pre / g``; for ``REVERSE=False`` at + fwd frame 0 we hardcode ``ddecay[0] = 0`` (state_in is exactly 0, but + the division would amplify any rounding noise). For very small ``g`` + the existing ``+ 1e-12`` epsilon is matched verbatim from + ``_fused_gdn_bwd_kernel``. + """ + pid = tl.program_id(0) + pid_b = pid // H + pid_h = pid % H + bh = pid_b * H + pid_h + + base_bh = pid_b * (H * D * N) + pid_h * (D * N) + q_bh = q_ptr + base_bh + k_bh = k_ptr + base_bh + v_bh = v_ptr + base_bh + do_bh = grad_out_ptr + base_bh + dq_bh = dq_ptr + base_bh + dk_bh = dk_ptr + base_bh + dv_bh = dv_ptr + base_bh + + beta_bh = beta_ptr + bh * F * S + decay_bh = decay_ptr + bh * F + dbeta_bh = dbeta_ptr + bh * F * S + ddecay_bh = ddecay_ptr + bh * F + + spre_bh = state_pre_ptr + bh * F * BLOCK_D * BLOCK_D + spost_bh = state_post_ptr + bh * F * BLOCK_D * BLOCK_D + + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < D + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + mask_dd = mask_d[:, None] & mask_d[None, :] + + # bf16 operands for tl.dot keep shared-memory pressure manageable for large + # BLOCK_D (e.g., 128 in reference). fp32 accumulators preserve precision. + grad_dtype = tl.bfloat16 + grad_ip: tl.constexpr = "tf32" + + # Reverse-time accumulator: gradient w.r.t. ``state_post`` for the iter + # currently being processed. + ds_post = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + + for f_rev in range(F): + # Walk fwd iters in reverse: F-1, F-2, ..., 0. + f_iter = F - 1 - f_rev + + if REVERSE: + q_frame = F - 1 - f_iter + kv_frame = F - f_iter if f_iter > 0 else 0 + skip_update = f_iter == 0 + else: + q_frame = f_iter + kv_frame = f_iter + skip_update = 0 + + if REVERSE and f_iter == 0: + g = 1.0 + else: + g = tl.load(decay_bh + kv_frame).to(tl.float32) + + # ---- Pass 2: dQ + ds_post += Q.T @ d_out ---------------------- + # Load state_post[q_frame] (zero in the REVERSE skip-update slot — + # the fwd save still writes the all-zero state at that index). + state = tl.load( + spost_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, + mask=mask_dd, + other=0.0, + ) + + q_n_base = q_frame * 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 = q_n_base + offs_s + + q_ptrs = q_bh + offs_d[None, :] * N + n_idx[:, None] + Q = tl.load(q_ptrs, mask=mask_sd, other=0.0) + + do_ptrs = do_bh + offs_d[None, :] * N + n_idx[:, None] + dO = tl.load(do_ptrs, mask=mask_sd, other=0.0) + + ds_post += tl.dot( + tl.trans(Q.to(grad_dtype)), + dO.to(grad_dtype), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + dQ = tl.dot( + dO.to(grad_dtype), + tl.trans(state.to(grad_dtype)), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + dq_ptrs = dq_bh + offs_d[:, None] * N + n_idx[None, :] + mask_ds = mask_d[:, None] & mask_s[None, :] + tl.store(dq_ptrs, tl.trans(dQ), mask=mask_ds) + + if skip_update == 0: + # ---- Reload state with state_pre[q_frame] for Pass 1 ---- + state = tl.load( + spre_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, + mask=mask_dd, + other=0.0, + ) + + # ds_pre starts equal to ds_post (direct pass-through term). + ds_pre = ds_post + + kv_n_base = kv_frame * 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 = kv_n_base + offs_s + + k_ptrs = k_bh + offs_d[None, :] * N + n_idx[:, None] + v_ptrs = v_bh + offs_d[None, :] * N + n_idx[:, None] + K = tl.load(k_ptrs, mask=mask_sd, other=0.0) + V = tl.load(v_ptrs, mask=mask_sd, other=0.0) + + bt = tl.load(beta_bh + kv_frame * S + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + V_pred = tl.dot( + K.to(grad_dtype), + state.to(grad_dtype), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + r = V - V_pred + delta = r * bt[:, None] + + ddelta = tl.dot( + K.to(grad_dtype), + ds_post.to(grad_dtype), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + dV = ddelta * bt[:, None] + dv_ptrs = dv_bh + offs_d[:, None] * N + n_idx[None, :] + mask_ds = mask_d[:, None] & mask_s[None, :] + tl.store(dv_ptrs, tl.trans(dV), mask=mask_ds) + + dbeta_st = tl.sum(ddelta * r, axis=1) + tl.store(dbeta_bh + kv_frame * S + offs_s, dbeta_st, mask=mask_s) + + dV_pred = -ddelta * bt[:, None] + + dK_part1 = tl.dot( + delta.to(grad_dtype), + tl.trans(ds_post.to(grad_dtype)), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + dK_part2 = tl.dot( + dV_pred.to(grad_dtype), + tl.trans(state.to(grad_dtype)), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + dK = dK_part1 + dK_part2 + dk_ptrs = dk_bh + offs_d[:, None] * N + n_idx[None, :] + tl.store(dk_ptrs, tl.trans(dK), mask=mask_ds) + + ds_pre += tl.dot( + tl.trans(K.to(grad_dtype)), + dV_pred.to(grad_dtype), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + # ---- ddecay[kv_frame] ---- + # state_pre = state_in * g, so state_in = state_pre / g. + # ddecay = sum(ds_pre * state_in) = sum(ds_pre * state_pre) / g. + # For REVERSE=False fwd frame 0, state_in is exactly 0 — hardcode + # 0 to avoid amplifying rounding noise via 1/g. + if (REVERSE == 0) and (f_iter == 0): + ddecay_f = 0.0 + else: + inv_g = 1.0 / (g + 1e-12) + ddecay_f = tl.sum(ds_pre * state) * inv_g + tl.store(ddecay_bh + kv_frame, ddecay_f) + + # Propagate to next bwd iter (which is fwd's previous iter). + ds_post = ds_pre * g + # else (skip_update branch): state_post == state_pre == 0, no + # ddecay write, no kv-side writes; ds_post passes through unchanged + # since ∂state_post/∂state_in = I when the update is skipped and g=1. + # In REVERSE=True this is the LAST bwd iter (f_iter=0) so the + # carried-over ds_post is discarded. + + +# ============================================================================= +# Python wrappers +# ============================================================================= + + +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. + """ + 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 + + +def _run_cam_prep_fwd_save( + 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, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + k_scale: float, + norm_eps: float, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Run :func:`cam_prep_func` while exposing intermediates needed by the bwd kernel. + + Mirrors :func:`cam_prep_func` exactly (same kernel launch, same outputs) + but additionally returns the per-token ``inv_rms`` for both Q and K, plus + the raw ``||k_pre_ucpe||^2`` and ``||k_post_ucpe||^2`` per ``(B, H, N)``. + These are required by :func:`_cam_prep_bwd_dispatch` to (a) replay the + ReLU/RMSNorm chain in fp32 without re-summing over the full ``C`` + channels and (b) chain ``grad_inflation_sq`` back through ``k_out`` and + the pre-UCPE ``||k||^2`` term with the correct ``clamp_min`` indicators. + + Returns: + ``(q_out, k_out, v_out, inflation_sq, q_inv_rms, k_inv_rms, + k_pre_sq, k_post_sq)``. The first four match :func:`cam_prep_func`; + the rest are fp32 contiguous saved-state tensors for the backward. + """ + 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 + + 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) + + _cam_prep_kernel[(B * N * H,)]( + 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 = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12) + return q_out, k_out, v_out, inflation_sq, q_inv_rms, k_inv_rms, k_pre_sq, k_post_sq + + +def _cam_prep_bwd_dispatch( + q_raw: torch.Tensor, + k_raw: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + proj_q: torch.Tensor, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + q_inv_rms: torch.Tensor, + k_inv_rms: torch.Tensor, + k_pre_sq: torch.Tensor, + k_post_sq: torch.Tensor, + k_out: torch.Tensor, + *, + grad_q: torch.Tensor | None, + grad_k: torch.Tensor | None, + grad_v: torch.Tensor | None, + grad_inflation_sq: torch.Tensor | None, + k_scale: float, +) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, +]: + """Hybrid Triton + torch backward for :func:`cam_prep_func`. + + Pipeline: + + 1. **(torch)** Chain ``grad_inflation_sq`` through ``k_post_sq`` and + ``k_pre_sq`` to produce + (a) ``eff_d_k_out = grad_k + 2 * k_out * d_post_k_sq`` and + (b) ``d_pre_k_sq`` — a per-(B, H, N) scalar fed into the kernel as a + direct contribution to ``d_k_post_kscale``. Both ``clamp_min(1e-12)`` + indicators are honored so the gradient is exactly 0 in the (rare) + saturating regime, matching :func:`_torch_cam_prep_reference`. + + 2. **(Triton)** Launch :func:`_cam_prep_bwd_kernel` per ``(b, n, h)`` + to apply RoPE^T, UCPE^T, K-scale^T and the ReLU mask. Emits + ``d_q_post_norm`` / ``d_k_post_norm`` ``(B, N, H, D)`` fp32 + intermediates (the post-norm pre-RMSNorm grad slots) and writes + ``dv_raw`` directly (V skips RMSNorm/ReLU/K-scale). + + 3. **(torch)** Apply the full-channel RMSNorm bwd. The cross-head + coupling means ``S_q[b, n] = sum_{h,d} d_q_post_norm * q_raw * + q_norm_w`` is reduced over ``H*D``; we then form + ``dq_raw = d_q_post_norm * inv_rms_q * q_norm_w + - inv_rms_q^3 / C * q_raw * S_q`` + elementwise, and ``dq_norm_weight[c] = sum_{b,n} + d_q_post_norm[b,n,c] * q_raw[b,n,c] * inv_rms_q[b,n]``. + + Returns: + ``(dq_raw, dk_raw, dv_raw, dq_norm_weight, dk_norm_weight)``. Each + slot is ``None`` if upstream did not request that grad — handled + by the caller via ``ctx.needs_input_grad``. ``dq_raw`` / ``dk_raw`` + / ``dv_raw`` are returned in the same dtype as ``q_raw``; + norm-weight grads are fp32 (matching the input dtype). + """ + B, N, H, D = q_raw.shape + assert k_raw.shape == q_raw.shape + assert q_raw.is_contiguous() and k_raw.is_contiguous() + assert q_inv_rms.shape == (B, N) and k_inv_rms.shape == (B, N) + assert k_pre_sq.shape == (B, H, N) and k_post_sq.shape == (B, H, N) + D_half = D // 2 + N_groups = D_half // 4 + C = H * D + + # ---- prepare inflation_sq grad chain (in fp32) ---- + eps_floor = 1e-12 + pre_clamped = k_pre_sq.clamp_min(eps_floor) + if grad_inflation_sq is not None: + gis = grad_inflation_sq.to(torch.float32) + post_clamped = k_post_sq.clamp_min(eps_floor) + pre_indicator = (k_pre_sq >= eps_floor).to(torch.float32) + post_indicator = (k_post_sq >= eps_floor).to(torch.float32) + # d(inflation_sq)/d(post_k_sq) = (1 / pre_clamped) * post_indicator + # d(inflation_sq)/d(pre_k_sq) = -post_clamped / pre_clamped^2 * pre_indicator + d_post_k_sq = (post_indicator / pre_clamped) * gis # (B, H, N) + d_pre_k_sq = (-post_clamped / (pre_clamped * pre_clamped) * pre_indicator) * gis # (B, H, N) + else: + d_post_k_sq = torch.zeros_like(k_pre_sq) + d_pre_k_sq = torch.zeros_like(k_pre_sq) + + # eff_d_k_out: (B, H, D, N) fp32, contiguous + if grad_k is None: + grad_k_f32 = torch.zeros((B, H, D, N), dtype=torch.float32, device=q_raw.device) + else: + grad_k_f32 = grad_k.to(torch.float32) + if grad_inflation_sq is not None: + # k_out: (B, H, D, N), d_post_k_sq: (B, H, N) → broadcast over D dim. + eff_d_k_out = (grad_k_f32 + 2.0 * k_out.to(torch.float32) * d_post_k_sq.unsqueeze(2)).contiguous() + else: + eff_d_k_out = grad_k_f32.contiguous() + d_pre_k_sq = d_pre_k_sq.contiguous() + + # grad_q / grad_v as fp32 (B, H, D, N) contiguous (zero-fill if absent) + if grad_q is None: + grad_q_f32 = torch.zeros((B, H, D, N), dtype=torch.float32, device=q_raw.device) + else: + grad_q_f32 = grad_q.to(torch.float32).contiguous() + if grad_v is None: + grad_v_f32 = torch.zeros((B, H, D, N), dtype=torch.float32, device=q_raw.device) + else: + grad_v_f32 = grad_v.to(torch.float32).contiguous() + + # ---- allocate outputs / intermediates ---- + d_q_post_norm = torch.empty((B, N, H, D), dtype=torch.float32, device=q_raw.device) + d_k_post_norm = torch.empty((B, N, H, D), dtype=torch.float32, device=q_raw.device) + dv_raw_f32 = torch.empty((B, N, H, D), 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) + + _cam_prep_bwd_kernel[(B * N * H,)]( + q_raw, + k_raw, + q_norm_weight, + k_norm_weight, + q_inv_rms, + k_inv_rms, + proj_q, + proj_kv, + rope_cos, + rope_sin, + grad_q_f32, + eff_d_k_out, + grad_v_f32, + d_pre_k_sq, + d_q_post_norm, + d_k_post_norm, + dv_raw_f32, + 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, + ) + + # ---- torch RMSNorm bwd over the saved post-norm grads ---- + # Cast raw inputs to fp32 for the cross-head reduction (matches kernel + # numerics — the kernel uses fp32 internally as well). + q_raw_f32 = q_raw.to(torch.float32) + k_raw_f32 = k_raw.to(torch.float32) + q_inv_rms_view = q_inv_rms.view(B, N, 1, 1) + k_inv_rms_view = k_inv_rms.view(B, N, 1, 1) + q_nw_view = q_norm_weight.view(1, 1, H, D) + k_nw_view = k_norm_weight.view(1, 1, H, D) + + # S_q[b, n] = sum_{h, d} d_q_post_norm[b, n, h, d] * q_raw[b, n, h, d] * q_norm_w[h, d] + weighted_q = d_q_post_norm * q_raw_f32 # reused for dq_norm_weight reduction + weighted_k = d_k_post_norm * k_raw_f32 + S_q = (weighted_q * q_nw_view).sum(dim=(2, 3)) # (B, N) + S_k = (weighted_k * k_nw_view).sum(dim=(2, 3)) + + # dq_raw[b, n, h, d] = d_q_post_norm * inv_rms_q * q_norm_w + # - inv_rms_q^3 / C * q_raw * S_q + inv_q3 = (q_inv_rms**3).view(B, N, 1, 1) + inv_k3 = (k_inv_rms**3).view(B, N, 1, 1) + inv_C = 1.0 / float(C) + dq_raw_f32 = d_q_post_norm * q_inv_rms_view * q_nw_view - inv_q3 * inv_C * q_raw_f32 * S_q.view(B, N, 1, 1) + dk_raw_f32 = d_k_post_norm * k_inv_rms_view * k_nw_view - inv_k3 * inv_C * k_raw_f32 * S_k.view(B, N, 1, 1) + + # dq_norm_weight[h, d] = sum_{b, n} d_q_post_norm[b, n, h, d] * q_raw[b, n, h, d] * inv_rms_q[b, n] + dq_norm_weight = (weighted_q * q_inv_rms_view).sum(dim=(0, 1)).reshape(-1).contiguous() + dk_norm_weight = (weighted_k * k_inv_rms_view).sum(dim=(0, 1)).reshape(-1).contiguous() + + # Cast Q/K/V grads back to input dtype to match torch.autograd convention. + dq_raw = dq_raw_f32.to(q_raw.dtype) + dk_raw = dk_raw_f32.to(q_raw.dtype) + dv_raw = dv_raw_f32.to(q_raw.dtype) + + return dq_raw, dk_raw, dv_raw, dq_norm_weight, dk_norm_weight + + +def cam_scan_func( + 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, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Fused numerator-only single-path delta-rule scan for the cam branch. + + Args: + q, k, v: ``(B, H, D, N)`` fp32 contiguous tensors. + beta: ``(B, H, F, S)`` fp32 contiguous. + decay: ``(B, H, F)`` fp32 contiguous. + reverse: If ``True``, run the scan as the backward pass (equivalent + to ``flip_and_shift``-ing the inputs along the frame axis and + running forward). + init_state: optional ``(B*H, BLOCK_D, BLOCK_D)`` fp32 contiguous + tensor holding the forward-scan KV state at the END of a prefix + sequence (i.e., AFTER the prefix's last update, BEFORE any + further decay applied by this call). When provided, the kernel + resumes the scan from this state instead of zero. ``BLOCK_D = + next_pow2(D)`` and only the top-left ``D x D`` submatrix is + read. Forward direction only — raises ``NotImplementedError`` + if combined with ``reverse=True``. + save_final_state: when True, allocate a fresh fp32 zero buffer for + the final KV state (after the last frame's update) and pass it + to the kernel for write-out. Returned as the second tuple slot. + Forward direction only. + + Returns: + ``out`` of shape ``(B, H, D, N)`` fp32 matching + ``torch_chunk_cam_single_path_delta_rule`` with ``chunk_size >= T``. + + When ``save_final_state=True``, returns ``(out, final_state)`` where + ``final_state`` is fp32 ``(B*H, BLOCK_D, BLOCK_D)``. + + Raises: + NotImplementedError: if ``reverse=True`` is combined with state + passing. The cam branch's anti-causal scan resets per chunk in + the reference block, so there is no global cross-prefix state + to cache for the reverse direction. + """ + # Chunkwise integration (2026-05-06): dispatch all paths (fwd + reverse) + # to `cam_scan_chunkwise`. Reverse uses chunkwise's existing direction=2 + # mode in phase_b_triton, which has the same flip-and-shift semantics as + # cam's REVERSE=1 path. Bypass via FUSED_GDN_FORCE_LEGACY=1. + if os.environ.get("FUSED_GDN_FORCE_LEGACY", "0") != "1": + return cam_scan_chunkwise( + q, + k, + v, + beta, + decay, + reverse=reverse, + init_state=init_state, + save_final_state=save_final_state, + ) + + assert q.shape == k.shape == v.shape + B, H, D, N = q.shape + assert beta.shape[0] == B and beta.shape[1] == H + F_frames = beta.shape[2] + assert N % F_frames == 0 + S = N // F_frames + assert beta.shape == (B, H, F_frames, S), f"beta shape {beta.shape}" + assert decay.shape == (B, H, F_frames), f"decay shape {decay.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 + + BLOCK_D = triton.next_power_of_2(D) + BLOCK_S = _DEFAULT_BLOCK_S + num_warps = 4 + num_stages = 1 + + if reverse and (init_state is not None or save_final_state): + raise NotImplementedError( + "cam_scan_func: 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." + ) + + if init_state is not None: + expected_shape = (B * H, BLOCK_D, BLOCK_D) + if tuple(init_state.shape) != expected_shape: + raise ValueError( + f"cam_scan_func: init_state shape {tuple(init_state.shape)} " + f"does not match expected {expected_shape} (BLOCK_D=next_pow2(D)={BLOCK_D})." + ) + if init_state.dtype != torch.float32: + raise ValueError(f"cam_scan_func: init_state must be fp32 (got {init_state.dtype}).") + if not init_state.is_contiguous(): + raise ValueError("cam_scan_func: init_state must be contiguous.") + if init_state.device != q.device: + raise ValueError("cam_scan_func: init_state must be on the same device as q.") + load_init = 1 + else: + load_init = 0 + + if save_final_state: + final_state = torch.zeros(B * H, BLOCK_D, BLOCK_D, device=q.device, dtype=torch.float32) + save_final = 1 + else: + final_state = None + save_final = 0 + + out = torch.empty_like(q) + + dummy_state = torch.empty(1, device=q.device, dtype=torch.float32) + init_state_ptr = init_state if load_init else dummy_state + final_state_ptr = final_state if save_final else dummy_state + + _cam_scan_kernel[(B * H,)]( + q, + k, + v, + beta, + decay, + out, + dummy_state, + dummy_state, + init_state_ptr, + final_state_ptr, + H=H, + F=F_frames, + S=S, + D=D, + N=N, + REVERSE=1 if reverse else 0, + SAVE_STATES=0, + LOAD_INIT_STATE=load_init, + SAVE_FINAL_STATE=save_final, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_warps=num_warps, + num_stages=num_stages, + ) + if save_final_state: + return out, final_state + return out + + +def _run_cam_scan_fwd_save( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + reverse: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the forward scan with per-frame state snapshots saved. + + Used by :class:`CamScanFunction` to preserve the ``(state_pre, state_post)`` + snapshots that the Triton bwd kernel consumes. Snapshots are indexed by + ``q_frame`` (matching the existing fwd-kernel save logic), so the bwd + kernel can load them with the same ``q_frame`` derived in its + ``REVERSE``-aware iteration. + + Returns: + (out, state_pre, state_post). ``state_pre`` and ``state_post`` are + ``(B, H, F, BLOCK_D, BLOCK_D)`` fp32 with ``BLOCK_D = next_pow2(D)``. + Padding columns/rows past ``D`` are zero-masked on store. + """ + assert q.shape == k.shape == v.shape + B, H, D, N = q.shape + F_frames = beta.shape[2] + S = N // F_frames + assert beta.shape == (B, H, F_frames, S) + assert decay.shape == (B, H, F_frames) + 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 + + BLOCK_D = triton.next_power_of_2(D) + BLOCK_S = _DEFAULT_BLOCK_S + num_warps = 4 + num_stages = 1 + + out = torch.empty_like(q) + state_pre = torch.zeros(B * H, F_frames, BLOCK_D, BLOCK_D, device=q.device, dtype=torch.float32) + state_post = torch.zeros(B * H, F_frames, BLOCK_D, BLOCK_D, device=q.device, dtype=torch.float32) + dummy_state = torch.empty(1, device=q.device, dtype=torch.float32) + + _cam_scan_kernel[(B * H,)]( + q, + k, + v, + beta, + decay, + out, + state_pre, + state_post, + dummy_state, + dummy_state, + H=H, + F=F_frames, + S=S, + D=D, + N=N, + REVERSE=1 if reverse else 0, + SAVE_STATES=1, + LOAD_INIT_STATE=0, + SAVE_FINAL_STATE=0, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_warps=num_warps, + num_stages=num_stages, + ) + return out, state_pre, state_post + + +def _cam_scan_bwd_dispatch( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + state_pre: torch.Tensor, + state_post: torch.Tensor, + grad_out: torch.Tensor, + *, + reverse: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Launch ``_cam_scan_bwd_kernel`` and return ``(dq, dk, dv, dbeta, ddecay)``. + + All gradient outputs are fp32 contiguous, matching the dtype of the + forward inputs. ``grad_out`` is cast to fp32 before launching. + + When ``reverse=True``, ``kv_frame=0`` is never visited (only used in the + skipped first iter), so ``dk[..., 0, :]``, ``dv[..., 0, :]``, + ``dbeta[..., 0, :]`` and ``ddecay[..., 0]`` must remain zero. We pre-zero + every output buffer here so the kernel only needs to write the live slots. + """ + assert q.shape == k.shape == v.shape + B, H, D, N = q.shape + F_frames = beta.shape[2] + S = N // F_frames + + grad_out_f32 = grad_out.to(torch.float32).contiguous() + dq = torch.zeros_like(q) + dk = torch.zeros_like(k) + dv = torch.zeros_like(v) + dbeta = torch.zeros_like(beta) + ddecay = torch.zeros_like(decay) + + BLOCK_D = triton.next_power_of_2(D) + # For small S, ``next_pow2(S) < _DEFAULT_BLOCK_S`` — using the smaller value + # avoids zero-padding huge unused tiles into shared memory. + BLOCK_S = min(_DEFAULT_BLOCK_S, max(triton.next_power_of_2(S), 16)) + num_stages = 1 + REVERSE = 1 if reverse else 0 + + last_err: Exception | None = None + for num_warps in (4, 2, 1): + try: + _cam_scan_bwd_kernel[(B * H,)]( + q, + k, + v, + beta, + decay, + state_pre, + state_post, + grad_out_f32, + dq, + dk, + dv, + dbeta, + ddecay, + H=H, + F=F_frames, + S=S, + D=D, + N=N, + REVERSE=REVERSE, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_warps=num_warps, + num_stages=num_stages, + ) + return dq, dk, dv, dbeta, ddecay + except triton.runtime.errors.OutOfResources as exc: + last_err = exc + continue + raise RuntimeError("_cam_scan_bwd_kernel exhausted all num_warps choices: " + str(last_err)) + + +# ============================================================================= +# Section: Torch reference implementations used by fallback backward paths +# ============================================================================= +# These references replicate the Triton-kernel math (full-channel RMSNorm + +# ReLU + K-scale + 4x4 UCPE projmat + interleaved-pair real-valued RoPE, then +# numerator-only single-path delta-rule scan). They run in fp32 internally and +# cast outputs back to the input dtype, matching the kernels. + + +def _flip_and_shift(x: torch.Tensor, dim: int, shift_val: float) -> torch.Tensor: + """Flip ``x`` along ``dim`` and right-shift by one (pad with ``shift_val``). + + Matches the reference ``sana_gdn_blocks.flip_and_shift`` semantics. + """ + 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) + + +def _torch_cam_scan_single_chunk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + init_state: torch.Tensor | None = None, + return_final_state: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Pure-torch single-chunk delta-rule scan (numerator-only). + + Algebraically equivalent to ``torch_chunk_cam_single_path_delta_rule`` with + ``chunk_size >= T``; matches the Triton ``_cam_scan_kernel`` math exactly + (which also runs as a single chunk over all F frames). + + Args: + q, k, v: ``(B, H, D, N)`` fp32 contiguous. + beta: ``(B, H, F, S)`` or ``(B, H, F)`` fp32 contiguous. + decay: ``(B, H, F)`` fp32 contiguous. + + Returns: + out: ``(B, H, D, N)`` fp32 contiguous, ``N = F * S``. + """ + B, H, D, N = q.shape + if beta.ndim == 4: + T = beta.shape[2] + elif beta.ndim == 3: + T = beta.shape[2] + else: + raise ValueError(f"beta must be (B,H,F[,S]); got ndim={beta.ndim}") + if N % T != 0: + raise ValueError(f"N ({N}) must be divisible by 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) # (B, H, T, D, S) + + q_t = to_frame_seq(q) + k_t = to_frame_seq(k) + v_t = to_frame_seq(v) + + if beta.ndim == 4: + beta_view = beta.unsqueeze(3) # (B, H, T, 1, S) + else: + beta_view = beta.view(B, H, T, 1, 1) + decay_view = decay.view(B, H, T, 1, 1) + + eye = torch.eye(D, device=q.device, dtype=q.dtype).view(1, 1, 1, D, D) + + k_beta = k_t * beta_view + W = decay_view * (eye - torch.matmul(k_beta, k_t.transpose(-1, -2))) + U = torch.matmul(v_t * beta_view, k_t.transpose(-1, -2)) + + state = ( + torch.zeros(B, H, D, D, device=q.device, dtype=q.dtype) + if init_state is None + else init_state.to(device=q.device, dtype=q.dtype) + ) + s_kv_list: list[torch.Tensor] = [] + for t in range(T): + state = torch.matmul(state, W[:, :, t]) + U[:, :, t] + s_kv_list.append(state) + s_all = torch.stack(s_kv_list, dim=2) # (B, H, T, D, D) + + out_t = torch.matmul(s_all, q_t) # (B, H, T, D, S) + out = out_t.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + return (out, state) if return_final_state else out + + +def _torch_cam_scan_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + reverse: bool = False, +) -> torch.Tensor: + """Pure-torch reference for ``cam_scan_func`` supporting ``reverse=True``. + + For ``reverse=False`` this is the standard forward delta-rule scan. + + For ``reverse=True`` we emulate the Triton kernel's per-chunk + ``flip_and_shift`` semantics (q is flipped only; k/v/beta are + flip-and-shifted with pad value 0; decay is flip-and-shifted with pad + value 1; output is then flipped back along the time axis). + """ + if not reverse: + return _torch_cam_scan_single_chunk(q, k, v, beta, decay) + + B, H, D, N = q.shape + if beta.ndim == 4: + T = beta.shape[2] + elif beta.ndim == 3: + T = beta.shape[2] + else: + raise ValueError(f"beta must be (B,H,F[,S]); got ndim={beta.ndim}") + S = N // T + + def to_frame(x: torch.Tensor) -> torch.Tensor: + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) # (B, H, T, D, S) + + def from_frame(x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + + q_bwd = torch.flip(to_frame(q), dims=[2]) + k_bwd = _flip_and_shift(to_frame(k), dim=2, shift_val=0.0) + v_bwd = _flip_and_shift(to_frame(v), 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 = _torch_cam_scan_single_chunk( + from_frame(q_bwd), + from_frame(k_bwd), + from_frame(v_bwd), + beta_bwd, + decay_bwd, + ) + out_bwd_t = out_bwd.view(B, H, D, T, S) # already in (B, H, D, T, S) + return torch.flip(out_bwd_t, dims=[3]).reshape(B, H, D, N) + + +def _torch_cam_prep_reference( + 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, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + k_scale: float, + norm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure-torch reference for ``cam_prep_func`` matching ``_cam_prep_kernel``. + + Replicates exactly: + - Full-channel (over ``H*D``) RMSNorm + per-channel weight on Q, K. + - ReLU on Q, K. + - K-scale on K. + - 4x4 UCPE projmat on first ``D/2`` dims (Q via ``proj_q``, + K and V via ``proj_kv``). + - Interleaved-pair real-valued RoPE on second ``D/2`` dims using + ``rope_cos`` / ``rope_sin`` (the same tables passed to the kernel). + - ``inflation_sq = ||k_post_ucpe||^2 / ||k_pre_ucpe||^2`` per (B, H, N), + with the same ``clamp_min(1e-12)`` floor as ``cam_prep_func``. + + All math runs in fp32 internally; outputs ``(q, k, v)`` are cast back to + ``q_raw.dtype`` and ``inflation_sq`` is fp32. + """ + B, N, H, D = q_raw.shape + if D % 2 != 0: + raise ValueError(f"D ({D}) must be even.") + if (D // 2) % 4 != 0: + raise ValueError(f"D/2 ({D // 2}) must be divisible by 4 (UCPE projmat).") + C = H * D + D_half = D // 2 + n_groups = D_half // 4 + + q32 = q_raw.float() + k32 = k_raw.float() + v32 = v_raw.float() + + # ---- Full-channel RMSNorm + per-channel weight (Q, K only) ---- + q_inv_rms = torch.rsqrt((q32 * q32).sum(dim=(-1, -2)) / C + norm_eps) # (B, N) + k_inv_rms = torch.rsqrt((k32 * k32).sum(dim=(-1, -2)) / C + norm_eps) + q_nw = q_norm_weight.float().view(1, 1, H, D) + k_nw = k_norm_weight.float().view(1, 1, H, D) + q_normed = q32 * q_inv_rms.view(B, N, 1, 1) * q_nw + k_normed = k32 * k_inv_rms.view(B, N, 1, 1) * k_nw + + # ---- ReLU + K-scale ---- + q_normed = torch.relu(q_normed) + k_normed = torch.relu(k_normed) * k_scale + + # ---- Pre-UCPE ||k||^2 over the full D dim ---- + pre_k_sq_BNH = (k_normed * k_normed).sum(dim=-1) # (B, N, H) + + # ---- UCPE 4x4 projmat on first half ---- + q_first = q_normed[..., :D_half].reshape(B, N, H, n_groups, 4) + k_first = k_normed[..., :D_half].reshape(B, N, H, n_groups, 4) + v_first = v32[..., :D_half].reshape(B, N, H, n_groups, 4) + + # out[b,n,h,g,i] = sum_j P[b,n,i,j] * x[b,n,h,g,j] + # einsum: 'bnij,bnhgj->bnhgi' + proj_q_f = proj_q.float() + proj_kv_f = proj_kv.float() + q_first_proj = torch.einsum("bnij,bnhgj->bnhgi", proj_q_f, q_first).reshape(B, N, H, D_half) + k_first_proj = torch.einsum("bnij,bnhgj->bnhgi", proj_kv_f, k_first).reshape(B, N, H, D_half) + v_first_proj = torch.einsum("bnij,bnhgj->bnhgi", proj_kv_f, v_first).reshape(B, N, H, D_half) + + # ---- Interleaved-pair real-valued RoPE on second half ---- + # Kernel form: y[d] = x[d]*rope_cos[d] + x[d^1]*rope_sin[d] + # where rope_cos/rope_sin come from _prepare_ucpe_rope_tables. + q_second = q_normed[..., D_half:] + k_second = k_normed[..., D_half:] + v_second = v32[..., D_half:] + + def _pair_swap(x: torch.Tensor) -> torch.Tensor: + # Swap consecutive pairs along the last dim: (..., D_half) where D_half is even. + # x[..., 2i] <-> x[..., 2i+1]. + x_pairs = x.unflatten(-1, (D_half // 2, 2)) + x_swapped = x_pairs.flip(-1) + return x_swapped.flatten(-2) + + cos_b = rope_cos.float().view(1, N, 1, D_half) + sin_b = rope_sin.float().view(1, N, 1, D_half) + q_rope = q_second * cos_b + _pair_swap(q_second) * sin_b + k_rope = k_second * cos_b + _pair_swap(k_second) * sin_b + v_rope = v_second * cos_b + _pair_swap(v_second) * sin_b + + # ---- Reassemble (B, N, H, D) and post-UCPE k norm ---- + q_out_BNHD = torch.cat([q_first_proj, q_rope], dim=-1) + k_out_BNHD = torch.cat([k_first_proj, k_rope], dim=-1) + v_out_BNHD = torch.cat([v_first_proj, v_rope], dim=-1) + + post_k_sq_BNH = (k_out_BNHD * k_out_BNHD).sum(dim=-1) # (B, N, H) + + out_dtype = q_raw.dtype + q_out = q_out_BNHD.to(out_dtype).permute(0, 2, 3, 1).contiguous() + k_out = k_out_BNHD.to(out_dtype).permute(0, 2, 3, 1).contiguous() + v_out = v_out_BNHD.to(out_dtype).permute(0, 2, 3, 1).contiguous() + + pre_k_sq = pre_k_sq_BNH.permute(0, 2, 1).contiguous() # (B, H, N) + post_k_sq = post_k_sq_BNH.permute(0, 2, 1).contiguous() + inflation_sq = post_k_sq.clamp_min(1e-12) / pre_k_sq.clamp_min(1e-12) + return q_out, k_out, v_out, inflation_sq + + +# ============================================================================= +# Section: Autograd-enabled wrappers +# ============================================================================= + + +def _cam_scan_torch_fallback_backward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + needs: tuple[bool, bool, bool, bool, bool], + grad_out: torch.Tensor, + reverse: bool, +) -> list[torch.Tensor | None]: + """Recompute the cam-branch scan via the torch reference and return grads. + + Used when ``CAM_SCAN_BWD_FALLBACK=1`` forces the torch-recompute backward + path. + """ + detached = [] + for tensor, need in zip((q, k, v, beta, decay), needs): + t = tensor.detach() + if need: + t = t.requires_grad_(True) + detached.append(t) + active = [t for t in detached if t.requires_grad] + + with torch.enable_grad(): + q_d, k_d, v_d, beta_d, decay_d = detached + ref_out = _torch_cam_scan_reference(q_d, k_d, v_d, beta_d, decay_d, reverse=reverse) + if active: + active_grads = torch.autograd.grad( + outputs=ref_out, + inputs=tuple(active), + grad_outputs=grad_out.to(ref_out.dtype), + allow_unused=True, + ) + else: + active_grads = [] + + grads: list[torch.Tensor | None] = [] + active_iter = iter(active_grads) + for tensor in detached: + grads.append(next(active_iter) if tensor.requires_grad else None) + return grads + + +class CamScanFunction(torch.autograd.Function): + """Autograd ``Function`` wrapping ``cam_scan_func``. + + Forward calls the Triton ``_cam_scan_kernel`` with ``SAVE_STATES=1`` so + per-frame state snapshots (``state_pre[q_frame]``, ``state_post[q_frame]``) + are kept for the backward pass. Backward runs the true Triton bwd + kernel (``_cam_scan_bwd_kernel``) for both ``reverse=False`` and + ``reverse=True``, replaying the recurrence in reverse time using the + saved snapshots. + + Set ``CAM_SCAN_BWD_FALLBACK=1`` to force the torch-recompute backward + validation path. + """ + + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + reverse: bool, + ) -> torch.Tensor: + ctx.set_materialize_grads(False) + ctx.reverse = bool(reverse) + + force_torch_fallback = os.environ.get("CAM_SCAN_BWD_FALLBACK", "0") == "1" + ctx.use_triton_bwd = not force_torch_fallback + + if ctx.use_triton_bwd: + out, state_pre, state_post = _run_cam_scan_fwd_save(q, k, v, beta, decay, reverse=ctx.reverse) + ctx.save_for_backward(q, k, v, beta, decay, state_pre, state_post) + return out + + # Torch-fallback backward path: don't bother saving state snapshots. + ctx.save_for_backward(q, k, v, beta, decay) + return cam_scan_func(q, k, v, beta, decay, reverse=reverse) + + @staticmethod + def backward(ctx, grad_out): # type: ignore[override] + if grad_out is None: + return (None, None, None, None, None, None) + + if ctx.use_triton_bwd: + q, k, v, beta, decay, state_pre, state_post = ctx.saved_tensors + needs = ctx.needs_input_grad[:5] # q, k, v, beta, decay + dq, dk, dv, dbeta, ddecay = _cam_scan_bwd_dispatch( + q, + k, + v, + beta, + decay, + state_pre, + state_post, + grad_out, + reverse=ctx.reverse, + ) + grads: list[torch.Tensor | None] = [ + dq if needs[0] else None, + dk if needs[1] else None, + dv if needs[2] else None, + dbeta if needs[3] else None, + ddecay if needs[4] else None, + ] + return (*grads, None) + + # Env-var torch-recompute backward. + q, k, v, beta, decay = ctx.saved_tensors + needs = ctx.needs_input_grad[:5] + grads = _cam_scan_torch_fallback_backward(q, k, v, beta, decay, needs, grad_out, ctx.reverse) + return (*grads, None) + + +def cam_scan_func_with_grad( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + reverse: bool = False, +) -> torch.Tensor: + """Autograd-enabled wrapper around :func:`cam_scan_func`. + + Forward is identical to :func:`cam_scan_func`; backward is computed via + a torch reference (``_torch_cam_scan_reference``). Use this in training + paths where any of ``q, k, v, beta, decay`` may require gradients. + + Inference paths can keep calling :func:`cam_scan_func` directly to avoid + the small autograd bookkeeping overhead. + """ + return CamScanFunction.apply(q, k, v, beta, decay, reverse) + + +def _cam_prep_torch_fallback_backward( + 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, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + needs: tuple[bool, ...], + grad_q: torch.Tensor | None, + grad_k: torch.Tensor | None, + grad_v: torch.Tensor | None, + grad_inflation_sq: torch.Tensor | None, + k_scale: float, + norm_eps: float, +) -> list[torch.Tensor | None]: + """Recompute the cam-branch prep via the torch reference and return grads. + + Used when any of ``proj_q / proj_kv / rope_cos / rope_sin`` requests a + gradient (the Triton kernel does not produce those grads), or when + ``CAM_PREP_BWD_FALLBACK=1`` forces the torch-recompute backward path. + + Args: + q_raw, k_raw, ..., rope_sin: the nine tensor inputs of + :func:`cam_prep_func` (in the same order as + :class:`CamPrepFunction.forward`'s arg list). + needs: ``ctx.needs_input_grad[:9]`` — boolean per-input flags. + grad_q, grad_k, grad_v, grad_inflation_sq: upstream gradients. + k_scale, norm_eps: scalar fwd args. + + Returns: + A 9-element list of ``torch.Tensor | None`` aligned with the + ``saved`` tuple. Entries that didn't request a gradient are ``None``. + """ + saved = ( + q_raw, + k_raw, + v_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + ) + detached: list[torch.Tensor] = [] + for tensor, need in zip(saved, needs): + t = tensor.detach() + if need: + t = t.requires_grad_(True) + detached.append(t) + active = [t for t in detached if t.requires_grad] + + with torch.enable_grad(): + (q_d, k_d, v_d, qnw_d, knw_d, pq_d, pkv_d, rc_d, rs_d) = detached + ref_q, ref_k, ref_v, ref_inf = _torch_cam_prep_reference( + q_d, + k_d, + v_d, + q_norm_weight=qnw_d, + k_norm_weight=knw_d, + proj_q=pq_d, + proj_kv=pkv_d, + rope_cos=rc_d, + rope_sin=rs_d, + k_scale=k_scale, + norm_eps=norm_eps, + ) + + outputs = [] + grad_outputs = [] + if grad_q is not None: + outputs.append(ref_q) + grad_outputs.append(grad_q.to(ref_q.dtype)) + if grad_k is not None: + outputs.append(ref_k) + grad_outputs.append(grad_k.to(ref_k.dtype)) + if grad_v is not None: + outputs.append(ref_v) + grad_outputs.append(grad_v.to(ref_v.dtype)) + if grad_inflation_sq is not None: + outputs.append(ref_inf) + grad_outputs.append(grad_inflation_sq.to(ref_inf.dtype)) + + if active and outputs: + active_grads = torch.autograd.grad( + outputs=tuple(outputs), + inputs=tuple(active), + grad_outputs=tuple(grad_outputs), + allow_unused=True, + ) + else: + active_grads = [] + + grads: list[torch.Tensor | None] = [] + active_iter = iter(active_grads) + for tensor in detached: + grads.append(next(active_iter) if tensor.requires_grad else None) + return grads + + +class CamPrepFunction(torch.autograd.Function): + """Autograd ``Function`` wrapping ``cam_prep_func``. + + Forward calls the fused Triton ``_cam_prep_kernel`` via + :func:`_run_cam_prep_fwd_save` so the per-token ``inv_rms`` / + ``k_pre_sq`` / ``k_post_sq`` snapshots required by the bwd kernel are + preserved alongside the standard outputs. + + Backward runs the true Triton bwd kernel via + :func:`_cam_prep_bwd_dispatch` for the standard training path + (``q_raw``, ``k_raw``, ``v_raw``, ``q_norm_weight``, ``k_norm_weight`` + only request grads). The Triton path implements: + + * RoPE^T, UCPE^T, K-scale^T, and ReLU mask in a single fused kernel + (one program per ``(b, n, h)``); + * ``grad_inflation_sq`` chain through ``k_post_sq`` (added into + ``eff_dO_k``) and ``k_pre_sq`` (direct contribution to + ``d_k_post_kscale``), with ``clamp_min(1e-12)`` indicators honored; + * The full-channel RMSNorm bwd (per-token cross-head reduction) is + done in PyTorch on the kernel's ``d_q_post_norm`` / + ``d_k_post_norm`` intermediates — see + :func:`_cam_prep_bwd_dispatch` for details. + + The torch-recompute fallback (running :func:`_torch_cam_prep_reference` + under autograd) is selected when any of ``proj_q / proj_kv / rope_cos / + rope_sin`` requests a gradient (the Triton path emits ``None`` for those + slots) or when ``CAM_PREP_BWD_FALLBACK=1`` is set. + """ + + @staticmethod + def forward( + ctx, + 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, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + k_scale: float, + norm_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ctx.set_materialize_grads(False) + ctx.k_scale = float(k_scale) + ctx.norm_eps = float(norm_eps) + + force_torch_fallback = os.environ.get("CAM_PREP_BWD_FALLBACK", "0") == "1" + ctx.use_triton_bwd = not force_torch_fallback + + if ctx.use_triton_bwd: + ( + q_out, + k_out, + v_out, + inflation_sq, + q_inv_rms, + k_inv_rms, + k_pre_sq, + k_post_sq, + ) = _run_cam_prep_fwd_save( + q_raw, + k_raw, + v_raw, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + proj_q=proj_q, + proj_kv=proj_kv, + rope_cos=rope_cos, + rope_sin=rope_sin, + k_scale=k_scale, + norm_eps=norm_eps, + ) + ctx.save_for_backward( + q_raw, + k_raw, + v_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + q_inv_rms, + k_inv_rms, + k_pre_sq, + k_post_sq, + k_out, + ) + else: + q_out, k_out, v_out, inflation_sq = cam_prep_func( + q_raw, + k_raw, + v_raw, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + proj_q=proj_q, + proj_kv=proj_kv, + rope_cos=rope_cos, + rope_sin=rope_sin, + k_scale=k_scale, + norm_eps=norm_eps, + ) + ctx.save_for_backward( + q_raw, + k_raw, + v_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + ) + return q_out, k_out, v_out, inflation_sq + + @staticmethod + def backward(ctx, grad_q, grad_k, grad_v, grad_inflation_sq): # type: ignore[override] + if grad_q is None and grad_k is None and grad_v is None and grad_inflation_sq is None: + return tuple([None] * 11) + + needs = ctx.needs_input_grad[:9] # nine tensor inputs + # If anyone outside (q_raw, k_raw, v_raw, q_norm_weight, k_norm_weight) + # requests a grad, the Triton bwd cannot handle it — fall back to the + # torch reference. + triton_bwd_supported_needs = needs[:5] + proj_or_rope_needs_grad = any(needs[5:]) + + if ctx.use_triton_bwd and not proj_or_rope_needs_grad: + ( + q_raw, + k_raw, + v_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + q_inv_rms, + k_inv_rms, + k_pre_sq, + k_post_sq, + k_out, + ) = ctx.saved_tensors + + ( + dq_raw, + dk_raw, + dv_raw, + dq_norm_weight, + dk_norm_weight, + ) = _cam_prep_bwd_dispatch( + q_raw, + k_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + q_inv_rms, + k_inv_rms, + k_pre_sq, + k_post_sq, + k_out, + grad_q=grad_q, + grad_k=grad_k, + grad_v=grad_v, + grad_inflation_sq=grad_inflation_sq, + k_scale=ctx.k_scale, + ) + grads: list[torch.Tensor | None] = [ + dq_raw if triton_bwd_supported_needs[0] else None, + dk_raw if triton_bwd_supported_needs[1] else None, + dv_raw if triton_bwd_supported_needs[2] else None, + dq_norm_weight if triton_bwd_supported_needs[3] else None, + dk_norm_weight if triton_bwd_supported_needs[4] else None, + None, # proj_q + None, # proj_kv + None, # rope_cos + None, # rope_sin + ] + return (*grads, None, None) + + # Torch fallback path. ``ctx.saved_tensors`` holds either 9 (legacy + # forward) or 14 (Triton fwd save) tensors — slice the leading nine. + saved = ctx.saved_tensors[:9] + ( + q_raw, + k_raw, + v_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + ) = saved + grads = _cam_prep_torch_fallback_backward( + q_raw, + k_raw, + v_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + needs=needs, + grad_q=grad_q, + grad_k=grad_k, + grad_v=grad_v, + grad_inflation_sq=grad_inflation_sq, + k_scale=ctx.k_scale, + norm_eps=ctx.norm_eps, + ) + # Two trailing None for non-tensor scalars (k_scale, norm_eps). + return (*grads, None, None) + + +def cam_prep_func_with_grad( + 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, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + k_scale: float, + norm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Autograd-enabled wrapper around :func:`cam_prep_func`. + + Forward is identical to :func:`cam_prep_func`; backward is computed via a + torch reference (``_torch_cam_prep_reference``). Use this in training + paths so gradients flow back through Q/K/V projection inputs and through + the RMSNorm weights. + + Inference paths can keep calling :func:`cam_prep_func` directly. + """ + return CamPrepFunction.apply( + q_raw, + k_raw, + v_raw, + q_norm_weight, + k_norm_weight, + proj_q, + proj_kv, + rope_cos, + rope_sin, + k_scale, + norm_eps, + ) + + +__all__ = [ + "CamPrepFunction", + "CamScanFunction", + "_cam_prep_bwd_dispatch", + "_cam_prep_bwd_kernel", + "_cam_prep_kernel", + "_cam_prep_torch_fallback_backward", + "_cam_scan_bwd_dispatch", + "_cam_scan_bwd_kernel", + "_cam_scan_kernel", + "_invert_SE3", + "_precompute_cam_inv_rms", + "_prepare_ucpe_rope_tables", + "_process_camera_conditions_raymats_only", + "_run_cam_prep_fwd_save", + "_run_cam_scan_fwd_save", + "_torch_cam_prep_reference", + "cam_prep_func", + "cam_prep_func_with_grad", + "cam_scan_func", + "cam_scan_func_with_grad", +] diff --git a/integrations/sana/sana_wm/ops/fused_gdn.py b/integrations/sana/sana_wm/ops/fused_gdn.py new file mode 100644 index 000000000..447e0f5ca --- /dev/null +++ b/integrations/sana/sana_wm/ops/fused_gdn.py @@ -0,0 +1,2249 @@ +"""Fused-BiGDN Triton kernels used by SANA-WM GDN attention blocks. + +Includes the unified forward kernel, backward kernels, RoPE/RMS helpers, and +autograd wrappers used by the Triton GDN attention blocks. + +Precision knob: env var ``FUSED_GDN_PRECISION`` or ``PRECISION_OVERRIDE``: + 0=IEEE fp32 dots, 1=TF32, 2=bf16 TC + fp32 state [default], 3=bf16 TC + bf16 state. +""" + +# ruff: noqa: E501 + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +# ===================================================================== +# 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_launch(D: int, beta: torch.Tensor, decay: torch.Tensor) -> tuple: + """Shared launcher preamble. + + Returns (BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta_c, decay_c). + ``beta_c`` / ``decay_c`` are the contiguous copies the kernel needs. + """ + BLOCK_D = triton.next_power_of_2(D) + cfg = _kcfg() + BLOCK_S = cfg["BLOCK_S"] + _, dot_prec, state_fp32, nw = _resolve_launch_config() + return BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta.contiguous(), decay.contiguous() + + +# ===================================================================== +# Unified forward Triton Mega-Kernel (inference-only variant) +# ===================================================================== +# Fuses: RMSNorm + ReLU + k_scale + RoPE + BiGDN recurrence. +# +# Inputs: +# qkv (B, N, 3, H, D) interleaved — strides passed explicitly. +# beta (B, H, F, S), decay (B, H, F) contiguous. +# q_norm_w, k_norm_w (H*D,) full-channel — only read when QK_NORM=1. +# rope_cos, rope_sin (N, D) contiguous. +# q_inv_rms, k_inv_rms (B, N) full-channel — only read when USE_PRECOMPUTED_RMS=1. +# +# Outputs: +# out (B, N, H, D) = num / (den + eps) — unused by BiGDN wrappers. +# num (B, N, H, D) = numerator before divide (summed across directions). +# den (B, H, N) = denominator before divide (summed across directions). +# +# NOTE (inference-only build): upstream also supports SAVE_STATE, +# LOAD_INIT_STATE, SAVE_FINAL_STATE for training backward / state caching. +# Those constexpr branches are preserved in the kernel so the source stays +# 1-for-1 with upstream (they compile away when launched with flags=0). + + +@triton.jit +def _fused_gdn_kernel( + # ---- interleaved QKV : (B, N, 3, H, D) ---- + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + # ---- gates ---- + beta_ptr, + decay_ptr, + # ---- inv-RMS (B, N) — only read when USE_PRECOMPUTED_RMS=1 ---- + q_inv_rms_ptr, + k_inv_rms_ptr, + # ---- norm weights (H*D,) full-channel — only read when QK_NORM=1 ---- + q_norm_w_ptr, + k_norm_w_ptr, + # ---- RoPE tables (N, D) contiguous ---- + rope_cos_ptr, + rope_sin_ptr, + # ---- outputs ---- + out_ptr, # (B, N, H, D) + num_ptr, # (B, N, H, D) + den_ptr, # (B, H, N) + # ---- saved-state dummies (unused in this build but kept for signature parity) ---- + saved_state_ptr, + saved_z_ptr, + saved_state_curr_ptr, + saved_z_curr_ptr, + init_state_kv_ptr, + init_state_z_ptr, + final_state_kv_ptr, + final_state_z_ptr, + # ---- scalars / dims ---- + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + K_SCALE, + NORM_EPS: tl.constexpr, + EPS: tl.constexpr, + QK_NORM: tl.constexpr, + USE_PRECOMPUTED_RMS: tl.constexpr, + STATE_FP32: tl.constexpr, + DOT_PRECISION: tl.constexpr, + REVERSE: tl.constexpr, + SAVE_STATE: tl.constexpr, + LOAD_INIT_STATE: tl.constexpr, + SAVE_FINAL_STATE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + # ---- dot product precision / operand dtype ---- + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + # ---- program → (batch, head) ---- + pid = tl.program_id(0) + pid_b = pid // H + pid_h = pid % H + N = F * S + bh = pid_b * H + pid_h + + # ---- base pointers ---- + qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h + out_bh = out_ptr + pid_b * (N * H * D) + pid_h * D + num_bh = num_ptr + pid_b * (N * H * D) + pid_h * D + den_bh = den_ptr + bh * N + beta_bh = beta_ptr + bh * (F * S) + decay_bh = decay_ptr + bh * F + if SAVE_STATE: + st_bh = saved_state_ptr + bh * F * BLOCK_D * BLOCK_D + sz_bh = saved_z_ptr + bh * F * BLOCK_D + stc_bh = saved_state_curr_ptr + bh * F * BLOCK_D * BLOCK_D + szc_bh = saved_z_curr_ptr + bh * F * BLOCK_D + + # ---- D-index helpers ---- + 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 + D_inv = 1.0 / D + + # ---- full-channel norm weights (only when QK_NORM=1) ---- + nw_offset = pid_h * D + if QK_NORM: + q_nw = tl.load(q_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) + k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) + q_nw_pair = tl.load(q_norm_w_ptr + nw_offset + offs_d_pair, mask=mask_d_pair, 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) + + k_scale = K_SCALE + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + mask_dd = mask_d[:, None] & mask_d[None, :] + + # ---- double-buffer state ---- + if LOAD_INIT_STATE: + init_kv_bh = init_state_kv_ptr + bh * BLOCK_D * BLOCK_D + state_curr = tl.load(init_kv_bh + offs_dd, mask=mask_dd, other=0.0).to(tl.float32) + init_z_bh = init_state_z_ptr + bh * BLOCK_D + state_z_curr = tl.load(init_z_bh + offs_d, mask=mask_d, other=0.0).to(tl.float32) + else: + state_curr = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + state_z_curr = tl.zeros([BLOCK_D], dtype=tl.float32) + if STATE_FP32: + state_prev = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + else: + state_prev = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.bfloat16) + state_z_prev = tl.zeros([BLOCK_D], dtype=tl.float32) + + # ======================================================== + # Temporal loop — serial over F + # ======================================================== + for f_iter in range(F): + if REVERSE: + q_frame = F - 1 - f_iter + kv_frame = F - f_iter if f_iter > 0 else 0 # unused at f=0 + skip_update = f_iter == 0 + else: + q_frame = f_iter + kv_frame = f_iter + skip_update = False + + # ---- decay + state snapshot ---- + if REVERSE and f_iter == 0: + g = 1.0 + else: + g = tl.load(decay_bh + kv_frame).to(tl.float32) + state_curr = state_curr * g + state_z_curr = state_z_curr * g + if STATE_FP32: + state_prev = state_curr + 0.0 + else: + state_prev = state_curr.to(tl.bfloat16) + state_z_prev = state_z_curr + + if SAVE_STATE: + st_f = st_bh + q_frame * BLOCK_D * BLOCK_D + tl.store(st_f + offs_dd, state_prev, mask=mask_dd) + tl.store(sz_bh + q_frame * BLOCK_D + offs_d, state_z_prev, mask=mask_d) + + # ------------------------------------------ + # Pass 1 — State Accumulation + # ------------------------------------------ + if skip_update == False: + kv_n_base = kv_frame * S + f_beta = beta_bh + kv_frame * 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, :] + mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] + n_idx = kv_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) + + if QK_NORM: + if USE_PRECOMPUTED_RMS: + k_inv_rms = tl.load(k_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0).to(tl.float32) + else: + k_var = tl.sum(K_raw * K_raw, axis=1) * D_inv + k_inv_rms = 1.0 / tl.sqrt(k_var + NORM_EPS) + K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] + else: + K_normed = K_raw + K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale + + k_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d_pair[None, :] * stride_d + K_pair_raw = tl.load(k_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) + if QK_NORM: + K_pair_normed = K_pair_raw * k_inv_rms[:, None] * k_nw_pair[None, :] + else: + K_pair_normed = K_pair_raw + 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 + + bt = tl.load(f_beta + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + K_rot_dc = K_rot.to(dot_dtype) + V_pred = tl.dot(K_rot_dc, state_prev.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + dv = (V_raw - V_pred) * bt[:, None] + state_curr += tl.dot(tl.trans(K_rot), dv, out_dtype=tl.float32, input_precision="tf32") + + z_hat = tl.sum(K * state_z_prev[None, :], axis=1) + dz = (1.0 - z_hat) * bt + state_z_curr += tl.sum(K * dz[:, None], axis=0) + + if SAVE_STATE: + stc_f = stc_bh + q_frame * BLOCK_D * BLOCK_D + tl.store(stc_f + offs_dd, state_curr, mask=mask_dd) + tl.store(szc_bh + q_frame * BLOCK_D + offs_d, state_z_curr, mask=mask_d) + + # ------------------------------------------ + # Pass 2 — Output (reads state_curr, inclusive) + # ------------------------------------------ + state_out = state_curr.to(dot_dtype) + state_z_out = state_z_curr + q_n_base = q_frame * 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, :] + mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] + n_idx = q_n_base + offs_s + + q_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d[None, :] * stride_d + q_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d_pair[None, :] * stride_d + Q_raw = tl.load(q_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + Q_pair_raw = tl.load(q_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) + + if QK_NORM: + if USE_PRECOMPUTED_RMS: + q_inv_rms = tl.load(q_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0).to(tl.float32) + else: + q_var = tl.sum(Q_raw * Q_raw, axis=1) * D_inv + q_inv_rms = 1.0 / tl.sqrt(q_var + NORM_EPS) + Q_normed = Q_raw * q_inv_rms[:, None] * q_nw[None, :] + Q_pair_normed = Q_pair_raw * q_inv_rms[:, None] * q_nw_pair[None, :] + else: + Q_normed = Q_raw + Q_pair_normed = Q_pair_raw + Q = tl.where(Q_normed > 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), state_out, out_dtype=tl.float32, input_precision=dot_ip) + den = tl.sum(Q * state_z_out[None, :], axis=1) + + result = num / (den[:, None] + EPS) + out_ptrs = out_bh + n_idx[:, None] * (H * D) + offs_d[None, :] + num_ptrs = num_bh + n_idx[:, None] * (H * D) + offs_d[None, :] + tl.store(out_ptrs, result.to(tl.bfloat16), mask=mask_sd) + tl.store(num_ptrs, num.to(tl.bfloat16), mask=mask_sd) + tl.store(den_bh + n_idx, den.to(tl.bfloat16), mask=mask_s) + + if SAVE_FINAL_STATE: + final_kv_bh = final_state_kv_ptr + bh * BLOCK_D * BLOCK_D + tl.store(final_kv_bh + offs_dd, state_curr, mask=mask_dd) + final_z_bh = final_state_z_ptr + bh * BLOCK_D + tl.store(final_z_bh + offs_d, state_z_curr, mask=mask_d) + + +# ===================================================================== +# Python wrappers +# ===================================================================== + + +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. + """ + 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 + + +@triton.jit +def _fused_bidi_merge_kernel( + num_fwd_ptr, + num_bwd_ptr, + den_fwd_ptr, + den_bwd_ptr, + gate_ptr, + out_ptr, + B, + N, + H, + D, + eps, + snum_b, + snum_n, + snum_h, + snum_d, + sden_b, + sden_h, + sden_n, + APPLY_GATE: tl.constexpr, + PRE_SUMMED: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_bh = tl.program_id(0) + pid_n = tl.program_id(1) + b = pid_bh // H + h = pid_bh % H + + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_D) + mask_n = offs_n < N + mask_d = offs_d < D + mask_nd = mask_n[:, None] & mask_d[None, :] + + num_base = b * snum_b + offs_n[:, None] * snum_n + h * snum_h + offs_d[None, :] * snum_d + nf = tl.load(num_fwd_ptr + num_base, mask=mask_nd, other=0.0).to(tl.float32) + den_base = b * sden_b + h * sden_h + offs_n * sden_n + df = tl.load(den_fwd_ptr + den_base, mask=mask_n, other=0.0).to(tl.float32) + + if PRE_SUMMED: + num_total = nf + den_total = df + eps + else: + nb = tl.load(num_bwd_ptr + num_base, mask=mask_nd, other=0.0).to(tl.float32) + db = tl.load(den_bwd_ptr + den_base, mask=mask_n, other=0.0).to(tl.float32) + num_total = nf + nb + den_total = df + db + eps + out_val = num_total / den_total[:, None] + + if APPLY_GATE: + g = tl.load(gate_ptr + num_base, mask=mask_nd, other=0.0).to(tl.float32) + silu_g = g * (1.0 / (1.0 + tl.exp(-g))) + out_val = out_val * silu_g + + tl.store(out_ptr + num_base, out_val.to(tl.bfloat16), mask=mask_nd) + + +def fused_bidi_merge( + num_fwd: torch.Tensor, + num_bwd: torch.Tensor | None, + den_fwd: torch.Tensor, + den_bwd: torch.Tensor | None, + eps: float, + gate: torch.Tensor | None = None, +) -> torch.Tensor: + pre_summed = num_bwd is None + assert (num_bwd is None) == (den_bwd is None), "num_bwd/den_bwd must both be None or both provided" + if not pre_summed: + assert num_fwd.shape == num_bwd.shape and den_fwd.shape == den_bwd.shape + assert num_fwd.dtype == num_bwd.dtype and den_fwd.dtype == den_bwd.dtype + B, N, H, D = num_fwd.shape + out = torch.empty( + B, N, H, D, device=num_fwd.device, dtype=(torch.float32 if num_fwd.dtype == torch.float32 else torch.bfloat16) + ) + BLOCK_D = triton.next_power_of_2(D) + BLOCK_N = 64 + grid = (B * H, triton.cdiv(N, BLOCK_N)) + if gate is not None: + assert gate.shape == (B, N, H, D), f"gate shape {gate.shape} != {(B, N, H, D)}" + gate_arg = gate + apply_gate = 1 + else: + gate_arg = num_fwd + apply_gate = 0 + num_bwd_arg = num_bwd if num_bwd is not None else num_fwd + den_bwd_arg = den_bwd if den_bwd is not None else den_fwd + _fused_bidi_merge_kernel[grid]( + num_fwd, + num_bwd_arg, + den_fwd, + den_bwd_arg, + gate_arg, + out, + B, + N, + H, + D, + float(eps), + num_fwd.stride(0), + num_fwd.stride(1), + num_fwd.stride(2), + num_fwd.stride(3), + den_fwd.stride(0), + den_fwd.stride(1), + den_fwd.stride(2), + APPLY_GATE=apply_gate, + PRE_SUMMED=1 if pre_summed else 0, + BLOCK_N=BLOCK_N, + BLOCK_D=BLOCK_D, + ) + return out + + +# ===================================================================== +# Single-direction GDN entry point (delegates to chunkwise) +# ===================================================================== + + +def fused_gdn_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,) = (H*D,) 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, + reverse: bool = False, + init_state_kv: torch.Tensor | None = None, + init_state_z: torch.Tensor | None = None, + save_final_state: bool = False, +) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """One direction of fused BiGDN via the unified kernel. + + Args: + qkv .. eps: see kernel signature. + reverse: forward (False) or anti-causal (True) scan. + init_state_kv: optional ``(B*H, BLOCK_D, BLOCK_D)`` fp32 contiguous + tensor holding the forward-scan KV state at the END of a prefix + sequence (i.e., AFTER the prefix's last update, BEFORE any further + decay applied by this call). When provided, the kernel resumes the + scan from this state instead of zero. ``BLOCK_D = next_pow2(D)``. + Only the top-left ``D x D`` submatrix of the tile is read. + init_state_z: optional ``(B*H, BLOCK_D)`` fp32 contiguous companion + for the Z denominator state. Must be provided iff ``init_state_kv`` + is provided. + save_final_state: when True, allocate fresh fp32 zero buffers for the + final KV / Z state (after the last frame's update) and pass them to + the kernel for write-out. Returns the buffers as additional outputs. + + Returns: + ``(num, den)`` — bf16 numerator ``(B, N, H, D)`` and denominator + ``(B, H, N)`` before divide. + + When ``save_final_state=True``, also returns + ``(final_state_kv, final_state_z)`` fp32 with shapes + ``(B*H, BLOCK_D, BLOCK_D)`` and ``(B*H, BLOCK_D)``. + + Raises: + NotImplementedError: if any state I/O argument is set together with + ``reverse=True``. The kernel supports state passing in both + directions, but state I/O is only defined for the forward direction + here to avoid silent misuse. + """ + # Dispatch both the stateless bidi case and the stateful forward path to + # chunkwise so split-equivalence uses one numeric implementation. + # Bypass via env: FUSED_GDN_FORCE_LEGACY=1. + if os.environ.get("FUSED_GDN_FORCE_LEGACY", "0") != "1": + from .fused_gdn_chunkwise import ( + fused_gdn_func_chunkwise, + fused_gdn_stateful_chunkwise, + ) + + # Validate state I/O args upfront — preserves the legacy fused_gdn_func's + # validation contract (callers depend on these specific ValueError / + # NotImplementedError signatures, e.g., test_state_validation). + if (init_state_kv is None) != (init_state_z is None): + raise ValueError( + "fused_gdn_func: init_state_kv and init_state_z must be provided together " + "(both None or both fp32 tensors)." + ) + if reverse and (init_state_kv is not None or save_final_state): + raise NotImplementedError( + "fused_gdn_func: state passing (init_state_kv / init_state_z / " + "save_final_state) is only supported for the forward direction " + "(reverse=False)." + ) + if init_state_kv is not None: + B_q, _N, _three, H_q, D_q = qkv.shape + BLOCK_D_q = triton.next_power_of_2(D_q) + expected_kv = (B_q * H_q, BLOCK_D_q, BLOCK_D_q) + expected_z = (B_q * H_q, BLOCK_D_q) + if tuple(init_state_kv.shape) != expected_kv: + raise ValueError( + f"fused_gdn_func: init_state_kv shape {tuple(init_state_kv.shape)} != " f"expected {expected_kv}." + ) + if tuple(init_state_z.shape) != expected_z: + raise ValueError( + f"fused_gdn_func: init_state_z shape {tuple(init_state_z.shape)} != " f"expected {expected_z}." + ) + if init_state_kv.dtype != torch.float32 or init_state_z.dtype != torch.float32: + raise ValueError( + f"fused_gdn_func: init_state_kv/init_state_z must be fp32 " + f"(got {init_state_kv.dtype}, {init_state_z.dtype})." + ) + if not init_state_kv.is_contiguous() or not init_state_z.is_contiguous(): + raise ValueError("fused_gdn_func: init_state_kv / init_state_z must be contiguous.") + + # Stateless path + if init_state_kv is None and init_state_z is None and not save_final_state: + return fused_gdn_func_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, + reverse=reverse, + ) + + # Stateful path: shape-adapt state I/O. + # state_kv: (B*H, BLOCK_D, BLOCK_D) row-major as M[K_feat, V_feat] + # chunkwise stateful: takes user-facing (B, H, D_in, D_out) and transposes + # internally to (B*H, D_out, D_in) for kernel storage. + # state_z: (B*H, BLOCK_D) + # chunkwise stateful: (B, H, D, 1) or (B, H, D) + B, N, _three, H, D = qkv.shape + BLOCK_D = triton.next_power_of_2(D) + + ck_init_kv = None + ck_init_z = None + if init_state_kv is not None: + # (B*H, BLOCK_D, BLOCK_D) → (B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D] + # then transpose so chunkwise's internal `.transpose(-1, -2)` undoes it. + ck_init_kv = init_state_kv.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D].transpose(-1, -2).contiguous() + if init_state_z is not None: + # (B*H, BLOCK_D) → (B, H, BLOCK_D)[:, :, :D] → (B, H, D, 1) + ck_init_z = init_state_z.view(B, H, BLOCK_D)[:, :, :D].unsqueeze(-1).contiguous() + + result = fused_gdn_stateful_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, + reverse=reverse, + init_state_kv=ck_init_kv, + init_state_z=ck_init_z, + return_final_state=save_final_state, + ) + + if not save_final_state: + return result # (num, den) + + num, den, ck_state_kv, ck_state_z = result + # chunkwise returns state_kv as (B, H, D, D), [K_feat, V_feat] (post its + # internal back-transpose). Convert to stateful (B*H, BLOCK_D, BLOCK_D) + # by transposing back to internal storage and padding to BLOCK_D. + out_state_kv = torch.zeros(B * H, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) + out_state_kv[:, :D, :D] = ck_state_kv.transpose(-1, -2).reshape(B * H, D, D) + out_state_z = torch.zeros(B * H, BLOCK_D, device=qkv.device, dtype=torch.float32) + out_state_z[:, :D] = ck_state_z.squeeze(-1).reshape(B * H, D) + return num, den, out_state_kv, out_state_z + + B, N, three, H, D = qkv.shape + assert three == 3 + + BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta, decay = _prepare_launch(D, beta, decay) + + has_init_state = init_state_kv is not None or init_state_z is not None + if reverse and (has_init_state or save_final_state): + raise NotImplementedError( + "fused_gdn_func: state passing (init_state_kv / init_state_z / " + "save_final_state) is only supported for the forward direction " + "(reverse=False). The chunk-causal anti-causal pass resets state " + "per chunk and has no global cross-prefix state to cache." + ) + + if has_init_state: + if init_state_kv is None or init_state_z is None: + raise ValueError( + "fused_gdn_func: init_state_kv and init_state_z must be " + "provided together (got " + f"init_state_kv={'set' if init_state_kv is not None else 'None'}, " + f"init_state_z={'set' if init_state_z is not None else 'None'})." + ) + expected_kv_shape = (B * H, BLOCK_D, BLOCK_D) + expected_z_shape = (B * H, BLOCK_D) + if tuple(init_state_kv.shape) != expected_kv_shape: + raise ValueError( + f"fused_gdn_func: init_state_kv shape {tuple(init_state_kv.shape)} " + f"does not match expected {expected_kv_shape} (BLOCK_D=next_pow2(D)={BLOCK_D})." + ) + if tuple(init_state_z.shape) != expected_z_shape: + raise ValueError( + f"fused_gdn_func: init_state_z shape {tuple(init_state_z.shape)} " + f"does not match expected {expected_z_shape}." + ) + if init_state_kv.dtype != torch.float32 or init_state_z.dtype != torch.float32: + raise ValueError( + "fused_gdn_func: init_state_kv and init_state_z must be fp32 " + f"(got {init_state_kv.dtype}, {init_state_z.dtype})." + ) + if not init_state_kv.is_contiguous() or not init_state_z.is_contiguous(): + raise ValueError("fused_gdn_func: init_state_kv and init_state_z must be contiguous.") + if init_state_kv.device != qkv.device or init_state_z.device != qkv.device: + raise ValueError("fused_gdn_func: init_state_* must live on the same device as qkv.") + load_init = 1 + init_kv_arg = init_state_kv + init_z_arg = init_state_z + else: + load_init = 0 + init_kv_arg = None # placeholder set below + + if save_final_state: + final_state_kv = torch.zeros(B * H, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) + final_state_z = torch.zeros(B * H, BLOCK_D, device=qkv.device, dtype=torch.float32) + save_final = 1 + else: + final_state_kv = None + final_state_z = None + save_final = 0 + + num = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) + den = torch.empty(B, H, N, device=qkv.device, dtype=qkv.dtype) + dummy = torch.empty(1, device=qkv.device, dtype=torch.float32) + + # Resolve pointer args for the unused slots to a shared scratch tensor; + # the kernel compiles the corresponding load/store away when the + # constexpr flag is 0. + init_kv_ptr = init_kv_arg if load_init else dummy + init_z_ptr = init_z_arg if load_init else dummy + final_kv_ptr = final_state_kv if save_final else dummy + final_z_ptr = final_state_z if save_final else dummy + + _fused_gdn_kernel[(B * H,)]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta, + decay, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + num, # `out_ptr` reuses `num` buffer (result immediately overwritten below) + num, + den, + dummy, + dummy, + dummy, + dummy, # saved-state dummies (SAVE_STATE=0) + init_kv_ptr, + init_z_ptr, + final_kv_ptr, + final_z_ptr, + H=H, + F=F, + S=S, + D=D, + K_SCALE=k_scale, + NORM_EPS=1e-5, # unused with USE_PRECOMPUTED_RMS=1 + EPS=eps, + QK_NORM=1, + USE_PRECOMPUTED_RMS=1, + STATE_FP32=1 if state_fp32 else 0, + DOT_PRECISION=dot_prec, + REVERSE=1 if reverse else 0, + SAVE_STATE=0, + LOAD_INIT_STATE=load_init, + SAVE_FINAL_STATE=save_final, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_stages=cfg["num_stages"], + num_warps=nw, + ) + if save_final_state: + return num, den, final_state_kv, final_state_z + return num, den + + +def fused_bigdn_func( + qkv: torch.Tensor, # (B, N, 3, H, D) + q_inv_rms: torch.Tensor, # (B, N) — pre-computed via `_precompute_inv_rms` + k_inv_rms: torch.Tensor, # (B, N) + q_norm_weight: torch.Tensor, # (C,) float32 + k_norm_weight: torch.Tensor, # (C,) + rope_cos: torch.Tensor, # (N, D) + rope_sin: torch.Tensor, # (N, D) + beta: torch.Tensor, # (B, H, F, S) + decay: torch.Tensor, # (B, H, F) + F: int, + S: int, + k_scale: float, + eps: float = 1e-6, + # -- chunk-causal extensions (not in upstream; see adapter notes below) -- + qkv_bwd: torch.Tensor | None = None, + beta_bwd: torch.Tensor | None = None, + decay_bwd: torch.Tensor | None = None, + q_inv_rms_bwd: torch.Tensor | None = None, + k_inv_rms_bwd: torch.Tensor | None = None, +) -> torch.Tensor: + """Full bidirectional fused GDN. + + Returns: out (B, N, H, D) bf16 = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps). + + Chunk-causal extensions (optional): + For chunk-causal GDN we need to zero state at chunk boundaries in the + BACKWARD direction only. Pass separately pre-processed backward tensors + (decay_bwd with zeros at boundary frames, and optionally qkv_bwd / + beta_bwd with K/V or beta zeroed at boundary frames). If any `*_bwd` + argument is None, the forward tensor is reused. + """ + if ( + os.environ.get("FUSED_GDN_FORCE_LEGACY", "0") != "1" + and qkv_bwd is None + and beta_bwd is None + and decay_bwd is None + and q_inv_rms_bwd is None + and k_inv_rms_bwd is None + ): + from .fused_gdn_chunkwise import fused_bigdn_bidi_chunkwise + + 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, + ) + + num_fwd, den_fwd = fused_gdn_func( + 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, + reverse=False, + ) + num_bwd, den_bwd = fused_gdn_func( + qkv if qkv_bwd is None else qkv_bwd, + q_inv_rms if q_inv_rms_bwd is None else q_inv_rms_bwd, + k_inv_rms if k_inv_rms_bwd is None else k_inv_rms_bwd, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta if beta_bwd is None else beta_bwd, + decay if decay_bwd is None else decay_bwd, + F=F, + S=S, + k_scale=k_scale, + eps=eps, + reverse=True, + ) + # num: (B, N, H, D), den: (B, H, N). Fuse then divide. + total_num = num_fwd + num_bwd + total_den = (den_fwd + den_bwd).permute(0, 2, 1).unsqueeze(-1) # (B, N, H, 1) + return total_num / (total_den + eps) + + +# ===================================================================== +# Backward / autograd Functions +# ===================================================================== +# Adds: +# 1. ``_fused_gdn_bwd_kernel`` -- Triton jit kernel that replays the +# forward recurrence in reverse time using per-frame state snapshots +# written by the forward kernel under ``SAVE_STATE=1``. +# 2. ``_run_fwd_save`` -- helper that runs the existing forward +# ``_fused_gdn_kernel`` with ``SAVE_STATE=1``. Adapted to pass our +# extra ``init_state_kv_ptr / init_state_z_ptr / final_state_kv_ptr / +# final_state_z_ptr`` pointers + ``LOAD_INIT_STATE / SAVE_FINAL_STATE`` +# constexpr flags (all unused on the autograd path -> dummy / 0). +# 3. ``FusedGDNFunction`` -- autograd Function for unidirectional GDN +# with ``QK_NORM=1`` (in-kernel per-head RMSNorm). +# 4. ``FusedBiGDNFunction`` -- autograd Function for bidirectional BiGDN. +# Pre-normalizes Q/K in PyTorch with full-channel RMSNorm, runs the +# forward kernel twice with ``QK_NORM=0 + SAVE_STATE=1``, fuses +# ``(num_fwd + num_bwd) / (den_fwd + den_bwd + eps)``. Backward +# computes ``dnum / dden`` from upstream ``dout`` and runs the bwd +# kernel twice with ``BIDI_MODE=1``. +# 5. Python wrappers ``fused_gdn_forward_with_grad`` / +# ``fused_bigdn_forward_with_grad`` -- drop-in autograd-enabled +# replacements for ``fused_gdn_func`` / ``fused_bigdn_func``. +# +# Chunk-causal autograd support: ``FusedBiGDNFunction`` (and the public +# wrapper ``fused_bigdn_forward_with_grad``) accepts optional +# ``beta_bwd`` / ``decay_bwd`` overrides for the reverse-direction +# kernel call -- exactly the same masking convention used by the +# inference path ``fused_bigdn_func``. When provided, the reverse +# direction's forward and backward kernels both run on these masked +# tensors, and the backward returns separate gradient tensors +# (``dbeta_bwd`` / ``ddecay_bwd``) so autograd can route them back +# through any ``clone() + index = 0`` masking the caller applied. + + +@triton.jit +def _fused_gdn_bwd_kernel( + # ---- original inputs ---- + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + beta_ptr, + decay_ptr, + q_norm_w_ptr, + k_norm_w_ptr, + rope_cos_ptr, + rope_sin_ptr, + # ---- saved from forward ---- + saved_state_ptr, # (B*H, F, BLOCK_D, BLOCK_D) -- state_prev snapshots + saved_z_ptr, # (B*H, F, BLOCK_D) + saved_state_curr_ptr, # (B*H, F, BLOCK_D, BLOCK_D) -- state_curr (after update) + saved_z_curr_ptr, # (B*H, F, BLOCK_D) + # ---- upstream gradient / pre-computed dnum ---- + dout_ptr, # GDN mode: (B, N, H, D) upstream grad. BiDI mode: pre-computed dnum + # ---- BiDI mode: external dden ---- + dden_ext_ptr, # BiDI mode: (B, H, N) pre-computed dden. GDN mode: unused + # ---- output gradients ---- + dqkv_ptr, # (B, N, 3, H, D) -- same layout as qkv + dbeta_ptr, # (B, H, F, S) + ddecay_ptr, # (B, H, F) + # ---- dims ---- + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + K_SCALE, + NORM_EPS: tl.constexpr, + EPS: tl.constexpr, + QK_NORM: tl.constexpr, + STATE_FP32: tl.constexpr, + REVERSE_BWD: tl.constexpr, # 0=backward of forward GDN, 1=backward of reversed GDN + BIDI_MODE: tl.constexpr, # 0=GDN (compute dnum/dden), 1=BiGDN (use provided) + DOT_PRECISION: tl.constexpr, # 0=bf16 TC, 1=TF32 TC, 2=IEEE fp32 + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + pid = tl.program_id(0) + pid_b = pid // H + pid_h = pid % H + N: tl.constexpr = F * S + bh = pid_b * H + pid_h + + qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h + dqkv_bh = dqkv_ptr + pid_b * stride_b + pid_h * stride_h + dout_bh = dout_ptr + pid_b * (N * H * D) + pid_h * D + beta_bh = beta_ptr + bh * (F * S) + decay_bh = decay_ptr + bh * F + dbeta_bh = dbeta_ptr + bh * (F * S) + ddecay_bh = ddecay_ptr + bh * F + st_bh = saved_state_ptr + bh * F * BLOCK_D * BLOCK_D + sz_bh = saved_z_ptr + bh * F * BLOCK_D + stc_bh = saved_state_curr_ptr + bh * F * BLOCK_D * BLOCK_D + szc_bh = saved_z_curr_ptr + bh * F * BLOCK_D + if BIDI_MODE: + dden_ext_bh = dden_ext_ptr + bh * N + + 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 + if QK_NORM: + q_nw = tl.load(q_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) + k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) + q_nw_pair = tl.load(q_norm_w_ptr + nw_offset + offs_d_pair, mask=mask_d_pair, 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) + + D_inv = 1.0 / D + k_scale = K_SCALE + + # Dot precision: mirror forward kernel + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + # Gradient matmuls: always use bf16 TC + TF32 input precision (matching PyTorch backward) + grad_dtype = tl.bfloat16 + grad_ip: tl.constexpr = "tf32" + + # ---- Gradient state accumulators (reverse time) ---- + dstate = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + dstate_z = tl.zeros([BLOCK_D], dtype=tl.float32) + + for f_rev in range(F): + # Backward iterates in reverse of forward direction. + if REVERSE_BWD: + f = f_rev # backward of reversed GDN: iterate 0..F-1 + # In fwd_save REVERSE, q_frame=f had kv_frame=f+1 (or skip at f=F-1). + kv_frame_bwd = f + 1 if f < F - 1 else f + skip_bwd = f == F - 1 # f=F-1 was dummy step (f_iter=0 in fwd) + else: + f = F - 1 - f_rev # backward of forward GDN: iterate F-1..0 + kv_frame_bwd = f + skip_bwd = False + q_n_base = f * S + kv_n_base = kv_frame_bwd * S + f_beta = beta_bh + kv_frame_bwd * S + + # ---- Load state_curr for Pass 2 output (both directions use inclusive) ---- + st_f = st_bh + f * BLOCK_D * BLOCK_D + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + mask_dd = mask_d[:, None] & mask_d[None, :] + + stc_f = stc_bh + f * BLOCK_D * BLOCK_D + P_state = tl.load(stc_f + offs_dd, mask=mask_dd, other=0.0) + Pz_state = tl.load(szc_bh + f * BLOCK_D + offs_d, mask=mask_d, other=0.0) + if STATE_FP32 == 0: + P_state = P_state.to(tl.float32) + + # Decay: for REVERSE_BWD, use decay[kv_frame] matching fwd_save. + if REVERSE_BWD and skip_bwd: + g = 1.0 + elif REVERSE_BWD: + g = tl.load(decay_bh + kv_frame_bwd).to(tl.float32) + else: + g = tl.load(decay_bh + f).to(tl.float32) + + # ======================================================== + # Pass 2 backward: Output gradients -> dQ, dstate, dstate_z + # ======================================================== + 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, :] + mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] + n_idx = q_n_base + offs_s # Q data from q_frame + + # Load dout; recompute Q, Q_pair, Q_rot, num, den from saved P_f/Pz_f. + dout_ptrs = dout_bh + n_idx[:, None] * (H * D) + offs_d[None, :] + d_out = tl.load(dout_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + + # Recompute Q, Q_pair, Q_rot (same as forward). + q_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d[None, :] * stride_d + Q_raw = tl.load(q_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + q_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d_pair[None, :] * stride_d + Q_pair_raw = tl.load(q_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) + + if QK_NORM: + q_var = tl.sum(Q_raw * Q_raw, axis=1) * D_inv + q_inv_rms = 1.0 / tl.sqrt(q_var + NORM_EPS) + Q_normed = Q_raw * q_inv_rms[:, None] * q_nw[None, :] + Q_pair_normed = Q_pair_raw * q_inv_rms[:, None] * q_nw_pair[None, :] + else: + Q_normed = Q_raw + Q_pair_normed = Q_pair_raw + Q = tl.where(Q_normed > 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 + + # Compute dnum and dden. + if BIDI_MODE: + # BiGDN: dnum and dden pre-computed externally from total num/den. + dnum = d_out # dout_ptr already contains pre-computed dnum + dden = tl.load(dden_ext_bh + n_idx, mask=mask_s, other=0.0).to(tl.float32) + else: + # GDN: recompute num/den using direction-appropriate state. + num_tile = tl.dot( + Q_rot.to(dot_dtype), P_state.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip + ) + den_tile = tl.sum(Q * Pz_state[None, :], axis=1) + inv_den = 1.0 / (den_tile + EPS) + dnum = d_out * inv_den[:, None] + dden = -tl.sum(d_out * num_tile, axis=1) * inv_den * inv_den + + # dstate += Q_rot^T @ dnum (state contribution from num = Q_rot @ P_state). + dstate = dstate + tl.dot( + tl.trans(Q_rot.to(grad_dtype)), + dnum.to(grad_dtype), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + # dstate_z += sum(dden * Q, axis=0) (Pz contribution from den = Q . Pz). + dstate_z += tl.sum(dden[:, None] * Q, axis=0) + + # dQ_rot = dnum @ P_state^T (uses state that forward's output read). + dQ_rot = tl.dot( + dnum.to(grad_dtype), + tl.trans(P_state.to(grad_dtype)), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + # dQ_from_den = dden * Pz_state. + dQ_from_den = dden[:, None] * Pz_state[None, :] + + # RoPE inverse for Q: store dQ_rot, reload at paired indices. + # Store dQ_rot temporarily to dqkv[Q] at normal d positions. + dq_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d[None, :] * stride_d + tl.store(dq_ptrs, dQ_rot.to(tl.bfloat16), mask=mask_sd) + # The XOR-paired channel can be owned by another warp. Synchronize + # both sides of the scratch roundtrip before dqkv is overwritten. + tl.debug_barrier() + + # Load dQ_rot at paired positions. + dq_pair_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d_pair[None, :] * stride_d + dQ_rot_pair = tl.load(dq_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) + tl.debug_barrier() + + # RoPE inverse: dQ = dQ_rot * Cos - dQ_rot_pair * Sin. + dQ = dQ_rot * Cos - dQ_rot_pair * Sin + dQ_from_den + + # ReLU backward. + relu_mask_q = (Q_normed > 0).to(tl.float32) + dQ_normed = dQ * relu_mask_q + + # Norm backward (QK_NORM) or direct (no norm). + if QK_NORM: + gw = dQ_normed * q_nw[None, :] + corr = tl.sum(gw * Q_raw, axis=1) * D_inv * q_inv_rms * q_inv_rms + dQ_raw = q_inv_rms[:, None] * (gw - Q_raw * corr[:, None]) + else: + dQ_raw = dQ_normed + + # Store final dQ_raw to dqkv[Q]. + tl.store(dq_ptrs, dQ_raw.to(tl.bfloat16), mask=mask_sd) + + # Both directions use inclusive output (state_curr), so capture dDelta AFTER Pass 2. + dDelta = dstate + dDelta_z = dstate_z + + # ======================================================== + # Reload state_prev for Pass 1 backward (reuse P_state variable) + # ======================================================== + P_state = tl.load(st_f + offs_dd, mask=mask_dd, other=0.0) + if STATE_FP32 == 0: + P_state = P_state.to(tl.float32) + Pz_state = tl.load(sz_bh + f * BLOCK_D + offs_d, mask=mask_d, other=0.0) + + # ======================================================== + # Pass 1 backward: State update gradients -> dK, dV, dbeta, dstate + # Skip for REVERSE_BWD dummy frame (skip_bwd=True) to avoid clobbering. + # ======================================================== + if skip_bwd == False: + 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, :] + mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] + n_idx = kv_n_base + offs_s # K/V from kv_frame + + # Recompute K, K_pair, K_rot, V. + 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) + + k_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d_pair[None, :] * stride_d + K_pair_raw = tl.load(k_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) + + if QK_NORM: + k_var = tl.sum(K_raw * K_raw, axis=1) * D_inv + k_inv_rms = 1.0 / tl.sqrt(k_var + NORM_EPS) + K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] + K_pair_normed = K_pair_raw * k_inv_rms[:, None] * k_nw_pair[None, :] + else: + K_normed = K_raw + K_pair_normed = K_pair_raw + K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale + 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 + + bt = tl.load(f_beta + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + # Recompute V_pred and delta_v. + K_rot_dc = K_rot.to(dot_dtype) + V_pred = tl.dot(K_rot_dc, P_state.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) + delta_v = (V_raw - V_pred) * bt[:, None] + + # ---- KV stream backward ---- + ddelta_v = tl.dot( + K_rot.to(grad_dtype), + dDelta.to(grad_dtype), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + dK_rot_from_delta = tl.dot( + delta_v.to(grad_dtype), + tl.trans(dDelta.to(grad_dtype)), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + dV = ddelta_v * bt[:, None] + dbeta_kv = tl.sum(ddelta_v * (V_raw - V_pred), axis=1) + + dV_pred = -ddelta_v * bt[:, None] + dK_rot_from_vpred = tl.dot( + dV_pred.to(grad_dtype), + tl.trans(P_state.to(grad_dtype)), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + dstate = dstate + tl.dot( + tl.trans(K_rot.to(grad_dtype)), + dV_pred.to(grad_dtype), + out_dtype=tl.float32, + input_precision=grad_ip, + ) + + dK_rot = dK_rot_from_delta + dK_rot_from_vpred + + # ---- Z stream backward ---- + z_hat = tl.sum(K * Pz_state[None, :], axis=1) + dz = (1.0 - z_hat) * bt + + ddz = tl.sum(K * dDelta_z[None, :], axis=1) + dz_hat = -ddz * bt + dK_z = dDelta_z[None, :] * dz[:, None] + dz_hat[:, None] * Pz_state[None, :] + dstate_z = dstate_z + tl.sum(dz_hat[:, None] * K, axis=0) + + dbeta_z = ddz * (1.0 - z_hat) + dbeta_total = dbeta_kv + dbeta_z + tl.store(dbeta_bh + kv_frame_bwd * S + offs_s, dbeta_total.to(tl.bfloat16), mask=mask_s) + + # ---- RoPE inverse for K ---- + dk_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d[None, :] * stride_d + tl.store(dk_ptrs, dK_rot.to(tl.bfloat16), mask=mask_sd) + # Match the dQ synchronization: all temporary values must be + # visible before paired loads and consumed before overwrites. + tl.debug_barrier() + dk_pair_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d_pair[None, :] * stride_d + dK_rot_pair = tl.load(dk_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) + tl.debug_barrier() + + dK_from_kv = dK_rot * Cos - dK_rot_pair * Sin + dK_total = dK_from_kv + dK_z + + relu_mask_k = (K_normed > 0).to(tl.float32) + dK_normed = dK_total * k_scale * relu_mask_k + + if QK_NORM: + gw_k = dK_normed * k_nw[None, :] + corr_k = tl.sum(gw_k * K_raw, axis=1) * D_inv * k_inv_rms * k_inv_rms + dK_raw = k_inv_rms[:, None] * (gw_k - K_raw * corr_k[:, None]) + else: + dK_raw = dK_normed + + tl.store(dk_ptrs, dK_raw.to(tl.bfloat16), mask=mask_sd) + dv_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 2 * stride_3 + offs_d[None, :] * stride_d + tl.store(dv_ptrs, dV.to(tl.bfloat16), mask=mask_sd) + + # ======================================================== + # Decay backward (inside skip_bwd guard) + # ======================================================== + is_first_frame = f_rev == F - 1 + if is_first_frame: + ddecay_f = 0.0 + else: + inv_g = 1.0 / (g + 1e-12) + ddecay_kv = tl.sum(dstate * P_state) * inv_g + ddecay_z_val = tl.sum(dstate_z * Pz_state) * inv_g + ddecay_f = ddecay_kv + ddecay_z_val + tl.store(ddecay_bh + kv_frame_bwd, ddecay_f) + + # Propagate gradient through decay: dS_{f-1} = g[f] * dP_f. + dstate = dstate * g + dstate_z = dstate_z * g + + +# ===================================================================== +# Forward-with-state-save helper (for autograd Functions) +# ===================================================================== + + +def _run_fwd_save( + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F: int, + S: int, + k_scale: float, + norm_eps: float, + eps: float, + qk_norm: bool, + reverse: bool, + cfg, +): + """Run forward kernel for one direction with ``SAVE_STATE=1``. + + Returns ``(num, den, saved_state, saved_z, saved_state_curr, saved_z_curr)``. + The forward kernel writes ``out = num/(den+eps)`` first and then overwrites + the same buffer with raw ``num``, so the returned ``num`` tensor holds raw + numerator values (matching the BiGDN combine-then-divide convention). + """ + B, N, three, H, D = qkv.shape + BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, _, beta, decay = _prepare_launch(D, beta, decay) + + num_out = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) + den_out = torch.empty(B, H, N, device=qkv.device, dtype=qkv.dtype) + state_dtype = torch.float32 if state_fp32 else torch.bfloat16 + saved_state = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=state_dtype) + saved_z = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) + saved_state_curr = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) + saved_z_curr = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) + # The kernel writes ``out = num/(den+eps)`` first then overwrites with raw num + # in the same buffer. Reuse num_out as the (discarded) ``out`` slot so the + # final contents end up being raw num. + out_discard = num_out + dummy_inv = torch.empty(1, device=qkv.device, dtype=torch.float32) + + _fused_gdn_kernel[(B * H,)]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta, + decay, + dummy_inv, + dummy_inv, # unused inv_rms ptrs (USE_PRECOMPUTED_RMS=0) + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + out_discard, + num_out, + den_out, + saved_state, + saved_z, + saved_state_curr, + saved_z_curr, + dummy_inv, + dummy_inv, + dummy_inv, + dummy_inv, # init/final-state dummies + H=H, + F=F, + S=S, + D=D, + K_SCALE=k_scale, + NORM_EPS=norm_eps, + EPS=eps, + QK_NORM=1 if qk_norm else 0, + USE_PRECOMPUTED_RMS=0, + STATE_FP32=1 if state_fp32 else 0, + DOT_PRECISION=dot_prec, + REVERSE=1 if reverse else 0, + SAVE_STATE=1, + LOAD_INIT_STATE=0, + SAVE_FINAL_STATE=0, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_stages=cfg["num_stages"], + num_warps=nw, + ) + return num_out, den_out, saved_state, saved_z, saved_state_curr, saved_z_curr + + +# ===================================================================== +# Unidirectional GDN autograd Function +# ===================================================================== + + +class FusedGDNFunction(torch.autograd.Function): + """Autograd Function for unidirectional fused GDN with in-kernel RMSNorm. + + Forward runs ``_fused_gdn_kernel`` with ``QK_NORM=1`` and ``SAVE_STATE=1``, + saving per-frame state snapshots for backward. Backward runs + ``_fused_gdn_bwd_kernel`` with ``BIDI_MODE=0`` (kernel computes + ``dnum``/``dden`` from upstream ``dout``). + """ + + @staticmethod + def forward( + ctx, + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F: int, + S: int, + k_scale: float = 1.0, + norm_eps: float = 1e-6, + eps: float = 1e-6, + qk_norm: bool = True, + ): + B, N, three, H, D = qkv.shape + assert three == 3 and N == F * S + + BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta, decay = _prepare_launch(D, beta, decay) + + if q_norm_weight is None: + q_norm_weight = torch.ones(D, device=qkv.device, dtype=torch.float32) + if k_norm_weight is None: + k_norm_weight = torch.ones(D, device=qkv.device, dtype=torch.float32) + + out = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) + + # Saved states for backward. + state_dtype = torch.float32 if state_fp32 else torch.bfloat16 + saved_state = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=state_dtype) + saved_z = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) + saved_state_curr = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) + saved_z_curr = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) + + # Dummy num/den for forward kernel (still writes them but we discard). + num_out = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) + den_out = torch.empty(B, H, N, device=qkv.device, dtype=qkv.dtype) + dummy_inv = torch.empty(1, device=qkv.device, dtype=torch.float32) + + _fused_gdn_kernel[(B * H,)]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta, + decay, + dummy_inv, + dummy_inv, # unused inv_rms ptrs (USE_PRECOMPUTED_RMS=0) + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + out, + num_out, + den_out, + saved_state, + saved_z, + saved_state_curr, + saved_z_curr, + dummy_inv, + dummy_inv, + dummy_inv, + dummy_inv, # init/final-state dummies + H=H, + F=F, + S=S, + D=D, + K_SCALE=k_scale, + NORM_EPS=norm_eps, + EPS=eps, + QK_NORM=1 if qk_norm else 0, + USE_PRECOMPUTED_RMS=0, + STATE_FP32=1 if state_fp32 else 0, + DOT_PRECISION=dot_prec, + REVERSE=0, + SAVE_STATE=1, + LOAD_INIT_STATE=0, + SAVE_FINAL_STATE=0, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_stages=cfg["num_stages"], + num_warps=nw, + ) + del num_out, den_out + + ctx.save_for_backward( + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + saved_state, + saved_z, + saved_state_curr, + saved_z_curr, + ) + ctx.F = F + ctx.S = S + ctx.k_scale = k_scale + ctx.norm_eps = norm_eps + ctx.eps = eps + ctx.qk_norm = qk_norm + ctx.dot_prec = dot_prec + return out + + @staticmethod + def backward(ctx, dout): + ( + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + saved_state, + saved_z, + saved_state_curr, + saved_z_curr, + ) = ctx.saved_tensors + + B, N, three, H, D = qkv.shape + F_val = ctx.F + S = ctx.S + + BLOCK_D, BLOCK_S_BWD, _, _, _, cfg, beta, decay = _prepare_launch(D, beta, decay) + dqkv = torch.zeros_like(qkv) + dbeta = torch.zeros_like(beta) + ddecay = torch.zeros_like(decay) + + # Dummy dden_ext (unused in GDN mode). + dden_ext = torch.empty(1, device=qkv.device, dtype=torch.float32) + + # Progressive num_warps reduction on tmem overflow. + nw = cfg["num_warps"] + if ctx.dot_prec >= 1: + nw = min(nw, 4) + while nw >= 1: + try: + _fused_gdn_bwd_kernel[(B * H,)]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + saved_state, + saved_z, + saved_state_curr, + saved_z_curr, + dout.contiguous(), + dden_ext, + dqkv, + dbeta, + ddecay, + H=H, + F=F_val, + S=S, + D=D, + K_SCALE=ctx.k_scale, + NORM_EPS=ctx.norm_eps, + EPS=ctx.eps, + QK_NORM=1 if ctx.qk_norm else 0, + STATE_FP32=1 if cfg["STATE_FP32"] else 0, + REVERSE_BWD=0, + BIDI_MODE=0, + DOT_PRECISION=ctx.dot_prec, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S_BWD, + num_stages=cfg["num_stages"], + num_warps=nw, + ) + break + except Exception as e: + if "OutOfResources" in str(type(e).__name__) or "out of resource" in str(e).lower(): + nw = nw // 2 + if nw < 1: + raise RuntimeError( + "FusedGDN backward: Triton kernel OutOfResources at all warp " + f"counts (8, 4, 2, 1). Most recent error: {e}" + ) from e + else: + raise + + return dqkv, dbeta, ddecay, None, None, None, None, None, None, None, None, None, None + + +def fused_gdn_forward_with_grad( + qkv: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + q_norm_weight: torch.Tensor | None, + k_norm_weight: torch.Tensor | None, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + F: int, + S: int, + k_scale: float = 1.0, + norm_eps: float = 1e-6, + eps: float = 1e-6, + qk_norm: bool = True, +) -> torch.Tensor: + """Drop-in autograd-enabled replacement for the unidirectional GDN path. + + Unlike ``fused_gdn_func`` (which expects pre-computed ``q_inv_rms``/ + ``k_inv_rms``), this wrapper computes per-head RMSNorm inside the + Triton kernel (``QK_NORM=1`` / ``USE_PRECOMPUTED_RMS=0``) so the + backward kernel can reproduce the exact same normed Q/K when + replaying the recurrence. + """ + return FusedGDNFunction.apply( + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F, + S, + k_scale, + norm_eps, + eps, + qk_norm, + ) + + +# ===================================================================== +# Bidirectional BiGDN autograd Function (full-channel RMSNorm in Python) +# ===================================================================== + + +class FusedBiGDNFunction(torch.autograd.Function): + """Autograd Function for bidirectional fused BiGDN. + + Full-channel RMSNorm is applied in Python (so the norm backward can + couple all heads correctly), then the kernel runs with ``QK_NORM=0`` + on the pre-normed QKV. Forward and reverse directions are run + separately with ``SAVE_STATE=1``, and combined as + ``out = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps)``. + + Backward computes ``dnum`` and ``dden`` from upstream ``dout`` and + runs the bwd kernel twice (forward + reverse) with ``BIDI_MODE=1``. + Norm backward is computed in Python (full-channel RMSNorm couples + all heads). + + Chunk-causal masking (optional): + Pass ``beta_bwd`` and/or ``decay_bwd`` to override the beta/decay + tensors used by the **reverse-direction** kernel calls (forward + save + backward). The forward direction always uses the + unmasked ``beta`` / ``decay``. This mirrors the inference path + in :func:`fused_bigdn_func` and unlocks chunk-causal autograd + training: callers typically build ``beta_bwd`` / ``decay_bwd`` + as ``beta.clone()`` / ``decay.clone()`` with interior chunk + boundaries zeroed, so the anti-causal scan resets state at + every chunk boundary. + + When ``beta_bwd`` is ``None``, the kernel-emitted reverse- + direction beta gradient is summed into the forward-direction + gradient (returned via the ``beta`` slot) and the ``beta_bwd`` + gradient slot returns ``None``. When ``beta_bwd`` is provided, + the two gradient streams are kept separate: the forward- + direction gradient flows through the ``beta`` slot and the + reverse-direction gradient flows through the ``beta_bwd`` slot + so autograd can route them through any ``clone() + index = 0`` + masking applied by the caller. ``decay_bwd`` is handled + identically. + """ + + @staticmethod + def forward( + ctx, + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F: int, + S: int, + k_scale: float = 1.0, + norm_eps: float = 1e-5, + eps: float = 1e-6, + beta_bwd: torch.Tensor | None = None, + decay_bwd: torch.Tensor | None = None, + ): + B, N, three, H, D = qkv.shape + C = H * D + assert three == 3 and N == F * S + cfg = _kcfg() + + if q_norm_weight is None: + q_norm_weight = torch.ones(C, device=qkv.device, dtype=torch.float32) + if k_norm_weight is None: + k_norm_weight = torch.ones(C, device=qkv.device, dtype=torch.float32) + + # Full-channel RMSNorm: inv_rms over all H*D dims. + q_raw = qkv[:, :, 0].float() # (B, N, H, D) + k_raw = qkv[:, :, 1].float() + q_inv_rms = torch.rsqrt((q_raw * q_raw).sum(dim=(-2, -1)) / C + norm_eps) # (B, N) + k_inv_rms = torch.rsqrt((k_raw * k_raw).sum(dim=(-2, -1)) / C + norm_eps) + + # Apply norm to Q and K: Q_normed = Q_raw * inv_rms * weight. + q_nw_hd = q_norm_weight.reshape(H, D) + k_nw_hd = k_norm_weight.reshape(H, D) + qkv_normed = qkv.clone() + qkv_normed[:, :, 0] = (q_raw * q_inv_rms[:, :, None, None] * q_nw_hd[None, None]).to(qkv.dtype) + qkv_normed[:, :, 1] = (k_raw * k_inv_rms[:, :, None, None] * k_nw_hd[None, None]).to(qkv.dtype) + + # Reverse-direction beta/decay overrides for chunk-causal masking. + # When the caller supplies ``beta_bwd`` / ``decay_bwd`` (typically + # ``beta.clone()`` / ``decay.clone()`` with interior chunk-boundary + # frames zeroed), the reverse-direction kernel reads them instead of + # the unmasked tensors so the anti-causal scan resets state at chunk + # boundaries. The forward (causal) direction always uses the + # unmasked ``beta`` / ``decay``. + beta_for_bwd_dir = beta_bwd if beta_bwd is not None else beta + decay_for_bwd_dir = decay_bwd if decay_bwd is not None else decay + + # Run forward-save with QK_NORM=0 on pre-normed data. + dummy_nw = torch.ones(D, device=qkv.device, dtype=torch.float32) + num_fwd, den_fwd, sv_fwd, sz_fwd, svc_fwd, szc_fwd = _run_fwd_save( + qkv_normed, + beta, + decay, + dummy_nw, + dummy_nw, + rope_cos, + rope_sin, + F, + S, + k_scale, + norm_eps, + eps, + False, + False, + cfg, + ) + num_bwd, den_bwd, sv_bwd, sz_bwd, svc_bwd, szc_bwd = _run_fwd_save( + qkv_normed, + beta_for_bwd_dir, + decay_for_bwd_dir, + dummy_nw, + dummy_nw, + rope_cos, + rope_sin, + F, + S, + k_scale, + norm_eps, + eps, + False, + True, + cfg, + ) + + # Combine: out = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps). + total_num = num_fwd.float() + num_bwd.float() + total_den = den_fwd.float() + den_bwd.float() + total_den_exp = total_den.permute(0, 2, 1).unsqueeze(-1) # (B, N, H, 1) + out = (total_num / (total_den_exp + eps)).to(qkv.dtype) + + # Save ``beta_bwd`` / ``decay_bwd`` (possibly ``None``) so the + # backward pass can (a) replay the reverse-direction kernel against + # the same masked inputs, and (b) decide whether to keep the + # reverse-direction beta/decay gradients separate (caller-supplied + # override) or fold them into the forward-direction gradient + # (no override, legacy behaviour). + ctx.save_for_backward( + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + q_inv_rms, + k_inv_rms, + rope_cos, + rope_sin, + sv_fwd, + sz_fwd, + svc_fwd, + szc_fwd, + sv_bwd, + sz_bwd, + svc_bwd, + szc_bwd, + out, + total_den.to(qkv.dtype), + beta_bwd, + decay_bwd, + ) + _, _dot_prec, _, _ = _resolve_launch_config() + ctx.dot_prec = _dot_prec + ctx.F = F + ctx.S = S + ctx.k_scale = k_scale + ctx.norm_eps = norm_eps + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, dout): + ( + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + q_inv_rms, + k_inv_rms, + rope_cos, + rope_sin, + sv_fwd, + sz_fwd, + svc_fwd, + szc_fwd, + sv_bwd, + sz_bwd, + svc_bwd, + szc_bwd, + out, + total_den_saved, + beta_bwd_saved, + decay_bwd_saved, + ) = ctx.saved_tensors + + # Track whether the caller supplied separate ``beta_bwd`` / + # ``decay_bwd`` overrides; this controls whether the reverse- + # direction kernel gradients are summed into the forward-direction + # slot (legacy behaviour) or routed back through dedicated grad + # slots so autograd can flow through the caller's masking ops + # (``clone() + index = 0``). + has_beta_bwd = beta_bwd_saved is not None + has_decay_bwd = decay_bwd_saved is not None + + B, N, three, H, D = qkv.shape + C = H * D + + # Recompute qkv_normed (avoid saving B*N*3*H*D extra tensor). + q_raw = qkv[:, :, 0].float() + k_raw = qkv[:, :, 1].float() + q_nw_hd = q_norm_weight.reshape(H, D) + k_nw_hd = k_norm_weight.reshape(H, D) + qkv_normed = qkv.clone() + qkv_normed[:, :, 0] = (q_raw * q_inv_rms[:, :, None, None] * q_nw_hd[None, None]).to(qkv.dtype) + qkv_normed[:, :, 1] = (k_raw * k_inv_rms[:, :, None, None] * k_nw_hd[None, None]).to(qkv.dtype) + F_val = ctx.F + S = ctx.S + eps = ctx.eps + BLOCK_D, _, _, _, _, cfg, beta, decay = _prepare_launch(D, beta, decay) + + # Reverse-direction beta/decay actually fed to the reverse kernel. + # When the caller supplied an override, we replay against the + # masked tensor; otherwise we reuse the unmasked beta/decay so the + # legacy summing path is bit-identical to the pre-extension + # behaviour. + if has_beta_bwd: + beta_for_bwd_dir = beta_bwd_saved.contiguous() + else: + beta_for_bwd_dir = beta + if has_decay_bwd: + decay_for_bwd_dir = decay_bwd_saved.contiguous() + else: + decay_for_bwd_dir = decay + + # ---- Pre-compute dnum and dden ---- + total_den_exp = total_den_saved.float().permute(0, 2, 1).unsqueeze(-1) + inv_total_den = 1.0 / (total_den_exp + eps) + dnum = (dout.float() * inv_total_den).to(qkv.dtype).contiguous() + dden = ( + (-(dout.float() * out.float()).sum(dim=-1) * inv_total_den.squeeze(-1)) + .permute(0, 2, 1) + .to(qkv.dtype) + .contiguous() + ) + del out, total_den_saved, total_den_exp, inv_total_den + + dummy_nw = torch.ones(D, device=qkv.device, dtype=torch.float32) + + # ---- Backward for forward direction (QK_NORM=0, operates on normed QKV) ---- + dqkv_fwd = torch.zeros_like(qkv) + dbeta_fwd = torch.zeros_like(beta) + ddecay_fwd = torch.zeros_like(decay) + + def _run_triton_bwd(sv, sz, svc, szc, dqkv_out, dbeta_out, ddecay_out, reverse_bwd, beta_kernel, decay_kernel): + """Try backward kernel with progressively fewer warps on tmem overflow. + + ``beta_kernel`` / ``decay_kernel`` are passed as explicit + arguments (instead of closing over the outer-scope ``beta`` / + ``decay``) so the reverse-direction call can replay against the + chunk-causal-masked tensors (``beta_bwd`` / ``decay_bwd``) when + present, while the forward-direction call always uses the + unmasked ``beta`` / ``decay``. + """ + nw = cfg["num_warps"] + if ctx.dot_prec >= 1: + nw = min(nw, 4) + while nw >= 1: + try: + _fused_gdn_bwd_kernel[(B * H,)]( + qkv_normed, + qkv_normed.stride(0), + qkv_normed.stride(1), + qkv_normed.stride(2), + qkv_normed.stride(3), + qkv_normed.stride(4), + beta_kernel, + decay_kernel, + dummy_nw, + dummy_nw, + rope_cos, + rope_sin, + sv, + sz, + svc, + szc, + dnum, + dden, + dqkv_out, + dbeta_out, + ddecay_out, + H=H, + F=F_val, + S=S, + D=D, + K_SCALE=ctx.k_scale, + NORM_EPS=ctx.norm_eps, + EPS=eps, + QK_NORM=0, + STATE_FP32=1 if cfg["STATE_FP32"] else 0, + REVERSE_BWD=reverse_bwd, + BIDI_MODE=1, + DOT_PRECISION=ctx.dot_prec, + BLOCK_D=BLOCK_D, + BLOCK_S=cfg["BLOCK_S"], + num_stages=cfg["num_stages"], + num_warps=nw, + ) + return # success + except Exception as e: + if "OutOfResources" in str(type(e).__name__) or "out of resource" in str(e).lower(): + nw = nw // 2 + if nw >= 1: + continue + raise RuntimeError( + "FusedBiGDN backward: Triton kernel OutOfResources at all warp counts " + f"(8, 4, 2, 1). Most recent error: {e}" + ) from e + else: + raise + + _run_triton_bwd(sv_fwd, sz_fwd, svc_fwd, szc_fwd, dqkv_fwd, dbeta_fwd, ddecay_fwd, 0, beta, decay) + del sv_fwd, sz_fwd, svc_fwd, szc_fwd + + # ---- Backward for reversed direction (replays against masked beta/decay if any) ---- + # Allocate kernel-output gradients with the exact shape the kernel + # writes — these always match the input ``beta_for_bwd_dir`` / + # ``decay_for_bwd_dir`` shapes (override or fall-back). + dqkv_bwd = torch.zeros_like(qkv) + dbeta_bwd_kernel = torch.zeros_like(beta_for_bwd_dir) + ddecay_bwd_kernel = torch.zeros_like(decay_for_bwd_dir) + + _run_triton_bwd( + sv_bwd, + sz_bwd, + svc_bwd, + szc_bwd, + dqkv_bwd, + dbeta_bwd_kernel, + ddecay_bwd_kernel, + 1, + beta_for_bwd_dir, + decay_for_bwd_dir, + ) + del sv_bwd, sz_bwd, svc_bwd, szc_bwd + del qkv_normed, dnum, dden + + # Q/K/V gradient is always summed: qkv is shared by both directions. + dqkv_fwd += dqkv_bwd + del dqkv_bwd + + # Beta gradient: route depends on whether the caller supplied an + # override. With override -> keep separate (so autograd routes the + # reverse-direction grad through the caller's clone+mask op). + # Without override -> sum into the forward-direction grad + # (legacy behaviour, bit-identical to pre-extension code). + if has_beta_bwd: + dbeta = dbeta_fwd + dbeta_bwd_out: torch.Tensor | None = dbeta_bwd_kernel + else: + dbeta_fwd += dbeta_bwd_kernel + dbeta = dbeta_fwd + dbeta_bwd_out = None + del dbeta_bwd_kernel + + # Decay gradient: same routing logic, independent of beta override. + if has_decay_bwd: + ddecay = ddecay_fwd + ddecay_bwd_out: torch.Tensor | None = ddecay_bwd_kernel + else: + ddecay_fwd += ddecay_bwd_kernel + ddecay = ddecay_fwd + ddecay_bwd_out = None + del ddecay_bwd_kernel + + dqkv_normed = dqkv_fwd + + # ---- Full-channel RMSNorm backward for Q and K ---- + # y = x * inv_rms * w -> dL/dx = inv_rms*w*dL/dy - inv_rms^3/C * x * sum(w*dL/dy*x) + # Process Q and K sequentially to reduce peak fp32 memory. + + # Q norm backward. + q_irms = q_inv_rms[:, :, None, None] + dq_normed = dqkv_normed[:, :, 0].float() + gw_q = dq_normed * q_nw_hd[None, None] + dq_nw = (dq_normed * q_raw * q_irms).sum(dim=(0, 1)).reshape(-1) + corr_q = (gw_q * q_raw).sum(dim=(-2, -1), keepdim=True) + dqkv_normed[:, :, 0] = (q_irms * gw_q - (q_irms**3) / C * q_raw * corr_q).to(qkv.dtype) + del dq_normed, gw_q, corr_q, q_raw, q_irms + + # K norm backward. + k_irms = k_inv_rms[:, :, None, None] + dk_normed = dqkv_normed[:, :, 1].float() + gw_k = dk_normed * k_nw_hd[None, None] + dk_nw = (dk_normed * k_raw * k_irms).sum(dim=(0, 1)).reshape(-1) + corr_k = (gw_k * k_raw).sum(dim=(-2, -1), keepdim=True) + dqkv_normed[:, :, 1] = (k_irms * gw_k - (k_irms**3) / C * k_raw * corr_k).to(qkv.dtype) + del dk_normed, gw_k, corr_k, k_raw, k_irms + + return ( + dqkv_normed, + dbeta, + ddecay, + dq_nw.to(q_norm_weight.dtype), + dk_nw.to(k_norm_weight.dtype), + None, # rope_cos + None, # rope_sin + None, # F + None, # S + None, # k_scale + None, # norm_eps + None, # eps + dbeta_bwd_out, # beta_bwd + ddecay_bwd_out, # decay_bwd + ) + + +def fused_bigdn_forward_with_grad( + qkv: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + q_norm_weight: torch.Tensor | None, + k_norm_weight: torch.Tensor | None, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + F: int, + S: int, + k_scale: float = 1.0, + norm_eps: float = 1e-5, + eps: float = 1e-6, + beta_bwd: torch.Tensor | None = None, + decay_bwd: torch.Tensor | None = None, +) -> torch.Tensor: + """Bidirectional fused BiGDN with autograd support (full-channel RMSNorm). + + Unlike ``fused_bigdn_func`` (which expects pre-computed ``q_inv_rms`` / + ``k_inv_rms``), this wrapper computes the full-channel inv-RMS in Python + so the norm backward can flow through the autograd graph naturally. + + Chunk-causal masking (optional): + Pass ``beta_bwd`` and/or ``decay_bwd`` to override the beta/decay + tensors used by the **reverse-direction** kernel only. These are + typically built by the caller as ``beta.clone()`` / ``decay.clone()`` + with interior chunk-boundary frames zeroed, so the anti-causal scan + resets state at chunk boundaries while the causal scan keeps full + context. The reverse-direction beta/decay gradients are routed + back through the ``beta_bwd`` / ``decay_bwd`` slots (instead of + being summed into the forward-direction grad), which lets autograd + flow the reverse-direction gradient through the caller's + ``clone() + index = 0`` masking op. + + When ``beta_bwd`` / ``decay_bwd`` is ``None`` (default), behaviour + is bit-identical to the pre-extension full-sequence-bidirectional + path: the reverse-direction kernel uses the unmasked ``beta`` / + ``decay`` and its kernel-emitted gradient is summed into the + forward-direction gradient before being returned. + """ + return FusedBiGDNFunction.apply( + qkv, + beta, + decay, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + F, + S, + k_scale, + norm_eps, + eps, + beta_bwd, + decay_bwd, + ) diff --git a/integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py b/integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py new file mode 100644 index 000000000..715231860 --- /dev/null +++ b/integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py @@ -0,0 +1,2269 @@ +# Copyright 2024 NVIDIA CORPORATION & AFFILIATES +# +# 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Fused GDN — Chunkwise-parallel forward (v2). + +V2 changes vs v1: + 1. Phase A is split into TWO kernels along the GDN data streams (KV and Z; + these are the two gating sub-paths within the GDN block, not CUDA + streams — both kernels are launched on the same CUDA stream): + _phase_a_kv_kernel: P_kv (with K_rot) + A (with K_rot, V) — uses RoPE + _phase_a_z_kernel : P_z (with K) + B (with K) — no RoPE + Z block is genuinely lighter (no V/Cos/Sin loads, no K_pair flip). + On H100 this enables 2 blocks/SM resident → multi-tenancy / latency hiding. + 2. Phase A stores (I - P_kv) and (I - P_z) instead of P_kv/P_z. Phase B then + uses these directly: M = g · (I-P_kv)·M + A_f. The MMA `(I-P_kv) @ M` folds + the identity-add into the matmul (no separate M-PM elementwise pass). + +BiGDN inference path: QK_NORM=1, USE_PRECOMPUTED_RMS=1, SAVE_STATE=0. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +_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) + full_M = lambda: torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype) + full_z = lambda: 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).""" + from .fused_gdn import _resolve_launch_config + + _, 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. + """ + 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) diff --git a/integrations/sana/sana_wm/refiner.py b/integrations/sana/sana_wm/refiner.py index 0d7c61107..8a1246281 100644 --- a/integrations/sana/sana_wm/refiner.py +++ b/integrations/sana/sana_wm/refiner.py @@ -48,6 +48,70 @@ ) +def _move_tensor_attr(module: nn.Module, name: str, device: torch.device | str) -> None: + tensor = getattr(module, name, None) + if isinstance(tensor, nn.Parameter): + setattr(module, name, nn.Parameter(tensor.to(device), requires_grad=tensor.requires_grad)) + elif isinstance(tensor, Tensor): + setattr(module, name, tensor.to(device)) + + +def _offload_video_unused_audio_modules( + transformer: nn.Module, + device: torch.device | str = "cpu", +) -> None: + for name in ( + "audio_proj_in", + "audio_caption_projection", + "audio_time_embed", + "av_cross_attn_video_scale_shift", + "av_cross_attn_audio_scale_shift", + "av_cross_attn_video_a2v_gate", + "av_cross_attn_audio_v2a_gate", + "audio_rope", + "cross_attn_rope", + "cross_attn_audio_rope", + "audio_norm_out", + "audio_proj_out", + ): + child = getattr(transformer, name, None) + if isinstance(child, nn.Module): + child.to(device) + for block in getattr(transformer, "transformer_blocks", ()): + for name in ( + "audio_norm1", + "audio_attn1", + "audio_norm2", + "audio_attn2", + "audio_to_video_norm", + "audio_to_video_attn", + "video_to_audio_norm", + "video_to_audio_attn", + "audio_norm3", + "audio_ff", + ): + child = getattr(block, name, None) + if isinstance(child, nn.Module): + child.to(device) + + +def _move_ltx2_video_modules_to( + transformer: nn.Module, + device: torch.device | str, +) -> None: + for name in ("proj_in", "caption_projection", "time_embed", "rope", "norm_out", "proj_out"): + child = getattr(transformer, name, None) + if isinstance(child, nn.Module): + child.to(device) + _move_tensor_attr(transformer, "scale_shift_table", device) + for block in getattr(transformer, "transformer_blocks", ()): + _move_tensor_attr(block, "scale_shift_table", device) + for name in ("norm1", "attn1", "norm2", "attn2", "norm3", "ff"): + child = getattr(block, name, None) + if isinstance(child, nn.Module): + child.to(device) + + class SanaWMLTX2Refiner(nn.Module): """Run the SANA-WM LTX-2 latent refiner without importing a Sana checkout.""" @@ -61,6 +125,7 @@ def __init__( precision: Precision = "bf16", quant_backend: QuantBackend = "torch", text_max_sequence_length: int = 1024, + cache_text_encoder: bool = False, ) -> None: super().__init__() self.refiner_root = Path(refiner_root) @@ -70,6 +135,7 @@ def __init__( self.precision = precision self.quant_backend = quant_backend self.text_max_sequence_length = int(text_max_sequence_length) + self.cache_text_encoder = bool(cache_text_encoder) self._quantized = False self._text_encoder_built = False self.transformer, self.connectors = self._load_diffusers_components() @@ -94,7 +160,9 @@ def refine_latents( ) prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt) - self.transformer.to(self.device) + + _move_ltx2_video_modules_to(self.transformer, self.device) + _offload_video_unused_audio_modules(self.transformer, "cpu") self.transformer.eval() self._prepare_quantization() @@ -205,14 +273,11 @@ def _prepare_quantization(self) -> None: def _ensure_text_encoder(self) -> None: """Load the Gemma tokenizer + encoder once and cache them in CPU RAM. - The encoder is ~20 GB; reloading it from disk on every refine call - dominated repeated-use latency, so the built module is kept cached and - a reused refiner pays the load once. It stays on the CPU between calls - and is moved to the GPU only for the encode forward (see - :meth:`_encode_prompt`) so it never inflates peak memory during the - denoise/decode phases. When ``offload_refiner`` is set the whole - refiner (and this encoder) is released between runs by - ``release_runtime``. + The encoder is ~20 GB. By default it is released after the one-shot + prompt encode to match upstream's single-generation path and avoid a + large GPU-to-CPU copy. Set ``cache_text_encoder=True`` for repeated + pipeline reuse, where paying the copy once can beat reloading Gemma on + every call. """ if self._text_encoder_built: return @@ -255,6 +320,7 @@ def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: # Pull the cached encoder onto the GPU only for the forward, then park # it back on the CPU so it is absent during denoise/decode. self.text_encoder.to(self.device) + text_backbone: nn.Module | None = None try: text_backbone = getattr(self.text_encoder, "model", self.text_encoder) outputs = text_backbone( @@ -266,9 +332,17 @@ def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: sequence_lengths = attention_mask.sum(dim=-1) del outputs finally: - self.text_encoder.to("cpu") - if torch.cuda.is_available(): - torch.cuda.empty_cache() + release_text_encoder = not self.cache_text_encoder + if release_text_encoder: + del text_backbone + del self.text_encoder + self._text_encoder_built = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + else: + self.text_encoder.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() prompt_embeds = _pack_text_embeds( hidden_states, diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py index bce53ea96..ec7c73161 100644 --- a/integrations/sana/sana_wm/runner.py +++ b/integrations/sana/sana_wm/runner.py @@ -294,23 +294,24 @@ def run(self) -> None: "sink_size": cfg.sink_size, } ) - decoded = pipeline.generate( - 0, - cache, - input=SanaWMI2VConditioningRequest( - image=image, - prompt=prompt, - poses_c2w=c2w, - intrinsics_vec4=intrinsics_vec4, - num_frames=num_frames, - fps=cfg.fps, - steps=cfg.step, - cfg_scale=cfg.cfg_scale, - flow_shift=cfg.flow_shift, - seed=cfg.seed, - negative_prompt=cfg.negative_prompt, - ), - ) + with torch.inference_mode(): + decoded = pipeline.generate( + 0, + cache, + input=SanaWMI2VConditioningRequest( + image=image, + prompt=prompt, + poses_c2w=c2w, + intrinsics_vec4=intrinsics_vec4, + num_frames=num_frames, + fps=cfg.fps, + steps=cfg.step, + cfg_scale=cfg.cfg_scale, + flow_shift=cfg.flow_shift, + seed=cfg.seed, + negative_prompt=cfg.negative_prompt, + ), + ) pipeline.finalize(0, cache) if not isinstance(decoded, SanaWMDecodedVideo): raise TypeError( diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index 522b2b44c..c376a14c5 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -24,12 +24,33 @@ from collections.abc import Callable from dataclasses import dataclass import math +import os import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor +try: + from .ops.fused_cam_gdn import ( + _prepare_ucpe_rope_tables as _fused_prepare_ucpe_rope_tables, + cam_prep_func as _fused_cam_prep_func, + ) + from .ops.fused_gdn import ( + fused_bigdn_func as _fused_bigdn_func, + fused_qk_inv_rms as _fused_qk_inv_rms, + prepare_rope_tables as _fused_prepare_rope_tables, + ) + from .ops.fused_gdn_chunkwise import ( + cam_scan_bidi_chunkwise as _fused_cam_scan_bidi_chunkwise, + ) +except ImportError: + _fused_bigdn_func = None + _fused_cam_prep_func = None + _fused_cam_scan_bidi_chunkwise = None + _fused_prepare_rope_tables = None + _fused_prepare_ucpe_rope_tables = None + @dataclass(frozen=True) class SanaWMStage1Spec: @@ -50,6 +71,8 @@ class SanaWMStage1Spec: plucker_channels: int = 48 raymap_channels: int = 3 softmax_every_n: int = 4 + chunk_size: int | None = None + chunk_split_strategy: str = "first_chunk_plus_one" @property def mlp_inner_size(self) -> int: @@ -70,6 +93,17 @@ def block_uses_gdn(self, index: int) -> bool: """Architecture spec for the public SANA-WM bidirectional Stage-1 checkpoint.""" +@dataclass(frozen=True) +class _CameraProjectionCache: + raymats: Tensor + proj: Tensor + proj_q: Tensor + proj_kv: Tensor + rope_cam: Tensor | None + rope_cos: Tensor + rope_sin: Tensor + + class RMSNorm(nn.Module): """RMSNorm parameter container matching the SANA checkpoint schema.""" @@ -242,10 +276,11 @@ def forward(self, x: Tensor, *, frames: int, height: int, width: int) -> Tensor: """Run spatial GLU plus temporal aggregation.""" batch, tokens, channels = x.shape x_2d = x.reshape(batch * frames, height, width, channels).permute(0, 3, 1, 2) - x_2d = F.silu(self.inverted_conv(x_2d)) + x_2d = F.silu(self.inverted_conv(x_2d), inplace=True) x_2d = self.depth_conv(x_2d) value, gate = x_2d.chunk(2, dim=1) - x_2d = self.point_conv(value * F.silu(gate)) + gate = F.silu(gate, inplace=not torch.is_grad_enabled()) + x_2d = self.point_conv(value * gate) x_time = x_2d.view(batch, frames, channels, height * width).permute(0, 2, 1, 3) x_time = x_time + self.t_conv(x_time) @@ -360,6 +395,7 @@ def forward( HW: tuple[int, int, int] | None = None, rotary_emb: Tensor | None = None, camera_conditions: Tensor | None = None, + camera_cache: _CameraProjectionCache | None = None, apply_output_gate: bool = True, **kwargs: object, ) -> Tensor: @@ -372,7 +408,9 @@ def forward( f"channels={channels} != heads*dim={self.heads * self.dim}" ) - precomputed_gates = self._compute_frame_gates(x, HW) + precomputed_gates = ( + self._compute_frame_gates(x, HW) if self.use_gdn_convs else None + ) if self.use_gdn_convs: main_raw = self._forward_gdn_main( x, @@ -385,6 +423,9 @@ def forward( x, HW=HW, rotary_emb=rotary_emb, + chunk_size=kwargs.get("chunk_size"), + chunk_split_strategy=str(kwargs.get("chunk_split_strategy", "uniform")), + chunk_index=kwargs.get("chunk_index"), ) cam_contrib: Tensor | int = 0 @@ -395,6 +436,7 @@ def forward( HW=HW, rotary_emb=rotary_emb, camera_conditions=camera_conditions, + camera_cache=camera_cache, precomputed_gates=precomputed_gates, ) else: @@ -403,12 +445,21 @@ def forward( HW=HW, rotary_emb=rotary_emb, camera_conditions=camera_conditions, + camera_cache=camera_cache, + chunk_size=kwargs.get("chunk_size"), + chunk_split_strategy=str(kwargs.get("chunk_split_strategy", "uniform")), + chunk_index=kwargs.get("chunk_index"), ) cam_contrib = self.out_proj_cam(cam_raw) combined = main_raw + cam_contrib if apply_output_gate: - combined = _apply_output_gate(combined, x, self.output_gate) + combined = _apply_output_gate( + combined, + x, + self.output_gate.weight, + self.output_gate.bias, + ) combined = self.proj(combined.to(dtype=self.proj.weight.dtype)) del kwargs return combined @@ -421,6 +472,14 @@ def _forward_gdn_main( rotary_emb: Tensor | None, precomputed_gates: tuple[Tensor, Tensor], ) -> Tensor: + if _use_fused_gdn(x): + return self._forward_fused_gdn_main( + x, + HW=HW, + rotary_emb=rotary_emb, + precomputed_gates=precomputed_gates, + ) + batch, tokens, channels = x.shape qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim) if hasattr(self, "conv_k"): @@ -460,14 +519,67 @@ def _forward_gdn_main( ) return out.reshape(batch, tokens, channels).to(dtype=x.dtype) + def _forward_fused_gdn_main( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + precomputed_gates: tuple[Tensor, Tensor], + ) -> Tensor: + if ( + _fused_bigdn_func is None + or _fused_qk_inv_rms is None + or _fused_prepare_rope_tables is None + ): + raise RuntimeError("fused SANA-WM GDN kernels are not available.") + + batch, tokens, channels = x.shape + frames, height, width = HW + spatial = height * width + if tokens != frames * spatial: + raise ValueError(f"tokens={tokens} != T*H*W={frames * spatial}") + qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim).contiguous() + if hasattr(self, "conv_k"): + k_raw = qkv[:, :, 1].contiguous().reshape(batch, tokens, channels) + k_conv = _apply_bidirectional_temporal_conv(k_raw, self.conv_k, HW) + qkv[:, :, 1].copy_(k_conv.reshape(batch, tokens, self.heads, self.dim)) + + beta, decay = precomputed_gates + q_inv_rms, k_inv_rms = _fused_qk_inv_rms(qkv, eps=self.q_norm.eps) + rope_cos, rope_sin = _fused_prepare_rope_tables( + rotary_emb, + tokens, + self.dim, + x.device, + ) + out = _fused_bigdn_func( + qkv, + q_inv_rms, + k_inv_rms, + self.q_norm.weight.float().contiguous(), + self.k_norm.weight.float().contiguous(), + rope_cos, + rope_sin, + beta.contiguous(), + decay.contiguous(), + F=frames, + S=spatial, + k_scale=_gdn_key_scale(self.dim, HW), + eps=self.eps, + ) + return out.reshape(batch, tokens, channels).to(dtype=x.dtype) + def _forward_softmax_main( self, x: Tensor, *, HW: tuple[int, int, int], rotary_emb: Tensor | None, + chunk_size: int | None, + chunk_split_strategy: str, + chunk_index: list[int] | None, ) -> Tensor: - del HW batch, tokens, channels = x.shape qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim) q, k, v = qkv.unbind(dim=2) @@ -485,12 +597,14 @@ def _forward_softmax_main( ) q = _apply_complex_rope(q, rotary_emb) k = _apply_complex_rope(k, rotary_emb) - out = F.scaled_dot_product_attention( + out = _scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), - dropout_p=0.0, - is_causal=False, + HW=HW, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, ) return out.transpose(1, 2).reshape(batch, tokens, channels).to(dtype=x.dtype) @@ -501,8 +615,19 @@ def _forward_gdn_camera( HW: tuple[int, int, int], rotary_emb: Tensor | None, camera_conditions: Tensor, + camera_cache: _CameraProjectionCache | None, precomputed_gates: tuple[Tensor, Tensor], ) -> Tensor: + if _use_fused_gdn(x): + return self._forward_fused_gdn_camera( + x, + HW=HW, + rotary_emb=rotary_emb, + camera_conditions=camera_conditions, + camera_cache=camera_cache, + precomputed_gates=precomputed_gates, + ) + q_cam, k_cam, v_cam = _camera_qkv(self, x, HW) q_trans, k_trans, v_trans, inflation_sq, output_projector = _prepare_ucpe_qkv( q_cam, @@ -511,6 +636,7 @@ def _forward_gdn_camera( camera_conditions=camera_conditions, HW=HW, rotary_emb=rotary_emb, + camera_cache=camera_cache, q_norm_weight=self.q_norm_cam.weight, k_norm_weight=self.k_norm_cam.weight, norm_eps=self.q_norm_cam.eps, @@ -534,6 +660,76 @@ def _forward_gdn_camera( out = output_projector(out) return out.reshape(x.shape[0], x.shape[1], -1).to(dtype=x.dtype) + def _forward_fused_gdn_camera( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + camera_conditions: Tensor, + camera_cache: _CameraProjectionCache | None, + precomputed_gates: tuple[Tensor, Tensor], + ) -> Tensor: + if ( + _fused_cam_prep_func is None + or _fused_cam_scan_bidi_chunkwise is None + or _fused_prepare_ucpe_rope_tables is None + ): + raise RuntimeError("fused SANA-WM camera GDN kernels are not available.") + + q_raw, k_raw, v_raw = _camera_qkv(self, x, HW) + batch, tokens, heads, dim = q_raw.shape + frames, height, width = HW + spatial = height * width + if camera_cache is None: + camera_cache = _prepare_camera_projection_cache( + camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + head_dim=dim, + ) + if camera_cache.rope_cam is None: + rope_cos = camera_cache.rope_cos + rope_sin = camera_cache.rope_sin + else: + rope_cos = camera_cache.rope_cos + rope_sin = camera_cache.rope_sin + + q_trans, k_trans, v_trans, inflation_sq = _fused_cam_prep_func( + q_raw.contiguous(), + k_raw.contiguous(), + v_raw.contiguous(), + q_norm_weight=self.q_norm_cam.weight.float().contiguous(), + k_norm_weight=self.k_norm_cam.weight.float().contiguous(), + proj_q=camera_cache.proj_q, + proj_kv=camera_cache.proj_kv, + rope_cos=rope_cos, + rope_sin=rope_sin, + k_scale=_gdn_key_scale(dim, HW), + norm_eps=self.q_norm_cam.eps, + ) + + beta, decay = precomputed_gates + frame_inflation = inflation_sq.reshape(batch, heads, frames, spatial).mean( + dim=-1, + ) + beta = beta / frame_inflation.unsqueeze(-1).clamp_min(1.0) + out = _fused_cam_scan_bidi_chunkwise( + q_trans.float().contiguous(), + k_trans.float().contiguous(), + v_trans.float().contiguous(), + beta.float().contiguous(), + decay.float().contiguous(), + ) + out = out.permute(0, 3, 1, 2).contiguous() + out = _ucpe_transform_apply( + out.to(dtype=x.dtype), + camera_cache.proj, + camera_cache.rope_cam, + inverse_rope=True, + ) + return out.reshape(batch, tokens, heads * dim).to(dtype=x.dtype) + def _forward_softmax_camera( self, x: Tensor, @@ -541,6 +737,10 @@ def _forward_softmax_camera( HW: tuple[int, int, int], rotary_emb: Tensor | None, camera_conditions: Tensor, + camera_cache: _CameraProjectionCache | None, + chunk_size: int | None, + chunk_split_strategy: str, + chunk_index: list[int] | None, ) -> Tensor: q_cam, k_cam, v_cam = _camera_qkv(self, x, HW) q_trans, k_trans, v_trans, output_projector = _prepare_ucpe_qkv_softmax( @@ -550,16 +750,19 @@ def _forward_softmax_camera( camera_conditions=camera_conditions, HW=HW, rotary_emb=rotary_emb, + camera_cache=camera_cache, q_norm_weight=self.q_norm_cam.weight, k_norm_weight=self.k_norm_cam.weight, norm_eps=self.q_norm_cam.eps, ) - out = F.scaled_dot_product_attention( + out = _scaled_dot_product_attention( q_trans.transpose(1, 2), k_trans.transpose(1, 2), v_trans.transpose(1, 2), - dropout_p=0.0, - is_causal=False, + HW=HW, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, ) return output_projector(out.transpose(1, 2)).reshape( x.shape[0], @@ -577,13 +780,18 @@ def _compute_frame_gates( spatial = height * width if tokens != frames * spatial: raise ValueError(f"tokens={tokens} != T*H*W={frames * spatial}") - beta = self.beta_proj(x).sigmoid() - beta = beta.reshape(batch, frames, spatial, self.heads).permute(0, 3, 1, 2) - x_frame = x.reshape(batch, frames, spatial, channels).mean(dim=2) - gate = self.gate_proj(x_frame).float() - dt = self.dt_bias.float().view(1, 1, -1) - a_val = self.A_log.float().exp().view(1, 1, -1) - decay = (-a_val * F.softplus(gate + dt)).exp().transpose(1, 2) + beta, decay = _compute_frame_gates( + x, + frames, + spatial, + self.heads, + self.beta_proj.weight, + self.beta_proj.bias, + self.gate_proj.weight, + self.gate_proj.bias, + self.dt_bias, + self.A_log, + ) return beta.float(), decay.float() @@ -627,7 +835,8 @@ def forward( attn_mask = None if mask is not None: attn_mask = (1 - mask.to(dtype=q.dtype)) * -10000.0 - attn_mask = attn_mask[:, None, None, :] + if attn_mask.ndim == 2: + attn_mask = attn_mask[:, None, None].repeat(1, self.num_heads, 1, 1) out = F.scaled_dot_product_attention( q, k, @@ -748,6 +957,21 @@ def forward( x.device, ) camera_conditions = kwargs.get("camera_conditions") + camera_conditions = ( + camera_conditions.to(device=x.device, dtype=x.dtype) + if isinstance(camera_conditions, Tensor) + else None + ) + camera_cache = ( + _prepare_camera_projection_cache( + camera_conditions, + HW=(frames, height, width), + rotary_emb=rotary_emb, + head_dim=self.spec.head_dim, + ) + if camera_conditions is not None + else None + ) block_kwargs = { key: value for key, value in kwargs.items() if key != "camera_conditions" } @@ -759,7 +983,12 @@ def forward( timestep_embed = self.t_embedder(timestep.flatten()) timestep_embed = timestep_embed.unflatten(0, timestep.shape) - block_t = self.t_block(timestep_embed).reshape(batch, 1, frames, 6 * self.spec.hidden_size) + block_t = self.t_block(timestep_embed).reshape( + batch, + 1, + frames, + 6 * self.spec.hidden_size, + ) for block in self.blocks: x = block( @@ -773,11 +1002,10 @@ def forward( plucker_emb=plucker_emb, HW=(frames, height, width), rotary_emb=rotary_emb, - camera_conditions=( - camera_conditions.to(device=x.device, dtype=x.dtype) - if isinstance(camera_conditions, Tensor) - else None - ), + camera_conditions=camera_conditions, + camera_cache=camera_cache, + chunk_size=self.spec.chunk_size, + chunk_split_strategy=self.spec.chunk_split_strategy, **block_kwargs, ) @@ -807,11 +1035,224 @@ def _gdn_key_scale(head_dim: int, HW: tuple[int, int, int]) -> float: return (head_dim**-0.5) * ((HW[1] * HW[2]) ** -0.5) -def _apply_output_gate(out: Tensor, gate_x: Tensor, gate: nn.Module) -> Tensor: - gate_values = F.silu(gate(gate_x).float()) +def _sdpa_needs_head_pad(head_dim: int) -> bool: + return head_dim not in (32, 64, 128, 256) and head_dim < 256 + + +def _scaled_dot_product_attention( + xq: Tensor, + xk: Tensor, + xv: Tensor, + *, + HW: tuple[int, int, int] | None = None, + chunk_size: int | None = None, + chunk_split_strategy: str = "uniform", + chunk_index: list[int] | None = None, +) -> Tensor: + if HW is not None and chunk_size is not None and chunk_size < HW[0]: + frames, height, width = HW + spatial = height * width + boundaries = _normalize_chunk_index( + chunk_index, + frames, + chunk_size, + chunk_split_strategy, + ) + out_chunks = [] + for start, end in zip(boundaries[:-1], boundaries[1:]): + q_chunk = xq[:, :, start * spatial : end * spatial] + out_chunks.append( + _scaled_dot_product_attention_full( + q_chunk, + xk[:, :, : end * spatial], + xv[:, :, : end * spatial], + ), + ) + return torch.cat(out_chunks, dim=2) + return _scaled_dot_product_attention_full(xq, xk, xv) + + +def _scaled_dot_product_attention_full(xq: Tensor, xk: Tensor, xv: Tensor) -> Tensor: + head_dim = xq.shape[-1] + if not _sdpa_needs_head_pad(head_dim): + return F.scaled_dot_product_attention( + xq, + xk, + xv, + dropout_p=0.0, + is_causal=False, + ) + pad_to = 128 if head_dim <= 128 else 256 + pad_size = pad_to - head_dim + xq_padded = F.pad(xq, (0, pad_size)) + xk_padded = F.pad(xk, (0, pad_size)) + xv_padded = F.pad(xv, (0, pad_size)) + out = F.scaled_dot_product_attention( + xq_padded, + xk_padded, + xv_padded, + dropout_p=0.0, + is_causal=False, + scale=head_dim**-0.5, + ) + return out[..., :head_dim] + + +def _normalize_chunk_index( + chunk_index: list[int] | None, + frames: int, + chunk_size: int, + chunk_split_strategy: str, +) -> list[int]: + if chunk_index is None: + chunk_index = _chunk_index_from_chunk_size( + frames, + chunk_size, + chunk_split_strategy, + ) + else: + chunk_index = list(chunk_index) + if not chunk_index or chunk_index[0] != 0: + chunk_index = [0] + [idx for idx in chunk_index if idx > 0] + chunk_index = [idx for idx in chunk_index if idx < frames] + if not chunk_index: + chunk_index = [0] + if chunk_index[-1] != frames: + chunk_index.append(frames) + return chunk_index + + +def _chunk_index_from_chunk_size( + frames: int, + chunk_size: int, + chunk_split_strategy: str, +) -> list[int]: + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}.") + if frames <= 0: + raise ValueError(f"frames must be > 0, got {frames}.") + strategy = "uniform" if chunk_split_strategy is None else str(chunk_split_strategy) + strategy = strategy.lower() + if strategy in ("uniform", "default"): + indices = list(range(0, frames, chunk_size)) + if len(indices) > 1 and (frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if strategy in ("first_frame", "first_frame_alone", "first_frame_only"): + if frames <= 1: + return [0] + indices = [0] + list(range(1, frames, chunk_size)) + if len(indices) > 2 and (frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if strategy in ("first_plus_one", "first_chunk_plus_one"): + if frames <= chunk_size + 1: + return [0] + indices = [0] + list(range(chunk_size + 1, frames, chunk_size)) + if len(indices) > 1 and (frames - indices[-1]) < chunk_size: + indices.pop() + return indices + raise ValueError( + "Unknown chunk_split_strategy " + f"{chunk_split_strategy!r}; expected uniform, first_frame, or first_plus_one." + ) + + +def _use_fused_gdn(x: Tensor) -> bool: + return ( + x.is_cuda + and _fused_bigdn_func is not None + and os.environ.get("SANA_WM_DISABLE_FUSED_GDN", "0") != "1" + ) + + +def _prepare_camera_projection_cache( + camera_conditions: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + head_dim: int, +) -> _CameraProjectionCache: + batch, frames = camera_conditions.shape[:2] + latent_frames, height, width = HW + if frames != latent_frames: + raise ValueError( + f"camera_conditions has {frames} frames but latent grid has {latent_frames}." + ) + tokens = frames * height * width + raymats = _camera_ray_mats(camera_conditions.float(), HW).to( + dtype=camera_conditions.dtype, + ) + proj = raymats.reshape(batch, tokens, 4, 4) + proj_q = proj.transpose(-1, -2).contiguous() + proj_kv = _invert_se3(proj).contiguous() + rope_cam = _slice_rope_for_camera(rotary_emb, head_dim) + if rope_cam is not None and _fused_prepare_ucpe_rope_tables is not None: + rope_cos, rope_sin = _fused_prepare_ucpe_rope_tables( + rope_cam, + tokens, + head_dim // 2, + camera_conditions.device, + ) + else: + rope_cos = torch.ones( + tokens, + head_dim // 2, + device=camera_conditions.device, + dtype=torch.float32, + ) + rope_sin = torch.zeros( + tokens, + head_dim // 2, + device=camera_conditions.device, + dtype=torch.float32, + ) + return _CameraProjectionCache( + raymats=raymats, + proj=proj, + proj_q=proj_q, + proj_kv=proj_kv, + rope_cam=rope_cam, + rope_cos=rope_cos, + rope_sin=rope_sin, + ) + + +@torch.compile +def _apply_output_gate( + out: Tensor, + gate_x: Tensor, + gate_weight: Tensor, + gate_bias: Tensor, +) -> Tensor: + gate_values = F.silu(F.linear(gate_x, gate_weight, gate_bias).float()) return out * gate_values.to(dtype=out.dtype) +@torch.compile +def _compute_frame_gates( + x: Tensor, + frames: int, + spatial: int, + heads: int, + beta_weight: Tensor, + beta_bias: Tensor, + gate_weight: Tensor, + gate_bias: Tensor, + dt_bias: Tensor, + A_log: Tensor, +) -> tuple[Tensor, Tensor]: + batch, tokens, channels = x.shape + beta = F.linear(x, beta_weight, beta_bias).sigmoid() + beta = beta.reshape(batch, frames, spatial, heads).permute(0, 3, 1, 2) + x_frame = x.reshape(batch, frames, spatial, channels).mean(dim=2) + gate = 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(gate + dt)).exp().transpose(1, 2) + return beta, decay + + def _apply_bidirectional_temporal_conv( x: Tensor, conv: nn.Conv1d, @@ -1088,12 +1529,29 @@ def _camera_qkv( HW: tuple[int, int, int], ) -> tuple[Tensor, Tensor, Tensor]: batch, tokens, _channels = x.shape - q_raw = module.q_proj_cam(x).reshape(batch, tokens, module.heads, module.dim) - k_raw_flat = module.k_proj_cam(x) + qkv_weight = torch.cat( + [ + module.q_proj_cam.weight, + module.k_proj_cam.weight, + module.v_proj_cam.weight, + ], + ) + qkv_bias = torch.cat( + [ + module.q_proj_cam.bias, + module.k_proj_cam.bias, + module.v_proj_cam.bias, + ], + ) + q_raw_flat, k_raw_flat, v_raw_flat = F.linear(x, qkv_weight, qkv_bias).chunk( + 3, + dim=-1, + ) if hasattr(module, "conv_k_cam"): k_raw_flat = _apply_bidirectional_temporal_conv(k_raw_flat, module.conv_k_cam, HW) + q_raw = q_raw_flat.reshape(batch, tokens, module.heads, module.dim) k_raw = k_raw_flat.reshape(batch, tokens, module.heads, module.dim) - v_raw = module.v_proj_cam(x).reshape(batch, tokens, module.heads, module.dim) + v_raw = v_raw_flat.reshape(batch, tokens, module.heads, module.dim) return q_raw, k_raw, v_raw @@ -1105,6 +1563,7 @@ def _prepare_ucpe_qkv( camera_conditions: Tensor, HW: tuple[int, int, int], rotary_emb: Tensor | None, + camera_cache: _CameraProjectionCache | None, q_norm_weight: Tensor, k_norm_weight: Tensor, norm_eps: float, @@ -1118,36 +1577,39 @@ def _prepare_ucpe_qkv( k_norm = F.relu(k_raw.float() * k_inv[:, :, None, None] * k_weight[None, None]) k_norm = k_norm * _gdn_key_scale(dim, HW) v_float = v_raw.float() - raymats = _camera_ray_mats(camera_conditions.float(), HW) - matrix_q = raymats.reshape(batch, tokens, 4, 4).transpose(-1, -2).contiguous() - matrix_kv = _invert_se3(raymats.reshape(batch, tokens, 4, 4)).contiguous() - rope_cam = _slice_rope_for_camera(rotary_emb, dim) + if camera_cache is None: + camera_cache = _prepare_camera_projection_cache( + camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + head_dim=dim, + ) - q_trans, _q_pre, _q_post = _ucpe_transform( + q_trans = _ucpe_transform_apply( q_norm, - matrix_q, - rope_cam, + camera_cache.proj_q, + camera_cache.rope_cam, inverse_rope=False, ) k_trans, k_pre_sq, k_post_sq = _ucpe_transform( k_norm, - matrix_kv, - rope_cam, + camera_cache.proj_kv, + camera_cache.rope_cam, inverse_rope=False, ) - v_trans, _v_pre, _v_post = _ucpe_transform( + v_trans = _ucpe_transform_apply( v_float, - matrix_kv, - rope_cam, + camera_cache.proj_kv, + camera_cache.rope_cam, inverse_rope=False, ) inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12) def output_projector(out: Tensor) -> Tensor: - projected, _pre, _post = _ucpe_transform( + projected = _ucpe_transform_apply( out.float(), - raymats.reshape(batch, tokens, 4, 4), - rope_cam, + camera_cache.proj, + camera_cache.rope_cam, inverse_rope=True, ) return projected.to(dtype=out.dtype) @@ -1163,6 +1625,7 @@ def _prepare_ucpe_qkv_softmax( camera_conditions: Tensor, HW: tuple[int, int, int], rotary_emb: Tensor | None, + camera_cache: _CameraProjectionCache | None, q_norm_weight: Tensor, k_norm_weight: Tensor, norm_eps: float, @@ -1172,37 +1635,38 @@ def _prepare_ucpe_qkv_softmax( k_inv = _inv_rms(k_raw, norm_eps) q_weight = q_norm_weight.float().view(heads, dim) k_weight = k_norm_weight.float().view(heads, dim) - q_norm = q_raw.float() * q_inv[:, :, None, None] * q_weight[None, None] - k_norm = k_raw.float() * k_inv[:, :, None, None] * k_weight[None, None] - v_float = v_raw.float() - raymats = _camera_ray_mats(camera_conditions.float(), HW) - matrix_q = raymats.reshape(batch, tokens, 4, 4).transpose(-1, -2).contiguous() - matrix_kv = _invert_se3(raymats.reshape(batch, tokens, 4, 4)).contiguous() - rope_cam = _slice_rope_for_camera(rotary_emb, dim) - q_trans, _q_pre, _q_post = _ucpe_transform( + q_norm = ( + q_raw.float() * q_inv[:, :, None, None] * q_weight[None, None] + ).to(dtype=q_raw.dtype) + k_norm = ( + k_raw.float() * k_inv[:, :, None, None] * k_weight[None, None] + ).to(dtype=k_raw.dtype) + if camera_cache is None: + camera_cache = _prepare_camera_projection_cache( + camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + head_dim=dim, + ) + q_trans = _ucpe_transform_apply( q_norm, - matrix_q, - rope_cam, - inverse_rope=False, - ) - k_trans, _k_pre, _k_post = _ucpe_transform( - k_norm, - matrix_kv, - rope_cam, + camera_cache.proj_q, + camera_cache.rope_cam, inverse_rope=False, - ) - v_trans, _v_pre, _v_post = _ucpe_transform( - v_float, - matrix_kv, - rope_cam, + ).to(dtype=q_raw.dtype) + kv_trans = _ucpe_transform_apply( + torch.cat([k_norm, v_raw], dim=2), + camera_cache.proj_kv, + camera_cache.rope_cam, inverse_rope=False, - ) + ).to(dtype=q_raw.dtype) + k_trans, v_trans = kv_trans.chunk(2, dim=2) def output_projector(out: Tensor) -> Tensor: - projected, _pre, _post = _ucpe_transform( - out.float(), - raymats.reshape(batch, tokens, 4, 4), - rope_cam, + projected = _ucpe_transform_apply( + out.to(dtype=q_raw.dtype), + camera_cache.proj, + camera_cache.rope_cam, inverse_rope=True, ) return projected.to(dtype=out.dtype) @@ -1228,7 +1692,7 @@ def _ucpe_transform( if half % 4 != 0: raise ValueError(f"UCPE requires head_dim/2 divisible by 4, got {half}.") first = x[..., :half].reshape(batch, tokens, heads, half // 4, 4) - first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), first.float()) + first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix, first) first_out = first_out.reshape(batch, tokens, heads, half) second = x[..., half:] if inverse_rope and rotary_emb is not None: @@ -1241,6 +1705,28 @@ def _ucpe_transform( return out, pre_sq, post_sq +def _ucpe_transform_apply( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> Tensor: + batch, tokens, heads, dim = x.shape + half = dim // 2 + if half % 4 != 0: + raise ValueError(f"UCPE requires head_dim/2 divisible by 4, got {half}.") + first = x[..., :half].reshape(batch, tokens, heads, half // 4, 4) + first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), first.float()) + first_out = first_out.reshape(batch, tokens, heads, half) + second = x[..., half:] + if inverse_rope and rotary_emb is not None: + second_out = _apply_complex_rope(second, rotary_emb.conj()) + else: + second_out = _apply_complex_rope(second, rotary_emb) + return torch.cat([first_out, second_out], dim=-1) + + def _camera_ray_mats( camera_conditions: Tensor, HW: tuple[int, int, int], diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py index 0b07b2810..dd744772c 100644 --- a/integrations/sana/sana_wm/transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -337,10 +337,7 @@ def _predict_conditioned( torch.cat([conditioning.uncondition, conditioning.condition], dim=0), _batched_cfg_model_kwargs(conditioning.model_kwargs), ) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0) - flow = noise_pred_uncond + conditioning.cfg_scale * ( - noise_pred_text - noise_pred_uncond - ) + flow = _cfg_guidance(noise_pred, conditioning.cfg_scale) return _zero_conditioned_frame_flow(flow, conditioning) def _predict_with_prompt( @@ -491,6 +488,11 @@ def _batched_cfg_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, obje return kwargs +def _cfg_guidance(noise_pred: Tensor, cfg_scale: float) -> Tensor: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0) + return noise_pred_uncond + cfg_scale * (noise_pred_text - noise_pred_uncond) + + def _conditioned_frame_timestep( *, noisy_latent: Tensor, From a9f56a42a35831002c9d0d11d687cdcf83007aa0 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 24 Jul 2026 13:44:52 -0700 Subject: [PATCH 26/36] Split benchmarks by GPU, chnage metrics, update model card --- .../performance/sana_wm/perf-bf16-0724.md | 5 ++ docs/source/models/sana_wm.rst | 31 +++++++++ integrations/sana/README.md | 14 ++-- .../sana/tests/parity_check/README.md | 31 ++++++--- integrations/sana/tests/parity_check/bench.sh | 11 ++-- .../sana/tests/parity_check/bench_summary.py | 38 ++++++++--- .../tests/parity_check/bench_sweep_summary.py | 15 +++-- .../sana/tests/parity_check/run_native.py | 37 ++++++----- .../tests/test_parity_benchmark_summary.py | 64 +++++++++++++------ integrations/sana/tests/test_smoke.py | 30 +++++---- 10 files changed, 190 insertions(+), 86 deletions(-) create mode 100644 docs/source/_static/performance/sana_wm/perf-bf16-0724.md diff --git a/docs/source/_static/performance/sana_wm/perf-bf16-0724.md b/docs/source/_static/performance/sana_wm/perf-bf16-0724.md new file mode 100644 index 000000000..95ada7025 --- /dev/null +++ b/docs/source/_static/performance/sana_wm/perf-bf16-0724.md @@ -0,0 +1,5 @@ +# SANA-WM BF16 Benchmark Data (ms) + +| device | official | flashdreams | +| --- | ---: | ---: | +| GB202 | 77826.74 | 76938.72 | diff --git a/docs/source/models/sana_wm.rst b/docs/source/models/sana_wm.rst index b58ccb795..735fa64a1 100644 --- a/docs/source/models/sana_wm.rst +++ b/docs/source/models/sana_wm.rst @@ -124,6 +124,37 @@ What to expect See :doc:`/developer_guides/inference_pipeline_overview` for what one pass does end-to-end. +Profiling benchmark +------------------- + +Here is the BF16 profiling benchmark on post-load generation latency per +generated clip for FlashDreams SANA-WM compared to the +`official SANA-WM implementation `_ under +matched settings. + +.. raw:: html + +
+
+
+

+ This chart shows post-load generation latency per generated clip in milliseconds for a 121-frame + Stage-1-only BF16 run. The measured GB202 row used an NVIDIA RTX PRO 6000 Blackwell + Workstation Edition. + Model construction, checkpoint loading, video writing, and frame dumps are outside the timing boundary. + For the official SANA-WM implementation, see + this instruction. +

+
+
+ + Citation -------- diff --git a/integrations/sana/README.md b/integrations/sana/README.md index 16f64236b..2690a0e4e 100644 --- a/integrations/sana/README.md +++ b/integrations/sana/README.md @@ -146,10 +146,16 @@ DEVICE_LABEL="RTX PRO 6000 Blackwell" bash bench.sh ``` `bench.sh` writes `outputs/bench/bench.md` for review and `outputs/bench/perf.md` -for model-card chart data. Both use the same benchmark metric: post-load -generation latency per generated frame. Stage-1 DiT, conditioning/encode, VAE -decode, optional refiner, and memory rows are reported as the breakdown of that -metric. +for chart-ready data. Both use the same benchmark metric: post-load generation +latency per generated clip. SANA-WM renders each requested bidirectional clip in +one generation pass, rather than as independently timed frames. Stage-1 DiT, +conditioning/encode, VAE decode, optional refiner, memory, and frame-normalized +diagnostic rows are reported as breakdowns of that metric. + +For precision sweeps, copy the per-precision chart data from +`outputs/bench//perf.md` into the docs. Each model-card chart should +cover one precision and use GPU/device as the differentiating row, matching the +other FlashDreams model-card benchmark charts. ## Tests diff --git a/integrations/sana/tests/parity_check/README.md b/integrations/sana/tests/parity_check/README.md index d0119a828..5171f7376 100644 --- a/integrations/sana/tests/parity_check/README.md +++ b/integrations/sana/tests/parity_check/README.md @@ -65,18 +65,33 @@ Outputs are under `outputs/bench/`: - `bench.json` - machine-readable inputs, medians, p90s, memory, and stage timings. - `bench.md` - human-readable report. -- `perf.md` - chart-ready model-card data using the same benchmark metric as - `bench.md`. +- `perf.md` - chart-ready data using the same benchmark metric as `bench.md`. -The benchmark metric is post-load generation latency per generated frame. Model +The benchmark metric is post-load generation latency per generated clip. Model construction, checkpoint loading, video writing, and frame dumps are outside the -timing boundary. With the default `NO_REFINER=1`, the timed work covers -conditioning, Stage-1 DiT, and SANA VAE decode. Set `NO_REFINER=0` to benchmark -the full Stage-1 + LTX-2 refiner path with the same timing boundary. +timing boundary. SANA-WM renders each requested bidirectional clip in one +generation pass, rather than as independently timed frames. With the default +`NO_REFINER=1`, the timed work covers conditioning, Stage-1 DiT, and SANA VAE +decode. Set `NO_REFINER=0` to benchmark the full Stage-1 + LTX-2 refiner path +with the same timing boundary. `bench.md` also reports Stage-1 DiT, conditioning/encode, VAE decode, optional -refiner, and memory breakdowns. Those rows explain the benchmark result; they -are not a second benchmark metric. +refiner, memory, and frame-normalized diagnostic breakdowns. Those rows explain +the benchmark result; they are not a second benchmark metric. + +The model card should use one chart per precision. Each chart file should have +GPU/device as the first column and implementation as the series columns, for +example: + +```markdown +| device | official | flashdreams | +| --- | ---: | ---: | +| GB202 | 77826.74 | 76938.72 | +``` + +In precision-sweep mode, the chart-ready files are +`outputs/bench//perf.md`. The top-level `outputs/bench/perf.md` is a +precision summary, not the model-card chart data. Set `FORCE_CUDNN_SDPA=1` only for backend-isolation probes. It is not the default SANA-WM benchmark setting. diff --git a/integrations/sana/tests/parity_check/bench.sh b/integrations/sana/tests/parity_check/bench.sh index 2769d7e59..c7ec2c296 100644 --- a/integrations/sana/tests/parity_check/bench.sh +++ b/integrations/sana/tests/parity_check/bench.sh @@ -15,7 +15,7 @@ # limitations under the License. # Stack-matched SANA-WM upstream vs FlashDreams benchmark. The chart metric is -# post-load generation latency per generated frame. Defaults to Stage-1-only so +# post-load generation latency per generated clip. Defaults to Stage-1-only so # conditioning, Stage-1 DiT, and SANA VAE decode are measured without the LTX-2 # refiner. @@ -90,7 +90,7 @@ if [[ -n "${BENCH_PRECISIONS}" ]]; then STAGE1_PRECISION="${PRECISION}" \ REFINER_PRECISION="${PRECISION}" \ QUANT_BACKEND="${QUANT_BACKEND}" \ - CHART_LABEL="${PRECISION_LABEL}" \ + CHART_LABEL="${CHART_LABEL}" \ bash "${SCRIPT_DIR}/bench.sh" SWEEP_ITEMS+=(--item "${PRECISION_LABEL}:${PRECISION_OUTPUT_DIR}/bench.json") done @@ -98,7 +98,7 @@ if [[ -n "${BENCH_PRECISIONS}" ]]; then echo "[bench] ERROR: BENCH_PRECISIONS produced no precision rows." >&2 exit 1 fi - echo "[bench] aggregating precision sweep -> ${OUTPUT_DIR}/perf.md" + echo "[bench] aggregating precision sweep -> ${OUTPUT_DIR}/bench.md" ( cd "${SCRIPT_DIR}" && \ uv run python "${SCRIPT_DIR}/bench_sweep_summary.py" \ "${SWEEP_ITEMS[@]}" \ @@ -106,8 +106,9 @@ if [[ -n "${BENCH_PRECISIONS}" ]]; then --output-md "${OUTPUT_DIR}/bench.md" \ --output-chart-md "${OUTPUT_DIR}/perf.md" ) echo "[bench] done." - echo " summary: ${OUTPUT_DIR}/bench.md" - echo " chart data: ${OUTPUT_DIR}/perf.md" + echo " precision summary: ${OUTPUT_DIR}/bench.md" + echo " precision data: ${OUTPUT_DIR}/perf.md" + echo " chart data: ${OUTPUT_DIR}/{bf16,fp8,fp4}/perf.md" exit 0 fi diff --git a/integrations/sana/tests/parity_check/bench_summary.py b/integrations/sana/tests/parity_check/bench_summary.py index c0d8bbf22..52b40d8a9 100644 --- a/integrations/sana/tests/parity_check/bench_summary.py +++ b/integrations/sana/tests/parity_check/bench_summary.py @@ -128,11 +128,26 @@ def _ms_per_frame(wall_s: float | None, num_frames: int) -> float | None: return wall_s * 1000.0 / num_frames +def _ms_per_clip(wall_s: float | None) -> float | None: + if wall_s is None: + return None + return wall_s * 1000.0 + + def _metric_value(summary: dict[str, Any], side: str, key: str) -> float | None: value = summary[side].get(key) return float(value) if isinstance(value, (int, float)) else None +def _generation_ms_per_clip( + summary: dict[str, Any], + side: str, + *, + percentile: str = "median", +) -> float | None: + return _ms_per_clip(_metric_value(summary, side, f"wall_{percentile}_s")) + + def _generation_ms_per_frame( summary: dict[str, Any], side: str, @@ -176,15 +191,18 @@ def _render_markdown(summary: dict[str, Any]) -> str: "", "## Benchmark metric", "", - "The chart metric is post-load generation latency per generated frame.", + "The chart metric is post-load generation latency per generated clip.", "Model construction, checkpoint loading, video writing, and frame dumps are outside this timing boundary.", + "SANA-WM renders each requested bidirectional clip in one generation pass, not as independently timed frames.", "With the default `NO_REFINER=1`, the timed work covers conditioning, Stage-1 DiT, and SANA VAE decode.", "", "| metric | upstream | FlashDreams |", "| --- | ---: | ---: |", f"| measured runs | {upstream['runs_measured']} | {native['runs_measured']} |", - f"| generation median / frame | {_fmt(_generation_ms_per_frame(summary, 'upstream'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams'), ' ms')} |", - f"| generation p90 / frame | {_fmt(_generation_ms_per_frame(summary, 'upstream', percentile='p90'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams', percentile='p90'), ' ms')} |", + f"| generation median / clip | {_fmt(_generation_ms_per_clip(summary, 'upstream'), ' ms')} | {_fmt(_generation_ms_per_clip(summary, 'flashdreams'), ' ms')} |", + f"| generation p90 / clip | {_fmt(_generation_ms_per_clip(summary, 'upstream', percentile='p90'), ' ms')} | {_fmt(_generation_ms_per_clip(summary, 'flashdreams', percentile='p90'), ' ms')} |", + f"| generation median / frame, diagnostic | {_fmt(_generation_ms_per_frame(summary, 'upstream'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams'), ' ms')} |", + f"| generation p90 / frame, diagnostic | {_fmt(_generation_ms_per_frame(summary, 'upstream', percentile='p90'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams', percentile='p90'), ' ms')} |", f"| wall median | {_fmt(upstream['wall_median_s'], ' s')} | {_fmt(native['wall_median_s'], ' s')} |", f"| wall p90 | {_fmt(upstream['wall_p90_s'], ' s')} | {_fmt(native['wall_p90_s'], ' s')} |", "", @@ -209,13 +227,13 @@ def _render_markdown(summary: dict[str, Any]) -> str: def _render_chart_markdown(summary: dict[str, Any], device_label: str) -> str: - official = _generation_ms_per_frame(summary, "upstream") - flashdreams = _generation_ms_per_frame(summary, "flashdreams") + official = _generation_ms_per_clip(summary, "upstream") + flashdreams = _generation_ms_per_clip(summary, "flashdreams") if official is None or flashdreams is None: raise ValueError("cannot render chart data without wall-clock stats") return "\n".join( [ - "# SANA-WM Benchmark Data (ms/frame)", + "# SANA-WM Benchmark Data (ms)", "", "| device | official | flashdreams |", "| --- | ---: | ---: |", @@ -298,16 +316,16 @@ def main(argv: list[str] | None = None) -> None: ), } summary["benchmark"] = { - "metric": "generation_ms_per_frame", - "unit": "ms/frame", + "metric": "generation_ms_per_clip", + "unit": "ms", "timing_boundary": ( "pipeline.generate after model setup; excludes model construction, " "checkpoint loading, video writing, and frame dumps" ), "device_label": args.device_label, "chart_label": args.chart_label or args.device_label, - "official": _generation_ms_per_frame(summary, "upstream"), - "flashdreams": _generation_ms_per_frame(summary, "flashdreams"), + "official": _generation_ms_per_clip(summary, "upstream"), + "flashdreams": _generation_ms_per_clip(summary, "flashdreams"), } args.output_json.parent.mkdir(parents=True, exist_ok=True) args.output_json.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") diff --git a/integrations/sana/tests/parity_check/bench_sweep_summary.py b/integrations/sana/tests/parity_check/bench_sweep_summary.py index c2343325f..258b752ff 100644 --- a/integrations/sana/tests/parity_check/bench_sweep_summary.py +++ b/integrations/sana/tests/parity_check/bench_sweep_summary.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Aggregate SANA-WM precision benchmark summaries into chart data.""" +"""Aggregate SANA-WM precision benchmark summaries.""" from __future__ import annotations @@ -58,7 +58,7 @@ def _load_item(raw: str) -> dict[str, Any]: def _render_chart(rows: list[dict[str, Any]]) -> str: lines = [ - "# SANA-WM Benchmark Data (ms/frame)", + "# SANA-WM Precision Sweep Summary (ms)", "", "| precision | official | flashdreams |", "| --- | ---: | ---: |", @@ -75,15 +75,16 @@ def _render_report(rows: list[dict[str, Any]]) -> str: lines = [ "# SANA-WM precision benchmark sweep", "", - "The chart metric is post-load generation latency per generated frame.", + "The chart metric is post-load generation latency per generated clip.", + "Model-card chart data is grouped by GPU/device in each precision subdirectory's `perf.md`.", "", "| precision | official | FlashDreams | source |", "| --- | ---: | ---: | --- |", ] for row in rows: lines.append( - f"| {row['label']} | {row['official']:.2f} ms/frame | " - f"{row['flashdreams']:.2f} ms/frame | `{row['path']}` |" + f"| {row['label']} | {row['official']:.2f} ms | " + f"{row['flashdreams']:.2f} ms | `{row['path']}` |" ) lines.append("") return "\n".join(lines) @@ -105,8 +106,8 @@ def main(argv: list[str] | None = None) -> None: payload = { "benchmark": { - "metric": "generation_ms_per_frame", - "unit": "ms/frame", + "metric": "generation_ms_per_clip", + "unit": "ms", "timing_boundary": ( "pipeline.generate after model setup; excludes model construction, " "checkpoint loading, video writing, and frame dumps" diff --git a/integrations/sana/tests/parity_check/run_native.py b/integrations/sana/tests/parity_check/run_native.py index c93903fb9..d49601ab7 100644 --- a/integrations/sana/tests/parity_check/run_native.py +++ b/integrations/sana/tests/parity_check/run_native.py @@ -249,6 +249,20 @@ def main() -> None: if vae_decoder is not None: vae_decoder._ensure_vae() + request = SanaWMI2VConditioningRequest( + image=image, + prompt=prompt, + poses_c2w=c2w, + intrinsics_vec4=intrinsics_vec4, + num_frames=num_frames, + fps=args.fps, + steps=args.step, + cfg_scale=args.cfg_scale, + flow_shift=args.flow_shift, + seed=args.seed, + negative_prompt=args.negative_prompt, + ) + cache = pipeline.initialize_cache( decoder_context={ "prompt": prompt, @@ -262,23 +276,12 @@ def main() -> None: torch.cuda.reset_peak_memory_stats(device) torch.cuda.synchronize(device) generation_start = time.perf_counter() - decoded = pipeline.generate( - 0, - cache, - input=SanaWMI2VConditioningRequest( - image=image, - prompt=prompt, - poses_c2w=c2w, - intrinsics_vec4=intrinsics_vec4, - num_frames=num_frames, - fps=args.fps, - steps=args.step, - cfg_scale=args.cfg_scale, - flow_shift=args.flow_shift, - seed=args.seed, - negative_prompt=args.negative_prompt, - ), - ) + with torch.inference_mode(): + decoded = pipeline.generate( + 0, + cache, + input=request, + ) if torch.cuda.is_available(): torch.cuda.synchronize(device) wall_s = time.perf_counter() - generation_start diff --git a/integrations/sana/tests/test_parity_benchmark_summary.py b/integrations/sana/tests/test_parity_benchmark_summary.py index cebdc2d5f..2a10e18b6 100644 --- a/integrations/sana/tests/test_parity_benchmark_summary.py +++ b/integrations/sana/tests/test_parity_benchmark_summary.py @@ -47,7 +47,7 @@ def _load_bench_sweep_summary() -> ModuleType: return module -def test_benchmark_summary_uses_generation_ms_per_frame_for_chart() -> None: +def test_benchmark_summary_uses_generation_ms_per_clip_for_chart() -> None: module = _load_bench_summary() summary = { "inputs": { @@ -90,17 +90,18 @@ def test_benchmark_summary_uses_generation_ms_per_frame_for_chart() -> None: assert "## Benchmark metric" in report assert "- stage1_precision: `fp8`" in report - assert "generation median / frame | 100.00 ms | 75.00 ms" in report + assert "generation median / clip | 2000.00 ms | 1500.00 ms" in report + assert "generation median / frame, diagnostic | 100.00 ms | 75.00 ms" in report assert "## Timing breakdown" in report assert "| Stage-1 DiT median | 1500.00 ms | 1100.00 ms |" in report assert "Stage-1 DiT metric" not in report assert "Generation after model load" not in report assert chart == ( - "# SANA-WM Benchmark Data (ms/frame)\n" + "# SANA-WM Benchmark Data (ms)\n" "\n" "| device | official | flashdreams |\n" "| --- | ---: | ---: |\n" - "| Test GPU | 100.00 | 75.00 |\n" + "| Test GPU | 2000.00 | 1500.00 |\n" ) @@ -164,8 +165,6 @@ def test_benchmark_summary_writes_chart_data(tmp_path: Path) -> None: "42", "--device-label", "Test GPU", - "--chart-label", - "BF16", "--stage1-precision", "bf16", "--refiner-precision", @@ -179,35 +178,57 @@ def test_benchmark_summary_writes_chart_data(tmp_path: Path) -> None: ] ) - assert "| BF16 | 100.00 | 50.00 |" in out_chart.read_text(encoding="utf-8") - assert "generation median / frame | 100.00 ms | 50.00 ms" in out_md.read_text( - encoding="utf-8" - ) + assert "| Test GPU | 4000.00 | 2000.00 |" in out_chart.read_text(encoding="utf-8") + report = out_md.read_text(encoding="utf-8") + assert "generation median / clip | 4000.00 ms | 2000.00 ms" in report + assert "generation median / frame, diagnostic | 100.00 ms | 50.00 ms" in report + assert "generation median / frame |" not in report + assert "per generated clip" in report + assert "not as independently timed frames" in report + assert "per generated frame" not in report summary = json.loads(out_json.read_text(encoding="utf-8")) assert summary["benchmark"] == { - "metric": "generation_ms_per_frame", - "unit": "ms/frame", + "metric": "generation_ms_per_clip", + "unit": "ms", "timing_boundary": ( "pipeline.generate after model setup; excludes model construction, " "checkpoint loading, video writing, and frame dumps" ), "device_label": "Test GPU", - "chart_label": "BF16", - "official": 100.0, - "flashdreams": 50.0, + "chart_label": "Test GPU", + "official": 4000.0, + "flashdreams": 2000.0, } +def test_benchmark_summary_keeps_frame_normalized_diagnostics() -> None: + module = _load_bench_summary() + assert module._generation_ms_per_frame( + { + "inputs": {"num_frames": 40}, + "upstream": {"wall_median_s": 4.0}, + }, + "upstream", + ) == 100.0 + assert module._generation_ms_per_clip( + { + "inputs": {"num_frames": 40}, + "upstream": {"wall_median_s": 4.0}, + }, + "upstream", + ) == 4000.0 + + def test_benchmark_sweep_summary_writes_precision_chart_data(tmp_path: Path) -> None: module = _load_bench_sweep_summary() bf16_json = tmp_path / "bf16.json" fp8_json = tmp_path / "fp8.json" bf16_json.write_text( - json.dumps({"benchmark": {"official": 100.0, "flashdreams": 90.0}}), + json.dumps({"benchmark": {"official": 1000.0, "flashdreams": 900.0}}), encoding="utf-8", ) fp8_json.write_text( - json.dumps({"benchmark": {"official": 80.0, "flashdreams": 70.0}}), + json.dumps({"benchmark": {"official": 800.0, "flashdreams": 700.0}}), encoding="utf-8", ) out_json = tmp_path / "bench.json" @@ -230,13 +251,14 @@ def test_benchmark_sweep_summary_writes_precision_chart_data(tmp_path: Path) -> ) assert out_chart.read_text(encoding="utf-8") == ( - "# SANA-WM Benchmark Data (ms/frame)\n" + "# SANA-WM Precision Sweep Summary (ms)\n" "\n" "| precision | official | flashdreams |\n" "| --- | ---: | ---: |\n" - "| BF16 | 100.00 | 90.00 |\n" - "| FP8 | 80.00 | 70.00 |\n" + "| BF16 | 1000.00 | 900.00 |\n" + "| FP8 | 800.00 | 700.00 |\n" ) payload = json.loads(out_json.read_text(encoding="utf-8")) - assert payload["benchmark"]["metric"] == "generation_ms_per_frame" + assert payload["benchmark"]["metric"] == "generation_ms_per_clip" + assert payload["benchmark"]["unit"] == "ms" assert [row["label"] for row in payload["rows"]] == ["BF16", "FP8"] diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index e079e56ed..38c08f97d 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -671,6 +671,8 @@ def test_stage1_model_matches_checkpoint_schema() -> None: """Pin the Stage-1 module to the public checkpoint schema.""" state = SanaWMStage1Model().state_dict() + assert SANA_WM_STAGE1_SPEC.chunk_size is None + assert SANA_WM_STAGE1_SPEC.chunk_split_strategy == "first_chunk_plus_one" assert len(state) == 872 assert tuple(state["x_embedder.proj.weight"].shape) == (2240, 128, 1, 1, 1) assert tuple(state["raymap_embedder.proj.weight"].shape) == (2240, 3, 1, 1, 1) @@ -843,8 +845,8 @@ def test_postprocess_releases_stage1_when_offloading() -> None: assert cache.conditioning is None -def test_vae_tiling_uses_low_memory_tiles() -> None: - """Keep VAE decode on the configured low-memory tiles.""" +def test_vae_tiling_defaults_match_upstream_fast_path() -> None: + """Keep VAE decode on the upstream-style fast tiles by default.""" class DummyVAE: def __init__(self) -> None: @@ -869,20 +871,20 @@ def enable_tiling(self, **kwargs: int) -> None: assert decoder.vae.calls == [ { - "tile_sample_min_height": 256, - "tile_sample_stride_height": 192, - "tile_sample_min_width": 256, - "tile_sample_stride_width": 224, - "tile_sample_min_num_frames": 24, - "tile_sample_stride_num_frames": 8, + "tile_sample_min_height": 512, + "tile_sample_stride_height": 448, + "tile_sample_min_width": 512, + "tile_sample_stride_width": 448, + "tile_sample_min_num_frames": 96, + "tile_sample_stride_num_frames": 64, } ] - assert decoder.vae.tile_sample_min_height == 256 - assert decoder.vae.tile_sample_stride_height == 192 - assert decoder.vae.tile_sample_min_width == 256 - assert decoder.vae.tile_sample_stride_width == 224 - assert decoder.vae.tile_sample_min_num_frames == 24 - assert decoder.vae.tile_sample_stride_num_frames == 8 + assert decoder.vae.tile_sample_min_height == 512 + assert decoder.vae.tile_sample_stride_height == 448 + assert decoder.vae.tile_sample_min_width == 512 + assert decoder.vae.tile_sample_stride_width == 448 + assert decoder.vae.tile_sample_min_num_frames == 96 + assert decoder.vae.tile_sample_stride_num_frames == 64 assert decoder.vae.use_framewise_encoding is True assert decoder.vae.use_framewise_decoding is True From 9372087096d068df2eeff2331028effbd5e7aac3 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Fri, 24 Jul 2026 15:45:51 -0700 Subject: [PATCH 27/36] Remove uupstream TE dep --- integrations/sana/tests/parity_check/pyproject.toml | 10 +++++++++- integrations/sana/tests/parity_check/uv.lock | 9 +++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/integrations/sana/tests/parity_check/pyproject.toml b/integrations/sana/tests/parity_check/pyproject.toml index 0f2fbdf09..f92ac2738 100644 --- a/integrations/sana/tests/parity_check/pyproject.toml +++ b/integrations/sana/tests/parity_check/pyproject.toml @@ -30,11 +30,19 @@ dependencies = [ "torch>=2.11", "torchvision>=0.26", "tqdm>=4.60", - "transformer-engine[pytorch,core-cu13]>=2.12; sys_platform != 'win32'", "transformers>=5.0,<6", "triton>=3.6; sys_platform == 'linux'", ] +[project.optional-dependencies] +# Transformer Engine is only needed for the FP8/FP4 quant backends. It is gated +# behind an opt-in extra so the default `uv sync` (BF16 benchmarking) does not +# try to build/import it. Install with `uv sync --extra quant` once a +# transformer-engine-torch build matching the local torch/CUDA ABI is available. +quant = [ + "transformer-engine[pytorch,core-cu13]>=2.12; sys_platform != 'win32'", +] + [tool.uv.sources] flashdreams = { path = "../../../../flashdreams", editable = true } flashdreams-sana-wm = { path = "../..", editable = true } diff --git a/integrations/sana/tests/parity_check/uv.lock b/integrations/sana/tests/parity_check/uv.lock index baa519ed3..81c9f57a0 100644 --- a/integrations/sana/tests/parity_check/uv.lock +++ b/integrations/sana/tests/parity_check/uv.lock @@ -1239,11 +1239,15 @@ dependencies = [ { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, { name = "torchvision" }, { name = "tqdm" }, - { name = "transformer-engine", extra = ["core-cu13", "pytorch"], marker = "sys_platform != 'win32'" }, { name = "transformers" }, { name = "triton", marker = "sys_platform == 'linux'" }, ] +[package.optional-dependencies] +quant = [ + { name = "transformer-engine", extra = ["core-cu13", "pytorch"], marker = "sys_platform != 'win32'" }, +] + [package.metadata] requires-dist = [ { name = "accelerate", specifier = ">=1.3" }, @@ -1272,10 +1276,11 @@ requires-dist = [ { name = "torch", specifier = ">=2.11" }, { name = "torchvision", specifier = ">=0.26" }, { name = "tqdm", specifier = ">=4.60" }, - { name = "transformer-engine", extras = ["pytorch", "core-cu13"], marker = "sys_platform != 'win32'", specifier = ">=2.12" }, + { name = "transformer-engine", extras = ["pytorch", "core-cu13"], marker = "sys_platform != 'win32' and extra == 'quant'", specifier = ">=2.12" }, { name = "transformers", specifier = ">=5.0,<6" }, { name = "triton", marker = "sys_platform == 'linux'", specifier = ">=3.6" }, ] +provides-extras = ["quant"] [[package]] name = "scipy" From 1435fc034b2790612f60287b891522284cf61d98 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 27 Jul 2026 10:23:45 -0700 Subject: [PATCH 28/36] Update data after Claude review --- .../sana_wm/{perf-bf16-0724.md => perf-bf16-0727.md} | 2 +- docs/source/models/sana_wm.rst | 5 +++-- integrations/sana/sana_wm/stage1_model.py | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) rename docs/source/_static/performance/sana_wm/{perf-bf16-0724.md => perf-bf16-0727.md} (74%) diff --git a/docs/source/_static/performance/sana_wm/perf-bf16-0724.md b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md similarity index 74% rename from docs/source/_static/performance/sana_wm/perf-bf16-0724.md rename to docs/source/_static/performance/sana_wm/perf-bf16-0727.md index 95ada7025..9e9004f8e 100644 --- a/docs/source/_static/performance/sana_wm/perf-bf16-0724.md +++ b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md @@ -2,4 +2,4 @@ | device | official | flashdreams | | --- | ---: | ---: | -| GB202 | 77826.74 | 76938.72 | +| GB202 | 76503.01 | 77109.93 | diff --git a/docs/source/models/sana_wm.rst b/docs/source/models/sana_wm.rst index 735fa64a1..ac297d626 100644 --- a/docs/source/models/sana_wm.rst +++ b/docs/source/models/sana_wm.rst @@ -130,7 +130,8 @@ Profiling benchmark Here is the BF16 profiling benchmark on post-load generation latency per generated clip for FlashDreams SANA-WM compared to the `official SANA-WM implementation `_ under -matched settings. +matched settings. On this metric the two implementations are at parity: the +measured difference is within run-to-run variance (under 1%). .. raw:: html @@ -138,7 +139,7 @@ matched settings.
diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index c376a14c5..4d44eab5f 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -1692,7 +1692,7 @@ def _ucpe_transform( if half % 4 != 0: raise ValueError(f"UCPE requires head_dim/2 divisible by 4, got {half}.") first = x[..., :half].reshape(batch, tokens, heads, half // 4, 4) - first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix, first) + first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), first.float()) first_out = first_out.reshape(batch, tokens, heads, half) second = x[..., half:] if inverse_rope and rotary_emb is not None: From edf91668864cd9abc7cb6733c2179825725e6f79 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 27 Jul 2026 10:53:18 -0700 Subject: [PATCH 29/36] Uaw GB300 numbers for model card --- docs/source/_static/performance/sana_wm/perf-bf16-0727.md | 2 +- docs/source/models/sana_wm.rst | 7 +++---- integrations/sana/tests/parity_check/bench.sh | 2 +- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/docs/source/_static/performance/sana_wm/perf-bf16-0727.md b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md index 9e9004f8e..8ceb9ae16 100644 --- a/docs/source/_static/performance/sana_wm/perf-bf16-0727.md +++ b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md @@ -2,4 +2,4 @@ | device | official | flashdreams | | --- | ---: | ---: | -| GB202 | 76503.01 | 77109.93 | +| GB300 | 33051.56 | 31640.88 | diff --git a/docs/source/models/sana_wm.rst b/docs/source/models/sana_wm.rst index ac297d626..8df606d75 100644 --- a/docs/source/models/sana_wm.rst +++ b/docs/source/models/sana_wm.rst @@ -130,8 +130,8 @@ Profiling benchmark Here is the BF16 profiling benchmark on post-load generation latency per generated clip for FlashDreams SANA-WM compared to the `official SANA-WM implementation `_ under -matched settings. On this metric the two implementations are at parity: the -measured difference is within run-to-run variance (under 1%). +matched settings. On GB300, FlashDreams is about 4% faster than the official +implementation at median clip latency. .. raw:: html @@ -146,8 +146,7 @@ measured difference is within run-to-run variance (under 1%).

This chart shows post-load generation latency per generated clip in milliseconds for a 121-frame - Stage-1-only BF16 run. The measured GB202 row used an NVIDIA RTX PRO 6000 Blackwell - Workstation Edition. + Stage-1-only BF16 run. The measured row used an NVIDIA GB300. Model construction, checkpoint loading, video writing, and frame dumps are outside the timing boundary. For the official SANA-WM implementation, see this instruction. diff --git a/integrations/sana/tests/parity_check/bench.sh b/integrations/sana/tests/parity_check/bench.sh index c7ec2c296..10df4260c 100644 --- a/integrations/sana/tests/parity_check/bench.sh +++ b/integrations/sana/tests/parity_check/bench.sh @@ -30,7 +30,7 @@ _abspath() { esac } -SANA_REPO="$(_abspath "${SANA_REPO:-${HOME}/dev/Sana}")" +SANA_REPO="$(_abspath "${SANA_REPO:-${SCRIPT_DIR}/Sana}")" PATCH_FILE="${SCRIPT_DIR}/changes.patch" REPO_URL="https://github.com/NVlabs/Sana.git" PIN_COMMIT="6298508" From fd6faa6787223bc2f83296e6cfed7083709f391c Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 27 Jul 2026 11:20:53 -0700 Subject: [PATCH 30/36] Fix bench labels --- .../sana/tests/parity_check/bench_summary.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/integrations/sana/tests/parity_check/bench_summary.py b/integrations/sana/tests/parity_check/bench_summary.py index 52b40d8a9..fa76ed17f 100644 --- a/integrations/sana/tests/parity_check/bench_summary.py +++ b/integrations/sana/tests/parity_check/bench_summary.py @@ -213,15 +213,25 @@ def _render_markdown(summary: dict[str, Any]) -> str: f"| Stage-1 incl. conditioning median | {_fmt(upstream['stage1_total_median_ms'], ' ms')} | {_fmt(_sum_optional(native['encode_median_ms'], native['dit_median_ms']), ' ms')} |", f"| Stage-1 DiT median | {_fmt(upstream['dit_median_ms'], ' ms')} | {_fmt(native['dit_median_ms'], ' ms')} |", f"| conditioning/encode median | n/a | {_fmt(native['encode_median_ms'], ' ms')} |", - f"| VAE decode median | {_fmt(upstream['vae_decode_median_ms'], ' ms')} | {_fmt(native['vae_decode_median_ms'], ' ms')} |", - f"| peak memory median | {_fmt(upstream['mem_peak_median_gib'], ' GiB')} | {_fmt(native['mem_peak_median_gib'], ' GiB')} |", ] - if upstream.get("refiner_median_ms") is not None: - rows.extend( - [ - f"| refiner median | {_fmt(upstream['refiner_median_ms'], ' ms')} | n/a |", - ] + if summary["inputs"]["no_refiner"]: + # No refiner: FlashDreams `decode_ms` and upstream `vae_decode_s` are both + # the pure SANA VAE decode, so they compare directly. + rows.append( + f"| VAE decode median | {_fmt(upstream['vae_decode_median_ms'], ' ms')} | {_fmt(native['vae_decode_median_ms'], ' ms')} |" ) + else: + # Refiner enabled: both sides bundle the refiner denoise together with its + # VAE decode into a single measurement (upstream `refiner_s`, FlashDreams + # `decode_ms`). They are only apples-to-apples as one combined row; the + # standalone upstream `vae_decode_s` is the Stage-1 decode and is NOT + # comparable to FlashDreams `decode_ms`. + rows.append( + f"| refiner + VAE decode median | {_fmt(upstream.get('refiner_median_ms'), ' ms')} | {_fmt(native['vae_decode_median_ms'], ' ms')} |" + ) + rows.append( + f"| peak memory median | {_fmt(upstream['mem_peak_median_gib'], ' GiB')} | {_fmt(native['mem_peak_median_gib'], ' GiB')} |" + ) rows.append("") return "\n".join(rows) From 67ed4093ceaf102ac7fe4b810c03a1fa5ab8b52a Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 27 Jul 2026 12:39:17 -0700 Subject: [PATCH 31/36] Update result to use refiner --- docs/source/_static/performance/sana_wm/perf-bf16-0727.md | 2 +- docs/source/models/sana_wm.rst | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/source/_static/performance/sana_wm/perf-bf16-0727.md b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md index 8ceb9ae16..64ba5dca3 100644 --- a/docs/source/_static/performance/sana_wm/perf-bf16-0727.md +++ b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md @@ -2,4 +2,4 @@ | device | official | flashdreams | | --- | ---: | ---: | -| GB300 | 33051.56 | 31640.88 | +| GB300 | 45823.92 | 42399.31 | diff --git a/docs/source/models/sana_wm.rst b/docs/source/models/sana_wm.rst index 8df606d75..a25299bda 100644 --- a/docs/source/models/sana_wm.rst +++ b/docs/source/models/sana_wm.rst @@ -130,8 +130,9 @@ Profiling benchmark Here is the BF16 profiling benchmark on post-load generation latency per generated clip for FlashDreams SANA-WM compared to the `official SANA-WM implementation `_ under -matched settings. On GB300, FlashDreams is about 4% faster than the official -implementation at median clip latency. +matched settings. On GB300, FlashDreams is about 7% faster than the official +implementation at median clip latency (full pipeline: Stage-1 DiT + LTX-2 +refiner, 10 measured runs). .. raw:: html @@ -146,7 +147,7 @@ implementation at median clip latency.

This chart shows post-load generation latency per generated clip in milliseconds for a 121-frame - Stage-1-only BF16 run. The measured row used an NVIDIA GB300. + full-pipeline BF16 run (Stage-1 DiT + LTX-2 refiner). The measured row used an NVIDIA GB300. Model construction, checkpoint loading, video writing, and frame dumps are outside the timing boundary. For the official SANA-WM implementation, see this instruction. From 3af585f1f48e45360d0d296c5cf5d53873a0856d Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 27 Jul 2026 13:27:00 -0700 Subject: [PATCH 32/36] Add some logging for debug --- integrations/sana/sana_wm/stage1_model.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index 4d44eab5f..3ad510d00 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -1158,12 +1158,22 @@ def _chunk_index_from_chunk_size( ) +_GDN_PATH_LOGGED = False + + def _use_fused_gdn(x: Tensor) -> bool: - return ( - x.is_cuda - and _fused_bigdn_func is not None - and os.environ.get("SANA_WM_DISABLE_FUSED_GDN", "0") != "1" - ) + available = _fused_bigdn_func is not None + disabled_by_env = os.environ.get("SANA_WM_DISABLE_FUSED_GDN", "0") == "1" + use = x.is_cuda and available and not disabled_by_env + global _GDN_PATH_LOGGED + if not _GDN_PATH_LOGGED and x.is_cuda: + _GDN_PATH_LOGGED = True + print( + f"[sana-wm] GDN path={'fused' if use else 'eager'} " + f"(fused_available={available}, disabled_by_env={disabled_by_env})", + flush=True, + ) + return use def _prepare_camera_projection_cache( From bbfec5c114ba3374735f1f1b0e8b8bd29adb9d38 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 27 Jul 2026 13:49:31 -0700 Subject: [PATCH 33/36] Clean old bench runs --- integrations/sana/tests/parity_check/bench.sh | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/integrations/sana/tests/parity_check/bench.sh b/integrations/sana/tests/parity_check/bench.sh index 10df4260c..c6f1a105f 100644 --- a/integrations/sana/tests/parity_check/bench.sh +++ b/integrations/sana/tests/parity_check/bench.sh @@ -147,6 +147,11 @@ fi mkdir -p "${UPSTREAM_ROOT}" "${NATIVE_ROOT}" +# Remove stale run_* dirs from prior (possibly longer) benchmarks. bench_summary +# reads every run_* dir, so leftovers from an earlier run with more MEASURED_RUNS +# would silently pollute this invocation's aggregate. +rm -rf "${UPSTREAM_ROOT:?}"/run_* "${NATIVE_ROOT:?}"/run_* + UPSTREAM_REFINER_ARGS=() NATIVE_REFINER_ARGS=() if _is_true "${NO_REFINER}"; then From 51767eada49ddec5d6aac0b7a4cc80e3c517ebbb Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Mon, 27 Jul 2026 16:42:36 -0700 Subject: [PATCH 34/36] Strip out upstream fused kernels --- .../sana/sana_wm/ops/fused_cam_gdn.py | 2448 ----------------- integrations/sana/sana_wm/ops/fused_gdn.py | 2249 --------------- .../sana/sana_wm/ops/fused_gdn_chunkwise.py | 2269 --------------- integrations/sana/sana_wm/stage1_model.py | 312 +-- 4 files changed, 39 insertions(+), 7239 deletions(-) delete mode 100644 integrations/sana/sana_wm/ops/fused_cam_gdn.py delete mode 100644 integrations/sana/sana_wm/ops/fused_gdn.py delete mode 100644 integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py diff --git a/integrations/sana/sana_wm/ops/fused_cam_gdn.py b/integrations/sana/sana_wm/ops/fused_cam_gdn.py deleted file mode 100644 index 6f6229637..000000000 --- a/integrations/sana/sana_wm/ops/fused_cam_gdn.py +++ /dev/null @@ -1,2448 +0,0 @@ -"""Triton-fused camera-branch UCPE single-path delta rule. - -Companion to :mod:`diffusion.model.ops.fused_gdn` (main GDN -branch). This module fuses the *camera* branch of -:class:`ChunkCausalGDNUCPESinglePathLiteLA` through two Triton kernels: - -1. ``_cam_prep_kernel`` — fuses, per ``(batch, token, head)``, the RMSNorm - over the full ``C`` channels, ReLU on Q/K, K-scale on K, the UCPE 4x4 - block-diagonal projection matrix on the first ``D/2`` dims, and the - interleaved-pair complex RoPE on the second ``D/2`` dims. Q, K, V are - processed in one pass. The kernel also emits per-token pre-UCPE and - post-UCPE ``||k||^2`` so the caller can compute the inflation-squared - factor used for Dynamic Beta Discounting. - -2. ``_cam_scan_kernel`` — fuses the numerator-only single-path delta-rule - scan per ``(batch, head)``. ``REVERSE=1`` implements the - ``flip_and_shift`` backward pass semantics directly, avoiding the - torch-side flips in the per-chunk backward loop. - -The prep and scan kernels are runtime paths. Torch implementations below are -kept only for fallback backward paths and focused validation. - -Notes: - - V skips RMSNorm / ReLU / K-scale but receives the same UCPE 4x4 + - RoPE transforms as K (apply_fn_kv in reference). - - The short convolution on K and the inverse UCPE output transform - (``apply_fn_o``) stay in PyTorch — they are single lightweight ops - not on the critical path. -""" - -# ruff: noqa: E501 - -from __future__ import annotations - -import os - -import torch -import triton -import triton.language as tl - -from .fused_gdn_chunkwise import cam_scan_chunkwise - -# ============================================================================= -# Scalar helpers -# ============================================================================= - - -def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: - """Invert a 4x4 SE(3) matrix batch (closed-form). - - Mirrors the reference ``_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 reference 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 -# ============================================================================= - - -_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) - - -@triton.jit -def _cam_prep_bwd_kernel( - # --- forward inputs (replayed for ReLU mask + k_post_kscale recompute) --- - q_raw_ptr, # (B, N, H, D) contiguous, any fp dtype - k_raw_ptr, # (B, N, H, D) contiguous (post short-conv on K) - q_norm_w_ptr, # (C,) = (H*D,) float32 - k_norm_w_ptr, # (C,) float32 - q_inv_rms_ptr, # (B, N) float32 — saved from forward - k_inv_rms_ptr, # (B, N) 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 - # --- upstream gradients (B, H, D, N) layout matching forward outputs --- - d_q_out_ptr, # grad of q_out (any dtype, cast to fp32 on load) - eff_d_k_out_ptr, # grad of k_out + inflation_sq contribution through k_out (fp32) - d_v_out_ptr, # grad of v_out - # --- inflation_sq direct grad to k_post_kscale^2 sum (B, H, N) fp32 --- - d_pre_k_sq_ptr, - # --- outputs --- - d_q_post_norm_ptr, # (B, N, H, D) fp32 — grad after RoPE^T+UCPE^T+Kscale+ReLU; consumed by torch RMSNorm bwd - d_k_post_norm_ptr, # (B, N, H, D) fp32 — same for K - dv_raw_ptr, # (B, N, H, D) fp32 — final dv_raw (V skips norm/ReLU/Kscale) - # --- dims --- - H: tl.constexpr, - N: tl.constexpr, - D: tl.constexpr, - D_HALF: tl.constexpr, - N_GROUPS: tl.constexpr, - 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) — matches the forward kernel's parallelism. - - Implements the bwd of the fused fwd: - forward order: RMSNorm -> ReLU -> [K-scale on K] -> UCPE (first D/2) - -> RoPE (second D/2) -> output (B, H, D, N). - backward order: RoPE^T (second D/2) -> UCPE^T (first D/2) - -> K-scale (only K) -> ReLU mask -> emit d_post_norm - intermediates for the cross-head RMSNorm bwd handled - outside (in :func:`_cam_prep_bwd_dispatch`). - - The full-channel RMSNorm bwd's outer-product term and per-channel weight - grad both require a sum over ``H*D`` (the full ``C``) per token, which - couples heads. We deliberately leave that step in PyTorch (a couple of - fused element-wise + reduction ops) — see :func:`_cam_prep_bwd_dispatch`. - - Inflation handling: the kernel takes an *effective* ``dO_k`` that already - includes the contribution from ``grad_inflation_sq`` flowing through - ``k_out`` (i.e. ``eff_dO_k = grad_k + 2 * k_out * d_post_k_sq``), plus a - per-(b, h, n) scalar ``d_pre_k_sq`` that is the chain-rule contribution - of ``grad_inflation_sq`` into the pre-UCPE ``||k_post_kscale||^2`` sum. - Inside the kernel we recompute ``k_post_kscale`` (= post-norm * ReLU * - K_SCALE) and add ``2 * k_post_kscale[d] * d_pre_k_sq`` as a direct - contribution to ``d_k_post_kscale``. - """ - pid = tl.program_id(0) - h_idx = pid % H - bn_idx = pid // H - b_idx = bn_idx // N - n_idx = bn_idx % N - - # ---- load saved scalars: inv-RMS (per b, n) and d_pre_k_sq (per b, h, n) ---- - 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) - bhn_idx = (b_idx * H + h_idx) * N + n_idx - d_pre_k_sq = tl.load(d_pre_k_sq_ptr + bhn_idx).to(tl.float32) - - # ---- load per-token P matrices (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) - - # ---- layout offsets ---- - row_base = b_idx * (N * H * D) + n_idx * (H * D) + h_idx * D # (B, N, H, D) - nw_off = h_idx * D - out_base_BHDN = b_idx * (H * D * N) + h_idx * (D * N) + n_idx # (B, H, D, N) for dO_* - norm_base = row_base # d_*_post_norm and dv_raw share the (B, N, H, D) layout - - # ============================================================ - # First half — UCPE region - # ============================================================ - 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] - - # Recompute post-norm (pre-ReLU) Q/K for the ReLU mask + k_post_kscale. - q_half_raw = tl.load(q_raw_ptr + row_base + offs_gj, mask=mask_gj, other=0.0).to(tl.float32) - k_half_raw = tl.load(k_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_post_norm_half = q_half_raw * q_inv_rms * q_nw_half - k_post_norm_half = k_half_raw * k_inv_rms * k_nw_half - q_relu_mask_half = q_post_norm_half > 0 - k_relu_mask_half = k_post_norm_half > 0 - - # k_post_kscale = relu(k_post_norm) * K_SCALE (used for direct inflation contribution) - k_post_relu_half = tl.where(k_relu_mask_half, k_post_norm_half, 0.0) - k_post_kscale_half = k_post_relu_half * K_SCALE - - # Load upstream gradients for first half. - # In (B, H, D, N): ptr[b, h, d, n] = out_base_BHDN + d * N. d = g*4 + i. - offs_d_half = offs_g[:, None] * 4 + offs_i[None, :] # (BLOCK_GROUPS, 4) - mask_d_half = mask_g[:, None] - dO_q_half = tl.load(d_q_out_ptr + out_base_BHDN + offs_d_half * N, mask=mask_d_half, other=0.0).to(tl.float32) - dO_k_eff_half = tl.load(eff_d_k_out_ptr + out_base_BHDN + offs_d_half * N, mask=mask_d_half, other=0.0).to( - tl.float32 - ) - dO_v_half = tl.load(d_v_out_ptr + out_base_BHDN + offs_d_half * N, mask=mask_d_half, other=0.0).to(tl.float32) - - # UCPE^T: din[g, j] = sum_i P[i, j] * dout[g, i] - # forward: out[g, i] = sum_j q[g, j] * P[i, j] -- so bwd sums over i - d_q_post_relu_half = tl.sum(dO_q_half[:, :, None] * P_q[None, :, :], axis=1) - d_k_post_kscale_via_ucpe_half = tl.sum(dO_k_eff_half[:, :, None] * P_kv[None, :, :], axis=1) - d_v_first_half = tl.sum(dO_v_half[:, :, None] * P_kv[None, :, :], axis=1) - - # K direct inflation contribution: 2 * k_post_kscale * d_pre_k_sq - d_k_post_kscale_half = d_k_post_kscale_via_ucpe_half + 2.0 * k_post_kscale_half * d_pre_k_sq - - # K-scale bwd (multiply by K_SCALE) - d_k_post_relu_half = d_k_post_kscale_half * K_SCALE - - # ReLU mask - d_q_post_norm_half = tl.where(q_relu_mask_half, d_q_post_relu_half, 0.0) - d_k_post_norm_half = tl.where(k_relu_mask_half, d_k_post_relu_half, 0.0) - - # Mask out-of-bounds groups to 0 explicitly (for safety on uneven N_GROUPS). - d_q_post_norm_half = tl.where(mask_d_half, d_q_post_norm_half, 0.0) - d_k_post_norm_half = tl.where(mask_d_half, d_k_post_norm_half, 0.0) - d_v_first_half = tl.where(mask_d_half, d_v_first_half, 0.0) - - # Store at (B, N, H, D), d = g*4+i - tl.store(d_q_post_norm_ptr + norm_base + offs_d_half, d_q_post_norm_half, mask=mask_d_half) - tl.store(d_k_post_norm_ptr + norm_base + offs_d_half, d_k_post_norm_half, mask=mask_d_half) - tl.store(dv_raw_ptr + norm_base + offs_d_half, d_v_first_half, mask=mask_d_half) - - # ============================================================ - # Second half — RoPE region - # ============================================================ - 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_pair = tl.load(rope_sin_ptr + rope_row + offs_r_pair, mask=mask_r_pair, other=0.0).to(tl.float32) - - # Recompute post-norm (pre-ReLU) Q/K for the ReLU mask. - rope_base = row_base + D_HALF - q_r_raw = tl.load(q_raw_ptr + rope_base + offs_r, mask=mask_r, other=0.0).to(tl.float32) - k_r_raw = tl.load(k_raw_ptr + rope_base + offs_r, mask=mask_r, 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_post_norm_r = q_r_raw * q_inv_rms * q_nw_r - k_post_norm_r = k_r_raw * k_inv_rms * k_nw_r - q_relu_mask_r = q_post_norm_r > 0 - k_relu_mask_r = k_post_norm_r > 0 - - k_post_relu_r = tl.where(k_relu_mask_r, k_post_norm_r, 0.0) - k_post_kscale_r = k_post_relu_r * K_SCALE - - # Load upstream gradients (second half) — direct + pair. - offs_d_r = D_HALF + offs_r - offs_d_r_pair = D_HALF + offs_r_pair - dO_q_r = tl.load(d_q_out_ptr + out_base_BHDN + offs_d_r * N, mask=mask_r, other=0.0).to(tl.float32) - dO_k_eff_r = tl.load(eff_d_k_out_ptr + out_base_BHDN + offs_d_r * N, mask=mask_r, other=0.0).to(tl.float32) - dO_v_r = tl.load(d_v_out_ptr + out_base_BHDN + offs_d_r * N, mask=mask_r, other=0.0).to(tl.float32) - dO_q_r_pair = tl.load(d_q_out_ptr + out_base_BHDN + offs_d_r_pair * N, mask=mask_r_pair, other=0.0).to(tl.float32) - dO_k_eff_r_pair = tl.load(eff_d_k_out_ptr + out_base_BHDN + offs_d_r_pair * N, mask=mask_r_pair, other=0.0).to( - tl.float32 - ) - dO_v_r_pair = tl.load(d_v_out_ptr + out_base_BHDN + offs_d_r_pair * N, mask=mask_r_pair, other=0.0).to(tl.float32) - - # RoPE^T: forward y[r] = x[r]*cos[r] + x[r^1]*sin[r] - # bwd dx[r] = dy[r]*cos[r] + dy[r^1]*sin[r^1] - d_q_post_relu_r = dO_q_r * cos_v + dO_q_r_pair * sin_v_pair - d_k_post_kscale_via_rope_r = dO_k_eff_r * cos_v + dO_k_eff_r_pair * sin_v_pair - d_v_second_r = dO_v_r * cos_v + dO_v_r_pair * sin_v_pair - - # K direct inflation contribution - d_k_post_kscale_r = d_k_post_kscale_via_rope_r + 2.0 * k_post_kscale_r * d_pre_k_sq - - # K-scale bwd - d_k_post_relu_r = d_k_post_kscale_r * K_SCALE - - # ReLU mask - d_q_post_norm_r = tl.where(q_relu_mask_r, d_q_post_relu_r, 0.0) - d_k_post_norm_r = tl.where(k_relu_mask_r, d_k_post_relu_r, 0.0) - - # Out-of-bound mask - d_q_post_norm_r = tl.where(mask_r, d_q_post_norm_r, 0.0) - d_k_post_norm_r = tl.where(mask_r, d_k_post_norm_r, 0.0) - d_v_second_r = tl.where(mask_r, d_v_second_r, 0.0) - - norm_offs_r = D_HALF + offs_r - tl.store(d_q_post_norm_ptr + norm_base + norm_offs_r, d_q_post_norm_r, mask=mask_r) - tl.store(d_k_post_norm_ptr + norm_base + norm_offs_r, d_k_post_norm_r, mask=mask_r) - tl.store(dv_raw_ptr + norm_base + norm_offs_r, d_v_second_r, mask=mask_r) - - -@triton.jit -def _cam_scan_kernel( - # --- inputs (B, H, D, N) contiguous, fp32 --- - q_ptr, - k_ptr, - v_ptr, - # --- gates --- - beta_ptr, # (B, H, F, S) contiguous - decay_ptr, # (B, H, F) contiguous - # --- output (B, H, D, N) fp32 --- - out_ptr, - # --- saved state snapshots (used when SAVE_STATES=1) --- - state_pre_ptr, # (B, H, F, BLOCK_D, BLOCK_D) fp32 — state after decay, before update - state_post_ptr, # (B, H, F, BLOCK_D, BLOCK_D) fp32 — state after update - # --- forward-direction cache state (used when LOAD_INIT_STATE / SAVE_FINAL_STATE) --- - init_state_ptr, # (B*H, BLOCK_D, BLOCK_D) fp32 — state at end of prefix - final_state_ptr, # (B*H, BLOCK_D, BLOCK_D) fp32 — state after last frame's update - # --- dims --- - H: tl.constexpr, - F: tl.constexpr, - S: tl.constexpr, - D: tl.constexpr, - N: tl.constexpr, # F * S - REVERSE: tl.constexpr, - SAVE_STATES: tl.constexpr, - LOAD_INIT_STATE: tl.constexpr, - SAVE_FINAL_STATE: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_S: tl.constexpr, -): - """One program per (b, h) — runs the full numerator-only delta-rule scan. - - When ``SAVE_STATES=1`` the kernel additionally writes per-frame snapshots - of ``state_prev`` (state after applying ``decay``, before the K@delta - update) to ``state_pre_ptr`` and ``state_curr`` (state after the update) - to ``state_post_ptr``, both indexed by ``q_frame``. The bwd kernel - (``_cam_scan_bwd_kernel``) consumes these snapshots for both - ``REVERSE=0`` and ``REVERSE=1``. In ``REVERSE=1`` the slot at - ``q_frame=F-1`` always holds the all-zero state (skip-update, decay=1 - on the zero initial state), and the bwd kernel reads exactly that — - no special-case load is needed. - - When ``LOAD_INIT_STATE=1`` (forward direction only — wrapper enforces - ``REVERSE=0``) the per-program ``state_curr`` is initialized from - ``init_state_ptr`` instead of zero. The convention is: the loaded value - is the state AT THE END of a prefix sequence (i.e., AFTER the prefix's - last update, BEFORE any further decay applied here). On the very first - frame, the kernel's own ``state_curr *= g`` then applies ``decay[0]`` - to this loaded state — which is exactly the decay that the global - sequence's f=K-th frame would have applied. This keeps split/resume - state trajectories identical from frame K onwards. - - When ``SAVE_FINAL_STATE=1`` (forward direction only) the final - ``state_curr`` (after the last frame's update) is written to - ``final_state_ptr``. This is the state to be loaded with - ``LOAD_INIT_STATE`` for a downstream segment. - """ - pid = tl.program_id(0) - pid_b = pid // H - pid_h = pid % H - bh = pid_b * H + pid_h - - offs_d = tl.arange(0, BLOCK_D) - mask_d = offs_d < D - offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] - mask_dd = mask_d[:, None] & mask_d[None, :] - - base_bh = pid_b * (H * D * N) + pid_h * (D * N) - q_bh = q_ptr + base_bh - k_bh = k_ptr + base_bh - v_bh = v_ptr + base_bh - out_bh = out_ptr + base_bh - beta_bh = beta_ptr + bh * F * S - decay_bh = decay_ptr + bh * F - if SAVE_STATES: - spre_bh = state_pre_ptr + bh * F * BLOCK_D * BLOCK_D - spost_bh = state_post_ptr + bh * F * BLOCK_D * BLOCK_D - - # State: (D_k, D_v) in the upstream convention. Here we call rows "k-dim" - # (input dim of state) and cols "v-dim" (output dim of state). - if LOAD_INIT_STATE: - init_bh = init_state_ptr + bh * BLOCK_D * BLOCK_D - state_curr = tl.load(init_bh + offs_dd, mask=mask_dd, other=0.0).to(tl.float32) - else: - state_curr = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) - - for f_iter in range(F): - if REVERSE: - q_frame = F - 1 - f_iter - kv_frame = F - f_iter if f_iter > 0 else 0 - skip_update = f_iter == 0 - else: - q_frame = f_iter - kv_frame = f_iter - skip_update = False - - if REVERSE and f_iter == 0: - g = 1.0 - else: - g = tl.load(decay_bh + kv_frame).to(tl.float32) - state_curr = state_curr * g - state_prev = state_curr # fp32 snapshot (same tensor, kept for clarity) - - if SAVE_STATES: - tl.store( - spre_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, - state_prev, - mask=mask_dd, - ) - - if skip_update == 0: - kv_n_base = kv_frame * S - f_beta = beta_bh + kv_frame * 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 = kv_n_base + offs_s - - # Load K, V tiles (BLOCK_S, BLOCK_D) from (B, H, D, N) layout: - # ptr[s, d] = base_bh + offs_d[d] * N + n_idx[s] - k_ptrs = k_bh + offs_d[None, :] * N + n_idx[:, None] - v_ptrs = v_bh + offs_d[None, :] * N + n_idx[:, None] - K = tl.load(k_ptrs, mask=mask_sd, other=0.0) - V = tl.load(v_ptrs, mask=mask_sd, other=0.0) - - bt = tl.load(f_beta + offs_s, mask=mask_s, other=0.0).to(tl.float32) - - # V_pred = K @ state_prev : (BLOCK_S, BLOCK_D) - V_pred = tl.dot( - K, - state_prev, - out_dtype=tl.float32, - input_precision="tf32", - ) - dv = (V - V_pred) * bt[:, None] - state_curr += tl.dot( - tl.trans(K), - dv, - out_dtype=tl.float32, - input_precision="tf32", - ) - - if SAVE_STATES: - tl.store( - spost_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, - state_curr, - mask=mask_dd, - ) - - # --- Pass 2: output --- - state_out = state_curr - q_n_base = q_frame * 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 = q_n_base + offs_s - - q_ptrs = q_bh + offs_d[None, :] * N + n_idx[:, None] - Q = tl.load(q_ptrs, mask=mask_sd, other=0.0) - - # num = Q @ state_out : (BLOCK_S, BLOCK_D), rows=S, cols=D_v - num = tl.dot( - Q, - state_out, - out_dtype=tl.float32, - input_precision="tf32", - ) - - # Store transposed into (B, H, D, N): - # ptr[d, s] = out_bh + offs_d[d] * N + n_idx[s] - out_ptrs = out_bh + offs_d[:, None] * N + n_idx[None, :] - mask_ds = mask_d[:, None] & mask_s[None, :] - tl.store(out_ptrs, tl.trans(num), mask=mask_ds) - - if SAVE_FINAL_STATE: - final_bh = final_state_ptr + bh * BLOCK_D * BLOCK_D - tl.store(final_bh + offs_dd, state_curr, mask=mask_dd) - - -# ============================================================================= -# Backward Triton kernel -# ============================================================================= - - -@triton.jit -def _cam_scan_bwd_kernel( - # --- forward inputs (B, H, D, N) fp32 contiguous --- - q_ptr, - k_ptr, - v_ptr, - # --- gates --- - beta_ptr, # (B, H, F, S) fp32 - decay_ptr, # (B, H, F) fp32 - # --- saved state snapshots (B*H, F, BLOCK_D, BLOCK_D) fp32, indexed by q_frame --- - state_pre_ptr, # state after decay, before update - state_post_ptr, # state after update - # --- upstream gradient (B, H, D, N) fp32 --- - grad_out_ptr, - # --- output gradients --- - dq_ptr, # (B, H, D, N) fp32 - dk_ptr, - dv_ptr, - dbeta_ptr, # (B, H, F, S) fp32 - ddecay_ptr, # (B, H, F) fp32 - # --- dims --- - H: tl.constexpr, - F: tl.constexpr, - S: tl.constexpr, - D: tl.constexpr, - N: tl.constexpr, # F * S - REVERSE: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_S: tl.constexpr, -): - """Reverse-time backward of the numerator-only delta-rule scan. - - One program per ``(b, h)``. Walks the same ``f_iter`` index space as the - forward kernel — but in reverse time order — replaying the recurrence - using the per-``q_frame`` state snapshots saved by the forward pass with - ``SAVE_STATES=1``. Accumulates gradients into ``dq``, ``dk``, ``dv``, - ``dbeta``, ``ddecay``. - - Forward indexing (matched here exactly):: - - REVERSE=False: q_frame = kv_frame = f_iter, skip_update = False - REVERSE=True : q_frame = F-1-f_iter, - kv_frame = F-f_iter (skip when f_iter==0) - g = decay[kv_frame] (1.0 when f_iter==0) - - Per-iteration backward derivation (when ``not skip_update``): - - ds_post += Q.T @ d_out # accumulate output grad - dQ[q] = d_out @ s_post.T - - ddelta = K @ ds_post # via K.T @ delta term - dV[kv] = ddelta * beta[kv,:] - dbeta[kv,:] = sum_d (ddelta * (V - K @ s_pre)) - dV_pred = -ddelta * beta[kv,:] - dK[kv] = delta @ ds_post.T + dV_pred @ s_pre.T - ds_pre = ds_post + K.T @ dV_pred # direct + V_pred path - - ddecay[kv] = sum(ds_pre * s_pre[q]) / g (= 0 when state_in is 0) - ds_post[next-bwd-iter] = ds_pre * g # propagate through decay - - For the ``skip_update`` branch (``REVERSE=True`` and ``f_iter==0``) the - update / decay path is bypassed: ``state_post == state_pre == 0`` so - ``dQ == 0``, ``ds_post`` is left unchanged across the iter, and - no writes to ``dK``, ``dV``, ``dbeta`` or ``ddecay`` happen at - ``kv_frame == 0``. Caller MUST pre-zero those output buffers (the - dispatch wrapper uses ``torch.zeros_like``). - - The ``ddecay`` formula uses ``state_pre / g``; for ``REVERSE=False`` at - fwd frame 0 we hardcode ``ddecay[0] = 0`` (state_in is exactly 0, but - the division would amplify any rounding noise). For very small ``g`` - the existing ``+ 1e-12`` epsilon is matched verbatim from - ``_fused_gdn_bwd_kernel``. - """ - pid = tl.program_id(0) - pid_b = pid // H - pid_h = pid % H - bh = pid_b * H + pid_h - - base_bh = pid_b * (H * D * N) + pid_h * (D * N) - q_bh = q_ptr + base_bh - k_bh = k_ptr + base_bh - v_bh = v_ptr + base_bh - do_bh = grad_out_ptr + base_bh - dq_bh = dq_ptr + base_bh - dk_bh = dk_ptr + base_bh - dv_bh = dv_ptr + base_bh - - beta_bh = beta_ptr + bh * F * S - decay_bh = decay_ptr + bh * F - dbeta_bh = dbeta_ptr + bh * F * S - ddecay_bh = ddecay_ptr + bh * F - - spre_bh = state_pre_ptr + bh * F * BLOCK_D * BLOCK_D - spost_bh = state_post_ptr + bh * F * BLOCK_D * BLOCK_D - - offs_d = tl.arange(0, BLOCK_D) - mask_d = offs_d < D - offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] - mask_dd = mask_d[:, None] & mask_d[None, :] - - # bf16 operands for tl.dot keep shared-memory pressure manageable for large - # BLOCK_D (e.g., 128 in reference). fp32 accumulators preserve precision. - grad_dtype = tl.bfloat16 - grad_ip: tl.constexpr = "tf32" - - # Reverse-time accumulator: gradient w.r.t. ``state_post`` for the iter - # currently being processed. - ds_post = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) - - for f_rev in range(F): - # Walk fwd iters in reverse: F-1, F-2, ..., 0. - f_iter = F - 1 - f_rev - - if REVERSE: - q_frame = F - 1 - f_iter - kv_frame = F - f_iter if f_iter > 0 else 0 - skip_update = f_iter == 0 - else: - q_frame = f_iter - kv_frame = f_iter - skip_update = 0 - - if REVERSE and f_iter == 0: - g = 1.0 - else: - g = tl.load(decay_bh + kv_frame).to(tl.float32) - - # ---- Pass 2: dQ + ds_post += Q.T @ d_out ---------------------- - # Load state_post[q_frame] (zero in the REVERSE skip-update slot — - # the fwd save still writes the all-zero state at that index). - state = tl.load( - spost_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, - mask=mask_dd, - other=0.0, - ) - - q_n_base = q_frame * 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 = q_n_base + offs_s - - q_ptrs = q_bh + offs_d[None, :] * N + n_idx[:, None] - Q = tl.load(q_ptrs, mask=mask_sd, other=0.0) - - do_ptrs = do_bh + offs_d[None, :] * N + n_idx[:, None] - dO = tl.load(do_ptrs, mask=mask_sd, other=0.0) - - ds_post += tl.dot( - tl.trans(Q.to(grad_dtype)), - dO.to(grad_dtype), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - dQ = tl.dot( - dO.to(grad_dtype), - tl.trans(state.to(grad_dtype)), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - dq_ptrs = dq_bh + offs_d[:, None] * N + n_idx[None, :] - mask_ds = mask_d[:, None] & mask_s[None, :] - tl.store(dq_ptrs, tl.trans(dQ), mask=mask_ds) - - if skip_update == 0: - # ---- Reload state with state_pre[q_frame] for Pass 1 ---- - state = tl.load( - spre_bh + q_frame * BLOCK_D * BLOCK_D + offs_dd, - mask=mask_dd, - other=0.0, - ) - - # ds_pre starts equal to ds_post (direct pass-through term). - ds_pre = ds_post - - kv_n_base = kv_frame * 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 = kv_n_base + offs_s - - k_ptrs = k_bh + offs_d[None, :] * N + n_idx[:, None] - v_ptrs = v_bh + offs_d[None, :] * N + n_idx[:, None] - K = tl.load(k_ptrs, mask=mask_sd, other=0.0) - V = tl.load(v_ptrs, mask=mask_sd, other=0.0) - - bt = tl.load(beta_bh + kv_frame * S + offs_s, mask=mask_s, other=0.0).to(tl.float32) - - V_pred = tl.dot( - K.to(grad_dtype), - state.to(grad_dtype), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - r = V - V_pred - delta = r * bt[:, None] - - ddelta = tl.dot( - K.to(grad_dtype), - ds_post.to(grad_dtype), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - dV = ddelta * bt[:, None] - dv_ptrs = dv_bh + offs_d[:, None] * N + n_idx[None, :] - mask_ds = mask_d[:, None] & mask_s[None, :] - tl.store(dv_ptrs, tl.trans(dV), mask=mask_ds) - - dbeta_st = tl.sum(ddelta * r, axis=1) - tl.store(dbeta_bh + kv_frame * S + offs_s, dbeta_st, mask=mask_s) - - dV_pred = -ddelta * bt[:, None] - - dK_part1 = tl.dot( - delta.to(grad_dtype), - tl.trans(ds_post.to(grad_dtype)), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - dK_part2 = tl.dot( - dV_pred.to(grad_dtype), - tl.trans(state.to(grad_dtype)), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - dK = dK_part1 + dK_part2 - dk_ptrs = dk_bh + offs_d[:, None] * N + n_idx[None, :] - tl.store(dk_ptrs, tl.trans(dK), mask=mask_ds) - - ds_pre += tl.dot( - tl.trans(K.to(grad_dtype)), - dV_pred.to(grad_dtype), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - # ---- ddecay[kv_frame] ---- - # state_pre = state_in * g, so state_in = state_pre / g. - # ddecay = sum(ds_pre * state_in) = sum(ds_pre * state_pre) / g. - # For REVERSE=False fwd frame 0, state_in is exactly 0 — hardcode - # 0 to avoid amplifying rounding noise via 1/g. - if (REVERSE == 0) and (f_iter == 0): - ddecay_f = 0.0 - else: - inv_g = 1.0 / (g + 1e-12) - ddecay_f = tl.sum(ds_pre * state) * inv_g - tl.store(ddecay_bh + kv_frame, ddecay_f) - - # Propagate to next bwd iter (which is fwd's previous iter). - ds_post = ds_pre * g - # else (skip_update branch): state_post == state_pre == 0, no - # ddecay write, no kv-side writes; ds_post passes through unchanged - # since ∂state_post/∂state_in = I when the update is skipped and g=1. - # In REVERSE=True this is the LAST bwd iter (f_iter=0) so the - # carried-over ds_post is discarded. - - -# ============================================================================= -# Python wrappers -# ============================================================================= - - -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. - """ - 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 - - -def _run_cam_prep_fwd_save( - 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, - proj_kv: torch.Tensor, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - k_scale: float, - norm_eps: float, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Run :func:`cam_prep_func` while exposing intermediates needed by the bwd kernel. - - Mirrors :func:`cam_prep_func` exactly (same kernel launch, same outputs) - but additionally returns the per-token ``inv_rms`` for both Q and K, plus - the raw ``||k_pre_ucpe||^2`` and ``||k_post_ucpe||^2`` per ``(B, H, N)``. - These are required by :func:`_cam_prep_bwd_dispatch` to (a) replay the - ReLU/RMSNorm chain in fp32 without re-summing over the full ``C`` - channels and (b) chain ``grad_inflation_sq`` back through ``k_out`` and - the pre-UCPE ``||k||^2`` term with the correct ``clamp_min`` indicators. - - Returns: - ``(q_out, k_out, v_out, inflation_sq, q_inv_rms, k_inv_rms, - k_pre_sq, k_post_sq)``. The first four match :func:`cam_prep_func`; - the rest are fp32 contiguous saved-state tensors for the backward. - """ - 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 - - 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) - - _cam_prep_kernel[(B * N * H,)]( - 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 = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12) - return q_out, k_out, v_out, inflation_sq, q_inv_rms, k_inv_rms, k_pre_sq, k_post_sq - - -def _cam_prep_bwd_dispatch( - q_raw: torch.Tensor, - k_raw: torch.Tensor, - q_norm_weight: torch.Tensor, - k_norm_weight: torch.Tensor, - proj_q: torch.Tensor, - proj_kv: torch.Tensor, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - q_inv_rms: torch.Tensor, - k_inv_rms: torch.Tensor, - k_pre_sq: torch.Tensor, - k_post_sq: torch.Tensor, - k_out: torch.Tensor, - *, - grad_q: torch.Tensor | None, - grad_k: torch.Tensor | None, - grad_v: torch.Tensor | None, - grad_inflation_sq: torch.Tensor | None, - k_scale: float, -) -> tuple[ - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, -]: - """Hybrid Triton + torch backward for :func:`cam_prep_func`. - - Pipeline: - - 1. **(torch)** Chain ``grad_inflation_sq`` through ``k_post_sq`` and - ``k_pre_sq`` to produce - (a) ``eff_d_k_out = grad_k + 2 * k_out * d_post_k_sq`` and - (b) ``d_pre_k_sq`` — a per-(B, H, N) scalar fed into the kernel as a - direct contribution to ``d_k_post_kscale``. Both ``clamp_min(1e-12)`` - indicators are honored so the gradient is exactly 0 in the (rare) - saturating regime, matching :func:`_torch_cam_prep_reference`. - - 2. **(Triton)** Launch :func:`_cam_prep_bwd_kernel` per ``(b, n, h)`` - to apply RoPE^T, UCPE^T, K-scale^T and the ReLU mask. Emits - ``d_q_post_norm`` / ``d_k_post_norm`` ``(B, N, H, D)`` fp32 - intermediates (the post-norm pre-RMSNorm grad slots) and writes - ``dv_raw`` directly (V skips RMSNorm/ReLU/K-scale). - - 3. **(torch)** Apply the full-channel RMSNorm bwd. The cross-head - coupling means ``S_q[b, n] = sum_{h,d} d_q_post_norm * q_raw * - q_norm_w`` is reduced over ``H*D``; we then form - ``dq_raw = d_q_post_norm * inv_rms_q * q_norm_w - - inv_rms_q^3 / C * q_raw * S_q`` - elementwise, and ``dq_norm_weight[c] = sum_{b,n} - d_q_post_norm[b,n,c] * q_raw[b,n,c] * inv_rms_q[b,n]``. - - Returns: - ``(dq_raw, dk_raw, dv_raw, dq_norm_weight, dk_norm_weight)``. Each - slot is ``None`` if upstream did not request that grad — handled - by the caller via ``ctx.needs_input_grad``. ``dq_raw`` / ``dk_raw`` - / ``dv_raw`` are returned in the same dtype as ``q_raw``; - norm-weight grads are fp32 (matching the input dtype). - """ - B, N, H, D = q_raw.shape - assert k_raw.shape == q_raw.shape - assert q_raw.is_contiguous() and k_raw.is_contiguous() - assert q_inv_rms.shape == (B, N) and k_inv_rms.shape == (B, N) - assert k_pre_sq.shape == (B, H, N) and k_post_sq.shape == (B, H, N) - D_half = D // 2 - N_groups = D_half // 4 - C = H * D - - # ---- prepare inflation_sq grad chain (in fp32) ---- - eps_floor = 1e-12 - pre_clamped = k_pre_sq.clamp_min(eps_floor) - if grad_inflation_sq is not None: - gis = grad_inflation_sq.to(torch.float32) - post_clamped = k_post_sq.clamp_min(eps_floor) - pre_indicator = (k_pre_sq >= eps_floor).to(torch.float32) - post_indicator = (k_post_sq >= eps_floor).to(torch.float32) - # d(inflation_sq)/d(post_k_sq) = (1 / pre_clamped) * post_indicator - # d(inflation_sq)/d(pre_k_sq) = -post_clamped / pre_clamped^2 * pre_indicator - d_post_k_sq = (post_indicator / pre_clamped) * gis # (B, H, N) - d_pre_k_sq = (-post_clamped / (pre_clamped * pre_clamped) * pre_indicator) * gis # (B, H, N) - else: - d_post_k_sq = torch.zeros_like(k_pre_sq) - d_pre_k_sq = torch.zeros_like(k_pre_sq) - - # eff_d_k_out: (B, H, D, N) fp32, contiguous - if grad_k is None: - grad_k_f32 = torch.zeros((B, H, D, N), dtype=torch.float32, device=q_raw.device) - else: - grad_k_f32 = grad_k.to(torch.float32) - if grad_inflation_sq is not None: - # k_out: (B, H, D, N), d_post_k_sq: (B, H, N) → broadcast over D dim. - eff_d_k_out = (grad_k_f32 + 2.0 * k_out.to(torch.float32) * d_post_k_sq.unsqueeze(2)).contiguous() - else: - eff_d_k_out = grad_k_f32.contiguous() - d_pre_k_sq = d_pre_k_sq.contiguous() - - # grad_q / grad_v as fp32 (B, H, D, N) contiguous (zero-fill if absent) - if grad_q is None: - grad_q_f32 = torch.zeros((B, H, D, N), dtype=torch.float32, device=q_raw.device) - else: - grad_q_f32 = grad_q.to(torch.float32).contiguous() - if grad_v is None: - grad_v_f32 = torch.zeros((B, H, D, N), dtype=torch.float32, device=q_raw.device) - else: - grad_v_f32 = grad_v.to(torch.float32).contiguous() - - # ---- allocate outputs / intermediates ---- - d_q_post_norm = torch.empty((B, N, H, D), dtype=torch.float32, device=q_raw.device) - d_k_post_norm = torch.empty((B, N, H, D), dtype=torch.float32, device=q_raw.device) - dv_raw_f32 = torch.empty((B, N, H, D), 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) - - _cam_prep_bwd_kernel[(B * N * H,)]( - q_raw, - k_raw, - q_norm_weight, - k_norm_weight, - q_inv_rms, - k_inv_rms, - proj_q, - proj_kv, - rope_cos, - rope_sin, - grad_q_f32, - eff_d_k_out, - grad_v_f32, - d_pre_k_sq, - d_q_post_norm, - d_k_post_norm, - dv_raw_f32, - 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, - ) - - # ---- torch RMSNorm bwd over the saved post-norm grads ---- - # Cast raw inputs to fp32 for the cross-head reduction (matches kernel - # numerics — the kernel uses fp32 internally as well). - q_raw_f32 = q_raw.to(torch.float32) - k_raw_f32 = k_raw.to(torch.float32) - q_inv_rms_view = q_inv_rms.view(B, N, 1, 1) - k_inv_rms_view = k_inv_rms.view(B, N, 1, 1) - q_nw_view = q_norm_weight.view(1, 1, H, D) - k_nw_view = k_norm_weight.view(1, 1, H, D) - - # S_q[b, n] = sum_{h, d} d_q_post_norm[b, n, h, d] * q_raw[b, n, h, d] * q_norm_w[h, d] - weighted_q = d_q_post_norm * q_raw_f32 # reused for dq_norm_weight reduction - weighted_k = d_k_post_norm * k_raw_f32 - S_q = (weighted_q * q_nw_view).sum(dim=(2, 3)) # (B, N) - S_k = (weighted_k * k_nw_view).sum(dim=(2, 3)) - - # dq_raw[b, n, h, d] = d_q_post_norm * inv_rms_q * q_norm_w - # - inv_rms_q^3 / C * q_raw * S_q - inv_q3 = (q_inv_rms**3).view(B, N, 1, 1) - inv_k3 = (k_inv_rms**3).view(B, N, 1, 1) - inv_C = 1.0 / float(C) - dq_raw_f32 = d_q_post_norm * q_inv_rms_view * q_nw_view - inv_q3 * inv_C * q_raw_f32 * S_q.view(B, N, 1, 1) - dk_raw_f32 = d_k_post_norm * k_inv_rms_view * k_nw_view - inv_k3 * inv_C * k_raw_f32 * S_k.view(B, N, 1, 1) - - # dq_norm_weight[h, d] = sum_{b, n} d_q_post_norm[b, n, h, d] * q_raw[b, n, h, d] * inv_rms_q[b, n] - dq_norm_weight = (weighted_q * q_inv_rms_view).sum(dim=(0, 1)).reshape(-1).contiguous() - dk_norm_weight = (weighted_k * k_inv_rms_view).sum(dim=(0, 1)).reshape(-1).contiguous() - - # Cast Q/K/V grads back to input dtype to match torch.autograd convention. - dq_raw = dq_raw_f32.to(q_raw.dtype) - dk_raw = dk_raw_f32.to(q_raw.dtype) - dv_raw = dv_raw_f32.to(q_raw.dtype) - - return dq_raw, dk_raw, dv_raw, dq_norm_weight, dk_norm_weight - - -def cam_scan_func( - 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, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Fused numerator-only single-path delta-rule scan for the cam branch. - - Args: - q, k, v: ``(B, H, D, N)`` fp32 contiguous tensors. - beta: ``(B, H, F, S)`` fp32 contiguous. - decay: ``(B, H, F)`` fp32 contiguous. - reverse: If ``True``, run the scan as the backward pass (equivalent - to ``flip_and_shift``-ing the inputs along the frame axis and - running forward). - init_state: optional ``(B*H, BLOCK_D, BLOCK_D)`` fp32 contiguous - tensor holding the forward-scan KV state at the END of a prefix - sequence (i.e., AFTER the prefix's last update, BEFORE any - further decay applied by this call). When provided, the kernel - resumes the scan from this state instead of zero. ``BLOCK_D = - next_pow2(D)`` and only the top-left ``D x D`` submatrix is - read. Forward direction only — raises ``NotImplementedError`` - if combined with ``reverse=True``. - save_final_state: when True, allocate a fresh fp32 zero buffer for - the final KV state (after the last frame's update) and pass it - to the kernel for write-out. Returned as the second tuple slot. - Forward direction only. - - Returns: - ``out`` of shape ``(B, H, D, N)`` fp32 matching - ``torch_chunk_cam_single_path_delta_rule`` with ``chunk_size >= T``. - - When ``save_final_state=True``, returns ``(out, final_state)`` where - ``final_state`` is fp32 ``(B*H, BLOCK_D, BLOCK_D)``. - - Raises: - NotImplementedError: if ``reverse=True`` is combined with state - passing. The cam branch's anti-causal scan resets per chunk in - the reference block, so there is no global cross-prefix state - to cache for the reverse direction. - """ - # Chunkwise integration (2026-05-06): dispatch all paths (fwd + reverse) - # to `cam_scan_chunkwise`. Reverse uses chunkwise's existing direction=2 - # mode in phase_b_triton, which has the same flip-and-shift semantics as - # cam's REVERSE=1 path. Bypass via FUSED_GDN_FORCE_LEGACY=1. - if os.environ.get("FUSED_GDN_FORCE_LEGACY", "0") != "1": - return cam_scan_chunkwise( - q, - k, - v, - beta, - decay, - reverse=reverse, - init_state=init_state, - save_final_state=save_final_state, - ) - - assert q.shape == k.shape == v.shape - B, H, D, N = q.shape - assert beta.shape[0] == B and beta.shape[1] == H - F_frames = beta.shape[2] - assert N % F_frames == 0 - S = N // F_frames - assert beta.shape == (B, H, F_frames, S), f"beta shape {beta.shape}" - assert decay.shape == (B, H, F_frames), f"decay shape {decay.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 - - BLOCK_D = triton.next_power_of_2(D) - BLOCK_S = _DEFAULT_BLOCK_S - num_warps = 4 - num_stages = 1 - - if reverse and (init_state is not None or save_final_state): - raise NotImplementedError( - "cam_scan_func: 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." - ) - - if init_state is not None: - expected_shape = (B * H, BLOCK_D, BLOCK_D) - if tuple(init_state.shape) != expected_shape: - raise ValueError( - f"cam_scan_func: init_state shape {tuple(init_state.shape)} " - f"does not match expected {expected_shape} (BLOCK_D=next_pow2(D)={BLOCK_D})." - ) - if init_state.dtype != torch.float32: - raise ValueError(f"cam_scan_func: init_state must be fp32 (got {init_state.dtype}).") - if not init_state.is_contiguous(): - raise ValueError("cam_scan_func: init_state must be contiguous.") - if init_state.device != q.device: - raise ValueError("cam_scan_func: init_state must be on the same device as q.") - load_init = 1 - else: - load_init = 0 - - if save_final_state: - final_state = torch.zeros(B * H, BLOCK_D, BLOCK_D, device=q.device, dtype=torch.float32) - save_final = 1 - else: - final_state = None - save_final = 0 - - out = torch.empty_like(q) - - dummy_state = torch.empty(1, device=q.device, dtype=torch.float32) - init_state_ptr = init_state if load_init else dummy_state - final_state_ptr = final_state if save_final else dummy_state - - _cam_scan_kernel[(B * H,)]( - q, - k, - v, - beta, - decay, - out, - dummy_state, - dummy_state, - init_state_ptr, - final_state_ptr, - H=H, - F=F_frames, - S=S, - D=D, - N=N, - REVERSE=1 if reverse else 0, - SAVE_STATES=0, - LOAD_INIT_STATE=load_init, - SAVE_FINAL_STATE=save_final, - BLOCK_D=BLOCK_D, - BLOCK_S=BLOCK_S, - num_warps=num_warps, - num_stages=num_stages, - ) - if save_final_state: - return out, final_state - return out - - -def _run_cam_scan_fwd_save( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - *, - reverse: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Run the forward scan with per-frame state snapshots saved. - - Used by :class:`CamScanFunction` to preserve the ``(state_pre, state_post)`` - snapshots that the Triton bwd kernel consumes. Snapshots are indexed by - ``q_frame`` (matching the existing fwd-kernel save logic), so the bwd - kernel can load them with the same ``q_frame`` derived in its - ``REVERSE``-aware iteration. - - Returns: - (out, state_pre, state_post). ``state_pre`` and ``state_post`` are - ``(B, H, F, BLOCK_D, BLOCK_D)`` fp32 with ``BLOCK_D = next_pow2(D)``. - Padding columns/rows past ``D`` are zero-masked on store. - """ - assert q.shape == k.shape == v.shape - B, H, D, N = q.shape - F_frames = beta.shape[2] - S = N // F_frames - assert beta.shape == (B, H, F_frames, S) - assert decay.shape == (B, H, F_frames) - 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 - - BLOCK_D = triton.next_power_of_2(D) - BLOCK_S = _DEFAULT_BLOCK_S - num_warps = 4 - num_stages = 1 - - out = torch.empty_like(q) - state_pre = torch.zeros(B * H, F_frames, BLOCK_D, BLOCK_D, device=q.device, dtype=torch.float32) - state_post = torch.zeros(B * H, F_frames, BLOCK_D, BLOCK_D, device=q.device, dtype=torch.float32) - dummy_state = torch.empty(1, device=q.device, dtype=torch.float32) - - _cam_scan_kernel[(B * H,)]( - q, - k, - v, - beta, - decay, - out, - state_pre, - state_post, - dummy_state, - dummy_state, - H=H, - F=F_frames, - S=S, - D=D, - N=N, - REVERSE=1 if reverse else 0, - SAVE_STATES=1, - LOAD_INIT_STATE=0, - SAVE_FINAL_STATE=0, - BLOCK_D=BLOCK_D, - BLOCK_S=BLOCK_S, - num_warps=num_warps, - num_stages=num_stages, - ) - return out, state_pre, state_post - - -def _cam_scan_bwd_dispatch( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - state_pre: torch.Tensor, - state_post: torch.Tensor, - grad_out: torch.Tensor, - *, - reverse: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Launch ``_cam_scan_bwd_kernel`` and return ``(dq, dk, dv, dbeta, ddecay)``. - - All gradient outputs are fp32 contiguous, matching the dtype of the - forward inputs. ``grad_out`` is cast to fp32 before launching. - - When ``reverse=True``, ``kv_frame=0`` is never visited (only used in the - skipped first iter), so ``dk[..., 0, :]``, ``dv[..., 0, :]``, - ``dbeta[..., 0, :]`` and ``ddecay[..., 0]`` must remain zero. We pre-zero - every output buffer here so the kernel only needs to write the live slots. - """ - assert q.shape == k.shape == v.shape - B, H, D, N = q.shape - F_frames = beta.shape[2] - S = N // F_frames - - grad_out_f32 = grad_out.to(torch.float32).contiguous() - dq = torch.zeros_like(q) - dk = torch.zeros_like(k) - dv = torch.zeros_like(v) - dbeta = torch.zeros_like(beta) - ddecay = torch.zeros_like(decay) - - BLOCK_D = triton.next_power_of_2(D) - # For small S, ``next_pow2(S) < _DEFAULT_BLOCK_S`` — using the smaller value - # avoids zero-padding huge unused tiles into shared memory. - BLOCK_S = min(_DEFAULT_BLOCK_S, max(triton.next_power_of_2(S), 16)) - num_stages = 1 - REVERSE = 1 if reverse else 0 - - last_err: Exception | None = None - for num_warps in (4, 2, 1): - try: - _cam_scan_bwd_kernel[(B * H,)]( - q, - k, - v, - beta, - decay, - state_pre, - state_post, - grad_out_f32, - dq, - dk, - dv, - dbeta, - ddecay, - H=H, - F=F_frames, - S=S, - D=D, - N=N, - REVERSE=REVERSE, - BLOCK_D=BLOCK_D, - BLOCK_S=BLOCK_S, - num_warps=num_warps, - num_stages=num_stages, - ) - return dq, dk, dv, dbeta, ddecay - except triton.runtime.errors.OutOfResources as exc: - last_err = exc - continue - raise RuntimeError("_cam_scan_bwd_kernel exhausted all num_warps choices: " + str(last_err)) - - -# ============================================================================= -# Section: Torch reference implementations used by fallback backward paths -# ============================================================================= -# These references replicate the Triton-kernel math (full-channel RMSNorm + -# ReLU + K-scale + 4x4 UCPE projmat + interleaved-pair real-valued RoPE, then -# numerator-only single-path delta-rule scan). They run in fp32 internally and -# cast outputs back to the input dtype, matching the kernels. - - -def _flip_and_shift(x: torch.Tensor, dim: int, shift_val: float) -> torch.Tensor: - """Flip ``x`` along ``dim`` and right-shift by one (pad with ``shift_val``). - - Matches the reference ``sana_gdn_blocks.flip_and_shift`` semantics. - """ - 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) - - -def _torch_cam_scan_single_chunk( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - init_state: torch.Tensor | None = None, - return_final_state: bool = False, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Pure-torch single-chunk delta-rule scan (numerator-only). - - Algebraically equivalent to ``torch_chunk_cam_single_path_delta_rule`` with - ``chunk_size >= T``; matches the Triton ``_cam_scan_kernel`` math exactly - (which also runs as a single chunk over all F frames). - - Args: - q, k, v: ``(B, H, D, N)`` fp32 contiguous. - beta: ``(B, H, F, S)`` or ``(B, H, F)`` fp32 contiguous. - decay: ``(B, H, F)`` fp32 contiguous. - - Returns: - out: ``(B, H, D, N)`` fp32 contiguous, ``N = F * S``. - """ - B, H, D, N = q.shape - if beta.ndim == 4: - T = beta.shape[2] - elif beta.ndim == 3: - T = beta.shape[2] - else: - raise ValueError(f"beta must be (B,H,F[,S]); got ndim={beta.ndim}") - if N % T != 0: - raise ValueError(f"N ({N}) must be divisible by 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) # (B, H, T, D, S) - - q_t = to_frame_seq(q) - k_t = to_frame_seq(k) - v_t = to_frame_seq(v) - - if beta.ndim == 4: - beta_view = beta.unsqueeze(3) # (B, H, T, 1, S) - else: - beta_view = beta.view(B, H, T, 1, 1) - decay_view = decay.view(B, H, T, 1, 1) - - eye = torch.eye(D, device=q.device, dtype=q.dtype).view(1, 1, 1, D, D) - - k_beta = k_t * beta_view - W = decay_view * (eye - torch.matmul(k_beta, k_t.transpose(-1, -2))) - U = torch.matmul(v_t * beta_view, k_t.transpose(-1, -2)) - - state = ( - torch.zeros(B, H, D, D, device=q.device, dtype=q.dtype) - if init_state is None - else init_state.to(device=q.device, dtype=q.dtype) - ) - s_kv_list: list[torch.Tensor] = [] - for t in range(T): - state = torch.matmul(state, W[:, :, t]) + U[:, :, t] - s_kv_list.append(state) - s_all = torch.stack(s_kv_list, dim=2) # (B, H, T, D, D) - - out_t = torch.matmul(s_all, q_t) # (B, H, T, D, S) - out = out_t.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) - return (out, state) if return_final_state else out - - -def _torch_cam_scan_reference( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - *, - reverse: bool = False, -) -> torch.Tensor: - """Pure-torch reference for ``cam_scan_func`` supporting ``reverse=True``. - - For ``reverse=False`` this is the standard forward delta-rule scan. - - For ``reverse=True`` we emulate the Triton kernel's per-chunk - ``flip_and_shift`` semantics (q is flipped only; k/v/beta are - flip-and-shifted with pad value 0; decay is flip-and-shifted with pad - value 1; output is then flipped back along the time axis). - """ - if not reverse: - return _torch_cam_scan_single_chunk(q, k, v, beta, decay) - - B, H, D, N = q.shape - if beta.ndim == 4: - T = beta.shape[2] - elif beta.ndim == 3: - T = beta.shape[2] - else: - raise ValueError(f"beta must be (B,H,F[,S]); got ndim={beta.ndim}") - S = N // T - - def to_frame(x: torch.Tensor) -> torch.Tensor: - return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) # (B, H, T, D, S) - - def from_frame(x: torch.Tensor) -> torch.Tensor: - return x.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) - - q_bwd = torch.flip(to_frame(q), dims=[2]) - k_bwd = _flip_and_shift(to_frame(k), dim=2, shift_val=0.0) - v_bwd = _flip_and_shift(to_frame(v), 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 = _torch_cam_scan_single_chunk( - from_frame(q_bwd), - from_frame(k_bwd), - from_frame(v_bwd), - beta_bwd, - decay_bwd, - ) - out_bwd_t = out_bwd.view(B, H, D, T, S) # already in (B, H, D, T, S) - return torch.flip(out_bwd_t, dims=[3]).reshape(B, H, D, N) - - -def _torch_cam_prep_reference( - 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, - proj_kv: torch.Tensor, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - k_scale: float, - norm_eps: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Pure-torch reference for ``cam_prep_func`` matching ``_cam_prep_kernel``. - - Replicates exactly: - - Full-channel (over ``H*D``) RMSNorm + per-channel weight on Q, K. - - ReLU on Q, K. - - K-scale on K. - - 4x4 UCPE projmat on first ``D/2`` dims (Q via ``proj_q``, - K and V via ``proj_kv``). - - Interleaved-pair real-valued RoPE on second ``D/2`` dims using - ``rope_cos`` / ``rope_sin`` (the same tables passed to the kernel). - - ``inflation_sq = ||k_post_ucpe||^2 / ||k_pre_ucpe||^2`` per (B, H, N), - with the same ``clamp_min(1e-12)`` floor as ``cam_prep_func``. - - All math runs in fp32 internally; outputs ``(q, k, v)`` are cast back to - ``q_raw.dtype`` and ``inflation_sq`` is fp32. - """ - B, N, H, D = q_raw.shape - if D % 2 != 0: - raise ValueError(f"D ({D}) must be even.") - if (D // 2) % 4 != 0: - raise ValueError(f"D/2 ({D // 2}) must be divisible by 4 (UCPE projmat).") - C = H * D - D_half = D // 2 - n_groups = D_half // 4 - - q32 = q_raw.float() - k32 = k_raw.float() - v32 = v_raw.float() - - # ---- Full-channel RMSNorm + per-channel weight (Q, K only) ---- - q_inv_rms = torch.rsqrt((q32 * q32).sum(dim=(-1, -2)) / C + norm_eps) # (B, N) - k_inv_rms = torch.rsqrt((k32 * k32).sum(dim=(-1, -2)) / C + norm_eps) - q_nw = q_norm_weight.float().view(1, 1, H, D) - k_nw = k_norm_weight.float().view(1, 1, H, D) - q_normed = q32 * q_inv_rms.view(B, N, 1, 1) * q_nw - k_normed = k32 * k_inv_rms.view(B, N, 1, 1) * k_nw - - # ---- ReLU + K-scale ---- - q_normed = torch.relu(q_normed) - k_normed = torch.relu(k_normed) * k_scale - - # ---- Pre-UCPE ||k||^2 over the full D dim ---- - pre_k_sq_BNH = (k_normed * k_normed).sum(dim=-1) # (B, N, H) - - # ---- UCPE 4x4 projmat on first half ---- - q_first = q_normed[..., :D_half].reshape(B, N, H, n_groups, 4) - k_first = k_normed[..., :D_half].reshape(B, N, H, n_groups, 4) - v_first = v32[..., :D_half].reshape(B, N, H, n_groups, 4) - - # out[b,n,h,g,i] = sum_j P[b,n,i,j] * x[b,n,h,g,j] - # einsum: 'bnij,bnhgj->bnhgi' - proj_q_f = proj_q.float() - proj_kv_f = proj_kv.float() - q_first_proj = torch.einsum("bnij,bnhgj->bnhgi", proj_q_f, q_first).reshape(B, N, H, D_half) - k_first_proj = torch.einsum("bnij,bnhgj->bnhgi", proj_kv_f, k_first).reshape(B, N, H, D_half) - v_first_proj = torch.einsum("bnij,bnhgj->bnhgi", proj_kv_f, v_first).reshape(B, N, H, D_half) - - # ---- Interleaved-pair real-valued RoPE on second half ---- - # Kernel form: y[d] = x[d]*rope_cos[d] + x[d^1]*rope_sin[d] - # where rope_cos/rope_sin come from _prepare_ucpe_rope_tables. - q_second = q_normed[..., D_half:] - k_second = k_normed[..., D_half:] - v_second = v32[..., D_half:] - - def _pair_swap(x: torch.Tensor) -> torch.Tensor: - # Swap consecutive pairs along the last dim: (..., D_half) where D_half is even. - # x[..., 2i] <-> x[..., 2i+1]. - x_pairs = x.unflatten(-1, (D_half // 2, 2)) - x_swapped = x_pairs.flip(-1) - return x_swapped.flatten(-2) - - cos_b = rope_cos.float().view(1, N, 1, D_half) - sin_b = rope_sin.float().view(1, N, 1, D_half) - q_rope = q_second * cos_b + _pair_swap(q_second) * sin_b - k_rope = k_second * cos_b + _pair_swap(k_second) * sin_b - v_rope = v_second * cos_b + _pair_swap(v_second) * sin_b - - # ---- Reassemble (B, N, H, D) and post-UCPE k norm ---- - q_out_BNHD = torch.cat([q_first_proj, q_rope], dim=-1) - k_out_BNHD = torch.cat([k_first_proj, k_rope], dim=-1) - v_out_BNHD = torch.cat([v_first_proj, v_rope], dim=-1) - - post_k_sq_BNH = (k_out_BNHD * k_out_BNHD).sum(dim=-1) # (B, N, H) - - out_dtype = q_raw.dtype - q_out = q_out_BNHD.to(out_dtype).permute(0, 2, 3, 1).contiguous() - k_out = k_out_BNHD.to(out_dtype).permute(0, 2, 3, 1).contiguous() - v_out = v_out_BNHD.to(out_dtype).permute(0, 2, 3, 1).contiguous() - - pre_k_sq = pre_k_sq_BNH.permute(0, 2, 1).contiguous() # (B, H, N) - post_k_sq = post_k_sq_BNH.permute(0, 2, 1).contiguous() - inflation_sq = post_k_sq.clamp_min(1e-12) / pre_k_sq.clamp_min(1e-12) - return q_out, k_out, v_out, inflation_sq - - -# ============================================================================= -# Section: Autograd-enabled wrappers -# ============================================================================= - - -def _cam_scan_torch_fallback_backward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - needs: tuple[bool, bool, bool, bool, bool], - grad_out: torch.Tensor, - reverse: bool, -) -> list[torch.Tensor | None]: - """Recompute the cam-branch scan via the torch reference and return grads. - - Used when ``CAM_SCAN_BWD_FALLBACK=1`` forces the torch-recompute backward - path. - """ - detached = [] - for tensor, need in zip((q, k, v, beta, decay), needs): - t = tensor.detach() - if need: - t = t.requires_grad_(True) - detached.append(t) - active = [t for t in detached if t.requires_grad] - - with torch.enable_grad(): - q_d, k_d, v_d, beta_d, decay_d = detached - ref_out = _torch_cam_scan_reference(q_d, k_d, v_d, beta_d, decay_d, reverse=reverse) - if active: - active_grads = torch.autograd.grad( - outputs=ref_out, - inputs=tuple(active), - grad_outputs=grad_out.to(ref_out.dtype), - allow_unused=True, - ) - else: - active_grads = [] - - grads: list[torch.Tensor | None] = [] - active_iter = iter(active_grads) - for tensor in detached: - grads.append(next(active_iter) if tensor.requires_grad else None) - return grads - - -class CamScanFunction(torch.autograd.Function): - """Autograd ``Function`` wrapping ``cam_scan_func``. - - Forward calls the Triton ``_cam_scan_kernel`` with ``SAVE_STATES=1`` so - per-frame state snapshots (``state_pre[q_frame]``, ``state_post[q_frame]``) - are kept for the backward pass. Backward runs the true Triton bwd - kernel (``_cam_scan_bwd_kernel``) for both ``reverse=False`` and - ``reverse=True``, replaying the recurrence in reverse time using the - saved snapshots. - - Set ``CAM_SCAN_BWD_FALLBACK=1`` to force the torch-recompute backward - validation path. - """ - - @staticmethod - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - reverse: bool, - ) -> torch.Tensor: - ctx.set_materialize_grads(False) - ctx.reverse = bool(reverse) - - force_torch_fallback = os.environ.get("CAM_SCAN_BWD_FALLBACK", "0") == "1" - ctx.use_triton_bwd = not force_torch_fallback - - if ctx.use_triton_bwd: - out, state_pre, state_post = _run_cam_scan_fwd_save(q, k, v, beta, decay, reverse=ctx.reverse) - ctx.save_for_backward(q, k, v, beta, decay, state_pre, state_post) - return out - - # Torch-fallback backward path: don't bother saving state snapshots. - ctx.save_for_backward(q, k, v, beta, decay) - return cam_scan_func(q, k, v, beta, decay, reverse=reverse) - - @staticmethod - def backward(ctx, grad_out): # type: ignore[override] - if grad_out is None: - return (None, None, None, None, None, None) - - if ctx.use_triton_bwd: - q, k, v, beta, decay, state_pre, state_post = ctx.saved_tensors - needs = ctx.needs_input_grad[:5] # q, k, v, beta, decay - dq, dk, dv, dbeta, ddecay = _cam_scan_bwd_dispatch( - q, - k, - v, - beta, - decay, - state_pre, - state_post, - grad_out, - reverse=ctx.reverse, - ) - grads: list[torch.Tensor | None] = [ - dq if needs[0] else None, - dk if needs[1] else None, - dv if needs[2] else None, - dbeta if needs[3] else None, - ddecay if needs[4] else None, - ] - return (*grads, None) - - # Env-var torch-recompute backward. - q, k, v, beta, decay = ctx.saved_tensors - needs = ctx.needs_input_grad[:5] - grads = _cam_scan_torch_fallback_backward(q, k, v, beta, decay, needs, grad_out, ctx.reverse) - return (*grads, None) - - -def cam_scan_func_with_grad( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - *, - reverse: bool = False, -) -> torch.Tensor: - """Autograd-enabled wrapper around :func:`cam_scan_func`. - - Forward is identical to :func:`cam_scan_func`; backward is computed via - a torch reference (``_torch_cam_scan_reference``). Use this in training - paths where any of ``q, k, v, beta, decay`` may require gradients. - - Inference paths can keep calling :func:`cam_scan_func` directly to avoid - the small autograd bookkeeping overhead. - """ - return CamScanFunction.apply(q, k, v, beta, decay, reverse) - - -def _cam_prep_torch_fallback_backward( - 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, - proj_kv: torch.Tensor, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - needs: tuple[bool, ...], - grad_q: torch.Tensor | None, - grad_k: torch.Tensor | None, - grad_v: torch.Tensor | None, - grad_inflation_sq: torch.Tensor | None, - k_scale: float, - norm_eps: float, -) -> list[torch.Tensor | None]: - """Recompute the cam-branch prep via the torch reference and return grads. - - Used when any of ``proj_q / proj_kv / rope_cos / rope_sin`` requests a - gradient (the Triton kernel does not produce those grads), or when - ``CAM_PREP_BWD_FALLBACK=1`` forces the torch-recompute backward path. - - Args: - q_raw, k_raw, ..., rope_sin: the nine tensor inputs of - :func:`cam_prep_func` (in the same order as - :class:`CamPrepFunction.forward`'s arg list). - needs: ``ctx.needs_input_grad[:9]`` — boolean per-input flags. - grad_q, grad_k, grad_v, grad_inflation_sq: upstream gradients. - k_scale, norm_eps: scalar fwd args. - - Returns: - A 9-element list of ``torch.Tensor | None`` aligned with the - ``saved`` tuple. Entries that didn't request a gradient are ``None``. - """ - saved = ( - q_raw, - k_raw, - v_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - ) - detached: list[torch.Tensor] = [] - for tensor, need in zip(saved, needs): - t = tensor.detach() - if need: - t = t.requires_grad_(True) - detached.append(t) - active = [t for t in detached if t.requires_grad] - - with torch.enable_grad(): - (q_d, k_d, v_d, qnw_d, knw_d, pq_d, pkv_d, rc_d, rs_d) = detached - ref_q, ref_k, ref_v, ref_inf = _torch_cam_prep_reference( - q_d, - k_d, - v_d, - q_norm_weight=qnw_d, - k_norm_weight=knw_d, - proj_q=pq_d, - proj_kv=pkv_d, - rope_cos=rc_d, - rope_sin=rs_d, - k_scale=k_scale, - norm_eps=norm_eps, - ) - - outputs = [] - grad_outputs = [] - if grad_q is not None: - outputs.append(ref_q) - grad_outputs.append(grad_q.to(ref_q.dtype)) - if grad_k is not None: - outputs.append(ref_k) - grad_outputs.append(grad_k.to(ref_k.dtype)) - if grad_v is not None: - outputs.append(ref_v) - grad_outputs.append(grad_v.to(ref_v.dtype)) - if grad_inflation_sq is not None: - outputs.append(ref_inf) - grad_outputs.append(grad_inflation_sq.to(ref_inf.dtype)) - - if active and outputs: - active_grads = torch.autograd.grad( - outputs=tuple(outputs), - inputs=tuple(active), - grad_outputs=tuple(grad_outputs), - allow_unused=True, - ) - else: - active_grads = [] - - grads: list[torch.Tensor | None] = [] - active_iter = iter(active_grads) - for tensor in detached: - grads.append(next(active_iter) if tensor.requires_grad else None) - return grads - - -class CamPrepFunction(torch.autograd.Function): - """Autograd ``Function`` wrapping ``cam_prep_func``. - - Forward calls the fused Triton ``_cam_prep_kernel`` via - :func:`_run_cam_prep_fwd_save` so the per-token ``inv_rms`` / - ``k_pre_sq`` / ``k_post_sq`` snapshots required by the bwd kernel are - preserved alongside the standard outputs. - - Backward runs the true Triton bwd kernel via - :func:`_cam_prep_bwd_dispatch` for the standard training path - (``q_raw``, ``k_raw``, ``v_raw``, ``q_norm_weight``, ``k_norm_weight`` - only request grads). The Triton path implements: - - * RoPE^T, UCPE^T, K-scale^T, and ReLU mask in a single fused kernel - (one program per ``(b, n, h)``); - * ``grad_inflation_sq`` chain through ``k_post_sq`` (added into - ``eff_dO_k``) and ``k_pre_sq`` (direct contribution to - ``d_k_post_kscale``), with ``clamp_min(1e-12)`` indicators honored; - * The full-channel RMSNorm bwd (per-token cross-head reduction) is - done in PyTorch on the kernel's ``d_q_post_norm`` / - ``d_k_post_norm`` intermediates — see - :func:`_cam_prep_bwd_dispatch` for details. - - The torch-recompute fallback (running :func:`_torch_cam_prep_reference` - under autograd) is selected when any of ``proj_q / proj_kv / rope_cos / - rope_sin`` requests a gradient (the Triton path emits ``None`` for those - slots) or when ``CAM_PREP_BWD_FALLBACK=1`` is set. - """ - - @staticmethod - def forward( - ctx, - 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, - proj_kv: torch.Tensor, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - k_scale: float, - norm_eps: float, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - ctx.set_materialize_grads(False) - ctx.k_scale = float(k_scale) - ctx.norm_eps = float(norm_eps) - - force_torch_fallback = os.environ.get("CAM_PREP_BWD_FALLBACK", "0") == "1" - ctx.use_triton_bwd = not force_torch_fallback - - if ctx.use_triton_bwd: - ( - q_out, - k_out, - v_out, - inflation_sq, - q_inv_rms, - k_inv_rms, - k_pre_sq, - k_post_sq, - ) = _run_cam_prep_fwd_save( - q_raw, - k_raw, - v_raw, - q_norm_weight=q_norm_weight, - k_norm_weight=k_norm_weight, - proj_q=proj_q, - proj_kv=proj_kv, - rope_cos=rope_cos, - rope_sin=rope_sin, - k_scale=k_scale, - norm_eps=norm_eps, - ) - ctx.save_for_backward( - q_raw, - k_raw, - v_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - q_inv_rms, - k_inv_rms, - k_pre_sq, - k_post_sq, - k_out, - ) - else: - q_out, k_out, v_out, inflation_sq = cam_prep_func( - q_raw, - k_raw, - v_raw, - q_norm_weight=q_norm_weight, - k_norm_weight=k_norm_weight, - proj_q=proj_q, - proj_kv=proj_kv, - rope_cos=rope_cos, - rope_sin=rope_sin, - k_scale=k_scale, - norm_eps=norm_eps, - ) - ctx.save_for_backward( - q_raw, - k_raw, - v_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - ) - return q_out, k_out, v_out, inflation_sq - - @staticmethod - def backward(ctx, grad_q, grad_k, grad_v, grad_inflation_sq): # type: ignore[override] - if grad_q is None and grad_k is None and grad_v is None and grad_inflation_sq is None: - return tuple([None] * 11) - - needs = ctx.needs_input_grad[:9] # nine tensor inputs - # If anyone outside (q_raw, k_raw, v_raw, q_norm_weight, k_norm_weight) - # requests a grad, the Triton bwd cannot handle it — fall back to the - # torch reference. - triton_bwd_supported_needs = needs[:5] - proj_or_rope_needs_grad = any(needs[5:]) - - if ctx.use_triton_bwd and not proj_or_rope_needs_grad: - ( - q_raw, - k_raw, - v_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - q_inv_rms, - k_inv_rms, - k_pre_sq, - k_post_sq, - k_out, - ) = ctx.saved_tensors - - ( - dq_raw, - dk_raw, - dv_raw, - dq_norm_weight, - dk_norm_weight, - ) = _cam_prep_bwd_dispatch( - q_raw, - k_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - q_inv_rms, - k_inv_rms, - k_pre_sq, - k_post_sq, - k_out, - grad_q=grad_q, - grad_k=grad_k, - grad_v=grad_v, - grad_inflation_sq=grad_inflation_sq, - k_scale=ctx.k_scale, - ) - grads: list[torch.Tensor | None] = [ - dq_raw if triton_bwd_supported_needs[0] else None, - dk_raw if triton_bwd_supported_needs[1] else None, - dv_raw if triton_bwd_supported_needs[2] else None, - dq_norm_weight if triton_bwd_supported_needs[3] else None, - dk_norm_weight if triton_bwd_supported_needs[4] else None, - None, # proj_q - None, # proj_kv - None, # rope_cos - None, # rope_sin - ] - return (*grads, None, None) - - # Torch fallback path. ``ctx.saved_tensors`` holds either 9 (legacy - # forward) or 14 (Triton fwd save) tensors — slice the leading nine. - saved = ctx.saved_tensors[:9] - ( - q_raw, - k_raw, - v_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - ) = saved - grads = _cam_prep_torch_fallback_backward( - q_raw, - k_raw, - v_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - needs=needs, - grad_q=grad_q, - grad_k=grad_k, - grad_v=grad_v, - grad_inflation_sq=grad_inflation_sq, - k_scale=ctx.k_scale, - norm_eps=ctx.norm_eps, - ) - # Two trailing None for non-tensor scalars (k_scale, norm_eps). - return (*grads, None, None) - - -def cam_prep_func_with_grad( - 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, - proj_kv: torch.Tensor, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - k_scale: float, - norm_eps: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Autograd-enabled wrapper around :func:`cam_prep_func`. - - Forward is identical to :func:`cam_prep_func`; backward is computed via a - torch reference (``_torch_cam_prep_reference``). Use this in training - paths so gradients flow back through Q/K/V projection inputs and through - the RMSNorm weights. - - Inference paths can keep calling :func:`cam_prep_func` directly. - """ - return CamPrepFunction.apply( - q_raw, - k_raw, - v_raw, - q_norm_weight, - k_norm_weight, - proj_q, - proj_kv, - rope_cos, - rope_sin, - k_scale, - norm_eps, - ) - - -__all__ = [ - "CamPrepFunction", - "CamScanFunction", - "_cam_prep_bwd_dispatch", - "_cam_prep_bwd_kernel", - "_cam_prep_kernel", - "_cam_prep_torch_fallback_backward", - "_cam_scan_bwd_dispatch", - "_cam_scan_bwd_kernel", - "_cam_scan_kernel", - "_invert_SE3", - "_precompute_cam_inv_rms", - "_prepare_ucpe_rope_tables", - "_process_camera_conditions_raymats_only", - "_run_cam_prep_fwd_save", - "_run_cam_scan_fwd_save", - "_torch_cam_prep_reference", - "cam_prep_func", - "cam_prep_func_with_grad", - "cam_scan_func", - "cam_scan_func_with_grad", -] diff --git a/integrations/sana/sana_wm/ops/fused_gdn.py b/integrations/sana/sana_wm/ops/fused_gdn.py deleted file mode 100644 index 447e0f5ca..000000000 --- a/integrations/sana/sana_wm/ops/fused_gdn.py +++ /dev/null @@ -1,2249 +0,0 @@ -"""Fused-BiGDN Triton kernels used by SANA-WM GDN attention blocks. - -Includes the unified forward kernel, backward kernels, RoPE/RMS helpers, and -autograd wrappers used by the Triton GDN attention blocks. - -Precision knob: env var ``FUSED_GDN_PRECISION`` or ``PRECISION_OVERRIDE``: - 0=IEEE fp32 dots, 1=TF32, 2=bf16 TC + fp32 state [default], 3=bf16 TC + bf16 state. -""" - -# ruff: noqa: E501 - -from __future__ import annotations - -import os - -import torch -import triton -import triton.language as tl - -# ===================================================================== -# 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_launch(D: int, beta: torch.Tensor, decay: torch.Tensor) -> tuple: - """Shared launcher preamble. - - Returns (BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta_c, decay_c). - ``beta_c`` / ``decay_c`` are the contiguous copies the kernel needs. - """ - BLOCK_D = triton.next_power_of_2(D) - cfg = _kcfg() - BLOCK_S = cfg["BLOCK_S"] - _, dot_prec, state_fp32, nw = _resolve_launch_config() - return BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta.contiguous(), decay.contiguous() - - -# ===================================================================== -# Unified forward Triton Mega-Kernel (inference-only variant) -# ===================================================================== -# Fuses: RMSNorm + ReLU + k_scale + RoPE + BiGDN recurrence. -# -# Inputs: -# qkv (B, N, 3, H, D) interleaved — strides passed explicitly. -# beta (B, H, F, S), decay (B, H, F) contiguous. -# q_norm_w, k_norm_w (H*D,) full-channel — only read when QK_NORM=1. -# rope_cos, rope_sin (N, D) contiguous. -# q_inv_rms, k_inv_rms (B, N) full-channel — only read when USE_PRECOMPUTED_RMS=1. -# -# Outputs: -# out (B, N, H, D) = num / (den + eps) — unused by BiGDN wrappers. -# num (B, N, H, D) = numerator before divide (summed across directions). -# den (B, H, N) = denominator before divide (summed across directions). -# -# NOTE (inference-only build): upstream also supports SAVE_STATE, -# LOAD_INIT_STATE, SAVE_FINAL_STATE for training backward / state caching. -# Those constexpr branches are preserved in the kernel so the source stays -# 1-for-1 with upstream (they compile away when launched with flags=0). - - -@triton.jit -def _fused_gdn_kernel( - # ---- interleaved QKV : (B, N, 3, H, D) ---- - qkv_ptr, - stride_b: tl.constexpr, - stride_n: tl.constexpr, - stride_3: tl.constexpr, - stride_h: tl.constexpr, - stride_d: tl.constexpr, - # ---- gates ---- - beta_ptr, - decay_ptr, - # ---- inv-RMS (B, N) — only read when USE_PRECOMPUTED_RMS=1 ---- - q_inv_rms_ptr, - k_inv_rms_ptr, - # ---- norm weights (H*D,) full-channel — only read when QK_NORM=1 ---- - q_norm_w_ptr, - k_norm_w_ptr, - # ---- RoPE tables (N, D) contiguous ---- - rope_cos_ptr, - rope_sin_ptr, - # ---- outputs ---- - out_ptr, # (B, N, H, D) - num_ptr, # (B, N, H, D) - den_ptr, # (B, H, N) - # ---- saved-state dummies (unused in this build but kept for signature parity) ---- - saved_state_ptr, - saved_z_ptr, - saved_state_curr_ptr, - saved_z_curr_ptr, - init_state_kv_ptr, - init_state_z_ptr, - final_state_kv_ptr, - final_state_z_ptr, - # ---- scalars / dims ---- - H: tl.constexpr, - F: tl.constexpr, - S: tl.constexpr, - D: tl.constexpr, - K_SCALE, - NORM_EPS: tl.constexpr, - EPS: tl.constexpr, - QK_NORM: tl.constexpr, - USE_PRECOMPUTED_RMS: tl.constexpr, - STATE_FP32: tl.constexpr, - DOT_PRECISION: tl.constexpr, - REVERSE: tl.constexpr, - SAVE_STATE: tl.constexpr, - LOAD_INIT_STATE: tl.constexpr, - SAVE_FINAL_STATE: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_S: tl.constexpr, -): - # ---- dot product precision / operand dtype ---- - if DOT_PRECISION >= 1: - dot_dtype = tl.float32 - else: - dot_dtype = tl.bfloat16 - dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" - - # ---- program → (batch, head) ---- - pid = tl.program_id(0) - pid_b = pid // H - pid_h = pid % H - N = F * S - bh = pid_b * H + pid_h - - # ---- base pointers ---- - qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h - out_bh = out_ptr + pid_b * (N * H * D) + pid_h * D - num_bh = num_ptr + pid_b * (N * H * D) + pid_h * D - den_bh = den_ptr + bh * N - beta_bh = beta_ptr + bh * (F * S) - decay_bh = decay_ptr + bh * F - if SAVE_STATE: - st_bh = saved_state_ptr + bh * F * BLOCK_D * BLOCK_D - sz_bh = saved_z_ptr + bh * F * BLOCK_D - stc_bh = saved_state_curr_ptr + bh * F * BLOCK_D * BLOCK_D - szc_bh = saved_z_curr_ptr + bh * F * BLOCK_D - - # ---- D-index helpers ---- - 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 - D_inv = 1.0 / D - - # ---- full-channel norm weights (only when QK_NORM=1) ---- - nw_offset = pid_h * D - if QK_NORM: - q_nw = tl.load(q_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) - k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) - q_nw_pair = tl.load(q_norm_w_ptr + nw_offset + offs_d_pair, mask=mask_d_pair, 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) - - k_scale = K_SCALE - offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] - mask_dd = mask_d[:, None] & mask_d[None, :] - - # ---- double-buffer state ---- - if LOAD_INIT_STATE: - init_kv_bh = init_state_kv_ptr + bh * BLOCK_D * BLOCK_D - state_curr = tl.load(init_kv_bh + offs_dd, mask=mask_dd, other=0.0).to(tl.float32) - init_z_bh = init_state_z_ptr + bh * BLOCK_D - state_z_curr = tl.load(init_z_bh + offs_d, mask=mask_d, other=0.0).to(tl.float32) - else: - state_curr = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) - state_z_curr = tl.zeros([BLOCK_D], dtype=tl.float32) - if STATE_FP32: - state_prev = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) - else: - state_prev = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.bfloat16) - state_z_prev = tl.zeros([BLOCK_D], dtype=tl.float32) - - # ======================================================== - # Temporal loop — serial over F - # ======================================================== - for f_iter in range(F): - if REVERSE: - q_frame = F - 1 - f_iter - kv_frame = F - f_iter if f_iter > 0 else 0 # unused at f=0 - skip_update = f_iter == 0 - else: - q_frame = f_iter - kv_frame = f_iter - skip_update = False - - # ---- decay + state snapshot ---- - if REVERSE and f_iter == 0: - g = 1.0 - else: - g = tl.load(decay_bh + kv_frame).to(tl.float32) - state_curr = state_curr * g - state_z_curr = state_z_curr * g - if STATE_FP32: - state_prev = state_curr + 0.0 - else: - state_prev = state_curr.to(tl.bfloat16) - state_z_prev = state_z_curr - - if SAVE_STATE: - st_f = st_bh + q_frame * BLOCK_D * BLOCK_D - tl.store(st_f + offs_dd, state_prev, mask=mask_dd) - tl.store(sz_bh + q_frame * BLOCK_D + offs_d, state_z_prev, mask=mask_d) - - # ------------------------------------------ - # Pass 1 — State Accumulation - # ------------------------------------------ - if skip_update == False: - kv_n_base = kv_frame * S - f_beta = beta_bh + kv_frame * 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, :] - mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] - n_idx = kv_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) - - if QK_NORM: - if USE_PRECOMPUTED_RMS: - k_inv_rms = tl.load(k_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0).to(tl.float32) - else: - k_var = tl.sum(K_raw * K_raw, axis=1) * D_inv - k_inv_rms = 1.0 / tl.sqrt(k_var + NORM_EPS) - K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] - else: - K_normed = K_raw - K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale - - k_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d_pair[None, :] * stride_d - K_pair_raw = tl.load(k_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) - if QK_NORM: - K_pair_normed = K_pair_raw * k_inv_rms[:, None] * k_nw_pair[None, :] - else: - K_pair_normed = K_pair_raw - 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 - - bt = tl.load(f_beta + offs_s, mask=mask_s, other=0.0).to(tl.float32) - - K_rot_dc = K_rot.to(dot_dtype) - V_pred = tl.dot(K_rot_dc, state_prev.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) - dv = (V_raw - V_pred) * bt[:, None] - state_curr += tl.dot(tl.trans(K_rot), dv, out_dtype=tl.float32, input_precision="tf32") - - z_hat = tl.sum(K * state_z_prev[None, :], axis=1) - dz = (1.0 - z_hat) * bt - state_z_curr += tl.sum(K * dz[:, None], axis=0) - - if SAVE_STATE: - stc_f = stc_bh + q_frame * BLOCK_D * BLOCK_D - tl.store(stc_f + offs_dd, state_curr, mask=mask_dd) - tl.store(szc_bh + q_frame * BLOCK_D + offs_d, state_z_curr, mask=mask_d) - - # ------------------------------------------ - # Pass 2 — Output (reads state_curr, inclusive) - # ------------------------------------------ - state_out = state_curr.to(dot_dtype) - state_z_out = state_z_curr - q_n_base = q_frame * 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, :] - mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] - n_idx = q_n_base + offs_s - - q_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d[None, :] * stride_d - q_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d_pair[None, :] * stride_d - Q_raw = tl.load(q_ptrs, mask=mask_sd, other=0.0).to(tl.float32) - Q_pair_raw = tl.load(q_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) - - if QK_NORM: - if USE_PRECOMPUTED_RMS: - q_inv_rms = tl.load(q_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0).to(tl.float32) - else: - q_var = tl.sum(Q_raw * Q_raw, axis=1) * D_inv - q_inv_rms = 1.0 / tl.sqrt(q_var + NORM_EPS) - Q_normed = Q_raw * q_inv_rms[:, None] * q_nw[None, :] - Q_pair_normed = Q_pair_raw * q_inv_rms[:, None] * q_nw_pair[None, :] - else: - Q_normed = Q_raw - Q_pair_normed = Q_pair_raw - Q = tl.where(Q_normed > 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), state_out, out_dtype=tl.float32, input_precision=dot_ip) - den = tl.sum(Q * state_z_out[None, :], axis=1) - - result = num / (den[:, None] + EPS) - out_ptrs = out_bh + n_idx[:, None] * (H * D) + offs_d[None, :] - num_ptrs = num_bh + n_idx[:, None] * (H * D) + offs_d[None, :] - tl.store(out_ptrs, result.to(tl.bfloat16), mask=mask_sd) - tl.store(num_ptrs, num.to(tl.bfloat16), mask=mask_sd) - tl.store(den_bh + n_idx, den.to(tl.bfloat16), mask=mask_s) - - if SAVE_FINAL_STATE: - final_kv_bh = final_state_kv_ptr + bh * BLOCK_D * BLOCK_D - tl.store(final_kv_bh + offs_dd, state_curr, mask=mask_dd) - final_z_bh = final_state_z_ptr + bh * BLOCK_D - tl.store(final_z_bh + offs_d, state_z_curr, mask=mask_d) - - -# ===================================================================== -# Python wrappers -# ===================================================================== - - -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. - """ - 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 - - -@triton.jit -def _fused_bidi_merge_kernel( - num_fwd_ptr, - num_bwd_ptr, - den_fwd_ptr, - den_bwd_ptr, - gate_ptr, - out_ptr, - B, - N, - H, - D, - eps, - snum_b, - snum_n, - snum_h, - snum_d, - sden_b, - sden_h, - sden_n, - APPLY_GATE: tl.constexpr, - PRE_SUMMED: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_D: tl.constexpr, -): - pid_bh = tl.program_id(0) - pid_n = tl.program_id(1) - b = pid_bh // H - h = pid_bh % H - - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_D) - mask_n = offs_n < N - mask_d = offs_d < D - mask_nd = mask_n[:, None] & mask_d[None, :] - - num_base = b * snum_b + offs_n[:, None] * snum_n + h * snum_h + offs_d[None, :] * snum_d - nf = tl.load(num_fwd_ptr + num_base, mask=mask_nd, other=0.0).to(tl.float32) - den_base = b * sden_b + h * sden_h + offs_n * sden_n - df = tl.load(den_fwd_ptr + den_base, mask=mask_n, other=0.0).to(tl.float32) - - if PRE_SUMMED: - num_total = nf - den_total = df + eps - else: - nb = tl.load(num_bwd_ptr + num_base, mask=mask_nd, other=0.0).to(tl.float32) - db = tl.load(den_bwd_ptr + den_base, mask=mask_n, other=0.0).to(tl.float32) - num_total = nf + nb - den_total = df + db + eps - out_val = num_total / den_total[:, None] - - if APPLY_GATE: - g = tl.load(gate_ptr + num_base, mask=mask_nd, other=0.0).to(tl.float32) - silu_g = g * (1.0 / (1.0 + tl.exp(-g))) - out_val = out_val * silu_g - - tl.store(out_ptr + num_base, out_val.to(tl.bfloat16), mask=mask_nd) - - -def fused_bidi_merge( - num_fwd: torch.Tensor, - num_bwd: torch.Tensor | None, - den_fwd: torch.Tensor, - den_bwd: torch.Tensor | None, - eps: float, - gate: torch.Tensor | None = None, -) -> torch.Tensor: - pre_summed = num_bwd is None - assert (num_bwd is None) == (den_bwd is None), "num_bwd/den_bwd must both be None or both provided" - if not pre_summed: - assert num_fwd.shape == num_bwd.shape and den_fwd.shape == den_bwd.shape - assert num_fwd.dtype == num_bwd.dtype and den_fwd.dtype == den_bwd.dtype - B, N, H, D = num_fwd.shape - out = torch.empty( - B, N, H, D, device=num_fwd.device, dtype=(torch.float32 if num_fwd.dtype == torch.float32 else torch.bfloat16) - ) - BLOCK_D = triton.next_power_of_2(D) - BLOCK_N = 64 - grid = (B * H, triton.cdiv(N, BLOCK_N)) - if gate is not None: - assert gate.shape == (B, N, H, D), f"gate shape {gate.shape} != {(B, N, H, D)}" - gate_arg = gate - apply_gate = 1 - else: - gate_arg = num_fwd - apply_gate = 0 - num_bwd_arg = num_bwd if num_bwd is not None else num_fwd - den_bwd_arg = den_bwd if den_bwd is not None else den_fwd - _fused_bidi_merge_kernel[grid]( - num_fwd, - num_bwd_arg, - den_fwd, - den_bwd_arg, - gate_arg, - out, - B, - N, - H, - D, - float(eps), - num_fwd.stride(0), - num_fwd.stride(1), - num_fwd.stride(2), - num_fwd.stride(3), - den_fwd.stride(0), - den_fwd.stride(1), - den_fwd.stride(2), - APPLY_GATE=apply_gate, - PRE_SUMMED=1 if pre_summed else 0, - BLOCK_N=BLOCK_N, - BLOCK_D=BLOCK_D, - ) - return out - - -# ===================================================================== -# Single-direction GDN entry point (delegates to chunkwise) -# ===================================================================== - - -def fused_gdn_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,) = (H*D,) 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, - reverse: bool = False, - init_state_kv: torch.Tensor | None = None, - init_state_z: torch.Tensor | None = None, - save_final_state: bool = False, -) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """One direction of fused BiGDN via the unified kernel. - - Args: - qkv .. eps: see kernel signature. - reverse: forward (False) or anti-causal (True) scan. - init_state_kv: optional ``(B*H, BLOCK_D, BLOCK_D)`` fp32 contiguous - tensor holding the forward-scan KV state at the END of a prefix - sequence (i.e., AFTER the prefix's last update, BEFORE any further - decay applied by this call). When provided, the kernel resumes the - scan from this state instead of zero. ``BLOCK_D = next_pow2(D)``. - Only the top-left ``D x D`` submatrix of the tile is read. - init_state_z: optional ``(B*H, BLOCK_D)`` fp32 contiguous companion - for the Z denominator state. Must be provided iff ``init_state_kv`` - is provided. - save_final_state: when True, allocate fresh fp32 zero buffers for the - final KV / Z state (after the last frame's update) and pass them to - the kernel for write-out. Returns the buffers as additional outputs. - - Returns: - ``(num, den)`` — bf16 numerator ``(B, N, H, D)`` and denominator - ``(B, H, N)`` before divide. - - When ``save_final_state=True``, also returns - ``(final_state_kv, final_state_z)`` fp32 with shapes - ``(B*H, BLOCK_D, BLOCK_D)`` and ``(B*H, BLOCK_D)``. - - Raises: - NotImplementedError: if any state I/O argument is set together with - ``reverse=True``. The kernel supports state passing in both - directions, but state I/O is only defined for the forward direction - here to avoid silent misuse. - """ - # Dispatch both the stateless bidi case and the stateful forward path to - # chunkwise so split-equivalence uses one numeric implementation. - # Bypass via env: FUSED_GDN_FORCE_LEGACY=1. - if os.environ.get("FUSED_GDN_FORCE_LEGACY", "0") != "1": - from .fused_gdn_chunkwise import ( - fused_gdn_func_chunkwise, - fused_gdn_stateful_chunkwise, - ) - - # Validate state I/O args upfront — preserves the legacy fused_gdn_func's - # validation contract (callers depend on these specific ValueError / - # NotImplementedError signatures, e.g., test_state_validation). - if (init_state_kv is None) != (init_state_z is None): - raise ValueError( - "fused_gdn_func: init_state_kv and init_state_z must be provided together " - "(both None or both fp32 tensors)." - ) - if reverse and (init_state_kv is not None or save_final_state): - raise NotImplementedError( - "fused_gdn_func: state passing (init_state_kv / init_state_z / " - "save_final_state) is only supported for the forward direction " - "(reverse=False)." - ) - if init_state_kv is not None: - B_q, _N, _three, H_q, D_q = qkv.shape - BLOCK_D_q = triton.next_power_of_2(D_q) - expected_kv = (B_q * H_q, BLOCK_D_q, BLOCK_D_q) - expected_z = (B_q * H_q, BLOCK_D_q) - if tuple(init_state_kv.shape) != expected_kv: - raise ValueError( - f"fused_gdn_func: init_state_kv shape {tuple(init_state_kv.shape)} != " f"expected {expected_kv}." - ) - if tuple(init_state_z.shape) != expected_z: - raise ValueError( - f"fused_gdn_func: init_state_z shape {tuple(init_state_z.shape)} != " f"expected {expected_z}." - ) - if init_state_kv.dtype != torch.float32 or init_state_z.dtype != torch.float32: - raise ValueError( - f"fused_gdn_func: init_state_kv/init_state_z must be fp32 " - f"(got {init_state_kv.dtype}, {init_state_z.dtype})." - ) - if not init_state_kv.is_contiguous() or not init_state_z.is_contiguous(): - raise ValueError("fused_gdn_func: init_state_kv / init_state_z must be contiguous.") - - # Stateless path - if init_state_kv is None and init_state_z is None and not save_final_state: - return fused_gdn_func_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, - reverse=reverse, - ) - - # Stateful path: shape-adapt state I/O. - # state_kv: (B*H, BLOCK_D, BLOCK_D) row-major as M[K_feat, V_feat] - # chunkwise stateful: takes user-facing (B, H, D_in, D_out) and transposes - # internally to (B*H, D_out, D_in) for kernel storage. - # state_z: (B*H, BLOCK_D) - # chunkwise stateful: (B, H, D, 1) or (B, H, D) - B, N, _three, H, D = qkv.shape - BLOCK_D = triton.next_power_of_2(D) - - ck_init_kv = None - ck_init_z = None - if init_state_kv is not None: - # (B*H, BLOCK_D, BLOCK_D) → (B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D] - # then transpose so chunkwise's internal `.transpose(-1, -2)` undoes it. - ck_init_kv = init_state_kv.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D].transpose(-1, -2).contiguous() - if init_state_z is not None: - # (B*H, BLOCK_D) → (B, H, BLOCK_D)[:, :, :D] → (B, H, D, 1) - ck_init_z = init_state_z.view(B, H, BLOCK_D)[:, :, :D].unsqueeze(-1).contiguous() - - result = fused_gdn_stateful_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, - reverse=reverse, - init_state_kv=ck_init_kv, - init_state_z=ck_init_z, - return_final_state=save_final_state, - ) - - if not save_final_state: - return result # (num, den) - - num, den, ck_state_kv, ck_state_z = result - # chunkwise returns state_kv as (B, H, D, D), [K_feat, V_feat] (post its - # internal back-transpose). Convert to stateful (B*H, BLOCK_D, BLOCK_D) - # by transposing back to internal storage and padding to BLOCK_D. - out_state_kv = torch.zeros(B * H, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) - out_state_kv[:, :D, :D] = ck_state_kv.transpose(-1, -2).reshape(B * H, D, D) - out_state_z = torch.zeros(B * H, BLOCK_D, device=qkv.device, dtype=torch.float32) - out_state_z[:, :D] = ck_state_z.squeeze(-1).reshape(B * H, D) - return num, den, out_state_kv, out_state_z - - B, N, three, H, D = qkv.shape - assert three == 3 - - BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta, decay = _prepare_launch(D, beta, decay) - - has_init_state = init_state_kv is not None or init_state_z is not None - if reverse and (has_init_state or save_final_state): - raise NotImplementedError( - "fused_gdn_func: state passing (init_state_kv / init_state_z / " - "save_final_state) is only supported for the forward direction " - "(reverse=False). The chunk-causal anti-causal pass resets state " - "per chunk and has no global cross-prefix state to cache." - ) - - if has_init_state: - if init_state_kv is None or init_state_z is None: - raise ValueError( - "fused_gdn_func: init_state_kv and init_state_z must be " - "provided together (got " - f"init_state_kv={'set' if init_state_kv is not None else 'None'}, " - f"init_state_z={'set' if init_state_z is not None else 'None'})." - ) - expected_kv_shape = (B * H, BLOCK_D, BLOCK_D) - expected_z_shape = (B * H, BLOCK_D) - if tuple(init_state_kv.shape) != expected_kv_shape: - raise ValueError( - f"fused_gdn_func: init_state_kv shape {tuple(init_state_kv.shape)} " - f"does not match expected {expected_kv_shape} (BLOCK_D=next_pow2(D)={BLOCK_D})." - ) - if tuple(init_state_z.shape) != expected_z_shape: - raise ValueError( - f"fused_gdn_func: init_state_z shape {tuple(init_state_z.shape)} " - f"does not match expected {expected_z_shape}." - ) - if init_state_kv.dtype != torch.float32 or init_state_z.dtype != torch.float32: - raise ValueError( - "fused_gdn_func: init_state_kv and init_state_z must be fp32 " - f"(got {init_state_kv.dtype}, {init_state_z.dtype})." - ) - if not init_state_kv.is_contiguous() or not init_state_z.is_contiguous(): - raise ValueError("fused_gdn_func: init_state_kv and init_state_z must be contiguous.") - if init_state_kv.device != qkv.device or init_state_z.device != qkv.device: - raise ValueError("fused_gdn_func: init_state_* must live on the same device as qkv.") - load_init = 1 - init_kv_arg = init_state_kv - init_z_arg = init_state_z - else: - load_init = 0 - init_kv_arg = None # placeholder set below - - if save_final_state: - final_state_kv = torch.zeros(B * H, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) - final_state_z = torch.zeros(B * H, BLOCK_D, device=qkv.device, dtype=torch.float32) - save_final = 1 - else: - final_state_kv = None - final_state_z = None - save_final = 0 - - num = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) - den = torch.empty(B, H, N, device=qkv.device, dtype=qkv.dtype) - dummy = torch.empty(1, device=qkv.device, dtype=torch.float32) - - # Resolve pointer args for the unused slots to a shared scratch tensor; - # the kernel compiles the corresponding load/store away when the - # constexpr flag is 0. - init_kv_ptr = init_kv_arg if load_init else dummy - init_z_ptr = init_z_arg if load_init else dummy - final_kv_ptr = final_state_kv if save_final else dummy - final_z_ptr = final_state_z if save_final else dummy - - _fused_gdn_kernel[(B * H,)]( - qkv, - qkv.stride(0), - qkv.stride(1), - qkv.stride(2), - qkv.stride(3), - qkv.stride(4), - beta, - decay, - q_inv_rms, - k_inv_rms, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - num, # `out_ptr` reuses `num` buffer (result immediately overwritten below) - num, - den, - dummy, - dummy, - dummy, - dummy, # saved-state dummies (SAVE_STATE=0) - init_kv_ptr, - init_z_ptr, - final_kv_ptr, - final_z_ptr, - H=H, - F=F, - S=S, - D=D, - K_SCALE=k_scale, - NORM_EPS=1e-5, # unused with USE_PRECOMPUTED_RMS=1 - EPS=eps, - QK_NORM=1, - USE_PRECOMPUTED_RMS=1, - STATE_FP32=1 if state_fp32 else 0, - DOT_PRECISION=dot_prec, - REVERSE=1 if reverse else 0, - SAVE_STATE=0, - LOAD_INIT_STATE=load_init, - SAVE_FINAL_STATE=save_final, - BLOCK_D=BLOCK_D, - BLOCK_S=BLOCK_S, - num_stages=cfg["num_stages"], - num_warps=nw, - ) - if save_final_state: - return num, den, final_state_kv, final_state_z - return num, den - - -def fused_bigdn_func( - qkv: torch.Tensor, # (B, N, 3, H, D) - q_inv_rms: torch.Tensor, # (B, N) — pre-computed via `_precompute_inv_rms` - k_inv_rms: torch.Tensor, # (B, N) - q_norm_weight: torch.Tensor, # (C,) float32 - k_norm_weight: torch.Tensor, # (C,) - rope_cos: torch.Tensor, # (N, D) - rope_sin: torch.Tensor, # (N, D) - beta: torch.Tensor, # (B, H, F, S) - decay: torch.Tensor, # (B, H, F) - F: int, - S: int, - k_scale: float, - eps: float = 1e-6, - # -- chunk-causal extensions (not in upstream; see adapter notes below) -- - qkv_bwd: torch.Tensor | None = None, - beta_bwd: torch.Tensor | None = None, - decay_bwd: torch.Tensor | None = None, - q_inv_rms_bwd: torch.Tensor | None = None, - k_inv_rms_bwd: torch.Tensor | None = None, -) -> torch.Tensor: - """Full bidirectional fused GDN. - - Returns: out (B, N, H, D) bf16 = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps). - - Chunk-causal extensions (optional): - For chunk-causal GDN we need to zero state at chunk boundaries in the - BACKWARD direction only. Pass separately pre-processed backward tensors - (decay_bwd with zeros at boundary frames, and optionally qkv_bwd / - beta_bwd with K/V or beta zeroed at boundary frames). If any `*_bwd` - argument is None, the forward tensor is reused. - """ - if ( - os.environ.get("FUSED_GDN_FORCE_LEGACY", "0") != "1" - and qkv_bwd is None - and beta_bwd is None - and decay_bwd is None - and q_inv_rms_bwd is None - and k_inv_rms_bwd is None - ): - from .fused_gdn_chunkwise import fused_bigdn_bidi_chunkwise - - 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, - ) - - num_fwd, den_fwd = fused_gdn_func( - 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, - reverse=False, - ) - num_bwd, den_bwd = fused_gdn_func( - qkv if qkv_bwd is None else qkv_bwd, - q_inv_rms if q_inv_rms_bwd is None else q_inv_rms_bwd, - k_inv_rms if k_inv_rms_bwd is None else k_inv_rms_bwd, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - beta if beta_bwd is None else beta_bwd, - decay if decay_bwd is None else decay_bwd, - F=F, - S=S, - k_scale=k_scale, - eps=eps, - reverse=True, - ) - # num: (B, N, H, D), den: (B, H, N). Fuse then divide. - total_num = num_fwd + num_bwd - total_den = (den_fwd + den_bwd).permute(0, 2, 1).unsqueeze(-1) # (B, N, H, 1) - return total_num / (total_den + eps) - - -# ===================================================================== -# Backward / autograd Functions -# ===================================================================== -# Adds: -# 1. ``_fused_gdn_bwd_kernel`` -- Triton jit kernel that replays the -# forward recurrence in reverse time using per-frame state snapshots -# written by the forward kernel under ``SAVE_STATE=1``. -# 2. ``_run_fwd_save`` -- helper that runs the existing forward -# ``_fused_gdn_kernel`` with ``SAVE_STATE=1``. Adapted to pass our -# extra ``init_state_kv_ptr / init_state_z_ptr / final_state_kv_ptr / -# final_state_z_ptr`` pointers + ``LOAD_INIT_STATE / SAVE_FINAL_STATE`` -# constexpr flags (all unused on the autograd path -> dummy / 0). -# 3. ``FusedGDNFunction`` -- autograd Function for unidirectional GDN -# with ``QK_NORM=1`` (in-kernel per-head RMSNorm). -# 4. ``FusedBiGDNFunction`` -- autograd Function for bidirectional BiGDN. -# Pre-normalizes Q/K in PyTorch with full-channel RMSNorm, runs the -# forward kernel twice with ``QK_NORM=0 + SAVE_STATE=1``, fuses -# ``(num_fwd + num_bwd) / (den_fwd + den_bwd + eps)``. Backward -# computes ``dnum / dden`` from upstream ``dout`` and runs the bwd -# kernel twice with ``BIDI_MODE=1``. -# 5. Python wrappers ``fused_gdn_forward_with_grad`` / -# ``fused_bigdn_forward_with_grad`` -- drop-in autograd-enabled -# replacements for ``fused_gdn_func`` / ``fused_bigdn_func``. -# -# Chunk-causal autograd support: ``FusedBiGDNFunction`` (and the public -# wrapper ``fused_bigdn_forward_with_grad``) accepts optional -# ``beta_bwd`` / ``decay_bwd`` overrides for the reverse-direction -# kernel call -- exactly the same masking convention used by the -# inference path ``fused_bigdn_func``. When provided, the reverse -# direction's forward and backward kernels both run on these masked -# tensors, and the backward returns separate gradient tensors -# (``dbeta_bwd`` / ``ddecay_bwd``) so autograd can route them back -# through any ``clone() + index = 0`` masking the caller applied. - - -@triton.jit -def _fused_gdn_bwd_kernel( - # ---- original inputs ---- - qkv_ptr, - stride_b: tl.constexpr, - stride_n: tl.constexpr, - stride_3: tl.constexpr, - stride_h: tl.constexpr, - stride_d: tl.constexpr, - beta_ptr, - decay_ptr, - q_norm_w_ptr, - k_norm_w_ptr, - rope_cos_ptr, - rope_sin_ptr, - # ---- saved from forward ---- - saved_state_ptr, # (B*H, F, BLOCK_D, BLOCK_D) -- state_prev snapshots - saved_z_ptr, # (B*H, F, BLOCK_D) - saved_state_curr_ptr, # (B*H, F, BLOCK_D, BLOCK_D) -- state_curr (after update) - saved_z_curr_ptr, # (B*H, F, BLOCK_D) - # ---- upstream gradient / pre-computed dnum ---- - dout_ptr, # GDN mode: (B, N, H, D) upstream grad. BiDI mode: pre-computed dnum - # ---- BiDI mode: external dden ---- - dden_ext_ptr, # BiDI mode: (B, H, N) pre-computed dden. GDN mode: unused - # ---- output gradients ---- - dqkv_ptr, # (B, N, 3, H, D) -- same layout as qkv - dbeta_ptr, # (B, H, F, S) - ddecay_ptr, # (B, H, F) - # ---- dims ---- - H: tl.constexpr, - F: tl.constexpr, - S: tl.constexpr, - D: tl.constexpr, - K_SCALE, - NORM_EPS: tl.constexpr, - EPS: tl.constexpr, - QK_NORM: tl.constexpr, - STATE_FP32: tl.constexpr, - REVERSE_BWD: tl.constexpr, # 0=backward of forward GDN, 1=backward of reversed GDN - BIDI_MODE: tl.constexpr, # 0=GDN (compute dnum/dden), 1=BiGDN (use provided) - DOT_PRECISION: tl.constexpr, # 0=bf16 TC, 1=TF32 TC, 2=IEEE fp32 - BLOCK_D: tl.constexpr, - BLOCK_S: tl.constexpr, -): - pid = tl.program_id(0) - pid_b = pid // H - pid_h = pid % H - N: tl.constexpr = F * S - bh = pid_b * H + pid_h - - qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h - dqkv_bh = dqkv_ptr + pid_b * stride_b + pid_h * stride_h - dout_bh = dout_ptr + pid_b * (N * H * D) + pid_h * D - beta_bh = beta_ptr + bh * (F * S) - decay_bh = decay_ptr + bh * F - dbeta_bh = dbeta_ptr + bh * (F * S) - ddecay_bh = ddecay_ptr + bh * F - st_bh = saved_state_ptr + bh * F * BLOCK_D * BLOCK_D - sz_bh = saved_z_ptr + bh * F * BLOCK_D - stc_bh = saved_state_curr_ptr + bh * F * BLOCK_D * BLOCK_D - szc_bh = saved_z_curr_ptr + bh * F * BLOCK_D - if BIDI_MODE: - dden_ext_bh = dden_ext_ptr + bh * N - - 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 - if QK_NORM: - q_nw = tl.load(q_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) - k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to(tl.float32) - q_nw_pair = tl.load(q_norm_w_ptr + nw_offset + offs_d_pair, mask=mask_d_pair, 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) - - D_inv = 1.0 / D - k_scale = K_SCALE - - # Dot precision: mirror forward kernel - if DOT_PRECISION >= 1: - dot_dtype = tl.float32 - else: - dot_dtype = tl.bfloat16 - dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" - - # Gradient matmuls: always use bf16 TC + TF32 input precision (matching PyTorch backward) - grad_dtype = tl.bfloat16 - grad_ip: tl.constexpr = "tf32" - - # ---- Gradient state accumulators (reverse time) ---- - dstate = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) - dstate_z = tl.zeros([BLOCK_D], dtype=tl.float32) - - for f_rev in range(F): - # Backward iterates in reverse of forward direction. - if REVERSE_BWD: - f = f_rev # backward of reversed GDN: iterate 0..F-1 - # In fwd_save REVERSE, q_frame=f had kv_frame=f+1 (or skip at f=F-1). - kv_frame_bwd = f + 1 if f < F - 1 else f - skip_bwd = f == F - 1 # f=F-1 was dummy step (f_iter=0 in fwd) - else: - f = F - 1 - f_rev # backward of forward GDN: iterate F-1..0 - kv_frame_bwd = f - skip_bwd = False - q_n_base = f * S - kv_n_base = kv_frame_bwd * S - f_beta = beta_bh + kv_frame_bwd * S - - # ---- Load state_curr for Pass 2 output (both directions use inclusive) ---- - st_f = st_bh + f * BLOCK_D * BLOCK_D - offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] - mask_dd = mask_d[:, None] & mask_d[None, :] - - stc_f = stc_bh + f * BLOCK_D * BLOCK_D - P_state = tl.load(stc_f + offs_dd, mask=mask_dd, other=0.0) - Pz_state = tl.load(szc_bh + f * BLOCK_D + offs_d, mask=mask_d, other=0.0) - if STATE_FP32 == 0: - P_state = P_state.to(tl.float32) - - # Decay: for REVERSE_BWD, use decay[kv_frame] matching fwd_save. - if REVERSE_BWD and skip_bwd: - g = 1.0 - elif REVERSE_BWD: - g = tl.load(decay_bh + kv_frame_bwd).to(tl.float32) - else: - g = tl.load(decay_bh + f).to(tl.float32) - - # ======================================================== - # Pass 2 backward: Output gradients -> dQ, dstate, dstate_z - # ======================================================== - 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, :] - mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] - n_idx = q_n_base + offs_s # Q data from q_frame - - # Load dout; recompute Q, Q_pair, Q_rot, num, den from saved P_f/Pz_f. - dout_ptrs = dout_bh + n_idx[:, None] * (H * D) + offs_d[None, :] - d_out = tl.load(dout_ptrs, mask=mask_sd, other=0.0).to(tl.float32) - - # Recompute Q, Q_pair, Q_rot (same as forward). - q_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d[None, :] * stride_d - Q_raw = tl.load(q_ptrs, mask=mask_sd, other=0.0).to(tl.float32) - q_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d_pair[None, :] * stride_d - Q_pair_raw = tl.load(q_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) - - if QK_NORM: - q_var = tl.sum(Q_raw * Q_raw, axis=1) * D_inv - q_inv_rms = 1.0 / tl.sqrt(q_var + NORM_EPS) - Q_normed = Q_raw * q_inv_rms[:, None] * q_nw[None, :] - Q_pair_normed = Q_pair_raw * q_inv_rms[:, None] * q_nw_pair[None, :] - else: - Q_normed = Q_raw - Q_pair_normed = Q_pair_raw - Q = tl.where(Q_normed > 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 - - # Compute dnum and dden. - if BIDI_MODE: - # BiGDN: dnum and dden pre-computed externally from total num/den. - dnum = d_out # dout_ptr already contains pre-computed dnum - dden = tl.load(dden_ext_bh + n_idx, mask=mask_s, other=0.0).to(tl.float32) - else: - # GDN: recompute num/den using direction-appropriate state. - num_tile = tl.dot( - Q_rot.to(dot_dtype), P_state.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip - ) - den_tile = tl.sum(Q * Pz_state[None, :], axis=1) - inv_den = 1.0 / (den_tile + EPS) - dnum = d_out * inv_den[:, None] - dden = -tl.sum(d_out * num_tile, axis=1) * inv_den * inv_den - - # dstate += Q_rot^T @ dnum (state contribution from num = Q_rot @ P_state). - dstate = dstate + tl.dot( - tl.trans(Q_rot.to(grad_dtype)), - dnum.to(grad_dtype), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - # dstate_z += sum(dden * Q, axis=0) (Pz contribution from den = Q . Pz). - dstate_z += tl.sum(dden[:, None] * Q, axis=0) - - # dQ_rot = dnum @ P_state^T (uses state that forward's output read). - dQ_rot = tl.dot( - dnum.to(grad_dtype), - tl.trans(P_state.to(grad_dtype)), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - # dQ_from_den = dden * Pz_state. - dQ_from_den = dden[:, None] * Pz_state[None, :] - - # RoPE inverse for Q: store dQ_rot, reload at paired indices. - # Store dQ_rot temporarily to dqkv[Q] at normal d positions. - dq_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d[None, :] * stride_d - tl.store(dq_ptrs, dQ_rot.to(tl.bfloat16), mask=mask_sd) - # The XOR-paired channel can be owned by another warp. Synchronize - # both sides of the scratch roundtrip before dqkv is overwritten. - tl.debug_barrier() - - # Load dQ_rot at paired positions. - dq_pair_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 0 * stride_3 + offs_d_pair[None, :] * stride_d - dQ_rot_pair = tl.load(dq_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) - tl.debug_barrier() - - # RoPE inverse: dQ = dQ_rot * Cos - dQ_rot_pair * Sin. - dQ = dQ_rot * Cos - dQ_rot_pair * Sin + dQ_from_den - - # ReLU backward. - relu_mask_q = (Q_normed > 0).to(tl.float32) - dQ_normed = dQ * relu_mask_q - - # Norm backward (QK_NORM) or direct (no norm). - if QK_NORM: - gw = dQ_normed * q_nw[None, :] - corr = tl.sum(gw * Q_raw, axis=1) * D_inv * q_inv_rms * q_inv_rms - dQ_raw = q_inv_rms[:, None] * (gw - Q_raw * corr[:, None]) - else: - dQ_raw = dQ_normed - - # Store final dQ_raw to dqkv[Q]. - tl.store(dq_ptrs, dQ_raw.to(tl.bfloat16), mask=mask_sd) - - # Both directions use inclusive output (state_curr), so capture dDelta AFTER Pass 2. - dDelta = dstate - dDelta_z = dstate_z - - # ======================================================== - # Reload state_prev for Pass 1 backward (reuse P_state variable) - # ======================================================== - P_state = tl.load(st_f + offs_dd, mask=mask_dd, other=0.0) - if STATE_FP32 == 0: - P_state = P_state.to(tl.float32) - Pz_state = tl.load(sz_bh + f * BLOCK_D + offs_d, mask=mask_d, other=0.0) - - # ======================================================== - # Pass 1 backward: State update gradients -> dK, dV, dbeta, dstate - # Skip for REVERSE_BWD dummy frame (skip_bwd=True) to avoid clobbering. - # ======================================================== - if skip_bwd == False: - 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, :] - mask_sd_pair = mask_s[:, None] & mask_d_pair[None, :] - n_idx = kv_n_base + offs_s # K/V from kv_frame - - # Recompute K, K_pair, K_rot, V. - 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) - - k_pair_ptrs = qkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d_pair[None, :] * stride_d - K_pair_raw = tl.load(k_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) - - if QK_NORM: - k_var = tl.sum(K_raw * K_raw, axis=1) * D_inv - k_inv_rms = 1.0 / tl.sqrt(k_var + NORM_EPS) - K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] - K_pair_normed = K_pair_raw * k_inv_rms[:, None] * k_nw_pair[None, :] - else: - K_normed = K_raw - K_pair_normed = K_pair_raw - K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale - 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 - - bt = tl.load(f_beta + offs_s, mask=mask_s, other=0.0).to(tl.float32) - - # Recompute V_pred and delta_v. - K_rot_dc = K_rot.to(dot_dtype) - V_pred = tl.dot(K_rot_dc, P_state.to(dot_dtype), out_dtype=tl.float32, input_precision=dot_ip) - delta_v = (V_raw - V_pred) * bt[:, None] - - # ---- KV stream backward ---- - ddelta_v = tl.dot( - K_rot.to(grad_dtype), - dDelta.to(grad_dtype), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - dK_rot_from_delta = tl.dot( - delta_v.to(grad_dtype), - tl.trans(dDelta.to(grad_dtype)), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - dV = ddelta_v * bt[:, None] - dbeta_kv = tl.sum(ddelta_v * (V_raw - V_pred), axis=1) - - dV_pred = -ddelta_v * bt[:, None] - dK_rot_from_vpred = tl.dot( - dV_pred.to(grad_dtype), - tl.trans(P_state.to(grad_dtype)), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - dstate = dstate + tl.dot( - tl.trans(K_rot.to(grad_dtype)), - dV_pred.to(grad_dtype), - out_dtype=tl.float32, - input_precision=grad_ip, - ) - - dK_rot = dK_rot_from_delta + dK_rot_from_vpred - - # ---- Z stream backward ---- - z_hat = tl.sum(K * Pz_state[None, :], axis=1) - dz = (1.0 - z_hat) * bt - - ddz = tl.sum(K * dDelta_z[None, :], axis=1) - dz_hat = -ddz * bt - dK_z = dDelta_z[None, :] * dz[:, None] + dz_hat[:, None] * Pz_state[None, :] - dstate_z = dstate_z + tl.sum(dz_hat[:, None] * K, axis=0) - - dbeta_z = ddz * (1.0 - z_hat) - dbeta_total = dbeta_kv + dbeta_z - tl.store(dbeta_bh + kv_frame_bwd * S + offs_s, dbeta_total.to(tl.bfloat16), mask=mask_s) - - # ---- RoPE inverse for K ---- - dk_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d[None, :] * stride_d - tl.store(dk_ptrs, dK_rot.to(tl.bfloat16), mask=mask_sd) - # Match the dQ synchronization: all temporary values must be - # visible before paired loads and consumed before overwrites. - tl.debug_barrier() - dk_pair_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 1 * stride_3 + offs_d_pair[None, :] * stride_d - dK_rot_pair = tl.load(dk_pair_ptrs, mask=mask_sd_pair, other=0.0).to(tl.float32) - tl.debug_barrier() - - dK_from_kv = dK_rot * Cos - dK_rot_pair * Sin - dK_total = dK_from_kv + dK_z - - relu_mask_k = (K_normed > 0).to(tl.float32) - dK_normed = dK_total * k_scale * relu_mask_k - - if QK_NORM: - gw_k = dK_normed * k_nw[None, :] - corr_k = tl.sum(gw_k * K_raw, axis=1) * D_inv * k_inv_rms * k_inv_rms - dK_raw = k_inv_rms[:, None] * (gw_k - K_raw * corr_k[:, None]) - else: - dK_raw = dK_normed - - tl.store(dk_ptrs, dK_raw.to(tl.bfloat16), mask=mask_sd) - dv_ptrs = dqkv_bh + n_idx[:, None] * stride_n + 2 * stride_3 + offs_d[None, :] * stride_d - tl.store(dv_ptrs, dV.to(tl.bfloat16), mask=mask_sd) - - # ======================================================== - # Decay backward (inside skip_bwd guard) - # ======================================================== - is_first_frame = f_rev == F - 1 - if is_first_frame: - ddecay_f = 0.0 - else: - inv_g = 1.0 / (g + 1e-12) - ddecay_kv = tl.sum(dstate * P_state) * inv_g - ddecay_z_val = tl.sum(dstate_z * Pz_state) * inv_g - ddecay_f = ddecay_kv + ddecay_z_val - tl.store(ddecay_bh + kv_frame_bwd, ddecay_f) - - # Propagate gradient through decay: dS_{f-1} = g[f] * dP_f. - dstate = dstate * g - dstate_z = dstate_z * g - - -# ===================================================================== -# Forward-with-state-save helper (for autograd Functions) -# ===================================================================== - - -def _run_fwd_save( - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - F: int, - S: int, - k_scale: float, - norm_eps: float, - eps: float, - qk_norm: bool, - reverse: bool, - cfg, -): - """Run forward kernel for one direction with ``SAVE_STATE=1``. - - Returns ``(num, den, saved_state, saved_z, saved_state_curr, saved_z_curr)``. - The forward kernel writes ``out = num/(den+eps)`` first and then overwrites - the same buffer with raw ``num``, so the returned ``num`` tensor holds raw - numerator values (matching the BiGDN combine-then-divide convention). - """ - B, N, three, H, D = qkv.shape - BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, _, beta, decay = _prepare_launch(D, beta, decay) - - num_out = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) - den_out = torch.empty(B, H, N, device=qkv.device, dtype=qkv.dtype) - state_dtype = torch.float32 if state_fp32 else torch.bfloat16 - saved_state = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=state_dtype) - saved_z = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) - saved_state_curr = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) - saved_z_curr = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) - # The kernel writes ``out = num/(den+eps)`` first then overwrites with raw num - # in the same buffer. Reuse num_out as the (discarded) ``out`` slot so the - # final contents end up being raw num. - out_discard = num_out - dummy_inv = torch.empty(1, device=qkv.device, dtype=torch.float32) - - _fused_gdn_kernel[(B * H,)]( - qkv, - qkv.stride(0), - qkv.stride(1), - qkv.stride(2), - qkv.stride(3), - qkv.stride(4), - beta, - decay, - dummy_inv, - dummy_inv, # unused inv_rms ptrs (USE_PRECOMPUTED_RMS=0) - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - out_discard, - num_out, - den_out, - saved_state, - saved_z, - saved_state_curr, - saved_z_curr, - dummy_inv, - dummy_inv, - dummy_inv, - dummy_inv, # init/final-state dummies - H=H, - F=F, - S=S, - D=D, - K_SCALE=k_scale, - NORM_EPS=norm_eps, - EPS=eps, - QK_NORM=1 if qk_norm else 0, - USE_PRECOMPUTED_RMS=0, - STATE_FP32=1 if state_fp32 else 0, - DOT_PRECISION=dot_prec, - REVERSE=1 if reverse else 0, - SAVE_STATE=1, - LOAD_INIT_STATE=0, - SAVE_FINAL_STATE=0, - BLOCK_D=BLOCK_D, - BLOCK_S=BLOCK_S, - num_stages=cfg["num_stages"], - num_warps=nw, - ) - return num_out, den_out, saved_state, saved_z, saved_state_curr, saved_z_curr - - -# ===================================================================== -# Unidirectional GDN autograd Function -# ===================================================================== - - -class FusedGDNFunction(torch.autograd.Function): - """Autograd Function for unidirectional fused GDN with in-kernel RMSNorm. - - Forward runs ``_fused_gdn_kernel`` with ``QK_NORM=1`` and ``SAVE_STATE=1``, - saving per-frame state snapshots for backward. Backward runs - ``_fused_gdn_bwd_kernel`` with ``BIDI_MODE=0`` (kernel computes - ``dnum``/``dden`` from upstream ``dout``). - """ - - @staticmethod - def forward( - ctx, - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - F: int, - S: int, - k_scale: float = 1.0, - norm_eps: float = 1e-6, - eps: float = 1e-6, - qk_norm: bool = True, - ): - B, N, three, H, D = qkv.shape - assert three == 3 and N == F * S - - BLOCK_D, BLOCK_S, dot_prec, state_fp32, nw, cfg, beta, decay = _prepare_launch(D, beta, decay) - - if q_norm_weight is None: - q_norm_weight = torch.ones(D, device=qkv.device, dtype=torch.float32) - if k_norm_weight is None: - k_norm_weight = torch.ones(D, device=qkv.device, dtype=torch.float32) - - out = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) - - # Saved states for backward. - state_dtype = torch.float32 if state_fp32 else torch.bfloat16 - saved_state = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=state_dtype) - saved_z = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) - saved_state_curr = torch.empty(B * H, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=torch.float32) - saved_z_curr = torch.empty(B * H, F, BLOCK_D, device=qkv.device, dtype=torch.float32) - - # Dummy num/den for forward kernel (still writes them but we discard). - num_out = torch.empty(B, N, H, D, device=qkv.device, dtype=qkv.dtype) - den_out = torch.empty(B, H, N, device=qkv.device, dtype=qkv.dtype) - dummy_inv = torch.empty(1, device=qkv.device, dtype=torch.float32) - - _fused_gdn_kernel[(B * H,)]( - qkv, - qkv.stride(0), - qkv.stride(1), - qkv.stride(2), - qkv.stride(3), - qkv.stride(4), - beta, - decay, - dummy_inv, - dummy_inv, # unused inv_rms ptrs (USE_PRECOMPUTED_RMS=0) - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - out, - num_out, - den_out, - saved_state, - saved_z, - saved_state_curr, - saved_z_curr, - dummy_inv, - dummy_inv, - dummy_inv, - dummy_inv, # init/final-state dummies - H=H, - F=F, - S=S, - D=D, - K_SCALE=k_scale, - NORM_EPS=norm_eps, - EPS=eps, - QK_NORM=1 if qk_norm else 0, - USE_PRECOMPUTED_RMS=0, - STATE_FP32=1 if state_fp32 else 0, - DOT_PRECISION=dot_prec, - REVERSE=0, - SAVE_STATE=1, - LOAD_INIT_STATE=0, - SAVE_FINAL_STATE=0, - BLOCK_D=BLOCK_D, - BLOCK_S=BLOCK_S, - num_stages=cfg["num_stages"], - num_warps=nw, - ) - del num_out, den_out - - ctx.save_for_backward( - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - saved_state, - saved_z, - saved_state_curr, - saved_z_curr, - ) - ctx.F = F - ctx.S = S - ctx.k_scale = k_scale - ctx.norm_eps = norm_eps - ctx.eps = eps - ctx.qk_norm = qk_norm - ctx.dot_prec = dot_prec - return out - - @staticmethod - def backward(ctx, dout): - ( - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - saved_state, - saved_z, - saved_state_curr, - saved_z_curr, - ) = ctx.saved_tensors - - B, N, three, H, D = qkv.shape - F_val = ctx.F - S = ctx.S - - BLOCK_D, BLOCK_S_BWD, _, _, _, cfg, beta, decay = _prepare_launch(D, beta, decay) - dqkv = torch.zeros_like(qkv) - dbeta = torch.zeros_like(beta) - ddecay = torch.zeros_like(decay) - - # Dummy dden_ext (unused in GDN mode). - dden_ext = torch.empty(1, device=qkv.device, dtype=torch.float32) - - # Progressive num_warps reduction on tmem overflow. - nw = cfg["num_warps"] - if ctx.dot_prec >= 1: - nw = min(nw, 4) - while nw >= 1: - try: - _fused_gdn_bwd_kernel[(B * H,)]( - qkv, - qkv.stride(0), - qkv.stride(1), - qkv.stride(2), - qkv.stride(3), - qkv.stride(4), - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - saved_state, - saved_z, - saved_state_curr, - saved_z_curr, - dout.contiguous(), - dden_ext, - dqkv, - dbeta, - ddecay, - H=H, - F=F_val, - S=S, - D=D, - K_SCALE=ctx.k_scale, - NORM_EPS=ctx.norm_eps, - EPS=ctx.eps, - QK_NORM=1 if ctx.qk_norm else 0, - STATE_FP32=1 if cfg["STATE_FP32"] else 0, - REVERSE_BWD=0, - BIDI_MODE=0, - DOT_PRECISION=ctx.dot_prec, - BLOCK_D=BLOCK_D, - BLOCK_S=BLOCK_S_BWD, - num_stages=cfg["num_stages"], - num_warps=nw, - ) - break - except Exception as e: - if "OutOfResources" in str(type(e).__name__) or "out of resource" in str(e).lower(): - nw = nw // 2 - if nw < 1: - raise RuntimeError( - "FusedGDN backward: Triton kernel OutOfResources at all warp " - f"counts (8, 4, 2, 1). Most recent error: {e}" - ) from e - else: - raise - - return dqkv, dbeta, ddecay, None, None, None, None, None, None, None, None, None, None - - -def fused_gdn_forward_with_grad( - qkv: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - q_norm_weight: torch.Tensor | None, - k_norm_weight: torch.Tensor | None, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - F: int, - S: int, - k_scale: float = 1.0, - norm_eps: float = 1e-6, - eps: float = 1e-6, - qk_norm: bool = True, -) -> torch.Tensor: - """Drop-in autograd-enabled replacement for the unidirectional GDN path. - - Unlike ``fused_gdn_func`` (which expects pre-computed ``q_inv_rms``/ - ``k_inv_rms``), this wrapper computes per-head RMSNorm inside the - Triton kernel (``QK_NORM=1`` / ``USE_PRECOMPUTED_RMS=0``) so the - backward kernel can reproduce the exact same normed Q/K when - replaying the recurrence. - """ - return FusedGDNFunction.apply( - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - F, - S, - k_scale, - norm_eps, - eps, - qk_norm, - ) - - -# ===================================================================== -# Bidirectional BiGDN autograd Function (full-channel RMSNorm in Python) -# ===================================================================== - - -class FusedBiGDNFunction(torch.autograd.Function): - """Autograd Function for bidirectional fused BiGDN. - - Full-channel RMSNorm is applied in Python (so the norm backward can - couple all heads correctly), then the kernel runs with ``QK_NORM=0`` - on the pre-normed QKV. Forward and reverse directions are run - separately with ``SAVE_STATE=1``, and combined as - ``out = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps)``. - - Backward computes ``dnum`` and ``dden`` from upstream ``dout`` and - runs the bwd kernel twice (forward + reverse) with ``BIDI_MODE=1``. - Norm backward is computed in Python (full-channel RMSNorm couples - all heads). - - Chunk-causal masking (optional): - Pass ``beta_bwd`` and/or ``decay_bwd`` to override the beta/decay - tensors used by the **reverse-direction** kernel calls (forward - save + backward). The forward direction always uses the - unmasked ``beta`` / ``decay``. This mirrors the inference path - in :func:`fused_bigdn_func` and unlocks chunk-causal autograd - training: callers typically build ``beta_bwd`` / ``decay_bwd`` - as ``beta.clone()`` / ``decay.clone()`` with interior chunk - boundaries zeroed, so the anti-causal scan resets state at - every chunk boundary. - - When ``beta_bwd`` is ``None``, the kernel-emitted reverse- - direction beta gradient is summed into the forward-direction - gradient (returned via the ``beta`` slot) and the ``beta_bwd`` - gradient slot returns ``None``. When ``beta_bwd`` is provided, - the two gradient streams are kept separate: the forward- - direction gradient flows through the ``beta`` slot and the - reverse-direction gradient flows through the ``beta_bwd`` slot - so autograd can route them through any ``clone() + index = 0`` - masking applied by the caller. ``decay_bwd`` is handled - identically. - """ - - @staticmethod - def forward( - ctx, - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - F: int, - S: int, - k_scale: float = 1.0, - norm_eps: float = 1e-5, - eps: float = 1e-6, - beta_bwd: torch.Tensor | None = None, - decay_bwd: torch.Tensor | None = None, - ): - B, N, three, H, D = qkv.shape - C = H * D - assert three == 3 and N == F * S - cfg = _kcfg() - - if q_norm_weight is None: - q_norm_weight = torch.ones(C, device=qkv.device, dtype=torch.float32) - if k_norm_weight is None: - k_norm_weight = torch.ones(C, device=qkv.device, dtype=torch.float32) - - # Full-channel RMSNorm: inv_rms over all H*D dims. - q_raw = qkv[:, :, 0].float() # (B, N, H, D) - k_raw = qkv[:, :, 1].float() - q_inv_rms = torch.rsqrt((q_raw * q_raw).sum(dim=(-2, -1)) / C + norm_eps) # (B, N) - k_inv_rms = torch.rsqrt((k_raw * k_raw).sum(dim=(-2, -1)) / C + norm_eps) - - # Apply norm to Q and K: Q_normed = Q_raw * inv_rms * weight. - q_nw_hd = q_norm_weight.reshape(H, D) - k_nw_hd = k_norm_weight.reshape(H, D) - qkv_normed = qkv.clone() - qkv_normed[:, :, 0] = (q_raw * q_inv_rms[:, :, None, None] * q_nw_hd[None, None]).to(qkv.dtype) - qkv_normed[:, :, 1] = (k_raw * k_inv_rms[:, :, None, None] * k_nw_hd[None, None]).to(qkv.dtype) - - # Reverse-direction beta/decay overrides for chunk-causal masking. - # When the caller supplies ``beta_bwd`` / ``decay_bwd`` (typically - # ``beta.clone()`` / ``decay.clone()`` with interior chunk-boundary - # frames zeroed), the reverse-direction kernel reads them instead of - # the unmasked tensors so the anti-causal scan resets state at chunk - # boundaries. The forward (causal) direction always uses the - # unmasked ``beta`` / ``decay``. - beta_for_bwd_dir = beta_bwd if beta_bwd is not None else beta - decay_for_bwd_dir = decay_bwd if decay_bwd is not None else decay - - # Run forward-save with QK_NORM=0 on pre-normed data. - dummy_nw = torch.ones(D, device=qkv.device, dtype=torch.float32) - num_fwd, den_fwd, sv_fwd, sz_fwd, svc_fwd, szc_fwd = _run_fwd_save( - qkv_normed, - beta, - decay, - dummy_nw, - dummy_nw, - rope_cos, - rope_sin, - F, - S, - k_scale, - norm_eps, - eps, - False, - False, - cfg, - ) - num_bwd, den_bwd, sv_bwd, sz_bwd, svc_bwd, szc_bwd = _run_fwd_save( - qkv_normed, - beta_for_bwd_dir, - decay_for_bwd_dir, - dummy_nw, - dummy_nw, - rope_cos, - rope_sin, - F, - S, - k_scale, - norm_eps, - eps, - False, - True, - cfg, - ) - - # Combine: out = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps). - total_num = num_fwd.float() + num_bwd.float() - total_den = den_fwd.float() + den_bwd.float() - total_den_exp = total_den.permute(0, 2, 1).unsqueeze(-1) # (B, N, H, 1) - out = (total_num / (total_den_exp + eps)).to(qkv.dtype) - - # Save ``beta_bwd`` / ``decay_bwd`` (possibly ``None``) so the - # backward pass can (a) replay the reverse-direction kernel against - # the same masked inputs, and (b) decide whether to keep the - # reverse-direction beta/decay gradients separate (caller-supplied - # override) or fold them into the forward-direction gradient - # (no override, legacy behaviour). - ctx.save_for_backward( - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - q_inv_rms, - k_inv_rms, - rope_cos, - rope_sin, - sv_fwd, - sz_fwd, - svc_fwd, - szc_fwd, - sv_bwd, - sz_bwd, - svc_bwd, - szc_bwd, - out, - total_den.to(qkv.dtype), - beta_bwd, - decay_bwd, - ) - _, _dot_prec, _, _ = _resolve_launch_config() - ctx.dot_prec = _dot_prec - ctx.F = F - ctx.S = S - ctx.k_scale = k_scale - ctx.norm_eps = norm_eps - ctx.eps = eps - return out - - @staticmethod - def backward(ctx, dout): - ( - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - q_inv_rms, - k_inv_rms, - rope_cos, - rope_sin, - sv_fwd, - sz_fwd, - svc_fwd, - szc_fwd, - sv_bwd, - sz_bwd, - svc_bwd, - szc_bwd, - out, - total_den_saved, - beta_bwd_saved, - decay_bwd_saved, - ) = ctx.saved_tensors - - # Track whether the caller supplied separate ``beta_bwd`` / - # ``decay_bwd`` overrides; this controls whether the reverse- - # direction kernel gradients are summed into the forward-direction - # slot (legacy behaviour) or routed back through dedicated grad - # slots so autograd can flow through the caller's masking ops - # (``clone() + index = 0``). - has_beta_bwd = beta_bwd_saved is not None - has_decay_bwd = decay_bwd_saved is not None - - B, N, three, H, D = qkv.shape - C = H * D - - # Recompute qkv_normed (avoid saving B*N*3*H*D extra tensor). - q_raw = qkv[:, :, 0].float() - k_raw = qkv[:, :, 1].float() - q_nw_hd = q_norm_weight.reshape(H, D) - k_nw_hd = k_norm_weight.reshape(H, D) - qkv_normed = qkv.clone() - qkv_normed[:, :, 0] = (q_raw * q_inv_rms[:, :, None, None] * q_nw_hd[None, None]).to(qkv.dtype) - qkv_normed[:, :, 1] = (k_raw * k_inv_rms[:, :, None, None] * k_nw_hd[None, None]).to(qkv.dtype) - F_val = ctx.F - S = ctx.S - eps = ctx.eps - BLOCK_D, _, _, _, _, cfg, beta, decay = _prepare_launch(D, beta, decay) - - # Reverse-direction beta/decay actually fed to the reverse kernel. - # When the caller supplied an override, we replay against the - # masked tensor; otherwise we reuse the unmasked beta/decay so the - # legacy summing path is bit-identical to the pre-extension - # behaviour. - if has_beta_bwd: - beta_for_bwd_dir = beta_bwd_saved.contiguous() - else: - beta_for_bwd_dir = beta - if has_decay_bwd: - decay_for_bwd_dir = decay_bwd_saved.contiguous() - else: - decay_for_bwd_dir = decay - - # ---- Pre-compute dnum and dden ---- - total_den_exp = total_den_saved.float().permute(0, 2, 1).unsqueeze(-1) - inv_total_den = 1.0 / (total_den_exp + eps) - dnum = (dout.float() * inv_total_den).to(qkv.dtype).contiguous() - dden = ( - (-(dout.float() * out.float()).sum(dim=-1) * inv_total_den.squeeze(-1)) - .permute(0, 2, 1) - .to(qkv.dtype) - .contiguous() - ) - del out, total_den_saved, total_den_exp, inv_total_den - - dummy_nw = torch.ones(D, device=qkv.device, dtype=torch.float32) - - # ---- Backward for forward direction (QK_NORM=0, operates on normed QKV) ---- - dqkv_fwd = torch.zeros_like(qkv) - dbeta_fwd = torch.zeros_like(beta) - ddecay_fwd = torch.zeros_like(decay) - - def _run_triton_bwd(sv, sz, svc, szc, dqkv_out, dbeta_out, ddecay_out, reverse_bwd, beta_kernel, decay_kernel): - """Try backward kernel with progressively fewer warps on tmem overflow. - - ``beta_kernel`` / ``decay_kernel`` are passed as explicit - arguments (instead of closing over the outer-scope ``beta`` / - ``decay``) so the reverse-direction call can replay against the - chunk-causal-masked tensors (``beta_bwd`` / ``decay_bwd``) when - present, while the forward-direction call always uses the - unmasked ``beta`` / ``decay``. - """ - nw = cfg["num_warps"] - if ctx.dot_prec >= 1: - nw = min(nw, 4) - while nw >= 1: - try: - _fused_gdn_bwd_kernel[(B * H,)]( - qkv_normed, - qkv_normed.stride(0), - qkv_normed.stride(1), - qkv_normed.stride(2), - qkv_normed.stride(3), - qkv_normed.stride(4), - beta_kernel, - decay_kernel, - dummy_nw, - dummy_nw, - rope_cos, - rope_sin, - sv, - sz, - svc, - szc, - dnum, - dden, - dqkv_out, - dbeta_out, - ddecay_out, - H=H, - F=F_val, - S=S, - D=D, - K_SCALE=ctx.k_scale, - NORM_EPS=ctx.norm_eps, - EPS=eps, - QK_NORM=0, - STATE_FP32=1 if cfg["STATE_FP32"] else 0, - REVERSE_BWD=reverse_bwd, - BIDI_MODE=1, - DOT_PRECISION=ctx.dot_prec, - BLOCK_D=BLOCK_D, - BLOCK_S=cfg["BLOCK_S"], - num_stages=cfg["num_stages"], - num_warps=nw, - ) - return # success - except Exception as e: - if "OutOfResources" in str(type(e).__name__) or "out of resource" in str(e).lower(): - nw = nw // 2 - if nw >= 1: - continue - raise RuntimeError( - "FusedBiGDN backward: Triton kernel OutOfResources at all warp counts " - f"(8, 4, 2, 1). Most recent error: {e}" - ) from e - else: - raise - - _run_triton_bwd(sv_fwd, sz_fwd, svc_fwd, szc_fwd, dqkv_fwd, dbeta_fwd, ddecay_fwd, 0, beta, decay) - del sv_fwd, sz_fwd, svc_fwd, szc_fwd - - # ---- Backward for reversed direction (replays against masked beta/decay if any) ---- - # Allocate kernel-output gradients with the exact shape the kernel - # writes — these always match the input ``beta_for_bwd_dir`` / - # ``decay_for_bwd_dir`` shapes (override or fall-back). - dqkv_bwd = torch.zeros_like(qkv) - dbeta_bwd_kernel = torch.zeros_like(beta_for_bwd_dir) - ddecay_bwd_kernel = torch.zeros_like(decay_for_bwd_dir) - - _run_triton_bwd( - sv_bwd, - sz_bwd, - svc_bwd, - szc_bwd, - dqkv_bwd, - dbeta_bwd_kernel, - ddecay_bwd_kernel, - 1, - beta_for_bwd_dir, - decay_for_bwd_dir, - ) - del sv_bwd, sz_bwd, svc_bwd, szc_bwd - del qkv_normed, dnum, dden - - # Q/K/V gradient is always summed: qkv is shared by both directions. - dqkv_fwd += dqkv_bwd - del dqkv_bwd - - # Beta gradient: route depends on whether the caller supplied an - # override. With override -> keep separate (so autograd routes the - # reverse-direction grad through the caller's clone+mask op). - # Without override -> sum into the forward-direction grad - # (legacy behaviour, bit-identical to pre-extension code). - if has_beta_bwd: - dbeta = dbeta_fwd - dbeta_bwd_out: torch.Tensor | None = dbeta_bwd_kernel - else: - dbeta_fwd += dbeta_bwd_kernel - dbeta = dbeta_fwd - dbeta_bwd_out = None - del dbeta_bwd_kernel - - # Decay gradient: same routing logic, independent of beta override. - if has_decay_bwd: - ddecay = ddecay_fwd - ddecay_bwd_out: torch.Tensor | None = ddecay_bwd_kernel - else: - ddecay_fwd += ddecay_bwd_kernel - ddecay = ddecay_fwd - ddecay_bwd_out = None - del ddecay_bwd_kernel - - dqkv_normed = dqkv_fwd - - # ---- Full-channel RMSNorm backward for Q and K ---- - # y = x * inv_rms * w -> dL/dx = inv_rms*w*dL/dy - inv_rms^3/C * x * sum(w*dL/dy*x) - # Process Q and K sequentially to reduce peak fp32 memory. - - # Q norm backward. - q_irms = q_inv_rms[:, :, None, None] - dq_normed = dqkv_normed[:, :, 0].float() - gw_q = dq_normed * q_nw_hd[None, None] - dq_nw = (dq_normed * q_raw * q_irms).sum(dim=(0, 1)).reshape(-1) - corr_q = (gw_q * q_raw).sum(dim=(-2, -1), keepdim=True) - dqkv_normed[:, :, 0] = (q_irms * gw_q - (q_irms**3) / C * q_raw * corr_q).to(qkv.dtype) - del dq_normed, gw_q, corr_q, q_raw, q_irms - - # K norm backward. - k_irms = k_inv_rms[:, :, None, None] - dk_normed = dqkv_normed[:, :, 1].float() - gw_k = dk_normed * k_nw_hd[None, None] - dk_nw = (dk_normed * k_raw * k_irms).sum(dim=(0, 1)).reshape(-1) - corr_k = (gw_k * k_raw).sum(dim=(-2, -1), keepdim=True) - dqkv_normed[:, :, 1] = (k_irms * gw_k - (k_irms**3) / C * k_raw * corr_k).to(qkv.dtype) - del dk_normed, gw_k, corr_k, k_raw, k_irms - - return ( - dqkv_normed, - dbeta, - ddecay, - dq_nw.to(q_norm_weight.dtype), - dk_nw.to(k_norm_weight.dtype), - None, # rope_cos - None, # rope_sin - None, # F - None, # S - None, # k_scale - None, # norm_eps - None, # eps - dbeta_bwd_out, # beta_bwd - ddecay_bwd_out, # decay_bwd - ) - - -def fused_bigdn_forward_with_grad( - qkv: torch.Tensor, - beta: torch.Tensor, - decay: torch.Tensor, - q_norm_weight: torch.Tensor | None, - k_norm_weight: torch.Tensor | None, - rope_cos: torch.Tensor, - rope_sin: torch.Tensor, - F: int, - S: int, - k_scale: float = 1.0, - norm_eps: float = 1e-5, - eps: float = 1e-6, - beta_bwd: torch.Tensor | None = None, - decay_bwd: torch.Tensor | None = None, -) -> torch.Tensor: - """Bidirectional fused BiGDN with autograd support (full-channel RMSNorm). - - Unlike ``fused_bigdn_func`` (which expects pre-computed ``q_inv_rms`` / - ``k_inv_rms``), this wrapper computes the full-channel inv-RMS in Python - so the norm backward can flow through the autograd graph naturally. - - Chunk-causal masking (optional): - Pass ``beta_bwd`` and/or ``decay_bwd`` to override the beta/decay - tensors used by the **reverse-direction** kernel only. These are - typically built by the caller as ``beta.clone()`` / ``decay.clone()`` - with interior chunk-boundary frames zeroed, so the anti-causal scan - resets state at chunk boundaries while the causal scan keeps full - context. The reverse-direction beta/decay gradients are routed - back through the ``beta_bwd`` / ``decay_bwd`` slots (instead of - being summed into the forward-direction grad), which lets autograd - flow the reverse-direction gradient through the caller's - ``clone() + index = 0`` masking op. - - When ``beta_bwd`` / ``decay_bwd`` is ``None`` (default), behaviour - is bit-identical to the pre-extension full-sequence-bidirectional - path: the reverse-direction kernel uses the unmasked ``beta`` / - ``decay`` and its kernel-emitted gradient is summed into the - forward-direction gradient before being returned. - """ - return FusedBiGDNFunction.apply( - qkv, - beta, - decay, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - F, - S, - k_scale, - norm_eps, - eps, - beta_bwd, - decay_bwd, - ) diff --git a/integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py b/integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py deleted file mode 100644 index 715231860..000000000 --- a/integrations/sana/sana_wm/ops/fused_gdn_chunkwise.py +++ /dev/null @@ -1,2269 +0,0 @@ -# Copyright 2024 NVIDIA CORPORATION & AFFILIATES -# -# 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. -# -# SPDX-License-Identifier: Apache-2.0 - -""" -Fused GDN — Chunkwise-parallel forward (v2). - -V2 changes vs v1: - 1. Phase A is split into TWO kernels along the GDN data streams (KV and Z; - these are the two gating sub-paths within the GDN block, not CUDA - streams — both kernels are launched on the same CUDA stream): - _phase_a_kv_kernel: P_kv (with K_rot) + A (with K_rot, V) — uses RoPE - _phase_a_z_kernel : P_z (with K) + B (with K) — no RoPE - Z block is genuinely lighter (no V/Cos/Sin loads, no K_pair flip). - On H100 this enables 2 blocks/SM resident → multi-tenancy / latency hiding. - 2. Phase A stores (I - P_kv) and (I - P_z) instead of P_kv/P_z. Phase B then - uses these directly: M = g · (I-P_kv)·M + A_f. The MMA `(I-P_kv) @ M` folds - the identity-add into the matmul (no separate M-PM elementwise pass). - -BiGDN inference path: QK_NORM=1, USE_PRECOMPUTED_RMS=1, SAVE_STATE=0. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -import torch -import triton -import triton.language as tl - -_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) - full_M = lambda: torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype) - full_z = lambda: 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).""" - from .fused_gdn import _resolve_launch_config - - _, 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. - """ - 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) diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index 3ad510d00..f67e76154 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -24,33 +24,12 @@ from collections.abc import Callable from dataclasses import dataclass import math -import os import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor -try: - from .ops.fused_cam_gdn import ( - _prepare_ucpe_rope_tables as _fused_prepare_ucpe_rope_tables, - cam_prep_func as _fused_cam_prep_func, - ) - from .ops.fused_gdn import ( - fused_bigdn_func as _fused_bigdn_func, - fused_qk_inv_rms as _fused_qk_inv_rms, - prepare_rope_tables as _fused_prepare_rope_tables, - ) - from .ops.fused_gdn_chunkwise import ( - cam_scan_bidi_chunkwise as _fused_cam_scan_bidi_chunkwise, - ) -except ImportError: - _fused_bigdn_func = None - _fused_cam_prep_func = None - _fused_cam_scan_bidi_chunkwise = None - _fused_prepare_rope_tables = None - _fused_prepare_ucpe_rope_tables = None - @dataclass(frozen=True) class SanaWMStage1Spec: @@ -100,8 +79,6 @@ class _CameraProjectionCache: proj_q: Tensor proj_kv: Tensor rope_cam: Tensor | None - rope_cos: Tensor - rope_sin: Tensor class RMSNorm(nn.Module): @@ -472,14 +449,6 @@ def _forward_gdn_main( rotary_emb: Tensor | None, precomputed_gates: tuple[Tensor, Tensor], ) -> Tensor: - if _use_fused_gdn(x): - return self._forward_fused_gdn_main( - x, - HW=HW, - rotary_emb=rotary_emb, - precomputed_gates=precomputed_gates, - ) - batch, tokens, channels = x.shape qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim) if hasattr(self, "conv_k"): @@ -519,57 +488,6 @@ def _forward_gdn_main( ) return out.reshape(batch, tokens, channels).to(dtype=x.dtype) - def _forward_fused_gdn_main( - self, - x: Tensor, - *, - HW: tuple[int, int, int], - rotary_emb: Tensor | None, - precomputed_gates: tuple[Tensor, Tensor], - ) -> Tensor: - if ( - _fused_bigdn_func is None - or _fused_qk_inv_rms is None - or _fused_prepare_rope_tables is None - ): - raise RuntimeError("fused SANA-WM GDN kernels are not available.") - - batch, tokens, channels = x.shape - frames, height, width = HW - spatial = height * width - if tokens != frames * spatial: - raise ValueError(f"tokens={tokens} != T*H*W={frames * spatial}") - qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim).contiguous() - if hasattr(self, "conv_k"): - k_raw = qkv[:, :, 1].contiguous().reshape(batch, tokens, channels) - k_conv = _apply_bidirectional_temporal_conv(k_raw, self.conv_k, HW) - qkv[:, :, 1].copy_(k_conv.reshape(batch, tokens, self.heads, self.dim)) - - beta, decay = precomputed_gates - q_inv_rms, k_inv_rms = _fused_qk_inv_rms(qkv, eps=self.q_norm.eps) - rope_cos, rope_sin = _fused_prepare_rope_tables( - rotary_emb, - tokens, - self.dim, - x.device, - ) - out = _fused_bigdn_func( - qkv, - q_inv_rms, - k_inv_rms, - self.q_norm.weight.float().contiguous(), - self.k_norm.weight.float().contiguous(), - rope_cos, - rope_sin, - beta.contiguous(), - decay.contiguous(), - F=frames, - S=spatial, - k_scale=_gdn_key_scale(self.dim, HW), - eps=self.eps, - ) - return out.reshape(batch, tokens, channels).to(dtype=x.dtype) - def _forward_softmax_main( self, x: Tensor, @@ -618,16 +536,6 @@ def _forward_gdn_camera( camera_cache: _CameraProjectionCache | None, precomputed_gates: tuple[Tensor, Tensor], ) -> Tensor: - if _use_fused_gdn(x): - return self._forward_fused_gdn_camera( - x, - HW=HW, - rotary_emb=rotary_emb, - camera_conditions=camera_conditions, - camera_cache=camera_cache, - precomputed_gates=precomputed_gates, - ) - q_cam, k_cam, v_cam = _camera_qkv(self, x, HW) q_trans, k_trans, v_trans, inflation_sq, output_projector = _prepare_ucpe_qkv( q_cam, @@ -660,76 +568,6 @@ def _forward_gdn_camera( out = output_projector(out) return out.reshape(x.shape[0], x.shape[1], -1).to(dtype=x.dtype) - def _forward_fused_gdn_camera( - self, - x: Tensor, - *, - HW: tuple[int, int, int], - rotary_emb: Tensor | None, - camera_conditions: Tensor, - camera_cache: _CameraProjectionCache | None, - precomputed_gates: tuple[Tensor, Tensor], - ) -> Tensor: - if ( - _fused_cam_prep_func is None - or _fused_cam_scan_bidi_chunkwise is None - or _fused_prepare_ucpe_rope_tables is None - ): - raise RuntimeError("fused SANA-WM camera GDN kernels are not available.") - - q_raw, k_raw, v_raw = _camera_qkv(self, x, HW) - batch, tokens, heads, dim = q_raw.shape - frames, height, width = HW - spatial = height * width - if camera_cache is None: - camera_cache = _prepare_camera_projection_cache( - camera_conditions, - HW=HW, - rotary_emb=rotary_emb, - head_dim=dim, - ) - if camera_cache.rope_cam is None: - rope_cos = camera_cache.rope_cos - rope_sin = camera_cache.rope_sin - else: - rope_cos = camera_cache.rope_cos - rope_sin = camera_cache.rope_sin - - q_trans, k_trans, v_trans, inflation_sq = _fused_cam_prep_func( - q_raw.contiguous(), - k_raw.contiguous(), - v_raw.contiguous(), - q_norm_weight=self.q_norm_cam.weight.float().contiguous(), - k_norm_weight=self.k_norm_cam.weight.float().contiguous(), - proj_q=camera_cache.proj_q, - proj_kv=camera_cache.proj_kv, - rope_cos=rope_cos, - rope_sin=rope_sin, - k_scale=_gdn_key_scale(dim, HW), - norm_eps=self.q_norm_cam.eps, - ) - - beta, decay = precomputed_gates - frame_inflation = inflation_sq.reshape(batch, heads, frames, spatial).mean( - dim=-1, - ) - beta = beta / frame_inflation.unsqueeze(-1).clamp_min(1.0) - out = _fused_cam_scan_bidi_chunkwise( - q_trans.float().contiguous(), - k_trans.float().contiguous(), - v_trans.float().contiguous(), - beta.float().contiguous(), - decay.float().contiguous(), - ) - out = out.permute(0, 3, 1, 2).contiguous() - out = _ucpe_transform_apply( - out.to(dtype=x.dtype), - camera_cache.proj, - camera_cache.rope_cam, - inverse_rope=True, - ) - return out.reshape(batch, tokens, heads * dim).to(dtype=x.dtype) - def _forward_softmax_camera( self, x: Tensor, @@ -1158,24 +996,6 @@ def _chunk_index_from_chunk_size( ) -_GDN_PATH_LOGGED = False - - -def _use_fused_gdn(x: Tensor) -> bool: - available = _fused_bigdn_func is not None - disabled_by_env = os.environ.get("SANA_WM_DISABLE_FUSED_GDN", "0") == "1" - use = x.is_cuda and available and not disabled_by_env - global _GDN_PATH_LOGGED - if not _GDN_PATH_LOGGED and x.is_cuda: - _GDN_PATH_LOGGED = True - print( - f"[sana-wm] GDN path={'fused' if use else 'eager'} " - f"(fused_available={available}, disabled_by_env={disabled_by_env})", - flush=True, - ) - return use - - def _prepare_camera_projection_cache( camera_conditions: Tensor, *, @@ -1197,34 +1017,12 @@ def _prepare_camera_projection_cache( proj_q = proj.transpose(-1, -2).contiguous() proj_kv = _invert_se3(proj).contiguous() rope_cam = _slice_rope_for_camera(rotary_emb, head_dim) - if rope_cam is not None and _fused_prepare_ucpe_rope_tables is not None: - rope_cos, rope_sin = _fused_prepare_ucpe_rope_tables( - rope_cam, - tokens, - head_dim // 2, - camera_conditions.device, - ) - else: - rope_cos = torch.ones( - tokens, - head_dim // 2, - device=camera_conditions.device, - dtype=torch.float32, - ) - rope_sin = torch.zeros( - tokens, - head_dim // 2, - device=camera_conditions.device, - dtype=torch.float32, - ) return _CameraProjectionCache( raymats=raymats, proj=proj, proj_q=proj_q, proj_kv=proj_kv, rope_cam=rope_cam, - rope_cos=rope_cos, - rope_sin=rope_sin, ) @@ -1457,79 +1255,47 @@ def _gdn_histories( .float() ) v = v.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4).float() - eye = torch.eye(dim, dtype=torch.float32, device=k.device).view(1, 1, dim, dim) - m = torch.zeros(batch, heads, dim, dim, dtype=torch.float32, device=k.device) - z = torch.zeros(batch, heads, dim, dtype=torch.float32, device=k.device) - m_hist = torch.empty( - batch, - heads, - frames, - dim, - dim, - dtype=torch.float32, - device=k.device, - ) - z_hist = ( - torch.empty(batch, heads, frames, dim, dtype=torch.float32, device=k.device) - if include_denominator - else None - ) + eye = torch.eye(dim, dtype=torch.float32, device=k.device).view(1, 1, 1, dim, dim) + + # The per-frame GDN transition is linear in the running state, so the + # spatial reductions that build each frame's transition are batched over all + # frames in one einsum, leaving only the (light) sequential state scan as a + # loop — run forward and reverse to form the bidirectional history. + p_kv = torch.einsum("bhfsd,bhfse->bhfde", k_rot, beta[..., None] * k_rot) + a_kv = torch.einsum("bhfsd,bhfse->bhfde", k_rot, beta[..., None] * v) + trans_m = decay[..., None, None] * (eye - p_kv) + m_state = torch.zeros(batch, heads, dim, dim, dtype=torch.float32, device=k.device) + m_forward: list[Tensor] = [] for frame in range(frames): - beta_f = beta[:, :, frame] - decay_f = decay[:, :, frame] - p_kv = torch.einsum( - "bhsd,bhse->bhde", - k_rot[:, :, frame], - beta_f[..., None] * k_rot[:, :, frame], - ) - a_val = torch.einsum( - "bhsd,bhse->bhde", - k_rot[:, :, frame], - beta_f[..., None] * v[:, :, frame], - ) - m = decay_f[..., None, None] * torch.einsum("bhde,bhef->bhdf", eye - p_kv, m) - m = m + a_val - m_hist[:, :, frame] = m - if include_denominator and z_hist is not None: - p_z = torch.einsum( - "bhsd,bhse->bhde", - k[:, :, frame], - beta_f[..., None] * k[:, :, frame], - ) - b_z = (beta_f[..., None] * k[:, :, frame]).sum(dim=2) - z = decay_f[..., None] * torch.einsum("bhde,bhe->bhd", eye - p_z, z) - z = z + b_z - z_hist[:, :, frame] = z - - m = torch.zeros_like(m) - z = torch.zeros_like(z) + m_state = torch.einsum("bhde,bhef->bhdf", trans_m[:, :, frame], m_state) + m_state = m_state + a_kv[:, :, frame] + m_forward.append(m_state) + m_hist = torch.stack(m_forward, dim=2) + m_state = torch.zeros(batch, heads, dim, dim, dtype=torch.float32, device=k.device) for src in range(frames - 1, 0, -1): - dst = src - 1 - beta_f = beta[:, :, src] - decay_f = decay[:, :, src] - p_kv = torch.einsum( - "bhsd,bhse->bhde", - k_rot[:, :, src], - beta_f[..., None] * k_rot[:, :, src], - ) - a_val = torch.einsum( - "bhsd,bhse->bhde", - k_rot[:, :, src], - beta_f[..., None] * v[:, :, src], - ) - m = decay_f[..., None, None] * torch.einsum("bhde,bhef->bhdf", eye - p_kv, m) - m = m + a_val - m_hist[:, :, dst] += m - if include_denominator and z_hist is not None: - p_z = torch.einsum( - "bhsd,bhse->bhde", - k[:, :, src], - beta_f[..., None] * k[:, :, src], - ) - b_z = (beta_f[..., None] * k[:, :, src]).sum(dim=2) - z = decay_f[..., None] * torch.einsum("bhde,bhe->bhd", eye - p_z, z) - z = z + b_z - z_hist[:, :, dst] += z + m_state = torch.einsum("bhde,bhef->bhdf", trans_m[:, :, src], m_state) + m_state = m_state + a_kv[:, :, src] + m_hist[:, :, src - 1] = m_hist[:, :, src - 1] + m_state + + z_hist = None + if include_denominator: + beta_k = beta[..., None] * k + p_z = torch.einsum("bhfsd,bhfse->bhfde", k, beta_k) + b_z = beta_k.sum(dim=3) + trans_z = decay[..., None, None] * (eye - p_z) + z_state = torch.zeros(batch, heads, dim, dtype=torch.float32, device=k.device) + z_forward: list[Tensor] = [] + for frame in range(frames): + z_state = torch.einsum("bhde,bhe->bhd", trans_z[:, :, frame], z_state) + z_state = z_state + b_z[:, :, frame] + z_forward.append(z_state) + z_hist = torch.stack(z_forward, dim=2) + z_state = torch.zeros(batch, heads, dim, dtype=torch.float32, device=k.device) + for src in range(frames - 1, 0, -1): + z_state = torch.einsum("bhde,bhe->bhd", trans_z[:, :, src], z_state) + z_state = z_state + b_z[:, :, src] + z_hist[:, :, src - 1] = z_hist[:, :, src - 1] + z_state + return m_hist, z_hist From 36d2d8aa2cc86ee5596bb2590ddbf513b4f1e610 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 28 Jul 2026 13:40:37 -0700 Subject: [PATCH 35/36] Optimizations --- integrations/sana/sana_wm/stage1_model.py | 220 +++++++++++++++++++- integrations/sana/sana_wm/transformer.py | 21 +- integrations/sana/tests/test_smoke.py | 69 ++++++ integrations/sana/tests/test_stage1_cuda.py | 79 +++++++ 4 files changed, 382 insertions(+), 7 deletions(-) create mode 100644 integrations/sana/tests/test_stage1_cuda.py diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index f67e76154..388eb656a 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -30,6 +30,13 @@ import torch.nn.functional as F from torch import Tensor +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - exercised in minimal CPU environments. + triton = None + tl = None + @dataclass(frozen=True) class SanaWMStage1Spec: @@ -782,10 +789,20 @@ def forward( x = self.x_embedder(x) x = x.permute(0, 2, 3, 4, 1).reshape(batch, frames * height * width, -1) - plucker_emb = None - if chunk_plucker is not None: - plucker_emb = self.plucker_embedder(chunk_plucker) - plucker_emb = plucker_emb.permute(0, 2, 3, 4, 1).reshape_as(x) + plucker_emb = kwargs.get("chunk_plucker_emb") + if isinstance(plucker_emb, Tensor): + if ( + plucker_emb.ndim != 3 + or plucker_emb.shape[1:] != x.shape[1:] + or plucker_emb.shape[0] not in (1, batch) + ): + raise ValueError( + "chunk_plucker_emb must have shape [1|B, T*H*W, hidden], " + f"got {tuple(plucker_emb.shape)} for x={tuple(x.shape)}." + ) + plucker_emb = plucker_emb.to(device=x.device, dtype=x.dtype) + elif chunk_plucker is not None: + plucker_emb = self.prepare_plucker_embedding(chunk_plucker) rotary_emb = _wan_rope_complex( self.spec.head_dim, @@ -811,7 +828,13 @@ def forward( else None ) block_kwargs = { - key: value for key, value in kwargs.items() if key != "camera_conditions" + key: value + for key, value in kwargs.items() + if key + not in { + "camera_conditions", + "chunk_plucker_emb", + } } y = self.y_embedder(y) @@ -850,6 +873,16 @@ def forward( x = self.final_layer(x, timestep_embed, frames=frames) return x.reshape(batch, frames, height, width, -1).permute(0, 4, 1, 2, 3) + def prepare_plucker_embedding(self, chunk_plucker: Tensor) -> Tensor: + """Project static Plucker conditioning into Stage-1 token space.""" + batch = chunk_plucker.shape[0] + plucker_emb = self.plucker_embedder(chunk_plucker) + return plucker_emb.permute(0, 2, 3, 4, 1).reshape( + batch, + -1, + self.spec.hidden_size, + ) + def _modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor: return x * (1 + scale) + shift @@ -1061,7 +1094,155 @@ def _compute_frame_gates( return beta, decay -def _apply_bidirectional_temporal_conv( +if triton is not None and tl is not None: + + @triton.jit + def _bidirectional_temporal_conv_kernel( + x_ptr, + weight_ptr, + out_ptr, + x_stride_b, + x_stride_t, + x_stride_c, + weight_stride_c, + weight_stride_k, + out_stride_b, + out_stride_t, + out_stride_c, + frames: tl.constexpr, + spatial: tl.constexpr, + channels, + kernel_size: tl.constexpr, + block_f: tl.constexpr, + block_s: tl.constexpr, + block_c: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + spatial_offsets = tl.program_id(1) * block_s + tl.arange(0, block_s) + channel_offsets = tl.program_id(2) * block_c + tl.arange(0, block_c) + frame_offsets = tl.arange(0, block_f) + + frame_mask = frame_offsets < frames + spatial_mask = spatial_offsets < spatial + channel_mask = channel_offsets < channels + element_mask = ( + frame_mask[:, None, None] + & spatial_mask[None, :, None] + & channel_mask[None, None, :] + ) + + token_offsets = frame_offsets[:, None, None] * spatial + spatial_offsets[ + None, + :, + None, + ] + x_offsets = ( + batch_idx * x_stride_b + + token_offsets * x_stride_t + + channel_offsets[None, None, :] * x_stride_c + ) + out_offsets = ( + batch_idx * out_stride_b + + token_offsets * out_stride_t + + channel_offsets[None, None, :] * out_stride_c + ) + + center = tl.load( + weight_ptr + + channel_offsets * weight_stride_c + + (kernel_size - 1) * weight_stride_k, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + value = tl.load(x_ptr + x_offsets, mask=element_mask, other=0.0).to( + tl.float32 + ) + acc = value * center[None, None, :] + + for offset in range(1, kernel_size): + coeff = tl.load( + weight_ptr + + channel_offsets * weight_stride_c + + (kernel_size - 1 - offset) * weight_stride_k, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + + prev_frame_offsets = frame_offsets - offset + prev_token_offsets = prev_frame_offsets[:, None, None] * spatial + ( + spatial_offsets[None, :, None] + ) + prev_mask = element_mask & (frame_offsets[:, None, None] >= offset) + prev_offsets = ( + batch_idx * x_stride_b + + prev_token_offsets * x_stride_t + + channel_offsets[None, None, :] * x_stride_c + ) + prev = tl.load(x_ptr + prev_offsets, mask=prev_mask, other=0.0).to( + tl.float32 + ) + + next_frame_offsets = frame_offsets + offset + next_token_offsets = next_frame_offsets[:, None, None] * spatial + ( + spatial_offsets[None, :, None] + ) + next_mask = element_mask & (next_frame_offsets[:, None, None] < frames) + next_offsets = ( + batch_idx * x_stride_b + + next_token_offsets * x_stride_t + + channel_offsets[None, None, :] * x_stride_c + ) + next_value = tl.load(x_ptr + next_offsets, mask=next_mask, other=0.0).to( + tl.float32 + ) + acc += (prev + next_value) * coeff[None, None, :] + + tl.store(out_ptr + out_offsets, acc, mask=element_mask) + +else: + _bidirectional_temporal_conv_kernel = None + + +def _apply_bidirectional_temporal_conv_fast( + x: Tensor, + conv: nn.Conv1d, + HW: tuple[int, int, int], +) -> Tensor: + batch, tokens, channels = x.shape + frames, height, width = HW + spatial = height * width + out = torch.empty_like(x) + assert triton is not None + assert _bidirectional_temporal_conv_kernel is not None + block_f = triton.next_power_of_2(frames) + block_s = 4 + block_c = 128 + grid = (batch, triton.cdiv(spatial, block_s), triton.cdiv(channels, block_c)) + weight = conv.weight[:, 0, :] + _bidirectional_temporal_conv_kernel[grid]( + x, + weight, + out, + x.stride(0), + x.stride(1), + x.stride(2), + weight.stride(0), + weight.stride(1), + out.stride(0), + out.stride(1), + out.stride(2), + frames, + spatial, + channels, + int(conv.kernel_size[0]), + block_f, + block_s, + block_c, + ) + return out + + +def _apply_bidirectional_temporal_conv_eager( x: Tensor, conv: nn.Conv1d, HW: tuple[int, int, int], @@ -1088,6 +1269,33 @@ def _apply_bidirectional_temporal_conv( ) +def _apply_bidirectional_temporal_conv( + x: Tensor, + conv: nn.Conv1d, + HW: tuple[int, int, int], +) -> Tensor: + if ( + x.is_cuda + and conv.weight.is_cuda + and conv.weight.dtype == x.dtype + and conv.bias is None + and conv.in_channels == x.shape[-1] + and conv.out_channels == x.shape[-1] + and conv.groups == x.shape[-1] + and conv.stride == (1,) + and conv.padding == (0,) + and conv.dilation == (1,) + and not torch.is_grad_enabled() + and triton is not None + and _bidirectional_temporal_conv_kernel is not None + ): + batch, tokens, _channels = x.shape + frames, height, width = HW + if tokens == frames * height * width: + return _apply_bidirectional_temporal_conv_fast(x, conv, HW) + return _apply_bidirectional_temporal_conv_eager(x, conv, HW) + + def _causal_depthwise_conv(x: Tensor, conv: nn.Conv1d) -> Tensor: kernel = int(conv.kernel_size[0]) padded = F.pad(x, (kernel - 1, 0)) diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py index dd744772c..fb4c61092 100644 --- a/integrations/sana/sana_wm/transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -314,6 +314,7 @@ def _predict_conditioned( timestep: Tensor, conditioning: SanaWMStage1Conditioning, ) -> Tensor: + self._prepare_static_model_kwargs(conditioning) model_timestep = _conditioned_frame_timestep( noisy_latent=noisy_latent, timestep=timestep, @@ -340,6 +341,24 @@ def _predict_conditioned( flow = _cfg_guidance(noise_pred, conditioning.cfg_scale) return _zero_conditioned_frame_flow(flow, conditioning) + def _prepare_static_model_kwargs( + self, + conditioning: SanaWMStage1Conditioning, + ) -> None: + """Cache model projections for conditioning that is constant per rollout.""" + model_kwargs = conditioning.model_kwargs + self._ensure_model() + chunk_plucker = model_kwargs.get("chunk_plucker") + prepare_plucker = getattr(self.model, "prepare_plucker_embedding", None) + if ( + "chunk_plucker_emb" not in model_kwargs + and isinstance(chunk_plucker, Tensor) + and callable(prepare_plucker) + ): + with torch.inference_mode(): + model_kwargs["chunk_plucker_emb"] = prepare_plucker(chunk_plucker) + del model_kwargs["chunk_plucker"] + def _predict_with_prompt( self, noisy_latent: Tensor, @@ -469,7 +488,7 @@ def _condition_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object return { key: value for key, value in model_kwargs.items() - if key != "negative_mask" + if key != "negative_mask" and not key.startswith("_") } diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 38c08f97d..7b8f23398 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -465,6 +465,75 @@ def forward( assert call["data_info"] == {"condition_frame_info": {0: 0.0}} +def test_transformer_predict_flow_caches_static_plucker_embedding() -> None: + """Avoid re-projecting rollout-constant Plucker conditioning each step.""" + + class DummyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.prepare_calls = 0 + self.forward_calls: list[dict[str, torch.Tensor]] = [] + + def prepare_plucker_embedding( + self, + chunk_plucker: torch.Tensor, + ) -> torch.Tensor: + self.prepare_calls += 1 + return chunk_plucker.flatten(2).transpose(1, 2) + + def forward( + self, + noisy_latent: torch.Tensor, + timestep: torch.Tensor, + prompt_embeds: torch.Tensor, + *, + chunk_plucker_emb: torch.Tensor, + **_kwargs: object, + ) -> torch.Tensor: + del timestep, prompt_embeds + assert "chunk_plucker" not in _kwargs + self.forward_calls.append({"chunk_plucker_emb": chunk_plucker_emb}) + return torch.ones_like(noisy_latent) + + transformer = SanaWMTransformerConfig().setup() + dummy_model = DummyModel() + transformer.model = dummy_model + transformer._model_built = True + chunk_plucker = torch.ones((1, 2, 2, 1, 1)) + conditioning = SanaWMStage1Conditioning( + condition=torch.ones((1, 1, 1)), + uncondition=None, + model_kwargs={ + "chunk_plucker": chunk_plucker, + "data_info": {"condition_frame_info": {0: 0.0}}, + }, + first_latent=torch.empty((1, 1, 1, 1, 1)), + latent_shape=(1, 1, 2, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + cache = SanaWMTransformerCache() + + for _ in range(2): + transformer.predict_flow( + noisy_latent=torch.zeros((1, 1, 2, 1, 1)), + timestep=torch.full((1, 1, 2), 1000.0), + cache=cache, + input=conditioning, + ) + + assert dummy_model.prepare_calls == 1 + assert "chunk_plucker" not in conditioning.model_kwargs + assert len(dummy_model.forward_calls) == 2 + cached = conditioning.model_kwargs["chunk_plucker_emb"] + assert isinstance(cached, torch.Tensor) + assert all( + call["chunk_plucker_emb"] is cached for call in dummy_model.forward_calls + ) + + def test_inference_config_loads_yaml( tmp_path: Path, ) -> None: diff --git a/integrations/sana/tests/test_stage1_cuda.py b/integrations/sana/tests/test_stage1_cuda.py new file mode 100644 index 000000000..a72c40ada --- /dev/null +++ b/integrations/sana/tests/test_stage1_cuda.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""CUDA tests for SANA-WM Stage-1 kernels.""" + +from __future__ import annotations + +import pytest +import torch + +import sana_wm.stage1_model as stage1 + +pytestmark = pytest.mark.ci_gpu + + +def _require_stage1_cuda() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for Stage-1 kernel tests") + if stage1.triton is None: + pytest.skip("Triton is required for the Stage-1 temporal conv CUDA kernel") + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bidirectional_temporal_conv_fast_path_matches_eager( + dtype: torch.dtype, +) -> None: + """Compare the CUDA temporal-conv fast path against the eager conv path.""" + _require_stage1_cuda() + torch.manual_seed(1234) + batch, frames, height, width, channels = 2, 5, 3, 4, 16 + x = torch.randn( + batch, + frames * height * width, + channels, + device="cuda", + dtype=dtype, + ) + conv = torch.nn.Conv1d( + channels, + channels, + kernel_size=4, + groups=channels, + bias=False, + device="cuda", + dtype=dtype, + ) + + reference = stage1._apply_bidirectional_temporal_conv_eager( + x, + conv, + (frames, height, width), + ) + with torch.inference_mode(): + actual = stage1._apply_bidirectional_temporal_conv( + x, + conv, + (frames, height, width), + ) + + assert actual.shape == x.shape + assert actual.dtype == dtype + torch.testing.assert_close( + actual.float(), + reference.float(), + rtol=2e-2, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) From 4c1d857e1ac6e4f3e09b81d5fb3b3066b04df3d6 Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Tue, 28 Jul 2026 18:11:59 -0700 Subject: [PATCH 36/36] Optimize SANA-WM Stage 1 kernels Add guarded clean-room Triton kernels for UCPE transforms, UCPE norm reductions, GLU activation multiply, and RMSNorm/ReLU head prep. Precompute static RoPE and camera projection tensors for singleton camera conditioning so diffusion steps avoid rebuilding those tensors. Correct Stage-1 main attention routing to use the GDN path consistently while preserving the hybrid camera branch behavior. Add CPU and CUDA coverage for the cache behavior, routing, and fused kernel fallbacks. Final BF16 benchmark with refiner enabled, one warmup and three measured runs on GB202: FlashDreams generation median 92.39s versus upstream 107.55s, with Stage-1 DiT median 69.13s versus upstream 75.30s. Signed-off-by: Aidan Foster --- integrations/sana/sana_wm/stage1_model.py | 842 ++++++++++++++++++-- integrations/sana/sana_wm/transformer.py | 32 +- integrations/sana/tests/test_smoke.py | 187 ++++- integrations/sana/tests/test_stage1_cuda.py | 201 +++++ 4 files changed, 1179 insertions(+), 83 deletions(-) diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py index 388eb656a..2f93cecce 100644 --- a/integrations/sana/sana_wm/stage1_model.py +++ b/integrations/sana/sana_wm/stage1_model.py @@ -24,6 +24,7 @@ from collections.abc import Callable from dataclasses import dataclass import math +import os import torch import torch.nn as nn @@ -79,6 +80,15 @@ def block_uses_gdn(self, index: int) -> bool: """Architecture spec for the public SANA-WM bidirectional Stage-1 checkpoint.""" +def _env_flag_enabled(name: str) -> bool: + return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"} + + +_DISABLE_UCPE_FAST = _env_flag_enabled("SANA_WM_STAGE1_DISABLE_UCPE_FAST") +_DISABLE_GLU_FAST = _env_flag_enabled("SANA_WM_STAGE1_DISABLE_GLU_FAST") +_DISABLE_RMS_RELU_FAST = _env_flag_enabled("SANA_WM_STAGE1_DISABLE_RMS_RELU_FAST") + + @dataclass(frozen=True) class _CameraProjectionCache: raymats: Tensor @@ -263,8 +273,9 @@ def forward(self, x: Tensor, *, frames: int, height: int, width: int) -> Tensor: x_2d = F.silu(self.inverted_conv(x_2d), inplace=True) x_2d = self.depth_conv(x_2d) value, gate = x_2d.chunk(2, dim=1) - gate = F.silu(gate, inplace=not torch.is_grad_enabled()) - x_2d = self.point_conv(value * gate) + x_2d = self.point_conv( + _silu_multiply(value, gate, inplace=not torch.is_grad_enabled()) + ) x_time = x_2d.view(batch, frames, channels, height * width).permute(0, 2, 1, 3) x_time = x_time + self.t_conv(x_time) @@ -392,25 +403,13 @@ def forward( f"channels={channels} != heads*dim={self.heads * self.dim}" ) - precomputed_gates = ( - self._compute_frame_gates(x, HW) if self.use_gdn_convs else None + precomputed_gates = self._compute_frame_gates(x, HW) + main_raw = self._forward_gdn_main( + x, + HW=HW, + rotary_emb=rotary_emb, + precomputed_gates=precomputed_gates, ) - if self.use_gdn_convs: - main_raw = self._forward_gdn_main( - x, - HW=HW, - rotary_emb=rotary_emb, - precomputed_gates=precomputed_gates, - ) - else: - main_raw = self._forward_softmax_main( - x, - HW=HW, - rotary_emb=rotary_emb, - chunk_size=kwargs.get("chunk_size"), - chunk_split_strategy=str(kwargs.get("chunk_split_strategy", "uniform")), - chunk_index=kwargs.get("chunk_index"), - ) cam_contrib: Tensor | int = 0 if camera_conditions is not None: @@ -463,21 +462,18 @@ def _forward_gdn_main( k_conv = _apply_bidirectional_temporal_conv(k_raw, self.conv_k, HW) qkv[:, :, 1] = k_conv.reshape(batch, tokens, self.heads, self.dim) - q = self.q_norm(qkv[:, :, 0].reshape(batch, tokens, channels)).reshape( - batch, - tokens, - self.heads, - self.dim, + q = _rmsnorm_relu_heads( + qkv[:, :, 0], + self.q_norm.weight, + self.q_norm.eps, ) - k = self.k_norm(qkv[:, :, 1].reshape(batch, tokens, channels)).reshape( - batch, - tokens, - self.heads, - self.dim, + k = _rmsnorm_relu_heads( + qkv[:, :, 1], + self.k_norm.weight, + self.k_norm.eps, + scale=_gdn_key_scale(self.dim, HW), ) v = qkv[:, :, 2] - q = F.relu(q.float()) - k = F.relu(k.float()) * _gdn_key_scale(self.dim, HW) v = v.float() q_rot = _apply_complex_rope(q, rotary_emb) k_rot = _apply_complex_rope(k, rotary_emb) @@ -804,28 +800,44 @@ def forward( elif chunk_plucker is not None: plucker_emb = self.prepare_plucker_embedding(chunk_plucker) - rotary_emb = _wan_rope_complex( - self.spec.head_dim, - frames, - height, - width, - x.device, - ) - camera_conditions = kwargs.get("camera_conditions") + rotary_emb = kwargs.get("rotary_emb") + if isinstance(rotary_emb, Tensor): + rotary_emb = rotary_emb.to(device=x.device) + else: + rotary_emb = _wan_rope_complex( + self.spec.head_dim, + frames, + height, + width, + x.device, + ) + raw_camera_conditions = kwargs.get("camera_conditions") camera_conditions = ( - camera_conditions.to(device=x.device, dtype=x.dtype) - if isinstance(camera_conditions, Tensor) + raw_camera_conditions.to(device=x.device, dtype=x.dtype) + if isinstance(raw_camera_conditions, Tensor) else None ) + camera_cache_value = kwargs.get("camera_cache") camera_cache = ( - _prepare_camera_projection_cache( - camera_conditions, - HW=(frames, height, width), - rotary_emb=rotary_emb, - head_dim=self.spec.head_dim, + _camera_projection_cache_to( + camera_cache_value, + device=x.device, + dtype=x.dtype, ) - if camera_conditions is not None - else None + if isinstance(camera_cache_value, _CameraProjectionCache) + else ( + _prepare_camera_projection_cache( + camera_conditions, + HW=(frames, height, width), + rotary_emb=rotary_emb, + head_dim=self.spec.head_dim, + ) + if camera_conditions is not None + else None + ) + ) + camera_signal = ( + camera_cache.raymats if camera_cache is not None else camera_conditions ) block_kwargs = { key: value @@ -834,6 +846,8 @@ def forward( not in { "camera_conditions", "chunk_plucker_emb", + "camera_cache", + "rotary_emb", } } @@ -863,7 +877,7 @@ def forward( plucker_emb=plucker_emb, HW=(frames, height, width), rotary_emb=rotary_emb, - camera_conditions=camera_conditions, + camera_conditions=camera_signal, camera_cache=camera_cache, chunk_size=self.spec.chunk_size, chunk_split_strategy=self.spec.chunk_split_strategy, @@ -883,6 +897,32 @@ def prepare_plucker_embedding(self, chunk_plucker: Tensor) -> Tensor: self.spec.hidden_size, ) + def prepare_camera_projection_cache( + self, + camera_conditions: Tensor, + *, + frames: int, + height: int, + width: int, + ) -> tuple[Tensor, _CameraProjectionCache]: + """Precompute static RoPE and camera projection tensors for a rollout.""" + param = next(self.parameters()) + camera_conditions = camera_conditions.to(device=param.device, dtype=param.dtype) + rotary_emb = _wan_rope_complex( + self.spec.head_dim, + frames, + height, + width, + param.device, + ) + camera_cache = _prepare_camera_projection_cache( + camera_conditions, + HW=(frames, height, width), + rotary_emb=rotary_emb, + head_dim=self.spec.head_dim, + ) + return rotary_emb, camera_cache + def _modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor: return x * (1 + scale) + shift @@ -1059,6 +1099,35 @@ def _prepare_camera_projection_cache( ) +def _camera_projection_cache_to( + cache: _CameraProjectionCache, + *, + device: torch.device, + dtype: torch.dtype, +) -> _CameraProjectionCache: + if ( + cache.raymats.device == device + and cache.proj.device == device + and cache.proj_q.device == device + and cache.proj_kv.device == device + and cache.raymats.dtype == dtype + and cache.proj.dtype == dtype + and cache.proj_q.dtype == dtype + and cache.proj_kv.dtype == dtype + and (cache.rope_cam is None or cache.rope_cam.device == device) + ): + return cache + return _CameraProjectionCache( + raymats=cache.raymats.to(device=device, dtype=dtype), + proj=cache.proj.to(device=device, dtype=dtype), + proj_q=cache.proj_q.to(device=device, dtype=dtype), + proj_kv=cache.proj_kv.to(device=device, dtype=dtype), + rope_cam=( + cache.rope_cam.to(device=device) if cache.rope_cam is not None else None + ), + ) + + @torch.compile def _apply_output_gate( out: Tensor, @@ -1199,8 +1268,359 @@ def _bidirectional_temporal_conv_kernel( tl.store(out_ptr + out_offsets, acc, mask=element_mask) + @triton.jit + def _ucpe_first_half_kernel( + x_ptr, + matrix_ptr, + out_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + matrix_stride_b, + matrix_stride_n, + matrix_stride_i, + matrix_stride_j, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + tokens: tl.constexpr, + groups4: tl.constexpr, + matrix_batch: tl.constexpr, + block_n: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + packed_idx = tl.program_id(2) + group_idx = packed_idx % groups4 + token_block = packed_idx // groups4 + token_offsets = token_block * block_n + tl.arange(0, block_n) + token_mask = token_offsets < tokens + + matrix_batch_idx = 0 if matrix_batch == 1 else batch_idx + col0 = group_idx * 4 + col1 = col0 + 1 + col2 = col0 + 2 + col3 = col0 + 3 + + x_base = ( + batch_idx * x_stride_b + + token_offsets * x_stride_n + + head_idx * x_stride_h + ) + x0 = tl.load( + x_ptr + x_base + col0 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + x1 = tl.load( + x_ptr + x_base + col1 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + x2 = tl.load( + x_ptr + x_base + col2 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + x3 = tl.load( + x_ptr + x_base + col3 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + + matrix_base = ( + matrix_batch_idx * matrix_stride_b + token_offsets * matrix_stride_n + ) + out_base = ( + batch_idx * out_stride_b + + token_offsets * out_stride_n + + head_idx * out_stride_h + ) + for row in range(4): + row_base = matrix_base + row * matrix_stride_i + m0 = tl.load( + matrix_ptr + row_base + 0 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + m1 = tl.load( + matrix_ptr + row_base + 1 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + m2 = tl.load( + matrix_ptr + row_base + 2 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + m3 = tl.load( + matrix_ptr + row_base + 3 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + value = (x0 * m0 + x1 * m1) + (x2 * m2 + x3 * m3) + tl.store( + out_ptr + out_base + (col0 + row) * out_stride_d, + value, + mask=token_mask, + ) + + @triton.jit + def _ucpe_rope_second_half_kernel( + x_ptr, + rope_ptr, + out_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + rope_stride_n, + rope_stride_p, + rope_stride_ri, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + tokens: tl.constexpr, + half: tl.constexpr, + rope_pairs: tl.constexpr, + inverse_rope: tl.constexpr, + block_n: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + packed_idx = tl.program_id(2) + pair_idx = packed_idx % rope_pairs + token_block = packed_idx // rope_pairs + token_offsets = token_block * block_n + tl.arange(0, block_n) + token_mask = token_offsets < tokens + + real_d = half + pair_idx * 2 + imag_d = real_d + 1 + x_base = ( + batch_idx * x_stride_b + + token_offsets * x_stride_n + + head_idx * x_stride_h + ) + real = tl.load( + x_ptr + x_base + real_d * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + imag = tl.load( + x_ptr + x_base + imag_d * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + rope_base = token_offsets * rope_stride_n + pair_idx * rope_stride_p + rope_real = tl.load( + rope_ptr + rope_base + 0 * rope_stride_ri, + mask=token_mask, + other=0.0, + ).to(tl.float32) + rope_imag = tl.load( + rope_ptr + rope_base + 1 * rope_stride_ri, + mask=token_mask, + other=0.0, + ).to(tl.float32) + if inverse_rope: + rope_imag = -rope_imag + + out_base = ( + batch_idx * out_stride_b + + token_offsets * out_stride_n + + head_idx * out_stride_h + ) + tl.store( + out_ptr + out_base + real_d * out_stride_d, + real * rope_real - imag * rope_imag, + mask=token_mask, + ) + tl.store( + out_ptr + out_base + imag_d * out_stride_d, + real * rope_imag + imag * rope_real, + mask=token_mask, + ) + + @triton.jit + def _ucpe_norms_kernel( + x_ptr, + out_ptr, + pre_ptr, + post_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + pre_stride_b, + pre_stride_h, + pre_stride_n, + post_stride_b, + post_stride_h, + post_stride_n, + tokens: tl.constexpr, + dim: tl.constexpr, + block_n: tl.constexpr, + block_d: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + token_block = tl.program_id(2) + token_offsets = token_block * block_n + tl.arange(0, block_n) + dim_offsets = tl.arange(0, block_d) + mask = (token_offsets[:, None] < tokens) & (dim_offsets[None, :] < dim) + + x_offsets = ( + batch_idx * x_stride_b + + token_offsets[:, None] * x_stride_n + + head_idx * x_stride_h + + dim_offsets[None, :] * x_stride_d + ) + out_offsets = ( + batch_idx * out_stride_b + + token_offsets[:, None] * out_stride_n + + head_idx * out_stride_h + + dim_offsets[None, :] * out_stride_d + ) + x_values = tl.load(x_ptr + x_offsets, mask=mask, other=0.0).to(tl.float32) + out_values = tl.load(out_ptr + out_offsets, mask=mask, other=0.0).to( + tl.float32 + ) + pre = tl.sum(x_values * x_values, axis=1) + post = tl.sum(out_values * out_values, axis=1) + token_mask = token_offsets < tokens + pre_offsets = ( + batch_idx * pre_stride_b + + head_idx * pre_stride_h + + token_offsets * pre_stride_n + ) + post_offsets = ( + batch_idx * post_stride_b + + head_idx * post_stride_h + + token_offsets * post_stride_n + ) + tl.store(pre_ptr + pre_offsets, pre, mask=token_mask) + tl.store(post_ptr + post_offsets, post, mask=token_mask) + + @triton.jit + def _silu_multiply_kernel( + value_ptr, + gate_ptr, + out_ptr, + value_stride_n, + value_stride_c, + value_stride_h, + value_stride_w, + gate_stride_n, + gate_stride_c, + gate_stride_h, + gate_stride_w, + channels: tl.constexpr, + height: tl.constexpr, + width: tl.constexpr, + total: tl.constexpr, + block: tl.constexpr, + ) -> None: + offsets = tl.program_id(0) * block + tl.arange(0, block) + mask = offsets < total + channel_offsets = offsets % channels + spatial_offsets = offsets // channels + width_offsets = spatial_offsets % width + height_offsets = (spatial_offsets // width) % height + batch_offsets = spatial_offsets // (height * width) + + value_offsets = ( + batch_offsets * value_stride_n + + channel_offsets * value_stride_c + + height_offsets * value_stride_h + + width_offsets * value_stride_w + ) + gate_offsets = ( + batch_offsets * gate_stride_n + + channel_offsets * gate_stride_c + + height_offsets * gate_stride_h + + width_offsets * gate_stride_w + ) + value = tl.load(value_ptr + value_offsets, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptr + gate_offsets, mask=mask, other=0.0).to(tl.float32) + silu_gate = gate / (1.0 + tl.exp(-gate)) + tl.store(out_ptr + offsets, value * silu_gate, mask=mask) + + @triton.jit + def _rmsnorm_relu_heads_kernel( + x_ptr, + weight_ptr, + out_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + tokens: tl.constexpr, + heads: tl.constexpr, + dim: tl.constexpr, + channels: tl.constexpr, + eps: tl.constexpr, + scale: tl.constexpr, + block_c: tl.constexpr, + ) -> None: + row = tl.program_id(0) + batch_idx = row // tokens + token_idx = row - batch_idx * tokens + channel_offsets = tl.arange(0, block_c) + mask = channel_offsets < channels + head_offsets = channel_offsets // dim + dim_offsets = channel_offsets - head_offsets * dim + + x_offsets = ( + batch_idx * x_stride_b + + token_idx * x_stride_n + + head_offsets * x_stride_h + + dim_offsets * x_stride_d + ) + values = tl.load(x_ptr + x_offsets, mask=mask, other=0.0).to(tl.float32) + mean_sq = tl.sum(values * values, axis=0) / channels + inv_rms = tl.rsqrt(mean_sq + eps) + weights = tl.load(weight_ptr + channel_offsets, mask=mask, other=0.0).to( + tl.float32 + ) + out_values = tl.maximum(values * inv_rms * weights * scale, 0.0) + out_offsets = ( + batch_idx * out_stride_b + + token_idx * out_stride_n + + head_offsets * out_stride_h + + dim_offsets * out_stride_d + ) + tl.store(out_ptr + out_offsets, out_values, mask=mask) + else: _bidirectional_temporal_conv_kernel = None + _ucpe_first_half_kernel = None + _ucpe_rope_second_half_kernel = None + _ucpe_norms_kernel = None + _silu_multiply_kernel = None + _rmsnorm_relu_heads_kernel = None def _apply_bidirectional_temporal_conv_fast( @@ -1553,13 +1973,13 @@ def _prepare_ucpe_qkv( norm_eps: float, ) -> tuple[Tensor, Tensor, Tensor, Tensor, Callable[[Tensor], Tensor]]: batch, tokens, heads, dim = q_raw.shape - q_inv = _inv_rms(q_raw, norm_eps) - k_inv = _inv_rms(k_raw, norm_eps) - q_weight = q_norm_weight.float().view(heads, dim) - k_weight = k_norm_weight.float().view(heads, dim) - q_norm = F.relu(q_raw.float() * q_inv[:, :, None, None] * q_weight[None, None]) - k_norm = F.relu(k_raw.float() * k_inv[:, :, None, None] * k_weight[None, None]) - k_norm = k_norm * _gdn_key_scale(dim, HW) + q_norm = _rmsnorm_relu_heads(q_raw, q_norm_weight, norm_eps) + k_norm = _rmsnorm_relu_heads( + k_raw, + k_norm_weight, + norm_eps, + scale=_gdn_key_scale(dim, HW), + ) v_float = v_raw.float() if camera_cache is None: camera_cache = _prepare_camera_projection_cache( @@ -1664,6 +2084,75 @@ def _inv_rms(x: Tensor, eps: float) -> Tensor: return torch.rsqrt(x.float().pow(2).sum(dim=(-1, -2)) / channels + eps) +def _rmsnorm_relu_heads( + x: Tensor, + weight: Tensor, + eps: float, + *, + scale: float = 1.0, +) -> Tensor: + fast = _rmsnorm_relu_heads_fast(x, weight, eps, scale=scale) + if fast is not None: + return fast + batch, tokens, heads, dim = x.shape + inv_rms = _inv_rms(x, eps) + norm_weight = weight.float().view(heads, dim) + out = F.relu(x.float() * inv_rms[:, :, None, None] * norm_weight[None, None]) + if scale != 1.0: + out = out * scale + return out + + +def _rmsnorm_relu_heads_fast( + x: Tensor, + weight: Tensor, + eps: float, + *, + scale: float, +) -> Tensor | None: + batch, tokens, heads, dim = x.shape + channels = heads * dim + block_c = 1 << (channels - 1).bit_length() + if ( + _DISABLE_RMS_RELU_FAST + or not x.is_cuda + or not weight.is_cuda + or torch.is_grad_enabled() + or triton is None + or _rmsnorm_relu_heads_kernel is None + or weight.numel() != channels + or block_c > 4096 + ): + return None + out = torch.empty( + (batch, tokens, heads, dim), + device=x.device, + dtype=torch.float32, + ) + grid = (batch * tokens,) + _rmsnorm_relu_heads_kernel[grid]( + x, + weight, + out, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + tokens, + heads, + dim, + channels, + float(eps), + float(scale), + block_c, + ) + return out + + def _ucpe_transform( x: Tensor, matrix: Tensor, @@ -1671,6 +2160,50 @@ def _ucpe_transform( *, inverse_rope: bool, ) -> tuple[Tensor, Tensor, Tensor]: + out = _ucpe_transform_apply(x, matrix, rotary_emb, inverse_rope=inverse_rope) + fast_norms = _ucpe_transform_norms_fast(x, out) + if fast_norms is not None: + pre_sq, post_sq = fast_norms + return out, pre_sq, post_sq + pre_sq = _ucpe_transform_norms_eager(x) + post_sq = _ucpe_transform_norms_eager(out) + return out, pre_sq, post_sq + + +def _ucpe_transform_norms_eager(x: Tensor) -> Tensor: + return x.float().pow(2).sum(dim=-1).transpose(1, 2).contiguous() + + +def _ucpe_transform_apply( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> Tensor: + fast = _ucpe_transform_apply_fast( + x, + matrix, + rotary_emb, + inverse_rope=inverse_rope, + ) + if fast is not None: + return fast + return _ucpe_transform_apply_eager( + x, + matrix, + rotary_emb, + inverse_rope=inverse_rope, + ) + + +def _ucpe_transform_apply_eager( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> Tensor: batch, tokens, heads, dim = x.shape half = dim // 2 if half % 4 != 0: @@ -1684,31 +2217,198 @@ def _ucpe_transform( else: second_out = _apply_complex_rope(second, rotary_emb) out = torch.cat([first_out, second_out], dim=-1) - pre_sq = x.float().pow(2).sum(dim=-1).transpose(1, 2).contiguous() - post_sq = out.float().pow(2).sum(dim=-1).transpose(1, 2).contiguous() - return out, pre_sq, post_sq + return out -def _ucpe_transform_apply( +def _ucpe_transform_apply_fast( x: Tensor, matrix: Tensor, rotary_emb: Tensor | None, *, inverse_rope: bool, -) -> Tensor: +) -> Tensor | None: batch, tokens, heads, dim = x.shape half = dim // 2 if half % 4 != 0: raise ValueError(f"UCPE requires head_dim/2 divisible by 4, got {half}.") - first = x[..., :half].reshape(batch, tokens, heads, half // 4, 4) - first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), first.float()) - first_out = first_out.reshape(batch, tokens, heads, half) - second = x[..., half:] - if inverse_rope and rotary_emb is not None: - second_out = _apply_complex_rope(second, rotary_emb.conj()) - else: - second_out = _apply_complex_rope(second, rotary_emb) - return torch.cat([first_out, second_out], dim=-1) + if ( + _DISABLE_UCPE_FAST + or not x.is_cuda + or not matrix.is_cuda + or rotary_emb is None + or not rotary_emb.is_cuda + or torch.is_grad_enabled() + or triton is None + or _ucpe_first_half_kernel is None + or _ucpe_rope_second_half_kernel is None + or matrix.shape[-2:] != (4, 4) + or matrix.shape[1] != tokens + or matrix.shape[0] not in (1, batch) + ): + return None + rope = torch.view_as_real(rotary_emb.squeeze(0).squeeze(0)) + if rope.shape != (tokens, half // 2, 2): + return None + + out = torch.empty_like(x, dtype=torch.float32) + groups4 = half // 4 + rope_pairs = half // 2 + block_n = 128 + first_grid = ( + batch, + heads, + groups4 * triton.cdiv(tokens, block_n), + ) + _ucpe_first_half_kernel[first_grid]( + x, + matrix, + out, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + matrix.stride(0), + matrix.stride(1), + matrix.stride(2), + matrix.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + tokens, + groups4, + matrix.shape[0], + block_n, + ) + rope_grid = ( + batch, + heads, + rope_pairs * triton.cdiv(tokens, block_n), + ) + _ucpe_rope_second_half_kernel[rope_grid]( + x, + rope, + out, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + rope.stride(0), + rope.stride(1), + rope.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + tokens, + half, + rope_pairs, + inverse_rope, + block_n, + ) + return out + + +def _ucpe_transform_norms_fast(x: Tensor, out: Tensor) -> tuple[Tensor, Tensor] | None: + batch, tokens, heads, dim = x.shape + block_d = 1 << (dim - 1).bit_length() + if ( + _DISABLE_UCPE_FAST + or not x.is_cuda + or not out.is_cuda + or torch.is_grad_enabled() + or triton is None + or _ucpe_norms_kernel is None + or out.shape != x.shape + or block_d > 256 + ): + return None + + pre_sq = torch.empty((batch, heads, tokens), device=x.device, dtype=torch.float32) + post_sq = torch.empty_like(pre_sq) + block_n = 16 + grid = ( + batch, + heads, + triton.cdiv(tokens, block_n), + ) + _ucpe_norms_kernel[grid]( + x, + out, + pre_sq, + post_sq, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + pre_sq.stride(0), + pre_sq.stride(1), + pre_sq.stride(2), + post_sq.stride(0), + post_sq.stride(1), + post_sq.stride(2), + tokens, + dim, + block_n, + block_d, + ) + return pre_sq, post_sq + + +def _silu_multiply(value: Tensor, gate: Tensor, *, inplace: bool) -> Tensor: + fast = _silu_multiply_fast(value, gate) + if fast is not None: + return fast + return value * F.silu(gate, inplace=inplace) + + +def _silu_multiply_fast(value: Tensor, gate: Tensor) -> Tensor | None: + if ( + _DISABLE_GLU_FAST + or value.shape != gate.shape + or value.dim() != 4 + or not value.is_cuda + or not gate.is_cuda + or torch.is_grad_enabled() + or triton is None + or _silu_multiply_kernel is None + ): + return None + batch, channels, height, width = value.shape + total = batch * channels * height * width + if total <= 0: + return torch.empty_like(value, memory_format=torch.channels_last) + out = torch.empty( + (batch, channels, height, width), + device=value.device, + dtype=value.dtype, + memory_format=torch.channels_last, + ) + block = 128 + grid = (triton.cdiv(total, block),) + _silu_multiply_kernel[grid]( + value, + gate, + out, + value.stride(0), + value.stride(1), + value.stride(2), + value.stride(3), + gate.stride(0), + gate.stride(1), + gate.stride(2), + gate.stride(3), + channels, + height, + width, + total, + block, + ) + return out def _camera_ray_mats( diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py index fb4c61092..c05a71c59 100644 --- a/integrations/sana/sana_wm/transformer.py +++ b/integrations/sana/sana_wm/transformer.py @@ -359,6 +359,25 @@ def _prepare_static_model_kwargs( model_kwargs["chunk_plucker_emb"] = prepare_plucker(chunk_plucker) del model_kwargs["chunk_plucker"] + camera_conditions = model_kwargs.get("camera_conditions") + prepare_camera = getattr(self.model, "prepare_camera_projection_cache", None) + if ( + "camera_cache" not in model_kwargs + and isinstance(camera_conditions, Tensor) + and camera_conditions.shape[0] == 1 + and callable(prepare_camera) + ): + _batch, _channels, frames, height, width = conditioning.latent_shape + with torch.inference_mode(): + rotary_emb, camera_cache = prepare_camera( + camera_conditions, + frames=frames, + height=height, + width=width, + ) + model_kwargs["rotary_emb"] = rotary_emb + model_kwargs["camera_cache"] = camera_cache + def _predict_with_prompt( self, noisy_latent: Tensor, @@ -500,10 +519,15 @@ def _batched_cfg_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, obje if isinstance(mask, Tensor): mask_uncond = negative_mask if isinstance(negative_mask, Tensor) else mask kwargs["mask"] = torch.cat([mask_uncond, mask], dim=0) - for key in ("camera_conditions", "chunk_plucker"): - value = kwargs.get(key) - if isinstance(value, Tensor): - kwargs[key] = torch.cat([value, value], dim=0) + camera_conditions = kwargs.get("camera_conditions") + if isinstance(camera_conditions, Tensor) and camera_conditions.shape[0] != 1: + kwargs["camera_conditions"] = torch.cat( + [camera_conditions, camera_conditions], + dim=0, + ) + chunk_plucker = kwargs.get("chunk_plucker") + if isinstance(chunk_plucker, Tensor): + kwargs["chunk_plucker"] = torch.cat([chunk_plucker, chunk_plucker], dim=0) return kwargs diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py index 7b8f23398..c04a30158 100644 --- a/integrations/sana/tests/test_smoke.py +++ b/integrations/sana/tests/test_smoke.py @@ -454,10 +454,7 @@ def forward( assert call["timestep"].shape == (2, 1, 2) assert call["prompt_embeds"].shape == (2, 1, 1, 1) torch.testing.assert_close(call["mask"], torch.cat([neg_mask, cond_mask], dim=0)) - torch.testing.assert_close( - call["camera_conditions"], - torch.cat([camera, camera], dim=0), - ) + torch.testing.assert_close(call["camera_conditions"], camera) torch.testing.assert_close( call["chunk_plucker"], torch.cat([chunk_plucker, chunk_plucker], dim=0), @@ -472,7 +469,9 @@ class DummyModel(torch.nn.Module): def __init__(self) -> None: super().__init__() self.prepare_calls = 0 - self.forward_calls: list[dict[str, torch.Tensor]] = [] + self.prepare_camera_calls = 0 + self.forward_calls: list[dict[str, object]] = [] + self.camera_cache = object() def prepare_plucker_embedding( self, @@ -481,6 +480,18 @@ def prepare_plucker_embedding( self.prepare_calls += 1 return chunk_plucker.flatten(2).transpose(1, 2) + def prepare_camera_projection_cache( + self, + camera_conditions: torch.Tensor, + *, + frames: int, + height: int, + width: int, + ) -> tuple[torch.Tensor, object]: + del camera_conditions + self.prepare_camera_calls += 1 + return torch.empty(frames, height, width), self.camera_cache + def forward( self, noisy_latent: torch.Tensor, @@ -488,11 +499,19 @@ def forward( prompt_embeds: torch.Tensor, *, chunk_plucker_emb: torch.Tensor, + rotary_emb: torch.Tensor, + camera_cache: object, **_kwargs: object, ) -> torch.Tensor: del timestep, prompt_embeds assert "chunk_plucker" not in _kwargs - self.forward_calls.append({"chunk_plucker_emb": chunk_plucker_emb}) + self.forward_calls.append( + { + "chunk_plucker_emb": chunk_plucker_emb, + "rotary_emb": rotary_emb, + "camera_cache": camera_cache, + } + ) return torch.ones_like(noisy_latent) transformer = SanaWMTransformerConfig().setup() @@ -505,6 +524,7 @@ def forward( uncondition=None, model_kwargs={ "chunk_plucker": chunk_plucker, + "camera_conditions": torch.zeros((1, 2, 20)), "data_info": {"condition_frame_info": {0: 0.0}}, }, first_latent=torch.empty((1, 1, 1, 1, 1)), @@ -525,13 +545,19 @@ def forward( ) assert dummy_model.prepare_calls == 1 + assert dummy_model.prepare_camera_calls == 1 assert "chunk_plucker" not in conditioning.model_kwargs + assert "camera_cache" in conditioning.model_kwargs assert len(dummy_model.forward_calls) == 2 cached = conditioning.model_kwargs["chunk_plucker_emb"] assert isinstance(cached, torch.Tensor) assert all( call["chunk_plucker_emb"] is cached for call in dummy_model.forward_calls ) + assert all( + call["camera_cache"] is dummy_model.camera_cache + for call in dummy_model.forward_calls + ) def test_inference_config_loads_yaml( @@ -776,8 +802,8 @@ def test_stage1_forward_preserves_latent_shape() -> None: text_dim=12, timestep_dim=8, depth=2, - num_heads=4, - head_dim=4, + num_heads=2, + head_dim=8, max_text_length=5, latent_grid_size=(2, 2), mlp_ratio=1, @@ -799,6 +825,67 @@ def test_stage1_forward_preserves_latent_shape() -> None: assert out.shape == latents.shape +def test_stage1_forward_accepts_precomputed_camera_cache() -> None: + """Precomputed camera cache must match the raw camera-conditioning path.""" + torch.manual_seed(0) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=1, + ) + model = SanaWMStage1Model(spec).eval() + with torch.no_grad(): + for param in model.parameters(): + param.normal_(mean=0.0, std=0.02) + latents = torch.randn(1, 4, 3, 2, 2) + timesteps = torch.ones(1, 1, 3) + text = torch.randn(1, 1, 5, 12) + mask = torch.ones(1, 5) + plucker = torch.randn(1, 6, 3, 2, 2) + camera = torch.zeros(1, 3, 20) + camera[..., :16] = torch.eye(4).flatten() + camera[..., 16:] = torch.tensor([1.0, 1.0, 0.5, 0.5]) + rotary_emb, camera_cache = model.prepare_camera_projection_cache( + camera, + frames=3, + height=2, + width=2, + ) + + raw = model( + latents, + timesteps, + text, + mask=mask, + chunk_plucker=plucker, + camera_conditions=camera, + ) + cached = model( + latents, + timesteps, + text, + mask=mask, + chunk_plucker=plucker, + camera_conditions=camera, + rotary_emb=rotary_emb, + camera_cache=camera_cache, + ) + + torch.testing.assert_close(cached, raw) + + def test_stage1_self_attention_uses_camera_conditions() -> None: """Changing camera conditions must affect camera attention.""" torch.manual_seed(0) @@ -834,6 +921,90 @@ def test_stage1_self_attention_uses_camera_conditions() -> None: assert not torch.allclose(base, shifted) +def test_stage1_self_attention_broadcasts_single_camera_batch() -> None: + """CFG can reuse one camera projection cache for both prompt branches.""" + torch.manual_seed(1) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + attn = Stage1SelfAttention(spec, use_gdn_convs=True).eval() + hidden = torch.randn(2, 8, 16) + camera = torch.zeros(1, 2, 20) + camera[..., :16] = torch.eye(4).flatten() + camera[..., 16:] = torch.tensor([1.0, 1.0, 0.5, 0.5]) + + broadcast = attn(hidden, HW=(2, 2, 2), camera_conditions=camera) + duplicated = attn( + hidden, + HW=(2, 2, 2), + camera_conditions=torch.cat([camera, camera], dim=0), + ) + + assert broadcast.shape == hidden.shape + torch.testing.assert_close(broadcast, duplicated) + + +def test_stage1_self_attention_uses_gdn_main_without_gdn_convs() -> None: + """Blocks without GDN conv weights still route main attention through GDN.""" + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=1, + ) + attn = Stage1SelfAttention(spec, use_gdn_convs=False).eval() + hidden = torch.randn(1, 4, 16) + called = {"gdn_main": False} + + def fake_gdn_main( + x: torch.Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: torch.Tensor | None, + precomputed_gates: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + del HW, rotary_emb + called["gdn_main"] = True + assert precomputed_gates[0].shape == (1, 2, 1, 4) + return torch.zeros_like(x) + + def fail_softmax_main(*_args: object, **_kwargs: object) -> torch.Tensor: + raise AssertionError("main attention should not use softmax") + + attn._forward_gdn_main = fake_gdn_main # type: ignore[method-assign] + attn._forward_softmax_main = fail_softmax_main # type: ignore[method-assign] + + out = attn(hidden, HW=(1, 2, 2), apply_output_gate=False) + + assert called["gdn_main"] + torch.testing.assert_close(out, torch.zeros_like(hidden)) + + def test_transformer_releases_stage1_runtime() -> None: """Free Stage-1-only modules and conditioning before decode/refine.""" transformer = SanaWMTransformerConfig().setup() diff --git a/integrations/sana/tests/test_stage1_cuda.py b/integrations/sana/tests/test_stage1_cuda.py index a72c40ada..d0f60ada6 100644 --- a/integrations/sana/tests/test_stage1_cuda.py +++ b/integrations/sana/tests/test_stage1_cuda.py @@ -77,3 +77,204 @@ def test_bidirectional_temporal_conv_fast_path_matches_eager( rtol=2e-2, atol=8e-2 if dtype is torch.bfloat16 else 1e-5, ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("inverse_rope", [False, True]) +@pytest.mark.parametrize("matrix_batch", [1, 2]) +def test_ucpe_transform_apply_fast_path_matches_eager( + dtype: torch.dtype, + inverse_rope: bool, + matrix_batch: int, +) -> None: + """Compare the CUDA UCPE transform fast path against the eager path.""" + _require_stage1_cuda() + torch.manual_seed(4321) + batch, frames, height, width, heads, dim = 2, 3, 2, 3, 2, 16 + tokens = frames * height * width + x = torch.randn( + batch, + tokens, + heads, + dim, + device="cuda", + dtype=dtype, + ) + matrix = torch.randn( + matrix_batch, + tokens, + 4, + 4, + device="cuda", + dtype=dtype, + ) + rope = stage1._slice_rope_for_camera( + stage1._wan_rope_complex(dim, frames, height, width, x.device), + dim, + ) + + reference = stage1._ucpe_transform_apply_eager( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + with torch.inference_mode(): + actual = stage1._ucpe_transform_apply( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + + assert actual.shape == x.shape + assert actual.dtype is torch.float32 + torch.testing.assert_close( + actual, + reference, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("inverse_rope", [False, True]) +@pytest.mark.parametrize("matrix_batch", [1, 2]) +def test_ucpe_transform_fast_path_matches_eager_norms( + dtype: torch.dtype, + inverse_rope: bool, + matrix_batch: int, +) -> None: + """Compare the UCPE fast path including GDN norm buffers.""" + _require_stage1_cuda() + torch.manual_seed(2468) + batch, frames, height, width, heads, dim = 2, 3, 2, 3, 2, 16 + tokens = frames * height * width + x = torch.randn( + batch, + tokens, + heads, + dim, + device="cuda", + dtype=dtype, + ) + matrix = torch.randn( + matrix_batch, + tokens, + 4, + 4, + device="cuda", + dtype=dtype, + ) + rope = stage1._slice_rope_for_camera( + stage1._wan_rope_complex(dim, frames, height, width, x.device), + dim, + ) + + reference = stage1._ucpe_transform_apply_eager( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + reference_pre = stage1._ucpe_transform_norms_eager(x) + reference_post = stage1._ucpe_transform_norms_eager(reference) + with torch.inference_mode(): + actual, actual_pre, actual_post = stage1._ucpe_transform( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + + assert actual.shape == x.shape + assert actual_pre.shape == (batch, heads, tokens) + assert actual_post.shape == (batch, heads, tokens) + torch.testing.assert_close( + actual, + reference, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + torch.testing.assert_close( + actual_pre, + reference_pre, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + torch.testing.assert_close( + actual_post, + reference_post, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_silu_multiply_fast_path_matches_eager(dtype: torch.dtype) -> None: + """Compare the fused GLU elementwise fast path against PyTorch eager.""" + _require_stage1_cuda() + torch.manual_seed(1357) + value = torch.randn( + 4, + 17, + 5, + 7, + device="cuda", + dtype=dtype, + ).contiguous(memory_format=torch.channels_last) + gate = torch.randn_like(value).contiguous(memory_format=torch.channels_last) + + reference = value * torch.nn.functional.silu(gate) + with torch.inference_mode(): + actual = stage1._silu_multiply(value, gate, inplace=False) + + assert actual.shape == value.shape + assert actual.dtype == value.dtype + assert actual.is_contiguous(memory_format=torch.channels_last) + torch.testing.assert_close( + actual.float(), + reference.float(), + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("scale", [1.0, 0.0375]) +def test_rmsnorm_relu_heads_fast_path_matches_eager( + dtype: torch.dtype, + scale: float, +) -> None: + """Compare fused GDN RMSNorm/ReLU preparation against eager math.""" + _require_stage1_cuda() + torch.manual_seed(9753) + batch, tokens, heads, dim = 2, 13, 4, 8 + x = torch.randn( + batch, + tokens, + heads, + dim, + device="cuda", + dtype=dtype, + ) + weight = torch.randn(heads * dim, device="cuda", dtype=dtype) + eps = 1e-6 + inv_rms = stage1._inv_rms(x, eps) + reference = torch.nn.functional.relu( + x.float() * inv_rms[:, :, None, None] * weight.float().view(heads, dim) + ) + if scale != 1.0: + reference = reference * scale + + with torch.inference_mode(): + actual = stage1._rmsnorm_relu_heads(x, weight, eps, scale=scale) + + assert actual.shape == x.shape + assert actual.dtype is torch.float32 + torch.testing.assert_close( + actual, + reference, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + )