Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,25 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+code/infer/example/eg_cn.wav filter=lfs diff=lfs merge=lfs -text
+code/infer/example/edit_cn.mp3 filter=lfs diff=lfs merge=lfs -text
+code/infer/example/edit_en.mp3 filter=lfs diff=lfs merge=lfs -text
+code/infer/example/eg_en.mp3 filter=lfs diff=lfs merge=lfs -text
+code/infer/example/output-windows/en-utf8-fixed.mp3 filter=lfs diff=lfs merge=lfs -text
+code/infer/example/output-windows/en-without-utf8-fix.mp3 filter=lfs diff=lfs merge=lfs -text
+code/infer/example/output/output.mp3 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/3c7495db974af283bccc2dee306a75da6491b2ca664d55597da926c4fd104dd5 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/0b4ec6da503185d389dcfeefc2c5d03ab13e0d76620b5fb6ce022f79bb64140f filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/556ec3c4c0e12c97f4c9b8dd0cd65ae321238c55438fbf498055f0ad7afc05bb filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/1f47921d53a3b8de72172e2fd5f0a0d27f1e8f4b77016c6b8a1a4b591bb83ae5 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/b9ddf00cad9d9a3bc91d09c8ee0f4fef42056861bb5ca1209fcc204b68950abf filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/2d16f4104fc6e34c556d42d6a35adedd5ffa396e71d59ef6dd7d5e88c01225e0 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/b7c2e1ec67bfd9d4578a81ca6987f095b74869db5f518bf2fc9d3ffeb30821e5 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/4ec2b03c373212fbe7fda97dbf9eaadb1bd42c2b0afc8b8a49806fe1dbc507f7 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/49353ead98e6a241dd015dda937bf8dcf335e4eed27fe732df8163dd955d042f filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/98c4455a98982b35380ec62f90d5cf0edff02eeb4b3bb18595d871fdeb0c0c67 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/3375e6602aa7264d6d67ae69ccb99690c8f40af7ab48d0af060dac1a2e9c24c9 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/e211ac567cceb6b5be62db5e135a1a5f3e77bcacebdac6a8cf9e0cdc24158c09 filter=lfs diff=lfs merge=lfs -text
+code/g2p/sources/chinese_lexicon.txt filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/672d584c094e51e29aa5e46a6dd22cbfb993f4f839609e954d2c299912ea5d77 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/428f99cfe94d9b1e42af2f3d8104598f8ca6e6105925c7687595a18c78d21011 filter=lfs diff=lfs merge=lfs -text

code/g2p/sources/chinese_lexicon.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3a7685d1c3e68eb2fa304bfc63e90c90c3c1a1948839a5b1b507b2131b3e2fb
+size 14779443

code/g2p/sources/g2p_chinese_model/poly_bert_model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8765d835ffdf9811c832d4dc7b6a552757aa8615c01d1184db716a50c20aebbc
+size 76583333

code/infer/example/edit_cn.lrc

@@ -0,0 +1,23 @@
+[00:13.55]亲爱的 人间爱情就像一朵花
+[00:20.29]想要种鲜花
+[00:22.51]手上难免沾泥巴
+[00:27.90]大红喜字铺开乡间田野
+[00:32.01]让我们住下
+[00:34.96]听微风吟唱吧
+[00:38.44]用光阴发新芽
+[00:41.04]裂缝钻出的新绿
+[00:43.97]正叩响整个盛夏
+[00:47.15]候鸟衔来远方的彩霞
+[00:52.03]荆棘里开出的花
+[00:55.67]每片花瓣都是回答
+[00:58.56]伤疤酿成琥珀时差
+[01:01.42]折射出星空的密码
+[01:04.36]蝴蝶吻过褪皮的痂
+[01:07.28]绽放成盔甲
+[01:10.11]牵你的手啊
+[01:12.90]在岁月种下鲜花
+[01:16.49]在寻常烟火人家
+[01:20.05]用漫长慢慢地回答
+[01:25.23]云会升起 雨会落下
+[01:28.86]不用催促啊
+[01:33.25]耐心的人啊才可以看见童话

code/infer/example/edit_cn.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fbc319c58215aff58b018a253abf6e6ec31673d6336947b932f1212964240d4
+size 8534237

code/infer/example/edit_en.lrc

@@ -0,0 +1,19 @@
+[00:12.56]后来
+[00:14.58]我总算学会了
+[00:17.81]如何去爱
+[00:19.37]可惜你 早已远去
+[00:22.21]消失在人海
+[00:25.42]Through tears I understood too late
+[00:32.11]Some souls slip away beyond fate's gate
+[00:39.86]White petals of gardenia
+[00:46.32]Drift onto my blue pleated skirt
+[00:51.95]I love you your whisper fell
+[00:58.96]I bowed my head
+[01:00.97]Breathing your lingering scent
+[01:05.40]That eternal summer night
+[01:09.35]Seventeen years young
+[01:12.20]When your lips found mine
+[01:18.23]Now through passing years
+[01:21.79]Each wistful sigh
+[01:24.91]Recalls starlight in your eyes
+[01:30.91]Why did young love's promise

code/infer/example/edit_en.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d72cb739c351b00a738f50e37a94cfc783dfa12b091c84ef9943b85a064015eb
+size 4771890

code/infer/example/eg_cn.lrc

@@ -0,0 +1,16 @@
+[00:18.23]让我掉下眼泪的
+[00:21.80]不止昨夜的酒
+[00:26.06]让我依依不舍的
+[00:29.91]不止你的温柔
+[00:33.79]余路还要走多久
+[00:37.92]你攥着我的手
+[00:41.81]让我感到为难的
+[00:45.78]是挣扎的自由
+[00:51.81]分别总是在九月
+[00:55.66]回忆是思念的愁
+[00:59.77]深秋嫩绿的垂柳
+[01:03.55]亲吻着我额头
+[01:07.55]在那座阴雨的小城里
+[01:11.47]我从未忘记你
+[01:15.37]成都 带不走的 只有你
+[01:23.51]和我在成都的街头走一走

code/infer/example/eg_cn.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67fc2163fa200544e5ecc1fcd7889124a60e6ff20fd70c53ad05a18859d79fd3
+size 882044

code/infer/example/eg_cn_full.lrc

@@ -0,0 +1,44 @@
+[00:18.23]让我掉下眼泪的
+[00:21.80]不止昨夜的酒
+[00:26.06]让我依依不舍的
+[00:29.91]不止你的温柔
+[00:33.79]余路还要走多久
+[00:37.92]你攥着我的手
+[00:41.81]让我感到为难的
+[00:45.78]是挣扎的自由
+[00:51.81]分别总是在九月
+[00:55.66]回忆是思念的愁
+[00:59.77]深秋嫩绿的垂柳
+[01:03.55]亲吻着我额头
+[01:07.55]在那座阴雨的小城里
+[01:11.47]我从未忘记你
+[01:15.37]成都 带不走的 只有你
+[01:23.51]和我在成都的街头走一走
+[01:31.38]直到所有的灯都熄灭了也不停留
+[01:39.23]你会挽着我的衣袖
+[01:43.12]我会把手揣进裤兜
+[01:47.13]走到玉林路的尽头
+[01:50.97]坐在小酒馆的门口
+[02:30.76]分别总是在九月
+[02:34.46]回忆是思念的愁
+[02:38.48]深秋嫩绿的垂柳
+[02:42.38]亲吻着我额头
+[02:46.40]在那座阴雨的小城里
+[02:50.35]我从未忘记你
+[02:54.21]成都 带不走的 只有你
+[03:02.18]和我在成都的街头走一走
+[03:10.02]直到所有的灯都熄灭了也不停留
+[03:18.16]你会挽着我的衣袖
+[03:21.90]我会把手揣进裤兜
+[03:25.80]走到玉林路的尽头
+[03:29.81]坐在小酒馆的门口
+[03:37.97]和我在成都的街头走一走
+[03:45.76]直到所有的灯都熄灭了也不停留
+[03:53.72]和我在成都的街头走一走
+[04:01.44]直到所有的灯都熄灭了也不停留
+[04:09.68]你会挽着我的衣袖
+[04:13.35]我会把手揣进裤兜
+[04:17.41]走到玉林路的尽头
+[04:21.27]走过小酒馆的门口
+[04:35.78]和我在成都的街头走一走
+[04:43.18]直到所有的灯都熄灭了也不停留

code/infer/example/eg_en.lrc

@@ -0,0 +1,26 @@
+[00:10.00]Moonlight spills through broken blinds
+[00:13.20]Your shadow dances on the dashboard shrine
+[00:16.85]Neon ghosts in gasoline rain
+[00:20.40]I hear your laughter down the midnight train
+[00:24.15]Static whispers through frayed wires
+[00:27.65]Guitar strings hum our cathedral choirs
+[00:31.30]Flicker screens show reruns of June
+[00:34.90]I'm drowning in this mercury lagoon
+[00:38.55]Electric veins pulse through concrete skies
+[00:42.10]Your name echoes in the hollow where my heartbeat lies
+[00:45.75]We're satellites trapped in parallel light
+[00:49.25]Burning through the atmosphere of endless night
+[01:00.00]Dusty vinyl spins reverse
+[01:03.45]Our polaroid timeline bleeds through the verse
+[01:07.10]Telescope aimed at dead stars
+[01:10.65]Still tracing constellations through prison bars
+[01:14.30]Electric veins pulse through concrete skies
+[01:17.85]Your name echoes in the hollow where my heartbeat lies
+[01:21.50]We're satellites trapped in parallel light
+[01:25.05]Burning through the atmosphere of endless night
+[02:10.00]Clockwork gears grind moonbeams to rust
+[02:13.50]Our fingerprint smudged by interstellar dust
+[02:17.15]Velvet thunder rolls through my veins
+[02:20.70]Chasing phantom trains through solar plane
+[02:24.35]Electric veins pulse through concrete skies
+[02:27.90]Your name echoes in the hollow where my heartbeat lies

code/infer/example/eg_en.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cdf3610c437715af336102a9bd34c63a3e882d9d99fe86b9fb193d8ebf9f2aa4
+size 10386551

code/infer/example/eg_en_full.lrc

@@ -0,0 +1,57 @@
+[00:00.52]Abracadabra abracadabra
+[00:03.97]Ha
+[00:04.66]Abracadabra abracadabra
+[00:12.02]Yeah
+[00:15.80]Pay the toll to the angels
+[00:19.08]Drawin' circles in the clouds
+[00:23.31]Keep your mind on the distance
+[00:26.67]When the devil turns around
+[00:30.95]Hold me in your heart tonight
+[00:34.11]In the magic of the dark moonlight
+[00:38.44]Save me from this empty fight
+[00:43.83]In the game of life
+[00:45.84]Like a poem said by a lady in red
+[00:49.45]You hear the last few words of your life
+[00:53.15]With a haunting dance now you're both in a trance
+[00:56.90]It's time to cast your spell on the night
+[01:01.40]Abracadabra ama-ooh-na-na
+[01:04.88]Abracadabra porta-ooh-ga-ga
+[01:08.92]Abracadabra abra-ooh-na-na
+[01:12.30]In her tongue she's sayin'
+[01:14.76]Death or love tonight
+[01:18.61]Abracadabra abracadabra
+[01:22.18]Abracadabra abracadabra
+[01:26.08]Feel the beat under your feet
+[01:27.82]The floor's on fire
+[01:29.90]Abracadabra abracadabra
+[01:33.78]Choose the road on the west side
+[01:37.09]As the dust flies watch it burn
+[01:41.45]Don't waste time on feeling
+[01:44.64]Your depression won't return
+[01:49.15]Hold me in your heart tonight
+[01:52.21]In the magic of the dark moonlight
+[01:56.54]Save me from this empty fight
+[02:01.77]In the game of life
+[02:03.94]Like a poem said by a lady in red
+[02:07.52]You hear the last few words of your life
+[02:11.19]With a haunting dance now you're both in a trance
+[02:14.95]It's time to cast your spell on the night
+[02:19.53]Abracadabra ama-ooh-na-na
+[02:22.71]Abracadabra porta-ooh-ga-ga
+[02:26.94]Abracadabra abra-ooh-na-na
+[02:30.42]In her tongue she's sayin'
+[02:32.83]Death or love tonight
+[02:36.55]Abracadabra abracadabra
+[02:40.27]Abracadabra abracadabra
+[02:44.19]Feel the beat under your feet
+[02:46.14]The floor's on fire
+[02:47.95]Abracadabra abracadabra
+[02:51.17]Phantom of the dance floor come to me
+[02:58.46]Sing for me a sinful melody
+[03:06.51]Ah-ah-ah-ah-ah ah-ah ah-ah
+[03:13.76]Ah-ah-ah-ah-ah ah-ah ah-ah
+[03:22.39]Abracadabra ama-ooh-na-na
+[03:25.66]Abracadabra porta-ooh-ga-ga
+[03:29.87]Abracadabra abra-ooh-na-na
+[03:33.16]In her tongue she's sayin'
+[03:35.55]Death or love tonight

code/infer/example/music_prompt_keywords.json

@@ -0,0 +1,67 @@
+{
+    "genre": [
+      "Pop",
+      "Rock",
+      "Electronic",
+      "Blues",
+      "R&B",
+      "Jazz",
+      "Soul",
+      "Country",
+      "Folk",
+      "Dance",
+      "Rap",
+      "Classical",
+      "Hip-Hop",
+      "Chinese",
+      "House",
+      "Punk",
+      "Metal",
+      "Funk",
+      "Reggae",
+      "Disco"
+    ],
+    "mood": [
+      "Romantic",
+      "Nostalgic",
+      "Heartfelt",
+      "Happy",
+      "Melancholic",
+      "Love",
+      "Upbeat",
+      "Energetic",
+      "Uplifting",
+      "Carefree",
+      "Fun",
+      "Emotional",
+      "Heartbroken",
+      "Reflective",
+      "Mellow",
+      "Hopeful",
+      "Relaxed",
+      "Passionate",
+      "Aggressive"
+    ],
+    "instrument": [
+      "Piano",
+      "Electric Guitar",
+      "Drums",
+      "Synthesizer",
+      "Bass",
+      "Strings",
+      "Bass Guitar",
+      "Acoustic Guitar",
+      "Guitar",
+      "Trumpet",
+      "Saxophone",
+      "Violin",
+      "Organ",
+      "Percussion",
+      "Vocal Sample",
+      "Flute",
+      "Harmonica",
+      "Banjo",
+      "Electric Bass",
+      "Cello"
+    ]
+  }

code/infer/example/output-windows/en-utf8-fixed.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10abbeefc22c1d03607cc074f825bbd4e0b7326e438fd13cb46985edf4a7ae6f
+size 1522668

code/infer/example/output-windows/en-without-utf8-fix.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:acf1fd072cf269745eaf2ee3c8674662a4c99441fc839c3bc0635f2dd89794ca
+size 1522668

code/infer/example/output/output.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6311c518821ee0d37d24c65d92c479f6c26da68b782687ab0ee69fb7a9ef97d3
+size 2832762

code/infer/example/vocal.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cb7d74eb7a8eda12acb8247b21d373928301db8a8cb0db480d341799fed3ce5
+size 2176

code/models/tt_dit/experimental/lora/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0

code/models/tt_dit/experimental/lora/adapter_loader.py

@@ -0,0 +1,392 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Adapter loader for the runtime-LoRA Wan 2.2 pipeline.
+
+Reads a LoRA safetensors file, parses A/B pairs using the same key
+normalization as ``pipeline_wan_lora.py``, and uploads them into the
+LoRA-aware Linear modules of a `WanTransformer3DModel` constructed with
+``lora_enabled=True``.
+
+Fused-QKV handling (the hard part):
+  Wan attention uses a fused ``to_qkv`` (self-attn) and ``to_kv``
+  (cross-attn) projection. The adapter ships separate ``to_q``,
+  ``to_k``, ``to_v`` LoRA pairs. We combine them into a single LoRA
+  pair on the fused module by:
+
+    - stacking A along the rank dim:  A_fused = [A_q; A_k; A_v]  → [3r, in]
+    - building a block-diagonal B with rank=3r and head-interleaving
+      it the same way the base weight is head-interleaved:
+        B_fused = head_interleave([B_q | 0 | 0,  0 | B_k | 0,  0 | 0 | B_v])
+
+  This triples the LoRA matmul cost on those layers (rank=3r) but keeps
+  the math exact and shippable in v0. A follow-up could split this into
+  3 independent adapters at the layer level for ~3x lower cost.
+
+`.diff` / `.diff_b` direct deltas are out of scope for v0 — the loader
+warns and skips them.
+"""
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from loguru import logger
+from safetensors.torch import load_file
+
+from models.tt_dit.experimental.utils.lightx2v_loader import wan_lightx2v_to_diffusers_key
+
+# Reused from pipeline_wan_lora.py — kept inline so this module is self-contained.
+_STRIP_PREFIXES = ("diffusion_model.", "transformer.", "unet.", "model.")
+_LOW_RANK_RE = re.compile(r"^(?P<base>.*)\.lora_(?P<slot>A|B|down|up)(?:\.[^.]+)?\.weight$")
+_SLOT_MAP = {"A": "A", "down": "A", "B": "B", "up": "B"}
+
+
+def _strip_known_prefixes(key: str) -> str:
+    for prefix in _STRIP_PREFIXES:
+        if key.startswith(prefix):
+            return key[len(prefix) :]
+    return key
+
+
+def _kohya_to_lightx2v(key: str) -> str:
+    if not key.startswith("lora_unet_"):
+        return key
+    parts = key.split(".", 1)
+    module_path = parts[0][len("lora_unet_") :]
+    suffix = f".{parts[1]}" if len(parts) > 1 else ""
+    m = re.match(r"blocks_(\d+)_(cross_attn|self_attn)_([a-z]+)", module_path)
+    if m:
+        return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
+    m = re.match(r"blocks_(\d+)_(ffn)_(\d+)", module_path)
+    if m:
+        return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
+    logger.warning(f"unrecognized kohya key structure: {key}")
+    return key
+
+
+def _normalize_key(raw: str) -> str:
+    return _kohya_to_lightx2v(_strip_known_prefixes(raw))
+
+
+# --------------------------------------------------------------------
+# Public API
+# --------------------------------------------------------------------
+@dataclass(frozen=True)
+class AdapterHandle:
+    name: str
+    rank: int  # the "expected bound rank" — may be 3r for fused-QKV targets
+    target_indices: dict[str, int]
+    """Map from a stable target identifier (e.g. ``'blocks.0.attn1.to_qkv'``)
+    to the bank index returned by `register_lora` on that module."""
+
+
+def load_adapter_into(
+    transformer,  # WanTransformer3DModel constructed with lora_enabled=True
+    path: str,
+    *,
+    scale: float = 1.0,
+    name: str = "",
+) -> AdapterHandle:
+    """Load a single LoRA safetensors file into `transformer`.
+
+    Returns an `AdapterHandle` recording the bank index assigned to each
+    LoRA-targeted Linear. Use it later with
+    ``pipeline.set_active_lora(handle)``.
+    """
+    raw = load_file(str(path))
+    if not any(_is_lora_key(k) for k in raw):
+        raise RuntimeError(f"no LoRA-style keys (lora_A/lora_B, lora_down/lora_up) in {path}")
+
+    pairs, alphas, skipped_direct = _collect_pairs(raw)
+    if skipped_direct:
+        logger.warning(
+            f"{path}: skipping {skipped_direct} direct (.diff/.diff_b) deltas "
+            "— runtime-LoRA pipeline does not support these yet."
+        )
+
+    # Group targets
+    fused_qkv: dict[tuple[int, str], dict[str, dict[str, torch.Tensor]]] = {}
+    singletons: list[tuple[str, dict[str, torch.Tensor]]] = []  # (diffusers_path, ab)
+    skipped_unmapped = 0
+
+    for base_path, ab in pairs.items():
+        diff_path = _diffusers_path(base_path)
+        if diff_path is None:
+            skipped_unmapped += 1
+            continue
+        m = re.match(r"blocks\.(\d+)\.(attn1|attn2)\.to_([qkv])$", diff_path)
+        if m:
+            block_idx, attn, qkv = int(m.group(1)), m.group(2), m.group(3)
+            # attn1 (self-attn): Q/K/V all fold into to_qkv.
+            # attn2 (cross-attn): Q is its own singleton; K/V fold into to_kv.
+            if attn == "attn1" or qkv in ("k", "v"):
+                fused_qkv.setdefault((block_idx, attn), {})[qkv] = ab
+                continue
+            # attn2.to_q is a real singleton — fall through.
+        # Everything else: 1:1 mapping
+        singletons.append((diff_path, ab))
+
+    if skipped_unmapped:
+        logger.warning(
+            f"{path}: {skipped_unmapped} pairs unmapped — adapter targets a module "
+            "that doesn't exist in the tt-dit name map."
+        )
+
+    target_indices: dict[str, int] = {}
+    fused_ranks: list[int] = []
+    singleton_ranks: list[int] = []
+
+    # Fused QKV / KV
+    for (block_idx, attn_name), qkvs in fused_qkv.items():
+        attn = getattr(transformer.blocks[block_idx], attn_name)
+        bank_idx, fused_rank = _register_fused_qkv(attn, block_idx, qkvs, alphas, scale, name)
+        if bank_idx is None:
+            continue
+        key = f"blocks.{block_idx}.{attn_name}.{'to_qkv' if attn.is_self else 'to_kv'}"
+        target_indices[key] = bank_idx
+        fused_ranks.append(fused_rank)
+
+    # Singletons
+    for diff_path, ab in singletons:
+        target, canonical = _resolve_singleton(transformer, diff_path)
+        if target is None:
+            logger.warning(f"unmapped singleton LoRA target: {diff_path}")
+            continue
+        A = ab["A"]
+        B = ab["B"]
+        rank = A.shape[0]
+        alpha = alphas.get(_lightx2v_for_path(diff_path), float(rank))
+        eff_scale = scale * (alpha / rank)
+        idx = target.register_lora(A, B, scale=eff_scale, name=name)
+        target_indices[canonical] = idx
+        singleton_ranks.append(rank)
+
+    # Sanity: every LoRA module in this adapter should have the same
+    # bound rank, but fused_QKV produces a 3x larger rank than singletons.
+    # The bound-tensor system handles per-module ranks independently, so
+    # we just record both for reporting.
+    all_ranks = fused_ranks + singleton_ranks
+    if not all_ranks:
+        raise RuntimeError(f"{path}: no LoRA targets registered")
+    representative_rank = max(all_ranks)
+    return AdapterHandle(name=name or Path(path).stem, rank=representative_rank, target_indices=target_indices)
+
+
+# --------------------------------------------------------------------
+# Helpers
+# --------------------------------------------------------------------
+def _is_lora_key(raw: str) -> bool:
+    key = _normalize_key(raw)
+    return bool(_LOW_RANK_RE.match(key)) or key.endswith((".diff", ".diff_b"))
+
+
+def _collect_pairs(raw: dict[str, torch.Tensor]):
+    """Returns ({base_path: {'A': T, 'B': T}}, {base_path: alpha}, num_skipped_direct)."""
+    pairs: dict[str, dict[str, torch.Tensor]] = {}
+    alphas: dict[str, float] = {}
+    skipped_direct = 0
+    unrecognized: list[tuple[str, str]] = []
+
+    for raw_key, tensor in raw.items():
+        key = _normalize_key(raw_key)
+        m = _LOW_RANK_RE.match(key)
+        if m:
+            pairs.setdefault(m.group("base"), {})[_SLOT_MAP[m.group("slot")]] = tensor
+        elif key.endswith(".alpha"):
+            alphas[key[: -len(".alpha")]] = float(tensor.item())
+        elif key.endswith(".diff") or key.endswith(".diff_b"):
+            skipped_direct += 1
+        else:
+            unrecognized.append((raw_key, key))
+
+    # Surface dropped keys — a non-standard naming convention silently
+    # producing plausible-but-wrong output is the failure mode we cannot
+    # afford here.
+    if unrecognized:
+        sample = ", ".join(f"{raw}→{norm}" for raw, norm in unrecognized[:5])
+        more = "" if len(unrecognized) <= 5 else f" (+{len(unrecognized) - 5} more)"
+        logger.warning(
+            f"adapter loader: {len(unrecognized)} key(s) did not match any known "
+            f"pattern (lora_A/B, lora_down/up, .alpha, .diff, .diff_b) and were "
+            f"dropped. Samples: {sample}{more}"
+        )
+
+    return pairs, alphas, skipped_direct
+
+
+def _diffusers_path(lightx2v_base: str) -> str | None:
+    """Map a lightx2v base path to its diffusers-canonical equivalent
+    (without ``.weight``). Returns None if mapping fails."""
+    try:
+        key = wan_lightx2v_to_diffusers_key(f"{lightx2v_base}.weight")
+    except Exception:  # noqa: BLE001
+        return None
+    if not key.endswith(".weight"):
+        return None
+    return key[: -len(".weight")]
+
+
+def _lightx2v_for_path(diffusers_path: str) -> str:
+    """Reverse-map (best-effort) a diffusers path back to lightx2v for alpha
+    lookups — alphas are keyed in the source naming."""
+    # Simple cases that the existing map covers:
+    rev = {
+        "attn1.to_q": "self_attn.q",
+        "attn1.to_k": "self_attn.k",
+        "attn1.to_v": "self_attn.v",
+        "attn1.to_out.0": "self_attn.o",
+        "attn2.to_q": "cross_attn.q",
+        "attn2.to_k": "cross_attn.k",
+        "attn2.to_v": "cross_attn.v",
+        "attn2.to_out.0": "cross_attn.o",
+        "ffn.net.0.proj": "ffn.0",
+        "ffn.net.2": "ffn.2",
+    }
+    for diff_seg, lx_seg in rev.items():
+        if diff_seg in diffusers_path:
+            return diffusers_path.replace(diff_seg, lx_seg)
+    return diffusers_path
+
+
+def _resolve_singleton(transformer, diff_path: str):
+    """Resolve a diffusers-style path to (module, tt_dit_dotted_path).
+
+    Returns (None, None) when no match. The tt-dit path is the one
+    consumed by the pipeline mixin's bind/unbind walker — it skips
+    diffusers-only segments like the trailing ``.0`` in ``to_out.0``.
+    """
+    m = re.match(r"blocks\.(\d+)\.(attn1|attn2)\.to_q$", diff_path)
+    if m:
+        block_idx, attn_name = int(m.group(1)), m.group(2)
+        attn = getattr(transformer.blocks[block_idx], attn_name)
+        if not attn.is_self:
+            return attn.to_q, f"blocks.{block_idx}.{attn_name}.to_q"
+        return None, None  # self-attn to_q shouldn't arrive here
+    m = re.match(r"blocks\.(\d+)\.(attn1|attn2)\.to_out\.0$", diff_path)
+    if m:
+        block_idx, attn_name = int(m.group(1)), m.group(2)
+        return getattr(transformer.blocks[block_idx], attn_name).to_out, (f"blocks.{block_idx}.{attn_name}.to_out")
+    m = re.match(r"blocks\.(\d+)\.ffn\.net\.0\.proj$", diff_path)
+    if m:
+        block_idx = int(m.group(1))
+        return transformer.blocks[block_idx].ffn.ff1, f"blocks.{block_idx}.ffn.ff1"
+    m = re.match(r"blocks\.(\d+)\.ffn\.net\.2$", diff_path)
+    if m:
+        block_idx = int(m.group(1))
+        return transformer.blocks[block_idx].ffn.ff2, f"blocks.{block_idx}.ffn.ff2"
+    return None, None
+
+
+def _register_fused_qkv(
+    attn,
+    block_idx: int,
+    qkvs: dict[str, dict[str, torch.Tensor]],
+    alphas: dict[str, float],
+    scale: float,
+    name: str,
+) -> tuple[int | None, int]:
+    """Stack Q/K/V (or K/V for cross-attn) LoRA pairs onto attn.to_qkv (or
+    attn.to_kv). Returns (bank_idx, fused_rank) or (None, 0) on skip."""
+    if attn.is_self:
+        required = ["q", "k", "v"]
+        target = attn.to_qkv
+    else:
+        required = ["k", "v"]
+        target = attn.to_kv
+
+    missing = [r for r in required if r not in qkvs]
+    if missing:
+        logger.warning(
+            f"fused QKV LoRA on {target.__class__.__name__} is missing {missing}; "
+            "skipping this group (need all of "
+            f"{required}, got {list(qkvs)})."
+        )
+        return None, 0
+
+    A_per = [qkvs[r]["A"] for r in required]
+    B_per = [qkvs[r]["B"] for r in required]
+    ranks = [A.shape[0] for A in A_per]
+    if len(set(ranks)) != 1:
+        raise ValueError(f"QKV LoRA ranks must match; got {ranks}")
+    r = ranks[0]
+    n = len(required)
+    in_dim = attn.dim
+    out_per = attn.dim  # each of Q/K/V has out=dim
+
+    # A stacked on rank dim → [n*r, in]
+    A_fused = torch.cat(A_per, dim=0)
+    assert A_fused.shape == (n * r, in_dim), f"A_fused {A_fused.shape}"
+
+    # B: each per-source B is [out_per=dim, r]. We embed each into a
+    # [out_per, n*r] tile (zeros for other-source rank columns). Then
+    # head-interleave the list of [out_per, n*r] tiles into a single
+    # [n*out_per, n*r] tensor — matching the head layout of the fused
+    # to_qkv / to_kv base weight.
+    n_dev = attn.parallel_config.tensor_parallel.factor
+    n_local_heads = attn.n_local_heads
+    head_dim = attn.head_dim
+
+    B_per_padded: list[torch.Tensor] = []
+    for i, B in enumerate(B_per):
+        pad = torch.zeros(out_per, n * r, dtype=B.dtype)
+        pad[:, i * r : (i + 1) * r] = B
+        B_per_padded.append(pad)
+
+    B_fused = _head_interleave_lora_B(B_per_padded, n_dev=n_dev, n_local_heads=n_local_heads, head_dim=head_dim)
+    assert B_fused.shape == (n * out_per, n * r), f"B_fused {B_fused.shape}"
+
+    # Effective scale: pick the first source's alpha (alphas are usually
+    # identical across Q/K/V in a single adapter; if not, the choice is
+    # safe-ish because we apply scale to all three the same way).
+    alpha_keys = [_lightx2v_qkv_key(attn, block_idx, r_name) for r_name in required]
+    alpha = next((alphas[k] for k in alpha_keys if k in alphas), float(r))
+    eff_scale = scale * (alpha / r)
+
+    bank_idx = target.register_lora(A_fused, B_fused, scale=eff_scale, name=name)
+    return bank_idx, n * r
+
+
+def _lightx2v_qkv_key(attn, block_idx: int, qkv: str) -> str:
+    """Build the lightx2v alpha key for a Q/K/V projection inside a
+    given WanAttention. Alphas in `_collect_pairs` are keyed by the full
+    normalized base path (``blocks.<i>.<side>.<qkv>``), so the block
+    index must be threaded in alongside ``attn.is_self``."""
+    side = "self_attn" if attn.is_self else "cross_attn"
+    return f"blocks.{block_idx}.{side}.{qkv}"
+
+
+def _head_interleave_lora_B(
+    tensors: Sequence[torch.Tensor],
+    *,
+    n_dev: int,
+    n_local_heads: int,
+    head_dim: int,
+) -> torch.Tensor:
+    """Mirror of WanAttention._interleave_heads for LoRA B tensors.
+
+    Each input is [out_per_source=num_heads*head_dim, rank_fused].
+    Returns a concatenated [N*out_per, rank_fused] tensor with the out
+    dim head-interleaved so column-parallel sharding lines up with the
+    fused base weight's layout.
+    """
+    out_per = n_dev * n_local_heads * head_dim
+    fused_rank = tensors[0].shape[1]
+    assert all(t.shape == (out_per, fused_rank) for t in tensors), (
+        f"all B tensors must be [{out_per}, {fused_rank}]; got " f"{[tuple(t.shape) for t in tensors]}"
+    )
+    n = len(tensors)
+    # Transpose each to [fused_rank, out_per]
+    transposed = [t.T for t in tensors]
+    # Reshape out dim to [fused_rank, n_dev, n_local_heads, head_dim]
+    reshaped = [t.reshape(fused_rank, n_dev, n_local_heads, head_dim) for t in transposed]
+    # Concat on heads sub-dim → [fused_rank, n_dev, N*n_local_heads, head_dim]
+    merged = torch.cat(reshaped, dim=2)
+    # Flatten back → [fused_rank, N*n_dev*n_local_heads*head_dim] = [fused_rank, N*out_per]
+    merged = merged.reshape(fused_rank, n * out_per)
+    # Transpose to [N*out_per, fused_rank]
+    return merged.T

code/models/tt_dit/experimental/lora/ltx_adapter_loader.py

@@ -0,0 +1,363 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Adapter loader for on-device fuse-mode LoRA on the LTX-2.3 transformer.
+
+Reads a LoRA safetensors file and registers its A/B pairs into the
+LoRA-aware Linear modules of an ``LTXTransformerModel`` built with
+``lora_enabled=True``. Swap is then a ``bind_active`` on each module
+(``weight.data += scale * A@B`` on device) — no host fuse, no reload.
+
+Why this differs from the Wan loader
+-------------------------------------
+LTX's ``LTXAttention._prepare_torch_state`` transforms the raw checkpoint
+Q/K/V weights before they reach the device:
+
+  - ``_permute_qk``: reorders each head's channels from the checkpoint's
+    SPLIT rotary layout to the INTERLEAVED layout the RoPE op expects.
+    Applied to Q and K (self-attn) and to the cross-attn ``to_q``.
+  - ``_interleave_heads``: folds separate Q/K/V (or K/V) into the fused
+    ``to_qkv`` / ``to_kv`` weight with heads interleaved so column-parallel
+    sharding lands the right head on each device.
+
+The host-fuse path (``utils.fuse_loras.fuse_loras_into``) is correct
+without any of this because it merges the delta into the *raw* checkpoint
+weight and lets ``_prepare_torch_state`` permute the merged result. Here
+the on-device weight is ALREADY permuted+interleaved, so the LoRA B must
+be pre-transformed identically or the merged delta lands in the wrong
+layout. ``to_out``/``ff1``/``ff2`` are untouched by ``_prepare_torch_state``,
+so their adapters need no transform (the mixin copies W's own sharding).
+
+Scope: every LoRA-adaptable Linear in the transformer — video/audio/cross-modal
+attention (``to_qkv``/``to_q``/``to_kv``/``to_out``/``to_gate_logits``), both FFNs
+(ff1/ff2), and the globals (``patchify_proj``, ``proj_out``, and the
+``adaln_single`` / timestep-embedder linears). ``promote_to_lora`` makes the plain
+globals bindable; attention/FFN are already LoRA via ``lora_enabled``. Only
+attention q/k/v/q-of-cross carry the rotary permute + head-interleave transform;
+everything else registers directly. ``.diff``/``.diff_b`` direct deltas remain
+unsupported; unmapped keys are surfaced loudly, never silently dropped.
+"""
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from loguru import logger
+from safetensors.torch import load_file
+
+from ...layers.lora import LoRAMixin
+from .promote import promote_to_lora
+
+_STRIP_PREFIXES = ("model.diffusion_model.", "diffusion_model.", "transformer.", "model.")
+_LOW_RANK_RE = re.compile(r"^(?P<base>.*)\.lora_(?P<slot>A|B|down|up)(?:\.[^.]+)?\.weight$")
+_SLOT_MAP = {"A": "A", "down": "A", "B": "B", "up": "B"}
+
+# Diffusers LTX checkpoints name blocks "transformer_blocks.<i>"; some exports
+# use the shorter "blocks.<i>". Accept both.
+_BLOCK_RE = r"(?:transformer_blocks|blocks)\.(\d+)"
+
+# All LTXAttention instances on a block share the same to_q/k/v/out/gate_logits
+# structure and fusion rules (self -> to_qkv, cross -> to_q + to_kv), so the loader
+# handles them uniformly, reading is_self/heads off each resolved module.
+_ATTN_NAMES = ("attn1", "attn2", "audio_attn1", "audio_attn2", "audio_to_video_attn", "video_to_audio_attn")
+# Diffusers FFN container name -> tt-dit attribute (ff.net.0.proj/net.2 -> ffX).
+_FF_NAME_TO_ATTR = {"ff": "ffn", "audio_ff": "audio_ff"}
+
+
+@dataclass(frozen=True)
+class LTXAdapterHandle:
+    name: str
+    target_indices: dict[str, int]
+    """Map from a canonical tt-dit module path (e.g. ``transformer_blocks.0.attn1.to_qkv``)
+    to the bank index returned by ``register_lora`` on that module."""
+
+
+def load_ltx_adapter_into(
+    transformer,  # LTXTransformerModel built with lora_enabled=True
+    path: str,
+    *,
+    scale: float = 1.0,
+    name: str = "",
+) -> LTXAdapterHandle:
+    """Register a single LoRA safetensors file into ``transformer``.
+
+    Returns a handle recording the bank index assigned per module; pass it
+    to the pipeline's ``set_active_lora`` to bind it on device.
+    """
+    raw = load_file(str(path))
+    pairs, alphas, skipped_direct, unrecognized = _collect_pairs(raw)
+    if not pairs:
+        raise RuntimeError(f"no LoRA A/B pairs (lora_A/B or lora_down/up) found in {path}")
+    if skipped_direct:
+        logger.warning(f"{path}: skipping {skipped_direct} direct (.diff/.diff_b) deltas — not supported.")
+    if unrecognized:
+        sample = ", ".join(unrecognized[:5])
+        more = "" if len(unrecognized) <= 5 else f" (+{len(unrecognized) - 5} more)"
+        logger.warning(
+            f"{path}: {len(unrecognized)} key(s) matched no known LoRA pattern and were dropped: {sample}{more}"
+        )
+
+    # Promote every plain Linear (adaln/proj/patchify/gate globals) to a LoRA
+    # variant so it can be bound; attn/ffn are already LoRA via lora_enabled.
+    promote_to_lora(transformer)
+
+    # Q/K/V fold into a fused projection; the rest are 1:1 block singletons or
+    # direct globals (patchify/proj_out/adaln/timestep-embedder).
+    fused: dict[tuple[int, str], dict[str, dict[str, torch.Tensor]]] = {}
+    singletons: list[tuple[int, str, str, dict[str, torch.Tensor]]] = []  # (block, container, sub, ab)
+    globals_: list[tuple[str, dict[str, torch.Tensor]]] = []
+
+    for base, ab in pairs.items():
+        parsed = _parse_target(base)
+        if parsed is None:
+            globals_.append((base, ab))
+            continue
+        block_idx, container, sub = parsed
+        if sub in ("to_q", "to_k", "to_v"):
+            attn = _get_attn(transformer, block_idx, container)
+            is_self = attn.is_self if attn is not None else container in ("attn1", "audio_attn1")
+            qkv = sub[-1]
+            # self-attn: q/k/v all fold into to_qkv. cross-attn: k/v fold into to_kv, q is a singleton.
+            if is_self or qkv in ("k", "v"):
+                fused.setdefault((block_idx, container), {})[qkv] = ab
+                continue
+        singletons.append((block_idx, container, sub, ab))
+
+    target_indices: dict[str, int] = {}
+
+    for (block_idx, container), qkvs in fused.items():
+        attn = _get_attn(transformer, block_idx, container)
+        if attn is None:
+            logger.warning(f"{path}: no attention module {container} in block {block_idx} — skipping fused group.")
+            continue
+        idx = _register_fused(attn, qkvs, alphas, scale, name, block_idx, container)
+        if idx is None:
+            continue
+        canonical = f"transformer_blocks.{block_idx}.{container}.{'to_qkv' if attn.is_self else 'to_kv'}"
+        target_indices[canonical] = idx
+
+    for block_idx, container, sub, ab in singletons:
+        if "A" not in ab or "B" not in ab:
+            logger.warning(f"{path}: {container}.{sub} in block {block_idx} missing one of A/B — skipping.")
+            continue
+        target, canonical, permute_dims = _resolve_singleton(transformer, block_idx, container, sub)
+        if target is None:
+            logger.warning(f"{path}: unmapped {container}.{sub} in block {block_idx} — skipping.")
+            continue
+        A = ab["A"]
+        B = ab["B"]
+        if permute_dims is not None:
+            B = _permute_qk_rows(B, *permute_dims)
+        rank = A.shape[0]
+        alpha = alphas.get(_alpha_key(block_idx, container, sub), float(rank))
+        idx = target.register_lora(A, B, scale=scale * (alpha / rank), name=name)
+        target_indices[canonical] = idx
+
+    unmapped_global = 0
+    for base, ab in globals_:
+        if "A" not in ab or "B" not in ab:
+            continue
+        target = _resolve_global(transformer, base)
+        if not isinstance(target, LoRAMixin):
+            unmapped_global += 1
+            continue
+        A = ab["A"]
+        B = ab["B"]
+        rank = A.shape[0]
+        alpha = alphas.get(base, float(rank))
+        idx = target.register_lora(A, B, scale=scale * (alpha / rank), name=name)
+        target_indices[base] = idx
+
+    if unmapped_global:
+        logger.warning(f"{path}: {unmapped_global} global key(s) matched no LTX module — dropped.")
+
+    if not target_indices:
+        raise RuntimeError(f"{path}: no LoRA targets registered — check the adapter's key naming against LTX modules.")
+
+    return LTXAdapterHandle(name=name or Path(path).stem, target_indices=target_indices)
+
+
+def iter_lora_modules(root, prefix: str = ""):
+    """Yield (dotted_path, module) for every LoRAMixin descendant of ``root``.
+    Paths match the loader's canonical keys (e.g. ``transformer_blocks.0.attn1.to_qkv``)."""
+    for name, child in root.named_children():
+        path = f"{prefix}.{name}" if prefix else name
+        if isinstance(child, LoRAMixin):
+            yield path, child
+        yield from iter_lora_modules(child, path)
+
+
+# --------------------------------------------------------------------
+# Key parsing
+# --------------------------------------------------------------------
+def _strip_prefixes(key: str) -> str:
+    for p in _STRIP_PREFIXES:
+        if key.startswith(p):
+            return key[len(p) :]
+    return key
+
+
+def _collect_pairs(raw: dict[str, torch.Tensor]):
+    """Return ({base: {'A','B'}}, {base: alpha}, num_skipped_direct, unrecognized_keys)."""
+    pairs: dict[str, dict[str, torch.Tensor]] = {}
+    alphas: dict[str, float] = {}
+    skipped_direct = 0
+    unrecognized: list[str] = []
+    for raw_key, tensor in raw.items():
+        key = _strip_prefixes(raw_key)
+        m = _LOW_RANK_RE.match(key)
+        if m:
+            pairs.setdefault(m.group("base"), {})[_SLOT_MAP[m.group("slot")]] = tensor
+        elif key.endswith(".alpha"):
+            alphas[key[: -len(".alpha")]] = float(tensor.item())
+        elif key.endswith((".diff", ".diff_b")):
+            skipped_direct += 1
+        else:
+            unrecognized.append(raw_key)
+    return pairs, alphas, skipped_direct, unrecognized
+
+
+def _parse_target(base: str) -> tuple[int, str, str] | None:
+    """Parse a stripped per-block LoRA base path into (block_idx, container, sub).
+
+    Covers every block attention family (video/audio/cross-modal) and both FFN
+    containers. Returns None for non-block paths (globals: patchify/proj_out/adaln),
+    which ``_resolve_global`` handles by direct path lookup."""
+    attn_alt = "|".join(_ATTN_NAMES)
+    m = re.match(rf"^{_BLOCK_RE}\.({attn_alt})\.to_(q|k|v|out|gate_logits)(?:\.0)?$", base)
+    if m:
+        raw = m.group(3)
+        sub = {"out": "to_out", "gate_logits": "to_gate_logits"}.get(raw, f"to_{raw}")
+        return int(m.group(1)), m.group(2), sub
+    ff_alt = "|".join(_FF_NAME_TO_ATTR)
+    m = re.match(rf"^{_BLOCK_RE}\.({ff_alt})\.net\.0\.proj$", base)
+    if m:
+        return int(m.group(1)), m.group(2), "ff1"
+    m = re.match(rf"^{_BLOCK_RE}\.({ff_alt})\.net\.2$", base)
+    if m:
+        return int(m.group(1)), m.group(2), "ff2"
+    return None
+
+
+def _alpha_key(block_idx: int, attn_name: str, sub: str) -> str:
+    return f"transformer_blocks.{block_idx}.{attn_name}.{sub}"
+
+
+# --------------------------------------------------------------------
+# Module resolution
+# --------------------------------------------------------------------
+def _get_attn(transformer, block_idx: int, attn_name: str):
+    blocks = transformer.transformer_blocks
+    if not (0 <= block_idx < len(blocks)):
+        return None
+    return getattr(blocks[block_idx], attn_name, None)
+
+
+def _resolve_singleton(transformer, block_idx: int, container: str, sub: str):
+    """Return (module, canonical_path, permute_dims) where permute_dims is
+    (num_heads, head_dim) if the base weight is _permute_qk'd, else None.
+    (None, None, None) when the target does not exist."""
+    blocks = transformer.transformer_blocks
+    if not (0 <= block_idx < len(blocks)):
+        return None, None, None
+    block = blocks[block_idx]
+    if sub in ("to_q", "to_out", "to_gate_logits"):
+        attn = getattr(block, container, None)
+        if attn is None:
+            return None, None, None
+        target = getattr(attn, sub, None)
+        if target is None:
+            return None, None, None
+        # Only cross-attn to_q is _permute_qk'd; to_out/to_gate_logits are not.
+        permute_dims = (attn.num_heads, attn.head_dim) if (sub == "to_q" and not attn.is_self) else None
+        return target, f"transformer_blocks.{block_idx}.{container}.{sub}", permute_dims
+    if sub in ("ff1", "ff2"):
+        ff_attr = _FF_NAME_TO_ATTR.get(container, "ffn")
+        ff = getattr(block, ff_attr, None)
+        if ff is None:
+            return None, None, None
+        target = getattr(ff, sub, None)
+        return target, f"transformer_blocks.{block_idx}.{ff_attr}.{sub}", None
+    return None, None, None
+
+
+def _resolve_global(transformer, base: str):
+    """Resolve a non-block LoRA key (patchify_proj, proj_out, adaln_single.*, ...)
+    to the module at that dotted path on the transformer root, or None if absent.
+    These are plain replicated/parallel Linears with no rotary layout transform."""
+    mod = transformer
+    for part in base.split("."):
+        mod = getattr(mod, part, None)
+        if mod is None:
+            return None
+    return mod
+
+
+# --------------------------------------------------------------------
+# Fused QKV / KV registration (permute_qk + head-interleave on B)
+# --------------------------------------------------------------------
+def _register_fused(attn, qkvs, alphas, scale, name, block_idx, attn_name) -> int | None:
+    required = ["q", "k", "v"] if attn.is_self else ["k", "v"]
+    missing = [r for r in required if r not in qkvs or "A" not in qkvs[r] or "B" not in qkvs[r]]
+    if missing:
+        logger.warning(f"fused LoRA on block {block_idx}.{attn_name} missing {missing}; skipping (got {list(qkvs)}).")
+        return None
+
+    ranks = [qkvs[r]["A"].shape[0] for r in required]
+    if len(set(ranks)) != 1:
+        raise ValueError(f"QKV LoRA ranks must match on block {block_idx}.{attn_name}; got {ranks}")
+    r = ranks[0]
+    n = len(required)
+
+    num_heads = attn.num_heads
+    head_dim = attn.head_dim
+    n_dev = attn.parallel_config.tensor_parallel.factor
+    n_local_heads = attn.n_local_heads
+    out_per = attn.dim  # each fused source projects to dim
+
+    # A stacked on rank -> [n*r, in]
+    A_fused = torch.cat([qkvs[x]["A"] for x in required], dim=0)
+
+    # B: permute Q/K rows (RoPE), pad each into a block-diagonal [out, n*r] tile,
+    # then head-interleave the tiles into [n*out, n*r] to match the fused weight.
+    B_padded = []
+    for i, x in enumerate(required):
+        B = qkvs[x]["B"]
+        if x in ("q", "k"):
+            B = _permute_qk_rows(B, num_heads, head_dim)
+        pad = torch.zeros(out_per, n * r, dtype=B.dtype)
+        pad[:, i * r : (i + 1) * r] = B
+        B_padded.append(pad)
+    B_fused = _interleave_heads_rows(B_padded, num_heads, head_dim, n_dev, n_local_heads)
+
+    alpha = alphas.get(_alpha_key(block_idx, attn_name, f"to_{required[0]}"), float(r))
+    eff_scale = scale * (alpha / r)
+    # register_lora lives on the fused projection Linear (LoRAMixin), not the attention wrapper.
+    fused_linear = attn.to_qkv if attn.is_self else attn.to_kv
+    return fused_linear.register_lora(A_fused, B_fused, scale=eff_scale, name=name)
+
+
+def _permute_qk_rows(t: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor:
+    """Reorder each head's output channels SPLIT -> INTERLEAVED, mirroring
+    ``LTXAttention._prepare_torch_state._permute_qk``. ``t`` is [num_heads*head_dim, ...]."""
+    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)
+
+
+def _interleave_heads_rows(tensors, num_heads: int, head_dim: int, n_dev: int, n_local_heads: int) -> torch.Tensor:
+    """Interleave the OUT (row) dim across sources, mirroring
+    ``LTXAttention._prepare_torch_state._interleave_heads``. Each input is
+    [out=num_heads*head_dim, C]; returns [n*out, C]."""
+    transposed = [t.T for t in tensors]  # [C, out]
+    reshaped = [t.reshape(t.shape[0], n_dev, n_local_heads, head_dim) for t in transposed]
+    merged = torch.cat(reshaped, dim=2)  # [C, n_dev, n*n_local_heads, head_dim]
+    merged = merged.reshape(merged.shape[0], len(tensors) * num_heads * head_dim)  # [C, n*out]
+    return merged.T  # [n*out, C]

