diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..48365d8fdfa6b23fcd03ba8b17ad32ff35a85e06 --- /dev/null +++ b/README.md @@ -0,0 +1,49 @@ +--- +tags: +- blackhole +- tt-dit-server +- tt-model-cache +- tt-model-container +--- + +# flux2-dev-qb2 + +A **tt-model container package**: the serving platform ships as a Docker image, so a consumer needs only Docker and a Tenstorrent card — no tt-metal, no vLLM, no venv on the host. + +## Serve it + +```bash +tt-model pull stisiTT/flux2-dev-qb2 +tt-model serve stisiTT/flux2-dev-qb2 +``` + +## Serve profiles + +One image serves every profile below; pick one with `--profile`. + +| profile | hardware | mesh | max_num_seqs | max_model_len | +| --- | --- | --- | --- | --- | +| `default` *(default)* | p300x2 | QB2 | | | + +## What is inside + +- **weights**: [`black-forest-labs/FLUX.2-dev`](https://huggingface.co/black-forest-labs/FLUX.2-dev) — downloaded to *your* HF cache at pull time, never baked into the image +- **arch**: blackhole +- **serving stack**: `tt-dit-server` + +## Provenance + +Everything below is pinned; the image was built from exactly these. + +| component | pinned to | +| --- | --- | +| tt-metal | `bb518e6548f4aa2ce45f7e803873711665724e11` *(dirty tree)* | +| code digest | `f97b5fc89d804602` | +| built | 2026-08-28T02:50:33+00:00 by tt-model 0.1.0 | + +## Shipped code + +`code/` in this repo is byte-identical to what runs inside the image. + +- `models/common/` +- `models/tt_dit/` diff --git a/code/models/common/README.md b/code/models/common/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f5167e8faee535b5274ecf3b3fa8471cd17e9f57 --- /dev/null +++ b/code/models/common/README.md @@ -0,0 +1,319 @@ +# TTNN Validation & Testing Utilities + +Helpers for validating TTNN computations against reference implementations and +for moving tensors between TTNN and PyTorch. The public API is implemented +across `models.common.validation_tools`, `models.common.metrics`, +`models.common.auto_compose`, and `models.common.distribute_as`, and is +exercised in: + +- `models/common/tests/test_validation_tools.py` +- `models/common/tests/test_metrics.py` +- `models/common/tests/test_auto_compose.py` +- `models/common/tests/test_distribute_as.py` +- `models/common/tests/host/test_metrics_pytorch_only.py` + +The examples in these tests are the most up‑to‑date reference for usage. + +## Quick Start – host reference (`compare_to_torch`) + +Use `compare_to_torch` when your reference implementation is a PyTorch function. +Inputs and outputs are automatically converted between TTNN and PyTorch. + +```python +import torch +import ttnn +from models.common.validation_tools import compare_to_torch, Metric, get_validation_registry + + +@compare_to_torch( + reference_fn=torch.matmul, + metric_tolerances={ + Metric.MAX_ABS_ERROR: 1e-1, + Metric.PCC: 0.99, + }, +) +def ttnn_matmul(a, b): + # a, b are TTNN tensors (possibly sharded) + return ttnn.matmul(a, b) + + +def run_example(device: ttnn.MeshDevice): + m, n, k = 16, 24, 12 + a = torch.randn(1, m, k, dtype=torch.bfloat16) + b = torch.randn(1, k, n, dtype=torch.bfloat16) + + a_tt = ttnn.from_torch(a.unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT) + b_tt = ttnn.from_torch(b.unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT) + + _ = ttnn_matmul(a_tt, b_tt) + + registry = get_validation_registry() + registry.print_report() +``` + +Note: +- When the signature of the decorated function is different from the reference function, `input_to_torch` and `output_to_torch` can be used to map the inputs and outputs between the decorated function and the reference function. +- See `models/common/tests/test_validation_tools.py::test_validation_matmul` for a real test using +this pattern. + +## Quick Start – TTNN reference (`compare_to_ttnn`) + +Use `compare_to_ttnn` when both your implementation and reference are TTNN‑based +and you want metrics computed directly on device. + +```python +import torch +import ttnn +from models.common.validation_tools import compare_to_ttnn + + +def torch_rms_norm(x, weight, eps=1e-6): + var = x.pow(2).mean(-1, keepdim=True) + return weight * x * torch.rsqrt(var + eps) + + +class DeviceValidatedRMSNorm: + def __init__(self, weight: torch.Tensor, eps: float, device: ttnn.MeshDevice): + self.eps = eps + self.device = device + self.weight_torch = weight + self.weight = ttnn.from_torch( + weight.unsqueeze(0).unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT + ) + + def _reference_impl(self, x): + x_torch = ttnn.to_torch(x).squeeze(0) + y_torch = torch_rms_norm(x_torch, self.weight_torch, self.eps) + return ttnn.from_torch( + y_torch.unsqueeze(0), device=self.device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT + ) + + @compare_to_ttnn(reference_fn=lambda self, x: self._reference_impl(x)) + def __call__(self, x): + x_sq = ttnn.mul(x, x) + mean_x_sq = ttnn.mean(x_sq, dim=-1, keepdim=True) + rms = ttnn.sqrt(ttnn.add(mean_x_sq, self.eps)) + x_norm = ttnn.mul(x, ttnn.reciprocal(rms)) + return ttnn.mul(x_norm, self.weight) +``` +Note: +- When the signature of the decorated function is different from the reference function, `input_to_ttnn` and `output_to_ttnn` can be used to map the inputs and outputs between the decorated function and the reference function. +- This mirrors the pattern used by `DeviceValidatedRMSNorm` in +`models/common/tests/test_validation_tools.py`. + +## Features + +- **Decorator‑based validation** – Wrap TTNN functions or methods and compare + them against PyTorch or TTNN references. +- **Host and device modes** – `compare_to_torch` (PyTorch reference) and + `compare_to_ttnn` (TTNN reference). +- **TTNN‑native metrics** – When both outputs are TTNN tensors, metrics are + computed on device with minimal host transfer. +- **Configurable tolerances** – Per‑metric tolerances via the `Metric` enum, + string keys, or `MetricSpec`. +- **Custom metrics** – Inject your own metric functions. +- **Global registry** – Collects all validation runs for reporting. +- **Easy disabling** – Turn validation on/off globally without changing call + sites. + +## Core Components + +### Validation decorators + +All decorators live in `models.common.validation_tools`: + +- `compare_to_torch(reference_fn, *, input_to_torch=None, output_to_torch=None, metric_tolerances=None, enabled=True, raise_exceptions=False, return_reference_output=False)` + - Use when `reference_fn` is a PyTorch implementation. + - By default, all TTNN tensors in the arguments/outputs are converted to + PyTorch via `to_torch_auto_compose`. + - Optional `input_to_torch(*args, **kwargs)` lets you override how inputs + are mapped to the reference. + - Optional `output_to_torch(output)` converts the implementation output + before metrics are computed. + +- `compare_to_ttnn(reference_fn, *, input_to_ttnn=None, output_to_ttnn=None, metric_tolerances=None, enabled=True, raise_exceptions=False, return_reference_output=False)` + - Use when `reference_fn` consumes and returns TTNN tensors. + - Optional `input_to_ttnn(*args, **kwargs)` lets you override how inputs + are mapped to the reference. + - Optional `output_to_ttnn(output)` converts the implementation output + before metrics are computed. + - If both implementation and reference return TTNN tensors, metrics run + entirely on device. + +In both cases, decorating a function records a `ValidationResult` in the global +`ValidationRegistry` every time the function is called (unless disabled). + +### Metrics + +Metric utilities are implemented in `models.common.metrics`: + +- `compute_max_abs_error(impl, ref)` – max absolute error. +- `compute_mean_abs_error(impl, ref)` – mean absolute error. +- `compute_pcc(impl, ref)` – Pearson correlation coefficient; uses TTNN + operations when possible and falls back to host. +- `comp_allclose(impl, ref, rtol=..., atol=...)` – allclose check plus a + detailed delta string. +- `DEFAULT_METRICS` – dict with built‑in metrics (`"max_abs_error"`, + `"mean_abs_error"`, `"pcc"`). + +Metrics support both TTNN and PyTorch tensors. + +### Registry and control functions + +From `models.common.validation_tools`: + +- `get_validation_registry() -> ValidationRegistry` + - Holds all `ValidationResult` objects. + - Provides `get_summary()` and `print_report(verbose: bool = False)`. + +- `enable_validation(enabled: bool = True)` + - Globally enable/disable validation; when disabled, decorators become + transparent wrappers. + +- `clear_validation_results()` + - Clear all accumulated validation results. + +`ValidationResult` includes: + +- `function_name` +- `passed` (bool) +- `metrics` – map of metric name → per‑metric result (value, passed, error) +- `execution_time_impl`, `execution_time_ref` +- `timestamp` +- `logs` – optional debug strings + +### Auto‑compose helper + +`to_torch_auto_compose` lives in `models.common.auto_compose`. + +It converts an arbitrary TTNN tensor (including sharded/replicated multi‑device +tensors) to a single PyTorch tensor by automatically choosing the appropriate +mesh composer. + +It is heavily used in: + +- `test_auto_compose.py` +- `test_distribute_as.py` +- all `compare_to_torch`‑based examples. + +## Usage Patterns + +High‑level patterns illustrated in the tests: + +1. **Host reference with explicit input mapping** + - See `HostValidatedRMSNorm` in `models/common/tests/test_validation_tools.py`. + - Uses `compare_to_torch` with `input_to_torch` to map TTNN inputs and + TTNN weights to a pure‑PyTorch reference function. + +2. **TTNN reference (on‑device metrics)** + - See `DeviceValidatedRMSNorm` in `models/common/tests/test_validation_tools.py`. + - Uses `compare_to_ttnn` where both implementation and reference return + TTNN tensors; metrics run on device. + +3. **Simple library calls** + - See `ttnn_matmul` and `ttnn_matmul_reverse` in `models/common/tests/test_validation_tools.py`. + - `compare_to_torch(reference_fn=torch.matmul, ...)` with optional + `input_to_torch` remapping. + +4. **Checkpoint / `from_torch` validation** + - See `from_torch_checkpoint` in `models/common/tests/test_validation_tools.py`. + - Validates a direct `ttnn.from_torch(...)` call using `compare_to_torch` + and `output_to_torch`. + +5. **Custom metric via `MetricSpec`** + - See `ttnn_matmul_metric_spec` in `models/common/tests/test_validation_tools.py` + and `MetricSpec` usage in `models/common/tests/host/test_metrics_pytorch_only.py`. + - Use `MetricSpec(tolerance=..., higher_is_better=..., compute_fn=...)` + in `metric_tolerances`. + +6. **Non‑decorator usage** + - `test_validation_non_decorator_class_vs_class_torch` demonstrates calling + `compare_to_torch` in a more manual, non‑decorator style between two + callable classes. + +## Default Metrics and Tolerances + +When `metric_tolerances` is omitted, the framework uses sensible defaults: + +- `Metric.MAX_ABS_ERROR` with tolerance `1e-2` +- `Metric.PCC` with tolerance `0.99` + +If you pass a `metric_tolerances` dict, keys can be: + +- `Metric` enum members (recommended), e.g. `Metric.MAX_ABS_ERROR` +- strings (`"max_abs_error"`, `"mean_abs_error"`, `"pcc"`) +- arbitrary names when used with `MetricSpec` + +Values can be: + +- a float tolerance (uses the built‑in metric) +- a `MetricSpec` instance to define a custom metric and tolerance + +Example: + +```python +from models.common.validation_tools import Metric, MetricSpec +from models.common.metrics import compute_pcc + + +@compare_to_torch( + reference_fn=torch.matmul, + metric_tolerances={ + Metric.PCC: MetricSpec(tolerance=0.99, higher_is_better=True, compute_fn=compute_pcc), + Metric.MAX_ABS_ERROR: 1.5e-1, + }, +) +def ttnn_matmul_metric_spec(a, b): + return ttnn.matmul(a, b) +``` + +## Testing + +The local test suite in `models/common/tests` shows end‑to‑end usage: + +- `test_validation_tools.py` + - Core decorator usage, registry behaviour, error handling, custom metrics. +- `test_metrics.py` + - Numerical correctness of device and host metric functions. +- `host/test_metrics_pytorch_only.py` + - Pure‑PyTorch metric tests. +- `test_auto_compose.py` + - Auto‑composition of sharded/replicated TTNN tensors into PyTorch. +- `test_distribute_as.py` + - Distribution helpers (`from_torch_dist_as`) that mirror an existing TTNN + tensor’s topology. + +Example commands (run from the repo root, with TTNN available): + +```bash +python -m pytest models/common/tests/test_validation_tools.py -v +python -m pytest models/common/tests/test_metrics.py -v +python -m pytest models/common/tests/host/test_metrics_pytorch_only.py -v +``` + +## API Reference (public surface) + +All symbols below are imported from `models.common.validation_tools` and `models.common.metrics`: + +- Decorators: + - `compare_to_torch` + - `compare_to_ttnn` +- Registry and control: + - `ValidationResult` + - `ValidationRegistry` + - `get_validation_registry` + - `enable_validation` + - `clear_validation_results` +- Metrics: + - `Metric` (enum: `MAX_ABS_ERROR`, `MEAN_ABS_ERROR`, `PCC`) + - `MetricSpec` + - `compute_max_abs_error` + - `compute_mean_abs_error` + - `compute_pcc` + - `comp_allclose` + - `DEFAULT_METRICS` +- Auto‑compose: + - `to_torch_auto_compose` + +For concrete, runnable examples of each API, see the tests listed at the top +of this document. diff --git a/code/models/common/auto_compose.py b/code/models/common/auto_compose.py new file mode 100644 index 0000000000000000000000000000000000000000..571fa0f114b567d4bb61734d36dbac010a29ab37 --- /dev/null +++ b/code/models/common/auto_compose.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Automatic composition of multi-device sharded tensors using TensorTopology. + +This module provides utilities to infer the correct MeshToTensor composer from a +sharded ttnn.Tensor's topology metadata and use it to compose shards on host. +""" + +from typing import Optional + +import torch +from loguru import logger + +import ttnn + +# ====================================================================================== +# Public API +# ====================================================================================== + + +def to_torch_auto_compose(tensor: ttnn.Tensor, device: Optional[ttnn.MeshDevice] = None) -> torch.Tensor: + """ + Convert a (possibly multi-device) TTNN tensor to torch, automatically + composing shards based on the tensor's topology. + + Args: + tensor: The distributed tensor to convert + device: Optional MeshDevice to use when the tensor lives on host + + Returns: + PyTorch tensor with shards composed + """ + composer = _infer_mesh_composer_from_topology(tensor, device=device) + try: + return ttnn.to_torch(tensor, mesh_composer=composer) + except Exception as e: + logger.error(f"Failed to convert tensor to torch with mesh_composer: {e}") + raise + + +def extract_tensor_topology_info( + tensor: ttnn.Tensor, +) -> tuple[list[object], list[int]]: + """ + Extract placements and distribution shape from a tensor's topology. + + Returns: + (placements, dist_shape) + """ + topology = tensor.tensor_topology() + placements = topology.placements() + dist_shape = list(topology.distribution_shape()) + return placements, dist_shape + + +def get_device_from_tensor(tensor: ttnn.Tensor) -> Optional[ttnn.MeshDevice]: + """Get device from tensor or fallback to provided mesh_device.""" + device = tensor.device() + # tensor.device() returns None if the tensor is on the host (ttnn/core/tensor/tensor.cpp --> Tensor::device()) + if device is None: + logger.debug("tensor.device() returns None, tensor is on the host") + else: + logger.debug(f"tensor.device() returns {device}") + + return device + + +# ====================================================================================== +# Private Implementation +# ====================================================================================== + + +def _infer_mesh_composer_from_topology( + tensor: ttnn.Tensor, *, device: Optional[ttnn.MeshDevice] = None +) -> Optional[ttnn.CppMeshToTensor]: + """ + Return a MeshToTensor composer inferred from the tensor's TensorTopology, + or None if no composition is needed (fully replicated, single-device). + + Note: For ND meshes with replicated dimensions, the composer will concatenate + all replicas, resulting in duplicated data. Callers may want to slice the + result if only one copy is desired. + + Args: + tensor: The distributed tensor to infer composer for + + Returns: + MeshToTensor composer or None if no composition needed + """ + placements, dist_shape = extract_tensor_topology_info(tensor) + + # No distribution or trivial 1-device case + if len(dist_shape) == 0 or (len(dist_shape) == 1 and dist_shape[0] == 1): + return None + + tensor_device = get_device_from_tensor(tensor) + mesh_device = tensor_device or device + if mesh_device is None: + # As a last resort, try default device for backward-compatibility + mesh_device = ttnn.GetDefaultDevice() + if mesh_device is None: + raise RuntimeError( + "Tensor is on host and no mesh_device provided. " + "Pass device=... to to_torch_auto_compose or set a default via ttnn.SetDefaultDevice(...)." + ) + + # Must match length (should be guaranteed by C++ TT_FATAL in ttnn/core/distributed/distributed_tensor.cpp) + assert len(dist_shape) == len(placements) + + if len(dist_shape) == 1 and mesh_device.shape.dims() == 1: + return _compose_1d_sharded(mesh_device, placements, dist_shape) + else: + # N >= 2 dimensions + return _compose_nd_sharded(mesh_device, placements, dist_shape) + + +def _compose_1d_sharded( + device: ttnn.MeshDevice, + placements: list[object], + dist_shape: list[int], +) -> Optional[ttnn.CppMeshToTensor]: + """Handle 1D case - returns None if fully replicated.""" + p = placements[0] + if isinstance(p, ttnn.PlacementShard): + # Use ND composer with shape override to match the tensor's distribution + composer_cfg = ttnn.MeshComposerConfig(dims=[p.dim], mesh_shape_override=ttnn.MeshShape(dist_shape)) + return ttnn.create_mesh_composer(device, composer_cfg) + # Fully replicated - no composition needed + return None + + +def _compose_nd_sharded( + device: ttnn.MeshDevice, + placements: list[object], + dist_shape: list[int], +) -> ttnn.CppMeshToTensor: + """ + Handle ND (N>=2) case. + + For replicated mesh dims, we use dim 0 as convention (the composed result + will include all replicas concatenated, which is typically not desired but + is how the C++ API works). + """ + dims = [] + shape_override = [] + for i, p in enumerate(placements): + if isinstance(p, ttnn.PlacementShard): + dims.append(p.dim) + shape_override.append(dist_shape[i]) + else: + assert isinstance(p, ttnn.PlacementReplicate) + # [INFO] steal from TensorDistribution2x4Test test case in test_distributed_tensor.cpp + # Replicated: use dim 0 as convention + dims.append(0) + # Replicated: use shape 1 to skip concatenation + shape_override.append(1) + + composer_cfg = ttnn.MeshComposerConfig(dims=dims, mesh_shape_override=ttnn.MeshShape(shape_override)) + return ttnn.create_mesh_composer(device, composer_cfg) diff --git a/code/models/common/device_utils.py b/code/models/common/device_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0066f8fe3123d24a1d14a25bd9345a6690b6470e --- /dev/null +++ b/code/models/common/device_utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Device topology naming helpers shared by TTTv2 modules.""" + +from __future__ import annotations + +import ttnn + + +def is_blackhole() -> bool: + return "blackhole" in ttnn.get_arch_name() + + +def get_device_name(mesh_device: ttnn.MeshDevice, num_devices: int | None = None) -> str: + """Return the product/topology name for a TT mesh device. + + By default, the full mesh device count is used. CCL callers can pass a + host-local device count when they need link-count tuning for the current + process rather than for the full mesh. + """ + num_devices = mesh_device.get_num_devices() if num_devices is None else num_devices + dram_grid_size = mesh_device.dram_grid_size() + + if ttnn.device.is_blackhole(mesh_device): + device_names = { + 1: "P100" if dram_grid_size and dram_grid_size.x == 7 else "P150", + 2: "P300", + 4: "P150x4", + 8: "P150x8", + 32: "BHGLX", + } + elif ttnn.device.is_wormhole_b0(mesh_device): + device_names = { + 1: "N150", + 2: "N300", + 4: "N150x4", + 8: "T3K", + 32: "TG", + } + else: + raise ValueError(f"Unsupported architecture: {ttnn.get_arch_name()}") + + if num_devices in device_names: + return device_names[num_devices] + raise ValueError(f"Unsupported number of devices: {num_devices} for {ttnn.get_arch_name()}") diff --git a/code/models/common/distribute_as.py b/code/models/common/distribute_as.py new file mode 100644 index 0000000000000000000000000000000000000000..fc309609103fc149c0388737236041fe4d7a4c59 --- /dev/null +++ b/code/models/common/distribute_as.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Distribute a torch.Tensor over a mesh using the same topology as a reference TTNN tensor. + +This mirrors the composition logic in `auto_compose.py` but in reverse: we infer a +TensorToMesh mapper from the reference tensor's TensorTopology and use it to distribute +the torch tensor accordingly. +""" + +from typing import Optional + +import torch + +import ttnn + +from .auto_compose import extract_tensor_topology_info, get_device_from_tensor + +# ====================================================================================== +# Public API +# ====================================================================================== + + +def from_torch_dist_as( + from_tensor_pt: torch.Tensor, as_tensor_tt: ttnn.Tensor, device: Optional[ttnn.MeshDevice] = None +) -> ttnn.Tensor: + """ + Distribute a torch.Tensor over a mesh using the same topology as an existing TTNN tensor. + + Args: + from_tensor_pt: Source PyTorch tensor on host. + as_tensor_tt: Reference TTNN tensor whose topology (placements + distribution shape) will be mirrored. + device: Optional mesh device. If omitted, inferred from ``as_tensor_tt`` when possible. + + Returns: + A TTNN tensor distributed according to ``as_tensor_tt``'s topology and memory configuration + (e.g. height-sharded decode heads). + """ + mapper, device = _infer_mesh_mapper_from_topology(as_tensor_tt, device=device) + + # Usage Patterns: unlike ttnn.to_torch, `device` is required here! + # Pattern 1: Using mesh_mapper without device (tensor stays in host memory) Programming_Mesh_of_Devices_with_TT-NN.md:370-375 + # Then transfer to device separately: Programming_Mesh_of_Devices_with_TT-NN.md:404-405 + # Pattern 2: Using both mesh_mapper and device together (direct to device) llms.md:1204-1218 + mem_cfg = as_tensor_tt.memory_config() + return ttnn.from_torch( + from_tensor_pt, + dtype=getattr(as_tensor_tt, "dtype", None), + layout=getattr(as_tensor_tt, "layout", None), + device=device, + memory_config=mem_cfg, + mesh_mapper=mapper, + ) + + +# ====================================================================================== +# Private Implementation +# ====================================================================================== + + +def _infer_mesh_mapper_from_topology( + tensor: ttnn.Tensor, *, device: Optional[ttnn.MeshDevice] = None +) -> Optional[ttnn.CppTensorToMesh]: + """ + Return a TensorToMesh mapper inferred from the tensor's TensorTopology, + or (None, mesh_device) if no distribution is needed (fully replicated, single-device). + """ + placements, dist_shape = extract_tensor_topology_info(tensor) + + tensor_device = get_device_from_tensor(tensor) + mesh_device = tensor_device or device + if mesh_device is None: + mesh_device = ttnn.GetDefaultDevice() + if mesh_device is None: + raise RuntimeError( + "Tensor is on host and no mesh_device provided. " "Set a default via ttnn.SetDefaultDevice(...)." + ) + + assert len(dist_shape) == len(placements) + + if len(dist_shape) == 1 and mesh_device.shape.dims() == 1: + return _map_1d(mesh_device, placements, dist_shape), mesh_device + else: + return _map_nd(mesh_device, placements, dist_shape), mesh_device + + +def _map_1d( + device: ttnn.MeshDevice, + placements: list[object], + dist_shape: list[int], +) -> Optional[ttnn.CppTensorToMesh]: + """ + Build a 1D TensorToMesh mapper. Returns None if fully trivial (handled earlier). + """ + p = placements[0] + if isinstance(p, ttnn.PlacementShard): + mapper_cfg = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(p.dim)], + mesh_shape_override=ttnn.MeshShape(dist_shape), + ) + return ttnn.create_mesh_mapper(device, mapper_cfg) + else: + # Replicate across the 1D mesh extent + mapper_cfg = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate()], + mesh_shape_override=ttnn.MeshShape(dist_shape), + ) + return ttnn.create_mesh_mapper(device, mapper_cfg) + + +def _map_nd( + device: ttnn.MeshDevice, + placements: list[object], + dist_shape: list[int], +) -> ttnn.CppTensorToMesh: + """ + Build an ND TensorToMesh mapper that mirrors the tensor's placements and distribution shape. + """ + mapper_placements = [] + for p in placements: + if isinstance(p, ttnn.PlacementShard): + mapper_placements.append(ttnn.PlacementShard(p.dim)) + else: + assert isinstance(p, ttnn.PlacementReplicate) + mapper_placements.append(ttnn.PlacementReplicate()) + + mapper_cfg = ttnn.MeshMapperConfig(placements=mapper_placements, mesh_shape_override=ttnn.MeshShape(dist_shape)) + return ttnn.create_mesh_mapper(device, mapper_cfg) diff --git a/code/models/common/generation_utils.py b/code/models/common/generation_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9f6d76a287d06960e8012ee42c1feb529e5424c3 --- /dev/null +++ b/code/models/common/generation_utils.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + +import torch +from loguru import logger +from transformers.generation.configuration_utils import GenerationConfig +from transformers.generation.logits_process import ( # ForceTokensLogitsProcessor, + EncoderNoRepeatNGramLogitsProcessor, + EncoderRepetitionPenaltyLogitsProcessor, + ExponentialDecayLengthPenalty, + ForcedBOSTokenLogitsProcessor, + ForcedEOSTokenLogitsProcessor, + InfNanRemoveLogitsProcessor, + LogitNormalization, + LogitsProcessorList, + MinLengthLogitsProcessor, + MinNewTokensLengthLogitsProcessor, + NoBadWordsLogitsProcessor, + NoRepeatNGramLogitsProcessor, + PrefixConstrainedLogitsProcessor, + RepetitionPenaltyLogitsProcessor, + SuppressTokensAtBeginLogitsProcessor, + SuppressTokensLogitsProcessor, +) + +# HammingDiversityLogitsProcessor (diverse beam search) was removed in +# transformers 5.x with no replacement. Import it optionally so this module +# still loads; it's only used when diversity_penalty > 0, which TT generation +# paths don't exercise. +try: + from transformers.generation.logits_process import HammingDiversityLogitsProcessor +except ImportError: # transformers >= 5.x + HammingDiversityLogitsProcessor = None + + +def _merge_criteria_processor_list( + default_list, # Union[LogitsProcessorList, StoppingCriteriaList], + custom_list, # Union[LogitsProcessorList, StoppingCriteriaList], +): # -> Union[LogitsProcessorList, StoppingCriteriaList]: + if len(custom_list) == 0: + return default_list + + for default in default_list: + for custom in custom_list: + if type(custom) is type(default): + object_type = "stopping criteria" if isinstance(custom, StoppingCriteria) else "logits processor" + raise ValueError( + f"A custom {object_type} of type {type(custom)} with values {custom} has been passed to" + f" `generate`, but it has already been created with the values {default}. {default} has been" + " created by passing the corresponding arguments to generate or by the model's config default" + f" values. If you just want to change the default values of {object_type} consider passing" + f" them as arguments to `generate` instead of using a custom {object_type}." + ) + default_list.extend(custom_list) + return default_list + + +def _get_logits_processor( + generation_config: GenerationConfig, + input_ids_seq_length: int, + encoder_input_ids, # torch.LongTensor + prefix_allowed_tokens_fn, # Callable[[int, torch.Tensor], List[int]], + logits_processor, # Optional[LogitsProcessorList] +): # -> LogitsProcessorList: + """ + This class returns a [`LogitsProcessorList`] list object that contains all relevant [`LogitsProcessor`] + instances used to modify the scores of the language model head. + """ + # instantiate processors list + processors = LogitsProcessorList() + + # the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files + # all samplers can be found in `generation_utils_samplers.py` + if generation_config.diversity_penalty is not None and generation_config.diversity_penalty > 0.0: + if HammingDiversityLogitsProcessor is None: + raise NotImplementedError( + "diversity_penalty > 0 (diverse beam search) requires HammingDiversityLogitsProcessor, " + "which was removed in transformers 5.x." + ) + processors.append( + HammingDiversityLogitsProcessor( + diversity_penalty=generation_config.diversity_penalty, + num_beams=generation_config.num_beams, + num_beam_groups=generation_config.num_beam_groups, + ) + ) + if generation_config.encoder_repetition_penalty is not None and generation_config.encoder_repetition_penalty != 1.0: + processors.append( + EncoderRepetitionPenaltyLogitsProcessor( + penalty=generation_config.encoder_repetition_penalty, + encoder_input_ids=encoder_input_ids, + ) + ) + if generation_config.repetition_penalty is not None and generation_config.repetition_penalty != 1.0: + processors.append(RepetitionPenaltyLogitsProcessor(penalty=generation_config.repetition_penalty)) + if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0: + processors.append(NoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size)) + if ( + generation_config.encoder_no_repeat_ngram_size is not None + and generation_config.encoder_no_repeat_ngram_size > 0 + ): + if len(encoder_input_ids.shape) == 2: + processors.append( + EncoderNoRepeatNGramLogitsProcessor(generation_config.encoder_no_repeat_ngram_size, encoder_input_ids) + ) + else: + raise ValueError("It's impossible to use `encoder_no_repeat_ngram_size` with decoder-only architecture") + if generation_config.bad_words_ids is not None: + processors.append(NoBadWordsLogitsProcessor(generation_config.bad_words_ids, generation_config.eos_token_id)) + if ( + generation_config.min_length is not None + and generation_config.eos_token_id is not None + and generation_config.min_length > 0 + ): + processors.append(MinLengthLogitsProcessor(generation_config.min_length, generation_config.eos_token_id)) + if ( + generation_config.min_new_tokens is not None + and generation_config.eos_token_id is not None + and generation_config.min_new_tokens > 0 + ): + processors.append( + MinNewTokensLengthLogitsProcessor( + input_ids_seq_length, + generation_config.min_new_tokens, + generation_config.eos_token_id, + ) + ) + if prefix_allowed_tokens_fn is not None: + processors.append( + PrefixConstrainedLogitsProcessor( + prefix_allowed_tokens_fn, + generation_config.num_beams // generation_config.num_beam_groups, + ) + ) + if generation_config.forced_bos_token_id is not None: + processors.append(ForcedBOSTokenLogitsProcessor(generation_config.forced_bos_token_id)) + if generation_config.forced_eos_token_id is not None: + processors.append( + ForcedEOSTokenLogitsProcessor(generation_config.max_length, generation_config.forced_eos_token_id) + ) + if generation_config.remove_invalid_values is True: + processors.append(InfNanRemoveLogitsProcessor()) + if generation_config.exponential_decay_length_penalty is not None: + processors.append( + ExponentialDecayLengthPenalty( + generation_config.exponential_decay_length_penalty, + generation_config.eos_token_id, + input_ids_seq_length, + ) + ) + if generation_config.suppress_tokens is not None: + processors.append(SuppressTokensLogitsProcessor(generation_config.suppress_tokens)) + if generation_config.begin_suppress_tokens is not None: + begin_index = input_ids_seq_length + begin_index = ( + begin_index + if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None) + else begin_index + 1 + ) + processors.append(SuppressTokensAtBeginLogitsProcessor(generation_config.begin_suppress_tokens, begin_index)) + processors = _merge_criteria_processor_list(processors, logits_processor) + # `LogitNormalization` should always be the last logit processor, when present + if generation_config.renormalize_logits is True: + processors.append(LogitNormalization()) + return processors + + +def get_logits_processor(input_ids, config): + generation_config = GenerationConfig.from_model_config(config) + input_ids_seq_length = input_ids.shape[-1] + + logits_processor = _get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_seq_length, + encoder_input_ids=input_ids, + prefix_allowed_tokens_fn=None, + logits_processor=LogitsProcessorList(), + ) + + return logits_processor + + +def pad_input_32(tensor, value): + len = tensor.shape[1] + + if len % 32 == 0: + return tensor + + padded_len = ((len // 32) + 1) * 32 + + pad_tensor = (value * torch.ones(tensor.shape[0], padded_len - len)).to(torch.long) + tensor = torch.cat([tensor, pad_tensor], dim=1) + + return tensor + + +def run_generate( + input_sentance, + tokenizer, + tt_model_constructor, + device, + run_tt_model=True, + log=True, + comp_pcc=None, +): + tt_model, hf_reference_model = tt_model_constructor(device) + + # Prepare input + tokenized = tokenizer(input_sentance, return_tensors="pt") # Batch size 1 + + input_ids = pad_input_32(tokenized.input_ids, hf_reference_model.generation_config.pad_token_id) + attention_mask = pad_input_32(tokenized.attention_mask, 0) + + if log: + logger.debug(f"input_ids {input_ids.shape} {input_ids}") + logger.debug(f"attention_mask {attention_mask.shape} {attention_mask}") + + logits_processor = get_logits_processor(input_ids, hf_reference_model.config) + + decoder_start_values = hf_reference_model.generation_config.pad_token_id * torch.ones(1, 32).to(torch.long) + decoder_input_ids = hf_reference_model.generation_config.pad_token_id * torch.ones(1, 64).to(torch.long) + + if log: + logger.debug(f"decoder_input_ids {decoder_input_ids}") + + encoder_outputs = None + use_cache = False + + for i in range(64): + # PyTorch forward pass + pt_out = hf_reference_model( + input_ids=input_ids, + decoder_input_ids=decoder_input_ids, + attention_mask=attention_mask, + ) + + if run_tt_model: + tt_out = tt_model( + input_ids=input_ids, + decoder_input_ids=decoder_input_ids, + attention_mask=attention_mask, + encoder_outputs=encoder_outputs, + return_dict=True, + use_cache=use_cache, + ) + encoder_outputs = tt_out.encoder_outputs + next_token_logits = tt_out.logits + + if comp_pcc is not None: + does_pass, pcc_message = comp_pcc(pt_out.logits, tt_out.logits, 0.98) + + if log: + logger.info(pcc_message) + else: + next_token_logits = pt_out.logits + + # pre-process distribution + next_tokens_scores = logits_processor(input_ids, next_token_logits) + + # argmax + next_tokens = torch.argmax(next_tokens_scores, dim=-1) + + if log: + logger.debug(f"next_tokens {next_tokens}") + + if next_tokens[0][i] == hf_reference_model.generation_config.eos_token_id: + break + + # We need to expand decoder_input_ids + if (i + 1) % 32 == 0: + decoder_input_ids = torch.cat([decoder_input_ids, decoder_start_values], dim=1) + + decoder_input_ids[0][i + 1] = next_tokens[0][i] + + if log: + logger.debug(f"decoder_input_ids {decoder_input_ids[0]}") + + return tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True) diff --git a/code/models/common/helper_funcs.py b/code/models/common/helper_funcs.py new file mode 100644 index 0000000000000000000000000000000000000000..c371e70e6a056dd7e0cff7ac0da8899322a1bac7 --- /dev/null +++ b/code/models/common/helper_funcs.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import ttnn + + +def Linear( + in_features: int, + out_features: int, + weight: ttnn.Tensor, + bias: Optional[ttnn.Tensor] = None, + output_mem_config=ttnn.DRAM_MEMORY_CONFIG, +): + """ + Returns a function that performs a Linear operation with optional bias. + + ``weight`` must be tt_tensor. + """ + assert weight.padded_shape == [ + 1, + 1, + out_features, + in_features, + ], "weight does not have the expected shape" + + if bias is not None: + assert bias.padded_shape[-1] == out_features, "bias does not have the expected shape" + + weight = weight + bias = bias + weight_T = ttnn.transpose(weight, -2, -1) + + def linear_(activation): + nonlocal bias + assert activation.padded_shape[-1] == in_features, "activation tensor do not have the expected shape" + if bias is not None and bias.get_layout() != ttnn.TILE_LAYOUT: + bias = ttnn.to_layout(bias, ttnn.TILE_LAYOUT) + return ttnn.linear(activation, weight_T, bias=bias, memory_config=output_mem_config) + + return linear_ diff --git a/code/models/common/lightweightmodule.py b/code/models/common/lightweightmodule.py new file mode 100644 index 0000000000000000000000000000000000000000..6743db9026f0547737bda649db0b3ab8c78289c0 --- /dev/null +++ b/code/models/common/lightweightmodule.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + + +class LightweightModule: + """Torch modules add a surprising amount of host overhead for attribute + access and method calls. This class is a lightweight alternative that + just wraps a forward function for now.""" + + def __call__(self, *args, **kwargs): + return self.forward(*args, **kwargs) diff --git a/code/models/common/llama_models.py b/code/models/common/llama_models.py new file mode 100644 index 0000000000000000000000000000000000000000..4e54dd25927c78fd8923f51ec7524ceb39250f55 --- /dev/null +++ b/code/models/common/llama_models.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + + +from dataclasses import dataclass +from enum import Enum +from typing import Dict, List, Optional, Union + +import torch +from PIL import Image +from pydantic import BaseModel, validator + +# ``AutoModelForImageTextToText`` (the replacement for ``AutoModelForVision2Seq``, +# which was removed in transformers 5.x) is only consumed by the +# ``GeneratorChat``/``GeneratorText`` constructors below — defer the imports +# so loading this module doesn't break every downstream import chain +# (e.g. ``tt_transformers.tt.generator``, used by every TT vLLM bridge) +# under transformers >= 5. + + +class Role(Enum): + system = "system" + user = "user" + assistant = "assistant" + ipython = "ipython" + + +class StopReason(Enum): + end_of_turn = "end_of_turn" + end_of_message = "end_of_message" + out_of_tokens = "out_of_tokens" + + +@dataclass +class TokenResult: + token: int + text: str + logprobs: Optional[List[float]] = None + + +@dataclass +class CompletionMessage: + content: str + role: Role = Role.assistant.value + + +class BuiltinTool(Enum): + brave_search = "brave_search" + wolfram_alpha = "wolfram_alpha" + photogen = "photogen" + code_interpreter = "code_interpreter" + + +Primitive = Union[str, int, float, bool, None] +RecursiveType = Union[Primitive, List[Primitive], Dict[str, Primitive]] + + +class ToolCall(BaseModel): + call_id: str + tool_name: Union[BuiltinTool, str] + arguments: Dict[str, RecursiveType] + + @validator("tool_name", pre=True) + @classmethod + def validate_field(cls, v): + if isinstance(v, str): + try: + return BuiltinTool(v) + except ValueError: + return v + return v + + +class ChatPrediction: + generation: CompletionMessage + decoded_tokens: Optional[List[str]] = None + logprobs: Optional[List[List[float]]] = None + + +class CompletionPrediction: + generation: str + decoded_tokens: Optional[List[str]] = None + logprobs: Optional[List[List[float]]] = None + + +def sample_top_p(probs, p): + """ + Perform top-p (nucleus) sampling on a probability distribution. + + Args: + probs (torch.Tensor): Probability distribution tensor. + p (float): Probability threshold for top-p sampling. + + Returns: + torch.Tensor: Sampled token indices. + + Note: + Top-p sampling selects the smallest set of tokens whose cumulative probability mass + exceeds the threshold p. The distribution is renormalized based on the selected tokens. + From: https://github.com/meta-llama/llama-models/blob/v0.1.5/models/llama3/reference_impl/generation.py#L450-L472 + """ + probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True) + probs_sum = torch.cumsum(probs_sort, dim=-1) + mask = probs_sum - probs_sort > p + probs_sort[mask] = 0.0 + probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True)) + next_token = torch.multinomial(probs_sort, num_samples=1) + next_token = torch.gather(probs_idx, -1, next_token) + return next_token + + +def extract_images_from_messages(messages): + images = [] + for message in messages: + if "content" in message: + contents = message["content"] + for content in contents: + if (content["type"] == "image") and ("image" in content): + images.append(content["image"]) + return images + + +def create_vision_mask( + tokens: List[int], + vision_token: int, +) -> List[List[int]]: + """From: https://github.com/meta-llama/llama-models/blob/v0.1.5/models/llama3/api/chat_format.py#L253-L276""" + vision_token_locations = [i for i, token in enumerate(tokens) if token == vision_token] + if len(vision_token_locations) == 0: + return [] + + if len(vision_token_locations) == 1: + # only one image present, unmask until end of sequence + return [[vision_token_locations[0], -1]] + vision_masks = [[loc1, loc2] for loc1, loc2 in zip(vision_token_locations[:-1], vision_token_locations[1:])] + # last image will attend to all subsequent text + vision_masks.append([vision_token_locations[-1], len(tokens)]) + + # if there are two or more consecutive vision tokens, + # they should all attend to all subsequent + # text present + last_mask_end = vision_masks[-1][1] + for vision_mask in vision_masks[::-1]: + if vision_mask[0] == vision_mask[1] - 1: + vision_mask[1] = last_mask_end + last_mask_end = vision_mask[1] + return vision_masks + + +def encode_content(content, images, image_token): + if isinstance(content, Image): + images.append(content) + assert image_token is not None + return image_token + if isinstance(content, str): + return content + if isinstance(content, (list, tuple)): + return "\n".join(encode_content(item, images) for item in content) + if isinstance(content, dict): + content_type = content.get("type") + if content_type == "text": + return content["text"] + if content_type == "image": + # TBD: support url + images.append(content["image"]) + assert image_token is not None + return image_token + raise ValueError(f"Unknown content format: {content}") + + +class GeneratorChat: + def __init__(self, model_name, max_batch_size=1): + from transformers import pipeline + + self.pipe = pipeline("image-text-to-text", model=model_name, batch_size=max_batch_size) + + def chat_completion( + self, + messages, + temperature=0.6, + top_p: float = 0.9, + max_gen_len=None, + ): + generation_output = self.pipe( + text=messages, temperature=temperature, top_p=top_p, max_new_tokens=max_gen_len, return_full_text=False + ) + + if len(generation_output) == 1: + return CompletionMessage(content=generation_output[0]["generated_text"]) + return [CompletionMessage(content=output[0]["generated_text"]) for output in generation_output] + + +class GeneratorText: + def __init__(self, model_name): + from transformers import AutoModelForImageTextToText, AutoProcessor + + self.processor = AutoProcessor.from_pretrained(model_name) + self.model = AutoModelForImageTextToText.from_pretrained(model_name) + + def text_completion( + self, + content: Union[str, Image.Image, Dict, List[Dict]], + temperature: float = 0.6, + top_p: float = 0.9, + max_gen_len=None, + ): + images = [] + text = encode_content(content, images, self.processor.image_token) + model_input = self.processor(text=text, images=images or None, return_tensors="pt", add_special_tokens=False) + tokens = self.model.generate(**model_input, temperature=temperature, top_p=top_p, max_new_tokens=max_gen_len)[0] + tokens = tokens[model_input["input_ids"].shape[-1] :] + return self.processor.decode(tokens, skip_special_tokens=True) diff --git a/code/models/common/metrics.py b/code/models/common/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..852f199626702f99517c5ce176b3fb7c8d300e6b --- /dev/null +++ b/code/models/common/metrics.py @@ -0,0 +1,488 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTNN Metric Functions + +Metric functions for comparing TTNN implementations against reference implementations. +All metrics support both PyTorch and TTNN tensors, with TTNN-native computation +that stays on device until the final scalar result. + +Key Features: +- Automatic detection of tensor type (PyTorch vs TTNN) +- TTNN-native computation using device operations +- Minimal host transfers (only final scalar) +- Graceful fallback to PyTorch for non-tensor inputs +""" + +import numpy as np +import torch +from loguru import logger + +import ttnn + +from .auto_compose import to_torch_auto_compose + +# ====================================================================================== +# Public API +# ====================================================================================== + + +def compute_max_abs_error(impl, ref): + """ + Compute maximum absolute error between two tensors. + + Supports both PyTorch and TTNN tensors. For TTNN tensors, computation + stays on device until the final scalar value. + + Args: + impl: Implementation output (PyTorch or TTNN tensor) + ref: Reference output (PyTorch or TTNN tensor) + + Returns: + float: Maximum absolute difference between tensors + + Examples: + >>> a = torch.tensor([1.0, 2.0, 3.0]) + >>> b = torch.tensor([1.1, 2.0, 2.9]) + >>> _compute_max_abs_error(a, b) + 0.10000002384185791 + """ + try: + if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref): + # TTNN path - stay on device + diff = ttnn.subtract(impl, ref) + abs_diff = _ttnn_op_layout_invariant(diff, ttnn.abs) + return _ttnn_max_scalar_all_dtype(abs_diff) + elif _is_ttnn_tensor(impl): + return (to_torch_auto_compose(impl) - ref).abs().max().item() + elif _is_ttnn_tensor(ref): + return (impl - to_torch_auto_compose(ref)).abs().max().item() + else: + # PyTorch path - both must be torch tensors + return (impl - ref).abs().max().item() + except Exception as e: + return float("inf") + + +def compute_mean_abs_error(impl, ref): + """ + Compute mean absolute error between two tensors. + + Supports both PyTorch and TTNN tensors. For TTNN tensors, computation + stays on device until the final scalar value. + + Args: + impl: Implementation output (PyTorch or TTNN tensor) + ref: Reference output (PyTorch or TTNN tensor) + + Returns: + float: Mean absolute difference between tensors + + Examples: + >>> a = torch.tensor([1.0, 2.0, 3.0]) + >>> b = torch.tensor([1.1, 2.0, 2.9]) + >>> _compute_mean_abs_error(a, b) + 0.06666667014360428 + """ + try: + if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref): + # TTNN path - stay on device + diff = ttnn.subtract(impl, ref) + abs_diff = _ttnn_op_layout_invariant(diff, ttnn.abs) + return _ttnn_mean_scalar_all_dtype(abs_diff) + elif _is_ttnn_tensor(impl): + return (to_torch_auto_compose(impl) - ref).abs().mean().item() + elif _is_ttnn_tensor(ref): + return (impl - to_torch_auto_compose(ref)).abs().mean().item() + else: + # PyTorch path - both must be torch tensors + return (impl - ref).abs().mean().item() + except Exception as e: + return float("inf") + + +def compute_pcc(impl, ref): + """ + Compute Pearson Correlation Coefficient (PCC) between two tensors. + + PCC measures the linear correlation between two tensors. Returns 1.0 for + perfect positive correlation, 0.0 for no correlation, and -1.0 for perfect + negative correlation. This is a common metric used in tt-metal for validating + tensor computations. + + For TTNN tensors: Uses on-device computation (100-1000× faster). If edge cases + produce non-finite values (NaN/inf), raises ValueError with instructions to use + Pattern 2 (PyTorch reference with output_map for robust CPU-based PCC). + + For PyTorch tensors: Uses robust numpy implementation (from tt-metal + comparison_funcs.py) that handles complex, NaN, inf, and constant tensors. + + Args: + impl: Implementation output (PyTorch or TTNN tensor) + ref: Reference output (PyTorch or TTNN tensor) + + Returns: + float: PCC value in range [-1.0, 1.0], or 0.0 on error + + Raises: + ValueError: If TTNN-native computation produces non-finite values + + Examples: + >>> a = torch.tensor([1.0, 2.0, 3.0]) + >>> b = torch.tensor([1.0, 2.0, 3.0]) + >>> _compute_pcc(a, b) + 1.0 + + >>> a = torch.tensor([1.0, 2.0, 3.0]) + >>> b = torch.tensor([3.0, 2.0, 1.0]) + >>> _compute_pcc(a, b) + -1.0 + """ + try: + # TTNN fast path - compute on device + if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref): + return compute_pcc_device(impl, ref) + elif _is_ttnn_tensor(impl): + return compute_pcc_host(to_torch_auto_compose(impl), ref) + elif _is_ttnn_tensor(ref): + return compute_pcc_host(impl, to_torch_auto_compose(ref)) + else: + return compute_pcc_host(impl, ref) + except Exception: + return 0.0 + + +# Default metrics dictionary for easy import +DEFAULT_METRICS = { + "max_abs_error": compute_max_abs_error, + "mean_abs_error": compute_mean_abs_error, + "pcc": compute_pcc, +} + + +# Allclose comparison with detailed delta string +def comp_allclose(impl, ref, rtol=1e-05, atol=1e-08): + """ + Compare two tensors using an allclose criterion and return (passing, details). + + Provides both a TTNN-native on-device implementation and a PyTorch fallback. + Mirrors semantics of torch.allclose(..., equal_nan=True) and reports the + maximum absolute and relative deltas similar to comparison_funcs.py. + + Args: + impl: Implementation output (PyTorch or TTNN tensor) + ref: Reference output (PyTorch or TTNN tensor) + rtol (float): Relative tolerance + atol (float): Absolute tolerance + + Returns: + tuple[bool, str]: (passing, "Max ATOL Delta: x, Max RTOL Delta: y[,...]") + """ + try: + # TTNN-native path: compute deltas and allclose on device, then transfer final scalars + if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref): + # Compute deltas (device) + diff = _ttnn_op_layout_invariant(ttnn.subtract(impl, ref), ttnn.abs) + cal_atol = _ttnn_max_scalar_all_dtype(diff) + # For rtol delta, divide by abs(ref) (may produce inf for zeros; acceptable for reporting) + divided_by_ref = ttnn.divide(diff, _ttnn_op_layout_invariant(ref, ttnn.abs)) + cal_rtol = _ttnn_max_scalar_all_dtype(divided_by_ref) + + # equal_nan=True semantics and finite/infinite handling + isnan_impl = _ttnn_op_layout_invariant(impl, ttnn.isnan) + isnan_ref = _ttnn_op_layout_invariant(ref, ttnn.isnan) + both_nan = ttnn.logical_and(isnan_impl, isnan_ref) + + isinf_impl = _ttnn_op_layout_invariant(impl, ttnn.isinf) + isinf_ref = _ttnn_op_layout_invariant(ref, ttnn.isinf) + impl_sign = _ttnn_op_layout_invariant(impl, ttnn.sign) + ref_sign = _ttnn_op_layout_invariant(ref, ttnn.sign) + same_sign_inf = ttnn.eq(impl_sign, ref_sign) + both_inf_same_sign = ttnn.logical_and(ttnn.logical_and(isinf_impl, isinf_ref), same_sign_inf) + + # Finite elements where numeric closeness applies + any_nan = ttnn.logical_or(isnan_impl, isnan_ref) + any_inf = ttnn.logical_or(isinf_impl, isinf_ref) + finite_both = _ttnn_op_layout_invariant(ttnn.logical_or(any_nan, any_inf), ttnn.logical_not) + + # |impl - ref| <= atol + rtol * |ref| + bound = ttnn.add(ttnn.mul(_ttnn_op_layout_invariant(ref, ttnn.abs), rtol, dtype=ttnn.bfloat16), atol) + close_numeric = ttnn.le(diff, bound) + finite_and_close = ttnn.logical_and(finite_both, close_numeric) + + ok_mask = ttnn.logical_or(ttnn.logical_or(both_nan, both_inf_same_sign), finite_and_close) + fail_mask = _ttnn_op_layout_invariant(ok_mask, ttnn.logical_not) + + # Reduce to scalar: any failure -> 1.0 else 0.0 + fail_indicator = _ttnn_op_layout_invariant(fail_mask, ttnn.where, true_value=1.0, false_value=0.0) + any_fail = _ttnn_max_scalar_all_dtype(fail_indicator) + passing = any_fail == 0.0 + + output_str = f"Max ATOL Delta: {cal_atol}, Max RTOL Delta: {cal_rtol}" + if not passing: + output_str += ", Allclose check failed" + return passing, output_str + + # Fallback: compute with PyTorch (handles mixed inputs by converting TTNN -> torch) + impl_torch = to_torch_auto_compose(impl) if _is_ttnn_tensor(impl) else impl + ref_torch = to_torch_auto_compose(ref) if _is_ttnn_tensor(ref) else ref + + if torch.is_tensor(impl_torch) and torch.is_tensor(ref_torch): + # Match dtype for fair comparison + if impl_torch.dtype != ref_torch.dtype: + ref_torch = ref_torch.to(impl_torch.dtype) + + atol_delta = torch.max(torch.abs(impl_torch - ref_torch)).item() + # May produce inf where ref == 0; this mirrors comparison_funcs.py behavior + rtol_delta = torch.max(torch.abs(impl_torch - ref_torch) / torch.abs(ref_torch)).item() + passing = torch.allclose(impl_torch, ref_torch, rtol, atol, True) + output_str = f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}" + if not passing: + output_str += ", Allclose check failed" + return passing, output_str + + # Unsupported types + return False, "Unsupported input types for comp_allclose" + except Exception as e: + return False, f"Error computing comp_allclose: {e}" + + +def compute_pcc_device(impl, ref): + """Fast on-device PCC for TTNN tensors only.""" + try: + # Early edge-case handling to mirror CPU semantics + # - All NaNs → 1.0; mixed NaNs → 0.0 + # - One tensor all zero and the other not → 0.0 + # - Both constant → 1.0 if equal, else 0.0 + + # Any nonzero check (all-zero detection) + impl_abs_max = _ttnn_max_scalar_all_dtype(_ttnn_op_layout_invariant(impl, ttnn.abs)) + ref_abs_max = _ttnn_max_scalar_all_dtype(_ttnn_op_layout_invariant(ref, ttnn.abs)) + + impl_has_any = impl_abs_max != 0.0 + ref_has_any = ref_abs_max != 0.0 + if impl_has_any != ref_has_any: + return 0.0 + + # Min/Max scalars for constant and NaN detection + impl_min = _ttnn_min_scalar_all_dtype(impl) + impl_max = _ttnn_max_scalar_all_dtype(impl) + ref_min = _ttnn_min_scalar_all_dtype(ref) + ref_max = _ttnn_max_scalar_all_dtype(ref) + + impl_min_finite = np.isfinite(impl_min) + impl_max_finite = np.isfinite(impl_max) + ref_min_finite = np.isfinite(ref_min) + ref_max_finite = np.isfinite(ref_max) + + impl_all_nan = (not impl_min_finite) and (not impl_max_finite) + ref_all_nan = (not ref_min_finite) and (not ref_max_finite) + if impl_all_nan and ref_all_nan: + return 1.0 + if impl_all_nan != ref_all_nan: + return 0.0 + + # Constant tensors + if impl_min_finite and impl_max_finite and ref_min_finite and ref_max_finite: + if impl_min == impl_max and ref_min == ref_max: + return ( + 1.0 + if torch.isclose( + torch.tensor(impl_max, dtype=torch.float32), torch.tensor(ref_max, dtype=torch.float32) + ) + else 0.0 + ) + + # Standard PCC formula on device + mean_impl = _ttnn_mean_scalar_all_dtype(impl) + mean_ref = _ttnn_mean_scalar_all_dtype(ref) + + impl_centered = ttnn.subtract(impl, mean_impl) + ref_centered = ttnn.subtract(ref, mean_ref) + + # todo)) ttnn.sum() does local reduction only; need CCL reduction for global sum when adding support for multiple-devices + # [INFO] we cast to float32 to avoid overflow when impl and ref are in bfloat8_b or bfloat4_b + numerator = ttnn.sum(ttnn.mul(impl_centered, ref_centered, dtype=ttnn.float32)) + impl_sq_sum = ttnn.sum(ttnn.mul(impl_centered, impl_centered, dtype=ttnn.float32)) + ref_sq_sum = ttnn.sum(ttnn.mul(ref_centered, ref_centered, dtype=ttnn.float32)) + denominator = ttnn.sqrt(ttnn.mul(impl_sq_sum, ref_sq_sum, dtype=ttnn.float32)) + + # Safe divide + denom_scalar = denominator.item() + if denom_scalar == 0.0 or not np.isfinite(denom_scalar): + return 0.0 + + pcc = numerator.item() / denom_scalar + if not np.isfinite(pcc): + return 0.0 + return pcc + except Exception as e: + # todo)) maybe return a string for logging: f"Error computing PCC on device: {e}, impl: {impl}, ref: {ref}" + return 0.0 + + +# code stolen from tests/tt_eager/python_api_testing/sweep_tests/comparison_funcs.py +# and models/common/utility_functions.py +def compute_pcc_host(impl, ref): + """Robust CPU PCC for PyTorch tensors only.""" + try: + calculated = impl + golden = ref + if golden.dtype != calculated.dtype: + calculated = calculated.type(golden.dtype) + + # Handle complex tensors + if golden.is_complex() and calculated.is_complex(): + golden = torch.view_as_real(golden.clone()) + calculated = torch.view_as_real(calculated.clone()) + + # Convert to float if needed + if not (golden.is_floating_point() or calculated.is_floating_point()): + golden = golden.to(torch.float) + calculated = calculated.to(torch.float) + + # Both tensors are nan + if torch.all(torch.isnan(golden)) and torch.all(torch.isnan(calculated)): + return 1.0 + + # One tensor is all nan, the other is not + if torch.all(torch.isnan(golden)) or torch.all(torch.isnan(calculated)): + return 0.0 + + # One tensor is all zero, the other is not — also a zero-variance case. + if torch.any(golden.bool()) != torch.any(calculated.bool()): + logger.warning("One tensor is all zero. PCC undefined; falling back to allclose.") + return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04)) + + # Mask all infs and nans + golden = golden.clone() + golden[ + torch.logical_or( + torch.isnan(golden), + torch.logical_or(torch.isinf(golden), torch.isneginf(golden)), + ) + ] = 0 + calculated = calculated.clone() + calculated[ + torch.logical_or( + torch.isnan(calculated), + torch.logical_or(torch.isinf(calculated), torch.isneginf(calculated)), + ) + ] = 0 + + if torch.equal(golden, calculated): + return 1.0 + + if golden.dtype == torch.bfloat16: + golden = golden.type(torch.float32) + calculated = calculated.type(torch.float32) + + # Single element or constant tensor: PCC is undefined. + if golden.numel() == 1: + return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04)) + + if torch.max(golden) == torch.min(golden) or torch.max(calculated) == torch.min(calculated): + logger.warning("One or both tensors are constant (zero std dev). PCC undefined; falling back to allclose.") + return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04)) + + # Compute PCC using numpy's corrcoef + cal_pcc = np.ma.corrcoef( + np.ma.masked_invalid(torch.squeeze(golden).detach().numpy()).flatten(), + np.ma.masked_invalid(torch.squeeze(calculated).detach().numpy()).flatten(), + ) + # Read off-diagonal directly to avoid diagonal contamination. + cal_pcc = cal_pcc[0, 1] + + if isinstance(cal_pcc, np.ma.core.MaskedConstant) or np.isnan(float(cal_pcc)): + logger.warning("PCC returned NaN/masked. Falling back to allclose.") + return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04)) + + return float(cal_pcc) + except Exception: + return 0.0 + + +# ====================================================================================== +# Private Implementation +# ====================================================================================== + + +def _is_ttnn_tensor(x): + """Safely detect TTNN tensors even if ttnn.Tensor is not defined in this environment.""" + return isinstance(x, ttnn.Tensor) + + +def _ttnn_op_layout_invariant(x, op_func, **kwargs): + """ + Generic helper for TTNN operations that require TILE layout for non-sharded tensors. + + Args: + x: TTNN tensor + op_func: TTNN operation function to call (e.g., ttnn.sign, ttnn.abs, ttnn.typecast) + **kwargs: Additional keyword arguments to pass to op_func (e.g., dtype=ttnn.bfloat16 for typecast) + + Returns: + Result of op_func applied to x, with layout preserved. + """ + # [ttnn contract] all TTNN operations require tensors to be in TILE layout when working with non-sharded tensors: + # ttnn.sign(), ttnn.abs(), ttnn.isinf(), ttnn.isnan(), ttnn.typecast(), ttnn.where(), ttnn.logical_not() + layout = x.get_layout() + if layout == ttnn.TILE_LAYOUT or x.is_sharded(): # sharded tensors can use either layout + return op_func(x, **kwargs) + else: + return ttnn.to_layout(op_func(ttnn.to_layout(x, ttnn.TILE_LAYOUT), **kwargs), layout) + + +def _ttnn_max_scalar_all_dtype(x): + x_bf16 = _ttnn_op_layout_invariant(x, ttnn.typecast, dtype=ttnn.bfloat16) + + # [ttnn contract] ttnn.max() internally calls a FillPad operation, which only supports the following dtypes: + # BFLOAT16 + # FLOAT32 + # UINT16, UINT32, INT32 + # UINT8 + # see ttnn/cpp/ttnn/operations/data_movement/fill_pad/device/fill_pad_program_factory.hpp for more details + # [INFO] we cast to bfloat16 to avoid the limitation of the FillPad operation when dealing with e.g., bfloat8_b and bfloat4_b dtypes + + # [ttnn contract] When called without a dim parameter, ttnn.max() returns the maximum value across the entire tensor as a scalar. + max_val_tensor = ttnn.max(x_bf16) + + # [ttnn contract] When called without a dim parameter, ttnn.max() returns the maximum value across the entire tensor as a scalar. + # The method supports multiple data types: + # FLOAT32 → Python float + # BFLOAT16 → Python float (cast from bfloat16) + # BFLOAT8_B and BFLOAT4_B → Python float + return max_val_tensor.item() + + +def _ttnn_mean_scalar_all_dtype(x): + x_bf16 = _ttnn_op_layout_invariant(x, ttnn.typecast, dtype=ttnn.bfloat16) + + # [ttnn contract] ttnn.mean() internally calls a FillPad operation, which only supports the following dtypes: + # BFLOAT16 + # FLOAT32 + # UINT16, UINT32, INT32 + # UINT8 + # see ttnn/cpp/ttnn/operations/data_movement/fill_pad/device/fill_pad_program_factory.hpp for more details + # [INFO] we cast to bfloat16 to avoid the limitation of the FillPad operation when dealing with e.g., bfloat8_b and bfloat4_b dtypes + + # [ttnn contract] When called without a dim parameter, ttnn.mean() returns the mean value across the entire tensor as a scalar. + mean_val_tensor = ttnn.mean(x_bf16) + return mean_val_tensor.item() + + +def _ttnn_min_scalar_all_dtype(x): + x_bf16 = _ttnn_op_layout_invariant(x, ttnn.typecast, dtype=ttnn.bfloat16) + + # [ttnn contract] ttnn.mean() internally calls a FillPad operation, which only supports the following dtypes: + # BFLOAT16 + # FLOAT32 + # UINT16, UINT32, INT32 + # UINT8 + # see ttnn/cpp/ttnn/operations/data_movement/fill_pad/device/fill_pad_program_factory.hpp for more details + # [INFO] we cast to bfloat16 to avoid the limitation of the FillPad operation when dealing with e.g., bfloat8_b and bfloat4_b dtypes + + # [ttnn contract] When called without a dim parameter, ttnn.mean() returns the mean value across the entire tensor as a scalar. + min_val_tensor = ttnn.min(x_bf16) + return min_val_tensor.item() diff --git a/code/models/common/model_capabilities.py b/code/models/common/model_capabilities.py new file mode 100644 index 0000000000000000000000000000000000000000..cb8300277e8b13803ce79713daeeec793229bfc7 --- /dev/null +++ b/code/models/common/model_capabilities.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +# Maximum number of tokens that can simultaneously occupy the KV cache +# across all concurrent users. This fallback applies to model/device +# configurations not covered by a model-specific override. +# Derived from the default branch of the per-model KV-cache rules in +# the TT vLLM worker (tenstorrent/vllm#315). +# See also: https://github.com/tenstorrent/vllm/issues/315 +FALLBACK_MAX_TOKENS_ALL_USERS = 131_072 + + +class ModelCapabilitiesMixin: + """Defines interface for hardware- or model-specific configurations. + + NOTE: The default values here and per-model overrides will eventually be + unified with the corresponding vLLM scheduler configuration so that both + paths derive from the same source of truth. + """ + + @classmethod + def get_max_tokens_all_users(cls, **kwargs) -> int: + """Returns the fallback all-user KV-cache token capacity. + + Used when no model- or device-specific override applies. + """ + return FALLBACK_MAX_TOKENS_ALL_USERS diff --git a/code/models/common/rmsnorm.py b/code/models/common/rmsnorm.py new file mode 100644 index 0000000000000000000000000000000000000000..3e4e766e299cf33e454b72850868d12ed225cf6f --- /dev/null +++ b/code/models/common/rmsnorm.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.utility_functions import copy_to_buffer +from models.tt_transformers.tt.common import Mode + +TILE = 32 +SHARD_HEIGHT = TILE # Current ttnn.rms_norm implementation requires shard height to be a single tile + + +class RMSNorm(LightweightModule): + """ + RMSNorm supporting replication over a MeshDevice and sharding within devices. + + This class implements a Root Mean Square Normalization (RMSNorm) that can be + distributed across multiple devices and cores. If the `device` parameter is a + MeshDevice, the weights and computations are replicated across all devices in + the mesh. Expects an interleaved input tensor, can optionally output a sharded tensor. + + Args: + device: The device or MeshDevice on which to perform the computations. + state_dict: The state dictionary containing the model parameters. + dim: Input dimension (e.g. model hidden dimension size). + layer_num: The layer number to determine the weight key in the state dictionary. + weight_key: The key for retrieving the weight from the state dictionary. + weight_cache_path: Optional path for caching the tilized weights. + weight_memory_config: Configuration for the weight memory, default is DRAM_MEMORY_CONFIG. + weight_dtype: The data type for the tensors, bfp8_b hits >0.999 PCC in the models we tested. + model_config: Optional configuration dictionary for the model. + eps (float): Small value to avoid division by zero in normalization, default is 1e-05. + + If model_config is provided, it must specify SHARDED_NORM_INPUT_MEMCFG, SHARDED_NORM_PRGM_CFG + and SHARDED_NORM_OUTPUT_MEMCFG. If not provided, default configurations will be generated. + """ + + def __init__( + self, + device, + dim, + state_dict, + weight_key, + layer_num=None, + state_dict_prefix=None, + weight_cache_path=None, + weight_memory_config=ttnn.DRAM_MEMORY_CONFIG, + weight_dtype=ttnn.bfloat16, + is_distributed=None, + eps: float = 1e-05, + add_unit_offset=False, + sharded_program_config=None, + sharded_output_config=None, + output_mem_config=None, + ccl_topology=ttnn.Topology.Ring, + tt_ccl=None, + fp32_dest_acc_en=True, + ): + super().__init__() + self.device = device + self.eps = eps + self.is_distributed = is_distributed + self.ccl_topology = ccl_topology + self.tt_ccl = tt_ccl + self.add_unit_offset = add_unit_offset + + if state_dict_prefix: + weight_name = f"{state_dict_prefix}{weight_key}.weight" + else: + if layer_num is None: + weight_name = f"{weight_key}.weight" + else: + weight_name = f"layers.{layer_num}.{weight_key}.weight" + + torch_weight = ( + state_dict[weight_name].unsqueeze(0).view(1, 1, dim).reshape([1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT]) + ) + + # Add offset before caching + if add_unit_offset: + torch_weight = torch_weight + 1.0 + + # Compatibility with models that don't use mesh devices (e.g. single-chip Mistral-7b) + is_mesh_device = device.__class__.__name__ == "MeshDevice" + + self.weight = ttnn.as_tensor( + torch_weight, + device=device, + dtype=weight_dtype, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=weight_memory_config, + cache_file_name=None if weight_cache_path is None else weight_cache_path / weight_name, + mesh_mapper=ttnn.ReplicateTensorToMesh(device) if is_mesh_device else None, + ) + + if self.is_distributed: + self.weight_distributed = ttnn.as_tensor( + torch_weight, + device=device, + dtype=weight_dtype, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=weight_memory_config, + cache_file_name=( + None if weight_cache_path is None else weight_cache_path / (weight_name + "_distributed") + ), + mesh_mapper=( + ttnn.ShardTensor2dMesh(device, dims=(None, 2), mesh_shape=list(device.shape)) + if is_mesh_device + else None + ), + ) + + self.sharded_output_config = sharded_output_config + self.sharded_program_config = sharded_program_config + self.output_mem_config = output_mem_config + + self.compute_kernel_config_hifi2 = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=fp32_dest_acc_en, + packer_l1_acc=True, + ) + + def update(self, *, weight: ttnn.Tensor) -> None: + """In-place replace the RMSNorm gamma via ``ttnn.copy``. + + HF-format input: ``weight`` is HF ``...norm.weight``, shape + ``(1, 1, 1, dim)``, bf16, TILE, DRAM-interleaved, replicated. + + ``copy_to_buffer`` reshapes to the storage shape + ``(1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT)`` and TILE -> ROW_MAJOR to + match ``self.weight``. ``add_unit_offset`` is not supported (see + assert): the caller must ship a gamma that already includes the +1. + + When ``self.weight_distributed`` (the column-sharded mirror) exists it's + kept in sync on device: project ``self.weight`` into the sharded layout + via ``ttnn.mesh_partition`` (the inverse of the constructor's + ``ShardTensor2dMesh(dims=(None, 2))``, hence ``dim=2, cluster_axis=1``) + and ``ttnn.copy`` into it. Both buffers keep their address, so captured + traces and the prefetcher's recorded addresses stay valid. + """ + assert not self.add_unit_offset, "RMSNorm.update does not support add_unit_offset=True" + copy_to_buffer(weight, self.weight, self.weight.dtype) + + if getattr(self, "weight_distributed", None) is not None: + partitioned = ttnn.mesh_partition( + self.weight, + memory_config=self.weight_distributed.memory_config(), + dim=2, + cluster_axis=1, + ) + copy_to_buffer(partitioned, self.weight_distributed, self.weight_distributed.dtype) + + def forward( + self, + x: ttnn.Tensor, + mode: Mode | str, + in_sharded=False, + out_sharded=False, + norm_config=None, + ) -> ttnn.Tensor: + if isinstance(mode, str): + try: + mode = Mode(mode) + except ValueError: + raise ValueError(f"Invalid mode: {mode}") + elif not isinstance(mode, Mode): + raise ValueError(f"Invalid mode: {mode}") + + sharded_program_config = norm_config.get("sharded_program_config") if norm_config else None + sharded_output_config = norm_config.get("sharded_output_config") if norm_config else None + output_mem_config = norm_config.get("output_mem_config") if norm_config else None + # Optional L1 placement for the distributed 3-op outputs (pre/gather/post); None -> DRAM default. + distributed_out_mc = norm_config.get("distributed_output_mem_config") if norm_config else None + + # If input is sharded do sharded RMSNorm and optionally return sharded output + program_config = sharded_program_config if in_sharded else None + memory_config = sharded_output_config if out_sharded else None + distributed = self.is_distributed and self.is_distributed(mode) + weight = self.weight_distributed if distributed else self.weight + + if in_sharded: + assert not distributed, "Distributed RMSNorm does not support sharded inputs" + else: + assert not out_sharded, "Non-sharded version of RMSNorm cannot output a sharded tensor" + + if distributed: + x = self._distributed_rmsnorm( + x, + epsilon=self.eps, + weight=weight, + compute_kernel_config=self.compute_kernel_config_hifi2, + output_memory_config=distributed_out_mc, + ) + else: + x = ttnn.rms_norm( + x, + epsilon=self.eps, + weight=weight, + program_config=program_config, + memory_config=memory_config, + compute_kernel_config=self.compute_kernel_config_hifi2, + ) + + if in_sharded and not out_sharded: + return ttnn.sharded_to_interleaved(x) + else: + if output_mem_config is not None: + x = ttnn.to_memory_config(x, output_mem_config) + return x + + def _distributed_rmsnorm( + self, + inp, + epsilon=None, + weight=None, + program_config=None, + memory_config=None, + compute_kernel_config=None, + output_memory_config=None, + ): + assert program_config is None, "Distributed RMSNorm does not support sharded inputs" + assert memory_config is None, "Distributed RMSNorm does not support sharded outputs" + assert self.tt_ccl is not None, "Distributed RMSNorm requires tt_ccl" + + # Interleaved output placement for the 3 ops; default DRAM (matches the prior hardcoded behavior). + mc = output_memory_config if output_memory_config is not None else ttnn.DRAM_MEMORY_CONFIG + + # Run distributed rmsnorm part 1 + tt_stats = ttnn.rms_norm_pre_all_gather( + inp, compute_kernel_config=compute_kernel_config, dtype=ttnn.bfloat16, memory_config=mc + ) + # AllGather stats + tt_stats = ttnn.experimental.all_gather_async( + tt_stats, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=self.ccl_topology, + memory_config=mc, + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + # Run distributed rmsnorm part 2 + tt_out = ttnn.rms_norm_post_all_gather( + inp, + tt_stats, + epsilon=epsilon, + weight=weight, + compute_kernel_config=compute_kernel_config, + memory_config=mc, + ) + tt_stats.deallocate(True) + + return tt_out diff --git a/code/models/common/tensor_utils.py b/code/models/common/tensor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..382e92e35d38fb80d60b80e2cd9206538b8fe417 --- /dev/null +++ b/code/models/common/tensor_utils.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Tensor utility functions for TTTv2 modules. +""" + +import json +import math +import re + +import torch + +import ttnn + +# Standard tile size - hardware constant +TILE_SIZE = ttnn.TILE_SIZE # 32 + + +def nearest_multiple(value: int, multiple: int) -> int: + return math.ceil(value / multiple) * multiple + + +def nearest_32(value: int) -> int: + return nearest_multiple(value, TILE_SIZE) + + +def num_to_core_range_set(num_cores: int): + assert num_cores < 8 or num_cores % 8 == 0 + num_x = min(num_cores, 8) + num_y = num_cores // num_x + assert num_x * num_y == num_cores + return ttnn.CoreRangeSet( + { + ttnn.CoreRange( + ttnn.CoreCoord(0, 0), + ttnn.CoreCoord(num_x - 1, num_y - 1), + ) + } + ) + + +def get_out_subblock_w(per_core_n: int, out_subblock_h: int = 1) -> int: + out_subblock_w = 4 + while out_subblock_w > 1: + if out_subblock_w * out_subblock_h <= 4 and per_core_n % out_subblock_w == 0: + break + out_subblock_w -= 1 + return out_subblock_w + + +def align_shape_to_tile(shape, tile_size: int = TILE_SIZE): + """Round up the last two dimensions of *shape* to multiples of *tile_size*. + + This is the recommended replacement for the deprecated ``ttnn.pad_to_tile_shape``. + + Args: + shape: An iterable of dimension sizes (list, tuple, or ttnn.Shape). + tile_size: Tile dimension to align to (default 32). + + Returns: + List[int]: A new shape with the last two dims tile-aligned. + """ + import math + + result = list(shape) + if len(result) >= 1: + result[-1] = math.ceil(result[-1] / tile_size) * tile_size + if len(result) >= 2: + result[-2] = math.ceil(result[-2] / tile_size) * tile_size + return result + + +def get_rot_transformation_mat(dhead: int = TILE_SIZE) -> torch.Tensor: + """ + Create rotation transformation matrix for RoPE. + + Constructs a permutation matrix that pairs adjacent dimensions with + signs (+1, -1) for the RoPE rotation: + [0, 1] → +1 at (0,1), -1 at (1,0) + [2, 3] → +1 at (2,3), -1 at (3,2) + ... + + Used by ttnn.experimental.rotary_embedding_llama. + + Args: + dhead: Matrix dimension. Must equal TILE_SIZE. Use TILE_SIZE for decode. + + Returns: + torch.Tensor of shape [1, 1, dhead, dhead]. + """ + rot_emb_matrix = torch.zeros(1, 1, dhead, dhead) + rot_emb_matrix[..., torch.arange(0, dhead, 2), torch.arange(1, dhead, 2)] = 1 + rot_emb_matrix[..., torch.arange(1, dhead, 2), torch.arange(0, dhead, 2)] = -1 + return rot_emb_matrix + + +def zeros_like_kv_cache(batch_size: int, n_kv_heads: int, max_seq_len: int, head_dim: int) -> torch.Tensor: + """Create zeros tensor for standard KV cache.""" + return torch.zeros((batch_size, n_kv_heads, max_seq_len, head_dim)) + + +def zeros_like_paged_cache(paged_config, n_kv_heads: int, head_dim: int) -> torch.Tensor: + """Create zeros tensor for paged KV cache.""" + return torch.zeros((paged_config.max_num_blocks, n_kv_heads, paged_config.block_size, head_dim)) + + +# todo)) add a on-device pad_dim_to_size function? +def pad_dim_to_size(x: "torch.Tensor", dim: int, size: int) -> "torch.Tensor": + """Pads the specified dimension of the input tensor with zeros.""" + if dim < 0: + dim = x.dim() + dim + current_size = x.size(dim) + pad_size = size - current_size + + if pad_size < 0: + raise ValueError(f"Target size {size} is smaller than current size {current_size} on dim {dim}") + + if pad_size == 0: + return x + + pad = [0] * (2 * x.dim()) + pad_index = 2 * (x.dim() - dim - 1) + pad[pad_index + 1] = pad_size + + return torch.nn.functional.pad(x, pad, mode="constant", value=0) + + +def pad_to_shape(x: "torch.Tensor", target_shape: tuple[int, ...], pad_value: float = 0.0) -> "torch.Tensor": + """Pad tensor to target_shape in a single F.pad call (more efficient than per-dim padding).""" + if x.shape == target_shape: + return x + + # F.pad expects: (left_last, right_last, left_second_last, right_second_last, ...) + pad = [] + for orig, target in zip(reversed(x.shape), reversed(target_shape)): + if target < orig: + raise ValueError(f"Target size {target} is smaller than current size {orig}") + pad.extend([0, target - orig]) + + return torch.nn.functional.pad(x, pad, mode="constant", value=pad_value) + + +def get_padded_hidden_dim(hidden_dim: int, num_devices: int, tile_size: int = 32) -> int: + """ + Compute padded hidden_dim to satisfy ttnn.from_torch's tile alignment constraint. + + ttnn.from_torch requires physical shard shapes to be tile-aligned. When sharding + a tensor across devices, each shard_dim = hidden_dim / num_devices must be + divisible by tile_size. + + We pad the global tensor first, then shard evenly so only the last shard has padding. + """ + shard_dim = hidden_dim // num_devices + padded_shard = ((shard_dim + tile_size - 1) // tile_size) * tile_size + return padded_shard * num_devices + + +def parse_shard_dims_from_mesh_mapper_config(mesh_mapper_config: ttnn.MeshMapperConfig) -> list[int]: + """ + Parse shard dimensions from MeshMapperConfig's repr. + + MeshMapperConfig doesn't expose .placements directly, but repr shows them: + 'MeshMapperConfig(placements: [PlacementShard(-1)], mesh_shape_override=MeshShape([8]))' + + This parses out the shard dimensions (e.g., [-1]) from PlacementShard entries. + Returns empty list if no PlacementShard found (e.g., replicated). + + Note: This is a workaround until TTNN exposes .placements directly. + """ + config_repr = repr(mesh_mapper_config) + matches = re.findall(r"PlacementShard\((-?\d+)\)", config_repr) + return [int(d) for d in matches] + + +def memory_config_to_dict(memory_config: ttnn.MemoryConfig): + # Convert to plain types for deterministic serialization. + return { + "memory_layout": str(memory_config.memory_layout), + "buffer_type": str(memory_config.buffer_type), + "shard_spec": str(memory_config.shard_spec), + "is_sharded": bool(memory_config.is_sharded()), + "interleaved": bool(memory_config.interleaved), + "hash": int(memory_config.__hash__()), + } + + +def compute_kernel_config_to_str(compute_kernel_config: ttnn.WormholeComputeKernelConfig): + # Backward compat shim; prefer compute_kernel_config_to_dict + serialize_config. + cfg = compute_kernel_config_to_dict(compute_kernel_config) + return serialize_config(cfg) + + +def compute_kernel_config_to_dict(compute_kernel_config: ttnn.WormholeComputeKernelConfig): + return { + "math_fidelity": str(compute_kernel_config.math_fidelity), + "math_approx_mode": str(compute_kernel_config.math_approx_mode), + "fp32_dest_acc_en": bool(compute_kernel_config.fp32_dest_acc_en), + "packer_l1_acc": bool(compute_kernel_config.packer_l1_acc), + "dst_full_sync_en": bool(compute_kernel_config.dst_full_sync_en), + "throttle_level": str(compute_kernel_config.throttle_level), + } + + +def program_config_to_str(program_config: ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig): + # Backward compat shim; prefer program_config_to_dict + serialize_config. + cfg = program_config_to_dict(program_config) + return serialize_config(cfg) + + +def program_config_to_dict(program_config): + if hasattr(program_config, "to_json"): + d = json.loads(program_config.to_json()) + d["type"] = type(program_config).__name__ + return d + else: + return {"type": type(program_config).__name__, "repr": repr(program_config)} + + +def serialize_config(cfg_dict: dict, fmt: str = "json") -> str: + if fmt == "json": + return json.dumps(cfg_dict, sort_keys=True) + if fmt == "yaml": + try: + import yaml + except ImportError as exc: # pragma: no cover - optional dependency + raise RuntimeError("PyYAML is required for yaml serialization") from exc + return yaml.safe_dump(cfg_dict, sort_keys=True) + raise ValueError(f"Unsupported format: {fmt}") diff --git a/code/models/common/utility_functions.py b/code/models/common/utility_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..9011d2df6425fafc66226087af0d9851b0dbed6d --- /dev/null +++ b/code/models/common/utility_functions.py @@ -0,0 +1,1270 @@ +# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + +import math +import os +import struct +import time +from typing import Union + +import numpy as np +import torch +from loguru import logger +from ttnn.device import Arch +from typing_extensions import deprecated + +import ttnn + + +def get_mesh_device(): + """Fixture to provide mesh device configuration.""" + mesh_device = os.environ.get("MESH_DEVICE", "N150") + mesh_config = { + "N150": (1, 1), + "N300": (2, 1), + "T3K": (8, 1), + "TG": (8, 4), + }.get(mesh_device, (ttnn.get_num_devices(), 1)) + return mesh_config + + +### Math operations ### +def _nearest_32(x): + return math.ceil(x / 32) * 32 + + +def nearest_32( + x, +): # needs refctoring; to match alias called in some scripts (e.g. test_padding_test in unit tests) + return _nearest_32(x) + + +def _nearest_y(x, y): + return math.ceil(x / y) * y + + +def nearest_y(x, y): + return _nearest_y(x, y) + + +def divup(a, b): + return (a + b - 1) // b + + +def roundup(a, b): + result = divup(a, b) * b + return result + + +def roundup32(a): + return roundup(a, 32) + + +def float_to_bits(x): + s = struct.pack(">f", x) + return struct.unpack(">l", s)[0] + + +def torch_random(shape, low, high, dtype): + if dtype in [torch.int64, torch.int32, torch.int16, torch.int8]: + return torch.randint(low, high, shape, dtype=dtype) + return torch.zeros(shape, dtype=dtype).uniform_(low, high) + + +def torch_random_with_zeros(shape, low, high, dtype, zero_fraction=0.1): + total_elements = torch.prod(torch.tensor(shape)).item() + num_zeros = int(total_elements * zero_fraction) + num_random = total_elements - num_zeros + + # Generate random values between low and high + random_values = torch.empty(num_random).uniform_(low, high) + zeros = torch.zeros(num_zeros) + + # Combine zeros and random values + combined = torch.cat([zeros, random_values]) + + # Shuffle the tensor + shuffled = combined[torch.randperm(combined.size(0))] + + # Reshape to the desired shape + result_tensor = shuffled.view(shape) + result_tensor.to(dtype) + return result_tensor + + +### Profiling ### +class Profiler: + def __init__(self): + self.start_times = dict() + self.times = dict() + self.disabled = False + + def clear(self): + self.start_times = dict() + self.times = dict() + self.disabled = False + + def enable(self): + self.disabled = False + + def disable(self): + self.disabled = True + + def start(self, key, force_enable=False): + if self.disabled and not force_enable: + return + + self.start_times[key] = time.time() + + def end(self, key, PERF_CNT=1, force_enable=False): + if self.disabled and not force_enable: + return + + if key not in self.start_times: + return + + diff = time.time() - self.start_times[key] + + if key not in self.times: + self.times[key] = [] + + self.times[key].append(diff / PERF_CNT) + + def get(self, key): + if key not in self.times: + return 0 + + return sum(self.times[key]) / len(self.times[key]) + + def print(self, units="s"): + for key in self.times: + average = self.get(key) + if units == "s": + pass + elif units == "ms": + average *= 1000 + elif units == "us": + average *= 1000000 + elif units == "ns": + average *= 1000000000 + else: + raise ValueError(f"Invalid units: {units}") + print(f"{key}: {average:.3f}{units}") + + +profiler = Profiler() + + +### Turn flags on/off ### +def enable_memory_reports(): + """ + Enables generating reports of memory allocation statistics in .reports/tt_metal dir + """ + return ttnn.device.EnableMemoryReports() + + +def disable_memory_reports(): + """ + Disables generating reports of memory allocation statistics + """ + return ttnn.device.DisableMemoryReports() + + +### Tensor conversion ### +def torch2tt_tensor( + py_tensor: torch.Tensor, + tt_device, + tt_layout=ttnn.TILE_LAYOUT, + tt_memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED), + tt_dtype=ttnn.bfloat16, +): + size = list(py_tensor.size()) + + while len(size) < 4: + size.insert(0, 1) + + tt_tensor = ttnn.Tensor(py_tensor.reshape(size), tt_dtype) + tt_tensor = tt_tensor.to(tt_layout) + + if tt_device is not None: + tt_tensor = tt_tensor.to(tt_device, tt_memory_config) + else: + tt_tensor = tt_tensor.cpu() + + return tt_tensor + + +def tt_tensors_to_torch_tensors( + tt_tensors_device: ttnn.Tensor, mesh_device: Union[ttnn.MeshDevice, ttnn.Device], concat_dim: int = 0 +): + # Convert tensors to interleaved + if tt_tensors_device.is_sharded(): + tt_tensors_device = ttnn.sharded_to_interleaved(tt_tensors_device) + + # Convert tensors to RM layout + if tt_tensors_device.layout == ttnn.TILE_LAYOUT: + # Convert to bfloat16 to ensure untilize works + if tt_tensors_device.dtype != ttnn.bfloat16: + tt_tensors_device = ttnn.clone( + tt_tensors_device, dtype=ttnn.bfloat16, memory_config=ttnn.DRAM_MEMORY_CONFIG + ) + # Untilize using singlecore since multicore version runs out of l1 memory (Issue #9022) + tt_tensors_device = ttnn.untilize(tt_tensors_device, use_multicore=False) + + return torch.cat([t.to_torch() for t in ttnn.get_device_tensors(tt_tensors_device.cpu())], dim=concat_dim) + + +def tt2torch_tensor(tt_tensor): + tt_output = tt_tensor.cpu() + if tt_output.get_layout() != ttnn.ROW_MAJOR_LAYOUT: + tt_output = tt_output.to(ttnn.ROW_MAJOR_LAYOUT) + return tt_output.to_torch() + + +def tt_to_torch_tensor(tt_tensor): + tt_output = tt_tensor.cpu().to(ttnn.ROW_MAJOR_LAYOUT) + return tt_output.to_torch() + + +def torch_to_tt_tensor_rm(py_tensor, device, shape=None, put_on_device=True): + if shape is None: + shape = list(py_tensor.size()) + while len(shape) < 4: + shape.insert(0, 1) + + tt_tensor = ttnn.Tensor(py_tensor.reshape(shape), ttnn.bfloat16) + if put_on_device: + tt_tensor = tt_tensor.to(device) + return tt_tensor + + +def torch_to_tt_tensor(py_tensor, device): + shape = list(py_tensor.size()) + while len(shape) < 4: + shape.insert(0, 1) + + tt_tensor = ( + ttnn.Tensor(py_tensor.reshape(shape), ttnn.bfloat16) + .to( + ttnn.TILE_LAYOUT + ) # change memory layout of TT Tensor to TILE (as operation that will use it expects TILE layout) + .to(device) # move TT Tensor from host to TT accelerator device (device is of type ttnn.device.Device) + ) + + return tt_tensor + + +def unpad_from_zero(x, desired_shape): + if x.padded_shape[-1] == desired_shape[-1] and x.padded_shape[-2] == desired_shape[-2]: + x = tt2torch_tensor(x) + else: + x = x.cpu() + if x.get_layout() != ttnn.ROW_MAJOR_LAYOUT: + x = x.to(ttnn.ROW_MAJOR_LAYOUT) + x = x.unpad( + (0, 0, 0, 0), + ( + desired_shape[0], + desired_shape[1], + desired_shape[2], + desired_shape[3], + ), + ) + + x = x.to_torch() + return x + + +def pad_activation(x): + """ + This function pads an activation with 0s as a pre-preprocessing step to tilization. + + In the 2d case, it pads a vector to the right with 0s, and in the 2+d case, + it pads the bottom and right corners of the last two dimensions. + + :param x: Input PyTorch Tensor + :type x: class:`torch.Tensor` + + WARNING: This function should eventually be retired in favour of padding on device + """ + nearest_32 = _nearest_32 + + assert isinstance(x, torch.Tensor), "Input to this function must be an instance of torch.Tensor" + assert len(x.shape) >= 1 and len(x.shape) <= 4, "Only tensors with dimension 1-4 supported" + if len(x.shape) == 1: # (num_features,) + padded_tensor = torch.zeros(1, 1, 32, nearest_32(x.shape[0])) + padded_tensor[:, 0, 0, : x.shape[0]] = x + elif len(x.shape) == 2: # (batch, num features) + padded_tensor = torch.zeros(x.shape[0], 1, 32, nearest_32(x.shape[1])) + padded_tensor[:, 0, 0, : x.shape[1]] = x + elif len(x.shape) == 3: # (batch, num features y, num features x) + padded_tensor = torch.zeros(x.shape[0], 1, nearest_32(x.shape[-2]), nearest_32(x.shape[-1])) + padded_tensor[..., 0, : x.shape[-2], : x.shape[-1]] = x + else: # (batch, num channels, num features y, num features x) + padded_tensor = torch.zeros(*x.shape[:-2], nearest_32(x.shape[-2]), nearest_32(x.shape[-1])) + padded_tensor[..., : x.shape[-2], : x.shape[-1]] = x + return padded_tensor + + +def pad_weight(x): + """ + This function pads a weight/bias with 0s as a pre-preprocessing step to tilization. + + tt_tensor = ttnn.Tensor( + py_tensor.reshape(shape), ttnn.bfloat16 + In the 2d case, it pads a vector to the right with 0s, and in the 2+d case, + it pads the bottom and right corners of the last two dimensions. + + :param x: Input PyTorch Tensor + :type x: class:`torch.Tensor` + + WARNING: This function should eventually be retired in favour of padding on device + """ + nearest_32 = _nearest_32 + + assert isinstance(x, torch.Tensor), "Input to this function must be an instance of torch.Tensor" + assert len(x.shape) >= 1 and len(x.shape) <= 4, "Only tensors with dimension 1-4 supported" + + if len(x.shape) == 1: # (num_features,) + padded_tensor = torch.zeros(1, 1, 32, nearest_32(x.shape[0])) + padded_tensor[:, 0, 0, : x.shape[0]] = x + elif len(x.shape) == 2: # (r_features, c_features) + padded_tensor = torch.zeros(1, 1, nearest_32(x.shape[0]), nearest_32(x.shape[1])) + padded_tensor[:, 0, : x.shape[0], : x.shape[1]] = x + else: + padded_tensor = torch.zeros(*x.shape[:-2], nearest_32(x.shape[-2]), nearest_32(x.shape[-1])) + padded_tensor[..., : x.shape[-2], : x.shape[-1]] = x + + return padded_tensor + + +def convert_weights_2d_matrix(weights, w_shape): + """ + :param weights: Input PyTorch Tensor + :type weights: class:`torch.Tensor` + """ + ret_shape = [1, 1, w_shape[0], w_shape[1] * w_shape[2] * w_shape[3]] + if isinstance(weights, torch.Tensor): + ret = torch.zeros(np.prod(ret_shape)) + else: + ret = np.zeros(np.prod(ret_shape)) + idx = 0 + for k in range(w_shape[0]): + for r in range(w_shape[2]): + for s in range(w_shape[3]): + for c in range(w_shape[1]): + ret[idx] = weights[k][c][r][s] + idx += 1 + assert idx == np.prod(ret_shape) + return ret.reshape(ret_shape).transpose(2, 3) + + +def convert_act_2d_matrix(activation, kernel_y, kernel_x, stride_y, stride_x, pad_y, pad_x): + """ + :param activation: Input PyTorch Tensor + :type activation: class:`torch.Tensor` + """ + N = activation.shape[0] + C = activation.shape[1] + H = activation.shape[2] + W = activation.shape[3] + + OH = (int)((H - kernel_y + 2 * pad_y) // stride_y) + 1 + OW = ((W - kernel_x + 2 * pad_x) // stride_x) + 1 + nrows = OH * OW + ncols = C * kernel_x * kernel_y + ret_shape = [1, N, nrows, ncols] + if isinstance(activation, torch.Tensor): + ret = torch.zeros(np.prod(ret_shape)) + else: + ret = np.zeros(np.prod(ret_shape)) + idx = 0 + for n in range(N): + for h in range(-1 * pad_y, H + pad_y - kernel_y + 1, stride_y): + for w in range(-1 * pad_x, W + pad_x - kernel_x + 1, stride_x): + for r in range(kernel_y): + for s in range(kernel_x): + for c in range(C): + h_offs = h + r + w_offs = w + s + pad = h_offs < 0 or h_offs >= H or w_offs < 0 or w_offs >= W + ret[idx] = 0 if pad else activation[n][c][h_offs][w_offs] + idx += 1 + assert idx == np.prod(ret_shape) + return ret.reshape(ret_shape) + + +### Tilizing / Untilizing ### +@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:") +def tilize(x): + return x + + +@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:") +def tilize_to_list(x): + """ + Returns a flattened list of the tensor + """ + return tilize(x).reshape(-1).tolist() + + +@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:") +def untilize(x): + return x + + +### Measuring accuracy and other metrics ### +def is_close(a, b, rtol=1e-2, atol=1e-2, max_mag=2.0, max_mag_fraction=0.02): + """ + A variant of np.isclose with logging. + """ + absdiff = (a - b).abs() + reldiff1 = (a.abs() / b.abs()) - 1.0 + reldiff2 = (a.abs() + 1.0) / (b.abs() + 1.0) - 1.0 # in case b.abs() is 0 + reldiff_or = torch.logical_or(reldiff1.abs() < rtol, reldiff2.abs() < rtol) + max_mag_ok = absdiff < max_mag * max_mag_fraction + + or_abs_rel = torch.logical_or(absdiff < atol, reldiff_or) + or_abs_rel = torch.logical_or(or_abs_rel, max_mag_ok) + debug_index = or_abs_rel.to(torch.int32).argmin().item() + + if not or_abs_rel.reshape(-1)[debug_index]: + logger.info(f"isclose mismatch at index={debug_index}") + logger.info(a.reshape(-1)[debug_index]) + logger.info(b.reshape(-1)[debug_index]) + logger.info(f"reldiff1={reldiff1.reshape(-1)[debug_index]}") + logger.info(f"reldiff2={reldiff2.reshape(-1)[debug_index]}") + logger.info(f"absdiff={absdiff.reshape(-1)[debug_index]}") + + HT = a.shape[-2] // 32 + WT = a.shape[-1] // 32 + hwt = debug_index // 1024 + wt = hwt % WT + ht = hwt // WT + h = (debug_index % 1024) // 32 + w = (debug_index % 1024) % 32 + + logger.info(f"**** at {debug_index} --- HTWT={ht} {wt} HW={h} {w}") + + return torch.all(or_abs_rel) + + +def _comp_nonfinite(golden, calculated): + """ + Returns True if tensors contain the same non-finite values (nan, inf, -inf) at the same positions. Also returns True if all elements are finite. + Returns False if non-finite values differ between both tensors. + """ + + # torch.equal(['nan'], ['nan']] => False + # For this reason, we check for nan and inf separately + if torch.not_equal(torch.isnan(golden), torch.isnan(calculated)).any(): + return False + + golden_inf_mask = torch.isinf(golden) + calculated_inf_mask = torch.isinf(calculated) + + if torch.not_equal(golden_inf_mask, calculated_inf_mask).any(): + return False + + golden_inf = golden[golden_inf_mask] + calculated_inf = calculated[calculated_inf_mask] + return torch.equal(golden_inf, calculated_inf) + + +def comp_allclose(golden, calculated, rtol=1e-05, atol=1e-08): + if golden.dtype != calculated.dtype: + calculated = calculated.type(golden.dtype) + + atol_delta = torch.max(torch.abs(golden - calculated)).item() + rtol_delta = torch.max(torch.abs(golden - calculated) / torch.abs(calculated)).item() + return ( + torch.allclose(golden, calculated, rtol, atol, True), + f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}", + ) + + +def comp_pcc(golden, calculated, pcc=0.99, rtol=1e-05, atol=1e-04): + golden = torch.Tensor(golden) + calculated = torch.Tensor(calculated) + + if golden.dtype != calculated.dtype: + calculated = calculated.type(golden.dtype) + + if torch.all(torch.isnan(golden)) and torch.all(torch.isnan(calculated)): + logger.warning("Both tensors are 'nan'") + return True, 1.0 + + if torch.all(torch.isnan(golden)) or torch.all(torch.isnan(calculated)): + logger.error("One tensor is all nan, the other is not.") + return False, 0.0 + + # Test if either is completely zero — but a zero tensor is also a constant tensor, + # so fall back to allclose instead of a hard 0.0: zero-vs-small-constant may be + # within the caller's tolerances. + if torch.any(golden.bool()) != torch.any(calculated.bool()): + logger.warning("One tensor is all zero. PCC undefined; falling back to allclose.") + result = torch.allclose(golden, calculated, rtol=rtol, atol=atol) + return result, float(result) + + golden = torch.squeeze(golden).flatten() + calculated = torch.squeeze(calculated).flatten() + + # For now, mask all infs and nans (to zero) so that we check the rest... TODO + # Skip this for integer types which don't have NaN/Inf values. + if golden.dtype.is_floating_point: + # FP8 doesn't support isfinite/nan_to_num and bfloat16 products lose precision, + # so correlate these in float32. + if golden.dtype in (torch.float8_e4m3fn, torch.float8_e5m2, torch.bfloat16): + golden = golden.to(torch.float32) + calculated = calculated.to(torch.float32) + + # Zero out NaN/Inf, preserving the historical PCC values. nan_to_num allocates a + # full-size copy of each tensor, so only do it when invalid values are actually + # present; on the common all-finite path the tensors stay as views and no copy is + # made (this short-circuit is what keeps peak memory near 1x of one input). + if not bool((torch.isfinite(golden) & torch.isfinite(calculated)).all()): + golden = torch.nan_to_num(golden, nan=0.0, posinf=0.0, neginf=0.0) + calculated = torch.nan_to_num(calculated, nan=0.0, posinf=0.0, neginf=0.0) + + if torch.equal(golden, calculated): + return True, 1.0 + + # Integer tensors must be correlated in floating point (centering/products would + # otherwise truncate/overflow). float32 keeps the working set small. + if not golden.dtype.is_floating_point: + golden = golden.to(torch.float32) + calculated = calculated.to(torch.float32) + + # Pearson r with float64 *accumulation* (dtype= on the reductions) over the float32 + # data: no float64 copy of either tensor is materialized, so peak memory stays near + # 1x of one input on large tensors while matching a full-float64 correlation to + # |Δ|<1e-9 across the high-PCC (>=0.999) range. + n = golden.numel() + g_centered = golden - (golden.sum(dtype=torch.float64) / n).to(golden.dtype) + c_centered = calculated - (calculated.sum(dtype=torch.float64) / n).to(calculated.dtype) + cov = (g_centered * c_centered).sum(dtype=torch.float64) + g_sq_sum = g_centered.pow(2).sum(dtype=torch.float64) + c_sq_sum = c_centered.pow(2).sum(dtype=torch.float64) + denom = torch.sqrt(g_sq_sum * c_sq_sum) + # pow/sum stay in float32 before the reduction; large-magnitude tensors (e.g. ldexp) + # can overflow to inf here even though float64 accumulation would be finite. + if not math.isfinite(denom.item()) or not math.isfinite(cov.item()): + g_centered64 = g_centered.to(torch.float64) + c_centered64 = c_centered.to(torch.float64) + cov = (g_centered64 * c_centered64).sum() + denom = torch.sqrt(g_centered64.pow(2).sum() * c_centered64.pow(2).sum()) + cal_pcc = (cov / denom).item() + + # Zero variance -> denom == 0 -> cal_pcc is nan: PCC is undefined for constant tensors. + # Fall back to allclose rather than returning a misleading 1.0. + if math.isnan(cal_pcc): + logger.warning("PCC is NaN (zero variance / constant tensor). Falling back to allclose check.") + result = torch.allclose(golden, calculated, rtol=rtol, atol=atol) + return result, float(result) + + return cal_pcc >= pcc, cal_pcc + + +def ulp(x: Union[ttnn.Tensor, torch.Tensor]) -> Union[ttnn.Tensor, torch.Tensor]: + "Return Unit of Least Precision for each element of a given tensor" + + received_ttnn_input = False + if isinstance(x, ttnn.Tensor): + x = ttnn.to_torch(x) + received_ttnn_input = True + + # Notes: + # - This should be identical to the definition of ULP by Goldberg + # "What every computer scientist should know about floating-point arithmetic" + # https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html + # - We use torch.abs(x) to ensure symmetry ULP(-x) == ULP(x) + # - For x powers of 2, x + ULP(x) is not closest number but second closest (previous number is 2x closer) + # However, this avoids rounding-to-nearest-tie-to-even issues on addition (i.e. x + ULP(x) != x) + abs_x = torch.abs(x) + next = torch.nextafter( + abs_x, torch.tensor(math.inf, dtype=x.dtype) + ) # 1 ULP ~ Difference between two consecutive floating point numbers + ulp_value = next - abs_x + + # Special case: if abs_x == torch.finfo(x.dtype).max, then next == math.inf, which leads to ULP(x) == inf rather than finite number + # We fix this problem by manually calculating ULP at max value, and masking tensor when input == max + dtype_max = torch.finfo(x.dtype).max + max_epsilon = dtype_max - torch.nextafter( + torch.tensor(dtype_max, dtype=x.dtype), torch.tensor(-math.inf, dtype=x.dtype) + ) + ulp_value = torch.where(abs_x == dtype_max, max_epsilon, ulp_value) + + if received_ttnn_input: # Ensures that type(input) == type(output) + ulp_value = ttnn.from_torch(ulp_value) + + return ulp_value + + +def comp_ulp(golden, calculated, ulp_threshold, allow_nonfinite=False): + """ + Compute absolute error between two tensors in Units of Least Precision (ULP) + """ + + # If both tensors are empty, then we can return True + if torch.numel(golden) == 0 and torch.numel(calculated) == 0: + return True, "Both tensors are empty" + + if not allow_nonfinite and not torch.all(torch.isfinite(calculated)): + return False, "Calculated tensor contains non-finite values" + + if not _comp_nonfinite(golden, calculated): + return False, "Tensors are not finite at the same positions" + # nonfinite elements can interfere with ULP error calculation + # To avoid this, replace nan, +inf, -inf with 0 + # (we have already checked that both tensors have the same nonfinite elements) + mask_finite = ~torch.isfinite(golden) + golden = golden.clone() + calculated = calculated.clone() + golden[mask_finite] = 0 + calculated[mask_finite] = 0 + + # ULP is measured according to the golden tensor + # In most cases, data type of golden tensor should be the same as calculated tensor. + # However, in some cases, we may want to measure < 1 ULP differences, which requires golden tensor + # to have higher precision than calculated tensor. + # If we passed golden tensor to ulp() as is, we would get ULP of higher precision. + # e.g. ulp of float32 rather bfloat16 calculation, which would give us a wrong value. + ulp_value = ulp(golden.type(calculated.dtype)) + + if golden.dtype != calculated.dtype: # Note: assumes that golden has higher precision than calculated tensor + calculated = calculated.type(golden.dtype) + ulp_value = ulp_value.type(golden.dtype) # Convert ULP to higher precision (for sub-1 ULP measurements) + + ulp_tensor = torch.abs(calculated - golden) / ulp_value + ulp_delta = torch.max(ulp_tensor) + within_threshold = ulp_delta <= ulp_threshold + message = f"Max ULP Delta: {ulp_delta}" + if not within_threshold: + ulp_index = torch.argmax(ulp_tensor) + ulp_index_tuple = tuple(int(idx) for idx in torch.unravel_index(ulp_index, golden.shape)) + message += ( + f" @ {list(ulp_index_tuple)} = " + f"|calculated {calculated[ulp_index_tuple]} - golden {golden[ulp_index_tuple]}| " + f"/ ULP(golden) {ulp_value[ulp_index_tuple]}" + ) + return (within_threshold, message) + + +def calculate_detailed_ulp_stats(expected, actual): + """ + Calculate detailed ULP statistics for analysis. + + Returns: + dict: Dictionary with ULP statistics including max, mean, std, and percentiles + """ + if isinstance(actual, ttnn.Tensor): + actual = ttnn.to_torch(actual) + if isinstance(expected, ttnn.Tensor): + expected = ttnn.to_torch(expected) + + # Convert to bfloat16 if not already + expected = expected.to(torch.bfloat16) + actual = actual.to(torch.bfloat16) + + # Handle special cases + if torch.allclose(expected, actual, rtol=0, atol=0, equal_nan=True): + return { + "max_ulp": 0.0, + "mean_ulp": 0.0, + "median_ulp": 0.0, + "std_ulp": 0.0, + "p95_ulp": 0.0, + "p99_ulp": 0.0, + "perfect_matches": 1.0, + } + + # Convert bfloat16 to uint16 representation for bit manipulation + expected_bits = expected.view(torch.int16).to(torch.int32) + actual_bits = actual.view(torch.int16).to(torch.int32) + + # Handle sign differences + expected_sign = expected_bits < 0 + actual_sign = actual_bits < 0 + same_sign = expected_sign == actual_sign + + # Calculate ULP differences + expected_abs_bits = torch.where(expected_sign, -expected_bits, expected_bits) + actual_abs_bits = torch.where(actual_sign, -actual_bits, actual_bits) + + ulp_diff = torch.where(same_sign, torch.abs(expected_bits - actual_bits), expected_abs_bits + actual_abs_bits) + + # Handle non-finite values + expected_finite = torch.isfinite(expected) + actual_finite = torch.isfinite(actual) + both_finite = expected_finite & actual_finite + + ulp_diff = torch.where(both_finite, ulp_diff, torch.tensor(float("inf"))) + + # Handle same non-finite values + both_nan = torch.isnan(expected) & torch.isnan(actual) + both_posinf = torch.isposinf(expected) & torch.isposinf(actual) + both_neginf = torch.isneginf(expected) & torch.isneginf(actual) + same_nonfinite = both_nan | both_posinf | both_neginf + + ulp_diff = torch.where(same_nonfinite, torch.tensor(0.0), ulp_diff) + + # Calculate statistics only on finite ULP differences + finite_ulp = ulp_diff[torch.isfinite(ulp_diff)] + + if len(finite_ulp) == 0: + return { + "max_ulp": float("inf"), + "mean_ulp": float("inf"), + "median_ulp": float("inf"), + "std_ulp": float("inf"), + "p95_ulp": float("inf"), + "p99_ulp": float("inf"), + "perfect_matches": 0.0, + } + + finite_ulp_float = finite_ulp.float() + perfect_matches = (finite_ulp == 0).float().mean().item() + + return { + "max_ulp": torch.max(finite_ulp).item(), + "mean_ulp": torch.mean(finite_ulp_float).item(), + "median_ulp": torch.median(finite_ulp_float).item(), + "std_ulp": torch.std(finite_ulp_float).item(), + "p95_ulp": torch.quantile(finite_ulp_float, 0.95).item(), + "p99_ulp": torch.quantile(finite_ulp_float, 0.99).item(), + "perfect_matches": perfect_matches, + } + + +def comp_allclose_and_pcc(golden, calculated, rtol=1e-05, atol=1e-08, pcc=0.99): + # 0-volume tensors are special because they don't have elements, so we can't compute PCC, etc. + # If one of the tensors is a 0-volume tensor, simply call torch.equal to check if they are equal + # (i.e. that both are 0-volume tensors and they have equal shapes). + if golden.numel() == 0 or calculated.numel() == 0: + return torch.equal(golden, calculated), f"{golden} != {calculated}" + + if golden.dtype != calculated.dtype: + calculated = calculated.type(golden.dtype) + + passing = True + output = "" + passing_allclose, output_allclose = comp_allclose(golden, calculated, rtol, atol) + passing &= passing_allclose + output += output_allclose + if torch.numel(golden) != 1: + passing_pcc, output_pcc = comp_pcc(golden, calculated, pcc, rtol=rtol, atol=atol) + passing &= passing_pcc + output += f", pcc={output_pcc}" + + return passing, output + + +def comp_equal(golden, calculated): + if golden.dtype != calculated.dtype: + calculated = calculated.type(golden.dtype) + + # If either tensor is zero-volume, broadcasting can still yield an empty delta and + # crash torch.max(); defer entirely to torch.equal (False on shape mismatch). + if golden.numel() == 0 or calculated.numel() == 0: + return torch.equal(golden, calculated), f"{golden} != {calculated}" + + atol_delta = torch.max(torch.abs(golden - calculated)).item() + rtol_delta = torch.max(torch.abs(golden - calculated) / torch.abs(calculated)).item() + return ( + torch.equal(golden, calculated), + f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}", + ) + + +def get_oom_of_float(float_lst): + """ + Given a list of floats, returns a list of the order or magnitudes + of the floats. Useful when you want to make sure that even if your + tt outputs don't match pytorch all that well, they are at least + on the same order of magnitude + """ + ooms = [] + for el in float_lst: + str_el = str(el) + if "e" in str_el: + oom = int(str_el.split("e")[1]) + elif str_el[:2] == "0.": + str_el = str_el.split(".")[1] + + oom = -1 + for e in str_el: + if e != "0": + break + oom -= 1 + else: + oom = len(str_el.split(".")[0]) + + ooms.append(oom) + + return ooms + + +def print_diff_argmax(a, b, annotation=""): + """ + Prints out the value of both tensors at a point where the absolute difference is the largest. + """ + absdiff = (a - b).abs() + argmax = absdiff.argmax().item() + diff = absdiff.reshape(-1)[argmax] + rela = a.abs() / (torch.max(a.abs(), b.abs())) + relb = b.abs() / (torch.max(a.abs(), b.abs())) + HT = a.shape[-2] // 32 + WT = a.shape[-1] // 32 + hwt = argmax // 1024 + wt = hwt % WT + ht = hwt // WT + h = (argmax % 1024) // 32 + w = (argmax % 1024) % 32 + print( + "Abs diff=", + diff, + " at ", + argmax, + " --- ", + annotation, + "HTWT=", + ht, + wt, + "HW=", + h, + w, + ) + print(" (a=", a.reshape(-1)[argmax].item(), ")") + print(" (b=", b.reshape(-1)[argmax].item(), ")") + print(" Rel a=", rela.reshape(-1)[argmax], " at ", argmax) + print(" Rel b=", relb.reshape(-1)[argmax], " at ", argmax) + return diff.item() + + +def print_diff_tt_pyt(a, b, annotation=""): + # first convert a pytorch tensor argument b to tt + padded_b = pad_weight(b) + pyt_a = tt2torch(a) # untilizes also + return print_diff_argmax(pyt_a, padded_b, annotation) + + +def ttP(x, count=4, offset=0, stride=1): + if type(x) == torch.Tensor: + t1 = x.reshape(-1) + else: + tt_out = x.cpu() + torch_out = untilize(tt_out.to_torch()) + t1 = torch_out.reshape(-1) + print("Tensor vals: (", end="") + for j in range(offset, offset + count * stride, stride): + print(t1[j].item(), " ", end="") + print(")") + + +### Conv related helpers ### +def read_conv_act_into_mm_act_block( + conv_act, + act_address_map_index, + address_map, + address_map_this_block_size, + act_block_h, + act_block_w, +): + mm_act_block_shape = [1, 1, act_block_h * 32, act_block_w * 32] + mm_act_block_size = act_block_h * act_block_w * 1024 + mm_act_block = torch.zeros(mm_act_block_size, dtype=torch.bfloat16).float() + for i in range(0, address_map_this_block_size, 4): + src_address = address_map[act_address_map_index] + dst_address = address_map[act_address_map_index + 1] + read_size = address_map[act_address_map_index + 2] + pad = address_map[act_address_map_index + 3] + for s in range(read_size): + assert dst_address + s < mm_act_block_size + if pad: + mm_act_block[dst_address + s] = 0 + else: + assert src_address + s < len(conv_act) + mm_act_block[dst_address + s] = conv_act[src_address + s] + act_address_map_index += 4 + return (mm_act_block.reshape(mm_act_block_shape), act_address_map_index) + + +def read_conv_weight_into_mm_weight_block( + conv_weight, + weight_address_map_index, + weight_address_map, + weight_address_map_this_block_size, + weight_block_h, + weight_block_w, +): + mm_weight_block_shape = [1, 1, weight_block_h * 32, weight_block_w * 32] + mm_weight_block_size = weight_block_h * weight_block_w * 1024 + mm_weight_block = torch.zeros(mm_weight_block_size, dtype=torch.bfloat16).float() + for i in range(0, weight_address_map_this_block_size, 4): + src_address = weight_address_map[weight_address_map_index] + dst_address = weight_address_map[weight_address_map_index + 1] + read_size = weight_address_map[weight_address_map_index + 2] + pad = weight_address_map[weight_address_map_index + 3] + for s in range(read_size): + assert dst_address + s < mm_weight_block_size + if pad: + mm_weight_block[dst_address + s] = 0 + else: + assert src_address + s < len(conv_weight) + mm_weight_block[dst_address + s] = conv_weight[src_address + s] + weight_address_map_index += 4 + return (mm_weight_block.reshape(mm_weight_block_shape), weight_address_map_index) + + +def blocked_mm_with_conv_act( + conv_act, + mm_weight, + act_address_map, + weight_address_map, + num_blocks_act_h, + num_blocks_act_w, + num_blocks_weight_w, + act_block_h, + act_block_w, + weight_block_w, +): + # act refers to conv activation tensor + # weight refers to conv weight tensor + mm_output_shape = [ + 1, + 1, + num_blocks_act_h * act_block_h * 32, + num_blocks_weight_w * weight_block_w * 32, + ] + ret = torch.zeros(mm_output_shape, dtype=torch.bfloat16).float() + mm_output_block_shape = [1, 1, act_block_h * 32, weight_block_w * 32] + act_address_map_index = 0 + weight_address_map_index = 0 + weight_block_h = act_block_w + num_groups = act_address_map[act_address_map_index] + assert num_groups == num_blocks_act_h * num_blocks_act_w * num_blocks_weight_w + weight_num_groups = act_address_map[weight_address_map_index] + assert weight_num_groups == num_groups + act_address_map_index += 1 + weight_address_map_index += 1 + for block_act_h in range(num_blocks_act_h): + # Reset weight (weight) to the starting tile in this column + for block_weight_w in range(num_blocks_weight_w): + output_block = torch.zeros(mm_output_block_shape, dtype=torch.bfloat16).float() + for block_act_w in range(num_blocks_act_w): + address_map_this_block_size = act_address_map[act_address_map_index] + act_address_map_index += 1 + weight_address_map_this_block_size = weight_address_map[weight_address_map_index] + weight_address_map_index += 1 + (mm_act_block, act_address_map_index) = read_conv_act_into_mm_act_block( + conv_act, + act_address_map_index, + act_address_map, + address_map_this_block_size, + act_block_h, + act_block_w, + ) + ( + mm_weight_block, + weight_address_map_index, + ) = read_conv_weight_into_mm_weight_block( + mm_weight, + weight_address_map_index, + weight_address_map, + weight_address_map_this_block_size, + weight_block_h, + weight_block_w, + ) + # Untilize weight block (this CPU reference does matmul on untilized blocks) + mm_weight_block = untilize(mm_weight_block) + for out_h_block in range(act_block_h * 32): + for out_w_block in range(weight_block_w * 32): + output_block[0][0][out_h_block][out_w_block] += torch.dot( + mm_act_block[0, 0, out_h_block, :].reshape(-1), + mm_weight_block[0, 0, :, out_w_block].reshape(-1), + ) + start_oh = block_act_h * act_block_h * 32 + start_ow = block_weight_w * weight_block_w * 32 + end_oh = start_oh + (act_block_h * 32) + end_ow = start_ow + (weight_block_w * 32) + ret[0, 0, start_oh:end_oh, start_ow:end_ow] = output_block + + return ret + + +def is_conv_supported_on_device(conv_params): + K, C, R, S, U, V, P_H, P_W, dilation, groups = [conv_params[i] for i in range(10)] + + if K % 32 != 0 or dilation != 1 or groups != 1: + logger.warning("DOES NOT HAVE SUPPORT FOR Conv with following parameters -") + logger.warning( + "K=" + + str(K) + + " C=" + + str(C) + + " R=" + + str(R) + + " S=" + + str(S) + + " U=" + + str(U) + + " V=" + + str(V) + + " PH=" + + str(P_H) + + " PW=" + + str(P_W) + + " dilation=" + + str(dilation) + + " groups=" + + str(groups) + ) + return False + + return True + + +def is_x2_harvested(device): + grid = device.compute_with_storage_grid_size() + return device.arch() == Arch.WORMHOLE_B0 and (grid.x, grid.y) == (8, 7) + + +def is_single_chip(): + return ttnn.GetNumAvailableDevices() == 1 + + +def is_quasar(): + ARCH_NAME = ttnn.get_arch_name() + return "quasar" in ARCH_NAME + + +def is_blackhole(): + ARCH_NAME = ttnn.get_arch_name() + return "blackhole" in ARCH_NAME + + +def is_wormhole_b0(): + ARCH_NAME = ttnn.get_arch_name() + return "wormhole_b0" in ARCH_NAME + + +def is_watcher_enabled(): + watcher = os.environ.get("TT_METAL_WATCHER") + lightweight_asserts = os.environ.get("TT_METAL_LIGHTWEIGHT_KERNEL_ASSERTS") + return (watcher is not None and watcher != "") or lightweight_asserts == "1" + + +def is_llk_assert_enabled(): + llk_assert = os.environ.get("TT_METAL_LLK_ASSERTS") + return llk_assert == "1" + + +def is_n300(): + return os.environ.get("MESH_DEVICE", "N150") == "N300" + + +def is_slow_dispatch(): + return os.environ.get("TT_METAL_SLOW_DISPATCH_MODE") == "1" + + +def ti_skip(condition, reason="Invalid test parameters"): + # Imported here, not at module scope: this module sits on the model path (tt_dit's + # attention imports it), so a top-level pytest import makes every deployment that + # ships models/common — a release container, for instance — need a test framework. + import pytest + + return pytest.mark.skipif(condition, reason="Skipping unsupported case: " + reason) + + +def skip_for_blackhole(reason_str="not a blackhole test"): + return ti_skip(is_blackhole(), reason=reason_str) + + +def skip_for_wormhole_b0(reason_str="not a wormhole test"): + return ti_skip(is_wormhole_b0(), reason=reason_str) + + +def skip_with_watcher(reason_str="Test is not passing with watcher enabled"): + return ti_skip(is_watcher_enabled(), reason=reason_str) + + +def skip_with_llk_assert(reason_str="Test is not passing with LLK asserts enabled"): + return ti_skip(is_llk_assert_enabled(), reason=reason_str) + + +def run_for_blackhole(reason_str="only runs for Blackhole"): + return ti_skip(not is_blackhole(), reason=reason_str) + + +def run_for_wormhole_b0(reason_str="only runs for Wormhole B0"): + return ti_skip(not is_wormhole_b0(), reason=reason_str) + + +def run_for_wormhole_b0_or_blackhole(reason_str="only runs for Wormhole B0 or Blackhole"): + return ti_skip(not (is_wormhole_b0() or is_blackhole()), reason=reason_str) + + +def run_for_n_dev(n, reason_str="Test is not meant for this number of devices"): + return ti_skip(ttnn.get_num_devices() != n, reason=reason_str) + + +def skip_for_n_dev(n, reason_str="Test is not meant for this number of devices"): + return ti_skip(ttnn.get_num_devices() == n, reason=reason_str) + + +def skip_for_n_or_less_dev(n, reason_str="Test is not meant for this number of devices"): + return ti_skip(ttnn.get_num_devices() <= n, reason=reason_str) + + +def skip_for_slow_dispatch(reason_str="not working for slow dispatch"): + return ti_skip(is_slow_dispatch(), reason=reason_str) + + +def ttl_complex_2_torch_complex(tt_tensor): + torch_tensor = tt2torch_tensor(tt_tensor) + + # extract real and imag parts of the complex tensor + real = torch_tensor[:, :, :, : torch_tensor.shape[-1] // 2].to(torch.bfloat16).to(torch.float) + imag = torch_tensor[:, :, :, torch_tensor.shape[-1] // 2 :].to(torch.bfloat16).to(torch.float) + + # create torch complex tensor + result = torch.complex(real, imag) + return result + + +def pad_and_fold_conv_filters_for_unity_stride(filter_pyt_nchw_tensor, stride_h, stride_w, align_c=4): + assert stride_h == stride_w + assert filter_pyt_nchw_tensor.shape[2] == filter_pyt_nchw_tensor.shape[3] + assert isinstance(align_c, int) and align_c > 0 + # Fold activation for unity stride + # Pad channel size to align_c. This keeps L1 read addresses aligned; extra channels become + # zero-valued weights that contribute nothing to the convolution. align_c=4 is the WH/BH default + # (16B alignment for bf16 gives C a multiple of 4 with a tiled conv reader). Quasar's row-major + # fold needs align_c=8 (bf16 row-major shard width must be a multiple of 8) so the first conv + # folds to groups*8 input channels and consumes the aligned output without per-group padding strip. + C = _nearest_y(filter_pyt_nchw_tensor.shape[1], align_c) + # Pad filter to nearest stride + Padded_filter_height = _nearest_y(filter_pyt_nchw_tensor.shape[2], stride_h) + Padded_filter_width = _nearest_y(filter_pyt_nchw_tensor.shape[3], stride_w) + filter_pyt_padded = torch.nn.functional.pad( + filter_pyt_nchw_tensor, + ( + 0, + Padded_filter_width - filter_pyt_nchw_tensor.shape[3], + 0, + Padded_filter_height - filter_pyt_nchw_tensor.shape[2], + 0, + C - filter_pyt_nchw_tensor.shape[1], + ), + ) + # Fold filter for unity stride. + filter_pyt_padded_folded = torch.zeros( + [ + filter_pyt_padded.shape[0], + C * stride_h * stride_w, + (int)(filter_pyt_padded.shape[2] / stride_h), + (int)(filter_pyt_padded.shape[3] / stride_w), + ] + ) + for h in range(0, filter_pyt_padded.shape[2], stride_h): + for w in range(0, filter_pyt_padded.shape[3], stride_w): + folded_h = (int)(h / stride_h) + folded_w = (int)(w / stride_w) + for i in range(4): + start_c = i * C + filter_pyt_padded_folded[:, start_c : start_c + C, folded_h, folded_w] = filter_pyt_padded[ + :, :, h + (int)(i / stride_w), w + (int)(i % stride_w) + ] + return filter_pyt_padded_folded + + +# produces a tensor where each element in a page is the page number +# this tensor is easy to debug and visualize +def get_debug_tensor(num_pages_width, num_pages_height, dtype, page_width=32, page_height=32): + torch_tensor = None + for row_idx in range(0, int(num_pages_height)): + tile_row = None + for col_idx in range(0, int(num_pages_width)): + tile_idx = col_idx + num_pages_width * row_idx + tile = torch.full((1, 1, page_width, page_height), tile_idx + 1, dtype=dtype) + if tile_row == None: + tile_row = tile + else: + tile_row = torch.cat((tile_row, tile), 3) + if torch_tensor == None: + torch_tensor = tile_row + else: + torch_tensor = torch.cat((torch_tensor, tile_row), 2) + + return torch_tensor + + +# ── transformers 5.x Cache API compatibility ──────────────────────────────── +# transformers 5.x removed the legacy Cache API: DynamicCache no longer exposes +# from_legacy_cache / to_legacy_cache / key_cache / value_cache (per-layer KV now +# lives at cache.layers[i].keys/.values). These helpers work on both 4.x and 5.x. +def hf_cache_layer_kv(cache, layer_idx): + """Return (key, value) tensors for a layer of a transformers Cache. + + Handles the legacy tuple-of-tuples past_key_values, transformers <5 Cache + (key_cache/value_cache), and transformers >=5 Cache (layers[i].keys/.values). + """ + if isinstance(cache, (tuple, list)): # legacy tuple-of-tuples past_key_values + return cache[layer_idx][0], cache[layer_idx][1] + if hasattr(cache, "key_cache"): # transformers < 5.x Cache + return cache.key_cache[layer_idx], cache.value_cache[layer_idx] + layer = cache.layers[layer_idx] # transformers >= 5.x Cache + return layer.keys, layer.values + + +def hf_cache_to_legacy(cache): + """Export a transformers Cache to the legacy tuple-of-(key, value) format.""" + if hasattr(cache, "to_legacy_cache"): # transformers < 5.x + return cache.to_legacy_cache() + return tuple((layer.keys, layer.values) for layer in cache.layers) # transformers >= 5.x + + +def hf_dynamic_cache_from_legacy(layer_kvs): + """Build a transformers DynamicCache from per-layer (key, value) tuples.""" + from transformers import DynamicCache + + layer_kvs = tuple(layer_kvs) + if hasattr(DynamicCache, "from_legacy_cache"): # transformers < 5.x + return DynamicCache.from_legacy_cache(layer_kvs) + return DynamicCache(layer_kvs) # transformers >= 5.x + + +def hf_cache_num_layers(cache): + """Number of populated layers in a transformers Cache (version-tolerant).""" + return len(cache.key_cache) if hasattr(cache, "key_cache") else len(cache.layers) + + +def hf_empty_encoder_decoder_cache(): + """Create an empty transformers EncoderDecoderCache (version-tolerant).""" + from transformers import DynamicCache, EncoderDecoderCache + + if hasattr(EncoderDecoderCache, "from_legacy_cache"): # transformers < 5.x + return EncoderDecoderCache.from_legacy_cache(None) + return EncoderDecoderCache(DynamicCache(), DynamicCache()) # transformers >= 5.x + + +def copy_to_buffer(src: "ttnn.Tensor", dst: "ttnn.Tensor", target_dtype) -> None: + """Convert ``src`` to ``dst``'s layout/dtype/shape/memcfg and write it into + ``dst``. ``dst``'s device buffer is preserved (no reallocation) so any + captured trace and the DRAM prefetcher's recorded buffer addresses remain + valid. The final ``ttnn.to_memory_config`` with ``output_tensor=dst`` both + reshards to ``dst``'s memory config and copies into ``dst``'s buffer. + """ + converted = src + + if converted.layout != dst.layout: + converted = ttnn.to_layout(converted, layout=dst.layout) + + if converted.dtype != target_dtype: + converted = ttnn.typecast(converted, dtype=target_dtype) + + if tuple(converted.shape) != tuple(dst.shape): + converted = ttnn.reshape(converted, list(dst.shape)) + + ttnn.to_memory_config(converted, dst.memory_config(), output_tensor=dst) diff --git a/code/models/common/utils.py b/code/models/common/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a999b20aab2de3e0f38635af0d5f162a665cf563 --- /dev/null +++ b/code/models/common/utils.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn.functional as F +from torch import Tensor + +# Backward compatibility: filter_none moved to models.common.sampling._utils +from models.common.sampling._utils import filter_none # noqa: F401 + +# Backward compatibility: LogProbsCalculator moved to models.common.sampling.tt_log_probs +from models.common.sampling.tt_log_probs import LogProbsCalculator # noqa: F401 + + +def top_k_top_p_filtering( + logits: Tensor, + top_k: int = 0, + top_p: float = 1.0, + filter_value: float = -float("Inf"), + min_tokens_to_keep: int = 1, +) -> Tensor: + """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering + Args: + logits: logits distribution shape (batch size, vocabulary size) + if top_k > 0: keep only top k tokens with highest probability (top-k filtering). + if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering). + Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751) + Make sure we keep at least min_tokens_to_keep per batch example in the output + From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317 + """ + if top_k > 0: + top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1)) # Safety check + # Remove all tokens with a probability less than the last token of the top-k + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits[indices_to_remove] = filter_value + + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + + # Remove tokens with cumulative probability above the threshold (token with 0 are kept) + sorted_indices_to_remove = cumulative_probs > top_p + if min_tokens_to_keep > 1: + # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below) + sorted_indices_to_remove[..., :min_tokens_to_keep] = 0 + # Shift the indices to the right to keep also the first token above the threshold + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 0] = 0 + + # scatter sorted tensors to original indexing + indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + logits[indices_to_remove] = filter_value + return logits diff --git a/code/models/common/validation_tools.py b/code/models/common/validation_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..d9ea69e5c8a04c0e17d35446e2fc85909c80f1fc --- /dev/null +++ b/code/models/common/validation_tools.py @@ -0,0 +1,719 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTNN Validation Framework + +A decorator-based validation system for comparing TTNN implementations against +reference implementations (in PyTorch). Supports automatic input/output +mapping, metric computation, and result collection. + +Key Features: +- Automatic comparison of TTNN vs reference implementations +- TTNN-native metric computation (stays on device until final scalar) +- Flexible input/output mapping +- Built-in metrics: max_abs_error, mean_abs_error, cosine_similarity +- Performance tracking +- Result registry for batch reporting +""" + +import time +from dataclasses import dataclass, field +from enum import Enum +from functools import wraps +from typing import Any, Callable, Dict, List, Optional + +import torch + +import ttnn + +from .auto_compose import to_torch_auto_compose +from .distribute_as import from_torch_dist_as +from .metrics import DEFAULT_METRICS + +# ============================================================================ +# Public API +# ============================================================================ + +# Module exports are defined at the package level in __init__.py + + +def get_validation_registry() -> "ValidationRegistry": + """Get the global validation registry""" + return _validation_registry + + +def enable_validation(enabled: bool = True): + """Enable or disable validation globally""" + _validation_registry.enabled = enabled + + +def clear_validation_results(): + """Clear all validation results""" + _validation_registry.results.clear() + + +def compare_to_ttnn( + reference_fn: Callable, + *, + input_to_ttnn: Optional[Callable] = None, + output_to_ttnn: Optional[Callable] = None, + metric_tolerances: Optional[Dict[Any, Any]] = None, + enabled: bool = True, + raise_exceptions: bool = False, + return_reference_output: bool = False, +): + """ + Convenience wrapper for TTNN-on-device comparison. Provides useful visual cue to users that the reference function is a TTNN-native function. + + Args: + reference_fn: Reference function to compare against + input_to_ttnn: Maps decorated function inputs to reference function inputs + output_to_ttnn: Maps decorated function outputs to reference function outputs + metric_tolerances: Dictionary specifying tolerances and optionally custom metrics. + enabled: Whether validation is enabled (can disable globally via registry) + raise_exceptions: When True, re-raise any exceptions encountered during + reference execution, output mapping, or metric computation instead + of logging them into validation results. + + Examples: + @compare_to_ttnn( + reference_fn=lambda self, x: ttnn.matmul(x, self.weight), + input_to_ttnn=lambda self, x: (self, x), + ) + def __call__(self, x): + return torch.matmul(x, self.torch_weight) + # alternatively, the decorated function can return a TTNN tensor: return ttnn.from_torch(x) @ self.weight + + NOTES: + - The reference function is expected to accepts TTNN tensors and returns a TTNN tensor + - The decorated function inputs/outputs TTNN tensors, Torch tensors, or mixed TTNN and Torch tensors + - When decorated function returns torch tensors: + - the reference function's inputs will be constructed through either input_to_ttnn or from_torch(decorated function inputs, device=ttnn.GetDefaultDevice()) + - the metric on output tensor will be computed on the host + - Experimental support for on-device metric computation is provided and used when both the decorated function and the reference function return TTNN tensors + """ + + # Default converters: recursively convert any TTNN tensors to torch, auto-compose shards. + # Non-tensor objects are passed through unchanged. + + def _to_ttnn_auto(x: Any) -> Any: + if torch.is_tensor(x): + # Use auto-compose; relies on tensor.device() or a globally-set default device + assert ( + ttnn.GetDefaultDevice() is not None + ), "Default device is not set. It is required by compare_to_ttnn. Please set it via ttnn.SetDefaultDevice(...)." + return ttnn.from_torch(x, device=ttnn.GetDefaultDevice()) + return x + + def _default_input_map(*args, **kwargs): + ref_args = _map_structure(args, _to_ttnn_auto) + ref_kwargs = _map_structure(kwargs, _to_ttnn_auto) + return ref_args, ref_kwargs + + map_fn_to_match_sig = lambda tt_tensor, filler: to_torch_auto_compose(tt_tensor) + + return __validate_against( + reference_fn=reference_fn, + input_map=input_to_ttnn or _default_input_map, + output_map=output_to_ttnn, + metric_tolerances=metric_tolerances, + enabled=enabled, + raise_exceptions=raise_exceptions, + reference_output_map_fn=map_fn_to_match_sig if return_reference_output else None, + ) + + +def compare_to_torch( + reference_fn: Callable, + *, + input_to_torch: Optional[Callable] = None, + output_to_torch: Optional[Callable] = None, + metric_tolerances: Optional[Dict[Any, Any]] = None, + enabled: bool = True, + raise_exceptions: bool = False, + return_reference_output: Optional[Callable[..., bool] | bool] = False, +): + """ + Convenience wrapper for host/CPU comparison using torch. + + # Args: + # reference_fn: Reference function to compare against + # input_to_torch: Maps decorated function inputs to reference function inputs + # output_to_torch: Maps decorated function outputs to reference function outputs + # metric_tolerances: Dictionary specifying tolerances and optionally custom metrics. + # enabled: Whether validation is enabled (can disable globally via registry) + # raise_exceptions: When True, re-raise any exceptions encountered during + # reference execution, output mapping, or metric computation instead + # of logging them into validation results. + # + # Notes: + # - compare_to_torch is used when the reference function is a PyTorch function + # - the reference function takes as inputs to_torch_auto_compose(decorated function inputs) and compares the outputs with to_torch_auto_compose(decorated function outputs) + # - the decorated function inputs/outputs TTNN tensors, Torch tensors, or mixed TTNN and Torch tensors + """ + + # Default converters: recursively convert any TTNN tensors to torch, auto-compose shards. + # Non-tensor objects are passed through unchanged. + + def _to_torch_auto(x: Any) -> Any: + if isinstance(x, ttnn.Tensor): + # Use auto-compose; relies on tensor.device() or a globally-set default device + return to_torch_auto_compose(x) + return x + + def _default_input_map(*args, **kwargs): + ref_args = _map_structure(args, _to_torch_auto) + ref_kwargs = _map_structure(kwargs, _to_torch_auto) + return ref_args, ref_kwargs + + def _default_output_map(output): + return _map_structure(output, _to_torch_auto) + + return __validate_against( + reference_fn=reference_fn, + input_map=input_to_torch or _default_input_map, + output_map=output_to_torch or _default_output_map, + metric_tolerances=metric_tolerances, + enabled=enabled, + raise_exceptions=raise_exceptions, + reference_output_map_fn=from_torch_dist_as if return_reference_output else None, + ) + + +# ============================================================================ +# Data Structures +# ============================================================================ + + +@dataclass +class MetricResult: + """Per-metric validation outcome""" + + value: float = float("inf") + passed: bool = False + error: str = "" + + +@dataclass +class ValidationResult: + """Results from a single validation run""" + + function_name: str + passed: bool + # Map of metric name to its result (value/pass/fail/error) + metrics: Dict[Any, MetricResult] = field(default_factory=dict) + execution_time_impl: float = 0.0 + execution_time_ref: float = 0.0 + timestamp: float = field(default_factory=time.time) + logs: List[str] = field(default_factory=list) + + +class ValidationRegistry: + """Global registry for validation results""" + + def __init__(self): + self.results: List[ValidationResult] = [] + self.enabled = True + + def add_result(self, result: ValidationResult): + self.results.append(result) + + def get_summary(self) -> Dict[str, Any]: + """Get summary statistics of all validations""" + if not self.results: + return {"total": 0, "passed": 0, "failed": 0} + + passed = sum(1 for r in self.results if r.passed) + failed = len(self.results) - passed + + return { + "total": len(self.results), + "passed": passed, + "failed": failed, + "pass_rate": passed / len(self.results) if self.results else 0.0, + "avg_speedup": ( + sum(r.execution_time_ref / r.execution_time_impl for r in self.results if r.execution_time_impl > 0) + / len(self.results) + if self.results + else 0.0 + ), + } + + def print_report(self, verbose: bool = False): + """Print detailed validation report""" + summary = self.get_summary() + print("\n" + "=" * 80) + print("VALIDATION REPORT") + print("=" * 80) + print() + + for result in self.results: + status = "✓ PASS" if result.passed else "✗ FAIL" + print(f"{status} - {result.function_name}") + print( + f" Execution time: impl={result.execution_time_impl*1000:.2f}ms, ref={result.execution_time_ref*1000:.2f}ms" + ) + + if result.metrics: + print(f" Metrics:") + for metric_name, mres in result.metrics.items(): + # Use enum value for readability if metric is an Enum + name_str = metric_name.value if hasattr(metric_name, "value") else str(metric_name) + if mres.value is not None: + try: + val_str = f"{mres.value:.6f}" + except Exception: + val_str = str(mres.value) + else: + val_str = "-" + status = "PASS" if mres.passed else "FAIL" + print(f" {name_str}: {val_str} — {status}") + if mres.error: + print(f" error: {mres.error}") + + # Print any collected logs for this validation + if result.logs and verbose: + print(" Logs:") + for entry in result.logs: + try: + msg = str(entry) + except Exception: + msg = "" + print(f" {msg}") + + # All errors are reported via per-metric entries + print() + + print("-" * 36 + "Summary:" + "-" * 36) + print(f"Total validations: {summary['total']}") + print(f"Passed: {summary['passed']} ({summary['pass_rate']*100:.1f}%)") + print(f"Failed: {summary['failed']}") + print(f"Average speedup: {summary['avg_speedup']:.2f}x") + print() + print("=" * 80 + "\n") + + +# Global validation registry +_validation_registry = ValidationRegistry() + + +# ============================================================================ +# Validation Decorator +# ============================================================================ + + +class Metric(str, Enum): + """Enumeration of supported metric names, values match current string keys.""" + + MAX_ABS_ERROR = "max_abs_error" + MEAN_ABS_ERROR = "mean_abs_error" + PCC = "pcc" + + +@dataclass +class MetricSpec: + """Metric specification: name, tolerance, direction, and compute function.""" + + tolerance: float + higher_is_better: bool + compute_fn: Callable[[Any, Any], float] + name: str = field(default="") + + +# Registry of built-in metrics with defaults. Tolerances here are sensible +# defaults; callers can override per-validation via `tolerances`. +METRIC_SPECS: Dict[Metric, MetricSpec] = { + Metric.MAX_ABS_ERROR: MetricSpec( + name=Metric.MAX_ABS_ERROR.value, + tolerance=0.0, + higher_is_better=False, + compute_fn=DEFAULT_METRICS[Metric.MAX_ABS_ERROR.value], + ), + Metric.MEAN_ABS_ERROR: MetricSpec( + name=Metric.MEAN_ABS_ERROR.value, + tolerance=0.0, + higher_is_better=False, + compute_fn=DEFAULT_METRICS[Metric.MEAN_ABS_ERROR.value], + ), + Metric.PCC: MetricSpec( + name=Metric.PCC.value, + tolerance=0.0, + higher_is_better=True, + compute_fn=DEFAULT_METRICS[Metric.PCC.value], + ), +} + +# Convenience groupings for quick checks +HIGHER_IS_BETTER_METRICS = {m.value for m, spec in METRIC_SPECS.items() if spec.higher_is_better} +LOWER_IS_BETTER_METRICS = {m.value for m, spec in METRIC_SPECS.items() if not spec.higher_is_better} + + +# Helper: prefer Metric enum as dict key when possible +def _metric_key(key: Any) -> Any: + try: + return Metric(key) + except Exception: + return key + + +# Helper: Build active metrics map (name -> compute fn). Accept Metric enum keys for tolerances. +def _normalize_key(k: Any) -> str: + try: + # Enum or similar objects with .value as canonical string + return k.value if hasattr(k, "value") else str(k) + except Exception: + return str(k) + + +# Helper: Prepare metrics, tolerances, and directionality +def _prepare_metric_config(metric_tolerances_input): + metrics_map = {name: fn for name, fn in DEFAULT_METRICS.items()} + hib = set(HIGHER_IS_BETTER_METRICS) + logs_local: List[str] = [] + tol_map: Dict[str, float] = {} + + if not isinstance(metric_tolerances_input, dict): + logs_local.append(f"metric_tolerances_input must be a dict, got {type(metric_tolerances_input)}") + metric_tolerances_input = dict() + + if not metric_tolerances_input: + logs_local.append("no metric tolerances provided") + metric_tolerances_input = dict() + + for raw_key, spec in metric_tolerances_input.items(): + name = _normalize_key(raw_key) + if isinstance(spec, MetricSpec): + tol_map[name] = float(spec.tolerance) + metrics_map[name] = spec.compute_fn + spec.name = name if spec.name == "" else spec.name + if spec.higher_is_better: + hib.add(name) + else: + hib.discard(name) + continue + try: + tol_map[name] = float(spec) + except Exception: + logs_local.append(f"unrecognized tolerance: {raw_key}: {spec}") + + return metrics_map, hib, tol_map, logs_local + + +# todo)) also allow raise an exception from the a failed metric! + +# todo)) add support for multiple outputs from the reference function and the decorated function! +# e.g., return logits, past_key_values, etc. + +# todo)) make sure the dtypes are taken care of in the validate_against decorator! +# e.g., if the decorated function is of dtype bfp4, what is the dtype of the to_torch_auto_compose output? + +# todo)) add file line number to the validation results! + +# todo)) add function to export the validation results to a csv file! + +# todo)) enhance report to use file line number as index to summarize the validation results +# e.g., ✗ FAIL - __main__.Attention.__call__ (line 100) -> 100 failed validations + +# todo)) remove compile time from speed up calculation -- e.g., 9118.15ms should be removed in the example below: +# ================================================================================ +# VALIDATION REPORT +# ================================================================================ +# Total validations: 1400 +# Passed: 1400 (100.0%) +# Failed: 0 +# Average speedup: 0.97x + +# ✓ PASS - __main__.TransformerBlock.__call__ +# Execution time: impl=9118.15ms, ref=14.84ms +# Metrics: +# pcc: 0.999743 — PASS + +# ✓ PASS - __main__.TransformerBlock.__call__ +# Execution time: impl=3.05ms, ref=12.73ms +# Metrics: +# pcc: 0.999913 — PASS + +# ✓ PASS - __main__.TransformerBlock.__call__ +# Execution time: impl=3.31ms, ref=12.48ms +# Metrics: +# pcc: 0.999962 — PASS + +# ✓ PASS - __main__.TransformerBlock.__call__ +# Execution time: impl=3.11ms, ref=12.89ms +# Metrics: +# pcc: 1.000000 — PASS + +# ✓ PASS - __main__.TransformerBlock.__call__ +# Execution time: impl=3.16ms, ref=12.97ms +# Metrics: +# pcc: 0.999998 — PASS + + +# todo)) stretch goals: +# - generate unit test automatically from the failed validations +def __validate_against( + reference_fn: Callable, + *, + input_map: Optional[Callable] = None, + output_map: Optional[Callable] = None, + metric_tolerances: Optional[Dict[Any, Any]] = None, + enabled: bool = True, + raise_exceptions: bool = False, + reference_output_map_fn: Optional[Callable] = None, +): + """ + Decorator to validate a function against a reference implementation. + + Args: + reference_fn: Reference function to compare against + input_map: Maps decorated function inputs to reference function inputs + Signature: (args, kwargs) -> (ref_args, ref_kwargs) + If None, inputs are passed as-is + output_map: Converts impl output to match ref output's type + Signature: (output) -> comparable_output + Applied ONLY to impl_output to convert it to ref_output's type + Common use: lambda x: ttnn.to_torch(x).squeeze() to convert ttnn → torch + If None, outputs are used as-is (both must already be same type) + metric_tolerances: Dictionary specifying tolerances and optionally custom metrics. + Accepts the following per metric key (str or Metric): + - float: tolerance only (uses built-in compute + direction) + - MetricSpec instance + Validation fails if any metric exceeds its tolerance + enabled: Whether validation is enabled (can disable globally via registry) + raise_exceptions: When True, re-raise any exceptions encountered during + reference execution, output mapping, or metric computation instead + of logging them into validation results. + + Examples: + # Pattern 1: TTNN-native metrics (recommended, 100-1000× faster!) + # Both impl and ref return ttnn.Tensor, no output_map needed + def _reference_impl(self, x): + x_torch = ttnn.to_torch(x).squeeze(0) + result_torch = torch.matmul(x_torch, self.weight_torch) + # Convert back to TTNN for on-device metrics! + return ttnn.from_torch(result_torch.unsqueeze(0), device=self.device, ...) + + @validate_against( + reference_fn=lambda self, x: self._reference_impl(x), + tolerances={'max_abs_error': 1e-3} + ) + def __call__(self, x): + return ttnn.matmul(x, self.weight) + + # Pattern 2: PyTorch metrics (when reference returns torch.Tensor) + # Use output_map to convert impl output (ttnn.Tensor) to match ref (torch.Tensor) + @validate_against( + reference_fn=torch.nn.functional.rms_norm, + input_map=lambda args, kwargs: ( + (ttnn.to_torch(args[1]).squeeze(),), + {'eps': args[0].eps} + ), + output_map=lambda x: ttnn.to_torch(x).squeeze(), # Convert impl: ttnn → torch + tolerances={'max_abs_error': 1e-3} + ) + def __call__(self, x): + return ttnn.rms_norm(x, self.weight, self.eps) # Returns ttnn.Tensor + """ + + if metric_tolerances is None: + metric_tolerances = { + Metric.MAX_ABS_ERROR: 1e-2, + Metric.PCC: 0.99, + } + + metrics_to_use, higher_is_better_effective, tolerances_map, pre_logs = _prepare_metric_config(metric_tolerances) + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + # Check if validation is enabled + if not enabled or not _validation_registry.enabled: + return func(*args, **kwargs) + + # Execute implementation + start_time = time.perf_counter() + impl_output = func(*args, **kwargs) + impl_time = time.perf_counter() - start_time + logs: List[str] = pre_logs.copy() + + # Map inputs for reference function: prefer input_map, else pass-through + if input_map: + _nm = getattr(input_map, "__name__", None) or type(input_map).__name__ + logs.append(f"input_map={_nm}") + try: + mapped = input_map(*args, **kwargs) + except Exception as e: + # If input mapping fails, log error, record result, and return impl output + logs.append(f"input_mapping_error={str(e)}") + result = ValidationResult( + function_name=f"{func.__module__}.{func.__qualname__}", + passed=False, + metrics={ + "input_mapping": MetricResult( + value=None, passed=False, error=f"Input mapping failed: {str(e)}" + ) + }, + execution_time_impl=impl_time, + execution_time_ref=0.0, + logs=logs, + ) + _validation_registry.add_result(result) + # Re-raise exception if raise_exceptions is True + if raise_exceptions: + raise + return impl_output + # Normalize mapper output: + # - If (ref_args, ref_kwargs) with kwargs as dict, use directly + # - Otherwise, treat return as positional args and use empty kwargs + if isinstance(mapped, tuple) and len(mapped) == 2 and isinstance(mapped[1], dict): + ref_args, ref_kwargs = mapped + else: + ref_args = mapped if isinstance(mapped, (list, tuple)) else (mapped,) + ref_kwargs = {} + else: + logs.append("input_map=pass-through") + ref_args, ref_kwargs = args, kwargs + + # Execute reference + try: + start_time = time.perf_counter() + ref_output = reference_fn(*ref_args, **ref_kwargs) + ref_time = time.perf_counter() - start_time + except Exception as e: + # If reference fails, just return impl output and log error via metrics + logs.append(f"reference_execution_error={str(e)}") + # Record elapsed time until failure + ref_time = time.perf_counter() - start_time + result = ValidationResult( + function_name=f"{func.__module__}.{func.__qualname__}", + passed=False, + metrics={ + "reference_execution": MetricResult( + value=None, passed=False, error=f"Reference execution failed: {str(e)}" + ) + }, + execution_time_impl=impl_time, + execution_time_ref=ref_time, + logs=logs, + ) + _validation_registry.add_result(result) + # Re-raise exception if raise_exceptions is True + if raise_exceptions: + raise + return impl_output + + # Map outputs for comparison + # Note: output_map only applies to impl_output to convert it to match ref_output's type + try: + _nm = getattr(output_map, "__name__", None) or type(output_map).__name__ + logs.append(f"output_map={_nm}") + impl_comparable = output_map(impl_output) if output_map else impl_output + ref_comparable = ref_output # Reference output is always used as-is + except Exception as e: + logs.append(f"output_mapping_error={str(e)}") + result = ValidationResult( + function_name=f"{func.__module__}.{func.__qualname__}", + passed=False, + metrics={ + "output_mapping": MetricResult( + value=None, passed=False, error=f"Output mapping failed: {str(e)}" + ) + }, + execution_time_impl=impl_time, + execution_time_ref=ref_time, + logs=logs, + ) + _validation_registry.add_result(result) + # Re-raise exception if raise_exceptions is True + if raise_exceptions: + raise + return impl_output + + # Compute metrics + computed_metrics: Dict[Any, MetricResult] = {} + passed = True + + for metric_name, threshold in tolerances_map.items(): + try: + metric_fn = metrics_to_use.get(metric_name) + + # Store results keyed by enum when available + metric_key = _metric_key(metric_name) + # If metric function isn't known, record an error + if metric_fn is None: + computed_metrics[metric_key] = MetricResult( + value=None, passed=False, error=f"Unknown metric: {metric_name}" + ) + passed = False + continue + + value = metric_fn(impl_comparable, ref_comparable) + + # Determine direction using registry when available + if metric_name in higher_is_better_effective: + ok = value >= threshold + err = None + if not ok: + passed = False + err = f"{metric_name}={value:.6e} below threshold {threshold:.6e}" + computed_metrics[metric_key] = MetricResult(value=value, passed=ok, error=err) + else: + ok = value <= threshold + err = None + if not ok: + passed = False + err = f"{metric_name}={value:.6e} exceeds tolerance {threshold:.6e}" + computed_metrics[metric_key] = MetricResult(value=value, passed=ok, error=err) + except Exception as e: + msg = f"Metric {metric_name} failed: {str(e)}" + computed_metrics[metric_key] = MetricResult(value=None, passed=False, error=msg) + passed = False + if raise_exceptions: + raise + + # Optionally return the (aligned) reference output instead of impl output + backup_impl_output = impl_output + try: + if reference_output_map_fn: + impl_output = reference_output_map_fn(ref_output, impl_output) + logs.append(f"reference_output_mapping_fn={reference_output_map_fn.__name__}") + except Exception as e: + # If alignment fails, fall back to impl output + impl_output = backup_impl_output + # Re-raise exception if raise_exceptions is True after logging the error + logs.append(f"reference_output_mapping_error={str(e)}") + if raise_exceptions: + raise + + # Record results + pass_count = sum(1 for v in computed_metrics.values() if v.passed) + fail_count = sum(1 for v in computed_metrics.values() if not v.passed) + logs.append(f"metrics={pass_count}_pass,{fail_count}_fail") + result = ValidationResult( + function_name=f"{func.__module__}.{func.__qualname__}", + passed=passed, + metrics=computed_metrics, + execution_time_impl=impl_time, + execution_time_ref=ref_time, + logs=logs, + ) + _validation_registry.add_result(result) + + return impl_output + + return wrapper + + return decorator + + +def _map_structure(obj: Any, fn: Callable[[Any], Any]) -> Any: + """ + Map a structure of objects to a new structure using a function. + """ + if isinstance(obj, (list, tuple)): + mapped = [_map_structure(x, fn) for x in obj] + return type(obj)(mapped) + if isinstance(obj, dict): + return {k: _map_structure(v, fn) for k, v in obj.items()} + return fn(obj) diff --git a/code/models/common/weight_cache.py b/code/models/common/weight_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..6dd08f59b00e36e2f6e4be420dd1c4e00b578566 --- /dev/null +++ b/code/models/common/weight_cache.py @@ -0,0 +1,426 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Shared, ModelArgs-agnostic warm ttnn weight-cache helpers (generalizes PR #50550 / #48531 +to forked model loaders — issue #45400 follow-up). + +On a warm cache, a model can be built from its on-disk ``.tensorbin`` files without the +expensive host-side HF ``from_pretrained`` load (the load that OOMs/hangs during prefill, +#48509): ``ttnn.as_tensor(torch_weight, cache_file_name=...)`` loads the cached tensor and +ignores ``torch_weight`` on a cache hit (see ttnn/operations/core.py). So most weights only +need a dataless placeholder (``torch.empty`` of the right shape/dtype) to satisfy the modules' +host-side reshape ops before ``as_tensor``. + +Some forks (e.g. gemma4) additionally consume a *small* set of weights on the host — token +embeddings used via ``F.embedding``, per-layer scalars read via ``.item()``, etc. Those must be +real. ``mark_weight_cache_complete`` persists exactly those tensors to a sidecar at cold-build +time (write-access run), and ``build_cached_state_dict`` serves them real on later warm runs +while placeholdering the rest — a HYBRID state_dict. The host subset is a tiny fraction of the +weight bytes, so the full from_pretrained (and its OOM) is still avoided. +""" + +import collections.abc +import hashlib +import json +import os +from pathlib import Path + +import torch +from loguru import logger + +WEIGHT_CACHE_MARKER = ".weights_complete" +HOST_WEIGHTS_SIDECAR = ".host_weights.pt" +# Bump when the set/naming/layout of cached weights, or this marker schema, changes such that an +# existing cache would not satisfy a new build. A marker written by an older format is rejected, +# so the run cold-loads and regenerates rather than building from an incompatible cache. +# v2: model/n_layers/mesh_shape validation + a {key: [shape, dtype]} manifest. +# v3: canonical mesh_shape encoding shared with ModelArgs (the two writers previously encoded it +# differently and each rejected the other's marker), a `components` field so a text-only seed +# cannot certify a cache for a build that also needs the vision tower, and `cache_files` -- +# the recursive list of .tensorbin files the completed build actually produced, verified +# per-file on read. That last one is load-bearing: ttnn.as_tensor PERSISTS whatever tensor it +# is handed on a cache miss, so a marker that outlives some of its tensorbins would otherwise +# dump placeholders to disk as real cache entries -- silent, permanent corruption. Verifying +# the recorded file set turns every such case back into a plain cold load. Also `build_variant` +# -- the build options (prefetcher, precision) that change an as_tensor cache FILENAME, matched +# exactly, because a different variant needs different files rather than fewer. +WEIGHT_CACHE_FORMAT_VERSION = 3 + +DEFAULT_FORCE_ENV = "TT_TRANSFORMERS_FORCE_MODEL_LOAD" + + +def _variant_digest(build_variant): + """Stable short digest of a build_variant dict ("none" for None).""" + if build_variant is None: + return "none" + return hashlib.sha1(json.dumps(build_variant, sort_keys=True, default=str).encode()).hexdigest()[:12] + + +def _variant_unverifiable(build_variant): + return bool(build_variant) and bool(build_variant.get("unverifiable")) + + +def marker_path(cache_path, build_variant=None): + """The marker file for one (cache dir, build variant). + + The variant digest is part of the FILENAME, not just a field compared inside one shared + marker. A cache dir legitimately serves several build variants (the Llama CI job runs + eval-32 with and without the DRAM prefetcher against the same instruct cache), and a single + marker matched exactly would make each variant's seed evict the other's on every run -- both + then cold-load forever with nothing going red. One marker per variant lets them coexist. + (#45400 review, finding B3)""" + return Path(cache_path) / f"{WEIGHT_CACHE_MARKER}.{_variant_digest(build_variant)}" + + +def _dtype_from_str(s): + return getattr(torch, s.rsplit(".", 1)[-1]) + + +def normalize_mesh_shape(mesh_shape): + """Canonical marker encoding for a mesh shape. + + ``ttnn.MeshShape`` stringifies as ``MeshShape([1, 8])`` while callers that pass a plain tuple + stringify as ``(1, 8)``. Both writers must agree or each rejects the other's marker and the + model cold-loads forever (gemma3 inherits ModelArgs but its demos call this module). Normalize + everything to a plain tuple-of-ints string.""" + try: + return str(tuple(int(d) for d in mesh_shape)) + except TypeError: + return str(mesh_shape) + + +def _normalize_components(components): + """Canonical component list. ``None`` means "the whole model as this loader builds it" and is + encoded as a single implicit component so old-style callers stay self-consistent.""" + if components is None: + return ["all"] + if isinstance(components, str): + return [components] + return sorted(str(c) for c in components) + + +def list_cache_files(cache_path): + """Every ``.tensorbin`` under ``cache_path``, recursively, as sorted relative POSIX paths. + + Recursive because forked loaders nest per-layer weights in subdirectories (qwen36 + ``layers.{n}/``, gemma4 ``layer_{i}/``); a top-level ``glob`` would call a cache complete when + only the root-level ``output.weight`` survived an interrupted seed.""" + cache_path = Path(cache_path) + return sorted(p.relative_to(cache_path).as_posix() for p in cache_path.rglob("*.tensorbin")) + + +# One-entry cache so the completeness gate's validation load is reused by +# build_cached_state_dict instead of torch.load-ing the same multi-GB file twice per warm run +# (gemma-4-31b's embedding alone is ~2.8 GB). Keyed on (path, mtime, size) so a republished +# sidecar is never served stale; the builder consumes the entry so the tensors are not pinned +# past the build. (#45400 review, finding R1) +_SIDECAR_CACHE = {} + + +def load_host_sidecar(cache_path, *, consume=False): + """Load the host-weights sidecar, or None if absent/unreadable. + + ``consume=True`` drops the memoized entry after returning it (the caller takes ownership).""" + sidecar = Path(cache_path) / HOST_WEIGHTS_SIDECAR + if not sidecar.is_file(): + return None + try: + st = sidecar.stat() + key = (str(sidecar), st.st_mtime_ns, st.st_size) + host = _SIDECAR_CACHE.get(key) + if host is None: + host = torch.load(sidecar, map_location="cpu", weights_only=True) + _SIDECAR_CACHE.clear() + _SIDECAR_CACHE[key] = host + if consume: + _SIDECAR_CACHE.pop(key, None) + return host + except Exception: + return None + + +def weight_cache_is_complete( + cache_path, + *, + model_name, + n_layers, + mesh_shape, + components=None, + build_variant=None, + force_env=DEFAULT_FORCE_ENV, +): + """True when the on-disk ttnn weight cache at ``cache_path`` was fully built by a previous run + for this exact build, and every tensorbin that build produced is still present. + + ``components`` names the model parts this build will construct (e.g. ``"text"`` vs + ``"text+vision"``); a marker written by a narrower build does not satisfy a wider one, because + the wider build needs tensorbins the narrower one never wrote. ``force_env=...=1`` forces a + cold load.""" + if force_env and os.getenv(force_env) == "1": + return False + # A variant we could not compute is a variant we cannot verify: accepting it could hand a + # placeholder to a build whose cache-filename set we did not check, and as_tensor would + # persist that placeholder to disk. Fail closed, loudly. (#45400 review, finding R3) + if _variant_unverifiable(build_variant): + logger.warning( + f"Warm-cache check for {cache_path}: build_variant could not be computed " + f"({build_variant.get('error', 'unknown error')}); forcing a cold load." + ) + return False + cache_path = Path(cache_path) + marker = marker_path(cache_path, build_variant) + if not marker.is_file(): + return False + try: + meta = json.loads(marker.read_text()) + except (ValueError, OSError): + return False + if meta.get("format_version") != WEIGHT_CACHE_FORMAT_VERSION: + return False + if meta.get("model_name") != model_name or meta.get("n_layers") != n_layers: + return False + if meta.get("mesh_shape") != normalize_mesh_shape(mesh_shape): + return False + # The recorded build must cover every component this build needs. Superset is fine (a + # text+vision seed wrote the text tensorbins too, so it satisfies a text-only build); a subset + # is not (a text-only seed never wrote the vision tower's tensorbins, and accepting it would + # make as_tensor dump placeholders for them). + if not set(_normalize_components(components)).issubset(set(meta.get("components") or [])): + return False + # Build options that change an as_tensor cache FILENAME (prefetcher, precision) must match + # exactly. A superset rule is wrong here: a different variant does not need fewer files, it + # needs DIFFERENT ones, and any it is missing would be regenerated from the placeholder. + if meta.get("build_variant") != build_variant: + return False + if not meta.get("weights"): + return False + # Every tensorbin the completed build produced must still be on disk. Any missing file would + # otherwise be regenerated by as_tensor FROM THE PLACEHOLDER we are about to hand it, writing + # garbage into the cache permanently. Missing file => cold load, which rebuilds it correctly. + recorded = meta.get("cache_files") + if not recorded: + return False + present = set(list_cache_files(cache_path)) + if not all(f in present for f in recorded): + return False + # If host weights were captured, the sidecar must be present AND loadable. A torn/corrupt + # sidecar (interrupted or racing seed) must fall back to a cold load -- the way a torn marker + # already does via the except above -- rather than pass this gate and then crash torch.load on + # every subsequent run, bricking the cache dir. Checked LAST so the load it performs is + # memoized only when the gate is about to pass, for build_cached_state_dict to consume. + # (#45400 review) + if meta.get("host_weights") and load_host_sidecar(cache_path) is None: + return False + return True + + +def mark_weight_cache_complete( + cache_path, + state_dict, + *, + model_name, + n_layers, + mesh_shape, + components=None, + build_variant=None, + is_moe=False, + is_host_weight=None, +): + """Record that the ttnn weight cache at ``cache_path`` is fully built. + + Writes a ``.weights_complete`` marker holding a ``{key: [shape, dtype]}`` manifest of every + weight plus the recursive list of ``.tensorbin`` files this build produced (verified per-file + on read). If ``is_host_weight(key)`` is provided, the (real) tensors it matches are also saved + to a ``.host_weights.pt`` sidecar so a later warm run can serve them for real (hybrid). + + Call this only AFTER the model has been constructed, so the tensorbins exist to be recorded.""" + if _variant_unverifiable(build_variant): + # Never certify a cache under an identity we could not compute -- a later run computing + # the same error string would otherwise warm-match it. (#45400 review, finding R3) + logger.warning( + f"Not marking weight cache complete at {cache_path}: build_variant could not be " + f"computed ({build_variant.get('error', 'unknown error')})." + ) + return + cache_path = Path(cache_path) + marker = marker_path(cache_path, build_variant) + weights = {} + host = {} + for k, v in state_dict.items(): + shape = getattr(v, "shape", None) + dt = getattr(v, "dtype", None) + if shape is None or dt is None: + continue # skip non-tensor entries + weights[k] = [list(shape), str(dt)] + if is_host_weight is not None and is_host_weight(k): + host[k] = v + try: + cache_path.mkdir(parents=True, exist_ok=True) + cache_files = list_cache_files(cache_path) + if not cache_files: + logger.warning(f"Not marking weight cache complete: no .tensorbin files under {cache_path}") + return + # Write both the sidecar and the marker atomically (temp file + os.replace, atomic on + # POSIX). Two jobs can seed the same (model, dtype, mesh) dir on one host concurrently, and + # an interrupted write must never leave a torn file that a later run picks up: a half-written + # sidecar would otherwise pass the is_file() gate and crash torch.load on every subsequent + # run. The temp name is pid-unique so two concurrent seeders cannot write the SAME temp + # inode -- with a fixed name, B could publish the file while A was still writing into it. + # Sidecar first, then marker, so the completeness gate only appears once its sidecar is + # fully in place. (#45400 review) + uniq = os.getpid() + if host: + sidecar = cache_path / HOST_WEIGHTS_SIDECAR + sidecar_tmp = sidecar.with_suffix(sidecar.suffix + f".tmp.{uniq}") + torch.save(host, sidecar_tmp) + os.replace(sidecar_tmp, sidecar) + marker_body = json.dumps( + { + "format_version": WEIGHT_CACHE_FORMAT_VERSION, + "model_name": model_name, + "n_layers": n_layers, + "mesh_shape": normalize_mesh_shape(mesh_shape), + "components": _normalize_components(components), + "build_variant": build_variant, + "cache_files": cache_files, + "is_moe": bool(is_moe), + "host_weights": sorted(host.keys()), + "weights": weights, + } + ) + marker_tmp = marker.with_suffix(marker.suffix + f".tmp.{uniq}") + marker_tmp.write_text(marker_body) + os.replace(marker_tmp, marker) + logger.info(f"Marked ttnn weight cache complete: {marker} ({len(weights)} weights, {len(host)} host-loaded)") + except Exception as e: + # Deliberately broad: this function only RECORDS completion -- failing to record must + # never kill a build that already succeeded. The concrete case: on a read-only + # /mnt/MLPerf, torch.save of the host sidecar raises RuntimeError from torch's C++ + # serializer (inline_container.cc "Read-only file system"), not OSError, and the narrow + # except crashed every read-only cold run of the sidecar models (gemma4/gemma3) right + # after a successful build. (#45400 review, finding R5; seen on Gemma-4-E2B bh_p150, + # run 32511945147) + logger.warning(f"Could not write weight-cache completion marker {marker}: {e}") + + +class CachedStateDict(collections.abc.MutableMapping): + """A stand-in ``state_dict`` for warm-cache builds. + + Serves the real tensor for keys captured in the host-weights sidecar; for every other key it + returns a fresh dataless ``torch.empty`` of the manifest shape/dtype (which ``ttnn.as_tensor`` + discards on the guaranteed cache hit). Mutable (some loaders ``setdefault`` missing KV-shared + weights) and truthy (some loaders gate real-weight loading on ``if state_dict:``).""" + + # Explicit marker that this is a warm-cache stand-in, NOT real weights. Callers that must tell + # "warm-cache placeholder" apart from "real weights" MUST branch on this attribute, never on + # truthiness: this mapping is truthy (non-zero __len__) but tt_transformers' _PlaceholderStateDict + # is falsy (__bool__ -> False), so a truthiness test silently means opposite things for the two. + # If tt_transformers is ever collapsed onto this class (a listed follow-up), the attribute keeps + # `if is_placeholder(...)` reload sites (e.g. test_model_prefill) correct. (#45400 review) + is_placeholder = True + + def __init__(self, manifest, host): + self._manifest = manifest # key -> (shape, dtype_str) + self._host = dict(host or {}) # key -> real torch.Tensor + self._overrides = {} # keys set by the caller at build time + self._deleted = set() + + def __getitem__(self, key): + if key in self._deleted: + raise KeyError(key) + if key in self._overrides: + return self._overrides[key] + if key in self._host: + return self._host[key] + spec = self._manifest.get(key) + if spec is None: + raise KeyError(key) + shape, dt = spec + return torch.empty(tuple(shape), dtype=_dtype_from_str(dt)) + + def __setitem__(self, key, value): + self._deleted.discard(key) + self._overrides[key] = value + + def __delitem__(self, key): + if key not in self: + raise KeyError(key) + self._overrides.pop(key, None) + if key in self._host or key in self._manifest: + self._deleted.add(key) + + def __iter__(self): + seen = set() + for k in list(self._overrides) + list(self._host) + list(self._manifest): + if k in self._deleted or k in seen: + continue + seen.add(k) + yield k + + def __len__(self): + return sum(1 for _ in self) + + # Mapping's default __contains__/get/items route through __getitem__, which allocates a + # full-size torch.empty for EVERY key touched -- including multi-GB ones like lm_head.weight. + # substate() (models/tt_dit/utils/substate.py) iterates .items() and filters by prefix, so a + # 62-layer gemma4 build would allocate the entire model once per layer just to discard it. + # Answer membership from the key sets, and make items() lazy so only matching keys materialize. + def __contains__(self, key): + if key in self._deleted: + return False + return key in self._overrides or key in self._host or key in self._manifest + + def keys(self): + return list(self) + + def items(self): + for k in self: + yield k, self[k] + + def get(self, key, default=None): + if key not in self: + return default + return self[key] + + +def build_cached_state_dict(cache_path, host=None, args=None, build_variant=None): + """Build the warm-cache stand-in ``state_dict`` from the marker manifest + host sidecar. + + ``host`` may be a sidecar dict already loaded by ``weight_cache_is_complete``'s validation, to + avoid a second multi-GB ``torch.load`` of the same file on every warm run (gemma-4-31b's + embedding alone is ~2.8 GB). + + ``args`` (a ModelArgs-like) has ``is_mixture_of_experts`` restored from the marker. That flag is + normally set as a side effect of ``load_state_dict`` (by sniffing for ``.experts.`` keys), which + the warm path skips -- so without this a MoE checkpoint would build a dense decoder and die on a + missing ``feed_forward.w1.weight``. (#45400 review)""" + cache_path = Path(cache_path) + meta = json.loads(marker_path(cache_path, build_variant).read_text()) + manifest = meta["weights"] + if args is not None and hasattr(args, "__dict__"): + args.is_mixture_of_experts = bool(meta.get("is_moe", False)) + # fuse_qkv / fuse_mlp are normally sniffed from the checkpoint keys inside load_state_dict, + # which the warm path skips -- leaving them at their __init__ defaults and silently changing + # how the decoder is built. The manifest holds the same key set, so derive them identically. + keys = manifest.keys() + args.fuse_qkv = any("qkv" in k for k in keys) + args.fuse_mlp = any("gate_up" in k for k in keys) + if args.is_mixture_of_experts: + args.moe = True + expert_indices = [int(k[-11]) + 1 for k in keys if "block_sparse_moe.experts" in k] + if expert_indices: + args.num_experts = max(expert_indices) + elif hasattr(args, "num_local_experts"): + args.num_experts = args.num_local_experts + if host is None and meta.get("host_weights"): + # consume=True: reuse the load the completeness gate just performed and release the + # memoized entry, so the sidecar is read from NAS once per warm run, not twice. (R1) + host = load_host_sidecar(cache_path, consume=True) + host = host or {} + logger.info( + f"Warm ttnn weight cache: built state_dict for {len(manifest)} weights " + f"({len(host)} real host weights, no full HF load)." + ) + return CachedStateDict(manifest, host) diff --git a/code/models/tt_dit/README.md b/code/models/tt_dit/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cb3e1d0cd71041c5f3ed18e253d0d7f32848a994 --- /dev/null +++ b/code/models/tt_dit/README.md @@ -0,0 +1,147 @@ +# TT-DiT: Tenstorrent Diffusion Transformers + +This directory contains the implementation of the Tenstorrent Diffusion Transformer (TT-DiT) architecture, designed for optimized parallelism of DiTs for image and video generation. + +## Supported Models + +For detailed information about each model including performance metrics, usage instructions, and specific requirements, see: + +- **[Stable Diffusion 3.5 Large](models/StableDiffusion35.md)** - Text-to-image generation +- **[Flux 1](models/Flux1.md)** - Text-to-image generation (schnell & dev variants) +- **[Motif](models/Motif.md)** - Text-to-image generation model +- **[Qwen-Image](models/QwenImage.md)** - Text-to-image generation model +- **[Ideogram 4.0](models/Ideogram4.md)** - Text-to-image generation model (best-in-class typography; Blackhole) +- **[Mochi-1](models/Mochi_1.md)** - Video generation model +- **[Wan2.2-T2V-A14B](models/Wan2_2.md)** - Text-to-video generation model +- See **[experimental/](experimental/)** for in-progress model variants (AniSora, Wan2.2-Distill, LoRA adapters) +- **[LTX-2.3](models/LTX2.md)** - Joint audio-video generation model (Pro and Fast variants) + +## Directory Structure + +``` +tt_dit/ +├── layers/ # Core neural network layers +├── models/ # Model architectures and documentation +│ ├── transformers/ # Transformer implementations (SD35, Mochi, Wan, LTX, Flux1, Motif, QwenImage) +│ ├── vae/ # VAE/Autoencoder implementations +│ ├── StableDiffusion35.md # SD3.5 model documentation +│ ├── Flux1.md # Flux 1 model documentation +│ ├── Motif.md # Motif model documentation +│ ├── QwenImage.md # Qwen-Image model documentation +│ ├── Mochi_1.md # Mochi-1 model documentation +│ ├── Wan2_2.md # Wan2.2 model documentation +│ └── LTX2.md # LTX-2.3 model documentation +├── encoders/ # Text encoder implementations +│ ├── clip/ # CLIP encoder +│ ├── gemma/ # Gemma encoder +│ └── t5/ # T5 encoder +├── parallel/ # Parallelization utilities +│ ├── config.py # Parallel configuration +│ └── manager.py # Parallel execution management +├── pipelines/ # End-to-end model pipelines +│ ├── stable_diffusion_35_large/ +│ ├── mochi/ +│ ├── wan/ +│ ├── ltx/ +│ ├── flux1/ +│ ├── motif/ +│ └── qwenimage/ +├── tests/ # Test suite +│ ├── models/ # Model-level tests (sd35, mochi, wan2_2, ltx, flux1, motif, qwenimage) +│ ├── encoders/ # Encoder tests +│ ├── blocks/ # Block-level tests +│ └── unit/ # Unit tests for layers +└── utils/ # Utility functions + ├── check.py # Validation utilities + ├── padding.py # Padding operations + ├── substate.py # State management + └── tensor.py # Tensor operations +``` + +## Core Components + +### Layers +- **Embeddings**: Position and token embedding implementations +- **Feedforward**: MLP and feed-forward network components +- **Linear**: Optimized linear transformation layers +- **Normalization**: Layer normalization implementations + +### Models +- **Transformers**: DiT transformer implementations for various generative models + - Support for multiple architectures including SD3.5, Mochi, Wan2.2, LTX-2, Flux1, Motif, and Qwen-Image + - Model-specific attention mechanisms and transformer architectures +- **Autoencoders**: VAE implementations for different models + +### Parallel Processing +- **Config**: Parallel configuration management with `DiTParallelConfig` +- **Manager**: Parallel execution and device management + +### Pipelines +End-to-end pipeline implementations for multiple generative models: +- **Stable Diffusion 3.5 Large**: Text-to-image generation (1024x1024px) +- **Flux 1**: Text-to-image generation (schnell & dev variants, 1024x1024px) +- **Motif**: Text-to-image generation (6B model, 1024x1024px) +- **Qwen-Image**: Text-to-image generation (1024x1024px) +- **Mochi-1**: Video generation model (824x480px, 168 frames) +- **Wan2.2-T2V-A14B**: Text-to-video generation +- **LTX-2.3**: Joint audio-video generation (Pro one-stage and Fast distilled two-stage) + +Each pipeline includes: +- Automatic parallel configuration for different device meshes +- Optimized execution on Wormhole systems +- Comprehensive timing collection and profiling + +## Testing + +The test suite is organized into three main categories: + +1. **Unit Tests** (`tests/unit/`): + - Layer-level testing + - Individual component validation + - Coverage for embeddings, feedforward, linear, and normalization layers + +2. **Model Tests** (`tests/models/`): + - Smoke tests for model components + - Performance testing + - Pipeline validation + +Running tests: +```bash +# Run tests from tt-metal/ + +# Run unit tests +python -m pytest models/tt_dit/tests/unit/ + +# Run all model tests +python -m pytest models/tt_dit/tests/models/ + +# Run specific model pipeline tests +python -m pytest models/tt_dit/tests/models/sd35/test_pipeline_sd35.py -v +python -m pytest models/tt_dit/tests/models/flux1/test_pipeline_flux1.py -v +python -m pytest models/tt_dit/tests/models/motif/test_pipeline_motif.py -v +python -m pytest models/tt_dit/tests/models/qwenimage/test_pipeline_qwenimage.py -v +python -m pytest models/tt_dit/tests/models/mochi/test_pipeline_mochi.py -v +python -m pytest models/tt_dit/tests/models/wan2_2/test_pipeline_wan.py -v +python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py -v +python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py -v +``` + +## Key Features + +- **Modular Architecture**: Clean separation between layers, models, and parallel processing +- **Optimized Parallelism**: + - Automatic configuration of tensor and sequence parallelism + - Efficient device mesh management + - Dedicated parallel managers for different model components +- **Performance Monitoring**: + - Built-in timing collection + - Performance profiling support + - Tracing capabilities for optimization + +## Contributing + +When adding new components: +1. Place implementation in appropriate subdirectory +2. Add corresponding unit tests +3. Update model tests if necessary +4. Ensure parallel processing compatibility diff --git a/code/models/tt_dit/blocks/attention_opt.py b/code/models/tt_dit/blocks/attention_opt.py new file mode 100644 index 0000000000000000000000000000000000000000..5f12bb2433f4814b421084a96a676adcf1d6d65c --- /dev/null +++ b/code/models/tt_dit/blocks/attention_opt.py @@ -0,0 +1,702 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from loguru import logger + +import ttnn +from models.common.utility_functions import is_blackhole + +from ..layers.linear import ColParallelLinear +from ..layers.module import Module, Parameter, UnregisteredModule +from ..layers.normalization import DistributedRMSNorm +from ..utils.matmul import get_matmul_config, get_matmul_core_grid +from ..utils.mochi import get_rot_transformation_mat +from ..utils.padding import PaddingConfig, pad_weight_tensor +from ..utils.substate import pop_substate + +if TYPE_CHECKING: + pass + + from ..parallel.config import DiTParallelConfig + from ..parallel.manager import CCLManager + + +# Best compute core grid per shape from flux2 dual-grid sweep (2026-05-27). +# See flux2_dual_grid_mm_sweep.md "Best grid per shape". +_FLUX2_MATMUL_CORE_GRIDS: dict[tuple[int, int, int], tuple[int, int]] = { + (1024, 128, 768): ttnn.CoreCoord(11, 10), + (512, 15360, 768): ttnn.CoreCoord(12, 9), + (1024, 6144, 768): ttnn.CoreCoord(11, 10), + (512, 6144, 768): ttnn.CoreCoord(12, 9), + (1024, 6144, 4608): ttnn.CoreCoord(12, 9), + (512, 6144, 4608): ttnn.CoreCoord(12, 9), + (1024, 6144, 2304): ttnn.CoreCoord(12, 9), + (512, 6144, 2304): ttnn.CoreCoord(12, 9), +} + + +# adapted from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/models/attention_processor.py +class Attention(Module): + # SDPA chunk sizes keyed by (is_blackhole, sp_factor, tp_factor). Resolution priority for + # non-ring: caller overrides > class map > constructor args > class default. + sdpa_chunk_size_map: dict[tuple, tuple[int, int]] = {} + default_sdpa_chunk_size: tuple[int, int] = (128, 512) + + ring_sdpa_chunk_size_map: dict[tuple, tuple[int, int]] = { + # -1 is the default resolution. + (False, 2, 4): {-1: (256, 256)}, + (False, 8, 4): {-1: (256, 256)}, + (True, 2, 2): {-1: (128, 512)}, + (True, 4, 8): { + -1: (256, 512), # default + 4096: (128, 256), # 1024×1024 — 18.0% util (ring_sdpa_sweep.md) + 4608: (128, 256), # 1024×1024 — 18.0% util (ring_sdpa_sweep.md) + 4096 * 4: (256, 512), # 2048×2048 — 55.4% util (ring_sdpa_sweep_2048.md) + (4096 + 128) * 4: (256, 512), # 2048×2048 — 55.4% util (ring_sdpa_sweep_2048.md) + 4096 * 4: (256, 512), # 2048×2048 — 55.4% util (ring_sdpa_sweep_2048.md) + 4096 * 16: (192, 512), # 4096×4096 — 67.4% util (ring_sdpa_sweep_4096.md) + }, + (True, 8, 4): { + -1: (256, 512), # default + 16384: (320, 384), # 2048×2048 — 43.3% util (ring_sdpa_sweep_8x4_2048.md) + 65536: (256, 512), # 4096×4096 — 65.7% util (ring_sdpa_sweep_8x4_4096.md) + }, + } + default_ring_sdpa_chunk_size: tuple[int, int] = {-1: (256, 256)} + + def __init__( + self, + *, + query_dim: int, + head_dim: int, + heads: int, + out_dim: int, + added_kv_proj_dim: int, + context_pre_only: bool = False, + pre_only: bool = False, + use_spatial_weights_for_prompt: bool = False, + context_head_scaling: bool = False, + proj_bias: bool = True, + eps: float, + mesh_device: ttnn.MeshDevice, + ccl_manager: CCLManager | None, + parallel_config: DiTParallelConfig, + padding_config: PaddingConfig | None, + k_chunk_size: int | None = None, + q_chunk_size: int | None = None, + sdpa_chunk_size_overrides: dict | None = None, + per_head_norm: bool = False, + is_fsdp: bool = False, + shard_prompt: bool = False, + ) -> None: + super().__init__() + + self.head_dim = head_dim + self.pre_only = pre_only + self.mesh_device = mesh_device + self.ccl_manager = ccl_manager + self.parallel_config = parallel_config + self.shard_prompt = shard_prompt + self.padding_config = padding_config + self.use_spatial_weights_for_prompt = use_spatial_weights_for_prompt + # True = per-head RMS (head_dim divisor, stats stay local); False = global RMS over + # padded_inner_dim (AG of stats). + self.per_head_norm = per_head_norm + + self.padded_heads = padding_config.target_heads if padding_config is not None else heads + self.n_local_heads = self.padded_heads // self.parallel_config.tensor_parallel.factor + + tp_axis = parallel_config.tensor_parallel.mesh_axis + padded_inner_dim = head_dim * self.padded_heads + + # FSDP: shard weights on sequence parallel axis to reduce memory + fsdp_mesh_axis = parallel_config.sequence_parallel.mesh_axis if is_fsdp else None + + common_args = dict(mesh_device=mesh_device, ccl_manager=ccl_manager, fsdp_mesh_axis=fsdp_mesh_axis) + + full_grid = self.mesh_device.compute_with_storage_grid_size() + + # Reserve last row for CCL. + self.sdpa_worker_grid = (full_grid.x, full_grid.y - 1) + + chunk_lookup = {**self.sdpa_chunk_size_map, **(sdpa_chunk_size_overrides or {})} + resolved_q_chunk, resolved_k_chunk = chunk_lookup.get( + ( + is_blackhole(), + parallel_config.sequence_parallel.factor, + parallel_config.tensor_parallel.factor, + ), + ( + q_chunk_size if q_chunk_size is not None else self.default_sdpa_chunk_size[0], + k_chunk_size if k_chunk_size is not None else self.default_sdpa_chunk_size[1], + ), + ) + + self.sdpa_program_config = ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=self.sdpa_worker_grid, + q_chunk_size=resolved_q_chunk, + k_chunk_size=resolved_k_chunk, + exp_approx_mode=False, # NOTE: False is more correct + ) + + self.ring_sdpa_worker_grid = None # (full_grid.x, full_grid.y - 5) + self.ring_sdpa_program_config = {} + self.sdpa_compute_kernel_config = ttnn.init_device_compute_kernel_config( + self.mesh_device.arch(), + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, # NOTE: Set to True if there's a correctness issue + ) + self.mm_compute_kernel_config = ttnn.init_device_compute_kernel_config( + self.mesh_device.arch(), + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=True, + fp32_dest_acc_en=True, + packer_l1_acc=True, + ) + + # chunks=3 splits QKV inside minimal_matmul_split (faster than matmul + ttnn.chunk). + self.to_qkv = ColParallelLinear( + query_dim, 3 * padded_inner_dim, bias=proj_bias, mesh_axis=tp_axis, chunks=3, **common_args + ) + + # Fuses head-split + RMSNorm + RoPE. Weight is tiled per-head → [padded_inner_dim] in + # _prepare_torch_state. + self.norm_q = DistributedRMSNorm( + embedding_dim=padded_inner_dim, + norm_eps=eps, + mesh_axis=tp_axis, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + ) + self.norm_k = DistributedRMSNorm( + embedding_dim=padded_inner_dim, + norm_eps=eps, + mesh_axis=tp_axis, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + ) + + self.to_out = ( + ColParallelLinear(padded_inner_dim, out_dim, bias=proj_bias, mesh_axis=tp_axis, **common_args) + if not self.pre_only + else None + ) + + if use_spatial_weights_for_prompt: + self.add_qkv_proj = UnregisteredModule(self.to_qkv) + self.norm_added_q = UnregisteredModule(self.norm_q) + self.norm_added_k = UnregisteredModule(self.norm_k) + self.to_add_out = UnregisteredModule(self.to_out) if self.to_out is not None else None + elif added_kv_proj_dim > 0: + self.add_qkv_proj = ColParallelLinear( + added_kv_proj_dim, 3 * padded_inner_dim, bias=proj_bias, mesh_axis=tp_axis, chunks=3, **common_args + ) + + self.norm_added_q = DistributedRMSNorm( + embedding_dim=padded_inner_dim, + norm_eps=eps, + mesh_axis=tp_axis, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + ) + self.norm_added_k = DistributedRMSNorm( + embedding_dim=padded_inner_dim, + norm_eps=eps, + mesh_axis=tp_axis, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + ) + + self.to_add_out = ( + ColParallelLinear(padded_inner_dim, out_dim, bias=proj_bias, mesh_axis=tp_axis, **common_args) + if not context_pre_only + else None + ) + else: + self.add_qkv_proj = None + self.norm_added_q = None + self.norm_added_k = None + self.to_add_out = None + + self.context_head_factors = ( + Parameter(total_shape=[self.padded_heads, 1, 1], device=mesh_device, mesh_axes=[tp_axis, None, None]) + if context_head_scaling and self.add_qkv_proj is not None + else None + ) + + assert not (use_spatial_weights_for_prompt and context_head_scaling), ( + "context_head_scaling is not supported with use_spatial_weights_for_prompt=True " + "(concatenated path keeps q/k/v combined; per-head scaling on add_q is never applied)" + ) + + # Empty joint input for ring SDPA when there is no prompt stream. + self.dummy_joint_input = ttnn.zeros( + [1, self.n_local_heads, 0, self.head_dim], + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + dtype=ttnn.bfloat16, + ) + + self.trans_mat = ttnn.from_torch( + get_rot_transformation_mat(), + device=mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + ) + + @staticmethod + def get_core_grid(M, full_grid) -> ttnn.CoreCoord: + if M <= 1152: + return ttnn.CoreCoord(full_grid.x, full_grid.y - 2) + else: + return ttnn.CoreCoord(full_grid.x, full_grid.y - 1) + + def get_ring_sdpa_program_config(self, per_device_seq_len: int) -> ttnn.SDPAProgramConfig: + if per_device_seq_len not in self.ring_sdpa_program_config: + device_parallel_key = ( + is_blackhole(), + self.parallel_config.sequence_parallel.factor, + self.parallel_config.tensor_parallel.factor, + ) + ring_chunk_size = self.ring_sdpa_chunk_size_map.get(device_parallel_key, self.default_ring_sdpa_chunk_size) + if per_device_seq_len in ring_chunk_size: + ring_chunk_size = ring_chunk_size[per_device_seq_len] + else: + logger.warning( + f"No ring SDPA chunk size found for resolution {per_device_seq_len}, using default {ring_chunk_size[-1]}" + ) + ring_chunk_size = ring_chunk_size[-1] + + self.ring_sdpa_program_config[per_device_seq_len] = ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=self.ring_sdpa_worker_grid, + q_chunk_size=ring_chunk_size[0], + k_chunk_size=ring_chunk_size[1], + exp_approx_mode=False, # NOTE: False is more correct + ) + return self.ring_sdpa_program_config[per_device_seq_len] + + def get_ring_sdpa_core_grid(self, resolved_seq_len: int) -> ttnn.CoreCoord: + full_grid = self.mesh_device.compute_with_storage_grid_size() + if resolved_seq_len == 4096: + return (full_grid.x, full_grid.y - 5) + elif resolved_seq_len == 4608: + return (full_grid.x, full_grid.y - 5) + elif resolved_seq_len == 4096 * 4: + return (full_grid.x, full_grid.y - 1) + elif resolved_seq_len == (4096 + 128) * 4: + return (full_grid.x, full_grid.y - 1) + elif resolved_seq_len == 4096 * 16: + return (full_grid.x, full_grid.y - 1) + else: + return (full_grid.x, full_grid.y - 1) + + def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None: + weight, bias = self._reshape_and_merge_qkv( + pop_substate(state, "to_q"), + pop_substate(state, "to_k"), + pop_substate(state, "to_v"), + ) + if weight is not None: + state["to_qkv.weight"] = weight + if bias is not None: + state["to_qkv.bias"] = bias + + weight, bias = self._reshape_and_merge_qkv( + pop_substate(state, "add_q_proj"), + pop_substate(state, "add_k_proj"), + pop_substate(state, "add_v_proj"), + ) + if weight is not None: + state["add_qkv_proj.weight"] = weight + if bias is not None: + state["add_qkv_proj.bias"] = bias + + if "to_out.0.weight" in state: + state["to_out.weight"] = state.pop("to_out.0.weight") + if "to_out.0.bias" in state: + state["to_out.bias"] = state.pop("to_out.0.bias") + + if self.padding_config is not None: + if "to_out.weight" in state: + weight = state["to_out.weight"].T + weight = pad_weight_tensor(weight, self.padding_config, pad_input_dim=True) + state["to_out.weight"] = weight.T + if "to_add_out.weight" in state: + weight = state["to_add_out.weight"].T + weight = pad_weight_tensor(weight, self.padding_config, pad_input_dim=True) + state["to_add_out.weight"] = weight.T + + if "context_head_factors" in state: + factors = state["context_head_factors"] + if self.padding_config is not None: + pad = (0, self.padding_config.head_padding) + factors = torch.nn.functional.pad(factors, pad) + state["context_head_factors"] = factors.reshape([-1, 1, 1]) + + # HF stores per-head weight [head_dim]; DistributedRMSNorm expects [padded_inner_dim]. + for key in ["norm_q.weight", "norm_k.weight", "norm_added_q.weight", "norm_added_k.weight"]: + if key in state and state[key].shape[0] == self.head_dim: + state[key] = state[key].repeat(self.padded_heads) + + def _reshape_and_merge_qkv( + self, + q_state: dict[str, torch.Tensor], + k_state: dict[str, torch.Tensor], + v_state: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + # Rearrange QKV projections such column-fracturing shards the heads + def _merge_tensors(q, k, v): + n_dev = self.parallel_config.tensor_parallel.factor + q, k, v = q.T, k.T, v.T + # Pad QKV weights and biases to match the padded heads + if self.padding_config is not None: + q = pad_weight_tensor(q, self.padding_config, pad_output_dim=True) + k = pad_weight_tensor(k, self.padding_config, pad_output_dim=True) + v = pad_weight_tensor(v, self.padding_config, pad_output_dim=True) + q = q.reshape(q.shape[0], n_dev, self.n_local_heads, self.head_dim) + k = k.reshape(k.shape[0], n_dev, self.n_local_heads, self.head_dim) + v = v.reshape(v.shape[0], n_dev, self.n_local_heads, self.head_dim) + qkv = torch.cat([q, k, v], dim=2) + qkv = qkv.reshape(qkv.shape[0], 3 * self.padded_heads * self.head_dim) + qkv = qkv.T + return qkv + + if "weight" in q_state and "weight" in k_state and "weight" in v_state: + weight = _merge_tensors(q_state["weight"], k_state["weight"], v_state["weight"]) + else: + weight = None + + if "bias" in q_state and "bias" in k_state and "bias" in v_state: + bias = _merge_tensors( + q_state["bias"].unsqueeze(-1), k_state["bias"].unsqueeze(-1), v_state["bias"].unsqueeze(-1) + ) + bias = bias.squeeze(-1) + else: + bias = None + + return weight, bias + + def _to_out_fused_addcmul( + self, + x: ttnn.Tensor, + addcmul_residual: ttnn.Tensor | None, + addcmul_gate: ttnn.Tensor | None, + to_out_module: ColParallelLinear, + parallel_config, + tp_axis: int | None = None, + ) -> ttnn.Tensor: + """Fused projection + addcmul: output = residual + proj(x) * gate. + + Ring (parallel_config, tp>1): all_gather_minimal_matmul_async fuses AG+matmul+addcmul. + Linear (parallel_config=None): gathers x on tp_axis if K-dim mismatches, then + dit_minimal_matmul_addcmul_fused. + """ + if ( + to_out_module.fsdp_mesh_axis is not None + and to_out_module.mesh_device.shape[to_out_module.fsdp_mesh_axis] > 1 + ): + unsqueezed_weight = ttnn.unsqueeze_to_4D(to_out_module.weight.data) + weight = self.ccl_manager.all_gather_persistent_buffer( + unsqueezed_weight, dim=2, mesh_axis=to_out_module.fsdp_mesh_axis + ) + weight = ttnn.reshape(weight, (weight.shape[-2], weight.shape[-1])) + else: + weight = to_out_module.weight.data + + bias = to_out_module.bias.data if to_out_module.bias is not None else None + + if parallel_config is not None and parallel_config.tensor_parallel.factor > 1: + M = x.padded_shape[-2] + K = weight.padded_shape[-2] + N = weight.padded_shape[-1] + full_grid = self.mesh_device.compute_with_storage_grid_size() + core_grid = self.get_core_grid(M, full_grid) + # core_grid = _FLUX2_MATMUL_CORE_GRIDS.get((M, K, N)) or ttnn.CoreCoord(full_grid.x, full_grid.y - 1) + # print(f"Using core grid {core_grid} for matmul config (M, K, N) = ({M}, {K}, {N})", flush=True) + matmul_config = get_matmul_config( + M, K, N, core_grid, use_heuristic=True, num_k_shards=self.parallel_config.tensor_parallel.factor + ) + + ag_persistent_buffer = self.ccl_manager.get_ag_ping_pong_buffer( + x.shape, -1, parallel_config.tensor_parallel.mesh_axis, dtype=x.get_dtype() + ) + ag_global_semaphores = self.ccl_manager.get_ag_ping_pong_semaphore( + parallel_config.tensor_parallel.mesh_axis + ) + return ttnn.experimental.all_gather_minimal_matmul_async( + input_tensor=x, + weight_tensor=weight, + bias_tensor=bias, + config=matmul_config, + compute_kernel_config=self.mm_compute_kernel_config, + persistent_output_buffer=ag_persistent_buffer, + multi_device_global_semaphore=ag_global_semaphores, + num_links=self.ccl_manager.num_links, + topology=self.ccl_manager.topology, + cluster_axis=parallel_config.tensor_parallel.mesh_axis, + barrier_semaphore=None, + force_transpose=True, + num_workers_per_link=full_grid.x // self.ccl_manager.num_links, + num_buffers_per_channel=48 if not is_blackhole() else 24, + scalar=1.0, + addcmul_input_tensor1=addcmul_residual, + addcmul_input_tensor2=addcmul_gate, + )[0] + else: + if tp_axis is not None and x.padded_shape[-1] != weight.padded_shape[-2]: + x = self.ccl_manager.all_gather_persistent_buffer(x, dim=-1, mesh_axis=tp_axis, use_hyperparams=True) + M = x.padded_shape[-2] + K = x.padded_shape[-1] + N = weight.padded_shape[-1] + # core_grid = _FLUX2_MATMUL_CORE_GRIDS.get((M, K, N)) or self.mesh_device.compute_with_storage_grid_size() + # core_grid = self.mesh_device.compute_with_storage_grid_size() + core_grid = get_matmul_core_grid(self.mesh_device) + matmul_config = get_matmul_config( + M, K, N, core_grid, use_heuristic=True, num_k_shards=self.parallel_config.tensor_parallel.factor + ) + return ttnn.experimental.dit_minimal_matmul_addcmul_fused( + x, + weight, + 1.0, + addcmul_residual, + addcmul_gate, + bias_tensor=bias, + config=matmul_config, + compute_kernel_config=self.mm_compute_kernel_config, + ) + + def _project_out( + self, + input: ttnn.Tensor, + proj_op: ColParallelLinear | None, + addcmul_residual: ttnn.Tensor | None, + addcmul_gate: ttnn.Tensor | None, + is_ring: bool, + tp_axis: int, + ) -> ttnn.Tensor: + """Project with optional fused addcmul. Returns input unchanged if proj_op is None.""" + if input is None or proj_op is None: + return input + return self._to_out_fused_addcmul( + input, + addcmul_residual, + addcmul_gate, + proj_op, + self.parallel_config if is_ring else None, + None if is_ring else tp_axis, + ) + + def forward( + self, + *, + sequence_1: ttnn.Tensor, + sequence_1_length: int, + sequence_2: ttnn.Tensor | None = None, + sequence_2_length: int = 0, + sequence_1_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + sequence_2_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + addcmul_spatial_residual: ttnn.Tensor | None = None, + addcmul_spatial_gate: ttnn.Tensor | None = None, + addcmul_prompt_residual: ttnn.Tensor | None = None, + addcmul_prompt_gate: ttnn.Tensor | None = None, + ) -> tuple[ttnn.Tensor, ttnn.Tensor | None]: + """Forward with optional fused to_out + addcmul. + + Ring callers pass spatial/prompt fractured on TP (fused AG+matmul inside); + Linear callers must all-gather first. addcmul_*_residual+gate fuse the + out projection's residual+gate when both are provided. + """ + + assert len(sequence_1.shape) == 3 + if sequence_2 is not None: + assert len(sequence_2.shape) == 3 + for t in sequence_1_rope or (): + assert len(t.shape) == 4, f"rope tensors must be pre-shaped to (1, 1, seq, head_dim), got {t.shape}" + for t in sequence_2_rope or (): + assert len(t.shape) == 4, f"rope tensors must be pre-shaped to (1, 1, seq, head_dim), got {t.shape}" + + sp_axis = self.parallel_config.sequence_parallel.mesh_axis + + shard_prompt = self.shard_prompt and self.parallel_config.sequence_parallel.factor > 1 + + def _split_heads(x: ttnn.Tensor) -> ttnn.Tensor: + # V only — Q/K head-split is fused inside DistributedRMSNorm. + out, _, _ = ttnn.experimental.nlp_create_qkv_heads( + ttnn.unsqueeze(x, 0), num_heads=self.n_local_heads, num_kv_heads=0, transpose_k_heads=False + ) + return out + + # Sequence 1 processing Typically spatial or spatial+prompt. + q_flat_c, k_flat_c, v_flat_c = self.to_qkv( + sequence_1, + compute_kernel_config=self.mm_compute_kernel_config, + parallel_config=self.parallel_config, + core_grid=self.get_core_grid(sequence_1.shape[-2], self.mesh_device.compute_with_storage_grid_size()), + use_heuristic_mmcfg=True, + ) + + # Combined RoPE: concat spatial and prompt tensors along the seq dim (dim=2). + sequence_1_cos, sequence_1_sin = sequence_1_rope + trans_mat = self.trans_mat if sequence_1_rope is not None else None + + # Single norm_q/norm_k/v on the full combined [spatial+prompt] shard. + q = self.norm_q( + ttnn.unsqueeze(q_flat_c, 0), + num_heads_per_device=self.n_local_heads, + rope_cos=sequence_1_cos, + rope_sin=sequence_1_sin, + trans_mat=trans_mat, + per_head_norm=self.per_head_norm, + ) + + k = self.norm_k( + ttnn.unsqueeze(k_flat_c, 0), + num_heads_per_device=self.n_local_heads, + rope_cos=sequence_1_cos, + rope_sin=sequence_1_sin, + trans_mat=trans_mat, + per_head_norm=self.per_head_norm, + ) + + v = _split_heads(v_flat_c) # [1, H, M_s+M_p, D] + + # No joint stream — dummies handled in the post-processing block below. + add_q = add_k = add_v = self.dummy_joint_input + persistent_joint_k = persistent_joint_v = None + + if self.add_qkv_proj is not None and sequence_2 is not None: + # Rope tensors arrive pre-shaped as (1, 1, seq, head_dim); no per-call reshape needed. + prompt_cos = sequence_2_rope[0] if sequence_2_rope is not None else None + prompt_sin = sequence_2_rope[1] if sequence_2_rope is not None else None + prompt_trans_mat = self.trans_mat if sequence_2_rope is not None else None + + add_q_flat, add_k_flat, add_v_flat = self.add_qkv_proj( + sequence_2, + compute_kernel_config=self.mm_compute_kernel_config, + parallel_config=self.parallel_config, + core_grid=self.get_core_grid(sequence_2.shape[-2], self.mesh_device.compute_with_storage_grid_size()), + use_heuristic_mmcfg=True, + ) + add_q = self.norm_added_q( + ttnn.unsqueeze(add_q_flat, 0), + num_heads_per_device=self.n_local_heads, + rope_cos=prompt_cos, + rope_sin=prompt_sin, + trans_mat=prompt_trans_mat, + per_head_norm=self.per_head_norm, + ) + add_k = self.norm_added_k( + ttnn.unsqueeze(add_k_flat, 0), + num_heads_per_device=self.n_local_heads, + rope_cos=prompt_cos, + rope_sin=prompt_sin, + trans_mat=prompt_trans_mat, + per_head_norm=self.per_head_norm, + ) + add_v = _split_heads(add_v_flat) + + if self.context_head_factors is not None: + add_q = add_q * self.context_head_factors.data + + if shard_prompt: + persistent_joint_k = self.ccl_manager.get_ag_ping_pong_buffer(add_k.shape, 2, sp_axis) + persistent_joint_v = self.ccl_manager.get_ag_ping_pong_buffer(add_v.shape, 2, sp_axis) + + if self.ring_sdpa_worker_grid is None: + self.ring_sdpa_worker_grid = self.get_ring_sdpa_core_grid(sequence_1_length) + + if self.parallel_config.sequence_parallel.factor > 1: + spatial, sequence_2, _lse = ttnn.transformer.ring_joint_scaled_dot_product_attention( + q, + k, + v, + add_q, + add_k, + add_v, + persistent_output_buffer_k=self.ccl_manager.get_ag_ping_pong_buffer( + k.shape, 2, self.parallel_config.sequence_parallel.mesh_axis + ), + persistent_output_buffer_v=self.ccl_manager.get_ag_ping_pong_buffer( + v.shape, 2, self.parallel_config.sequence_parallel.mesh_axis + ), + joint_strategy="rear", + logical_n=sequence_1_length, + logical_l=sequence_2_length, + program_config=self.get_ring_sdpa_program_config(sequence_1_length), + compute_kernel_config=self.sdpa_compute_kernel_config, + dim=2, + multi_device_global_semaphore=self.ccl_manager.get_ag_ping_pong_semaphore( + self.parallel_config.sequence_parallel.mesh_axis + ), + num_links=self.ccl_manager.num_links, + cluster_axis=self.parallel_config.sequence_parallel.mesh_axis, + mesh_device=self.mesh_device, + topology=self.ccl_manager.topology, + subdevice_id=self.ccl_manager.ccl_sub_device_id, + ccl_core_grid_offset=(0, self.ring_sdpa_worker_grid[1]), + persistent_output_buffer_joint_k=persistent_joint_k, + persistent_output_buffer_joint_v=persistent_joint_v, + ) + else: + assert ( + sequence_1_length == spatial.shape[1] + ), "spatial sequence must not be padded without sequence parallelism" + + spatial, sequence_2 = ttnn.transformer.joint_scaled_dot_product_attention( + q, + k, + v, + add_q, + add_k, + add_v, + joint_strategy="rear", + program_config=self.sdpa_program_config, + compute_kernel_config=self.sdpa_compute_kernel_config, + ) + + spatial = ttnn.transformer.concatenate_heads(spatial) + if self.to_out is not None: + spatial = self.to_out( + x=spatial, + addcmul_a=addcmul_spatial_residual, + addcmul_b=addcmul_spatial_gate, + parallel_config=self.parallel_config, + core_grid=self.get_core_grid(spatial.shape[-2], self.mesh_device.compute_with_storage_grid_size()), + ) + + if sequence_2 is not None and sequence_2.shape[2] > 0: + sequence_2 = ttnn.transformer.concatenate_heads(sequence_2) + if self.to_add_out is not None: + sequence_2 = self.to_add_out( + x=sequence_2, + addcmul_a=addcmul_prompt_residual, + addcmul_b=addcmul_prompt_gate, + parallel_config=self.parallel_config, + core_grid=self.get_core_grid( + sequence_2.shape[-2], self.mesh_device.compute_with_storage_grid_size() + ), + ) + + return spatial, sequence_2 + + @classmethod + def spatial_sequence_padding_length(cls, *, length: int, sp_factor: int, k_chunk_size: int = 512) -> int: + if sp_factor == 1: + return 0 + + divisor = k_chunk_size * sp_factor + return -length % divisor + + @classmethod + def pad_spatial_sequence(cls, x: torch.Tensor, /, *, sp_factor: int, k_chunk_size: int = 512) -> torch.Tensor: + padding_len = cls.spatial_sequence_padding_length( + length=x.shape[-2], sp_factor=sp_factor, k_chunk_size=k_chunk_size + ) + return torch.nn.functional.pad(x, (0, 0, 0, padding_len)) diff --git a/code/models/tt_dit/parallel/config.py b/code/models/tt_dit/parallel/config.py new file mode 100644 index 0000000000000000000000000000000000000000..980c585a66c842dbecf89c1fe64e63571a3fada9 --- /dev/null +++ b/code/models/tt_dit/parallel/config.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass +from typing import NamedTuple + +import ttnn + + +class ParallelFactor(NamedTuple): + factor: int + mesh_axis: int + + +class DiTParallelConfig(NamedTuple): + cfg_parallel: ParallelFactor + tensor_parallel: ParallelFactor + sequence_parallel: ParallelFactor + + @classmethod + def from_tuples(cls, *, cfg: tuple[int, int], sp: tuple[int, int], tp: tuple[int, int]) -> DiTParallelConfig: + return cls( + cfg_parallel=ParallelFactor(*cfg), + sequence_parallel=ParallelFactor(*sp), + tensor_parallel=ParallelFactor(*tp), + ) + + +class DiTGParallelConfigNoCFG(NamedTuple): + tensor_parallel: ParallelFactor + sequence_parallel: ParallelFactor + + +class EncoderParallelConfig(NamedTuple): + tensor_parallel: ParallelFactor + sequence_parallel: ParallelFactor | None = None + + @classmethod + def from_tuple(cls, tp: tuple[int, int]) -> EncoderParallelConfig: + return cls(tensor_parallel=ParallelFactor(*tp)) + + +class VAEParallelConfig(NamedTuple): + tensor_parallel: ParallelFactor + + @classmethod + def from_tuple(cls, tp: tuple[int, int]) -> VAEParallelConfig: + return cls(tensor_parallel=ParallelFactor(*tp)) + + +class VaeHWParallelConfig(NamedTuple): + height_parallel: ParallelFactor + width_parallel: ParallelFactor + + @classmethod + def from_tuples(cls, *, height: tuple[int, int], width: tuple[int, int]) -> VaeHWParallelConfig: + return cls( + height_parallel=ParallelFactor(*height), + width_parallel=ParallelFactor(*width), + ) + + +class AudioTParallelConfig(NamedTuple): + axis0: ParallelFactor + axis1: ParallelFactor + + @property + def factor(self) -> int: + return self.axis0.factor * self.axis1.factor + + +class AudioTCParallelConfig(NamedTuple): + time_parallel: ParallelFactor + channel_parallel: ParallelFactor + + @property + def factor(self) -> int: + return self.time_parallel.factor + + @property + def mesh_axis(self) -> int: + return self.time_parallel.mesh_axis + + +class MochiVAEParallelConfig(NamedTuple): + time_parallel: ParallelFactor + h_parallel: ParallelFactor + w_parallel: ParallelFactor + + @classmethod + def from_tuples( + cls, + *, + time: tuple[int, int], + h: tuple[int, int], + w: tuple[int, int], + ) -> MochiVAEParallelConfig: + return cls( + time_parallel=ParallelFactor(*time), + h_parallel=ParallelFactor(*h), + w_parallel=ParallelFactor(*w), + ) + + +@dataclass(frozen=True) +class Flux2VaeParallelConfig: + """Parallel config for the Flux2 VAE decoder. + + Any two of the three axes may be set; the third must be None (2-D mesh constraint). + All non-None entries must use distinct ``mesh_axis`` values. + """ + + tp_parallel: ParallelFactor | None = None + h_parallel: ParallelFactor | None = None + w_parallel: ParallelFactor | None = None + + def __post_init__(self) -> None: + active = [p for p in (self.tp_parallel, self.h_parallel, self.w_parallel) if p is not None] + if len(active) > 2: + raise ValueError("At most 2 of tp_parallel/h_parallel/w_parallel may be set on a 2-D mesh") + axes = [p.mesh_axis for p in active] + if len(axes) != len(set(axes)): + raise ValueError(f"Parallel axes must be distinct; got mesh_axes={axes}") + + def _asdict(self) -> dict[str, ParallelFactor | None]: + return {"tp_parallel": self.tp_parallel, "h_parallel": self.h_parallel, "w_parallel": self.w_parallel} + + @classmethod + def from_axes( + cls, + mesh_device: ttnn.MeshDevice, + *, + tp_axis: int | None = None, + h_axis: int | None = None, + w_axis: int | None = None, + ) -> Flux2VaeParallelConfig: + def _pf(axis: int | None) -> ParallelFactor | None: + return ParallelFactor(factor=int(mesh_device.shape[axis]), mesh_axis=axis) if axis is not None else None + + return cls(tp_parallel=_pf(tp_axis), h_parallel=_pf(h_axis), w_parallel=_pf(w_axis)) + + +class OldParallelConfig(NamedTuple): + mesh_shape: tuple[int, int] + factor: int + mesh_axis: int + + +def vae_all_gather( + ccl_manager, + x: ttnn.Tensor, + cluster_axis: int = 1, + dim: int = 3, + reshape: bool = True, + use_barrier: bool = True, +) -> ttnn.Tensor: + if x.device().shape[cluster_axis] == 1: + return x + + global_semaphores = ccl_manager.get_ag_ping_pong_semaphore(cluster_axis) + + if reshape: + # reshape to b,1,h*w,c. This was tested to be faster. Need to verify overhead. TODO: Cleanup + b, h, w, c = x.shape + if h != 1: # Check if its already in desired shape. E.g group norm already reshaped to 1,1,h*w,c + x = x.reshape(b, 1, h * w, c) + + if x.layout != ttnn.TILE_LAYOUT: + x = ttnn.to_layout(x, ttnn.TILE_LAYOUT) # All gather requires tile layout + + # NOTE: We can't use ping-pong persistent buffers because we run out of memory. + # Single-buffered persistent buffers is a potential correctness issue, so we can't do that. + # barrier_semaphore is required for correctness on repeated all_gathers (prevents + # cross-dispatch races where fast devices start a new all_gather while slow devices + # are still processing the previous one). + # However, barrier_semaphore causes hangs in some pipelines (e.g. SD3.5 large). + # Those callers pass use_barrier=False and rely on synchronize_device instead. + if use_barrier: + barrier_semaphore = ccl_manager.get_barrier_semaphore(cluster_axis) + else: + barrier_semaphore = None + ttnn.synchronize_device(x.device()) + + x_g = ttnn.experimental.all_gather_async( + input_tensor=x, + dim=dim, + persistent_output_buffer=None, + barrier_semaphore=barrier_semaphore, + multi_device_global_semaphore=global_semaphores, + topology=ttnn.Topology.Linear, + cluster_axis=cluster_axis, + num_links=ccl_manager.num_links, + num_workers_per_link=4, + chunks_per_sync=80, + num_buffers_per_channel=4, + ) + + if reshape: + # reshape back to original expected shape + if h != 1: + x_g = x_g.reshape(b, h, w, -1) + return x_g + + +def vae_neighbor_pad( + ccl_manager, + x: ttnn.Tensor, + cluster_axis: int = 1, + dim: int = 0, + padding_left: int = 0, + padding_right: int = 0, + padding_mode: str = "replicate", +) -> ttnn.Tensor: + neighbor_semaphore = ccl_manager.get_np_ping_pong_semaphore(cluster_axis) + barrier_semaphore = ccl_manager.get_barrier_semaphore(cluster_axis) + + x_pad = ttnn.experimental.neighbor_pad_async( + x, + [dim], + [padding_left], + [padding_right], + padding_mode, + [cluster_axis], + [neighbor_semaphore], + [barrier_semaphore], + num_links=[ccl_manager.num_links], + topology=ttnn.Topology.Linear, + ) + + return x_pad + + +def vae_slice_reshard( + ccl_manager, x: ttnn.Tensor, cluster_axis: int = 1, dim: int = 0, output_shape: int = 88, output_offset: int = 0 +) -> ttnn.Tensor: + global_semaphore = ccl_manager.get_sr_ping_pong_semaphore(cluster_axis) + barrier_semaphore = ccl_manager.get_barrier_semaphore(cluster_axis) + + x_sr = ttnn.experimental.slice_reshard_async( + x, + dim=dim, + output_dim_shape=output_shape, + output_dim_offset=output_offset, + cluster_axis=cluster_axis, + final_semaphore=global_semaphore, + barrier_semaphore=barrier_semaphore, + num_links=ccl_manager.num_links, + topology=ttnn.Topology.Linear, + ) + + return x_sr diff --git a/code/models/tt_dit/tests/models/ltx/ltx_mesh_params.py b/code/models/tt_dit/tests/models/ltx/ltx_mesh_params.py new file mode 100644 index 0000000000000000000000000000000000000000..bc3ce5b529783cd85350f3bdadc6796148ad05ba --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/ltx_mesh_params.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +import os + +import pytest + +import ttnn +from models.tt_dit.utils.test import ( + line_params_req_exact_devices, + ring_params_8k_req_exact_devices, + ring_params_req_exact_devices, +) + +_line = line_params_req_exact_devices +_ring = ring_params_req_exact_devices + +_1x1sp0tp1nl1_line_is_fsdp0 = pytest.param( + (1, 1), 0, 1, 1, {}, ttnn.Topology.Linear, False, id="1x1sp0tp1nl1_line_is_fsdp0" +) +_2x2sp0tp1nl2_line_is_fsdp1 = pytest.param( + (2, 2), 0, 1, 2, _line, ttnn.Topology.Linear, True, id="2x2sp0tp1nl2_line_is_fsdp1" +) +_2x4sp0tp1nl1_line_is_fsdp1 = pytest.param( + (2, 4), 0, 1, 1, _line, ttnn.Topology.Linear, True, id="2x4sp0tp1nl1_line_is_fsdp1" +) +_2x4sp1tp0nl1_line_is_fsdp1 = pytest.param( + (2, 4), 1, 0, 1, _line, ttnn.Topology.Linear, True, id="2x4sp1tp0nl1_line_is_fsdp1" +) +# is_fsdp=False mirrors production BH 2x4 (dynamic_load path, not per-layer FSDP gathers) +_2x4sp1tp0nl2_line_is_fsdp0 = pytest.param( + (2, 4), 1, 0, 2, _line, ttnn.Topology.Linear, False, id="2x4sp1tp0nl2_line_is_fsdp0" +) + +# WH has 4 links. +_4x8sp1tp0nl4_ring_is_fsdp1 = pytest.param( + (4, 8), 1, 0, 4, _ring, ttnn.Topology.Ring, True, id="4x8sp1tp0nl4_ring_is_fsdp1" +) +_4x8sp1tp0nl2_line_is_fsdp0 = pytest.param( + (4, 8), 1, 0, 2, _line, ttnn.Topology.Linear, False, id="4x8sp1tp0nl2_line_is_fsdp0" +) +_4x8sp1tp0nl2_ring_is_fsdp0 = pytest.param( + (4, 8), 1, 0, 2, _ring, ttnn.Topology.Ring, False, id="4x8sp1tp0nl2_ring_is_fsdp0" +) +_4x32sp1tp0nl2_ring_is_fsdp0 = pytest.param( + (4, 32), 1, 0, 2, _ring, ttnn.Topology.Ring, False, id="4x32sp1tp0nl2_ring_is_fsdp0" +) + +LTX_PIPELINE_MESH_PARAMS = [ + _2x2sp0tp1nl2_line_is_fsdp1, + _2x4sp0tp1nl1_line_is_fsdp1, + _2x4sp1tp0nl2_line_is_fsdp0, + _4x8sp1tp0nl4_ring_is_fsdp1, + _4x8sp1tp0nl2_line_is_fsdp0, + _4x8sp1tp0nl2_ring_is_fsdp0, + _4x32sp1tp0nl2_ring_is_fsdp0, +] + + +def _with_dynamic_load(param, dynamic_load): + """Append dynamic_node arg to a 7-field mesh param, yielding the pipeline signature + (mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load). + Some tests use 7-field mesh_device params, some need an 8th dynamic_load. This function + allows us concatenate them together""" + v = param.values # (mesh_shape, sp, tp, num_links, device_params, topology, is_fsdp) + assert len(v) == 7, f"expected a 7-field mesh param, got {len(v)}: {param.id}" + return pytest.param(*v, dynamic_load, marks=param.marks, id=param.id) + + +# dynamic_load=True only for the 2x4 (BH-like) configs — mirrors production, where 2x4 pages +# weights in/out (dynamic_load path) to avoid init-time DRAM OOM rather than per-layer FSDP gathers. +_PIPELINE_DL_TRUE = {"2x4sp0tp1nl1_line_is_fsdp1", "2x4sp1tp0nl2_line_is_fsdp0"} + +LTX_PIPELINE_MESH_PARAMS_DL = [_with_dynamic_load(p, p.id in _PIPELINE_DL_TRUE) for p in LTX_PIPELINE_MESH_PARAMS] + + +def _override_base_device_params(base, device_params, *, id=None): + """Swap the device_params (element 4) of a shared building block, keeping its geometry. + Optional id override for variants that need a distinct test id (e.g. the i2v prefix).""" + v = base.values # (mesh_shape, sp, tp, num_links, device_params, topology, is_fsdp) + return pytest.param(*v[:4], device_params, *v[5:], marks=base.marks, id=id or base.id) + + +# One-stage (Pro) AV needs the L1_SMALL pool for the audio-vocoder conv1d taps. The base pipeline +# params leave l1_small_size=0, which OOMs the vocoder on the 4x8 galaxy ("L1_SMALL ... bank size is +# 0 B" in bank_manager). Reserve 32 KB L1_SMALL on the BH 4x8 configs (untraced -> no trace region), +# matching the distilled AV configs. Other geometries (deselected by the CI -k) pass through unchanged. +_LTX_ONE_STAGE_L1SMALL_IDS = {"4x8sp1tp0nl2_ring_is_fsdp0", "4x8sp1tp0nl2_line_is_fsdp0"} +LTX_ONE_STAGE_MESH_PARAMS_DL = [ + _override_base_device_params(p, {**p.values[4], "l1_small_size": 32768}, id=p.id) + if p.id in _LTX_ONE_STAGE_L1SMALL_IDS + else p + for p in LTX_PIPELINE_MESH_PARAMS_DL +] + +# Two-stages (Pro) drives the same audio vocoder, so it needs the same pool for the same reason. +LTX_TWO_STAGES_MESH_PARAMS_DL = LTX_ONE_STAGE_MESH_PARAMS_DL + + +# --------------------------------------------------------------------------- +# Distilled AV pipeline mesh params. Same geometry as the pipeline configs, but the audio +# decode chain needs bigger device pools on a few configs, and the WH 4x8 ring config runs +# dynamic_load. Reuse the shared building blocks for geometry (single-sourced); override +# device_params only where audio decode diverges, then append the distilled dynamic_load. +# --------------------------------------------------------------------------- +# Override device_params for the audio decode chain. Each dict augments the base line/ring +# params only with the pool(s) that config actually needs: +# l1_small_size: the audio vocoder's native conv1d/conv2d taps (depthwise audio filters) run an +# UntilizeWithHalo gather whose sharding/config tensors allocate from the dedicated L1_SMALL +# pool. It defaults to 0, which OOMs the vocoder in decode; 32 KB matches the audio component +# tests. +# worker_l1_size: the WH 4x8 ring config's RingAttention otherwise hits a kernel code-size error; +# the larger worker L1 gives its command stream room. +# trace_region_size: under LTX_TRACED=1 both stage traces' command streams (stage-1 + the +# larger-sequence stage-2) live here; measured need is ~236 MB at 1080p (get_trace_buffers_size), +# so 500 MB leaves headroom. +_line_l1small = {**_line, "l1_small_size": 32768} +_ring_worker_l1 = {"worker_l1_size": 1344544, **_ring} +_line_trace = {**_line, "trace_region_size": 500_000_000, "l1_small_size": 32768} +# fabric_router_config (8 KB payload): the strided all-gather packs up to 4 bf16 tiles per fabric packet +_ring_trace = {**ring_params_8k_req_exact_devices, "trace_region_size": 500_000_000, "l1_small_size": 32768} + +LTX_DISTILLED_MESH_PARAMS_DL = [ + _with_dynamic_load(_2x2sp0tp1nl2_line_is_fsdp1, False), + _with_dynamic_load(_2x4sp0tp1nl1_line_is_fsdp1, True), + # BH on 2x4: L1_SMALL scratch for the vocoder conv taps. + _with_dynamic_load(_override_base_device_params(_2x4sp1tp0nl2_line_is_fsdp0, _line_l1small), True), + # WH (ring) on 4x8: bigger worker L1 for RingAttention. + _with_dynamic_load(_override_base_device_params(_4x8sp1tp0nl4_ring_is_fsdp1, _ring_worker_l1), True), + # BH (linear) on 4x8. + _with_dynamic_load(_4x8sp1tp0nl2_line_is_fsdp0, False), + # BH (ring) on 4x8: trace region + L1_SMALL for the traced decode. + _with_dynamic_load(_override_base_device_params(_4x8sp1tp0nl2_ring_is_fsdp0, _ring_trace), False), + _with_dynamic_load(_4x32sp1tp0nl2_ring_is_fsdp0, False), +] + +# I2V chained t2v->i2v; both configs are traced (trace region + L1_SMALL). 4x8 Galaxy (ring): +# full-res 1088x1920 latent shards unevenly on the 4x8 mesh (s1 cond latent 17x30, full 34x60), +# so the VAE encoder fold + even-shard padding must handle non-mesh-aligned dims here — the 2x4 +# loudbox shards evenly and never hits this, so it's kept as a distinct id. +LTX_DISTILLED_I2V_MESH_PARAMS_DL = [ + _with_dynamic_load(_override_base_device_params(_2x4sp1tp0nl2_line_is_fsdp0, _line_trace), True), + _with_dynamic_load( + _override_base_device_params(_4x8sp1tp0nl2_ring_is_fsdp0, _ring_trace, id="i2v_4x8sp1tp0nl2_ring_is_fsdp0"), + False, + ), +] + +# Two-stages (Pro) I2V smoke reuses the same geometry + device params as the distilled i2v chain +# (2x4 loudbox + 4x8 galaxy, traced + L1_SMALL for the audio vocoder), so the Pro-I2V leg can run on +# the single 4x8 galaxy in addition to the 2x4 loudbox. +LTX_TWO_STAGES_I2V_MESH_PARAMS_DL = LTX_DISTILLED_I2V_MESH_PARAMS_DL + +# Audio-decode-only profiling: both configs use the traced line params (trace region + L1_SMALL) +# so LTX_TRACED=1 can capture the vocoder. +LTX_DISTILLED_AUDIO_MESH_PARAMS_DL = [ + _with_dynamic_load(_override_base_device_params(_2x4sp1tp0nl2_line_is_fsdp0, _line_trace), True), + _with_dynamic_load(_override_base_device_params(_4x8sp1tp0nl2_line_is_fsdp0, _line_trace), False), +] + +# No 1x1 config: real-grid shapes require SP padding, and video self-attention only masks +# padded keys via ring SDPA's logical_n (sp>1). Production never runs sp=1. +LTX_TRANSFORMER_MESH_PARAMS = [ + _2x4sp0tp1nl1_line_is_fsdp1, + _2x4sp1tp0nl2_line_is_fsdp0, + _4x8sp1tp0nl4_ring_is_fsdp1, + _4x8sp1tp0nl2_ring_is_fsdp0, + _4x8sp1tp0nl2_line_is_fsdp0, +] + +LTX_ATTENTION_MESH_PARAMS = [ + _1x1sp0tp1nl1_line_is_fsdp0, + _2x4sp0tp1nl1_line_is_fsdp1, + _2x4sp1tp0nl1_line_is_fsdp1, + _4x8sp1tp0nl4_ring_is_fsdp1, + _4x8sp1tp0nl2_line_is_fsdp0, +] + +# Upsampler uses VaeHWParallelConfig (h_axis, w_axis) instead of DiT sp/tp axes. +_2x4h0w1nl1_line = pytest.param( + (2, 4), 0, 1, 1, {**_line, "trace_region_size": 23887872}, ttnn.Topology.Linear, id="2x4h0w1nl1_line" +) + +_4x8h0w1nl2_line = pytest.param( + (4, 8), 0, 1, 2, {**_line, "trace_region_size": 23887872}, ttnn.Topology.Linear, id="4x8h0w1nl2_line" +) + +LTX_UPSAMPLER_MESH_PARAMS = [ + _2x4h0w1nl1_line, + _4x8h0w1nl2_line, +] + +# VAE decoder mesh params — multi-device configs for full LTXVideoDecoder parity. +# Format: (mesh_shape, h_axis, w_axis, num_links, device_params, topology) +_LTX_SKIP_SMALL_MESH_FABRIC = os.environ.get("LTX_VAE_FORCE_SMALL_MESH_FABRIC", "0") != "1" + +_line_fabric = {"fabric_config": ttnn.FabricConfig.FABRIC_1D, "require_exact_physical_num_devices": True} + +_1x1h0w1nl1_line_vae = pytest.param( + (1, 1), + 0, + 1, + 1, + _line_fabric, + ttnn.Topology.Linear, + marks=pytest.mark.skipif( + _LTX_SKIP_SMALL_MESH_FABRIC, + reason="Known flaky fabric handshake on 1x1 mesh; set LTX_VAE_FORCE_SMALL_MESH_FABRIC=1 to force-run.", + ), + id="1x1h0w1nl1_line", +) + +_2x4h0w1nl1_line_vae = pytest.param((2, 4), 0, 1, 1, _line_fabric, ttnn.Topology.Linear, id="2x4h0w1nl1_line") + +_2x4h1w0nl1_line_vae = pytest.param((2, 4), 1, 0, 1, _line_fabric, ttnn.Topology.Linear, id="2x4h1w0nl1_line") + +_1x8h0w1nl1_line_vae = pytest.param((1, 8), 0, 1, 1, _line_fabric, ttnn.Topology.Linear, id="1x8h0w1nl1_line") + +_1x4h1w0nl1_line_vae = pytest.param( + (1, 4), + 1, + 0, + 1, + _line_fabric, + ttnn.Topology.Linear, + marks=pytest.mark.skipif( + _LTX_SKIP_SMALL_MESH_FABRIC, + reason="Known flaky fabric handshake on 1x4 mesh; set LTX_VAE_FORCE_SMALL_MESH_FABRIC=1 to force-run.", + ), + id="1x4h1w0nl1_line", +) + +_4x8h0w1nl2_line_vae = pytest.param((4, 8), 0, 1, 2, _line_fabric, ttnn.Topology.Linear, id="4x8h0w1nl2_line") + +LTX_VAE_DECODER_MESH_PARAMS = [ + _1x1h0w1nl1_line_vae, + _2x4h0w1nl1_line_vae, + _2x4h1w0nl1_line_vae, + _1x8h0w1nl1_line_vae, + _1x4h1w0nl1_line_vae, + _4x8h0w1nl2_line_vae, +] + +# Multi-device only (2K decode doesn't fit on a single device). +LTX_VAE_DECODER_MULTI_ONLY_MESH_PARAMS = [ + _2x4h0w1nl1_line_vae, + _2x4h1w0nl1_line_vae, + _4x8h0w1nl2_line_vae, +] + +# Encoder production mesh (mirrors the I2V pipeline layout). +LTX_VAE_ENCODER_PROD_MESH_PARAMS = [ + _4x8h0w1nl2_line_vae, +] + +# --------------------------------------------------------------------------- +# Audio mesh params — (mesh_device, mesh_shape, sp_axis, tp_axis, num_links, +# dynamic_load, device_params, topology, is_fsdp) +# Audio tests create a submesh from mesh_device using mesh_shape, hence the +# duplicate shape field. +# --------------------------------------------------------------------------- +_line_audio = {**_line, "l1_small_size": 32768} +_ring_audio = {**_ring, "l1_small_size": 32768} +_ring_trace_audio = {**_ring_audio, "trace_region_size": 300_000_000} + +LTX_AUDIO_MESH_PARAMS_FULL = [ + pytest.param( + (2, 2), (2, 2), 0, 1, 2, False, _line_audio, ttnn.Topology.Linear, True, id="2x2sp0tp1nl2_line_is_fsdp1" + ), + pytest.param( + (2, 4), (2, 4), 0, 1, 1, True, _line_audio, ttnn.Topology.Linear, True, id="2x4sp0tp1nl1_line_is_fsdp1" + ), + pytest.param( + (2, 4), (2, 4), 1, 0, 2, True, _line_audio, ttnn.Topology.Linear, False, id="2x4sp1tp0nl2_line_is_fsdp0" + ), + pytest.param( + (4, 8), (4, 8), 1, 0, 4, False, _ring_audio, ttnn.Topology.Ring, True, id="4x8sp1tp0nl4_ring_is_fsdp1" + ), + pytest.param( + (4, 8), (4, 8), 1, 0, 2, False, _line_audio, ttnn.Topology.Linear, False, id="4x8sp1tp0nl2_line_is_fsdp0" + ), + pytest.param( + (4, 8), (4, 8), 1, 0, 2, False, _ring_trace_audio, ttnn.Topology.Ring, False, id="4x8sp1tp0nl2_ring_is_fsdp0" + ), + pytest.param( + (4, 32), (4, 32), 1, 0, 2, False, _ring_audio, ttnn.Topology.Ring, False, id="4x32sp1tp0nl2_ring_is_fsdp0" + ), +] diff --git a/code/models/tt_dit/tests/models/ltx/test_attention_ltx.py b/code/models/tt_dit/tests/models/ltx/test_attention_ltx.py new file mode 100644 index 0000000000000000000000000000000000000000..5c0597ae96ee9169f3c6a9496ff4ad82468d62dd --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_attention_ltx.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""LTX-2 attention unit tests; mesh/fabric parametrization mirrors test_wan_attention.""" + +import pytest +import torch +from loguru import logger + +import ttnn +from models.tt_dit.models.transformers.ltx.attention_ltx import LTXAttention +from models.tt_dit.parallel.config import DiTParallelConfig, ParallelFactor +from models.tt_dit.parallel.manager import CCLManager +from models.tt_dit.utils.check import assert_quality +from models.tt_dit.utils.mochi import get_rot_transformation_mat +from models.tt_dit.utils.tensor import bf16_tensor, bf16_tensor_2dshard +from models.tt_dit.utils.test import skip_if_unsupported_num_links + +from .ltx_mesh_params import LTX_ATTENTION_MESH_PARAMS + + +def _diffusers_qk_to_split(t: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor: + """diffusers (interleaved-rotation) Q/K → Lightricks SPLIT convention the TT loader expects.""" + D = head_dim + D_half = D // 2 + inv = torch.empty(D, dtype=torch.long) + inv[:D_half] = torch.arange(0, D, 2) + inv[D_half:] = torch.arange(1, D, 2) + rest = t.shape[1:] + return t.reshape(num_heads, D, *rest).index_select(1, inv).reshape(num_heads * D, *rest) + + +def _convert_diffusers_attn_state(state: dict, num_heads: int, head_dim: int) -> dict: + """diffusers LTX2Attention state_dict → TT LTXAttention loader input (split Q/K convention).""" + out = dict(state) + for base in ("to_q", "to_k"): + for suffix in ("weight", "bias"): + k = f"{base}.{suffix}" + if k in out: + out[k] = _diffusers_qk_to_split(out[k], num_heads, head_dim) + for k in ("norm_q.weight", "norm_k.weight"): + if k in out: + out[k] = _diffusers_qk_to_split(out[k], num_heads, head_dim) + return out + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + LTX_ATTENTION_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_self_attention( + mesh_device: ttnn.MeshDevice, + sp_axis: int, + tp_axis: int, + num_links: int, + topology: ttnn.Topology, + is_fsdp: bool, +) -> None: + """ + Test LTX-2 self-attention: compare TT LTXAttention vs PyTorch Attention. + """ + from diffusers.models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoRotaryPosEmbed + + skip_if_unsupported_num_links(mesh_device, num_links) + + dim = 4096 + num_heads = 32 + head_dim = dim // num_heads + B = 1 + seq_len = 256 # Small for fast test + + # interleaved matches the (B, N, dim) cos/sin layout the device kernel expects + torch_model = LTX2Attention( + query_dim=dim, heads=num_heads, kv_heads=num_heads, dim_head=head_dim, norm_eps=1e-6, rope_type="interleaved" + ) + torch_model.eval() + torch_state = torch_model.state_dict() + + ccl_manager = CCLManager(mesh_device=mesh_device, num_links=num_links, topology=topology) + parallel_config = DiTParallelConfig( + cfg_parallel=ParallelFactor(factor=1, mesh_axis=0), + sequence_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[sp_axis], mesh_axis=sp_axis), + tensor_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[tp_axis], mesh_axis=tp_axis), + ) + + tt_model = LTXAttention( + dim=dim, + num_heads=num_heads, + eps=1e-6, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + is_self=True, + ) + tt_model.load_torch_state_dict(_convert_diffusers_attn_state(torch_state, num_heads, head_dim)) + + # seq_len=256 → N_local=32 for 4x8 (sp=8); exactly 1 tile, the minimum + torch.manual_seed(42) + x = torch.randn(B, seq_len, dim, dtype=torch.float32) + + # Build RoPE (small grid for test) with the diffusers reference, like test_rope_ltx. + F, H, W = 4, 8, 8 # 4*8*8=256 = seq_len + t_ids = torch.arange(F) + h_ids = torch.arange(H) + w_ids = torch.arange(W) + grid_t, grid_h, grid_w = torch.meshgrid(t_ids, h_ids, w_ids, indexing="ij") + indices_grid = torch.stack([grid_t.flatten(), grid_h.flatten(), grid_w.flatten()], dim=0).float().unsqueeze(0) + + rope = LTX2AudioVideoRotaryPosEmbed( + dim=dim, + base_num_frames=20, + base_height=2048, + base_width=2048, + theta=10000.0, + modality="video", + double_precision=False, + rope_type="interleaved", + num_attention_heads=num_heads, + ) + # Interleaved cos/sin: (B, seq_len, dim) + cos_freq, sin_freq = rope(indices_grid) + # Reshape for the device kernel: (B, seq_len, num_heads, head_dim) + cos_apply = cos_freq.reshape(B, seq_len, num_heads, head_dim) + sin_apply = sin_freq.reshape(B, seq_len, num_heads, head_dim) + + # diffusers applies interleaved RoPE to q/k on (B, T, dim) before the head split. + with torch.no_grad(): + torch_out = torch_model(x, query_rotary_emb=(cos_freq, sin_freq)) + + logger.info(f"PyTorch output shape: {torch_out.shape}") + + # Prepare TT tensors + # spatial: (1, B, N, D) fractured on SP/TP + spatial = x.unsqueeze(0) # (1, B, N, D) + tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3}) + + # rotary_embedding_llama expects per-head (B, H, N, head_dim) + rope_cos_bhnd = cos_apply.permute(0, 2, 1, 3) # (B, H, N, D) + rope_sin_bhnd = sin_apply.permute(0, 2, 1, 3) + tt_cos = bf16_tensor_2dshard(rope_cos_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}) + tt_sin = bf16_tensor_2dshard(rope_sin_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + # TT forward + tt_out = tt_model( + spatial_1BND=tt_spatial, + N=seq_len, + rope_cos=tt_cos, + rope_sin=tt_sin, + trans_mat=tt_trans_mat, + ) + + # Gather and compare + concat_dims = [None, None] + concat_dims[sp_axis] = 2 + concat_dims[tp_axis] = 3 + tt_out_torch = ttnn.to_torch( + tt_out, + mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)), + ).squeeze(0) + + # 8-way SP-ring all-gathers on 4x8 loosen the bound; <=8 devices hold the single-op bar. + pcc = 0.988 if mesh_device.get_num_devices() > 8 else 0.999 + assert_quality(torch_out, tt_out_torch, pcc=pcc) + logger.info("PASSED: LTX self-attention matches PyTorch reference") + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + LTX_ATTENTION_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_cross_attention( + mesh_device: ttnn.MeshDevice, + sp_axis: int, + tp_axis: int, + num_links: int, + topology: ttnn.Topology, + is_fsdp: bool, +) -> None: + """ + Test LTX-2 cross-attention: compare TT LTXAttention vs PyTorch Attention. + """ + from diffusers.models.transformers.transformer_ltx2 import LTX2Attention + + skip_if_unsupported_num_links(mesh_device, num_links) + + dim = 4096 + context_dim = 4096 + num_heads = 32 + head_dim = dim // num_heads + B = 1 + seq_len = 256 + prompt_len = 32 + + # diffusers reference cross-attention (cross_attention_dim set; no RoPE on text K/V) + torch_model = LTX2Attention( + query_dim=dim, + cross_attention_dim=context_dim, + heads=num_heads, + kv_heads=num_heads, + dim_head=head_dim, + norm_eps=1e-6, + ) + torch_model.eval() + torch_state = torch_model.state_dict() + + ccl_manager = CCLManager(mesh_device=mesh_device, num_links=num_links, topology=topology) + parallel_config = DiTParallelConfig( + cfg_parallel=ParallelFactor(factor=1, mesh_axis=0), + sequence_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[sp_axis], mesh_axis=sp_axis), + tensor_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[tp_axis], mesh_axis=tp_axis), + ) + + tt_model = LTXAttention( + dim=dim, + num_heads=num_heads, + eps=1e-6, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + is_self=False, + context_dim=context_dim, + ) + tt_model.load_torch_state_dict(torch_state) + + torch.manual_seed(42) + x = torch.randn(B, seq_len, dim, dtype=torch.float32) + context = torch.randn(B, prompt_len, context_dim, dtype=torch.float32) + + with torch.no_grad(): + torch_out = torch_model(x, encoder_hidden_states=context) + + logger.info(f"PyTorch cross-attn output shape: {torch_out.shape}") + + spatial = x.unsqueeze(0) + prompt = context.unsqueeze(0) + tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3}) + tt_prompt = bf16_tensor(prompt, device=mesh_device) + + tt_out = tt_model(spatial_1BND=tt_spatial, N=seq_len, prompt_1BLP=tt_prompt) + + concat_dims = [None, None] + concat_dims[sp_axis] = 2 + concat_dims[tp_axis] = 3 + tt_out_torch = ttnn.to_torch( + tt_out, + mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)), + ).squeeze(0) + + # 8-way SP-ring all-gathers on 4x8 loosen the bound; <=8 devices hold the single-op bar. + pcc = 0.988 if mesh_device.get_num_devices() > 8 else 0.999 + assert_quality(torch_out, tt_out_torch, pcc=pcc) + logger.info("PASSED: LTX cross-attention matches PyTorch reference") diff --git a/code/models/tt_dit/tests/models/ltx/test_audio_ltx.py b/code/models/tt_dit/tests/models/ltx/test_audio_ltx.py new file mode 100644 index 0000000000000000000000000000000000000000..efa66a87662b4f384277f2e852d5402d05e8a639 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_audio_ltx.py @@ -0,0 +1,1207 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 + +"""Audio tests for LTX-2: per-component stage checks + post-denoise decode e2e. + +Component stage tests (random weights vs diffusers LTX-2 torch reference): +- Stage A: ``audio_decoder`` (mel-VAE) +- Stage B: ``vocoder`` (24 kHz waveform) +- Stage C: ``vocoder_with_bwe`` (48 kHz BWE path) + +End-to-end decode (mel-VAE → vocoder → BWE, requires real checkpoints): +- on-device decode runs on representative meshes +- warm-path decode stays functional after cold compile/build +- on-device decode quality remains close to host reference (PSNR) + +Denoise-path audio regressions live in sibling files. +""" + +from __future__ import annotations + +import math +import os +import time +from glob import glob + +import einops +import pytest +import torch +from loguru import logger + +import ttnn +from models.tt_dit.models.audio_vae.audio_decoder_ltx import LTXAudioDecoderAdapter +from models.tt_dit.models.audio_vae.bwe_ltx import MelSTFT, VocoderWithBWE +from models.tt_dit.models.audio_vae.mel_decoder_ltx import AudioUpsample, MelDecoder, ResnetBlock +from models.tt_dit.models.audio_vae.vocoder_ltx import AMPBlock1, Vocoder +from models.tt_dit.parallel.config import AudioTCParallelConfig, DiTParallelConfig, ParallelFactor +from models.tt_dit.parallel.manager import CCLManager +from models.tt_dit.pipelines.ltx.pipeline_ltx import LTXPipeline +from models.tt_dit.utils.check import assert_quality +from models.tt_dit.utils.test import skip_if_unsupported_num_links +from models.tt_dit.utils.video import Audio + +from .ltx_mesh_params import LTX_AUDIO_MESH_PARAMS_FULL + +_WARM_ITERS = 3 + + +# --------------------------------------------------------------------------- +# Checkpoint discovery +# --------------------------------------------------------------------------- +# Default checkpoint filenames used by the e2e/perf decode tests. Any of these can +# be overridden at runtime with LTX_CHECKPOINT, which is honored regardless of its +# basename (see _resolve_checkpoint). The tuples below only drive *fallback* +# discovery when LTX_CHECKPOINT is unset. +_DEV_CHECKPOINT = "ltx-2.3-22b-dev.safetensors" +_DISTILLED_CHECKPOINT = "ltx-2.3-22b-distilled-1.1.safetensors" +_LOCAL_CHECKPOINT_DIR = "~/.cache/ltx-checkpoints" +# HF cache snapshot glob for the LTX-2.3 repo; {name} is the checkpoint filename. +_HF_SNAPSHOT_GLOB = "~/.cache/huggingface/hub/models--Lightricks--LTX-2.3/snapshots/*/{name}" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _resolve_checkpoint(candidates: list[str], expected_filenames: tuple[str, ...]) -> str | None: + # An explicit LTX_CHECKPOINT override wins unconditionally as long as the file + # exists — that's the whole point of the env var. The expected_filenames tuple + # only gates the fallback discovery below, never the override. + explicit = os.environ.get("LTX_CHECKPOINT") + if explicit and os.path.exists(explicit): + return explicit + for path in candidates: + expanded = os.path.expanduser(path) + if os.path.exists(expanded): + return expanded + # Reuse default HF cache location if present; avoid network dependency. + for name in expected_filenames: + hits = glob(os.path.expanduser(_HF_SNAPSHOT_GLOB.format(name=name))) + if hits: + return hits[0] + return None + + +def _build_pipeline( + mesh_device: ttnn.MeshDevice, + *, + sp_axis: int, + tp_axis: int, + checkpoint: str, + num_links: int, + topology: ttnn.Topology, +) -> tuple[LTXPipeline, DiTParallelConfig]: + mesh_shape = tuple(mesh_device.shape) + parallel_config = DiTParallelConfig( + cfg_parallel=ParallelFactor(factor=1, mesh_axis=0), + sequence_parallel=ParallelFactor(factor=mesh_shape[sp_axis], mesh_axis=sp_axis), + tensor_parallel=ParallelFactor(factor=mesh_shape[tp_axis], mesh_axis=tp_axis), + ) + ccl_manager = CCLManager(mesh_device, num_links=num_links, topology=topology) + pipeline = LTXPipeline( + mesh_device=mesh_device, + parallel_config=parallel_config, + ccl_manager=ccl_manager, + checkpoint_name=None, + ) + pipeline.checkpoint_name = checkpoint + pipeline._audio_adapter = LTXAudioDecoderAdapter( + pipeline.checkpoint_name, + mesh_device=pipeline.mesh_device, + vae_ccl_manager=pipeline.vae_ccl_manager, + dit_parallel_config=pipeline.parallel_config, + traced=pipeline._traced, + ) + return pipeline, parallel_config + + +def _psnr(ref: torch.Tensor, test: torch.Tensor) -> float: + ref = ref.float() + test = test.float() + mse = torch.mean((ref - test) ** 2).item() + if mse == 0.0: + return float("inf") + peak = ref.abs().max().item() + if peak == 0.0: + return float("inf") + return 20.0 * math.log10(peak) - 10.0 * math.log10(mse) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ( + "mesh_device", + "mesh_shape", + "sp_axis", + "tp_axis", + "num_links", + "dynamic_load", + "device_params", + "topology", + "is_fsdp", + ), + LTX_AUDIO_MESH_PARAMS_FULL, + indirect=["mesh_device", "device_params"], +) +def test_audio_decode_perf( + mesh_device: ttnn.MeshDevice, + mesh_shape: tuple[int, int], + sp_axis: int, + tp_axis: int, + num_links: int, + dynamic_load: bool, + topology: ttnn.Topology, + is_fsdp: bool, +): + ckpt = _resolve_checkpoint( + [f"{_LOCAL_CHECKPOINT_DIR}/{_DEV_CHECKPOINT}"], + (_DEV_CHECKPOINT,), + ) + if ckpt is None: + pytest.skip("checkpoint not found (set LTX_CHECKPOINT)") + + skip_if_unsupported_num_links(mesh_device, num_links) + _ = (dynamic_load, is_fsdp) # parity with fast_av config schema + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + torch.manual_seed(0) + pipeline, parallel_config = _build_pipeline( + mesh_device, + sp_axis=sp_axis, + tp_axis=tp_axis, + checkpoint=ckpt, + num_links=num_links, + topology=topology, + ) + mesh_shape = tuple(mesh_device.shape) + + audio_latent = torch.randn(1, 256, 128, dtype=torch.float32) + num_frames, fps = 145, 24.0 + + t0 = time.time() + dev_audio = pipeline.decode_audio(audio_latent, num_frames, fps=fps) + t_cold = time.time() - t0 + assert dev_audio is not None + + warm_times = [] + for _ in range(_WARM_ITERS): + t0 = time.time() + out = pipeline.decode_audio(audio_latent, num_frames, fps=fps) + warm_times.append(time.time() - t0) + assert out is not None + + logger.info( + f"Audio decode perf (mesh {mesh_shape}, T-shard={parallel_config.sequence_parallel.factor}): " + f"cold={t_cold:.2f}s warm_min={min(warm_times):.2f}s warm_mean={sum(warm_times)/len(warm_times):.2f}s " + f"wave={tuple(dev_audio.waveform.shape)} sr={dev_audio.sampling_rate}" + ) + + +@pytest.mark.parametrize( + ( + "mesh_device", + "mesh_shape", + "sp_axis", + "tp_axis", + "num_links", + "dynamic_load", + "device_params", + "topology", + "is_fsdp", + ), + LTX_AUDIO_MESH_PARAMS_FULL, + indirect=["mesh_device", "device_params"], +) +def test_audio_decode_e2e_psnr( + mesh_device: ttnn.MeshDevice, + mesh_shape: tuple[int, int], + sp_axis: int, + tp_axis: int, + num_links: int, + dynamic_load: bool, + topology: ttnn.Topology, + is_fsdp: bool, +): + ckpt = _resolve_checkpoint( + [f"{_LOCAL_CHECKPOINT_DIR}/{_DISTILLED_CHECKPOINT}"], + (_DISTILLED_CHECKPOINT,), + ) + if ckpt is None: + pytest.skip("distilled checkpoint not found (set LTX_CHECKPOINT)") + + skip_if_unsupported_num_links(mesh_device, num_links) + _ = (dynamic_load, is_fsdp) # parity with fast_av config schema + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + torch.manual_seed(0) + mesh_shape = tuple(mesh_device.shape) + pipeline, _ = _build_pipeline( + mesh_device, + sp_axis=sp_axis, + tp_axis=tp_axis, + checkpoint=ckpt, + num_links=num_links, + topology=topology, + ) + + audio_latent = torch.randn(1, 256, 128, dtype=torch.float32) + num_frames, fps = 145, 24.0 + + ref_audio = _decode_audio_reference(ckpt, audio_latent, num_frames, fps=fps) + assert ref_audio is not None, "host reference audio decode returned None" + + dev_audio = pipeline.decode_audio(audio_latent, num_frames, fps=fps) + assert dev_audio is not None, "on-device audio decode returned None" + assert dev_audio.sampling_rate == ref_audio.sampling_rate + + n = min(ref_audio.waveform.shape[-1], dev_audio.waveform.shape[-1]) + ref_w = ref_audio.waveform[..., :n].float() + dev_w = dev_audio.waveform[..., :n].float() + psnr = _psnr(ref_w, dev_w) + + logger.info( + f"Audio decode PSNR (mesh {mesh_shape}): {psnr:.2f} dB " + f"(ref {tuple(ref_audio.waveform.shape)} -> dev {tuple(dev_audio.waveform.shape)})" + ) + assert psnr >= 28.0, f"audio decode PSNR {psnr:.2f} dB < 28 dB" + + +# =========================================================================== +# Component stage tests (pipeline_ltx.py audio decode chain) +# +# Torch references use diffusers LTX-2 audio modules (same as transformer/rope tests). +# Stage/component IDs match ``pipeline_ltx.py`` audio decode chain: +# - Stage A: ``audio_decoder`` (mel-VAE) +# - Stage B: ``vocoder`` (24 kHz waveform) +# - Stage C: ``vocoder_with_bwe`` (48 kHz BWE path) +# =========================================================================== +_AUDIO_DECODER_CFG = dict( + ch=128, + out_ch=2, + ch_mult=(1, 2, 4), + num_res_blocks=2, + attn_resolutions=(), + mid_block_add_attention=False, + z_channels=8, + resolution=64, + mel_bins=64, + sample_rate=16000, + mel_hop_length=160, + is_causal=True, +) + +_MAIN_VOCODER_CFG = dict( + resblock_kernel_sizes=[3, 7, 11], + upsample_rates=[5, 2, 2, 2, 2, 2], + upsample_kernel_sizes=[11, 4, 4, 4, 4, 4], + resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + upsample_initial_channel=1536, + resblock="AMP1", + output_sampling_rate=24000, + activation="snakebeta", + use_tanh_at_final=False, + apply_final_activation=True, + use_bias_at_final=False, +) + +_BWE_VOCODER_CFG = dict( + resblock_kernel_sizes=[3, 7, 11], + upsample_rates=[6, 5, 2, 2, 2], + upsample_kernel_sizes=[12, 11, 4, 4, 4], + resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + upsample_initial_channel=512, + resblock="AMP1", + output_sampling_rate=48000, + activation="snakebeta", + use_tanh_at_final=False, + apply_final_activation=False, + use_bias_at_final=False, +) + +_MEL_STFT_CFG = dict( + filter_length=512, + hop_length=80, + win_length=512, + n_mel_channels=64, +) + +_INPUT_SR = 16000 +_OUTPUT_SR = 48000 +_HOP_LENGTH = 80 + + +# --------------------------------------------------------------------------- +# Component helpers +# --------------------------------------------------------------------------- +def _vocoder_mel(t_frames: int = 120) -> torch.Tensor: + """Stereo mel input. t_frames=120 with T-shard factor 4 leaves t_pad=8 (128-aligned), + so the tile-align pad fits inside the last shard and the sharded path's per-op tail + materialization is exercised — the boundary the click came from.""" + return torch.randn(1, 2, t_frames, 64, dtype=torch.float32) * 0.5 + + +def _assert_sharded_matches_unsharded(unsharded: torch.Tensor, sharded: torch.Tensor, *, name: str) -> None: + """The T-sharded vocoder must reproduce the unsharded device path up to the bf16 storage + floor. A t_pad boundary regression (the audible click) shows up as a localized tail spike + of order 1e-2 — far above the ~1e-4 floor — so a tight max|Δ| catches it where aggregate + PCC against the torch reference (with its own fp32-kernel floor) does not.""" + max_abs = (unsharded - sharded).abs().max().item() + logger.info(f"{name}: sharded vs unsharded device max|Δ| = {max_abs:.3e}") + assert max_abs < 5e-3, f"{name}: sharded vs unsharded device max|Δ| {max_abs:.3e} — t_pad boundary drift" + + +def _require_diffusers(): + pytest.importorskip("diffusers") + from diffusers.models.autoencoders.autoencoder_kl_ltx2_audio import AutoencoderKLLTX2Audio + from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder, LTX2VocoderWithBWE + + return AutoencoderKLLTX2Audio, LTX2Vocoder, LTX2VocoderWithBWE + + +def _diffusers_vocoder_state_to_tt(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Invert diffusers LTX-2.3 vocoder key names to the TT/Lightricks layout.""" + out: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + new_key = ( + key.replace("conv_in.", "conv_pre.") + .replace("conv_out.", "conv_post.") + .replace("act_out.", "act_post.") + .replace("upsamplers.", "ups.") + .replace("resnets.", "resblocks.") + ) + # TT derives downsample filter taps at init; these are not loadable Parameters. + if new_key.endswith("downsample.filter"): + continue + out[new_key] = value + return out + + +def _diffusers_vocoder_with_bwe_state_to_tt(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + for prefix in ("vocoder.", "bwe_generator."): + sub = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} + for k, v in _diffusers_vocoder_state_to_tt(sub).items(): + out[f"{prefix}{k}"] = v + for key, value in state_dict.items(): + if key.startswith("mel_stft."): + out[key] = value + return out + + +def _audio_decoder_state_from_diffusers( + vae, + *, + stats_std: torch.Tensor, + stats_mean: torch.Tensor, +) -> dict[str, torch.Tensor]: + state = {k[len("decoder.") :]: v for k, v in vae.state_dict().items() if k.startswith("decoder.")} + state["per_channel_statistics.std-of-means"] = stats_std + state["per_channel_statistics.mean-of-means"] = stats_mean + return state + + +def _decode_latent_ref( + vae, + latent: torch.Tensor, + *, + stats_std: torch.Tensor, + stats_mean: torch.Tensor, +) -> torch.Tensor: + """Match the Lightricks MelDecoder / TT denormalize + decoder forward.""" + b, channels, _t, mel_bins = latent.shape + patched = einops.rearrange(latent, "b c t f -> b t (c f)") + std = stats_std.to(dtype=patched.dtype, device=patched.device) + mean = stats_mean.to(dtype=patched.dtype, device=patched.device) + last_dim = patched.shape[-1] + if std.shape[0] != last_dim: + if std.shape[0] == channels and last_dim == channels * mel_bins: + std = std.repeat_interleave(mel_bins) + mean = mean.repeat_interleave(mel_bins) + else: + raise ValueError(f"per-channel stats {tuple(std.shape)} do not broadcast to last dim {last_dim}") + sample = einops.rearrange(patched * std + mean, "b t (c f) -> b c t f", c=channels, f=mel_bins) + return vae.decoder(sample) + + +def _compute_mel_ref(vocoder_with_bwe, audio_bct: torch.Tensor) -> torch.Tensor: + batch, n_channels, _ = audio_bct.shape + flat = audio_bct.reshape(batch * n_channels, -1).unsqueeze(1) + log_mel, _, _, _ = vocoder_with_bwe.mel_stft(flat) + return log_mel.reshape(batch, n_channels, log_mel.shape[1], log_mel.shape[2]) + + +def _tt_vocoder_cfg(cfg: dict) -> dict: + # Torch reference config carries keys that TT constructor does not accept. + return {k: v for k, v in cfg.items() if k != "output_sampling_rate"} + + +def _randomize_snake_alphas(module: torch.nn.Module) -> None: + with torch.no_grad(): + for submodule in module.modules(): + if hasattr(submodule, "alpha") and isinstance(submodule.alpha, torch.nn.Parameter): + submodule.alpha.data = torch.randn_like(submodule.alpha.data) * 0.1 + if hasattr(submodule, "beta") and isinstance(submodule.beta, torch.nn.Parameter): + submodule.beta.data = torch.randn_like(submodule.beta.data) * 0.1 + + +def _build_torch_stage_b(seed: int = 42): + _, LTX2Vocoder, _ = _require_diffusers() + + torch.manual_seed(seed) + model = LTX2Vocoder( + in_channels=128, + hidden_channels=_MAIN_VOCODER_CFG["upsample_initial_channel"], + out_channels=2, + upsample_kernel_sizes=_MAIN_VOCODER_CFG["upsample_kernel_sizes"], + upsample_factors=_MAIN_VOCODER_CFG["upsample_rates"], + resnet_kernel_sizes=_MAIN_VOCODER_CFG["resblock_kernel_sizes"], + resnet_dilations=_MAIN_VOCODER_CFG["resblock_dilation_sizes"], + act_fn=_MAIN_VOCODER_CFG["activation"], + antialias=True, + final_act_fn=None, + final_bias=_MAIN_VOCODER_CFG["use_bias_at_final"], + output_sampling_rate=_MAIN_VOCODER_CFG["output_sampling_rate"], + ).eval() + _randomize_snake_alphas(model) + return model + + +def _audio_parallel_config(mesh_shape: tuple[int, int]) -> AudioTCParallelConfig | ParallelFactor | None: + t_axis = 0 if mesh_shape[0] >= mesh_shape[1] else 1 + t_factor = mesh_shape[t_axis] + c_axis = 1 - t_axis + c_factor = mesh_shape[c_axis] + if t_factor > 1 and c_factor > 1: + return AudioTCParallelConfig( + time_parallel=ParallelFactor(factor=t_factor, mesh_axis=t_axis), + channel_parallel=ParallelFactor(factor=c_factor, mesh_axis=c_axis), + ) + if t_factor > 1: + return ParallelFactor(factor=t_factor, mesh_axis=t_axis) + return None + + +def _build_tt_stage_b(mesh_device: ttnn.MeshDevice, *, parallel_config, ccl_manager: CCLManager | None) -> Vocoder: + return Vocoder( + mesh_device=mesh_device, + dtype=ttnn.float32, + in_channels=128, + out_channels=2, + parallel_config=parallel_config, + ccl_manager=ccl_manager, + **_tt_vocoder_cfg(_MAIN_VOCODER_CFG), + ) + + +def _build_torch_stage_c(seed: int = 42): + _, LTX2Vocoder, LTX2VocoderWithBWE = _require_diffusers() + + torch.manual_seed(seed) + full = LTX2VocoderWithBWE( + in_channels=128, + hidden_channels=_MAIN_VOCODER_CFG["upsample_initial_channel"], + out_channels=2, + upsample_kernel_sizes=_MAIN_VOCODER_CFG["upsample_kernel_sizes"], + upsample_factors=_MAIN_VOCODER_CFG["upsample_rates"], + resnet_kernel_sizes=_MAIN_VOCODER_CFG["resblock_kernel_sizes"], + resnet_dilations=_MAIN_VOCODER_CFG["resblock_dilation_sizes"], + act_fn=_MAIN_VOCODER_CFG["activation"], + antialias=True, + final_act_fn=None, + final_bias=_MAIN_VOCODER_CFG["use_bias_at_final"], + bwe_in_channels=128, + bwe_hidden_channels=_BWE_VOCODER_CFG["upsample_initial_channel"], + bwe_out_channels=2, + bwe_upsample_kernel_sizes=_BWE_VOCODER_CFG["upsample_kernel_sizes"], + bwe_upsample_factors=_BWE_VOCODER_CFG["upsample_rates"], + bwe_resnet_kernel_sizes=_BWE_VOCODER_CFG["resblock_kernel_sizes"], + bwe_resnet_dilations=_BWE_VOCODER_CFG["resblock_dilation_sizes"], + bwe_act_fn=_BWE_VOCODER_CFG["activation"], + bwe_antialias=True, + bwe_final_act_fn=None, + bwe_final_bias=_BWE_VOCODER_CFG["use_bias_at_final"], + filter_length=_MEL_STFT_CFG["filter_length"], + hop_length=_MEL_STFT_CFG["hop_length"], + window_length=_MEL_STFT_CFG["win_length"], + num_mel_channels=_MEL_STFT_CFG["n_mel_channels"], + input_sampling_rate=_INPUT_SR, + output_sampling_rate=_OUTPUT_SR, + ) + n_freqs = _MEL_STFT_CFG["filter_length"] // 2 + 1 + with torch.no_grad(): + full.mel_stft.stft_fn.forward_basis.copy_(torch.randn(n_freqs * 2, 1, _MEL_STFT_CFG["win_length"]) * 0.05) + full.mel_stft.stft_fn.inverse_basis.copy_(torch.randn(n_freqs * 2, 1, _MEL_STFT_CFG["win_length"]) * 0.05) + full.mel_stft.mel_basis.copy_(torch.randn(_MEL_STFT_CFG["n_mel_channels"], n_freqs).abs() * 0.1) + _randomize_snake_alphas(full) + return full.eval() + + +def _build_torch_stage_c_real(checkpoint_name: str): + """diffusers ``LTX2VocoderWithBWE`` with the *real* checkpoint weights — the torch oracle + for the vocoder+BWE. The checkpoint uses the original BigVGAN key names; invert the + diffusers→TT rename (see ``_diffusers_vocoder_state_to_tt``) to load them.""" + from safetensors import safe_open + + full = _build_torch_stage_c(seed=0) # correct architecture; weights overwritten below + + def ckpt_to_diffusers(s: str) -> str: + return ( + s.replace("conv_pre.", "conv_in.") + .replace("conv_post.", "conv_out.") + .replace("act_post.", "act_out.") + .replace("ups.", "upsamplers.") + .replace("resblocks.", "resnets.") + .replace("downsample.lowpass.filter", "downsample.filter") + ) + + state: dict[str, torch.Tensor] = {} + with safe_open(checkpoint_name, framework="pt") as f: + for k in f.keys(): + if k.startswith("vocoder."): + state[ckpt_to_diffusers(k[len("vocoder.") :])] = f.get_tensor(k).float() + res = full.load_state_dict(state, strict=False) + assert ( + not res.missing_keys and not res.unexpected_keys + ), f"torch oracle weight load mismatch: {len(res.missing_keys)} missing, {len(res.unexpected_keys)} unexpected" + return full.eval() + + +def _build_torch_audio_decoder_real(checkpoint_name: str): + """diffusers ``AutoencoderKLLTX2Audio`` (mel-VAE decoder) with the *real* checkpoint weights — + the host oracle for the mel-VAE stage. Architecture dims are read from the checkpoint's + embedded config, the same source ``LTXAudioDecoderAdapter`` uses, so this tracks the + real audio VAE rather than the toy component-test config. + + Returns ``(vae, stats_std, stats_mean, z_channels)`` ready for ``_decode_latent_ref``. + """ + import json + + from safetensors import safe_open + + AutoencoderKLLTX2Audio, _, _ = _require_diffusers() + + state: dict[str, torch.Tensor] = {} + stats_std = stats_mean = None + with safe_open(checkpoint_name, framework="pt") as f: + config = json.loads(f.metadata()["config"]) + ad = config["audio_vae"]["model"]["params"] + ddconfig = ad["ddconfig"] + stft_cfg = config["audio_vae"].get("preprocessing", {}).get("stft", {}) + mel_cfg = config["audio_vae"].get("preprocessing", {}).get("mel", {}) + mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") + for k in f.keys(): + if k.startswith("audio_vae.decoder."): + state[k[len("audio_vae.") :]] = f.get_tensor(k).float() # -> "decoder.*" + elif k == "audio_vae.per_channel_statistics.std-of-means": + stats_std = f.get_tensor(k).float() + elif k == "audio_vae.per_channel_statistics.mean-of-means": + stats_mean = f.get_tensor(k).float() + + z_channels = ddconfig.get("z_channels", 8) + vae = AutoencoderKLLTX2Audio( + base_channels=ddconfig.get("ch", 128), + output_channels=ddconfig.get("out_ch", 2), + ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), + num_res_blocks=ddconfig.get("num_res_blocks", 2), + attn_resolutions=tuple(ddconfig.get("attn_resolutions", ())), + resolution=ddconfig.get("resolution", 256), + latent_channels=z_channels, + norm_type="pixel", + causality_axis="height", + mid_block_add_attention=ddconfig.get("mid_block_add_attention", False), + sample_rate=ad.get("sampling_rate", 16000), + mel_hop_length=stft_cfg.get("hop_length", 160), + is_causal=stft_cfg.get("causal", True), + mel_bins=mel_bins, + ).eval() + + # Decoder-only load: the checkpoint carries no encoder for this path, so encoder keys are + # legitimately missing — only decoder keys must be satisfied, with nothing left unexpected. + res = vae.load_state_dict(state, strict=False) + missing_decoder = [k for k in res.missing_keys if k.startswith("decoder.")] + assert not missing_decoder and not res.unexpected_keys, ( + f"audio decoder oracle weight load mismatch: {len(missing_decoder)} decoder keys missing, " + f"{len(res.unexpected_keys)} unexpected" + ) + + if stats_std is None or stats_mean is None: + # Checkpoint without denormalize stats → identity (matches the component-test default). + stats_std = torch.ones(z_channels * mel_bins) + stats_mean = torch.zeros(z_channels * mel_bins) + return vae, stats_std, stats_mean, z_channels + + +def _decode_audio_reference(checkpoint_name: str, audio_latent: torch.Tensor, num_frames: int, fps: float) -> Audio: + """Host (torch) reference for the full audio decode chain — mel-VAE → vocoder+BWE — using the + real checkpoint weights. Mirrors ``LTXPipeline.decode_audio`` on CPU so the e2e PSNR test can + compare the on-device decode against a torch oracle instead of skipping when the pipeline + revision doesn't expose its own reference decode. + + Returns an ``Audio(waveform=(C, T), sampling_rate=...)`` matching the device output layout. + """ + vae, stats_std, stats_mean, z = _build_torch_audio_decoder_real(checkpoint_name) + vocoder_with_bwe = _build_torch_stage_c_real(checkpoint_name) + + # Unpatchify exactly as decode_audio does: (1, audio_N, 128) -> (1, z, audio_N, 128 // z). + audio_N = audio_latent.shape[1] + audio_spatial = audio_latent.reshape(1, audio_N, z, audio_latent.shape[2] // z).permute(0, 2, 1, 3).float() + + with torch.no_grad(): + mel = _decode_latent_ref(vae, audio_spatial, stats_std=stats_std, stats_mean=stats_mean) + waveform = vocoder_with_bwe(mel).squeeze(0).float() + sampling_rate = _OUTPUT_SR + + # Trim to the clip duration, matching decode_audio. + target_samples = int((num_frames / fps) * sampling_rate) + if waveform.shape[-1] > target_samples: + waveform = waveform[..., :target_samples] + return Audio(waveform=waveform, sampling_rate=sampling_rate) + + +def _build_tt_stage_c( + mesh_device: ttnn.MeshDevice, *, parallel_config, ccl_manager: CCLManager | None +) -> VocoderWithBWE: + main_voc = Vocoder( + mesh_device=mesh_device, + dtype=ttnn.float32, + in_channels=128, + out_channels=2, + parallel_config=parallel_config, + ccl_manager=ccl_manager, + **_tt_vocoder_cfg(_MAIN_VOCODER_CFG), + ) + # Mirror the pipeline (LTXAudioDecoderAdapter): channel-TP the BWE generator only + # where the channel axis pays for the gather (factor > 2); single-axis on 2x4 (factor 2). + if isinstance(parallel_config, AudioTCParallelConfig): + bwe_pc = parallel_config if parallel_config.channel_parallel.factor > 2 else parallel_config.time_parallel + else: + bwe_pc = parallel_config + bwe_voc = Vocoder( + mesh_device=mesh_device, + dtype=ttnn.float32, + in_channels=128, + out_channels=2, + parallel_config=bwe_pc, + ccl_manager=ccl_manager, + **_tt_vocoder_cfg(_BWE_VOCODER_CFG), + ) + mel_stft = MelSTFT(mesh_device=mesh_device, dtype=ttnn.float32, **_MEL_STFT_CFG) + return VocoderWithBWE( + vocoder=main_voc, + bwe_generator=bwe_voc, + mel_stft=mel_stft, + input_sampling_rate=_INPUT_SR, + output_sampling_rate=_OUTPUT_SR, + hop_length=_HOP_LENGTH, + mesh_device=mesh_device, + dtype=ttnn.float32, + ) + + +# --------------------------------------------------------------------------- +# Component stage tests +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ( + "mesh_device", + "mesh_shape", + "sp_axis", + "tp_axis", + "num_links", + "dynamic_load", + "device_params", + "topology", + "is_fsdp", + ), + LTX_AUDIO_MESH_PARAMS_FULL, + indirect=["mesh_device", "device_params"], +) +def test_stage_a_audio_decoder( + mesh_device: ttnn.MeshDevice, + mesh_shape: tuple[int, int], + sp_axis: int, + tp_axis: int, + num_links: int, + dynamic_load: bool, + topology: ttnn.Topology, + is_fsdp: bool, +): + AutoencoderKLLTX2Audio, _, _ = _require_diffusers() + + skip_if_unsupported_num_links(mesh_device, num_links) + _ = (sp_axis, tp_axis, num_links, dynamic_load, topology, is_fsdp) # parity with fast_av config schema + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + torch.manual_seed(42) + ref_vae = AutoencoderKLLTX2Audio( + base_channels=_AUDIO_DECODER_CFG["ch"], + output_channels=_AUDIO_DECODER_CFG["out_ch"], + ch_mult=_AUDIO_DECODER_CFG["ch_mult"], + num_res_blocks=_AUDIO_DECODER_CFG["num_res_blocks"], + attn_resolutions=_AUDIO_DECODER_CFG["attn_resolutions"], + resolution=_AUDIO_DECODER_CFG["resolution"], + latent_channels=_AUDIO_DECODER_CFG["z_channels"], + norm_type="pixel", + causality_axis="height", + mid_block_add_attention=_AUDIO_DECODER_CFG["mid_block_add_attention"], + sample_rate=_AUDIO_DECODER_CFG["sample_rate"], + mel_hop_length=_AUDIO_DECODER_CFG["mel_hop_length"], + is_causal=_AUDIO_DECODER_CFG["is_causal"], + mel_bins=_AUDIO_DECODER_CFG["mel_bins"], + ).eval() + + z_times_f = _AUDIO_DECODER_CFG["z_channels"] * _AUDIO_DECODER_CFG["mel_bins"] + stats_std = torch.ones(z_times_f) + stats_mean = torch.zeros(z_times_f) + + tt_decoder = MelDecoder(mesh_device=mesh_device, **_AUDIO_DECODER_CFG) + tt_decoder.load_torch_state_dict( + _audio_decoder_state_from_diffusers(ref_vae, stats_std=stats_std, stats_mean=stats_mean) + ) + + latent = torch.randn(1, _AUDIO_DECODER_CFG["z_channels"], 64, _AUDIO_DECODER_CFG["mel_bins"], dtype=torch.float32) + with torch.no_grad(): + torch_out = _decode_latent_ref(ref_vae, latent, stats_std=stats_std, stats_mean=stats_mean) + tt_out = tt_decoder(latent) + + assert torch_out.shape == tt_out.shape + assert_quality(torch_out, tt_out, pcc=0.998, relative_rmse=0.05) + + +@pytest.mark.parametrize( + ( + "mesh_device", + "mesh_shape", + "sp_axis", + "tp_axis", + "num_links", + "dynamic_load", + "device_params", + "topology", + "is_fsdp", + ), + LTX_AUDIO_MESH_PARAMS_FULL, + indirect=["mesh_device", "device_params"], +) +def test_stage_b_vocoder( + mesh_device: ttnn.MeshDevice, + mesh_shape: tuple[int, int], + sp_axis: int, + tp_axis: int, + num_links: int, + dynamic_load: bool, + topology: ttnn.Topology, + is_fsdp: bool, +): + skip_if_unsupported_num_links(mesh_device, num_links) + _ = (sp_axis, tp_axis, dynamic_load, is_fsdp) # parity with fast_av config schema + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + pc = _audio_parallel_config(mesh_shape) + ccl = CCLManager(mesh_device, num_links=num_links, topology=ttnn.Topology.Linear) if pc is not None else None + + torch_voc = _build_torch_stage_b(seed=42) + tt_voc = _build_tt_stage_b(mesh_device, parallel_config=pc, ccl_manager=ccl) + tt_voc.load_torch_state_dict(_diffusers_vocoder_state_to_tt(torch_voc.state_dict())) + + mel = _vocoder_mel() + with torch.no_grad(): + ref_out = torch_voc(mel) + tt_out = tt_voc(mel) + + assert tt_out.shape == ref_out.shape + assert_quality(ref_out, tt_out, pcc=0.99) + + # When T-sharded, assert the sharded path reproduces the unsharded device path: the t_pad + # boundary click is localized to the tail and slips past aggregate PCC, but breaks this. + if pc is not None: + tt_voc_un = _build_tt_stage_b(mesh_device, parallel_config=None, ccl_manager=None) + tt_voc_un.load_torch_state_dict(_diffusers_vocoder_state_to_tt(torch_voc.state_dict())) + _assert_sharded_matches_unsharded(tt_voc_un(mel), tt_out, name="stage_b_vocoder") + + +@pytest.mark.parametrize( + ( + "mesh_device", + "mesh_shape", + "sp_axis", + "tp_axis", + "num_links", + "dynamic_load", + "device_params", + "topology", + "is_fsdp", + ), + LTX_AUDIO_MESH_PARAMS_FULL, + indirect=["mesh_device", "device_params"], +) +def test_stage_c_vocoder_with_bwe( + mesh_device: ttnn.MeshDevice, + mesh_shape: tuple[int, int], + sp_axis: int, + tp_axis: int, + num_links: int, + dynamic_load: bool, + topology: ttnn.Topology, + is_fsdp: bool, +): + skip_if_unsupported_num_links(mesh_device, num_links) + _ = (sp_axis, tp_axis, dynamic_load, is_fsdp) # parity with fast_av config schema + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + pc = _audio_parallel_config(mesh_shape) + ccl = CCLManager(mesh_device, num_links=num_links, topology=ttnn.Topology.Linear) if pc is not None else None + + torch_full = _build_torch_stage_c(seed=42) + tt_full = _build_tt_stage_c(mesh_device, parallel_config=pc, ccl_manager=ccl) + + incompatible = tt_full.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict())) + assert not incompatible.missing_keys, f"missing keys: {incompatible.missing_keys}" + + mel = _vocoder_mel() + + # Stage B inside Stage C. + with torch.no_grad(): + ref_main = torch_full.vocoder(mel) + tt_main = tt_full.vocoder(mel) + assert tt_main.shape == ref_main.shape + assert_quality(ref_main, tt_main, pcc=0.99) + if pc is not None: + tt_full_un = _build_tt_stage_c(mesh_device, parallel_config=None, ccl_manager=None) + tt_full_un.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict())) + _assert_sharded_matches_unsharded(tt_full_un.vocoder(mel), tt_main, name="stage_c_vocoder") + + # Log-mel path (apples-to-apples same waveform). + remainder = ref_main.shape[-1] % _HOP_LENGTH + if remainder: + pad_right = _HOP_LENGTH - remainder + ref_main = torch.nn.functional.pad(ref_main, (0, pad_right)) + with torch.no_grad(): + ref_mel = _compute_mel_ref(torch_full, ref_main) + tt_mel = tt_full._compute_mel_device(ref_main) + assert tt_mel.shape == ref_mel.shape + assert_quality(ref_mel, tt_mel, pcc=0.99) + + # BWE residual and skip-resample paths. + ref_mel_for_bwe = ref_mel.transpose(2, 3).contiguous() + with torch.no_grad(): + ref_residual = torch_full.bwe_generator(ref_mel_for_bwe) + ref_skip = torch_full.resampler(ref_main) + tt_residual = tt_full.bwe_generator(ref_mel_for_bwe) + tt_skip = tt_full._resample_device(ref_main) + assert tt_residual.shape == ref_residual.shape + assert tt_skip.shape == ref_skip.shape + assert_quality(ref_residual, tt_residual, pcc=0.98) + assert_quality(ref_skip, tt_skip, pcc=0.99) + + # Full Stage C forward: the muxed waveform is where the sharded tail click surfaced. + with torch.no_grad(): + ref_out = torch_full(mel) + tt_out = tt_full(mel) + assert tt_out.shape == ref_out.shape + if pc is not None: + _assert_sharded_matches_unsharded(tt_full_un(mel), tt_out, name="stage_c_full") + + +# Localized audio regression: aggregate PCC over a full clip masks a brief burst, so this +# drives the real ~6s length and gates each 0.5s window vs the unsharded eager baseline. +_STATIC_WIN_S = 0.5 + + +def _window_error_rows(ref: torch.Tensor, tt: torch.Tensor, *, sr: int, win_s: float): + """Per-window (start_s, max|Δ|, rmse/σ_ref) over the time axis of (B, C, T).""" + T = ref.shape[-1] + n = max(1, int(win_s * sr)) + rows = [] + for start in range(0, T, n): + r = ref[..., start : start + n] + t = tt[..., start : start + n] + max_abs = (r - t).abs().max().item() + rmse = ((r - t) ** 2).mean().sqrt().item() + sigma = ((r**2).mean().sqrt().item()) + 1e-9 + rows.append((start / sr, max_abs, rmse / sigma)) + return rows + + +@pytest.mark.parametrize( + "device_params", + [ + { + "l1_small_size": 32768, + "fabric_config": ttnn.FabricConfig.FABRIC_1D, + "trace_region_size": 300000000, + "require_exact_physical_num_devices": True, + } + ], + indirect=True, +) +@pytest.mark.parametrize("mesh_device", [(2, 4)], indirect=True) +def test_stage_c_audio_decode(mesh_device, device_params): + """End-to-end audio synthesis (VocoderWithBWE): channel-TP main vocoder + conv1d depthwise, + gated against the unsharded eager baseline per 0.5s window so a localized burst (the demo + static) can't hide behind an aggregate PCC. Drives the real ~6s clip length (601 mel + frames).""" + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(2, 4)) + pc = AudioTCParallelConfig( + time_parallel=ParallelFactor(factor=4, mesh_axis=1), + channel_parallel=ParallelFactor(factor=2, mesh_axis=0), + ) + ccl = CCLManager(mesh_device, num_links=1, topology=ttnn.Topology.Linear) + + torch_full = _build_torch_stage_c(seed=42) + tt_full = _build_tt_stage_c(mesh_device, parallel_config=pc, ccl_manager=ccl) + tt_full.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict())) + tt_un = _build_tt_stage_c(mesh_device, parallel_config=None, ccl_manager=None) + tt_un.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict())) + + torch.manual_seed(7) + _tf = int(os.environ.get("AUDIO_T_FRAMES", "601")) # 601 mel frames ≈ the 6s girl clip + mel = torch.randn(1, 2, _tf, 64, dtype=torch.float32) * 0.5 + tt = tt_full(mel) + out_un = tt_un(mel) + + sr = tt_full.output_sampling_rate + logger.info(f"[audio] output {tuple(tt.shape)} = {tt.shape[-1]/sr:.2f}s @ {sr}Hz") + su = _window_error_rows(out_un, tt, sr=sr, win_s=_STATIC_WIN_S) + for s, su_max, su_rmse in su: + logger.info(f"[audio] t={s:4.2f}s shard-vs-un max|Δ|={su_max:.3e} rmse/σ={su_rmse:.3e}") + worst_su_max = max(r[1] for r in su) + worst_su_rmse = max(r[2] for r in su) + print( + f"\nAUDIO_DECODE shard-vs-un worst_max|Δ|={worst_su_max:.3e} worst_rmse/σ={worst_su_rmse:.3e}", + flush=True, + ) + # Regression guard: channel-TP audio must match the unsharded eager TT baseline. Catches + # broken sharding (the actual risk); torch-vs-TT on the full muxed BWE output is inherently + # loose with random weights and is checked per-component in test_stage_c_vocoder_with_bwe. + assert worst_su_max < 5e-3, f"audio window shard-vs-un max|Δ| {worst_su_max:.3e} — divergence" + assert worst_su_rmse < 1e-2, f"audio window shard-vs-un rmse/σ {worst_su_rmse:.3e} — divergence" + + +# =========================================================================== +# Tracy per-op device-time profile of the on-device audio decode chain. +# +# Goal: emit every op in the audio decode (mel-VAE -> vocoder -> BWE) with its +# device duration, and have the summed device time land on the ~0.49s traced +# steady-state audio decode (within a small delta). +# +# Why eager (traced=False): each op dispatches as its own program, so (a) the +# on-device profiler can be flushed per-block to stay under the 125-scope/core +# buffer, and (b) every op gets a named DEVICE FW / KERNEL DURATION. The device +# kernel durations are identical to the traced replay — trace only removes host +# dispatch — so the summed device time reproduces the traced ~0.49s while the +# host wall-clock (printed) is far larger; that gap is pure host dispatch. +# +# Two entry points: +# * test_audio_decode_profile — the workload (run under `python -m tracy`) +# * test_audio_decode_perf_table — driver that runs the profiler on the above +# and prints the per-op table + device total +# +# Manual one-liner (equivalent to the driver): +# LTX_AUDIO_PROF=1 \ +# LTX_CHECKPOINT=~/.cache/ltx-checkpoints/ltx-2.3-22b-distilled-1.1.safetensors \ +# python -m tracy -p -r -v -o generated/profiler/ltx_audio_decode -m \ +# pytest 'models/tt_dit/tests/models/ltx/test_audio_ltx.py::test_audio_decode_profile' -s +# then sum "DEVICE FW DURATION [ns]" from +# generated/profiler/ltx_audio_decode/reports//ops_perf_results_.csv +# =========================================================================== + +# The ~0.49s traced steady-state audio decode (4x8 ring, distilled, 145f@24fps), +# expressed in ms; the summed device time should land within _PROF_TOL_MS of it. +_PROF_TARGET_MS = 490.0 +_PROF_TOL_MS = 150.0 + +# Compute-heavy leaf blocks; a ReadDeviceProfiler flush after each keeps every +# Tracy zone window well under the per-core 125-scope device buffer. +_PROF_FLUSH_TYPES = (ResnetBlock, AudioUpsample, AMPBlock1, MelSTFT) + + +def _walk_tt(module): + """Depth-first walk over a tt_dit Module tree (root included).""" + yield module + for _, child in module.named_children(): + yield from _walk_tt(child) + + +def _flush_forward_after(module, mesh_device): + """Wrap ``module.forward`` so the on-device profiler is drained after every call, + keeping each Tracy zone window under the per-core scope buffer (see gemma harness).""" + orig = module.forward + + def timed(*a, _orig=orig, **k): + r = _orig(*a, **k) + ttnn.ReadDeviceProfiler(mesh_device) + return r + + module.forward = timed + + +@pytest.mark.skipif( + os.environ.get("LTX_AUDIO_PROF") != "1", + reason="Tracy profiling workload — run via test_audio_decode_perf_table, or manually with " + "LTX_AUDIO_PROF=1 python -m tracy -p -r -o -m pytest '::test_audio_decode_profile' -s", +) +# Profiler-instrumented kernels compile fresh on the first profiling run and the eager decode +# runs op-by-op, so this comfortably exceeds the default 300s per-test cap. +@pytest.mark.timeout(3600) +@pytest.mark.parametrize( + "device_params", + [ + { + "l1_small_size": 32768, + "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING, + "trace_region_size": 300_000_000, + "require_exact_physical_num_devices": True, + } + ], + indirect=True, +) +@pytest.mark.parametrize("mesh_device", [(4, 8)], indirect=True) +def test_audio_decode_profile(mesh_device, device_params): + """Eager on-device audio decode wrapped in Tracy signposts, one profiled iteration. + + Mirrors the 4x8 ring distilled pipeline's audio-decode config (T-shard=8, replicated + channel; the pipeline default with LTX_AUDIO_CHANNEL_TP off). Emits ``signpost("start")`` + / ``signpost("stop")`` around a single decode so ``post_process_ops_log(has_signposts=True)`` + slices out warmup/compile; per-block ReadDeviceProfiler flushes keep the device buffer live. + Pure timing — no PCC (correctness is covered by the sibling e2e/component tests).""" + ckpt = _resolve_checkpoint( + [f"{_LOCAL_CHECKPOINT_DIR}/{_DISTILLED_CHECKPOINT}"], + (_DISTILLED_CHECKPOINT,), + ) + if ckpt is None: + pytest.skip("distilled checkpoint not found (set LTX_CHECKPOINT)") + + from tracy import signpost + + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(4, 8)) + + torch.manual_seed(0) + # traced=False (LTXPipeline default): decode_audio runs eager, one program per op. + pipeline, _ = _build_pipeline( + mesh_device, + sp_axis=1, + tp_axis=0, + checkpoint=ckpt, + num_links=2, + topology=ttnn.Topology.Ring, + ) + assert not pipeline._traced, "profile expects the eager decode path (per-op programs)" + + audio_latent = torch.randn(1, 256, 128, dtype=torch.float32) + num_frames, fps = 145, 24.0 + + # Warmup: compile kernels + one-time device weight prep, unprofiled. Two iters so the + # program cache and any lazy device state are fully warm before the measured pass. + for _ in range(2): + assert pipeline.decode_audio(audio_latent, num_frames, fps=fps) is not None + ttnn.synchronize_device(mesh_device) + ttnn.ReadDeviceProfiler(mesh_device) # drain warmup zones (they precede the "start" marker) + + # Install per-block flushes only for the measured pass. + for root in (pipeline.tt_mel_decoder, pipeline.tt_vocoder_with_bwe): + for mod in _walk_tt(root): + if isinstance(mod, _PROF_FLUSH_TYPES): + _flush_forward_after(mod, mesh_device) + + # Capture the device-op-id range around the single forward. GLOBAL CALL COUNT in the C++ + # device digest (cpp_device_perf_report.csv) == this counter, so [op0, op1) slices the digest + # to exactly this forward — no dependence on the fragile Tracy -r ops report or signpost rows. + from ttnn import _ttnn + + op0 = _ttnn.get_device_operation_id() + signpost("start") + t0 = time.perf_counter() + out = pipeline.decode_audio(audio_latent, num_frames, fps=fps) + ttnn.synchronize_device(mesh_device) + host_wall_ms = (time.perf_counter() - t0) * 1000.0 + ttnn.ReadDeviceProfiler(mesh_device) # flush the tail ops while still inside the region + signpost("stop") + op1 = _ttnn.get_device_operation_id() + + assert out is not None + logger.info(f"Audio decode profiled: host_wall={host_wall_ms:.1f}ms wave={tuple(out.waveform.shape)}") + print(f"\nAUDIO_DECODE_HOST_WALL_MS={host_wall_ms:.2f}", flush=True) + print(f"AUDIO_DECODE_OPID_RANGE={op0},{op1}", flush=True) + + +@pytest.mark.skipif( + os.environ.get("LTX_AUDIO_PROF") != "1", + reason="Tracy driver — run: LTX_AUDIO_PROF=1 LTX_CHECKPOINT= pytest " + "'::test_audio_decode_perf_table' -s", +) +@pytest.mark.timeout(3600) +def test_audio_decode_perf_table(): + """Profile ``test_audio_decode_profile`` and print the per-op device-time table. + + Spawns ``python -m tracy`` on the workload (opening the device itself; this driver does not + take the mesh_device fixture, matching test_ring_joint_sdpa_perf_table), then reads the + signpost-sliced ops log. Prints one row per op (device FW/kernel duration) plus the summed + device time, and asserts the total lands within ``_PROF_TOL_MS`` of the ~0.49s traced target. + """ + import pandas as pd + from tracy.process_model_log import post_process_ops_log, run_device_profiler + + subdir = "ltx_audio_decode" + test_id = "models/tt_dit/tests/models/ltx/test_audio_ltx.py::test_audio_decode_profile" + run_device_profiler( + f"-m 'pytest {test_id} -s'", + subdir, + check_test_return_code=True, + is_command_binary_exe=True, + ) + + df = post_process_ops_log(subdir, has_signposts=True) # DataFrame sliced between start/stop + fw_col, k_col = "DEVICE FW DURATION [ns]", "DEVICE KERNEL DURATION [ns]" + + def _num(series): + return pd.to_numeric(series[series != "-"], errors="coerce").fillna(0.0) + + fw_ns = _num(df[fw_col]) + k_ns = _num(df[k_col]) + total_fw_ms = fw_ns.sum() / 1e6 + total_k_ms = k_ns.sum() / 1e6 + + # Per-op table, ops in dispatch order. + print("\n" + "=" * 78, flush=True) + print(f"{'#':>3} {'OP CODE':<34}{'FW [µs]':>12}{'KERNEL [µs]':>14}", flush=True) + print("-" * 78, flush=True) + codes = df["OP CODE"].tolist() + fw_list = _num(df[fw_col]).tolist() + k_list = _num(df[k_col]).tolist() + for i, (code, fw, kd) in enumerate(zip(codes, fw_list, k_list)): + print(f"{i:>3} {str(code):<34}{fw / 1e3:>12.2f}{kd / 1e3:>14.2f}", flush=True) + print("-" * 78, flush=True) + + # Per-op-code aggregate (sorted by total FW time). + agg = ( + pd.DataFrame({"code": codes, "fw_ns": fw_list, "k_ns": k_list}) + .groupby("code") + .agg(count=("fw_ns", "size"), fw_us=("fw_ns", lambda s: s.sum() / 1e3), k_us=("k_ns", lambda s: s.sum() / 1e3)) + .sort_values("fw_us", ascending=False) + ) + print("\nBy op code (FW-descending):", flush=True) + print(f"{'OP CODE':<34}{'count':>7}{'FW [µs]':>12}{'KERNEL [µs]':>14}", flush=True) + for code, row in agg.iterrows(): + print(f"{str(code):<34}{int(row['count']):>7}{row['fw_us']:>12.1f}{row['k_us']:>14.1f}", flush=True) + print("=" * 78, flush=True) + + print( + f"\nAUDIO_DECODE_DEVICE_FW_MS={total_fw_ms:.2f} " + f"AUDIO_DECODE_DEVICE_KERNEL_MS={total_k_ms:.2f} " + f"(target ~{_PROF_TARGET_MS:.0f}ms, tol ±{_PROF_TOL_MS:.0f}ms, n_ops={len(df)})", + flush=True, + ) + logger.info(f"Audio decode device time: FW={total_fw_ms:.1f}ms kernel={total_k_ms:.1f}ms over {len(df)} ops") + assert len(df) > 0, "no ops captured between signposts — check flushing / signpost naming" + assert abs(total_fw_ms - _PROF_TARGET_MS) <= _PROF_TOL_MS, ( + f"summed device FW {total_fw_ms:.1f}ms is >{_PROF_TOL_MS:.0f}ms from the ~{_PROF_TARGET_MS:.0f}ms " + f"traced target — device-bound audio-decode time drifted" + ) diff --git a/code/models/tt_dit/tests/models/ltx/test_embeddings_ltx.py b/code/models/tt_dit/tests/models/ltx/test_embeddings_ltx.py new file mode 100644 index 0000000000000000000000000000000000000000..8c0cd747928443de4d5dfc65cc3f1844c8cd28e5 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_embeddings_ltx.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from loguru import logger + +import ttnn +from models.tt_dit.layers.embeddings import LTXAdaLayerNormSingle +from models.tt_dit.utils.check import assert_quality + + +@pytest.mark.parametrize( + "mesh_device", + [(1, 1)], + ids=["1x1"], + indirect=["mesh_device"], +) +@pytest.mark.parametrize("embedding_dim", [4096], ids=["dim4096"]) +def test_ltx_adaln_single(mesh_device: ttnn.MeshDevice, embedding_dim: int): + """ + Test LTXAdaLayerNormSingle: compare TT output vs diffusers LTX2AdaLayerNormSingle. + """ + from diffusers.models.transformers.transformer_ltx2 import LTX2AdaLayerNormSingle as TorchAdaLayerNormSingle + + B = 2 + + # Create PyTorch reference model + torch_model = TorchAdaLayerNormSingle( + embedding_dim=embedding_dim, num_mod_params=6, use_additional_conditions=False + ) + torch_model.eval() + torch_state = torch_model.state_dict() + logger.info(f"PyTorch AdaLayerNormSingle state keys: {list(torch_state.keys())}") + + # Create TT model and load weights + tt_model = LTXAdaLayerNormSingle( + embedding_dim=embedding_dim, + embedding_coefficient=6, + mesh_device=mesh_device, + dtype=ttnn.float32, + ) + tt_model.load_torch_state_dict(torch_state) + + # Create timestep input + torch.manual_seed(42) + timestep = torch.rand(B) * 1000 # random timesteps in [0, 1000) + + # PyTorch reference forward + with torch.no_grad(): + torch_modulation, torch_embedded = torch_model(timestep) + logger.info(f"PyTorch modulation shape: {torch_modulation.shape}, embedded: {torch_embedded.shape}") + + # TT forward + # Timesteps module expects (1, 1, B, 1) so B broadcasts with the 128-dim factor. + # Use float32 for timestep input — sinusoidal embeddings are precision-sensitive. + tt_timestep = ttnn.from_torch( + timestep.reshape(1, 1, B, 1), + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + dtype=ttnn.float32, + ) + tt_modulation, tt_embedded = tt_model(tt_timestep) + + # Convert back and compare — output is (1, 1, B, dim) in ttnn + tt_modulation_torch = ttnn.to_torch(tt_modulation).squeeze(0).squeeze(0) + tt_embedded_torch = ttnn.to_torch(tt_embedded).squeeze(0).squeeze(0) + + logger.info(f"TT modulation shape: {tt_modulation_torch.shape}, embedded: {tt_embedded_torch.shape}") + + assert_quality(torch_modulation, tt_modulation_torch, pcc=0.999) + assert_quality(torch_embedded, tt_embedded_torch, pcc=0.999) + logger.info("PASSED: LTXAdaLayerNormSingle matches PyTorch reference") diff --git a/code/models/tt_dit/tests/models/ltx/test_latent_upsampler.py b/code/models/tt_dit/tests/models/ltx/test_latent_upsampler.py new file mode 100644 index 0000000000000000000000000000000000000000..6e14605dd16b4e28748a69248e1c372cd64a9664 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_latent_upsampler.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 + +"""PCC tests for ``LTXLatentUpsampler`` vs the diffusers ``LTX2LatentUpsamplerModel`` reference. + +All shapes match the production LTX-2 Fast pipeline (mid_channels=1024, num_frames=19, +2x4 BH mesh, half-res latent 17x30 / 18x32).""" + +import pytest +import torch +from loguru import logger + +import ttnn +from models.tt_dit.models.upsampler.latent_upsampler_ltx import LTXLatentUpsampler, LTXUpsamplerResBlock +from models.tt_dit.parallel.config import ParallelFactor, VaeHWParallelConfig +from models.tt_dit.parallel.manager import CCLManager +from models.tt_dit.utils.check import assert_quality +from models.tt_dit.utils.conv3d import ConvDims, conv_pad_height, conv_pad_width +from models.tt_dit.utils.tensor import fast_device_to_host, typed_tensor_2dshard +from models.tt_dit.utils.test import skip_if_unsupported_num_links + +from .ltx_mesh_params import LTX_UPSAMPLER_MESH_PARAMS + + +def _parallel_kwargs( + mesh_device: ttnn.MeshDevice, h_axis: int, w_axis: int, num_links: int, topology: ttnn.Topology +) -> dict: + mesh_shape = tuple(mesh_device.shape) + parallel_config = VaeHWParallelConfig( + height_parallel=ParallelFactor(factor=mesh_shape[h_axis], mesh_axis=h_axis), + width_parallel=ParallelFactor(factor=mesh_shape[w_axis], mesh_axis=w_axis), + ) + ccl_manager = CCLManager(mesh_device, num_links=num_links, topology=topology) + return {"parallel_config": parallel_config, "ccl_manager": ccl_manager} + + +@pytest.mark.parametrize( + "H, W", + [(17, 30), (18, 32)], + ids=["nondiv_17x30", "div_18x32"], +) +@pytest.mark.parametrize( + ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"), + LTX_UPSAMPLER_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_upsampler_resblock( + mesh_device: ttnn.MeshDevice, h_axis: int, w_axis: int, num_links: int, topology: ttnn.Topology, H: int, W: int +): + """Single LTXUpsamplerResBlock at the production shape (channels=1024, 2x4 sharded, + T=19). ``nondiv_17x30`` exercises the sharded pad/mask/crop path; ``div_18x32`` is + the no-mask control.""" + from diffusers.pipelines.ltx2.latent_upsampler import ResBlock as TorchResBlock + + skip_if_unsupported_num_links(mesh_device, num_links) + channels, T, B = 1024, 19, 1 + torch.manual_seed(0xABBA) + pk = _parallel_kwargs(mesh_device, h_axis, w_axis, num_links, topology) + pc = pk["parallel_config"] + hf, wf = pc.height_parallel.factor, pc.width_parallel.factor + padded_h = ((H + hf - 1) // hf) * hf + padded_w = ((W + wf - 1) // wf) * wf + + torch_block = TorchResBlock(channels=channels, dims=3) + torch_block.eval() + + tt_block = LTXUpsamplerResBlock( + channels=channels, + gn_input_nhw=T * H * W, + mesh_device=mesh_device, + conv_dims=ConvDims(T=T + 2, H=padded_h // hf, W=padded_w // wf), + **pk, + ) + tt_block.load_torch_state_dict(torch_block.state_dict()) + + x = torch.randn(B, channels, T, H, W, dtype=torch.float32) + with torch.no_grad(): + torch_out = torch_block(x) # (B, C, T, H, W) + + x_bthwc = x.permute(0, 2, 3, 4, 1).contiguous() # (B, T, H, W, C) + x_bthwc, logical_h = conv_pad_height(x_bthwc, hf) + x_bthwc, logical_w = conv_pad_width(x_bthwc, wf) + x_tt = typed_tensor_2dshard( + x_bthwc, + mesh_device, + shard_mapping={pc.height_parallel.mesh_axis: 2, pc.width_parallel.mesh_axis: 3}, + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=ttnn.bfloat16, + ) + + out = tt_block(x_tt, logical_h, logical_w) + + concat_dims = [None, None] + concat_dims[pc.height_parallel.mesh_axis] = 2 + concat_dims[pc.width_parallel.mesh_axis] = 3 + out_torch = fast_device_to_host(out, mesh_device, concat_dims, ccl_manager=pk["ccl_manager"]) + out_torch = out_torch[:, :, :logical_h, :logical_w, :channels].permute(0, 4, 1, 2, 3) + + assert out_torch.shape == torch_out.shape, f"shape mismatch: TT {out_torch.shape} vs torch {torch_out.shape}" + assert_quality(torch_out, out_torch, pcc=0.99) + logger.info(f"PASSED: LTXUpsamplerResBlock(1024 ch, {T}x{H}x{W}, 2x4) matches reference") + + +@pytest.mark.parametrize( + "H, W", + [(17, 30), (18, 32)], + ids=["nondiv_17x30", "div_18x32"], +) +@pytest.mark.parametrize( + ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"), + LTX_UPSAMPLER_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_latent_upsampler( + mesh_device: ttnn.MeshDevice, h_axis: int, w_axis: int, num_links: int, topology: ttnn.Topology, H: int, W: int +): + """Full ``LTXLatentUpsampler`` PCC vs reference, random weights at production shape + (in=128, mid=1024, T=19, 2x4). No HF checkpoint required.""" + from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel + + skip_if_unsupported_num_links(mesh_device, num_links) + in_c, mid_c, T, B = 128, 1024, 19, 1 + torch.manual_seed(0xC0FFEE) + + torch_model = LTX2LatentUpsamplerModel( + in_channels=in_c, + mid_channels=mid_c, + num_blocks_per_stage=4, + dims=3, + spatial_upsample=True, + temporal_upsample=False, + rational_spatial_scale=2.0, + use_rational_resampler=False, + ) + torch_model.eval() + + tt_model = LTXLatentUpsampler( + input_hw=(H, W), + in_channels=in_c, + mid_channels=mid_c, + num_blocks_per_stage=4, + spatial_upsample=True, + temporal_upsample=False, + spatial_scale=2.0, + rational_resampler=False, + mesh_device=mesh_device, + num_frames=T, + **_parallel_kwargs(mesh_device, h_axis, w_axis, num_links, topology), + ) + tt_model.load_torch_state_dict(torch_model.state_dict()) + + latent = torch.randn(B, in_c, T, H, W, dtype=torch.float32) + + with torch.no_grad(): + torch_out = torch_model(latent) + tt_out = tt_model(latent) + + assert tt_out.shape == torch_out.shape, f"shape mismatch: TT {tt_out.shape} vs torch {torch_out.shape}" + assert_quality(torch_out, tt_out, pcc=0.99) + logger.info(f"PASSED: LTXLatentUpsampler({in_c}->{mid_c}, {T}x{H}x{W} → {T}x{H * 2}x{W * 2}) matches reference") + + +def _resolve_upsampler_checkpoint() -> str | None: + """Resolve the real spatial-upscaler safetensors (local HF cache, no download).""" + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import LocalEntryNotFoundError + + repo_id, filename = "Lightricks/LTX-2.3", "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" + try: + return hf_hub_download(repo_id, filename, local_files_only=True) + except (LocalEntryNotFoundError, Exception): # noqa: BLE001 + return None + + +@pytest.mark.parametrize( + "T, H, W", + [(19, 17, 30), (19, 18, 32)], + ids=["nondiv_17x30", "div_18x32"], +) +@pytest.mark.parametrize( + ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"), + LTX_UPSAMPLER_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_latent_upsampler_real_checkpoint( + mesh_device: ttnn.MeshDevice, + h_axis: int, + w_axis: int, + num_links: int, + topology: ttnn.Topology, + T: int, + H: int, + W: int, +): + """PCC vs reference with the real ``ltx-2.3-spatial-upscaler-x2`` weights (mid_channels=1024) + on the production 2x4 mesh. Input is a smooth, per-channel-scaled latent: white noise + is unrepresentative and lower-bounds bf16 conv PCC through the 8-block residual chain.""" + import json + + from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel + from safetensors import safe_open + from safetensors.torch import load_file + + skip_if_unsupported_num_links(mesh_device, num_links) + ckpt = _resolve_upsampler_checkpoint() + if ckpt is None: + pytest.skip("real upsampler checkpoint not in local HF cache") + + with safe_open(ckpt, framework="pt") as f: + cfg = json.loads(f.metadata()["config"]) + in_c, mid_c = cfg["in_channels"], cfg["mid_channels"] + n_blocks = cfg["num_blocks_per_stage"] + B = 1 + torch.manual_seed(0xC0FFEE) + + torch_model = LTX2LatentUpsamplerModel( + in_channels=in_c, + mid_channels=mid_c, + num_blocks_per_stage=n_blocks, + dims=3, + spatial_upsample=True, + temporal_upsample=False, + rational_spatial_scale=2.0, + use_rational_resampler=False, + ) + torch_model.load_state_dict(load_file(ckpt)) + torch_model.eval() + + tt_model = LTXLatentUpsampler( + input_hw=(H, W), + in_channels=in_c, + mid_channels=mid_c, + num_blocks_per_stage=n_blocks, + spatial_upsample=True, + temporal_upsample=False, + spatial_scale=2.0, + rational_resampler=False, + mesh_device=mesh_device, + num_frames=T, + **_parallel_kwargs(mesh_device, h_axis, w_axis, num_links, topology), + ) + tt_model.load_torch_state_dict(torch_model.state_dict()) + + # Smooth, spatially-correlated latent with per-channel scale (mimics real VAE latents). + noise = torch.randn(B, in_c, T, H, W, dtype=torch.float32) + latent = torch.nn.functional.avg_pool3d(noise, kernel_size=(1, 3, 3), stride=1, padding=(0, 1, 1)) + latent = latent * (torch.rand(in_c) * 4 + 0.5).view(1, in_c, 1, 1, 1) + + with torch.no_grad(): + torch_out = torch_model(latent) + tt_out = tt_model(latent) + + assert tt_out.shape == torch_out.shape, f"shape mismatch: TT {tt_out.shape} vs torch {torch_out.shape}" + assert_quality(torch_out, tt_out, pcc=0.99) + logger.info(f"PASSED: real-checkpoint LTXLatentUpsampler({in_c}->{mid_c}, {n_blocks} blocks) matches reference") diff --git a/code/models/tt_dit/tests/models/ltx/test_ltx_adapter_loader.py b/code/models/tt_dit/tests/models/ltx/test_ltx_adapter_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..7321c3dedc4bdd9eeca800b9e0725c1511f133fb --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_ltx_adapter_loader.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +"""Host-only unit tests for the LTX LoRA adapter loader's module resolution. + +No device, no checkpoint: these pin the loader's warn-and-skip contract for +targets that do not exist on the transformer variant being loaded. A LoRA file +authored against a deeper variant references block indices past this variant's +depth; resolution must return "absent" so the caller skips, not raise IndexError. +""" + +from types import SimpleNamespace + +from models.tt_dit.experimental.lora.ltx_adapter_loader import _get_attn, _resolve_singleton + + +def _fake_transformer(num_blocks: int): + """A transformer stub with ``num_blocks`` blocks, each carrying an ``attn1``.""" + blocks = [SimpleNamespace(attn1=SimpleNamespace(name=f"attn1_{i}")) for i in range(num_blocks)] + return SimpleNamespace(transformer_blocks=blocks) + + +def test_get_attn_out_of_range_returns_none(): + tf = _fake_transformer(2) + assert _get_attn(tf, 5, "attn1") is None + assert _get_attn(tf, -1, "attn1") is None + + +def test_get_attn_in_range_resolves(): + tf = _fake_transformer(2) + assert _get_attn(tf, 1, "attn1") is tf.transformer_blocks[1].attn1 + # Missing attribute on an in-range block still degrades to None, not AttributeError. + assert _get_attn(tf, 1, "attn2") is None + + +def test_resolve_singleton_out_of_range_returns_none_triple(): + tf = _fake_transformer(2) + assert _resolve_singleton(tf, 9, "attn1", "to_out") == (None, None, None) + assert _resolve_singleton(tf, -1, "attn1", "ff1") == (None, None, None) diff --git a/code/models/tt_dit/tests/models/ltx/test_ltx_lora_ondevice.py b/code/models/tt_dit/tests/models/ltx/test_ltx_lora_ondevice.py new file mode 100644 index 0000000000000000000000000000000000000000..e706c09d471831dd5ec5122f1da50817854d72f3 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_ltx_lora_ondevice.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +"""On-device fuse-mode LoRA swap for the LTX-2.3 pipeline. + +Unlike the host-fuse path (``fuse_loras_into`` + full weight reload), swapping here is a +``bind_active`` on each LoRA Linear: ``weight.data += scale * A@B`` on device, no reload. + +Correctness anchor: the host-layout algebra (loader A/B == ``fuse_loras_into``) is already +proven CPU-side. This test closes the remaining gap — that ``LoRAMixin._apply_delta`` applies +that delta correctly on the *sharded device* weight — by comparing the on-device fused weight +against a host recompute of the same registered A/B. + +The device cases need the real 22B checkpoint + a LoRA file + a mesh; they mirror the opt-in +style of ``test_ltx_lora_runtime.py``. Set ``RUN_LORA_GEN=1`` for an end-to-end base/lora/unloaded +generate A/B. +""" + +import os + +import pytest +import torch +from loguru import logger + +import ttnn +from models.common.utility_functions import comp_pcc +from models.tt_dit.experimental.lora.ltx_adapter_loader import iter_lora_modules +from models.tt_dit.pipelines.ltx.pipeline_ltx import LTXPipeline +from models.tt_dit.utils import tensor +from models.tt_dit.utils.test import line_params + + +def _default_checkpoint() -> str: + explicit = os.environ.get("LTX_CHECKPOINT") + if explicit: + return explicit + local = os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-dev.safetensors") + if os.path.exists(local): + return local + return "Lightricks/LTX-2.3:ltx-2.3-22b-dev.safetensors" + + +def _default_lora() -> str: + explicit = os.environ.get("LORA_PATH") + if explicit: + return explicit + local = os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-distilled-lora-384-1.1.safetensors") + if os.path.exists(local): + return local + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id="Lightricks/LTX-2.3", filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors") + + +def _read_weight(param) -> torch.Tensor: + """Reconstruct a sharded Parameter's full weight on host.""" + return tensor.to_torch(param.data, mesh_axes=param.mesh_axes) + + +def _first_qkv_lora_module(transformer): + """First self-attn to_qkv LoRA module (path, module) — the fused/permuted/interleaved + path, i.e. the one most sensitive to a layout bug.""" + for path, mod in iter_lora_modules(transformer): + if path.endswith("attn1.to_qkv"): + return path, mod + return None, None + + +@pytest.mark.timeout(3600) +@pytest.mark.parametrize( + "mesh_device, mesh_shape, sp_axis, tp_axis, num_links, dynamic_load, device_params, topology, is_fsdp", + [ + [(2, 2), (2, 2), 0, 1, 2, False, line_params, ttnn.Topology.Linear, True], + [(4, 8), (4, 8), 1, 0, 2, False, line_params, ttnn.Topology.Linear, False], + ], + ids=["2x2sp0tp1", "bh_4x8sp1tp0"], + indirect=["mesh_device", "device_params"], +) +def test_ondevice_lora_bind_matches_host_delta( + mesh_device, + mesh_shape, + sp_axis, + tp_axis, + num_links, + dynamic_load, + topology, + is_fsdp, +): + """bind_active changes the weight to match base + host(scale*A@B); unbind restores the base + (within bf16 add/subtract drift); the transformer Module is reused throughout.""" + ckpt = _default_checkpoint() + lora = _default_lora() + strength = float(os.environ.get("LORA_STRENGTH", "1.0")) + + parent_mesh = mesh_device + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + pipeline = LTXPipeline.create_pipeline( + mesh_device=mesh_device, + checkpoint_name=ckpt, + gemma_path=None, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + dynamic_load=dynamic_load, + topology=topology, + is_fsdp=is_fsdp, + mode="av", + num_frames=int(os.environ.get("NUM_FRAMES", "25")), + height=int(os.environ.get("HEIGHT", "256")), + width=int(os.environ.get("WIDTH", "256")), + lora_enabled=True, + ) + + if int(ttnn.distributed_context_get_rank()) != 0: + logger.info("Skipping assertions on non-zero rank") + return + + transformer_id = id(pipeline.transformer) + assert pipeline._active_lora is None, "fresh pipeline must have no active LoRA" + + path, module = _first_qkv_lora_module(pipeline.transformer) + assert module is not None, "no attn1.to_qkv LoRA module found — is lora_enabled wired?" + base_w = _read_weight(module.weight).clone() + + # --- bind (on-device fuse) --------------------------------------------- + pipeline.load_lora_weights(lora, strength=strength) + assert id(pipeline.transformer) == transformer_id, "model object was rebuilt, not reused" + assert module.is_lora_active, f"{path} was not bound" + fused_w = _read_weight(module.weight) + assert not torch.allclose(fused_w, base_w), "bind did not change the weight" + + # host recompute of the SAME registered delta; scale folded in exactly as _apply_delta does. + adapter = module.lora_bank[module.active_idx] + delta_oi = module.active_scale * (adapter.B.to(torch.float32) @ adapter.A.to(torch.float32)) # [out,in] + delta = delta_oi if delta_oi.shape == base_w.shape else delta_oi.T + assert delta.shape == base_w.shape, f"delta {tuple(delta.shape)} vs weight {tuple(base_w.shape)}" + expected = base_w.to(torch.float32) + delta + + ok, pcc = comp_pcc(expected, fused_w.to(torch.float32), pcc=0.999) + logger.info(f"[{path}] on-device bind vs host delta: {pcc}") + assert ok, f"on-device fused weight disagrees with host delta recompute: {pcc}" + + # --- unbind ------------------------------------------------------------ + pipeline.unload_lora_weights() + assert pipeline._active_lora is None + assert id(pipeline.transformer) == transformer_id, "model object was rebuilt, not reused" + restored_w = _read_weight(module.weight) + # fuse-mode unbind subtracts in bf16 — close, not bit-exact (see lora.py drift note). + assert torch.allclose( + restored_w.to(torch.float32), base_w.to(torch.float32), atol=1e-2 + ), "unbind did not restore the base weight within bf16 drift" + logger.info("On-device LoRA bind/unbind weight checks passed.") + + if os.environ.get("RUN_LORA_GEN", "0") in ("1", "true", "True"): + steps = int(os.environ.get("STEPS", "6")) + num_frames = int(os.environ.get("NUM_FRAMES", "25")) + height = int(os.environ.get("HEIGHT", "256")) + width = int(os.environ.get("WIDTH", "256")) + seq = pipeline.gemma_encoder_pair.sequence_length + v_p = torch.zeros(1, seq, pipeline.gemma_encoder_pair.video_dim) + a_p = torch.zeros(1, seq, pipeline.gemma_encoder_pair.audio_dim) + for tag, spec in (("base", None), ("lora", lora), ("unloaded", None)): + if spec is not None: + pipeline.load_lora_weights(spec, strength=strength) + else: + pipeline.unload_lora_weights() + pipeline._prepare_transformer(0) + v_lat, a_lat = pipeline.call_av( + video_prompt_embeds=v_p, + audio_prompt_embeds=a_p, + neg_video_prompt_embeds=v_p, + neg_audio_prompt_embeds=a_p, + num_frames=num_frames, + height=height, + width=width, + num_inference_steps=steps, + seed=0, + ge_gamma=0.0, + ) + logger.info(f"[{tag}] generate ran; video latent {tuple(v_lat.shape)}") + assert id(pipeline.transformer) == transformer_id, "model object was rebuilt during generation" diff --git a/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py new file mode 100644 index 0000000000000000000000000000000000000000..d72316f0407555a3fb4fe205ca4916b1603b87e5 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py @@ -0,0 +1,739 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +import functools +import itertools +import math +import os +import shutil +import subprocess +import sys +import tempfile +import time + +import numpy as np +import pytest +import torch +from loguru import logger + +import ttnn +from models.tt_dit.models.audio_vae.audio_decoder_ltx import LTXAudioDecoderAdapter +from models.tt_dit.pipelines.ltx.pipeline_ltx_distilled import LTXDistilledPipeline +from models.tt_dit.utils.ltx import ( + DEFAULT_LTX_PROMPT, + STEADY_STATE_LTX_PROMPT, + STEADY_STATE_REPLAY_LTX_PROMPT, + default_ltx_checkpoint, + default_ltx_gemma, + print_ltx_timing_table, +) +from models.tt_dit.utils.patchifiers import AudioLatentShape, VideoPixelShape +from models.tt_dit.utils.test import skip_if_unsupported_num_links +from models.tt_dit.utils.vbench import assert_vbench_quality + +from .ltx_mesh_params import ( + LTX_DISTILLED_AUDIO_MESH_PARAMS_DL, + LTX_DISTILLED_I2V_MESH_PARAMS_DL, + LTX_DISTILLED_MESH_PARAMS_DL, +) + + +def _ltx_checkpoint_cached(filename: str) -> bool: + """True if the LTX checkpoint is available -- either a real local file (LTX_CHECKPOINT / flat + ~/.cache/ltx-checkpoints) OR already in the HF hub cache. The old guard checked + os.path.exists(default_ltx_checkpoint(...)), but when the checkpoint resolves from the hub that + returns a "repo:file" string (never a path), so it always skipped in CI even with the weights + cached. Checking the hub cache too (local_files_only) lets the offline CI run actually execute.""" + ref = default_ltx_checkpoint(filename) + if os.path.exists(ref): + return True + if ":" in ref: # "repo_id:filename" -> resolvable from the local hub cache without a network call + repo_id, fname = ref.split(":", 1) + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError, LocalEntryNotFoundError + + try: + hf_hub_download(repo_id=repo_id, filename=fname, local_files_only=True) + return True + except (LocalEntryNotFoundError, EntryNotFoundError, FileNotFoundError): + return False + return False + + +# Default-off: full AV gen needs the real LTX checkpoint + Gemma, so it skips without one. +@pytest.mark.skipif( + not _ltx_checkpoint_cached("ltx-2.3-22b-distilled-1.1.safetensors"), + reason="needs the LTX checkpoint (set LTX_CHECKPOINT to a local .safetensors)", +) +@pytest.mark.parametrize( + "no_prompt", + [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), True)], +) +@pytest.mark.parametrize( + "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load", + LTX_DISTILLED_MESH_PARAMS_DL, + indirect=["mesh_device", "device_params"], +) +def test_pipeline_distilled( + mesh_device, + sp_axis, + tp_axis, + num_links, + dynamic_load, + topology, + is_fsdp, + no_prompt, +): + """LTX-2.3 distilled 2-stage AV pipeline.""" + skip_if_unsupported_num_links(mesh_device, num_links) + ckpt = default_ltx_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors") + gemma = default_ltx_gemma() + + parent_mesh = mesh_device + mesh_shape = tuple(parent_mesh.shape) + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + # Stage-2 output target: 1080p @ 24fps, ~6s. Constraints from the pipeline: + # - height/width must be divisible by 64 (asserted at the top of generate(), + # so half-res stage-1 is divisible by 32 for the VAE). + # 1920 is already % 64; 1080 is not → round up to 1088 (next multiple of 64). + # If you need exactly 1080 lines, post-crop the 8 extra rows. + # - num_frames must satisfy (num_frames - 1) % 8 == 0 so the VAE decoder maps + # latent_frames → num_frames frames exactly. 144 (== 6s × 24fps) doesn't + # satisfy this; 145 does (19 latent frames → 145 decoded frames ≈ 6.04s). + num_frames = int(os.environ.get("NUM_FRAMES", "145")) + height = int(os.environ.get("HEIGHT", "1088")) + width = int(os.environ.get("WIDTH", "1920")) + + run_warmup = os.environ.get("RUN_WARMUP", "0") in ("1", "true", "True") + traced = os.environ.get("LTX_TRACED", "0") in ("1", "true", "True") + + # Conditioning image (I2V). Its mere presence drives image_conditioning: with a path the + # transformer builds the per-token video-timestep (I2V) modulation; without one pure T2V keeps + # the fast scalar-AdaLN path. Resolved before create_pipeline so the bool can gate the build. + image_path = os.environ.get("LTX_I2V_IMAGE") + images = None + if image_path: + if not os.path.exists(image_path): + pytest.skip(f"LTX_I2V_IMAGE set but file not found: {image_path}") + strength = float(os.environ.get("LTX_I2V_STRENGTH", "1.0")) + images = [(image_path, 0, strength)] + + pipeline = LTXDistilledPipeline.create_pipeline( + mesh_device=mesh_device, + checkpoint_name=ckpt, + gemma_path=gemma, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + dynamic_load=dynamic_load, + topology=topology, + is_fsdp=is_fsdp, + run_warmup=run_warmup, + traced=traced, + num_frames=num_frames, + height=height, + width=width, + image_conditioning=bool(image_path), + ) + + prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT) + # Gen #1 is the steady-state measurement, and on console every request encodes a prompt the + # cache has never seen — so it gets its own prompt to keep the encoder in the measured path. + # Under dynamic_load the encoder is coresident-excluded with the DiT: a second encode would + # evict the DiT and clobber the captured traces, so that path reuses the cached embedding. + steady_state_prompt = prompt if dynamic_load else os.environ.get("PROMPT_STEADY_STATE", STEADY_STATE_LTX_PROMPT) + replay_prompt = prompt if dynamic_load else STEADY_STATE_REPLAY_LTX_PROMPT + + def run(*, prompt, number, seed): + output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_fast_{width}x{height}_{number}.mp4") + logger.info(f"Running LTX AV Fast: '{prompt[:80]}...'") + logger.info(f"Config: {height}x{width}, {num_frames} frames") + if images: + logger.info(f"I2V: conditioning image {images[0][0]} (strength={images[0][2]})") + + if int(ttnn.distributed_context_get_rank()) != 0: + logger.info(f"Skipping generation on rank {ttnn.distributed_context_get_rank()}") + return + + pipeline.generate( + prompt, + output_path=output_filename, + images=images, + num_frames=num_frames, + height=height, + width=width, + seed=seed, + ) + logger.info(f"Saved video to: {output_filename}") + print_ltx_timing_table( + pipeline, + label="LTX DISTILLED", + num_frames=num_frames, + height=height, + width=width, + mesh_shape=mesh_shape, + sp_axis=sp_axis, + tp_axis=tp_axis, + topology=topology, + output_path=output_filename, + prompt=prompt, + ) + + vbench_thresholds_by_height = { + 1088: { + "subject_consistency": 0.92, + "background_consistency": 0.93, + "motion_smoothness": 0.955, + # Averaged over VBENCH_SEEDS clips: dynamic_degree is near-binary per clip and this content + # sits at ~0.8-1.0, so the floor requires ~4/5 seeds dynamic (1.0 would demand every seed). + "dynamic_degree": 0.8, + "imaging_quality": 0.645, + }, + } + + # RUN_VBENCH=0 skips the quality gate (e.g. perf-only iteration); defaults on so CI gates. + run_vbench = os.environ.get("RUN_VBENCH", "1") in ("1", "true", "True") + # RUN_CLIP=0 skips the CLIP prompt-alignment gate; defaults on (mirrors the wan2.2 test). + run_clip = os.environ.get("RUN_CLIP", "1") in ("1", "true", "True") + + # An enabled quality gate whose deps are absent must report SKIPPED (visible), never a silent + # green pass. Unguarded by rank on purpose: dep availability is per-process (same venv on every + # rank), and a rank-divergent skip would hang the collective generate(). + if run_vbench: + pytest.importorskip("vbench", reason="RUN_VBENCH=1 but vbench not installed (set RUN_VBENCH=0)") + if run_clip: + pytest.importorskip("decord", reason="RUN_CLIP=1 but decord not installed (set RUN_CLIP=0)") + + def check_output_with_vbench(prompt, number, seed=None): + if not run_vbench: + logger.info("RUN_VBENCH=0, skipping VBench quality gate") + return + if int(ttnn.distributed_context_get_rank()) != 0: + return + thresholds = vbench_thresholds_by_height.get(height) + if thresholds is None: + pytest.skip(f"no VBench thresholds calibrated for height {height}") + + # Gate on the average of VBENCH_SEEDS traced replays (default 5): a single clip's VBench score + # is too noisy to gate on — dynamic_degree is near-binary per clip — so score the set together + # (VBench averages over a directory). The primary clip (`number`, `seed`) is already rendered; + # the rest are cheap warm replays. A single-clip caller (no seed) or OUTPUT_PATH (one pinned + # filename) keeps the one-clip path. + num_seeds = int(os.environ.get("VBENCH_SEEDS", "5")) + if seed is None or num_seeds <= 1 or os.environ.get("OUTPUT_PATH"): + output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_fast_{width}x{height}_{number}.mp4") + assert_vbench_quality(output_filename, prompt=prompt, thresholds=thresholds) + return + + cwd = os.path.abspath(os.getcwd()) + + with tempfile.TemporaryDirectory() as vbench_dir: + + def _link(idx: int, mp4_number: int) -> None: + # The render always lands directly in CWD; resolve the path and confirm it stays there so + # the number threaded into the filename can't escape it, and give the link a plain integer + # name so nothing dynamic reaches the destination path either. + src = os.path.abspath(os.path.join(cwd, f"ltx_av_fast_{width}x{height}_{mp4_number}.mp4")) + if os.path.dirname(src) != cwd: + raise ValueError(f"render path escaped the working dir: {src!r}") + os.symlink(src, os.path.join(vbench_dir, f"seed_{idx}.mp4")) + + _link(0, number) + for k in range(1, num_seeds): + run(prompt=prompt, number=1000 + k, seed=seed + k) + _link(k, 1000 + k) + logger.info(f"VBench gate averaged over {num_seeds} seeds (base seed {seed})") + assert_vbench_quality(vbench_dir, prompt=prompt, thresholds=thresholds) + + def check_output_with_clip(prompt, number, clip_threshold=None): + # Mirrors wan2.2's check_output_with_clip: sample ~8 evenly-spaced frames, score each + # against the prompt with CLIP, assert the mean clears a floor. LTX writes the video to + # disk (generate() returns only the path), so frames are read back from the mp4. + if not run_clip: + logger.info("RUN_CLIP=0, skipping CLIP score check") + return + if int(ttnn.distributed_context_get_rank()) != 0: + return + try: + from decord import VideoReader + from PIL import Image + + from models.tt_dit.tests.dataset_eval.clip_encoder import CLIPEncoder + except ImportError as e: + pytest.skip(f"CLIP deps unavailable ({e})") + + # LTX baseline mean ~31.3 (vs wan2.2's ~37): the LTX prompt exceeds CLIP's 77-token limit + # so only its head is scored, and frames are read back re-encoded (CRF=25). 28.0 leaves + # ~3pt headroom for run-to-run variance. Override with CLIP_THRESHOLD. + threshold = float(os.environ.get("CLIP_THRESHOLD", "28.0")) if clip_threshold is None else clip_threshold + output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_fast_{width}x{height}_{number}.mp4") + + vr = VideoReader(output_filename) + total_frames = len(vr) + indices = np.linspace(0, total_frames - 1, min(8, total_frames), dtype=int) + + clip_encoder = CLIPEncoder() + scores = [] + for i in indices: + frame = vr[int(i)].asnumpy() # (H, W, 3) uint8 RGB + pil_img = Image.fromarray(frame.astype(np.uint8)) + score = clip_encoder.get_clip_score(prompt, pil_img).item() * 100.0 + scores.append(score) + + clip_min = min(scores) + clip_max = max(scores) + clip_mean = sum(scores) / len(scores) + logger.info(f"CLIP scores: min={clip_min:.2f}, max={clip_max:.2f}, mean={clip_mean:.2f}") + + assert clip_mean >= threshold, ( + f"Mean CLIP score {clip_mean:.2f} is below threshold {threshold:.2f}. " + f"Per-frame scores: {[f'{s:.2f}' for s in scores]}" + ) + + if no_prompt: + seed = int(os.environ.get("SEED", "10")) + run(prompt=prompt, number=0, seed=seed) + # Traced: gen #0 captures the denoise/VAE/audio traces (lazily, on first step of each stage). + if traced: + # Gen #1 captures the encode trace (the pipeline opens that gate once its own captures + # are done), so gen #2 is the first pass where every trace replays — the one whose + # timings are the steady state. + logger.info("=== gen #1: encode trace capture ===") + run(prompt=steady_state_prompt, number=1, seed=seed) + logger.info("=== traced steady-state pass (gen #2, pure replay) ===") + run(prompt=replay_prompt, number=2, seed=seed) + # Gate gen #2: a corrupted full-replay pass is the failure this structure exists to + # catch, and gen #1 still has a capture in it. + check_output_with_clip(replay_prompt, 2) + check_output_with_vbench(replay_prompt, 2, seed=seed) + else: + check_output_with_clip(prompt, 0) + check_output_with_vbench(prompt, 0) + else: + for i in itertools.count(): + new_prompt = input("Enter the input prompt, or q to exit: ") + if new_prompt: + prompt = new_prompt + if prompt[0] == "q": + break + run(prompt=prompt, number=i, seed=i) + check_output_with_clip(prompt, i) + check_output_with_vbench(prompt, i) + + if traced: + pipeline.release_traces() + + +@functools.lru_cache(maxsize=1) +def _ffmpeg(): + """An ffmpeg executable. Prefer the system binary; fall back to imageio-ffmpeg's + bundled static build (pip-installed on demand) so a host without ffmpeg still runs + the frame-decode checks. The static build is real ffmpeg, so decoded pixels — and + thus the calibrated seam thresholds — are unchanged.""" + exe = shutil.which("ffmpeg") + if exe: + return exe + try: + import imageio_ffmpeg + except ImportError: + pip_install = [sys.executable, "-m", "pip", "install", "-q", "imageio-ffmpeg"] + if subprocess.run(pip_install).returncode != 0: + # uv-managed venvs ship without pip; bootstrap it from the stdlib, then retry. + subprocess.run([sys.executable, "-m", "ensurepip", "--upgrade"], check=True) + subprocess.run(pip_install, check=True) + import imageio_ffmpeg + return imageio_ffmpeg.get_ffmpeg_exe() + + +def _ffmpeg_frames(path, n): + """Sample ``n`` luma frames evenly across ``path`` (via ffmpeg) as float ndarrays.""" + import io + + from PIL import Image + + frames = [] + for t in np.linspace(0.1, 5.5, n): + raw = subprocess.run( + [ + _ffmpeg(), + "-v", + "error", + "-ss", + f"{t}", + "-i", + path, + "-vframes", + "1", + "-f", + "image2pipe", + "-vcodec", + "png", + "-", + ], + capture_output=True, + ).stdout + if raw: + frames.append(np.asarray(Image.open(io.BytesIO(raw)).convert("L")).astype(float)) + return frames + + +def _temporal_seam_score(path): + """Grid-seam strength at the 2x4 mesh boundaries (W/4,W/2,3W/4 vertical; H/2 horizontal). + Per-column/row TEMPORAL MEDIAN of the gradient isolates a static seam from moving content. + Each boundary is scored against its LOCAL neighbour gradient (a ring +/-12 px, skipping + +/-1 px for seam width) — not the global median, which a centred subject over a dark, + low-gradient background inflates (the subject straddles W/2, not a mesh line). A baked seam + is a thin gradient spike pinned to the boundary, so it stands well above its flat local + ring; clean content scores ~1.""" + fs = _ffmpeg_frames(path, 16) + assert fs, f"no frames decoded from {path}" + h, w = fs[0].shape + gx = np.median(np.stack([np.abs(np.diff(f, axis=1)).mean(0) for f in fs]), 0) + gy = np.median(np.stack([np.abs(np.diff(f, axis=0)).mean(1) for f in fs]), 0) + + def _ring_ratio(g, idx, skip=1, span=12): + lo, hi = max(0, idx - span), min(len(g), idx + span + 1) + ring = np.concatenate([g[lo : idx - skip], g[idx + skip + 1 : hi]]) + return float(g[idx] / (np.median(ring) + 1e-6)) + + v = float(np.mean([_ring_ratio(gx, round(w * i / 4)) for i in (1, 2, 3)])) + hh = _ring_ratio(gy, round(h / 2)) + return v, hh + + +@pytest.mark.skipif( + not _ltx_checkpoint_cached("ltx-2.3-22b-distilled-1.1.safetensors"), + reason="needs the LTX checkpoint (set LTX_CHECKPOINT to a local .safetensors)", +) +@pytest.mark.parametrize( + "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load", + LTX_DISTILLED_I2V_MESH_PARAMS_DL, + indirect=["mesh_device", "device_params"], +) +def test_pipeline_distilled_i2v(mesh_device, sp_axis, tp_axis, num_links, dynamic_load, topology, is_fsdp, tmp_path): + """E2E chained t2v->i2v: generate a t2v clip, then i2v continuing from its LAST frame + (same DEFAULT_LTX_PROMPT), and splice both into one continuous ~12s clip. Asserts the + I2V output (a) reproduces the conditioning frame at frame-0 (conditioning works) and + (b) is free of the VAE 2x4 grid seam (guards the non-mesh-aligned i2v fix at 1088x1920, + whose s1 cond latent is the uneven 17x30 that used to seam). An external LTX_I2V_IMAGE + overrides the t2v pre-gen (single i2v, no splice).""" + import subprocess + + from PIL import Image + + skip_if_unsupported_num_links(mesh_device, num_links) + + ckpt = default_ltx_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors") + parent_mesh = mesh_device + mesh_shape = tuple(parent_mesh.shape) + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + num_frames = int(os.environ.get("NUM_FRAMES", "145")) + height = int(os.environ.get("HEIGHT", "1088")) # 1088x1920 -> s1 latent 17x30 (uneven, the fixed case) + width = int(os.environ.get("WIDTH", "1920")) + traced = os.environ.get("LTX_TRACED", "0") in ("1", "true", "True") + seed = int(os.environ.get("SEED", "10")) + prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT) + + pipeline = LTXDistilledPipeline.create_pipeline( + mesh_device=mesh_device, + checkpoint_name=ckpt, + gemma_path=default_ltx_gemma(), + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + dynamic_load=dynamic_load, + topology=topology, + is_fsdp=is_fsdp, + run_warmup=traced, # trace capture needs warmup; eager doesn't + traced=traced, + num_frames=num_frames, + height=height, + width=width, + # This test conditions on a frame produced in-process (t2v last frame), so there's no + # conditioning-image path in hand at construction — force the I2V path on explicitly. + image_conditioning=True, + ) + + if int(ttnn.distributed_context_get_rank()) != 0: + return + + def _gen(out, images): + pipeline.generate( + prompt, output_path=str(out), images=images, num_frames=num_frames, height=height, width=width, seed=seed + ) + + # Conditioning image: an explicit LTX_I2V_IMAGE conditions directly on it and SKIPS the t2v + # pre-gen (an external/user image, or a t2v last frame produced by a prior process). Else this + # generates a t2v clip and conditions on its LAST FRAME so the i2v continues the shot. + cond_override = os.environ.get("LTX_I2V_IMAGE") + strength = float(os.environ.get("LTX_I2V_STRENGTH", "1.0")) + t2v = None + if cond_override: + cond = cond_override + else: + t2v = tmp_path / "t2v.mp4" + _gen(t2v, None) + # -sseof seeks near the end and -update rewrites the same file per decoded frame, so it + # lands on the true final frame without buffering the whole clip into memory. + cond = tmp_path / "cond_lastframe.png" + subprocess.run( + [_ffmpeg(), "-v", "error", "-sseof", "-3", "-i", str(t2v), "-update", "1", "-y", str(cond)], + check=True, + ) + if os.environ.get("LTX_T2V_OUT"): + shutil.copy(str(t2v), os.environ["LTX_T2V_OUT"]) + if os.environ.get("LTX_COND_OUT"): + shutil.copy(str(cond), os.environ["LTX_COND_OUT"]) + # On a 4x8 mesh the image encoder and DiT coresident-evict, so a t2v->i2v chain in one + # process reloads the encoder over t2v's captured traces and wedges. LTX_T2V_ONLY stops + # here so a fresh process runs i2v on LTX_COND_OUT (via LTX_I2V_IMAGE) and host-splices. + if os.environ.get("LTX_T2V_ONLY"): + print( + f"I2V_E2E t2v-only -> clip={os.environ.get('LTX_T2V_OUT')} cond={os.environ.get('LTX_COND_OUT')}", + flush=True, + ) + if traced: + pipeline.release_traces() + return + + # i2v conditioned on that frame + i2v = tmp_path / "i2v.mp4" + _gen(i2v, [(str(cond), 0, strength)]) + if os.environ.get("LTX_I2V_OUT"): + shutil.copy(str(i2v), os.environ["LTX_I2V_OUT"]) + + # In-process splice of t2v + i2v into one continuous clip (each gen is 145f@24fps -> ~12s). + # Only reached where encoder+DiT+traces fit one process (not the 4x8 mesh — use LTX_T2V_ONLY + # plus a second i2v process there). An external LTX_I2V_IMAGE has no preceding t2v to splice. + if t2v is not None: + chained = os.environ.get("LTX_CHAINED_OUT") or str(tmp_path / "chained.mp4") + subprocess.run( + [ + _ffmpeg(), + "-v", + "error", + "-i", + str(t2v), + "-i", + str(i2v), + "-filter_complex", + "[0:v:0][0:a:0][1:v:0][1:a:0]concat=n=2:v=1:a=1[v][a]", + "-map", + "[v]", + "-map", + "[a]", + "-y", + chained, + ], + check=True, + ) + print(f"I2V_E2E chained t2v+i2v -> {chained}", flush=True) + + # (a) conditioning works: i2v frame-0 reproduces the conditioning frame (VAE roundtrip + CRF, + # so not identity — but a far tighter correlation than an unconditioned gen of the same prompt). + def _luma0(path): + raw = subprocess.run( + [_ffmpeg(), "-v", "error", "-i", path, "-vframes", "1", "-f", "image2pipe", "-vcodec", "png", "-"], + capture_output=True, + ).stdout + return Image.open(__import__("io").BytesIO(raw)).convert("L") + + def _vec(im): + return torch.from_numpy(np.asarray(im).astype("float32")).flatten() + + f0_img = _luma0(str(i2v)) + cond_img = _luma0(str(cond)) + # Compare a matched field of view: an external LTX_I2V_IMAGE can differ in resolution/aspect + # from the generated frame (the pipeline center-crops the conditioning image to the target + # aspect before encoding). Mirror that — center-crop cond to the frame's aspect, then resize. + # No-op when cond is a same-size/same-aspect frame (the self-conditioned t2v-last-frame case). + fw, fh = f0_img.size + cw, ch = cond_img.size + aspect = fw / fh + if cw / ch > aspect: + nw = round(ch * aspect) + x0 = (cw - nw) // 2 + cond_img = cond_img.crop((x0, 0, x0 + nw, ch)) + elif cw / ch < aspect: + nh = round(cw / aspect) + y0 = (ch - nh) // 2 + cond_img = cond_img.crop((0, y0, cw, y0 + nh)) + if cond_img.size != f0_img.size: + cond_img = cond_img.resize(f0_img.size) + + c, f0 = _vec(cond_img), _vec(f0_img) + pcc = torch.corrcoef(torch.stack([c, f0]))[0, 1].item() + print(f"\nI2V_E2E frame0-vs-cond PCC={pcc:.4f}", flush=True) + + # (b) seam-free at the uneven 1088x1920 (the i2v grid-seam fix). Local-ring ratios sit at + # ~1.0 on clean content (measured 0.85-1.16 across all boundaries); a baked seam spikes + # well above its flat neighbourhood. Threshold 1.5 clears the clean band with margin. + v, hh = _temporal_seam_score(str(i2v)) + print(f"I2V_E2E seam V={v:.2f} H={hh:.2f} (clean~1.0, seam>>1)", flush=True) + + if traced: + pipeline.release_traces() + + assert pcc > 0.85, f"i2v frame-0 does not reproduce the conditioning frame (PCC={pcc:.4f}) — conditioning broken" + assert v < 1.5 and hh < 1.5, f"i2v grid seam present (V={v:.2f}, H={hh:.2f}) at 1088x1920" + + +@pytest.mark.skipif( + not _ltx_checkpoint_cached("ltx-2.3-22b-distilled-1.1.safetensors"), + reason="needs the LTX checkpoint (set LTX_CHECKPOINT to a local .safetensors)", +) +@pytest.mark.parametrize( + "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load", + LTX_DISTILLED_AUDIO_MESH_PARAMS_DL, + indirect=["mesh_device", "device_params"], +) +def test_audio_decode_girl(mesh_device, sp_axis, tp_axis, num_links, dynamic_load, topology, is_fsdp): + """Just the audio section of the AV pipeline: real checkpoint weights, real girl-clip + latent shape (num_frames -> als.frames), no transformer/gemma (decode_audio never + encodes prompts). Profiles cold (weight load + compile) vs warm (steady-state per-gen) + decode wall. LTX_TRACED=1 traces the main vocoder (cold/warm timing only — the eager + torch-oracle quality block is skipped under trace; run LTX_TRACED=0 for the stats).""" + skip_if_unsupported_num_links(mesh_device, num_links) + + parent_mesh = mesh_device + mesh_shape = tuple(parent_mesh.shape) + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + num_frames = int(os.environ.get("NUM_FRAMES", "145")) # ~6.04s @ 24fps + height = int(os.environ.get("HEIGHT", "1088")) + width = int(os.environ.get("WIDTH", "1920")) + traced = os.environ.get("LTX_TRACED", "0") in ("1", "true", "True") + + ckpt = default_ltx_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors") + # Audio-only build: create the pipeline WITHOUT a checkpoint so __init__ builds/primes + # nothing (skipping the ~10-min video warmup and the 22B transformer / VAE / upsampler + # prime that dominate a cold run — decode_audio only touches the audio decoder + vocoder). + # Then point it at the checkpoint and construct + prime just the audio decode chain; + # decode_audio compiles + captures its own trace lazily on its first call. + pipeline = LTXDistilledPipeline.create_pipeline( + mesh_device=mesh_device, + checkpoint_name=None, + gemma_path=default_ltx_gemma(), # built as a lazy shell only; never loaded here + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + dynamic_load=dynamic_load, + topology=topology, + is_fsdp=is_fsdp, + run_warmup=False, + traced=traced, + num_frames=num_frames, + height=height, + width=width, + ) + pipeline.checkpoint_name = LTXDistilledPipeline._resolve_checkpoint_file(ckpt) + pipeline._audio_adapter = LTXAudioDecoderAdapter( + pipeline.checkpoint_name, + mesh_device=pipeline.mesh_device, + vae_ccl_manager=pipeline.vae_ccl_manager, + dit_parallel_config=pipeline.parallel_config, + traced=pipeline._traced, + ) + pipeline._prepare_audio_decoder() + + vps = VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=24) + als = AudioLatentShape.from_video_pixel_shape(vps) + # Seeded synthetic latent scaled to a real girl-clip audio latent's statistics (mean≈0.2, + # std≈0.92). The PCC gate below is TT-vs-torch with both legs fed the same mel, so the input + # need not be "real" content — only AUDIO_OUT .wav listening would differ. AUDIO_LATENT + # (.npy or .pt) still overrides for anyone who wants to decode a real dumped latent. + _lat = os.environ.get("AUDIO_LATENT") + if _lat and os.path.exists(_lat): + latent = (torch.from_numpy(np.load(_lat)) if _lat.endswith(".npy") else torch.load(_lat)).float() + else: + torch.manual_seed(0) + latent = torch.randn(1, als.frames, pipeline.in_channels, dtype=torch.float32) * 0.92 + 0.2 + + t0 = time.perf_counter() + audio = pipeline.decode_audio(latent, num_frames, fps=24.0) # cold: weight load + compile (+ capture) + cold_ms = (time.perf_counter() - t0) * 1000 + + # One untimed decode to absorb a late first-replay cost (the BWE trace captures on the + # first post-cold call, not during the cold decode), so warm_ms is true steady state. + pipeline.decode_audio(latent, num_frames, fps=24.0) + # WARM_REPS shrinks the steady-state loop for fast dev iteration (the PCC oracle below is + # unaffected); default 5 keeps warm_ms a stable average for reported timings. + N = int(os.environ.get("WARM_REPS", "5")) + t0 = time.perf_counter() + for _ in range(N): + audio = pipeline.decode_audio(latent, num_frames, fps=24.0) + warm_ms = (time.perf_counter() - t0) * 1000 / N + wav = audio.waveform + sr = audio.sampling_rate + out_dir = os.environ.get("AUDIO_OUT") + + def _save(w, name): + if not out_dir: + return + import soundfile as sf + + os.makedirs(out_dir, exist_ok=True) + a = w[0] if w.dim() == 3 else w # (2, T) -> (T, 2) + path = os.path.join(out_dir, name) + sf.write(path, a.transpose(0, 1).cpu().numpy(), int(sr)) + logger.info(f"wrote {path}") + + _save(wav, f"girl_audio_conv1d{'_traced' if traced else ''}.wav") + + # Torch-oracle quality gate vs the diffusers vocoder+BWE (real checkpoint, CPU) — the ground + # truth. The conv1d path runs under BOTH eager and traced (gating the traced vocoder/BWE + # output, which is otherwise unverified — exactly the gap that let a broken trace ship). + # Per-1s-interval error-to-signal (rmse/σ, dB) + overall PCC; the absolute level (~−18 dB + # conv1d) is the fp32/bf16 floor, gated only against a gross spike. + from models.tt_dit.tests.models.ltx.test_audio_ltx import _build_torch_stage_c_real + + z = pipeline.tt_mel_decoder.z_channels + audio_spatial = latent.reshape(1, latent.shape[1], z, latent.shape[2] // z).permute(0, 2, 1, 3).float() + mel = pipeline.tt_mel_decoder(audio_spatial) # TT mel, fed to both TT and torch + w_conv = pipeline.tt_vocoder_with_bwe(mel).squeeze(0).float() + with torch.no_grad(): + w_torch = _build_torch_stage_c_real(ckpt)(mel.float()).squeeze(0).float() + _save(w_conv, f"girl_audio_conv1d{'_traced' if traced else ''}.wav") + _save(w_torch, "girl_audio_torch.wav") + + nn = min(w_torch.shape[-1], w_conv.shape[-1]) + + def _vs_torch(test, name): + rows = [] + for start in range(0, nn, sr): + r, t = w_torch[..., start : start + sr], test[..., start : start + sr] + sig = r.pow(2).mean().sqrt().item() + 1e-9 + rr = ((r - t) ** 2).mean().sqrt().item() / sig + rows.append(rr) + logger.info( + f"[oracle] t={start//sr}s {name}-vs-torch rmse/σ={rr:.3e} ({20*math.log10(rr+1e-12):+.1f} dB) " + f"max|Δ|={(r - t).abs().max().item():.3e}" + ) + pcc = torch.corrcoef(torch.stack([w_torch[..., :nn].flatten(), test[..., :nn].flatten()]))[0, 1].item() + worst_, med_ = max(rows), sorted(rows)[len(rows) // 2] + print( + f"\nAUDIO_GIRL {name}-vs-torch traced={traced} worst_rmse/σ={worst_:.3e} " + f"({20*math.log10(worst_+1e-12):+.1f} dB) median={med_:.3e} ratio={worst_/(med_+1e-9):.2f} PCC={pcc:.5f}", + flush=True, + ) + return worst_, med_, pcc + + cw, cm, c_pcc = _vs_torch(w_conv, "conv1d") + # Gross localized spike only: an interval >4x the median AND past a clearly-audible floor + # (rmse/σ 0.5 = −6 dB). Also gate overall PCC vs torch — a broken (e.g. mis-traced) vocoder + # drops PCC far below the ~0.99 fp32/bf16 floor. + assert cw < 0.5 and cw < 4 * cm + 1e-6, f"conv1d-vs-torch localized spike: worst {cw:.3e} vs median {cm:.3e}" + assert c_pcc > 0.95, f"conv1d-vs-torch PCC {c_pcc:.4f} below floor (traced={traced}) — vocoder output is wrong" + + pipeline.release_traces() + dur = wav.shape[-1] / sr + print( + f"\nAUDIO_GIRL traced={traced} " + f"latent_frames={als.frames} out={tuple(wav.shape)} {dur:.2f}s@{sr}Hz cold={cold_ms:.1f}ms warm={warm_ms:.1f}ms", + flush=True, + ) + assert torch.isfinite(wav).all(), "decoded waveform has non-finite samples" + assert abs(dur - num_frames / 24.0) < 0.2, f"duration {dur:.2f}s != ~{num_frames/24.0:.2f}s" diff --git a/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_one_stage.py b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_one_stage.py new file mode 100644 index 0000000000000000000000000000000000000000..1e3794a6424ec3d73354f5ed60ee4d6cde8df5b0 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_one_stage.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +""" +LTX-2 one-stage pipeline tests: +- Euler step scheduler primitive (no device needed) +- End-to-end one-stage AV generation (LTXOneStagePipeline) on a multi-chip mesh +""" + +import itertools +import os + +import pytest +import torch +from loguru import logger + +import ttnn +from models.tt_dit.pipelines.ltx.pipeline_ltx import euler_step +from models.tt_dit.pipelines.ltx.pipeline_ltx_one_stage import LTXOneStagePipeline +from models.tt_dit.utils.ltx import ( + DEFAULT_LTX_PROMPT, + default_ltx_checkpoint, + default_ltx_gemma, + print_ltx_timing_table, +) +from models.tt_dit.utils.test import skip_if_unsupported_num_links + +from .ltx_mesh_params import LTX_ONE_STAGE_MESH_PARAMS_DL + + +def test_euler_step(): + """Test Euler step: x_{t+1} = x_t + velocity * dt.""" + torch.manual_seed(42) + sample = torch.randn(1, 256, 128) + denoised = torch.randn(1, 256, 128) + sigma, sigma_next = 0.8, 0.6 + + result = euler_step(sample, denoised, sigma=sigma, sigma_next=sigma_next) + + # Manual computation + velocity = (sample.float() - denoised.float()) / sigma + dt = sigma_next - sigma + expected = (sample.float() + velocity * dt).to(sample.dtype) + + assert torch.allclose( + result, expected, atol=1e-6 + ), f"Euler step mismatch: max diff {(result - expected).abs().max():.2e}" + logger.info("Euler step matches manual computation") + + +# --------------------------------------------------------------------------- +# End-to-end one-stage AV generation (requires device + checkpoint) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "no_prompt", + [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), True)], +) +@pytest.mark.parametrize( + "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load", + LTX_ONE_STAGE_MESH_PARAMS_DL, + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize( + "width, height", + [ + (1920, 1088), + ], + ids=[ + "resolution_1080p", + ], +) +def test_pipeline_one_stage( + mesh_device, + sp_axis, + tp_axis, + num_links, + dynamic_load, + topology, + width, + height, + is_fsdp, + no_prompt, +): + skip_if_unsupported_num_links(mesh_device, num_links) + ckpt = default_ltx_checkpoint("ltx-2.3-22b-dev.safetensors") + gemma = default_ltx_gemma() + # ckpt / gemma always resolve (env var → local → HF repo string fallback). The pipeline's + # resolver downloads from HF if needed. + + parent_mesh = mesh_device + mesh_shape = tuple(parent_mesh.shape) + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + num_frames = int(os.environ.get("NUM_FRAMES", "145")) + num_inference_steps = int(os.environ.get("NUM_STEPS", "30")) + width = int(os.environ.get("WIDTH", width)) + height = int(os.environ.get("HEIGHT", height)) + + run_warmup = os.environ.get("RUN_WARMUP", "0") in ("1", "true", "True") + + pipeline = LTXOneStagePipeline.create_pipeline( + mesh_device=mesh_device, + checkpoint_name=ckpt, + gemma_path=gemma, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + dynamic_load=dynamic_load, + topology=topology, + is_fsdp=is_fsdp, + run_warmup=run_warmup, + num_frames=num_frames, + height=height, + width=width, + ) + + prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT) + + def run(*, prompt, number, seed): + logger.info(f"Running inference with prompt: '{prompt}'") + logger.info(f"Parameters: {height}x{width}, {num_frames} frames, {num_inference_steps} steps") + + output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_pro_{width}x{height}_{number}.mp4") + + pipeline.generate( + prompt, + output_path=output_filename, + num_frames=num_frames, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + ) + + if int(ttnn.distributed_context_get_rank()) == 0: + # Minimal artifact gate: the leg must not pass on missing/empty output. A calibrated + # VBench/CLIP quality check (like the distilled path) is the stronger follow-up. + assert os.path.exists(output_filename), f"pipeline produced no video at {output_filename}" + assert os.path.getsize(output_filename) > 0, f"pipeline produced an empty video at {output_filename}" + logger.info(f"Saved video to: {output_filename}") + print_ltx_timing_table( + pipeline, + label="LTX ONE-STAGE", + num_frames=num_frames, + height=height, + width=width, + mesh_shape=mesh_shape, + sp_axis=sp_axis, + tp_axis=tp_axis, + topology=topology, + output_path=output_filename, + prompt=prompt, + ) + else: + logger.info(f"Skipping video export on rank {ttnn.distributed_context_get_rank()}") + + if no_prompt: + run(prompt=prompt, number=0, seed=42) + else: + for i in itertools.count(): + new_prompt = input("Enter the input prompt, or q to exit: ") + if new_prompt: + prompt = new_prompt + if prompt[0] == "q": + break + run(prompt=prompt, number=i, seed=i) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py new file mode 100644 index 0000000000000000000000000000000000000000..56eb790ef190f73eb0cfdd9d5125dfd3284a98ea --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +import itertools +import os + +import pytest +from loguru import logger + +import ttnn +from models.tt_dit.pipelines.ltx.pipeline_ltx_two_stages import LTXTwoStagesPipeline +from models.tt_dit.utils.ltx import ( + DEFAULT_LTX_PROMPT, + default_ltx_checkpoint, + default_ltx_gemma, + print_ltx_timing_table, +) +from models.tt_dit.utils.test import skip_if_unsupported_num_links + +from .ltx_mesh_params import LTX_TWO_STAGES_I2V_MESH_PARAMS_DL, LTX_TWO_STAGES_MESH_PARAMS_DL + + +def _default_distilled_lora() -> str: + """Resolve distilled stage-2 LoRA: env var > local file > HF (auto-downloaded here).""" + explicit = os.environ.get("DISTILLED_LORA_PATH") + if explicit: + return explicit + local = os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-distilled-lora-384-1.1.safetensors") + if os.path.exists(local): + return local + from huggingface_hub import hf_hub_download + + logger.info( + "Resolving HuggingFace distilled LoRA Lightricks/LTX-2.3:ltx-2.3-22b-distilled-lora-384-1.1.safetensors" + ) + return hf_hub_download(repo_id="Lightricks/LTX-2.3", filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors") + + +@pytest.mark.parametrize( + "no_prompt", + [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), True)], +) +@pytest.mark.parametrize( + "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load", + LTX_TWO_STAGES_MESH_PARAMS_DL, + indirect=["mesh_device", "device_params"], +) +def test_pipeline_two_stages( + mesh_device, + sp_axis, + tp_axis, + num_links, + dynamic_load, + topology, + is_fsdp, + no_prompt, +): + """LTX-2.3 22B 2-stage AV pipeline: full-guidance s1 + distilled-LoRA s2 refine.""" + skip_if_unsupported_num_links(mesh_device, num_links) + ckpt = default_ltx_checkpoint("ltx-2.3-22b-dev.safetensors") + distilled_lora = _default_distilled_lora() + gemma = default_ltx_gemma() + + parent_mesh = mesh_device + mesh_shape = tuple(parent_mesh.shape) + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + num_frames = int(os.environ.get("NUM_FRAMES", "145")) + height = int(os.environ.get("HEIGHT", "1088")) + width = int(os.environ.get("WIDTH", "1920")) + num_inference_steps = int(os.environ.get("NUM_INFERENCE_STEPS", "30")) + + run_warmup = os.environ.get("RUN_WARMUP", "0") in ("1", "true", "True") + pipeline = LTXTwoStagesPipeline.create_pipeline( + mesh_device=mesh_device, + checkpoint_name=ckpt, + gemma_path=gemma, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + dynamic_load=dynamic_load, + topology=topology, + is_fsdp=is_fsdp, + distilled_lora_path=distilled_lora, + run_warmup=run_warmup, + num_frames=num_frames, + height=height, + width=width, + ) + + prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT) + + def run(*, prompt, number, seed): + output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_two_stages_{width}x{height}_{number}.mp4") + logger.info(f"Running LTX AV Two-Stages: '{prompt[:80]}...'") + logger.info(f"Config: {height}x{width}, {num_frames} frames, {num_inference_steps} stage-1 steps") + + if int(ttnn.distributed_context_get_rank()) != 0: + logger.info(f"Skipping generation on rank {ttnn.distributed_context_get_rank()}") + return + + pipeline.generate( + prompt, + output_path=output_filename, + num_frames=num_frames, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + ) + logger.info(f"Saved video to: {output_filename}") + print_ltx_timing_table( + pipeline, + label="LTX TWO-STAGES", + num_frames=num_frames, + height=height, + width=width, + mesh_shape=mesh_shape, + sp_axis=sp_axis, + tp_axis=tp_axis, + topology=topology, + output_path=output_filename, + prompt=prompt, + ) + + if no_prompt: + run(prompt=prompt, number=0, seed=int(os.environ.get("SEED", "10"))) + else: + for i in itertools.count(): + new_prompt = input("Enter the input prompt, or q to exit: ") + if new_prompt: + prompt = new_prompt + if prompt[0] == "q": + break + run(prompt=prompt, number=i, seed=i) + + +# --------------------------------------------------------------------------- +# I2V smoke test (image + text -> video) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load", + LTX_TWO_STAGES_I2V_MESH_PARAMS_DL, + indirect=["mesh_device", "device_params"], +) +def test_pipeline_two_stages_i2v_smoke( + mesh_device, + sp_axis, + tp_axis, + num_links, + dynamic_load, + topology, + is_fsdp, +): + """End-to-end I2V smoke: generate(images=[(img, 0, 1.0)]) produces a non-empty MP4 whose + first frame correlates with the conditioning image. + + Gated behind LTX_I2V_IMAGE (path to a conditioning image); skips otherwise so it never runs + in the default CI sweep without an input image / checkpoint present. + """ + image_path = os.environ.get("LTX_I2V_IMAGE") + if not image_path or not os.path.exists(image_path): + pytest.skip("set LTX_I2V_IMAGE to a conditioning image path to run the I2V smoke test") + + skip_if_unsupported_num_links(mesh_device, num_links) + + ckpt = default_ltx_checkpoint("ltx-2.3-22b-dev.safetensors") + distilled_lora = _default_distilled_lora() + gemma = default_ltx_gemma() + + parent_mesh = mesh_device + mesh_shape = tuple(parent_mesh.shape) + mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape)) + + num_frames = int(os.environ.get("NUM_FRAMES", "25")) + height = int(os.environ.get("HEIGHT", "512")) + width = int(os.environ.get("WIDTH", "768")) + num_inference_steps = int(os.environ.get("NUM_INFERENCE_STEPS", "8")) + strength = float(os.environ.get("LTX_I2V_STRENGTH", "1.0")) + + pipeline = LTXTwoStagesPipeline.create_pipeline( + mesh_device=mesh_device, + checkpoint_name=ckpt, + gemma_path=gemma, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + dynamic_load=dynamic_load, + topology=topology, + is_fsdp=is_fsdp, + distilled_lora_path=distilled_lora, + run_warmup=False, + num_frames=num_frames, + height=height, + width=width, + ) + + if int(ttnn.distributed_context_get_rank()) != 0: + logger.info(f"Skipping generation on rank {ttnn.distributed_context_get_rank()}") + return + + output_filename = os.environ.get("OUTPUT_PATH", f"ltx_i2v_two_stages_{width}x{height}.mp4") + prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT) + pipeline.generate( + prompt, + output_path=output_filename, + images=[(image_path, 0, strength)], + num_frames=num_frames, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=int(os.environ.get("SEED", "10")), + ) + + assert os.path.exists(output_filename), f"no MP4 produced at {output_filename}" + assert os.path.getsize(output_filename) > 0, "produced MP4 is empty" + logger.info(f"I2V smoke: saved {output_filename} ({os.path.getsize(output_filename)} bytes)") + + # Optional first-frame correlation check (only if imageio + PIL are available). + try: + import imageio.v3 as iio + import numpy as np + from PIL import Image + + frames = iio.imread(output_filename, index=None, plugin="pyav") # (T, H, W, 3), in-process + first = frames[0].astype(np.float32) / 255.0 + cond = Image.open(image_path).convert("RGB").resize((width, height)) + cond_np = np.asarray(cond, dtype=np.float32) / 255.0 + a = first.reshape(-1) - first.mean() + b = cond_np.reshape(-1) - cond_np.mean() + pcc = float((a @ b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8)) + logger.info(f"I2V smoke: first-frame vs conditioning-image PCC = {pcc:.3f}") + if strength >= 1.0: + assert pcc > 0.3, f"first frame does not resemble the conditioning image (PCC={pcc:.3f})" + except ImportError: + logger.info("imageio/PIL unavailable — skipping first-frame correlation check") diff --git a/code/models/tt_dit/tests/models/ltx/test_rope_ltx.py b/code/models/tt_dit/tests/models/ltx/test_rope_ltx.py new file mode 100644 index 0000000000000000000000000000000000000000..2fa9718e8d7588efeca38973af0fcd1d1275d9f1 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_rope_ltx.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from loguru import logger + +import ttnn +from models.tt_dit.utils.check import assert_quality +from models.tt_dit.utils.mochi import get_rot_transformation_mat +from models.tt_dit.utils.tensor import bf16_tensor, bf16_tensor_2dshard + + +def make_indices_grid(F: int, H: int, W: int) -> torch.Tensor: + """3D position indices grid (1, n_dims=3, seq_len) in the layout diffusers' RoPE expects.""" + t_ids = torch.arange(F) + h_ids = torch.arange(H) + w_ids = torch.arange(W) + grid_t, grid_h, grid_w = torch.meshgrid(t_ids, h_ids, w_ids, indexing="ij") + indices_grid = torch.stack([grid_t.flatten(), grid_h.flatten(), grid_w.flatten()], dim=0).float() + return indices_grid.unsqueeze(0) # (1, n_dims=3, seq_len) + + +@pytest.mark.parametrize( + "mesh_device, sp_axis, tp_axis", + [ + [(1, 1), 0, 1], + [(1, 2), 0, 1], + [(2, 1), 0, 1], + [(2, 2), 0, 1], + [(2, 4), 0, 1], + [(2, 4), 1, 0], + ], + ids=[ + "1x1sp0tp1", + "1x2sp0tp1", + "2x1sp0tp1", + "2x2sp0tp1", + "2x4sp0tp1", + "2x4sp1tp0", + ], + indirect=["mesh_device"], +) +@pytest.mark.parametrize( + "F, H, W", + [ + (5, 16, 28), # Small: 2240 tokens + (5, 30, 52), # 480p-ish: 7800 tokens + (5, 34, 64), # 2K spatial grid (1088x2048 latent H×W): 10880 tokens + ], + ids=["small", "480p", "2k"], +) +def test_ltx_rope_interleaved(mesh_device: ttnn.MeshDevice, sp_axis: int, tp_axis: int, F: int, H: int, W: int): + """Test LTX-2 interleaved RoPE against the diffusers reference (mirrors test_wan_rotary_pos_embed).""" + # double_precision=False keeps the freq grid fp32 to match the device path. + from diffusers.models.transformers.transformer_ltx2 import ( + LTX2AudioVideoRotaryPosEmbed, + apply_interleaved_rotary_emb, + ) + + dim = 4096 + num_heads = 32 + head_dim = dim // num_heads # 128 + theta = 10000.0 + max_pos = [20, 2048, 2048] + B = 1 + + indices_grid = make_indices_grid(F, H, W) + seq_len = F * H * W + logger.info(f"indices_grid shape: {indices_grid.shape}, seq_len: {seq_len}") + + # Reference cos/sin on CPU; raw indices (ndim==3) normalize by base_{frames,height,width}. + rope = LTX2AudioVideoRotaryPosEmbed( + dim=dim, + base_num_frames=max_pos[0], + base_height=max_pos[1], + base_width=max_pos[2], + theta=theta, + modality="video", + double_precision=False, + rope_type="interleaved", + num_attention_heads=num_heads, + ) + # Interleaved cos/sin: (B, seq_len, dim) + cos_freq, sin_freq = rope(indices_grid) + logger.info(f"cos_freq shape: {cos_freq.shape}, sin_freq shape: {sin_freq.shape}") + + # Create random input on the flat (B, seq_len, dim) layout the diffusers apply expects. + torch.manual_seed(42) + input_flat = torch.randn(B, seq_len, dim, dtype=torch.float32) + + # PyTorch reference: diffusers applies interleaved RoPE on (B, N, dim). + output_ref_flat = apply_interleaved_rotary_emb(input_flat, (cos_freq, sin_freq)) + + # Reshape into (B, seq_len, num_heads, head_dim) for the device kernel + comparison. + input_tensor = input_flat.reshape(B, seq_len, num_heads, head_dim) + output_ref = output_ref_flat.reshape(B, seq_len, num_heads, head_dim) + cos_for_apply = cos_freq.reshape(B, seq_len, num_heads, head_dim) + sin_for_apply = sin_freq.reshape(B, seq_len, num_heads, head_dim) + logger.info(f"Reference output shape: {output_ref.shape}") + + # rotary_embedding_llama expects (B, num_heads, seq_len, head_dim) + input_bhnd = input_tensor.permute(0, 2, 1, 3) # (B, H, N, D) + cos_bhnd = cos_for_apply.permute(0, 2, 1, 3) # (B, H, N, D) + sin_bhnd = sin_for_apply.permute(0, 2, 1, 3) # (B, H, N, D) + + tt_input = bf16_tensor_2dshard(input_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}) + # cos/sin sharded same as input: SP axis on sequence, TP axis on heads + tt_cos = bf16_tensor_2dshard(cos_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}) + tt_sin = bf16_tensor_2dshard(sin_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + # Apply RoPE on device + tt_output = ttnn.experimental.rotary_embedding_llama(tt_input, tt_cos, tt_sin, tt_trans_mat) + + # Gather output back to host + concat_dims = [None, None] + concat_dims[sp_axis] = 2 + concat_dims[tp_axis] = 1 + tt_output_torch = ttnn.to_torch( + tt_output, + mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)), + ) + tt_output_torch = tt_output_torch.permute(0, 2, 1, 3) # Back to (B, N, H, D) + + # Compare + assert_quality(output_ref, tt_output_torch, pcc=0.99) + logger.info("PASSED: LTX interleaved RoPE matches PyTorch reference") diff --git a/code/models/tt_dit/tests/models/ltx/test_transformer_ltx.py b/code/models/tt_dit/tests/models/ltx/test_transformer_ltx.py new file mode 100644 index 0000000000000000000000000000000000000000..c811a9407c01211b4ad0ff53a002be861dfbaab6 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_transformer_ltx.py @@ -0,0 +1,2113 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""LTX-2.3 transformer unit tests; structure mirrors test_transformer_wan.py. + +Torch reference is diffusers' LTX2VideoTransformer{Block,3DModel}. The TT runtime fuses RoPE +into rotary_embedding_llama (INTERLEAVED + trans_mat) with a load-time Q/K permute, so every +TT freq feed is INTERLEAVED. Modality axis selects video vs av; set LTX_SKIP_PCC=1 to assert +shape/finiteness only. +""" + +import os +import time + +import pytest +import torch +from loguru import logger +from safetensors.torch import load_file +from tracy import signpost + +import ttnn +from models.tt_dit.models.transformers.ltx.quant_config import LtxQuantProfile +from models.tt_dit.models.transformers.ltx.rope_ltx import LTXRopeType, precompute_freqs_cis +from models.tt_dit.models.transformers.ltx.transformer_ltx import ( + LTXTransformerBlock, + LTXTransformerModel, + build_audio_masks, + build_video_pad_mask, +) +from models.tt_dit.parallel.config import DiTParallelConfig, ParallelFactor +from models.tt_dit.parallel.manager import CCLManager +from models.tt_dit.utils.check import assert_quality +from models.tt_dit.utils.mochi import get_rot_transformation_mat +from models.tt_dit.utils.patchifiers import AudioLatentShape, VideoPixelShape +from models.tt_dit.utils.tensor import bf16_tensor, bf16_tensor_2dshard +from models.tt_dit.utils.test import ( + line_params_req_exact_devices, + ring_params_4k, + ring_params_8k, + skip_if_unsupported_num_links, +) +from models.tt_dit.utils.tracing import Tracer + +from .ltx_mesh_params import LTX_TRANSFORMER_MESH_PARAMS, _4x8sp1tp0nl2_ring_is_fsdp0 + +# --------------------------------------------------------------------------- +# LTX-2.3-22B distilled transformer configuration +# --------------------------------------------------------------------------- +DIM = 4096 +NUM_HEADS = 32 +HEAD_DIM = DIM // NUM_HEADS # 128 +IN_CHANNELS = 128 +OUT_CHANNELS = 128 +CTX_DIM = 4096 +AUDIO_DIM = 2048 +AUDIO_HEAD_DIM = AUDIO_DIM // NUM_HEADS # 64 +AUDIO_IN_CHANNELS = 128 +AUDIO_CTX_DIM = 2048 +EPS = 1e-6 +PROMPT_LEN = 32 + +# RoPE +ROPE_THETA = 10000.0 +VIDEO_ROPE_MAX_POS = [20, 2048, 2048] +AUDIO_ROPE_MAX_POS = [20] +CROSS_PE_MAX_POS = [20] + +# AV-mode auxiliary +AUDIO_N = 256 # sp-aligned for sp ∈ {1, 2, 4, 8} +# 22B checkpoint variants (AV mode only): "fast" = distilled-1.1 (production), "dev" = base 22B. +_CHECKPOINT_22B_VARIANTS = { + "fast": os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-distilled-1.1.safetensors"), + "dev": os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-dev.safetensors"), +} + + +# Fallback: the same checkpoints as cached by huggingface_hub (Lightricks/LTX-2.3 repo). +_CHECKPOINT_22B_HF_FILENAMES = { + "fast": "ltx-2.3-22b-distilled-1.1.safetensors", + "dev": "ltx-2.3-22b-dev.safetensors", +} + + +def _resolve_checkpoint_22b(variant: str) -> str: + """Path for the 22B variant: LTX_CHECKPOINT env, then local cache, then huggingface_hub snapshot.""" + env = os.environ.get("LTX_CHECKPOINT") + if env: + return env + local = _CHECKPOINT_22B_VARIANTS[variant] + if os.path.exists(local): + return local + snap_root = os.path.expanduser("~/.cache/huggingface/hub/models--Lightricks--LTX-2.3/snapshots") + if os.path.isdir(snap_root): + for snap in sorted(os.listdir(snap_root)): + cand = os.path.join(snap_root, snap, _CHECKPOINT_22B_HF_FILENAMES[variant]) + if os.path.exists(cand): + return cand + return local + + +# Block-test reference uses small scaled random weights; default init blows up the adaln chains. +WEIGHT_SEED = 1234 +INPUT_SEED = 42 +TIMESTEP_VAL = 0.01 + +# Toggle PCC verification via env (mirrors Wan's `dit_unit_test`). Default ON. +_RUN_PCC_DEFAULT = {"1": False, "0": True}.get(os.environ.get("LTX_SKIP_PCC"), True) + + +# --------------------------------------------------------------------------- +# Parametrize lists +# --------------------------------------------------------------------------- + +# 1080p fast-pipeline real latent grids (F, H_lat, W_lat) — the exact (latent_frames, h//32, +# w//32) production feeds RoPE. These are NOT tile/SP-aligned, so the test SP-pads the sequence +# (round F·H·W up to TILE * sp_factor) exactly like LTXPipeline._sp_pad_len, then crops back. +_LTX_TRANSFORMER_SHAPE_PARAMS = [ + pytest.param(19, 17, 30, id="stage_1"), # F·H·W = 9690 (real); SP-padded on device + pytest.param(19, 34, 60, id="stage_2"), # F·H·W = 38760 (real); SP-padded on device +] + +_LTX_TRANSFORMER_MODALITY_PARAMS = [ + pytest.param(False, id="video"), + pytest.param(True, id="av"), +] + +_LTX_TRANSFORMER_RUN_PCC_PARAMS = [pytest.param(_RUN_PCC_DEFAULT, id="pcc" if _RUN_PCC_DEFAULT else "nopcc")] + +# Checkpoint variant for AV mode ("fast" = production default); ignored in video mode. +_LTX_TRANSFORMER_CKPT_PARAMS = [ + pytest.param("fast", id="ckpt_fast"), + pytest.param("dev", id="ckpt_dev"), +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _interleaved_to_bhnd(t: torch.Tensor, num_heads: int) -> torch.Tensor: + """(B, N, dim) INTERLEAVED freqs → (B, num_heads, N, head_dim) for rotary_embedding_llama.""" + B, N, dim = t.shape + head_dim = dim // num_heads + return t.reshape(B, N, num_heads, head_dim).permute(0, 2, 1, 3).contiguous() + + +def _sp_pad_len(n_real: int, sp_factor: int) -> int: + """Round a seq len up to ttnn.TILE_SIZE * sp_factor (mirrors LTXPipeline._sp_pad_len).""" + divisor = ttnn.TILE_SIZE * sp_factor + return ((n_real + divisor - 1) // divisor) * divisor + + +def _pad_seq_dim(t: torch.Tensor, n_pad: int, dim: int = -2) -> torch.Tensor: + """Zero-pad a torch tensor's sequence dim up to n_pad (no-op if already ≥ n_pad).""" + n = t.shape[dim] + if n_pad <= n: + return t + shape = list(t.shape) + shape[dim] = n_pad + out = torch.zeros(shape, dtype=t.dtype) + idx = [slice(None)] * t.ndim + idx[dim] = slice(0, n) + out[tuple(idx)] = t + return out + + +def _pad_rope_bhnd(cos_i: torch.Tensor, sin_i: torch.Tensor, pad_to: int | None): + """Right-pad BHND INTERLEAVED freqs on the seq dim (2) with identity rotation (cos=1, sin=0). + + Same convention as rope_ltx.pad_video_rope_sp: padded slots are no-op rotations; SDPA + still masks them out via logical_n / padding masks. + """ + if pad_to is None or pad_to <= cos_i.shape[2]: + return cos_i, sin_i + pad = pad_to - cos_i.shape[2] + cos_i = torch.nn.functional.pad(cos_i, (0, 0, 0, pad), value=1.0) + sin_i = torch.nn.functional.pad(sin_i, (0, 0, 0, pad), value=0.0) + return cos_i, sin_i + + +def _diffusers_qk_to_split(t: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor: + """diffusers (interleaved-rotation) Q/K → Lightricks SPLIT convention the TT loader expects.""" + D = head_dim + D_half = D // 2 + inv = torch.empty(D, dtype=torch.long) + inv[:D_half] = torch.arange(0, D, 2) + inv[D_half:] = torch.arange(1, D, 2) + rest = t.shape[1:] + return t.reshape(num_heads, D, *rest).index_select(1, inv).reshape(num_heads * D, *rest) + + +def _convert_diffusers_video_block_to_tt(state: dict, *, num_heads: int, head_dim: int) -> dict: + """diffusers LTX2VideoTransformerBlock state_dict → TT video block loader input (video subset, split Q/K).""" + keep_exact = {"scale_shift_table", "prompt_scale_shift_table"} + out = {k: v for k, v in state.items() if k in keep_exact or k.startswith(("attn1.", "attn2.", "ff."))} + for base in ("attn1", "attn2"): + for proj in ("to_q", "to_k"): + for suf in ("weight", "bias"): + kk = f"{base}.{proj}.{suf}" + if kk in out: + out[kk] = _diffusers_qk_to_split(out[kk], num_heads, head_dim) + for nk in ("norm_q.weight", "norm_k.weight"): + kk = f"{base}.{nk}" + if kk in out: + out[kk] = _diffusers_qk_to_split(out[kk], num_heads, head_dim) + return out + + +def _make_diffusers_video_block(): + """Build the diffusers LTX-2 block in the video config used as the PCC reference.""" + from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformerBlock + + return LTX2VideoTransformerBlock( + dim=DIM, + num_attention_heads=NUM_HEADS, + attention_head_dim=HEAD_DIM, + cross_attention_dim=CTX_DIM, + audio_dim=AUDIO_DIM, + audio_num_attention_heads=NUM_HEADS, + audio_attention_head_dim=AUDIO_HEAD_DIM, + audio_cross_attention_dim=AUDIO_CTX_DIM, + video_gated_attn=False, + video_cross_attn_adaln=True, + audio_gated_attn=False, + audio_cross_attn_adaln=False, + eps=EPS, + rope_type="interleaved", + ) + + +def _diffusers_video_block_ref(block, *, x, context, temb, prompt_temb, cos_i, sin_i): + """Run the diffusers block video-only (a2v/v2a disabled) and return the video output.""" + B, N, _ = x.shape + a_N = 32 + audio_hidden = torch.zeros(B, a_N, AUDIO_DIM) + audio_enc = torch.zeros(B, context.shape[1], AUDIO_CTX_DIM) + temb_audio = torch.zeros(B, 1, 6 * AUDIO_DIM) # audio_cross_attn_adaln=False → 6 mod params + temb_prompt_audio = torch.zeros(B, 1, 2 * AUDIO_DIM) + with torch.no_grad(): + out_v, _ = block( + hidden_states=x, + audio_hidden_states=audio_hidden, + encoder_hidden_states=context, + audio_encoder_hidden_states=audio_enc, + temb=temb, + temb_audio=temb_audio, + temb_ca_scale_shift=None, + temb_ca_audio_scale_shift=None, + temb_ca_gate=None, + temb_ca_audio_gate=None, + temb_prompt=prompt_temb, + temb_prompt_audio=temb_prompt_audio, + video_rotary_emb=(cos_i, sin_i), + audio_rotary_emb=None, + use_a2v_cross_attention=False, + use_v2a_cross_attention=False, + ) + return out_v + + +def _make_diffusers_video_model(num_layers: int = 1): + """diffusers LTX2VideoTransformer3DModel configured to match the TT video model.""" + from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel + + return LTX2VideoTransformer3DModel( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_attention_heads=NUM_HEADS, + attention_head_dim=HEAD_DIM, + cross_attention_dim=CTX_DIM, + audio_in_channels=AUDIO_IN_CHANNELS, + audio_out_channels=AUDIO_IN_CHANNELS, + audio_num_attention_heads=NUM_HEADS, + audio_attention_head_dim=AUDIO_HEAD_DIM, + audio_cross_attention_dim=AUDIO_CTX_DIM, + num_layers=num_layers, + cross_attn_mod=True, + audio_cross_attn_mod=False, + norm_eps=EPS, + rope_theta=10000.0, + rope_double_precision=False, + use_prompt_embeddings=False, + rope_type="interleaved", + ) + + +def _convert_diffusers_video_model_to_tt(state: dict, *, num_heads: int, head_dim: int) -> dict: + """diffusers LTX2VideoTransformer3DModel state_dict → TT LTXTransformerModel loader input.""" + drop_top = ("audio_", "av_cross_attn_", "caption_projection") + rename_top = {"proj_in.": "patchify_proj.", "time_embed.": "adaln_single.", "prompt_adaln.": "prompt_adaln_single."} + out: dict = {} + for k, v in state.items(): + if k.startswith("transformer_blocks."): + _, _idx, rest = k.split(".", 2) + if rest in ("scale_shift_table", "prompt_scale_shift_table") or rest.startswith( + ("attn1.", "attn2.", "ff.") + ): + out[k] = v + continue + if k.startswith(drop_top): + continue + nk = k + for a, b in rename_top.items(): + if nk.startswith(a): + nk = b + nk[len(a) :] + break + out[nk] = v + qk_suffixes = ("to_q.weight", "to_q.bias", "to_k.weight", "to_k.bias", "norm_q.weight", "norm_k.weight") + for k in list(out.keys()): + if (".attn1." in k or ".attn2." in k) and k.endswith(qk_suffixes): + out[k] = _diffusers_qk_to_split(out[k], num_heads, head_dim) + return out + + +def _diffusers_video_model_ref(model, *, video_lat, video_prompt, sigma_val, F, H, W): + """Run the diffusers 3D model video-only (isolate_modalities=True) and return video output.""" + B, N, _ = video_lat.shape + gt, gh, gw, _ = _video_grid(F, H, W) + video_coords = torch.stack([gt, gh, gw], dim=0).float().unsqueeze(0) # (1, 3, N) + a_N = 64 + audio_lat = torch.zeros(B, a_N, AUDIO_IN_CHANNELS) + audio_prompt = torch.zeros(B, video_prompt.shape[1], AUDIO_CTX_DIM) + audio_coords = torch.arange(a_N).reshape(1, 1, a_N).float() + ts = torch.tensor([sigma_val * 1000.0]) + with torch.no_grad(): + out = model( + hidden_states=video_lat, + audio_hidden_states=audio_lat, + encoder_hidden_states=video_prompt, + audio_encoder_hidden_states=audio_prompt, + timestep=ts, + audio_timestep=ts, + sigma=ts, + audio_sigma=ts, + video_coords=video_coords, + audio_coords=audio_coords, + isolate_modalities=True, + return_dict=False, + ) + return out[0] if isinstance(out, (tuple, list)) else out + + +def _diffusers_video_model_ref_pertoken(model, *, video_lat, video_prompt, video_ts_real, sigma_scalar, F, H, W): + """Per-token oracle: same as ``_diffusers_video_model_ref`` but ``timestep`` is ``(1, N)``. + + The diffusers LTX2 model flattens ``timestep`` through ``time_embed`` then reshapes to + ``(B, N, ...)`` modulation, so a per-token timestep yields genuine per-token modulation. + ``sigma`` (prompt/cross modulation) stays scalar, matching the TT scalar-timestep path. + """ + B, N, _ = video_lat.shape + gt, gh, gw, _ = _video_grid(F, H, W) + video_coords = torch.stack([gt, gh, gw], dim=0).float().unsqueeze(0) # (1, 3, N) + a_N = 64 + audio_lat = torch.zeros(B, a_N, AUDIO_IN_CHANNELS) + audio_prompt = torch.zeros(B, video_prompt.shape[1], AUDIO_CTX_DIM) + audio_coords = torch.arange(a_N).reshape(1, 1, a_N).float() + ts_pertoken = video_ts_real.reshape(1, -1) * 1000.0 # (1, N) + ts_scalar = torch.tensor([sigma_scalar * 1000.0]) + with torch.no_grad(): + out = model( + hidden_states=video_lat, + audio_hidden_states=audio_lat, + encoder_hidden_states=video_prompt, + audio_encoder_hidden_states=audio_prompt, + timestep=ts_pertoken, + audio_timestep=ts_scalar, + sigma=ts_scalar, + audio_sigma=ts_scalar, + video_coords=video_coords, + audio_coords=audio_coords, + isolate_modalities=True, + return_dict=False, + ) + return out[0] if isinstance(out, (tuple, list)) else out + + +def _lightricks_qk_to_interleaved(t: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor: + """Forward per-head SPLIT→INTERLEAVED Q/K permute (same as LTXAttention._permute_qk).""" + D = head_dim + D_half = D // 2 + perm = torch.empty(D, dtype=torch.long) + perm[0::2] = torch.arange(D_half) + perm[1::2] = torch.arange(D_half, D) + rest = t.shape[1:] + return t.reshape(num_heads, D, *rest).index_select(1, perm).reshape(num_heads * D, *rest) + + +# Per-attention head_dim for the AV block (video self/text-cross use 128; audio + cross-modal use 64). +_AV_ATTN_HEAD_DIM = { + "attn1": HEAD_DIM, + "attn2": HEAD_DIM, + "audio_attn1": AUDIO_HEAD_DIM, + "audio_attn2": AUDIO_HEAD_DIM, + "audio_to_video_attn": AUDIO_HEAD_DIM, + "video_to_audio_attn": AUDIO_HEAD_DIM, +} + + +def _convert_lightricks_av_to_diffusers(state: dict, *, num_heads: int) -> dict: + """Raw Lightricks AV state_dict → diffusers LTX2VideoTransformer3DModel (AV) names.""" + rename_top = { + "patchify_proj.": "proj_in.", + "adaln_single.": "time_embed.", + "prompt_adaln_single.": "prompt_adaln.", + "audio_patchify_proj.": "audio_proj_in.", + "audio_adaln_single.": "audio_time_embed.", + "audio_prompt_adaln_single.": "audio_prompt_adaln.", + "av_ca_video_scale_shift_adaln_single.": "av_cross_attn_video_scale_shift.", + "av_ca_audio_scale_shift_adaln_single.": "av_cross_attn_audio_scale_shift.", + "av_ca_a2v_gate_adaln_single.": "av_cross_attn_video_a2v_gate.", + "av_ca_v2a_gate_adaln_single.": "av_cross_attn_audio_v2a_gate.", + } + block_table_rename = { + "scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table", + "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table", + } + out: dict = {} + for k, v in state.items(): + if k.startswith("transformer_blocks."): + prefix, idx, rest = k.split(".", 2) + if rest in block_table_rename: + rest = block_table_rename[rest] + rest = rest.replace("q_norm", "norm_q").replace("k_norm", "norm_k") + out[f"{prefix}.{idx}.{rest}"] = v + else: + nk = k + for a, b in rename_top.items(): + if nk.startswith(a): + nk = b + nk[len(a) :] + break + out[nk] = v + qk_suffixes = ("to_q.weight", "to_q.bias", "to_k.weight", "to_k.bias", "norm_q.weight", "norm_k.weight") + for k in list(out.keys()): + if not k.startswith("transformer_blocks."): + continue + attn = k.split(".")[2] + if attn in _AV_ATTN_HEAD_DIM and k.endswith(qk_suffixes): + out[k] = _lightricks_qk_to_interleaved(out[k], num_heads, _AV_ATTN_HEAD_DIM[attn]) + return out + + +def _make_diffusers_av_model(num_layers: int = 1): + """diffusers ``LTX2VideoTransformer3DModel`` configured to match the TT AV model (gated + attention, 9-param video+audio modulation, cross-modal a2v/v2a, no caption projection).""" + from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel + + return LTX2VideoTransformer3DModel( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_attention_heads=NUM_HEADS, + attention_head_dim=HEAD_DIM, + cross_attention_dim=CTX_DIM, + audio_in_channels=AUDIO_IN_CHANNELS, + audio_out_channels=AUDIO_IN_CHANNELS, + audio_num_attention_heads=NUM_HEADS, + audio_attention_head_dim=AUDIO_HEAD_DIM, + audio_cross_attention_dim=AUDIO_CTX_DIM, + num_layers=num_layers, + gated_attn=True, + audio_gated_attn=True, + cross_attn_mod=True, + audio_cross_attn_mod=True, + norm_eps=EPS, + rope_theta=ROPE_THETA, + rope_double_precision=False, + use_prompt_embeddings=False, + rope_type="interleaved", + # =1 matches Lightricks' av_ca_timestep_scale_multiplier (a2v/v2a gate adaln sees raw sigma). + cross_attn_timestep_scale_multiplier=1, + ) + + +def _diffusers_av_model_ref(model, *, video_lat, video_prompt, audio_lat, audio_prompt, sigma_val, F, H, W, audio_N): + """Run the diffusers AV model (cross-modal enabled) and return (video, audio) outputs.""" + gt, gh, gw, _ = _video_grid(F, H, W) + video_coords = torch.stack([gt, gh, gw], dim=0).float().unsqueeze(0) # (1, 3, N) + audio_coords = torch.arange(audio_N).reshape(1, 1, audio_N).float() + ts = torch.tensor([sigma_val * 1000.0]) + with torch.no_grad(): + out = model( + hidden_states=video_lat, + audio_hidden_states=audio_lat, + encoder_hidden_states=video_prompt, + audio_encoder_hidden_states=audio_prompt, + timestep=ts, + audio_timestep=ts, + sigma=ts, + audio_sigma=ts, + video_coords=video_coords, + audio_coords=audio_coords, + isolate_modalities=False, + return_dict=False, + ) + return out[0], out[1] + + +def _make_parallel_config(mesh_device, sp_axis, tp_axis): + return DiTParallelConfig( + cfg_parallel=ParallelFactor(factor=1, mesh_axis=0), + sequence_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[sp_axis], mesh_axis=sp_axis), + tensor_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[tp_axis], mesh_axis=tp_axis), + ) + + +def _make_ccl_manager(mesh_device, num_links, topology): + return CCLManager(mesh_device=mesh_device, num_links=num_links, topology=topology) + + +def _video_grid(F, H, W): + """Return (grid_t, grid_h, grid_w) flat tensors and stacked 3D positions (1, N, 3).""" + t_ids, h_ids, w_ids = torch.arange(F), torch.arange(H), torch.arange(W) + gt, gh, gw = torch.meshgrid(t_ids, h_ids, w_ids, indexing="ij") + gt_f, gh_f, gw_f = gt.flatten(), gh.flatten(), gw.flatten() + positions = torch.stack([gt_f, gh_f, gw_f], dim=-1).float().unsqueeze(0) + return gt_f, gh_f, gw_f, positions + + +def _video_rope_freqs(F, H, W, rope_type=LTXRopeType.SPLIT): + _, _, _, positions = _video_grid(F, H, W) + return precompute_freqs_cis( + positions, + dim=DIM, + out_dtype=torch.float32, + theta=ROPE_THETA, + max_pos=VIDEO_ROPE_MAX_POS, + num_attention_heads=NUM_HEADS, + rope_type=rope_type, + ) + + +def _video_cross_pe_freqs(F, H, W, rope_type=LTXRopeType.SPLIT): + """T-only cross-PE for V↔A, dim matches audio (AUDIO_DIM).""" + _, _, _, positions = _video_grid(F, H, W) + return precompute_freqs_cis( + positions[..., 0:1], + dim=AUDIO_DIM, + out_dtype=torch.float32, + theta=ROPE_THETA, + max_pos=CROSS_PE_MAX_POS, + num_attention_heads=NUM_HEADS, + rope_type=rope_type, + ) + + +def _audio_rope_freqs(audio_N, rope_type=LTXRopeType.SPLIT): + a_pos = torch.arange(audio_N).float().unsqueeze(0).unsqueeze(-1) + return precompute_freqs_cis( + a_pos, + dim=AUDIO_DIM, + out_dtype=torch.float32, + theta=ROPE_THETA, + max_pos=AUDIO_ROPE_MAX_POS, + num_attention_heads=NUM_HEADS, + rope_type=rope_type, + ) + + +def _audio_cross_pe_freqs(audio_N, rope_type=LTXRopeType.SPLIT): + a_pos = torch.arange(audio_N).float().unsqueeze(0).unsqueeze(-1) + return precompute_freqs_cis( + a_pos, + dim=AUDIO_DIM, + out_dtype=torch.float32, + theta=ROPE_THETA, + max_pos=CROSS_PE_MAX_POS, + num_attention_heads=NUM_HEADS, + rope_type=rope_type, + ) + + +def _audio_seq_lens(F: int, sp_factor: int) -> tuple[int, int]: + """Fast-pipeline audio seq lengths from the latent frame count (audio_N_real, padded to 32*sp).""" + num_frames = (F - 1) * 8 + 1 + vps = VideoPixelShape(batch=1, frames=num_frames, height=64, width=64, fps=24) + audio_N_real = AudioLatentShape.from_video_pixel_shape(vps).frames + audio_N = ((audio_N_real + 32 * sp_factor - 1) // (32 * sp_factor)) * (32 * sp_factor) + return audio_N, audio_N_real + + +def _tt_rope(freqs_fn, *args, mesh_device, sp_axis, tp_axis, pad_to=None): + """Build INTERLEAVED freqs for the TT runtime: precompute → BHND reshape → SP-pad → 2D shard.""" + cos_i, sin_i = freqs_fn(*args, rope_type=LTXRopeType.INTERLEAVED) + cos_i = _interleaved_to_bhnd(cos_i, NUM_HEADS) + sin_i = _interleaved_to_bhnd(sin_i, NUM_HEADS) + cos_i, sin_i = _pad_rope_bhnd(cos_i, sin_i, pad_to) + tt_cos = bf16_tensor_2dshard(cos_i, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}) + tt_sin = bf16_tensor_2dshard(sin_i, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}) + return tt_cos, tt_sin + + +def _tt_rope_full(freqs_fn, *args, mesh_device, tp_axis, pad_to=None): + """INTERLEAVED freqs replicated full-seq (cross-attn K side): BHND reshape, SP-pad, no SP shard.""" + cos_i, sin_i = freqs_fn(*args, rope_type=LTXRopeType.INTERLEAVED) + cos_i = _interleaved_to_bhnd(cos_i, NUM_HEADS) + sin_i = _interleaved_to_bhnd(sin_i, NUM_HEADS) + cos_i, sin_i = _pad_rope_bhnd(cos_i, sin_i, pad_to) + tt_cos = bf16_tensor(cos_i, device=mesh_device, mesh_axis=tp_axis, shard_dim=1) + tt_sin = bf16_tensor(sin_i, device=mesh_device, mesh_axis=tp_axis, shard_dim=1) + return tt_cos, tt_sin + + +def _make_tt_block(*, mesh_device, ccl_manager, parallel_config, is_fsdp, has_audio, quant_config=None): + return LTXTransformerBlock( + video_dim=DIM, + video_ffn_dim=DIM * 4, + video_num_heads=NUM_HEADS, + video_cross_attention_dim=CTX_DIM, + audio_dim=AUDIO_DIM, + audio_ffn_dim=AUDIO_DIM * 4, + audio_num_heads=NUM_HEADS, + audio_cross_attention_dim=AUDIO_CTX_DIM, + eps=EPS, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=has_audio, + apply_gated_attention=has_audio, + cross_attention_adaln=True, + quant_config=quant_config, + ) + + +def _make_tt_model( + *, mesh_device, ccl_manager, parallel_config, is_fsdp, has_audio, num_layers, image_conditioning=False +): + return LTXTransformerModel( + num_attention_heads=NUM_HEADS, + attention_head_dim=HEAD_DIM, + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_layers=num_layers, + cross_attention_dim=CTX_DIM, + audio_num_attention_heads=NUM_HEADS, + audio_attention_head_dim=AUDIO_HEAD_DIM, + audio_in_channels=AUDIO_IN_CHANNELS, + audio_out_channels=AUDIO_IN_CHANNELS, + audio_cross_attention_dim=AUDIO_CTX_DIM, + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=has_audio, + apply_gated_attention=has_audio, + cross_attention_adaln=True, + image_conditioning=image_conditioning, + ) + + +def _load_22b_state_dict(num_layers: int, checkpoint_path: str) -> dict | None: + """Load LTX-2.3-22B weights filtered to the first num_layers blocks; None if checkpoint missing.""" + if not os.path.exists(checkpoint_path): + return None + raw = load_file(checkpoint_path) + prefix = "model.diffusion_model." + sd = {k[len(prefix) :]: v for k, v in raw.items() if k.startswith(prefix)} + return {k: v for k, v in sd.items() if not k.startswith("transformer_blocks.") or int(k.split(".")[1]) < num_layers} + + +def _scale_init_(module: torch.nn.Module, seed: int = WEIGHT_SEED) -> None: + """Deterministic N(0, 0.1²) reinit in-place. See block test comment for why.""" + torch.manual_seed(seed) + with torch.no_grad(): + for p in module.parameters(): + p.copy_(torch.randn(p.shape, dtype=p.dtype) * 0.1) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + LTX_TRANSFORMER_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize(("F", "H", "W"), _LTX_TRANSFORMER_SHAPE_PARAMS) +@pytest.mark.parametrize("has_audio", _LTX_TRANSFORMER_MODALITY_PARAMS) +@pytest.mark.parametrize("run_pcc", _LTX_TRANSFORMER_RUN_PCC_PARAMS) +@pytest.mark.parametrize("checkpoint_variant", _LTX_TRANSFORMER_CKPT_PARAMS) +def test_ltx_transformer_block( + mesh_device: ttnn.MeshDevice, + sp_axis: int, + tp_axis: int, + num_links: int, + topology: ttnn.Topology, + is_fsdp: bool, + F: int, + H: int, + W: int, + has_audio: bool, + run_pcc: bool, + checkpoint_variant: str, + reset_seeds, +) -> None: + """Test LTXTransformerBlock: TT forward, with optional PCC vs the diffusers LTX-2 block.""" + # Checkpoint variant only affects AV weight loading; skip the redundant copy in video mode. + skip_if_unsupported_num_links(mesh_device, num_links) + if not has_audio and checkpoint_variant != "fast": + pytest.skip("checkpoint_variant only affects AV mode (video uses random scaled weights)") + checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant) + sp_factor = tuple(mesh_device.shape)[sp_axis] + # Real latent grid → SP-padded sequence (mirrors LTXPipeline). video_N is the padded device + # length; video_N_real is the logical token count fed to SDPA's logical_n / padding masks. + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + # Audio seq length mirrors the fast pipeline so op sizes / padding match production. + audio_N, audio_N_real = _audio_seq_lens(F, sp_factor) + assert video_N % (32 * sp_factor) == 0, f"video_N={video_N} not sp/tile-aligned for sp={sp_factor}" + assert audio_N % (32 * sp_factor) == 0, f"audio_N={audio_N} not sp/tile-aligned for sp={sp_factor}" + + # AV mode lacks the LTXModel-level cross-PE plumbing, so it asserts shape/finiteness only. + do_pcc = run_pcc and not has_audio + torch_block = None + torch_out = None + if do_pcc: + # Reference is the diffusers LTX-2 block run video-only (audio + a2v/v2a disabled). + torch_block = _make_diffusers_video_block() + torch_block.eval() + _scale_init_(torch_block) + + # TT block + ccl_manager = _make_ccl_manager(mesh_device, num_links, topology) + parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis) + tt_block = _make_tt_block( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=has_audio, + ) + + if do_pcc: + # diffusers→TT: keep video subset, convert attn Q/K to split. Clone to avoid aliasing. + conv = _convert_diffusers_video_block_to_tt(torch_block.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM) + tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()}) + elif has_audio: + sd = _load_22b_state_dict(num_layers=1, checkpoint_path=checkpoint_22b) + if sd is None: + pytest.skip(f"22B checkpoint not found at {checkpoint_22b}") + block_sd = { + k[len("transformer_blocks.0.") :]: v for k, v in sd.items() if k.startswith("transformer_blocks.0.") + } + tt_block.load_torch_state_dict(block_sd, strict=False) + else: + # video mode, PCC disabled — harvest a shape-matching state dict from a throwaway block. + dummy = _make_diffusers_video_block() + _scale_init_(dummy) + conv = _convert_diffusers_video_block_to_tt(dummy.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM) + tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()}) + del dummy + + # Inputs (real-grid token count; the TT side gets a zero-padded copy below). + torch.manual_seed(INPUT_SEED) + x = torch.randn(1, video_N_real, DIM, dtype=torch.float32) + context = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + temb = torch.randn(1, 1, 9 * DIM, dtype=torch.float32) # 9 adaln params + prompt_temb = torch.randn(1, 1, 2 * DIM, dtype=torch.float32) # 2 adaln params for prompt + + # Torch forward (video-only branch): diffusers block consumes INTERLEAVED cos/sin, same + # convention as the TT runtime (tt_dit interleaved freqs ≈ diffusers rope, PCC ~1.0). + if do_pcc: + cos_int, sin_int = _video_rope_freqs(F, H, W, rope_type=LTXRopeType.INTERLEAVED) + torch_out = _diffusers_video_block_ref( + torch_block, x=x, context=context, temb=temb, prompt_temb=prompt_temb, cos_i=cos_int, sin_i=sin_int + ) + logger.info(f"torch block output {tuple(torch_out.shape)}") + + # TT video tensors. Runtime rope is INTERLEAVED + trans_mat. Sequence is SP-padded to video_N; + # rope pads with identity rotation; video_N (logical) goes to the block for SDPA masking. + spatial = _pad_seq_dim(x, video_N, dim=1).unsqueeze(0) + tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3}) + tt_prompt = bf16_tensor(context.unsqueeze(0), device=mesh_device) + # Outer-param layout (coeff, B=1, 1, D) so the block's chunk(dim=0) is a free tile-aligned slice. + tt_temb = bf16_tensor( + temb.reshape(9, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3 + ) + tt_prompt_temb = bf16_tensor(prompt_temb.reshape(2, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device) + tt_cos, tt_sin = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + forward_kwargs = dict( + video_1BND=tt_spatial, + video_prompt=tt_prompt, + video_temb=tt_temb, + video_N=video_N_real, + video_rope_cos=tt_cos, + video_rope_sin=tt_sin, + trans_mat=tt_trans_mat, + video_prompt_temb=tt_prompt_temb, + ) + if has_audio: + # Real tokens in [:audio_N_real], zeros in the padded tail (matches the padded audio latent). + a_x = torch.zeros(1, audio_N, AUDIO_DIM, dtype=torch.float32) + a_x[:, :audio_N_real, :] = torch.randn(1, audio_N_real, AUDIO_DIM, dtype=torch.float32) + a_ctx = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32) + a_temb = torch.randn(1, 1, 9 * AUDIO_DIM, dtype=torch.float32) + a_prompt_temb = torch.randn(1, 1, 2 * AUDIO_DIM, dtype=torch.float32) + av_ca_v = torch.randn(1, 1, 5 * DIM, dtype=torch.float32) # 4 scale-shift + 1 gate + av_ca_a = torch.randn(1, 1, 5 * AUDIO_DIM, dtype=torch.float32) + + a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis) + vx_cos, vx_sin = _tt_rope( + _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + ax_cos, ax_sin = _tt_rope( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis + ) + ax_cos_full, ax_sin_full = _tt_rope_full( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis + ) + + # Padding masks, same construction as the fast pipeline (audio padded, video aligned). + a_attn_mask, a_pad_sp, a_pad_full = build_audio_masks( + audio_N, audio_N_real, mesh_device=mesh_device, sp_axis=sp_axis + ) + v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis) + + forward_kwargs.update( + audio_1BND=bf16_tensor_2dshard( + a_x.unsqueeze(0), device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3} + ), + audio_prompt=bf16_tensor(a_ctx.unsqueeze(0), device=mesh_device), + audio_temb=bf16_tensor( + a_temb.reshape(9, AUDIO_DIM).unsqueeze(1).unsqueeze(1), + device=mesh_device, + mesh_axis=tp_axis, + shard_dim=3, + ), + audio_prompt_temb=bf16_tensor( + a_prompt_temb.reshape(2, AUDIO_DIM).unsqueeze(1).unsqueeze(1), device=mesh_device + ), + av_ca_temb=bf16_tensor( + av_ca_v.reshape(5, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3 + ), + av_ca_audio_temb=bf16_tensor( + av_ca_a.reshape(5, AUDIO_DIM).unsqueeze(1).unsqueeze(1), + device=mesh_device, + mesh_axis=tp_axis, + shard_dim=3, + ), + audio_N=audio_N, + audio_rope_cos=a_cos, + audio_rope_sin=a_sin, + video_cross_pe_cos=vx_cos, + video_cross_pe_sin=vx_sin, + audio_cross_pe_cos=ax_cos, + audio_cross_pe_sin=ax_sin, + audio_cross_pe_cos_full=ax_cos_full, + audio_cross_pe_sin_full=ax_sin_full, + audio_attn_mask=a_attn_mask, + audio_padding_mask=a_pad_sp, + audio_padding_mask_full=a_pad_full, + video_padding_mask=v_pad_sp, + ) + + # Signpost-bracket one warm forward for tt-perf-report --start/end-signpost. + for _ in range(2): + tt_out = tt_block(**forward_kwargs) + ttnn.synchronize_device(mesh_device) + signpost("start") + tt_out = tt_block(**forward_kwargs) + ttnn.synchronize_device(mesh_device) + signpost("stop") + + if has_audio: + tt_v, tt_a = tt_out + else: + tt_v = tt_out + + # Gather video output + concat_dims = [None, None] + concat_dims[sp_axis] = 2 + concat_dims[tp_axis] = 3 + tt_v_torch = ttnn.to_torch( + tt_v, + mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)), + ).squeeze(0) + + # Crop the SP-padding tail off the gathered video before checking/comparing real tokens. + tt_v_torch = tt_v_torch[:, :video_N_real, :] + assert tt_v_torch.shape == (1, video_N_real, DIM), f"video shape {tt_v_torch.shape}" + assert torch.isfinite(tt_v_torch).all(), "video output NaN/Inf" + + if has_audio: + tt_a_torch = ttnn.to_torch( + tt_a, + mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)), + ).squeeze(0)[:, :audio_N_real, :] + assert tt_a_torch.shape == (1, audio_N_real, AUDIO_DIM), f"audio shape {tt_a_torch.shape}" + assert torch.isfinite(tt_a_torch).all(), "audio output NaN/Inf" + + if do_pcc: + # 4x8 mesh has 8-way SP ring all-gathers — looser tolerance. + pcc = 0.988 if mesh_device.get_num_devices() > 8 else 0.999 + rmse = 0.10 if mesh_device.get_num_devices() > 8 else 0.032 + assert_quality(torch_out, tt_v_torch, pcc=pcc, relative_rmse=rmse) + logger.info(f"PASSED block PCC: video {tuple(tt_v_torch.shape)}") + else: + logger.info(f"PASSED block (no PCC): video {tuple(tt_v_torch.shape)}") + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + [_4x8sp1tp0nl2_ring_is_fsdp0], + indirect=["mesh_device", "device_params"], +) +def test_ltx_transformer_block_bf8( + mesh_device: ttnn.MeshDevice, + sp_axis: int, + tp_axis: int, + num_links: int, + topology: ttnn.Topology, + is_fsdp: bool, + reset_seeds, +) -> None: + """Shipped bf8 DiT-linear quant path (``all_bf8_lofi``) on the production 4x8 ring. + + Covers two things a bf16-only run cannot: (1) the profile actually produced bf8 device weights with + the ``to_out`` carve-out kept bf16 — a mis-wired role or a silent no-quant fails the dtype asserts; + (2) the bf8 forward still tracks the fp32 diffusers oracle. Video-only, matching the block test's + PCC path; the gate-role carve-out only exists under gated (audio) attention, so it is out of scope. + """ + skip_if_unsupported_num_links(mesh_device, num_links) + F, H, W = 19, 17, 30 # stage_1 latent grid + sp_factor = tuple(mesh_device.shape)[sp_axis] + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + + torch_block = _make_diffusers_video_block() + torch_block.eval() + _scale_init_(torch_block) + + ccl_manager = _make_ccl_manager(mesh_device, num_links, topology) + parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis) + tt_block = _make_tt_block( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=False, + quant_config=LtxQuantProfile.all_bf8_lofi(), + ) + conv = _convert_diffusers_video_block_to_tt(torch_block.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM) + tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()}) + + # Assert the loaded device-tensor dtypes the profile is supposed to map: projections bf8, the + # fused-addcmul to_out carve-out bf16. A wrong role or a silent no-quant (bf16 everywhere) fails here. + def wdtype(linear): + return linear.weight.data.dtype + + assert wdtype(tt_block.attn1.to_qkv) == ttnn.bfloat8_b, "self-attn qkv should be bf8" + assert wdtype(tt_block.attn1.to_out) == ttnn.bfloat16, "self-attn to_out carve-out should stay bf16" + assert wdtype(tt_block.attn2.to_q) == ttnn.bfloat8_b, "cross-attn q should be bf8" + assert wdtype(tt_block.attn2.to_kv) == ttnn.bfloat8_b, "cross-attn kv should be bf8" + assert wdtype(tt_block.attn2.to_out) == ttnn.bfloat16, "cross-attn to_out carve-out should stay bf16" + assert wdtype(tt_block.ffn.ff1) == ttnn.bfloat8_b, "ff1 should be bf8" + assert wdtype(tt_block.ffn.ff2) == ttnn.bfloat8_b, "ff2 should be bf8" + + # Activation casts ride LTX_QUANT_ACTIVATIONS (default on): the projections and the bf16-weight + # to_out both narrow their activation to bf8. PCC can't catch a missing cast (bf16 activations only + # raise it), so assert it directly. + assert tt_block.attn1.to_qkv.activation_dtype == ttnn.bfloat8_b, "self-attn qkv activation should be bf8" + assert tt_block.attn1.to_out.activation_dtype == ttnn.bfloat8_b, "to_out consumes the shared bf8 activation" + assert tt_block.attn2.to_q.activation_dtype == ttnn.bfloat8_b, "cross-attn q activation should be bf8" + + torch.manual_seed(INPUT_SEED) + x = torch.randn(1, video_N_real, DIM, dtype=torch.float32) + context = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + temb = torch.randn(1, 1, 9 * DIM, dtype=torch.float32) + prompt_temb = torch.randn(1, 1, 2 * DIM, dtype=torch.float32) + cos_int, sin_int = _video_rope_freqs(F, H, W, rope_type=LTXRopeType.INTERLEAVED) + torch_out = _diffusers_video_block_ref( + torch_block, x=x, context=context, temb=temb, prompt_temb=prompt_temb, cos_i=cos_int, sin_i=sin_int + ) + + spatial = _pad_seq_dim(x, video_N, dim=1).unsqueeze(0) + tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3}) + tt_prompt = bf16_tensor(context.unsqueeze(0), device=mesh_device) + tt_temb = bf16_tensor( + temb.reshape(9, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3 + ) + tt_prompt_temb = bf16_tensor(prompt_temb.reshape(2, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device) + tt_cos, tt_sin = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + tt_out = tt_block( + video_1BND=tt_spatial, + video_prompt=tt_prompt, + video_temb=tt_temb, + video_N=video_N_real, + video_rope_cos=tt_cos, + video_rope_sin=tt_sin, + trans_mat=tt_trans_mat, + video_prompt_temb=tt_prompt_temb, + ) + + concat_dims = [None, None] + concat_dims[sp_axis] = 2 + concat_dims[tp_axis] = 3 + tt_v = ttnn.to_torch( + tt_out, + mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)), + ).squeeze(0)[:, :video_N_real, :] + assert tt_v.shape == (1, video_N_real, DIM), f"video shape {tt_v.shape}" + assert torch.isfinite(tt_v).all(), "bf8 video output NaN/Inf" + + # bf8 measures PCC 99.93% / RMSE 4.2% on BH 4x8. Floor PCC at the bf16 block test's ring-noise + # tolerance (0.988), not tighter — the RMSE gate (0.08 vs observed 0.042) is the real quant- + # regression catch; a broken quant path craters PCC or blows up RMSE and trips here regardless. + assert_quality(torch_out, tt_v, pcc=0.988, relative_rmse=0.08) + logger.info(f"PASSED bf8 block PCC: video {tuple(tt_v.shape)}") + + +def _run_inner_step( + *, + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + has_audio, + run_pcc, + checkpoint_variant: str, + use_forward_alias: bool, +): + """Shared body for the model forward tests; use_forward_alias picks tt_model(...) vs .forward(...).""" + skip_if_unsupported_num_links(mesh_device, num_links) + # Checkpoint variant only affects AV weight loading; skip the redundant copy in video mode. + if not has_audio and checkpoint_variant != "fast": + pytest.skip("checkpoint_variant only affects AV mode (video uses random scaled weights)") + checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant) + sp_factor = tuple(mesh_device.shape)[sp_axis] + # Real latent grid → SP-padded sequence (mirrors LTXPipeline). video_N_real is the logical + # token count (SDPA logical_n / pad masks); video_N is the padded on-device length. + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + audio_N = AUDIO_N + assert video_N % (32 * sp_factor) == 0 + if has_audio: + assert audio_N % (32 * sp_factor) == 0 + + # === Reference model === + do_pcc = run_pcc + ref_video = None + ref_audio = None + state_dict = None + + if has_audio: + # AV PCC path: 22B with strict=False. + state_dict = _load_22b_state_dict(num_layers=1, checkpoint_path=checkpoint_22b) + if state_dict is None: + pytest.skip(f"22B checkpoint not found at {checkpoint_22b}") + else: + # Video PCC path: random scaled weights (1 layer) from the diffusers 3D model. + torch_model = _make_diffusers_video_model(num_layers=1) + torch_model.eval() + _scale_init_(torch_model) + state_dict = _convert_diffusers_video_model_to_tt( + torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM + ) + state_dict = {k: v.detach().clone() for k, v in state_dict.items()} + + # === Inputs === + # Real-grid latent for the torch reference; the TT side gets a zero-padded copy (video_N). + torch.manual_seed(INPUT_SEED) + video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32) + video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1) + video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + sigma_val = TIMESTEP_VAL if not has_audio else 0.5 # AV path uses production sigma=0.5 + timestep_torch = torch.tensor([sigma_val]) + + audio_lat = None + audio_prompt = None + if has_audio: + audio_lat = torch.randn(1, audio_N, AUDIO_IN_CHANNELS, dtype=torch.float32) + audio_prompt = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32) + + # === Reference forward (before TT to avoid weight aliasing) === + # The torch LTXModel computes rope internally from positions; only the TT side builds freqs. + if do_pcc and not has_audio: + ref_video = _diffusers_video_model_ref( + torch_model, video_lat=video_lat_real, video_prompt=video_prompt, sigma_val=sigma_val, F=F, H=H, W=W + ) + logger.info(f"ref video {tuple(ref_video.shape)} range=[{ref_video.min():.3f}, {ref_video.max():.3f}]") + del torch_model + + if do_pcc and has_audio: + # AV reference: load the raw 22B into the diffusers AV model and run the cross-modal forward. + ref_model = _make_diffusers_av_model(num_layers=1) + ref_model.load_state_dict(_convert_lightricks_av_to_diffusers(state_dict, num_heads=NUM_HEADS), strict=False) + ref_model.eval() + t0 = time.time() + ref_video, ref_audio = _diffusers_av_model_ref( + ref_model, + video_lat=video_lat_real, + video_prompt=video_prompt, + audio_lat=audio_lat, + audio_prompt=audio_prompt, + sigma_val=sigma_val, + F=F, + H=H, + W=W, + audio_N=audio_N, + ) + logger.info( + f"AV ref forward: {time.time() - t0:.1f}s — video {tuple(ref_video.shape)} audio {tuple(ref_audio.shape)}" + ) + del ref_model + import gc + + gc.collect() + + # === TT model === + ccl_manager = _make_ccl_manager(mesh_device, num_links, topology) + parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis) + tt_model = _make_tt_model( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=has_audio, + num_layers=1, + ) + t0 = time.time() + tt_model.load_torch_state_dict(state_dict, strict=not has_audio) + logger.info(f"state-dict load: {time.time() - t0:.1f}s") + + # === Build TT-side RoPE / cross-PE / prompt tensors (INTERLEAVED + trans_mat) === + tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device) + tt_vc, tt_vs = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + call_kwargs = dict( + video_1BNI_torch=video_lat.unsqueeze(0), + video_prompt_1BLP=tt_video_prompt, + video_rope_cos=tt_vc, + video_rope_sin=tt_vs, + video_N=video_N_real, + trans_mat=tt_trans_mat, + timestep_torch=timestep_torch, + ) + if has_audio: + a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis) + vx_cos, vx_sin = _tt_rope( + _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + ax_cos, ax_sin = _tt_rope( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis + ) + ax_cos_full, ax_sin_full = _tt_rope_full( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis + ) + # Zero padded video tokens as V→A cross-attention keys (matches build_video_pad_mask). + v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis) + call_kwargs.update( + audio_1BNI_torch=audio_lat.unsqueeze(0), + audio_prompt_1BLP=bf16_tensor(audio_prompt.unsqueeze(0), device=mesh_device), + audio_rope_cos=a_cos, + audio_rope_sin=a_sin, + audio_N=audio_N, + video_cross_pe_cos=vx_cos, + video_cross_pe_sin=vx_sin, + audio_cross_pe_cos=ax_cos, + audio_cross_pe_sin=ax_sin, + audio_cross_pe_cos_full=ax_cos_full, + audio_cross_pe_sin_full=ax_sin_full, + video_padding_mask=v_pad_sp, + ) + + # === Forward === + t0 = time.time() + result = tt_model(**call_kwargs) if use_forward_alias else tt_model.forward(**call_kwargs) + logger.info(f"TT forward: {time.time() - t0:.1f}s") + + # Crop the SP-padding tail off the gathered video output before checking/comparing real tokens. + if has_audio: + tt_v_dev, tt_a_dev = result + tt_video = LTXTransformerModel.device_to_host(tt_v_dev).squeeze(0)[:, :video_N_real, :] + tt_audio = LTXTransformerModel.device_to_host(tt_a_dev).squeeze(0) + assert tt_video.shape == (1, video_N_real, OUT_CHANNELS), f"video shape {tt_video.shape}" + assert tt_audio.shape == (1, audio_N, AUDIO_IN_CHANNELS), f"audio shape {tt_audio.shape}" + assert torch.isfinite(tt_video).all() and torch.isfinite(tt_audio).all(), "NaN/Inf in TT output" + else: + tt_video = LTXTransformerModel.device_to_host(result).squeeze(0)[:, :video_N_real, :] + assert tt_video.shape == (1, video_N_real, OUT_CHANNELS), f"video shape {tt_video.shape}" + assert torch.isfinite(tt_video).all(), "NaN/Inf in TT output" + + del tt_model + + if do_pcc: + assert_quality(ref_video, tt_video, pcc=0.992, relative_rmse=0.15) + if has_audio: + assert_quality(ref_audio, tt_audio, pcc=0.992, relative_rmse=0.15) + logger.info("PASSED PCC") + else: + logger.info("PASSED (no PCC)") + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + LTX_TRANSFORMER_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize(("F", "H", "W"), _LTX_TRANSFORMER_SHAPE_PARAMS) +@pytest.mark.parametrize("has_audio", _LTX_TRANSFORMER_MODALITY_PARAMS) +@pytest.mark.parametrize("run_pcc", _LTX_TRANSFORMER_RUN_PCC_PARAMS) +@pytest.mark.parametrize("checkpoint_variant", _LTX_TRANSFORMER_CKPT_PARAMS) +def test_ltx_transformer_model( + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + has_audio, + run_pcc, + checkpoint_variant, + reset_seeds, +) -> None: + """Full LTXTransformerModel forward via ``__call__``.""" + _run_inner_step( + mesh_device=mesh_device, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + topology=topology, + is_fsdp=is_fsdp, + F=F, + H=H, + W=W, + has_audio=has_audio, + run_pcc=run_pcc, + checkpoint_variant=checkpoint_variant, + use_forward_alias=True, + ) + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + LTX_TRANSFORMER_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize(("F", "H", "W"), _LTX_TRANSFORMER_SHAPE_PARAMS) +@pytest.mark.parametrize("has_audio", _LTX_TRANSFORMER_MODALITY_PARAMS) +@pytest.mark.parametrize("run_pcc", _LTX_TRANSFORMER_RUN_PCC_PARAMS) +@pytest.mark.parametrize("checkpoint_variant", _LTX_TRANSFORMER_CKPT_PARAMS) +def test_ltx_transformer_inner_step( + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + has_audio, + run_pcc, + checkpoint_variant, + reset_seeds, +) -> None: + """LTXTransformerModel.forward — explicit call, denoising-loop path.""" + _run_inner_step( + mesh_device=mesh_device, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + topology=topology, + is_fsdp=is_fsdp, + F=F, + H=H, + W=W, + has_audio=has_audio, + run_pcc=run_pcc, + checkpoint_variant=checkpoint_variant, + use_forward_alias=False, + ) + + +# --------------------------------------------------------------------------- +# Per-token video timestep (I2V) — equivalence to the scalar path +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + [pytest.param((2, 4), 1, 0, 2, line_params_req_exact_devices, ttnn.Topology.Linear, False, id="2x4sp1tp0")], + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 17, 30, id="stage_1")]) +def test_ltx_per_token_timestep_equivalence( + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + reset_seeds, +) -> None: + """A uniform per-token video timestep must reproduce the scalar-timestep path. + + This is the weight-agnostic invariant behind I2V: with ``image_conditioning=True`` and a + per-token ``video_timestep = sigma`` for every token, the model output must match the + ``image_conditioning=False`` scalar path (T2V/AV stays bit-similar). The frame-0 pinning + behaviour itself is covered end-to-end by the pipeline; here we isolate the DiT plumbing. + """ + skip_if_unsupported_num_links(mesh_device, num_links) + sp_factor = tuple(mesh_device.shape)[sp_axis] + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + + # Shared random-scaled weights (1 layer) from the diffusers 3D video model. + torch_model = _make_diffusers_video_model(num_layers=1) + torch_model.eval() + _scale_init_(torch_model) + state_dict = _convert_diffusers_video_model_to_tt(torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM) + state_dict = {k: v.detach().clone() for k, v in state_dict.items()} + del torch_model + + torch.manual_seed(INPUT_SEED) + video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32) + video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1) + video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + sigma_val = TIMESTEP_VAL + timestep_torch = torch.tensor([sigma_val]) + + ccl_manager = _make_ccl_manager(mesh_device, num_links, topology) + parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis) + + # Shared TT-side RoPE / prompt tensors. + tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device) + tt_vc, tt_vs = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + base_kwargs = dict( + video_1BNI_torch=video_lat.unsqueeze(0), + video_prompt_1BLP=tt_video_prompt, + video_rope_cos=tt_vc, + video_rope_sin=tt_vs, + video_N=video_N_real, + trans_mat=tt_trans_mat, + timestep_torch=timestep_torch, + ) + + # Scalar path (image_conditioning=False). + scalar_model = _make_tt_model( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=False, + num_layers=1, + image_conditioning=False, + ) + scalar_model.load_torch_state_dict(state_dict, strict=True) + out_scalar = LTXTransformerModel.device_to_host(scalar_model.forward(**base_kwargs)).squeeze(0)[:, :video_N_real, :] + del scalar_model + + # Per-token path (image_conditioning=True) with a uniform timestep = sigma over the padded grid. + pertoken_model = _make_tt_model( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=False, + num_layers=1, + image_conditioning=True, + ) + pertoken_model.load_torch_state_dict(state_dict, strict=True) + video_timestep_torch = torch.full((video_N,), sigma_val, dtype=torch.float32) + out_pertoken = LTXTransformerModel.device_to_host( + pertoken_model.forward(video_timestep_torch=video_timestep_torch, **base_kwargs) + ).squeeze(0)[:, :video_N_real, :] + del pertoken_model + + assert out_pertoken.shape == out_scalar.shape, f"{out_pertoken.shape} vs {out_scalar.shape}" + assert torch.isfinite(out_pertoken).all(), "NaN/Inf in per-token output" + # Same weights + uniform sigma: the per-token MLP collapses to the scalar broadcast (bf16-equal). + assert_quality(out_scalar, out_pertoken, pcc=0.999, relative_rmse=0.02) + logger.info("PASSED: uniform per-token timestep reproduces the scalar path") + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + [pytest.param((2, 4), 1, 0, 2, line_params_req_exact_devices, ttnn.Topology.Linear, False, id="2x4sp1tp0")], + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 17, 30, id="stage_1")]) +def test_ltx_per_token_timestep_nonuniform( + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + reset_seeds, +) -> None: + """NON-uniform per-token timestep (the real I2V case) must match the diffusers per-token oracle. + + This is the case the existing equivalence test does NOT cover: frame-0 tokens at sigma=0 + (the pinned image anchor) and every other token at a high sigma. If the per-token gate does + not reach the fused attn/FFN epilogues, non-frame-0 tokens get the wrong (frame-0) gate and + the output diverges from the oracle — the DiT-level signature of "frame 0 ok, rest collapses". + """ + skip_if_unsupported_num_links(mesh_device, num_links) + sp_factor = tuple(mesh_device.shape)[sp_axis] + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + frame0_tokens = H * W # token order is f*H*W + h*W + w, so frame 0 == first H*W tokens + sigma_high = float(os.environ.get("LTX_TEST_SIGMA", "0.7")) + + # Shared random-scaled weights (1 layer) from the diffusers 3D video model. + torch_model = _make_diffusers_video_model(num_layers=1) + torch_model.eval() + _scale_init_(torch_model) + state_dict = _convert_diffusers_video_model_to_tt(torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM) + state_dict = {k: v.detach().clone() for k, v in state_dict.items()} + + torch.manual_seed(INPUT_SEED) + video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32) + video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1) + video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + + # Per-token timestep: frame-0 tokens pinned at sigma=0, everything else (incl. SP padding) at sigma_high. + video_ts_real = torch.full((video_N_real,), sigma_high, dtype=torch.float32) + video_ts_real[:frame0_tokens] = 0.0 + video_timestep_torch = torch.full((video_N,), sigma_high, dtype=torch.float32) + video_timestep_torch[:frame0_tokens] = 0.0 + + # === Reference (per-token oracle) — before TT to avoid weight aliasing === + ref_video = _diffusers_video_model_ref_pertoken( + torch_model, + video_lat=video_lat_real, + video_prompt=video_prompt, + video_ts_real=video_ts_real, + sigma_scalar=sigma_high, + F=F, + H=H, + W=W, + ) + del torch_model + + ccl_manager = _make_ccl_manager(mesh_device, num_links, topology) + parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis) + tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device) + tt_vc, tt_vs = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + model = _make_tt_model( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=False, + num_layers=1, + image_conditioning=True, + ) + model.load_torch_state_dict(state_dict, strict=True) + out_tt = LTXTransformerModel.device_to_host( + model.forward( + video_1BNI_torch=video_lat.unsqueeze(0), + video_prompt_1BLP=tt_video_prompt, + video_rope_cos=tt_vc, + video_rope_sin=tt_vs, + video_N=video_N_real, + trans_mat=tt_trans_mat, + timestep_torch=torch.tensor([sigma_high]), + video_timestep_torch=video_timestep_torch, + ) + ).squeeze(0)[:, :video_N_real, :] + del model + + assert out_tt.shape == ref_video.shape, f"{out_tt.shape} vs {ref_video.shape}" + assert torch.isfinite(out_tt).all(), "NaN/Inf in per-token output" + + # Localize frame-0 (pinned, sigma=0) vs the rest (sigma_high): a per-token gate that doesn't + # reach the fused epilogues makes the "rest" track the oracle far worse than frame 0. + def _pcc(a, b): + a32, b32 = a.float().flatten(), b.float().flatten() + return torch.corrcoef(torch.stack([a32, b32]))[0, 1].item() + + pcc_all = _pcc(ref_video, out_tt) + pcc_f0 = _pcc(ref_video[:, :frame0_tokens, :], out_tt[:, :frame0_tokens, :]) + pcc_rest = _pcc(ref_video[:, frame0_tokens:, :], out_tt[:, frame0_tokens:, :]) + logger.info(f"non-uniform per-token (TT vs oracle): all={pcc_all:.5f} frame0={pcc_f0:.5f} rest={pcc_rest:.5f}") + + assert_quality(ref_video, out_tt, pcc=0.99, relative_rmse=0.03) + + +# --------------------------------------------------------------------------- Trace-based perf: VIDEO stage-2 ring +def _build_video_trace_setup( + *, + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + has_audio=False, + checkpoint_variant="fast", + num_layers=1, +): + """Build the TT model + device-only ``inner_step`` kwargs for the trace-perf test. + + Mirrors the TT setup in ``_run_inner_step`` and reproduces ``LTXTransformerModel.forward``'s + host→device upload preamble (which ``inner_step`` normally receives), so the returned ttnn + kwargs can be captured/replayed directly under ttnn trace. Returns + ``(tt_model, inner_kwargs, video_N_real)``. Perf-only; no PCC. + + Video mode uses random scaled weights. AV mode (``has_audio=True``) loads the 22B checkpoint + (``strict=False``) like the AV PCC path and adds the audio device tensors (audio latent, + audio/cross-modal rope, video pad mask) so a2v / audio shapes actually fire under the trace. + """ + sp_factor = tuple(mesh_device.shape)[sp_axis] + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + assert video_N % (32 * sp_factor) == 0 + audio_N = AUDIO_N + if has_audio: + assert audio_N % (32 * sp_factor) == 0 + + if has_audio: + # AV weights come from the 22B checkpoint (strict=False), same as the AV PCC path. + checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant) + state_dict = _load_22b_state_dict(num_layers=num_layers, checkpoint_path=checkpoint_22b) + if state_dict is None: + pytest.skip(f"22B checkpoint not found at {checkpoint_22b}") + else: + # Random scaled weights from the diffusers 3D video model — the video PCC path. + torch_model = _make_diffusers_video_model(num_layers=num_layers) + torch_model.eval() + _scale_init_(torch_model) + state_dict = _convert_diffusers_video_model_to_tt( + torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM + ) + state_dict = {k: v.detach().clone() for k, v in state_dict.items()} + del torch_model + + # Real-grid latent; the TT side gets a zero-padded copy (video_N). + torch.manual_seed(INPUT_SEED) + video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32) + video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1) + video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + sigma_val = 0.5 if has_audio else TIMESTEP_VAL # AV path uses production sigma=0.5 + timestep_torch = torch.tensor([sigma_val]) + + audio_lat = None + audio_prompt = None + if has_audio: + audio_lat = torch.randn(1, audio_N, AUDIO_IN_CHANNELS, dtype=torch.float32) + audio_prompt = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32) + + ccl_manager = _make_ccl_manager(mesh_device, num_links, topology) + parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis) + tt_model = _make_tt_model( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=has_audio, + num_layers=num_layers, + ) + t0 = time.time() + tt_model.load_torch_state_dict(state_dict, strict=not has_audio) + logger.info(f"state-dict load: {time.time() - t0:.1f}s") + + # Stage-constant device tensors (INTERLEAVED rope + trans_mat + prompt). + tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device) + tt_vc, tt_vs = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + # Per-step device inputs — reproduce LTXTransformerModel.forward's upload for the video path. + video_1BNI_torch = video_lat.unsqueeze(0) # (1, B=1, video_N, IN_CHANNELS) + B_size = video_1BNI_torch.shape[1] + video_1BNI = bf16_tensor(video_1BNI_torch, device=mesh_device, mesh_axis=sp_axis, shard_dim=-2) + timestep = bf16_tensor(timestep_torch.reshape(1, 1, B_size, 1) * 1000.0, device=mesh_device) + + # inner_step is keyword-only; these are exactly the device tensors it consumes. + inner_kwargs = dict( + video_1BNI=video_1BNI, + timestep=timestep, + video_prompt_1BLP=tt_video_prompt, + video_rope_cos=tt_vc, + video_rope_sin=tt_vs, + video_N=video_N_real, + trans_mat=tt_trans_mat, + ) + + if has_audio: + # Audio + cross-modal device tensors, mirroring the AV branch of _run_inner_step. + a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis) + vx_cos, vx_sin = _tt_rope( + _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + ax_cos, ax_sin = _tt_rope( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis + ) + ax_cos_full, ax_sin_full = _tt_rope_full( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis + ) + v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis) + audio_1BNI = bf16_tensor(audio_lat.unsqueeze(0), device=mesh_device, mesh_axis=sp_axis, shard_dim=-2) + inner_kwargs.update( + audio_1BNI=audio_1BNI, + audio_prompt_1BLP=bf16_tensor(audio_prompt.unsqueeze(0), device=mesh_device), + audio_rope_cos=a_cos, + audio_rope_sin=a_sin, + audio_N=audio_N, + video_cross_pe_cos=vx_cos, + video_cross_pe_sin=vx_sin, + audio_cross_pe_cos=ax_cos, + audio_cross_pe_sin=ax_sin, + audio_cross_pe_cos_full=ax_cos_full, + audio_cross_pe_sin_full=ax_sin_full, + video_padding_mask=v_pad_sp, + ) + return tt_model, inner_kwargs, video_N_real + + +def _build_trace_inner_kwargs(*, mesh_device, sp_axis, tp_axis, F, H, W, has_audio): + """Build ONLY the device inner_step kwargs for a given (F,H,W) — no model / CCLManager. + + Split out so ONE model + CCLManager can be driven at multiple shapes (e.g. stage-1 then + stage-2) to reproduce the e2e's cross-stage device/semaphore-pool state. Mirrors the per-shape + tensor setup inside _build_video_trace_setup. Returns (inner_kwargs, video_N_real). + """ + sp_factor = tuple(mesh_device.shape)[sp_axis] + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + assert video_N % (32 * sp_factor) == 0 + audio_N = AUDIO_N + if has_audio: + assert audio_N % (32 * sp_factor) == 0 + + torch.manual_seed(INPUT_SEED) + video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32) + video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1) + video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + sigma_val = 0.5 if has_audio else TIMESTEP_VAL + timestep_torch = torch.tensor([sigma_val]) + + audio_lat = None + audio_prompt = None + if has_audio: + audio_lat = torch.randn(1, audio_N, AUDIO_IN_CHANNELS, dtype=torch.float32) + audio_prompt = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32) + + tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device) + tt_vc, tt_vs = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + video_1BNI_torch = video_lat.unsqueeze(0) + B_size = video_1BNI_torch.shape[1] + video_1BNI = bf16_tensor(video_1BNI_torch, device=mesh_device, mesh_axis=sp_axis, shard_dim=-2) + timestep = bf16_tensor(timestep_torch.reshape(1, 1, B_size, 1) * 1000.0, device=mesh_device) + + inner_kwargs = dict( + video_1BNI=video_1BNI, + timestep=timestep, + video_prompt_1BLP=tt_video_prompt, + video_rope_cos=tt_vc, + video_rope_sin=tt_vs, + video_N=video_N_real, + trans_mat=tt_trans_mat, + ) + if has_audio: + a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis) + vx_cos, vx_sin = _tt_rope( + _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + ax_cos, ax_sin = _tt_rope( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis + ) + ax_cos_full, ax_sin_full = _tt_rope_full( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis + ) + v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis) + audio_1BNI = bf16_tensor(audio_lat.unsqueeze(0), device=mesh_device, mesh_axis=sp_axis, shard_dim=-2) + inner_kwargs.update( + audio_1BNI=audio_1BNI, + audio_prompt_1BLP=bf16_tensor(audio_prompt.unsqueeze(0), device=mesh_device), + audio_rope_cos=a_cos, + audio_rope_sin=a_sin, + audio_N=audio_N, + video_cross_pe_cos=vx_cos, + video_cross_pe_sin=vx_sin, + audio_cross_pe_cos=ax_cos, + audio_cross_pe_sin=ax_sin, + audio_cross_pe_cos_full=ax_cos_full, + audio_cross_pe_sin_full=ax_sin_full, + video_padding_mask=v_pad_sp, + ) + return inner_kwargs, video_N_real + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + [ + pytest.param((4, 8), 1, 0, 2, ring_params_4k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_4k"), + # 8k fabric payload variant (where stage-2 wedges at depth). + pytest.param((4, 8), 1, 0, 2, ring_params_8k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_8k"), + ], + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize("has_audio", [pytest.param(False, id="video"), pytest.param(True, id="av")]) +@pytest.mark.parametrize("checkpoint_variant", [pytest.param("fast", id="ckpt_fast")]) +def test_ltx_transformer_crossstage_hang( + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + has_audio, + checkpoint_variant, + reset_seeds, +) -> None: + """FAST cross-stage repro of the e2e stage-2 deadlock: drive ONE model+CCLManager through + stage-1 then stage-2 (eager, no VAE/Gemma), so stage-2 reuses the shared strided-AGMM pool + that stage-1 left mid-rotation. PASSES -> not cross-stage; HANGS -> confirmed. + """ + # Stage-2 shape; full depth (env-overridable) to stack many back-to-back strided ops per forward. + tt_model, kwargs_s2, vN2 = _build_video_trace_setup( + mesh_device=mesh_device, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + topology=topology, + is_fsdp=is_fsdp, + F=19, + H=34, + W=60, + has_audio=has_audio, + checkpoint_variant=checkpoint_variant, + num_layers=int(os.environ.get("CROSSSTAGE_NUM_LAYERS", "48")), + ) + # Stage-1 shape (half spatial res, like the e2e's height//2,width//2) on the SAME model. + kwargs_s1, vN1 = _build_trace_inner_kwargs( + mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, F=19, H=17, W=30, has_audio=has_audio + ) + logger.info(f"cross-stage: stage-1 vN={vN1} then stage-2 vN={vN2} (shared CCLManager)") + + # CROSSSTAGE_SKIP_S1=1 -> run only stage-2, isolating the stage transition from full-depth big-M. + skip_s1 = os.environ.get("CROSSSTAGE_SKIP_S1", "0") != "0" + if not skip_s1: + # Stage 1 first (eager) — populate/rotate the shared strided-AGMM semaphore pool at stage-1 M. + for i in range(3): + _ = tt_model.inner_step(**kwargs_s1) + ttnn.synchronize_device(mesh_device) + logger.info(f"stage-1 forward {i} done") + logger.info("stage-1 complete; entering STAGE-2 (the e2e hang point) ...") + else: + logger.info("SKIP_S1: entering STAGE-2 directly (48 layers, no stage-1) ...") + + # Stage 2 (eager) — reuses the same pool. This is where the e2e deadlocks. + for i in range(3): + _ = tt_model.inner_step(**kwargs_s2) + ttnn.synchronize_device(mesh_device) + logger.info(f"stage-2 forward {i} done") + + logger.info("PASSED cross-stage: stage-1 -> stage-2 completed without deadlock") + del tt_model + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + # VIDEO stage-2 ring config only (8k fabric payload); do NOT add AV/other-mesh params here. + [pytest.param((4, 8), 1, 0, 2, ring_params_8k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_8k")], + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 34, 60, id="stage_2")]) +@pytest.mark.parametrize("has_audio", [pytest.param(False, id="video"), pytest.param(True, id="av")]) +@pytest.mark.parametrize("checkpoint_variant", [pytest.param("fast", id="ckpt_fast")]) +def test_ltx_transformer_trace_perf( + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + has_audio, + checkpoint_variant, + reset_seeds, +) -> None: + """Trace-based perf for LTXTransformerModel.inner_step (VIDEO stage-2 ring). + + Measures StridedAllGatherMinimalMatmulAsync device time and block wall-clock at steady state + under a ttnn trace (gap-free replay), which the eager op-by-op path can't expose. Perf-only: + asserts shape/finiteness, no PCC. inner_step is @traced_function(prep_run=False), so kernels + are precompiled with a couple of eager warm calls before trace capture. + """ + tt_model, inner_kwargs, video_N_real = _build_video_trace_setup( + mesh_device=mesh_device, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + topology=topology, + is_fsdp=is_fsdp, + F=F, + H=H, + W=W, + has_audio=has_audio, + checkpoint_variant=checkpoint_variant, + ) + + # AV emits far more markers than video + n_warm = int(os.environ.get("LTX_PERF_N_WARM", "1" if has_audio else "2")) + n_replay = int(os.environ.get("LTX_PERF_N_REPLAY", "3" if has_audio else "5")) + + # Warm / compile: eager (untraced) calls precompile kernels (prep_run=False on inner_step). + t0 = time.time() + for _ in range(n_warm): + _ = tt_model.inner_step(**inner_kwargs) + ttnn.synchronize_device(mesh_device) + logger.info(f"eager warm/compile ({n_warm}x): {time.time() - t0:.1f}s") + + # Capture: first traced call captures the trace (and executes it once). + t0 = time.time() + _ = tt_model.inner_step(**inner_kwargs, traced=True) + ttnn.synchronize_device(mesh_device) + logger.info(f"trace capture: {time.time() - t0:.1f}s") + + # Drain warm+capture markers so only the signposted replays occupy the profiler buffer. + ttnn.ReadDeviceProfiler(mesh_device) + + # Steady-state replay under the signpost bracket for tt-perf-report --start/end-signpost. + signpost("start") + t0 = time.time() + result = None + for _ in range(n_replay): + result = tt_model.inner_step(**inner_kwargs, traced=True) + ttnn.synchronize_device(mesh_device) + replay_s = time.time() - t0 + signpost("stop") + logger.info(f"trace replay ({n_replay}x): {replay_s:.3f}s → {replay_s / n_replay * 1e3:.2f} ms/step") + + # Perf-only: crop SP-padding tail, assert shape + finiteness (no PCC) + video_out = result[0] if isinstance(result, tuple) else result + tt_video = LTXTransformerModel.device_to_host(video_out).squeeze(0)[:, :video_N_real, :] + assert tt_video.shape == (1, video_N_real, OUT_CHANNELS), f"video shape {tt_video.shape}" + assert torch.isfinite(tt_video).all(), "NaN/Inf in traced video output" + logger.info(f"PASSED trace-perf: video {tuple(tt_video.shape)} has_audio={has_audio}") + + del tt_model + + +# Trace-based perf for a SINGLE transformer block (isolates one block's program order). +def _build_block_trace_setup( + *, mesh_device, sp_axis, tp_axis, num_links, topology, is_fsdp, F, H, W, has_audio, checkpoint_variant="fast" +): + """Build an ``LTXTransformerBlock`` + device ``forward`` kwargs for the block trace-perf test. + + Mirrors the non-PCC setup of ``test_ltx_transformer_block``: AV loads the 22B checkpoint + (``strict=False``) and builds the full video+audio+cross-modal device tensors; video harvests + random scaled weights from a throwaway diffusers block. Returns + ``(tt_block, forward_kwargs, video_N_real, audio_N_real)``. Perf-only; no PCC. + """ + checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant) + sp_factor = tuple(mesh_device.shape)[sp_axis] + # Real latent grid → SP-padded sequence (mirrors LTXPipeline) + video_N_real = F * H * W + video_N = _sp_pad_len(video_N_real, sp_factor) + audio_N, audio_N_real = _audio_seq_lens(F, sp_factor) + assert video_N % (32 * sp_factor) == 0, f"video_N={video_N} not sp/tile-aligned for sp={sp_factor}" + assert audio_N % (32 * sp_factor) == 0, f"audio_N={audio_N} not sp/tile-aligned for sp={sp_factor}" + + ccl_manager = _make_ccl_manager(mesh_device, num_links, topology) + parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis) + tt_block = _make_tt_block( + mesh_device=mesh_device, + ccl_manager=ccl_manager, + parallel_config=parallel_config, + is_fsdp=is_fsdp, + has_audio=has_audio, + ) + + if has_audio: + sd = _load_22b_state_dict(num_layers=1, checkpoint_path=checkpoint_22b) + if sd is None: + pytest.skip(f"22B checkpoint not found at {checkpoint_22b}") + block_sd = { + k[len("transformer_blocks.0.") :]: v for k, v in sd.items() if k.startswith("transformer_blocks.0.") + } + tt_block.load_torch_state_dict(block_sd, strict=False) + else: + # video mode — harvest a shape-matching state dict from a throwaway block (perf-only). + dummy = _make_diffusers_video_block() + _scale_init_(dummy) + conv = _convert_diffusers_video_block_to_tt(dummy.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM) + tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()}) + del dummy + + # Inputs (real-grid token count; the TT side gets a zero-padded copy below). + torch.manual_seed(INPUT_SEED) + x = torch.randn(1, video_N_real, DIM, dtype=torch.float32) + context = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32) + temb = torch.randn(1, 1, 9 * DIM, dtype=torch.float32) # 9 adaln params + prompt_temb = torch.randn(1, 1, 2 * DIM, dtype=torch.float32) # 2 adaln params for prompt + + # TT video tensors + spatial = _pad_seq_dim(x, video_N, dim=1).unsqueeze(0) + tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3}) + tt_prompt = bf16_tensor(context.unsqueeze(0), device=mesh_device) + tt_temb = bf16_tensor( + temb.reshape(9, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3 + ) + tt_prompt_temb = bf16_tensor(prompt_temb.reshape(2, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device) + tt_cos, tt_sin = _tt_rope( + _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device) + + forward_kwargs = dict( + video_1BND=tt_spatial, + video_prompt=tt_prompt, + video_temb=tt_temb, + video_N=video_N_real, + video_rope_cos=tt_cos, + video_rope_sin=tt_sin, + trans_mat=tt_trans_mat, + video_prompt_temb=tt_prompt_temb, + ) + if has_audio: + # Real tokens in [:audio_N_real], zeros in the padded tail (matches the padded audio latent). + a_x = torch.zeros(1, audio_N, AUDIO_DIM, dtype=torch.float32) + a_x[:, :audio_N_real, :] = torch.randn(1, audio_N_real, AUDIO_DIM, dtype=torch.float32) + a_ctx = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32) + a_temb = torch.randn(1, 1, 9 * AUDIO_DIM, dtype=torch.float32) + a_prompt_temb = torch.randn(1, 1, 2 * AUDIO_DIM, dtype=torch.float32) + av_ca_v = torch.randn(1, 1, 5 * DIM, dtype=torch.float32) # 4 scale-shift + 1 gate + av_ca_a = torch.randn(1, 1, 5 * AUDIO_DIM, dtype=torch.float32) + + a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis) + vx_cos, vx_sin = _tt_rope( + _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N + ) + ax_cos, ax_sin = _tt_rope( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis + ) + ax_cos_full, ax_sin_full = _tt_rope_full( + _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis + ) + + # Padding masks, same construction as the fast pipeline (audio padded, video aligned). + a_attn_mask, a_pad_sp, a_pad_full = build_audio_masks( + audio_N, audio_N_real, mesh_device=mesh_device, sp_axis=sp_axis + ) + v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis) + + forward_kwargs.update( + audio_1BND=bf16_tensor_2dshard( + a_x.unsqueeze(0), device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3} + ), + audio_prompt=bf16_tensor(a_ctx.unsqueeze(0), device=mesh_device), + audio_temb=bf16_tensor( + a_temb.reshape(9, AUDIO_DIM).unsqueeze(1).unsqueeze(1), + device=mesh_device, + mesh_axis=tp_axis, + shard_dim=3, + ), + audio_prompt_temb=bf16_tensor( + a_prompt_temb.reshape(2, AUDIO_DIM).unsqueeze(1).unsqueeze(1), device=mesh_device + ), + av_ca_temb=bf16_tensor( + av_ca_v.reshape(5, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3 + ), + av_ca_audio_temb=bf16_tensor( + av_ca_a.reshape(5, AUDIO_DIM).unsqueeze(1).unsqueeze(1), + device=mesh_device, + mesh_axis=tp_axis, + shard_dim=3, + ), + audio_N=audio_N, + audio_rope_cos=a_cos, + audio_rope_sin=a_sin, + video_cross_pe_cos=vx_cos, + video_cross_pe_sin=vx_sin, + audio_cross_pe_cos=ax_cos, + audio_cross_pe_sin=ax_sin, + audio_cross_pe_cos_full=ax_cos_full, + audio_cross_pe_sin_full=ax_sin_full, + audio_attn_mask=a_attn_mask, + audio_padding_mask=a_pad_sp, + audio_padding_mask_full=a_pad_full, + video_padding_mask=v_pad_sp, + ) + return tt_block, forward_kwargs, video_N_real, audio_N_real + + +@pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + # Same ring stage-2 8k config as test_ltx_transformer_trace_perf + [pytest.param((4, 8), 1, 0, 2, ring_params_8k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_8k")], + indirect=["mesh_device", "device_params"], +) +@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 34, 60, id="stage_2")]) +@pytest.mark.parametrize("has_audio", [pytest.param(False, id="video"), pytest.param(True, id="av")]) +@pytest.mark.parametrize("checkpoint_variant", [pytest.param("fast", id="ckpt_fast")]) +def test_ltx_transformer_block_trace_perf( + mesh_device, + sp_axis, + tp_axis, + num_links, + topology, + is_fsdp, + F, + H, + W, + has_audio, + checkpoint_variant, + reset_seeds, +) -> None: + """Trace-based perf for a SINGLE ``LTXTransformerBlock`` forward (AV stage-2 ring). + + Isolates one block's program order — no model-level patch-embed / proj-out / tilize / + untilize / mesh-partition ops, and no multi-block scrambling — so tt-perf-report shows a clean + one-block, program-ordered view. Captures the block forward under a ttnn trace (via the same + ``Tracer`` machinery that backs ``inner_step``'s ``traced=True`` path) and replays it under the + same "start"/"stop" signpost bracket used by ``test_ltx_transformer_trace_perf``. Perf-only: + asserts shape/finiteness, no PCC. + """ + tt_block, forward_kwargs, video_N_real, _audio_N_real = _build_block_trace_setup( + mesh_device=mesh_device, + sp_axis=sp_axis, + tp_axis=tp_axis, + num_links=num_links, + topology=topology, + is_fsdp=is_fsdp, + F=F, + H=H, + W=W, + has_audio=has_audio, + checkpoint_variant=checkpoint_variant, + ) + + # Trace the block forward with the same Tracer settings inner_step uses + tracer = Tracer(tt_block.forward, device=mesh_device, prep_run=False, clone_prep_inputs=False) + + # AV runs many more ops than video (audio self/cross + a2v/v2a + audio ff/RS) + n_warm = 1 if has_audio else 2 + n_replay = 3 if has_audio else 5 + + # Warm / compile: eager (untraced) calls precompile kernels before trace capture. + t0 = time.time() + for _ in range(n_warm): + _ = tt_block(**forward_kwargs) + ttnn.synchronize_device(mesh_device) + logger.info(f"eager warm/compile ({n_warm}x): {time.time() - t0:.1f}s") + + # Capture: first traced call captures the trace (and executes it once). + t0 = time.time() + _ = tracer(**forward_kwargs, traced=True) + ttnn.synchronize_device(mesh_device) + logger.info(f"trace capture: {time.time() - t0:.1f}s") + + # Drain warm+capture markers so ONLY the signposted replays occupy the profiler DRAM buffer + ttnn.ReadDeviceProfiler(mesh_device) + + # Steady-state replay under the signpost bracket for tt-perf-report --start/end-signpost. + signpost("start") + t0 = time.time() + result = None + for _ in range(n_replay): + result = tracer(**forward_kwargs, traced=True) + ttnn.synchronize_device(mesh_device) + replay_s = time.time() - t0 + signpost("stop") + logger.info(f"trace replay ({n_replay}x): {replay_s:.3f}s → {replay_s / n_replay * 1e3:.2f} ms/step") + + # Perf-only: crop SP-padding tail, assert shape + finiteness (no PCC) + tt_v = result[0] if isinstance(result, tuple) else result + concat_dims = [None, None] + concat_dims[sp_axis] = 2 + concat_dims[tp_axis] = 3 + tt_v_torch = ttnn.to_torch( + tt_v, + mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)), + ).squeeze(0)[:, :video_N_real, :] + assert tt_v_torch.shape == (1, video_N_real, DIM), f"video shape {tt_v_torch.shape}" + assert torch.isfinite(tt_v_torch).all(), "NaN/Inf in traced video output" + logger.info(f"PASSED block trace-perf: video {tuple(tt_v_torch.shape)} has_audio={has_audio}") + + tracer.release_trace() + del tt_block diff --git a/code/models/tt_dit/tests/models/ltx/test_vae_ltx.py b/code/models/tt_dit/tests/models/ltx/test_vae_ltx.py new file mode 100644 index 0000000000000000000000000000000000000000..edbd3ba341923daf12e1da77ecdd034058384123 --- /dev/null +++ b/code/models/tt_dit/tests/models/ltx/test_vae_ltx.py @@ -0,0 +1,1151 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +""" +Tests for LTX-2 Video VAE components. + +Torch references use diffusers LTX-2 VAE modules (mirrors test_vae_wan2_1.py / test_audio_ltx.py). +""" + +from __future__ import annotations + +import os +from typing import Any + +import pytest +import torch +import torch.nn as nn +from loguru import logger +from tracy import signpost + +import ttnn +from models.tt_dit.models.vae.vae_ltx import ( + LTXCausalConv3d, + LTXDepthToSpaceUpsample, + LTXResnetBlock3D, + LTXVideoDecoder, + LTXVideoEncoder, +) +from models.tt_dit.parallel.config import ParallelFactor, VaeHWParallelConfig +from models.tt_dit.parallel.manager import CCLManager +from models.tt_dit.utils.check import assert_quality +from models.tt_dit.utils.conv3d import conv_pad_in_channels + +from .ltx_mesh_params import ( + LTX_VAE_DECODER_MESH_PARAMS, + LTX_VAE_DECODER_MULTI_ONLY_MESH_PARAMS, + LTX_VAE_ENCODER_PROD_MESH_PARAMS, +) + + +def _require_diffusers_ltx_vae(): + """Return diffusers LTX-2 VAE building blocks (preferred) or LTX-Video fallbacks.""" + pytest.importorskip("diffusers") + try: + from diffusers.models.autoencoders.autoencoder_kl_ltx2 import ( + LTX2VideoCausalConv3d, + LTX2VideoDownsampler3d, + LTX2VideoResnetBlock3d, + LTX2VideoUpsampler3d, + PerChannelRMSNorm, + ) + + return { + "causal_conv": LTX2VideoCausalConv3d, + "resnet": LTX2VideoResnetBlock3d, + "upsample": LTX2VideoUpsampler3d, + "downsample": LTX2VideoDownsampler3d, + "pixel_norm": PerChannelRMSNorm, + "ltx2": True, + } + except ImportError: + from diffusers.models.autoencoders.autoencoder_kl_ltx import ( + LTXVideoCausalConv3d, + LTXVideoDownsampler3d, + LTXVideoResnetBlock3d, + LTXVideoUpsampler3d, + ) + from diffusers.models.normalization import RMSNorm + + return { + "causal_conv": LTXVideoCausalConv3d, + "resnet": LTXVideoResnetBlock3d, + "upsample": LTXVideoUpsampler3d, + "downsample": LTXVideoDownsampler3d, + "pixel_norm": RMSNorm, + "ltx2": False, + } + + +def _diffusers_resnet_state_to_tt(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Remap diffusers LTX-2 resnet keys to the TT/Lightricks layout.""" + out = dict(state) + if "conv_shortcut.weight" in out: + out["conv_shortcut.conv.weight"] = out.pop("conv_shortcut.weight") + if "conv_shortcut.bias" in out: + out["conv_shortcut.conv.bias"] = out.pop("conv_shortcut.bias") + return out + + +def _run_diffusers_causal_conv(conv: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor: + if ltx2: + return conv(x, causal=causal) + return conv(x) + + +def _run_diffusers_resnet(block: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor: + if ltx2: + return block(x, causal=causal) + return block(x) + + +def _run_diffusers_upsample(block: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor: + if ltx2: + return block(x, causal=causal) + return block(x) + + +def _run_diffusers_downsample(block: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor: + if ltx2: + return block(x, causal=causal) + return block(x) + + +class _PerChannelStatistics(nn.Module): + """Minimal per-channel latent stats (Lightricks key names for TT weight load).""" + + def __init__(self, latent_channels: int) -> None: + super().__init__() + self.register_buffer("mean-of-means", torch.zeros(latent_channels)) + self.register_buffer("std-of-means", torch.ones(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1) + std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1) + return x * std + mean + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1) + std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1) + return (x - mean) / std + + +def _diffusers_decoder_state_to_tt(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Remap diffusers LTX-2 decoder keys (incl. nested resnets) to TT layout.""" + out: dict[str, torch.Tensor] = {} + for key, value in state.items(): + if key.endswith("conv_shortcut.weight"): + out[key.replace("conv_shortcut.weight", "conv_shortcut.conv.weight")] = value + elif key.endswith("conv_shortcut.bias"): + out[key.replace("conv_shortcut.bias", "conv_shortcut.conv.bias")] = value + else: + out[key] = value + return out + + +class _TorchUNetMidBlock3D(nn.Module): + def __init__( + self, + *, + in_channels: int, + num_layers: int, + resnet_cls: type, + ltx2: bool, + spatial_padding_mode: str = "zeros", + ) -> None: + super().__init__() + kwargs: dict[str, Any] = {"in_channels": in_channels, "out_channels": in_channels} + if ltx2 and spatial_padding_mode != "zeros": + kwargs["spatial_padding_mode"] = spatial_padding_mode + self.res_blocks = nn.ModuleList([resnet_cls(**kwargs) for _ in range(num_layers)]) + + def forward(self, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor: + for block in self.res_blocks: + x = _run_diffusers_resnet(block, x, causal=causal, ltx2=ltx2) + return x + + +def _unpatchify(x: torch.Tensor, *, patch_size_hw: int = 4, patch_size_t: int = 1) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = x.shape + p = patch_size_hw + p_t = patch_size_t + x = x.reshape(batch_size, -1, p_t, p, p, num_frames, height, width) + return x.permute(0, 1, 5, 2, 6, 4, 7, 3).flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +def _patchify(x: torch.Tensor, *, patch_size_hw: int = 4, patch_size_t: int = 1) -> torch.Tensor: + """Space-to-depth patchify (exact inverse of ``_unpatchify``): folds a + ``(p_t, p_hw, p_hw)`` patch into channels with order ``(c, p_t, p_width, p_height)``. + Mirrors the Lightricks/TT encoder patchify so both references agree. + """ + batch_size, num_channels, num_frames, height, width = x.shape + p = patch_size_hw + p_t = patch_size_t + x = x.reshape(batch_size, num_channels, num_frames // p_t, p_t, height // p, p, width // p, p) + # (b, c, f', p_t, h', p_h, w', p_w) -> (b, c, p_t, p_w, p_h, f', h', w') + x = x.permute(0, 1, 3, 7, 5, 2, 4, 6) + return x.reshape(batch_size, num_channels * p_t * p * p, num_frames // p_t, height // p, width // p) + + +class _TorchLTXVideoDecoder(nn.Module): + """Block-list LTX decoder built from diffusers VAE primitives (matches TT/Lightricks layout).""" + + def __init__( + self, + *, + decoder_blocks: list[tuple[str, dict | int]], + in_channels: int = 128, + out_channels: int = 3, + patch_size: int = 4, + base_channels: int = 128, + causal: bool = False, + spatial_padding_mode: str = "zeros", + vae_mods: dict, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.causal = causal + self.ltx2 = vae_mods["ltx2"] + causal_conv = vae_mods["causal_conv"] + resnet_cls = vae_mods["resnet"] + upsample_cls = vae_mods["upsample"] + pixel_norm = vae_mods["pixel_norm"] + + out_channels_patched = out_channels * patch_size**2 + feature_channels = base_channels * 8 + + conv_kwargs: dict[str, Any] = {} + if self.ltx2: + if spatial_padding_mode != "zeros": + conv_kwargs["spatial_padding_mode"] = spatial_padding_mode + else: + conv_kwargs["is_causal"] = causal + + self.per_channel_statistics = _PerChannelStatistics(in_channels) + self.conv_in = causal_conv( + in_channels=in_channels, + out_channels=feature_channels, + kernel_size=3, + stride=1, + **conv_kwargs, + ) + + stride_map = { + "compress_all": (2, 2, 2), + "compress_space": (1, 2, 2), + "compress_time": (2, 1, 1), + } + + self.up_blocks = nn.ModuleList() + ch = feature_channels + for block_name, block_params in reversed(decoder_blocks): + block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params + if block_name == "res_x": + self.up_blocks.append( + _TorchUNetMidBlock3D( + in_channels=ch, + num_layers=block_config["num_layers"], + resnet_cls=resnet_cls, + ltx2=self.ltx2, + spatial_padding_mode=spatial_padding_mode, + ) + ) + elif block_name in stride_map: + multiplier = block_config.get("multiplier", 1) + upsample_kwargs: dict[str, Any] = { + "in_channels": ch, + "stride": stride_map[block_name], + "residual": block_config.get("residual", False), + "upscale_factor": multiplier, + } + self.up_blocks.append(upsample_cls(**upsample_kwargs)) + ch = ch // multiplier + else: + raise ValueError(f"Unknown decoder block: {block_name}") + + self.conv_norm_out = pixel_norm() + self.conv_act = nn.SiLU() + self.conv_out = causal_conv( + in_channels=ch, + out_channels=out_channels_patched, + kernel_size=3, + stride=1, + **conv_kwargs, + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample = self.per_channel_statistics.un_normalize(sample) + sample = _run_diffusers_causal_conv(self.conv_in, sample, causal=self.causal, ltx2=self.ltx2) + for up_block in self.up_blocks: + if isinstance(up_block, _TorchUNetMidBlock3D): + sample = up_block(sample, causal=self.causal, ltx2=self.ltx2) + else: + sample = _run_diffusers_upsample(up_block, sample, causal=self.causal, ltx2=self.ltx2) + if self.ltx2: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample.movedim(1, -1)).movedim(-1, 1) + sample = self.conv_act(sample) + sample = _run_diffusers_causal_conv(self.conv_out, sample, causal=self.causal, ltx2=self.ltx2) + return _unpatchify(sample, patch_size_hw=self.patch_size) + + +def _vae_parallel_kwargs(mesh_device: ttnn.MeshDevice) -> dict: + """Build a VaeHWParallelConfig + CCLManager for a test mesh. + + Maps mesh axis 0 → height shard, axis 1 → width shard. For a (1, 1) mesh + both factors are 1 and ``LTXCausalConv3d`` skips halo exchange. + """ + mesh_shape = tuple(mesh_device.shape) + parallel_config = VaeHWParallelConfig( + height_parallel=ParallelFactor(factor=mesh_shape[0], mesh_axis=0), + width_parallel=ParallelFactor(factor=mesh_shape[1], mesh_axis=1), + ) + ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear) + return {"parallel_config": parallel_config, "ccl_manager": ccl_manager} + + +# Single-chip: no fabric (avoids BH fabric router handshake when only one device is opened). +# Stand-alone conv/res/upsample tests use raw (unsharded) tensors — 1x1 only. +_LTX_VAE_MESH_DEVICE_PARAMS = [ + ((1, 1), {}), +] +_LTX_VAE_MESH_DEVICE_IDS = ["1x1"] + +# (num_frames, height, width): H/W divisible by 64; (num_frames - 1) % 8 == 0 for VAE. +_LTX_DECODER_SHAPE_PARAMS = [ + pytest.param(145, 1088, 1920, id="145f_1088x1920"), # 1080p production (19 latent frames) +] + +_LTX_DECODER_2K_SHAPE_PARAMS = [ + pytest.param(9, 1088, 2048, id="9f_1088x2048"), # 2K (1088x2048, latent 34x64) +] + +# Smoke resolution — only shape exercised on a 1x1 mesh in test_ltx_video_decoder. +_LTX_DECODER_1X1_MAX_HW = (256, 256) + + +def _skip_ltx_decoder_if_single_chip_too_large(mesh_device: ttnn.MeshDevice, height: int, width: int) -> None: + if tuple(mesh_device.shape) == (1, 1) and ( + height > _LTX_DECODER_1X1_MAX_HW[0] or width > _LTX_DECODER_1X1_MAX_HW[1] + ): + pytest.skip(f"{height}x{width} LTX decode requires a multi-device mesh") + + +def _unwrap_tt_output(tt_out): + if isinstance(tt_out, (tuple, list)): + # Newer TT path may return auxiliary tensors alongside the primary output. + return tt_out[0] + return tt_out + + +# Full-decoder test — exercises the production decoder_blocks under sharding. +_LTX_PROD_DECODER_BLOCKS = [ + ("res_x", {"num_layers": 4}), + ("compress_space", {"multiplier": 2}), + ("res_x", {"num_layers": 6}), + ("compress_time", {"multiplier": 2}), + ("res_x", {"num_layers": 4}), + ("compress_all", {"multiplier": 1}), + ("res_x", {"num_layers": 2}), + ("compress_all", {"multiplier": 2}), + ("res_x", {"num_layers": 2}), +] + + +@pytest.mark.parametrize( + "in_c, out_c, kernel_size, stride, T, H, W", + [ + (128, 128, 3, 1, 3, 16, 16), # small residual block conv + (128, 256, 3, 1, 3, 16, 16), # channel expansion + (48, 128, 3, 1, 5, 32, 32), # conv_in (after patchify) + (128, 1024, 3, 1, 21, 17, 15), # ltx_s0_conv_in @ 2x4 per-device (1080p path) + (1024, 1024, 3, 1, 21, 17, 15), # ltx_s0_res + (1024, 4096, 3, 1, 21, 17, 15), # ltx_s0_up (upsample projection) + (512, 512, 3, 1, 39, 34, 30), # ltx_s1_res + (128, 48, 3, 1, 147, 136, 120), # ltx_s4_out + ], + ids=[ + "res_128_128", + "expand_128_256", + "conv_in_48_128", + "ltx_s0_conv_in", + "ltx_s0_res", + "ltx_s0_up", + "ltx_s1_res", + "ltx_s4_out", + ], +) +@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"]) +@pytest.mark.parametrize( + "mesh_device, device_params", + _LTX_VAE_MESH_DEVICE_PARAMS, + ids=_LTX_VAE_MESH_DEVICE_IDS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_causal_conv3d( + mesh_device: ttnn.MeshDevice, + in_c: int, + out_c: int, + kernel_size: int, + stride: int, + T: int, + H: int, + W: int, + mean: float, + std: float, +): + """ + Test LTXCausalConv3d against diffusers LTX causal conv reference. + """ + vae_mods = _require_diffusers_ltx_vae() + TorchCausalConv3d = vae_mods["causal_conv"] + ltx2 = vae_mods["ltx2"] + + B = 1 + torch.manual_seed(42) + + conv_kwargs: dict[str, Any] = { + "in_channels": in_c, + "out_channels": out_c, + "kernel_size": kernel_size, + "stride": stride if ltx2 else (stride, stride, stride), + } + if not ltx2: + conv_kwargs["is_causal"] = True + + torch_model = TorchCausalConv3d(**conv_kwargs) + torch_model.eval() + + # TT model + tt_model = LTXCausalConv3d( + in_channels=in_c, + out_channels=out_c, + kernel_size=kernel_size, + stride=stride, + mesh_device=mesh_device, + **_vae_parallel_kwargs(mesh_device), + ) + tt_model.load_torch_state_dict(torch_model.state_dict()) + + # Input + x = torch.randn(B, in_c, T, H, W, dtype=torch.float32) * std + mean + + # PyTorch forward (BCTHW format) + with torch.no_grad(): + torch_out = _run_diffusers_causal_conv(torch_model, x, causal=True, ltx2=ltx2) + + # TT forward (BTHWC format) + x_bthwc = x.permute(0, 2, 3, 4, 1) # (B, T, H, W, C) + x_bthwc = conv_pad_in_channels(x_bthwc) # Pad C to alignment + x_tt = ttnn.from_torch(x_bthwc, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=ttnn.bfloat16) + + tt_out = _unwrap_tt_output(tt_model(x_tt)) + tt_out_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0]) # (B, T_out, H_out, W_out, C_out) + tt_out_torch = tt_out_torch[:, :, :, :, :out_c] # Trim padded channels + tt_out_torch = tt_out_torch.permute(0, 4, 1, 2, 3) # Back to BCTHW + + logger.info(f"PyTorch out: {torch_out.shape}, TT out: {tt_out_torch.shape}") + assert_quality(torch_out, tt_out_torch, pcc=0.999) + logger.info(f"PASSED: LTXCausalConv3d ({in_c}->{out_c}) matches reference") + + +@pytest.mark.parametrize( + "in_c, out_c, T, H, W", + [ + (128, 128, 3, 16, 16), # same channels (no shortcut) + (128, 256, 3, 16, 16), # channel expansion (with shortcut) + (1024, 1024, 21, 17, 15), # ltx_s0_res @ 2x4 per-device + (512, 512, 39, 34, 30), # ltx_s1_res + (256, 256, 147, 70, 66), # ltx_s3_res @ 2K per-device + (128, 128, 147, 138, 130), # ltx_s4_res @ 2K per-device + ], + ids=[ + "same_channels", + "expand_channels", + "ltx_s0_res", + "ltx_s1_res", + "ltx_s3_res_2k", + "ltx_s4_res_2k", + ], +) +@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"]) +@pytest.mark.parametrize( + "mesh_device, device_params", + _LTX_VAE_MESH_DEVICE_PARAMS, + ids=_LTX_VAE_MESH_DEVICE_IDS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_resnet_block( + mesh_device: ttnn.MeshDevice, in_c: int, out_c: int, T: int, H: int, W: int, mean: float, std: float +): + """ + Test LTXResnetBlock3D against diffusers LTX ResnetBlock3D reference. + """ + vae_mods = _require_diffusers_ltx_vae() + TorchResnetBlock3D = vae_mods["resnet"] + ltx2 = vae_mods["ltx2"] + + B = 1 + torch.manual_seed(42) + + torch_model = TorchResnetBlock3D(in_channels=in_c, out_channels=out_c) + torch_model.eval() + + # TT model + tt_model = LTXResnetBlock3D( + in_channels=in_c, + out_channels=out_c, + mesh_device=mesh_device, + **_vae_parallel_kwargs(mesh_device), + ) + tt_model.load_torch_state_dict(_diffusers_resnet_state_to_tt(torch_model.state_dict())) + + # Input + x = torch.randn(B, in_c, T, H, W, dtype=torch.float32) * std + mean + + # PyTorch forward + with torch.no_grad(): + torch_out = _run_diffusers_resnet(torch_model, x, causal=True, ltx2=ltx2) + + # TT forward + x_bthwc = x.permute(0, 2, 3, 4, 1) + x_bthwc = conv_pad_in_channels(x_bthwc) + x_tt = ttnn.from_torch(x_bthwc, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=ttnn.bfloat16) + + tt_out = _unwrap_tt_output(tt_model(x_tt)) + tt_out_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0]) + tt_out_torch = tt_out_torch[:, :, :, :, :out_c] + tt_out_torch = tt_out_torch.permute(0, 4, 1, 2, 3) + + logger.info(f"PyTorch out: {torch_out.shape}, TT out: {tt_out_torch.shape}") + # Lower threshold for channel-expansion blocks (GroupNorm(1) vs layer_norm precision) + min_pcc = 0.995 if in_c != out_c else 0.999 + assert_quality(torch_out, tt_out_torch, pcc=min_pcc) + logger.info(f"PASSED: LTXResnetBlock3D ({in_c}->{out_c}) matches reference") + + +@pytest.mark.parametrize( + "in_c, stride, T, H, W", + [ + (128, (2, 2, 2), 4, 8, 8), # full 3D upsample (compress_all) + (128, (1, 2, 2), 3, 8, 8), # spatial only (compress_space) + (128, (2, 1, 1), 4, 8, 8), # temporal only (compress_time) + (512, (2, 2, 2), 39, 34, 30), # compress_all @ ltx_s1 stage + (1024, (1, 2, 2), 21, 17, 15), # compress_space @ ltx_s0 stage + (256, (2, 1, 1), 75, 68, 60), # compress_time @ ltx_s2 stage + ], + ids=[ + "upsample_all", + "upsample_space", + "upsample_time", + "upsample_all_s1", + "upsample_space_s0", + "upsample_time_s2", + ], +) +@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"]) +@pytest.mark.parametrize( + "mesh_device, device_params", + _LTX_VAE_MESH_DEVICE_PARAMS, + ids=_LTX_VAE_MESH_DEVICE_IDS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_depth_to_space_upsample( + mesh_device: ttnn.MeshDevice, in_c: int, stride: tuple, T: int, H: int, W: int, mean: float, std: float +): + """Test LTXDepthToSpaceUpsample against diffusers LTX upsampler reference.""" + vae_mods = _require_diffusers_ltx_vae() + TorchUpsample = vae_mods["upsample"] + ltx2 = vae_mods["ltx2"] + + B = 1 + torch.manual_seed(42) + + torch_model = TorchUpsample(in_channels=in_c, stride=stride, upscale_factor=1) + torch_model.eval() + + tt_model = LTXDepthToSpaceUpsample( + in_channels=in_c, + stride=stride, + mesh_device=mesh_device, + **_vae_parallel_kwargs(mesh_device), + ) + tt_model.load_torch_state_dict(torch_model.state_dict()) + + x = torch.randn(B, in_c, T, H, W, dtype=torch.float32) * std + mean + with torch.no_grad(): + torch_out = _run_diffusers_upsample(torch_model, x, causal=True, ltx2=ltx2) + + x_bthwc = x.permute(0, 2, 3, 4, 1) + x_bthwc = conv_pad_in_channels(x_bthwc) + x_tt = ttnn.from_torch(x_bthwc, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=ttnn.bfloat16) + + tt_out = _unwrap_tt_output(tt_model(x_tt)) + tt_out_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0]) + out_c = torch_out.shape[1] + tt_out_torch = tt_out_torch[:, :, :, :, :out_c].permute(0, 4, 1, 2, 3) + + logger.info(f"Upsample stride={stride}: {x.shape} -> torch {torch_out.shape}, tt {tt_out_torch.shape}") + assert_quality(torch_out, tt_out_torch, pcc=0.999) + logger.info(f"PASSED: DepthToSpaceUpsample stride={stride}") + + +def _run_ltx_decoder_parity( + mesh_device: ttnn.MeshDevice, + h_axis: int, + w_axis: int, + num_links: int, + num_frames: int, + height: int, + width: int, + *, + mean: float = 0.0, + std: float = 1.0, + pcc: float = 0.99, +): + """Build torch + TT production decoders at ``(num_frames, height, width)`` and + assert numerical parity. Shared by the standard multi-resolution decoder test + and the 2K decoder test. The TT decoder is given num_frames/height/width so the + dim walker fires and ``_BLOCKINGS`` gets an exact-match lookup for every conv3d. + """ + vae_mods = _require_diffusers_ltx_vae() + if not vae_mods["ltx2"]: + pytest.skip("Full decoder parity requires diffusers autoencoder_kl_ltx2 (runtime causal= support)") + + B = 1 + torch.manual_seed(42) + spatial_compression = 32 + latent_frames = (num_frames - 1) // 8 + 1 + latent_h = height // spatial_compression + latent_w = width // spatial_compression + + # PyTorch reference — full production decoder shape (diffusers primitives). + torch_decoder = _TorchLTXVideoDecoder( + decoder_blocks=_LTX_PROD_DECODER_BLOCKS, + in_channels=128, + out_channels=3, + patch_size=4, + base_channels=128, + causal=False, + spatial_padding_mode="zeros", + vae_mods=vae_mods, + ) + torch_decoder.eval() + + # Per-channel stats set to identity so denormalization is a no-op (random + # weights produce garbage stats otherwise). + torch_state = torch_decoder.state_dict() + torch_state["per_channel_statistics.mean-of-means"] = torch.zeros(128) + torch_state["per_channel_statistics.std-of-means"] = torch.ones(128) + torch_decoder.load_state_dict(torch_state) + + # Build VAE parallel config from h_axis / w_axis (mirrors test_wan_decoder). + ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear, num_links=num_links) + parallel_config = VaeHWParallelConfig( + height_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[h_axis], mesh_axis=h_axis), + width_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[w_axis], mesh_axis=w_axis), + ) + + # TT decoder — pass num_frames/height/width so the dim walker fires and + # _BLOCKINGS gets an exact-match lookup for every conv3d. + tt_decoder = LTXVideoDecoder( + decoder_blocks=_LTX_PROD_DECODER_BLOCKS, + in_channels=128, + out_channels=3, + patch_size=4, + base_channels=128, + causal=False, + num_frames=num_frames, + height=height, + width=width, + mesh_device=mesh_device, + parallel_config=parallel_config, + ccl_manager=ccl_manager, + ) + tt_decoder.load_torch_state_dict(_diffusers_decoder_state_to_tt(torch_decoder.state_dict())) + + # Latent shape matches what the pipeline would produce at this resolution. + latent = torch.randn(B, 128, latent_frames, latent_h, latent_w, dtype=torch.float32) * std + mean + + with torch.no_grad(): + torch_out = torch_decoder(latent) + + tt_out = tt_decoder(latent) + + logger.info(f"Decoder: {latent.shape} -> torch {torch_out.shape}, tt {tt_out.shape}") + assert_quality(torch_out, tt_out, pcc=pcc) + logger.info("PASSED: LTXVideoDecoder matches PyTorch reference") + + +@pytest.mark.parametrize("num_frames, height, width", _LTX_DECODER_SHAPE_PARAMS) +@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"]) +@pytest.mark.parametrize( + ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"), + LTX_VAE_DECODER_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_video_decoder( + mesh_device: ttnn.MeshDevice, + num_frames: int, + height: int, + width: int, + mean: float, + std: float, + h_axis: int, + w_axis: int, + num_links: int, + topology: ttnn.Topology, +): + """ + Test full LTXVideoDecoder against PyTorch VideoDecoder reference using the + production LTX-2.3 22B decoder_blocks. Random weights — checks numerical + parity, not absolute output quality. + + For 2x4 this validates the parallel halo-exchange path end-to-end, which + is otherwise only covered by the slow pipeline-level test. + """ + _skip_ltx_decoder_if_single_chip_too_large(mesh_device, height, width) + _run_ltx_decoder_parity(mesh_device, h_axis, w_axis, num_links, num_frames, height, width, mean=mean, std=std) + + +_NUM_FRAMES = 145 +_HEIGHT = 1088 +_WIDTH = 1920 + + +def _build_tt_decoder(mesh): + """Production decoder with random weights; torch module built for its + state_dict only (no host forward).""" + vae_mods = _require_diffusers_ltx_vae() + if not vae_mods["ltx2"]: + pytest.skip("LTX-2 decoder profiling requires diffusers autoencoder_kl_ltx2") + + torch.manual_seed(42) + torch_decoder = _TorchLTXVideoDecoder( + decoder_blocks=_LTX_PROD_DECODER_BLOCKS, + in_channels=128, + out_channels=3, + patch_size=4, + base_channels=128, + causal=False, + spatial_padding_mode="zeros", + vae_mods=vae_mods, + ) + torch_decoder.eval() + state = torch_decoder.state_dict() + # Identity per-channel stats so denorm is a no-op (random weights give garbage stats). + state["per_channel_statistics.mean-of-means"] = torch.zeros(128) + state["per_channel_statistics.std-of-means"] = torch.ones(128) + torch_decoder.load_state_dict(state) + + # num_links=2 mirrors the BH production device config; NP_LINKS probes the W-transport bound. + ccl_manager = CCLManager(mesh, topology=ttnn.Topology.Linear, num_links=int(os.environ.get("NP_LINKS", "2"))) + parallel_config = VaeHWParallelConfig( + height_parallel=ParallelFactor(factor=tuple(mesh.shape)[0], mesh_axis=0), + width_parallel=ParallelFactor(factor=tuple(mesh.shape)[1], mesh_axis=1), + ) + tt_decoder = LTXVideoDecoder( + decoder_blocks=_LTX_PROD_DECODER_BLOCKS, + in_channels=128, + out_channels=3, + patch_size=4, + base_channels=128, + causal=False, + num_frames=_NUM_FRAMES, + height=_HEIGHT, + width=_WIDTH, + mesh_device=mesh, + parallel_config=parallel_config, + ccl_manager=ccl_manager, + ) + tt_decoder.load_torch_state_dict(_diffusers_decoder_state_to_tt(torch_decoder.state_dict())) + return tt_decoder + + +def _latent(): + latent_frames = (_NUM_FRAMES - 1) // 8 + 1 + return torch.randn(1, 128, latent_frames, _HEIGHT // 32, _WIDTH // 32, dtype=torch.float32) + + +@pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True) +@pytest.mark.parametrize("mesh_device", [(4, 8)], indirect=True) +def test_prof_vae_ltx_devicetime(mesh_device, device_params): + mesh = mesh_device.create_submesh(ttnn.MeshShape(4, 8)) + tt_decoder = _build_tt_decoder(mesh) + latent = _latent() + + # One forward only: the decoder emits ~10x a transformer block's ops, so warm-ups + # would overflow the profiler DRAM buffer. Device kernel durations are warm-independent. + signpost("start") + _ = tt_decoder(latent) + ttnn.synchronize_device(mesh) + signpost("stop") + ttnn.ReadDeviceProfiler(mesh) + + +@pytest.mark.parametrize("num_frames, height, width", _LTX_DECODER_2K_SHAPE_PARAMS) +@pytest.mark.parametrize("mean, std", [(0, 1)]) +@pytest.mark.parametrize( + ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"), + LTX_VAE_DECODER_MULTI_ONLY_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_video_decoder_2k( + mesh_device: ttnn.MeshDevice, + num_frames: int, + height: int, + width: int, + mean: float, + std: float, + h_axis: int, + w_axis: int, + num_links: int, + topology: ttnn.Topology, +): + """Full LTXVideoDecoder at the 2K production resolution (2048x1080 -> 1088x2048). + + Multi-device only (2x4 is the production layout; 4x8 also exercised). The two + 512->4096 / 1024->4096 upsample-projection convs are not in the 2K ``_BLOCKINGS`` + table (they OOM a 2x4 mesh at full production T), so at 2K they fall back to the + generic blockings — correct, but not yet perf-tuned. See conv3d.py. + """ + _run_ltx_decoder_parity(mesh_device, h_axis, w_axis, num_links, num_frames, height, width, mean=mean, std=std) + + +# --------------------------------------------------------------------------- +# Encoder parity (LTXVideoEncoder vs a diffusers-primitives reference encoder) +# --------------------------------------------------------------------------- +def _diffusers_encoder_state_to_tt(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Remap diffusers LTX-2 encoder keys (incl. nested resnets) to the TT/Lightricks layout.""" + out: dict[str, torch.Tensor] = {} + for key, value in state.items(): + if key.endswith("conv_shortcut.weight"): + out[key.replace("conv_shortcut.weight", "conv_shortcut.conv.weight")] = value + elif key.endswith("conv_shortcut.bias"): + out[key.replace("conv_shortcut.bias", "conv_shortcut.conv.bias")] = value + else: + out[key] = value + return out + + +class _TorchLTXVideoEncoder(nn.Module): + """Block-list LTX encoder built from diffusers VAE primitives (matches TT/Lightricks layout). + + Mirrors the reference ``VideoEncoder`` encode path: patchify -> conv_in -> down_blocks + (``res_x`` mid blocks + ``compress_*_res`` space-to-depth downsamplers) -> PixelNorm + SiLU + -> conv_out (latent means) -> per-channel normalize. Always causal. + """ + + _STRIDE_MAP = { + "compress_time_res": (2, 1, 1), + "compress_space_res": (1, 2, 2), + "compress_all_res": (2, 2, 2), + } + + def __init__( + self, + *, + encoder_blocks: list[tuple[str, dict | int]], + in_channels: int = 3, + out_channels: int = 128, + patch_size: int = 4, + base_channels: int = 128, + causal: bool = True, + spatial_padding_mode: str = "zeros", + vae_mods: dict, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.causal = causal + self.ltx2 = vae_mods["ltx2"] + causal_conv = vae_mods["causal_conv"] + resnet_cls = vae_mods["resnet"] + downsample_cls = vae_mods["downsample"] + pixel_norm = vae_mods["pixel_norm"] + + in_channels_patched = in_channels * patch_size**2 + feature_channels = base_channels + + conv_kwargs: dict[str, Any] = {} + if self.ltx2: + if spatial_padding_mode != "zeros": + conv_kwargs["spatial_padding_mode"] = spatial_padding_mode + else: + conv_kwargs["is_causal"] = causal + + self.conv_in = causal_conv( + in_channels=in_channels_patched, + out_channels=feature_channels, + kernel_size=3, + stride=1, + **conv_kwargs, + ) + + self.down_blocks = nn.ModuleList() + ch = feature_channels + for block_name, block_params in encoder_blocks: + block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params + if block_name == "res_x": + self.down_blocks.append( + _TorchUNetMidBlock3D( + in_channels=ch, + num_layers=block_config["num_layers"], + resnet_cls=resnet_cls, + ltx2=self.ltx2, + spatial_padding_mode=spatial_padding_mode, + ) + ) + elif block_name in self._STRIDE_MAP: + multiplier = block_config.get("multiplier", 2) + new_ch = ch * multiplier + downsample_kwargs: dict[str, Any] = { + "in_channels": ch, + "out_channels": new_ch, + "stride": self._STRIDE_MAP[block_name], + } + if self.ltx2 and spatial_padding_mode != "zeros": + downsample_kwargs["spatial_padding_mode"] = spatial_padding_mode + self.down_blocks.append(downsample_cls(**downsample_kwargs)) + ch = new_ch + else: + raise ValueError(f"Unknown encoder block: {block_name}") + + self.conv_norm_out = pixel_norm() + self.conv_act = nn.SiLU() + self.conv_out = causal_conv( + in_channels=ch, + out_channels=out_channels, + kernel_size=3, + stride=1, + **conv_kwargs, + ) + self.per_channel_statistics = _PerChannelStatistics(out_channels) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + # Crop trailing frames so F == 1 + 8*k (mirrors the reference; no-op for the test shapes). + frames_count = sample.shape[2] + if (frames_count - 1) % 8 != 0: + sample = sample[:, :, : frames_count - ((frames_count - 1) % 8)] + + sample = _patchify(sample, patch_size_hw=self.patch_size) + sample = _run_diffusers_causal_conv(self.conv_in, sample, causal=self.causal, ltx2=self.ltx2) + for down_block in self.down_blocks: + if isinstance(down_block, _TorchUNetMidBlock3D): + sample = down_block(sample, causal=self.causal, ltx2=self.ltx2) + else: + sample = _run_diffusers_downsample(down_block, sample, causal=self.causal, ltx2=self.ltx2) + if self.ltx2: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample.movedim(1, -1)).movedim(-1, 1) + sample = self.conv_act(sample) + sample = _run_diffusers_causal_conv(self.conv_out, sample, causal=self.causal, ltx2=self.ltx2) + return self.per_channel_statistics.normalize(sample) + + +# Small encoder config: conv_in (48->128) -> res_x -> compress_space_res -> compress_time_res. +# Avoids resnet channel-shortcut paths (res_x keeps channels; compress_*_res change channels via +# space-to-depth), so the parity check isolates the new encoder primitives. +_LTX_ENCODER_SMALL_BLOCKS = [ + ("res_x", {"num_layers": 1}), + ("compress_space_res", {"multiplier": 2}), + ("compress_time_res", {"multiplier": 2}), +] + + +@pytest.mark.parametrize( + "num_frames, height, width", + [(9, 64, 64)], + ids=["9f_64x64"], +) +@pytest.mark.parametrize( + "mesh_device, device_params", + _LTX_VAE_MESH_DEVICE_PARAMS, + ids=_LTX_VAE_MESH_DEVICE_IDS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_video_encoder(mesh_device: ttnn.MeshDevice, num_frames: int, height: int, width: int): + """LTXVideoEncoder parity vs a diffusers-primitives reference encoder. + + Random weights — checks the encode path (patchify -> conv_in -> down_blocks -> PixelNorm+SiLU + -> conv_out (means) -> per-channel normalize) matches the reference numerically. + """ + vae_mods = _require_diffusers_ltx_vae() + if not vae_mods["ltx2"]: + pytest.skip("Encoder parity requires diffusers autoencoder_kl_ltx2 (runtime causal= support)") + + torch.manual_seed(42) + latent_channels = 128 + + ref = _TorchLTXVideoEncoder( + encoder_blocks=_LTX_ENCODER_SMALL_BLOCKS, + in_channels=3, + out_channels=latent_channels, + patch_size=4, + base_channels=latent_channels, + causal=True, + spatial_padding_mode="zeros", + vae_mods=vae_mods, + ) + # Set known per-channel stats so normalize() = (x - mean) / std is deterministic/well-conditioned. + with torch.no_grad(): + ref.per_channel_statistics.get_buffer("mean-of-means").copy_(torch.randn(latent_channels)) + ref.per_channel_statistics.get_buffer("std-of-means").copy_(torch.rand(latent_channels) + 0.5) + ref.eval() + + x = torch.randn(1, 3, num_frames, height, width, dtype=torch.float32) + with torch.no_grad(): + ref_out = ref(x.clone()) # (1, 128, F', H', W') + + tt_model = LTXVideoEncoder( + encoder_blocks=_LTX_ENCODER_SMALL_BLOCKS, + in_channels=3, + out_channels=latent_channels, + patch_size=4, + mesh_device=mesh_device, + num_frames=num_frames, + height=height, + width=width, + **_vae_parallel_kwargs(mesh_device), + ) + tt_model.load_torch_state_dict(_diffusers_encoder_state_to_tt(ref.state_dict())) + + tt_out = _unwrap_tt_output(tt_model(x.clone())) # host (B, C, F', H', W') + + logger.info(f"ref out {tuple(ref_out.shape)}, TT out {tuple(tt_out.shape)}") + assert tt_out.shape == ref_out.shape, f"shape mismatch: TT {tt_out.shape} vs ref {ref_out.shape}" + assert torch.isfinite(tt_out).all(), "NaN/Inf in TT encoder output" + assert_quality(ref_out, tt_out, pcc=0.99) + logger.info("PASSED: LTXVideoEncoder matches diffusers reference encoder") + + +# --------------------------------------------------------------------------- +# Production encoder parity under sharding (regression for the I2V VAE-encoder bug) +# --------------------------------------------------------------------------- +# Full LTX-2.3 22B encoder block list (from the checkpoint config metadata). The two +# ``compress_all_res`` blocks are NOT exercised by ``test_ltx_video_encoder`` above, and +# only matter under spatial sharding — they fold a (2,2,2) patch into channels per device. +_LTX_PROD_ENCODER_BLOCKS = [ + ("res_x", {"num_layers": 4}), + ("compress_space_res", {"multiplier": 2}), + ("res_x", {"num_layers": 6}), + ("compress_time_res", {"multiplier": 2}), + ("res_x", {"num_layers": 4}), + ("compress_all_res", {"multiplier": 2}), + ("res_x", {"num_layers": 2}), + ("compress_all_res", {"multiplier": 1}), + ("res_x", {"num_layers": 2}), +] + + +def _run_ltx_encoder_parity( + mesh_device: ttnn.MeshDevice, + h_axis: int, + w_axis: int, + num_links: int, + num_frames: int, + height: int, + width: int, + *, + pcc: float = 0.99, +): + """Build torch (diffusers reference) + TT production encoders at ``(num_frames, height, width)`` + and assert numerical parity. The TT encoder shards H/W across the mesh exactly like the I2V + pipeline; this is the path that fails when a per-device spatial dim is not divisible by the + downsample stride (patch groups would straddle a shard boundary).""" + vae_mods = _require_diffusers_ltx_vae() + if not vae_mods["ltx2"]: + pytest.skip("Encoder parity requires diffusers autoencoder_kl_ltx2 (runtime causal= support)") + + torch.manual_seed(42) + latent_channels = 128 + + ref = _TorchLTXVideoEncoder( + encoder_blocks=_LTX_PROD_ENCODER_BLOCKS, + in_channels=3, + out_channels=latent_channels, + patch_size=4, + base_channels=latent_channels, + causal=True, + spatial_padding_mode="zeros", + vae_mods=vae_mods, + ) + with torch.no_grad(): + ref.per_channel_statistics.get_buffer("mean-of-means").copy_(torch.randn(latent_channels)) + ref.per_channel_statistics.get_buffer("std-of-means").copy_(torch.rand(latent_channels) + 0.5) + ref.eval() + + x = torch.randn(1, 3, num_frames, height, width, dtype=torch.float32) + with torch.no_grad(): + ref_out = ref(x.clone()) # (1, 128, F', H', W') + + ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear, num_links=num_links) + parallel_config = VaeHWParallelConfig( + height_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[h_axis], mesh_axis=h_axis), + width_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[w_axis], mesh_axis=w_axis), + ) + + tt_model = LTXVideoEncoder( + encoder_blocks=_LTX_PROD_ENCODER_BLOCKS, + in_channels=3, + out_channels=latent_channels, + patch_size=4, + mesh_device=mesh_device, + num_frames=num_frames, + height=height, + width=width, + parallel_config=parallel_config, + ccl_manager=ccl_manager, + ) + tt_model.load_torch_state_dict(_diffusers_encoder_state_to_tt(ref.state_dict())) + + tt_out = _unwrap_tt_output(tt_model(x.clone())) + + logger.info(f"prod encoder: ref out {tuple(ref_out.shape)}, TT out {tuple(tt_out.shape)}") + assert tt_out.shape == ref_out.shape, f"shape mismatch: TT {tt_out.shape} vs ref {ref_out.shape}" + assert torch.isfinite(tt_out).all(), "NaN/Inf in TT encoder output" + assert_quality(ref_out, tt_out, pcc=pcc) + logger.info("PASSED: production LTXVideoEncoder matches diffusers reference encoder") + + +# (num_frames, H, W). On a 4x8 mesh (H//4, W//8) the per-device spatial dim after patchify and +# the early downsamples decides whether the last compress_all_res can fold locally: +# - "clean": per-device dims stay even through every fold (fast local path). +# - "boundary": per-device dim becomes odd at the last compress_all_res, so the space-to-depth +# patch would straddle a shard boundary — the case that corrupts the I2V latent. +_LTX_ENCODER_PROD_SHAPE_PARAMS = [ + pytest.param(1, 128, 256, id="1f_128x256_clean"), + pytest.param(1, 192, 384, id="1f_192x384_boundary"), + # The two real I2V conditioning-image resolutions (single frame): stage-0 (544x960 -> latent + # 17x30) and stage-1 (1088x1920 -> latent 34x60). Both put an odd per-device spatial dim into + # the last compress_all_res on 4x8 — this is the exact case that corrupted the I2V latent. + pytest.param(1, 544, 960, id="1f_544x960_stage0"), + pytest.param(1, 1088, 1920, id="1f_1088x1920_stage1"), +] + + +@pytest.mark.parametrize("num_frames, height, width", _LTX_ENCODER_PROD_SHAPE_PARAMS) +@pytest.mark.parametrize( + ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"), + LTX_VAE_ENCODER_PROD_MESH_PARAMS, + indirect=["mesh_device", "device_params"], +) +def test_ltx_video_encoder_prod( + mesh_device: ttnn.MeshDevice, + num_frames: int, + height: int, + width: int, + h_axis: int, + w_axis: int, + num_links: int, + topology: ttnn.Topology, +): + """Production-config LTXVideoEncoder parity vs the diffusers reference, sharded like the I2V pipeline. + + The "boundary" resolution reproduces the I2V conditioning-latent corruption (device-vs-host + PCC ~0.3): a per-device spatial dim becomes odd at the second ``compress_all_res``, so the + per-device space-to-depth folds a patch that should straddle two shards. + """ + _run_ltx_encoder_parity(mesh_device, h_axis, w_axis, num_links, num_frames, height, width) diff --git a/code/models/tt_dit/tests/models/minimax_h3/common.py b/code/models/tt_dit/tests/models/minimax_h3/common.py new file mode 100644 index 0000000000000000000000000000000000000000..6b6bcbcd9ac67954c6f4ded78960a970841c0eb7 --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/common.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + +"""Shared helpers for the MiniMax-H3 bringup tests.""" + +import json +import math +import os + +import numpy as np +import pytest +import torch +from PIL import Image + +import ttnn + +from ....utils.tensor import from_torch +from ....utils.test import ring_params_8k_req_exact_devices, ring_params_req_exact_devices + +# Fixed VAE work units: encoder (17, 256, 256) tiles, decoder (7, 16, 16) latent chunks. +TILE = 256 +CLIP_FRAMES = 17 +LATENT_TILE = 16 +DECODE_LATENT_FRAMES = 7 + + +def weights_subdir(subfolder: str) -> str | None: + base = os.environ.get("MINIMAX_H3_MODEL_PATH") + if not base: + return None + candidate = os.path.join(base, subfolder) + return candidate if os.path.isfile(os.path.join(candidate, "config.json")) else None + + +def load_config(weights_dir: str) -> dict: + return { + k: v + for k, v in json.loads(open(os.path.join(weights_dir, "config.json")).read()).items() + if not k.startswith("_") + } + + +def _reference_class(name: str): + pytest.importorskip("diffusers", reason="pinned diffusers reference not installed") + from diffusers.models.autoencoders import autoencoder_kl_minimax_h3 as ref + + return getattr(ref, name) + + +def random_encoder_state(config: dict) -> dict: + cls = _reference_class("MiniMaxH3VideoEncoder3d") + module = cls( + in_channels=3, + out_channels=2 * config["latent_channels"], + block_out_channels=tuple(config["block_out_channels"]), + layers_per_block=config["layers_per_block"], + spatial_downsample_factors=tuple(config["spatial_downsample_factors"]), + temporal_downsample_factors=tuple(config["temporal_downsample_factors"]), + norm_num_groups=config["norm_num_groups"], + norm_eps=config["norm_eps"], + spatial_padding_mode=config["spatial_padding_mode"], + ) + return dict(module.state_dict()) + + +def random_decoder_state(config: dict, *, num_layers: int | None = None) -> dict: + cls = _reference_class("MiniMaxH3VideoViTDecoder3d") + module = cls( + in_channels=config["latent_channels"], + out_channels=config["out_channels"], + patch_size=16, + patch_size_t=4, + num_layers=config["decoder_num_layers"] if num_layers is None else num_layers, + num_attention_heads=config["decoder_num_attention_heads"], + attention_head_dim=config["decoder_attention_head_dim"], + num_register_tokens=config["decoder_num_register_tokens"], + ffn_mult=config["decoder_ffn_mult"], + rope_theta=config["decoder_rope_theta"], + rope_dim_ratio=config["decoder_rope_dim_ratio"], + norm_eps=config["decoder_norm_eps"], + ) + return dict(module.state_dict()) + + +def psnr(reference: torch.Tensor, test: torch.Tensor) -> float: + """PSNR in dB, with the peak taken from the reference's own range.""" + mse = torch.mean((reference.float() - test.float()) ** 2).item() + if mse == 0.0: + return float("inf") + peak = reference.abs().max().item() + return float("inf") if peak == 0.0 else 20.0 * math.log10(peak) - 10.0 * math.log10(mse) + + +def create_fractal_image(width: int, height: int) -> Image.Image: + """Mandelbrot I2V seed: content the model would never generate, so a pipeline that ignores the keyframe cannot pass.""" + c = np.linspace(-2.0, 1.0, width)[None, :] + 1j * np.linspace(-1.5, 1.5, height)[:, None] + z = np.zeros_like(c) + img = np.zeros(c.shape, dtype=np.uint8) + for i in range(32): + z = z * z + c + img[(img == 0) & (np.abs(z) > 2)] = 255 - 8 * i + return Image.fromarray(np.dstack((img, np.roll(img, width // 10, 1), np.roll(img, height // 10, 0))), "RGB") + + +# The mesh shapes MiniMax-H3 is tuned for, named individually and composed into a list -- the shape +# `ltx_mesh_params.py` uses. Both carry `require_exact_physical_num_devices`, so exactly one runs on +# any given cluster and the other skips. +# +# The fabric must be FABRIC_1D_RING: `CCLManager` runs ring collectives, and on a line fabric one +# cannot resolve a forwarding direction (`TT_FATAL fabric.cpp:174 forwarding_direction.has_value()`). +# 4x32 additionally takes the 8 KB router payload, matching Wan's 4x32 rows, and a trace region for +# the quad's `trace_denoise`; the region is only reserved, so 4x8 pays nothing but address space. +_L1_SMALL = 65536 +_ring = {**ring_params_req_exact_devices, "l1_small_size": _L1_SMALL} +_ring_8k_trace = {**ring_params_8k_req_exact_devices, "trace_region_size": 150_000_000, "l1_small_size": _L1_SMALL} + +MESH_4X8_RING = pytest.param((4, 8), _ring, id="4x8") +MESH_4X32_RING = pytest.param((4, 32), _ring_8k_trace, id="4x32") + +GALAXY_MESHES = [MESH_4X8_RING, MESH_4X32_RING] + + +def randomize_norm_weights(module: torch.nn.Module, *, scale: float = 0.5) -> torch.nn.Module: + """Randomize every `nn.RMSNorm` affine in place, BEFORE taking `state_dict`: all-ones norms make norm-weight loading invisible to PCC.""" + for submodule in module.modules(): + if isinstance(submodule, torch.nn.RMSNorm) and submodule.weight is not None: + submodule.weight.data = 1.0 + scale * torch.randn_like(submodule.weight.data) + return module + + +# ---- transformer test fixtures, shared by the correctness and perf tests so they cannot drift ---- + +# The transformer tests take the axes explicitly: TP stays on axis 0 at factor 4 and SP absorbs the +# rest, so 4x8 -> 4x32 moves only `sp_factor`, which every test body derives from `mesh_device.shape`. +# `device_params` travels inside the tuple because the router payload differs per shape; crossing them +# independently would pair a 4x8 mesh with the 4x32 router config. +# The transformer tests take the axes explicitly: TP stays on axis 0 at factor 4 and SP absorbs the +# rest, so 4x8 -> 4x32 moves only `sp_factor`, which every test body derives from `mesh_device.shape`. +# `device_params` travels inside the tuple because the router payload differs per shape; crossing them +# independently would pair a 4x8 mesh with the 4x32 router config. +GALAXY_RING = pytest.mark.parametrize( + ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"), + [ + pytest.param((4, 8), 1, 0, 2, _ring, ttnn.Topology.Ring, False, id="4x8sp1tp0nl2_ring_is_fsdp0"), + pytest.param((4, 32), 1, 0, 2, _ring_8k_trace, ttnn.Topology.Ring, False, id="4x32sp1tp0nl2_ring_is_fsdp0"), + ], + indirect=["mesh_device", "device_params"], +) + +REAL_BLOCK_CONFIG = dict( + hidden_size=5376, + num_attention_heads=56, + attention_head_dim=128, + ffn_dim=14336, + time_embed_dim=2688, + norm_eps=1e-5, + qk_norm_eps=1e-5, +) +ROPE_FREQ_DIM = 16 +ROPE_THETA = 10000.0 + +TT_BLOCK_CONFIG = dict( + hidden_size=REAL_BLOCK_CONFIG["hidden_size"], + num_heads=REAL_BLOCK_CONFIG["num_attention_heads"], + head_dim=REAL_BLOCK_CONFIG["attention_head_dim"], + ffn_dim=REAL_BLOCK_CONFIG["ffn_dim"], + time_embed_dim=REAL_BLOCK_CONFIG["time_embed_dim"], + norm_eps=REAL_BLOCK_CONFIG["norm_eps"], + qk_norm_eps=REAL_BLOCK_CONFIG["qk_norm_eps"], +) + +# Reference token tags: 0 video, 1 text, 2 audio. +TAG_VIDEO, TAG_TEXT, TAG_AUDIO = 0, 1, 2 + + +def packed_layout( + num_text: int, + num_audio: int, + num_video: int, + grid_hw: tuple[int, int] = (8, 8), + padded_len: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build one packed layout `(position_ids, token_tags, timestep_indices)`: text, audio, video rows over two timesteps, zero-padded to `padded_len`.""" + grid_h, grid_w = grid_hw + frame = grid_h * grid_w + assert num_video % frame == 0, "num_video must fill whole (h, w) frames" + grid_t = num_video // frame + assert grid_t >= 2, "need at least one conditioning frame and one target frame" + + tags = torch.cat( + [ + torch.full((num_text,), TAG_TEXT, dtype=torch.long), + torch.full((num_audio,), TAG_AUDIO, dtype=torch.long), + torch.full((num_video,), TAG_VIDEO, dtype=torch.long), + ] + ) + timestep_indices = torch.cat( + [ + torch.zeros(num_text, dtype=torch.long), + torch.ones(num_audio, dtype=torch.long), + torch.zeros(frame, dtype=torch.long), + torch.ones(num_video - frame, dtype=torch.long), + ] + ) + + vt, vh, vw = torch.meshgrid(torch.arange(grid_t), torch.arange(grid_h), torch.arange(grid_w), indexing="ij") + video_pos = torch.stack([vt.reshape(-1), vh.reshape(-1), vw.reshape(-1)], dim=-1) + + def clock_pos(n: int) -> torch.Tensor: + return torch.stack([torch.arange(n), torch.zeros(n, dtype=torch.long), torch.zeros(n, dtype=torch.long)], -1) + + position_ids = torch.cat([clock_pos(num_text), clock_pos(num_audio), video_pos], dim=0) + + if padded_len is not None: + pad = padded_len - position_ids.shape[0] + if pad: + position_ids = torch.cat([position_ids, torch.zeros((pad, 3), dtype=position_ids.dtype)]) + tags = torch.cat([tags, torch.zeros(pad, dtype=tags.dtype)]) + timestep_indices = torch.cat([timestep_indices, torch.zeros(pad, dtype=timestep_indices.dtype)]) + return position_ids, tags, timestep_indices + + +def upload_rope(rope_cos: torch.Tensor, rope_sin: torch.Tensor, *, mesh_device, sp_axis: int): + """Upload prepared `(cos, sin)` tables: fp32, fractured on SP along the sequence, replicated on TP.""" + + def _upload(table: torch.Tensor) -> ttnn.Tensor: + return from_torch( + table.reshape(1, 1, *table.shape), + device=mesh_device, + dtype=ttnn.float32, + mesh_axes=[..., sp_axis, None], + ) + + return _upload(rope_cos), _upload(rope_sin) + + +def build_visual_encoder(config: dict, mesh_device, num_frames: int, **overrides): + from ....models.vae.minimax_h3.encoder_minimax_h3 import MiniMaxH3Encoder3d + + kwargs = dict( + num_frames=num_frames, + height=TILE, + width=TILE, + in_channels=3, + out_channels=2 * config["latent_channels"], + block_out_channels=tuple(config["block_out_channels"]), + layers_per_block=config["layers_per_block"], + spatial_downsample_factors=tuple(config["spatial_downsample_factors"]), + temporal_downsample_factors=tuple(config["temporal_downsample_factors"]), + mesh_device=mesh_device, + ) + kwargs.update(overrides) + return MiniMaxH3Encoder3d(**kwargs) + + +def build_visual_decoder(config: dict, mesh_device, **overrides): + from ....models.vae.minimax_h3.decoder_minimax_h3 import MiniMaxH3ViTDecoder3d + + kwargs = dict( + num_frames=DECODE_LATENT_FRAMES, + height=LATENT_TILE, + width=LATENT_TILE, + in_channels=config["latent_channels"], + out_channels=config["out_channels"], + patch_size=16, + patch_size_t=4, + num_layers=config["decoder_num_layers"], + num_heads=config["decoder_num_attention_heads"], + head_dim=config["decoder_attention_head_dim"], + num_register_tokens=config["decoder_num_register_tokens"], + ffn_mult=config["decoder_ffn_mult"], + rope_theta=config["decoder_rope_theta"], + rope_dim_ratio=config["decoder_rope_dim_ratio"], + eps=config["decoder_norm_eps"], + mesh_device=mesh_device, + ) + kwargs.update(overrides) + return MiniMaxH3ViTDecoder3d(**kwargs) + + +def build_audio_decoder(config: dict, mesh_device, **overrides): + from ....models.audio_vae.minimax_h3.decoder_minimax_h3_audio import MiniMaxH3AudioDecoder + + kwargs = dict( + latent_channels=config["latent_channels"], + latent_dim=config["latent_dim"], + decoder_dim=config["decoder_dim"], + decoder_rates=tuple(config["decoder_rates"]), + decoder_kernel_sizes=tuple(config["decoder_kernel_sizes"]), + resblock_kernel_sizes=tuple(config["resblock_kernel_sizes"]), + resblock_dilation_sizes=tuple(tuple(d) for d in config["resblock_dilation_sizes"]), + mesh_device=mesh_device, + ) + kwargs.update(overrides) + return MiniMaxH3AudioDecoder(**kwargs) + + +CONDITIONER_SUBFOLDER = "text_encoder" + + +def conditioner_checkpoint_dir(patterns: list[str]) -> str: + """Resolve the conditioner directory from $MINIMAX_H3_MODEL_PATH; a missing checkpoint skips. + + `patterns` are the shard patterns the test needs present (validation only, nothing is fetched). + """ + import glob + + root = os.environ.get("MINIMAX_H3_MODEL_PATH", "") + if not root or not os.path.isdir(root): + pytest.skip("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot") + missing = [pattern for pattern in patterns if not glob.glob(os.path.join(root, pattern))] + if missing: + pytest.skip(f"MiniMax-H3 conditioner checkpoint at {root} is missing {missing}") + return os.path.join(root, CONDITIONER_SUBFOLDER) + + +def load_reference_conditioner(path: str): + """Load the conditioner with load-info checked: a silently partial load would pass parity while testing the fresh init.""" + import transformers + + hf, info = transformers.Qwen3VLForConditionalGeneration.from_pretrained( + path, dtype=torch.bfloat16, output_loading_info=True + ) + bad = {k: sorted(info[k])[:5] for k in ("missing_keys", "unexpected_keys", "mismatched_keys") if info[k]} + assert not bad, f"conditioner load key mismatch: {bad}" + return hf diff --git a/code/models/tt_dit/tests/models/minimax_h3/common_av.py b/code/models/tt_dit/tests/models/minimax_h3/common_av.py new file mode 100644 index 0000000000000000000000000000000000000000..320a5496047ec0a01bf7741bdf166d15f21d8e0a --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/common_av.py @@ -0,0 +1,563 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""Reference-free A/V sanity checks plus the scaffolding shared by the MiniMax-H3 e2e gates. +A/V sync is checked structurally (duration/ordering); envelope-vs-motion correlation is diagnostic only.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +from pathlib import Path + +import numpy as np +import pytest +import torch +from loguru import logger + +import ttnn + +from ....pipelines.minimax_h3.packing import MINIMAX_H3_FPS + + +def check_audio_sanity(audio, *, sampling_rate, expected_seconds, tolerance_seconds=0.05): + """Guard against a soundtrack that is silent, clipped, constant, or the wrong length.""" + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu().numpy() + audio = np.asarray(audio) + if audio.ndim == 3: + assert audio.shape[0] == 1, f"expected a single generation, got batch {audio.shape[0]}" + audio = audio[0] + assert audio.ndim == 2, f"expected (channels, samples), got shape {audio.shape}" + + channels, samples = audio.shape + assert channels == 2, f"H3 generates stereo; got {channels} channel(s)" + assert np.isfinite(audio).all(), "soundtrack contains NaN/Inf" + + seconds = samples / sampling_rate + assert abs(seconds - expected_seconds) <= tolerance_seconds, ( + f"soundtrack is {seconds:.3f} s against {expected_seconds:.3f} s of video " + f"(off by {seconds - expected_seconds:+.3f} s, tolerance {tolerance_seconds:.3f} s)" + ) + + peak = float(np.abs(audio).max()) + rms = float(np.sqrt((audio.astype(np.float64) ** 2).mean())) + assert peak > 1e-3, f"soundtrack is silent (peak {peak:.2e})" + assert rms > 1e-4, f"soundtrack is near-silent (rms {rms:.2e})" + for index in range(channels): + channel_std = float(audio[index].std()) + assert channel_std > 1e-4, f"channel {index} is constant (std {channel_std:.2e})" + # Widespread clipping means the denormalization is wrong, not that the mix is loud. + clipped = float((np.abs(audio) >= 0.999).mean()) + assert clipped < 0.01, f"{clipped:.1%} of samples are at full scale; suspect a scaling error" + + logger.info( + f"Audio sanity OK: {channels}ch {seconds:.3f} s @ {sampling_rate} Hz, " + f"peak={peak:.3f}, rms={rms:.4f}, clipped={clipped:.3%}" + ) + + +def check_av_sync(frames, audio, *, sampling_rate, fps, tolerance_seconds=0.05): + """Durations must agree and stereo channels must differ -- identical channels mean + `unpack_audio_tokens` collapsed the channel-major layout. Envelope-motion lag is logged only.""" + if isinstance(frames, torch.Tensor): + frames = frames.detach().cpu().numpy() + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu().numpy() + frames = np.asarray(frames) + audio = np.asarray(audio) + if audio.ndim == 3: + audio = audio[0] + + num_frames = frames.shape[0] + video_seconds = num_frames / fps + audio_seconds = audio.shape[-1] / sampling_rate + assert num_frames > 1, f"need more than one frame to talk about sync, got {num_frames}" + assert audio.shape[-1] > sampling_rate // fps, "soundtrack is shorter than a single frame" + assert abs(video_seconds - audio_seconds) <= tolerance_seconds, ( + f"video is {video_seconds:.3f} s but audio is {audio_seconds:.3f} s " + f"(off by {audio_seconds - video_seconds:+.3f} s)" + ) + + if audio.shape[0] == 2: + assert not np.allclose(audio[0], audio[1]), "stereo channels are identical; suspect audio row unpacking" + + # Diagnostic only. + motion = np.abs(np.diff(frames.astype(np.float32), axis=0)).mean(axis=(1, 2, 3)) + samples_per_frame = audio.shape[-1] / num_frames + envelope = np.array( + [ + np.abs(audio[:, int(i * samples_per_frame) : int((i + 1) * samples_per_frame)]).mean() + for i in range(1, num_frames) + ] + ) + lag_frames = 0 + if motion.std() > 1e-8 and envelope.std() > 1e-8: + m = (motion - motion.mean()) / motion.std() + e = (envelope - envelope.mean()) / envelope.std() + correlation = np.correlate(m, e, mode="full") / len(m) + lag_frames = int(np.argmax(correlation)) - (len(m) - 1) + logger.info( + f"A/V envelope-motion best lag: {lag_frames:+d} frames ({lag_frames / fps:+.3f} s), " + f"peak r={correlation.max():.3f} (diagnostic, not asserted)" + ) + + logger.info( + f"A/V sync OK: video {video_seconds:.3f} s / {num_frames} frames @ {fps} fps, " + f"audio {audio_seconds:.3f} s @ {sampling_rate} Hz, delta {audio_seconds - video_seconds:+.4f} s" + ) + return {"video_seconds": video_seconds, "audio_seconds": audio_seconds, "lag_frames": lag_frames} + + +def check_spatial_seams(frames, *, vertical_boundaries, horizontal_boundaries, max_ratio=2.0): + """Gradient energy at the VAE's tile boundaries vs everywhere else (~1.0 = no seam; whole-frame means hide seams).""" + if isinstance(frames, torch.Tensor): + frames = frames.detach().cpu().numpy() + frames = np.asarray(frames).astype(np.float32) + if frames.ndim == 4: + frames = frames.mean(axis=-1) + + def ratio(gradient, boundaries): + inside = np.array([b for b in boundaries if 1 <= b < len(gradient) - 1], dtype=int) + if not len(inside): + return float("nan") + # Exclude +-2 px around each boundary so a smeared seam cannot inflate its own baseline. + baseline = np.ones(len(gradient), dtype=bool) + for b in inside: + baseline[max(0, b - 2) : b + 3] = False + if not baseline.any() or gradient[baseline].mean() == 0: + return float("nan") + return float(gradient[inside].mean() / gradient[baseline].mean()) + + column_gradient = np.abs(np.diff(frames, axis=2)).mean(axis=(0, 1)) + row_gradient = np.abs(np.diff(frames, axis=1)).mean(axis=(0, 2)) + vertical = ratio(column_gradient, vertical_boundaries) + horizontal = ratio(row_gradient, horizontal_boundaries) + + logger.info( + f"Spatial seam ratios (1.0 = no seam): vertical {vertical:.3f} at x={list(vertical_boundaries)}, " + f"horizontal {horizontal:.3f} at y={list(horizontal_boundaries)}" + ) + for name, value in (("vertical", vertical), ("horizontal", horizontal)): + if np.isfinite(value): + assert value < max_ratio, ( + f"{name} tile-boundary gradient is {value:.2f}x the surrounding image; " + "suspect the tile blend extent or per-tile normalization (artifact rubric: seams)" + ) + return {"vertical": vertical, "horizontal": horizontal} + + +def log_spectral_flatness(audio, *, sampling_rate, num_bands=64): + """Log-spectrum shape, logged only: white noise drives flatness toward 1.0, a stuck tone toward 0.""" + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu().numpy() + audio = np.asarray(audio) + if audio.ndim == 3: + audio = audio[0] + mono = audio.mean(axis=0).astype(np.float64) + + window = 2048 + hop = window // 2 + frames = [mono[i : i + window] for i in range(0, max(1, len(mono) - window), hop)] + if not frames: + return {"flatness": float("nan")} + spectrum = np.abs(np.fft.rfft(np.stack(frames) * np.hanning(window), axis=-1)) ** 2 + power = spectrum.mean(axis=0)[1:] + flatness = float(np.exp(np.log(power + 1e-20).mean()) / (power.mean() + 1e-20)) + band_edges = np.linspace(0, len(power), num_bands + 1).astype(int) + bands = np.array([power[a:b].mean() for a, b in zip(band_edges[:-1], band_edges[1:]) if b > a]) + logger.info( + f"Audio log-spectrum: flatness={flatness:.4f}, " + f"band dB range=[{10 * np.log10(bands.min() + 1e-20):.1f}, {10 * np.log10(bands.max() + 1e-20):.1f}]" + ) + return {"flatness": flatness} + + +def _ffmpeg(): + exe = shutil.which("ffmpeg") + if exe: + return exe + try: + import imageio_ffmpeg + except ImportError: + return None + return imageio_ffmpeg.get_ffmpeg_exe() + + +def write_artifacts(frames, audio, sampling_rate, directory: Path, stem: str = "t2va"): + import wave + + paths = {} + wav_path = directory / f"{stem}.wav" + interleaved = np.ascontiguousarray(audio[0].T if audio.ndim == 3 else audio.T) + pcm = np.clip(interleaved, -1.0, 1.0) + with wave.open(str(wav_path), "wb") as handle: + handle.setnchannels(pcm.shape[1]) + handle.setsampwidth(2) + handle.setframerate(sampling_rate) + handle.writeframes((pcm * 32767.0).astype(" dict: + exe = shutil.which("ffprobe") + if exe is None: + return {} + import json as _json + + result = subprocess.run( + [ + exe, + "-v", + "error", + "-show_entries", + "stream=codec_type,duration,nb_frames,sample_rate,channels,width,height", + "-of", + "json", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + streams = {} + for stream in _json.loads(result.stdout).get("streams", []): + streams[stream.get("codec_type", "?")] = stream + return streams + + +# Only reached when ffprobe is unavailable. +_FALLBACK_HEIGHT = 768 +_FALLBACK_WIDTH = 1344 + + +def decoded_frames(path: Path, count: int, height: int = 0, width: int = 0) -> np.ndarray: + """Sample `count` frames evenly out of the written file, as uint8 luma. + + `height`/`width` are the caller's known canvas. They matter: the frame count is derived by + dividing the raw byte stream by the frame size, and the only other source for that size is + `probe_streams`, which needs `ffprobe` -- absent wherever ffmpeg comes from imageio_ffmpeg. + Without them this fell back to a hardcoded 768x1344, which silently miscounts every canvas + with a different pixel total (4:3 read 94 frames for a correct 124-frame file). 16:9 and 21:9 + both happen to be 1032192 px, so the constant looked fine until a sweep left that budget. + """ + exe = _ffmpeg() + if exe is None: + return np.empty((0,)) + result = subprocess.run( + [ + exe, + "-v", + "error", + # Count what the file STORES. Muxing AAC writes an mp4 edit list (the codec's priming + # shifts the timeline), and a decoder honouring it drops frames at the edges: a 362-frame + # write reads back as 362 stored / 361 in-container / 359 after the edit list. The gate + # asks "were the frames written", so the edit list must not be applied here. + "-ignore_editlist", + "1", + "-i", + str(path), + # `select` with count<=1 is a no-op (`not(mod(n,1))` is always true) but combined with + # `-vsync 0` on an edit-list-shifted timeline it silently drops frames at the edges -- + # 2 of them on a muxed 15 s clip. Only ask for the filter when actually subsampling. + *(["-vf", f"select='not(mod(n\\,{count}))'", "-vsync", "0"] if count > 1 else []), + "-f", + "rawvideo", + "-pix_fmt", + "gray", + "-", + ], + check=True, + capture_output=True, + ) + probe = probe_streams(path).get("video", {}) + height = int(probe.get("height") or height or _FALLBACK_HEIGHT) + width = int(probe.get("width") or width or _FALLBACK_WIDTH) + buffer = np.frombuffer(result.stdout, dtype=np.uint8) + usable = (buffer.size // (height * width)) * height * width + return buffer[:usable].reshape(-1, height, width) + + +VBENCH_PYTHON = Path.home() / "vbench_env" / "bin" / "python" + + +def run_vbench(video: Path, prompt: str, dimensions) -> dict[str, float]: + """Run VBench in its own interpreter (~/vbench_env): it pins numpy < 2 / transformers 4.33.""" + if not VBENCH_PYTHON.is_file(): + pytest.skip( + f"no VBench interpreter at {VBENCH_PYTHON}; create it with " + "`uv venv --python 3.10 ~/vbench_env && uv pip install --python ~/vbench_env/bin/python " + "vbench decord 'numpy==1.26.4' 'opencv-python-headless<4.11' 'setuptools<81'`" + ) + interpreter = str(VBENCH_PYTHON) + runner = Path(__file__).with_name("tools") / "vbench_runner.py" + result = subprocess.run( + [interpreter, str(runner), str(video), ",".join(dimensions), "--prompt", prompt], + capture_output=True, + text=True, + timeout=5400, + ) + marker = "VBENCH_JSON " + line = next((l for l in result.stdout.splitlines() if l.startswith(marker)), None) + if line is None: + raise AssertionError( + f"VBench produced no scores (exit {result.returncode}).\n" + f"stdout tail:\n{result.stdout[-2000:]}\nstderr tail:\n{result.stderr[-2000:]}" + ) + import json as _json + + return _json.loads(line[len(marker) :]) + + +def clip_prompt_alignment(frames: np.ndarray, prompt: str, num_frames: int = 8) -> dict[str, float]: + """Mean CLIP similarity between the prompt and evenly-spaced frames, x100.""" + from PIL import Image + + from ...dataset_eval.clip_encoder import CLIPEncoder + + indices = np.linspace(0, frames.shape[0] - 1, num_frames).astype(int) + encoder = CLIPEncoder() + scores = [encoder.get_clip_score(prompt, Image.fromarray(frames[i])).item() * 100.0 for i in indices] + return {"mean": float(np.mean(scores)), "min": float(min(scores)), "max": float(max(scores))} + + +def temporal_seam_score(frames: np.ndarray, period: int) -> float: + """Inter-frame delta at temporal chunk boundaries vs elsewhere (~1.0 = no seam).""" + deltas = np.abs(np.diff(frames.astype(np.float32), axis=0)).mean(axis=(1, 2)) + if len(deltas) < 2 * period: + return float("nan") + index = np.arange(1, len(deltas) + 1) + at_boundary = deltas[index % period == 0] + elsewhere = deltas[index % period != 0] + if not len(at_boundary) or not elsewhere.mean(): + return float("nan") + return float(at_boundary.mean() / elsewhere.mean()) + + +# ------------------------------------------------------------------ shared e2e gate scaffolding + +# Matched pair with the tier-6 bars (CLIP 37.37, imaging_quality 0.6896); imported by fl2va so it cannot drift. +CALIBRATED_FOX_PROMPT = ( + "A red fox trots across a snowy field at dawn, its breath visible in the cold air. " + "The low sun throws long blue shadows behind it, and loose snow lifts from each footfall." +) + + +def weights_dir(*required_subdirs: str) -> Path: + """The snapshot dir from MINIMAX_H3_MODEL_PATH; skips when it or a required partition is missing.""" + root = os.environ.get("MINIMAX_H3_MODEL_PATH", "") + if not root or not Path(root).is_dir(): + pytest.skip("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot") + directory = Path(root) + missing = [name for name in required_subdirs if not (directory / name).is_dir()] + if missing: + pytest.skip(f"MiniMax-H3 snapshot at {directory} is missing {missing}") + return directory + + +def artifact_dir(name: str) -> Path: + """The gate's artifact directory `~/{name}`, created if absent.""" + directory = Path.home() / name + directory.mkdir(parents=True, exist_ok=True) + return directory + + +def run_warm_generation(pipeline, prompt: str, *, seed: int, **gen_kwargs): + """Warmup then the timed generation with identical kwargs; asserts padded-length agreement (programs are keyed on it).""" + pipeline.warmup(prompt=prompt, **gen_kwargs) + warm_padded_len = pipeline.last_padded_len + + # The warmup is a full generation too, so its stage timings are the *cold* numbers -- program + # compilation, per-shape conv3d blocking and buffer allocation all land in them. Log them + # before the measured call overwrites `last_timings`. Not gated: cold cost depends on what the + # weight/JIT caches already held, so it is a diagnostic, not a bar. + cold_rows = list(pipeline.last_timings) + if cold_rows: + cold_total = sum(seconds for _, seconds in cold_rows) + logger.info("WARMUP (cold: program compile + buffer alloc included, not a perf target)") + for row_label, seconds in cold_rows: + share = 100 * seconds / cold_total if cold_total else 0.0 + logger.info(f" {row_label:<18} {seconds:8.1f} s ({share:4.1f} %)") + logger.info(f" {'Total (warmup)':<18} {cold_total:8.1f} s") + + ttnn.synchronize_device(pipeline.mesh_device) + if ttnn.using_distributed_env(): + ttnn.distributed_context_barrier() + + output = pipeline(prompt, seed=seed, **gen_kwargs) + + assert pipeline.last_padded_len == warm_padded_len, ( + f"warmup ran at padded_len {warm_padded_len} but the measured call ran at " + f"{pipeline.last_padded_len}; this number is not warm" + ) + return output + + +def log_timing_table(pipeline, label: str, num_forwards: int, video_seconds: float, expected_total_s=None, extra=""): + """The MEASUREMENT block; `expected_total_s`, when given, asserts the total. Returns the total.""" + rows = pipeline.last_timings + total = sum(seconds for _, seconds in rows) + shape = tuple(pipeline.mesh_device.shape) + logger.info( + f"MEASUREMENT {label} fully warm | mesh {shape[0]}x{shape[1]} Blackhole, " + f"TP={pipeline.tp_factor} axis {pipeline.tp_axis} / SP={pipeline.sp_factor} axis {pipeline.sp_axis}, " + f"{pipeline.ccl_manager.topology}, {pipeline.ccl_manager.num_links} links{extra} " + f"| warm window: one full warmup generation at this shape, prepares and export excluded" + ) + for row_label, seconds in rows: + logger.info(f" {row_label:<18} {seconds:8.1f} s ({100 * seconds / total:4.1f} %)") + logger.info(f" {'Total (compute)':<18} {total:8.1f} s") + denoise = dict(rows).get("Denoise") + if denoise: + logger.info( + f" per forward {denoise / num_forwards * 1000:8.1f} ms " + f"({num_forwards} forwards over {denoise:.1f} s)" + ) + logger.info(f" realtime factor {total / video_seconds:8.1f} x (compute / video seconds)") + if expected_total_s is not None: + assert ( + total < expected_total_s + ), f"fully-warm total {total:.1f} s exceeds the {expected_total_s:.0f} s floor bar" + return total + + +def to_uint8_frames(output) -> np.ndarray: + """`(1, 3, F, H, W)` [0, 1] -> `(F, H, W, 3)` uint8; clamp + round first -- a bare astype wraps out-of-range values.""" + video = output.video + assert video.ndim == 5 and video.shape[0] == 1, f"unexpected video shape {tuple(video.shape)}" + frames = video[0].permute(1, 2, 3, 0).clamp(0, 1).mul(255).round().to(torch.uint8) + return frames.cpu().numpy() + + +def check_written_file(paths: dict, expected_frames: int, seam_period: int = 17, height: int = 0, width: int = 0): + """Gate the written *file* (streams, A/V skew, frame count, temporal seam); silently no-ops without an mp4.""" + if "mp4" not in paths: + return + streams = probe_streams(paths["mp4"]) + if streams: + logger.info(f"container streams: { {k: v.get('duration') for k, v in streams.items()} }") + assert "video" in streams and "audio" in streams, f"muxed file is missing a stream: {list(streams)}" + durations = {k: float(v["duration"]) for k, v in streams.items() if v.get("duration")} + if {"video", "audio"} <= set(durations): + skew = durations["audio"] - durations["video"] + # AAC pads to a frame boundary, so allow a little more than the tensor-level check. + assert abs(skew) < 0.15, f"muxed A/V skew {skew:+.3f} s" + logger.info(f"muxed A/V skew: {skew:+.4f} s") + + decoded = decoded_frames(paths["mp4"], count=1, height=height, width=width) + if decoded.size: + assert ( + decoded.shape[0] >= expected_frames - 1 + ), f"the written mp4 decodes to {decoded.shape[0]} frames, expected ~{expected_frames}" + # The VAE's temporal chunk covers clip_length (17) pixel frames. + seam = temporal_seam_score(decoded, period=seam_period) + logger.info(f"temporal seam score at the {seam_period}-frame chunk period: {seam:.3f} (1.0 = no seam)") + if np.isfinite(seam): + assert seam < 3.0, ( + f"inter-frame delta at chunk boundaries is {seam:.2f}x the delta elsewhere; " + "suspect temporal stitching (see the artifact rubric)" + ) + + +def gate_clip(frames: np.ndarray, prompt: str, threshold: float, label: str): + """CLIP prompt-alignment gate; skips only if `open_clip` is missing.""" + pytest.importorskip("open_clip", reason="the CLIP gate needs open_clip, which is not installed") + alignment = clip_prompt_alignment(frames, prompt) + logger.info( + f"{label} CLIP prompt alignment: mean={alignment['mean']:.2f} " + f"min={alignment['min']:.2f} max={alignment['max']:.2f} (bar {threshold})" + ) + assert alignment["mean"] >= threshold, ( + f"{label} CLIP prompt alignment {alignment['mean']:.2f} is below the {threshold} bar; " + "the video does not match the prompt" + ) + return alignment + + +def gate_vbench(paths: dict, prompt: str, thresholds: dict, label: str, skip_without_mp4=False): + """VBench gate (skips only when its interpreter is missing); a requested dimension with no returned score FAILS, never silently passes.""" + if "mp4" not in paths: + message = "the VBench gate needs the muxed mp4, which ffmpeg did not produce" + if skip_without_mp4: + pytest.skip(message) + raise AssertionError(message) + scores = run_vbench(paths["mp4"], prompt, tuple(thresholds)) + for dimension, bar in thresholds.items(): + assert dimension in scores, f"VBench returned no score for {dimension}" + logger.info(f"{label} VBench {dimension} = {scores[dimension]:.4f} (bar {bar})") + failures = [f"{d} = {scores[d]:.4f} < {bar:.4f}" for d, bar in thresholds.items() if scores[d] < bar] + assert not failures, "VBench below threshold: " + "; ".join(failures) + return scores diff --git a/code/models/tt_dit/tests/models/minimax_h3/decode_bench.py b/code/models/tt_dit/tests/models/minimax_h3/decode_bench.py new file mode 100644 index 0000000000000000000000000000000000000000..3775e9aa1f1da6bd01fe4fcf88db74091d4473dd --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/decode_bench.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""Median decode time over N runs, building the decoder once. + +Build once, warm once, then time N decodes and report median plus spread. Timing a single call per +process is where an earlier script's run-to-run noise came from. + +Usage -- the label names the saved waveform: + + export MINIMAX_H3_MODEL_PATH=/path/to/MiniMax-H3-diffusers + python models/tt_dit/tests/models/minimax_h3/decode_bench.py baseline + +Env: BENCH_N (default 5), BENCH_OUT (default `generated/minimax_h3_audio/bench`). + +For latency *with* PSNR against a CPU reference, and for mesh/sharding/trace configurations, use +`cpu_vs_device.py` instead; this script is single-device and reports timing only. +""" +import json +import os +import re +import statistics +import sys +import time + +import torch +from safetensors.torch import load_file + +import ttnn +from models.tt_dit.models.audio_vae.minimax_h3.convert_minimax_h3_audio import convert_minimax_h3_audio_state_dict +from models.tt_dit.models.audio_vae.minimax_h3.decoder_minimax_h3_audio import MiniMaxH3AudioDecoder + +# Reaches a filename below, so keep it to one path-safe component. +LABEL = re.sub(r"[^A-Za-z0-9._-]", "_", os.path.basename(sys.argv[1])) or "bench" +N = int(os.environ.get("BENCH_N", "5")) +OUT = os.environ.get("BENCH_OUT") or os.path.join( + os.environ.get("TT_METAL_HOME", os.getcwd()), "generated", "minimax_h3_audio", "bench" +) +# `MINIMAX_H3_MODEL_PATH` matches the test suite; the older var is still accepted. +_root = os.environ.get("MINIMAX_H3_MODEL_PATH") or os.environ.get("MINIMAX_H3_DIFFUSERS_DIR", "") +if not _root: + raise SystemExit("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot") +wd = os.path.join(_root, "audio_vae") +cfg = {k: v for k, v in json.load(open(os.path.join(wd, "config.json"))).items() if not k.startswith("_")} + +d = ttnn.open_mesh_device(ttnn.MeshShape(1, 1), l1_small_size=65536) +try: + dec = MiniMaxH3AudioDecoder( + latent_channels=cfg["latent_channels"], + latent_dim=cfg["latent_dim"], + decoder_dim=cfg["decoder_dim"], + decoder_rates=tuple(cfg["decoder_rates"]), + decoder_kernel_sizes=tuple(cfg["decoder_kernel_sizes"]), + resblock_kernel_sizes=tuple(cfg["resblock_kernel_sizes"]), + resblock_dilation_sizes=tuple(tuple(x) for x in cfg["resblock_dilation_sizes"]), + mesh_device=d, + ) + dec.load_torch_state_dict( + convert_minimax_h3_audio_state_dict(load_file(os.path.join(wd, "diffusion_pytorch_model.safetensors"))), + strict=False, + ) + torch.manual_seed(2) + lat = torch.randn(2, cfg["latent_channels"], 207) * 0.1 + out = dec(lat) + ttnn.synchronize_device(d) + s = [] + for _ in range(N): + t0 = time.perf_counter() + out = dec(lat) + ttnn.synchronize_device(d) + s.append(time.perf_counter() - t0) + os.makedirs(OUT, exist_ok=True) + # Both halves come from outside the process (BENCH_OUT, argv[1]), so contain the result here + # rather than relying on the sanitising above. + _out_dir = os.path.realpath(OUT) + _dest = os.path.realpath(os.path.join(_out_dir, LABEL + ".pt")) + if os.path.commonpath([_out_dir, _dest]) != _out_dir: + raise SystemExit(f"refusing to write outside {_out_dir}: {_dest}") + torch.save({"wav": torch.as_tensor(out).float().cpu(), "samples": s}, _dest) + print( + f"RESULT {LABEL}: median {statistics.median(s):.4f}s min {min(s):.4f} max {max(s):.4f} n={N} " + f"[{', '.join(f'{v:.3f}' for v in s)}]", + flush=True, + ) +finally: + ttnn.close_mesh_device(d) diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_adaln_precompute_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_adaln_precompute_minimax_h3.py new file mode 100644 index 0000000000000000000000000000000000000000..bb6b18ec6007c609e1a77c81f9b04df9090299c7 --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/test_adaln_precompute_minimax_h3.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""Host-only gate for the MiniMax-H3 AdaLN precompute, on a small synthetic checkpoint. +The design rests on ``adaln_proj`` being batch-stable while ``time_embedder`` is not. The full +26 GB build was verified once against the real checkpoint (git history: tools/build_adaln_table.py).""" + +import pytest +import torch +from safetensors.torch import save_file + +from ....pipelines.minimax_h3 import adaln_precompute as ap +from ....pipelines.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG +from ....pipelines.minimax_h3.scheduler import MiniMaxH3Scheduler + +HIDDEN = 64 +TIME_EMBED_DIM = 32 +FREQ_DIM = 16 +NUM_LAYERS = 3 +STEPS = 6 + + +def _schedules(steps=STEPS): + video = MiniMaxH3Scheduler(12.0) + audio = MiniMaxH3Scheduler(3.0) + video.set_timesteps(steps) + audio.set_timesteps(steps) + return video, audio + + +def _step_timesteps(steps=STEPS): + video, audio = _schedules(steps) + return ap.request_step_timesteps(video.sigmas, audio.sigmas, MINIMAX_H3_KEYFRAME_NOISE_AUG) + + +@pytest.fixture +def synthetic_checkpoint(tmp_path): + torch.manual_seed(0) + tensors = { + "time_embedder.proj_in.weight": torch.randn(HIDDEN, FREQ_DIM), + "time_embedder.proj_in.bias": torch.randn(HIDDEN), + "time_embedder.proj_out.weight": torch.randn(TIME_EMBED_DIM, HIDDEN), + "time_embedder.proj_out.bias": torch.randn(TIME_EMBED_DIM), + "final_layer.adaln_proj.linear.weight": torch.randn(2 * HIDDEN, TIME_EMBED_DIM).bfloat16(), + "final_layer.adaln_proj.linear.bias": torch.randn(2 * HIDDEN).bfloat16(), + } + for layer in range(NUM_LAYERS): + out = ap.MINIMAX_H3_ADALN_PARAMS * HIDDEN * ap.MINIMAX_H3_MODALITY_NUM + tensors[f"blocks.{layer}.adaln_proj.linear.weight"] = torch.randn(out, TIME_EMBED_DIM).bfloat16() + tensors[f"blocks.{layer}.adaln_proj.linear.bias"] = torch.randn(out).bfloat16() + # Two shards, so the multi-shard key index is exercised. + keys = sorted(tensors) + save_file({k: tensors[k] for k in keys[: len(keys) // 2]}, tmp_path / "model-00001-of-00002.safetensors") + save_file({k: tensors[k] for k in keys[len(keys) // 2 :]}, tmp_path / "model-00002-of-00002.safetensors") + return tmp_path, tensors + + +def _table(checkpoint, steps=STEPS): + return ap.precompute_adaln_table( + checkpoint, + _step_timesteps(steps), + num_layers=NUM_LAYERS, + hidden_size=HIDDEN, + freq_dim=FREQ_DIM, + ) + + +def test_frequency_embedding_is_cosine_then_sine(): + diffusers_embeddings = pytest.importorskip("diffusers.models.embeddings") + timesteps = torch.tensor([0.0, 0.02, 0.5, 0.999, 1.0], dtype=torch.float32) + ours = ap.timestep_frequency_embedding(timesteps, 256) + theirs = diffusers_embeddings.get_timestep_embedding(timesteps, 256, flip_sin_to_cos=True, downscale_freq_shift=0) + assert torch.equal(ours, theirs) + + +def test_silu_must_run_before_the_bfloat16_cast(): + """bf16 activation shifts modulation ~7.8e-3, coherently across every block and step.""" + torch.manual_seed(1) + weight = torch.randn(6 * HIDDEN * 3, TIME_EMBED_DIM).bfloat16() + bias = torch.randn(6 * HIDDEN * 3).bfloat16() + temb = torch.randn(4, TIME_EMBED_DIM, dtype=torch.float32) + + correct = ap.project_block_adaln(temb, weight, bias, HIDDEN) + hoisted = torch.nn.functional.linear(torch.nn.functional.silu(temb.bfloat16()), weight, bias) + hoisted = torch.stack(hoisted.view(-1, 6 * HIDDEN).chunk(6, dim=-1), dim=1) + + assert not torch.equal(correct, hoisted) + assert (correct.float() - hoisted.float()).abs().max() > 1e-3 + + +def test_adaln_projection_is_batch_stable(): + torch.manual_seed(2) + weight = torch.randn(6 * HIDDEN * 3, TIME_EMBED_DIM).bfloat16() + bias = torch.randn(6 * HIDDEN * 3).bfloat16() + temb = torch.randn(16, TIME_EMBED_DIM, dtype=torch.float32) + + batched = ap.project_block_adaln(temb, weight, bias, HIDDEN) + for row in (0, 5, 15): + alone = ap.project_block_adaln(temb[row : row + 1], weight, bias, HIDDEN) + assert torch.equal(batched[row * 3 : (row + 1) * 3], alone) + + +def test_time_embedding_is_not_batch_stable(synthetic_checkpoint): + """Why ``temb`` is computed per step, asserted rather than silently depended on.""" + _, tensors = synthetic_checkpoint + args = ( + tensors["time_embedder.proj_in.weight"], + tensors["time_embedder.proj_in.bias"], + tensors["time_embedder.proj_out.weight"], + tensors["time_embedder.proj_out.bias"], + ) + steps = _step_timesteps() + everything = torch.unique(torch.cat(steps), sorted=True) + + batched = ap.time_embedding(everything, *args, freq_dim=FREQ_DIM) + per_step = ap.time_embedding(steps[0], *args, freq_dim=FREQ_DIM) + rows = [int((everything == value).nonzero()[0, 0]) for value in steps[0]] + + assert torch.allclose(batched[rows], per_step, atol=1e-4) + + +def test_step_timesteps_pin_conditioning_at_the_noise_aug_floor(): + video, audio = _schedules() + steps = ap.request_step_timesteps(video.sigmas, audio.sigmas, MINIMAX_H3_KEYFRAME_NOISE_AUG) + for index, levels in enumerate(steps): + expected = max(float(video.timesteps[index]), MINIMAX_H3_KEYFRAME_NOISE_AUG) + assert pytest.approx(float(levels.max())) == expected + + +def test_table_matches_per_step_recompute(synthetic_checkpoint): + checkpoint, tensors = synthetic_checkpoint + steps = _step_timesteps() + table = _table(checkpoint) + + total_rows = sum(int(levels.numel()) for levels in steps) + assert len(steps) == _schedules()[0].num_inference_steps + assert table.num_layers == NUM_LAYERS + assert table.hidden_size == HIDDEN + assert table.num_steps == len(steps) + assert tuple(table.block_params.shape) == ( + NUM_LAYERS, + total_rows * ap.MINIMAX_H3_MODALITY_NUM, + ap.MINIMAX_H3_ADALN_PARAMS, + HIDDEN, + ) + assert table.block_params.dtype == torch.bfloat16 + assert tuple(table.final_shift.shape) == (total_rows, HIDDEN) + assert int(table.step_offsets[0]) == 0 + assert int(table.step_offsets[-1]) == total_rows + + embed_args = ( + tensors["time_embedder.proj_in.weight"], + tensors["time_embedder.proj_in.bias"], + tensors["time_embedder.proj_out.weight"], + tensors["time_embedder.proj_out.bias"], + ) + + for step, levels in enumerate(steps): + assert levels.numel() in (2, 3) + assert bool((levels[1:] > levels[:-1]).all()) + assert torch.equal(table.step_timesteps(step), levels) + + temb = ap.time_embedding(levels, *embed_args, freq_dim=FREQ_DIM) + rows = table.step_rows(step, torch.arange(levels.numel())) + low = int(rows[0]) * ap.MINIMAX_H3_MODALITY_NUM + high = (int(rows[-1]) + 1) * ap.MINIMAX_H3_MODALITY_NUM + + for layer in range(NUM_LAYERS): + expected = ap.project_block_adaln( + temb, + tensors[f"blocks.{layer}.adaln_proj.linear.weight"], + tensors[f"blocks.{layer}.adaln_proj.linear.bias"], + HIDDEN, + ) + assert torch.equal(table.block_params[layer, low:high], expected) + + shift, scale = ap.project_final_adaln( + temb, + tensors["final_layer.adaln_proj.linear.weight"], + tensors["final_layer.adaln_proj.linear.bias"], + ) + assert torch.equal(table.final_shift[rows[0] : rows[-1] + 1], shift) + assert torch.equal(table.final_scale[rows[0] : rows[-1] + 1], scale) + + +def test_adaln_indices_address_timestep_and_modality(synthetic_checkpoint): + checkpoint, _ = synthetic_checkpoint + table = _table(checkpoint) + step = 3 + offset = int(table.step_offsets[step]) + + timestep_indices = torch.tensor([0, 0, 1, 1]) + token_tags = torch.tensor([0, 2, 1, -1]) + indices = table.adaln_indices(step, timestep_indices, token_tags) + + expected = torch.tensor( + [ + (offset + 0) * 3 + 0, + (offset + 0) * 3 + 2, + (offset + 1) * 3 + 1, + # tag -1 padding rows clamp to the video slot; their output is discarded + (offset + 1) * 3 + 0, + ] + ) + assert torch.equal(indices, expected) diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_audio_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_audio_minimax_h3.py new file mode 100644 index 0000000000000000000000000000000000000000..7fd68ae44504992b1dd5f52017ad9f511a464b9c --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/test_audio_minimax_h3.py @@ -0,0 +1,666 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""Gate M8d: the MiniMax-H3 audio VAE (BigVGAN) at production shapes; stereo rides as batch 2. + +T-parallel decode works: ``test_audio_decode_t_parallel`` at the bottom of this file is resurrected from +git history, as the previous version of this docstring asked for. The 8-way shard layout it called +known-broken was `Vocoder.conv_pre` returning uninitialized memory when sharded, plus +``_forward_tap_matmul`` padding shard boundaries locally instead of through the halo; both are fixed, and +the test's PSNR assert is enforced for every factor rather than excused. ``test_neighbor_pad_t_minimax_h3.py`` +is still in git history if the halo itself ever needs its own gate again -- it was verified correct here +(exact at pad 1/3/25) while hunting the conv_pre bug.""" + +from __future__ import annotations + +import copy +import os +import time + +import pytest +import torch +from loguru import logger + +import ttnn + +from ....layers.audio_ops import Conv1dViaConv3d, depthwise_tap_filter +from ....models.audio_vae.minimax_h3.blockings_minimax_h3_audio import register_h3_audio_blockings +from ....models.audio_vae.minimax_h3.convert_minimax_h3_audio import ( + assert_weight_norm_axes_consistent, + convert_minimax_h3_audio_state_dict, + fuse_attention_biases, + fuse_weight_norm, + remap_amp_activations, +) +from ....models.audio_vae.minimax_h3.decoder_minimax_h3_audio import MiniMaxH3AudioDecoder +from ....parallel.config import ParallelFactor +from ....parallel.manager import CCLManager +from ....utils.check import assert_quality +from .common import build_audio_decoder, load_config, psnr, weights_subdir + +# The vocoder needs extra L1 scratch, as the LTX audio tests do. +SINGLE_DEVICE = [pytest.param((1, 1), {"l1_small_size": 65536}, id="single_device")] + +# Encode floor, calibrated before accurate mode became the constructed default (measured ~0.105 with +# fp32 conv multipliers over ~130 convs); the accurate defaults only improve on it. +AUDIO_RELATIVE_RMSE = 0.12 + +CONV_PRE_SHAPE = (2048, 1024, 7) # conv_pre: the decoder's widest reduction, Cin 2048 x k 7 +CONV_PRE_LATENT_FRAMES = 207 + +LATENT_CHANNELS = 32 +HOP_LENGTH = 800 +SAMPLING_RATE = 32000 + +PRODUCTION_LATENT_FRAMES = [pytest.param(207, id="5s")] # 40 Hz latent; 10 s exercises no new path + + +def _build_reference(load_weights: bool = True): + """The reference audio VAE, with weight norm still attached (the converter removes it).""" + weights_dir = weights_subdir("audio_vae") + if weights_dir is None: + pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH") + pytest.importorskip("diffusers", reason="pinned diffusers reference not installed") + from diffusers import AutoencoderKLMiniMaxH3Audio + + config = load_config(weights_dir) + reference = AutoencoderKLMiniMaxH3Audio(**config).eval() + if load_weights: + from safetensors.torch import load_file + + reference.load_state_dict(load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors"))) + return reference, config + + +def _golden_latents(reference, num_latent_frames: int): + """``(waveform, latents, expected_decode)``; latents come from the reference encoder, not ``randn``.""" + waveform = torch.randn(2, 1, num_latent_frames * HOP_LENGTH) * 0.1 + with torch.no_grad(): + posterior = reference.encode(waveform).latent_dist + latents = posterior.mode()[..., :num_latent_frames] + expected = reference.decode(latents).sample + return waveform, latents, expected + + +def _log_mel_distance(a: torch.Tensor, b: torch.Tensor, *, n_fft: int = 1024, hop: int = 256) -> float: + """Mean absolute log-magnitude-spectrogram difference; catches spectral drift an aggregate PCC hides.""" + window = torch.hann_window(n_fft) + spectra = [] + for signal in (a, b): + flat = signal.reshape(-1, signal.shape[-1]).float() + stft = torch.stft(flat, n_fft=n_fft, hop_length=hop, window=window, return_complex=True) + spectra.append(torch.log(stft.abs().clamp_min(1e-5))) + return (spectra[0] - spectra[1]).abs().mean().item() + + +def _tt_decoder(config: dict, mesh_device) -> MiniMaxH3AudioDecoder: + return build_audio_decoder(config, mesh_device) + + +@pytest.mark.parametrize("num_latent_frames", PRODUCTION_LATENT_FRAMES) +@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"]) +def test_decode(mesh_device, num_latent_frames): + """The whole decode path against the reference, at a production duration, stereo. + + Constructor defaults are accurate mode (split_mode='full', tap_matmul, prefer_mac), so the bars + are the accurate-mode ones: measured 0.0045 rel RMSE / 99.9990% PCC / 67.5 dB PSNR. + """ + reference, config = _build_reference() + torch.manual_seed(1) + + _, latents, expected = _golden_latents(reference, num_latent_frames) + with torch.no_grad(): + expected_proj = reference.dec_in_proj(latents) + + tt_decoder = _tt_decoder(config, mesh_device) + # The precision levers are the constructed defaults; assert they landed where they matter. + assert tt_decoder.dec_in_proj.split_mode == "full", "split_mode='full' did not land on dec_in_proj" + assert tt_decoder.dec_in_proj.tap_matmul, "tap_matmul=True did not land on dec_in_proj" + assert tt_decoder.decoder.conv_post.split_mode == "full", "split_mode='full' did not land on conv_post" + assert tt_decoder.decoder.act_post.downsample.lowpass.prefer_mac, "prefer_mac=True did not land on act_post" + + tt_decoder.load_torch_state_dict(convert_minimax_h3_audio_state_dict(dict(reference.state_dict())), strict=False) + + projected = tt_decoder._project_latents_device(latents) + assert projected.shape == expected_proj.shape, f"proj {tuple(projected.shape)} != {tuple(expected_proj.shape)}" + assert_quality(expected_proj, projected, pcc=0.999) + + actual = tt_decoder(latents) + + assert actual.shape == expected.shape, f"shape {tuple(actual.shape)} != reference {tuple(expected.shape)}" + assert_quality(expected, actual, pcc=0.9999, relative_rmse=0.006) + + psnr_db = psnr(expected, actual) + mel_distance = _log_mel_distance(expected, actual) + assert psnr_db >= 60.0, f"decode PSNR {psnr_db:.2f} dB < 60 dB" + assert mel_distance <= 5.0, f"log-spectrogram distance {mel_distance:.3f} > 5.0" + + left, right = actual[0, 0], actual[1, 0] + assert not torch.allclose(left, right, atol=1e-4), "stereo channels are identical -- a broadcast bug" + + +@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"]) +def test_depthwise_chunked_conv1d_matches_torch(mesh_device): + """`depthwise_tap_filter`'s C-chunked `ttnn.conv1d` recovery, against a torch depthwise reference. + + Nothing else covers it: H3 defaults to `prefer_mac=True`, which returns before the conv1d fallback, + so no decode test touches the slicing, the per-chunk weight, or the concat. Shape is the one the + decode really chunks (T_pad=166, C=512, K=7, stride=1), where the activation block is C*K wide and + the slicer runs out of L1 at full C. + """ + torch.manual_seed(0) + B, T_pad, C, K, stride = 1, 166, 512, 7, 1 + + x = torch.randn(B, T_pad, C, dtype=torch.float32) + taps = torch.randn(K, dtype=torch.float32) + x_dev = ttnn.from_torch(x, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=mesh_device) + + cache: dict = {} + out = depthwise_tap_filter( + x_dev, + [float(t) for t in taps], + stride, + mesh_device=mesh_device, + dtype=ttnn.float32, + cache=cache, + prefer_mac=False, + ) + actual = ttnn.to_torch(ttnn.get_device_tensors(out)[0]).float() + + # Valid depthwise conv over T: one tap vector shared by every channel, groups=C. + weight = taps.view(1, 1, K).expand(C, 1, K).contiguous() + expected = torch.nn.functional.conv1d(x.transpose(1, 2), weight, stride=stride, groups=C).transpose(1, 2) + + assert actual.shape == expected.shape, f"{actual.shape} != {expected.shape}" + psnr_db = psnr(expected, actual) + logger.info(f"chunked depthwise conv1d vs torch: {psnr_db:.2f} dB") + assert psnr_db > 60.0, f"chunked depthwise conv1d diverges from torch: {psnr_db:.2f} dB" + + # Prove the chunked path ran: it keys its prepared weight on the chunk width where the unchunked + # path keys on C. Without this the test would pass on plain conv1d or a silent MAC fallback, i.e. + # assert torch equivalence while covering none of the code it exists for. + chunk_widths = [k[1] for k in cache if isinstance(k, tuple) and k and k[0] == "w" and k[1] < C] + assert chunk_widths, ( + f"expected the C-chunked conv1d recovery to run at C={C}, K={K}, T_pad={T_pad}, but the weight " + f"cache has no sub-C chunk key: {sorted(k for k in cache if isinstance(k, tuple))}. Either " + "conv1d now fits at full C (then this shape no longer covers the chunked path and the test needs " + "a new one), or it fell back to MAC (then the recovery is broken)." + ) + logger.info(f"chunked into widths {sorted(set(chunk_widths))} from C={C}") + + +@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"]) +def test_conv_operand_split_improves_precision(mesh_device): + """Operand splitting really does beat the fp32 conv floor, at ``conv_pre``'s production shape.""" + # Raw Conv1dViaConv3d does not register the H3 conv blockings (the H3 audio module constructors + # do); without this, every conv silently falls back to C_in_block=32 and the "off" bar fails. + register_h3_audio_blockings() + in_channels, out_channels, kernel = CONV_PRE_SHAPE + torch.manual_seed(0) + reference = torch.nn.Conv1d(in_channels, out_channels, kernel, padding=kernel // 2).float().eval() + x = torch.randn(2, in_channels, CONV_PRE_LATENT_FRAMES) * 0.1 + with torch.no_grad(): + golden = copy.deepcopy(reference).double()(x.double()).float() + + x_device = None + errors = {} + for mode in ("off", "weight", "full"): + layer = Conv1dViaConv3d( + in_channels, + out_channels, + kernel_size=kernel, + padding_mode="zeros", + bias=True, + mesh_device=mesh_device, + dtype=ttnn.float32, + split_mode=mode, + ) + assert layer.split_mode == mode, f"layer resolved split_mode {layer.split_mode!r}, expected {mode!r}" + assert (layer.weight_lo is None) == (mode == "off"), f"weight residual allocation wrong for {mode!r}" + layer.load_torch_state_dict( + { + "weight": reference.weight.detach().contiguous(), + "bias": reference.bias.detach().contiguous(), + }, + strict=False, + ) + if x_device is None: + x_device = ttnn.from_torch( + x.transpose(1, 2).contiguous(), dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=mesh_device + ) + actual = ttnn.to_torch(layer(x_device)).float().transpose(1, 2) + errors[mode] = float((actual.double() - golden.double()).pow(2).mean().sqrt() / golden.double().std()) + + logger.info("conv_pre rel_rmse by split mode: " + ", ".join(f"{k}={v:.3e}" for k, v in errors.items())) + + # 2.1e-03 sits between production's 1.86e-03 and the unregistered C_in_block=32 fallback's 2.40e-03. + assert errors["off"] <= 2.1e-3, ( + f"baseline conv_pre error {errors['off']:.3e} is above the production floor -- " + "C_in_block has probably fallen back to 32, meaning register_h3_audio_blockings() did not run" + ) + # Measured 1.86e-03 / 1.25e-03 / 1.00e-03; ratios carry ~1.3x margin. + assert errors["weight"] <= 0.85 * errors["off"], f"weight-split did not improve: {errors}" + assert errors["full"] <= 0.70 * errors["off"], f"full split did not improve enough: {errors}" + assert errors["full"] < errors["weight"], f"full split should beat weight-only: {errors}" + + +def _tt_encoder(config: dict, mesh_device): + from ....models.audio_vae.minimax_h3.encoder_minimax_h3_audio import MiniMaxH3AudioEncoder + + return MiniMaxH3AudioEncoder( + encoder_dim=config["encoder_dim"], + encoder_rates=tuple(config["encoder_rates"]), + latent_dim=config["latent_dim"], + latent_channels=config["latent_channels"], + num_attention_heads=config["num_attention_heads"], + mesh_device=mesh_device, + ) + + +@pytest.mark.parametrize("num_latent_frames", PRODUCTION_LATENT_FRAMES) +@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"]) +def test_encode(mesh_device, num_latent_frames): + """The whole encode path -- DAC trunk, ``pre_block``, posterior heads -- vs the reference.""" + reference, config = _build_reference() + torch.manual_seed(2) + + waveform = torch.randn(2, 1, num_latent_frames * HOP_LENGTH) * 0.1 + with torch.no_grad(): + posterior = reference.encode(waveform).latent_dist + expected_mean, expected_logs = posterior.mean, posterior.logs + + tt_encoder = _tt_encoder(config, mesh_device) + tt_encoder.load_torch_state_dict(convert_minimax_h3_audio_state_dict(dict(reference.state_dict())), strict=False) + mean, logs = tt_encoder(waveform) + + assert mean.shape == expected_mean.shape, f"mean shape {tuple(mean.shape)} != {tuple(expected_mean.shape)}" + assert mean.shape[2] == num_latent_frames, f"expected {num_latent_frames} latent frames, got {mean.shape[2]}" + assert_quality(expected_mean, mean, pcc=0.99, relative_rmse=AUDIO_RELATIVE_RMSE) + assert_quality(expected_logs, logs, pcc=0.98) + + +@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"]) +def test_roundtrip(mesh_device): + """End to end: encode then decode on device, against the reference's own round trip.""" + reference, config = _build_reference() + torch.manual_seed(3) + num_latent_frames = 207 + + waveform, _, expected = _golden_latents(reference, num_latent_frames) + + tt_encoder = _tt_encoder(config, mesh_device) + tt_decoder = _tt_decoder(config, mesh_device) + converted = convert_minimax_h3_audio_state_dict(dict(reference.state_dict())) + tt_encoder.load_torch_state_dict(converted, strict=False) + tt_decoder.load_torch_state_dict(converted, strict=False) + + mean, _ = tt_encoder(waveform) + actual = tt_decoder(mean) + + assert actual.shape == expected.shape, f"shape {tuple(actual.shape)} != {tuple(expected.shape)}" + psnr_db = psnr(expected, actual) + assert psnr_db >= 28.0, f"round-trip PSNR {psnr_db:.2f} dB < 28 dB" + assert _log_mel_distance(expected, actual) <= 5.0, "round-trip spectrum drifted" + + +# -------------------------------------------------------------------- weight-norm fusion and checkpoint conversion (host only) + + +def test_fuse_weight_norm_matches_torch_conv_transpose1d(): + """The load-bearing case: axis 0 of a ConvTranspose1d weight is ``in_channels``, not out.""" + torch.manual_seed(1) + conv = torch.nn.ConvTranspose1d(32, 16, kernel_size=4, stride=2) + assert conv.weight.shape == (32, 16, 4), "ConvTranspose1d weight is (in, out, k)" + conv = torch.nn.utils.weight_norm(conv) + with torch.no_grad(): + conv.weight_g.normal_(1.0, 0.2) + conv.weight_v.normal_() + assert conv.weight_g.shape == (32, 1, 1), "weight_g is per-in_channel, confirming dim=0" + + weight_g = conv.weight_g.detach().clone() + weight_v = conv.weight_v.detach().clone() + fused = fuse_weight_norm(weight_g, weight_v) + torch.nn.utils.remove_weight_norm(conv) + + relative = (fused - conv.weight).abs().max().item() / conv.weight.abs().max().item() + assert relative < 1e-6, f"ConvTranspose1d fusion differs from torch by {relative:.3e}" + + wrong_norm = weight_v.transpose(0, 1).flatten(1).norm(dim=1).view(1, -1, 1) + wrong = weight_g * weight_v / wrong_norm + assert wrong.shape == fused.shape, "the wrong axis still type-checks -- hence this test" + assert not torch.allclose(wrong, fused, atol=1e-4), "the two axes agree, so this test proves nothing" + + +def test_remap_amp_activations_interleaves_correctly(): + """``activations.{0,2,4}`` -> ``acts1.{0,1,2}`` and ``{1,3,5}`` -> ``acts2.{0,1,2}``.""" + state = {f"resblocks.0.activations.{i}.act.alpha": torch.tensor([float(i)]) for i in range(6)} + remapped = remap_amp_activations(state) + + for i in range(3): + assert remapped[f"resblocks.0.acts1.{i}.act.alpha"].item() == 2 * i + assert remapped[f"resblocks.0.acts2.{i}.act.alpha"].item() == 2 * i + 1 + assert not any("activations." in key for key in remapped), "an activations key survived" + + +def test_fuse_attention_biases_rejects_a_nonzero_k_bias(): + """A ``zero_k_bias`` that is not zero must fail loudly, not be dropped.""" + state = { + "pre_block.attn.q_bias": torch.ones(8), + "pre_block.attn.v_bias": torch.full((8,), 2.0), + "pre_block.attn.zero_k_bias": torch.zeros(8), + } + fused = fuse_attention_biases(state) + bias = fused["pre_block.attn.qkv.bias"] + assert bias.shape == (24,) + assert torch.equal(bias[:8], torch.ones(8)) + assert torch.equal(bias[8:16], torch.zeros(8)) + assert torch.equal(bias[16:], torch.full((8,), 2.0)) + + state["pre_block.attn.zero_k_bias"] = torch.ones(8) + with pytest.raises(AssertionError, match="not all zero"): # allow-pytest.raises: guards a silent data-loss path + fuse_attention_biases(state) + + +def test_real_checkpoint_fusion_matches_reference_module(): + """Real checkpoint: axes hold, every pair fuses, fused weights equal torch's ``remove_weight_norm``.""" + weights_dir = weights_subdir("audio_vae") + if weights_dir is None: + pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH") + pytest.importorskip("diffusers", reason="pinned diffusers reference not installed") + from diffusers import AutoencoderKLMiniMaxH3Audio + from safetensors.torch import load_file + + config = load_config(weights_dir) + reference = AutoencoderKLMiniMaxH3Audio(**config) + state = load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors")) + reference.load_state_dict(state) + + assert_weight_norm_axes_consistent(state) + num_pairs = len([k for k in state if k.endswith(".weight_g")]) + assert num_pairs > 100, f"expected ~172 weight-normed convs, found {num_pairs}" + + converted = convert_minimax_h3_audio_state_dict(state) + assert not [k for k in converted if k.endswith((".weight_g", ".weight_v"))], "a weight-norm pair survived" + assert not [k for k in converted if k.endswith(("q_bias", "v_bias", "zero_k_bias"))], "an attn bias survived" + assert not [k for k in converted if "activations." in k], "an activations key survived" + assert ( + len(converted) == len(state) - num_pairs - 2 + ), f"key count {len(converted)} does not match {len(state)} minus {num_pairs} g/v pairs and 2 folded biases" + + for module in reference.modules(): + if hasattr(module, "weight_g"): + torch.nn.utils.remove_weight_norm(module) + expected = dict(reference.state_dict()) + + checked = 0 + worst_key, worst = None, 0.0 + for key, value in expected.items(): + if not key.endswith(".weight") or key not in converted: + continue + scale = max(value.abs().max().item(), 1e-12) + relative = (converted[key] - value).abs().max().item() / scale + if relative > worst: + worst_key, worst = key, relative + checked += 1 + assert checked > 100, f"only compared {checked} fused weights" + assert worst < 1e-6, f"worst fused weight is {worst_key} at relative {worst:.3e}" + + +# -------------------------------------------------------------------- traced audio decode + +# 450 MB trace region: the accurate-mode (default) decode graph needs 375463936 B. +TRACED = [ + pytest.param((1, 1), {"l1_small_size": 65536, "trace_region_size": 450_000_000}, id="single_device"), +] +NUM_LATENT_FRAMES = 207 +ITERS = 5 + + +def _best(fn) -> float: + fn() + best = float("inf") + for _ in range(ITERS): + t0 = time.perf_counter() + fn() + best = min(best, time.perf_counter() - t0) + return best + + +@pytest.mark.parametrize(("mesh_device", "device_params"), TRACED, indirect=["mesh_device", "device_params"]) +def test_audio_decode_traced(mesh_device): + reference, config = _build_reference() + converted = convert_minimax_h3_audio_state_dict(dict(reference.state_dict())) + + torch.manual_seed(2) + decoder = _tt_decoder(config, mesh_device) + decoder.load_torch_state_dict(converted, strict=False) + + latents = torch.randn(2, config["latent_channels"], NUM_LATENT_FRAMES) * 0.1 + + plain = decoder(latents) + traced = decoder(latents, traced=True) + assert traced.shape == plain.shape, f"{traced.shape} != {plain.shape}" + + # PSNR inf would also appear if traced=True silently fell through, so confirm a trace was captured. + tracers = type(decoder.decoder)._forward_device._tracers_keyed.get(decoder.decoder, {}) + assert tracers, "traced=True captured no trace; the PSNR below would be meaningless" + assert plain.abs().max() > 1e-6, "decoder produced all-zero output; PSNR would be vacuous" + + psnr_db = psnr(plain, traced) + logger.info(f"traced vs untraced PSNR: {psnr_db:.2f} dB ({len(tracers)} trace(s) captured)") + + untraced_s = _best(lambda: decoder(latents)) + traced_s = _best(lambda: decoder(latents, traced=True)) + logger.info( + f"PERF audio_decode_5s untraced {untraced_s:.4f} s | traced {traced_s:.4f} s " + f"-> {untraced_s / traced_s:.2f}x" + ) + decoder.release_trace() + + assert psnr_db > 60.0, f"traced output diverges from untraced: PSNR {psnr_db:.2f} dB" + + +MESH = [ + pytest.param( + (4, 8), + { + "fabric_config": ttnn.FabricConfig.FABRIC_1D, + "require_exact_physical_num_devices": True, + "l1_small_size": 65536, + }, + id="mesh4x8", + ) +] +# (t_factor, mesh_axis). Axis 1 is the 8-wide axis of the 4x8 Galaxy, axis 0 the 4-wide one. +# +# The factor must equal the length of the axis it shards: factor=2 or factor=4 on the 8-wide axis 1 +# both die in `_partition_t` at slice_device_operation.cpp:164 ("height begin index aligned to tiles"), +# because the partition indexes by the device's coordinate along the axis and assumes it covers it. +# That is why this list is (4, axis 0) and (8, axis 1) rather than a scan. +# +# `KNOWN_BROKEN` is deliberately empty; an entry silences the PSNR assert, which is the only thing +# separating a speedup from a fast wrong answer. It formerly held (8, 1) at -6.3 dB, blamed on one tile +# per shard -- the wrong suspect, since every factor was wrong and the cause was `conv_pre` returning +# uninitialized memory when sharded. Both factors now measure 83.4 dB against the unsharded path +# (1.87x at axis 0, 2.20x at axis 1), and `cpu_vs_device.py` scores sharded at the same PSNR as single +# device (81.89 vs 81.99 dB at the constructed defaults): sharding buys latency, not accuracy. +FACTORS = [(1, 1), (4, 0), (8, 1)] +KNOWN_BROKEN: set[tuple[int, int]] = set() + + +def _build(mesh_device, config, converted, parallel_config, ccl_manager): + """The decoder at this file's shared defaults, plus a shard layout. + + Goes through `build_audio_decoder` rather than constructing directly so the precision levers + (`split_mode`, `tap_matmul`, `prefer_mac`, `max_c_in_block`) stay at whatever the shipping default + is -- a sharded run must measure the same configuration the unsharded gates do, or a divergence + could be a lever difference rather than a sharding bug. + """ + decoder = build_audio_decoder( + config, + mesh_device, + parallel_config=parallel_config, + ccl_manager=ccl_manager, + ) + decoder.load_torch_state_dict(converted, strict=False) + return decoder + + +def _localize_divergence(baseline, parallel, *, factor: int, logger) -> None: + """Say *where* a diverging shard layout diverges. "PSNR -10.2 dB" alone names no bug. + + Three shapes of answer, each pointing somewhere different: + + - concentrated at the internal shard boundaries -> the halo exchange is wrong, and each conv in the + stack needs `kernel_size - 1` samples of its neighbour that it is not getting; + - uniform across the whole signal -> the shard layout itself is wrong, not its edges; + - a prefix or suffix only -> the causal padding or the final trim is off by a shard. + + Also reports the best cross-correlation lag: a diverging-but-highly-correlated output at a nonzero + lag is a *shift*, which is a trim bug rather than a numerics one, and PSNR cannot distinguish those. + """ + import numpy as np + + error = (parallel - baseline).abs() + total = baseline.shape[-1] + logger.info( + f" divergence localization, t_factor={factor}: overall mean {error.mean():.6f} " + f"max {error.max():.4f} against baseline absmax {baseline.abs().max():.4f}" + ) + per_shard = [] + for shard in range(factor): + lo, hi = shard * total // factor, (shard + 1) * total // factor + per_shard.append(float(error[..., lo:hi].mean())) + logger.info(f" per-shard mean error: {[f'{v:.6f}' for v in per_shard]}") + + # Boundary vs interior. A halo bug puts the error in a narrow band at each internal boundary. + window = 128 + boundary, interior = [], [] + for shard in range(1, factor): + cut = shard * total // factor + boundary.append(float(error[..., max(0, cut - window) : cut + window].mean())) + mask = torch.ones(total, dtype=torch.bool) + for shard in range(1, factor): + cut = shard * total // factor + mask[max(0, cut - window) : cut + window] = False + interior = float(error[..., mask].mean()) + logger.info( + f" boundary bands (+-{window}): {[f'{v:.6f}' for v in boundary]} interior {interior:.6f} " + f"ratio {max(boundary) / max(interior, 1e-12):.2f}" + ) + + a = baseline[0, 0].numpy() + c = parallel[0, 0].numpy() + best, best_lag = -2.0, None + for lag in range(-4096, 4097, 32): + x, y = (a[-lag:], c[: len(c) + lag]) if lag < 0 else (a[: len(a) - lag], c[lag:]) + n = min(len(x), len(y)) + if n < 2048: + continue + r = float(np.corrcoef(x[:n], y[:n])[0, 1]) + if r > best: + best, best_lag = r, lag + logger.info( + f" correlation at lag 0: {float(np.corrcoef(a, c)[0, 1]):.4f}; " + f"best {best:.4f} at lag {best_lag} (nonzero lag => a shift, i.e. a trim bug)" + ) + + +# ~14 min: three decoder builds plus a decode per factor, against `pytest.ini`'s 300 s default. +@pytest.mark.timeout(2400) +@pytest.mark.parametrize(("mesh_device", "device_params"), MESH, indirect=["mesh_device", "device_params"]) +def test_audio_decode_t_parallel(mesh_device): + weights_dir = weights_subdir("audio_vae") + if weights_dir is None: + pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH") + pytest.importorskip("diffusers", reason="pinned diffusers reference not installed") + from diffusers import AutoencoderKLMiniMaxH3Audio + from loguru import logger + from safetensors.torch import load_file + + from ....models.audio_vae.minimax_h3.convert_minimax_h3_audio import convert_minimax_h3_audio_state_dict + + config = load_config(weights_dir) + reference = AutoencoderKLMiniMaxH3Audio(**config).eval() + reference.load_state_dict(load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors"))) + converted = convert_minimax_h3_audio_state_dict(dict(reference.state_dict())) + + torch.manual_seed(2) + latents = torch.randn(2, config["latent_channels"], NUM_LATENT_FRAMES) * 0.1 + + baseline_out = None + baseline_s = None + results = [] + for factor, axis in FACTORS: + pc = None if factor <= 1 else ParallelFactor(factor=factor, mesh_axis=axis) + ccl = None if pc is None else CCLManager(mesh_device, num_links=1, topology=ttnn.Topology.Linear) + try: + decoder = _build(mesh_device, config, converted, pc, ccl) + out = decoder(latents) + seconds = _best(lambda: decoder(latents)) + except Exception as exc: + # Every FACTORS entry is a supported layout, so a raise is a regression, not a "result": + # recording it as unsupported let factor 4 break while factor 8 kept the test green, since + # the asserts below need only *some* factor to run. Optional layouts go in KNOWN_BROKEN. + logger.warning(f"t_factor={factor} axis={axis} FAILED: {str(exc)[:160]}") + results.append((factor, axis, None, None)) + if (factor, axis) not in KNOWN_BROKEN: + raise + continue + + if baseline_out is None: + baseline_out, baseline_s = out, seconds + psnr_db = float("inf") + else: + assert out.shape == baseline_out.shape, f"factor {factor}: {out.shape} != {baseline_out.shape}" + # `psnr_db`, not `psnr`: binding the float to `psnr` shadows the imported helper, so the + # first factor's `float("inf")` makes the second factor's call a TypeError. + psnr_db = psnr(baseline_out, out) + results.append((factor, axis, seconds, psnr_db)) + logger.info( + f"PERF audio_decode t_factor={factor} axis={axis}: {seconds:.4f} s " + f"({baseline_s / seconds:.2f}x) PSNR vs 1-device {psnr_db:.1f} dB" + ) + if psnr_db < 40.0 and out is not baseline_out: + _localize_divergence(baseline_out.float(), out.float(), factor=factor, logger=logger) + del decoder + + logger.info("=== audio decode T-parallel summary ===") + for factor, axis, seconds, psnr_db in results: + if seconds is None: + logger.info(f" t_factor={factor:2d} axis={axis}: unsupported") + else: + logger.info( + f" t_factor={factor:2d} axis={axis}: {seconds:.4f} s {baseline_s / seconds:5.2f}x " + f"PSNR {psnr_db:6.1f} dB" + ) + + # The baseline must have run, or there is nothing to compare against and every other factor was + # skipped for want of a reference. Without this the test PASSES when the whole stack is broken: + # observed once on a device left dirty by an unrelated crash, where all three factors raised + # TT_FATAL, each was swallowed as "unsupported", and the run reported green. A gate that cannot + # fail is worse than no gate. + baseline_ran = any(seconds is not None and factor == 1 for factor, _, seconds, _ in results) + assert baseline_ran, ( + "the single-device baseline did not run, so nothing was compared. If this follows a crashed " + "run, reset the device (`tt-smi -glx_reset`) -- an allocator TT_FATAL here is usually a dirty " + "device, not a code failure" + ) + # `baseline_ran` above catches "everything failed" but not the subtler case: if factor 1 raises, + # the loop's `baseline_out is None` branch promotes the *next* factor to baseline, and that factor + # then scores PSNR inf against itself and reads as correct. That happened for real -- a fusion + # crash killed factor 1, factor 4 became the reference, and a -10.1 dB configuration reported inf + # until the baseline was made to run. The first result must be the unsharded one. + assert results and results[0][0] == 1 and results[0][2] is not None, ( + f"factor 1 must be the baseline, but the first result that ran was {results[0][:2] if results else None}; " + "a later factor has been promoted to baseline and is being compared against itself" + ) + ran = [(f, a) for f, a, seconds, _ in results if seconds is not None and f != 1] + assert ran, "no parallel factor ran at all; the T-parallel path is entirely unavailable" + + # Any factor that ran must agree with the single-device path; a fast wrong answer fails. + for factor, axis, seconds, psnr_db in results: + if seconds is None or (factor, axis) in KNOWN_BROKEN: + continue + assert psnr_db > 40.0, f"t_factor={factor} axis={axis} diverges from 1-device: PSNR {psnr_db:.1f} dB" diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py new file mode 100644 index 0000000000000000000000000000000000000000..ae9baf07011bcdb474d62d5e1cba2dd1c5ce433b --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""Gate M8f: MiniMax-H3 VAE performance baselines, gated against ``EXPECTED_METRICS`` budgets +(style of ``tests/models/wan2_2/test_performance_wan.py``). Tracy capture entry points live +in ``tools/tracy_decode_harness.py``; quality gates live in the vae/audio test files.""" + +import os +import time + +import pytest +import torch + +import ttnn + +from ....models.audio_vae.minimax_h3.convert_minimax_h3_audio import convert_minimax_h3_audio_state_dict +from ....models.audio_vae.minimax_h3.encoder_minimax_h3_audio import MiniMaxH3AudioEncoder +from ....models.vae.minimax_h3.vae_minimax_h3 import MiniMaxH3Vae, MiniMaxH3VaeConfig +from ....parallel.manager import CCLManager +from .common import ( + CLIP_FRAMES, + DECODE_LATENT_FRAMES, + LATENT_TILE, + TILE, + build_audio_decoder, + build_visual_decoder, + build_visual_encoder, + load_config, + random_decoder_state, + random_encoder_state, + weights_subdir, +) + +SINGLE_DEVICE = [pytest.param((1, 1), {"l1_small_size": 65536}, id="single_device")] + +HOP_LENGTH = 800 + +# Work units per working point: 768P = 4x7 tiles, 1440P = 8x13; clip/chunk counts per the 17n+5 -> 5n+2 rule. +WORK_UNITS = { + "768P_5s": {"tiles": 28, "encode_clips": 8, "decode_chunks": 7}, + "768P_10s": {"tiles": 28, "encode_clips": 15, "decode_chunks": 14}, + "768P_15s": {"tiles": 28, "encode_clips": 22, "decode_chunks": 21}, + "1440P_5s": {"tiles": 104, "encode_clips": 8, "decode_chunks": 7}, + "1440P_10s": {"tiles": 104, "encode_clips": 15, "decode_chunks": 14}, + "1440P_15s": {"tiles": 104, "encode_clips": 22, "decode_chunks": 21}, +} + +AUDIO_LATENT_FRAMES_5S = 207 # ~5 s at 40 latent fps; budgets were calibrated at 207, not the exact 200 + +# Seconds per measurement. Generous regression bars, not tuned targets; visual entries gate a 32-unit wave. +# Audio decode runs the accurate-mode defaults (~13 s eager, ~3x the retired fast path); 60 s stays generous. +EXPECTED_METRICS = { + "visual_encoder_clip_wave": 20.0, + "visual_encoder_keyframe_wave": 5.0, + "visual_decoder_wave": 20.0, + "audio_encode_5s": 60.0, + "audio_decode_5s": 60.0, +} + + +def _time_it(fn, *, warmup: int = 1, iterations: int = 2, mesh_device=None) -> float: + for _ in range(warmup): + fn() + if mesh_device is not None: + ttnn.synchronize_device(mesh_device) + best = float("inf") + for _ in range(iterations): + start = time.time() + fn() + if mesh_device is not None: + ttnn.synchronize_device(mesh_device) + best = min(best, time.time() - start) + return best + + +def _report(measurements: dict[str, float]) -> None: + from loguru import logger + + failures = [] + for key, seconds in measurements.items(): + budget = EXPECTED_METRICS[key] + logger.info(f"PERF {key}: {seconds:.3f} s (budget {budget:.1f} s)") + if seconds > budget: + failures.append(f"{key} took {seconds:.3f} s, budget {budget:.1f} s") + assert not failures, "\n".join(failures) + + +MESH_4X8 = [ + pytest.param( + (4, 8), + {"fabric_config": None, "require_exact_physical_num_devices": True, "l1_small_size": 65536}, + id="mesh4x8", + ) +] + + +@pytest.mark.parametrize(("mesh_device", "device_params"), MESH_4X8, indirect=["mesh_device", "device_params"]) +def test_visual_data_parallel_throughput(mesh_device): + """Per-wave time with one work unit per device, gated, plus full-video projections.""" + from loguru import logger + + weights_dir = weights_subdir("vae") + if weights_dir is None: + pytest.skip("MiniMax-H3 vae not found; set MINIMAX_H3_MODEL_PATH") + config = load_config(weights_dir) + devices = mesh_device.get_num_devices() + torch.manual_seed(0) + + measurements = {} + + for label, num_frames, taps in ( + ("visual_encoder_clip_wave", CLIP_FRAMES, 3), + ("visual_encoder_keyframe_wave", 1, 1), + ): + encoder = build_visual_encoder(config, mesh_device, num_frames, temporal_taps=taps) + encoder.load_torch_state_dict(random_encoder_state(config)) + x = torch.randn(devices, num_frames, TILE, TILE, encoder.conv_in.in_channels) + x_device = ttnn.from_torch( + x, + dtype=ttnn.float32, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=0), + ) + seconds = _time_it(lambda: encoder(x_device), iterations=8, mesh_device=mesh_device) + logger.info(f"PERF {label} of {devices} units: {seconds:.3f} s ({seconds / devices:.4f} s/unit)") + measurements[label] = seconds + + decoder = build_visual_decoder(config, mesh_device) + decoder.load_torch_state_dict(random_decoder_state(config)) + tokens = torch.randn(devices, DECODE_LATENT_FRAMES * LATENT_TILE * LATENT_TILE, config["latent_channels"]) + tokens_device = ttnn.from_torch( + tokens, + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=0), + ) + decoder_wave = _time_it(lambda: decoder(tokens_device), iterations=8, mesh_device=mesh_device) + # 10.6 = TFLOP per invocation at 1797 tokens x 36 layers x dim 2048. + logger.info( + f"PERF visual_decoder_wave of {devices} units: {decoder_wave:.3f} s " + f"({decoder_wave / devices:.4f} s/unit, {10.6 / (decoder_wave / devices):.1f} TFLOP/s effective per unit)" + ) + measurements["visual_decoder_wave"] = decoder_wave + + encoder_wave = measurements["visual_encoder_clip_wave"] + for name, units in WORK_UNITS.items(): + encode_units = units["tiles"] * units["encode_clips"] + decode_units = units["tiles"] * units["decode_chunks"] + encode_waves = -(-encode_units // devices) + decode_waves = -(-decode_units // devices) + encode_total = encode_waves * encoder_wave + decode_total = decode_waves * decoder_wave + logger.info( + f"PROJECTED {name}: encode {encode_units} units / {encode_waves} waves = {encode_total:.1f} s, " + f"decode {decode_units} units / {decode_waves} waves = {decode_total:.1f} s, " + f"total {encode_total + decode_total:.1f} s" + ) + + _report(measurements) + + +@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"]) +def test_audio_baselines(mesh_device): + """Audio encode/decode timing baselines at 5 s, against their budgets.""" + weights_dir = weights_subdir("audio_vae") + if weights_dir is None: + pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH") + from safetensors.torch import load_file + + config = load_config(weights_dir) + converted = convert_minimax_h3_audio_state_dict( + load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors")) + ) + + torch.manual_seed(2) + waveform = torch.randn(2, 1, AUDIO_LATENT_FRAMES_5S * HOP_LENGTH) * 0.1 + latents = torch.randn(2, config["latent_channels"], AUDIO_LATENT_FRAMES_5S) * 0.1 + + encoder = MiniMaxH3AudioEncoder( + encoder_dim=config["encoder_dim"], + encoder_rates=tuple(config["encoder_rates"]), + latent_dim=config["latent_dim"], + latent_channels=config["latent_channels"], + num_attention_heads=config["num_attention_heads"], + mesh_device=mesh_device, + ) + decoder = build_audio_decoder(config, mesh_device) + # `strict=False`: the converted dict carries both halves' tensors. + encoder.load_torch_state_dict(converted, strict=False) + decoder.load_torch_state_dict(converted, strict=False) + + measurements = { + "audio_encode_5s": _time_it(lambda: encoder(waveform), iterations=1, mesh_device=mesh_device), + "audio_decode_5s": _time_it(lambda: decoder(latents), iterations=1, mesh_device=mesh_device), + } + _report(measurements) + + +# ---- whole-stage `decode` timing, per output type ------------------------------------------- +# +# `test_visual_data_parallel_throughput` times one decoder wave; this times the stage the pipeline +# bills to "VAE decode" -- tiling, upload, readback, unpatchify, stitch and the temporal assembly -- +# and prints `_report_profile`'s breakdown, which is the point: a stage total alone cannot say which +# of device / readback / stitch moved. +# +# RING, not the `fabric_config=None` the throughput test uses: `_decode_clip_device_stitched` +# hardcodes `all_gather(topology=Ring)`, so `yuv420` cannot run on a line fabric. Built with a +# `CCLManager` because `_read_wave_units` picks its reader on `ccl_manager is None` -- without one a +# single-host run silently takes the `ConcatMeshToTensor` path and measures something production +# never runs. +# +# export MINIMAX_H3_MODEL_PATH=/path/to/MiniMax-H3 +# pytest models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py -k decode_stage -s + +MESH_4X8_RING = [ + pytest.param( + (4, 8), + { + "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING, + "require_exact_physical_num_devices": True, + "l1_small_size": 65536, + }, + id="4x8ring", + ) +] + +# 1344x768 at 16x spatial compression. Latent frame counts follow the 17n -> 5n rule less +# `token_drop`, so 124 frames (5 s) is 37 and 362 (15 s) is 107. +DECODE_STAGE_LATENT_HW = (48, 84) +DECODE_STAGE_FRAMES = {5: 37, 15: 107} +# 7 latent frames is one chunk, the cheapest warm-up that still exercises a whole wave. +DECODE_STAGE_WARM_FRAMES = 7 + +MINIMAX_H3_PIXEL_MEAN = (0.485, 0.456, 0.406) +MINIMAX_H3_PIXEL_STD = (0.229, 0.224, 0.225) + +# Each mode is one lever, so a row of this table measures that lever alone. +DECODE_STAGE_MODES = { + # Today's default: bf16 tiles in token space, blended and unpatchified on host. + "float": {}, + # Same path, uint8 across PCIe. Halves the transfer. + "uint8": {"pixel_denorm": (MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD), "readback_uint8": True}, + # Stitch, clamp and colour-convert on device; read one planar canvas at 1.5 bytes/pixel. + "yuv420": {"pixel_denorm": (MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD), "device_stitch": True}, +} + + +def _decode_stage_state(weights_dir: str) -> dict[str, torch.Tensor]: + """Just the decoder-side tensors; over half the checkpoint is the encoder.""" + import json + + from safetensors.torch import load_file + + index_path = os.path.join(weights_dir, "diffusion_pytorch_model.safetensors.index.json") + if not os.path.isfile(index_path): + return { + k: v + for k, v in load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors")).items() + if k.startswith(("decoder.", "post_quant_conv.")) + } + weight_map = json.loads(open(index_path).read())["weight_map"] + wanted = {k: f for k, f in weight_map.items() if k.startswith(("decoder.", "post_quant_conv."))} + state: dict[str, torch.Tensor] = {} + for shard in sorted(set(wanted.values())): + loaded = load_file(os.path.join(weights_dir, shard)) + state.update({k: loaded[k] for k in wanted if k in loaded}) + return state + + +@pytest.mark.timeout(3600) +@pytest.mark.parametrize(("mesh_device", "device_params"), MESH_4X8_RING, indirect=["mesh_device", "device_params"]) +@pytest.mark.parametrize("seconds", [5, 15], ids=["5s", "15s"]) +@pytest.mark.parametrize("mode", sorted(DECODE_STAGE_MODES), ids=sorted(DECODE_STAGE_MODES)) +def test_decode_stage(mesh_device, seconds, mode): + """Measurement driver, not a gate: reports the split, asserts nothing about wall time.""" + from loguru import logger + + weights_dir = weights_subdir("vae") + if weights_dir is None: + pytest.skip("MiniMax-H3 vae not found; set MINIMAX_H3_MODEL_PATH") + + config = MiniMaxH3VaeConfig.from_pretrained(weights_dir) + ccl_manager = CCLManager(mesh_device, num_links=2, topology=ttnn.Topology.Ring) + vae = MiniMaxH3Vae(config, mesh_device=mesh_device, ccl_manager=ccl_manager, **DECODE_STAGE_MODES[mode]) + output_type = "yuv420" if mode == "yuv420" else "float" + + vae.load_decoder_state(_decode_stage_state(weights_dir)) + # Build the per-shape decoder outside the timed region, as `_prepare_vae(decode_shape=...)` does: + # its weight upload is seconds and would otherwise land inside the first `decode`. + vae._decoder_for(config.tokens_chunk_size + config.token_overlap, 16, 16) + + torch.manual_seed(0) + latent_h, latent_w = DECODE_STAGE_LATENT_HW + warm = torch.randn(1, config.latent_channels, DECODE_STAGE_WARM_FRAMES, latent_h, latent_w) + latents = torch.randn(1, config.latent_channels, DECODE_STAGE_FRAMES[seconds], latent_h, latent_w) + + vae.decode(warm, output_type=output_type) + + started = time.time() + video = vae.decode(latents, output_type=output_type) + elapsed = time.time() - started + + profile = dict(vae.last_decode_profile) + logger.info(f"DECODE_STAGE mode={mode} seconds={seconds} total={elapsed:.3f}s video={tuple(video.shape)}") + logger.info( + f"DECODE_SPLIT mode={mode} seconds={seconds} " + + " ".join( + f"{k}={profile.get(k, 0.0):.3f}" + for k in ( + "device", + "readback", + "stitch", + "unpatchify", + "blend", + "concat", + "tiling", + "upload", + "host_prep", + "residual", + ) + ) + + f" waves={int(profile.get('waves', 0))} units={int(profile.get('units', 0))}" + + f" readback_gb={profile.get('readback_mb', 0.0) / 1000:.2f}" + ) diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_minimax_h3.py new file mode 100644 index 0000000000000000000000000000000000000000..439509d854183de69fad9eab382ced63a4a5815f --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_minimax_h3.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""End-to-end `t2va` perf + quality gate, swept over the published working points. + +Six aspect ratios (21:9 .. 9:16) x three durations (5 / 10 / 15 s), 50 steps. The canvas comes from +`resolve_canvas_size` -- short edge 768 from 16:9 through 9:16, ~1 MPix for wider -- and the frame +count from `align_num_frames`, so neither is tabulated here. Each case writes its own artifact stem, +so a sweep does not overwrite itself. +""" + +from __future__ import annotations + +import os + +import pytest +from loguru import logger + +import ttnn + +from ....pipelines.minimax_h3.packing import MINIMAX_H3_FPS, align_num_frames, resolve_canvas_size +from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline +from ..wan2_2.common import check_output_sanity +from .common import GALAXY_MESHES +from .common_av import ( + CALIBRATED_FOX_PROMPT, + artifact_dir, + check_audio_sanity, + check_av_sync, + check_spatial_seams, + check_written_file, + gate_clip, + gate_vbench, + log_spectral_flatness, + log_timing_table, + run_warm_generation, + to_uint8_frames, + weights_dir, + write_artifacts, +) + +NUM_INFERENCE_STEPS = 50 +SEED = 0 + +# The published working points: short edge 768 from 16:9 through 9:16, and ~1 MPix for anything +# wider. `resolve_canvas_size` already implements exactly that (768 short edge, area capped at +# 768*1344, both axes snapped to 32), so the canvas is derived rather than tabulated here -- +# 21:9 lands on 672x1536, which is the documented example. +ASPECT_RATIOS = [(21, 9), (16, 9), (4, 3), (1, 1), (3, 4), (9, 16)] +DURATIONS_S = [5, 10, 15] + +# Loose did-something-collapse bar, per second of video rather than absolute: the 400 s figure was +# set for one 5.17 s clip, and 10 s / 15 s at 1 MPix cost proportionally more. 77 s of compute per +# video second is the same allowance (400 / 5.167), i.e. still ~5x looser than the 13.4x realtime +# factor measured at 16:9 / 5 s. +EXPECTED_S_PER_VIDEO_SECOND = 77.0 + + +# calibrated 2026-08-04, fox prompt, seed 0 (single sample; margins are generous) +VBENCH_THRESHOLDS = { + "subject_consistency": 0.95, + "background_consistency": 0.95, + "motion_smoothness": 0.97, + "dynamic_degree": 1.0, + "imaging_quality": 0.64, +} +CLIP_THRESHOLD = 33.0 # measured mean 37.05 (2026-08-04, fox prompt, seed 0) + +# tier-6 thresholds are calibrated against this exact prompt; swapping it invalidates both bars. +PROMPT = CALIBRATED_FOX_PROMPT + +SWEEP = [ + pytest.param(ratio, seconds, id=f"{ratio[0]}x{ratio[1]}_{seconds}s") + for seconds in DURATIONS_S + for ratio in ASPECT_RATIOS +] + + +@pytest.mark.timeout(7200) +@pytest.mark.parametrize(("aspect_ratio", "duration_s"), SWEEP) +@pytest.mark.parametrize(("mesh_device", "device_params"), GALAXY_MESHES, indirect=["mesh_device", "device_params"]) +def test_t2va_end_to_end(mesh_device, reset_seeds, aspect_ratio, duration_s): + weights = weights_dir("transformer", "text_encoder", "vae", "audio_vae") + artifacts = artifact_dir("h3_t2va_artifacts") + prompt = PROMPT + + HEIGHT, WIDTH = resolve_canvas_size(*aspect_ratio) + NUM_FRAMES = align_num_frames(round(duration_s * MINIMAX_H3_FPS)) + # One artifact per working point, so a sweep does not overwrite itself. + stem = f"t2va_{aspect_ratio[0]}x{aspect_ratio[1]}_{WIDTH}x{HEIGHT}_{duration_s}s" + logger.info(f"working point: {aspect_ratio[0]}:{aspect_ratio[1]} -> {WIDTH}x{HEIGHT}, {NUM_FRAMES} frames, {stem}") + + # A missing dependency must report SKIPPED before the long generation, never silently pass as green. + pytest.importorskip("open_clip", reason="the CLIP gate needs open_clip, which is not installed") + + if not os.environ.get("TT_DIT_CACHE_DIR"): + logger.warning( + "TT_DIT_CACHE_DIR is unset, so every weight load reads safetensors. Prepares are excluded " + "from the total either way, but the run will take far longer than the reported compute." + ) + + pipeline = MiniMaxH3Pipeline.create_pipeline(mesh_device=mesh_device, weights_dir=weights) + + output = run_warm_generation( + pipeline, + prompt, + num_frames=NUM_FRAMES, + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_INFERENCE_STEPS, + seed=SEED, + ) + + expected_frames = align_num_frames(NUM_FRAMES) + logger.info( + f"generated {output.num_frames} frames ({output.video_seconds:.3f} s) and " + f"{output.audio_seconds:.3f} s of audio at {output.sampling_rate} Hz" + ) + + num_forwards = NUM_INFERENCE_STEPS - 1 + video_seconds = expected_frames / MINIMAX_H3_FPS + log_timing_table( + pipeline, + "t2va", + num_forwards=num_forwards, + video_seconds=video_seconds, + expected_total_s=EXPECTED_S_PER_VIDEO_SECOND * video_seconds, + extra=( + f" | {aspect_ratio[0]}:{aspect_ratio[1]} {WIDTH}x{HEIGHT}, {expected_frames} frames " + f"@ {MINIMAX_H3_FPS} fps ({video_seconds:.2f} s), {num_forwards} forwards, " + f"padded_len {pipeline.last_padded_len}" + ), + ) + + frames = to_uint8_frames(output) + + check_output_sanity(frames, num_frames=expected_frames, height=HEIGHT, width=WIDTH) + check_audio_sanity( + output.audio, + sampling_rate=output.sampling_rate, + expected_seconds=expected_frames / MINIMAX_H3_FPS, + ) + sync = check_av_sync( + frames, + output.audio, + sampling_rate=output.sampling_rate, + fps=MINIMAX_H3_FPS, + ) + log_spectral_flatness(output.audio, sampling_rate=output.sampling_rate) + + # Boundaries from the VAE's own tile grid; re-deriving them with a wrong overlap gates non-boundaries. + ratio = pipeline.vae_config.spatial_compression_ratio + (y_starts, _, _), (x_starts, _, _) = pipeline.vae.decode_tile_grid(HEIGHT // ratio, WIDTH // ratio) + check_spatial_seams(frames, vertical_boundaries=x_starts[1:], horizontal_boundaries=y_starts[1:]) + + # Rank 0 alone writes and scores, and its outcome is held rather than raised: every other rank + # goes straight to the barrier below, so a gate failure, a VBench `pytest.skip` or an artifact I/O + # error raised here would leave them blocked and skip the trace release. `pytest.skip.Exception` + # is named alongside `Exception` because `Skipped` does not derive from it. + outcome: BaseException | None = None + is_distributed = ttnn.using_distributed_env() + if not is_distributed or int(ttnn.distributed_context_get_rank()) == 0: + try: + paths = write_artifacts(frames, output.audio.cpu().numpy(), output.sampling_rate, artifacts, stem=stem) + check_written_file(paths, expected_frames, height=HEIGHT, width=WIDTH) + + # CLIP_THRESHOLD was measured at 16:9 / 5 s. Applying it across the sweep is an extrapolation, + # but a generous one: the calibrated point measures ~37 against a bar of 33, and the score is a + # prompt-alignment number rather than a resolution-dependent one. A prompt change would + # invalidate it outright -- recalibrate before swapping PROMPT. + gate_clip(frames, prompt, CLIP_THRESHOLD, stem) + # RUN_VBENCH=0 drops the VBench gate. It needs its own interpreter (~/vbench_env, pinned to + # numpy<2 / transformers 4.33), and without one `run_vbench` skips -- which marks the whole + # test SKIPPED *after* the full generation, hiding the perf and A/V results behind a + # non-result. Off means "not measured": everything above still gates. + if os.environ.get("RUN_VBENCH", "1") not in ("0", "false", "False"): + gate_vbench(paths, prompt, VBENCH_THRESHOLDS, stem) + else: + logger.warning("RUN_VBENCH=0, so the VBench gate did not run; generative quality is UNMEASURED") + + logger.info(f"artifacts in {artifacts}: {sorted(p.name for p in artifacts.iterdir())}") + except (Exception, pytest.skip.Exception) as exc: + outcome = exc + if is_distributed: + ttnn.distributed_context_barrier() + logger.info( + "REMINDER: read the artifact rubric against these frames -- the seam and flicker scores above " + "are statistics, and only looking at the output catches what they average away" + ) + # Every rank, and before the fixture closes the mesh: a captured trace holds device buffers + # for the request. Only the 4x32 preset traces, so this is a no-op at 4x8. + pipeline.release_traces() + # Now that every rank has resynchronised and released its trace, rank 0 reports. mpirun surfaces + # a single failing rank as a failing job, so the run fails rather than hanging. + if outcome is not None: + raise outcome + assert sync["video_seconds"] > 0 diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_references_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_references_minimax_h3.py new file mode 100644 index 0000000000000000000000000000000000000000..19a218444123dcd10851a48fcf2eb66402c74b71 --- /dev/null +++ b/code/models/tt_dit/tests/models/minimax_h3/test_references_minimax_h3.py @@ -0,0 +1,809 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax-H3 ``ref2va`` reference conditioning, gated against the reference implementation's +own path (``diffusers.modular_pipelines.minimax_h3`` at PR #14355): host gates are bit-exact +(drift is silent), the device VAE encode is PCC on real media. importorskips that branch.""" + +from __future__ import annotations + +import os +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from loguru import logger + +from ....pipelines.minimax_h3 import packing as p +from ....pipelines.minimax_h3 import packing_ref2va as rp +from ....pipelines.minimax_h3 import references as R +from ....utils.check import assert_quality +from ....utils.test import ring_params_req_exact_devices + +# Base `packing` too: the rotary-clock gate pins both `_temporal_position_span` implementations. +reference_before_encoder, reference_packing, reference_encoders, reference_base_packing = ( + pytest.importorskip(module, reason="requires the minimax-h3 diffusers branch") + for module in ( + "diffusers.modular_pipelines.minimax_h3.before_encoder", + "diffusers.modular_pipelines.minimax_h3.packing_ref2va", + "diffusers.modular_pipelines.minimax_h3.encoders", + "diffusers.modular_pipelines.minimax_h3.packing", + ) +) + +AUDIO_RATE = 32000 +TARGET_FRAMES = 124 +DURATION = TARGET_FRAMES / p.MINIMAX_H3_FPS # 5.1667 s + +TARGET_HEIGHT, TARGET_WIDTH = 768, 1344 +VAE_RATIO = 16 +AUDIO_HOP = 800 +PATCH_SIZE = (1, 2, 2) + + +def _components(): + return SimpleNamespace(audio_sampling_rate=AUDIO_RATE) + + +def _rng(): + return np.random.default_rng(0) + + +def _image(width: int, height: int) -> np.ndarray: + """Structured content, not flat: a resize of a constant field cannot show an error.""" + rng = _rng() + return (rng.random((height, width, 3)) * 255).astype(np.uint8) + + +def _video(width: int, height: int, num_frames: int) -> np.ndarray: + rng = _rng() + return (rng.random((num_frames, height, width, 3)) * 255).astype(np.uint8) + + +def _waveform(seconds: float, sample_rate: int = AUDIO_RATE, channels: int = 2) -> torch.Tensor: + torch.manual_seed(0) + return torch.randn(channels, int(seconds * sample_rate)) * 0.1 + + +def _pair(specs: list[dict]): + ours = [rp.MiniMaxH3Reference(**spec) for spec in specs] + theirs = [reference_packing.MiniMaxH3Reference(**spec) for spec in specs] + return ours, theirs + + +SPECS = { + "image": dict(image=_image(1024, 1024)), + "image_16to9": dict(image=_image(1920, 1080)), + "video_24fps_at_canvas": dict(video=_video(1344, 768, TARGET_FRAMES), fps=24.0), + "video_30fps_off_canvas": dict(video=_video(1920, 1080, 60), fps=30.0), + "video_with_sound": dict(video=_video(1344, 768, TARGET_FRAMES), fps=24.0, audio=_waveform(DURATION)), + "audio": dict(audio=_waveform(DURATION)), + "audio_44k_mono": dict(audio=_waveform(DURATION, 44100, 1), sample_rate=44100), +} + + +def test_reference_limits_match_reference_and_are_enforced(expect_error): + assert rp.MINIMAX_H3_MAX_REFERENCE_IMAGES == reference_packing.MINIMAX_H3_MAX_REFERENCE_IMAGES == 9 + assert rp.MINIMAX_H3_MAX_REFERENCE_VIDEOS == reference_packing.MINIMAX_H3_MAX_REFERENCE_VIDEOS == 3 + assert rp.MINIMAX_H3_MAX_REFERENCE_AUDIOS == reference_packing.MINIMAX_H3_MAX_REFERENCE_AUDIOS == 3 + assert rp.MINIMAX_H3_MAX_REFERENCES == reference_packing.MINIMAX_H3_MAX_REFERENCES == 12 + assert rp.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE == reference_packing.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE == 2048 + + image = rp.MiniMaxH3Reference(image=_image(256, 256)) + video = rp.MiniMaxH3Reference(video=_video(256, 256, 10), fps=24.0) + audio = rp.MiniMaxH3Reference(audio=_waveform(DURATION)) + + assert R.check_references([image, video, audio]) == ["image", "video", "audio"] + + with expect_error(ValueError, "at least one reference"): + R.check_references([]) + with expect_error(ValueError, "at most 9 image"): + R.check_references([image] * 10) + with expect_error(ValueError, "at most 3 video"): + R.check_references([video] * 4) + with expect_error(ValueError, "at most 3 audio"): + R.check_references([image, audio, audio, audio, audio]) + with expect_error(ValueError, "at most 12 references"): + R.check_references([image] * 9 + [video] * 3 + [audio]) + with expect_error(ValueError, "paired with at least one image or video"): + R.check_references([audio]) + + +@pytest.mark.parametrize("name", list(SPECS), ids=list(SPECS)) +def test_prepare_references_matches_reference(name): + spec = SPECS[name] + # audio may not be alone; pair with an image + specs = [spec] if "audio" not in name else [SPECS["image"], spec] + ours, theirs = _pair(specs) + + got, got_frames = R.prepare_references(ours, TARGET_FRAMES, AUDIO_RATE) + want, want_frames = reference_before_encoder.MiniMaxH3Ref2VASetupStep.prepare_references( + _components(), theirs, TARGET_FRAMES + ) + + assert got_frames == want_frames == TARGET_FRAMES + assert len(got) == len(want) + for index, (a, b) in enumerate(zip(got, want)): + assert a.kind == b.kind, index + assert a.has_audio == b.has_audio, index + if a.kind == "image": + assert np.array_equal(np.asarray(a.image), np.asarray(b.image)), f"prepared image {index} differs" + assert a.image.size == b.image.size + if a.kind == "video": + assert np.array_equal(a.frames, b.frames), f"prepared frames {index} differ" + if a.has_audio: + assert torch.equal(a.waveform, b.waveform), f"prepared waveform {index} differs" + assert a.waveform.shape[0] == p.MINIMAX_H3_AUDIO_CHANNELS + + +def test_prepare_references_uses_each_references_own_resolution(): + references = [ + rp.MiniMaxH3Reference(image=_image(1024, 1024)), + rp.MiniMaxH3Reference(video=_video(1920, 1080, 30), fps=24.0), + ] + prepared, _ = R.prepare_references(references, TARGET_FRAMES, AUDIO_RATE) + + assert prepared[0].image.size == (2048, 2048) + assert min(prepared[0].image.size) == rp.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE + assert prepared[1].frames.shape[1:3] == (768, 1344) + assert min(prepared[1].frames.shape[1:3]) == p.MINIMAX_H3_SHORT_EDGE + assert min(prepared[0].image.size) > min(prepared[1].frames.shape[1:3]) + + +def test_prepare_references_truncates_a_video_and_its_soundtrack_to_the_target(): + long_frames = _video(1344, 768, TARGET_FRAMES + 60) + references = [rp.MiniMaxH3Reference(video=long_frames, fps=24.0, audio=_waveform(10.0))] + prepared, num_frames = R.prepare_references(references, TARGET_FRAMES, AUDIO_RATE) + + assert num_frames == TARGET_FRAMES + assert prepared[0].frames.shape[0] == TARGET_FRAMES + assert prepared[0].waveform.shape[1] == int(DURATION * AUDIO_RATE) + + +def test_num_frames_may_be_derived_from_a_single_audio_bearing_reference(): + references = [rp.MiniMaxH3Reference(image=_image(512, 512)), rp.MiniMaxH3Reference(audio=_waveform(6.0))] + _, num_frames = R.prepare_references(references, None, AUDIO_RATE) + assert num_frames == p.align_num_frames(round(6.0 * p.MINIMAX_H3_FPS)) + assert num_frames % p.MINIMAX_H3_FRAMES_PER_CHUNK == p.MINIMAX_H3_LATENTS_PER_CHUNK + theirs = [ + reference_packing.MiniMaxH3Reference(image=_image(512, 512)), + reference_packing.MiniMaxH3Reference(audio=_waveform(6.0)), + ] + _, want = reference_before_encoder.MiniMaxH3Ref2VASetupStep.prepare_references(_components(), theirs, None) + assert num_frames == want + + +def test_num_frames_is_ambiguous_with_two_soundtracks(expect_error): + references = [ + rp.MiniMaxH3Reference(video=_video(768, 768, 24), fps=24.0, audio=_waveform(6.0)), + rp.MiniMaxH3Reference(audio=_waveform(7.0)), + ] + with expect_error(ValueError, "exactly one of them carries audio"): + R.prepare_references(references, None, AUDIO_RATE) + + +@pytest.mark.parametrize("seconds", [4.0, 16.0]) +def test_a_derived_duration_outside_the_models_range_is_rejected(seconds, expect_error): + references = [rp.MiniMaxH3Reference(image=_image(512, 512)), rp.MiniMaxH3Reference(audio=_waveform(seconds))] + with expect_error(ValueError, "seconds"): + R.prepare_references(references, None, AUDIO_RATE) + + +@pytest.mark.parametrize("seconds", [DURATION, 5.0, 8.0, 15.0]) +def test_pad_waveform_to_hop_reproduces_the_reference_latent_count(seconds): + waveform = _waveform(seconds) + padded = R.pad_waveform_to_hop(waveform) + + assert padded.shape[-1] % R.MINIMAX_H3_AUDIO_HOP == 0 + assert padded.shape[-1] - waveform.shape[-1] < R.MINIMAX_H3_AUDIO_HOP + assert torch.equal(padded[:, : waveform.shape[-1]], waveform), "padding must not disturb the samples" + assert (padded[:, waveform.shape[-1] :] == 0).all(), "the pad is zeros, as F.pad's default is" + assert padded.shape[-1] // R.MINIMAX_H3_AUDIO_HOP == int(np.ceil(waveform.shape[-1] / R.MINIMAX_H3_AUDIO_HOP)) + + +def test_production_soundtrack_is_not_a_whole_number_of_hops(): + """The case that makes the padding load-bearing rather than defensive.""" + samples = int(DURATION * AUDIO_RATE) + assert samples == 165333 + assert samples % R.MINIMAX_H3_AUDIO_HOP != 0 + assert R.pad_waveform_to_hop(_waveform(DURATION)).shape[-1] // R.MINIMAX_H3_AUDIO_HOP == 207 + + +def test_span_matches_both_reference_summation_orders(): + for n in range(1, 61): + assert p._temporal_position_span(n) == reference_base_packing._temporal_position_span( + n + ), f"pairwise span disagrees with the reference at n={n}" + assert rp._temporal_position_span(n) == reference_packing._temporal_position_span( + n + ), f"sequential span disagrees with the reference at n={n}" + + +def test_the_two_span_orders_actually_differ(): + assert p._temporal_position_span(15) == rp._temporal_position_span(15), "expected agreement below 16" + assert p._temporal_position_span(16) != rp._temporal_position_span(16), "expected divergence from 16 latent frames" + + production = p.video_latent_num_frames(TARGET_FRAMES) + assert production == 37 + pairwise, sequential = p._temporal_position_span(production), rp._temporal_position_span(production) + assert pairwise != sequential, "the two orders must differ at the production frame count" + assert rp._temporal_position_span(16) > p._temporal_position_span(16) + assert sequential < pairwise + + +def _target(): + return ( + p.video_latent_num_frames(TARGET_FRAMES), + TARGET_HEIGHT // VAE_RATIO, + TARGET_WIDTH // VAE_RATIO, + p.audio_latent_num_frames(TARGET_FRAMES), + ) + + +def _image_geometry(source_width: int, source_height: int): + height, width = rp.resolve_reference_image_size(source_width, source_height) + return dict(num_latent_frames=1, latent_height=height // VAE_RATIO, latent_width=width // VAE_RATIO) + + +def _video_geometry(source_width: int, source_height: int, num_frames: int = TARGET_FRAMES): + frames = np.zeros((num_frames, source_height, source_width, 3), dtype=np.uint8) + frames = rp.prepare_reference_frames(frames, TARGET_FRAMES) + trimmed = rp.trim_reference_num_frames(frames.shape[0]) + return dict( + num_latent_frames=p.video_latent_num_frames(trimmed), + latent_height=frames.shape[1] // VAE_RATIO, + latent_width=frames.shape[2] // VAE_RATIO, + ) + + +def _audio_latents(seconds: float) -> int: + return int(np.ceil(int(seconds * AUDIO_RATE) / AUDIO_HOP)) + + +def _spec_to_prepared(spec, dataclass): + kind = spec["kind"] + geometry = {} + if kind == "image": + geometry = _image_geometry(*spec["source"]) + elif kind == "video": + geometry = _video_geometry(*spec["source"], spec.get("num_frames", TARGET_FRAMES)) + audio_seconds = spec.get("audio_seconds") + return dataclass( + kind=kind, + has_audio=audio_seconds is not None, + num_audio_latents=0 if audio_seconds is None else _audio_latents(audio_seconds), + **geometry, + ) + + +def _both(spec_list): + ours = [_spec_to_prepared(spec, rp.MiniMaxH3PreparedReference) for spec in spec_list] + theirs = [_spec_to_prepared(spec, reference_packing.MiniMaxH3PreparedReference) for spec in spec_list] + return ours, theirs + + +IMAGE_1TO1 = dict(kind="image", source=(1024, 1024)) +IMAGE_16TO9 = dict(kind="image", source=(1920, 1080)) +VIDEO_SOUND = dict(kind="video", source=(1344, 768), audio_seconds=DURATION) +# 1:1 sounded video: its soundtrack rows must pin to its OWN width grid, not the target's. +VIDEO_SOUND_1TO1 = dict(kind="video", source=(768, 768), audio_seconds=DURATION) +VIDEO_SILENT = dict(kind="video", source=(1344, 768)) +AUDIO_ONLY = dict(kind="audio", audio_seconds=DURATION) + +# Awkward order on purpose: exercises the label counters, soundtrack-before-video order and the +# shared rotary clock; the reversed copy is the host order-sensitivity gate (order is semantic). +_NINE_MIXED = [ + VIDEO_SOUND, + AUDIO_ONLY, + VIDEO_SILENT, + IMAGE_1TO1, + AUDIO_ONLY, + IMAGE_16TO9, + VIDEO_SOUND, + IMAGE_1TO1, + AUDIO_ONLY, +] + +CASES = { + "one_image": [IMAGE_1TO1], + "one_audio": [IMAGE_1TO1, AUDIO_ONLY], # audio may never be the only reference + "video_with_sound": [VIDEO_SOUND], + "video_with_sound_1to1": [VIDEO_SOUND_1TO1], + "video_without_sound": [VIDEO_SILENT], + "nine_mixed_awkward": _NINE_MIXED, + "nine_mixed_reversed": list(reversed(_NINE_MIXED)), +} + +TEXT_LEN = 1024 +# A reference's vision block lands inside the text rows tagged as VIDEO, not text. +VISION_BLOCK = slice(30, 120) + + +def _text_tags(): + tags = torch.ones(TEXT_LEN, dtype=torch.long) + tags[VISION_BLOCK] = p.MINIMAX_H3_VIDEO_TAG + return tags + + +@pytest.mark.parametrize("case", list(CASES), ids=list(CASES)) +def test_layout_matches_reference(case): + ours, theirs = _both(CASES[case]) + num_latent_frames, latent_height, latent_width, num_audio_latents = _target() + args = (num_latent_frames, latent_height, latent_width, num_audio_latents, PATCH_SIZE) + + got = rp.build_ref2va_packed_sequence(_text_tags(), ours, *args) + want = reference_packing.build_ref2va_packed_sequence(_text_tags(), theirs, *args) + + assert got.sequence_length == want.sequence_length + assert got.position_ids.dtype == torch.float64 + assert torch.equal(got.position_ids, want.position_ids), "the fp64 rotary grid is not bit-exact" + assert torch.equal(got.token_tags, want.token_tags) + assert torch.equal(got.video_indices, want.video_indices) + assert torch.equal(got.audio_indices, want.audio_indices) + assert torch.equal(got.text_indices, want.text_indices) + assert got.num_condition_video_rows == want.num_condition_video_rows + assert got.num_condition_audio_rows == want.num_condition_audio_rows + + +# `(width, height)` in, `(height, width)` out; both orientations listed so a transpose cannot pass. +@pytest.mark.parametrize( + "source,expected", + [ + ((1024, 1024), (2048, 2048)), + ((1920, 1080), (2048, 3648)), + ((1080, 1920), (3648, 2048)), + ((4096, 1024), (2048, 8192)), + ], +) +def test_resolve_reference_image_size(source, expected): + assert rp.resolve_reference_image_size(*source) == expected + assert rp.resolve_reference_image_size(*source) == reference_packing.resolve_reference_image_size(*source) + height, width = expected + assert height % p.MINIMAX_H3_CANVAS_MULTIPLE == 0 and width % p.MINIMAX_H3_CANVAS_MULTIPLE == 0 + assert min(height, width) == rp.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE + + +def test_resolve_reference_image_size_rejects_out_of_range(expect_error): + with expect_error(ValueError, "within 1:4 and 4:1"): + rp.resolve_reference_image_size(5000, 1000) + with expect_error(ValueError, "positive size"): + rp.resolve_reference_image_size(0, 100) + + +@pytest.mark.parametrize("fps", [24.0, 30.0, 25.0, 12.0, 60.0, 23.976]) +def test_resample_reference_frames_matches_reference(fps): + frames = np.arange(50 * 2 * 2 * 3, dtype=np.uint8).reshape(50, 2, 2, 3) + got = rp.resample_reference_frames(frames, fps) + want = reference_packing.resample_reference_frames(frames, fps) + assert got.shape == want.shape + assert np.array_equal(got, want) + if fps == 24.0: + assert got is frames + + +def test_prepare_reference_frames_uses_its_own_aspect_canvas(): + frames = np.zeros((10, 480, 480, 3), dtype=np.uint8) + prepared = rp.prepare_reference_frames(frames, TARGET_FRAMES) + assert prepared.shape[1:3] == p.resolve_canvas_size(480, 480) == (768, 768) + assert np.array_equal(prepared, reference_packing.prepare_reference_frames(frames, TARGET_FRAMES)) + + # `shares_memory`, not `is`: the frame-count cap slices first, so a view comes back. + at_canvas = np.zeros((10, 768, 1344, 3), dtype=np.uint8) + passthrough = rp.prepare_reference_frames(at_canvas, TARGET_FRAMES) + assert np.shares_memory(passthrough, at_canvas) + + +@pytest.mark.parametrize("num_frames", [1, 5, 12, 24, 25, 124, 192]) +def test_sample_reference_video_frames_matches_reference(num_frames): + frames = np.zeros((num_frames, 4, 4, 3), dtype=np.uint8) + got_frames, got_ts = rp.sample_reference_video_frames(frames) + want_frames, want_ts = reference_packing.sample_reference_video_frames(frames) + assert len(got_frames) == len(want_frames) + assert got_ts == want_ts + assert len(got_ts) == -(-len(got_frames) // rp.MINIMAX_H3_QWEN_TEMPORAL_PATCH) + + +def test_block_timestamps_round_half_to_even(): + frames = np.zeros((124, 4, 4, 3), dtype=np.uint8) + _, timestamps = rp.sample_reference_video_frames(frames) + assert timestamps[0] == 0.25 + assert f"{timestamps[0]:.1f}" == "0.2" + + +@pytest.mark.parametrize("num_frames", [1, 5, 22, 39, 124, 125, 130]) +def test_trim_reference_num_frames_matches_reference(num_frames): + got = rp.trim_reference_num_frames(num_frames) + assert got == reference_packing.trim_reference_num_frames(num_frames) + assert got % p.MINIMAX_H3_FRAMES_PER_CHUNK == p.MINIMAX_H3_LATENTS_PER_CHUNK + + +def test_trim_reference_num_frames_has_a_22_frame_floor(): + """1..23 frames all map to 22 -- not a plain round-down; a "fixed" floor of 5 silently changes the latent count.""" + for num_frames in (1, 5, 10, 22, 23): + assert rp.trim_reference_num_frames(num_frames) == 22 + assert rp.trim_reference_num_frames(num_frames) == reference_packing.trim_reference_num_frames(num_frames) + assert rp.trim_reference_num_frames(39) == 39 + assert rp.trim_reference_num_frames(55) == 39 + assert rp.trim_reference_num_frames(125) == 124 + + +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("sample_rate", [32000, 44100]) +def test_prepare_reference_waveform_matches_reference(channels, sample_rate): + torch.manual_seed(0) + waveform = torch.randn(channels, int(6.0 * sample_rate)) + got = rp.prepare_reference_waveform(waveform, sample_rate, AUDIO_RATE, DURATION) + want = reference_packing.prepare_reference_waveform(waveform, sample_rate, AUDIO_RATE, DURATION) + assert got.shape == want.shape + assert torch.equal(got, want) + assert got.shape[0] == p.MINIMAX_H3_AUDIO_CHANNELS + if sample_rate == AUDIO_RATE: + assert got.shape[1] == int(DURATION * sample_rate) + + +def test_reference_media_to_uint8_matches_reference(): + torch.manual_seed(0) + tensor = torch.rand(3, 8, 6) + array = (np.random.default_rng(0).random((8, 6, 3)) * 255).astype(np.uint8) + for media in (tensor, array, [tensor, tensor]): + got = rp.reference_media_to_uint8(media) + want = reference_packing.reference_media_to_uint8(media) + assert got.dtype == np.uint8 + assert np.array_equal(got, want) + + +def test_reference_requires_exactly_one_medium(expect_error): + """A video may carry audio; nothing else may carry two media.""" + image = np.zeros((4, 4, 3), dtype=np.uint8) + video = np.zeros((2, 4, 4, 3), dtype=np.uint8) + waveform = torch.zeros(2, 800) + + assert rp.MiniMaxH3Reference(image=image).kind == "image" + assert rp.MiniMaxH3Reference(video=video).kind == "video" + assert rp.MiniMaxH3Reference(audio=waveform).kind == "audio" + sounded = rp.MiniMaxH3Reference(video=video, audio=waveform) + assert sounded.kind == "video" and sounded.has_audio + + with expect_error(ValueError, "exactly one"): + rp.MiniMaxH3Reference(image=image, video=video) + with expect_error(ValueError, "exactly one"): + rp.MiniMaxH3Reference(image=image, audio=waveform) + with expect_error(ValueError, "exactly one"): + rp.MiniMaxH3Reference() + + +def test_reference_refuses_a_path(expect_error): + with expect_error(ValueError, "never opens media files"): + rp.MiniMaxH3Reference(image="subject.png") + + +def test_reference_defaults_fps_to_the_models_own(): + video = np.zeros((2, 4, 4, 3), dtype=np.uint8) + assert rp.MiniMaxH3Reference(video=video).fps == float(p.MINIMAX_H3_FPS) + + +def _prepared(kind, *, frames=1, height=48, width=84, audio_latents=0): + return rp.MiniMaxH3PreparedReference( + kind=kind, + has_audio=audio_latents > 0, + num_latent_frames=frames, + latent_height=height if kind != "audio" else 0, + latent_width=width if kind != "audio" else 0, + num_audio_latents=audio_latents, + ) + + +def test_split_condition_blocks_is_in_packed_order(): + """Per-modality order would keep every row count right and desynchronize the whole region.""" + references = [ + _prepared("video", frames=37, audio_latents=207), + _prepared("audio", audio_latents=100), + _prepared("image"), + ] + video_rows = torch.arange(sum(r.num_video_rows for r in references if r.kind != "audio")).float()[:, None] + audio_rows = torch.arange(sum(r.num_audio_rows for r in references)).float()[:, None] + 10000 + + blocks = R.split_condition_blocks(references, video_rows, audio_rows) + assert [kind for _, kind in blocks] == ["audio", "video", "audio", "video"] + assert [block.shape[0] for block, _ in blocks] == [ + references[0].num_audio_rows, + references[0].num_video_rows, + references[1].num_audio_rows, + references[2].num_video_rows, + ] + + assert torch.equal(blocks[0][0], audio_rows[: references[0].num_audio_rows]) + assert torch.equal(blocks[1][0], video_rows[: references[0].num_video_rows]) + assert torch.equal(blocks[2][0], audio_rows[references[0].num_audio_rows :]) + assert torch.equal(blocks[3][0], video_rows[references[0].num_video_rows :]) + + +def test_split_condition_blocks_row_count_matches_the_layout(): + references = [ + _prepared("image", height=128, width=128), + _prepared("video", frames=37, audio_latents=207), + _prepared("audio", audio_latents=207), + ] + layout = rp.build_ref2va_packed_sequence(torch.ones(64, dtype=torch.long), references, 37, 48, 84, 207, (1, 2, 2)) + video_rows = torch.zeros(layout.num_condition_video_rows, 96) + audio_rows = torch.zeros(layout.num_condition_audio_rows, 32) + + blocks = R.split_condition_blocks(references, video_rows, audio_rows) + assert sum(block.shape[0] for block, _ in blocks) == ( + layout.num_condition_video_rows + layout.num_condition_audio_rows + ) + + +def test_split_condition_blocks_rejects_a_row_count_mismatch(expect_error): + references = [_prepared("image")] + with expect_error(ValueError, "consumed"): + R.split_condition_blocks(references, torch.zeros(references[0].num_video_rows + 7, 96), None) + + +def test_reference_condition_shapes_skips_audio_references(): + references = [_prepared("image"), _prepared("audio", audio_latents=50), _prepared("video", frames=37)] + shapes = R.reference_condition_shapes(references) + assert shapes == ((1, 48, 84), (37, 48, 84)) + + +def test_normalize_reference_pixels_uses_imagenet_statistics(): + """Not [-1, 1]: H3 conditions on VLM-style normalized pixels.""" + frames = np.full((1, 32, 32, 3), 128, dtype=np.uint8) + pixels = R.normalize_reference_pixels(frames) + assert pixels.shape == (1, 3, 1, 32, 32) + expected = [(128 / 255.0 - m) / s for m, s in zip(R.MINIMAX_H3_PIXEL_MEAN, R.MINIMAX_H3_PIXEL_STD)] + for channel, value in enumerate(expected): + assert torch.allclose(pixels[0, channel], torch.full((1, 32, 32), value), atol=1e-6) + assert R.normalize_reference_pixels(np.zeros((5, 8, 8, 3), dtype=np.uint8)).shape == (1, 3, 5, 8, 8) + + +REFERENCE_MEDIA = Path.home() / "h3_fl2va_artifacts" / "fl2va_first.mp4" + + +def _real_media() -> Path: + if not REFERENCE_MEDIA.is_file(): + pytest.skip(f"no reference media at {REFERENCE_MEDIA}; place a video with a soundtrack there") + return REFERENCE_MEDIA + + +def _weights_dir() -> Path: + directory = Path(os.environ.get("MINIMAX_H3_MODEL_PATH", "")) + if not directory.is_dir(): + pytest.skip("set MINIMAX_H3_MODEL_PATH to a diffusers snapshot of the checkpoint") + return directory + + +_L1_SMALL = 16384 # 65536 and 32768 clash with the taps=3 encoder CBs + +# The e2e ring fabric params: a LINE config happens to pass here but production never runs it. +MESH_4X8 = [ + pytest.param( + (4, 8), + {**ring_params_req_exact_devices, "l1_small_size": _L1_SMALL}, + id="4x8", + ) +] + +REFERENCE_PCC = 0.99 # floor set by the encoder's ttnn.group_norm calls (no fp32 path) +# 17*1 + 5: smallest multi-chunk count (7 latent frames); the reference video VAE encode is CPU-hours at 124. +VIDEO_FRAMES = 22 + + +def _reference_components(weights: Path, device: str = "cpu"): + from diffusers.models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio + + vae = AutoencoderKLMiniMaxH3.from_pretrained(str(weights), subfolder="vae").to(device).eval() + audio_vae = AutoencoderKLMiniMaxH3Audio.from_pretrained(str(weights), subfolder="audio_vae").to(device).eval() + return SimpleNamespace( + vae=vae, + audio_vae=audio_vae, + patch_size=(1, 2, 2), + audio_latent_channels=audio_vae.config.latent_channels, + _execution_device=torch.device(device), + ) + + +def _real_frame() -> np.ndarray: + """Not randn: the fp16 round trip is only exercised on natural image statistics.""" + frames, _, _ = rp.decode_reference_video(_real_media()) + return frames[0] + + +def _reference_case(case: str): + """Audio pairs with an image (audio-only is forbidden); the pairing is inert for the comparison.""" + media = _real_media() + frames, fps, soundtrack = rp.decode_reference_video(media) + waveform, sample_rate = soundtrack + + if case == "image": + specs = [dict(image=_real_frame())] + elif case == "video": + specs = [dict(video=frames[:VIDEO_FRAMES], fps=fps)] + elif case == "audio": + specs = [dict(image=_real_frame()), dict(audio=waveform, sample_rate=sample_rate)] + else: + raise ValueError(case) + + ours, _ = R.prepare_references([rp.MiniMaxH3Reference(**s) for s in specs], VIDEO_FRAMES, AUDIO_RATE) + theirs, _ = reference_before_encoder.MiniMaxH3Ref2VASetupStep.prepare_references( + _components(), [reference_packing.MiniMaxH3Reference(**s) for s in specs], VIDEO_FRAMES + ) + return ours, theirs + + +@pytest.mark.timeout(7200) +@pytest.mark.parametrize(("mesh_device", "device_params"), MESH_4X8, indirect=["mesh_device", "device_params"]) +@pytest.mark.parametrize("case", ["image", "audio", "video"]) +def test_encode_references_matches_reference(mesh_device, case, reset_seeds): + """Per modality (the three encode paths fail independently); geometry asserted too, since the layout is built from it.""" + from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + weights = _weights_dir() + ours, theirs = _reference_case(case) + + pipeline = MiniMaxH3Pipeline.create_pipeline(mesh_device=mesh_device, weights_dir=weights, task="ref2va") + needs_video = any(reference.kind == "video" for reference in ours) + vae = pipeline._prepare_vae() + vae = pipeline._prepare_vae(encode_shape=(1, vae.tile_size, vae.tile_size), encode_taps=1) + if needs_video: + vae = pipeline._prepare_vae( + encode_shape=(pipeline.vae_config.clip_length, vae.tile_size, vae.tile_size), encode_taps=3 + ) + audio_encoder = pipeline._prepare_audio_encoder() if any(r.has_audio for r in ours) else None + + got_video, got_audio = R.encode_references( + ours, + encode_clip=vae.encode_clip, + encode_video=vae.encode if needs_video else None, + encode_audio=(lambda waveform: audio_encoder(waveform)[0]) if audio_encoder else None, + latents_mean=pipeline.vae_config.latents_mean, + latents_std=pipeline.vae_config.latents_std, + audio_latents_mean=pipeline.audio_config["latents_mean"], + audio_latents_std=pipeline.audio_config["latents_std"], + patch_size=pipeline.patch_size, + audio_latent_channels=pipeline.audio_config["latent_channels"], + ) + + components = _reference_components(weights) + with torch.no_grad(): + want_video, want_audio = reference_encoders.MiniMaxH3Ref2VAReferenceEncoderStep.encode_references( + components, theirs, device=torch.device("cpu") + ) + + for index, (a, b) in enumerate(zip(ours, theirs)): + assert (a.num_latent_frames, a.latent_height, a.latent_width) == ( + b.num_latent_frames, + b.latent_height, + b.latent_width, + ), f"reference {index} resolved to different visual latent geometry" + assert a.num_audio_latents == b.num_audio_latents, f"reference {index} resolved a different audio length" + logger.info( + f"[{case}] resolved geometry: " + + ", ".join( + f"{r.kind}({r.num_latent_frames}x{r.latent_height}x{r.latent_width}, audio {r.num_audio_latents})" + for r in ours + ) + ) + + assert got_video.shape == want_video.shape, f"{tuple(got_video.shape)} != {tuple(want_video.shape)}" + logger.info(f"[{case}] video condition rows {tuple(got_video.shape)}") + assert_quality(want_video, got_video, pcc=REFERENCE_PCC) + + if got_audio is not None: + assert got_audio.shape == want_audio.shape, f"{tuple(got_audio.shape)} != {tuple(want_audio.shape)}" + logger.info(f"[{case}] audio condition rows {tuple(got_audio.shape)}") + # The soundtrack takes the posterior MEAN, so it carries no sampling noise at all. + assert_quality(want_audio, got_audio, pcc=REFERENCE_PCC) + else: + assert want_audio is None, "the reference produced audio rows where we produced none" + + +def _tokenizer(): + weights = os.environ.get("MINIMAX_H3_MODEL_PATH") + if not weights or not (Path(weights) / "tokenizer").is_dir(): + pytest.skip("needs MINIMAX_H3_MODEL_PATH pointing at a diffusers snapshot with tokenizer/") + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(str(weights), subfolder="tokenizer") + + +@pytest.mark.parametrize("case", list(CASES), ids=list(CASES)) +def test_presentation_matches_reference(case): + """The tags are the load-bearing half: H3 tags the WHOLE vision block video, sentinels included.""" + tokenizer = _tokenizer() + ours, theirs = _both(CASES[case]) + image_counts = [4096 for reference in ours if reference.kind == "image"] + video_counts = [1008 for reference in ours if reference.kind == "video"] + for prepared in (ours, theirs): + for reference in prepared: + if reference.kind == "video": + reference.block_timestamps = [0.25, 1.25, 2.25, 3.25, 4.25, 5.0] + + prompt = "a red fox stepping through wet grass at dawn" + got_ids, got_tags = rp.build_ref2va_presentation(tokenizer, prompt, ours, image_counts, video_counts) + want_ids, want_tags = reference_packing.build_ref2va_presentation( + tokenizer, prompt, theirs, image_counts, video_counts + ) + assert got_ids == want_ids + assert got_tags == want_tags + assert len(got_tags) == len(got_ids) + + +class _ProcessorStub: + def __init__(self, weights: Path): + from transformers import AutoImageProcessor, AutoTokenizer, AutoVideoProcessor + + self.tokenizer = AutoTokenizer.from_pretrained(str(weights), subfolder="tokenizer") + self.image_processor = AutoImageProcessor.from_pretrained(str(weights), subfolder="text_encoder") + self.video_processor = AutoVideoProcessor.from_pretrained(str(weights), subfolder="text_encoder") + + +def _build(stub, prompt, references): + from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + return MiniMaxH3Pipeline._build_ref2va_presentation(stub, prompt, references) + + +def test_presentation_orders_vision_patches_by_reference_not_by_batch(): + """`_scatter_rows` consumes tower rows in run order, so processor-batch order would swap blocks.""" + stub = _ProcessorStub(_weights_dir()) + video = rp.reference_from_video_file(_real_media(), with_audio=False) + image = rp.MiniMaxH3Reference(image=_image(1024, 1024)) + + prepared, _ = R.prepare_references([video, image], TARGET_FRAMES, AUDIO_RATE) + input_ids, tags, type_ids, pixel_values, grid_thw, kinds = _build(stub, "a prompt", prepared) + + assert kinds == ["video", "image"], f"vision entries are in {kinds}, not presentation order" + assert int(grid_thw[0][0]) == len(prepared[0].block_timestamps) > 1 + assert tuple(grid_thw[1].tolist()) == (1, 128, 128) + assert pixel_values.shape[0] == sum(int(grid.prod()) for grid in grid_thw) + + reversed_prepared, _ = R.prepare_references([image, video], TARGET_FRAMES, AUDIO_RATE) + _, _, _, _, reversed_grid, reversed_kinds = _build(stub, "a prompt", reversed_prepared) + assert reversed_kinds == ["image", "video"] + assert tuple(reversed_grid[0].tolist()) == (1, 128, 128) + + +def test_presentation_vision_runs_line_up_with_the_towers_rows(): + from ....encoders.qwen3vl.model_qwen3vl import vision_token_runs + + stub = _ProcessorStub(_weights_dir()) + references = [ + rp.MiniMaxH3Reference(image=_image(1024, 1024)), + rp.reference_from_video_file(_real_media(), with_audio=False), + rp.MiniMaxH3Reference(audio=_waveform(DURATION)), + ] + prepared, _ = R.prepare_references(references, TARGET_FRAMES, AUDIO_RATE) + input_ids, tags, type_ids, pixel_values, grid_thw, kinds = _build(stub, "a prompt", prepared) + + pad_ids = [stub.tokenizer.convert_tokens_to_ids(t) for t in ("<|image_pad|>", "<|video_pad|>")] + runs = vision_token_runs(input_ids, pad_ids) + merge = stub.image_processor.merge_size**2 + + assert len(runs) == sum(int(grid[0]) for grid in grid_thw) + assert sum(length for _, length in runs) == sum(int(grid.prod()) for grid in grid_thw) // merge + # Runs are sorted and disjoint, which `_scatter_rows` requires outright. + assert all(a[0] + a[1] <= b[0] for a, b in zip(runs, runs[1:])) + assert stub.tokenizer.decode(input_ids[0]).count("