code/models/tt_dit/experimental/lora/promote.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Generic in-place promotion of Linear-family modules to their LoRA variants.
+
+Rather than thread a ``lora_enabled`` flag through every module constructor,
+walk a built model once and upgrade each plain ``Linear``/``ColParallelLinear``/
+``RowParallelLinear`` to the matching ``LoRAMixin`` subclass. The subclasses are
+thin (mixin + base + ``_init_lora_state``), so swapping ``__class__`` and calling
+``_init_lora_state`` post-construction is equivalent to having built the LoRA
+variant directly. Modules already LoRA-aware are left untouched.
+
+This keeps the swap surface identical — ``bind_active`` / ``set_active_lora`` /
+``reapply_after_load`` act on the promoted modules exactly as before — while
+making every Linear in the model a LoRA target, not just the attn/ffn subset.
+"""
+from __future__ import annotations
+
+from ...layers.linear import ColParallelLinear, Linear, RowParallelLinear
+from ...layers.lora import LoRAMixin
+
+# Base types eligible for in-place promotion. Exact-type match (not isinstance)
+# so an unknown Linear subclass isn't silently mis-promoted.
+_PROMOTABLE = (Linear, ColParallelLinear, RowParallelLinear)
+
+# Per-base promotion classes, built lazily. Base FIRST, mixin SECOND: Python
+# forbids ``__class__`` assignment when the mixin comes first (the instance
+# "solid base" changes -> 'object layout differs'). With the base first the
+# layout is identical to the original, so the swap is allowed. The pre-defined
+# LoRA*Linear classes are mixin-first (so their forward override wins in runtime
+# mode) and can't be used here; fuse mode never needs that override, and every
+# bind method is LoRAMixin-exclusive, so ordering is otherwise irrelevant.
+_PROMOTED_CACHE: dict[type, type] = {}
+
+
+def _promoted_class(base_cls: type) -> type:
+    cls = _PROMOTED_CACHE.get(base_cls)
+    if cls is None:
+        cls = type(f"LoRA{base_cls.__name__}", (base_cls, LoRAMixin), {})
+        _PROMOTED_CACHE[base_cls] = cls
+    return cls
+
+
+def _iter_modules(root):
+    yield root
+    for _, child in root.named_children():
+        yield from _iter_modules(child)
+
+
+def promote_to_lora(root, *, mode: str = "fuse") -> int:
+    """Upgrade every plain Linear-family descendant of ``root`` to a LoRA-aware
+    class in place. Returns the number promoted. Idempotent: modules already
+    ``LoRAMixin`` (e.g. built via ``lora_enabled``) are skipped."""
+    promoted = 0
+    for module in _iter_modules(root):
+        if isinstance(module, LoRAMixin):
+            continue
+        if type(module) not in _PROMOTABLE:
+            continue
+        module.__class__ = _promoted_class(type(module))
+        module._init_lora_state(mode=mode)
+        promoted += 1
+    return promoted

code/models/tt_dit/experimental/models/AniSora.md

@@ -0,0 +1,136 @@
+# Index-AniSora V3.2 (I2V)
+
+## Introduction
+
+[IndexTeam/Index-anisora](https://huggingface.co/IndexTeam/Index-anisora) V3.2
+is an anime-domain image-to-video model derived from
+[Wan2.2-I2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers).
+It preserves the two-expert (high-noise / low-noise) MoE structure of Wan2.2
+and ships a finetuned weight pair specialized for anime-style motion.
+
+This is a **full checkpoint replacement** (not a LoRA adapter or distilled model).
+The model graph is identical to base Wan2.2 I2V — see [Wan2_2.md](Wan2_2.md) for
+the underlying architecture.
+
+## Details
+
+The two-stage MoE schedule is preserved: the high-noise expert handles the
+first 90% of the denoising trajectory and the low-noise expert handles the
+final 10% (`boundary_ratio=0.9`, vs 0.5 for base Wan2.2).
+
+The AniSora safetensors use the original-Wan key naming
+(`blocks.X.self_attn.q.weight`, etc), identical to the lightx2v distill
+checkpoints. The same `wan_lightx2v_to_diffusers_key` rename function in
+`models/tt_dit/utils/lightx2v_loader.py` is reused.
+
+`NUM_STEPS` is configurable (default 40). Lower step counts (16, 8) can produce video at decent quality.
+
+## Weights
+
+### AniSora V3.2 experts (DiT only)
+
+- **Source:** [IndexTeam/Index-anisora](https://huggingface.co/IndexTeam/Index-anisora)
+- **Files:**
+  - `V3.2/high_noise_model/diffusion_pytorch_model.safetensors` (~28 GB, BF16)
+  - `V3.2/low_noise_model/diffusion_pytorch_model.safetensors` (~28 GB, BF16)
+- **Total download:** ~57 GB (auto-downloaded on first run with `TT_DIT_ALLOW_HF_DOWNLOAD=1`)
+
+### Base Wan2.2 components (shared, not re-downloaded if already cached)
+
+- **Source:** [Wan-AI/Wan2.2-I2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers)
+- **Components:** tokenizer, UMT5 text encoder, VAE, scheduler config
+- **Total download:** ~12 GB (shared with base Wan2.2)
+
+### Download weights manually (optional)
+
+Weights are auto-downloaded via HuggingFace Hub on first run. To pre-download or use a local copy:
+
+```bash
+# Option 1: Pre-download to HF cache
+python -c "
+from huggingface_hub import hf_hub_download
+for f in ['V3.2/high_noise_model/diffusion_pytorch_model.safetensors',
+          'V3.2/low_noise_model/diffusion_pytorch_model.safetensors']:
+    hf_hub_download('IndexTeam/Index-anisora', f)
+"
+
+# Option 2: Point to a local directory
+# Layout must be:
+#   $ANISORA_LOCAL_DIR/V3.2/high_noise_model/diffusion_pytorch_model.safetensors
+#   $ANISORA_LOCAL_DIR/V3.2/low_noise_model/diffusion_pytorch_model.safetensors
+export ANISORA_LOCAL_DIR=/path/to/index-anisora-weights
+```
+
+## Inference defaults
+
+Mirrors `anisoraV3.2/wan/configs/wan_i2v_A14B.py`:
+
+| Parameter | Value | Note |
+|---|---|---|
+| `num_inference_steps` | 40 | UniPC sampler (configurable via `NUM_STEPS` env var) |
+| `boundary_ratio` | 0.9 | High-noise covers first 90% of timesteps |
+| `guidance_scale` | 3.5 | High-noise stage CFG |
+| `guidance_scale_2` | 3.5 | Low-noise stage CFG |
+| `sample_shift` | 5.0 | Inherited from base Wan2.2 scheduler config |
+
+## Performance
+
+BH Galaxy 4x8 Ring (sp=8, tp=4), 81 frames @ 16fps (~5s video):
+
+| Stage | Base Wan2.2 I2V 480p | AniSora V3.2 I2V 480p | AniSora V3.2 I2V 720p |
+|---|---|---|---|
+| Text Encoding | 0.13s | 0.14s | 0.14s |
+| Image Encoding | 5.03s | 5.02s | 12.84s |
+| Denoising | 48.92s (40 steps, 1.22s/step) | 48.29s (40 steps, 1.21s/step) | 133.61s (40 steps, 3.34s/step) |
+| VAE Decoding | 0.43s | 0.43s | 0.70s |
+| **Total** | **54.52s** | **53.90s** | **147.32s** |
+
+### Reduced-steps perf (AniSora V3.2)
+
+| Steps | 480p (832x480) | 720p (1280x720) |
+|---|---|---|
+| 40 | 53.90s | 147.32s |
+| 16 | 25.0s | 67.2s |
+| 8 | 15.3s | 40.4s |
+
+## Supported configurations
+
+| System | Mesh | SP | TP | Topology | Test ID |
+|---|---|---|---|---|---|
+| BH Galaxy | 4x8 | 8 (axis 1) | 4 (axis 0) | Ring | `bh_4x8sp1tp0_ring` |
+
+## Prerequisites
+- Cloned [tt-metal repository](https://github.com/tenstorrent/tt-metal) for source code
+- Installed: [TT-Metalium / TT-NN](https://github.com/tenstorrent/tt-metal/blob/main/INSTALLING.md)
+
+## How to Run
+
+```bash
+# Environment setup
+export TT_DIT_CACHE_DIR=/your/cache/path
+export TT_METAL_HOME=$(pwd)
+export PYTHONPATH=$(pwd)
+export TT_DIT_ALLOW_HF_DOWNLOAD=1
+
+# Set input image and prompt
+export PROMPT_IMAGE=$TT_METAL_HOME/dog.jpg
+export PROMPT="An anime girl smiling, soft lighting, cinematic"
+export NUM_STEPS=40       # 40 default, or 16/8 for faster previews
+export NO_PROMPT=1
+
+# Run AniSora inference (480p)
+pytest models/tt_dit/experimental/tests/test_pipeline_anisora.py \
+  -v -k "bh_4x8sp1tp0_ring and resolution_480p and not random_weights" \
+  --timeout 1800 -s
+# Output: ./wan_anisora_i2v_832x480_0.mp4
+
+# Run AniSora inference (720p)
+pytest models/tt_dit/experimental/tests/test_pipeline_anisora.py \
+  -v -k "bh_4x8sp1tp0_ring and resolution_720p and not random_weights" \
+  --timeout 1800 -s
+# Output: ./wan_anisora_i2v_1280x720_0.mp4
+
+# Random-weights smoke test (no large download needed)
+pytest models/tt_dit/experimental/tests/test_pipeline_anisora.py::test_pipeline_inference_random_weights \
+  -v --timeout 1500 -s
+```

code/models/tt_dit/experimental/models/Wan2_2_Distill.md

@@ -0,0 +1,89 @@
+# Wan2.2-Distill (lightx2v, 4-step I2V)
+
+## Introduction
+
+[lightx2v/Wan2.2-Distill-Models](https://huggingface.co/lightx2v/Wan2.2-Distill-Models)
+is a 4-step distillation of [Wan2.2-I2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers).
+It produces image-to-video outputs in roughly 1/10 the denoising steps of the
+base model, with classifier-free guidance baked in.
+
+This page documents the tt_dit integration. See [Wan2_2.md](Wan2_2.md) for the
+underlying architecture — the distill is weight-only; the model graph is identical.
+
+## Details
+
+The distill ships only the DiT weights as flat `.safetensors` files. The
+tokenizer, UMT5 text encoder, VAE, and scheduler are loaded from the base
+`Wan-AI/Wan2.2-I2V-A14B-Diffusers` repo.
+
+The two-stage MoE structure is preserved: a high-noise expert handles the first
+half of the denoising trajectory and a low-noise expert handles the rest. With
+4 inference steps and `boundary_ratio=0.5`, that's 2 steps per expert.
+
+## Files used
+
+From [lightx2v/Wan2.2-Distill-Models](https://huggingface.co/lightx2v/Wan2.2-Distill-Models):
+
+- `wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors` (BF16, 28.6 GB) — high-noise expert
+- `wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors` (BF16, 28.6 GB) — low-noise expert
+
+From [Wan-AI/Wan2.2-I2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers):
+tokenizer, text encoder, VAE, scheduler config.
+
+## Inference defaults
+
+| Parameter | Value | Note |
+|---|---|---|
+| `num_inference_steps` | 4 | Distill target |
+| `boundary_ratio` | 0.5 | 2 high-noise + 2 low-noise |
+| `guidance_scale` | 1.0 | CFG baked in — disables negative-prompt path |
+| `guidance_scale_2` | 1.0 | Same as above for the low-noise stage |
+
+## Supported configurations (PR1)
+
+| System | Mesh | SP | TP | Topology | Test ID |
+|---|---|---|---|---|---|
+| BH Galaxy | 4x8 | 8 (axis 1) | 4 (axis 0) | Ring | `bh_4x8sp1tp0_ring` |
+
+Additional mesh shapes will be added in follow-up work.
+
+## How to run
+
+```bash
+# Required environment
+export TT_DIT_CACHE_DIR=/your/cache/path
+export TT_METAL_HOME=$(pwd)
+export PYTHONPATH=$(pwd)
+
+# Allow first-time download of the lightx2v safetensors (~57 GB total).
+# Skip if you've already downloaded them and put them in the HF cache or
+# pointed LIGHTX2V_LOCAL_DIR at a directory that contains both files.
+export TT_DIT_ALLOW_HF_DOWNLOAD=1
+
+# Place a seed image at ./prompt_image.png in the working directory.
+
+NO_PROMPT=1 pytest \
+  models/tt_dit/experimental/tests/test_pipeline_wan_distill_i2v.py \
+  -v -k "bh_4x8sp1tp0_ring and resolution_480p" --timeout 1500
+```
+
+To use weights from a custom location:
+
+```bash
+export LIGHTX2V_LOCAL_DIR=/path/to/lightx2v/weights
+```
+
+## Performance
+
+BH Galaxy 4×8 Ring, 81 frames, single timed iteration (2026-05-08):
+
+| Stage | Base 480p (40 steps, CFG) | Base 720p (40 steps, CFG) | Distill 480p (4 steps, no CFG) | Distill 720p (4 steps, no CFG) |
+|---|---:|---:|---:|---:|
+| Text encoding (UMT5) | 0.136 s | 0.133 s | 0.134 s | 0.138 s |
+| Image encoding (CLIP+VAE encode) | 5.020 s | 13.024 s | 4.949 s | 12.577 s |
+| Denoising | 48.806 s | 133.243 s | 2.436 s | 6.553 s |
+| VAE decoding | 0.426 s | 0.711 s | 0.424 s | 0.746 s |
+| **Total** | **54.4 s** | **147.1 s** | **7.96 s** | **20.0 s** |
+| Speedup vs base | — | — | 6.84× | 7.34× |
+
+Per-forward denoise cost is identical between base and distill at each resolution; distill's win is "10× fewer steps × 2× from no-CFG" with kernel performance unchanged.

code/models/tt_dit/experimental/models/Wan2_2_LoRA.md

@@ -0,0 +1,197 @@
+# Wan2.2 I2V — LoRA Adapter Pipeline
+
+## Introduction
+
+`WanPipelineI2VLora` fuses [LoRA](https://arxiv.org/abs/2106.09685) adapters
+into the base [Wan2.2-I2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers)
+weights on CPU before pushing to device. Inference uses vanilla I2V
+machinery with no LoRA-specific runtime cost.
+
+Each expert transformer (high-noise + low-noise) accepts an ordered LoRA
+stack — multiple adapters can be combined on the same expert. Later entries
+see earlier ones' deltas already folded into the base, which matches the
+training-time expectation for adapters that were fine-tuned on top of
+another adapter (e.g. an SVI LoRA on top of a LightX2V LoRA).
+
+See [../../models/Wan2_2.md](../../models/Wan2_2.md) for the base model architecture and
+[Wan2_2_SVI.md](Wan2_2_SVI.md) for long-video generation that builds on this pipeline.
+
+## Details
+
+- **CPU-side fusion:** `W' = W + sum_i scale_i * (B_i @ A_i)` for low-rank
+  adapter pairs, with per-module alpha scaling applied when an `.alpha` key
+  is present. The `fp32` promote → add → cast-back uses a fresh allocation
+  per pair to avoid mutating the base when a stack folds in multiple
+  adapters touching the same target.
+- **Single aggregate verification:** after the full stack is fused, an L2
+  norm check confirms the base weights actually changed (catches silent
+  load failures from key mismatches).
+- **Cache namespacing:** the TT cache is keyed by a SHA1 hash of the ordered
+  `(resolved_path, scale)` tuples per expert. Two stacks with the same files
+  in different order get different namespaces.
+
+### Supported adapter key formats
+
+`fuse_lora_state_dict` auto-detects these key shapes per adapter entry:
+
+1. **Low-rank pairs** at one of:
+   - `<base>.lora_A.weight` / `<base>.lora_B.weight` (lightx2v native)
+   - `<base>.lora_down.weight` / `<base>.lora_up.weight` (legacy / kohya)
+   - PEFT-style `<base>.lora_A.default.weight` (the trailing `.default`
+     adapter-name segment is tolerated; SVI uses this convention)
+2. **Bias deltas** at `<base>.diff_b` — added to the base `.bias`.
+3. **Full-parameter deltas** at `<base>.diff` — added to the base
+   `.weight`. Some LoRAs ship RMSNorm gammas via this channel.
+4. **Per-module alpha** at `<base>.alpha` — applied as
+   `effective_scale = scale * alpha / rank` for the low-rank pair under
+   the same `<base>`.
+
+### Supported key namespaces (auto-stripped or remapped)
+
+- **lightx2v native:** `blocks.<i>.attn.q.lora_A.weight` — passes through.
+- **diffusers prefixes:** `diffusion_model.`, `transformer.`, `unet.`,
+  `model.` — stripped automatically before matching.
+- **kohya / A1111:** `lora_unet_blocks_<i>_cross_attn_k.lora_down.weight`
+  — remapped to lightx2v-style.
+
+Anything outside the above is logged as a skipped-unknown warning rather
+than erroring, so legitimate-but-unsupported keys don't break the load.
+
+## API
+
+The adapter paths are constructor arguments, so build the config with
+`default_config` (which pins the I2V model type and checkpoint) and construct the
+pipeline directly.
+
+```python
+from models.tt_dit.experimental.pipelines.pipeline_wan_lora import LoRASpec, WanPipelineI2VLora
+
+config = WanPipelineI2VLora.default_config(mesh_device=mesh_device, height=480, width=832)
+
+# Single-LoRA, both experts
+pipe = WanPipelineI2VLora(
+    device=mesh_device,
+    config=config,
+    lora_high="/path/high.safetensors",
+    lora_low="/path/low.safetensors",
+)
+
+# Multi-LoRA stack (LightX2V applied first, then SVI on top, half strength)
+pipe = WanPipelineI2VLora(
+    device=mesh_device,
+    config=config,
+    lora_high=[
+        LoRASpec("/path/lightx2v_high.safetensors", scale=1.0),
+        LoRASpec("/path/svi_high.safetensors", scale=0.5),
+    ],
+    lora_low=[
+        LoRASpec("/path/lightx2v_low.safetensors", scale=1.0),
+        LoRASpec("/path/svi_low.safetensors", scale=1.0),
+    ],
+)
+```
+
+Either `lora_high` or `lora_low` may be `None` — that expert uses unmodified
+base weights. At least one must be set.
+
+## Weights
+
+### Earth zoom-out LoRA
+
+- **Source:** [wangkanai/wan22-fp16-i2v-loras](https://huggingface.co/wangkanai/wan22-fp16-i2v-loras)
+- **File:** `loras/wan/wan22-camera-earthzoomout.safetensors` (293 MB)
+- **Format:** Single-file, kohya/A1111-style keys
+
+### Camera rotation LoRA
+
+- **Source:** [wangkanai/wan22-fp16-i2v-loras](https://huggingface.co/wangkanai/wan22-fp16-i2v-loras)
+- **File:** `loras/wan/wan22-camera-rotation-rank16-v2.safetensors` (293 MB)
+- **Format:** Single-file, kohya/A1111-style keys
+
+### Base Wan2.2 components (shared)
+
+- **Source:** [Wan-AI/Wan2.2-I2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers)
+- **Components:** tokenizer, UMT5 text encoder, VAE, scheduler config, transformer base weights (~70 GB)
+
+### Download
+
+```bash
+python -c "
+from huggingface_hub import hf_hub_download
+import os
+os.makedirs('/path/to/lora-camera', exist_ok=True)
+for f in ['loras/wan/wan22-camera-earthzoomout.safetensors',
+          'loras/wan/wan22-camera-rotation-rank16-v2.safetensors']:
+    hf_hub_download('wangkanai/wan22-fp16-i2v-loras', f, local_dir='/path/to/lora-camera')
+"
+```
+
+## Supported configurations
+
+| System | Mesh | SP | TP | Topology | Test ID |
+|---|---|---|---|---|---|
+| BH Loud Box | 2x4 | 4 (axis 1) | 2 (axis 0) | Linear | `bh_2x4sp1tp0` |
+| BH Galaxy | 4x8 | 8 (axis 1) | 4 (axis 0) | Ring | `bh_4x8sp1tp0_ring` |
+
+## How to Run
+
+### Single-LoRA (40 steps)
+
+```bash
+export TT_DIT_CACHE_DIR=/your/cache/path
+export TT_METAL_HOME=$(pwd)
+export PYTHONPATH=$(pwd)
+export TT_DIT_ALLOW_HF_DOWNLOAD=1
+
+export LORA_HIGH_PATH=/path/to/lora-camera/wan22-camera-earthzoomout.safetensors
+unset LORA_LOW_PATH          # single-file: low-noise expert uses base weights
+export LORA_SCALE=1.0
+
+export PROMPT_IMAGE=$TT_METAL_HOME/dog.jpg
+export PROMPT="A golden retriever running on a sandy beach, waves in the background"
+export NUM_STEPS=40
+export GUIDANCE_SCALE=3.5
+export BOUNDARY_RATIO=0.875
+export NO_PROMPT=1
+
+pytest models/tt_dit/experimental/tests/test_pipeline_lora.py \
+  -v -k "bh_4x8sp1tp0_ring and resolution_480p" \
+  --timeout 1800 -s
+```
+
+### Multi-LoRA stack
+
+```bash
+# Comma-separated path[:scale] entries. Order matters -- earlier entries are
+# fused first.
+export LORA_STACK_HIGH="/path/lightx2v_high.safetensors:1.0,/path/style_high.safetensors:0.5"
+export LORA_STACK_LOW="/path/lightx2v_low.safetensors:1.0,/path/style_low.safetensors:1.0"
+
+pytest models/tt_dit/experimental/tests/test_pipeline_lora.py \
+  -v -k "bh_4x8sp1tp0_ring and resolution_480p" \
+  --timeout 1800 -s
+```
+
+## Environment variables
+
+| Variable | Required | Default | Description |
+|---|---|---|---|
+| `LORA_HIGH_PATH` | Either this or stack | — | Path to high-noise expert LoRA `.safetensors` |
+| `LORA_LOW_PATH` | No | None | Path to low-noise expert LoRA |
+| `LORA_SCALE` | No | 1.0 | Blend strength applied to single-LoRA entries |
+| `LORA_STACK_HIGH` | Either this or single | — | Comma-separated `path[:scale]` list for the high-noise expert |
+| `LORA_STACK_LOW` | No | None | Comma-separated `path[:scale]` list for the low-noise expert |
+| `NUM_STEPS` | No | 40 | Inference steps |
+| `GUIDANCE_SCALE` | No | 3.5 | CFG scale |
+| `BOUNDARY_RATIO` | No | 0.875 | High/low expert split ratio |
+| `PROMPT_IMAGE` | No | `./prompt_image.png` | Seed image path |
+| `PROMPT` | No | golden retriever prompt | Text prompt |
+
+## Limitations / open items
+
+- LoRA fusion happens on CPU before device push — adds ~5–10 s per expert to
+  pipeline initialization (one-time cost, scales with stack depth).
+- Only `bh_4x8sp1tp0_ring` is exercised in CI.
+- Any community LoRA trained for Wan2.2 I2V should work if its keys follow
+  lightx2v native, diffusers, or kohya/A1111 naming. Untested LoRAs may need
+  key remap additions.

code/models/tt_dit/experimental/models/Wan2_2_SVI.md

@@ -0,0 +1,204 @@
+# Wan2.2 I2V — Stable-Video-Infinity (SVI 2.0 Pro)
+
+## Introduction
+
+`WanPipelineSVI` generates long videos by autoregressively chaining short
+Wan2.2 I2V clips. Per upstream's
+[SVI 2.0 Pro doc](https://github.com/vita-epfl/Stable-Video-Infinity/blob/svi_wan22/docs/svi/svi_2.0_pro.md),
+the conditioning latent for each clip is built as
+
+```
+y = concat([anchor_latent, motion_latent, padding], dim=temporal)
+```
+
+- **`anchor_latent`** — the VAE-encoded user-provided first frame. The same
+  anchor is used for every clip so the long video keeps a fixed visual
+  reference.
+- **`motion_latent`** — the last `num_motion_latent` latent time steps from
+  the previous clip's denoised output (`None` on the first clip). This is a
+  continuity-strength knob, not a value derived from `num_frames`.
+- **`padding`** — literal-zero latents for the remaining frame slots.
+
+The original SVI 2.0 last-frame handoff is explicitly **not** used in 2.0
+Pro — the decoded last frame of the previous clip is never re-encoded.
+Setting `num_motion_latent=0` collapses 2.0 Pro back to 2.0 behavior (anchor
+only, no latent handoff).
+
+SVI ships as a LoRA adapter pair (one per Wan2.2 MoE expert) trained with
+"Error-Recycling Fine-Tuning" so the model is robust to its own drift when
+chained. See [Wan2_2_LoRA.md](Wan2_2_LoRA.md) for the underlying LoRA
+fusion pipeline this builds on.
+
+Upstream: [vita-epfl/Stable-Video-Infinity](https://github.com/vita-epfl/Stable-Video-Infinity)
+(svi_wan22 branch).
+
+## Sampling regimes
+
+| Regime | Steps | CFG | Scheduler | flow_shift | LoRA stack (high / low) |
+|---|---|---|---|---|---|
+| `python` | 50 | 5.0 | `FlowMatchEulerDiscreteScheduler` → `EulerSolver` | 5.0 | SVI 1.0 / SVI 1.0 |
+| `comfyui` | 6 | 1.5 | `UniPCMultistepScheduler(use_flow_sigmas=True)` → `UniPCSolver` *(see note)* | 8.0 | LightX2V 0.5 + SVI 1.0 / LightX2V 1.0 + SVI 1.0 |
+
+The `python` regime matches upstream's `inference_svi_2.0_pro.py` —
+diffsynth's `FlowMatchScheduler("Wan")` is a plain Euler step on a
+flow-matching schedule, so `FlowMatchEulerDiscreteScheduler` + `EulerSolver`
+is semantically equivalent.
+
+**Scheduler caveat for `comfyui` regime.** The upstream ComfyUI workflow
+(`comfyui_workflow/SVI-Wan22-1210-{4,10}-Clips.json`) uses k-diffusion's
+`dpm++_sde` sampler. tt-metal has no on-device port of that stochastic
+2nd-order singlestep solver, so we substitute UniPC at the same
+`flow_shift=8`. Both are order-2 flow-aware solvers; output is visually
+close but not bit-exact to ComfyUI. The 4-clip variant additionally uses
+`cfg=2.0` for clips 1-2 and `cfg=1.5` for clips 3-4 — pass
+`GUIDANCE_SCALE="2.0;2.0;1.5;1.5"` (and matching `SVI_NUM_CLIPS=4`) to
+match that pattern; the 10-clip variant uses uniform `cfg=1.5` which is the
+default.
+
+## Weights
+
+### SVI 2.0 Pro LoRA
+
+- **Source:** [vita-video-gen/svi-model](https://huggingface.co/vita-video-gen/svi-model)
+- **Files (one per MoE expert):**
+  - `SVI_Wan2.2-I2V-A14B_high_noise_lora_v2.0_pro.safetensors`
+  - `SVI_Wan2.2-I2V-A14B_low_noise_lora_v2.0_pro.safetensors`
+- **Format:** PEFT-style — keys are `<base>.lora_A.default.weight` /
+  `<base>.lora_B.default.weight` (the trailing `.default` adapter-name
+  segment is tolerated by the loader; see
+  [Wan2_2_LoRA.md](Wan2_2_LoRA.md#supported-adapter-key-formats)).
+
+### LightX2V LoRA (required for `comfyui` regime only)
+
+- **Source:** [lightx2v/Wan2.2-Distill-Loras](https://huggingface.co/lightx2v/Wan2.2-Distill-Loras)
+- **Files:**
+  - `wan2.2_i2v_A14b_high_noise_lora_rank64_lightx2v_4step_1022.safetensors`
+  - `wan2.2_i2v_A14b_low_noise_lora_rank64_lightx2v_4step_1022.safetensors`
+
+### Download
+
+```bash
+python -c "
+from huggingface_hub import hf_hub_download
+for f in [
+    'SVI_Wan2.2-I2V-A14B_high_noise_lora_v2.0_pro.safetensors',
+    'SVI_Wan2.2-I2V-A14B_low_noise_lora_v2.0_pro.safetensors',
+]:
+    hf_hub_download('vita-video-gen/svi-model', f, local_dir='/path/to/svi')
+for f in [
+    'wan2.2_i2v_A14b_high_noise_lora_rank64_lightx2v_4step_1022.safetensors',
+    'wan2.2_i2v_A14b_low_noise_lora_rank64_lightx2v_4step_1022.safetensors',
+]:
+    hf_hub_download('lightx2v/Wan2.2-Distill-Loras', f, local_dir='/path/to/lightx2v')
+"
+```
+
+### Base Wan2.2 components (shared)
+
+- **Source:** [Wan-AI/Wan2.2-I2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers)
+
+## Supported configurations
+
+| System | Mesh | Topology | Test ID |
+|---|---|---|---|
+| BH Loud Box | 2x4 | Linear | `bh_2x4sp1tp0` |
+| BH Galaxy | 4x8 | Ring | `bh_4x8sp1tp0_ring` |
+
+## How to Run
+
+### ComfyUI regime (6 steps, LightX2V-accelerated — recommended)
+
+```bash
+export TT_DIT_CACHE_DIR=/your/cache/path
+export TT_METAL_HOME=$(pwd)
+export PYTHONPATH=$(pwd)
+export TT_DIT_ALLOW_HF_DOWNLOAD=1
+
+export SVI_HIGH_PATH=/path/to/svi/SVI_Wan2.2-I2V-A14B_high_noise_lora_v2.0_pro.safetensors
+export SVI_LOW_PATH=/path/to/svi/SVI_Wan2.2-I2V-A14B_low_noise_lora_v2.0_pro.safetensors
+export LIGHTX2V_HIGH_PATH=/path/to/lightx2v/wan2.2_i2v_A14b_high_noise_lora_rank64_lightx2v_4step_1022.safetensors
+export LIGHTX2V_LOW_PATH=/path/to/lightx2v/wan2.2_i2v_A14b_low_noise_lora_rank64_lightx2v_4step_1022.safetensors
+
+export PROMPT_IMAGE=$TT_METAL_HOME/example.png
+export PROMPT="A subject moving naturally through the scene"
+export SVI_NUM_CLIPS=10
+
+pytest models/tt_dit/experimental/tests/test_pipeline_wan_svi.py \
+  -v -k "bh_2x4sp1tp0 and resolution_480p and comfyui" \
+  --timeout 3600 -s
+```
+
+### Python regime (50 steps, no LightX2V needed)
+
+Same `SVI_HIGH_PATH` / `SVI_LOW_PATH` exports as above; omit the LightX2V
+ones. Replace the `-k` selector with `... and python`. Expect roughly 6×
+the wall time per clip.
+
+## Configuration
+
+### Required paths
+
+- `SVI_HIGH_PATH` / `SVI_LOW_PATH`: SVI 2.0 Pro LoRA files for the high-noise
+  and low-noise Wan2.2 experts.
+- `LIGHTX2V_HIGH_PATH` / `LIGHTX2V_LOW_PATH`: LightX2V LoRA files for the
+  high-noise and low-noise experts. Required only for `regime='comfyui'`.
+
+### Normal inputs
+
+- `PROMPT_IMAGE` (default `./prompt_image.png`): the anchor image. This is the
+  first-frame reference used for every clip.
+- `PROMPT` (default golden retriever prompt): one prompt string for all clips,
+  or a semicolon-separated list with exactly `SVI_NUM_CLIPS` entries for
+  per-clip prompts. Example: `PROMPT="wide shot;close-up;camera pans left"` with
+  `SVI_NUM_CLIPS=3`.
+- `SVI_NUM_CLIPS` (default `2`): number of 81-frame clips to generate and
+  chain.
+
+### Expert knobs
+
+These are useful for matching upstream workflows, ablations, or debugging. The
+defaults are the intended starting point.
+
+- `NUM_STEPS` (default `50` for `python`, `6` for `comfyui`): denoising steps
+  per clip. Higher values are slower and usually improve quality. The `comfyui`
+  default assumes the LightX2V acceleration LoRAs are present.
+- `GUIDANCE_SCALE` (default `5.0` for `python`, `1.5` for `comfyui`): CFG for
+  the high-noise expert. Accepts either one value for all clips or a
+  semicolon-separated list of length `SVI_NUM_CLIPS`.
+- `GUIDANCE_SCALE_2` (default: matches `GUIDANCE_SCALE`): CFG for the low-noise
+  expert. Uses the same scalar or semicolon per-clip format. For the upstream
+  ComfyUI 4-clip workflow, use
+  `GUIDANCE_SCALE="2.0;2.0;1.5;1.5"` and the same value for
+  `GUIDANCE_SCALE_2`.
+- `SVI_SEED` (default `0`): base random seed. Clip `i` uses
+  `SVI_SEED + i * 42`, so each clip gets deterministic but distinct noise.
+- `SVI_NUM_MOTION_LATENT` (default `1`): number of latent time steps copied
+  from the previous clip into the next clip's conditioning. Increasing this can
+  strengthen cross-clip continuity but may make the next clip follow the
+  previous motion more tightly. Setting it to `0` disables the SVI Pro latent
+  handoff and leaves only the anchor image.
+- `SVI_NUM_OVERLAP_FRAME` (default `4`): number of decoded frames dropped from
+  the start of every clip after the first when concatenating the final video.
+  This removes duplicate or unstable boundary frames; it does not affect latent
+  conditioning.
+
+## Limitations / open items
+
+- **`comfyui` regime is not bit-exact to upstream** — see the scheduler
+  caveat above. A native `DPMSolverSDESolver` is the follow-up that would
+  close the gap.
+- **Anchor encoding differs slightly from upstream.** Our implementation
+  reuses `WanPipelineI2V.prepare_latents`, which encodes an anchor-at-frame-0 /
+  zero-elsewhere video and then overwrites the non-anchor latents with the
+  motion-latent splice and zero-padding. Upstream's `wan_video_svi_pro`
+  encodes only the single anchor frame and concatenates. Both produce the
+  same anchor latent at frame 0 and zeros / motion at the rest, but the
+  encode path is not byte-identical.
+- **CPU-side LoRA fusion** (~5–10 s per expert at cold cache, scales with
+  stack depth). Switching between LoRA stacks across runs rebuilds the TT
+  cache; runtime LoRA scale tuning or mid-clip switching is not supported.
+- **Per-clip latency** (BH-LB 2x4, 480p, with anchor VAE-encode cache hit
+  after clip 1): `comfyui` regime ~63 s/clip; `python` regime not yet
+  measured end-to-end on this branch.
+- `bh_2x4sp1tp0` (BH-LB) is the primary tested config; `bh_4x8sp1tp0_ring`
+  (Galaxy) is parameterized but not yet exercised.

code/models/tt_dit/experimental/pipelines/pipeline_anisora.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""IndexTeam/Index-anisora V3.2 I2V pipeline.
+
+A thin subclass of :class:`WanPipelineI2V` that:
+
+1. Uses ``Wan-AI/Wan2.2-I2V-A14B-Diffusers`` for tokenizer, text encoder, VAE,
+   and scheduler — the AniSora HF repo only ships the DiT experts.
+2. Replaces both expert transformer state dicts with AniSora's ``V3.2``
+   ``diffusion_pytorch_model.safetensors`` files (high-noise + low-noise pair).
+3. Defaults to 40 sampling steps with ``boundary_ratio=0.9`` and
+   ``guidance_scale=3.5`` for both stages, matching the upstream config in
+   ``anisoraV3.2/wan/configs/wan_i2v_A14B.py``.
+
+AniSora's safetensors use the same original-Wan key naming as the lightx2v
+distill checkpoints (verified: 1095 keys, identical layout). The
+``wan_lightx2v_to_diffusers_key`` rename function in
+:mod:`models.tt_dit.utils.lightx2v_loader` is reused as-is.
+
+A ``random_weights`` mode is provided for smoke-testing without HuggingFace
+downloads. It (a) skips the two ~28 GB transformer subfolder downloads by
+constructing config-only random ``TorchWanTransformer3DModel`` instances and
+(b) feeds those random ``state_dict()``s into the TT model in place of the
+AniSora safetensors. Tokenizer, text encoder, VAE, and scheduler still come
+from the base diffusers repo (~12 GB total).
+"""
+from __future__ import annotations
+
+import contextlib
+import os
+from typing import TYPE_CHECKING
+
+import ttnn
+from models.tt_dit.experimental.pipelines.pipeline_wan_distill import _patch_torch_transformer_random
+from models.tt_dit.experimental.utils.lightx2v_loader import load_lightx2v_state_dict
+from models.tt_dit.pipelines.wan.pipeline_wan import WanPipelineConfig
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import WanPipelineI2V
+from models.tt_dit.utils import cache
+
+if TYPE_CHECKING:
+    from diffusers.schedulers import SchedulerMixin
+
+
+class AniSoraPipeline(WanPipelineI2V):
+    HF_REPO = "IndexTeam/Index-anisora"
+    HIGH_NOISE_FILE = "V3.2/high_noise_model/diffusion_pytorch_model.safetensors"
+    LOW_NOISE_FILE = "V3.2/low_noise_model/diffusion_pytorch_model.safetensors"
+    CACHE_NAMESPACE = "Index-anisora-V3.2"
+    RANDOM_CACHE_NAMESPACE = "Index-anisora-random"
+    ANISORA_BOUNDARY_RATIO = 0.9
+
+    @classmethod
+    def _config_overrides(cls) -> dict[str, object]:
+        # Matches upstream anisoraV3.2/wan/configs/wan_i2v_A14B.py.
+        return {**super()._config_overrides(), "boundary_ratio": cls.ANISORA_BOUNDARY_RATIO}
+
+    def __init__(
+        self,
+        *,
+        device: ttnn.MeshDevice,
+        config: WanPipelineConfig,
+        anisora_local_dir: str | None = None,
+        allow_download: bool | None = None,
+        random_weights: bool | None = None,
+        scheduler: SchedulerMixin | None = None,
+        run_warmup: bool = True,
+        lora_enabled: bool = False,
+    ) -> None:
+        if allow_download is None:
+            allow_download = os.environ.get("TT_DIT_ALLOW_HF_DOWNLOAD") == "1"
+        if anisora_local_dir is None:
+            anisora_local_dir = os.environ.get("ANISORA_LOCAL_DIR")
+        if random_weights is None:
+            random_weights = os.environ.get("TT_DIT_RANDOM_WEIGHTS") == "1"
+
+        self._anisora_local_dir = anisora_local_dir
+        self._allow_download = allow_download
+        self._random_weights = random_weights
+
+        ctx = _patch_torch_transformer_random() if random_weights else contextlib.nullcontext()
+        with ctx:
+            super().__init__(
+                device=device, config=config, scheduler=scheduler, run_warmup=run_warmup, lora_enabled=lora_enabled
+            )
+
+    def _prepare_transformer(self, idx: int):
+        state = self.transformer_states[idx]
+        if self._random_weights:
+            cache.load_model(
+                state.model,
+                model_name=self.RANDOM_CACHE_NAMESPACE,
+                subfolder=state.checkpoint.subfolder,
+                parallel_config=self.parallel_config,
+                mesh_shape=tuple(self.mesh_device.shape),
+                mesh_device=self.mesh_device,
+                is_fsdp=self.is_fsdp,
+                get_torch_state_dict=lambda s=state: s.checkpoint.state_dict(),
+            )
+            return
+
+        filename = self.HIGH_NOISE_FILE if idx == 0 else self.LOW_NOISE_FILE
+        cache.load_model(
+            state.model,
+            model_name=self.CACHE_NAMESPACE,
+            subfolder=state.checkpoint.subfolder,
+            parallel_config=self.parallel_config,
+            mesh_shape=tuple(self.mesh_device.shape),
+            mesh_device=self.mesh_device,
+            is_fsdp=self.is_fsdp,
+            get_torch_state_dict=lambda f=filename: load_lightx2v_state_dict(
+                self.HF_REPO,
+                f,
+                allow_download=self._allow_download,
+                local_dir=self._anisora_local_dir,
+            ),
+        )

code/models/tt_dit/experimental/pipelines/pipeline_wan_distill.py

@@ -0,0 +1,181 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Wan2.2 lightx2v distilled I2V pipeline.
+
+A thin subclass of :class:`WanPipelineI2V` that:
+
+1. Uses ``Wan-AI/Wan2.2-I2V-A14B-Diffusers`` for tokenizer, text encoder, VAE,
+   and scheduler — the lightx2v repo ships only the DiT.
+2. Replaces both expert transformer state dicts with lightx2v's flat
+   ``.safetensors`` files (4-step distill, high-noise + low-noise pair).
+3. Defaults ``boundary_ratio=0.5`` (2 high-noise + 2 low-noise steps).
+
+Test/caller passes ``num_inference_steps=4`` and ``guidance_scale=1.0`` for both
+stages since CFG is baked into the distill.
+
+A ``random_weights`` mode is provided for smoke-testing without HuggingFace
+downloads. It (a) skips the two ~28 GB transformer subfolder downloads by
+constructing config-only random ``TorchWanTransformer3DModel`` instances and
+(b) feeds those random ``state_dict()``s into the TT model in place of the
+lightx2v safetensors. Tokenizer, text encoder, VAE, and scheduler still come
+from the base diffusers repo (~12 GB total).
+"""
+from __future__ import annotations
+
+import contextlib
+import os
+from typing import TYPE_CHECKING
+
+import torch
+
+import ttnn
+from models.tt_dit.experimental.utils.lightx2v_loader import load_lightx2v_state_dict
+from models.tt_dit.models.transformers.wan2_2.transformer_wan import TorchWanTransformer3DModel
+from models.tt_dit.pipelines.wan.pipeline_wan import WanPipelineConfig
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import WanPipelineI2V
+from models.tt_dit.utils import cache
+
+if TYPE_CHECKING:
+    from diffusers.schedulers import SchedulerMixin
+
+# Hard-coded config for Wan2.2-I2V-A14B-Diffusers transformer subfolders. Used
+# only in random_weights mode so we don't have to fetch transformer/config.json
+# from HF. Both transformer and transformer_2 share this architecture; only
+# weights differ between high-noise and low-noise experts.
+_RANDOM_I2V_TRANSFORMER_CONFIG = dict(
+    patch_size=(1, 2, 2),
+    num_attention_heads=40,
+    attention_head_dim=128,
+    in_channels=36,  # I2V: 4 mask + 16 image latent + 16 noise latent
+    out_channels=16,
+    text_dim=4096,
+    freq_dim=256,
+    ffn_dim=13824,
+    num_layers=40,
+    cross_attn_norm=True,
+    qk_norm="rms_norm_across_heads",
+    eps=1e-6,
+    rope_max_seq_len=1024,
+    image_dim=None,
+    added_kv_proj_dim=None,
+    pos_embed_seq_len=None,
+)
+
+
+@contextlib.contextmanager
+def _patch_torch_transformer_random(seed: int = 0):
+    """Replace ``TorchWanTransformer3DModel.from_pretrained`` with a stub that
+    instantiates from a hard-coded I2V config and never touches the network or
+    disk. Restored on exit."""
+    cls = TorchWanTransformer3DModel
+    sentinel = object()
+    saved = cls.__dict__.get("from_pretrained", sentinel)
+
+    def _stub(_cls, *args, **kwargs):
+        with torch.random.fork_rng():
+            torch.manual_seed(seed)
+            return _cls(**_RANDOM_I2V_TRANSFORMER_CONFIG)
+
+    cls.from_pretrained = classmethod(_stub)
+    try:
+        yield
+    finally:
+        if saved is sentinel:
+            # `from_pretrained` was inherited; delete our override to restore inheritance.
+            try:
+                delattr(cls, "from_pretrained")
+            except AttributeError:
+                pass
+        else:
+            cls.from_pretrained = saved
+
+
+class WanDistillPipelineI2V(WanPipelineI2V):
+    LIGHTX2V_REPO = "lightx2v/Wan2.2-Distill-Models"
+    HIGH_NOISE_FILE = "wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors"
+    LOW_NOISE_FILE = "wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors"
+    CACHE_NAMESPACE = "Wan2.2-Distill-lightx2v-4step"
+    RANDOM_CACHE_NAMESPACE = "Wan2.2-Distill-random"
+    DISTILL_BOUNDARY_RATIO = 0.5
+
+    @classmethod
+    def _config_overrides(cls) -> dict[str, object]:
+        # The 4-step distill splits 2 high-noise + 2 low-noise steps.
+        return {**super()._config_overrides(), "boundary_ratio": cls.DISTILL_BOUNDARY_RATIO}
+
+    def __init__(
+        self,
+        *,
+        device: ttnn.MeshDevice,
+        config: WanPipelineConfig,
+        lightx2v_local_dir: str | None = None,
+        allow_download: bool | None = None,
+        random_weights: bool | None = None,
+        scheduler: SchedulerMixin | None = None,
+        run_warmup: bool = True,
+        lora_enabled: bool = False,
+    ) -> None:
+        if allow_download is None:
+            allow_download = os.environ.get("TT_DIT_ALLOW_HF_DOWNLOAD") == "1"
+        if lightx2v_local_dir is None:
+            lightx2v_local_dir = os.environ.get("LIGHTX2V_LOCAL_DIR")
+        if random_weights is None:
+            random_weights = os.environ.get("TT_DIT_RANDOM_WEIGHTS") == "1"
+
+        self._lightx2v_local_dir = lightx2v_local_dir
+        self._allow_download = allow_download
+        self._random_weights = random_weights
+
+        ctx = _patch_torch_transformer_random() if random_weights else contextlib.nullcontext()
+        with ctx:
+            super().__init__(
+                device=device, config=config, scheduler=scheduler, run_warmup=run_warmup, lora_enabled=lora_enabled
+            )
+
+    def prepare_text_conditioning(self, tt_model, prompt_embeds, buffer, traced=False):
+        # When CFG is baked in (guidance_scale=1.0), encode_prompt returns
+        # negative_prompt_embeds=None. The base loop still calls this for the
+        # negative buffer; forwarding None into the text embedder hits a
+        # NoneType.padded_shape in Linear. combined_step already short-circuits
+        # on do_classifier_free_guidance=False, so leaving the buffer as-is is
+        # safe.
+        if prompt_embeds is None:
+            return buffer
+        return super().prepare_text_conditioning(tt_model, prompt_embeds, buffer, traced)
+
+    def _prepare_transformer(self, idx: int):
+        state = self.transformer_states[idx]
+        if self._random_weights:
+            # Use the (random) state_dict the WanCheckpoint loaded under the
+            # _patch_torch_transformer_random monkey-patch. Cache under a
+            # separate namespace so a real-weights run doesn't reuse it.
+            cache.load_model(
+                state.model,
+                model_name=self.RANDOM_CACHE_NAMESPACE,
+                subfolder=state.checkpoint.subfolder,
+                parallel_config=self.parallel_config,
+                mesh_shape=tuple(self.mesh_device.shape),
+                mesh_device=self.mesh_device,
+                is_fsdp=self.is_fsdp,
+                get_torch_state_dict=lambda s=state: s.checkpoint.state_dict(),
+            )
+            return
+
+        filename = self.HIGH_NOISE_FILE if idx == 0 else self.LOW_NOISE_FILE
+        cache.load_model(
+            state.model,
+            model_name=self.CACHE_NAMESPACE,
+            subfolder=state.checkpoint.subfolder,
+            parallel_config=self.parallel_config,
+            mesh_shape=tuple(self.mesh_device.shape),
+            mesh_device=self.mesh_device,
+            is_fsdp=self.is_fsdp,
+            get_torch_state_dict=lambda f=filename: load_lightx2v_state_dict(
+                self.LIGHTX2V_REPO,
+                f,
+                allow_download=self._allow_download,
+                local_dir=self._lightx2v_local_dir,
+            ),
+        )

code/models/tt_dit/experimental/pipelines/pipeline_wan_lora.py

@@ -0,0 +1,414 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Wan2.2 I2V pipeline with LoRA adapters fused into the base weights.
+
+Each expert (high/low noise) takes an ordered LoRA stack; stacks are fused
+on CPU before TT conversion so inference has no LoRA-specific runtime cost.
+See ``experimental/models/Wan2_2_LoRA.md`` for the adapter-key formats
+detected by ``fuse_lora_state_dict`` and the supported namespaces.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+from loguru import logger
+from safetensors.torch import load_file
+
+import ttnn
+from models.tt_dit.experimental.utils.lightx2v_loader import wan_lightx2v_to_diffusers_key
+from models.tt_dit.pipelines.wan.pipeline_wan import WanPipelineConfig
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import WanPipelineI2V
+from models.tt_dit.solvers.factory import CustomSigmaScheduler
+from models.tt_dit.utils import cache
+
+if TYPE_CHECKING:
+    from diffusers.schedulers import SchedulerMixin
+
+
+@dataclass(frozen=True)
+class LoRASpec:
+    """A single LoRA adapter file and its blend strength."""
+
+    path: str
+    scale: float = 1.0
+
+
+LoRAArg = LoRASpec | str | Sequence[LoRASpec | str] | None
+
+
+@dataclass(frozen=True)
+class _FusionStats:
+    applied_pairs: int = 0
+    applied_direct: int = 0
+    skipped_unknown: int = 0
+    skipped_unmapped: int = 0
+    skipped_shape_mismatch: int = 0
+
+    @property
+    def applied(self) -> int:
+        return self.applied_pairs + self.applied_direct
+
+
+def _normalize_lora_arg(arg: LoRAArg) -> list[LoRASpec]:
+    """Coerce None / path / LoRASpec / sequence-of-either into LoRASpecs."""
+    if arg is None:
+        return []
+    if isinstance(arg, LoRASpec):
+        return [arg]
+    if isinstance(arg, str):
+        return [LoRASpec(arg)]
+    if not isinstance(arg, Sequence):
+        raise TypeError(f"Expected LoRASpec, str, sequence, or None; got {type(arg).__name__}")
+    out: list[LoRASpec] = []
+    for item in arg:
+        if isinstance(item, LoRASpec):
+            out.append(item)
+        elif isinstance(item, str):
+            out.append(LoRASpec(item))
+        else:
+            raise TypeError(f"Expected LoRASpec or str in LoRA list, got {type(item).__name__}")
+    return out
+
+
+_STRIP_PREFIXES = ("diffusion_model.", "transformer.", "unet.", "model.")
+
+# PEFT may insert an adapter name before ``.weight``, e.g. ``lora_A.default.weight``.
+_LOW_RANK_RE = re.compile(r"^(?P<base>.*)\.lora_(?P<slot>A|B|down|up)(?:\.[^.]+)?\.weight$")
+_SLOT_MAP = {"A": "A", "down": "A", "B": "B", "up": "B"}
+
+
+def _strip_known_prefixes(key: str) -> str:
+    for prefix in _STRIP_PREFIXES:
+        if key.startswith(prefix):
+            return key[len(prefix) :]
+    return key
+
+
+def _kohya_to_lightx2v(key: str) -> str:
+    """Convert kohya/A1111-style keys to lightx2v-style keys.
+
+    ``lora_unet_blocks_0_cross_attn_k.lora_down.weight``
+    → ``blocks.0.cross_attn.k.lora_down.weight``
+    """
+    if not key.startswith("lora_unet_"):
+        return key
+    parts = key.split(".", 1)
+    module_path = parts[0][len("lora_unet_") :]
+    suffix = f".{parts[1]}" if len(parts) > 1 else ""
+
+    m = re.match(r"blocks_(\d+)_(cross_attn|self_attn)_([a-z]+)", module_path)
+    if m:
+        return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
+
+    m = re.match(r"blocks_(\d+)_(ffn)_(\d+)", module_path)
+    if m:
+        return f"blocks.{m.group(1)}.{m.group(2)}.{m.group(3)}{suffix}"
+
+    logger.warning(f"Unrecognized kohya key structure: {key}")
+    return key
+
+
+def _diffusers_target(lightx2v_base_path: str, suffix: str) -> str:
+    """Map ``<lightx2v_base>.weight``/``.bias`` to the diffusers parameter key."""
+    weight_key = wan_lightx2v_to_diffusers_key(f"{lightx2v_base_path}.weight")
+    if suffix == ".weight":
+        return weight_key
+    return weight_key[: -len(".weight")] + suffix
+
+
+def _is_lora_key(raw_key: str) -> bool:
+    key = _kohya_to_lightx2v(_strip_known_prefixes(raw_key))
+    return bool(_LOW_RANK_RE.match(key)) or key.endswith((".diff", ".diff_b"))
+
+
+def _has_lora_keys(state_dict: dict[str, torch.Tensor]) -> bool:
+    """Return True iff the safetensors file contains at least one LoRA-style key."""
+    return any(_is_lora_key(k) for k in state_dict)
+
+
+def fuse_lora_state_dict(
+    base_state_dict: dict[str, torch.Tensor],
+    lora_state_dict: dict[str, torch.Tensor],
+    *,
+    scale: float = 1.0,
+    return_stats: bool = False,
+) -> dict[str, torch.Tensor] | tuple[dict[str, torch.Tensor], _FusionStats]:
+    """Return a fused state dict, optionally with stats describing what applied.
+
+    The base dict is not mutated; entries that were touched are fresh tensors
+    so the caller can chain fusions across a stack safely. See
+    ``experimental/models/Wan2_2_LoRA.md`` for the supported adapter key
+    formats. Raises ``KeyError`` when a low-rank pair is missing one half.
+    """
+    pairs: dict[str, dict[str, torch.Tensor]] = {}
+    direct_deltas: list[tuple[str, str, torch.Tensor]] = []  # (base_path, suffix, tensor)
+    alphas: dict[str, float] = {}
+    skipped_unknown: list[str] = []
+
+    for raw_key, tensor in lora_state_dict.items():
+        key = _kohya_to_lightx2v(_strip_known_prefixes(raw_key))
+        m = _LOW_RANK_RE.match(key)
+        if m:
+            pairs.setdefault(m.group("base"), {})[_SLOT_MAP[m.group("slot")]] = tensor
+            continue
+        if key.endswith(".diff_b"):
+            direct_deltas.append((key[: -len(".diff_b")], ".bias", tensor))
+        elif key.endswith(".diff"):
+            direct_deltas.append((key[: -len(".diff")], ".weight", tensor))
+        elif key.endswith(".alpha"):
+            alphas[key[: -len(".alpha")]] = tensor.item()
+        else:
+            skipped_unknown.append(raw_key)
+
+    if skipped_unknown:
+        logger.warning(
+            f"LoRA fusion: {len(skipped_unknown)} unrecognized keys ignored. Examples: {skipped_unknown[:5]}"
+        )
+
+    fused = dict(base_state_dict)
+    applied_pairs = 0
+    skipped_unmapped: list[str] = []
+
+    for base_path, ab in pairs.items():
+        if "A" not in ab or "B" not in ab:
+            raise KeyError(f"LoRA pair incomplete for '{base_path}': have {list(ab)}")
+        diffusers_key = _diffusers_target(base_path, ".weight")
+        if diffusers_key not in fused:
+            skipped_unmapped.append(diffusers_key)
+            continue
+        base_weight = fused[diffusers_key]
+        rank = ab["A"].shape[0]
+        alpha = alphas.get(base_path, float(rank))
+        effective_scale = scale * (alpha / rank)
+        delta = effective_scale * (ab["B"].to(torch.float32) @ ab["A"].to(torch.float32))
+        # ``+`` always allocates a new tensor; avoids the .to(fp32)+.add_()
+        # alias trap when base_weight is already fp32 (would mutate the
+        # caller's dict on a subsequent stacked-fusion pass).
+        fused[diffusers_key] = (base_weight.float() + delta).to(base_weight.dtype)
+        applied_pairs += 1
+
+    applied_direct = 0
+    skipped_shape_mismatch = 0
+    for base_path, suffix, tensor in direct_deltas:
+        diffusers_key = _diffusers_target(base_path, suffix)
+        if diffusers_key not in fused:
+            skipped_unmapped.append(diffusers_key)
+            continue
+        base = fused[diffusers_key]
+        if base.shape != tensor.shape:
+            logger.warning(
+                f"LoRA direct delta shape mismatch for '{diffusers_key}': "
+                f"base {tuple(base.shape)} vs delta {tuple(tensor.shape)}; skipping."
+            )
+            skipped_shape_mismatch += 1
+            continue
+        fused[diffusers_key] = (base.float() + scale * tensor.to(torch.float32)).to(base.dtype)
+        applied_direct += 1
+
+    if skipped_unmapped:
+        logger.warning(
+            f"LoRA fusion: {len(skipped_unmapped)} adapter targets not present in base; skipped. "
+            f"Examples: {skipped_unmapped[:5]}"
+        )
+
+    logger.info(f"Fused {applied_pairs} low-rank pairs and {applied_direct} direct deltas (scale={scale})")
+    stats = _FusionStats(
+        applied_pairs=applied_pairs,
+        applied_direct=applied_direct,
+        skipped_unknown=len(skipped_unknown),
+        skipped_unmapped=len(skipped_unmapped),
+        skipped_shape_mismatch=skipped_shape_mismatch,
+    )
+    return (fused, stats) if return_stats else fused
+
+
+def verify_fusion_changed_weights(
+    base_sd: dict[str, torch.Tensor],
+    fused_sd: dict[str, torch.Tensor],
+    *,
+    min_changed: int = 3,
+    label: str,
+) -> None:
+    """Sanity check that at least ``min_changed`` weights actually differ post-fusion.
+
+    Catches the canonical "LoRA silently failed to apply" failure mode where
+    a key-mapping bug leaves every fused target absent from the base. Logs a
+    sample of changed tensors with their L2 diff for visibility.
+
+    Raises ``RuntimeError`` if no weights changed at all, or if fewer than
+    ``min_changed`` weights changed (likely a partial load).
+    """
+    changed: list[tuple[str, float]] = []
+    max_diff = 0.0
+    for k, base in base_sd.items():
+        fused = fused_sd.get(k)
+        if fused is None or fused.shape != base.shape:
+            continue
+        if fused.data_ptr() == base.data_ptr():
+            continue
+        diff = (fused.to(torch.float32) - base.to(torch.float32)).norm().item()
+        if diff > 0.0:
+            changed.append((k, diff))
+            max_diff = max(max_diff, diff)
+
+    if max_diff == 0.0:
+        raise RuntimeError(
+            f"LoRA silently failed to apply -- weights are unchanged for '{label}'. "
+            f"Verified {len(base_sd)} keys against fused dict; max L2 diff is 0.0."
+        )
+
+    sample = changed[:5]
+    logger.info(
+        f"LoRA fusion verified for '{label}': {len(changed)} tensors changed, "
+        f"max L2 diff={max_diff:.4f}. Sample diffs:"
+    )
+    for k, d in sample:
+        logger.info(f"  {k}: L2={d:.4f}")
+
+    if len(changed) < min_changed:
+        raise RuntimeError(
+            f"LoRA fusion changed only {len(changed)} weights for '{label}' "
+            f"(require >= {min_changed}). Likely indicates partial LoRA load."
+        )
+
+
+def _lora_stack_cache_namespace(specs_by_expert: dict[int, list[LoRASpec]]) -> str:
+    """Hash ordered ``(resolved_path, scale)`` per expert so distinct stacks cache separately."""
+    h = hashlib.sha1()
+    for idx in sorted(specs_by_expert.keys()):
+        h.update(f"expert_{idx}\x00".encode())
+        for spec in specs_by_expert[idx]:
+            h.update(str(Path(spec.path).resolve()).encode())
+            h.update(b"\x00")
+            h.update(f"{spec.scale:.6f}".encode())
+            h.update(b"\x00")
+    return f"Wan2.2-I2V-LoRA-{h.hexdigest()[:12]}"
+
+
+class WanPipelineI2VLora(WanPipelineI2V):
+    """Wan2.2 I2V with LoRA stacks fused into the base PyTorch weights.
+
+    Adapters distilled to be guidance-free want ``cfg_enabled=False``, which halves the
+    denoising cost by dropping the unconditional pass. With it off the caller must keep both
+    guidance scales at 1.0, which ``__call__`` enforces.
+    """
+
+    def __init__(
+        self,
+        *,
+        device: ttnn.MeshDevice,
+        config: WanPipelineConfig,
+        lora_high: LoRAArg = None,
+        lora_low: LoRAArg = None,
+        scheduler: SchedulerMixin | None = None,
+        run_warmup: bool = True,
+        lora_enabled: bool = False,
+    ) -> None:
+        high_specs = _normalize_lora_arg(lora_high)
+        low_specs = _normalize_lora_arg(lora_low)
+
+        if not high_specs and not low_specs:
+            raise ValueError(
+                "WanPipelineI2VLora requires at least one LoRA. "
+                "Pass lora_high and/or lora_low as a path, LoRASpec, or list."
+            )
+
+        for label, specs in [("lora_high", high_specs), ("lora_low", low_specs)]:
+            for spec in specs:
+                if not Path(spec.path).is_file():
+                    raise FileNotFoundError(f"{label}: file does not exist: {spec.path}")
+
+        self._lora_specs: dict[int, list[LoRASpec]] = {0: high_specs, 1: low_specs}
+        self._cache_namespace = _lora_stack_cache_namespace(self._lora_specs)
+        # Cleared after TT cache handoff (see _prepare_transformer) to free CPU memory.
+        self._fused_state_dicts: dict[int, dict[str, torch.Tensor] | None] = {0: None, 1: None}
+        self.custom_sigmas = None
+
+        if scheduler is None:
+            scheduler = CustomSigmaScheduler()
+            self.custom_sigmas = [1.0, 0.9375, 0.8333333, 0.625, 0.0]  # default for Wan 2.2 I2V Lightning
+            logger.info(
+                f"WanPipelineI2VLora: default Euler scheduler with custom sigmas {self.custom_sigmas} -- (Wan 2.2 I2V Lightning)"
+            )
+
+        super().__init__(
+            device=device, config=config, scheduler=scheduler, run_warmup=run_warmup, lora_enabled=lora_enabled
+        )
+
+    def prepare_text_conditioning(self, tt_model, prompt_embeds, buffer, traced=False):
+        # guidance_scale=1.0 → encoder returns None for negative embeds; combined_step
+        # skips CFG so the untouched buffer is fine.
+        if prompt_embeds is None:
+            return buffer
+        return super().prepare_text_conditioning(tt_model, prompt_embeds, buffer, traced)
+
+    def prepare_schedule(self, num_inference_steps: int, *, flow_shift: float | None = None) -> None:
+        if self._solver.scheduler is None:
+            self._solver.set_schedule(sigmas=self.custom_sigmas)
+        else:
+            super().prepare_schedule(num_inference_steps, flow_shift=flow_shift)
+
+    def _build_fused_state_dict(self, idx: int) -> dict[str, torch.Tensor] | None:
+        specs = self._lora_specs[idx]
+        state = self.transformer_states[idx]
+        subfolder = state.checkpoint.subfolder
+        if not specs:
+            logger.info(f"No LoRA for expert idx={idx} ('{subfolder}') -- using base weights")
+            return None
+
+        base_sd = state.checkpoint.state_dict()
+        fused_sd = base_sd
+        for spec in specs:
+            logger.info(f"Loading LoRA for '{subfolder}' from {spec.path} (scale={spec.scale})")
+            lora_sd = load_file(str(spec.path))
+            if not _has_lora_keys(lora_sd):
+                raise RuntimeError(
+                    f"No LoRA-style keys (lora_A/lora_B, lora_down/lora_up, diff/diff_b) found in {spec.path}"
+                )
+            previous_sd = fused_sd
+            fused_sd, stats = fuse_lora_state_dict(previous_sd, lora_sd, scale=spec.scale, return_stats=True)
+            label = f"{subfolder}: {Path(spec.path).name}"
+            if stats.applied == 0:
+                raise RuntimeError(
+                    f"LoRA fusion applied no tensors for '{label}'. "
+                    f"Skipped unmapped={stats.skipped_unmapped}, unknown={stats.skipped_unknown}, "
+                    f"shape_mismatch={stats.skipped_shape_mismatch}."
+                )
+            verify_fusion_changed_weights(previous_sd, fused_sd, label=label)
+        return fused_sd
+
+    def _prepare_transformer(self, idx: int):
+        state = self.transformer_states[idx]
+
+        if not self._lora_specs[idx]:
+            super()._prepare_transformer(idx)
+            return
+
+        def _get_state_dict(idx_=idx):
+            cached = self._fused_state_dicts.get(idx_)
+            if cached is not None:
+                return cached
+            sd = self._build_fused_state_dict(idx_)
+            self._fused_state_dicts[idx_] = sd
+            return sd
+
+        cache.load_model(
+            state.model,
+            model_name=self._cache_namespace,
+            subfolder=state.checkpoint.subfolder,
+            parallel_config=self.parallel_config,
+            mesh_shape=tuple(self.mesh_device.shape),
+            mesh_device=self.mesh_device,
+            is_fsdp=self.is_fsdp,
+            get_torch_state_dict=_get_state_dict,
+        )
+        self._fused_state_dicts[idx] = None

code/models/tt_dit/experimental/pipelines/pipeline_wan_runtime_lora.py

@@ -0,0 +1,267 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Wan 2.2 pipeline with runtime-swappable LoRA adapters.
+
+Constructs the transformer with ``lora_enabled=True`` so every attention
+and FFN Linear is a LoRA-aware variant. Exposes a small API for
+registering and switching adapters at runtime:
+
+  pipeline.register_lora(name, high_path, low_path, scale)  → AdapterHandle
+  pipeline.set_active_lora(handle | None)
+  pipeline.unregister_lora(handle)
+  pipeline.list_loras()
+
+Adapter swap merges the active LoRA delta directly into the base weight:
+``W += scale * A.T @ B.T`` per LoRA Linear, on device, with W's sharding.
+Forward-time LoRA overhead is zero — the captured trace runs the
+un-modified base path. Swap cost is bounded to ~one small matmul + an
+in-place add per LoRA Linear; sub-second end to end. See
+``models/tt_dit/layers/lora.py`` for the merge details.
+
+Out of scope for v0:
+  - Stacking multiple adapters into a single bound slot (callers wanting
+    stacked LoRAs should pre-merge offline, or use the older CPU-fusion
+    pipeline at ``pipeline_wan_lora.py``).
+  - ``.diff`` / ``.diff_b`` direct deltas (the loader warns and skips).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from loguru import logger
+
+from models.tt_dit.experimental.lora.adapter_loader import AdapterHandle, load_adapter_into
+from models.tt_dit.layers.lora import LoRAMixin
+from models.tt_dit.pipelines.wan.pipeline_wan import WanPipeline
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import WanPipelineI2V
+
+
+@dataclass
+class _RegisteredLoRA:
+    """Per-pipeline record of a registered adapter pair."""
+
+    name: str
+    handle_high: AdapterHandle | None
+    handle_low: AdapterHandle | None
+
+
+@dataclass
+class CombinedAdapterHandle:
+    """Returned by `register_lora`. Knows about both experts."""
+
+    name: str
+    high: AdapterHandle | None
+    low: AdapterHandle | None
+    rank: int = field(init=False)
+
+    def __post_init__(self):
+        ranks = [h.rank for h in (self.high, self.low) if h is not None]
+        if not ranks:
+            raise ValueError(f"adapter '{self.name}' loaded for neither expert")
+        self.rank = max(ranks)
+
+
+def _iter_lora_modules(module):
+    """Yield every LoRAMixin-bearing submodule of `module`, recursively."""
+    if isinstance(module, LoRAMixin):
+        yield module
+    for _, child in module.named_children():
+        yield from _iter_lora_modules(child)
+
+
+class _LoRAPipelineMixin:
+    """Implements the LoRA swap API on top of a WanPipeline subclass.
+
+    Assumes the host pipeline was constructed with ``lora_enabled=True``
+    so both ``self.transformer`` and ``self.transformer_2`` have LoRA-aware
+    Linear modules throughout.
+    """
+
+    def __init__(self, *args, lora_enabled: bool = True, **kwargs):
+        # Initialize LoRA state *before* super().__init__() because the base
+        # pipeline init calls self._prepare_transformer() during warmup, and
+        # our override below reads self._active_adapter_name.
+        self._lora_registry: dict[str, _RegisteredLoRA] = {}
+        self._active_adapter_name: str | None = None
+        super().__init__(*args, lora_enabled=lora_enabled, **kwargs)
+
+        # Fuse-mode bind/unbind adds and subtracts the delta in bf16; under
+        # dynamic_load=True the cache restore re-seeds the base every reload
+        # and drift is bounded. With dynamic_load=False the base weight is
+        # never refreshed, so W drifts by ~O(N · eps_bf16) over N swap cycles.
+        # See models/tt_dit/layers/lora.py for the snapshot+reseed escape hatch.
+        if not getattr(self, "dynamic_load", True):
+            logger.warning(
+                "LoRA fuse-mode pipeline constructed with dynamic_load=False — "
+                "repeated bind/unbind cycles will accumulate bf16 quantization "
+                "drift in the base weight. For long-lived sessions, snapshot "
+                "the transformer state dict at construction and re-seed via "
+                "load_torch_state_dict when drift becomes material."
+            )
+
+    @classmethod
+    def create_pipeline(cls, **kwargs):
+        # The base factory defaults lora_enabled=False; runtime LoRA needs it on.
+        kwargs.setdefault("lora_enabled", True)
+        return super().create_pipeline(**kwargs)
+
+    def _prepare_transformer(self, idx: int):
+        """Re-apply any active LoRA delta after the transformer is reloaded.
+
+        Under dynamic_load the two transformers swap in and out of DRAM
+        between uses; each reload restores the cached *base* weights, so
+        the merged LoRA delta is wiped. We walk this transformer's LoRA
+        modules and re-merge their active adapter (no-op for modules that
+        already have the delta applied, e.g., the first load).
+        """
+        super()._prepare_transformer(idx)
+        if self._active_adapter_name is None:
+            return
+        transformer = self.transformer_states[idx].model
+        for mod in _iter_lora_modules(transformer):
+            mod.reapply_after_load()
+
+    # ---- public API ----
+    def register_lora(
+        self,
+        name: str,
+        *,
+        high_path: str | Path | None = None,
+        low_path: str | Path | None = None,
+        scale: float = 1.0,
+    ) -> CombinedAdapterHandle:
+        """Load a LoRA adapter for one or both experts.
+
+        Both ``high_path`` and ``low_path`` are independently optional — pass
+        either or both. Adapter weights are parsed from disk and stored on the
+        host LoRA bank of every applicable Linear in the respective transformer;
+        device-side merge happens later in ``set_active_lora`` (fuse mode) or
+        in the layer's forward (runtime mode). Device-OOM therefore surfaces
+        at ``set_active_lora`` time, not here.
+        """
+        if not high_path and not low_path:
+            raise ValueError("register_lora requires high_path and/or low_path")
+        if name in self._lora_registry:
+            raise ValueError(f"LoRA name '{name}' already registered")
+
+        high_handle = None
+        low_handle = None
+
+        if high_path is not None:
+            logger.info(f"loading LoRA '{name}' high-noise from {high_path} (scale={scale})")
+            high_handle = load_adapter_into(self.transformer, str(high_path), scale=scale, name=name)
+        if low_path is not None:
+            logger.info(f"loading LoRA '{name}' low-noise from {low_path} (scale={scale})")
+            low_handle = load_adapter_into(self.transformer_2, str(low_path), scale=scale, name=name)
+
+        combined = CombinedAdapterHandle(name=name, high=high_handle, low=low_handle)
+        self._lora_registry[name] = _RegisteredLoRA(name=name, handle_high=high_handle, handle_low=low_handle)
+        return combined
+
+    def unregister_lora(self, handle: CombinedAdapterHandle | str) -> None:
+        name = handle if isinstance(handle, str) else handle.name
+        rec = self._lora_registry.pop(name, None)
+        if rec is None:
+            return
+        if self._active_adapter_name == name:
+            self.set_active_lora(None)
+        # Walk and unregister bank entries
+        for transformer, handle_ in (
+            (self.transformer, rec.handle_high),
+            (self.transformer_2, rec.handle_low),
+        ):
+            if handle_ is None:
+                continue
+            self._unregister_handle(transformer, handle_)
+
+    def set_active_lora(self, handle: CombinedAdapterHandle | str | None) -> None:
+        """Bind the named adapter as the active LoRA for both experts.
+
+        Each LoRA Linear's bind merges the new delta into the base weight
+        in-place. If a different adapter is currently active, its delta is
+        subtracted first. Passing ``None`` removes the active delta (W is
+        restored to its base value).
+        """
+        if handle is None:
+            self._unbind_all()
+            self._active_adapter_name = None
+            return
+
+        name = handle if isinstance(handle, str) else handle.name
+        rec = self._lora_registry.get(name)
+        if rec is None:
+            raise KeyError(f"no LoRA registered with name '{name}'")
+
+        self._bind_per_module(self.transformer, rec.handle_high)
+        self._bind_per_module(self.transformer_2, rec.handle_low)
+        self._active_adapter_name = name
+
+    def list_loras(self) -> list[str]:
+        return list(self._lora_registry)
+
+    def deallocate_all_loras(self) -> None:
+        """Unbind every active adapter and clear the LoRA bank on every layer.
+
+        Releases host-side adapter records across both transformers. Use this
+        on a long-running server to reset state between unrelated jobs.
+        """
+        self._unbind_all()
+        self._active_adapter_name = None
+        self._lora_registry.clear()
+        for transformer in (self.transformer, self.transformer_2):
+            for mod in _iter_lora_modules(transformer):
+                mod.deallocate_lora()
+
+    # ---- internals ----
+    def _path_to_module(self, transformer, dotted: str):
+        """Resolve a dotted path like ``blocks.0.attn1.to_qkv`` to a child module."""
+        cur = transformer
+        for part in dotted.split("."):
+            if part.isdigit():
+                cur = cur[int(part)]
+            else:
+                cur = getattr(cur, part)
+        return cur
+
+    def _bind_per_module(self, transformer, handle: AdapterHandle | None):
+        addressed: set[int] = set()
+        if handle is not None:
+            for dotted, bank_idx in handle.target_indices.items():
+                mod = self._path_to_module(transformer, dotted)
+                mod.bind_active(bank_idx)
+                addressed.add(id(mod))
+
+        # Any LoRA module not touched by this adapter (including the all-orphan
+        # case where handle is None, e.g. switching to an adapter that only
+        # covers the other expert) but currently active from a previous bind
+        # must have its delta removed.
+        for mod in _iter_lora_modules(transformer):
+            if id(mod) in addressed:
+                continue
+            if mod.is_lora_active:
+                mod.unbind_active()
+
+    def _unbind_all(self):
+        for transformer in (self.transformer, self.transformer_2):
+            for mod in _iter_lora_modules(transformer):
+                if mod.is_lora_active:
+                    mod.unbind_active()
+
+    def _unregister_handle(self, transformer, handle: AdapterHandle):
+        for dotted, bank_idx in handle.target_indices.items():
+            mod = self._path_to_module(transformer, dotted)
+            mod.unregister_lora(bank_idx)
+
+
+# --------------------------------------------------------------------
+# Public pipeline classes
+# --------------------------------------------------------------------
+class WanPipelineRuntimeLoRA(_LoRAPipelineMixin, WanPipeline):
+    """T2V pipeline with runtime LoRA support."""
+
+
+class WanPipelineI2VRuntimeLoRA(_LoRAPipelineMixin, WanPipelineI2V):
+    """I2V pipeline with runtime LoRA support."""

code/models/tt_dit/experimental/pipelines/pipeline_wan_svi.py

@@ -0,0 +1,399 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Stable-Video-Infinity (SVI) 2.0 Pro pipeline for Wan2.2 I2V.
+
+Chains short I2V clips into long videos with latent-space continuity.
+See ``experimental/models/Wan2_2_SVI.md`` for the regime parameters,
+upstream-workflow comparison, and the documented scheduler gap.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+import torch
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
+from loguru import logger
+from PIL import Image
+
+import ttnn
+from models.tt_dit.experimental.pipelines.pipeline_wan_lora import LoRASpec, WanPipelineI2VLora
+from models.tt_dit.pipelines.wan.pipeline_wan import WanPipeline, WanPipelineConfig
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import ImagePrompt
+
+Regime = Literal["python", "comfyui"]
+
+
+_REGIMES = {
+    "python": {
+        "flow_shift": 5.0,
+        "call_defaults": dict(num_inference_steps=50, guidance_scale=5.0, guidance_scale_2=5.0),
+    },
+    "comfyui": {
+        "flow_shift": 8.0,
+        "call_defaults": dict(num_inference_steps=6, guidance_scale=1.5, guidance_scale_2=1.5),
+    },
+}
+
+
+@dataclass(frozen=True)
+class _ClipSpec:
+    prompt: str
+    seed: int
+    guidance_scale: float
+    guidance_scale_2: float
+
+
+class WanPipelineSVI(WanPipelineI2VLora):
+    """Wan2.2 I2V with SVI 2.0 Pro LoRA + autoregressive clip-chaining driver."""
+
+    def __init__(
+        self,
+        *,
+        device: ttnn.MeshDevice,
+        config: WanPipelineConfig,
+        svi_high: str,
+        svi_low: str,
+        lightx2v_high: Optional[str] = None,
+        lightx2v_low: Optional[str] = None,
+        regime: Regime = "python",
+        num_motion_latent: int = 1,
+        num_overlap_frame: int = 4,
+    ) -> None:
+        if regime not in _REGIMES:
+            raise ValueError(f"regime must be one of {sorted(_REGIMES)}, got {regime!r}")
+
+        for label, path in [("svi_high", svi_high), ("svi_low", svi_low)]:
+            if not Path(path).is_file():
+                raise FileNotFoundError(f"{label}: file does not exist: {path}")
+
+        self._regime: Regime = regime
+        self._num_motion_latent = int(num_motion_latent)
+        self._num_overlap_frame = int(num_overlap_frame)
+        # Lives for the duration of one generate_long_video call; cleared in
+        # that method's finally block.
+        self._anchor_latents_cache: Optional[tuple] = None
+
+        lora_high, lora_low, scheduler = self._configure_regime(
+            regime=regime,
+            svi_high=svi_high,
+            svi_low=svi_low,
+            lightx2v_high=lightx2v_high,
+            lightx2v_low=lightx2v_low,
+        )
+
+        super().__init__(
+            device=device,
+            config=config,
+            lora_high=lora_high,
+            lora_low=lora_low,
+            scheduler=scheduler,
+        )
+
+    @staticmethod
+    def _configure_regime(
+        *,
+        regime: "Regime",
+        svi_high: str,
+        svi_low: str,
+        lightx2v_high: Optional[str],
+        lightx2v_low: Optional[str],
+    ) -> tuple[list[LoRASpec], list[LoRASpec], object]:
+        """Return ``(lora_high, lora_low, scheduler)`` per regime."""
+        shift = _REGIMES[regime]["flow_shift"]
+
+        if regime == "python":
+            if lightx2v_high is not None or lightx2v_low is not None:
+                raise ValueError("lightx2v_* are only valid in regime='comfyui'; drop them for 'python'.")
+            # Matches diffsynth's FlowMatchScheduler("Wan") — plain Euler on a
+            # flow-matching schedule, dispatched to EulerSolver downstream.
+            return (
+                [LoRASpec(svi_high, 1.0)],
+                [LoRASpec(svi_low, 1.0)],
+                FlowMatchEulerDiscreteScheduler(shift=shift),
+            )
+
+        # comfyui regime: requires LightX2V LoRAs stacked under SVI.
+        if lightx2v_high is None or lightx2v_low is None:
+            raise ValueError(
+                "regime='comfyui' requires lightx2v_high and lightx2v_low LoRA paths. "
+                "See experimental/models/Wan2_2_SVI.md for the recommended HF repo."
+            )
+        for label, path in [("lightx2v_high", lightx2v_high), ("lightx2v_low", lightx2v_low)]:
+            if not Path(path).is_file():
+                raise FileNotFoundError(f"{label}: file does not exist: {path}")
+        # Per upstream docs/svi/comfyui.md: LightX2V at 1.0 on the high-noise
+        # expert hurts SVI; 0.5 is recommended.
+        # ComfyUI workflow uses k-diffusion dpm++_sde; tt-metal has no
+        # on-device port, so we substitute UniPC at the same flow_shift.
+        # Visually close, not bit-exact.
+        return (
+            [LoRASpec(lightx2v_high, 0.5), LoRASpec(svi_high, 1.0)],
+            [LoRASpec(lightx2v_low, 1.0), LoRASpec(svi_low, 1.0)],
+            UniPCMultistepScheduler(
+                use_flow_sigmas=True,
+                prediction_type="flow_prediction",
+                flow_shift=shift,
+            ),
+        )
+
+    def __call__(
+        self,
+        *args,
+        anchor_image: Optional[Image.Image] = None,
+        prev_last_latent: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        self._svi_anchor_image = anchor_image
+        self._svi_prev_last_latent = prev_last_latent
+        try:
+            return super().__call__(*args, **kwargs)
+        finally:
+            self._svi_anchor_image = None
+            self._svi_prev_last_latent = None
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        image_prompt,
+        num_channels_latents: int = 16,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 81,
+        dtype=None,
+        device=None,
+        anchor_image: Optional[Image.Image] = None,
+        prev_last_latent: Optional[torch.Tensor] = None,
+    ):
+        if anchor_image is None:
+            anchor_image = getattr(self, "_svi_anchor_image", None)
+        if prev_last_latent is None:
+            prev_last_latent = getattr(self, "_svi_prev_last_latent", None)
+
+        # WanPipelineI2V.__init__'s warmup_buffers invokes __call__ with no
+        # SVI state — defer to the plain I2V path for that one pass.
+        if anchor_image is None and prev_last_latent is None:
+            return super().prepare_latents(
+                batch_size=batch_size,
+                image_prompt=image_prompt,
+                num_channels_latents=num_channels_latents,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                dtype=dtype,
+                device=device,
+            )
+
+        anchor_pil = _as_anchor_image(anchor_image if anchor_image is not None else image_prompt)
+
+        # Cache the encoded anchor + mask across clips; the encoder is the
+        # dominant non-DiT cost per clip (~3-8s on BH 2x4) and the result is
+        # invariant under (anchor, H, W, T, dtype).
+        cache_key = (id(anchor_pil), height, width, num_frames, num_channels_latents, dtype)
+        cached = self._anchor_latents_cache
+        if cached is not None and cached[0] == cache_key:
+            tt_y_template = cached[1]
+            latents, _ = WanPipeline.prepare_latents(
+                self,
+                batch_size=batch_size,
+                num_channels_latents=num_channels_latents,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                dtype=dtype,
+                device=device,
+            )
+            tt_y = tt_y_template.clone()
+        else:
+            latents, tt_y = super().prepare_latents(
+                batch_size=batch_size,
+                image_prompt=[ImagePrompt(image=anchor_pil, frame_pos=0)],
+                num_channels_latents=num_channels_latents,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                dtype=dtype,
+                device=device,
+            )
+            # Cloned because the motion-latent splice below mutates tt_y in place.
+            self._anchor_latents_cache = (cache_key, tt_y.clone())
+
+        # tt_y layout: (B, 4 mask channels + 16 image-latent channels, T_lat, H_lat, W_lat).
+        t_lat = tt_y.shape[2]
+        n = 0
+        if prev_last_latent is not None and self._num_motion_latent > 0:
+            n = min(self._num_motion_latent, prev_last_latent.shape[2], t_lat - 1)
+            if n > 0:
+                motion = prev_last_latent[:, :, -n:, :, :].to(tt_y.device, tt_y.dtype)
+                if motion.shape[-2:] != tt_y.shape[-2:]:
+                    logger.warning(
+                        f"prev_last_latent spatial shape {tuple(motion.shape[-2:])} != "
+                        f"conditioning {tuple(tt_y.shape[-2:])}; skipping motion splice."
+                    )
+                    n = 0
+                else:
+                    tt_y[:, 4:, 1 : 1 + n, :, :] = motion
+
+        if 1 + n < t_lat:
+            tt_y[:, 4:, 1 + n :, :, :] = 0.0
+
+        return latents, tt_y
+
+    def generate_long_video(
+        self,
+        *,
+        prompt: Union[str, List[str]],
+        num_clips: int,
+        anchor_image: Optional[Image.Image] = None,
+        image_prompt=None,
+        base_seed: int = 0,
+        seed_stride: int = 42,
+        partial_output_path: Optional[str] = None,
+        **call_kwargs,
+    ):
+        """Run ``num_clips`` chained I2V generations and return the concat.
+
+        ``anchor_image`` is the persistent SVI reference frame. ``image_prompt``
+        is accepted for compatibility with the base I2V test harness.
+        """
+        if num_clips < 1:
+            raise ValueError("num_clips must be >= 1")
+        call_kwargs.setdefault("output_type", "pt_with_last_latent")
+
+        defaults = _REGIMES[self._regime]["call_defaults"]
+        for k, v in defaults.items():
+            call_kwargs.setdefault(k, v)
+
+        anchor_pil = _as_anchor_image(anchor_image if anchor_image is not None else image_prompt)
+        clip_specs = _make_clip_specs(
+            prompt=prompt,
+            guidance_scale=call_kwargs.pop("guidance_scale"),
+            guidance_scale_2=call_kwargs.pop("guidance_scale_2"),
+            num_clips=num_clips,
+            base_seed=base_seed,
+            seed_stride=seed_stride,
+        )
+
+        try:
+            return self._generate_clips_loop(
+                clip_specs=clip_specs,
+                anchor_image=anchor_pil,
+                partial_output_path=partial_output_path,
+                call_kwargs=call_kwargs,
+            )
+        finally:
+            # Drop the anchor-encode cache so the next generate_long_video
+            # call can re-encode with a different anchor / shape.
+            self._anchor_latents_cache = None
+
+    def _generate_clips_loop(
+        self,
+        *,
+        clip_specs: List[_ClipSpec],
+        anchor_image: Image.Image,
+        partial_output_path,
+        call_kwargs,
+    ):
+        prev_last_latent: Optional[torch.Tensor] = None
+        clips: List[torch.Tensor] = []
+
+        for clip_idx, spec in enumerate(clip_specs):
+            logger.info(
+                f"SVI clip {clip_idx + 1}/{len(clip_specs)} (seed={spec.seed}, "
+                f"num_motion_latent={self._num_motion_latent}, "
+                f"cfg={spec.guidance_scale}/{spec.guidance_scale_2}, "
+                f"prompt={spec.prompt[:80]!r}...)"
+            )
+
+            result = self(
+                prompts=[spec.prompt],
+                image_prompt=[ImagePrompt(image=anchor_image, frame_pos=0)],
+                seed=spec.seed,
+                guidance_scale=spec.guidance_scale,
+                guidance_scale_2=spec.guidance_scale_2,
+                anchor_image=anchor_image,
+                prev_last_latent=prev_last_latent,
+                **call_kwargs,
+            )
+
+            if not (isinstance(result, tuple) and len(result) == 2):
+                raise RuntimeError("SVI requires output_type='pt_with_last_latent' from the inner pipeline.")
+            frames, prev_last_latent = result
+
+            clips.append(frames[0] if frames.ndim == 5 else frames)
+
+            # Save partial concat so a mid-run hang still leaves a usable
+            # video covering the clips that completed.
+            if partial_output_path is not None:
+                partial = _concat_with_overlap(clips, overlap=self._num_overlap_frame)
+                torch.save(partial.detach(), partial_output_path)
+                logger.info(
+                    f"saved partial output after clip {clip_idx + 1}/{len(clip_specs)} "
+                    f"({partial.shape[0]} frames) to {partial_output_path}"
+                )
+
+        return _concat_with_overlap(clips, overlap=self._num_overlap_frame)
+
+
+def _make_clip_specs(
+    *,
+    prompt: Union[str, List[str]],
+    guidance_scale,
+    guidance_scale_2,
+    num_clips: int,
+    base_seed: int,
+    seed_stride: int,
+) -> List[_ClipSpec]:
+    prompts = _broadcast_prompt(prompt, num_clips)
+    cfgs = _broadcast_float(guidance_scale, num_clips, "guidance_scale")
+    cfg2s = _broadcast_float(guidance_scale_2, num_clips, "guidance_scale_2")
+    return [
+        _ClipSpec(
+            prompt=prompts[i],
+            seed=base_seed + i * seed_stride,
+            guidance_scale=cfgs[i],
+            guidance_scale_2=cfg2s[i],
+        )
+        for i in range(num_clips)
+    ]
+
+
+def _broadcast_prompt(prompt: Union[str, List[str]], num_clips: int) -> List[str]:
+    if isinstance(prompt, list):
+        if len(prompt) != num_clips:
+            raise ValueError(f"prompt list length {len(prompt)} != num_clips {num_clips}")
+        return prompt
+    return [prompt] * num_clips
+
+
+def _broadcast_float(value, num_clips: int, name: str) -> List[float]:
+    if isinstance(value, (list, tuple)):
+        if len(value) != num_clips:
+            raise ValueError(f"{name} list length {len(value)} != num_clips {num_clips}")
+        return [float(v) for v in value]
+    return [float(value)] * num_clips
+
+
+def _as_anchor_image(value) -> Image.Image:
+    """Accept a PIL image or single ImagePrompt/list wrapper from the I2V API."""
+    if isinstance(value, list):
+        if len(value) != 1:
+            raise ValueError("SVI supports a single anchor image; got list of length != 1")
+        return _as_anchor_image(value[0])
+    if isinstance(value, ImagePrompt):
+        return value.image
+    if not isinstance(value, Image.Image):
+        raise TypeError("SVI requires anchor_image or image_prompt to be a PIL image")
+    return value
+
+
+def _concat_with_overlap(clips: List[torch.Tensor], *, overlap: int) -> torch.Tensor:
+    """Temporal-concat decoded clips, dropping ``overlap`` boundary frames between adjacent clips."""
+    if not clips:
+        raise ValueError("no clips to concatenate")
+    pieces = [clips[0]] + [c[overlap:] if overlap > 0 else c for c in clips[1:]]
+    return torch.cat(pieces, dim=0)

code/models/tt_dit/experimental/tests/conftest.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import torch
+from loguru import logger
+
+num_torch_threads = max(1, os.cpu_count())
+logger.info(f"Setting torch num_threads to {num_torch_threads}")
+torch.set_num_threads(num_torch_threads)

code/models/tt_dit/experimental/tests/test_adapter_loader.py

@@ -0,0 +1,224 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Integration tests for the runtime-LoRA adapter loader.
+
+Builds a tiny WanTransformer3DModel (1 block) with ``lora_enabled=True``,
+writes a synthetic LoRA safetensors file, runs ``load_adapter_into``, and
+checks that:
+
+  - every LoRA-aware Linear in the block got an adapter registered
+  - ``bind_active`` marks the right modules active across all paths
+  - ``unbind_active`` clears the active state
+
+The forward math is covered by ``test_lora_variants.py``; this file
+exercises the loader's key-mapping logic and the pipeline-mixin bind
+flow against the default (``fuse``) mode.
+
+Run with:
+    pytest -xvs models/tt_dit/experimental/tests/test_adapter_loader.py
+"""
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+
+import pytest
+import torch
+
+import ttnn
+from models.tt_dit.experimental.lora.adapter_loader import load_adapter_into
+from models.tt_dit.layers.lora import LoRAMixin
+from models.tt_dit.models.transformers.wan2_2.transformer_wan import WanTransformer3DModel
+from models.tt_dit.parallel.config import DiTParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+
+
+def _save_synthetic_adapter(
+    path: Path,
+    *,
+    num_blocks: int = 1,
+    dim: int = 128,
+    ffn_dim: int = 256,
+    rank: int = 8,
+) -> None:
+    """Write a fake LoRA safetensors that covers all 6 LoRA-targeted Linears
+    per block, using the lightx2v naming convention."""
+    from safetensors.torch import save_file
+
+    state: dict[str, torch.Tensor] = {}
+    dtype = torch.bfloat16
+
+    def add_pair(base: str, in_dim: int, out_dim: int):
+        state[f"{base}.lora_A.weight"] = torch.randn(rank, in_dim, dtype=dtype) * 0.02
+        state[f"{base}.lora_B.weight"] = torch.randn(out_dim, rank, dtype=dtype) * 0.02
+
+    for i in range(num_blocks):
+        # self-attn Q, K, V, O
+        add_pair(f"blocks.{i}.self_attn.q", dim, dim)
+        add_pair(f"blocks.{i}.self_attn.k", dim, dim)
+        add_pair(f"blocks.{i}.self_attn.v", dim, dim)
+        add_pair(f"blocks.{i}.self_attn.o", dim, dim)
+        # cross-attn Q, K, V, O
+        add_pair(f"blocks.{i}.cross_attn.q", dim, dim)
+        add_pair(f"blocks.{i}.cross_attn.k", dim, dim)
+        add_pair(f"blocks.{i}.cross_attn.v", dim, dim)
+        add_pair(f"blocks.{i}.cross_attn.o", dim, dim)
+        # ffn ff1, ff2 (lightx2v naming: ffn.0 = ff1, ffn.2 = ff2)
+        add_pair(f"blocks.{i}.ffn.0", dim, ffn_dim)
+        add_pair(f"blocks.{i}.ffn.2", ffn_dim, dim)
+
+    save_file(state, str(path))
+
+
+def _count_registered(transformer) -> int:
+    """Count LoRA modules with at least one entry in their bank."""
+    n = 0
+    for mod in _iter_lora_modules(transformer):
+        if any(a is not None for a in mod.lora_bank):
+            n += 1
+    return n
+
+
+def _iter_lora_modules(module):
+    if isinstance(module, LoRAMixin):
+        yield module
+    for _, child in module.named_children():
+        yield from _iter_lora_modules(child)
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+def test_adapter_loader_registers_all_lora_linears(mesh_device: ttnn.MeshDevice) -> None:
+    """Loader picks up Q/K/V/O for both attn1/attn2 plus ff1/ff2 per block."""
+    num_heads = 4
+    head_dim = 32
+    dim = num_heads * head_dim
+    ffn_dim = 2 * dim
+    rank = 8
+    num_blocks = 1
+
+    parallel_config = DiTParallelConfig(
+        cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+        tensor_parallel=ParallelFactor(factor=1, mesh_axis=1),
+        sequence_parallel=ParallelFactor(factor=1, mesh_axis=0),
+    )
+    ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+
+    transformer = WanTransformer3DModel(
+        patch_size=(1, 2, 2),
+        num_heads=num_heads,
+        dim=dim,
+        in_channels=16,
+        out_channels=16,
+        text_dim=128,
+        freq_dim=64,
+        ffn_dim=ffn_dim,
+        num_layers=num_blocks,
+        mesh_device=mesh_device,
+        ccl_manager=ccl_manager,
+        parallel_config=parallel_config,
+        is_fsdp=False,
+        lora_enabled=True,
+    )
+
+    # Every block should have 5 LoRA-aware Linears wired in: to_qkv (self),
+    # to_q (cross), to_kv (cross), to_out for each attention (×2), ff1, ff2.
+    # That's 1 + 1 + 1 + 2 + 1 + 1 = 7 modules per block.
+    expected_per_block = 7
+    n_modules = sum(1 for _ in _iter_lora_modules(transformer))
+    assert (
+        n_modules == num_blocks * expected_per_block
+    ), f"expected {num_blocks * expected_per_block} LoRA modules; got {n_modules}"
+
+    with tempfile.TemporaryDirectory() as tmp:
+        adapter_path = Path(tmp) / "synthetic.safetensors"
+        _save_synthetic_adapter(adapter_path, num_blocks=num_blocks, dim=dim, ffn_dim=ffn_dim, rank=rank)
+
+        handle = load_adapter_into(transformer, str(adapter_path), scale=1.0, name="synthetic")
+
+    # All 7 LoRA modules per block should have a registered adapter.
+    registered = _count_registered(transformer)
+    assert (
+        registered == num_blocks * expected_per_block
+    ), f"expected {num_blocks * expected_per_block} registered adapters; got {registered}"
+
+    # Fused-QKV (to_qkv self-attn): rank should be 3*r. Fused KV (cross-attn
+    # to_kv): rank should be 2*r. Singletons: rank = r.
+    block0 = transformer.blocks[0]
+    assert block0.attn1.to_qkv.lora_bank[0].rank == 3 * rank
+    assert block0.attn2.to_kv.lora_bank[0].rank == 2 * rank
+    assert block0.attn2.to_q.lora_bank[0].rank == rank
+    assert block0.attn1.to_out.lora_bank[0].rank == rank
+    assert block0.attn2.to_out.lora_bank[0].rank == rank
+    assert block0.ffn.ff1.lora_bank[0].rank == rank
+    assert block0.ffn.ff2.lora_bank[0].rank == rank
+
+    # handle.rank tracks the max rank across targets (3*rank for the fused-QKV group).
+    assert handle.rank == 3 * rank
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+def test_pipeline_bind_unbind_walks_all_modules(mesh_device: ttnn.MeshDevice) -> None:
+    """The pipeline mixin's bind/unbind walks every LoRA module on each set_active_lora."""
+    from models.tt_dit.experimental.pipelines.pipeline_wan_runtime_lora import _iter_lora_modules as pipeline_iter
+
+    num_heads = 4
+    head_dim = 32
+    dim = num_heads * head_dim
+    ffn_dim = 2 * dim
+    rank = 8
+
+    parallel_config = DiTParallelConfig(
+        cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+        tensor_parallel=ParallelFactor(factor=1, mesh_axis=1),
+        sequence_parallel=ParallelFactor(factor=1, mesh_axis=0),
+    )
+    ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+
+    transformer = WanTransformer3DModel(
+        patch_size=(1, 2, 2),
+        num_heads=num_heads,
+        dim=dim,
+        in_channels=16,
+        out_channels=16,
+        text_dim=128,
+        freq_dim=64,
+        ffn_dim=ffn_dim,
+        num_layers=1,
+        mesh_device=mesh_device,
+        ccl_manager=ccl_manager,
+        parallel_config=parallel_config,
+        is_fsdp=False,
+        lora_enabled=True,
+    )
+
+    with tempfile.TemporaryDirectory() as tmp:
+        adapter_path = Path(tmp) / "synthetic.safetensors"
+        _save_synthetic_adapter(adapter_path, num_blocks=1, dim=dim, ffn_dim=ffn_dim, rank=rank)
+        handle = load_adapter_into(transformer, str(adapter_path), name="synthetic")
+
+    # Simulate the bind path the pipeline mixin runs.
+    for dotted, bank_idx in handle.target_indices.items():
+        cur = transformer
+        for part in dotted.split("."):
+            cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
+        cur.bind_active(bank_idx)
+
+    # All modules covered by the handle should be active (delta merged).
+    for dotted in handle.target_indices:
+        cur = transformer
+        for part in dotted.split("."):
+            cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
+        assert cur.is_lora_active, f"{dotted} not active after bind_active"
+
+    # Unbind everything → delta subtracted, is_lora_active False.
+    for mod in pipeline_iter(transformer):
+        if mod.is_lora_active:
+            mod.unbind_active()
+    for dotted in handle.target_indices:
+        cur = transformer
+        for part in dotted.split("."):
+            cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
+        assert not cur.is_lora_active
+        assert cur.active_idx is None

code/models/tt_dit/experimental/tests/test_lora_variants.py

@@ -0,0 +1,262 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Parity tests for the three LoRA-aware Linear subclasses, each exercised
+under both ``lora_mode`` options:
+
+  fuse     — bind merges ``scale * A.T @ B.T`` into the base weight in-place.
+             Forward runs the un-modified base path.
+  runtime  — bind uploads A, B to device; forward adds
+             ``scale * (x @ A) @ B`` on top of the base output.
+
+The unbind branch is checked in both modes — for ``fuse`` it subtracts
+the delta back out of W; for ``runtime`` it frees the A/B tensors so the
+forward path collapses to the base.
+
+Mesh shape is 1x1 so sharding collapses to the replicated case for all
+three variants. The tests still exercise the LoRA-specific upload and
+forward code paths and catch class-level bugs in shape handling,
+fused-activation fallback, and chunked-output handling.
+
+Run with:
+    pytest -xvs models/tt_dit/experimental/tests/test_lora_variants.py
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+import ttnn
+from models.tt_dit.layers.linear import LoRAColParallelLinear, LoRALinear, LoRARowParallelLinear
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.tensor import bf16_tensor
+
+
+def _make_pair(in_features: int, out_features: int, rank: int, dtype: torch.dtype):
+    A = torch.randn(rank, in_features, dtype=dtype) * 0.02
+    B = torch.randn(out_features, rank, dtype=dtype) * 0.02
+    return A, B
+
+
+def _ref_lora(x, W, b, A, B, scale):
+    """y = x @ (W + scale * B @ A).T + b"""
+    W_eff = W + scale * (B @ A)
+    return torch.nn.functional.linear(x, W_eff, b)
+
+
+# --------------------------------------------------------------------
+# LoRALinear (replicated)
+# --------------------------------------------------------------------
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+@pytest.mark.parametrize("lora_mode", ["fuse", "runtime"])
+@pytest.mark.parametrize(
+    ("seq_len", "in_features", "out_features", "rank"),
+    [
+        (128, 256, 256, 8),
+        (256, 512, 1024, 16),
+    ],
+)
+def test_lora_linear_variant(
+    mesh_device: ttnn.MeshDevice,
+    lora_mode: str,
+    seq_len: int,
+    in_features: int,
+    out_features: int,
+    rank: int,
+) -> None:
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    scale = 0.6
+
+    torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
+    torch_lin.eval()
+    A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
+    x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)
+
+    with torch.no_grad():
+        y_base = torch_lin(x_t)
+        y_lora = _ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale)
+
+    tt = LoRALinear(in_features, out_features, bias=True, mesh_device=mesh_device, lora_mode=lora_mode)
+    tt.load_torch_state_dict(torch_lin.state_dict())
+    idx = tt.register_lora(A_t, B_t, scale=scale)
+    x_tt = bf16_tensor(x_t, device=mesh_device)
+
+    # 1) no adapter bound → base behavior
+    y = tt(x_tt)
+    for t in ttnn.get_device_tensors(y):
+        assert_quality(y_base, ttnn.to_torch(t), pcc=0.999_5)
+
+    # 2) bind active → LoRA effect applied
+    tt.bind_active(idx, scale=scale)
+    y = tt(x_tt)
+    for t in ttnn.get_device_tensors(y):
+        assert_quality(y_lora, ttnn.to_torch(t), pcc=0.999_0)
+
+    # 3) unbind collapses back to base
+    tt.unbind_active()
+    y = tt(x_tt)
+    for t in ttnn.get_device_tensors(y):
+        assert_quality(y_base, ttnn.to_torch(t), pcc=0.999_5)
+
+    tt.deallocate_lora()
+
+
+# --------------------------------------------------------------------
+# LoRAColParallelLinear — single output and chunked output
+# --------------------------------------------------------------------
+# Chunked configs run fuse-only: the runtime delta path doesn't split
+# its delta into the same chunk layout as the base output.
+_COL_PARALLEL_PARAMS = [
+    # (label, lora_mode, seq, in, out, rank, chunks, activation_fn)
+    ("plain_fuse", "fuse", 128, 256, 512, 8, None, None),
+    ("plain_runtime", "runtime", 128, 256, 512, 8, None, None),
+    ("gelu_fuse", "fuse", 128, 256, 1024, 8, None, "gelu"),
+    ("gelu_runtime", "runtime", 128, 256, 1024, 8, None, "gelu"),
+    ("chunked_qkv_fuse", "fuse", 128, 256, 3 * 256, 8, 3, None),
+]
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+@pytest.mark.parametrize(
+    ("lora_mode", "seq_len", "in_features", "out_features", "rank", "chunks", "activation_fn"),
+    [params[1:] for params in _COL_PARALLEL_PARAMS],
+    ids=[params[0] for params in _COL_PARALLEL_PARAMS],
+)
+def test_lora_col_parallel_variant(
+    mesh_device: ttnn.MeshDevice,
+    lora_mode: str,
+    seq_len: int,
+    in_features: int,
+    out_features: int,
+    rank: int,
+    chunks: int | None,
+    activation_fn: str | None,
+) -> None:
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    scale = 0.5
+
+    torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
+    torch_lin.eval()
+    A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
+    x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)
+
+    def torch_act(y):
+        if activation_fn == "gelu":
+            return torch.nn.functional.gelu(y)
+        return y
+
+    with torch.no_grad():
+        y_base = torch_act(torch_lin(x_t))
+        y_lora = torch_act(_ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale))
+
+    tt = LoRAColParallelLinear(
+        in_features,
+        out_features,
+        bias=True,
+        mesh_device=mesh_device,
+        mesh_axis=0,
+        activation_fn=activation_fn,
+        chunks=chunks,
+        lora_mode=lora_mode,
+    )
+    tt.load_torch_state_dict(torch_lin.state_dict())
+    idx = tt.register_lora(A_t, B_t, scale=scale)
+    x_tt = bf16_tensor(x_t, device=mesh_device)
+
+    def _join(out):
+        """Concat chunked outputs along last dim; pass through single tensors."""
+        if isinstance(out, list):
+            torch_chunks = [ttnn.to_torch(ttnn.get_device_tensors(o)[0]) for o in out]
+            return torch.cat(torch_chunks, dim=-1)
+        return ttnn.to_torch(ttnn.get_device_tensors(out)[0])
+
+    # 1) no adapter bound → base behavior
+    assert_quality(y_base, _join(tt(x_tt)), pcc=0.999_5)
+
+    # 2) bind active → LoRA effect applied
+    tt.bind_active(idx, scale=scale)
+    assert_quality(y_lora, _join(tt(x_tt)), pcc=0.999_0)
+
+    # 3) unbind collapses to base
+    tt.unbind_active()
+    assert_quality(y_base, _join(tt(x_tt)), pcc=0.999_0)
+
+    tt.deallocate_lora()
+
+
+# --------------------------------------------------------------------
+# LoRARowParallelLinear — N_tp=1 (1x1 mesh) sanity
+# --------------------------------------------------------------------
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+@pytest.mark.parametrize("lora_mode", ["fuse", "runtime"])
+@pytest.mark.parametrize(
+    ("seq_len", "in_features", "out_features", "rank"),
+    [
+        (128, 512, 256, 8),
+        (256, 1024, 512, 16),
+    ],
+)
+def test_lora_row_parallel_variant(
+    mesh_device: ttnn.MeshDevice,
+    lora_mode: str,
+    seq_len: int,
+    in_features: int,
+    out_features: int,
+    rank: int,
+) -> None:
+    from models.tt_dit.parallel.manager import CCLManager
+
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    scale = 0.4
+
+    torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
+    torch_lin.eval()
+    A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
+    x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)
+
+    with torch.no_grad():
+        y_base = torch_lin(x_t)
+        y_lora = _ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale)
+
+    ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+    tt = LoRARowParallelLinear(
+        in_features,
+        out_features,
+        bias=True,
+        mesh_device=mesh_device,
+        mesh_axis=0,
+        ccl_manager=ccl_manager,
+        lora_mode=lora_mode,
+    )
+    tt.load_torch_state_dict(torch_lin.state_dict())
+    idx = tt.register_lora(A_t, B_t, scale=scale)
+    x_tt = bf16_tensor(x_t, device=mesh_device)
+
+    # baseline
+    y = tt(x_tt)
+    assert_quality(y_base, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_5)
+
+    # bind active
+    tt.bind_active(idx, scale=scale)
+    y = tt(x_tt)
+    assert_quality(y_lora, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_0)
+
+    # unbind
+    tt.unbind_active()
+    y = tt(x_tt)
+    assert_quality(y_base, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_5)
+
+    tt.deallocate_lora()
+
+
+# --------------------------------------------------------------------
+# Construction-time validation
+# --------------------------------------------------------------------
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+def test_invalid_lora_mode_raises(mesh_device: ttnn.MeshDevice, expect_error) -> None:
+    with expect_error(ValueError, "lora_mode"):
+        LoRALinear(64, 64, bias=False, mesh_device=mesh_device, lora_mode="bogus")

code/models/tt_dit/experimental/tests/test_pipeline_anisora.py

@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+import os
+
+import numpy as np
+import PIL
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.experimental.pipelines.pipeline_anisora import AniSoraPipeline
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import ImagePrompt
+from models.tt_dit.utils.test import ring_params
+
+
+@pytest.mark.parametrize(
+    "no_prompt",
+    [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), False)],
+)
+@pytest.mark.parametrize(
+    "mesh_device, mesh_shape, num_links, dynamic_load, device_params, topology, is_fsdp",
+    [
+        # BH Galaxy 4x8 Ring — sole supported config in the first AniSora PR.
+        [(4, 8), (4, 8), 2, False, ring_params, ttnn.Topology.Ring, False],
+    ],
+    ids=["bh_4x8sp1tp0_ring"],
+    indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (832, 480),
+        (1280, 720),
+    ],
+    ids=[
+        "resolution_480p",
+        "resolution_720p",
+    ],
+)
+def test_pipeline_inference(
+    mesh_device,
+    mesh_shape,
+    num_links,
+    dynamic_load,
+    topology,
+    width,
+    height,
+    is_fsdp,
+    no_prompt,
+):
+    parent_mesh = mesh_device
+    mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+    pil_image = PIL.Image.open(os.environ.get("PROMPT_IMAGE", "./prompt_image.png"))
+    image_prompt = [ImagePrompt(image=pil_image, frame_pos=0)]
+    negative_prompt = ""
+
+    # AniSora upstream defaults (anisoraV3.2/wan/configs/wan_i2v_A14B.py):
+    #   sample_steps=40, boundary=0.9, sample_guide_scale=(3.5, 3.5).
+    num_frames = 81
+    num_inference_steps = int(os.environ.get("NUM_STEPS", "40"))
+    guidance_scale = 3.5
+    guidance_scale_2 = 3.5
+
+    pipeline = AniSoraPipeline.create_pipeline(
+        mesh_device=mesh_device,
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        config_overrides={
+            "num_links": num_links,
+            "dynamic_load": dynamic_load,
+            "topology": topology,
+            "is_fsdp": is_fsdp,
+        },
+    )
+
+    prompt = "An anime girl smiling, soft lighting, cinematic."
+
+    def run(*, prompt, number, seed):
+        logger.info(f"Running AniSora inference with prompt: '{prompt}'")
+        logger.info(f"Parameters: {height}x{width}, {num_frames} frames, {num_inference_steps} steps")
+
+        with torch.no_grad():
+            result = pipeline(
+                prompts=[prompt],
+                image_prompt=image_prompt,
+                negative_prompts=[negative_prompt],
+                num_inference_steps=num_inference_steps,
+                seed=seed,
+                guidance_scale=guidance_scale,
+                guidance_scale_2=guidance_scale_2,
+                output_type="uint8",
+            )
+
+        if hasattr(result, "frames"):
+            frames = result.frames
+        else:
+            frames = result[0] if isinstance(result, tuple) else result
+
+        logger.info("AniSora inference completed successfully")
+        logger.info(f"  Output shape: {frames.shape if hasattr(frames, 'shape') else 'Unknown'}")
+        logger.info(f"  Output type: {type(frames)}")
+
+        if isinstance(frames, np.ndarray):
+            logger.info(f"  Video data range: [{frames.min():.3f}, {frames.max():.3f}]")
+        elif isinstance(frames, torch.Tensor):
+            logger.info(f"  Video data range: [{frames.min().item():.3f}, {frames.max().item():.3f}]")
+
+        frames = frames[0]
+        output_filename = f"wan_anisora_i2v_{width}x{height}_{number}.mp4"
+        try:
+            from models.tt_dit.utils.video import export_to_video
+
+            export_to_video(frames, output_filename, fps=16)
+            logger.info(f"Saved video to: {output_filename}")
+        except ImportError:
+            logger.info("Could not export video - imageio_ffmpeg not available")
+
+    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)
+
+
+# ---------------------------------------------------------------------------
+# Random-weight smoke test
+# ---------------------------------------------------------------------------
+# Runs the full pipeline end-to-end without downloading the AniSora safetensors
+# or the two ~28 GB transformer subfolders from Wan-AI. Tokenizer / UMT5 text
+# encoder / VAE / scheduler still come from Wan-AI/Wan2.2-I2V-A14B-Diffusers
+# (~12 GB total) and require TT_DIT_ALLOW_HF_DOWNLOAD=1 the first time.
+# Output frames are garbage by design — this only validates compile / mesh /
+# shape / CCL plumbing, not numeric quality. Uses 4 sampling steps to keep the
+# smoke test short; the 40-step real-weights config is exercised above.
+
+
+@pytest.mark.parametrize(
+    "mesh_device, mesh_shape, num_links, dynamic_load, device_params, topology, is_fsdp",
+    [
+        [(4, 8), (4, 8), 2, False, ring_params, ttnn.Topology.Ring, False],
+    ],
+    ids=["bh_4x8sp1tp0_ring"],
+    indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (832, 480),
+    ],
+    ids=["resolution_480p"],
+)
+def test_pipeline_inference_random_weights(
+    mesh_device,
+    mesh_shape,
+    num_links,
+    dynamic_load,
+    topology,
+    width,
+    height,
+    is_fsdp,
+):
+    parent_mesh = mesh_device
+    mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+    pil_image = PIL.Image.open(os.environ.get("PROMPT_IMAGE", "./prompt_image.png"))
+    image_prompt = [ImagePrompt(image=pil_image, frame_pos=0)]
+
+    num_frames = 81
+    num_inference_steps = 4
+
+    os.environ["TT_DIT_RANDOM_WEIGHTS"] = "1"
+    try:
+        pipeline = AniSoraPipeline.create_pipeline(
+            mesh_device=mesh_device,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            config_overrides={
+                "num_links": num_links,
+                "dynamic_load": dynamic_load,
+                "topology": topology,
+                "is_fsdp": is_fsdp,
+            },
+        )
+    finally:
+        os.environ.pop("TT_DIT_RANDOM_WEIGHTS", None)
+
+    prompt = "smoke test — random weights, output is meaningless"
+    logger.info(f"Random-weight smoke test: {height}x{width}, {num_frames} frames, {num_inference_steps} steps")
+
+    with torch.no_grad():
+        result = pipeline(
+            prompts=[prompt],
+            image_prompt=image_prompt,
+            negative_prompts=[""],
+            num_inference_steps=num_inference_steps,
+            seed=42,
+            guidance_scale=3.5,
+            guidance_scale_2=3.5,
+            output_type="uint8",
+        )
+
+    if hasattr(result, "frames"):
+        frames = result.frames
+    else:
+        frames = result[0] if isinstance(result, tuple) else result
+
+    if isinstance(frames, np.ndarray):
+        rng_str = f"[{frames.min():.3f}, {frames.max():.3f}]"
+    elif isinstance(frames, torch.Tensor):
+        rng_str = f"[{frames.min().item():.3f}, {frames.max().item():.3f}]"
+    else:
+        rng_str = "n/a"
+
+    shape = frames.shape if hasattr(frames, "shape") else "Unknown"
+    logger.info(f"Random-weight smoke test PASSED. Output shape={shape}, range={rng_str}")

code/models/tt_dit/experimental/tests/test_pipeline_lora.py

@@ -0,0 +1,300 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""End-to-end test for Wan2.2 I2V with LoRA adapters fused into the base.
+
+Reads adapter paths via env vars (compatible with the previous experimental
+test): ``LORA_HIGH_PATH``, ``LORA_LOW_PATH``, ``LORA_SCALE``. Set
+``LORA_STACK_HIGH`` and/or ``LORA_STACK_LOW`` to a comma-separated list of
+``path[:scale]`` entries to exercise multi-LoRA stacking. If both single-LoRA
+and stack env vars are set, the stack form wins.
+"""
+import itertools
+import os
+import statistics
+from typing import List, Tuple
+
+import numpy as np
+import PIL
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.perf.benchmarking_utils import BenchmarkProfiler
+from models.tt_dit.experimental.pipelines.pipeline_wan_lora import LoRASpec, WanPipelineI2VLora
+from models.tt_dit.pipelines.events import profiler_event_callback
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import ImagePrompt
+from models.tt_dit.utils.test import line_params, ring_params_8k
+
+# Only the 4x32 row traces, so only it needs the trace region carved out of DRAM.
+DEVICE_PARAMS = {"trace_region_size": 150000000}
+
+
+def _rank() -> int:
+    return int(ttnn.distributed_context_get_rank()) if ttnn.using_distributed_env() else 0
+
+
+def create_fractal_image(width: int, height: int) -> PIL.Image.Image:
+    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 PIL.Image.fromarray(np.dstack((img, np.roll(img, width // 10, 1), np.roll(img, height // 10, 0))), "RGB")
+
+
+def _parse_stack(env_val: str) -> List[LoRASpec]:
+    """Parse ``path[:scale],path[:scale]`` into a LoRASpec list."""
+    out: List[LoRASpec] = []
+    for entry in env_val.split(","):
+        entry = entry.strip()
+        if not entry:
+            continue
+        if ":" in entry:
+            path, scale_str = entry.rsplit(":", 1)
+            out.append(LoRASpec(path.strip(), float(scale_str)))
+        else:
+            out.append(LoRASpec(entry))
+    return out
+
+
+def _resolve_lora_args() -> Tuple[List[LoRASpec], List[LoRASpec], float]:
+    stack_high = os.environ.get("LORA_STACK_HIGH")
+    stack_low = os.environ.get("LORA_STACK_LOW")
+    single_high = os.environ.get("LORA_HIGH_PATH")
+    single_low = os.environ.get("LORA_LOW_PATH")
+    scale = float(os.environ.get("LORA_SCALE", "1.0"))
+
+    if stack_high or stack_low:
+        return (
+            _parse_stack(stack_high) if stack_high else [],
+            _parse_stack(stack_low) if stack_low else [],
+            scale,
+        )
+
+    high = [LoRASpec(single_high, scale)] if single_high else []
+    low = [LoRASpec(single_low, scale)] if single_low else []
+    return high, low, scale
+
+
+@pytest.mark.parametrize(
+    "no_prompt",
+    [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), False)],
+)
+@pytest.mark.parametrize(
+    "mesh_device, mesh_shape, num_links, dynamic_load, device_params, topology, is_fsdp",
+    [
+        [(2, 4), (2, 4), 2, True, line_params, ttnn.Topology.Linear, False],
+        [(4, 8), (4, 8), 2, False, ring_params_8k, ttnn.Topology.Ring, False],
+        # ring_params_8k adds the fabric router config test_performance_wan uses at 4x32.
+        # [(4, 32), (4, 32), 2, False, {**DEVICE_PARAMS, **ring_params}, ttnn.Topology.Ring, False],
+        [(4, 32), (4, 32), 2, False, {**DEVICE_PARAMS, **ring_params_8k}, ttnn.Topology.Ring, False],
+    ],
+    ids=["bh_2x4sp1tp0", "bh_4x8sp1tp0_ring", "bh_4x32sp1tp0_ring"],
+    indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (832, 480),
+        (1280, 720),
+    ],
+    ids=[
+        "resolution_480p",
+        "resolution_720p",
+    ],
+)
+def test_pipeline_inference(
+    mesh_device,
+    mesh_shape,
+    num_links,
+    dynamic_load,
+    topology,
+    width,
+    height,
+    is_fsdp,
+    no_prompt,
+):
+    lora_high, lora_low, scale = _resolve_lora_args()
+    if not lora_high and not lora_low:
+        pytest.skip(
+            "Set LORA_HIGH_PATH / LORA_LOW_PATH (single-LoRA) or "
+            "LORA_STACK_HIGH / LORA_STACK_LOW (multi-LoRA) to run."
+        )
+
+    # Trace on 4x32 only, matching test_performance_wan; TEST_TRACED overrides either way.
+    traced = True if mesh_shape == (4, 32) else False
+    logger.info(f"traced={traced} (mesh_shape={mesh_shape})")
+
+    parent_mesh = mesh_device
+    mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+    prompt_image_path = os.environ.get("PROMPT_IMAGE")
+    pil_image = PIL.Image.open(prompt_image_path) if prompt_image_path else create_fractal_image(width, height)
+    image_prompt = [ImagePrompt(image=pil_image, frame_pos=0)]
+    negative_prompt = ""
+
+    num_frames = 81
+    num_inference_steps = int(os.environ.get("NUM_STEPS", "4"))
+    guidance_scale = float(os.environ.get("GUIDANCE_SCALE", "1"))
+    guidance_scale_2 = float(os.environ.get("GUIDANCE_SCALE_2", str(guidance_scale)))
+    boundary_ratio = float(os.environ.get("BOUNDARY_RATIO", "0.875"))
+
+    # A scale of 1.0 is a no-op lerp, so the unconditional pass is pure waste: skip CFG
+    # entirely and halve the denoising cost. Guidance-distilled LoRAs want exactly this.
+    cfg_enabled = guidance_scale > 1 or guidance_scale_2 > 1
+    logger.info(f"cfg_enabled={cfg_enabled} (guidance_scale={guidance_scale}, guidance_scale_2={guidance_scale_2})")
+
+    pipeline = WanPipelineI2VLora(
+        device=mesh_device,
+        config=WanPipelineI2VLora.default_config(
+            mesh_device=mesh_device,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            cfg_enabled=cfg_enabled,
+            config_overrides={
+                "num_links": num_links,
+                "dynamic_load": dynamic_load,
+                "topology": topology,
+                "is_fsdp": is_fsdp,
+                "boundary_ratio": boundary_ratio,
+            },
+        ),
+        lora_high=lora_high if lora_high else None,
+        lora_low=lora_low if lora_low else None,
+    )
+
+    # prompt = os.environ.get("PROMPT", "A golden retriever running on a sandy beach, waves in the background")
+    prompt = os.environ.get("PROMPT", "A  cat running on a sandy beach, waves in the background")
+
+    profiler = BenchmarkProfiler()
+
+    def run(*, prompt, number, seed, profiler_iteration=None):
+        logger.info(f"Running LoRA inference with prompt: '{prompt}'")
+        logger.info(
+            f"Parameters: {height}x{width}, {num_frames} frames, {num_inference_steps} steps, "
+            f"scale={scale}, stack_high={len(lora_high)}, stack_low={len(lora_low)}"
+        )
+
+        call_kwargs = dict(
+            prompts=[prompt],
+            image_prompt=image_prompt,
+            negative_prompts=[negative_prompt],
+            num_inference_steps=num_inference_steps,
+            seed=seed,
+            guidance_scale=guidance_scale,
+            guidance_scale_2=guidance_scale_2,
+            output_type="uint8",
+            traced=traced,
+        )
+        if profiler_iteration is not None:
+            call_kwargs["on_event"] = profiler_event_callback(profiler, profiler_iteration)
+
+        with torch.no_grad():
+            result = pipeline(**call_kwargs)
+
+        if hasattr(result, "frames"):
+            frames = result.frames
+        else:
+            frames = result[0] if isinstance(result, tuple) else result
+
+        logger.info(f"  Output shape: {frames.shape if hasattr(frames, 'shape') else 'Unknown'}")
+        if isinstance(frames, np.ndarray):
+            logger.info(f"  Video data range: [{frames.min():.3f}, {frames.max():.3f}]")
+        elif isinstance(frames, torch.Tensor):
+            logger.info(f"  Video data range: [{frames.min().item():.3f}, {frames.max().item():.3f}]")
+
+        frames = frames[0]
+        # Every rank holds the same gathered video, so only rank 0 writes it out.
+        if _rank() == 0:
+            output_filename = f"wan_lora_i2v_{width}x{height}_{number}.mp4"
+            try:
+                from models.tt_dit.utils.video import export_to_video
+
+                export_to_video(frames, output_filename, fps=16)
+                logger.info(f"Saved video to: {output_filename}")
+            except ImportError:
+                logger.info("Could not export video - imageio_ffmpeg not available")
+        else:
+            logger.info(f"Skipping video export on rank {_rank()}")
+
+        return frames
+
+    if no_prompt:
+        # Warmup run (not included in stats)
+        logger.info("Running warmup iteration...")
+        with profiler("warmup", iteration=0):
+            run(prompt=prompt, number=-1, seed=0)
+        logger.info(f"Warmup completed in {profiler.get_duration('warmup', 0):.2f}s")
+
+        # Timed performance run
+        num_perf_runs = 1
+        ttnn.synchronize_device(mesh_device)
+        # Line up all hosts so the measured window is comparable across ranks.
+        ttnn.distributed_context_barrier()
+
+        for i in range(num_perf_runs):
+            logger.info(f"Performance run {i+1}/{num_perf_runs}...")
+            with profiler("run", iteration=i):
+                run(prompt=prompt, number=i, seed=42, profiler_iteration=i)
+            logger.info(f"  Run {i+1} completed in {profiler.get_duration('run', i):.2f}s")
+
+        # Collect timing stats. "prepare_latents" is the I2V image-conditioning
+        # section: host preprocess/upload of the conditioned frames plus the VAE encode.
+        encoder_times = [profiler.get_duration("encoder", i) for i in range(num_perf_runs)]
+        image_encode_times = [profiler.get_duration("prepare_latents", i) for i in range(num_perf_runs)]
+        denoising_times = [profiler.get_duration("denoising", i) for i in range(num_perf_runs)]
+        vae_times = [profiler.get_duration("vae", i) for i in range(num_perf_runs)]
+        total_times = [profiler.get_duration("run", i) for i in range(num_perf_runs)]
+
+        def fmt_stats(times):
+            if not times:
+                return "No data available"
+            mean = statistics.mean(times)
+            if len(times) > 1:
+                std = statistics.stdev(times)
+                return f"Mean: {mean:8.3f}s | Std: {std:7.3f}s | Min: {min(times):8.3f}s | Max: {max(times):8.3f}s"
+            return f"{mean:8.3f}s"
+
+        if _rank() == 0:
+            sp_factor = tuple(mesh_device.shape)[1]
+            tp_factor = tuple(mesh_device.shape)[0]
+            print("\n" + "=" * 80)
+            print("WAN 2.2 I2V LoRA — PERFORMANCE RESULTS")
+            print("=" * 80)
+            print(f"  Resolution:      {width}x{height}")
+            print(f"  Frames:          {num_frames}")
+            print(f"  Denoising steps: {num_inference_steps}")
+            print(f"  Mesh shape:      {tuple(mesh_device.shape)}")
+            print(f"  SP factor:       {sp_factor}  TP factor: {tp_factor}")
+            print(f"  Topology:        {topology}")
+            print(f"  FSDP:            {is_fsdp}")
+            print(f"  Num links:       {num_links}")
+            print(f"  LoRA high:       {len(lora_high)} adapter(s)")
+            print(f"  LoRA low:        {len(lora_low)} adapter(s)")
+            print(f"  Perf runs:       {num_perf_runs}")
+            print("-" * 80)
+            print(f"  {'Text Encoding':25s} | {fmt_stats(encoder_times)}")
+            print(f"  {'Image Encoding (total)':25s} | {fmt_stats(image_encode_times)}")
+            print(f"  {'Denoising':25s} | {fmt_stats(denoising_times)}")
+            print(f"  {'VAE Decoding':25s} | {fmt_stats(vae_times)}")
+            print(f"  {'Total Pipeline':25s} | {fmt_stats(total_times)}")
+            print("=" * 80)
+        else:
+            logger.info(f"Skipping performance summary on rank {_rank()}")
+    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 traced:
+        # Release before the mesh fixture tears down, as test_performance_wan does.
+        pipeline.release_traces()

code/models/tt_dit/experimental/tests/test_pipeline_wan_distill_i2v.py

@@ -0,0 +1,226 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+import os
+
+import numpy as np
+import PIL
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.experimental.pipelines.pipeline_wan_distill import WanDistillPipelineI2V
+from models.tt_dit.pipelines.wan.pipeline_wan_i2v import ImagePrompt
+from models.tt_dit.utils.test import ring_params
+
+
+@pytest.mark.parametrize(
+    "no_prompt",
+    [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), False)],
+)
+@pytest.mark.parametrize(
+    "mesh_device, mesh_shape, num_links, dynamic_load, device_params, topology, is_fsdp",
+    [
+        # BH Galaxy 4x8 Ring — sole supported config in the first Distill PR.
+        [(4, 8), (4, 8), 2, False, ring_params, ttnn.Topology.Ring, False],
+    ],
+    ids=["bh_4x8sp1tp0_ring"],
+    indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (832, 480),
+        (1280, 720),
+    ],
+    ids=[
+        "resolution_480p",
+        "resolution_720p",
+    ],
+)
+def test_pipeline_inference(
+    mesh_device,
+    mesh_shape,
+    num_links,
+    dynamic_load,
+    topology,
+    width,
+    height,
+    is_fsdp,
+    no_prompt,
+):
+    parent_mesh = mesh_device
+    mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+    pil_image = PIL.Image.open(os.environ.get("PROMPT_IMAGE", "./prompt_image.png"))
+    image_prompt = [ImagePrompt(image=pil_image, frame_pos=0)]
+    # Distill bakes CFG in; negative_prompt is unused when guidance_scale == 1.0,
+    # but we keep one for parity with the base I2V test.
+    negative_prompt = ""
+
+    num_frames = 81
+    num_inference_steps = 4
+    guidance_scale = 1.0
+    guidance_scale_2 = 1.0
+
+    pipeline = WanDistillPipelineI2V.create_pipeline(
+        mesh_device=mesh_device,
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        config_overrides={
+            "num_links": num_links,
+            "dynamic_load": dynamic_load,
+            "topology": topology,
+            "is_fsdp": is_fsdp,
+        },
+    )
+
+    prompt = "The cat in the hat runs up the hill to the house."
+
+    def run(*, prompt, number, seed):
+        logger.info(f"Running distill inference with prompt: '{prompt}'")
+        logger.info(f"Parameters: {height}x{width}, {num_frames} frames, {num_inference_steps} steps")
+
+        with torch.no_grad():
+            result = pipeline(
+                prompts=[prompt],
+                image_prompt=image_prompt,
+                negative_prompts=[negative_prompt],
+                num_inference_steps=num_inference_steps,
+                seed=seed,
+                guidance_scale=guidance_scale,
+                guidance_scale_2=guidance_scale_2,
+                output_type="uint8",
+            )
+
+        if hasattr(result, "frames"):
+            frames = result.frames
+        else:
+            frames = result[0] if isinstance(result, tuple) else result
+
+        logger.info("Distill inference completed successfully")
+        logger.info(f"  Output shape: {frames.shape if hasattr(frames, 'shape') else 'Unknown'}")
+        logger.info(f"  Output type: {type(frames)}")
+
+        if isinstance(frames, np.ndarray):
+            logger.info(f"  Video data range: [{frames.min():.3f}, {frames.max():.3f}]")
+        elif isinstance(frames, torch.Tensor):
+            logger.info(f"  Video data range: [{frames.min().item():.3f}, {frames.max().item():.3f}]")
+
+        frames = frames[0]
+        output_filename = f"wan_distill_i2v_{width}x{height}_{number}.mp4"
+        try:
+            from models.tt_dit.utils.video import export_to_video
+
+            export_to_video(frames, output_filename, fps=16)
+            logger.info(f"Saved video to: {output_filename}")
+        except ImportError:
+            logger.info("Could not export video - imageio_ffmpeg not available")
+
+    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)
+
+
+# ---------------------------------------------------------------------------
+# Random-weight smoke test
+# ---------------------------------------------------------------------------
+# Runs the full pipeline end-to-end without downloading the lightx2v safetensors
+# or the two ~28 GB transformer subfolders from Wan-AI. Tokenizer / UMT5 text
+# encoder / VAE / scheduler still come from Wan-AI/Wan2.2-I2V-A14B-Diffusers
+# (~12 GB total) and require TT_DIT_ALLOW_HF_DOWNLOAD=1 the first time.
+# Output frames are garbage by design — this only validates compile / mesh /
+# shape / CCL plumbing, not numeric quality.
+
+
+@pytest.mark.parametrize(
+    "mesh_device, mesh_shape, num_links, dynamic_load, device_params, topology, is_fsdp",
+    [
+        [(4, 8), (4, 8), 2, False, ring_params, ttnn.Topology.Ring, False],
+    ],
+    ids=["bh_4x8sp1tp0_ring"],
+    indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (832, 480),
+    ],
+    ids=["resolution_480p"],
+)
+def test_pipeline_inference_random_weights(
+    mesh_device,
+    mesh_shape,
+    num_links,
+    dynamic_load,
+    topology,
+    width,
+    height,
+    is_fsdp,
+):
+    parent_mesh = mesh_device
+    mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+    pil_image = PIL.Image.open(os.environ.get("PROMPT_IMAGE", "./prompt_image.png"))
+    image_prompt = [ImagePrompt(image=pil_image, frame_pos=0)]
+
+    num_frames = 81
+    num_inference_steps = 4
+
+    os.environ["TT_DIT_RANDOM_WEIGHTS"] = "1"
+    try:
+        pipeline = WanDistillPipelineI2V.create_pipeline(
+            mesh_device=mesh_device,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            config_overrides={
+                "num_links": num_links,
+                "dynamic_load": dynamic_load,
+                "topology": topology,
+                "is_fsdp": is_fsdp,
+            },
+        )
+    finally:
+        os.environ.pop("TT_DIT_RANDOM_WEIGHTS", None)
+
+    prompt = "smoke test — random weights, output is meaningless"
+    logger.info(f"Random-weight smoke test: {height}x{width}, {num_frames} frames, {num_inference_steps} steps")
+
+    with torch.no_grad():
+        result = pipeline(
+            prompts=[prompt],
+            image_prompt=image_prompt,
+            negative_prompts=[""],
+            num_inference_steps=num_inference_steps,
+            seed=42,
+            guidance_scale=1.0,
+            guidance_scale_2=1.0,
+            output_type="uint8",
+        )
+
+    if hasattr(result, "frames"):
+        frames = result.frames
+    else:
+        frames = result[0] if isinstance(result, tuple) else result
+
+    if isinstance(frames, np.ndarray):
+        rng_str = f"[{frames.min():.3f}, {frames.max():.3f}]"
+    elif isinstance(frames, torch.Tensor):
+        rng_str = f"[{frames.min().item():.3f}, {frames.max().item():.3f}]"
+    else:
+        rng_str = "n/a"
+
+    shape = frames.shape if hasattr(frames, "shape") else "Unknown"
+    logger.info(f"Random-weight smoke test PASSED. Output shape={shape}, range={rng_str}")

code/models/tt_dit/experimental/tests/test_pipeline_wan_svi.py

@@ -0,0 +1,273 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""End-to-end test for Wan2.2 I2V with SVI 2.0 Pro.
+
+Required env vars (or the test is skipped):
+- ``SVI_HIGH_PATH`` / ``SVI_LOW_PATH``: SVI LoRA safetensors (high/low noise)
+- ``LIGHTX2V_HIGH_PATH`` / ``LIGHTX2V_LOW_PATH``: required only when regime='comfyui'
+
+Normal optional env vars:
+- ``PROMPT_IMAGE``: anchor image for every clip (default: ``./prompt_image.png``)
+- ``PROMPT``: text prompt for generation. Pass one string for all clips, or a
+  semicolon-separated list with exactly ``SVI_NUM_CLIPS`` entries for per-clip
+  prompts.
+- ``SVI_NUM_CLIPS``: number of 81-frame clips to chain (default: ``2``)
+
+Example per-clip prompt format:
+``PROMPT="clip 1 prompt;clip 2 prompt;clip 3 prompt"`` requires
+``SVI_NUM_CLIPS=3``. Example invocation::
+
+    SVI_HIGH_PATH=/path/svi_high.safetensors \\
+    SVI_LOW_PATH=/path/svi_low.safetensors \\
+    LIGHTX2V_HIGH_PATH=/path/lx2v_high.safetensors \\
+    LIGHTX2V_LOW_PATH=/path/lx2v_low.safetensors \\
+    SVI_NUM_CLIPS=10 PROMPT_IMAGE=./example.png \\
+    pytest models/tt_dit/experimental/tests/test_pipeline_wan_svi.py \\
+      -v -k "comfyui and bh_2x4" -s
+
+See ``experimental/models/Wan2_2_SVI.md`` for advanced knobs (steps, CFG, seed,
+motion-latent handoff, overlap) and the upstream-workflow comparison.
+"""
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import pytest
+import torch
+from loguru import logger
+from PIL import Image
+
+import ttnn
+from models.tt_dit.experimental.pipelines.pipeline_wan_svi import WanPipelineSVI
+from models.tt_dit.utils.test import line_params, ring_params
+
+PromptArg = Union[str, List[str]]
+FloatArg = Union[float, List[float]]
+
+
+_DEFAULT_PROMPT = "A golden retriever running on a sandy beach, waves in the background"
+
+
+@dataclass(frozen=True)
+class _SVIRunConfig:
+    regime: str
+    svi_high: str
+    svi_low: str
+    lightx2v_high: Optional[str]
+    lightx2v_low: Optional[str]
+    anchor_image: Image.Image
+    prompt: PromptArg
+    num_clips: int
+    num_motion_latent: int
+    num_overlap_frame: int
+    base_seed: int
+    num_inference_steps: int
+    guidance_scale: FloatArg
+    guidance_scale_2: FloatArg
+    num_frames: int
+    width: int
+    height: int
+    output_basename: str
+
+
+def _parse_per_clip_env(env_val: str, num_clips: int, name: str, cast):
+    """Parse a scalar or semicolon-separated per-clip env var."""
+    parts = [p.strip() for p in env_val.split(";")]
+    if len(parts) == 1:
+        return cast(parts[0])
+
+    if len(parts) != num_clips:
+        pytest.fail(f"{name} has {len(parts)} entries but SVI_NUM_CLIPS={num_clips}")
+    return [cast(p) for p in parts]
+
+
+def _read_svi_config(regime: str, width: int, height: int) -> _SVIRunConfig:
+    svi_high = os.environ.get("SVI_HIGH_PATH")
+    svi_low = os.environ.get("SVI_LOW_PATH")
+    if not svi_high or not svi_low:
+        pytest.skip("SVI_HIGH_PATH and SVI_LOW_PATH env vars are required for this test.")
+
+    lightx2v_high = os.environ.get("LIGHTX2V_HIGH_PATH")
+    lightx2v_low = os.environ.get("LIGHTX2V_LOW_PATH")
+    if regime == "comfyui" and (not lightx2v_high or not lightx2v_low):
+        pytest.skip("regime='comfyui' requires LIGHTX2V_HIGH_PATH / LIGHTX2V_LOW_PATH env vars.")
+
+    num_clips = int(os.environ.get("SVI_NUM_CLIPS", "2"))
+    default_steps = "50" if regime == "python" else "6"
+    default_cfg = "5.0" if regime == "python" else "1.5"
+    guidance_scale_env = os.environ.get("GUIDANCE_SCALE", default_cfg)
+
+    prompt = _parse_per_clip_env(os.environ.get("PROMPT", _DEFAULT_PROMPT), num_clips, "PROMPT", str)
+    guidance_scale = _parse_per_clip_env(guidance_scale_env, num_clips, "GUIDANCE_SCALE", float)
+    guidance_scale_2 = _parse_per_clip_env(
+        os.environ.get("GUIDANCE_SCALE_2", guidance_scale_env),
+        num_clips,
+        "GUIDANCE_SCALE_2",
+        float,
+    )
+
+    return _SVIRunConfig(
+        regime=regime,
+        svi_high=svi_high,
+        svi_low=svi_low,
+        lightx2v_high=lightx2v_high,
+        lightx2v_low=lightx2v_low,
+        anchor_image=Image.open(os.environ.get("PROMPT_IMAGE", "./prompt_image.png")),
+        prompt=prompt,
+        num_clips=num_clips,
+        num_motion_latent=int(os.environ.get("SVI_NUM_MOTION_LATENT", "1")),
+        num_overlap_frame=int(os.environ.get("SVI_NUM_OVERLAP_FRAME", "4")),
+        base_seed=int(os.environ.get("SVI_SEED", "0")),
+        num_inference_steps=int(os.environ.get("NUM_STEPS", default_steps)),
+        guidance_scale=guidance_scale,
+        guidance_scale_2=guidance_scale_2,
+        num_frames=81,
+        width=width,
+        height=height,
+        output_basename=f"wan_svi_{regime}_{num_clips}clips_{width}x{height}",
+    )
+
+
+def _pipeline_kwargs(config: _SVIRunConfig):
+    """The SVI-specific half of the constructor arguments (the config is built separately)."""
+    kwargs = dict(
+        svi_high=config.svi_high,
+        svi_low=config.svi_low,
+        regime=config.regime,
+        num_motion_latent=config.num_motion_latent,
+        num_overlap_frame=config.num_overlap_frame,
+    )
+    if config.regime == "comfyui":
+        kwargs["lightx2v_high"] = config.lightx2v_high
+        kwargs["lightx2v_low"] = config.lightx2v_low
+    return kwargs
+
+
+def _assert_video(video: torch.Tensor, *, num_clips: int, num_frames: int, num_overlap_frame: int) -> None:
+    expected_t = num_clips * num_frames - (num_clips - 1) * num_overlap_frame
+    actual_t = video.shape[0]
+    logger.info(f"Concat shape: {tuple(video.shape)} (expected T={expected_t})")
+    assert actual_t == expected_t, f"Expected {expected_t} frames, got {actual_t}"
+    assert not torch.isnan(video.float()).any(), "video contains NaN"
+
+
+def _save_outputs(video: torch.Tensor, output_basename: str) -> None:
+    video = video.detach()
+    torch.save(video, f"{output_basename}.pt")
+    logger.info(f"Saved raw tensor to: {output_basename}.pt")
+
+    try:
+        from models.tt_dit.utils.video import export_to_video
+
+        # output_type='pt_with_last_latent' decodes frames as (T, C, H, W) float;
+        # export_to_video wants (T, H, W, C) uint8.
+        arr = video
+        if arr.ndim == 4 and arr.shape[1] in (1, 3):
+            arr = arr.permute(0, 2, 3, 1).contiguous()
+        if arr.dtype != torch.uint8:
+            arr = arr.float()
+            if arr.min().item() < 0:
+                arr = (arr + 1.0) * 127.5
+            elif arr.max().item() <= 1.0 + 1e-3:
+                arr = arr * 255.0
+            arr = arr.clamp(0, 255).to(torch.uint8)
+        export_to_video(arr.numpy(), f"{output_basename}.mp4", fps=16)
+        logger.info(f"Saved video to: {output_basename}.mp4")
+    except ImportError:
+        logger.info("Could not export video - imageio_ffmpeg not available")
+
+
+def _touch(path) -> str:
+    path.write_bytes(b"")
+    return str(path)
+
+
+@pytest.mark.parametrize(
+    "mesh_device, mesh_shape, num_links, dynamic_load, device_params, topology, is_fsdp",
+    [
+        [(2, 4), (2, 4), 2, True, line_params, ttnn.Topology.Linear, False],
+        [(4, 8), (4, 8), 2, False, ring_params, ttnn.Topology.Ring, False],
+    ],
+    ids=["bh_2x4sp1tp0", "bh_4x8sp1tp0_ring"],
+    indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (832, 480),
+    ],
+    ids=["resolution_480p"],
+)
+@pytest.mark.parametrize(
+    "regime",
+    ["python", "comfyui"],
+)
+def test_long_video(
+    mesh_device,
+    mesh_shape,
+    num_links,
+    dynamic_load,
+    topology,
+    width,
+    height,
+    is_fsdp,
+    regime,
+):
+    config = _read_svi_config(regime, width, height)
+    mesh_device = mesh_device.create_submesh(ttnn.MeshShape(*mesh_shape))
+    pipeline = WanPipelineSVI(
+        device=mesh_device,
+        config=WanPipelineSVI.default_config(
+            mesh_device=mesh_device,
+            height=config.height,
+            width=config.width,
+            num_frames=config.num_frames,
+            config_overrides={
+                "num_links": num_links,
+                "dynamic_load": dynamic_load,
+                "topology": topology,
+                "is_fsdp": is_fsdp,
+            },
+        ),
+        **_pipeline_kwargs(config),
+    )
+
+    logger.info(
+        f"SVI run: {config.num_clips} clips × {config.num_frames} frames at {width}x{height}, "
+        f"steps={config.num_inference_steps}, motion={config.num_motion_latent}, "
+        f"overlap={config.num_overlap_frame}"
+    )
+
+    with torch.no_grad():
+        video = pipeline.generate_long_video(
+            prompt=config.prompt,
+            anchor_image=config.anchor_image,
+            num_clips=config.num_clips,
+            base_seed=config.base_seed,
+            num_inference_steps=config.num_inference_steps,
+            guidance_scale=config.guidance_scale,
+            guidance_scale_2=config.guidance_scale_2,
+            partial_output_path=f"{config.output_basename}.partial.pt",
+        )
+
+    _assert_video(
+        video,
+        num_clips=config.num_clips,
+        num_frames=config.num_frames,
+        num_overlap_frame=config.num_overlap_frame,
+    )
+    _save_outputs(video, config.output_basename)
+
+
+def test_comfyui_regime_requires_lightx2v(tmp_path, expect_error):
+    """regime='comfyui' must demand LightX2V LoRA paths."""
+    with expect_error(ValueError, "lightx2v_high"):
+        WanPipelineSVI(
+            device=None,
+            config=None,
+            svi_high=_touch(tmp_path / "svi_high.safetensors"),
+            svi_low=_touch(tmp_path / "svi_low.safetensors"),
+            regime="comfyui",
+        )

code/models/tt_dit/experimental/utils/lightx2v_loader.py

@@ -0,0 +1,199 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Loader for lightx2v Wan2.2 distilled weights.
+
+The ``lightx2v/Wan2.2-Distill-Models`` HF repo ships flat ``.safetensors`` files
+containing only the DiT weights (no ``transformer/``, ``vae/``, ``text_encoder/``
+subfolders). The keys use the *original Wan* naming
+(``blocks.X.self_attn.q.weight``, ``blocks.X.modulation``, ``head.head.weight``,
+``time_embedding.0.weight``, ...), but tt_dit's ``WanTransformer3DModel``
+consumes the *diffusers* layout (``blocks.X.attn1.to_q.weight``,
+``blocks.X.scale_shift_table``, ``proj_out.weight``,
+``condition_embedder.time_embedder.linear_1.weight``, ...).
+
+:func:`load_lightx2v_state_dict` applies :func:`rename_lightx2v_to_diffusers`
+to the loaded state dict using :func:`rename_substate` from the DiT library,
+so the returned dict matches what ``WanTransformer3DModel._prepare_torch_state``
+expects.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Mapping
+
+import torch
+from loguru import logger
+from safetensors.torch import load_file
+
+from models.tt_dit.utils.substate import rename_substate
+
+# Per-block prefix renames (lightx2v → diffusers), applied under each
+# ``blocks.<i>.`` subtree. Uses rename_substate for prefix-based renaming.
+_BLOCK_PREFIX_REMAP: list[tuple[str, str]] = [
+    # self-attention
+    ("self_attn.q", "attn1.to_q"),
+    ("self_attn.k", "attn1.to_k"),
+    ("self_attn.v", "attn1.to_v"),
+    ("self_attn.o", "attn1.to_out.0"),
+    ("self_attn.norm_q", "attn1.norm_q"),
+    ("self_attn.norm_k", "attn1.norm_k"),
+    # cross-attention
+    ("cross_attn.q", "attn2.to_q"),
+    ("cross_attn.k", "attn2.to_k"),
+    ("cross_attn.v", "attn2.to_v"),
+    ("cross_attn.o", "attn2.to_out.0"),
+    ("cross_attn.norm_q", "attn2.norm_q"),
+    ("cross_attn.norm_k", "attn2.norm_k"),
+    # feedforward
+    ("ffn.0", "ffn.net.0.proj"),
+    ("ffn.2", "ffn.net.2"),
+    # cross-attn input norm
+    ("norm3", "norm2"),
+]
+
+# Single-key renames that don't fit the prefix pattern (leaf key → leaf key).
+_BLOCK_KEY_REMAP: dict[str, str] = {
+    "modulation": "scale_shift_table",
+}
+
+# Top-level prefix renames (lightx2v → diffusers).
+_TOP_LEVEL_PREFIX_REMAP: list[tuple[str, str]] = [
+    ("text_embedding.0", "condition_embedder.text_embedder.linear_1"),
+    ("text_embedding.2", "condition_embedder.text_embedder.linear_2"),
+    ("time_embedding.0", "condition_embedder.time_embedder.linear_1"),
+    ("time_embedding.2", "condition_embedder.time_embedder.linear_2"),
+    ("time_projection.1", "condition_embedder.time_proj"),
+    ("head.head", "proj_out"),
+]
+
+_TOP_LEVEL_KEY_REMAP: dict[str, str] = {
+    "head.modulation": "scale_shift_table",
+}
+
+
+def _detect_num_blocks(state: dict[str, torch.Tensor]) -> int:
+    max_idx = -1
+    for k in state:
+        if k.startswith("blocks."):
+            parts = k.split(".")
+            if len(parts) >= 2 and parts[1].isdigit():
+                max_idx = max(max_idx, int(parts[1]))
+    return max_idx + 1
+
+
+def rename_lightx2v_to_diffusers(state: dict[str, torch.Tensor]) -> None:
+    """Rename all keys in *state* from lightx2v/native-Wan to diffusers layout, in place.
+
+    Uses :func:`rename_substate` from the DiT library for prefix-based renames.
+    """
+    num_blocks = _detect_num_blocks(state)
+
+    for i in range(num_blocks):
+        for src, dst in _BLOCK_PREFIX_REMAP:
+            rename_substate(state, f"blocks.{i}.{src}", f"blocks.{i}.{dst}")
+        for src, dst in _BLOCK_KEY_REMAP.items():
+            old_key = f"blocks.{i}.{src}"
+            if old_key in state:
+                state[f"blocks.{i}.{dst}"] = state.pop(old_key)
+
+    for src, dst in _TOP_LEVEL_PREFIX_REMAP:
+        rename_substate(state, src, dst)
+
+    for src, dst in _TOP_LEVEL_KEY_REMAP.items():
+        if src in state:
+            state[dst] = state.pop(src)
+
+
+def wan_lightx2v_to_diffusers_key(key: str) -> str:
+    """Rename one lightx2v key to its diffusers-canonical equivalent.
+
+    Convenience wrapper kept for callers that need per-key conversion
+    (e.g. LoRA key targeting). Internally builds a single-key dict and
+    applies :func:`rename_lightx2v_to_diffusers`.
+    """
+    tmp = {key: torch.empty(0)}
+    rename_lightx2v_to_diffusers(tmp)
+    return next(iter(tmp))
+
+
+class WeightsNotFoundError(FileNotFoundError):
+    """Raised when a lightx2v weight file is not on disk and downloads are disabled."""
+
+
+def _hf_cache_root() -> Path:
+    return Path(os.environ.get("HF_HOME") or Path.home() / ".cache" / "huggingface")
+
+
+def _resolve_path(
+    repo_id: str,
+    filename: str,
+    *,
+    allow_download: bool,
+    local_dir: str | None,
+) -> Path:
+    if local_dir is not None:
+        candidate = Path(local_dir) / filename
+        if candidate.is_file():
+            return candidate
+
+    repo_cache = _hf_cache_root() / "hub" / f"models--{repo_id.replace('/', '--')}" / "snapshots"
+    if repo_cache.is_dir():
+        for snap in sorted(repo_cache.iterdir(), reverse=True):
+            cand = snap / filename
+            if cand.is_file():
+                return cand
+
+    if not allow_download:
+        raise WeightsNotFoundError(
+            f"lightx2v weights '{filename}' not found for repo '{repo_id}'.\n"
+            f"Searched: {repo_cache}\n"
+            f"To download from HuggingFace, re-run with TT_DIT_ALLOW_HF_DOWNLOAD=1 "
+            f"or pass allow_download=True to the pipeline.\n"
+            f"Alternatively, place the file under a snapshot directory in: {repo_cache}"
+        )
+
+    from huggingface_hub import hf_hub_download
+
+    logger.info(f"Downloading lightx2v weights from HuggingFace: {repo_id}/{filename}")
+    return Path(hf_hub_download(repo_id=repo_id, filename=filename))
+
+
+def load_lightx2v_state_dict(
+    repo_id: str,
+    filename: str,
+    *,
+    allow_download: bool = False,
+    local_dir: str | None = None,
+    key_remap: Mapping[str, str] | None = None,
+) -> dict[str, torch.Tensor]:
+    """Load a lightx2v safetensors file and return a torch state dict.
+
+    Args:
+        repo_id: HuggingFace repo, e.g. ``"lightx2v/Wan2.2-Distill-Models"``.
+        filename: File within the repo, e.g.
+            ``"wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors"``.
+        allow_download: When ``False`` (default), the file must already exist
+            in ``local_dir`` or the HF cache. When ``True``, missing files are
+            fetched via ``huggingface_hub.hf_hub_download``.
+        local_dir: Optional directory to check before the HF cache.
+        key_remap: Flat-dict overrides applied after the lightx2v→diffusers
+            rename. Useful for one-off renames in tests.
+
+    Returns:
+        State dict with diffusers-canonical keys ready for
+        ``WanTransformer3DModel.load_torch_state_dict``.
+    """
+    path = _resolve_path(repo_id, filename, allow_download=allow_download, local_dir=local_dir)
+    logger.info(f"Loading lightx2v state dict from '{path}'")
+    state = load_file(str(path))
+
+    rename_lightx2v_to_diffusers(state)
+
+    if key_remap:
+        for old_k, new_k in key_remap.items():
+            if old_k in state:
+                state[new_k] = state.pop(old_k)
+    return state

code/models/tt_dit/models/audio_vae/minimax_h3/blockings_minimax_h3_audio.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""conv3d blocking stubs for the MiniMax-H3 audio VAE.
+
+``_FP32_BLOCKINGS`` in ``utils/conv3d.py`` is tuned for the **LTX** vocoder's channel
+schedule (1536 -> 768 -> ... -> 24). H3's is different at both ends -- the decoder runs
+1024 -> 512 -> ... -> **8** and the encoder 64 -> ... -> 2048 -- so every H3 audio conv
+misses the table and falls back to ``(32, 32, 1, 1, 1)``. For a 1-D conv over a long time
+axis (a 10 s clip decodes to 324000 samples) a ``T_out_block`` of 1 means one output frame
+per work unit, and the conv3d DRAM slicer cannot find a workable configuration at all.
+
+These are **stubs**, shaped after the LTX audio entries which have the same ``(k, 1, 1)``
+form: ``C_out_block`` 32, ``H``/``W`` 1, and a ``T_out_block`` in the range LTX's tuned
+values occupy (4-64). ``bruteforce_conv3d_sweep.py`` is the tool for real tuning, and that
+is a performance-pass job -- these exist only so the correctness gates can run.
+"""
+
+from __future__ import annotations
+
+from ....layers.audio_ops import DEFAULT_MAX_C_IN_BLOCK
+from ....utils.conv3d import _FP32_BLOCKINGS, aligned_channels
+
+# 16 overshoots L1 by 1.26x (1979264 B against 1572864 B) at the widest audio convs.
+T_OUT_BLOCK = 8
+C_OUT_BLOCK = 32
+
+# ``C_in_block`` is not only a performance knob: conv3d accumulates one partial sum per block over the
+# reduction and those partials are rounded, so a narrower block is measurably less accurate *in isolation*
+# -- on `conv_pre` (Cin 2048 x k 7) the error falls monotonically, 2.40e-03 at 32, 1.86e-03 at 128,
+# 1.63e-03 at 512, and with operand splitting on, 512 is 1.48x better than 128.
+#
+# **That does not transfer end to end, so 128 stays.** Measured with every other lever on: 128 gives
+# 3.20 % chain RMSE, 256 gives 3.31 % (no better), and 512 fails outright
+# (`program.cpp:1706`). The reason is structural -- the chain is dominated by the 126 AMP convs at
+# Cin 8-512, where the block cannot widen anyway, so only `conv_pre` and `dec_in_proj` would gain and
+# there are two of them against 126. Pass ``max_c_in_block`` to re-sweep. A non-default cap changes
+# the prepared weight *bytes* (`prepare_conv3d_weight_state` blocks by ``C_in_block``) with an
+# unchanged file set, which is why `weights_variant` folds it into the device-weight cache key.
+
+
+def _c_in_block(in_channels: int, max_c_in_block: int = DEFAULT_MAX_C_IN_BLOCK) -> int:
+    """Largest 32-multiple <= ``max_c_in_block`` that divides ``in_channels`` evenly.
+
+    The kernel requires ``C_in_block`` to be a multiple of the tile width and to divide the
+    padded input channel count.
+    """
+    aligned = aligned_channels(in_channels)
+    for block in range(min(max_c_in_block, aligned) // 32 * 32, 0, -32):
+        if aligned % block == 0:
+            return block
+    return 32
+
+
+def h3_audio_channel_widths(
+    *,
+    encoder_dim: int = 64,
+    encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5),
+    latent_dim: int = 2048,
+    latent_channels: int = 32,
+    decoder_dim: int = 1024,
+    decoder_rates: tuple[int, ...] = (5, 5, 2, 2, 2, 2, 2),
+) -> set[tuple[int, int]]:
+    """Every ``(in_channels, out_channels)`` pair the audio VAE's convolutions use."""
+    pairs: set[tuple[int, int]] = set()
+
+    # Encoder: conv_in, then per stage three residual units at dim//2 plus a strided
+    # channel-doubling conv.
+    pairs.add((aligned_channels(1), encoder_dim))
+    dim = encoder_dim
+    for stride in encoder_rates:
+        inner = dim
+        dim *= 2
+        pairs.add((inner, inner))  # residual unit k7 and k1
+        pairs.add((inner, dim))  # strided conv
+    pairs.add((dim, latent_dim))
+
+    # Decoder: conv_pre, the transposed upsamplers' inner convs, the AMP blocks, conv_post.
+    pairs.add((latent_channels, latent_dim))  # dec_in_proj
+    pairs.add((latent_dim, decoder_dim))  # conv_pre
+    channels = decoder_dim
+    for _ in decoder_rates:
+        nxt = channels // 2
+        pairs.add((channels, nxt))  # upsampler inner conv
+        pairs.add((aligned_channels(nxt), aligned_channels(nxt)))  # AMP convs
+        channels = nxt
+    pairs.add((aligned_channels(channels), aligned_channels(1)))  # conv_post -> mono
+    return pairs
+
+
+def register_h3_audio_blockings(*, max_c_in_block: int = DEFAULT_MAX_C_IN_BLOCK, **config) -> int:
+    """Seed ``_FP32_BLOCKINGS`` for every H3 audio conv shape. Returns the number added.
+
+    ``setdefault``, so a swept value that later lands in ``conv3d.py`` wins over these.
+    Kernels cover every size the model uses: 1 and 3 (projections), 4/7/8/9/10/11 (AMP
+    blocks, strided encoder convs and transposed upsamplers).
+    """
+    kernels = (1, 3, 4, 7, 8, 9, 10, 11)
+    added = 0
+    for in_channels, out_channels in h3_audio_channel_widths(**config):
+        for kernel in kernels:
+            key = (aligned_channels(in_channels), max(32, out_channels), (kernel, 1, 1))
+            blocking = (_c_in_block(in_channels, max_c_in_block), C_OUT_BLOCK, T_OUT_BLOCK, 1, 1)
+            existing = _FP32_BLOCKINGS.get(key)
+            if existing is None:
+                _FP32_BLOCKINGS[key] = blocking
+                added += 1
+            elif existing[2] > T_OUT_BLOCK:
+                # A few of H3's (C_in, C_out, kernel) triples coincide with LTX's -- e.g.
+                # (128, 64, (4,1,1)) is LTX's ups[3] -- and LTX's tuned T_out_block of 32 is
+                # sized for its tensors, not H3's: at H3's widths it overshoots L1 (1979264 B
+                # against 1572864 B). Cap the temporal block for the shapes H3 actually uses
+                # rather than deferring to an entry tuned for a different model.
+                _FP32_BLOCKINGS[key] = (existing[0], existing[1], T_OUT_BLOCK, existing[3], existing[4])
+                added += 1
+    return added

code/models/tt_dit/models/audio_vae/minimax_h3/convert_minimax_h3_audio.py

@@ -0,0 +1,160 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Torch-side conversion of the MiniMax-H3 audio VAE checkpoint. No ``ttnn`` here.
+
+``tt_dit`` has no weight-norm support anywhere -- ``git grep weight_g`` over
+``models/tt_dit`` returns nothing -- because the contract is that weight norm is removed
+before a state dict reaches it. The H3 audio checkpoint stores ``weight_g``/``weight_v``
+for all ~172 of its convolutions, so that fusion has to happen somewhere, and a single
+torch-side conversion boundary is the right place: it is orthogonal to every conv class,
+it would otherwise be duplicated across four of them, and it is testable with **zero
+device time**.
+
+The fusion rule is uniform because ``torch.nn.utils.weight_norm`` is used at its default
+``dim=0`` everywhere::
+
+    weight = weight_g * weight_v / ||weight_v||   (L2 over all axes but axis 0)
+
+The trap is that for ``ConvTranspose1d`` axis 0 is **``in_channels``**, not
+``out_channels``, because its weight is stored ``(in, out, k)``. Fusing over the wrong
+axis yields a well-formed decoder that sounds subtly wrong, which is the kind of bug that
+gets misattributed to precision for days -- hence
+:func:`assert_weight_norm_axes_consistent`.
+
+Also folded here, so the device side stays simple:
+
+* ``attn.q_bias`` / ``attn.zero_k_bias`` / ``attn.v_bias`` -> one ``attn.qkv.bias``. The
+  reference passes ``cat(q_bias, zero_k_bias, v_bias)`` to a single ``F.linear``; the k
+  third is a frozen zero, which is asserted before being discarded.
+* ``activations.{2i}`` -> ``acts1.{i}`` and ``activations.{2i+1}`` -> ``acts2.{i}``, which
+  is the layout ``vocoder_ltx.AMPBlock1`` already expects.
+* ``ups.{i}.0`` -> ``ups.{i}`` and ``activation_post`` -> ``act_post``.
+"""
+
+from __future__ import annotations
+
+import re
+
+import torch
+
+
+def fuse_weight_norm(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor:
+    """``weight_g * weight_v / ||weight_v||``, L2-reduced over every axis but axis 0.
+
+    Matches ``torch.nn.utils.weight_norm`` at its default ``dim=0`` for both ``Conv1d``
+    (axis 0 = out_channels) and ``ConvTranspose1d`` (axis 0 = **in_channels**). The result
+    keeps the checkpoint's own layout; transposing is the conv class's job, not ours.
+    """
+    norm = weight_v.flatten(1).norm(dim=1).view(-1, *([1] * (weight_v.dim() - 1)))
+    return weight_g * weight_v / norm
+
+
+def assert_weight_norm_axes_consistent(state: dict[str, torch.Tensor]) -> None:
+    """Guard the axis-0 assumption on every weight-normed conv in the checkpoint.
+
+    For each ``weight_g``/``weight_v`` pair, ``weight_g`` must be a per-axis-0 scalar --
+    shape ``(axis0, 1, 1)`` -- and any sibling ``bias`` must be length ``out_channels``.
+    For a ``Conv1d`` that is ``weight_v.shape[0]``; for a ``ConvTranspose1d`` it is
+    ``weight_v.shape[1]``. A bias that matches neither means the layout is not what this
+    converter assumes.
+    """
+    for key in sorted(k for k in state if k.endswith(".weight_g")):
+        prefix = key[: -len(".weight_g")]
+        weight_g = state[key]
+        weight_v = state[f"{prefix}.weight_v"]
+        assert (
+            weight_g.shape[0] == weight_v.shape[0]
+        ), f"{prefix}: weight_g axis 0 is {weight_g.shape[0]} but weight_v axis 0 is {weight_v.shape[0]}"
+        assert all(
+            size == 1 for size in weight_g.shape[1:]
+        ), f"{prefix}: weight_g {tuple(weight_g.shape)} is not a per-axis-0 scalar, so dim=0 is the wrong axis"
+        bias = state.get(f"{prefix}.bias")
+        if bias is not None:
+            assert bias.shape[0] in (weight_v.shape[0], weight_v.shape[1]), (
+                f"{prefix}: bias {bias.shape[0]} matches neither weight_v axis "
+                f"({weight_v.shape[0]}, {weight_v.shape[1]})"
+            )
+
+
+def fuse_all_weight_norms(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Replace every ``weight_g``/``weight_v`` pair with a single fused ``weight``."""
+    assert_weight_norm_axes_consistent(state)
+    fused = {}
+    for key, value in state.items():
+        if key.endswith(".weight_v"):
+            prefix = key[: -len(".weight_v")]
+            fused[f"{prefix}.weight"] = fuse_weight_norm(state[f"{prefix}.weight_g"], value)
+        elif key.endswith(".weight_g"):
+            continue
+        else:
+            fused[key] = value
+    return fused
+
+
+def fuse_attention_biases(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Fold ``q_bias`` / ``zero_k_bias`` / ``v_bias`` into one ``qkv.bias``.
+
+    The reference builds ``cat(q_bias, zero_k_bias, v_bias)`` for a single fused
+    projection. ``zero_k_bias`` is a frozen zero buffer; that is asserted rather than
+    assumed, so a future checkpoint that unfreezes it fails loudly instead of silently
+    losing a bias term.
+    """
+    out = dict(state)
+    prefixes = {k[: -len(".q_bias")] for k in state if k.endswith(".q_bias")}
+    for prefix in sorted(prefixes):
+        q_bias = out.pop(f"{prefix}.q_bias")
+        k_bias = out.pop(f"{prefix}.zero_k_bias", None)
+        v_bias = out.pop(f"{prefix}.v_bias")
+        if k_bias is None:
+            k_bias = torch.zeros_like(q_bias)
+        else:
+            assert (
+                torch.count_nonzero(k_bias) == 0
+            ), f"{prefix}.zero_k_bias is not all zero -- it is a real parameter and cannot be dropped"
+        out[f"{prefix}.qkv.bias"] = torch.cat([q_bias, k_bias, v_bias], dim=0)
+    return out
+
+
+_ACTIVATION_RE = re.compile(r"^(?P<prefix>.*?)activations\.(?P<index>\d+)\.(?P<rest>.*)$")
+
+
+def remap_amp_activations(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """``activations.{2i}`` -> ``acts1.{i}``, ``activations.{2i+1}`` -> ``acts2.{i}``.
+
+    H3 stores one flat list of six activations per AMP block; ``vocoder_ltx.AMPBlock1``
+    holds them as two lists of three, interleaved. Getting this backwards produces a
+    plausible-looking decoder, so it is worth doing in one obvious place.
+    """
+    out = {}
+    for key, value in state.items():
+        match = _ACTIVATION_RE.match(key)
+        if match is None:
+            out[key] = value
+            continue
+        index = int(match["index"])
+        target = "acts1" if index % 2 == 0 else "acts2"
+        out[f"{match['prefix']}{target}.{index // 2}.{match['rest']}"] = value
+    return out
+
+
+def remap_decoder_names(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """``ups.{i}.0`` -> ``ups.{i}``, ``activation_post`` -> ``act_post``."""
+    out = {}
+    for key, value in state.items():
+        key = re.sub(r"(^|\.)ups\.(\d+)\.0\.", r"\1ups.\2.", key)
+        key = key.replace("activation_post.", "act_post.")
+        out[key] = value
+    return out
+
+
+def convert_minimax_h3_audio_state_dict(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """The whole conversion, in the order the steps depend on each other.
+
+    Weight norm is fused first so the later renames only ever move plain ``weight`` keys.
+    """
+    converted = fuse_all_weight_norms(state)
+    converted = fuse_attention_biases(converted)
+    converted = remap_amp_activations(converted)
+    return remap_decoder_names(converted)

code/models/tt_dit/models/audio_vae/minimax_h3/decoder_minimax_h3_audio.py

@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""MiniMax-H3 audio VAE decode: a 1x1 latent projection into a BigVGAN vocoder.
+
+Almost all of this already exists. ``audio_vae/vocoder_ltx.py``'s ``Vocoder`` is a
+BigVGAN-v2 AMP1 stack whose ``forward`` is *identical* to H3's reference
+``MiniMaxH3AudioBigVGANDecoder.forward`` -- ``conv_pre``, then per stage
+``ups[i]`` followed by the mean of three parallel AMP blocks, then ``act_post``,
+``conv_post`` and ``clamp(-1, 1)`` -- including the row-major
+``resblocks[i * num_kernels + j]`` ordering. Every loop in it is generic over
+``len(upsample_rates)``, so H3's **seven** stages need only configuration, not code.
+
+The channel schedule falls out of the same ``upsample_initial_channel // 2**i`` rule:
+1024 -> 512, 256, 128, 64, 32, 16, **8**, and ``conv_post`` takes that 8 to 1 mono
+channel. Stereo is carried as batch 2, so nothing here is stereo-aware.
+
+What this module adds is only:
+
+* ``dec_in_proj``, the 1x1 conv from 32 latent channels to the vocoder's 2048;
+* the wiring that keeps its output on device instead of round-tripping to host between
+  the projection and the vocoder.
+
+Naming is chosen so the converted checkpoint loads with no fixups: the vocoder is held as
+a child called ``decoder``, matching ``decoder.conv_pre`` / ``decoder.ups.N`` /
+``decoder.resblocks.N`` / ``decoder.act_post`` / ``decoder.conv_post``, and the projection
+is ``dec_in_proj``. The ``ups.{i}.0 -> ups.{i}`` and ``activations -> acts1/acts2`` remaps,
+and the weight-norm fusion, all happen in ``convert_minimax_h3_audio.py``.
+
+fp32 throughout: ``vocoder_ltx``'s docstring records that bf16 accumulation measurably
+degrades spectral metrics through its 108-conv chain, and H3's is longer still.
+"""
+
+from __future__ import annotations
+
+import torch
+
+import ttnn
+
+from ....layers.audio_ops import DEFAULT_MAX_C_IN_BLOCK, _AlignedOutConv1d
+from ....layers.module import Module
+from ....parallel.config import ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils.tensor import local_device_to_torch
+from ..vocoder_ltx import Vocoder
+from .blockings_minimax_h3_audio import register_h3_audio_blockings
+
+TILE_HEIGHT = 32
+
+
+class MiniMaxH3AudioDecoder(Module):
+    """``(B, 32, T)`` latents to ``(B, 1, T * 800)`` mono waveform at 32 kHz."""
+
+    def __init__(
+        self,
+        *,
+        latent_channels: int = 32,
+        latent_dim: int = 2048,
+        decoder_dim: int = 1024,
+        decoder_rates: tuple[int, ...] = (5, 5, 2, 2, 2, 2, 2),
+        decoder_kernel_sizes: tuple[int, ...] = (9, 9, 4, 4, 4, 4, 4),
+        resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11),
+        resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+        parallel_config: ParallelFactor | None = None,
+        ccl_manager: CCLManager | None = None,
+        split_mode: str = "full",
+        max_c_in_block: int = DEFAULT_MAX_C_IN_BLOCK,
+    ) -> None:
+        super().__init__()
+        self.mesh_device = mesh_device
+        self.dtype = dtype
+        self.parallel_config = parallel_config
+        self.latent_channels = latent_channels
+        self.latent_dim = latent_dim
+        self.hop_length = 1
+        for rate in decoder_rates:
+            self.hop_length *= rate
+
+        # The conv3d precision levers default to accurate. H3-only: LTX constructs the same conv
+        # classes with its own fast defaults. Kept as attributes so the pipeline's device-weight
+        # cache key (`weights_variant`) reads the exact values this module was built with. The
+        # depthwise resample filters need no lever: their conv1d kernel is exact in fp32
+        # (see compute_depthwise_conv1d.cpp).
+        self.split_mode = split_mode
+        self.max_c_in_block = max_c_in_block
+
+        # H3's audio channel schedule differs from LTX's at both ends, so every conv misses
+        # _FP32_BLOCKINGS. Seed stubs before any conv is built; see that module for why stubs.
+        register_h3_audio_blockings(max_c_in_block=max_c_in_block)
+
+        # k1 conv, so no padding mode to get wrong. _AlignedOutConv1d rather than the base
+        # class per its own docstring: a non-32-multiple out count reaches conv3d and
+        # produces a buffer whose page size does not divide its length.
+        self.dec_in_proj = _AlignedOutConv1d(
+            latent_channels,
+            latent_dim,
+            kernel_size=1,
+            bias=True,
+            mesh_device=mesh_device,
+            dtype=dtype,
+            parallel_config=parallel_config,
+            ccl_manager=ccl_manager,
+            split_mode=split_mode,
+        )
+        self.decoder = Vocoder(
+            resblock_kernel_sizes=list(resblock_kernel_sizes),
+            upsample_rates=list(decoder_rates),
+            upsample_kernel_sizes=list(decoder_kernel_sizes),
+            resblock_dilation_sizes=[list(d) for d in resblock_dilation_sizes],
+            upsample_initial_channel=decoder_dim,
+            in_channels=latent_dim,
+            out_channels=1,
+            use_tanh_at_final=False,
+            apply_final_activation=True,
+            use_bias_at_final=False,  # H3's conv_post has no bias
+            mesh_device=mesh_device,
+            dtype=dtype,
+            parallel_config=parallel_config,
+            ccl_manager=ccl_manager,
+            # H3-only opt-in: LTX's vocoder keeps its default single-conv weights.
+            split_mode=split_mode,
+        )
+
+    def _project_latents_device(self, latents_BCT: torch.Tensor) -> torch.Tensor:
+        """``(B, 32, T)`` -> ``(B, 2048, T)`` through ``dec_in_proj`` on device.
+
+        Torch in, torch out, matching the stage convention ``bwe_ltx.VocoderWithBWE``
+        uses (``_compute_mel_device``, ``_resample_device``): each stage owns its own
+        upload and readback, which keeps the stages independently testable.
+        """
+        x = latents_BCT.transpose(1, 2).float().contiguous()  # (B, T, C)
+        t_pad = self._t_padding(x.shape[1])
+        if t_pad:
+            x = torch.nn.functional.pad(x, (0, 0, 0, t_pad))
+
+        x_device = ttnn.from_torch(x, device=self.mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=self.dtype)
+        projected_device = self.dec_in_proj(x_device)
+        projected = local_device_to_torch(projected_device).float()
+        if t_pad:
+            # Crop the alignment padding back off before handing T to the vocoder, which
+            # applies its own padding for its own sharding.
+            projected = projected[:, : x.shape[1] - t_pad]
+        return projected.transpose(1, 2).contiguous()  # (B, 2048, T)
+
+    def forward(self, latents_BCT: torch.Tensor, *, traced: bool = False) -> torch.Tensor:
+        """``(B, latent_channels, T)`` torch in, ``(B, 1, T * hop_length)`` torch out.
+
+        ``traced`` replays a captured device graph for the vocoder instead of dispatching it
+        op by op. The vocoder is ~70 % host-bound, so this is its dominant lever -- unlike the
+        *visual* halves, which measure 1.00x traced because they are device-bound. Needs a
+        ``trace_region_size`` on the mesh device; the first call at a shape captures, later
+        calls replay.
+        """
+        _, channels, _ = latents_BCT.shape
+        assert channels == self.latent_channels, f"expected {self.latent_channels} latent channels, got {channels}"
+        projected = self._project_latents_device(latents_BCT)
+        return self.decoder.forward_BCT_traced(projected) if traced else self.decoder.forward_BCT(projected)
+
+    def _t_padding(self, num_frames: int) -> int:
+        """T padding needed for tile-aligned per-chip shards; zero when unsharded."""
+        if self.parallel_config is None or self.parallel_config.factor <= 1:
+            return 0
+        align = TILE_HEIGHT * self.parallel_config.factor
+        return (-num_frames) % align
+
+    def release_trace(self) -> None:
+        """Free the vocoder's captured traces; safe when none is active."""
+        self.decoder.release_trace()

code/models/tt_dit/models/audio_vae/minimax_h3/encoder_minimax_h3_audio.py

@@ -0,0 +1,405 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""MiniMax-H3 audio VAE encode: a DAC waveform encoder plus a causal-attention projection.
+
+``(B, 1, samples)`` at 32 kHz to a 40 Hz, 32-channel posterior -- a hop of
+``prod(encoder_rates) = 800``. Mono; stereo is batch 2.
+
+Structure mirrors the reference one-for-one so the converted checkpoint loads with no
+fixups. Everything nests under ``block`` ModuleLists because the reference uses
+``nn.Sequential``, which is why the checkpoint keys read
+``encoder.block.1.block.0.block.1.weight``.
+
+Reused rather than rewritten, all from ``layers/audio_ops.py``:
+
+* ``Snake`` with ``alpha_logscale=False`` -- the DAC activation stores ``alpha`` linearly,
+  unlike the decoder's ``SnakeBeta`` which is log-scale;
+* ``DilatedConv1d`` for the residual units' k=7 convs, whose ``eff_k // 2`` same-padding
+  already equals the reference's ``((7-1)*d)//2``;
+* ``_AlignedOutConv1d`` everywhere else, per its own docstring: a non-32-multiple channel
+  count reaching ``conv3d`` produces a buffer whose page size does not divide its length.
+
+The five strided channel-doubling convs pair ``kernel = 2 * stride`` with
+``padding = ceil(stride / 2)``, passed through ``Conv1dViaConv3d``'s optional ``padding``
+argument: the derived default ``eff_k // 2`` gives ``L/stride + 1`` instead of ``L/stride``.
+The default stands, so LTX call sites are unaffected.
+
+The reference's residual-unit shortcut crop is a no-op here -- the dilated conv's padding is
+exact-same, so the crop width is zero -- but the lengths are asserted rather than assumed.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+import ttnn
+
+from ....layers.audio_ops import DEFAULT_MAX_C_IN_BLOCK, Snake, _AlignedOutConv1d
+from ....layers.module import Module, ModuleList
+from ....layers.normalization import LayerNorm
+from ....parallel.config import ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils.tensor import local_device_to_torch
+from ..vocoder_ltx import DilatedConv1d
+from .blockings_minimax_h3_audio import register_h3_audio_blockings
+
+
+def _snake_row_major(snake: Snake, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
+    """Apply ``Snake`` and hand back ROW_MAJOR, which is what the next conv requires.
+
+    LTX only ever uses ``Snake``/``SnakeBeta`` inside ``Activation1d``, whose resamplers own
+    the layout. H3's DAC encoder applies a **bare** Snake between convolutions, and the
+    elementwise chain inside it returns TILE, so ``Conv1dViaConv3d``'s ROW_MAJOR assertion
+    fires without this.
+    """
+    return ttnn.to_layout(snake(x_BTC), ttnn.ROW_MAJOR_LAYOUT)
+
+
+class MiniMaxH3AudioResidualUnit(Module):
+    """``Snake -> dilated Conv1d(k=7) -> Snake -> Conv1d(k=1)``, plus a residual add."""
+
+    def __init__(
+        self,
+        dim: int,
+        dilation: int,
+        *,
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+        parallel_config: ParallelFactor | None = None,
+        ccl_manager: CCLManager | None = None,
+        split_mode: str = "off",
+    ) -> None:
+        super().__init__()
+        shared = dict(mesh_device=mesh_device, dtype=dtype, parallel_config=parallel_config, ccl_manager=ccl_manager)
+        # Conv-only levers: Snake takes neither, so they ride a separate dict from `shared`.
+        levers = dict(split_mode=split_mode)
+        self.block = ModuleList(
+            [
+                Snake(dim, alpha_logscale=False, **{k: v for k, v in shared.items() if k != "ccl_manager"}),
+                DilatedConv1d(dim, dim, kernel_size=7, dilation=dilation, **shared, **levers),
+                Snake(dim, alpha_logscale=False, **{k: v for k, v in shared.items() if k != "ccl_manager"}),
+                _AlignedOutConv1d(dim, dim, kernel_size=1, **shared, **levers),
+            ]
+        )
+
+    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
+        residual = x_BTC
+        for layer in self.block:
+            x_BTC = _snake_row_major(layer, x_BTC) if isinstance(layer, Snake) else layer(x_BTC)
+        assert x_BTC.shape[1] == residual.shape[1], (
+            f"residual unit changed T from {residual.shape[1]} to {x_BTC.shape[1]}; "
+            "the reference's centre crop is only a no-op when the dilated conv is exact-same"
+        )
+        return ttnn.add(residual, x_BTC)
+
+
+class MiniMaxH3AudioEncoderBlock(Module):
+    """Three residual units at dilations 1/3/9, then a strided channel-doubling conv."""
+
+    def __init__(
+        self,
+        dim: int,
+        stride: int,
+        *,
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+        parallel_config: ParallelFactor | None = None,
+        ccl_manager: CCLManager | None = None,
+        split_mode: str = "off",
+    ) -> None:
+        super().__init__()
+        shared = dict(mesh_device=mesh_device, dtype=dtype, parallel_config=parallel_config, ccl_manager=ccl_manager)
+        levers = dict(split_mode=split_mode)
+        inner = dim // 2
+        self.stride = stride
+        self.block = ModuleList(
+            [
+                MiniMaxH3AudioResidualUnit(inner, 1, **shared, **levers),
+                MiniMaxH3AudioResidualUnit(inner, 3, **shared, **levers),
+                MiniMaxH3AudioResidualUnit(inner, 9, **shared, **levers),
+                Snake(inner, alpha_logscale=False, **{k: v for k, v in shared.items() if k != "ccl_manager"}),
+                # kernel = 2 * stride with padding = ceil(stride / 2) is what makes the output
+                # exactly L / stride; the class's default eff_k // 2 would give one extra frame.
+                _AlignedOutConv1d(
+                    inner,
+                    dim,
+                    kernel_size=2 * stride,
+                    stride=stride,
+                    padding=math.ceil(stride / 2),
+                    **shared,
+                    **levers,
+                ),
+            ]
+        )
+
+    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
+        expected_out = x_BTC.shape[1] // self.stride
+        for layer in self.block:
+            x_BTC = _snake_row_major(layer, x_BTC) if isinstance(layer, Snake) else layer(x_BTC)
+        assert (
+            x_BTC.shape[1] == expected_out
+        ), f"strided conv produced T={x_BTC.shape[1]}, expected {expected_out} -- check the padding override"
+        return x_BTC
+
+
+class MiniMaxH3AudioDACEncoder(Module):
+    """``conv(1 -> encoder_dim, k7) -> 5 encoder blocks -> Snake -> conv(k3)``."""
+
+    def __init__(
+        self,
+        *,
+        encoder_dim: int = 64,
+        encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5),
+        latent_dim: int = 2048,
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+        parallel_config: ParallelFactor | None = None,
+        ccl_manager: CCLManager | None = None,
+        split_mode: str = "off",
+    ) -> None:
+        super().__init__()
+        shared = dict(mesh_device=mesh_device, dtype=dtype, parallel_config=parallel_config, ccl_manager=ccl_manager)
+        no_ccl = {k: v for k, v in shared.items() if k != "ccl_manager"}
+        levers = dict(split_mode=split_mode)
+
+        layers: list[Module] = [_AlignedOutConv1d(1, encoder_dim, kernel_size=7, **shared, **levers)]
+        dim = encoder_dim
+        for stride in encoder_rates:
+            dim *= 2
+            layers.append(MiniMaxH3AudioEncoderBlock(dim, stride, **shared, **levers))
+        layers.append(Snake(dim, alpha_logscale=False, **no_ccl))
+        layers.append(_AlignedOutConv1d(dim, latent_dim, kernel_size=3, **shared, **levers))
+        self.block = ModuleList(layers)
+        self.hop_length = math.prod(encoder_rates)
+
+    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
+        for layer in self.block:
+            x_BTC = _snake_row_major(layer, x_BTC) if isinstance(layer, Snake) else layer(x_BTC)
+        return x_BTC
+
+
+class MiniMaxH3AudioCausalAttention(Module):
+    """Causal attention that narrows ``in_dim`` to ``out_dim`` by pooling, not concatenating.
+
+    The heads are **mean-pooled away** rather than concatenated, and the surviving head
+    dimension is average-pooled from ``in_dim // num_heads`` down to ``out_dim``. Since
+    ``256 % 32 == 0``, that pool is exactly a mean over contiguous groups of 8, so it is a
+    reshape and a mean rather than anything adaptive.
+
+    The checkpoint stores one fused biasless ``qkv`` plus separate ``q_bias``/``v_bias`` and
+    a frozen-zero ``zero_k_bias``; the converter folds all three into ``qkv.bias``.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        num_heads: int,
+        *,
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+    ) -> None:
+        super().__init__()
+        from ....layers.linear import Linear
+
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.head_dim = in_dim // num_heads
+        assert (
+            self.head_dim % out_dim == 0
+        ), f"head_dim {self.head_dim} must be a whole multiple of out_dim {out_dim} for the pool to be a plain mean"
+        self.pool_group = self.head_dim // out_dim
+        self.mesh_device = mesh_device
+        self.dtype = dtype
+
+        self.qkv = Linear(in_dim, 3 * in_dim, bias=True, mesh_device=mesh_device, dtype=dtype)
+        self.proj = Linear(out_dim, out_dim, bias=True, mesh_device=mesh_device, dtype=dtype)
+
+    def forward(self, x_BSC: ttnn.Tensor) -> ttnn.Tensor:
+        batch, seq_len, _ = x_BSC.shape
+        # The qkv projection is 2048 -> 6144 over the whole latent sequence, and at a 10 s
+        # clip (405 frames, batch 2 for stereo) the default matmul blocking overshoots L1
+        # (1979264 B against 1572864 B). A smaller block trades a little throughput for a
+        # working shape; the real fix is a tuned config, which is a performance-pass job.
+        qkv = self.qkv(x_BSC, default_block_size=(4, 4, 4))
+        query, key, value = (
+            ttnn.permute(ttnn.reshape(part, (batch, seq_len, self.num_heads, self.head_dim)), (0, 2, 1, 3))
+            for part in ttnn.chunk(qkv, 3, dim=-1)
+        )
+        # SDPA is bf16-only (`sdpa_device_operation.cpp:43`), so this one op is a bf16 island
+        # inside the otherwise-fp32 audio path -- the same shape of compromise the visual
+        # encoder makes for GroupNorm.
+        query, key, value = (
+            ttnn.typecast(t, ttnn.bfloat16) if t.get_dtype() != ttnn.bfloat16 else t for t in (query, key, value)
+        )
+        attended = ttnn.transformer.scaled_dot_product_attention(query, key, value, is_causal=True)
+        if attended.get_dtype() != self.dtype:
+            attended = ttnn.typecast(attended, self.dtype)
+
+        # (B, heads, S, head_dim) -> mean over heads -> (B, S, head_dim)
+        pooled = ttnn.mean(ttnn.permute(attended, (0, 2, 1, 3)), dim=2)
+        pooled = ttnn.reshape(pooled, (batch, seq_len, self.out_dim, self.pool_group))
+        pooled = ttnn.mean(pooled, dim=3)
+        return self.proj(pooled)
+
+
+class MiniMaxH3AudioGeGluMlp(Module):
+    """Pre-norm GeGLU: ``w2(gelu_tanh(w0(norm(x))) * w1(norm(x)))``."""
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int,
+        *,
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+    ) -> None:
+        super().__init__()
+        from ....layers.linear import Linear
+
+        self.norm = LayerNorm(in_features, norm_eps=1e-5, bias=True, mesh_device=mesh_device)
+        self.w0 = Linear(in_features, hidden_features, bias=True, mesh_device=mesh_device, dtype=dtype)
+        self.w1 = Linear(in_features, hidden_features, bias=True, mesh_device=mesh_device, dtype=dtype)
+        self.w2 = Linear(hidden_features, in_features, bias=True, mesh_device=mesh_device, dtype=dtype)
+
+    def forward(self, x: ttnn.Tensor) -> ttnn.Tensor:
+        normed = self.norm(x)
+        # nn.GELU(approximate="tanh")
+        gated = ttnn.gelu(self.w0(normed), fast_and_approximate_mode=True)
+        return self.w2(ttnn.mul(gated, self.w1(normed)))
+
+
+class MiniMaxH3AudioAttnProjection(Module):
+    """``pre_block``: ``proj(norm3(x)) + attn(norm1(x))``, then ``+ mlp(norm2(.))``.
+
+    Note the residual is a **projected** bypass, not ``x`` itself -- the block changes width
+    from ``latent_dim`` to ``latent_channels`` inside the residual, which is unusual and
+    easy to get wrong.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        num_heads: int,
+        *,
+        mlp_ratio: int = 2,
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+    ) -> None:
+        super().__init__()
+        from ....layers.linear import Linear
+
+        self.norm1 = LayerNorm(in_dim, norm_eps=1e-5, bias=True, mesh_device=mesh_device)
+        self.attn = MiniMaxH3AudioCausalAttention(in_dim, out_dim, num_heads, mesh_device=mesh_device, dtype=dtype)
+        self.proj = Linear(in_dim, out_dim, bias=True, mesh_device=mesh_device, dtype=dtype)
+        self.norm3 = LayerNorm(in_dim, norm_eps=1e-5, bias=True, mesh_device=mesh_device)
+        self.norm2 = LayerNorm(out_dim, norm_eps=1e-5, bias=True, mesh_device=mesh_device)
+        self.mlp = MiniMaxH3AudioGeGluMlp(out_dim, out_dim * mlp_ratio, mesh_device=mesh_device, dtype=dtype)
+
+    def forward(self, x: ttnn.Tensor) -> ttnn.Tensor:
+        hidden = ttnn.add(self.proj(self.norm3(x)), self.attn(self.norm1(x)))
+        return ttnn.add(hidden, self.mlp(self.norm2(hidden)))
+
+
+class MiniMaxH3AudioEncoder(Module):
+    """Waveform to posterior: DAC trunk, ``pre_block``, then ``mean_proj`` / ``logs_proj``.
+
+    Returns ``(mean, logs)`` as torch tensors. Sampling stays on host, matching how the
+    visual encoder leaves its posterior draw to the caller.
+    """
+
+    def __init__(
+        self,
+        *,
+        encoder_dim: int = 64,
+        encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5),
+        latent_dim: int = 2048,
+        latent_channels: int = 32,
+        num_attention_heads: int = 8,
+        mesh_device: ttnn.MeshDevice,
+        dtype: ttnn.DataType = ttnn.float32,
+        parallel_config: ParallelFactor | None = None,
+        ccl_manager: CCLManager | None = None,
+        split_mode: str = "full",
+        max_c_in_block: int = DEFAULT_MAX_C_IN_BLOCK,
+    ) -> None:
+        super().__init__()
+        self.mesh_device = mesh_device
+        self.dtype = dtype
+        self.latent_channels = latent_channels
+        self.hop_length = math.prod(encoder_rates)
+
+        # The precision levers default to accurate, same rationale as the decoder. H3-only: LTX
+        # constructs the same conv classes with its own fast defaults. Kept as attributes so the
+        # pipeline's device-weight cache key (`weights_variant`) reads the exact values this module
+        # was built with.
+        self.split_mode = split_mode
+        self.max_c_in_block = max_c_in_block
+
+        # Every H3 audio conv shape misses _FP32_BLOCKINGS; seed stubs before any conv is built.
+        register_h3_audio_blockings(max_c_in_block=max_c_in_block)
+
+        self.encoder = MiniMaxH3AudioDACEncoder(
+            encoder_dim=encoder_dim,
+            encoder_rates=encoder_rates,
+            latent_dim=latent_dim,
+            mesh_device=mesh_device,
+            dtype=dtype,
+            parallel_config=parallel_config,
+            ccl_manager=ccl_manager,
+            split_mode=split_mode,
+        )
+        self.pre_block = MiniMaxH3AudioAttnProjection(
+            latent_dim, latent_channels, num_attention_heads, mesh_device=mesh_device, dtype=dtype
+        )
+        self.mean_proj = _AlignedOutConv1d(
+            latent_channels,
+            latent_channels,
+            kernel_size=1,
+            mesh_device=mesh_device,
+            dtype=dtype,
+            split_mode=split_mode,
+        )
+        self.logs_proj = _AlignedOutConv1d(
+            latent_channels,
+            latent_channels,
+            kernel_size=1,
+            mesh_device=mesh_device,
+            dtype=dtype,
+            split_mode=split_mode,
+        )
+
+    def forward(self, waveform_BCT: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """``(B, 1, samples)`` torch in, ``(mean, logs)`` each ``(B, 32, samples/800)`` torch."""
+        _, channels, num_samples = waveform_BCT.shape
+        assert channels == 1, f"the audio VAE is mono; stereo is batch 2. Got {channels} channels"
+        assert (
+            num_samples % self.hop_length == 0
+        ), f"{num_samples} samples is not a whole number of {self.hop_length}-sample hops"
+
+        x = waveform_BCT.transpose(1, 2).float().contiguous()  # (B, T, 1)
+        x_device = ttnn.from_torch(x, device=self.mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=self.dtype)
+
+        trunk = self.encoder(x_device)
+        # pre_block is a transformer block, so it wants TILE; the convs want ROW_MAJOR.
+        projected = self.pre_block(ttnn.to_layout(trunk, ttnn.TILE_LAYOUT))
+        projected = ttnn.to_layout(projected, ttnn.ROW_MAJOR_LAYOUT)
+
+        # The upload replicates and nothing here is fractured, so every device holds the same
+        # result: read back one. A bare ``ttnn.to_torch`` asserts ``buffers.size() == 1`` and so
+        # only works on a single-device mesh. Same fix as ``MiniMaxH3AudioDecoder.__call__``.
+        def read(tensor: ttnn.Tensor) -> torch.Tensor:
+            # See `MiniMaxH3AudioDecoder.__call__`: a storage slice keeps the parent's distribution
+            # metadata and the converter rejects it on a multi-host mesh. The helper reads a shard this
+            # host owns instead, which for a replicated tensor is the whole answer.
+            return local_device_to_torch(tensor).float()
+
+        mean = read(self.mean_proj(projected))
+        logs = read(self.logs_proj(projected))
+        return mean.transpose(1, 2).contiguous(), logs.transpose(1, 2).contiguous()

code/models/tt_dit/models/transformers/ltx/attention_ltx.py

@@ -0,0 +1,870 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+import ttnn
+from models.common.utility_functions import is_blackhole
+
+from ....layers.linear import ColParallelLinear, LoRAColParallelLinear, maybe_cast_activation, resolve_output_dtype
+from ....layers.module import Module
+from ....layers.normalization import DistributedRMSNorm
+from ....parallel.config import DiTParallelConfig
+from ....parallel.manager import CCLManager
+from ....utils.matmul import get_fabric_agmm_config, get_matmul_config
+from ....utils.substate import pop_substate, rename_substate
+from ....utils.tensor import bf16_tensor
+from .quant_config import LtxQuantProfile
+
+# to_gate_logits and to_q/to_qkv are both ColParallelLinear fed the SAME activation, and each fuses
+# its own TP all-gather of it (all_gather_minimal_matmul_async) — the activation crosses the fabric
+# twice. The fused gather barely overlaps its matmul, so gathering once explicitly and running both
+# projections as plain matmuls on the gathered tensor is the same math for one gather less.
+# Set to 0 to restore the double gather for an A/B.
+LTX_DEDUP_GATE_GATHER = os.environ.get("LTX_DEDUP_GATE_GATHER", "1") in ("1", "true", "True")
+
+
+class LTXAttention(Module):
+    # Map from (is_blackhole, sp_factor, tp_factor) -> (q_chunk_size, k_chunk_size)
+    sdpa_chunk_size_map = {
+        (False, 2, 4): (256, 256),
+        (False, 8, 4): (256, 256),
+        (True, 2, 2): (128, 512),
+        (True, 8, 4): (128, 512),
+    }
+
+    # V2A cross ring-SDPA q_chunk = the per-device audio Q (audio_N / sp_factor), keyed by
+    # (is_blackhole, sp, tp); assumes audio_N=256. A q_chunk wider than the Q shard pads the
+    # query rows and burns ~2x SDPA compute, so it must track sp, not be fixed per model.
+    # k_chunk reuses the self-attn ring value; misses fall back to the self-attn ring q_chunk.
+    # TODO: audio_N depends on video duration (ceil(round((num_frames/fps)*25), 32*sp)); derive
+    # q_chunk from the actual Q shard (q_BHNE.shape[2]) instead of hardcoding per mesh.
+    cross_ring_sdpa_q_chunk_map = {
+        (True, 4, 2): 64,  # BH 2x4
+        (True, 8, 4): 32,  # BH 4x8
+    }
+    default_sdpa_chunk_size = (256, 256)
+
+    # Per-stage ring-SDPA chunk, keyed by (is_blackhole, sp, tp, N); N is the SP-padded
+    # sequence length passed to the op. Misses fall back to sdpa_chunk_size_map.
+    ring_sdpa_chunk_by_n = {
+        (True, 8, 4, 9728): (96, 256),
+        (True, 8, 4, 38912): (192, 512),
+    }
+
+    # Per-shape cross-attn SDPA chunk, keyed by (is_blackhole, q_seq, kv_seq); seqs are
+    # the per-device Q shard and full K. Misses fall back to sdpa_program_config.
+    sdpa_chunk_by_shape = {
+        (True, 1216, 32): (128, 128),  # video text cross-attn, stage 1
+        (True, 4864, 32): (192, 128),  # video text cross-attn, stage 2
+        (True, 1216, 256): (128, 128),  # audio->video cross-attn, stage 1
+        (True, 4864, 256): (192, 256),  # audio->video cross-attn, stage 2
+    }
+
+    def __init__(
+        self,
+        *,
+        dim: int,
+        num_heads: int,
+        qk_norm: bool = True,
+        eps: float = 1e-6,
+        mesh_device: ttnn.MeshDevice,
+        ccl_manager: CCLManager | None = None,
+        parallel_config: DiTParallelConfig,
+        is_fsdp: bool = False,
+        is_self: bool = True,
+        context_dim: int | None = None,
+        query_input_dim: int | None = None,
+        output_dim: int | None = None,
+        apply_gated_attention: bool = False,
+        quant_config: LtxQuantProfile | None = None,
+        lora_enabled: bool = False,
+    ) -> None:
+        super().__init__()
+
+        assert dim % num_heads == 0
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.is_self = is_self
+        self.query_input_dim = query_input_dim or dim
+        self.output_dim = output_dim or dim
+
+        self.mesh_device = mesh_device
+        self.ccl_manager = ccl_manager
+        self.parallel_config = parallel_config
+
+        self.n_local_heads = self.num_heads // self.parallel_config.tensor_parallel.factor
+
+        fsdp_mesh_axis = self.parallel_config.sequence_parallel.mesh_axis if is_fsdp else None
+
+        rms_kwargs = {
+            "embedding_dim": dim,
+            "norm_eps": eps,
+            "norm_elementwise_affine": True,
+            "bias": False,
+            "mesh_device": mesh_device,
+            "mesh_axis": parallel_config.tensor_parallel.mesh_axis,
+            "ccl_manager": ccl_manager,
+        }
+
+        self.norm_q = DistributedRMSNorm(**rms_kwargs)
+        self.norm_k = DistributedRMSNorm(**rms_kwargs)
+
+        col_parallel_kwargs = {
+            "bias": True,
+            "mesh_device": mesh_device,
+            "mesh_axis": parallel_config.tensor_parallel.mesh_axis,
+            "fsdp_mesh_axis": fsdp_mesh_axis,
+            "ccl_manager": ccl_manager,
+        }
+
+        self.kv_input_dim = context_dim if (context_dim is not None and not is_self) else dim
+
+        # Per-linear precision comes entirely from the quant profile (None => bf16 everywhere, matching
+        # the unquantized model). The profile owns the to_out carve-out and the LTX_QUANT_ACTIVATIONS
+        # gating; here each role's linear just spreads the kwargs the profile hands it.
+        def qk(role):
+            return quant_config.linear_kwargs(role) if quant_config is not None else {}
+
+        # Fuse-mode LoRA lives in weight.data, so the chunked (to_qkv/to_kv) and
+        # fused-addcmul (to_out) paths work unchanged; runtime mode is unsupported here. The quant
+        # kwargs forward through LoRAColParallelLinear's **kwargs to the base linear.
+        ColCls = LoRAColParallelLinear if lora_enabled else ColParallelLinear
+
+        # Fuse the per-head gate into the QKV/Q matmul via variable-width chunks
+        tp_factor = parallel_config.tensor_parallel.factor
+        uses_fused_agmm = tp_factor > 1 and ccl_manager is not None and ccl_manager.topology == ttnn.Topology.Ring
+        # A fused linear carries a single weight dtype, but the quant profile holds the gate at bf16
+        # while quantizing qkv/q, so the two cannot share one matmul: fusion is bf16-only.
+        self.fuse_gate = apply_gated_attention and uses_fused_agmm and quant_config is None
+
+        # Gate is num_heads/TP columns per device (sub-tile); pad to a whole tile so it's a legal chunk.
+        self.gate_width_per_device = self.n_local_heads
+        self.gate_padded_per_device = ((self.n_local_heads + ttnn.TILE_SIZE - 1) // ttnn.TILE_SIZE) * ttnn.TILE_SIZE
+        gate_fused_width = self.gate_padded_per_device * tp_factor
+
+        if is_self:
+            if self.fuse_gate:
+                self.to_qkv = ColCls(
+                    dim,
+                    3 * dim + gate_fused_width,
+                    chunks=4,
+                    chunk_sizes=[dim, dim, dim, gate_fused_width],
+                    **col_parallel_kwargs,
+                    **qk("qkv"),
+                )
+            else:
+                self.to_qkv = ColCls(dim, 3 * dim, chunks=3, **col_parallel_kwargs, **qk("qkv"))
+        else:
+            if self.fuse_gate:
+                self.to_q = ColCls(
+                    self.query_input_dim,
+                    dim + gate_fused_width,
+                    chunks=2,
+                    chunk_sizes=[dim, gate_fused_width],
+                    **col_parallel_kwargs,
+                    **qk("q"),
+                )
+            else:
+                self.to_q = ColCls(self.query_input_dim, dim, **col_parallel_kwargs, **qk("q"))
+            self.to_kv = ColCls(self.kv_input_dim, 2 * dim, chunks=2, **col_parallel_kwargs, **qk("kv"))
+
+        self.to_out = ColCls(
+            dim,
+            self.output_dim,
+            bias=True,
+            mesh_device=mesh_device,
+            mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+            fsdp_mesh_axis=fsdp_mesh_axis,
+            ccl_manager=ccl_manager,
+            **qk("out"),
+        )
+
+        # Per-head gate, sharded on num_heads to match the SDPA-output head layout
+        self.apply_gated_attention = apply_gated_attention
+        if apply_gated_attention and not self.fuse_gate:
+            # Standalone gate weight stays bf16 (the model's working dtype) while consuming the shared
+            # bf8 activation, so its output is pinned back to bf16 like the quantized projections.
+            self.to_gate_logits = ColParallelLinear(
+                in_features=self.query_input_dim,
+                out_features=self.num_heads,
+                bias=True,
+                dtype=ttnn.bfloat16,
+                mesh_device=mesh_device,
+                mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+                fsdp_mesh_axis=fsdp_mesh_axis,
+                ccl_manager=ccl_manager,
+                **qk("gate"),
+            )
+
+        self.dummy_joint_input = bf16_tensor(torch.zeros((1, self.n_local_heads, 0, self.head_dim)), device=mesh_device)
+
+        full_grid = self.mesh_device.compute_with_storage_grid_size()
+        self.sdpa_program_config = ttnn.SDPAProgramConfig(
+            compute_with_storage_grid_size=full_grid,
+            q_chunk_size=256,
+            k_chunk_size=256,
+            exp_approx_mode=False,
+        )
+
+        self.sdpa_worker_grid = (full_grid.x - 1, full_grid.y)
+        mesh_key = (
+            is_blackhole(),
+            self.parallel_config.sequence_parallel.factor,
+            self.parallel_config.tensor_parallel.factor,
+        )
+        ring_sdpa_chunk_size = self.sdpa_chunk_size_map.get(mesh_key, self.default_sdpa_chunk_size)
+        self.ring_sdpa_program_config = ttnn.SDPAProgramConfig(
+            compute_with_storage_grid_size=self.sdpa_worker_grid,
+            q_chunk_size=ring_sdpa_chunk_size[0],
+            k_chunk_size=ring_sdpa_chunk_size[1],
+            exp_approx_mode=False,
+        )
+        self._ring_pc_by_n = {
+            n: ttnn.SDPAProgramConfig(
+                compute_with_storage_grid_size=self.sdpa_worker_grid,
+                q_chunk_size=chunk[0],
+                k_chunk_size=chunk[1],
+                exp_approx_mode=False,
+            )
+            for (b, sp, tp, n), chunk in self.ring_sdpa_chunk_by_n.items()
+            if (b, sp, tp) == mesh_key
+        }
+        self._sdpa_pc_by_shape = {
+            (q, kv): ttnn.SDPAProgramConfig(
+                compute_with_storage_grid_size=full_grid,
+                q_chunk_size=chunk[0],
+                k_chunk_size=chunk[1],
+                exp_approx_mode=False,
+            )
+            for (b, q, kv), chunk in self.sdpa_chunk_by_shape.items()
+            if b == mesh_key[0]
+        }
+
+        # V2A cross ring SDPA: q_chunk matched to the per-device audio Q; k_chunk reuses the
+        # self-attn ring value.
+        cross_ring_q_chunk = self.cross_ring_sdpa_q_chunk_map.get(mesh_key, ring_sdpa_chunk_size[0])
+        self.cross_ring_sdpa_program_config = ttnn.SDPAProgramConfig(
+            compute_with_storage_grid_size=self.sdpa_worker_grid,
+            q_chunk_size=cross_ring_q_chunk,
+            k_chunk_size=ring_sdpa_chunk_size[1],
+            exp_approx_mode=False,
+        )
+
+        # All SDPA (ring + cross) runs HiFi2, matching the Wan attention 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,
+        )
+
+        self.rope_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+            self.mesh_device.arch(),
+            math_fidelity=ttnn.MathFidelity.HiFi4,
+            math_approx_mode=False,
+            fp32_dest_acc_en=True,
+            packer_l1_acc=True,
+        )
+
+        device_grid = self.mesh_device.compute_with_storage_grid_size()
+        self.core_grid = ttnn.CoreGrid(x=device_grid.x, y=device_grid.y)
+
+        # Attention QKV/out matmuls run HiFi2, matching the Wan attention config.
+        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,
+        )
+
+        # Under a quant preset, override the attention compute configs to the profile's fidelity.
+        # Self-attn additionally swaps in the ring-SDPA compute and narrows its SDPA inputs; cross-attn
+        # leaves SDPA and _sdpa_input_dtype unset, so forward's getattr(self, "_sdpa_input_dtype", None)
+        # keeps cross SDPA at bf16.
+        if quant_config is not None:
+            arch = self.mesh_device.arch()
+            self.mm_compute_kernel_config = quant_config.mm_compute_config(arch)
+            if self.is_self:
+                self.sdpa_compute_kernel_config, self._sdpa_input_dtype = quant_config.sdpa_self_config(arch)
+
+    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+        rename_substate(state, "to_out.0", "to_out")
+        rename_substate(state, "q_norm", "norm_q")
+        rename_substate(state, "k_norm", "norm_k")
+
+        # Permute Q/K (and norm_q/norm_k) channels per head from checkpoint SPLIT rotation to
+        # the INTERLEAVED layout rotary_embedding_llama expects: even output lanes take the
+        # first half of each head's channels, odd lanes the second half.
+        D = self.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)
+
+        def _permute_qk(t: torch.Tensor) -> torch.Tensor:
+            rest = t.shape[1:]
+            t = t.reshape(self.num_heads, D, *rest).index_select(1, perm)
+            return t.reshape(self.num_heads * D, *rest)
+
+        for nk in ("norm_q.weight", "norm_k.weight"):
+            if nk in state:
+                state[nk] = _permute_qk(state[nk])
+
+        def _interleave_heads(tensors: list[torch.Tensor]):
+            n_dev = self.parallel_config.tensor_parallel.factor
+            tensors = [t.T for t in tensors]
+            tensors = [t.reshape(t.shape[0], n_dev, self.n_local_heads, self.head_dim) for t in tensors]
+            merged = torch.cat(tensors, dim=2)
+            merged = merged.reshape(merged.shape[0], len(tensors) * self.num_heads * self.head_dim)
+            return merged.T
+
+        def _pad_gate_to_tile(t: torch.Tensor) -> torch.Tensor:
+            """Zero-pad each device's gate block from n_local_heads up to a whole tile (device-major)."""
+            n_dev = self.parallel_config.tensor_parallel.factor
+            pad = self.gate_padded_per_device - self.n_local_heads
+            if pad == 0:
+                return t
+            blocks = []
+            for d in range(n_dev):
+                blocks.append(t[d * self.n_local_heads : (d + 1) * self.n_local_heads])
+                blocks.append(torch.zeros(pad, *t.shape[1:], dtype=t.dtype))
+            return torch.cat(blocks, dim=0)
+
+        def _gate_state_or_zeros(reference: torch.Tensor, want_bias: bool) -> dict[str, torch.Tensor]:
+            """Gate substate, zero-filled when the checkpoint has none (zeros -> identity gating)."""
+            g = pop_substate(state, "to_gate_logits")
+            if "weight" not in g:
+                g["weight"] = torch.zeros(self.num_heads, reference.shape[1], dtype=reference.dtype)
+            if want_bias and "bias" not in g:
+                g["bias"] = torch.zeros(self.num_heads, dtype=reference.dtype)
+            return g
+
+        def _interleave_device_major(tensors: list[torch.Tensor]):
+            """Like _interleave_heads but for per-device widths that differ (q/k/v head_dim vs gate)."""
+            n_dev = self.parallel_config.tensor_parallel.factor
+            reshaped = []
+            for t in tensors:
+                t = t.T  # [in, out]
+                assert t.shape[1] % n_dev == 0, f"projection width {t.shape[1]} not divisible by TP={n_dev}"
+                reshaped.append(t.reshape(t.shape[0], n_dev, t.shape[1] // n_dev))
+            merged = torch.cat(reshaped, dim=2)
+            return merged.reshape(merged.shape[0], -1).T
+
+        if self.is_self:
+            q_state = pop_substate(state, "to_q")
+            k_state = pop_substate(state, "to_k")
+            v_state = pop_substate(state, "to_v")
+
+            q_state["weight"] = _permute_qk(q_state["weight"])
+            k_state["weight"] = _permute_qk(k_state["weight"])
+            if "bias" in q_state:
+                q_state["bias"] = _permute_qk(q_state["bias"])
+            if "bias" in k_state:
+                k_state["bias"] = _permute_qk(k_state["bias"])
+
+            if self.fuse_gate:
+                # Fold the gate in as a 4th chunk, device-major: device d holds [q_d | k_d | v_d | gate_d].
+                g_state = _gate_state_or_zeros(q_state["weight"], "bias" in q_state)
+                g_state = {k: _pad_gate_to_tile(v) for k, v in g_state.items()}
+                state["to_qkv.weight"] = _interleave_device_major(
+                    [q_state["weight"], k_state["weight"], v_state["weight"], g_state["weight"]]
+                )
+                if "bias" in q_state:
+                    bias = _interleave_device_major(
+                        [
+                            q_state["bias"].unsqueeze(-1),
+                            k_state["bias"].unsqueeze(-1),
+                            v_state["bias"].unsqueeze(-1),
+                            g_state["bias"].unsqueeze(-1),
+                        ]
+                    )
+                    state["to_qkv.bias"] = bias.squeeze(-1)
+            else:
+                state["to_qkv.weight"] = _interleave_heads([q_state["weight"], k_state["weight"], v_state["weight"]])
+                if "bias" in q_state:
+                    bias = _interleave_heads(
+                        [q_state["bias"].unsqueeze(-1), k_state["bias"].unsqueeze(-1), v_state["bias"].unsqueeze(-1)]
+                    )
+                    state["to_qkv.bias"] = bias.squeeze(-1)
+        else:
+            k_state = pop_substate(state, "to_k")
+            v_state = pop_substate(state, "to_v")
+
+            k_state["weight"] = _permute_qk(k_state["weight"])
+            if "bias" in k_state:
+                k_state["bias"] = _permute_qk(k_state["bias"])
+
+            if "to_q.weight" in state:
+                state["to_q.weight"] = _permute_qk(state["to_q.weight"])
+            if "to_q.bias" in state:
+                state["to_q.bias"] = _permute_qk(state["to_q.bias"])
+
+            if self.fuse_gate and "to_q.weight" in state:
+                # Cross-attn: gate rides along with Q as a 2nd chunk.
+                g_state = _gate_state_or_zeros(state["to_q.weight"], "to_q.bias" in state)
+                g_state = {k: _pad_gate_to_tile(v) for k, v in g_state.items()}
+                state["to_q.weight"] = _interleave_device_major([state["to_q.weight"], g_state["weight"]])
+                if "to_q.bias" in state and "bias" in g_state:
+                    bias = _interleave_device_major([state["to_q.bias"].unsqueeze(-1), g_state["bias"].unsqueeze(-1)])
+                    state["to_q.bias"] = bias.squeeze(-1)
+
+            state["to_kv.weight"] = _interleave_heads([k_state["weight"], v_state["weight"]])
+            if "bias" in k_state:
+                bias = _interleave_heads([k_state["bias"].unsqueeze(-1), v_state["bias"].unsqueeze(-1)])
+                state["to_kv.bias"] = bias.squeeze(-1)
+
+    def _to_out_fused_addcmul(
+        self,
+        x: ttnn.Tensor,
+        addcmul_residual: ttnn.Tensor,
+        addcmul_gate: ttnn.Tensor,
+        compute_kernel_config=None,
+        parallel_config: DiTParallelConfig | None = None,
+        dtype=None,
+    ) -> ttnn.Tensor:
+        """Fused to_out projection + addcmul: output = residual + (matmul(x, W) + bias) * gate."""
+        to_out = self.to_out
+        # to_out inlines the AG-matmul rather than calling ColParallelLinear.forward, so it has to
+        # honour the activation cast itself. The addcmul residual/gate stay bf16 — the kernel ties
+        # their tile size to the weight's, not the activation's — and the output is the residual
+        # stream, so it must be pinned back to bf16 rather than inheriting the bf8 activation.
+        x = maybe_cast_activation(x, to_out.activation_dtype)
+        if to_out.pin_output_bf16:
+            dtype = resolve_output_dtype(dtype, x)
+
+        if to_out.fsdp_mesh_axis is not None and to_out.mesh_device.shape[to_out.fsdp_mesh_axis] > 1:
+            unsqueezed_weight = ttnn.unsqueeze_to_4D(to_out.weight.data)
+            weight = self.ccl_manager.all_gather_persistent_buffer(
+                unsqueezed_weight, dim=2, mesh_axis=to_out.fsdp_mesh_axis
+            )
+            weight = ttnn.reshape(weight, (weight.shape[-2], weight.shape[-1]))
+        else:
+            weight = to_out.weight.data
+
+        if parallel_config is not None and parallel_config.tensor_parallel.factor > 1:
+            M, K, N_out = x.padded_shape[-2], weight.padded_shape[-2], weight.padded_shape[-1]
+            full_grid = self.mesh_device.compute_with_storage_grid_size()
+
+            # Known shapes route the addcmul to_out to the strided AG-matmul (out = a + scalar*matmul*b)
+            fabric_cfg = get_fabric_agmm_config(M, K, N_out, 1, full_grid)
+            if fabric_cfg is not None:
+                tp_axis = parallel_config.tensor_parallel.mesh_axis
+                dram = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM)
+                fabric_matmul_config = ttnn.MinimalMatmulConfig(
+                    M_block_size=fabric_cfg.M_block_size,
+                    K_block_size=fabric_cfg.K_block_size,
+                    N_block_size=fabric_cfg.N_block_size,
+                    subblock_h=fabric_cfg.subblock_h,
+                    subblock_w=fabric_cfg.subblock_w,
+                    compute_with_storage_grid_size=fabric_cfg.mm_core_grid,
+                )
+                outputs = ttnn.experimental.strided_all_gather_minimal_matmul_async(
+                    x,
+                    weight,
+                    persistent_output_buffer=self.ccl_manager.get_ag_ping_pong_buffer(
+                        x.shape, 3, tp_axis, dtype=x.get_dtype()
+                    ),
+                    dim=3,
+                    multi_device_global_semaphore=self.ccl_manager.get_strided_ag_mm_semaphore(
+                        tp_axis, fabric_cfg.num_workers_per_link
+                    ),
+                    strided_all_gather_core_grid_offset=fabric_cfg.ag_core_grid_offset,
+                    num_links=self.ccl_manager.num_links,
+                    memory_config_ag=dram,
+                    topology=self.ccl_manager.topology,
+                    cluster_axis=tp_axis,
+                    bias=to_out.bias.data if to_out.bias is not None else None,
+                    config=fabric_matmul_config,
+                    memory_config_mm=dram,
+                    compute_kernel_config=compute_kernel_config or to_out.compute_config,
+                    num_workers_per_link=fabric_cfg.num_workers_per_link,
+                    num_buffers_per_channel=fabric_cfg.num_buffers_per_channel,
+                    read_local_slice_from_input=True,
+                    fused_ternary_input_a=addcmul_residual,
+                    fused_ternary_input_b=addcmul_gate,
+                    fused_ternary_scalar=1.0,
+                    chunks=1,
+                )
+                # Op returns [all_gather_output, matmul_chunk_0]; take the single matmul chunk.
+                return outputs[1]
+
+            core_grid = ttnn.CoreCoord(full_grid.x, full_grid.y - 1)
+            matmul_config = get_matmul_config(M, K, N_out, core_grid)
+
+            ag_persistent_buffer = self.ccl_manager.get_ag_ping_pong_buffer(
+                x.shape, 3, 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
+            )
+            output = ttnn.experimental.all_gather_minimal_matmul_async(
+                input_tensor=x,
+                weight_tensor=weight,
+                bias_tensor=to_out.bias.data if to_out.bias is not None else None,
+                config=matmul_config,
+                compute_kernel_config=compute_kernel_config or to_out.compute_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,
+                dtype=dtype,
+            )[0]
+        else:
+            M, K, N_out = x.padded_shape[-2], x.padded_shape[-1], weight.padded_shape[-1]
+            core_grid = self.mesh_device.compute_with_storage_grid_size()
+            matmul_config = get_matmul_config(M, K, N_out, core_grid)
+
+            output = ttnn.experimental.dit_minimal_matmul_addcmul_fused(
+                x,
+                weight,
+                1.0,
+                addcmul_residual,
+                addcmul_gate,
+                bias_tensor=to_out.bias.data if to_out.bias is not None else None,
+                config=matmul_config,
+                compute_kernel_config=compute_kernel_config or to_out.compute_config,
+                dtype=dtype,
+            )
+        return output
+
+    def _gate_is_live(self) -> bool:
+        """True when the gate projection will actually run (and so will gather its input).
+
+        A fused gate has no projection of its own — it falls out of the QKV/Q matmul — so it neither
+        runs nor gathers, and ``to_gate_logits`` is never built.
+        """
+        return self.apply_gated_attention and not self.fuse_gate and self.to_gate_logits.weight._data is not None
+
+    def _compute_gate(
+        self, spatial_1BND: ttnn.Tensor, qkv_parallel_config: DiTParallelConfig | None
+    ) -> ttnn.Tensor | None:
+        """Per-head gate 2 * sigmoid(to_gate_logits(x)); returns (B, H_local, N, 1) or None."""
+        if not self._gate_is_live():
+            return None
+
+        gate_logits = self.to_gate_logits(spatial_1BND, parallel_config=qkv_parallel_config)
+        return self._gate_from_logits(gate_logits)
+
+    def _unpad_gate_logits(self, gate_logits: ttnn.Tensor) -> ttnn.Tensor:
+        """Drop the tile padding from the fused gate chunk, leaving n_local_heads real columns."""
+        if self.gate_padded_per_device == self.gate_width_per_device:
+            return gate_logits
+        shape = gate_logits.shape
+        return ttnn.slice(
+            gate_logits,
+            [0, 0, 0, 0],
+            [shape[0], shape[1], shape[2], self.gate_width_per_device],
+        )
+
+    def _gate_from_logits(self, gate_logits: ttnn.Tensor) -> ttnn.Tensor:
+        """2 * sigmoid(logits) as (B, H_local, N, 1); shared by the fused and unfused paths."""
+        gate = ttnn.multiply(ttnn.sigmoid(gate_logits), 2.0)
+        return ttnn.permute(gate, (1, 3, 2, 0))
+
+    def forward(
+        self,
+        spatial_1BND: ttnn.Tensor,
+        N: int,
+        prompt_1BLP: ttnn.Tensor | None = None,
+        rope_cos: ttnn.Tensor | None = None,
+        rope_sin: ttnn.Tensor | None = None,
+        trans_mat: ttnn.Tensor | None = None,
+        addcmul_residual: ttnn.Tensor | None = None,
+        addcmul_gate: ttnn.Tensor | None = None,
+        k_rope_cos: ttnn.Tensor | None = None,
+        k_rope_sin: ttnn.Tensor | None = None,
+        attn_mask: ttnn.Tensor | None = None,
+        skip_qk: bool = False,
+        kv_replicated: bool = False,
+        kv_logical_n: int | None = None,
+    ) -> ttnn.Tensor:
+        """Same interface as WanAttention.forward(); pass k_rope_cos/sin for separate K RoPE
+        in A2V/V2A cross-attention."""
+        if rope_cos is not None:
+            assert rope_sin is not None
+            assert trans_mat is not None, "INTERLEAVED RoPE requires trans_mat (load-time Q/K permute assumes it)"
+
+        use_nonfused_agmm = (self.ccl_manager.topology == ttnn.Topology.Linear) and (
+            self.parallel_config.tensor_parallel.factor > 1
+        )
+        # With the gate on, the gate and Q/QKV projections would each fuse a gather of the same
+        # activation; hoisting one explicit gather feeds both and halves the fabric traffic here.
+        dedup_gate_gather = (
+            LTX_DEDUP_GATE_GATHER
+            and not use_nonfused_agmm
+            and self._gate_is_live()
+            and self.parallel_config.tensor_parallel.factor > 1
+        )
+        if use_nonfused_agmm or dedup_gate_gather:
+            # Cast BEFORE the gather, not inside the linears downstream of it: this path hoists the
+            # gather out so Q/QKV (and the gate) can share it, so the fabric payload is this tensor.
+            # Casting it at the linear would happen on the already-gathered result and shrink nothing.
+            qkv_linear = self.to_qkv if self.is_self else self.to_q
+            spatial_1BND = maybe_cast_activation(spatial_1BND, qkv_linear.activation_dtype)
+            spatial_1BND = self.ccl_manager.all_gather_persistent_buffer(
+                spatial_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+            )
+
+        qkv_parallel_config = None if (use_nonfused_agmm or dedup_gate_gather) else self.parallel_config
+
+        # When fused, the gate falls out of the QKV/Q matmul below; otherwise it's its own projection.
+        gate_bhne = None if self.fuse_gate else self._compute_gate(spatial_1BND, qkv_parallel_config)
+
+        if self.is_self:
+            if self.fuse_gate:
+                q_1BNF, k_1BNF, v_1BNF, gate_logits = self.to_qkv(
+                    spatial_1BND,
+                    compute_kernel_config=self.mm_compute_kernel_config,
+                    parallel_config=qkv_parallel_config,
+                )
+                gate_bhne = self._gate_from_logits(self._unpad_gate_logits(gate_logits))
+            else:
+                q_1BNF, k_1BNF, v_1BNF = self.to_qkv(
+                    spatial_1BND,
+                    compute_kernel_config=self.mm_compute_kernel_config,
+                    parallel_config=qkv_parallel_config,
+                )
+        else:
+            kv_input = prompt_1BLP if prompt_1BLP is not None else spatial_1BND
+            # Cross K/V: gather TP-sharded context for to_kv (replicated text prompt is already full).
+            kv_parallel_config = None
+            if prompt_1BLP is not None and self.parallel_config.tensor_parallel.factor > 1:
+                local_k = kv_input.shape[-1]
+                kv_is_tp_sharded = local_k * self.parallel_config.tensor_parallel.factor == self.kv_input_dim
+                if kv_is_tp_sharded:
+                    if use_nonfused_agmm:
+                        kv_input = self.ccl_manager.all_gather_persistent_buffer(
+                            kv_input, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+                        )
+                    else:
+                        kv_parallel_config = self.parallel_config
+            if self.fuse_gate:
+                q_1BNF, gate_logits = self.to_q(
+                    spatial_1BND,
+                    compute_kernel_config=self.mm_compute_kernel_config,
+                    parallel_config=qkv_parallel_config,
+                )
+                gate_bhne = self._gate_from_logits(self._unpad_gate_logits(gate_logits))
+            else:
+                q_1BNF = self.to_q(
+                    spatial_1BND,
+                    compute_kernel_config=self.mm_compute_kernel_config,
+                    parallel_config=qkv_parallel_config,
+                )
+            k_1BNF, v_1BNF = self.to_kv(
+                kv_input,
+                compute_kernel_config=self.mm_compute_kernel_config,
+                parallel_config=kv_parallel_config,
+            )
+
+        # RMSNorm on Q/K fused with the head split (emits BHNE via num_heads_per_device).
+        q_BHNE = self.norm_q(q_1BNF, num_heads_per_device=self.n_local_heads)
+        k_BHNE = self.norm_k(k_1BNF, num_heads_per_device=self.n_local_heads)
+
+        def create_heads(inp):
+            out, _, _ = ttnn.experimental.nlp_create_qkv_heads(
+                inp, num_heads=self.n_local_heads, num_kv_heads=0, transpose_k_heads=False
+            )
+            return out
+
+        # V still goes through the explicit reshape (no norm to fuse with).
+        v_BHNE = create_heads(v_1BNF)
+
+        # Cross-attn K/V must be full-seq for SDPA; gather across SP only when genuinely sharded.
+        is_cross = prompt_1BLP is not None
+        sp_factor = self.parallel_config.sequence_parallel.factor
+        _k_cos_pe = k_rope_cos if k_rope_cos is not None else rope_cos
+        # V2A cross: K/V stay SP-sharded (caller passes the sharded K-rope and kv_logical_n) so the
+        # ring SDPA fuses the gather instead of an explicit K/V all-gather + local SDPA.
+        use_ring_cross = is_cross and sp_factor > 1 and not kv_replicated and kv_logical_n is not None
+        if is_cross and sp_factor > 1 and not use_ring_cross:
+            sp_axis = self.parallel_config.sequence_parallel.mesh_axis
+            if kv_replicated:
+                need_gather = False
+            elif _k_cos_pe is not None:
+                need_gather = k_BHNE.shape[2] < _k_cos_pe.shape[2]
+            else:
+                # Unknown sharded context with no rope reference: gather conservatively.
+                need_gather = True
+            if need_gather:
+                k_BHNE = self.ccl_manager.all_gather_persistent_buffer(k_BHNE, dim=2, mesh_axis=sp_axis)
+                v_BHNE = self.ccl_manager.all_gather_persistent_buffer(v_BHNE, dim=2, mesh_axis=sp_axis)
+
+        if rope_cos is not None:
+            _k_cos = _k_cos_pe
+            _k_sin = k_rope_sin if k_rope_sin is not None else rope_sin
+            q_BHNE = ttnn.experimental.rotary_embedding_llama(
+                q_BHNE, rope_cos, rope_sin, trans_mat, compute_kernel_config=self.rope_compute_kernel_config
+            )
+            k_BHNE = ttnn.experimental.rotary_embedding_llama(
+                k_BHNE, _k_cos, _k_sin, trans_mat, compute_kernel_config=self.rope_compute_kernel_config
+            )
+
+        # SDPA input quant, applied after RoPE so the rotation still runs at full precision. On the
+        # ring paths K/V are the fabric payload (SDPA fuses their SP gather), so this shrinks a
+        # collective as well as the QK^T/PV matmuls; dummy_joint is a real SDPA input and must carry
+        # the same dtype. Kept separate from the linear activation cast: SDPA inputs have the widest
+        # dynamic range in the block and are the likeliest place for bf8 to break accuracy.
+        sdpa_input_dtype = getattr(self, "_sdpa_input_dtype", None)
+        dummy_joint = self.dummy_joint_input
+        if sdpa_input_dtype is not None:
+            q_BHNE = maybe_cast_activation(q_BHNE, sdpa_input_dtype)
+            k_BHNE = maybe_cast_activation(k_BHNE, sdpa_input_dtype)
+            v_BHNE = maybe_cast_activation(v_BHNE, sdpa_input_dtype)
+            dummy_joint = maybe_cast_activation(dummy_joint, sdpa_input_dtype)
+
+        if skip_qk:
+            # STG perturbation: skip Q/K attention, use V passthrough.
+            spatial_BHNE = v_BHNE
+        elif prompt_1BLP is None:
+            if sp_factor > 1 and attn_mask is None:
+                spatial_BHNE, _prompt_BHLE, _lse = ttnn.transformer.ring_joint_scaled_dot_product_attention(
+                    q_BHNE,
+                    k_BHNE,
+                    v_BHNE,
+                    dummy_joint,
+                    dummy_joint,
+                    dummy_joint,
+                    # The gather buffer must be allocated at the gathered tensor's dtype: the fabric
+                    # writes raw tiles into it, so a dtype mismatch is silent corruption, not a cast.
+                    persistent_output_buffer_k=self.ccl_manager.get_ag_ping_pong_buffer(
+                        k_BHNE.shape, 2, self.parallel_config.sequence_parallel.mesh_axis, dtype=k_BHNE.get_dtype()
+                    ),
+                    persistent_output_buffer_v=self.ccl_manager.get_ag_ping_pong_buffer(
+                        v_BHNE.shape, 2, self.parallel_config.sequence_parallel.mesh_axis, dtype=v_BHNE.get_dtype()
+                    ),
+                    joint_strategy="rear",
+                    logical_n=N,
+                    program_config=self._ring_pc_by_n.get(N, self.ring_sdpa_program_config),
+                    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=(self.sdpa_worker_grid[0], 0),
+                    use_column_major_ccl=True,
+                )
+            elif sp_factor > 1:
+                # Masked audio self-attn: gather K/V, keep Q sharded; gather+local SDPA beats ring-joint here.
+                sp_axis = self.parallel_config.sequence_parallel.mesh_axis
+                k_full = self.ccl_manager.all_gather_persistent_buffer(k_BHNE, dim=2, mesh_axis=sp_axis)
+                v_full = self.ccl_manager.all_gather_persistent_buffer(v_BHNE, dim=2, mesh_axis=sp_axis)
+                spatial_BHNE = ttnn.transformer.scaled_dot_product_attention(
+                    q_BHNE,
+                    k_full,
+                    v_full,
+                    attn_mask=attn_mask,
+                    is_causal=False,
+                    program_config=self.sdpa_program_config,
+                    compute_kernel_config=self.sdpa_compute_kernel_config,
+                )
+            else:
+                spatial_BHNE = ttnn.transformer.scaled_dot_product_attention(
+                    q_BHNE,
+                    k_BHNE,
+                    v_BHNE,
+                    attn_mask=attn_mask,
+                    is_causal=False,
+                    program_config=self.sdpa_program_config,
+                    compute_kernel_config=self.sdpa_compute_kernel_config,
+                )
+        elif use_ring_cross:
+            # Short audio Q attends non-causally to the SP-sharded video K/V; is_cross fuses the
+            # K/V gather into the ring SDPA. Output is the per-device Q shard (same as local SDPA).
+            sp_mesh_axis = self.parallel_config.sequence_parallel.mesh_axis
+            spatial_BHNE, _prompt_BHLE, _lse = ttnn.transformer.ring_joint_scaled_dot_product_attention(
+                q_BHNE,
+                k_BHNE,
+                v_BHNE,
+                dummy_joint,
+                dummy_joint,
+                dummy_joint,
+                persistent_output_buffer_k=self.ccl_manager.get_ag_ping_pong_buffer(
+                    k_BHNE.shape, 2, sp_mesh_axis, dtype=k_BHNE.get_dtype()
+                ),
+                persistent_output_buffer_v=self.ccl_manager.get_ag_ping_pong_buffer(
+                    v_BHNE.shape, 2, sp_mesh_axis, dtype=v_BHNE.get_dtype()
+                ),
+                joint_strategy="rear",
+                logical_n=kv_logical_n,
+                is_cross=True,
+                program_config=self.cross_ring_sdpa_program_config,
+                compute_kernel_config=self.sdpa_compute_kernel_config,
+                dim=2,
+                multi_device_global_semaphore=self.ccl_manager.get_ag_ping_pong_semaphore(sp_mesh_axis),
+                num_links=self.ccl_manager.num_links,
+                cluster_axis=sp_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=(self.sdpa_worker_grid[0], 0),
+                use_column_major_ccl=True,
+            )
+        else:
+            # Cross-attention: K/V full-seq, Q SP-sharded so local SDPA returns the local shard.
+            spatial_BHNE = ttnn.transformer.scaled_dot_product_attention(
+                q_BHNE,
+                k_BHNE,
+                v_BHNE,
+                attn_mask=attn_mask,
+                is_causal=False,
+                program_config=self._sdpa_pc_by_shape.get((q_BHNE.shape[2], k_BHNE.shape[2]), self.sdpa_program_config),
+                compute_kernel_config=self.sdpa_compute_kernel_config,
+            )
+
+        # Apply per-head gate in BHNE space.
+        if gate_bhne is not None:
+            spatial_BHNE = ttnn.multiply(spatial_BHNE, gate_bhne)
+
+        spatial_1BND = ttnn.transformer.concatenate_heads(spatial_BHNE)
+        spatial_1BND = ttnn.unsqueeze(spatial_1BND, 0)
+
+        # Ring fuses the TP all-gather into the to_out matmul; only Linear needs explicit AG.
+        addcmul_fused = addcmul_residual is not None and addcmul_gate is not None
+        to_out_explicit_ag = self.parallel_config.tensor_parallel.factor > 1 and use_nonfused_agmm
+        if to_out_explicit_ag:
+            # Same ordering rule as the QKV gather above: shrink the tensor before it crosses.
+            spatial_1BND = maybe_cast_activation(spatial_1BND, self.to_out.activation_dtype)
+            spatial_1BND = self.ccl_manager.all_gather_persistent_buffer(
+                spatial_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+            )
+
+        if addcmul_fused:
+            spatial_1BND = self._to_out_fused_addcmul(
+                spatial_1BND,
+                addcmul_residual,
+                addcmul_gate,
+                compute_kernel_config=self.mm_compute_kernel_config,
+                parallel_config=None if use_nonfused_agmm else self.parallel_config,
+            )
+        else:
+            spatial_1BND = self.to_out(
+                spatial_1BND,
+                compute_kernel_config=self.mm_compute_kernel_config,
+                parallel_config=None if to_out_explicit_ag else self.parallel_config,
+            )
+
+        return spatial_1BND

code/models/tt_dit/models/transformers/ltx/quant_config.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Quantization and compute-precision policy for the LTX DiT linears.
+
+A single :class:`LtxQuantProfile` holds the whole policy — weight dtypes, the ``to_out`` carve-out,
+activation/SDPA casts, and matmul/SDPA fidelities — and exposes it as construction kwargs. The
+transformer modules stay precision-agnostic: they spread ``profile.linear_kwargs(role)`` onto each
+linear and read ``profile.mm_compute_config`` / ``profile.sdpa_self_config`` for the op configs, with
+no dtype or fidelity literals of their own.
+
+Weight quantization happens at Parameter construction, not as a state_dict pass: there is no Torch
+bfloat8, so ``ttnn.from_torch(dtype=...)`` is the only quantizer and the cache enforces the declared
+``Parameter.dtype``. The profile therefore sets each linear's ``dtype`` at build time, so a cache miss
+loads weights direct-to-quant and the cache write holds the quantized tensorbins.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+import ttnn
+
+# When set, activations are narrowed to bf8 wherever they cross the fabric — both the linear activation
+# casts and the ring-SDPA Q/K/V inputs — so this is the one knob that changes what the collectives move
+# (the linear payloads and the SP-gathered K/V; SDPA itself stays HiFi2, only its inputs narrow). More
+# aggressive than the weight quant it rides on: weights are a fixed, well-conditioned distribution,
+# activations are not. On by default because the shipped 1080p tier is measured and VBench-gated with it
+# on; set to 0 to A/B against the weight-only preset. Keeping the SDPA input cast on this same knob means
+# a run can't narrow the linear activations while silently leaving SDPA at bf16 (or vice versa).
+LTX_QUANT_ACTIVATIONS = os.environ.get("LTX_QUANT_ACTIVATIONS", "1") in ("1", "true", "True")
+
+
+@dataclass(frozen=True)
+class LtxQuantProfile:
+    """One quant profile for every DiT linear, with the ``to_out`` carve-out.
+
+    Every quantized projection (qkv/q/kv, ff1/ff2) shares one weight dtype and one compute profile;
+    only ``to_out`` carves out to ``out_weight_dtype`` (bf16), because it feeds the fused-addcmul
+    epilogue whose bf16 ternary inputs must match the weight tile format. The activation and SDPA-input
+    casts (and the output pins that go with them) ride ``LTX_QUANT_ACTIVATIONS`` together, so a
+    weight-only run keeps activations bf16.
+    """
+
+    weight_dtype: ttnn.DataType  # qkv/q/kv, ff1/ff2, and any other quantized projection
+    out_weight_dtype: ttnn.DataType  # to_out carve-out
+    activation_dtype: ttnn.DataType  # linear activation cast (gated by LTX_QUANT_ACTIVATIONS)
+    mm_fidelity: ttnn.MathFidelity
+    mm_fp32_dest_acc: bool
+    sdpa_input_dtype: ttnn.DataType  # self-attn ring-SDPA input cast (gated by LTX_QUANT_ACTIVATIONS)
+    sdpa_fidelity: ttnn.MathFidelity
+    sdpa_fp32_dest_acc: bool
+
+    def linear_kwargs(self, role: str) -> dict:
+        """Construction kwargs for a DiT linear of ``role`` ("out" carves out to bf16, else quantized).
+
+        A ``gate`` keeps its bf16 weight but still consumes the shared bf8 activation, so it is pinned
+        like the projections but gets no weight-dtype or activation-cast change.
+        """
+        pin = LTX_QUANT_ACTIVATIONS
+        if role == "gate":
+            return {"pin_output_bf16": pin}
+        return {
+            "dtype": self.out_weight_dtype if role == "out" else self.weight_dtype,
+            "activation_dtype": self.activation_dtype if LTX_QUANT_ACTIVATIONS else None,
+            "pin_output_bf16": pin,
+        }
+
+    def ffn_kwargs(self) -> dict:
+        """Construction kwargs for the FFN (ParallelFeedForward), where ff1 and ff2 are both quantized."""
+        lk = self.linear_kwargs("ff")
+        return {
+            "ff1_dtype": lk["dtype"],
+            "ff2_dtype": lk["dtype"],
+            "activation_dtype": lk["activation_dtype"],
+            "pin_output_bf16": lk["pin_output_bf16"],
+        }
+
+    def mm_compute_config(self, arch):
+        """Compute-kernel config for the DiT-linear matmuls (attention QKV/out and the FFN).
+
+        ``packer_l1_acc`` stays on regardless of tier: it accumulates matmul partials in L1 at fp32,
+        which is near-free and strictly helps. It is orthogonal to ``fp32_dest_acc_en`` — the pricier
+        fp32 *destination* accumulation that the LoFi/bf8 tier deliberately drops (``mm_fp32_dest_acc``).
+        """
+        return ttnn.init_device_compute_kernel_config(
+            arch,
+            math_fidelity=self.mm_fidelity,
+            math_approx_mode=False,
+            fp32_dest_acc_en=self.mm_fp32_dest_acc,
+            packer_l1_acc=True,
+        )
+
+    def sdpa_self_config(self, arch) -> tuple:
+        """``(compute_config, input_dtype)`` for the self-attn ring SDPA; input cast rides the gate.
+
+        Cross-attention keeps its default bf16/HiFi2 SDPA, so callers apply this to self-attn only.
+        """
+        compute_config = ttnn.init_device_compute_kernel_config(
+            arch,
+            math_fidelity=self.sdpa_fidelity,
+            math_approx_mode=False,
+            fp32_dest_acc_en=self.sdpa_fp32_dest_acc,
+        )
+        return compute_config, (self.sdpa_input_dtype if LTX_QUANT_ACTIVATIONS else None)
+
+    @staticmethod
+    def all_bf8_lofi() -> LtxQuantProfile:
+        """Weights bf8, LoFi compute, bf8 activations. SDPA math stays HiFi2 with bf8 inputs.
+
+        SDPA stays fully unquantized in math (FastVideo kept attention higher precision); only its
+        inputs narrow to bf8 under ``LTX_QUANT_ACTIVATIONS``. ``ff2`` uses the RowParallel RS-fused
+        addcmul, which Wan runs at bf8 with no issue, so it is quantized; only ``to_out`` carves out.
+        """
+        return LtxQuantProfile(
+            weight_dtype=ttnn.bfloat8_b,
+            out_weight_dtype=ttnn.bfloat16,
+            activation_dtype=ttnn.bfloat8_b,
+            mm_fidelity=ttnn.MathFidelity.LoFi,
+            mm_fp32_dest_acc=False,
+            sdpa_input_dtype=ttnn.bfloat8_b,
+            sdpa_fidelity=ttnn.MathFidelity.HiFi2,
+            sdpa_fp32_dest_acc=False,
+        )
+
+    @staticmethod
+    def all_bf4_lofi() -> LtxQuantProfile:
+        """Like ``all_bf8_lofi`` but 4-bit weights and 4-bit linear activations. Opt-in probe tier, not
+        shipped: bf4 *activations* visibly destroy quality (bf4 weights alone hold up), so this is for
+        A/B measurement only. Carve-outs and the ring-SDPA bf8 input match all_bf8_lofi.
+        """
+        return LtxQuantProfile(
+            weight_dtype=ttnn.bfloat4_b,
+            out_weight_dtype=ttnn.bfloat16,
+            activation_dtype=ttnn.bfloat4_b,
+            mm_fidelity=ttnn.MathFidelity.LoFi,
+            mm_fp32_dest_acc=False,
+            sdpa_input_dtype=ttnn.bfloat8_b,
+            sdpa_fidelity=ttnn.MathFidelity.HiFi2,
+            sdpa_fp32_dest_acc=False,
+        )

code/models/tt_dit/models/transformers/ltx/rope_ltx.py

@@ -0,0 +1,393 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+LTX-2 Rotary Position Embeddings for tt_dit.
+
+Computes cos/sin on CPU using LTX-2's fractional position encoding with custom
+frequency grids, then converts to ttnn tensors for use with
+ttnn.experimental.rotary_embedding_llama.
+
+Reference: LTX-2/packages/ltx-core/src/ltx_core/model/transformer/rope.py
+"""
+
+from __future__ import annotations
+
+import functools
+import math
+from enum import Enum
+
+import numpy as np
+import torch
+
+import ttnn
+
+from ....parallel.config import DiTParallelConfig
+from ....utils.patchifiers import (
+    AudioLatentShape,
+    VideoLatentShape,
+    audio_get_patch_grid_bounds,
+    get_pixel_coords,
+    video_get_patch_grid_bounds,
+)
+from ....utils.tensor import bf16_tensor, bf16_tensor_2dshard
+
+
+class LTXRopeType(Enum):
+    INTERLEAVED = "interleaved"
+    SPLIT = "split"
+
+
+@functools.lru_cache(maxsize=5)
+def generate_freq_grid(
+    positional_embedding_theta: float,
+    positional_embedding_max_pos_count: int,
+    inner_dim: int,
+) -> torch.Tensor:
+    """Power-law frequency grid for LTX-2 RoPE: theta^linspace(...) * pi/2.
+
+    Kept fp32, not HF's fp64: upgrading the grid to fp64 measurably hurt TT-vs-HF
+    audio correlation because the downstream residual stream is bf16. Keep fp32
+    unless the whole RoPE + residual path moves to fp32.
+    """
+    theta = positional_embedding_theta
+    n_elem = 2 * positional_embedding_max_pos_count
+
+    pow_indices = np.power(
+        theta,
+        np.linspace(
+            np.log(1.0) / np.log(theta),
+            np.log(theta) / np.log(theta),
+            inner_dim // n_elem,
+            dtype=np.float64,
+        ),
+    )
+    return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32)
+
+
+def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor:
+    """Normalize position indices to [0, 1] by dividing by max_pos per dimension.
+
+    Args:
+        indices_grid: (B, n_dims, N) position indices after averaging start/end
+        max_pos: max position per dimension
+    """
+    n_pos_dims = indices_grid.shape[1]
+    assert n_pos_dims == len(
+        max_pos
+    ), f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})"
+    return torch.stack([indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)], dim=-1)
+
+
+def generate_freqs(
+    indices: torch.Tensor,
+    indices_grid: torch.Tensor,
+    max_pos: list[int],
+    use_middle_indices_grid: bool,
+) -> torch.Tensor:
+    """
+    Compute raw frequencies from position indices and frequency grid.
+
+    freqs = indices * (fractional_positions * 2 - 1)
+    This maps fractional positions from [0,1] to [-1,1] before multiplying by freq indices.
+    """
+    if use_middle_indices_grid:
+        assert len(indices_grid.shape) == 4
+        assert indices_grid.shape[-1] == 2
+        indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1]
+        indices_grid = (indices_grid_start + indices_grid_end) / 2.0
+    elif len(indices_grid.shape) == 4:
+        indices_grid = indices_grid[..., 0]
+    else:
+        # 3D (B, N, n_dims) → transpose to (B, n_dims, N) for get_fractional_positions
+        indices_grid = indices_grid.transpose(1, 2)
+
+    fractional_positions = get_fractional_positions(indices_grid, max_pos)
+    indices = indices.to(device=fractional_positions.device)
+    freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2)
+    return freqs
+
+
+def reshape_interleaved_to_bhnd(t: torch.Tensor, num_heads: int) -> torch.Tensor:
+    """Reshape interleaved (B, N, dim) to (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 interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute cos/sin with interleaved repeat for rotary embedding.
+
+    Each frequency value is duplicated: [f0, f0, f1, f1, ...] to match
+    the interleaved (x0, x1) pair structure expected by rotary_embedding_llama.
+    """
+    cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
+    sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
+    if pad_size != 0:
+        cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+        sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size])
+        cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
+        sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
+    return cos_freq, sin_freq
+
+
+def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute cos/sin with split layout for rotary_embedding_llama.
+
+    The llama kernel expects (B, H, N, head_dim) and applies rotation to pairs (x[i], x[i+D/2]).
+    We repeat the unique frequencies: [f0,f1,...,f_{d/2}, f0,f1,...,f_{d/2}] so both halves
+    are rotated correctly.
+
+    Output shape: (B, num_heads, T, head_dim) where head_dim = 2 * D_half.
+    """
+    cos_freq = freqs.cos()
+    sin_freq = freqs.sin()
+
+    if pad_size != 0:
+        cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+        sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
+        cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
+        sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
+
+    b, t = cos_freq.shape[0], cos_freq.shape[1]
+    # Reshape to per-head: (B, T, H, D_half) -> (B, H, T, D_half)
+    cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1).swapaxes(1, 2)
+    sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1).swapaxes(1, 2)
+    return cos_freq, sin_freq
+
+
+def precompute_freqs_cis(
+    indices_grid: torch.Tensor,
+    dim: int,
+    out_dtype: torch.dtype = torch.float32,
+    theta: float = 10000.0,
+    max_pos: list[int] | None = None,
+    use_middle_indices_grid: bool = False,
+    num_attention_heads: int = 32,
+    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Precompute (cos, sin) for LTX-2 RoPE.
+
+    Output shape depends on rope_type:
+    - INTERLEAVED: (B, seq_len, dim), each freq duplicated for the (x0, x1) pairs.
+    - SPLIT: (B, num_heads, seq_len, dim // (2 * num_heads)).
+    """
+    if max_pos is None:
+        max_pos = [20, 2048, 2048]
+
+    # n_pos_dims is the number of position dimensions (e.g. 3 for temporal + H + W)
+    # For 4D grids (B, n_dims, N, 2), n_pos_dims is shape[1]; for 3D (B, N, n_dims), it's shape[-1]
+    if use_middle_indices_grid and len(indices_grid.shape) == 4:
+        n_pos_dims = indices_grid.shape[1]  # (B, n_dims, N, 2)
+    else:
+        n_pos_dims = indices_grid.shape[-1]
+    indices = generate_freq_grid(theta, n_pos_dims, dim)
+    freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid)
+
+    if rope_type == LTXRopeType.SPLIT:
+        expected_freqs = dim // 2
+        current_freqs = freqs.shape[-1]
+        pad_size = expected_freqs - current_freqs
+        cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads)
+    else:
+        n_elem = 2 * n_pos_dims
+        cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
+
+    return cos_freq.to(out_dtype), sin_freq.to(out_dtype)
+
+
+# =============================================================================
+# Device-side RoPE builders (INTERLEAVED cos/sin, sharded onto the DiT mesh) for
+# ttnn.experimental.rotary_embedding_llama. All positions stay fp32 — bf16 introduced
+# catastrophic phase error in high-frequency RoPE channels (randomizing the top half of
+# head_dim), so only the final ttnn tensors are cast to bf16.
+# =============================================================================
+
+
+def pad_video_rope_sp(
+    cos_freq: torch.Tensor, sin_freq: torch.Tensor, sp_factor: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Right-pad video RoPE cos/sin on dim=2 to the SP boundary (``ttnn.TILE_SIZE * sp_factor``).
+
+    Padded slots use cos=1, sin=0 (identity rotation); SDPA still masks them via ``logical_n``.
+    Same convention as the audio RoPE padding in ``prepare_audio_rope`` / ``prepare_av_cross_pe``.
+    """
+    video_N_real = cos_freq.shape[2]
+    divisor = ttnn.TILE_SIZE * sp_factor
+    video_N = ((video_N_real + divisor - 1) // divisor) * divisor
+    if video_N == video_N_real:
+        return cos_freq, sin_freq
+    pad = video_N - video_N_real
+    H = cos_freq.shape[1]
+    d_half = cos_freq.shape[-1]
+    cos_pad = torch.ones(1, H, pad, d_half, dtype=cos_freq.dtype)
+    sin_pad = torch.zeros(1, H, pad, d_half, dtype=sin_freq.dtype)
+    cos_freq = torch.cat([cos_freq, cos_pad], dim=2)
+    sin_freq = torch.cat([sin_freq, sin_pad], dim=2)
+    return cos_freq, sin_freq
+
+
+def prepare_video_rope(
+    latent_frames: int,
+    latent_height: int,
+    latent_width: int,
+    *,
+    inner_dim: int,
+    num_attention_heads: int,
+    theta: float,
+    max_pos: list[int],
+    mesh_device: ttnn.MeshDevice,
+    parallel_config: DiTParallelConfig,
+    fps: float = 24.0,
+) -> tuple[ttnn.Tensor, ttnn.Tensor]:
+    """Compute video RoPE in INTERLEAVED layout, SP×TP sharded onto the mesh."""
+    v_shape = VideoLatentShape(batch=1, channels=128, frames=latent_frames, height=latent_height, width=latent_width)
+    v_coords = video_get_patch_grid_bounds(v_shape)
+    v_positions = get_pixel_coords(v_coords, scale_factors=(8, 32, 32), causal_fix=True).float()
+    v_positions[:, 0, ...] = v_positions[:, 0, ...] / fps
+
+    cos_freq, sin_freq = precompute_freqs_cis(
+        v_positions,
+        dim=inner_dim,
+        out_dtype=torch.float32,
+        theta=theta,
+        max_pos=max_pos,
+        use_middle_indices_grid=True,
+        num_attention_heads=num_attention_heads,
+        rope_type=LTXRopeType.INTERLEAVED,
+    )  # (1, N, dim)
+
+    cos_freq = reshape_interleaved_to_bhnd(cos_freq, num_attention_heads)
+    sin_freq = reshape_interleaved_to_bhnd(sin_freq, num_attention_heads)
+
+    # Pad seq dim to ttnn.TILE_SIZE * sp_factor; padded slots use cos=1, sin=0 (identity).
+    cos_freq, sin_freq = pad_video_rope_sp(cos_freq, sin_freq, parallel_config.sequence_parallel.factor)
+
+    sp_axis = parallel_config.sequence_parallel.mesh_axis
+    tp_axis = parallel_config.tensor_parallel.mesh_axis
+    tt_cos = bf16_tensor_2dshard(cos_freq, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+    tt_sin = bf16_tensor_2dshard(sin_freq, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+    return tt_cos, tt_sin
+
+
+def prepare_audio_rope(
+    audio_N: int,
+    audio_N_real: int,
+    *,
+    theta: float,
+    mesh_device: ttnn.MeshDevice,
+    parallel_config: DiTParallelConfig,
+) -> tuple[ttnn.Tensor, ttnn.Tensor]:
+    """Compute audio RoPE in INTERLEAVED layout, SP×TP sharded onto the mesh."""
+    a_shape = AudioLatentShape(batch=1, channels=8, frames=audio_N_real, mel_bins=16)
+    a_positions = audio_get_patch_grid_bounds(a_shape).float()  # (1, 1, N, 2)
+
+    a_cos, a_sin = precompute_freqs_cis(
+        a_positions,
+        dim=2048,
+        out_dtype=torch.float32,
+        theta=theta,
+        max_pos=[20],
+        use_middle_indices_grid=True,
+        num_attention_heads=32,
+        rope_type=LTXRopeType.INTERLEAVED,
+    )  # (1, N, 2048)
+    a_cos = reshape_interleaved_to_bhnd(a_cos, num_heads=32)  # (1, 32, N, 64)
+    a_sin = reshape_interleaved_to_bhnd(a_sin, num_heads=32)
+
+    if audio_N > audio_N_real:
+        head_dim = a_cos.shape[-1]
+        a_cos_padded = torch.ones(1, 32, audio_N, head_dim)
+        a_cos_padded[:, :, :audio_N_real, :] = a_cos
+        a_sin_padded = torch.zeros(1, 32, audio_N, head_dim)
+        a_sin_padded[:, :, :audio_N_real, :] = a_sin
+        a_cos, a_sin = a_cos_padded, a_sin_padded
+
+    sp_axis = parallel_config.sequence_parallel.mesh_axis
+    tp_axis = parallel_config.tensor_parallel.mesh_axis
+    return (
+        bf16_tensor_2dshard(a_cos, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}),
+        bf16_tensor_2dshard(a_sin, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1}),
+    )
+
+
+def prepare_av_cross_pe(
+    latent_frames: int,
+    latent_height: int,
+    latent_width: int,
+    audio_N: int,
+    audio_N_real: int,
+    *,
+    theta: float,
+    mesh_device: ttnn.MeshDevice,
+    parallel_config: DiTParallelConfig,
+    fps: float = 24.0,
+    cross_pe_max_pos: int = 20,
+) -> tuple[ttnn.Tensor, ttnn.Tensor, ttnn.Tensor, ttnn.Tensor, ttnn.Tensor, ttnn.Tensor]:
+    """Temporal-only cross positional embeddings for A↔V cross-attention.
+
+    Reference: ``MultiModalTransformerArgsPreprocessor.prepare`` builds ``cross_pe`` from the
+    temporal slice ``modality.positions[:, 0:1, :]`` at ``dim=audio_cross_attention_dim`` with
+    ``max_pos=[cross_pe_max_pos]``. Video and audio share the scheme so a video token and an
+    audio token at the same time get the same rotary phase (AV temporal sync).
+
+    Returns 6 device tensors used by inner_step:
+        (v_q_cos, v_q_sin)  — video Q in A→V cross-attn (SP×TP sharded). Also reused as the
+                              video K rope in V→A (ring SDPA gathers the SP-sharded K).
+        (a_q_cos, a_q_sin)  — audio Q in V→A cross-attn (SP×TP sharded).
+        (a_k_cos, a_k_sin)  — audio K in A→V cross-attn (TP-only; K side after AllGather).
+    """
+    v_shape = VideoLatentShape(batch=1, channels=128, frames=latent_frames, height=latent_height, width=latent_width)
+    v_coords = video_get_patch_grid_bounds(v_shape)
+    v_positions = get_pixel_coords(v_coords, scale_factors=(8, 32, 32), causal_fix=True).float()
+    v_positions[:, 0, ...] = v_positions[:, 0, ...] / fps  # temporal axis → seconds
+    v_temporal = v_positions[:, 0:1, :]  # (1, 1, video_N, 2)
+
+    a_shape = AudioLatentShape(batch=1, channels=8, frames=audio_N_real, mel_bins=16)
+    a_positions = audio_get_patch_grid_bounds(a_shape).float()  # (1, 1, audio_N_real, 2)
+
+    rope_kwargs = dict(
+        dim=2048,  # audio_cross_attention_dim — both sides share this
+        out_dtype=torch.float32,
+        theta=theta,
+        max_pos=[cross_pe_max_pos],
+        use_middle_indices_grid=True,
+        num_attention_heads=32,
+        rope_type=LTXRopeType.INTERLEAVED,
+    )
+
+    v_cos, v_sin = precompute_freqs_cis(v_temporal, **rope_kwargs)  # (1, video_N, 2048)
+    a_cos, a_sin = precompute_freqs_cis(a_positions, **rope_kwargs)  # (1, audio_N_real, 2048)
+    v_cos = reshape_interleaved_to_bhnd(v_cos, num_heads=32)  # (1, 32, video_N, 64)
+    v_sin = reshape_interleaved_to_bhnd(v_sin, num_heads=32)
+    a_cos = reshape_interleaved_to_bhnd(a_cos, num_heads=32)
+    a_sin = reshape_interleaved_to_bhnd(a_sin, num_heads=32)
+
+    v_cos, v_sin = pad_video_rope_sp(v_cos, v_sin, parallel_config.sequence_parallel.factor)
+
+    if audio_N > audio_N_real:
+        head_dim = a_cos.shape[-1]
+        a_cos_padded = torch.ones(1, 32, audio_N, head_dim)
+        a_cos_padded[:, :, :audio_N_real, :] = a_cos
+        a_sin_padded = torch.zeros(1, 32, audio_N, head_dim)
+        a_sin_padded[:, :, :audio_N_real, :] = a_sin
+        a_cos, a_sin = a_cos_padded, a_sin_padded
+
+    sp_axis = parallel_config.sequence_parallel.mesh_axis
+    tp_axis = parallel_config.tensor_parallel.mesh_axis
+
+    # Q-side: SP×TP sharded (matches the Q tensor layout post-attention QKV split).
+    v_q_cos = bf16_tensor_2dshard(v_cos, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+    v_q_sin = bf16_tensor_2dshard(v_sin, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+    a_q_cos = bf16_tensor_2dshard(a_cos, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+    a_q_sin = bf16_tensor_2dshard(a_sin, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+
+    # K-side: TP-only on heads (sequence is replicated after AllGather on K). Only the audio K
+    # rope is needed (A→V gathers audio K); video K in V→A reuses the SP-sharded v_q rope.
+    a_k_cos = bf16_tensor(a_cos, device=mesh_device, mesh_axis=tp_axis, shard_dim=1)
+    a_k_sin = bf16_tensor(a_sin, device=mesh_device, mesh_axis=tp_axis, shard_dim=1)
+
+    return (v_q_cos, v_q_sin, a_q_cos, a_q_sin, a_k_cos, a_k_sin)