diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..92a8316488e12eaac18432ed01b9b01bc1524a56 100644
--- a/.gitattributes
+++ b/.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
diff --git a/code/g2p/sources/chinese_lexicon.txt b/code/g2p/sources/chinese_lexicon.txt
new file mode 100644
index 0000000000000000000000000000000000000000..4d7dbf347a29d3b87c199d0e56ef7f1dbf28a6ee
--- /dev/null
+++ b/code/g2p/sources/chinese_lexicon.txt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3a7685d1c3e68eb2fa304bfc63e90c90c3c1a1948839a5b1b507b2131b3e2fb
+size 14779443
diff --git a/code/g2p/sources/g2p_chinese_model/poly_bert_model.onnx b/code/g2p/sources/g2p_chinese_model/poly_bert_model.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..6b952b9717eb71bb5a7aa2492478095f117858dd
--- /dev/null
+++ b/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
diff --git a/code/infer/example/edit_cn.lrc b/code/infer/example/edit_cn.lrc
new file mode 100644
index 0000000000000000000000000000000000000000..4f472bafeefe9866e9b58844ccc4627b5d5d820d
--- /dev/null
+++ b/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]耐心的人啊才可以看见童话
\ No newline at end of file
diff --git a/code/infer/example/edit_cn.mp3 b/code/infer/example/edit_cn.mp3
new file mode 100644
index 0000000000000000000000000000000000000000..b133befcae41eca8e5fb7668d0c88593a1102300
--- /dev/null
+++ b/code/infer/example/edit_cn.mp3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fbc319c58215aff58b018a253abf6e6ec31673d6336947b932f1212964240d4
+size 8534237
diff --git a/code/infer/example/edit_en.lrc b/code/infer/example/edit_en.lrc
new file mode 100644
index 0000000000000000000000000000000000000000..1046adc05100ad1e1c60cd19fcbf69a4b4a9da4c
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/code/infer/example/edit_en.mp3 b/code/infer/example/edit_en.mp3
new file mode 100644
index 0000000000000000000000000000000000000000..47893b7ac48ca73cdba29b0bbfc305aadcad742d
--- /dev/null
+++ b/code/infer/example/edit_en.mp3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d72cb739c351b00a738f50e37a94cfc783dfa12b091c84ef9943b85a064015eb
+size 4771890
diff --git a/code/infer/example/eg_cn.lrc b/code/infer/example/eg_cn.lrc
new file mode 100644
index 0000000000000000000000000000000000000000..1cf7ff0efbe7036319287c0f2e0eadaaa1eacd3e
--- /dev/null
+++ b/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]和我在成都的街头走一走
\ No newline at end of file
diff --git a/code/infer/example/eg_cn.wav b/code/infer/example/eg_cn.wav
new file mode 100644
index 0000000000000000000000000000000000000000..b817eae85b328c8d6de9ea9ea3e545b8b617d9d0
--- /dev/null
+++ b/code/infer/example/eg_cn.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67fc2163fa200544e5ecc1fcd7889124a60e6ff20fd70c53ad05a18859d79fd3
+size 882044
diff --git a/code/infer/example/eg_cn_full.lrc b/code/infer/example/eg_cn_full.lrc
new file mode 100644
index 0000000000000000000000000000000000000000..f05dfaa370736e69a856a6f7ba6c31ede1336ab5
--- /dev/null
+++ b/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]直到所有的灯都熄灭了也不停留
\ No newline at end of file
diff --git a/code/infer/example/eg_en.lrc b/code/infer/example/eg_en.lrc
new file mode 100644
index 0000000000000000000000000000000000000000..c097f2e2062e931d660d7446ec0d08e7015556e0
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/code/infer/example/eg_en.mp3 b/code/infer/example/eg_en.mp3
new file mode 100644
index 0000000000000000000000000000000000000000..b1c0d493008b31c6b2e7748950a432c05d3655af
--- /dev/null
+++ b/code/infer/example/eg_en.mp3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cdf3610c437715af336102a9bd34c63a3e882d9d99fe86b9fb193d8ebf9f2aa4
+size 10386551
diff --git a/code/infer/example/eg_en_full.lrc b/code/infer/example/eg_en_full.lrc
new file mode 100644
index 0000000000000000000000000000000000000000..673721a9b25e2242fc16d94ea94bec3a36971f50
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/code/infer/example/music_prompt_keywords.json b/code/infer/example/music_prompt_keywords.json
new file mode 100644
index 0000000000000000000000000000000000000000..3f8a4ab7e07a86a0e893674db69f13a6f7aee2c8
--- /dev/null
+++ b/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"
+ ]
+ }
\ No newline at end of file
diff --git a/code/infer/example/output-windows/en-utf8-fixed.mp3 b/code/infer/example/output-windows/en-utf8-fixed.mp3
new file mode 100644
index 0000000000000000000000000000000000000000..f4957cd9c49340256ea74ee25b2102e86b5f0152
--- /dev/null
+++ b/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
diff --git a/code/infer/example/output-windows/en-without-utf8-fix.mp3 b/code/infer/example/output-windows/en-without-utf8-fix.mp3
new file mode 100644
index 0000000000000000000000000000000000000000..188009ca6381cab2b9bf588d737cc25f251fe67f
--- /dev/null
+++ b/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
diff --git a/code/infer/example/output/output.mp3 b/code/infer/example/output/output.mp3
new file mode 100644
index 0000000000000000000000000000000000000000..431939aadcc010031c20723b4995d2deef25b648
--- /dev/null
+++ b/code/infer/example/output/output.mp3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6311c518821ee0d37d24c65d92c479f6c26da68b782687ab0ee69fb7a9ef97d3
+size 2832762
diff --git a/code/infer/example/vocal.npy b/code/infer/example/vocal.npy
new file mode 100644
index 0000000000000000000000000000000000000000..f3056dc285d16aa6abaf81ba70f6f0eb7038c304
--- /dev/null
+++ b/code/infer/example/vocal.npy
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cb7d74eb7a8eda12acb8247b21d373928301db8a8cb0db480d341799fed3ce5
+size 2176
diff --git a/code/models/tt_dit/experimental/lora/__init__.py b/code/models/tt_dit/experimental/lora/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e13459ada2877a67fad216d55a31379cb8db5e78
--- /dev/null
+++ b/code/models/tt_dit/experimental/lora/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/code/models/tt_dit/experimental/lora/adapter_loader.py b/code/models/tt_dit/experimental/lora/adapter_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..350de02989670bbae56c55c940dcf3b8d7339807
--- /dev/null
+++ b/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.*)\.lora_(?PA|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...``), 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
diff --git a/code/models/tt_dit/experimental/lora/ltx_adapter_loader.py b/code/models/tt_dit/experimental/lora/ltx_adapter_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..6351462d928af30d5b5dc8e6933aa9cc0d5c8f99
--- /dev/null
+++ b/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.*)\.lora_(?PA|B|down|up)(?:\.[^.]+)?\.weight$")
+_SLOT_MAP = {"A": "A", "down": "A", "B": "B", "up": "B"}
+
+# Diffusers LTX checkpoints name blocks "transformer_blocks."; some exports
+# use the shorter "blocks.". 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]
diff --git a/code/models/tt_dit/experimental/lora/promote.py b/code/models/tt_dit/experimental/lora/promote.py
new file mode 100644
index 0000000000000000000000000000000000000000..729d0ff60d94d8d2a46191c3c81e61d428c33160
--- /dev/null
+++ b/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
diff --git a/code/models/tt_dit/experimental/models/AniSora.md b/code/models/tt_dit/experimental/models/AniSora.md
new file mode 100644
index 0000000000000000000000000000000000000000..6e8ca7384ba06234dbb3c5d06a9b9ac59004f636
--- /dev/null
+++ b/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
+```
diff --git a/code/models/tt_dit/experimental/models/Wan2_2_Distill.md b/code/models/tt_dit/experimental/models/Wan2_2_Distill.md
new file mode 100644
index 0000000000000000000000000000000000000000..58c07b65d431d965ccc0e1c2f0e4090ddb66aba0
--- /dev/null
+++ b/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.
diff --git a/code/models/tt_dit/experimental/models/Wan2_2_LoRA.md b/code/models/tt_dit/experimental/models/Wan2_2_LoRA.md
new file mode 100644
index 0000000000000000000000000000000000000000..d9a9b2878c52b4e640a8255cbcf005cd83ee63fa
--- /dev/null
+++ b/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:
+ - `.lora_A.weight` / `.lora_B.weight` (lightx2v native)
+ - `.lora_down.weight` / `.lora_up.weight` (legacy / kohya)
+ - PEFT-style `.lora_A.default.weight` (the trailing `.default`
+ adapter-name segment is tolerated; SVI uses this convention)
+2. **Bias deltas** at `.diff_b` — added to the base `.bias`.
+3. **Full-parameter deltas** at `.diff` — added to the base
+ `.weight`. Some LoRAs ship RMSNorm gammas via this channel.
+4. **Per-module alpha** at `.alpha` — applied as
+ `effective_scale = scale * alpha / rank` for the low-rank pair under
+ the same ``.
+
+### Supported key namespaces (auto-stripped or remapped)
+
+- **lightx2v native:** `blocks..attn.q.lora_A.weight` — passes through.
+- **diffusers prefixes:** `diffusion_model.`, `transformer.`, `unet.`,
+ `model.` — stripped automatically before matching.
+- **kohya / A1111:** `lora_unet_blocks__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.
diff --git a/code/models/tt_dit/experimental/models/Wan2_2_SVI.md b/code/models/tt_dit/experimental/models/Wan2_2_SVI.md
new file mode 100644
index 0000000000000000000000000000000000000000..55229ad3e3c0ff464882506f89f34b5ab7cfd89a
--- /dev/null
+++ b/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 `.lora_A.default.weight` /
+ `.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.
diff --git a/code/models/tt_dit/experimental/pipelines/__init__.py b/code/models/tt_dit/experimental/pipelines/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/code/models/tt_dit/experimental/pipelines/pipeline_anisora.py b/code/models/tt_dit/experimental/pipelines/pipeline_anisora.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb50da20d7ec321331a8e40f3cd43b5866f1ca8c
--- /dev/null
+++ b/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,
+ ),
+ )
diff --git a/code/models/tt_dit/experimental/pipelines/pipeline_wan_distill.py b/code/models/tt_dit/experimental/pipelines/pipeline_wan_distill.py
new file mode 100644
index 0000000000000000000000000000000000000000..17681213901417f9a507a324810f5cc1afe40f68
--- /dev/null
+++ b/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,
+ ),
+ )
diff --git a/code/models/tt_dit/experimental/pipelines/pipeline_wan_lora.py b/code/models/tt_dit/experimental/pipelines/pipeline_wan_lora.py
new file mode 100644
index 0000000000000000000000000000000000000000..90d6e00e5c61d896da7f7c27570a8bc13650d0f7
--- /dev/null
+++ b/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.*)\.lora_(?PA|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 ``.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
diff --git a/code/models/tt_dit/experimental/pipelines/pipeline_wan_runtime_lora.py b/code/models/tt_dit/experimental/pipelines/pipeline_wan_runtime_lora.py
new file mode 100644
index 0000000000000000000000000000000000000000..870b330c5c9303d5f183967a986e16f9f473aca7
--- /dev/null
+++ b/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."""
diff --git a/code/models/tt_dit/experimental/pipelines/pipeline_wan_svi.py b/code/models/tt_dit/experimental/pipelines/pipeline_wan_svi.py
new file mode 100644
index 0000000000000000000000000000000000000000..2955eb045098fdea2d8c36d5711cc5bb1cb5631b
--- /dev/null
+++ b/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)
diff --git a/code/models/tt_dit/experimental/tests/__init__.py b/code/models/tt_dit/experimental/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/code/models/tt_dit/experimental/tests/conftest.py b/code/models/tt_dit/experimental/tests/conftest.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f0ca06eded788a9fc20cd2f612eb183e603d670
--- /dev/null
+++ b/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)
diff --git a/code/models/tt_dit/experimental/tests/test_adapter_loader.py b/code/models/tt_dit/experimental/tests/test_adapter_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..098dafe786e05742cdcba6693daaca38d7bcdc78
--- /dev/null
+++ b/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
diff --git a/code/models/tt_dit/experimental/tests/test_lora_variants.py b/code/models/tt_dit/experimental/tests/test_lora_variants.py
new file mode 100644
index 0000000000000000000000000000000000000000..a56f8620eedf8ebd3577698a0ad3a2ef76bb942e
--- /dev/null
+++ b/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")
diff --git a/code/models/tt_dit/experimental/tests/test_pipeline_anisora.py b/code/models/tt_dit/experimental/tests/test_pipeline_anisora.py
new file mode 100644
index 0000000000000000000000000000000000000000..205a739d7fbd68fe9eba48493eb2167dd96cd881
--- /dev/null
+++ b/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}")
diff --git a/code/models/tt_dit/experimental/tests/test_pipeline_lora.py b/code/models/tt_dit/experimental/tests/test_pipeline_lora.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d2723e88c58cfafff706509876c32adc5ad1086
--- /dev/null
+++ b/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()
diff --git a/code/models/tt_dit/experimental/tests/test_pipeline_wan_distill_i2v.py b/code/models/tt_dit/experimental/tests/test_pipeline_wan_distill_i2v.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca2f1be0a5bfff099af6fa0ffa5e874c98e27127
--- /dev/null
+++ b/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}")
diff --git a/code/models/tt_dit/experimental/tests/test_pipeline_wan_svi.py b/code/models/tt_dit/experimental/tests/test_pipeline_wan_svi.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3f0375c4ff962094209b9d572edf5fdc5b24eeb
--- /dev/null
+++ b/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",
+ )
diff --git a/code/models/tt_dit/experimental/utils/__init__.py b/code/models/tt_dit/experimental/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/code/models/tt_dit/experimental/utils/lightx2v_loader.py b/code/models/tt_dit/experimental/utils/lightx2v_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e0e2b002f18dc22de21e3a66d7d12c53a8f1aa0
--- /dev/null
+++ b/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..`` 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
diff --git a/code/models/tt_dit/models/audio_vae/minimax_h3/.gitkeep b/code/models/tt_dit/models/audio_vae/minimax_h3/.gitkeep
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/code/models/tt_dit/models/audio_vae/minimax_h3/blockings_minimax_h3_audio.py b/code/models/tt_dit/models/audio_vae/minimax_h3/blockings_minimax_h3_audio.py
new file mode 100644
index 0000000000000000000000000000000000000000..44e57a113aa7ba2267b1e9cc17880349b67c3017
--- /dev/null
+++ b/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
diff --git a/code/models/tt_dit/models/audio_vae/minimax_h3/convert_minimax_h3_audio.py b/code/models/tt_dit/models/audio_vae/minimax_h3/convert_minimax_h3_audio.py
new file mode 100644
index 0000000000000000000000000000000000000000..65fb911df56e2efc1b75aeb5105a2ce9447e5a9d
--- /dev/null
+++ b/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.*?)activations\.(?P\d+)\.(?P.*)$")
+
+
+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)
diff --git a/code/models/tt_dit/models/audio_vae/minimax_h3/decoder_minimax_h3_audio.py b/code/models/tt_dit/models/audio_vae/minimax_h3/decoder_minimax_h3_audio.py
new file mode 100644
index 0000000000000000000000000000000000000000..e40ee87a68a820da0517e9dd16988515bc38d798
--- /dev/null
+++ b/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()
diff --git a/code/models/tt_dit/models/audio_vae/minimax_h3/encoder_minimax_h3_audio.py b/code/models/tt_dit/models/audio_vae/minimax_h3/encoder_minimax_h3_audio.py
new file mode 100644
index 0000000000000000000000000000000000000000..d3c68909ada24e2ce2d75f90c5f0b92377a1ea70
--- /dev/null
+++ b/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()
diff --git a/code/models/tt_dit/models/transformers/ltx/attention_ltx.py b/code/models/tt_dit/models/transformers/ltx/attention_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b75bc4afbb6eaf99461e6d3c78e93d5847b060a
--- /dev/null
+++ b/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
diff --git a/code/models/tt_dit/models/transformers/ltx/quant_config.py b/code/models/tt_dit/models/transformers/ltx/quant_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..4b28a04c6bc95f36dff9e4b62f3b9e869004d2cb
--- /dev/null
+++ b/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,
+ )
diff --git a/code/models/tt_dit/models/transformers/ltx/rope_ltx.py b/code/models/tt_dit/models/transformers/ltx/rope_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..d210c0f9aba95cac91dcdc5ea86453e2ee0164a7
--- /dev/null
+++ b/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)
diff --git a/code/models/tt_dit/models/transformers/ltx/transformer_ltx.py b/code/models/tt_dit/models/transformers/ltx/transformer_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..895ee587e2c2db11d87c49c06474fed322d54cae
--- /dev/null
+++ b/code/models/tt_dit/models/transformers/ltx/transformer_ltx.py
@@ -0,0 +1,1249 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+
+import torch
+from loguru import logger
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+import ttnn
+
+from ....layers.embeddings import LTXAdaLayerNormSingle
+from ....layers.feedforward import ParallelFeedForward
+from ....layers.linear import ColParallelLinear, Linear
+from ....layers.module import Module, ModuleList, Parameter
+from ....layers.normalization import DistributedLayerNorm, DistributedRMSNorm
+from ....parallel.config import DiTParallelConfig
+from ....parallel.manager import CCLManager
+from ....utils import cache as cache_module
+from ....utils.fuse_loras import LoraSpec, fuse_loras_into
+from ....utils.substate import pop_substate, rename_substate
+from ....utils.tensor import bf16_tensor
+from ....utils.tracing import traced_function
+from .attention_ltx import LTXAttention
+from .quant_config import LtxQuantProfile
+
+
+def _tile_preserving_chunk0(x: ttnn.Tensor, n: int) -> list[ttnn.Tensor]:
+ """Split ``x`` into ``n`` size-1 slices along dim 0 WITHOUT leaving TILE layout.
+
+ ``ttnn.chunk`` is a host fallback that untilizes to ROW_MAJOR (and, for the
+ ``(coeff,B,1,D)`` AdaLN modulation, the ``(1,B,1,D)`` slices never re-tile), forcing
+ every downstream ``addcmul``/matmul epilogue to re-tilize the shift/scale/gate vectors
+ (the dominant BF16->BF16 tilize cost per block). Slicing dim 0 -- a non-tile outer dim --
+ on the already-TILE ``shifted`` tensor keeps every slice in TILE, so consumers take the
+ fused LLK path instead of the composite (multiply+add w/ per-input tilize) fallback.
+ Output is bit-identical to ``ttnn.chunk`` (pure layout, no value change).
+
+ ``x`` is already TILE here (``TILE scale_shift_table.data`` + temb; ttnn add promotes to
+ TILE), so no explicit re-tilize is needed -- ``ttnn.slice`` is a trace-safe device op that
+ preserves the input layout, whereas ``ttnn.chunk``'s host untilize is what forced RM."""
+ shape = list(x.shape)
+ out = []
+ for i in range(n):
+ starts = [0] * len(shape)
+ ends = list(shape)
+ starts[0] = i
+ ends[0] = i + 1
+ out.append(ttnn.slice(x, starts, ends))
+ return out
+
+
+def build_audio_masks(
+ audio_N: int, audio_N_real: int, *, mesh_device: ttnn.MeshDevice, sp_axis: int
+) -> tuple[ttnn.Tensor | None, ttnn.Tensor | None, ttnn.Tensor | None]:
+ """SDPA attn mask + padding masks for SP-sharded vs gathered audio tokens.
+
+ Returns ``(attn_mask, pad_mask_sp, pad_mask_full)`` — ``pad_mask_sp`` is sharded on the
+ sequence dim for multiply with local audio activations; ``pad_mask_full`` is replicated
+ for multiply after the all_gather on A→V keys. All ``None`` when no padding is needed.
+ """
+ if audio_N <= audio_N_real:
+ return None, None, None
+
+ # Column mask only: real/padded queries are barred from attending TO padded keys.
+ # Do NOT mask padded-query rows to -inf — that makes all attention scores in those
+ # rows -inf → softmax NaN → NaN propagates via padded-token outputs (which we then
+ # multiply by 0; IEEE 0*NaN = NaN, not 0). The pad mask already zeros the padded-query
+ # outputs after attention, so column-only masking is sufficient and numerically safer
+ # at high σ where activations have largest magnitude.
+ mask = torch.zeros(1, 1, audio_N, audio_N)
+ mask[:, :, :, audio_N_real:] = float("-inf")
+ tt_attn_mask = bf16_tensor(mask.to(torch.bfloat16), device=mesh_device, mesh_axis=sp_axis, shard_dim=2)
+
+ pad_mask = torch.ones(1, 1, audio_N, 1, dtype=torch.bfloat16)
+ pad_mask[:, :, audio_N_real:, :] = 0.0
+ tt_pad_mask_sp = bf16_tensor(pad_mask, device=mesh_device, mesh_axis=sp_axis, shard_dim=2)
+ tt_pad_mask_full = bf16_tensor(pad_mask, device=mesh_device)
+ return tt_attn_mask, tt_pad_mask_sp, tt_pad_mask_full
+
+
+def build_video_pad_mask(
+ video_N: int, video_N_real: int, *, mesh_device: ttnn.MeshDevice, sp_axis: int
+) -> ttnn.Tensor | None:
+ """SP-sharded video padding mask ``(1, 1, video_N, 1)``; ``None`` when no padding is needed.
+
+ Multiply the local (sharded) video activations by this to zero padded slots before they
+ propagate downstream (self-attn residual / cross-attn K / FF). No SDPA attn_mask is needed
+ (unlike audio) — video self-attention uses ring SDPA, which masks padded keys via its
+ ``logical_n=video_N_real`` arg.
+ """
+ if video_N <= video_N_real:
+ return None
+ pad_mask = torch.ones(1, 1, video_N, 1, dtype=torch.bfloat16)
+ pad_mask[:, :, video_N_real:, :] = 0.0
+ return bf16_tensor(pad_mask, device=mesh_device, mesh_axis=sp_axis, shard_dim=2)
+
+
+class LTXTransformerBlock(Module):
+ def __init__(
+ self,
+ *,
+ video_dim: int = 4096,
+ audio_dim: int = 2048,
+ video_ffn_dim: int = 16384,
+ audio_ffn_dim: int = 8192,
+ video_num_heads: int = 32,
+ audio_num_heads: int = 32,
+ video_cross_attention_dim: int = 4096,
+ audio_cross_attention_dim: int = 2048,
+ eps: float = 1e-6,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager | None = None,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool = False,
+ has_audio: bool = False,
+ apply_gated_attention: bool = False,
+ cross_attention_adaln: bool = True,
+ quant_config: LtxQuantProfile | None = None,
+ lora_enabled: bool = False,
+ ) -> None:
+ super().__init__()
+
+ self.video_dim = video_dim
+ self.audio_dim = audio_dim
+ self.has_audio = has_audio
+ self.cross_attention_adaln = cross_attention_adaln
+ self.mesh_device = mesh_device
+ self.ccl_manager = ccl_manager
+ self.parallel_config = parallel_config
+
+ rms_norm_kwargs = {
+ "norm_eps": eps,
+ "norm_elementwise_affine": False,
+ "bias": False,
+ "mesh_axis": parallel_config.tensor_parallel.mesh_axis,
+ "mesh_device": mesh_device,
+ "ccl_manager": ccl_manager,
+ }
+
+ attn_kwargs = {
+ "eps": eps,
+ "mesh_device": mesh_device,
+ "ccl_manager": ccl_manager,
+ "parallel_config": parallel_config,
+ "is_fsdp": is_fsdp,
+ "apply_gated_attention": apply_gated_attention,
+ "quant_config": quant_config,
+ "lora_enabled": lora_enabled,
+ }
+
+ # FFN precision: the profile supplies the ff dtypes + casts; no quant_config leaves the
+ # ParallelFeedForward ctor defaults, reproducing bf16. lora_enabled always threads through.
+ ffn_kwargs = {"lora_enabled": lora_enabled}
+ if quant_config is not None:
+ ffn_kwargs.update(quant_config.ffn_kwargs())
+
+ # FSDP fractures FFN weights across the SP axis (on top of the TP fracture);
+ # without it the FFN is only TP-sharded and replicated across every SP device.
+ fsdp_mesh_axis = parallel_config.sequence_parallel.mesh_axis if is_fsdp else None
+
+ self.norm1 = DistributedRMSNorm(embedding_dim=video_dim, **rms_norm_kwargs)
+ self.attn1 = LTXAttention(dim=video_dim, num_heads=video_num_heads, is_self=True, **attn_kwargs)
+ self.norm2 = DistributedRMSNorm(embedding_dim=video_dim, **rms_norm_kwargs)
+ self.attn2 = LTXAttention(
+ dim=video_dim,
+ num_heads=video_num_heads,
+ is_self=False,
+ context_dim=video_cross_attention_dim,
+ **attn_kwargs,
+ )
+ self.norm3 = DistributedRMSNorm(embedding_dim=video_dim, **rms_norm_kwargs)
+ self.ffn = ParallelFeedForward(
+ video_dim,
+ inner_dim=video_ffn_dim,
+ activation_fn="gelu_tanh",
+ bias=True,
+ mesh_device=mesh_device,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ ccl_manager=ccl_manager,
+ fsdp_mesh_axis=fsdp_mesh_axis,
+ **ffn_kwargs,
+ )
+ self.adaln_coeff = 9 if cross_attention_adaln else 6
+ # Outer-param layout (coeff, 1, 1, D): keeps each modulation parameter on the
+ # non-tiled dim 0 so per-block ttnn.chunk(dim=0) is a free tile-aligned slice
+ # (no untilize/slice/tilize). D stays on dim 3 for TP sharding.
+ self.scale_shift_table = Parameter(
+ total_shape=[self.adaln_coeff, 1, 1, video_dim],
+ mesh_axes=[None, None, None, parallel_config.tensor_parallel.mesh_axis],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ if cross_attention_adaln:
+ self.prompt_scale_shift_table = Parameter(
+ total_shape=[2, 1, 1, video_dim],
+ mesh_axes=[None, None, None, None],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+
+ if has_audio:
+ self.audio_norm1 = DistributedRMSNorm(embedding_dim=audio_dim, **rms_norm_kwargs)
+ self.audio_attn1 = LTXAttention(dim=audio_dim, num_heads=audio_num_heads, is_self=True, **attn_kwargs)
+ self.audio_norm2 = DistributedRMSNorm(embedding_dim=audio_dim, **rms_norm_kwargs)
+ self.audio_attn2 = LTXAttention(
+ dim=audio_dim,
+ num_heads=audio_num_heads,
+ is_self=False,
+ context_dim=audio_cross_attention_dim,
+ **attn_kwargs,
+ )
+ self.audio_norm3 = DistributedRMSNorm(embedding_dim=audio_dim, **rms_norm_kwargs)
+ self.audio_ff = ParallelFeedForward(
+ audio_dim,
+ inner_dim=audio_ffn_dim,
+ activation_fn="gelu_tanh",
+ bias=True,
+ mesh_device=mesh_device,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ ccl_manager=ccl_manager,
+ fsdp_mesh_axis=fsdp_mesh_axis,
+ **ffn_kwargs,
+ )
+ self.audio_scale_shift_table = Parameter(
+ total_shape=[self.adaln_coeff, 1, 1, audio_dim],
+ mesh_axes=[None, None, None, parallel_config.tensor_parallel.mesh_axis],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ if cross_attention_adaln:
+ self.audio_prompt_scale_shift_table = Parameter(
+ total_shape=[2, 1, 1, audio_dim],
+ mesh_axes=[None, None, None, None],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+
+ self.audio_to_video_attn = LTXAttention(
+ dim=audio_dim,
+ num_heads=audio_num_heads,
+ is_self=False,
+ context_dim=audio_dim,
+ query_input_dim=video_dim,
+ output_dim=video_dim,
+ **attn_kwargs,
+ )
+ self.video_to_audio_attn = LTXAttention(
+ dim=audio_dim,
+ num_heads=audio_num_heads,
+ is_self=False,
+ context_dim=video_dim,
+ **attn_kwargs,
+ )
+ self.scale_shift_table_a2v_ca_audio = Parameter(
+ total_shape=[5, 1, 1, audio_dim],
+ mesh_axes=[None, None, None, parallel_config.tensor_parallel.mesh_axis],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.scale_shift_table_a2v_ca_video = Parameter(
+ total_shape=[5, 1, 1, video_dim],
+ mesh_axes=[None, None, None, parallel_config.tensor_parallel.mesh_axis],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+
+ if quant_config is not None:
+ self.ff_compute_kernel_config = quant_config.mm_compute_config(mesh_device.arch())
+ else:
+ self.ff_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ mesh_device.arch(),
+ math_fidelity=ttnn.MathFidelity.HiFi2,
+ math_approx_mode=False,
+ fp32_dest_acc_en=True,
+ packer_l1_acc=True,
+ )
+
+ def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+ rename_substate(state, "ff.net.0.proj", "ffn.ff1")
+ rename_substate(state, "ff.net.2", "ffn.ff2")
+
+ if self.has_audio:
+ rename_substate(state, "audio_ff.net.0.proj", "audio_ff.ff1")
+ rename_substate(state, "audio_ff.net.2", "audio_ff.ff2")
+
+ # Scale slots pre-incremented (+1) and cast to bf16 so each runtime
+ # modulation fuses to one ttnn.addcmul(shift, x, scale_p1).
+ tables_scale_idxs: dict[str, list[int]] = {}
+ tables_scale_idxs["scale_shift_table"] = [1, 4, 7] if self.cross_attention_adaln else [1, 4]
+ if self.cross_attention_adaln:
+ tables_scale_idxs["prompt_scale_shift_table"] = [1]
+ else:
+ pop_substate(state, "prompt_scale_shift_table")
+ if self.has_audio:
+ tables_scale_idxs["audio_scale_shift_table"] = [1, 4, 7] if self.cross_attention_adaln else [1, 4]
+ if self.cross_attention_adaln:
+ tables_scale_idxs["audio_prompt_scale_shift_table"] = [1]
+ else:
+ pop_substate(state, "audio_prompt_scale_shift_table")
+ tables_scale_idxs["scale_shift_table_a2v_ca_audio"] = [0, 2]
+ tables_scale_idxs["scale_shift_table_a2v_ca_video"] = [0, 2]
+ for key, scale_idxs in tables_scale_idxs.items():
+ if key in state:
+ # Reference table is [coeff, D]; reshape to outer-param [coeff, 1, 1, D]
+ # so dim 0 indexes the parameters (matches the (coeff,1,1,D) Parameter).
+ t = state[key].unsqueeze(1).unsqueeze(1).clone()
+ t[scale_idxs, :, :, :] += 1.0
+ state[key] = t.to(dtype=torch.bfloat16)
+
+ def _modulated_ffn(self, ffn, norm, x_1BND, shift_ff, scale_ff_p1, gate_ff):
+ """norm -> AdaLN (shift + x * scale_p1) -> gated FFN residual.
+
+ Ring fuses ff1(AG) + ff2 + RS + addcmul; Linear needs explicit AG + plain ffn().
+ """
+ normed = norm(x_1BND)
+ normed = ttnn.addcmul(shift_ff, normed, scale_ff_p1)
+ if self.ccl_manager.topology == ttnn.Topology.Ring:
+ return ffn.forward_fused_addcmul(
+ normed,
+ x_1BND,
+ gate_ff,
+ scalar=1.0,
+ compute_kernel_config=self.ff_compute_kernel_config,
+ parallel_config=self.parallel_config,
+ )
+ else:
+ if self.parallel_config.tensor_parallel.factor > 1:
+ normed = self.ccl_manager.all_gather_persistent_buffer(
+ normed, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ ff_out = ffn(normed, compute_kernel_config=self.ff_compute_kernel_config)
+ return ttnn.addcmul(x_1BND, ff_out, gate_ff)
+
+ def forward(
+ self,
+ video_1BND: ttnn.Tensor,
+ video_prompt: ttnn.Tensor,
+ video_temb: ttnn.Tensor, # (adaln_coeff, B, 1, video_dim)
+ video_N: int,
+ video_rope_cos: ttnn.Tensor,
+ video_rope_sin: ttnn.Tensor,
+ trans_mat: ttnn.Tensor | None = None,
+ video_prompt_temb: ttnn.Tensor | None = None,
+ # Audio args (only used when has_audio=True)
+ audio_1BND: ttnn.Tensor | None = None,
+ audio_prompt: ttnn.Tensor | None = None,
+ audio_temb: ttnn.Tensor | None = None,
+ av_ca_temb: ttnn.Tensor | None = None,
+ audio_N: int = 0,
+ audio_rope_cos: ttnn.Tensor | None = None,
+ audio_rope_sin: ttnn.Tensor | None = None,
+ av_ca_audio_temb: ttnn.Tensor | None = None,
+ audio_prompt_temb: ttnn.Tensor | None = None,
+ video_cross_pe_cos: ttnn.Tensor | None = None,
+ video_cross_pe_sin: ttnn.Tensor | None = None,
+ audio_cross_pe_cos: ttnn.Tensor | None = None,
+ audio_cross_pe_sin: ttnn.Tensor | None = None,
+ audio_cross_pe_cos_full: ttnn.Tensor | None = None,
+ audio_cross_pe_sin_full: ttnn.Tensor | None = None,
+ skip_cross_attn: bool = False,
+ skip_self_attn: bool = False,
+ audio_attn_mask: ttnn.Tensor | None = None,
+ audio_padding_mask: ttnn.Tensor | None = None,
+ audio_padding_mask_full: ttnn.Tensor | None = None,
+ video_padding_mask: ttnn.Tensor | None = None,
+ ) -> ttnn.Tensor | tuple[ttnn.Tensor, ttnn.Tensor]:
+ # Video modulation; `_p1` chunks carry +1 baked into the scale slot (see _prepare_torch_state).
+ shifted_v = self.scale_shift_table.data + video_temb
+ chunks = _tile_preserving_chunk0(shifted_v, self.adaln_coeff)
+ v_shift_sa, v_scale_sa_p1, v_gate_sa = chunks[0], chunks[1], chunks[2]
+ v_shift_ff, v_scale_ff_p1, v_gate_ff = chunks[3], chunks[4], chunks[5]
+ if self.cross_attention_adaln:
+ v_shift_ca, v_scale_ca_p1, v_gate_ca = chunks[6], chunks[7], chunks[8]
+
+ # Video self-attention
+ video_normed = self.norm1(video_1BND)
+ video_normed = ttnn.addcmul(v_shift_sa, video_normed, v_scale_sa_p1)
+ video_1BND = self.attn1(
+ spatial_1BND=video_normed,
+ N=video_N,
+ rope_cos=video_rope_cos,
+ rope_sin=video_rope_sin,
+ trans_mat=trans_mat,
+ addcmul_residual=video_1BND,
+ addcmul_gate=v_gate_sa,
+ skip_qk=skip_self_attn,
+ )
+
+ # Video text cross-attention
+ if self.cross_attention_adaln:
+ video_ca_input = ttnn.addcmul(v_shift_ca, self.norm2(video_1BND), v_scale_ca_p1)
+ if video_prompt_temb is not None:
+ shifted_prompt_v = self.prompt_scale_shift_table.data + video_prompt_temb
+ v_kv_shift, v_kv_scale_p1 = _tile_preserving_chunk0(shifted_prompt_v, 2)
+ video_prompt_mod = ttnn.addcmul(v_kv_shift, video_prompt, v_kv_scale_p1)
+ else:
+ video_prompt_mod = video_prompt
+ # Fold the gated residual (video_1BND + video_ca_out * v_gate_ca) into to_out's epilogue,
+ # mirroring attn1's fused path — avoids a separate full-sequence addcmul pass.
+ video_1BND = self.attn2(
+ spatial_1BND=video_ca_input,
+ N=video_N,
+ prompt_1BLP=video_prompt_mod,
+ kv_replicated=True,
+ addcmul_residual=video_1BND,
+ addcmul_gate=v_gate_ca,
+ )
+ else:
+ # 6-output mode: cross-attention with fixed norm (no AdaLN shift/scale/gate)
+ video_ca_input = self.norm2(video_1BND)
+ video_ca_out = self.attn2(
+ spatial_1BND=video_ca_input, N=video_N, prompt_1BLP=video_prompt, kv_replicated=True
+ )
+ video_1BND = video_1BND + video_ca_out
+
+ if not self.has_audio:
+ # Video-only feed forward
+ return self._modulated_ffn(self.ffn, self.norm3, video_1BND, v_shift_ff, v_scale_ff_p1, v_gate_ff)
+
+ # Audio path (has_audio=True from here)
+ shifted_a = self.audio_scale_shift_table.data + audio_temb
+ a_chunks = ttnn.chunk(shifted_a, self.adaln_coeff, dim=0)
+ a_shift_sa, a_scale_sa_p1, a_gate_sa = a_chunks[0], a_chunks[1], a_chunks[2]
+ a_shift_ff, a_scale_ff_p1, a_gate_ff = a_chunks[3], a_chunks[4], a_chunks[5]
+ if self.cross_attention_adaln:
+ a_shift_ca, a_scale_ca_p1, a_gate_ca = a_chunks[6], a_chunks[7], a_chunks[8]
+
+ # Audio self-attention
+ audio_normed = self.audio_norm1(audio_1BND)
+ audio_normed = ttnn.addcmul(a_shift_sa, audio_normed, a_scale_sa_p1)
+ audio_1BND = self.audio_attn1(
+ spatial_1BND=audio_normed,
+ N=audio_N,
+ rope_cos=audio_rope_cos,
+ rope_sin=audio_rope_sin,
+ trans_mat=trans_mat,
+ addcmul_residual=audio_1BND,
+ addcmul_gate=a_gate_sa,
+ skip_qk=skip_self_attn,
+ attn_mask=audio_attn_mask,
+ )
+
+ # Audio text cross-attention
+ if self.cross_attention_adaln:
+ audio_ca_input = ttnn.addcmul(a_shift_ca, self.audio_norm2(audio_1BND), a_scale_ca_p1)
+ if audio_prompt_temb is not None:
+ shifted_prompt_a = self.audio_prompt_scale_shift_table.data + audio_prompt_temb
+ a_kv_shift, a_kv_scale_p1 = ttnn.chunk(shifted_prompt_a, 2, dim=0)
+ audio_prompt_mod = ttnn.addcmul(a_kv_shift, audio_prompt, a_kv_scale_p1)
+ else:
+ audio_prompt_mod = audio_prompt
+ audio_1BND = self.audio_attn2(
+ spatial_1BND=audio_ca_input,
+ N=audio_N,
+ prompt_1BLP=audio_prompt_mod,
+ kv_replicated=True,
+ addcmul_residual=audio_1BND,
+ addcmul_gate=a_gate_ca,
+ )
+ else:
+ audio_ca_input = self.audio_norm2(audio_1BND)
+ audio_ca_out = self.audio_attn2(
+ spatial_1BND=audio_ca_input, N=audio_N, prompt_1BLP=audio_prompt, kv_replicated=True
+ )
+ audio_1BND = audio_1BND + audio_ca_out
+
+ # Bidirectional A<->V cross-attention
+ if not skip_cross_attn:
+ # Chunk layout [scale, shift, scale, shift, gate]: scale slots (idx 0, 2) have +1 baked in.
+ shifted_av = self.scale_shift_table_a2v_ca_video.data + av_ca_temb
+ v_ca_scale_p1, v_ca_shift, a_ca_scale_v_p1, a_ca_shift_v, v_ca_gate = ttnn.chunk(shifted_av, 5, dim=0)
+
+ shifted_av_a = self.scale_shift_table_a2v_ca_audio.data + av_ca_audio_temb
+ a_scale_a2v_p1, a_shift_a2v, a_scale_v2a_p1, a_shift_v2a, a_ca_gate = ttnn.chunk(shifted_av_a, 5, dim=0)
+
+ video_normed_xattn = self.norm3(video_1BND)
+ audio_normed_xattn = self.audio_norm3(audio_1BND)
+
+ # A→V: audio provides context for video
+ video_q_a2v = ttnn.addcmul(v_ca_shift, video_normed_xattn, v_ca_scale_p1)
+ audio_kv_a2v = ttnn.addcmul(a_shift_a2v, audio_normed_xattn, a_scale_a2v_p1)
+ if self.parallel_config.sequence_parallel.factor > 1:
+ audio_kv_a2v = self.ccl_manager.all_gather_persistent_buffer(
+ audio_kv_a2v, dim=2, mesh_axis=self.parallel_config.sequence_parallel.mesh_axis
+ )
+ # Zero padded audio tokens before to_kv so they don't contribute to A->V cross-attention.
+ pad_mask_a2v = audio_padding_mask_full if audio_padding_mask_full is not None else audio_padding_mask
+ if pad_mask_a2v is not None:
+ audio_kv_a2v = ttnn.multiply(audio_kv_a2v, pad_mask_a2v)
+ a2v_output = self.audio_to_video_attn(
+ spatial_1BND=video_q_a2v,
+ N=video_N,
+ prompt_1BLP=audio_kv_a2v,
+ rope_cos=video_cross_pe_cos,
+ rope_sin=video_cross_pe_sin,
+ k_rope_cos=audio_cross_pe_cos_full,
+ k_rope_sin=audio_cross_pe_sin_full,
+ trans_mat=trans_mat,
+ addcmul_residual=video_1BND,
+ addcmul_gate=v_ca_gate,
+ )
+ video_1BND = a2v_output
+
+ # V→A: video provides context for audio
+ audio_q_v2a = ttnn.addcmul(a_shift_v2a, audio_normed_xattn, a_scale_v2a_p1)
+ video_kv_v2a = ttnn.addcmul(a_ca_shift_v, video_normed_xattn, a_ca_scale_v_p1)
+ # Zero padded video tokens (on the SP-local shard) after the affine, before to_kv.
+ if video_padding_mask is not None:
+ video_kv_v2a = ttnn.multiply(video_kv_v2a, video_padding_mask)
+ v2a_output = self.video_to_audio_attn(
+ spatial_1BND=audio_q_v2a,
+ N=audio_N,
+ prompt_1BLP=video_kv_v2a,
+ rope_cos=audio_cross_pe_cos,
+ rope_sin=audio_cross_pe_sin,
+ # Ring cross keeps video K/V SP-sharded: pass the sharded K-rope and kv_logical_n so
+ # the ring SDPA gathers internally instead of a separate K/V all-gather.
+ k_rope_cos=video_cross_pe_cos,
+ k_rope_sin=video_cross_pe_sin,
+ kv_logical_n=video_N,
+ trans_mat=trans_mat,
+ addcmul_residual=audio_1BND,
+ addcmul_gate=a_ca_gate,
+ )
+ audio_1BND = v2a_output
+
+ # Video feed forward
+ video_1BND = self._modulated_ffn(self.ffn, self.norm3, video_1BND, v_shift_ff, v_scale_ff_p1, v_gate_ff)
+
+ # Audio feed forward
+ audio_1BND = self._modulated_ffn(
+ self.audio_ff, self.audio_norm3, audio_1BND, a_shift_ff, a_scale_ff_p1, a_gate_ff
+ )
+
+ return video_1BND, audio_1BND
+
+
+class LTXTransformerModel(Module):
+ def __init__(
+ self,
+ *,
+ # Video config
+ num_attention_heads: int = 32,
+ attention_head_dim: int = 128,
+ in_channels: int = 128,
+ out_channels: int = 128,
+ num_layers: int = 48,
+ cross_attention_dim: int = 4096,
+ # Audio config (only used when has_audio=True)
+ audio_num_attention_heads: int = 32,
+ audio_attention_head_dim: int = 64,
+ audio_in_channels: int = 128,
+ audio_out_channels: int = 128,
+ audio_cross_attention_dim: int = 2048,
+ # Common
+ norm_eps: float = 1e-6,
+ ffn_mult: int = 4,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager | None = None,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool = False,
+ has_audio: bool = False,
+ apply_gated_attention: bool = False,
+ cross_attention_adaln: bool = True,
+ lora_enabled: bool = False,
+ image_conditioning: bool = False,
+ quant_config: LtxQuantProfile | None = None,
+ ) -> None:
+ super().__init__()
+
+ self.inner_dim = num_attention_heads * attention_head_dim
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.num_layers = num_layers
+ self.has_audio = has_audio
+ self.cross_attention_adaln = cross_attention_adaln
+ self.lora_enabled = lora_enabled
+ # I2V: video AdaLN modulation is per-token (denoise_mask * sigma) instead of batch-scalar.
+ # Audio / prompt / A<->V cross AdaLN stay batch-scalar regardless.
+ self.image_conditioning = image_conditioning
+ self.mesh_device = mesh_device
+ self.ccl_manager = ccl_manager
+ self.parallel_config = parallel_config
+
+ if has_audio:
+ self.audio_inner_dim = audio_num_attention_heads * audio_attention_head_dim
+
+ self.patchify_proj = ColParallelLinear(
+ in_channels,
+ self.inner_dim,
+ bias=True,
+ mesh_device=mesh_device,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ ccl_manager=ccl_manager,
+ )
+ adaln_coeff = 9 if cross_attention_adaln else 6
+ self.adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.inner_dim,
+ embedding_coefficient=adaln_coeff,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ if cross_attention_adaln:
+ self.prompt_adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.inner_dim,
+ embedding_coefficient=2,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.norm_out = DistributedLayerNorm(
+ self.inner_dim,
+ norm_eps=norm_eps,
+ norm_elementwise_affine=False,
+ bias=False,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+ self.scale_shift_table = Parameter(
+ total_shape=[1, 1, 2, self.inner_dim],
+ mesh_axes=[None, None, None, parallel_config.tensor_parallel.mesh_axis],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.proj_out = Linear(self.inner_dim, out_channels, bias=True, mesh_device=mesh_device)
+
+ if has_audio:
+ self.audio_patchify_proj = ColParallelLinear(
+ audio_in_channels,
+ self.audio_inner_dim,
+ bias=True,
+ mesh_device=mesh_device,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ ccl_manager=ccl_manager,
+ )
+ self.audio_adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.audio_inner_dim,
+ embedding_coefficient=adaln_coeff,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ if cross_attention_adaln:
+ self.audio_prompt_adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.audio_inner_dim,
+ embedding_coefficient=2,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.audio_norm_out = DistributedLayerNorm(
+ self.audio_inner_dim,
+ norm_eps=norm_eps,
+ norm_elementwise_affine=False,
+ bias=False,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+ self.audio_scale_shift_table = Parameter(
+ total_shape=[1, 1, 2, self.audio_inner_dim],
+ mesh_axes=[None, None, None, parallel_config.tensor_parallel.mesh_axis],
+ device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.audio_proj_out = Linear(self.audio_inner_dim, audio_out_channels, bias=True, mesh_device=mesh_device)
+
+ # Cross-attention timestep conditioning
+ self.av_ca_video_scale_shift_adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.inner_dim,
+ embedding_coefficient=4,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.av_ca_audio_scale_shift_adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.audio_inner_dim,
+ embedding_coefficient=4,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.av_ca_a2v_gate_adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.inner_dim,
+ embedding_coefficient=1,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+ self.av_ca_v2a_gate_adaln_single = LTXAdaLayerNormSingle(
+ embedding_dim=self.audio_inner_dim,
+ embedding_coefficient=1,
+ mesh_device=mesh_device,
+ dtype=ttnn.bfloat16,
+ )
+
+ video_ffn_dim = self.inner_dim * ffn_mult
+ audio_ffn_dim = self.audio_inner_dim * ffn_mult if has_audio else 0
+ self.transformer_blocks = ModuleList()
+ for _ in range(num_layers):
+ self.transformer_blocks.append(
+ LTXTransformerBlock(
+ video_dim=self.inner_dim,
+ audio_dim=self.audio_inner_dim if has_audio else 0,
+ video_ffn_dim=video_ffn_dim,
+ audio_ffn_dim=audio_ffn_dim,
+ video_num_heads=num_attention_heads,
+ audio_num_heads=audio_num_attention_heads if has_audio else 0,
+ video_cross_attention_dim=cross_attention_dim,
+ audio_cross_attention_dim=audio_cross_attention_dim if has_audio else 0,
+ eps=norm_eps,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ apply_gated_attention=apply_gated_attention,
+ cross_attention_adaln=cross_attention_adaln,
+ quant_config=quant_config,
+ lora_enabled=lora_enabled,
+ )
+ )
+
+ # Output-proj config matches WanTransformer's hifi4 config (packer_l1_acc=True).
+ self.hifi4_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ mesh_device.arch(),
+ math_fidelity=ttnn.MathFidelity.HiFi4,
+ math_approx_mode=False,
+ fp32_dest_acc_en=True,
+ packer_l1_acc=True,
+ )
+
+ def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+ # Output-proj table is [shift, scale]; bake +1 into the scale slot and cast to bf16.
+ tables = ["scale_shift_table"]
+ if self.has_audio:
+ tables.append("audio_scale_shift_table")
+ for key in tables:
+ if key in state:
+ t = state[key].unsqueeze(0).unsqueeze(0).clone()
+ t[:, :, 1, :] += 1.0
+ state[key] = t.to(dtype=torch.bfloat16)
+
+ # Remove keys not implemented in TTNN
+ pop_substate(state, "video_embeddings_connector")
+ pop_substate(state, "audio_embeddings_connector")
+ # Caption projection is a 19B-distilled-only artifact, not implemented in TTNN.
+ pop_substate(state, "caption_projection")
+ if self.has_audio:
+ pop_substate(state, "audio_caption_projection")
+ # 6-output mode: no prompt_adaln modules
+ if not self.cross_attention_adaln:
+ pop_substate(state, "prompt_adaln_single")
+ if self.has_audio:
+ pop_substate(state, "audio_prompt_adaln_single")
+
+ def weight_cache_subfolder(self) -> str:
+ """Cache subfolder naming this model's weight LAYOUT, not just its parallel config.
+
+ On a Ring topology the attention gate is fused into the QKV/Q projection
+ (``LTXAttention.fuse_gate``): ``to_qkv`` grows by one padded tile per device and the
+ separate gate tensor disappears, so the cached tensorbins have different shapes and a
+ different count than the Linear-topology layout. Both used to land under ``transformer/``
+ under an identical key, and the second topology to run died loading the first one's cache.
+
+ A ``transformer/`` cache written by a Ring run before this split is stale (nothing reads it any
+ more) and can be deleted; Ring now writes ``transformer_fusedgate/``.
+ """
+ fused = any(getattr(block.attn1, "fuse_gate", False) for block in self.transformer_blocks)
+ return "transformer_fusedgate" if fused else "transformer"
+
+ def forward(
+ self,
+ # Video
+ video_1BNI_torch: torch.Tensor,
+ video_prompt_1BLP: ttnn.Tensor,
+ video_rope_cos: ttnn.Tensor,
+ video_rope_sin: ttnn.Tensor,
+ video_N: int,
+ # Shared
+ trans_mat: ttnn.Tensor | None,
+ timestep_torch: torch.Tensor,
+ # I2V per-token video timestep (B, N); required when image_conditioning=True.
+ video_timestep_torch: torch.Tensor | None = None,
+ # Audio (only used when has_audio=True)
+ audio_1BNI_torch: torch.Tensor | None = None,
+ audio_prompt_1BLP: ttnn.Tensor | None = None,
+ audio_rope_cos: ttnn.Tensor | None = None,
+ audio_rope_sin: ttnn.Tensor | None = None,
+ audio_N: int = 0,
+ # Cross-modal positional embeddings (optional)
+ video_cross_pe_cos: ttnn.Tensor | None = None,
+ video_cross_pe_sin: ttnn.Tensor | None = None,
+ audio_cross_pe_cos: ttnn.Tensor | None = None,
+ audio_cross_pe_sin: ttnn.Tensor | None = None,
+ audio_cross_pe_cos_full: ttnn.Tensor | None = None,
+ audio_cross_pe_sin_full: ttnn.Tensor | None = None,
+ skip_cross_attn: bool = False,
+ skip_self_attn_blocks: list[int] | None = None,
+ audio_attn_mask: ttnn.Tensor | None = None,
+ audio_padding_mask: ttnn.Tensor | None = None,
+ audio_padding_mask_full: ttnn.Tensor | None = None,
+ video_padding_mask: ttnn.Tensor | None = None,
+ ) -> ttnn.Tensor | tuple[ttnn.Tensor, ttnn.Tensor]:
+ """Host entry: upload torch latents/timestep, then run the device-only inner_step."""
+ sp_axis = self.parallel_config.sequence_parallel.mesh_axis
+ video_1BNI = bf16_tensor(video_1BNI_torch, device=self.mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ B_size = video_1BNI_torch.shape[1]
+ # Scalar timestep (1,1,B,1) drives audio / prompt / A<->V cross modulation.
+ timestep = bf16_tensor(timestep_torch.reshape(1, 1, B_size, 1) * 1000.0, device=self.mesh_device)
+ # Per-token video timestep (1,1,B*N,1), SP-sharded on dim 2 to match video_1BNI (shard_dim=-2).
+ video_timestep = None
+ if self.image_conditioning:
+ assert video_timestep_torch is not None, "video_timestep_torch required when image_conditioning=True"
+ video_timestep = bf16_tensor(
+ video_timestep_torch.reshape(1, 1, -1, 1) * 1000.0,
+ device=self.mesh_device,
+ mesh_axis=sp_axis,
+ shard_dim=-2,
+ )
+ audio_1BNI = (
+ bf16_tensor(audio_1BNI_torch, device=self.mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ if self.has_audio
+ else None
+ )
+ return self.inner_step(
+ video_1BNI=video_1BNI,
+ timestep=timestep,
+ video_timestep=video_timestep,
+ audio_1BNI=audio_1BNI,
+ video_prompt_1BLP=video_prompt_1BLP,
+ video_rope_cos=video_rope_cos,
+ video_rope_sin=video_rope_sin,
+ video_N=video_N,
+ trans_mat=trans_mat,
+ audio_prompt_1BLP=audio_prompt_1BLP,
+ audio_rope_cos=audio_rope_cos,
+ audio_rope_sin=audio_rope_sin,
+ audio_N=audio_N,
+ video_cross_pe_cos=video_cross_pe_cos,
+ video_cross_pe_sin=video_cross_pe_sin,
+ audio_cross_pe_cos=audio_cross_pe_cos,
+ audio_cross_pe_sin=audio_cross_pe_sin,
+ audio_cross_pe_cos_full=audio_cross_pe_cos_full,
+ audio_cross_pe_sin_full=audio_cross_pe_sin_full,
+ skip_cross_attn=skip_cross_attn,
+ skip_self_attn_blocks=skip_self_attn_blocks,
+ audio_attn_mask=audio_attn_mask,
+ audio_padding_mask=audio_padding_mask,
+ audio_padding_mask_full=audio_padding_mask_full,
+ video_padding_mask=video_padding_mask,
+ )
+
+ @traced_function(device=lambda self: self.mesh_device, clone_prep_inputs=False, prep_run=False)
+ def inner_step(
+ self,
+ *,
+ video_1BNI: ttnn.Tensor,
+ timestep: ttnn.Tensor,
+ video_timestep: ttnn.Tensor | None = None,
+ video_ts_pair: ttnn.Tensor | None = None,
+ video_pin_mask: ttnn.Tensor | None = None,
+ audio_1BNI: ttnn.Tensor | None = None,
+ video_prompt_1BLP: ttnn.Tensor = None,
+ video_rope_cos: ttnn.Tensor = None,
+ video_rope_sin: ttnn.Tensor = None,
+ video_N: int = 0,
+ trans_mat: ttnn.Tensor | None = None,
+ audio_prompt_1BLP: ttnn.Tensor | None = None,
+ audio_rope_cos: ttnn.Tensor | None = None,
+ audio_rope_sin: ttnn.Tensor | None = None,
+ audio_N: int = 0,
+ video_cross_pe_cos: ttnn.Tensor | None = None,
+ video_cross_pe_sin: ttnn.Tensor | None = None,
+ audio_cross_pe_cos: ttnn.Tensor | None = None,
+ audio_cross_pe_sin: ttnn.Tensor | None = None,
+ audio_cross_pe_cos_full: ttnn.Tensor | None = None,
+ audio_cross_pe_sin_full: ttnn.Tensor | None = None,
+ skip_cross_attn: bool = False,
+ skip_self_attn_blocks: list[int] | None = None,
+ audio_attn_mask: ttnn.Tensor | None = None,
+ audio_padding_mask: ttnn.Tensor | None = None,
+ audio_padding_mask_full: ttnn.Tensor | None = None,
+ video_padding_mask: ttnn.Tensor | None = None,
+ gather_output: bool = True,
+ ) -> ttnn.Tensor | tuple[ttnn.Tensor, ttnn.Tensor]:
+ """Device-only, trace-capturable denoising step. All tensor args are ttnn (no torch).
+ ``video_1BNI``/``audio_1BNI``/``timestep`` are the per-step inputs; the rest are
+ constant across a denoise stage.
+
+ ``gather_output`` SP-gathers the velocity outputs to full sequence length (default,
+ for the host-Euler path). Pass ``False`` to keep them SP-sharded so an on-device
+ solver can step the latent without leaving the device (the traced WAN pattern)."""
+ # Video modulation (6 or 9 params depending on cross_attention_adaln).
+ adaln_coeff = 9 if self.cross_attention_adaln else 6
+ # I2V compact path: per-token timestep has 2 values (pinned frame-0 vs. sigma), passed as a
+ # (1,1,2,1) pair + {0,1} pin mask. Blend per token to avoid the dense (1,1,N,coeff*D) modulation.
+ compact_i2v = self.image_conditioning and video_ts_pair is not None and video_pin_mask is not None
+ if compact_i2v:
+ N = video_pin_mask.shape[2]
+ mod_pair, emb_pair = self.adaln_single(video_ts_pair) # (1,1,2,coeff*D), (1,1,2,D)
+ mod_pair = ttnn.reshape(mod_pair, (1, 2, adaln_coeff, self.inner_dim))
+ if self.parallel_config.tensor_parallel.factor > 1:
+ mod_pair = ttnn.mesh_partition(
+ mod_pair, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ mod_pair = ttnn.permute(mod_pair, (2, 0, 1, 3)) # (coeff,1,2,Dloc)
+ mod_pin, mod_base = ttnn.chunk(mod_pair, 2, dim=2) # each (coeff,1,1,Dloc)
+ ttnn.deallocate(mod_pair)
+ # video_mod = base + (pin - base) * mask, materialized per token via broadcast-safe ops.
+ mod_delta = ttnn.sub(mod_pin, mod_base)
+ mod_delta = ttnn.repeat(mod_delta, ttnn.Shape([1, 1, N, 1])) # (coeff,1,N,Dloc)
+ mod_delta = ttnn.mul(mod_delta, video_pin_mask) # mask broadcasts over coeff/Dloc
+ video_mod_CB1D = ttnn.add(mod_base, mod_delta) # base broadcasts over N
+ ttnn.deallocate(mod_delta)
+ ttnn.deallocate(mod_pin)
+ ttnn.deallocate(mod_base)
+ # Embedded timestep (for norm_out), same per-token blend, kept full-D.
+ emb_pin, emb_base = ttnn.chunk(emb_pair, 2, dim=2) # each (1,1,1,D)
+ ttnn.deallocate(emb_pair)
+ emb_delta = ttnn.sub(emb_pin, emb_base)
+ emb_delta = ttnn.repeat(emb_delta, ttnn.Shape([1, 1, N, 1])) # (1,1,N,D)
+ emb_delta = ttnn.mul(emb_delta, video_pin_mask)
+ video_emb_ts = ttnn.add(emb_base, emb_delta)
+ ttnn.deallocate(emb_delta)
+ ttnn.deallocate(emb_pin)
+ ttnn.deallocate(emb_base)
+ B = 1
+ else:
+ # I2V (dense): feed the per-token timestep so each token gets its own AdaLN modulation.
+ video_ts = video_timestep if self.image_conditioning else timestep
+ video_modulation, video_emb_ts = self.adaln_single(video_ts)
+ # dim 2 is B (scalar) or B*N (per-token, B=1 -> N_local on the SP shard).
+ X = video_modulation.shape[2]
+ video_mod_CB1D = ttnn.reshape(video_modulation, (1, X, adaln_coeff, self.inner_dim))
+ if self.parallel_config.tensor_parallel.factor > 1:
+ video_mod_CB1D = ttnn.mesh_partition(
+ video_mod_CB1D, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ # Move the coeff axis to dim 0 once per step so each block's chunk(dim=0) is free.
+ # Scalar: (1,B,coeff,D) -> (coeff,B,1,D). Per-token: (1,N,coeff,D) -> (coeff,1,N,D).
+ if self.image_conditioning:
+ video_mod_CB1D = ttnn.permute(video_mod_CB1D, (2, 0, 1, 3))
+ else:
+ video_mod_CB1D = ttnn.permute(video_mod_CB1D, (2, 1, 0, 3))
+ B = 1 if self.image_conditioning else X
+
+ # Video prompt modulation (2 params, only for 9-output mode)
+ video_prompt_2B1D = None
+ if self.cross_attention_adaln:
+ video_prompt_mod, _ = self.prompt_adaln_single(timestep)
+ video_prompt_2B1D = ttnn.reshape(video_prompt_mod, (1, B, 2, self.inner_dim))
+ video_prompt_2B1D = ttnn.permute(video_prompt_2B1D, (2, 1, 0, 3))
+
+ # Audio modulation (only when has_audio)
+ audio_mod_CB1D = None
+ audio_prompt_2B1D = None
+ av_ca_video_temb = None
+ av_ca_audio_temb = None
+ audio_emb_ts = None
+
+ if self.has_audio:
+ audio_modulation, audio_emb_ts = self.audio_adaln_single(timestep)
+ audio_mod_CB1D = ttnn.reshape(audio_modulation, (1, B, adaln_coeff, self.audio_inner_dim))
+ if self.parallel_config.tensor_parallel.factor > 1:
+ audio_mod_CB1D = ttnn.mesh_partition(
+ audio_mod_CB1D, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ audio_mod_CB1D = ttnn.permute(audio_mod_CB1D, (2, 1, 0, 3))
+
+ if self.cross_attention_adaln:
+ audio_prompt_mod, _ = self.audio_prompt_adaln_single(timestep)
+ audio_prompt_2B1D = ttnn.reshape(audio_prompt_mod, (1, B, 2, self.audio_inner_dim))
+ audio_prompt_2B1D = ttnn.permute(audio_prompt_2B1D, (2, 1, 0, 3))
+
+ # Cross-attention timestep conditioning; concat on the param axis (dim 0).
+ av_ca_video_ss, _ = self.av_ca_video_scale_shift_adaln_single(timestep)
+ av_ca_video_ss = ttnn.permute(ttnn.reshape(av_ca_video_ss, (1, B, 4, self.inner_dim)), (2, 1, 0, 3))
+ av_ca_a2v_gate, _ = self.av_ca_a2v_gate_adaln_single(timestep)
+ av_ca_a2v_gate = ttnn.permute(ttnn.reshape(av_ca_a2v_gate, (1, B, 1, self.inner_dim)), (2, 1, 0, 3))
+ av_ca_video_temb = ttnn.concat([av_ca_video_ss, av_ca_a2v_gate], dim=0)
+
+ av_ca_audio_ss, _ = self.av_ca_audio_scale_shift_adaln_single(timestep)
+ av_ca_audio_ss = ttnn.permute(ttnn.reshape(av_ca_audio_ss, (1, B, 4, self.audio_inner_dim)), (2, 1, 0, 3))
+ av_ca_v2a_gate, _ = self.av_ca_v2a_gate_adaln_single(timestep)
+ av_ca_v2a_gate = ttnn.permute(ttnn.reshape(av_ca_v2a_gate, (1, B, 1, self.audio_inner_dim)), (2, 1, 0, 3))
+ av_ca_audio_temb = ttnn.concat([av_ca_audio_ss, av_ca_v2a_gate], dim=0)
+
+ if self.parallel_config.tensor_parallel.factor > 1:
+ av_ca_video_temb = ttnn.mesh_partition(
+ av_ca_video_temb, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ av_ca_audio_temb = ttnn.mesh_partition(
+ av_ca_audio_temb, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+
+ # Patchify
+ video_1BND = self.patchify_proj(video_1BNI)
+ audio_1BND = self.audio_patchify_proj(audio_1BNI) if self.has_audio else None
+
+ skip_self_attn_set = frozenset(skip_self_attn_blocks) if skip_self_attn_blocks else frozenset()
+
+ # Transformer blocks
+ for block_idx, block in enumerate(self.transformer_blocks):
+ result = block(
+ video_1BND=video_1BND,
+ video_prompt=video_prompt_1BLP,
+ video_temb=video_mod_CB1D,
+ video_N=video_N,
+ video_rope_cos=video_rope_cos,
+ video_rope_sin=video_rope_sin,
+ trans_mat=trans_mat,
+ video_prompt_temb=video_prompt_2B1D,
+ audio_1BND=audio_1BND,
+ audio_prompt=audio_prompt_1BLP,
+ audio_temb=audio_mod_CB1D,
+ av_ca_temb=av_ca_video_temb,
+ audio_N=audio_N,
+ audio_rope_cos=audio_rope_cos,
+ audio_rope_sin=audio_rope_sin,
+ av_ca_audio_temb=av_ca_audio_temb,
+ audio_prompt_temb=audio_prompt_2B1D,
+ video_cross_pe_cos=video_cross_pe_cos,
+ video_cross_pe_sin=video_cross_pe_sin,
+ audio_cross_pe_cos=audio_cross_pe_cos,
+ audio_cross_pe_sin=audio_cross_pe_sin,
+ audio_cross_pe_cos_full=audio_cross_pe_cos_full,
+ audio_cross_pe_sin_full=audio_cross_pe_sin_full,
+ skip_cross_attn=skip_cross_attn,
+ skip_self_attn=block_idx in skip_self_attn_set,
+ audio_attn_mask=audio_attn_mask,
+ audio_padding_mask=audio_padding_mask,
+ audio_padding_mask_full=audio_padding_mask_full,
+ video_padding_mask=video_padding_mask,
+ )
+ if self.has_audio:
+ video_1BND, audio_1BND = result
+ else:
+ video_1BND = result
+
+ v_inner_local = video_emb_ts.shape[-1]
+ if self.image_conditioning:
+ # Per-token (video_emb_ts is (1,1,N,D)): split the (shift, scale) table, broadcast-add per token.
+ v_emb_1B1D = video_emb_ts
+ if self.parallel_config.tensor_parallel.factor > 1:
+ v_emb_1B1D = ttnn.mesh_partition(
+ v_emb_1B1D, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ sst_shift, sst_scale_p1 = ttnn.chunk(self.scale_shift_table.data, 2, dim=2) # each (1,1,1,Dloc)
+ v_shift_out = ttnn.add(sst_shift, v_emb_1B1D) # (1,1,N,Dloc)
+ v_scale_out_p1 = ttnn.add(sst_scale_p1, v_emb_1B1D)
+ else:
+ v_emb_1B1D = ttnn.reshape(video_emb_ts, (1, B, 1, v_inner_local))
+ if self.parallel_config.tensor_parallel.factor > 1:
+ v_emb_1B1D = ttnn.mesh_partition(
+ v_emb_1B1D, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ shifted_v = self.scale_shift_table.data + v_emb_1B1D
+ v_shift_out, v_scale_out_p1 = ttnn.chunk(shifted_v, 2, dim=2)
+ if self.image_conditioning:
+ # Per-token: fused norm_out gamma/beta can't carry per-token scale/shift, so plain
+ # layernorm then manual shift + normed * scale_p1, in fp32 to match the T2V branch.
+ video_1BND = self.norm_out(video_1BND, dtype=ttnn.float32)
+ v_shift_out = ttnn.typecast(v_shift_out, ttnn.float32)
+ v_scale_out_p1 = ttnn.typecast(v_scale_out_p1, ttnn.float32)
+ video_1BND = ttnn.addcmul(v_shift_out, video_1BND, v_scale_out_p1)
+ else:
+ # Fuse the AdaLN (1 + scale) * normed + shift modulation into norm_out (WAN pattern).
+ video_1BND = self.norm_out(
+ video_1BND, dynamic_weight=v_scale_out_p1, dynamic_bias=v_shift_out, dtype=ttnn.float32
+ )
+ if self.parallel_config.tensor_parallel.factor > 1:
+ video_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ video_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ video_out = self.proj_out(
+ video_1BND, compute_kernel_config=self.hifi4_compute_kernel_config, dtype=ttnn.float32
+ )
+
+ # Output projection — audio
+ if self.has_audio:
+ a_inner_local = audio_emb_ts.shape[-1]
+ a_emb_1B1D = ttnn.reshape(audio_emb_ts, (1, B, 1, a_inner_local))
+ if self.parallel_config.tensor_parallel.factor > 1:
+ a_emb_1B1D = ttnn.mesh_partition(
+ a_emb_1B1D, dim=3, cluster_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ shifted_a = self.audio_scale_shift_table.data + a_emb_1B1D
+ a_shift_out, a_scale_out_p1 = ttnn.chunk(shifted_a, 2, dim=2)
+ # Fuse the AdaLN (1 + scale) * normed + shift modulation into audio_norm_out (WAN pattern).
+ audio_1BND = self.audio_norm_out(
+ audio_1BND, dynamic_weight=a_scale_out_p1, dynamic_bias=a_shift_out, dtype=ttnn.float32
+ )
+ if self.parallel_config.tensor_parallel.factor > 1:
+ audio_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ audio_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+ audio_out = self.audio_proj_out(
+ audio_1BND, compute_kernel_config=self.hifi4_compute_kernel_config, dtype=ttnn.float32
+ )
+
+ # SP gather (skipped when an on-device solver steps the still-sharded latent)
+ if gather_output and self.parallel_config.sequence_parallel.factor > 1:
+ video_out = self.ccl_manager.all_gather_persistent_buffer(
+ video_out, dim=2, mesh_axis=self.parallel_config.sequence_parallel.mesh_axis
+ )
+ if self.has_audio:
+ audio_out = self.ccl_manager.all_gather_persistent_buffer(
+ audio_out, dim=2, mesh_axis=self.parallel_config.sequence_parallel.mesh_axis
+ )
+
+ if self.has_audio:
+ return video_out, audio_out
+ return video_out
+
+ @staticmethod
+ def device_to_host(
+ tt_tensor: ttnn.Tensor,
+ *,
+ ccl_manager: CCLManager | None = None,
+ parallel_config: DiTParallelConfig | None = None,
+ sp_already_gathered: bool = False,
+ tp_already_gathered: bool = False,
+ ) -> torch.Tensor:
+ """Assemble SP/TP shards on host (do not read device 0 alone on multi-device meshes).
+
+ ``inner_step`` already AllGathers velocities on dim=2 before return; pass
+ ``sp_already_gathered=True`` from the denoise loop to avoid concatenating the
+ full sequence twice (audible doubling / ghost vocals).
+ """
+ if ccl_manager is not None and parallel_config is not None:
+ mesh_dims: list[int | None] = [None, None]
+ if not tp_already_gathered and parallel_config.tensor_parallel.factor > 1:
+ mesh_dims[parallel_config.tensor_parallel.mesh_axis] = 3
+ if not sp_already_gathered and parallel_config.sequence_parallel.factor > 1:
+ mesh_dims[parallel_config.sequence_parallel.mesh_axis] = 2
+ return ccl_manager.device_to_host(tt_tensor, mesh_dims).float().clone()
+ return ttnn.to_torch(ttnn.get_device_tensors(tt_tensor)[0]).float()
+
+
+class LTXTransformerCheckpoint:
+ """An LTX transformer checkpoint: parses variant config, builds transformers, and loads weights.
+
+ Mirrors ``WanCheckpoint`` but for the single-file LTX safetensors layout. The pipeline-level
+ flags that vary per instance (``has_audio`` / ``image_conditioning``) are passed into ``build``
+ explicitly — the checkpoint never reads pipeline state. LoRA specs vary per variant, so they
+ are threaded through ``state_dict`` / ``cache_name`` / ``load`` rather than stored.
+ """
+
+ def __init__(self, checkpoint_path: str, *, inner_dim: int) -> None:
+ self._checkpoint_path = checkpoint_path
+ # Transformer + connector detection (key scan + one tensor-shape read). No tensor loads.
+ with safe_open(checkpoint_path, framework="pt") as f:
+ keys = list(f.keys())
+ adaln_key = "model.diffusion_model.adaln_single.linear.weight"
+ if adaln_key in keys:
+ self.cross_attention_adaln = f.get_tensor(adaln_key).shape[0] > 6 * inner_dim
+ else:
+ self.cross_attention_adaln = True
+ self.has_gate = any("to_gate_logits" in k for k in keys)
+ logger.info(f"Detected: has_gate={self.has_gate}, cross_attention_adaln={self.cross_attention_adaln}")
+
+ def state_dict(self, lora_specs: list[LoraSpec]) -> dict[str, torch.Tensor]:
+ """Load + LoRA-fuse the transformer state dict from safetensors. Only
+ invoked on cache miss by ``cache_module.load_model``."""
+ logger.info(f"Transformer cache miss — loading safetensors: {self._checkpoint_path}")
+ raw = load_file(self._checkpoint_path)
+ prefix = "model.diffusion_model."
+ sd = {k[len(prefix) :]: v for k, v in raw.items() if k.startswith(prefix)}
+ if lora_specs:
+ sd = fuse_loras_into(sd, lora_specs)
+ return sd
+
+ def cache_name(self, lora_specs: list[LoraSpec], quant_tag: str | None = None) -> str:
+ """Cache key for ``cache_module.load_model``. LoRA-tagged so fused and base weights don't
+ alias in ``TT_DIT_CACHE_DIR``; quant-tagged because cached tensorbins carry their dtype, so
+ a bf8 preset run and the bf16 baseline must live in separate dirs."""
+ base = os.path.basename(self._checkpoint_path).removesuffix(".safetensors")
+ if lora_specs:
+ tag = "+".join(f"{os.path.basename(s.path).removesuffix('.safetensors')}@{s.strength}" for s in lora_specs)
+ base = f"{base}.lora-{tag}"
+ if quant_tag:
+ base = f"{base}.q-{quant_tag}"
+ return base
+
+ def build(
+ self,
+ *,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ in_channels: int,
+ out_channels: int,
+ num_layers: int,
+ cross_attention_dim: int,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool,
+ has_audio: bool,
+ image_conditioning: bool,
+ quant_config: LtxQuantProfile | None = None,
+ lora_enabled: bool = False,
+ ) -> LTXTransformerModel:
+ """Construct an ``LTXTransformerModel`` for this checkpoint (weights NOT loaded).
+
+ Loading is deferred so the caller can manage the lifecycle (deallocate / reload).
+ ``quant_config`` bakes the preset dtypes into construction, so a cache miss loads
+ weights direct-to-quant and the cache write holds the quantized tensorbins.
+ """
+ return LTXTransformerModel(
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ num_layers=num_layers,
+ cross_attention_dim=cross_attention_dim,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ apply_gated_attention=self.has_gate,
+ cross_attention_adaln=self.cross_attention_adaln,
+ image_conditioning=image_conditioning,
+ quant_config=quant_config,
+ lora_enabled=lora_enabled,
+ )
+
+ def load(
+ self,
+ model: LTXTransformerModel,
+ *,
+ parallel_config: DiTParallelConfig,
+ mesh_shape: tuple[int, ...],
+ is_fsdp: bool,
+ lora_specs: list[LoraSpec],
+ quant_tag: str | None = None,
+ ) -> None:
+ """Load (or reload) weights for a previously-built transformer."""
+ cache_module.load_model(
+ model,
+ model_name=self.cache_name(lora_specs, quant_tag),
+ subfolder=model.weight_cache_subfolder(),
+ parallel_config=parallel_config,
+ mesh_shape=mesh_shape,
+ mesh_device=model.mesh_device,
+ is_fsdp=is_fsdp,
+ get_torch_state_dict=lambda: self.state_dict(lora_specs),
+ )
diff --git a/code/models/tt_dit/models/transformers/minimax_h3/attention_minimax_h3.py b/code/models/tt_dit/models/transformers/minimax_h3/attention_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..645a83e2438ae2aabe9f20b340f216ab9b71c3d2
--- /dev/null
+++ b/code/models/tt_dit/models/transformers/minimax_h3/attention_minimax_h3.py
@@ -0,0 +1,609 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+import ttnn
+from models.common.utility_functions import is_blackhole
+
+from ....layers.linear import ColParallelLinear
+from ....layers.module import Module
+from ....layers.normalization import DistributedRMSNorm
+from ....parallel.config import DiTParallelConfig
+from ....parallel.manager import CCLManager
+from ....utils.mochi import get_rot_transformation_mat
+from ....utils.substate import pop_substate, rename_substate
+from ....utils.tensor import bf16_tensor
+from .agmm_config import agmm_block_size
+
+
+def rope_channel_permutation(head_dim: int, rotary_dim: int) -> torch.Tensor:
+ """Reorder a head's channels from MiniMax-H3's half-split RoPE layout to the interleaved one.
+
+ The fused RoPE inside `dit_fused_distributed_rmsnorm` rotates by multiplying every 32-column tile
+ by one 32x32 matrix, which pairs *adjacent* channels: `out[2i] = -in[2i+1]`, `out[2i+1] = in[2i]`.
+ MiniMax-H3's reference instead pairs `i` with `i + rotary_dim/2`. The two are the same operation
+ under this permutation of the rotary channels -- `out[2i] = in[i]`, `out[2i+1] = in[i + rot/2]` --
+ with the `head_dim - rotary_dim` pass-through channels left where they are.
+
+ Applied identically to the Q and K projection output channels, the QK-norm affine weight and the
+ cos/sin tables. Attention sees Q and K only through `q . k`, which any *shared* permutation of the
+ channel axis leaves unchanged, and V and `to_out` are untouched -- so the relayout is numerically
+ neutral, not an approximation. Same trick as `transformer_ideogram4.rope_halfsplit_to_interleaved_perm`,
+ extended to a rotary_dim narrower than head_dim.
+ """
+ half = rotary_dim // 2
+ rotary = torch.stack([torch.arange(half), torch.arange(half) + half], dim=1).flatten()
+ return torch.cat([rotary, torch.arange(rotary_dim, head_dim)])
+
+
+def prepare_rope_tables(cos: torch.Tensor, sin: torch.Tensor, head_dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+ """Turn the reference's `[.., rotary_dim]` cos/sin into what the fused RoPE consumes.
+
+ Permutes the rotary channels into the interleaved layout (see `rope_channel_permutation`) and pads
+ out to `head_dim` with cos=1 / sin=0. Those pad channels are the ones MiniMax-H3 passes through
+ unrotated: because the fused rotate only ever mixes channels *within* a 32-column tile, and the
+ pass-through channels occupy whole tiles of their own, `sin=0` there makes the rotate an exact
+ identity. That is what lets a partial-head-dim RoPE run on an op that has no notion of one.
+ """
+ rotary_dim = cos.shape[-1]
+ perm = rope_channel_permutation(rotary_dim, rotary_dim) # permute the rotary block only
+ cos, sin = cos[..., perm], sin[..., perm]
+ pad = head_dim - rotary_dim
+ if pad:
+ ones = torch.ones(*cos.shape[:-1], pad, dtype=cos.dtype)
+ zeros = torch.zeros(*sin.shape[:-1], pad, dtype=sin.dtype)
+ cos, sin = torch.cat([cos, ones], dim=-1), torch.cat([sin, zeros], dim=-1)
+ return cos, sin
+
+
+class MiniMaxH3Attention(Module):
+ """Full self-attention over one packed sequence. MiniMax-H3 has no cross-attention.
+
+ Two things differ from `WanAttention` and drive the shape of this module:
+
+ * The attention inner dim (`num_heads * head_dim` = 7168) is *larger* than the residual stream
+ (`hidden_size` = 5376), so `to_q/k/v` widen 5376 -> 7168 and `to_out` narrows 7168 -> 5376.
+ Nothing here may assume `inner_dim == hidden_size`. Every projection is bias-free.
+ * The query/key norms are RMSNorms over `head_dim` (128), not over the TP-sharded residual
+ stream. They use `DistributedRMSNorm` in `per_head_norm` mode, which reduces over each head's
+ head_dim *locally* -- no all-gather, since a head's channels all live on one device -- and in
+ the same op splits the heads and applies RoPE. So one fused op replaces the norm, the head
+ split and the whole rotary sequence. See `rope_channel_permutation` for how MiniMax-H3's
+ partial, half-split rotary is made to fit an op that implements a full-width interleaved one.
+ """
+
+ # Per-device sequence length -> measured-best ring SDPA (q_chunk_size, k_chunk_size).
+ # See `_sdpa_program_config` for how these were obtained and why the optimum moves with length.
+ # 4768 / 9216 / 13632 are 768P at 5s / 10s / 15s, packed and padded, divided by SP=8.
+ measured_sdpa_chunk_sizes = {
+ 4768: (320, 384),
+ 9216: (256, 512),
+ 13632: (256, 512),
+ }
+
+ def __init__(
+ self,
+ *,
+ hidden_size: int,
+ num_heads: int,
+ head_dim: int,
+ rotary_dim: int | None = None,
+ qk_norm_eps: float = 1e-5,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool = False,
+ is_sequence_parallel: bool = True,
+ ) -> None:
+ super().__init__()
+
+ # is_sequence_parallel=False means the sequence is *replicated* on the SP axis rather than
+ # fractured across it, so attention runs locally with plain SDPA and no ring all-gather. The
+ # token refiner uses that: its text stream is short and every SP device holds all of it.
+ self.is_sequence_parallel = is_sequence_parallel
+
+ self.hidden_size = hidden_size
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+ self.inner_dim = num_heads * head_dim
+ # Channels [rotary_dim, head_dim) pass through the rotary embedding unrotated. None means
+ # the whole head rotates, which is what the token refiner (no RoPE at all) leaves unused.
+ self.rotary_dim = head_dim if rotary_dim is None else rotary_dim
+ self.qk_norm_eps = qk_norm_eps
+
+ self.mesh_device = mesh_device
+ self.ccl_manager = ccl_manager
+ self.parallel_config = parallel_config
+
+ tp_factor = parallel_config.tensor_parallel.factor
+ assert num_heads % tp_factor == 0, f"{num_heads} heads must divide across TP={tp_factor}"
+ self.n_local_heads = num_heads // tp_factor
+ self.tp_mesh_axis = parallel_config.tensor_parallel.mesh_axis
+ self.sp_mesh_axis = parallel_config.sequence_parallel.mesh_axis
+ # Fractured sequence means attention has to gather K/V around the ring.
+ self.use_ring = is_sequence_parallel and parallel_config.sequence_parallel.factor > 1
+
+ fsdp_mesh_axis = self.sp_mesh_axis if is_fsdp else None
+
+ # Fused QKV: one matmul, output split into three. The state dict is rearranged in
+ # `_prepare_torch_state` so that column-parallel fracturing hands each device the same
+ # 14 heads of q, k and v.
+ self.to_qkv = ColParallelLinear(
+ hidden_size,
+ 3 * self.inner_dim,
+ chunks=3,
+ bias=False,
+ mesh_device=mesh_device,
+ mesh_axis=self.tp_mesh_axis,
+ fsdp_mesh_axis=fsdp_mesh_axis,
+ ccl_manager=ccl_manager,
+ )
+ self.to_out = ColParallelLinear(
+ self.inner_dim,
+ hidden_size,
+ bias=False,
+ mesh_device=mesh_device,
+ mesh_axis=self.tp_mesh_axis,
+ fsdp_mesh_axis=fsdp_mesh_axis,
+ ccl_manager=ccl_manager,
+ )
+
+ # QK-norm + head split + RoPE in one fused op. embedding_dim is the *inner* dim so the
+ # per-device weight slice covers n_local_heads * head_dim; `per_head_norm=True` at the call
+ # site makes the reduction per head and device-local.
+ qk_norm_kwargs = dict(
+ embedding_dim=self.inner_dim,
+ norm_eps=qk_norm_eps,
+ norm_elementwise_affine=True,
+ mesh_axis=self.tp_mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+ self.norm_q = DistributedRMSNorm(**qk_norm_kwargs)
+ self.norm_k = DistributedRMSNorm(**qk_norm_kwargs)
+ self.rope_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ # all_gather_minimal_matmul_async folds the TP all-gather into the matmul. Ring only: on a
+ # line topology WanAttention measured the unfused path faster, so match that condition.
+ self.use_fused_agmm = ccl_manager.topology == ttnn.Topology.Ring and tp_factor > 1
+
+ # Ring SDPA reuses the joint-attention entry point with empty joint inputs, as WanAttention does.
+ self.dummy_joint_input = bf16_tensor(torch.zeros((1, self.n_local_heads, 0, head_dim)), device=mesh_device)
+
+ full_grid = mesh_device.compute_with_storage_grid_size()
+ self.full_grid = full_grid
+ self.sdpa_worker_grid = (full_grid.x - 1, full_grid.y) # reserve last column for CCL
+ self._sdpa_program_configs: dict[tuple[int, bool], ttnn.SDPAProgramConfig] = {}
+
+ # The exp ring op walks head-SEGMENTS (a head's Q chunks split over segs_per_head rows) as
+ # serial passes, ceil(n_local_heads * segs / rows) passes per row. Segmentation is what
+ # balances 14 local heads over 10 rows: segs=1 gives 2 passes of 10-tile chunks with 6 rows
+ # idle on the second pass, while segs=2 gives 3 passes of 5-tile chunks on every core --
+ # 15 Q tile-rows per core instead of 20 on the bottleneck cores.
+ self.exp_ring_max_passes = 3 # kMaxPasses in exp_ring_joint_sdpa_program_factory.cpp
+ self.exp_ring_num_passes = math.ceil(self.n_local_heads / full_grid.y)
+ self.exp_ring_max_k_chunk = 512 # largest k worth trying; `_exp_sdpa_l1_bytes` picks down from here
+ self.use_exp_ring_sdpa = (
+ self.use_ring
+ and is_blackhole()
+ and tp_factor == 4
+ and parallel_config.sequence_parallel.factor == 32
+ and self.exp_ring_num_passes <= self.exp_ring_max_passes
+ )
+ self._exp_sdpa_program_configs: dict[int, ttnn.SDPAProgramConfig | None] = {}
+
+ self.sdpa_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ mesh_device.arch(),
+ math_fidelity=ttnn.MathFidelity.HiFi2,
+ math_approx_mode=False,
+ fp32_dest_acc_en=False,
+ )
+ self.mm_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ mesh_device.arch(),
+ math_fidelity=ttnn.MathFidelity.HiFi2,
+ math_approx_mode=True,
+ fp32_dest_acc_en=True,
+ packer_l1_acc=True,
+ )
+
+ # ------------------------------------------------------------------ weights
+
+ def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+ rename_substate(state, "to_out.0", "to_out")
+
+ def _interleave_heads(tensors: list[torch.Tensor]) -> torch.Tensor:
+ """Reorder [out, in] weights so TP column-fracturing gives each device matching heads.
+
+ Out dim is `num_heads * head_dim`. Reshaping it to [n_dev, n_local_heads, head_dim] and
+ concatenating the tensors on the heads axis puts device `d`'s q, k and v heads
+ contiguously inside shard `d`, which is also the order `chunks=3` splits them back out in.
+ Device `d` therefore owns canonical heads `[d * n_local, (d + 1) * n_local)`, so simply
+ all-gathering the attention output on TP rebuilds the canonical 7168-channel order that
+ `to_out` expects.
+ """
+ n_dev = self.parallel_config.tensor_parallel.factor
+ tensors = [t.T for t in tensors] # -> [in, out]
+ 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.inner_dim)
+ return merged.T
+
+ q_state = pop_substate(state, "to_q")
+ k_state = pop_substate(state, "to_k")
+ v_state = pop_substate(state, "to_v")
+
+ # Relayout Q and K into the interleaved rotary channel order the fused RoPE consumes. Shared
+ # between Q and K and absent from V, so Q.K is unchanged (see `rope_channel_permutation`).
+ perm = rope_channel_permutation(self.head_dim, self.rotary_dim)
+
+ def _permute_rotary(weight: torch.Tensor) -> torch.Tensor:
+ # [out, in] with out == num_heads * head_dim; permute within each head.
+ return weight.reshape(self.num_heads, self.head_dim, -1)[:, perm].reshape(weight.shape)
+
+ state["to_qkv.weight"] = _interleave_heads(
+ [_permute_rotary(q_state["weight"]), _permute_rotary(k_state["weight"]), v_state["weight"]]
+ )
+
+ # The reference's QK-norm affine is one head_dim vector shared by every head. The fused op
+ # wants it spanning the whole inner dim, so permute it the same way and repeat per head.
+ for name in ("norm_q", "norm_k"):
+ sub = pop_substate(state, name)
+ if "weight" in sub:
+ state[f"{name}.weight"] = sub["weight"][perm].repeat(self.num_heads)
+
+ # ------------------------------------------------------------------ helpers
+
+ def _sdpa_program_config(self, seq_local: int, *, ring: bool) -> ttnn.SDPAProgramConfig:
+ """Ring SDPA chunk sizes for a given per-device sequence length.
+
+ Measured points come from the sweep in
+ `tests/nightly/blackhole/sdpa/test_ring_joint_sdpa.py::test_ring_joint_attention_create_perf_table`
+ (model configs `minimax_h3_{5s,10s,15s}_768p`) on 4x8 Blackhole Galaxy, TP=4 / SP=8. Anything
+ else falls back to the generic rule below.
+
+ The optimum depends on the *local* sequence length, which is why this is keyed on it rather
+ than on the mesh shape the way `WanAttention.sdpa_chunk_size_map` is. At long sequences there
+ is enough work that (256, 512) wins -- a larger k halves the ring's K-loop iterations. At 5s
+ there is too little work to fill 110 cores that way: 14 heads x ceil(4768/256) = 266 work items
+ over 110 cores rounds up to 3 per core and wastes ~19% of the slots, whereas q=320 gives 210
+ items, 2 per core and ~4.5% waste. That is worth more than the larger k.
+
+ L1 (1.57 MB) bounds the (q, k) product, and the bound is what makes k=384 interesting: k=1024
+ never fits, (320, 512) reaches 1.65 MB and does not either, but (320, 384) does -- combining
+ the good q with a k larger than 256. Neither chunk size has to be a power of two, only a
+ multiple of TILE, and restricting the search to {256, 512} misses this point entirely.
+
+ Padding is not the thing to tune. 4768 is 32 x 149 with 149 prime, so no sensible chunk size
+ divides it, yet at q=320 the Q padding is only 0.67%. Core slot efficiency dominates, and the
+ optimum in it is sharp rather than a plateau: at 5s q=288 measured 10.43 ms against q=320's
+ 7.81 ms. Slot efficiency is a good candidate generator but not a predictor -- q=416 at 10s has
+ the best slot efficiency of any k=256 point there (97.6%) and measured the worst (28.39 ms).
+ """
+ key = (seq_local, ring)
+ if key not in self._sdpa_program_configs:
+ tile = ttnn.TILE_SIZE
+ measured = self.measured_sdpa_chunk_sizes.get(seq_local)
+ if measured is not None:
+ q_chunk, k_chunk = measured
+ else:
+ q_chunk = max(tile, min(256, (seq_local // tile) * tile))
+ k_chunk = max(tile, min(512, (seq_local // tile) * tile))
+ grid = (
+ ttnn.CoreCoord(*self.sdpa_worker_grid) if ring else ttnn.CoreCoord(self.full_grid.x, self.full_grid.y)
+ )
+ self._sdpa_program_configs[key] = ttnn.SDPAProgramConfig(
+ compute_with_storage_grid_size=grid,
+ q_chunk_size=q_chunk,
+ k_chunk_size=k_chunk,
+ exp_approx_mode=False, # NOTE: False is more correct
+ )
+ return self._sdpa_program_configs[key]
+
+ # One accumulator entry and one Q chunk per pass, in tiles of `_EXP_L1_TILE_BYTES`. Mirrors the
+ # CB table in exp_ring_joint_sdpa_program_factory.cpp; reproduces its measured 1,302,528 B at
+ # (224, 512) exactly. Nothing in the op validates this, so an oversized shape would only surface
+ # as a CB allocation failure at program build.
+ _EXP_L1_TILE_BYTES = 2048 # bf16 and Float16_b tiles are both 2 KiB
+ # CB space measured IN THE PIPELINE, not bare L1: the op's CBs must end below the lowest live
+ # L1 buffer (global semaphores etc. occupy the top of L1), which a 15s run measured at
+ # 1,504,000 with the CB region starting at 191,360. The factory checks the live value at build;
+ # this constant only has to be a safe lower bound so the k search picks a buildable shape.
+ _EXP_USABLE_L1_BYTES = 1_312_640
+ # DEST tiles from `get_dest_reg_count`: 1024 * 16 / (32 * 32), halved because dst_full_sync_en is
+ # off, not halved again because `sdpa_compute_kernel_config` has fp32_dest_acc_en off.
+ _EXP_DST_TILES = 8
+ # `determine_largest_subblock_size` in sdpa_subblock_utils.hpp, in its search order.
+ _EXP_SUBBLOCKS = (
+ (2, 4), (4, 2), (1, 8), (8, 1), (1, 7), (7, 1), (2, 3), (3, 2), (1, 6), (6, 1),
+ (1, 5), (5, 1), (2, 2), (1, 4), (4, 1), (1, 3), (3, 1), (1, 2), (2, 1), (1, 1),
+ ) # fmt: skip
+
+ def _exp_streaming_compute_enabled(self, sq_t: int, sk_t: int) -> bool:
+ """Whether the op picks its streaming compute path for a chunk shape.
+
+ Mirrors `use_streaming_compute` in exp_ring_joint_sdpa_program_factory.cpp. The exp compute
+ kernel `static_assert`s on it, so a shape the factory judges ineligible does not fall back --
+ it fails to build the kernel. The binding term in practice is `sk_t % (dst / h) == 0`: it
+ rejects k=320 at q=320 (h=1, so sk_t must be a multiple of 8, and 10 is not).
+ """
+ dst = self._EXP_DST_TILES
+ for h, w in self._EXP_SUBBLOCKS:
+ if h * w <= dst and sq_t % h == 0 and sk_t % w == 0:
+ return h <= 2 and sk_t % (dst // h) == 0 and sq_t // h > 1
+ return False
+
+ def _exp_sdpa_l1_bytes(self, sq_t: int, sk_t: int, p: int, resident_q: bool = True) -> int:
+ """L1 the exp op's circular buffers need for a (q, k, passes) shape, in tiles of head_dim.
+
+ `p` is the candidate's pass count (head-segment scheduling makes it per-candidate:
+ ceil(n_local_heads * segs / rows)). `resident_q=False` models the op's streamed-Q fallback:
+ when the resident total does not fit, the factory sizes c_0 to a single chunk and the reader
+ re-reads each pass's Q every ring iteration. The op selects the mode itself from this same
+ arithmetic; the model only needs it to know which (q, k) shapes are buildable.
+ """
+ dh_t = self.head_dim // ttnn.TILE_SIZE
+ tiles = (
+ (p if resident_q else 1) * sq_t * dh_t # c_0 Q: one chunk per pass, or one when streamed
+ + 4 * sk_t * dh_t # c_1/c_14 K and c_2/c_15 V, double buffered
+ + 7 # c_3 mask, scalars, reciprocal scratch
+ + 2 * p * sq_t # c_6 / c_11 state FIFO running max and sum
+ + p * sq_t * dh_t # c_7 state FIFO partial output
+ + sq_t # c_17 stats out (c_10 is dead on this path and not allocated)
+ + 16 # c_16 streaming output ping-pong
+ + sq_t * sk_t # c_24 qk intermediate
+ + 2 * sq_t * dh_t # c_25/c_26 output scratch halves
+ + 4 * sq_t # c_27-c_30 max/sum scratch halves
+ + sq_t # c_31 exp max diff
+ )
+ return tiles * self._EXP_L1_TILE_BYTES
+
+ def _exp_sdpa_program_config(self, seq_local: int) -> ttnn.SDPAProgramConfig | None:
+ """Exp ring SDPA config for a per-device sequence length, or None if it cannot use the op.
+
+ The op gives Q chunk `x` to core column `x`, so a head's chunks must fill its row exactly:
+ `ceil(seq_local / q_chunk)` has to equal the SDPA column count. That pins q_chunk to the
+ window `[seq_local / cols, seq_local / (cols - 1))` for a given width, and only a TILE
+ multiple will do, so a width is usable only if one lands there. `measured_sdpa_chunk_sizes`
+ therefore does not apply on this path -- though at the H3 10s shape the window happens to
+ give q=224, the value WanAttention measured anyway.
+
+ Widest usable grid wins, so try `full_grid.x - 1` columns first and step down. A narrower
+ grid frees no L1 whatsoever -- every CB is sized from (q_chunk, k_chunk, passes), never from
+ the column count -- so it is never worth taking while a wider one fits. It is only worth
+ taking when no wider width admits a q_chunk at all, and there the comparison is not against
+ a full-grid exp config but against no exp config at all: at 5s, 11 columns admit nothing
+ (96 -> 13 chunks, 128 -> 10) while 10 columns take q=128 with L1 to spare.
+
+ k_chunk is the one genuinely free variable, and the only one that buys L1 headroom, so take
+ the largest that both fits L1 and keeps the op on its streaming compute path -- the kernel
+ static_asserts on the latter, so an ineligible k fails the build rather than falling back.
+ A shape fits if either Q mode does: resident Q (all passes' chunks stay in L1, read once) or
+ the op's streamed-Q fallback (one chunk resident, re-read per pass per ring iteration).
+ That gives 512 at q=224 resident, and 384 at q=320 streamed -- where the k=256 that resident
+ Q would force measured far slower (small k doubles the per-chunk flash overhead; see
+ exp_more_heads_per_row.md §9).
+ """
+ if not self.use_exp_ring_sdpa:
+ return None
+ if seq_local not in self._exp_sdpa_program_configs:
+ self._exp_sdpa_program_configs[seq_local] = self._build_exp_sdpa_program_config(seq_local)
+ return self._exp_sdpa_program_configs[seq_local]
+
+ def _build_exp_sdpa_program_config(self, seq_local: int) -> ttnn.SDPAProgramConfig | None:
+ """Search (cols, segs_per_head, q_chunk, k_chunk) and take the lightest bottleneck load.
+
+ The per-core matmul work per ring iteration is passes * q_chunk Q rows against the full
+ K/V stream, so `passes * q_chunk` is the primary score: at 14 heads on 10 rows, segs=1
+ gives 2 passes x 320 = 640 rows while segs=2 gives 3 passes x 160 = 480. Larger k_chunk
+ is the tie-break (fewer per-chunk overheads), then wider grids.
+ """
+ tile = ttnn.TILE_SIZE
+ rows = self.full_grid.y
+ best = None
+ for cols in range(self.full_grid.x - 1, 1, -1):
+ for segs in (1, 2, 3):
+ chunks = cols * segs
+ q_chunk = math.ceil(math.ceil(seq_local / chunks) / tile) * tile
+ if math.ceil(seq_local / q_chunk) != chunks:
+ continue # this (cols, segs) admits no tile-multiple q_chunk
+ passes = math.ceil(self.n_local_heads * segs / rows)
+ if passes > self.exp_ring_max_passes:
+ continue
+ for k_chunk in range(self.exp_ring_max_k_chunk, 0, -tile):
+ sq_t, sk_t = q_chunk // tile, k_chunk // tile
+ fits = (
+ self._exp_sdpa_l1_bytes(sq_t, sk_t, passes) <= self._EXP_USABLE_L1_BYTES
+ or self._exp_sdpa_l1_bytes(sq_t, sk_t, passes, resident_q=False) <= self._EXP_USABLE_L1_BYTES
+ )
+ if fits and self._exp_streaming_compute_enabled(sq_t, sk_t):
+ score = (passes * q_chunk, -k_chunk, -cols)
+ if best is None or score < best[0]:
+ best = (score, cols, q_chunk, k_chunk)
+ break
+ if best is None:
+ return None
+ _, cols, q_chunk, k_chunk = best
+ return ttnn.SDPAProgramConfig(
+ compute_with_storage_grid_size=ttnn.CoreCoord(cols + 1, self.full_grid.y),
+ q_chunk_size=q_chunk,
+ k_chunk_size=k_chunk,
+ exp_approx_mode=False, # NOTE: False is more correct
+ )
+
+ # ------------------------------------------------------------------ forward
+
+ def forward(
+ self,
+ spatial_1BND: ttnn.Tensor,
+ N: int | None = None,
+ rope_cos: ttnn.Tensor | None = None,
+ rope_sin: ttnn.Tensor | None = None,
+ addcmul_residual: ttnn.Tensor | None = None,
+ addcmul_gate: ttnn.Tensor | None = None,
+ ) -> ttnn.Tensor:
+ """
+ spatial_1BND: fractured hidden_size on TP; fractured N on SP when `is_sequence_parallel`,
+ otherwise replicated on SP.
+ rope_cos/rope_sin: [1, 1, N_local, rotary_dim], fractured N on SP, replicated on TP. Both
+ None skips the rotary embedding entirely, as the token refiner requires.
+ N: logical (unfractured) sequence length. Only needed for ring attention.
+ addcmul_residual/addcmul_gate: when both are given, the gated residual
+ `addcmul_residual + to_out(...) * addcmul_gate` is folded into the to_out matmul's
+ epilogue instead of running as separate ops. Both must be TP-fractured like the output.
+
+ Returns the attention output with the same distribution as the input.
+ """
+ assert (addcmul_residual is None) == (addcmul_gate is None), "addcmul residual/gate come as a pair"
+ assert (rope_cos is None) == (rope_sin is None), "rope_cos and rope_sin must be given together"
+ # The fused RoPE consumes head_dim-wide tables, not the reference's rotary_dim-wide ones: the
+ # pass-through channels must be present as cos=1 / sin=0. Passing the raw reference tables
+ # here is silently wrong rather than a shape error, so check it. See `prepare_rope_tables`.
+ if rope_cos is not None and rope_cos.shape[-1] != self.head_dim:
+ msg = (
+ f"rope tables must be head_dim ({self.head_dim}) wide, got {rope_cos.shape[-1]}; "
+ "build them with prepare_rope_tables()"
+ )
+ raise ValueError(msg)
+
+ tp_factor = self.parallel_config.tensor_parallel.factor
+ assert not (self.use_ring and N is None), "ring attention needs the logical sequence length N"
+
+ # Passing parallel_config puts ColParallelLinear on all_gather_minimal_matmul_async: the TP
+ # all-gather of the K-fractured input folds into the matmul that consumes it, instead of
+ # running as a separate op. Only on ring topologies -- on a line the unfused path is faster,
+ # which is the same condition WanAttention uses.
+ matmul_parallel_config = self.parallel_config if self.use_fused_agmm else None
+ if not self.use_fused_agmm and tp_factor > 1:
+ spatial_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_1BND, dim=3, mesh_axis=self.tp_mesh_axis
+ )
+
+ q_1BNF, k_1BNF, v_1BNF = self.to_qkv(
+ spatial_1BND,
+ compute_kernel_config=self.mm_compute_kernel_config,
+ parallel_config=matmul_parallel_config,
+ default_block_size=agmm_block_size(self.hidden_size, 3 * self.inner_dim // tp_factor),
+ )
+
+ def create_heads(inp: ttnn.Tensor) -> ttnn.Tensor:
+ out, _, _ = ttnn.experimental.nlp_create_qkv_heads(
+ inp,
+ num_heads=self.n_local_heads,
+ num_kv_heads=0,
+ transpose_k_heads=False,
+ )
+ return out
+
+ # One fused op per stream: per-head RMSNorm over head_dim, head split, and RoPE. It emits
+ # head-split [B, n_local_heads, N, head_dim] directly, so Q and K need no create_heads.
+ norm_kwargs = dict(
+ num_heads_per_device=self.n_local_heads,
+ per_head_norm=True,
+ rope_cos=rope_cos,
+ rope_sin=rope_sin,
+ trans_mat=self.rope_trans_mat if rope_cos is not None else None,
+ )
+ q_BHNE = self.norm_q(q_1BNF, **norm_kwargs)
+ k_BHNE = self.norm_k(k_1BNF, **norm_kwargs)
+ v_BHNE = create_heads(v_1BNF)
+
+ # Sequence is fractured across SP, so attention must gather K/V around the ring.
+ # The packed sequence is one attention document and logical_n masks the pad tail, so no mask.
+ exp_program_config = self._exp_sdpa_program_config(q_BHNE.shape[2])
+ if exp_program_config is not None:
+ spatial_BHNE, _prompt, _lse = ttnn.transformer.exp_ring_joint_scaled_dot_product_attention(
+ q_BHNE,
+ k_BHNE,
+ v_BHNE,
+ self.dummy_joint_input,
+ self.dummy_joint_input,
+ self.dummy_joint_input,
+ persistent_output_buffer_k=self.ccl_manager.get_ag_ping_pong_buffer(
+ k_BHNE.shape, 2, self.sp_mesh_axis, dtype=k_BHNE.dtype
+ ),
+ persistent_output_buffer_v=self.ccl_manager.get_ag_ping_pong_buffer(
+ v_BHNE.shape, 2, self.sp_mesh_axis, dtype=v_BHNE.dtype
+ ),
+ joint_strategy="rear",
+ logical_n=N,
+ program_config=exp_program_config,
+ compute_kernel_config=self.sdpa_compute_kernel_config,
+ dim=2,
+ multi_device_global_semaphore=self.ccl_manager.get_exp_ring_ping_pong_semaphore(self.sp_mesh_axis),
+ num_links=self.ccl_manager.num_links,
+ cluster_axis=self.sp_mesh_axis,
+ mesh_device=self.mesh_device,
+ topology=self.ccl_manager.topology,
+ subdevice_id=self.ccl_manager.ccl_sub_device_id,
+ num_workers_per_link=5,
+ num_buffers_per_channel=32,
+ )
+ elif self.use_ring:
+ spatial_BHNE, _prompt, _lse = ttnn.transformer.ring_joint_scaled_dot_product_attention(
+ q_BHNE,
+ k_BHNE,
+ v_BHNE,
+ self.dummy_joint_input,
+ self.dummy_joint_input,
+ self.dummy_joint_input,
+ persistent_output_buffer_k=self.ccl_manager.get_ag_ping_pong_buffer(
+ k_BHNE.shape, 2, self.sp_mesh_axis, dtype=k_BHNE.dtype
+ ),
+ persistent_output_buffer_v=self.ccl_manager.get_ag_ping_pong_buffer(
+ v_BHNE.shape, 2, self.sp_mesh_axis, dtype=v_BHNE.dtype
+ ),
+ joint_strategy="rear",
+ logical_n=N,
+ program_config=self._sdpa_program_config(q_BHNE.shape[2], ring=True),
+ compute_kernel_config=self.sdpa_compute_kernel_config,
+ dim=2,
+ multi_device_global_semaphore=self.ccl_manager.get_ag_ping_pong_semaphore(self.sp_mesh_axis),
+ num_links=self.ccl_manager.num_links,
+ cluster_axis=self.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:
+ spatial_BHNE = ttnn.transformer.scaled_dot_product_attention(
+ q_BHNE,
+ k_BHNE,
+ v_BHNE,
+ is_causal=False,
+ program_config=self._sdpa_program_config(q_BHNE.shape[2], ring=False),
+ compute_kernel_config=self.sdpa_compute_kernel_config,
+ )
+
+ spatial_1BND = ttnn.transformer.concatenate_heads(spatial_BHNE)
+ spatial_1BND = ttnn.unsqueeze(spatial_1BND, 0)
+
+ # Each device holds canonical heads [d * n_local, (d+1) * n_local), so gathering on TP
+ # rebuilds the full inner_dim in canonical order for to_out -- fused into the matmul when
+ # use_fused_agmm.
+ if not self.use_fused_agmm and tp_factor > 1:
+ spatial_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_1BND, dim=3, mesh_axis=self.tp_mesh_axis
+ )
+
+ # The gated residual rides along in the matmul epilogue on the fused path; on the unfused
+ # path the op has no addcmul, so apply it afterwards.
+ fuse_gate = addcmul_residual is not None and self.use_fused_agmm
+ out = self.to_out(
+ spatial_1BND,
+ compute_kernel_config=self.mm_compute_kernel_config,
+ parallel_config=matmul_parallel_config,
+ default_block_size=agmm_block_size(self.inner_dim, self.hidden_size // tp_factor),
+ addcmul_a=addcmul_residual if fuse_gate else None,
+ addcmul_b=addcmul_gate if fuse_gate else None,
+ )
+ if addcmul_residual is not None and not fuse_gate:
+ out = ttnn.addcmul(addcmul_residual, out, addcmul_gate)
+ return out
diff --git a/code/models/tt_dit/models/transformers/minimax_h3/mmrs_config.py b/code/models/tt_dit/models/transformers/minimax_h3/mmrs_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e18f1c4bfd06ef577982c43bdb84b8db23efc5f
--- /dev/null
+++ b/code/models/tt_dit/models/transformers/minimax_h3/mmrs_config.py
@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Fused matmul + reduce-scatter (+ addcmul) blockings for the MiniMax-H3 feed-forward ff2.
+
+`ff2` is row-parallel, so its reduce-scatter and the gated residual after it collapse into one
+`minimal_matmul_strided_reduce_scatter_async`. That fusion is only a win with a *swept* blocking:
+`get_fused_mmrs_config` falls back to `default_fused_mmrs_config`, which puts the matmul on an 8x7 =
+56-core grid with `M_block=2` and subblock 1x1 -- half the machine at the least efficient subblock,
+where the unfused path runs on 110 cores at subblock (2, 2). Measured, that fallback makes the fusion a
+**45% regression** on the stage (1.75 -> 2.55 ms).
+
+Unlike `agmm_config` these must be keyed on the full `(M, K, N)`, because that is what
+`get_fused_mmrs_config` looks up -- there is no `default_block_size` hook to key more loosely. M is the
+per-device packed sequence length, so it varies with the video duration; `has_mmrs_config` gates the
+fused path off entirely for shapes it cannot serve, and `register_mmrs_config` registers the blocking
+for whichever M the fused path is about to run at.
+
+`compute_with_storage_grid_size` is the *matmul* grid, and the reduce-scatter workers occupy the rows
+between it and the full device grid:
+
+ rs_zone_capacity = (device_grid.y - mm_grid.y) * device_grid.x
+ num_workers_per_link = rs_zone_capacity // (2 * num_links) - 1
+
+so the matmul grid and the reduce-scatter's bandwidth trade against each other directly, and that
+trade is the dominant axis. Swept on 4x8 Blackhole Galaxy at 2 links with
+`models/tt_dit/utils/sweep_mm_block_sizes.py` (`mmrs` use case), first across the three candidate
+grids at M=4768:
+
+ 12x7 84 mm cores, 8 RS workers/link M=6 K=4 N=16 sb(2,2) 1.373 ms
+ 12x8 96 mm cores, 5 RS workers/link M=4 K=8 N=14 sb(2,2) 1.313 ms <- best, -25.0%
+ 12x9 108 mm cores, 2 RS workers/link M=6 K=2 N=8 sb(2,2) 1.487 ms
+
+The optimum is interior: 12x9 starves the reduce-scatter and 12x7 starves the matmul. Ms without a
+swept entry are NOT registered here: they fall through to the v2.3 rule engine inside
+`get_fused_mmrs_config` (`utils/mmrs_rules.py`), which picks the 12x8 blocking per M. Measured
+across the 41 swept H3 Ms, the rules average +1.9% over each M's swept optimum (worst +5.3%),
+versus +6.6% (worst +26.7%, at small M) for the earlier policy of reusing the M=3424 blocking
+at every duration.
+"""
+
+from __future__ import annotations
+
+import ttnn
+
+from ....utils.matmul import FusedMMRSConfig, register_fused_mmrs_configs
+
+# ff2 per-device K = 14336 / tp 4 = 3584; N = 5376 is the full hidden size (the reduce-scatter
+# fractures it back to 1344 per device on the way out).
+_K = 3584
+_N = 5376
+
+# Per-M entries for the Ms that have been swept at their own shape; any other M is left
+# unregistered so `get_fused_mmrs_config` resolves it with the v2.3 rules instead. Earlier
+# revisions of this file keyed everything off M=4768 ("5s@768P"), then reused the nearest
+# swept blocking for every duration; both traded up to 27% on Ms far from the swept one.
+#
+# M=3424 (the current run's per-device ff2 M) swept 2026-08-25 under the windowed L1 handoff,
+# across all three candidate grids -- 12x8 still wins: 755.6 us vs 804.2 at 12x7 and 856.8 at
+# 12x9. Mt_per_core = 107/12 -> 9, so M_block=8 leaves 2 blocks and the window rotates.
+_SWEPT_BLOCKINGS = {
+ 3424: FusedMMRSConfig(ttnn.CoreCoord(12, 8), 8, 2, 8, 2, 2, None, 1),
+}
+
+_DEVICE_GRID = ttnn.CoreCoord(12, 10)
+
+_TILE = 32
+# The swept blocking's K divisibility holds for the fixed K: K_block 2 divides K/32 = 112. The
+# rule engine guarantees the same for its own picks (K_block is chosen among divisors of K/32).
+
+
+def has_mmrs_config(m: int, k: int, n: int) -> bool:
+ """Whether the fused MM+RS+addcmul path should be taken for this shape. Pure query, no side effects.
+
+ Gate the fused path on this, and call `register_mmrs_config` before running the fused op. A shape
+ this rejects would resolve neither a table entry nor a rule pick and land on
+ `default_fused_mmrs_config`, whose 56-core matmul grid at subblock 1x1 makes the fused op far
+ *slower* than not fusing -- measured as a 45% regression on this stage.
+
+ Rather than a whitelist of the shipped durations, any tile-aligned M is accepted, so an
+ arbitrary duration gets the fused path: a swept M resolves to its table entry, and every other
+ M falls through to the v2.3 rules inside `get_fused_mmrs_config` (blind-validated within 5% of
+ the swept optimum on 14/14 fresh H3 Ms). `M_block` does *not* have to divide the M tile count --
+ a partial trailing block along M is fine, unlike along K, where the ring delivers the gathered
+ input in fixed chunks; requiring divisibility here would silently disable the fused path for
+ most durations.
+ """
+ return k == _K and n == _N and m % _TILE == 0
+
+
+def register_mmrs_config(m: int, k: int, n: int) -> None:
+ """Register the swept blocking for this shape into the global fused-MMRS table, if one exists.
+
+ Call at the site that is about to take the fused path (idempotent, cheap), for a shape
+ `has_mmrs_config` accepts; registration is explicit so that merely *querying* a shape never
+ mutates the global table. An M without a swept entry is deliberately NOT registered:
+ `get_fused_mmrs_config` then resolves it with the v2.3 rule engine, which beats reusing a
+ neighboring swept blocking (see the module docstring for the measured comparison).
+ """
+ if not has_mmrs_config(m, k, n):
+ msg = f"No fused MMRS blocking for (M, K, N) = ({m}, {k}, {n}); gate on has_mmrs_config first"
+ raise ValueError(msg)
+ if m in _SWEPT_BLOCKINGS:
+ register_fused_mmrs_configs({_DEVICE_GRID: {(m, _K, _N): _SWEPT_BLOCKINGS[m]}})
diff --git a/code/models/tt_dit/models/transformers/minimax_h3/transformer_block_minimax_h3.py b/code/models/tt_dit/models/transformers/minimax_h3/transformer_block_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..1dbb1fb0d4fcb6bcec5456bc0d2c136dac263bcb
--- /dev/null
+++ b/code/models/tt_dit/models/transformers/minimax_h3/transformer_block_minimax_h3.py
@@ -0,0 +1,339 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+
+import ttnn
+
+from ....layers.feedforward import ParallelFeedForward
+from ....layers.linear import ColParallelLinear
+from ....layers.module import Module
+from ....layers.normalization import DistributedRMSNorm
+from ....parallel.config import DiTParallelConfig
+from ....parallel.manager import CCLManager
+from ....utils.substate import rename_substate
+from .agmm_config import agmm_block_size
+from .attention_minimax_h3 import MiniMaxH3Attention
+from .mmrs_config import has_mmrs_config, register_mmrs_config
+
+# Number of modalities the AdaLN table is indexed by: video (0), text (1), audio (2).
+# Mirrors `MINIMAX_H3_MODALITY_NUM` in the reference; padding rows (-1) are clamped to 0.
+MODALITY_NUM = 3
+
+# Modulation parameters per block, in the reference's chunk order.
+NUM_MODULATION_PARAMS = 6
+(
+ _SHIFT_MSA,
+ _SCALE_MSA,
+ _GATE_MSA,
+ _SHIFT_MLP,
+ _SCALE_MLP,
+ _GATE_MLP,
+) = range(NUM_MODULATION_PARAMS)
+
+
+class MiniMaxH3TransformerBlock(Module):
+ """One MiniMax-H3 block: pre-norm self-attention and SwiGLU feed-forward, each modulated by
+ AdaLN parameters selected *per row* of the packed sequence.
+
+ The per-row modulation is what sets this apart from Wan. A single projection turns the shared
+ timestep embedding into a `(num_timesteps * MODALITY_NUM, hidden_size)` table for each of the six
+ modulation parameters, and every row of the packed sequence indexes that table by
+ `timestep_indices * MODALITY_NUM + token_tags`. One forward therefore serves rows at different
+ noise levels and of different modalities. Wan, by contrast, broadcasts one modulation over the
+ whole sequence, so it needs no gather at all.
+
+ The projection is shared by `norm1` and `norm2` and by all three modalities, so -- unlike
+ `AdaLayerNormZero` -- it cannot be folded into either norm and lives as a block-level module,
+ named `adaln_proj` after the checkpoint.
+ """
+
+ def __init__(
+ self,
+ *,
+ hidden_size: int,
+ num_heads: int,
+ head_dim: int,
+ rotary_dim: int | None = None,
+ ffn_dim: int,
+ time_embed_dim: int,
+ norm_eps: float = 1e-5,
+ qk_norm_eps: float = 1e-5,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool = False,
+ precomputed_adaln: bool = False,
+ ) -> None:
+ super().__init__()
+
+ self.hidden_size = hidden_size
+ self.ffn_dim = ffn_dim
+ self.time_embed_dim = time_embed_dim
+
+ self.mesh_device = mesh_device
+ self.ccl_manager = ccl_manager
+ self.parallel_config = parallel_config
+
+ self.tp_mesh_axis = parallel_config.tensor_parallel.mesh_axis
+ self.tp_factor = parallel_config.tensor_parallel.factor
+ assert hidden_size % self.tp_factor == 0
+ self.hidden_local = hidden_size // self.tp_factor
+
+ fsdp_mesh_axis = parallel_config.sequence_parallel.mesh_axis if is_fsdp else None
+
+ self.norm1 = DistributedRMSNorm(
+ embedding_dim=hidden_size,
+ norm_eps=norm_eps,
+ norm_elementwise_affine=True,
+ mesh_axis=self.tp_mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+ self.attn = MiniMaxH3Attention(
+ hidden_size=hidden_size,
+ num_heads=num_heads,
+ head_dim=head_dim,
+ rotary_dim=rotary_dim,
+ qk_norm_eps=qk_norm_eps,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ )
+ self.norm2 = DistributedRMSNorm(
+ embedding_dim=hidden_size,
+ norm_eps=norm_eps,
+ norm_elementwise_affine=True,
+ mesh_axis=self.tp_mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+ # The reference's SwiGLU is `up, gate = proj(x).chunk(2); up * silu(gate)`, which is the
+ # `gate_is_first=False` layout `ColParallelLinear` already assumes for activation_fn="swiglu".
+ self.ff = ParallelFeedForward(
+ hidden_size,
+ inner_dim=ffn_dim,
+ activation_fn="swiglu",
+ bias=False,
+ mesh_device=mesh_device,
+ mesh_axis=self.tp_mesh_axis,
+ fsdp_mesh_axis=fsdp_mesh_axis,
+ ccl_manager=ccl_manager,
+ )
+ # With precomputed modulation the projection never exists on device: the caller passes the
+ # six tables into `forward` instead, and this block's `adaln_proj.*` checkpoint keys are
+ # dropped in `_prepare_torch_state`. See `adaln_cache_minimax_h3`.
+ self.precomputed_adaln = precomputed_adaln
+ self.adaln_proj = (
+ None
+ if precomputed_adaln
+ else ColParallelLinear(
+ time_embed_dim,
+ NUM_MODULATION_PARAMS * hidden_size * MODALITY_NUM,
+ bias=True,
+ mesh_device=mesh_device,
+ mesh_axis=self.tp_mesh_axis,
+ fsdp_mesh_axis=fsdp_mesh_axis,
+ ccl_manager=ccl_manager,
+ )
+ )
+
+ self.mm_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ mesh_device.arch(),
+ math_fidelity=ttnn.MathFidelity.HiFi2,
+ math_approx_mode=True,
+ fp32_dest_acc_en=True,
+ packer_l1_acc=True,
+ )
+ self.use_fused_agmm = ccl_manager.topology == ttnn.Topology.Ring and self.tp_factor > 1
+ # ff1 packs gate and up together for the fused SwiGLU, so its per-device N is 2 * ffn_dim / tp.
+ self.ff1_block_size = agmm_block_size(hidden_size, 2 * ffn_dim // self.tp_factor)
+
+ # ------------------------------------------------------------------ weights
+
+ def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+ rename_substate(state, "ff.net.0.proj", "ff.ff1")
+ rename_substate(state, "ff.net.2", "ff.ff2")
+
+ # The reference lays the AdaLN projection's output dim out as
+ # [modality][param][hidden], and `hidden` is what TP fractures. Reordering the output dim to
+ # [device][modality][param][hidden_local] makes ColParallelLinear's contiguous column split
+ # hand each device every (modality, param) pair restricted to its own hidden slice.
+ def _reorder_for_tp(t: torch.Tensor) -> torch.Tensor:
+ # t is [out, ...] in torch convention; out == 6 * hidden_size * MODALITY_NUM.
+ trailing = t.shape[1:]
+ t = t.reshape(MODALITY_NUM, NUM_MODULATION_PARAMS, self.tp_factor, self.hidden_local, *trailing)
+ t = t.permute(2, 0, 1, 3, *range(4, t.ndim))
+ return t.reshape(-1, *trailing)
+
+ weight = state.pop("adaln_proj.linear.weight", None)
+ bias = state.pop("adaln_proj.linear.bias", None)
+ if self.precomputed_adaln:
+ # Dropped: these are the 26 GB the precomputed table exists to keep off the
+ # device. The pops above already strip them from the state the loader checks.
+ return
+ if weight is not None:
+ state["adaln_proj.weight"] = _reorder_for_tp(weight)
+ if bias is not None:
+ state["adaln_proj.bias"] = _reorder_for_tp(bias.unsqueeze(-1)).squeeze(-1)
+
+ # ------------------------------------------------------------------ modulation
+
+ def _modulation_tables(self, temb: ttnn.Tensor) -> list[ttnn.Tensor]:
+ """Project `temb` into the six per-(timestep, modality) modulation tables.
+
+ `temb` is [1, 1, num_timesteps, time_embed_dim]. Returns six [num_timesteps * MODALITY_NUM,
+ hidden_local] tables, row `t * MODALITY_NUM + modality`, matching the row order that
+ `adaln_indices` addresses.
+
+ The SiLU runs at `temb`'s own (float32) precision and only its result is cast
+ down to the bfloat16 projection, as the reference is explicit about: every block reads the
+ same `temb`, so rounding applied before the activation biases every block's modulation
+ identically at every sampling step and accumulates over the denoising trajectory.
+ """
+ num_timesteps = temb.shape[2]
+ # silu at temb's precision, then cast to the projection's dtype -- the reference's
+ # `self.linear(silu(temb).to(self.linear.weight.dtype))`. Casting here also keeps the
+ # resulting tables bfloat16, which `ttnn.embedding` requires of its weights.
+ activated = ttnn.silu(temb)
+ if activated.dtype != ttnn.bfloat16:
+ activated = ttnn.typecast(activated, ttnn.bfloat16)
+
+ # [1, 1, num_timesteps, MODALITY_NUM * 6 * hidden_local] per device
+ projected = self.adaln_proj(activated, compute_kernel_config=self.mm_compute_kernel_config)
+
+ # Fold the modality axis into rows: [num_timesteps, mod, param, h] -> [t * mod, param * h].
+ # Done in row-major because the row count changes and num_timesteps is tiny (a handful of
+ # rows), so the layout round-trip is negligible.
+ rows = num_timesteps * MODALITY_NUM
+ projected = ttnn.to_layout(projected, ttnn.ROW_MAJOR_LAYOUT)
+ projected = ttnn.reshape(projected, (1, 1, rows, NUM_MODULATION_PARAMS * self.hidden_local))
+ projected = ttnn.to_layout(projected, ttnn.TILE_LAYOUT)
+
+ tables = []
+ for p in range(NUM_MODULATION_PARAMS):
+ table = ttnn.slice(projected, [0, 0, 0, p * self.hidden_local], [1, 1, rows, (p + 1) * self.hidden_local])
+ # AdaLN applies the scales as (1 + scale). Adding the 1 to the table costs one elementwise
+ # op on `num_timesteps * MODALITY_NUM` rows -- six, typically -- where doing it after the
+ # gather would cost one over the whole packed sequence, per scale, per block.
+ if p in (_SCALE_MSA, _SCALE_MLP):
+ table = ttnn.add(table, 1.0)
+ # ttnn.embedding wants a 2D [num_embeddings, embedding_dim] weight.
+ table = ttnn.to_layout(table, ttnn.ROW_MAJOR_LAYOUT)
+ table = ttnn.reshape(table, (rows, self.hidden_local))
+ tables.append(ttnn.to_layout(table, ttnn.TILE_LAYOUT))
+ return tables
+
+ def _gather_rows(self, table: ttnn.Tensor, adaln_indices: ttnn.Tensor) -> ttnn.Tensor:
+ """Select one table row per row of the local packed sequence -> [1, 1, S_local, hidden_local]."""
+ out = ttnn.embedding(adaln_indices, table, layout=ttnn.TILE_LAYOUT)
+ return ttnn.unsqueeze(out, 0)
+
+ # ------------------------------------------------------------------ forward
+
+ def forward(
+ self,
+ spatial_1BND: ttnn.Tensor,
+ N: int,
+ temb: ttnn.Tensor | None,
+ adaln_indices: ttnn.Tensor,
+ rope_cos: ttnn.Tensor,
+ rope_sin: ttnn.Tensor,
+ modulation_tables: list[ttnn.Tensor] | None = None,
+ ) -> ttnn.Tensor:
+ """
+ spatial_1BND: fractured N on SP, fractured hidden_size on TP
+ temb: [1, 1, num_timesteps, time_embed_dim], replicated, float32
+ adaln_indices: [1, 1, 1, N_local] integer row indices, fractured N on SP
+ rope_cos/rope_sin: [1, 1, N_local, rotary_dim], fractured N on SP, replicated on TP
+ N: logical (unfractured) packed sequence length
+
+ Returns the block output, fractured N on SP and hidden_size on TP.
+ """
+ # Precomputed: the six tables come from the host-built cache, addressed by the same absolute
+ # `adaln_indices`, so nothing downstream changes. Otherwise project `temb` on device.
+ if modulation_tables is not None:
+ if len(modulation_tables) != NUM_MODULATION_PARAMS:
+ raise ValueError(f"expected {NUM_MODULATION_PARAMS} modulation tables, got {len(modulation_tables)}")
+ tables = modulation_tables
+ else:
+ if self.precomputed_adaln:
+ raise ValueError("block was built with precomputed_adaln but forward got no modulation_tables")
+ tables = self._modulation_tables(temb)
+
+ # ttnn.embedding takes [batch, seq] indices; uint32 is the dtype it expects.
+ indices = ttnn.reshape(adaln_indices, (1, adaln_indices.shape[-1]))
+ if indices.dtype != ttnn.uint32:
+ indices = ttnn.typecast(indices, ttnn.uint32)
+
+ def modulation(param: int) -> ttnn.Tensor:
+ return self._gather_rows(tables[param], indices)
+
+ # 1. Modulated self-attention. The (1 + scale) and shift are handed to the norm as a per-token
+ # dynamic weight and bias, so the fused norm op applies the modulation itself rather than the
+ # block following it with a separate multiply and add. The gated residual is one addcmul
+ # (residual + gate * branch) rather than a multiply and an add.
+ residual = spatial_1BND
+ normed = self.norm1(
+ spatial_1BND,
+ dynamic_weight=modulation(_SCALE_MSA),
+ dynamic_bias=modulation(_SHIFT_MSA),
+ )
+ # The gated residual is fused into to_out's matmul epilogue, so `attn` returns
+ # `residual + attn_out * gate` directly rather than the block adding it afterwards.
+ spatial_1BND = self.attn(
+ normed,
+ N=N,
+ rope_cos=rope_cos,
+ rope_sin=rope_sin,
+ addcmul_residual=residual,
+ addcmul_gate=modulation(_GATE_MSA),
+ )
+
+ # 2. Modulated feed-forward.
+ residual = spatial_1BND
+ normed = self.norm2(
+ spatial_1BND,
+ dynamic_weight=modulation(_SCALE_MLP),
+ dynamic_bias=modulation(_SHIFT_MLP),
+ )
+ # ff1 gathers the TP-fractured input inside its matmul (all_gather_minimal_matmul_async) when
+ # parallel_config is passed; ff2 is row-parallel and reduce-scatters back to TP-fractured.
+ if not self.use_fused_agmm and self.tp_factor > 1:
+ normed = self.ccl_manager.all_gather_persistent_buffer(normed, dim=3, mesh_axis=self.tp_mesh_axis)
+ # ff2's reduce-scatter and the gated residual after it fuse into a single
+ # minimal_matmul_strided_reduce_scatter_async, computing residual + ff2(...) * gate in one op.
+ #
+ # Only take it for a shape with a swept blocking. The generic fallback config runs the matmul
+ # on 56 of the device's 120 cores at subblock 1x1, making the fused op a measured 45%
+ # regression on this stage (1.75 -> 2.55 ms) -- and silently, since an unknown shape does not
+ # warn. See mmrs_config for the sweep and the grid/bandwidth tradeoff behind it.
+ #
+ # Ring only: the op asserts `attributes.topology == Ring` ("MinimalMatmulStridedReduceScatter
+ # Async only supports Ring topology"), so a line-cabled mesh has to take the unfused path
+ # below. Gated here rather than left to fail, because the assert fires on the first denoise
+ # step of the first request -- long after warmup reports the model loaded.
+ ff2_shape = (normed.shape[2], self.ffn_dim // self.tp_factor, self.hidden_size)
+ if self.tp_factor > 1 and self.ccl_manager.topology == ttnn.Topology.Ring and has_mmrs_config(*ff2_shape):
+ # M is only known here (it tracks the packed sequence length), so the blocking is
+ # registered at the point of use rather than at construction. Idempotent and cheap.
+ register_mmrs_config(*ff2_shape)
+ return self.ff.forward_fused_addcmul(
+ normed,
+ residual,
+ modulation(_GATE_MLP),
+ compute_kernel_config=self.mm_compute_kernel_config,
+ parallel_config=self.parallel_config if self.use_fused_agmm else None,
+ default_block_size=self.ff1_block_size,
+ )
+ ff_out = self.ff(
+ normed,
+ compute_kernel_config=self.mm_compute_kernel_config,
+ parallel_config=self.parallel_config if self.use_fused_agmm else None,
+ default_block_size=self.ff1_block_size,
+ )
+ return ttnn.addcmul(residual, ff_out, modulation(_GATE_MLP))
diff --git a/code/models/tt_dit/models/transformers/wan2_2/attention_wan.py b/code/models/tt_dit/models/transformers/wan2_2/attention_wan.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c7f7e1cc1fc35e069f29aefff0c417b748c4b9e
--- /dev/null
+++ b/code/models/tt_dit/models/transformers/wan2_2/attention_wan.py
@@ -0,0 +1,529 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+
+import ttnn
+from models.common.utility_functions import is_blackhole
+
+from ....layers.linear import ColParallelLinear, LoRAColParallelLinear
+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_matmul_config
+from ....utils.substate import pop_substate, rename_substate
+from ....utils.tensor import bf16_tensor
+
+
+class WanAttention(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): (288, 512),
+ (True, 32, 4): (224, 512),
+ }
+ default_sdpa_chunk_size = (256, 256)
+
+ def __init__(
+ self,
+ *,
+ dim: int,
+ num_heads: int,
+ qk_norm: bool = True,
+ eps: float = 1e-5,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager | None = None,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool = False,
+ is_self: bool = True,
+ sdpa_chunk_size_overrides: dict | 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.mesh_device = mesh_device
+ self.ccl_manager = ccl_manager
+ self.parallel_config = parallel_config
+ self.is_self = is_self
+
+ 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,
+ }
+ ColCls = LoRAColParallelLinear if lora_enabled else ColParallelLinear
+
+ if is_self:
+ # Fused QKV for self-attention: single matmul split into 3 outputs
+ self.to_qkv = ColCls(
+ dim,
+ 3 * dim,
+ chunks=3,
+ **col_parallel_kwargs,
+ )
+ else:
+ # Cross-attention: Q from spatial, K/V from prompt
+ self.to_q = ColCls(dim, dim, **col_parallel_kwargs)
+ # Fused KV: single matmul split into 2 outputs
+ self.to_kv = ColCls(
+ dim,
+ 2 * dim,
+ chunks=2,
+ **col_parallel_kwargs,
+ )
+
+ self.to_out = ColCls(
+ dim,
+ 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,
+ )
+
+ 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, # NOTE: False is more correct
+ )
+
+ self.sdpa_worker_grid = (full_grid.x - 1, full_grid.y) # Reserve last column for CCL
+ chunk_lookup = {**self.sdpa_chunk_size_map, **(sdpa_chunk_size_overrides or {})}
+ ring_sdpa_chunk_size = chunk_lookup.get(
+ (
+ is_blackhole(),
+ self.parallel_config.sequence_parallel.factor,
+ self.parallel_config.tensor_parallel.factor,
+ ),
+ 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, # NOTE: False is more correct
+ )
+
+ self.use_exp_ring_sdpa = (
+ self.parallel_config.tensor_parallel.factor == 4 and self.parallel_config.sequence_parallel.factor == 32
+ )
+
+ if self.use_exp_ring_sdpa:
+ self.exp_ring_sdpa_program_config = ttnn.SDPAProgramConfig(
+ compute_with_storage_grid_size=full_grid,
+ q_chunk_size=ring_sdpa_chunk_size[0],
+ k_chunk_size=ring_sdpa_chunk_size[1],
+ exp_approx_mode=False,
+ )
+
+ self.sdpa_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ self.mesh_device.arch(),
+ math_fidelity=ttnn.MathFidelity.HiFi2,
+ math_approx_mode=False,
+ fp32_dest_acc_en=False, # NOTE: Set to True if there's a correctness issue
+ )
+
+ self.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)
+
+ 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,
+ )
+
+ def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+ rename_substate(state, "to_out.0", "to_out")
+
+ def _interleave_heads(tensors: list[torch.Tensor]):
+ """Interleave weight/bias tensors so column-parallel fracturing assigns correct heads per device.
+
+ Each tensor has out-dim = num_heads * head_dim. We interleave them on the heads
+ axis so that after TP column-sharding each device gets the right heads from every tensor.
+
+ Args:
+ tensors: list of tensors in [out, in] PyTorch convention (or [out, 1] for bias).
+ Returns:
+ Merged tensor in [out, in] PyTorch convention.
+ """
+ n_dev = self.parallel_config.tensor_parallel.factor
+ # Transpose to [in, out]
+ tensors = [t.T for t in tensors]
+ # Reshape out dim to [in, n_dev, n_local_heads, head_dim]
+ tensors = [t.reshape(t.shape[0], n_dev, self.n_local_heads, self.head_dim) for t in tensors]
+ # Concatenate on the heads dim so each device shard gets its own heads from every tensor
+ merged = torch.cat(tensors, dim=2) # [in, n_dev, len(tensors)*n_local_heads, head_dim]
+ merged = merged.reshape(merged.shape[0], len(tensors) * self.num_heads * self.head_dim)
+ # Transpose back to [out, in] PyTorch convention
+ return merged.T
+
+ if self.is_self:
+ # Merge separate to_q, to_k, to_v weights into fused to_qkv
+ q_state = pop_substate(state, "to_q")
+ k_state = pop_substate(state, "to_k")
+ v_state = pop_substate(state, "to_v")
+
+ 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:
+ # Merge separate to_k, to_v weights into fused to_kv
+ k_state = pop_substate(state, "to_k")
+ v_state = pop_substate(state, "to_v")
+
+ 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=None,
+ dtype=None,
+ ) -> ttnn.Tensor:
+ """Fused to_out projection + addcmul: output = residual + (matmul(x, W) + bias) * gate."""
+ to_out = self.to_out
+
+ # Reads to_out.weight.data directly, so runtime-mode LoRA (which keeps
+ # the delta in self._runtime_A/B) would silently no-op. Fuse mode is
+ # fine — the delta lives in weight.data.
+ if getattr(to_out, "lora_mode", None) == "runtime" and getattr(to_out, "is_lora_active", False):
+ raise RuntimeError(
+ "runtime LoRA mode is incompatible with WanAttention._to_out_fused_addcmul; "
+ "construct to_out with lora_mode='fuse'"
+ )
+
+ # Handle FSDP weight gathering (mirrors ColParallelLinear.forward)
+ 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 = x.padded_shape[-2], weight.padded_shape[-2], weight.padded_shape[-1]
+ full_grid = self.mesh_device.compute_with_storage_grid_size()
+ core_grid = ttnn.CoreCoord(full_grid.x, full_grid.y - 1)
+ matmul_config = get_matmul_config(M, K, N, 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, # scalar
+ 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 forward(
+ self,
+ spatial_1BND: ttnn.Tensor,
+ N: int,
+ prompt_1BLP: ttnn.Tensor | None = None,
+ cross_attn_mask: 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,
+ ) -> ttnn.Tensor:
+ """
+ spatial_1BND: fractured N on SP, fracturd D on TP
+ prompt_1BLP: replicated on SP, replicated D on TP (optional)
+ cross_attn_mask: (optional, cross-attn only) mask for SDPA, shape [B|1, nqh|1, Sq, Sk].
+ Must be TILE layout, bf16/bfp8/bfp4 dtype, on device in DRAM.
+ rope_cos: fractured on SP, TP
+ rope_sin: fractured on SP, TP
+ trans_mat: replicated
+ addcmul_residual: (optional) residual tensor for fused matmul+addcmul (self-attn only)
+ addcmul_gate: (optional) gate tensor for fused matmul+addcmul (self-attn only)
+
+ If prompt_1BLP is not provided, run self-attention.
+ Otherwise, run cross-attention on prompt.
+
+ When addcmul_residual and addcmul_gate are both provided (self-attention only),
+ the to_out projection and residual addcmul are fused into a single op:
+ output = addcmul_residual + to_out(attn_output) * addcmul_gate
+
+ Outputs:
+ spatial_1BND: fractured N on SP, fractured D on TP
+ """
+
+ if rope_cos is not None:
+ # If ROPE is given, this is self-attention
+ assert rope_sin is not None
+ assert trans_mat is not None
+ assert prompt_1BLP is None
+
+ use_nonfused_agmm = (self.ccl_manager.topology == ttnn.Topology.Linear) and (
+ self.parallel_config.tensor_parallel.factor > 1
+ )
+ if use_nonfused_agmm:
+ spatial_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+
+ if self.is_self:
+ # Fused QKV matmul with split output for self-attention
+ q_1BNF, k_1BNF, v_1BNF = self.to_qkv(
+ spatial_1BND,
+ compute_kernel_config=self.mm_compute_kernel_config,
+ parallel_config=None if use_nonfused_agmm else self.parallel_config,
+ )
+ else:
+ # Cross-attention: Q from spatial, fused KV from prompt
+ assert prompt_1BLP is not None
+ kv_input = prompt_1BLP
+ q_1BNF = self.to_q(
+ spatial_1BND,
+ compute_kernel_config=self.mm_compute_kernel_config,
+ parallel_config=None if use_nonfused_agmm else self.parallel_config,
+ )
+ k_1BNF, v_1BNF = self.to_kv(kv_input, compute_kernel_config=self.mm_compute_kernel_config)
+
+ # Set norm output dtype to the input dtype required for ring self-attn.
+ sdpa_input_dtype = getattr(self, "_sdpa_input_dtype", None)
+ use_ring_sdpa = self.parallel_config.sequence_parallel.factor > 1
+ norm_output_dtype = sdpa_input_dtype if (use_ring_sdpa and prompt_1BLP is None) else None
+
+ # Norm spatial before splitting heads
+ q_BHNE = self.norm_q(
+ q_1BNF,
+ num_heads_per_device=self.n_local_heads,
+ rope_cos=rope_cos,
+ rope_sin=rope_sin,
+ trans_mat=trans_mat,
+ dtype=norm_output_dtype,
+ )
+ k_BHNE = self.norm_k(
+ k_1BNF,
+ num_heads_per_device=self.n_local_heads,
+ rope_cos=rope_cos,
+ rope_sin=rope_sin,
+ trans_mat=trans_mat,
+ dtype=norm_output_dtype,
+ )
+
+ 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_BHNE = create_heads(v_1BNF)
+
+ if prompt_1BLP is None:
+ # Self attention
+ if self.parallel_config.sequence_parallel.factor > 1:
+ # Q and K already cast by norm kernel; cast V and dummy joint inputs to match
+ dummy_joint = self.dummy_joint_input
+ if sdpa_input_dtype is not None:
+ if v_BHNE.dtype != sdpa_input_dtype:
+ v_BHNE = ttnn.typecast(v_BHNE, sdpa_input_dtype)
+ if dummy_joint.dtype != sdpa_input_dtype:
+ dummy_joint = ttnn.typecast(dummy_joint, sdpa_input_dtype)
+
+ # HACK: pass null joint inputs to take advantage of ring attention, even though this is self-attention.
+ if self.use_exp_ring_sdpa:
+ spatial_BHNE, prompt_BHLE, _lse = ttnn.transformer.exp_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, self.parallel_config.sequence_parallel.mesh_axis, dtype=k_BHNE.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.dtype
+ ),
+ joint_strategy="rear",
+ logical_n=N,
+ program_config=self.exp_ring_sdpa_program_config,
+ compute_kernel_config=self.sdpa_compute_kernel_config,
+ dim=2,
+ multi_device_global_semaphore=self.ccl_manager.get_exp_ring_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,
+ num_workers_per_link=5,
+ num_buffers_per_channel=32,
+ )
+ else:
+ 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, self.parallel_config.sequence_parallel.mesh_axis, dtype=k_BHNE.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.dtype
+ ),
+ joint_strategy="rear",
+ logical_n=N,
+ program_config=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,
+ )
+ else:
+ spatial_BHNE = ttnn.transformer.scaled_dot_product_attention(
+ q_BHNE,
+ k_BHNE,
+ v_BHNE,
+ is_causal=False,
+ program_config=self.sdpa_program_config,
+ compute_kernel_config=self.sdpa_compute_kernel_config,
+ )
+ else:
+ # Cross attention
+ spatial_BHNE = ttnn.transformer.scaled_dot_product_attention(
+ q_BHNE,
+ k_BHNE,
+ v_BHNE,
+ is_causal=False,
+ attn_mask=cross_attn_mask,
+ program_config=self.sdpa_program_config,
+ compute_kernel_config=self.sdpa_compute_kernel_config,
+ )
+
+ spatial_1BND = ttnn.transformer.concatenate_heads(spatial_BHNE)
+ spatial_1BND = ttnn.unsqueeze(spatial_1BND, 0)
+
+ if use_nonfused_agmm:
+ # Gather spatial on TP axis before projection
+ spatial_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+
+ if addcmul_residual is not None and addcmul_gate is not None:
+ # Fused to_out projection + addcmul (self-attention only)
+ 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 use_nonfused_agmm else self.parallel_config,
+ )
+
+ return spatial_1BND
diff --git a/code/models/tt_dit/models/transformers/wan2_2/transformer_wan.py b/code/models/tt_dit/models/transformers/wan2_2/transformer_wan.py
new file mode 100644
index 0000000000000000000000000000000000000000..8668c15ec5966f47a195923ccdc20d66271f8d54
--- /dev/null
+++ b/code/models/tt_dit/models/transformers/wan2_2/transformer_wan.py
@@ -0,0 +1,779 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from diffusers import WanTransformer3DModel as TorchWanTransformer3DModel
+from diffusers.models.transformers.transformer_wan import WanRotaryPosEmbed as TorchWanRotaryPosEmbed
+from loguru import logger
+
+if not os.environ.get("TT_DIT_DEBUG"):
+ logger.disable(__name__)
+
+
+import ttnn
+
+from ....layers.embeddings import WanPatchEmbed, WanTimeTextImageEmbedding
+from ....layers.feedforward import ParallelFeedForward
+from ....layers.linear import Linear
+from ....layers.module import Module, ModuleList, Parameter
+from ....layers.normalization import DistributedLayerNorm
+from ....parallel.config import DiTParallelConfig
+from ....parallel.manager import CCLManager
+from ....utils import cache
+from ....utils.mochi import get_rot_transformation_mat
+from ....utils.padding import pad_vision_seq_parallel
+from ....utils.substate import pop_substate, rename_substate
+from ....utils.tensor import bf16_tensor, from_torch, unflatten
+from ....utils.tracing import traced_function
+from .attention_wan import WanAttention
+
+
+class WanTransformerBlock(Module):
+ def __init__(
+ self,
+ *,
+ dim: int,
+ ffn_dim: int,
+ num_heads: int,
+ cross_attention_norm: bool = True,
+ eps: float = 1e-6,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager | None = None,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool = False,
+ sdpa_chunk_size_overrides: dict | None = None,
+ lora_enabled: bool = False,
+ ) -> None:
+ super().__init__()
+
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.num_heads = num_heads
+ self.cross_attention_norm = cross_attention_norm
+ self.eps = eps
+
+ self.mesh_device = mesh_device
+ self.ccl_manager = ccl_manager
+ self.parallel_config = parallel_config
+
+ fsdp_mesh_axis = self.parallel_config.sequence_parallel.mesh_axis if is_fsdp else None
+
+ self.norm1 = DistributedLayerNorm(
+ dim,
+ norm_eps=eps,
+ norm_elementwise_affine=False,
+ bias=False,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+
+ self.attn1 = WanAttention(
+ dim=dim,
+ num_heads=num_heads,
+ eps=eps,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ is_self=True,
+ sdpa_chunk_size_overrides=sdpa_chunk_size_overrides,
+ lora_enabled=lora_enabled,
+ )
+
+ self.attn2 = WanAttention(
+ dim=dim,
+ num_heads=num_heads,
+ eps=eps,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ is_self=False,
+ sdpa_chunk_size_overrides=sdpa_chunk_size_overrides,
+ lora_enabled=lora_enabled,
+ )
+
+ self.norm2 = (
+ DistributedLayerNorm(
+ dim,
+ norm_eps=eps,
+ norm_elementwise_affine=True,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+ if cross_attention_norm
+ else None
+ )
+
+ self.ffn = ParallelFeedForward(
+ dim,
+ inner_dim=ffn_dim,
+ activation_fn="gelu_tanh",
+ bias=True,
+ mesh_device=mesh_device,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ ccl_manager=ccl_manager,
+ fsdp_mesh_axis=fsdp_mesh_axis,
+ lora_enabled=lora_enabled,
+ )
+
+ self.norm3 = DistributedLayerNorm(
+ dim,
+ norm_eps=eps,
+ norm_elementwise_affine=False,
+ bias=False,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+
+ self.scale_shift_table = Parameter(
+ total_shape=[1, 1, 6, dim],
+ mesh_axes=[None, None, None, parallel_config.tensor_parallel.mesh_axis],
+ device=mesh_device,
+ dtype=ttnn.float32,
+ )
+
+ self.ff_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ mesh_device.arch(),
+ math_fidelity=ttnn.MathFidelity.HiFi2,
+ 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)
+
+ def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+ rename_substate(state, "ffn.net.0.proj", "ffn.ff1")
+ rename_substate(state, "ffn.net.2", "ffn.ff2")
+
+ if "scale_shift_table" in state:
+ state["scale_shift_table"] = state["scale_shift_table"].unsqueeze(0)
+
+ def forward(
+ self,
+ spatial_1BND: ttnn.Tensor,
+ prompt_1BLP: ttnn.Tensor,
+ temb_1BTD: ttnn.Tensor,
+ N: int,
+ rope_cos: ttnn.Tensor,
+ rope_sin: ttnn.Tensor,
+ trans_mat: ttnn.Tensor,
+ cross_attn_mask: ttnn.Tensor | None = None,
+ ) -> ttnn.Tensor:
+ """
+ spatial_1BND: fractured N on SP, fractured D on TP
+ prompt_1BLP: replicated on SP, replicated D on TP
+ temb_1BTD: replicated on SP, fractured D on TP
+ N: logical sequence length of the spatial input
+ rope_cos_BANH: fractured N on SP, A (num_heads) on TP
+ rope_sin_BANH: fractured N on SP, A (num_heads) on TP
+ trans_mat: replicated on SP, replicated D on TP
+
+ Outputs:
+ spatial_1BND: fractured N on SP, fractured D on TP
+ """
+
+ assert temb_1BTD.shape[2] == 6, "wan2.2 14b expects 6 chunks in timestep embedding"
+
+ shifted_temb_1BTD = self.scale_shift_table.data + temb_1BTD
+ shift_msa_1B1D, scale_msa_1B1D, gate_msa_1B1D, c_shift_msa_1B1D, c_scale_msa_1B1D, c_gate_msa_1B1D = ttnn.chunk(
+ shifted_temb_1BTD, 6, dim=2
+ )
+
+ # NOTE: workaround - addcmul (fused and unfused) is less accurate with fp32 gate input
+ gate_msa_1B1D = ttnn.typecast(gate_msa_1B1D, dtype=ttnn.bfloat16)
+ c_gate_msa_1B1D = ttnn.typecast(c_gate_msa_1B1D, dtype=ttnn.bfloat16)
+
+ spatial_normed_1BND = self.norm1(
+ spatial_1BND, dynamic_weight=(1.0 + scale_msa_1B1D), dynamic_bias=shift_msa_1B1D
+ )
+
+ # Self attention on spatial with fused residual addcmul
+ # Fuses: spatial_1BND = spatial_1BND + to_out(attn_output) * gate_msa_1B1D
+ spatial_1BND = self.attn1(
+ spatial_1BND=spatial_normed_1BND,
+ N=N,
+ rope_cos=rope_cos,
+ rope_sin=rope_sin,
+ trans_mat=trans_mat,
+ addcmul_residual=spatial_1BND,
+ addcmul_gate=gate_msa_1B1D,
+ )
+
+ # Cross attention on prompt
+ spatial_normed_1BND = self.norm2(spatial_1BND)
+
+ attn_output_1BND = self.attn2(
+ spatial_1BND=spatial_normed_1BND,
+ N=N,
+ prompt_1BLP=prompt_1BLP,
+ cross_attn_mask=cross_attn_mask,
+ )
+ spatial_1BND = spatial_1BND + attn_output_1BND
+
+ # Feed Forward
+ spatial_normed_1BND = self.norm3(
+ spatial_1BND, dynamic_weight=(1.0 + c_scale_msa_1B1D), dynamic_bias=c_shift_msa_1B1D
+ )
+
+ if self.ccl_manager.topology == ttnn.Topology.Linear:
+ if self.parallel_config.tensor_parallel.factor > 1:
+ spatial_normed_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_normed_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+
+ spatial_ff_1BND = self.ffn(spatial_normed_1BND, compute_kernel_config=self.ff_compute_kernel_config)
+ # spatial_1BND = spatial_1BND + spatial_ff_1BND * c_gate_msa_1B1D
+ # NOTE: higher precision compute config in addcmul may be needed for correctness
+ spatial_1BND = ttnn.addcmul(spatial_1BND, spatial_ff_1BND, c_gate_msa_1B1D)
+ else:
+ # Fused FFN + addcmul at the RS final write step (Phase 2).
+ # Both 'a' (spatial_1BND) and 'b' (c_gate_msa_1B1D) are already at [D/tp] size
+ # on each TP device — no AllGather or scatter matmul needed.
+ spatial_1BND = self.ffn.forward_fused_addcmul(
+ spatial_normed_1BND,
+ spatial_1BND,
+ c_gate_msa_1B1D,
+ scalar=1.0,
+ compute_kernel_config=self.ff_compute_kernel_config,
+ parallel_config=self.parallel_config,
+ )
+
+ return spatial_1BND
+
+
+class WanTransformer3DModel(Module):
+ def __init__(
+ self,
+ *,
+ patch_size: tuple = (1, 2, 2),
+ num_heads: int = 40,
+ dim: int = 5120,
+ in_channels: int = 16,
+ out_channels: int = 16,
+ text_dim: int = 4096,
+ freq_dim: int = 256,
+ ffn_dim: int = 13824,
+ num_layers: int = 40,
+ cross_attn_norm: bool = True,
+ eps: float = 1e-6,
+ rope_max_seq_len: int = 1024,
+ mesh_device: ttnn.MeshDevice,
+ ccl_manager: CCLManager | None = None,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool = True,
+ model_type: str = "t2v",
+ output_dtype: ttnn.DataType = ttnn.float32,
+ lora_enabled: bool = False,
+ ) -> None:
+ super().__init__()
+
+ self.mesh_device = mesh_device
+ self.ccl_manager = ccl_manager
+ self.parallel_config = parallel_config
+ self.is_fsdp = is_fsdp
+ self.lora_enabled = lora_enabled
+ self.fsdp_mesh_axis = self.parallel_config.sequence_parallel.mesh_axis if is_fsdp else None
+ self.model_type = model_type
+ self.cached_rope_features = {}
+ self.output_dtype = output_dtype
+
+ assert model_type in ["t2v", "i2v"], "model_type must be either t2v or i2v"
+ if model_type == "i2v":
+ in_channels = 36
+ else:
+ assert in_channels == 16, "in_channels must be 16 for t2v"
+
+ self.patch_size = patch_size
+ self.dim = dim
+
+ self.out_channels = out_channels
+
+ # NOTE: Fallback
+ self.rope = TorchWanRotaryPosEmbed(
+ dim // num_heads,
+ patch_size,
+ rope_max_seq_len,
+ )
+
+ self.patch_embedding = WanPatchEmbed(
+ patch_size=patch_size,
+ in_channels=in_channels,
+ embed_dim=dim,
+ mesh_device=mesh_device,
+ tp_mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ )
+
+ self.condition_embedder = WanTimeTextImageEmbedding(
+ dim=dim,
+ time_freq_dim=freq_dim,
+ time_proj_dim=dim * 6,
+ text_embed_dim=text_dim,
+ mesh_device=self.mesh_device,
+ tp_mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ ccl_manager=ccl_manager,
+ )
+
+ self.blocks = ModuleList(
+ WanTransformerBlock(
+ dim=dim,
+ ffn_dim=ffn_dim,
+ num_heads=num_heads,
+ cross_attention_norm=cross_attn_norm,
+ eps=eps,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ lora_enabled=lora_enabled,
+ )
+ for i in range(num_layers)
+ )
+
+ self.norm_out = DistributedLayerNorm(
+ dim,
+ norm_eps=eps,
+ norm_elementwise_affine=False,
+ bias=False,
+ mesh_axis=parallel_config.tensor_parallel.mesh_axis,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ )
+
+ self.proj_out = Linear(
+ dim,
+ patch_size[0] * patch_size[1] * patch_size[2] * self.out_channels,
+ bias=True,
+ mesh_device=mesh_device,
+ )
+
+ self.scale_shift_table = Parameter(
+ total_shape=[1, 2, dim],
+ device=mesh_device,
+ mesh_axes=[None, None, parallel_config.tensor_parallel.mesh_axis],
+ dtype=ttnn.float32,
+ )
+
+ self.hifi4_compute_kernel_config = ttnn.init_device_compute_kernel_config(
+ 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)
+
+ def save(self, directory: str | Path, /, *, prefix: str = "") -> None:
+ super().save(directory, prefix=prefix)
+
+ directory = Path(directory)
+
+ # Torch fallbacks
+ torch.save(self.rope.state_dict(), directory / f"{prefix}rope.pt")
+
+ def load(self, directory: str | Path, /, *, prefix: str = "") -> None:
+ super().load(directory, prefix=prefix)
+
+ directory = Path(directory)
+
+ # Torch fallbacks
+ self.rope.load_state_dict(torch.load(directory / f"{prefix}rope.pt"))
+
+ def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
+ # Torch fallbacks
+ self.rope.load_state_dict(pop_substate(state, "rope"))
+
+ def get_rope_features(self, hidden_states):
+ if tuple(hidden_states.shape) not in self.cached_rope_features:
+ rope_features = self.prepare_rope_features(hidden_states)
+ self.cached_rope_features[tuple(hidden_states.shape)] = rope_features
+ return self.cached_rope_features[tuple(hidden_states.shape)]
+
+ def prepare_rope_features(self, hidden_states):
+ """
+ Given video input, compute RoPE features.
+ Return tensors on device.
+ """
+ logger.info(f"Preparing rope features for shape {hidden_states.shape}")
+ rope_cos, rope_sin = self.rope(hidden_states)
+
+ # Convert to TT tensors with proper sharding
+ rope_cos_1HND = rope_cos.permute(0, 2, 1, 3)
+ rope_sin_1HND = rope_sin.permute(0, 2, 1, 3)
+
+ rope_cos_1HND = pad_vision_seq_parallel(
+ rope_cos_1HND, num_devices=self.parallel_config.sequence_parallel.factor
+ )
+ rope_sin_1HND = pad_vision_seq_parallel(
+ rope_sin_1HND, num_devices=self.parallel_config.sequence_parallel.factor
+ )
+
+ trans_mat = get_rot_transformation_mat()
+
+ sp_axis = self.parallel_config.sequence_parallel.mesh_axis
+ tt_rope_cos_1HND = from_torch(
+ rope_cos_1HND,
+ device=self.mesh_device,
+ dtype=ttnn.float32,
+ mesh_axes=[..., sp_axis, None],
+ )
+ tt_rope_sin_1HND = from_torch(
+ rope_sin_1HND,
+ device=self.mesh_device,
+ dtype=ttnn.float32,
+ mesh_axes=[..., sp_axis, None],
+ )
+ tt_trans_mat = bf16_tensor(trans_mat, device=self.mesh_device)
+
+ logger.info(f"TT rope cos shape: {tt_rope_cos_1HND.shape}")
+ logger.info(f"TT rope sin shape: {tt_rope_sin_1HND.shape}")
+ logger.info(f"TT trans mat shape: {tt_trans_mat.shape}")
+
+ return tt_rope_cos_1HND, tt_rope_sin_1HND, tt_trans_mat
+
+ def prepare_text_conditioning(self, encoder_hidden_states):
+ tt_prompt_1BLP = self.condition_embedder.forward_text(encoder_hidden_states)
+
+ logger.info(f"TT prompt shape: {tt_prompt_1BLP.shape}")
+ return tt_prompt_1BLP
+
+ def prepare_timestep_conditioning(self, timestep):
+ tt_temb_11BD, tt_timestep_proj_1BTD = self.condition_embedder.forward_timestep(timestep, timestep_seq_len=None)
+ tt_timestep_proj_1BTD = unflatten(ttnn.squeeze(tt_timestep_proj_1BTD, -2), -1, (6, -1))
+ logger.info(f"TT temb shape: {tt_temb_11BD.shape}")
+ logger.info(f"TT timestep proj shape: {tt_timestep_proj_1BTD.shape}")
+ return tt_temb_11BD, tt_timestep_proj_1BTD
+
+ def prepare_conditioning(self, timestep, encoder_hidden_states):
+ """
+ Given inputs, execute the combined timestep and text embedding.
+ Return tensors on device.
+ """
+ tt_temb_11BD, tt_timestep_proj_1BTD = self.prepare_timestep_conditioning(timestep)
+ tt_prompt_1BLP = self.prepare_text_conditioning(encoder_hidden_states)
+
+ logger.info(f"TT temb shape: {tt_temb_11BD.shape}")
+ logger.info(f"TT timestep proj shape: {tt_timestep_proj_1BTD.shape}")
+ logger.info(f"TT prompt shape: {tt_prompt_1BLP.shape}")
+ return tt_temb_11BD, tt_timestep_proj_1BTD, tt_prompt_1BLP
+
+ def preprocess_spatial_input_host(self, spatial):
+ B, C, F, H, W = spatial.shape
+ logger.info(f"Preprocessing spatial input with shape {spatial.shape}")
+ assert B == 1, "Batch size must be 1"
+ pF, pH, pW = self.patch_size
+ patch_F, patch_H, patch_W = F // pF, H // pH, W // pW
+ N = patch_F * patch_H * patch_W
+
+ # Patchify video input
+ spatial = spatial.reshape(B, C, patch_F, pF, patch_H, pH, patch_W, pW)
+ spatial = spatial.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(1, B, N, pF * pH * pW * C)
+ logger.info(f"spatial input after patchifying: {spatial.shape}")
+
+ spatial = pad_vision_seq_parallel(spatial, num_devices=self.parallel_config.sequence_parallel.factor)
+ logger.info(f"spatial input after padding: {spatial.shape}")
+
+ return spatial, N
+
+ def preprocess_spatial_input(self, spatial):
+ spatial, N = self.preprocess_spatial_input_host(spatial)
+ spatial = bf16_tensor(
+ spatial, device=self.mesh_device, mesh_axis=self.parallel_config.sequence_parallel.mesh_axis, shard_dim=-2
+ )
+ logger.info(f"TT spatial shape: {spatial.shape}")
+ return spatial, N
+
+ def postprocess_spatial_output_host(self, spatial_1BND, F, H, W, N):
+ """
+ This is the reverse of preprocess_spatial_input
+ Input is of shape: 1 B (patch_F patch_H patch_W) (pF pH pW C)
+ returns shape: B C F H W
+ """
+ assert len(spatial_1BND.shape) == 4
+ assert spatial_1BND.shape[0] == 1
+ B = spatial_1BND.shape[1]
+ pF, pH, pW = self.patch_size
+ patch_F, patch_H, patch_W = F // pF, H // pH, W // pW
+ logger.info(f"Postprocessing spatial output with shape {spatial_1BND.shape}")
+ spatial_BND = spatial_1BND.squeeze(0)
+
+ spatial_BND = spatial_BND[:, :N] # Slice out sequence-parallel padding tokens
+ logger.info(f"Spatial output after slicing: {spatial_BND.shape}")
+
+ spatial_patches = spatial_BND.reshape(B, patch_F, patch_H, patch_W, pF, pH, pW, self.out_channels)
+ logger.info(f"Spatial output after reshaping: {spatial_patches.shape}")
+
+ spatial_BCFHW = spatial_patches.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, self.out_channels, F, H, W)
+ logger.info(f"Spatial output after permuting: {spatial_BCFHW.shape}")
+ return spatial_BCFHW
+
+ def postprocess_spatial_output(self, spatial_1BND, F, H, W, N):
+ """
+ This is the reverse of preprocess_spatial_input
+ Input is of shape: 1 B (patch_F patch_H patch_W) (pF pH pW C)
+ returns shape: B C F H W
+ """
+ assert len(spatial_1BND.shape) == 4
+ assert spatial_1BND.shape[0] == 1
+ # Gather sequence-parallel output
+ spatial_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_1BND, dim=2, mesh_axis=self.parallel_config.sequence_parallel.mesh_axis
+ )
+
+ spatial_1BND = ttnn.to_torch(ttnn.get_device_tensors(spatial_1BND)[0])
+
+ spatial_BCFHW = self.postprocess_spatial_output_host(spatial_1BND, F, H, W, N)
+ return spatial_BCFHW
+
+ def forward(
+ self,
+ spatial: torch.Tensor,
+ prompt: ttnn.Tensor,
+ timestep: torch.Tensor,
+ y: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """
+ Inputs are all torch tensors
+ y is an optional argument for image-to-video generation.
+ We assume that preprocessing has already been done for y and y has the same shape as spatial.
+ Output is torch tensor
+ """
+
+ if self.model_type == "i2v":
+ assert y is not None, "y must be provided for image-to-video generation"
+
+ B, C, F, H, W = spatial.shape
+ pF, pH, pW = self.patch_size
+ patch_F, patch_H, patch_W = F // pF, H // pH, W // pW
+ N = patch_F * patch_H * patch_W
+
+ rope_cos_1HND, rope_sin_1HND, trans_mat = self.get_rope_features(spatial)
+
+ temb_11BD, timestep_proj_1BTD, prompt_1BLP = self.prepare_conditioning(timestep, prompt)
+
+ # Concatenate spatial and y along the channel dimension
+ if self.model_type == "i2v":
+ print(spatial.shape, y.shape)
+ spatial = torch.cat([spatial, y], dim=1)
+
+ spatial_1BNI, N = self.preprocess_spatial_input(spatial)
+
+ spatial_1BND = self.patch_embedding(spatial_1BNI)
+
+ for idx, block in enumerate(self.blocks):
+ spatial_1BND = block(
+ spatial_1BND=spatial_1BND,
+ prompt_1BLP=prompt_1BLP,
+ temb_1BTD=timestep_proj_1BTD,
+ N=N,
+ rope_cos=rope_cos_1HND,
+ rope_sin=rope_sin_1HND,
+ trans_mat=trans_mat,
+ )
+
+ scale_shift_1BSD = self.scale_shift_table.data + temb_11BD
+ shift_11BD, scale_11BD = ttnn.chunk(scale_shift_1BSD, 2, -2)
+
+ spatial_norm_1BND = self.norm_out(
+ spatial_1BND, dynamic_weight=(1 + scale_11BD), dynamic_bias=shift_11BD, dtype=ttnn.float32
+ )
+
+ if self.parallel_config.tensor_parallel.factor > 1:
+ spatial_norm_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_norm_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+
+ proj_out_1BNI = self.proj_out(
+ spatial_norm_1BND, compute_kernel_config=self.hifi4_compute_kernel_config, dtype=ttnn.float32
+ )
+
+ spatial_out = self.postprocess_spatial_output(proj_out_1BNI, F, H, W, N)
+
+ return spatial_out
+
+ def inner_step(
+ self, spatial_1BNI, prompt_1BLP, rope_cos_1HND, rope_sin_1HND, trans_mat, N, timestep, gather_output=True
+ ):
+ """
+ Reduced forward function which assumes outer loop has cached certain inputs that are step independent:
+ - prompt_1BLP
+ - rope_cos_1HND
+ - rope_sin_1HND
+ - trans_mat
+ - N
+
+ Spatial input is a tensor with layout `1 B (patch_F patch_H patch_W) (pF pH pW C)`.
+ Spatial output is an fp32 ttnn.Tensor on device with same layout.
+ """
+ temb_11BD, timestep_proj_1BTD = self.prepare_timestep_conditioning(timestep)
+
+ spatial_1BND = self.patch_embedding(spatial_1BNI)
+
+ for block in self.blocks:
+ spatial_1BND = block(
+ spatial_1BND=spatial_1BND,
+ prompt_1BLP=prompt_1BLP,
+ temb_1BTD=timestep_proj_1BTD,
+ N=N,
+ rope_cos=rope_cos_1HND,
+ rope_sin=rope_sin_1HND,
+ trans_mat=trans_mat,
+ )
+ scale_shift_1BSD = self.scale_shift_table.data + temb_11BD
+ shift_11BD, scale_11BD = ttnn.chunk(scale_shift_1BSD, 2, -2)
+
+ spatial_norm_1BND = self.norm_out(
+ spatial_1BND, dynamic_weight=(1 + scale_11BD), dynamic_bias=shift_11BD, dtype=ttnn.float32
+ )
+
+ if self.parallel_config.tensor_parallel.factor > 1:
+ spatial_norm_1BND = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_norm_1BND, dim=3, mesh_axis=self.parallel_config.tensor_parallel.mesh_axis
+ )
+
+ spatial_1BNI = self.proj_out(
+ spatial_norm_1BND, compute_kernel_config=self.hifi4_compute_kernel_config, dtype=self.output_dtype
+ )
+
+ # Gather fp32 spatial output across sequence parallel devices (remains on device)
+ if gather_output:
+ spatial_1BNI = self.ccl_manager.all_gather_persistent_buffer(
+ spatial_1BNI, dim=2, mesh_axis=self.parallel_config.sequence_parallel.mesh_axis
+ )
+
+ return spatial_1BNI
+
+ # Prep run is False because we warmup the entire pipeline first. Remove if this is not desired.
+ @traced_function(device=lambda self: self.mesh_device, clone_prep_inputs=False, prep_run=False)
+ def combined_step(
+ self,
+ do_classifier_free_guidance: bool,
+ spatial_1BNI: ttnn.Tensor,
+ prompt_1BLP: ttnn.Tensor,
+ negative_prompt_1BLP: ttnn.Tensor,
+ N: int,
+ rope_cos_1HND: ttnn.Tensor,
+ rope_sin_1HND: ttnn.Tensor,
+ trans_mat: ttnn.Tensor,
+ timestep: ttnn.Tensor,
+ guidance_scale: ttnn.Tensor,
+ *,
+ gather_output: bool = True,
+ ) -> ttnn.Tensor:
+ cond = self.inner_step(
+ spatial_1BNI,
+ prompt_1BLP,
+ rope_cos_1HND,
+ rope_sin_1HND,
+ trans_mat,
+ N,
+ timestep,
+ gather_output=gather_output,
+ )
+ if not do_classifier_free_guidance:
+ return cond
+
+ uncond = self.inner_step(
+ spatial_1BNI,
+ negative_prompt_1BLP,
+ rope_cos_1HND,
+ rope_sin_1HND,
+ trans_mat,
+ N,
+ timestep,
+ gather_output=gather_output,
+ )
+
+ combined = ttnn.lerp(uncond, cond, guidance_scale)
+
+ return combined
+
+
+class WanCheckpoint:
+ """A Wan transformer-subfolder checkpoint: fetches weights and builds loaded transformers."""
+
+ def __init__(self, name: str, subfolder: str) -> None:
+ self._name = name
+ self._subfolder = subfolder
+ torch_transformer = TorchWanTransformer3DModel.from_pretrained(
+ name,
+ subfolder=subfolder,
+ trust_remote_code=True,
+ )
+ torch_transformer.eval()
+ self._config = torch_transformer.config
+ self._state_dict = torch_transformer.state_dict()
+
+ @property
+ def subfolder(self) -> str:
+ return self._subfolder
+
+ def state_dict(self) -> dict[str, torch.Tensor]:
+ return dict(self._state_dict)
+
+ def build(
+ self,
+ *,
+ ccl_manager: CCLManager,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool,
+ model_type: str,
+ lora_enabled: bool = False,
+ ) -> WanTransformer3DModel:
+ """Construct a ``WanTransformer3DModel`` for this checkpoint (weights NOT loaded).
+
+ Loading is deferred so the caller can manage the lifecycle (deallocate / reload).
+ """
+ c = self._config
+ return WanTransformer3DModel(
+ patch_size=c.patch_size,
+ num_heads=c.num_attention_heads,
+ dim=c.num_attention_heads * c.attention_head_dim,
+ in_channels=c.in_channels,
+ out_channels=c.out_channels,
+ text_dim=c.text_dim,
+ freq_dim=c.freq_dim,
+ ffn_dim=c.ffn_dim,
+ cross_attn_norm=c.cross_attn_norm,
+ eps=c.eps,
+ rope_max_seq_len=c.rope_max_seq_len,
+ mesh_device=ccl_manager.mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ model_type=model_type,
+ lora_enabled=lora_enabled,
+ )
+
+ def load(
+ self,
+ model: WanTransformer3DModel,
+ *,
+ mesh_device: ttnn.MeshDevice,
+ parallel_config: DiTParallelConfig,
+ is_fsdp: bool,
+ ) -> None:
+ """Load (or reload) weights for a previously-built transformer."""
+ cache.load_model(
+ tt_model=model,
+ get_torch_state_dict=lambda: self._state_dict,
+ model_name=os.path.basename(self._name),
+ subfolder=self._subfolder,
+ parallel_config=parallel_config,
+ mesh_shape=tuple(mesh_device.shape),
+ mesh_device=mesh_device,
+ is_fsdp=is_fsdp,
+ )
diff --git a/code/models/tt_dit/tests/dataset_eval/clip_encoder.py b/code/models/tt_dit/tests/dataset_eval/clip_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b082724bf02005f851e1b47af88e9e18ef34699
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/clip_encoder.py
@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List, Optional, Union
+
+import open_clip
+import torch
+import torch.nn as nn
+from PIL import Image
+
+
+class CLIPEncoder(nn.Module):
+ """
+ A class for encoding images and texts using a specified CLIP model and computing the similarity between them.
+
+ Attributes:
+ -----------
+ clip_version: str
+ The version of the CLIP model to be used.
+ pretrained: str
+ The pre-trained weights to load.
+ model: nn.Module
+ The CLIP model.
+ preprocess: Callable
+ The preprocessing transform to apply to the input image.
+ device: str
+ The device to which the model is moved.
+ """
+
+ def __init__(
+ self,
+ clip_version: str = "ViT-B/32",
+ pretrained: Optional[str] = "",
+ cache_dir: Optional[str] = None,
+ ):
+ """
+ Initializes the CLIPEncoder with the specified CLIP model version and pre-trained weights.
+
+ Parameters:
+ -----------
+ clip_version: str, optional
+ The version of the CLIP model to be used. Defaults to 'ViT-B/32'.
+ pretrained: str, optional
+ The pre-trained weights to load. If not provided, it defaults based on clip_version.
+ cache_dir: str, optional
+ The directory to cache the model. Defaults to None.
+ device: str, optional
+ The device to which the model is moved. Defaults to 'cuda'.
+ """
+ super().__init__()
+
+ self.clip_version = clip_version
+ self.pretrained = pretrained if pretrained else self._get_default_pretrained()
+
+ self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+ self.clip_version, pretrained=self.pretrained, cache_dir=cache_dir
+ )
+
+ self.model.eval()
+
+ def _get_default_pretrained(self) -> str:
+ """Returns the default pretrained weights based on the clip_version."""
+ return "openai"
+
+ @torch.no_grad()
+ def get_clip_score(self, text: Union[str, List[str]], image: Union[Image.Image, torch.Tensor]) -> torch.Tensor:
+ """
+ Computes the similarity score between the given text(s) and image using the CLIP model.
+
+ Parameters:
+ -----------
+ text: Union[str, List[str]]
+ The text or list of texts to compare with the image.
+ image: Image.Image
+ The input image.
+
+ Returns:
+ --------
+ torch.Tensor
+ The similarity score between the text(s) and image.
+ """
+ # Preprocess the image and move it to the specified device
+ image = self.preprocess(image).unsqueeze(0)
+
+ # Normalize the image features
+ image_features = self.model.encode_image(image).float()
+ image_features /= image_features.norm(dim=-1, keepdim=True)
+
+ # If a single text string is provided, convert it to a list
+ if not isinstance(text, (list, tuple)):
+ text = [text]
+
+ # Tokenize the text and move it to the specified device
+ text = open_clip.tokenize(text)
+
+ # Normalize the text features
+ text_features = self.model.encode_text(text).float()
+ text_features /= text_features.norm(dim=-1, keepdim=True)
+
+ # Compute the similarity between the image and text features
+ similarity = image_features @ text_features.T
+
+ return similarity
diff --git a/code/models/tt_dit/tests/dataset_eval/conftest.py b/code/models/tt_dit/tests/dataset_eval/conftest.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c325baff574b41eacb15aa3493e34e23fe3bea4
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/conftest.py
@@ -0,0 +1,208 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+import ttnn
+
+from ...pipelines.flux1.pipeline_flux1 import Flux1Pipeline
+from ...pipelines.mochi.pipeline_mochi import MochiPipeline
+from ...pipelines.motif.pipeline_motif import MotifPipeline
+from ...pipelines.stable_diffusion_35_large.pipeline_stable_diffusion_35_large import StableDiffusion3Pipeline
+from ...pipelines.wan.pipeline_wan import WanPipeline
+
+targets_setup = {
+ "stable-diffusion-3.5-large": {
+ "hf_id": "stabilityai/stable-diffusion-3.5-large",
+ "num_inference_steps": 28,
+ "pipeline": StableDiffusion3Pipeline,
+ "device_params": {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "l1_small_size": 32768,
+ "trace_region_size": 25000000,
+ },
+ "accuracy": {
+ "5000": {"fid_valid_range": [22.63592, 24.0361780112], "clip_valid_range": [30.38525, 32.264881565]},
+ "500": {"fid_valid_range": [79.734, 84.66634524], "clip_valid_range": [30.423468, 32.30546373048]},
+ "100": {
+ "fid_valid_range": [175.35951, 186.2072492886],
+ "clip_valid_range": [30.162635, 32.028495601100005],
+ },
+ },
+ },
+ "flux.1-dev": {
+ "hf_id": "black-forest-labs/FLUX.1-dev",
+ "num_inference_steps": 28,
+ "pipeline": Flux1Pipeline,
+ "device_params": {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "l1_small_size": 32768,
+ "trace_region_size": 50000000,
+ },
+ "accuracy": {
+ "5000": {"fid_valid_range": [28.14067, 29.8814518462], "clip_valid_range": [29.7014, 31.538728604]},
+ "500": {"fid_valid_range": [82.37143, 87.46692665980001], "clip_valid_range": [29.79258, 31.6355489988]},
+ "100": {"fid_valid_range": [174.97054, 185.79421760440002], "clip_valid_range": [29.86339, 31.7107393054]},
+ },
+ },
+ "flux.1-schnell": {
+ "hf_id": "black-forest-labs/FLUX.1-schnell",
+ "num_inference_steps": 4,
+ "pipeline": Flux1Pipeline,
+ "device_params": {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "l1_small_size": 32768,
+ "trace_region_size": 50000000,
+ },
+ "accuracy": {
+ "5000": {
+ "fid_valid_range": [22.195637, 23.568659104820004],
+ "clip_valid_range": [30.010539, 31.866990942540003],
+ },
+ "500": {"fid_valid_range": [77.21976, 81.9965743536], "clip_valid_range": [30.19901, 32.067120758600005]},
+ "100": {"fid_valid_range": [171.065611, 181.64772969646], "clip_valid_range": [30.41629, 32.2978416994]},
+ },
+ },
+ "motif-image-6b-preview": {
+ "hf_id": "Motif-Technologies/Motif-Image-6B-Preview",
+ "num_inference_steps": 28,
+ "pipeline": MotifPipeline,
+ "device_params": {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "l1_small_size": 32768,
+ "trace_region_size": 31000000,
+ },
+ "accuracy": {
+ "5000": {"fid_valid_range": [33.85203, 35.9461165758], "clip_valid_range": [29.6529, 31.487228394]},
+ "500": {"fid_valid_range": [86.02833, 91.3500424938], "clip_valid_range": [29.65581, 31.4903184066]},
+ "100": {"fid_valid_range": [180.97872, 192.1740636192], "clip_valid_range": [29.65581, 31.4903184066]},
+ },
+ },
+ "mochi-1-preview": {
+ "hf_id": "genmo/mochi-1-preview",
+ "num_inference_steps": 50,
+ "pipeline": MochiPipeline,
+ "device_params": {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ },
+ "accuracy": {"16": {"clip_valid_range": [26.20977, 27.83099]}},
+ "is_video": True,
+ "extra_args": {"num_frames": 168},
+ },
+ "wan2.2": {
+ "hf_id": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+ "num_inference_steps": 40,
+ "pipeline": WanPipeline,
+ "device_params": {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ },
+ "accuracy": {"16": {"clip_valid_range": [28.25276, 30.00036]}},
+ "is_video": True,
+ "extra_args": {"output_type": "pil"},
+ },
+}
+
+
+def pytest_addoption(parser):
+ parser.addoption(
+ "--start-from",
+ action="store",
+ default=0,
+ help="Start from prompt number (0-4999)",
+ )
+ parser.addoption(
+ "--num-prompts",
+ action="store",
+ default=5000,
+ help="Number of prompts to process (default: 5000)",
+ )
+
+ model_id_choices = targets_setup.keys()
+ parser.addoption(
+ "--model-id",
+ action="store",
+ required=True,
+ choices=model_id_choices,
+ help=f"Model ID to use for evaluation. Options:{','.join(model_id_choices)}",
+ )
+
+ parser.addoption(
+ "--num-inference-steps",
+ action="store",
+ default=None,
+ type=int,
+ help="Number of inference steps (default: 28)",
+ )
+
+ parser.addoption(
+ "--24hrs-test",
+ action="store_true",
+ default=False,
+ help="Whether to run a 24-hour test",
+ )
+
+
+@pytest.fixture
+def evaluation_range(request):
+ start_from = request.config.getoption("--start-from")
+ num_prompts = request.config.getoption("--num-prompts")
+ if start_from is not None:
+ start_from = int(start_from)
+ else:
+ start_from = 0
+ if num_prompts is not None:
+ num_prompts = int(num_prompts)
+ else:
+ num_prompts = 5
+ return start_from, num_prompts
+
+
+@pytest.fixture
+def model_id(request):
+ return request.config.getoption("--model-id")
+
+
+@pytest.fixture
+def is_24hrs_test(request):
+ return request.config.getoption("--24hrs-test")
+
+
+@pytest.fixture
+def num_inference_steps(request):
+ return (
+ request.config.getoption("--num-inference-steps")
+ or targets_setup[request.config.getoption("--model-id")]["num_inference_steps"]
+ )
+
+
+@pytest.fixture
+def dit_pipeline(model_id):
+ return targets_setup[model_id]["pipeline"]
+
+
+@pytest.fixture
+def model_metadata(model_id):
+ return targets_setup[model_id]
+
+
+@pytest.fixture
+def device_params(request, model_id):
+ """Return device_params based on model_id, with optional parametrization override."""
+
+ # If parametrized (indirect=True), use that value as base
+ params = getattr(request, "param", {}).copy()
+
+ # If no parametrization, use model-specific params
+ if not params:
+ params = targets_setup[model_id]["device_params"]
+ # If parametrized, merge with model-specific (parametrized takes precedence)
+ else:
+ model_defaults = targets_setup[model_id]["device_params"]
+ params = {**model_defaults, **params} # Parametrized overrides model defaults
+
+ return params
+
+
+@pytest.fixture
+def is_video(model_id):
+ return targets_setup[model_id].get("is_video", False)
diff --git a/code/models/tt_dit/tests/dataset_eval/fid_score.py b/code/models/tt_dit/tests/dataset_eval/fid_score.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f210c8e3257f3cc973063736b2b5179182af7e5
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/fid_score.py
@@ -0,0 +1,214 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import urllib
+
+import numpy as np
+import torch
+from loguru import logger
+from PIL import Image
+from scipy import linalg
+from threadpoolctl import threadpool_limits
+from torchvision import transforms as TF
+
+from .inception import InceptionV3
+
+COCO_STATISTICS_DOWNLOAD_PATH = "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/tools/val2014.npz"
+
+
+def get_activations(files, model, batch_size=1, dims=2048, device="cpu", num_workers=1):
+ """Calculates the activations of the pool_3 layer for all images.
+
+ Params:
+ -- files : List of image files paths
+ -- model : Instance of inception model
+ -- batch_size : Batch size of images for the model to process at once.
+ Make sure that the number of samples is a multiple of
+ the batch size, otherwise some samples are ignored. This
+ behavior is retained to match the original FID score
+ implementation.
+ -- dims : Dimensionality of features returned by Inception
+ -- device : Device to run calculations
+ -- num_workers : Number of parallel dataloader workers
+
+ Returns:
+ -- A numpy array of dimension (num images, dims) that contains the
+ activations of the given tensor when feeding inception with the
+ query tensor.
+ """
+ model.eval()
+
+ if batch_size > len(files):
+ logger.warning("Warning: batch size is bigger than the data size. Setting batch size to data size")
+ batch_size = len(files)
+
+ dataset = ImagesDataset(files, transforms=TF.ToTensor())
+ dataloader = torch.utils.data.DataLoader(
+ dataset,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ )
+
+ pred_arr = np.empty((len(files), dims))
+
+ start_idx = 0
+
+ for batch in dataloader:
+ batch = batch.to(device)
+
+ with torch.no_grad():
+ pred = model(batch)[0]
+
+ pred = pred.squeeze(3).squeeze(2).cpu().numpy()
+
+ pred_arr[start_idx : start_idx + pred.shape[0]] = pred
+
+ start_idx = start_idx + pred.shape[0]
+
+ return pred_arr
+
+
+def calculate_activation_statistics(files, model):
+ """Calculation of the statistics used by the FID.
+ Params:
+ -- files : List of images
+ -- model : Instance of inception model
+
+ Returns:
+ -- mu : The mean over samples of the activations of the pool_3 layer of
+ the inception model.
+ -- sigma : The covariance matrix of the activations of the pool_3 layer of
+ the inception model.
+ """
+ act = get_activations(files, model)
+ mu = np.mean(act, axis=0)
+ sigma = np.cov(act, rowvar=False)
+ return mu, sigma
+
+
+def compute_statistics_of_path(path):
+ if not os.path.isfile(path):
+ logger.info(f"File {path} not found. Downloading...")
+ os.makedirs(os.path.dirname(path), exist_ok=True)
+ urllib.request.urlretrieve(COCO_STATISTICS_DOWNLOAD_PATH, path)
+ logger.info("Download complete.")
+
+ with np.load(path) as f:
+ m, s = f["mu"][:], f["sigma"][:]
+
+ return m, s
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+ """Numpy implementation of the Frechet Distance.
+ The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+ and X_2 ~ N(mu_2, C_2) is
+ d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+ Stable version by Dougal J. Sutherland.
+
+ Params:
+ -- mu1 : Numpy array containing the activations of a layer of the
+ inception net (like returned by the function 'get_predictions')
+ for generated samples.
+ -- mu2 : The sample mean over activations, precalculated on an
+ representative data set.
+ -- sigma1: The covariance matrix over activations for generated samples.
+ -- sigma2: The covariance matrix over activations, precalculated on an
+ representative data set.
+
+ Returns:
+ -- : The Frechet Distance.
+ """
+
+ mu1 = np.atleast_1d(mu1)
+ mu2 = np.atleast_1d(mu2)
+
+ sigma1 = np.atleast_2d(sigma1)
+ sigma2 = np.atleast_2d(sigma2)
+
+ assert mu1.shape == mu2.shape, "Training and test mean vectors have different lengths"
+ assert sigma1.shape == sigma2.shape, "Training and test covariances have different dimensions"
+
+ diff = mu1 - mu2
+
+ # save sigma matrices in case of hangs
+ np.save("sigma1.npy", sigma1.astype(np.float32))
+ np.save("sigma2.npy", sigma2.astype(np.float32))
+ logger.info("Saved sigma1.npy and sigma2.npy for manual continuation.")
+
+ # ensure float32 to avoid massive slowdowns
+ sigma1 = sigma1.astype(np.float32)
+ sigma2 = sigma2.astype(np.float32)
+
+ # compute dot product
+ logger.info("calculating dot")
+ s_dot_s = sigma1 @ sigma2
+
+ # compute matrix square root with threadpool limit
+ logger.info("calculating sqrtm...")
+ with threadpool_limits(limits=1): # limit mkl/openblas threads only here
+ covmean, _ = linalg.sqrtm(s_dot_s, disp=False)
+ logger.info("done sqrtm")
+
+ # handle numerical issues
+ if not np.isfinite(covmean).all():
+ logger.info(f"fid calculation produces singular product; adding {eps} to diagonal of cov estimates")
+ offset = np.eye(sigma1.shape[0], dtype=np.float32) * eps
+ with threadpool_limits(limits=1):
+ covmean, _ = linalg.sqrtm((sigma1 + offset) @ (sigma2 + offset), disp=False)
+
+ # handle imaginary parts due to sqrtm precision issues
+ if np.iscomplexobj(covmean):
+ if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+ m = np.max(np.abs(covmean.imag))
+ raise ValueError(f"Imaginary component {m}")
+ covmean = covmean.real
+
+ covmean = covmean.astype(np.float64)
+ tr_covmean = np.trace(covmean)
+
+ return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
+
+
+def calculate_fid_score(images, coco_statistics_path, inception_dims=2048):
+ """Calculates the activations of the pool_3 layer for all images.
+
+ Params:
+ -- images : List of images
+ -- coco_statistics_path : Path to the precomputed statistics of the COCO dataset
+ -- inception_dims : Dimensionality of features returned by Inception
+ """
+ block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[inception_dims]
+
+ model = InceptionV3([block_idx])
+ model.eval()
+
+ m1, s1 = compute_statistics_of_path(coco_statistics_path)
+
+ m2, s2 = calculate_activation_statistics(images, model)
+
+ # Ensure dimensions match before calculating FID
+ assert s1.shape == s2.shape, f"Covariance shapes mismatch: {s1.shape} vs {s2.shape}"
+ assert m1.shape == m2.shape, f"Mean shapes mismatch: {m1.shape} vs {m2.shape}"
+
+ return calculate_frechet_distance(m1, s1, m2, s2)
+
+
+class ImagesDataset(torch.utils.data.Dataset):
+ def __init__(self, imgs, transforms=None):
+ self.imgs = imgs
+ self.transforms = transforms
+
+ def __len__(self):
+ return len(self.imgs)
+
+ def __getitem__(self, i):
+ img = self.imgs[i]
+ if not isinstance(img, Image.Image):
+ raise TypeError(f"Unsupported image type: {type(img)}")
+ if self.transforms is not None:
+ img = self.transforms(img)
+ return img
diff --git a/code/models/tt_dit/tests/dataset_eval/inception.py b/code/models/tt_dit/tests/dataset_eval/inception.py
new file mode 100644
index 0000000000000000000000000000000000000000..3cde165db84bf2e87de0ac556368bcffeab635b6
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/inception.py
@@ -0,0 +1,341 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+try:
+ from torchvision.models.utils import load_state_dict_from_url
+except ImportError:
+ from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# Inception weights ported to Pytorch from
+FID_WEIGHTS_URL = (
+ "https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth"
+)
+
+
+class InceptionV3(nn.Module):
+ """Pretrained InceptionV3 network returning feature maps"""
+
+ # Index of default block of inception to return,
+ # corresponds to output of final average pooling
+ DEFAULT_BLOCK_INDEX = 3
+
+ # Maps feature dimensionality to their output blocks indices
+ BLOCK_INDEX_BY_DIM = {
+ 64: 0, # First max pooling features
+ 192: 1, # Second max pooling features
+ 768: 2, # Pre-aux classifier features
+ 2048: 3, # Final average pooling features
+ }
+
+ def __init__(
+ self,
+ output_blocks=(DEFAULT_BLOCK_INDEX,),
+ resize_input=True,
+ normalize_input=True,
+ requires_grad=False,
+ use_fid_inception=True,
+ ):
+ """Build pretrained InceptionV3
+
+ Parameters
+ ----------
+ output_blocks : list of int
+ Indices of blocks to return features of. Possible values are:
+ - 0: corresponds to output of first max pooling
+ - 1: corresponds to output of second max pooling
+ - 2: corresponds to output which is fed to aux classifier
+ - 3: corresponds to output of final average pooling
+ resize_input : bool
+ If true, bilinearly resizes input to width and height 299 before
+ feeding input to model. As the network without fully connected
+ layers is fully convolutional, it should be able to handle inputs
+ of arbitrary size, so resizing might not be strictly needed
+ normalize_input : bool
+ If true, scales the input from range (0, 1) to the range the
+ pretrained Inception network expects, namely (-1, 1)
+ requires_grad : bool
+ If true, parameters of the model require gradients. Possibly useful
+ for finetuning the network
+ use_fid_inception : bool
+ If true, uses the pretrained Inception model used in Tensorflow's
+ FID implementation. If false, uses the pretrained Inception model
+ available in torchvision. The FID Inception model has different
+ weights and a slightly different structure from torchvision's
+ Inception model. If you want to compute FID scores, you are
+ strongly advised to set this parameter to true to get comparable
+ results.
+ """
+ super(InceptionV3, self).__init__()
+
+ self.resize_input = resize_input
+ self.normalize_input = normalize_input
+ self.output_blocks = sorted(output_blocks)
+ self.last_needed_block = max(output_blocks)
+
+ assert self.last_needed_block <= 3, "Last possible output block index is 3"
+
+ self.blocks = nn.ModuleList()
+
+ if use_fid_inception:
+ inception = fid_inception_v3()
+ else:
+ inception = _inception_v3(weights="DEFAULT")
+
+ # Block 0: input to maxpool1
+ block0 = [
+ inception.Conv2d_1a_3x3,
+ inception.Conv2d_2a_3x3,
+ inception.Conv2d_2b_3x3,
+ nn.MaxPool2d(kernel_size=3, stride=2),
+ ]
+ self.blocks.append(nn.Sequential(*block0))
+
+ # Block 1: maxpool1 to maxpool2
+ if self.last_needed_block >= 1:
+ block1 = [
+ inception.Conv2d_3b_1x1,
+ inception.Conv2d_4a_3x3,
+ nn.MaxPool2d(kernel_size=3, stride=2),
+ ]
+ self.blocks.append(nn.Sequential(*block1))
+
+ # Block 2: maxpool2 to aux classifier
+ if self.last_needed_block >= 2:
+ block2 = [
+ inception.Mixed_5b,
+ inception.Mixed_5c,
+ inception.Mixed_5d,
+ inception.Mixed_6a,
+ inception.Mixed_6b,
+ inception.Mixed_6c,
+ inception.Mixed_6d,
+ inception.Mixed_6e,
+ ]
+ self.blocks.append(nn.Sequential(*block2))
+
+ # Block 3: aux classifier to final avgpool
+ if self.last_needed_block >= 3:
+ block3 = [
+ inception.Mixed_7a,
+ inception.Mixed_7b,
+ inception.Mixed_7c,
+ nn.AdaptiveAvgPool2d(output_size=(1, 1)),
+ ]
+ self.blocks.append(nn.Sequential(*block3))
+
+ for param in self.parameters():
+ param.requires_grad = requires_grad
+
+ def forward(self, inp):
+ """Get Inception feature maps
+
+ Parameters
+ ----------
+ inp : torch.autograd.Variable
+ Input tensor of shape Bx3xHxW. Values are expected to be in
+ range (0, 1)
+
+ Returns
+ -------
+ List of torch.autograd.Variable, corresponding to the selected output
+ block, sorted ascending by index
+ """
+ outp = []
+ x = inp
+
+ if self.resize_input:
+ x = F.interpolate(x, size=(299, 299), mode="bilinear", align_corners=False)
+
+ if self.normalize_input:
+ x = 2 * x - 1 # Scale from range (0, 1) to range (-1, 1)
+
+ for idx, block in enumerate(self.blocks):
+ x = block(x)
+ if idx in self.output_blocks:
+ outp.append(x)
+
+ if idx == self.last_needed_block:
+ break
+
+ return outp
+
+
+def _inception_v3(*args, **kwargs):
+ """Wraps `torchvision.models.inception_v3`"""
+ try:
+ version = tuple(map(int, torchvision.__version__.split(".")[:2]))
+ except ValueError:
+ # Just a caution against weird version strings
+ version = (0,)
+
+ # Skips default weight initialization if supported by torchvision
+ # version. See https://github.com/mseitzer/pytorch-fid/issues/28.
+ if version >= (0, 6):
+ kwargs["init_weights"] = False
+
+ # Backwards compatibility: `weights` argument was handled by `pretrained`
+ # argument prior to version 0.13.
+ if version < (0, 13) and "weights" in kwargs:
+ if kwargs["weights"] == "DEFAULT":
+ kwargs["pretrained"] = True
+ elif kwargs["weights"] is None:
+ kwargs["pretrained"] = False
+ else:
+ raise ValueError(
+ "weights=={} not supported in torchvision {}".format(kwargs["weights"], torchvision.__version__)
+ )
+ del kwargs["weights"]
+
+ return torchvision.models.inception_v3(*args, **kwargs)
+
+
+def fid_inception_v3():
+ """Build pretrained Inception model for FID computation
+
+ The Inception model for FID computation uses a different set of weights
+ and has a slightly different structure than torchvision's Inception.
+
+ This method first constructs torchvision's Inception and then patches the
+ necessary parts that are different in the FID Inception model.
+ """
+ inception = _inception_v3(num_classes=1008, aux_logits=False, weights=None)
+ inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
+ inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
+ inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
+ inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
+ inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
+ inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
+ inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
+ inception.Mixed_7b = FIDInceptionE_1(1280)
+ inception.Mixed_7c = FIDInceptionE_2(2048)
+
+ state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
+ inception.load_state_dict(state_dict)
+ return inception
+
+
+class FIDInceptionA(torchvision.models.inception.InceptionA):
+ """InceptionA block patched for FID computation"""
+
+ def __init__(self, in_channels, pool_features):
+ super(FIDInceptionA, self).__init__(in_channels, pool_features)
+
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch5x5 = self.branch5x5_1(x)
+ branch5x5 = self.branch5x5_2(branch5x5)
+
+ branch3x3dbl = self.branch3x3dbl_1(x)
+ branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+ branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+ # Patch: Tensorflow's average pool does not use the padded zero's in
+ # its average calculation
+ branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
+ branch_pool = self.branch_pool(branch_pool)
+
+ outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+ return torch.cat(outputs, 1)
+
+
+class FIDInceptionC(torchvision.models.inception.InceptionC):
+ """InceptionC block patched for FID computation"""
+
+ def __init__(self, in_channels, channels_7x7):
+ super(FIDInceptionC, self).__init__(in_channels, channels_7x7)
+
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch7x7 = self.branch7x7_1(x)
+ branch7x7 = self.branch7x7_2(branch7x7)
+ branch7x7 = self.branch7x7_3(branch7x7)
+
+ branch7x7dbl = self.branch7x7dbl_1(x)
+ branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+ branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+ branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+ branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+ # Patch: Tensorflow's average pool does not use the padded zero's in
+ # its average calculation
+ branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
+ branch_pool = self.branch_pool(branch_pool)
+
+ outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+ return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_1(torchvision.models.inception.InceptionE):
+ """First InceptionE block patched for FID computation"""
+
+ def __init__(self, in_channels):
+ super(FIDInceptionE_1, self).__init__(in_channels)
+
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch3x3 = self.branch3x3_1(x)
+ branch3x3 = [
+ self.branch3x3_2a(branch3x3),
+ self.branch3x3_2b(branch3x3),
+ ]
+ branch3x3 = torch.cat(branch3x3, 1)
+
+ branch3x3dbl = self.branch3x3dbl_1(x)
+ branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+ branch3x3dbl = [
+ self.branch3x3dbl_3a(branch3x3dbl),
+ self.branch3x3dbl_3b(branch3x3dbl),
+ ]
+ branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+ # Patch: Tensorflow's average pool does not use the padded zero's in
+ # its average calculation
+ branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
+ branch_pool = self.branch_pool(branch_pool)
+
+ outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+ return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_2(torchvision.models.inception.InceptionE):
+ """Second InceptionE block patched for FID computation"""
+
+ def __init__(self, in_channels):
+ super(FIDInceptionE_2, self).__init__(in_channels)
+
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch3x3 = self.branch3x3_1(x)
+ branch3x3 = [
+ self.branch3x3_2a(branch3x3),
+ self.branch3x3_2b(branch3x3),
+ ]
+ branch3x3 = torch.cat(branch3x3, 1)
+
+ branch3x3dbl = self.branch3x3dbl_1(x)
+ branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+ branch3x3dbl = [
+ self.branch3x3dbl_3a(branch3x3dbl),
+ self.branch3x3dbl_3b(branch3x3dbl),
+ ]
+ branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+ # Patch: The FID Inception model uses max pooling instead of average
+ # pooling. This is likely an error in this specific Inception
+ # implementation, as other Inception models use average pooling here
+ # (which matches the description in the paper).
+ branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+ branch_pool = self.branch_pool(branch_pool)
+
+ outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+ return torch.cat(outputs, 1)
diff --git a/code/models/tt_dit/tests/dataset_eval/reference/README.md b/code/models/tt_dit/tests/dataset_eval/reference/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..fd8da6f972369274d39ae485c8e438dfba222315
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/reference/README.md
@@ -0,0 +1,103 @@
+# Reference Implementations for Dataset Evaluation
+
+This directory contains reference implementations for evaluating models on H100 hardware. These implementations use standard diffusers pipelines and provide baseline metrics for comparison with TT-Metal implementations.
+
+## Setup
+
+1. `pip install -r requirements.txt`
+2. Download COCO captions:
+ - `wget https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/coco2014/captions/captions_source.tsv`
+ - Rename file to `captions.tsv`
+3. Download COCO validation statistics:
+ - `wget https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/tools/val2014.npz`
+4. Login to HuggingFace: `huggingface-cli login`
+ - Enter your token
+ - Accept license agreements as required for each model
+
+## Flux Evaluation
+
+**File**: `flux.py`
+
+**Steps**:
+
+1. Load weights locally:
+ - `pip install accelerate`
+ - ```
+ huggingface-cli download black-forest-labs/FLUX.1-schnell \
+ --repo-type model \
+ --local-dir ./flux_schnell \
+ --include "*"
+ ```
+
+2. Run eval:
+ ```
+ pytest flux.py::test_accuracy_model -v -s --start-from=0 --num-prompts=5
+ ```
+
+**Notes**:
+- Default settings: 1024x1024, guidance_scale=1.0, 4 inference steps (schnell model)
+- Results are saved to `test_reports/flux_test_results.json`
+
+## Stable Diffusion 3.5 Evaluation
+
+**File**: `sd35.py`
+
+**Steps**:
+
+1. Accept the license agreement for SD3.5 at: https://huggingface.co/stabilityai/stable-diffusion-3.5-large
+
+2. Load weights locally:
+ - `pip install accelerate`
+ - ```
+ huggingface-cli download stabilityai/stable-diffusion-3.5-large \
+ --repo-type model \
+ --local-dir ./sd35_large \
+ --include "*"
+ ```
+
+3. Run eval:
+ ```
+ pytest sd35.py::test_accuracy_sd35 -v -s --start-from=0 --num-prompts=5
+ ```
+
+**Notes**:
+- SD3.5 uses three text encoders: CLIP-L, CLIP-G, and T5-XXL
+- The T5 encoder provides enhanced text understanding and longer prompt support (up to 256 tokens)
+- Default settings: 1024x1024, guidance_scale=3.5, 40 inference steps
+- Requires ~40GB VRAM on H100 for full resolution inference
+- Results are saved to `test_reports/sd35_test_results.json`
+
+## Motif Evaluation
+
+**File**: `motif.py`
+
+**Note**: There isn't an existing diffusers pipeline for Motif yet, so this uses the official Motif inference script.
+
+**Steps**:
+
+1. Clone the Motif repository:
+ ```
+ git lfs install
+ GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/Motif-Technologies/Motif-Image-6B-Preview
+ cd Motif-Image-6B-Preview
+ huggingface-cli download Motif-Technologies/Motif-Image-6B-Preview --local-dir ./ --include "checkpoints/*"
+ ```
+
+2. Download captions for eval:
+ - `wget https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/coco2014/captions/captions_source.tsv`
+ - `mv captions_source.tsv captions.tsv`
+
+3. Accept the license agreement for Motif if required
+
+4. Run eval:
+ ```
+ python motif.py \
+ --images-dir outputs/seed_7777/scale_4.0 \
+ --prompts-file captions.tsv \
+ --start-from 0 \
+ --num-prompts 500
+ ```
+
+**Notes**:
+- This script uses the official Motif inference.py script to generate images
+- Results are saved to `dataset_eval_results.json`
diff --git a/code/models/tt_dit/tests/dataset_eval/reference/__init__.py b/code/models/tt_dit/tests/dataset_eval/reference/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..41cc875f93477a9c1e59a76f0c7d5bf101f6621c
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/reference/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/code/models/tt_dit/tests/dataset_eval/reference/flux.py b/code/models/tt_dit/tests/dataset_eval/reference/flux.py
new file mode 100644
index 0000000000000000000000000000000000000000..aac55c621b738c80e762ce566d325561932c7624
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/reference/flux.py
@@ -0,0 +1,295 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import csv
+import json
+
+# must be the very first lines!
+import os
+import statistics
+import time
+
+# os.environ["TRANSFORMERS_NO_TF"] = "1" # block tensorflow backend
+# os.environ["TRANSFORMERS_NO_FLAX"] = "1" # block flax/jax backend
+# os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" # silence tf if it somehow loads
+import urllib
+
+import pytest
+import torch
+from diffusers import FluxPipeline
+from loguru import logger
+
+from ..clip_encoder import CLIPEncoder
+from ..fid_score import calculate_fid_score
+
+
+class SimpleProfiler:
+ def __init__(self):
+ self.times = {}
+ self.start_times = {}
+
+ def start(self, name):
+ self.start_times[name] = time.time()
+ if name not in self.times:
+ self.times[name] = []
+
+ def end(self, name):
+ if name in self.start_times:
+ elapsed = time.time() - self.start_times[name]
+ self.times[name].append(elapsed)
+ del self.start_times[name]
+
+ def get(self, name):
+ if name in self.times and self.times[name]:
+ return sum(self.times[name]) / len(self.times[name])
+ return 0.0
+
+
+profiler = SimpleProfiler()
+
+# currently same as sdxl
+COCO_CAPTIONS_DOWNLOAD_PATH = "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/coco2014/captions/captions_source.tsv"
+OUT_ROOT, RESULTS_FILE_NAME = "test_reports", "flux_test_results.json"
+
+
+@pytest.mark.parametrize(
+ "device_params",
+ [{"device": "cuda" if torch.cuda.is_available() else "cpu"}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "model_name, image_w, image_h, guidance_scale, num_inference_steps",
+ [
+ # ("dev", 1024, 1024, 3.5, 28), # Full resolution on H100
+ ("schnell", 1024, 1024, 1.0, 4)
+ ],
+)
+@pytest.mark.parametrize("captions_path", ["captions.tsv"])
+@pytest.mark.parametrize("coco_statistics_path", ["val2014.npz"])
+def test_accuracy_model(
+ device_params,
+ model_name,
+ image_w,
+ image_h,
+ guidance_scale,
+ num_inference_steps,
+ captions_path,
+ coco_statistics_path,
+ evaluation_range,
+):
+ start_from, num_prompts = evaluation_range
+ prompts = flux_get_prompts(captions_path, start_from, num_prompts)
+ logger.info(f"start inference from prompt index: {start_from} to {start_from + num_prompts}")
+
+ device = device_params.get("device", "cuda" if torch.cuda.is_available() else "cpu")
+ logger.info(f"Using device: {device}")
+
+ if device == "cuda":
+ logger.info(f"GPU: {torch.cuda.get_device_name(0)}")
+ logger.info(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
+
+ # Initialize Diffusers pipeline
+ model_id = "./flux_schnell"
+ logger.info(f"Loading model: {model_id}")
+
+ pipeline = FluxPipeline.from_pretrained(
+ model_id,
+ dtype=torch.bfloat16 if device == "cuda" else torch.float32, # Use bfloat16 for H100
+ use_safetensors=True,
+ )
+
+ # Move to device manually
+ pipeline = pipeline.to(device)
+ logger.info(f"Pipeline loaded on {device}")
+
+ if device == "cuda":
+ # Enable optimizations for H100
+ try:
+ pipeline.enable_attention_slicing()
+ logger.info("✓ Attention slicing enabled")
+ except Exception as e:
+ logger.warning(f"Could not enable attention slicing: {e}")
+
+ # Optional: Enable VAE slicing for memory efficiency
+ try:
+ if hasattr(pipeline, "enable_vae_slicing"):
+ pipeline.enable_vae_slicing()
+ logger.info("✓ VAE slicing enabled")
+ except Exception as e:
+ logger.warning(f"Could not enable VAE slicing: {e}")
+
+ # Enable memory efficient attention for Flux
+ try:
+ pipeline.enable_memory_efficient_attention()
+ logger.info("✓ Memory efficient attention enabled")
+ except Exception as e:
+ logger.warning(f"Could not enable memory efficient attention: {e}")
+
+ # Set generator for reproducibility
+ generator = torch.Generator(device=device).manual_seed(0)
+
+ images = []
+ total_times = []
+
+ logger.info(f"Starting generation of {len(prompts)} images at {image_w}x{image_h}...")
+
+ for i, prompt in enumerate(prompts):
+ logger.info(f"Generating image {i+1}/{len(prompts)}: {prompt[:50]}...")
+
+ start_total = time.time()
+ profiler.start("denoising_loop")
+
+ # Generate image
+ with torch.no_grad():
+ generated_images = pipeline(
+ prompt=prompt,
+ height=image_h,
+ width=image_w,
+ num_inference_steps=num_inference_steps,
+ max_sequence_length=256,
+ generator=generator,
+ output_type="pil",
+ ).images
+
+ profiler.end("denoising_loop")
+ end_total = time.time()
+
+ total_time = end_total - start_total
+ total_times.append(total_time)
+
+ images.append(generated_images[0])
+ logger.info(f"Image {i+1} completed in {total_time:.2f}s")
+
+ # Optional: Save image for inspection
+ os.makedirs("generated_images", exist_ok=True)
+ generated_images[0].save(f"generated_images/image_{i+1}.png")
+
+ # Calculate metrics
+ logger.info("Calculating CLIP scores...")
+ clip = CLIPEncoder() # Remove device parameter, it auto-detects
+ clip_scores = [100 * clip.get_clip_score(prompts[i], img).item() for i, img in enumerate(images)]
+ average_clip_score = sum(clip_scores) / len(clip_scores)
+
+ deviation_clip_score = "N/A"
+ fid_score = "N/A"
+
+ # Calculate CLIP standard deviation if we have multiple prompts
+ if num_prompts >= 2:
+ deviation_clip_score = statistics.stdev(clip_scores)
+ logger.info(f"CLIP scores individual: {[f'{score:.2f}' for score in clip_scores]}")
+
+ # Calculate FID score if we have the COCO statistics file
+ if num_prompts >= 2 and os.path.isfile(coco_statistics_path):
+ logger.info("Calculating FID score...")
+ fid_score = calculate_fid_score(images, coco_statistics_path)
+ elif num_prompts >= 2 and not os.path.isfile(coco_statistics_path):
+ logger.warning(f"FID calculation skipped: COCO stats file not found at {coco_statistics_path}")
+ logger.info("To enable FID calculation, download COCO validation statistics:")
+ logger.info("wget http://bioinf.jku.at/research/ttur/ttur_stats/fid_stats_imagenet_256_hdf5.npz -O val2014.npz")
+ elif num_prompts < 2:
+ logger.info("FID calculation requires at least 2 images")
+
+ # Performance metrics
+ average_inference_time = sum(total_times) / len(total_times) if total_times else 0
+ min_inference_time = min(total_times) if total_times else 0
+ max_inference_time = max(total_times) if total_times else 0
+
+ # Results
+ logger.info("=== RESULTS ===")
+ logger.info(f"Average CLIP Score: {average_clip_score:.2f}")
+ logger.info(f"CLIP Score Std Dev: {deviation_clip_score}")
+ logger.info(f"Average Inference Time: {average_inference_time:.2f}s")
+ logger.info(f"Min/Max Inference Time: {min_inference_time:.2f}s / {max_inference_time:.2f}s")
+ logger.info(f"FID Score: {fid_score}")
+
+ print(f"FID score: {fid_score}")
+ print(f"Average CLIP Score: {average_clip_score}")
+ print(f"Standard Deviation of CLIP Scores: {deviation_clip_score}")
+
+ # Ensure values are JSON serializable
+ clip_std_serializable = float(deviation_clip_score) if deviation_clip_score != "N/A" else "N/A"
+ fid_serializable = float(fid_score) if fid_score != "N/A" else "N/A"
+
+ data = {
+ "model": "flux",
+ "metadata": {
+ "device": device.upper(),
+ "model_name": model_name,
+ "actual_model": model_id,
+ "start_from": start_from,
+ "num_prompts": num_prompts,
+ "image_width": image_w,
+ "image_height": image_h,
+ "guidance_scale": guidance_scale,
+ "num_inference_steps": num_inference_steps,
+ "backend": "diffusers",
+ "dtype": str(torch.bfloat16 if device == "cuda" else torch.float32),
+ "optimizations": ["attention_slicing", "vae_slicing", "memory_efficient_attention"],
+ },
+ "benchmarks_summary": [
+ {
+ "device": device.upper(),
+ "model": "flux",
+ "average_denoising_time": profiler.get("denoising_loop"),
+ "average_vae_time": 0.0, # VAE time is included in denoising for diffusers
+ "average_inference_time": average_inference_time,
+ "min_inference_time": min_inference_time,
+ "max_inference_time": max_inference_time,
+ "average_clip": average_clip_score,
+ "deviation_clip": clip_std_serializable,
+ "fid_score": fid_serializable,
+ "individual_clip_scores": [float(score) for score in clip_scores], # Add individual scores
+ }
+ ],
+ }
+
+ os.makedirs(OUT_ROOT, exist_ok=True)
+ with open(f"{OUT_ROOT}/{RESULTS_FILE_NAME}", "w") as f:
+ json.dump(data, f, indent=4)
+ logger.info(f"test results saved to {OUT_ROOT}/{RESULTS_FILE_NAME}")
+
+ # Cleanup
+ if device == "cuda":
+ torch.cuda.empty_cache()
+
+ print(f"\n🎉 Test completed successfully!")
+ print(f"📊 Average CLIP Score: {average_clip_score:.2f}")
+ print(f"⚡ Average Inference Time: {average_inference_time:.2f}s")
+ print(f"🖼️ Images saved to: generated_images/")
+
+
+def flux_get_prompts(captions_path, start_from, num_prompts):
+ assert (
+ 0 <= start_from < 5000 and start_from + num_prompts <= 5000
+ ), "start_from must be between 0 and 4999, and start_from + num_prompts must not exceed 5000."
+
+ prompts = []
+ if not os.path.isfile(captions_path):
+ logger.info(f"file {captions_path} not found. downloading...")
+ os.makedirs(os.path.dirname(captions_path), exist_ok=True)
+ urllib.request.urlretrieve(COCO_CAPTIONS_DOWNLOAD_PATH, captions_path)
+ logger.info("download complete.")
+
+ with open(captions_path, "r") as tsv_file:
+ reader = csv.reader(tsv_file, delimiter="\t")
+ next(reader) # skip header
+ for index, row in enumerate(reader):
+ if index < start_from:
+ continue
+ if index >= start_from + num_prompts:
+ break
+ prompts.append(row[2])
+ return prompts
+
+
+# Fixture for device parameters
+@pytest.fixture
+def device_params(request):
+ return request.param
+
+
+# Fixture for evaluation range
+# @pytest.fixture
+# def evaluation_range():
+# # Default evaluation range - can be overridden
+# return (0, 5) # start_from=0, num_prompts=5 for better statistics
diff --git a/code/models/tt_dit/tests/dataset_eval/reference/motif.py b/code/models/tt_dit/tests/dataset_eval/reference/motif.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab16176d4074b53d8f90358148c8a0c9f37d6f66
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/reference/motif.py
@@ -0,0 +1,192 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import csv
+import json
+import os
+import subprocess
+import time
+import urllib
+from pathlib import Path
+
+from loguru import logger
+from PIL import Image
+
+from ..clip_encoder import CLIPEncoder
+from ..fid_score import calculate_fid_score
+
+COCO_CAPTIONS_DOWNLOAD_PATH = (
+ "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/"
+ "text_to_image/coco2014/captions/captions_source.tsv"
+)
+
+
+def load_coco_prompts(captions_path, start_from=0, num_prompts=16):
+ """loads prompts from COCO captions tsv, downloading if missing"""
+ prompts = []
+ if not os.path.isfile(captions_path):
+ logger.info(f"file {captions_path} not found. downloading...")
+ urllib.request.urlretrieve(COCO_CAPTIONS_DOWNLOAD_PATH, captions_path)
+ logger.info("download complete.")
+
+ with open(captions_path, "r") as tsv_file:
+ reader = csv.reader(tsv_file, delimiter="\t")
+ next(reader) # skip header
+ for index, row in enumerate(reader):
+ if index < start_from:
+ continue
+ if index >= start_from + num_prompts:
+ break
+ if len(row) >= 3:
+ prompts.append(row[2].strip())
+ logger.info(f"loaded {len(prompts)} coco prompts from {captions_path}")
+ return prompts
+
+
+def generate_images_with_inference(prompts_file, output_dir, num_prompts):
+ """forces motif inference to regenerate only the requested number of prompts"""
+ logger.info("clearing any previous images...")
+ os.makedirs(output_dir, exist_ok=True)
+ for f in Path(output_dir).rglob("*.png"):
+ f.unlink()
+
+ # create temporary subset prompt file
+ tmp_prompts_file = os.path.join(output_dir, "subset_prompts.txt")
+ with open(prompts_file, "r") as f_in, open(tmp_prompts_file, "w") as f_out:
+ lines = [line.strip() for line in f_in if line.strip()]
+ subset = lines[:num_prompts]
+ for l in subset:
+ f_out.write(l + "\n")
+
+ logger.info(f"generating {len(subset)} images using motif inference.py ...")
+
+ cmd = [
+ "python",
+ "inference.py",
+ "--model-config",
+ "configs/mmdit_xlarge_hq.json",
+ "--model-ckpt",
+ "checkpoints/pytorch_model_fsdp.bin",
+ "--seed",
+ "7777",
+ "--steps",
+ "20",
+ "--resolution",
+ "1024",
+ "--prompt-file",
+ tmp_prompts_file, # use subset
+ "--guidance-scales",
+ "4.0",
+ "--output-dir",
+ output_dir,
+ "--batch-size",
+ "1",
+ ]
+ subprocess.run(cmd, check=True)
+ logger.info("image generation complete.")
+
+
+def compute_clip_scores(images_dir, prompts):
+ """computes average clip similarity across generated images and corresponding prompts"""
+ clip_encoder = CLIPEncoder(clip_version="ViT-B/32", pretrained="openai")
+ scores = []
+
+ # find all pngs recursively
+ all_images = sorted(Path(images_dir).rglob("*.png"))
+ if len(all_images) < len(prompts):
+ logger.warning(f"found only {len(all_images)} images, expected {len(prompts)}")
+
+ for idx, (prompt, img_path) in enumerate(zip(prompts, all_images)):
+ try:
+ image = Image.open(img_path).convert("RGB")
+ score = clip_encoder.get_clip_score(prompt, image)
+ scores.append(score.item())
+ except Exception as e:
+ logger.error(f"error processing {img_path}: {e}")
+
+ if not scores:
+ return 0.0
+
+ avg_score = sum(scores) / len(scores)
+ logger.info(f"average clip score = {avg_score:.4f} ({avg_score * 100:.2f} CLIPScore)")
+ return avg_score
+
+
+def compute_fid(images_dir, coco_stats="val2014.npz"):
+ """computes fid score against coco statistics"""
+ if not os.path.exists(coco_stats):
+ logger.info(f"COCO stats not found at {coco_stats}, downloading...")
+ os.makedirs(os.path.dirname(coco_stats) or ".", exist_ok=True)
+ urllib.request.urlretrieve(
+ "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/"
+ "text_to_image/tools/val2014.npz",
+ coco_stats,
+ )
+ logger.info("download complete.")
+
+ all_images = [Image.open(f).convert("RGB") for f in sorted(Path(images_dir).rglob("*.png"))]
+ if not all_images:
+ logger.error("no images found for fid computation")
+ return None
+
+ logger.info(f"computing fid over {len(all_images)} images ...")
+ fid = calculate_fid_score(all_images, coco_stats)
+ logger.info(f"fid score = {fid:.4f}")
+ return fid
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Motif Dataset Evaluation")
+ parser.add_argument("--images-dir", type=str, required=True, help="directory containing generated images")
+ parser.add_argument("--prompts-file", type=str, required=True, help="path to prompts file (.tsv or .txt)")
+ parser.add_argument("--output-json", type=str, default="dataset_eval_results.json", help="output results json")
+ parser.add_argument("--start-from", type=int, default=0, help="starting index for captions.tsv")
+ parser.add_argument("--num-prompts", type=int, default=16, help="number of prompts to evaluate")
+ args = parser.parse_args()
+
+ start_time = time.time()
+
+ # load prompts dynamically depending on file type
+ if args.prompts_file.endswith(".tsv"):
+ prompts = load_coco_prompts(args.prompts_file, args.start_from, args.num_prompts)
+ else:
+ with open(args.prompts_file, "r") as f:
+ prompts = [line.strip() for line in f if line.strip()]
+
+ # always regenerate fresh images
+ generate_images_with_inference(args.prompts_file, args.images_dir, args.num_prompts)
+
+ # automatically detect where the actual images were saved
+ def find_deepest_image_dir(root):
+ """walks the output dir and returns the first subdir containing pngs"""
+ for dirpath, _, filenames in os.walk(root):
+ if any(fn.lower().endswith(".png") for fn in filenames):
+ return dirpath
+ return root
+
+ image_root = find_deepest_image_dir(args.images_dir)
+ logger.info(f"found generated images under: {image_root}")
+
+ # compute metrics at the detected path
+ clip_score = compute_clip_scores(image_root, prompts)
+ fid_score = compute_fid(image_root, coco_stats="val2014.npz")
+ total_time = time.time() - start_time
+
+ results = {
+ "num_images": len(prompts),
+ "clip_score": clip_score,
+ "fid_score": fid_score,
+ "total_time_sec": total_time,
+ "images_dir": image_root,
+ }
+
+ with open(args.output_json, "w") as f:
+ json.dump(results, f, indent=2)
+
+ logger.info(f"saved results to {args.output_json}")
+ logger.info(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/code/models/tt_dit/tests/dataset_eval/reference/requirements.txt b/code/models/tt_dit/tests/dataset_eval/reference/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..becf4d3a83deee981ec278af3cf0e626b8d2522e
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/reference/requirements.txt
@@ -0,0 +1,16 @@
+# core dependencies
+torch
+numpy
+scipy
+loguru
+pytest
+
+# model + tokenizer
+sentencepiece
+transformers==4.46.2
+# latest dev version required for SD3.5 model architecture (qk_norm, norm_added_k, etc.)
+git+https://github.com/huggingface/diffusers.git
+open_clip_torch
+
+# vision + utilities
+Pillow>=12.1.1
diff --git a/code/models/tt_dit/tests/dataset_eval/reference/sd35.py b/code/models/tt_dit/tests/dataset_eval/reference/sd35.py
new file mode 100644
index 0000000000000000000000000000000000000000..83792c9404218f3cb462ca4205b49ab43b033761
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/reference/sd35.py
@@ -0,0 +1,291 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import csv
+import json
+
+# must be the very first lines!
+import os
+import statistics
+import time
+import urllib
+
+import pytest
+import torch
+from diffusers import StableDiffusion3Pipeline
+from loguru import logger
+
+from ..clip_encoder import CLIPEncoder
+from ..fid_score import calculate_fid_score
+
+
+class SimpleProfiler:
+ def __init__(self):
+ self.times = {}
+ self.start_times = {}
+
+ def start(self, name):
+ self.start_times[name] = time.time()
+ if name not in self.times:
+ self.times[name] = []
+
+ def end(self, name):
+ if name in self.start_times:
+ elapsed = time.time() - self.start_times[name]
+ self.times[name].append(elapsed)
+ del self.start_times[name]
+
+ def get(self, name):
+ if name in self.times and self.times[name]:
+ return sum(self.times[name]) / len(self.times[name])
+ return 0.0
+
+
+profiler = SimpleProfiler()
+
+COCO_CAPTIONS_DOWNLOAD_PATH = "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/coco2014/captions/captions_source.tsv"
+OUT_ROOT, RESULTS_FILE_NAME = "test_reports", "sd35_test_results.json"
+
+
+@pytest.mark.parametrize(
+ "device_params",
+ [{"device": "cuda" if torch.cuda.is_available() else "cpu"}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "model_name, image_w, image_h, guidance_scale, num_inference_steps, max_sequence_length",
+ [
+ ("stable-diffusion-3.5-large", 1024, 1024, 3.5, 40, 256),
+ ],
+)
+@pytest.mark.parametrize("captions_path", ["captions.tsv"])
+@pytest.mark.parametrize("coco_statistics_path", ["val2014.npz"])
+def test_accuracy_sd35(
+ device_params,
+ model_name,
+ image_w,
+ image_h,
+ guidance_scale,
+ num_inference_steps,
+ max_sequence_length,
+ captions_path,
+ coco_statistics_path,
+ evaluation_range,
+):
+ start_from, num_prompts = evaluation_range
+ prompts = sd35_get_prompts(captions_path, start_from, num_prompts)
+ logger.info(f"start inference from prompt index: {start_from} to {start_from + num_prompts}")
+
+ device = device_params.get("device", "cuda" if torch.cuda.is_available() else "cpu")
+ logger.info(f"Using device: {device}")
+
+ if device == "cuda":
+ logger.info(f"GPU: {torch.cuda.get_device_name(0)}")
+ logger.info(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
+
+ model_id = "./sd35_large"
+ logger.info(f"Loading model: {model_id}")
+ logger.info(f"T5 text encoder enabled with max_sequence_length={max_sequence_length}")
+
+ pipeline = StableDiffusion3Pipeline.from_pretrained(
+ model_id,
+ torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, # Use bfloat16 for H100
+ use_safetensors=True,
+ )
+
+ pipeline = pipeline.to(device)
+ logger.info(f"Pipeline loaded on {device}")
+
+ if device == "cuda":
+ # enable optimizations for H100
+ try:
+ pipeline.enable_attention_slicing()
+ logger.info("✓ Attention slicing enabled")
+ except Exception as e:
+ logger.warning(f"Could not enable attention slicing: {e}")
+
+ # enable VAE slicing for memory efficiency
+ try:
+ if hasattr(pipeline, "enable_vae_slicing"):
+ pipeline.enable_vae_slicing()
+ logger.info("✓ VAE slicing enabled")
+ except Exception as e:
+ logger.warning(f"Could not enable VAE slicing: {e}")
+
+ # set generator for reproducibility
+ generator = torch.Generator(device=device).manual_seed(0)
+
+ images = []
+ total_times = []
+
+ logger.info(f"Starting generation of {len(prompts)} images at {image_w}x{image_h}...")
+
+ for i, prompt in enumerate(prompts):
+ logger.info(f"Generating image {i+1}/{len(prompts)}: {prompt[:50]}...")
+
+ start_total = time.time()
+ profiler.start("denoising_loop")
+
+ # generate image with triple prompt (SD3.5 uses 3 text encoders)
+ with torch.no_grad():
+ generated_images = pipeline(
+ prompt=prompt,
+ prompt_2=prompt,
+ prompt_3=prompt,
+ negative_prompt="",
+ negative_prompt_2="",
+ negative_prompt_3="",
+ height=image_h,
+ width=image_w,
+ num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale,
+ max_sequence_length=max_sequence_length,
+ generator=generator,
+ output_type="pil",
+ ).images
+
+ profiler.end("denoising_loop")
+ end_total = time.time()
+
+ total_time = end_total - start_total
+ total_times.append(total_time)
+
+ images.append(generated_images[0])
+ logger.info(f"Image {i+1} completed in {total_time:.2f}s")
+
+ # optional: save image
+ os.makedirs("generated_images", exist_ok=True)
+ generated_images[0].save(f"generated_images/image_{i+1}.png")
+
+ # calculate metrics
+ logger.info("Calculating CLIP scores...")
+ clip = CLIPEncoder()
+ clip_scores = [100 * clip.get_clip_score(prompts[i], img).item() for i, img in enumerate(images)]
+ average_clip_score = sum(clip_scores) / len(clip_scores)
+
+ deviation_clip_score = "N/A"
+ fid_score = "N/A"
+
+ # Calculate CLIP standard deviation if we have multiple prompts
+ if num_prompts >= 2:
+ deviation_clip_score = statistics.stdev(clip_scores)
+ logger.info(f"CLIP scores individual: {[f'{score:.2f}' for score in clip_scores]}")
+
+ # Calculate FID score if we have the COCO statistics file
+ if num_prompts >= 2 and os.path.isfile(coco_statistics_path):
+ logger.info("Calculating FID score...")
+ fid_score = calculate_fid_score(images, coco_statistics_path)
+ elif num_prompts >= 2 and not os.path.isfile(coco_statistics_path):
+ logger.warning(f"FID calculation skipped: COCO stats file not found at {coco_statistics_path}")
+ logger.info("To enable FID calculation, download COCO validation statistics:")
+ logger.info(
+ "wget https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/tools/val2014.npz"
+ )
+ elif num_prompts < 2:
+ logger.info("FID calculation requires at least 2 images")
+
+ average_inference_time = sum(total_times) / len(total_times) if total_times else 0
+ min_inference_time = min(total_times) if total_times else 0
+ max_inference_time = max(total_times) if total_times else 0
+
+ # results
+ logger.info("=== RESULTS ===")
+ logger.info(f"Average CLIP Score: {average_clip_score:.2f}")
+ logger.info(f"CLIP Score Std Dev: {deviation_clip_score}")
+ logger.info(f"Average Inference Time: {average_inference_time:.2f}s")
+ logger.info(f"Min/Max Inference Time: {min_inference_time:.2f}s / {max_inference_time:.2f}s")
+ logger.info(f"FID Score: {fid_score}")
+
+ print(f"FID score: {fid_score}")
+ print(f"Average CLIP Score: {average_clip_score}")
+ print(f"Standard Deviation of CLIP Scores: {deviation_clip_score}")
+
+ # check if T5 encoder is being used
+ t5_enabled = hasattr(pipeline, "text_encoder_3") and pipeline.text_encoder_3 is not None
+ logger.info(f"T5 text encoder enabled: {t5_enabled}")
+
+ # ensure values are JSON serializable
+ clip_std_serializable = float(deviation_clip_score) if deviation_clip_score != "N/A" else "N/A"
+ fid_serializable = float(fid_score) if fid_score != "N/A" else "N/A"
+
+ data = {
+ "model": model_name,
+ "metadata": {
+ "device": device.upper(),
+ "model_name": model_name,
+ "actual_model": model_id,
+ "start_from": start_from,
+ "num_prompts": num_prompts,
+ "image_width": image_w,
+ "image_height": image_h,
+ "guidance_scale": guidance_scale,
+ "num_inference_steps": num_inference_steps,
+ "max_sequence_length": max_sequence_length,
+ "t5_enabled": t5_enabled,
+ "backend": "diffusers",
+ "dtype": str(torch.bfloat16 if device == "cuda" else torch.float32),
+ "optimizations": ["attention_slicing", "vae_slicing"],
+ },
+ "benchmarks_summary": [
+ {
+ "device": device.upper(),
+ "model": model_name,
+ "average_denoising_time": profiler.get("denoising_loop"),
+ "average_vae_time": 0.0, # VAE time is included in denoising for diffusers
+ "average_inference_time": average_inference_time,
+ "min_inference_time": min_inference_time,
+ "max_inference_time": max_inference_time,
+ "average_clip": average_clip_score,
+ "deviation_clip": clip_std_serializable,
+ "fid_score": fid_serializable,
+ "individual_clip_scores": [float(score) for score in clip_scores],
+ }
+ ],
+ }
+
+ os.makedirs(OUT_ROOT, exist_ok=True)
+ with open(f"{OUT_ROOT}/{RESULTS_FILE_NAME}", "w") as f:
+ json.dump(data, f, indent=4)
+ logger.info(f"test results saved to {OUT_ROOT}/{RESULTS_FILE_NAME}")
+
+ # cleanup
+ if device == "cuda":
+ torch.cuda.empty_cache()
+
+ print(f"\n🎉 Test completed successfully!")
+ print(f"📊 Average CLIP Score: {average_clip_score:.2f}")
+ print(f"⚡ Average Inference Time: {average_inference_time:.2f}s")
+ print(f"🖼️ Images saved to: generated_images/")
+
+
+def sd35_get_prompts(captions_path, start_from, num_prompts):
+ assert (
+ 0 <= start_from < 5000 and start_from + num_prompts <= 5000
+ ), "start_from must be between 0 and 4999, and start_from + num_prompts must not exceed 5000."
+
+ prompts = []
+ if not os.path.isfile(captions_path):
+ logger.info(f"file {captions_path} not found. downloading...")
+ os.makedirs(os.path.dirname(captions_path), exist_ok=True)
+ urllib.request.urlretrieve(COCO_CAPTIONS_DOWNLOAD_PATH, captions_path)
+ logger.info("download complete.")
+
+ with open(captions_path, "r") as tsv_file:
+ reader = csv.reader(tsv_file, delimiter="\t")
+ next(reader) # skip header
+ for index, row in enumerate(reader):
+ if index < start_from:
+ continue
+ if index >= start_from + num_prompts:
+ break
+ prompts.append(row[2])
+ return prompts
+
+
+# fixture for device parameters
+@pytest.fixture
+def device_params(request):
+ return request.param
+
+
+# python -m pytest model.py::test_accuracy_sd35 -v --start-from=0 --num-prompts=5 -s
diff --git a/code/models/tt_dit/tests/dataset_eval/test_tt_dit.py b/code/models/tt_dit/tests/dataset_eval/test_tt_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..57174714ce649812cf62e24647f0e0f5c90f19bb
--- /dev/null
+++ b/code/models/tt_dit/tests/dataset_eval/test_tt_dit.py
@@ -0,0 +1,264 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import csv
+import itertools
+import json
+import os
+import random
+import statistics
+import time
+import urllib
+
+import pytest
+from loguru import logger
+
+import ttnn
+
+# from .fid_score import calculate_fid_score
+from models.demos.stable_diffusion_xl_base.utils.fid_score import calculate_fid_score
+from models.perf.benchmarking_utils import BenchmarkProfiler
+from models.tt_dit.pipelines.events import profiler_event_callback
+
+from .clip_encoder import CLIPEncoder
+
+# COCO captions download path
+COCO_CAPTIONS_DOWNLOAD_PATH = "https://github.com/mlcommons/inference/raw/4b1d1156c23965172ae56eacdd8372f8897eb771/text_to_image/coco2014/captions/captions_source.tsv"
+
+
+def get_prompts(captions_path, start_from, num_prompts):
+ assert (
+ 0 <= start_from < 5000 and start_from + num_prompts <= 5000
+ ), "start_from must be between 0 and 4999, and start_from + num_prompts must not exceed 5000."
+
+ prompts = []
+ if not os.path.isfile(captions_path):
+ logger.info(f"file {captions_path} not found. downloading...")
+ os.makedirs(os.path.dirname(captions_path), exist_ok=True)
+ urllib.request.urlretrieve(COCO_CAPTIONS_DOWNLOAD_PATH, captions_path)
+ logger.info("download complete.")
+
+ with open(captions_path, "r") as tsv_file:
+ reader = csv.reader(tsv_file, delimiter="\t")
+ next(reader) # skip header
+ for index, row in enumerate(reader):
+ if index < start_from:
+ continue
+ if index >= start_from + num_prompts:
+ break
+ prompts.append(row[2])
+ return prompts
+
+
+def get_img_evals(model_id, model_setup, images, prompts, coco_statistics_path):
+ # Calculate metrics
+ logger.info("Calculating CLIP scores...")
+ num_prompts = len(prompts)
+ clip = CLIPEncoder()
+ clip_scores = [100 * clip.get_clip_score(prompts[i], img).item() for i, img in enumerate(images)]
+ average_clip_score = sum(clip_scores) / len(clip_scores)
+
+ deviation_clip_score = "N/A"
+ fid_score = "N/A"
+
+ # Calculate CLIP standard deviation if we have multiple prompts
+ if num_prompts >= 2:
+ deviation_clip_score = statistics.stdev(clip_scores)
+ logger.info(f"CLIP scores individual: {[f'{score:.2f}' for score in clip_scores]}")
+
+ # Calculate FID score if we have the COCO statistics file
+ logger.info("Calculating FID score...")
+ try:
+ fid_score = calculate_fid_score(images, coco_statistics_path)
+ except Exception as e:
+ logger.error(f"Error calculating FID score: {e}")
+ else:
+ logger.info("FID calculation requires at least 2 images")
+
+ evals = [
+ {
+ "model": model_id,
+ "average_clip": average_clip_score,
+ "clip_standard_deviation": deviation_clip_score,
+ "num_prompts": num_prompts,
+ "accuracy_check": check_accuracy(num_prompts, average_clip_score, "clip", model_setup),
+ },
+ {
+ "model": model_id,
+ "fid_score": fid_score,
+ "num_prompts": num_prompts,
+ "accuracy_check": check_accuracy(num_prompts, fid_score, "fid", model_setup),
+ },
+ ]
+ return evals
+
+
+def get_vid_evals(model_id, model_setup, videos, prompts):
+ num_prompts = len(prompts)
+ clip = CLIPEncoder()
+ clip_scores = [
+ [100 * clip.get_clip_score(prompts[i], img).item() for img in video] for i, video in enumerate(videos)
+ ]
+ clip_score_stats = [(min(c), max(c), statistics.mean(c), statistics.stdev(c)) for c in clip_scores]
+
+ for p, c in zip(prompts, clip_score_stats):
+ logger.info(f"Prompt: {p}, CLIP statistics (min, max, mean, stddev): {c}")
+
+ video_mean_scores = [c[2] for c in clip_score_stats]
+ min_clip = min(video_mean_scores)
+ max_clip = max(video_mean_scores)
+ mean_clip = statistics.mean(video_mean_scores)
+ stddev_clip = statistics.stdev(video_mean_scores)
+
+ evals = [
+ {
+ "model": model_id,
+ "average_clip": mean_clip,
+ "min_clip": min_clip,
+ "max_clip": max_clip,
+ "clip_standard_deviation": stddev_clip,
+ "num_prompts": num_prompts,
+ "accuracy_check": check_accuracy(num_prompts, mean_clip, "clip", model_setup),
+ }
+ ]
+ return evals
+
+
+def check_accuracy(num_prompts, score, score_key, model_info):
+ accuracy_status = 1
+ prompt_data = model_info["accuracy"].get(str(num_prompts), {})
+ range_key = f"{score_key}_valid_range"
+ if range_key in prompt_data:
+ score_range = prompt_data[range_key]
+ if score >= score_range[0] and score <= score_range[1]:
+ return 2
+ else:
+ return 3
+
+ return accuracy_status
+
+
+@pytest.mark.parametrize(
+ "mesh_device,device_name",
+ [
+ [(1, 2), "p300"],
+ [(2, 2), "qbge"],
+ [(2, 4), "t3K"],
+ [(4, 8), "galaxy"],
+ ],
+ ids=[
+ "1x2",
+ "2x2",
+ "2x4",
+ "4x8",
+ ],
+ indirect=["mesh_device"],
+)
+def test_tt_dit_accuracy(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ device_name: str,
+ num_inference_steps,
+ model_id,
+ model_location_generator,
+ dit_pipeline,
+ model_metadata,
+ evaluation_range,
+ is_video,
+ is_24hrs_test,
+ monkeypatch,
+) -> None:
+ """Accuracy test for TT-Metal DiT pipelines with CLIP and FID score evaluation.
+ It relies on the default configurations set in each pipeline.
+ """
+
+ if os.getenv("TT_DIT_CACHE_DIR") is None:
+ monkeypatch.setenv("TT_DIT_CACHE_DIR", "/tmp/TT_DIT_CACHE")
+ benchmark_profiler = BenchmarkProfiler()
+ if ttnn.device.is_blackhole() and device_name == "galaxy":
+ device_name = "bh_" + device_name
+ captions_path = "coco_captions/captions.tsv"
+ coco_statistics_path = "coco_statistics/val2014.npz"
+ start_from, num_prompts = evaluation_range
+ prompts = get_prompts(captions_path, start_from, num_prompts)
+ logger.info(f"start inference from prompt index: {start_from} to {start_from + num_prompts}")
+
+ logger.info(f" Mesh shape: {mesh_device.shape}")
+ logger.info(f" Model ID: {model_id}")
+ logger.info(f" Inference steps: {num_inference_steps}")
+
+ # Create pipeline similar to test_performance_flux1.py
+ pipeline = dit_pipeline.create_pipeline(
+ mesh_device=mesh_device, checkpoint_name=model_location_generator(model_metadata["hf_id"])
+ )
+
+ images = []
+ gen_type = "image" if not is_video else "video"
+ logger.info(f"Starting generation of {len(prompts)} {gen_type}s...")
+
+ # Generate images for each prompt
+ iter = itertools.cycle(prompts) if is_24hrs_test else prompts
+ start_time = time.time()
+ for i, prompt in enumerate(iter):
+ logger.info(
+ f"Generating {gen_type} {i+1} of {len(prompts) if not is_24hrs_test else '24hrs {gen_type}s'}: {prompt[:50]}..."
+ )
+
+ # Generate image using TT-Metal pipeline
+ with benchmark_profiler("inference", iteration=i):
+ generated_images = pipeline.run_single_prompt(
+ prompt=prompt,
+ num_inference_steps=num_inference_steps,
+ on_event=profiler_event_callback(benchmark_profiler, i),
+ **model_metadata.get("extra_args", {}),
+ )
+
+ images.append(generated_images[0])
+ logger.info(f"{gen_type} {i+1} completed in {benchmark_profiler.get_duration('inference', i):.2f}s")
+
+ if is_24hrs_test:
+ if (time.time() - start_time) > 24 * 60 * 60:
+ break
+ # Sleep for a random time between 0.1 and 10 seconds
+ time.sleep(random.uniform(0.1, 10))
+
+ # Performance metrics
+ average_inference_time = benchmark_profiler.get_duration_average("inference", 1) # Skip warmup iterations
+
+ data = {
+ "model": model_id,
+ "metadata": {
+ "device": f"{device_name}",
+ "model_name": model_id,
+ "start_from": start_from,
+ "num_prompts": num_prompts,
+ "num_inference_steps": num_inference_steps,
+ "mesh_shape": list(mesh_device.shape),
+ },
+ "benchmarks_summary": [
+ {
+ "device": f"{device_name}",
+ "model": model_id,
+ "average_inference_time": average_inference_time,
+ "average_vae_time": benchmark_profiler.get_duration_average("vae", 1),
+ "average_encoding_time": benchmark_profiler.get_duration_average("encoder", 1),
+ "average_denoising_time": benchmark_profiler.get_duration_average("denoising", 1),
+ }
+ ],
+ "evals": get_vid_evals(model_id, model_metadata, images, prompts)
+ if is_video
+ else get_img_evals(model_id, model_metadata, images, prompts, coco_statistics_path),
+ }
+
+ out_root = "test_reports"
+ os.makedirs(out_root, exist_ok=True)
+ file_path = f"{out_root}/sdxl_test_results.json"
+ with open(file_path, "w") as f:
+ json.dump(data, f, indent=4)
+ logger.info(f"test results saved to {file_path}")
+
+ # Synchronize devices
+ pipeline.synchronize_devices()
+
+ print(f"\n🎉 Test completed successfully!")
+ print(f"⚡ Average Inference Time: {average_inference_time:.2f}s")
diff --git a/code/models/tt_dit/tests/encoders/clip/test_clip_embeddings.py b/code/models/tt_dit/tests/encoders/clip/test_clip_embeddings.py
new file mode 100644
index 0000000000000000000000000000000000000000..03199d589e39bf566fc736c37e0e327adc9bd735
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/clip/test_clip_embeddings.py
@@ -0,0 +1,138 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer
+
+import ttnn
+from models.tt_dit.encoders.clip.model_clip import CLIPConfig, TextEmbeddings
+from models.tt_dit.utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "clip_path, tokenizer_path, expected_pcc",
+ [
+ ("text_encoder", "tokenizer", 0.99),
+ ("text_encoder_2", "tokenizer_2", 0.985),
+ ],
+ ids=["encoder_1", "encoder_2"],
+)
+@pytest.mark.parametrize("mesh_device", [(2, 4)], ids=["t3k"], indirect=True)
+@pytest.mark.parametrize("submesh_shape", [(1, 4), (2, 2)], ids=["1x4", "2x2"])
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_clip_text_embeddings(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ clip_path: str,
+ tokenizer_path: str,
+ expected_pcc: float,
+ topology: ttnn.Topology,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ hf_model = CLIPTextModelWithProjection.from_pretrained(
+ model_name_checkpoint, subfolder=clip_path, local_files_only=True
+ )
+ tokenizer = CLIPTokenizer.from_pretrained(model_name_checkpoint, subfolder=tokenizer_path, local_files_only=True)
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace CLIP Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.hidden_size}")
+ logger.info(f"intermediate_size: {hf_model.config.intermediate_size}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_attention_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_hidden_layers}")
+
+ test_text = "A coffee shop on Main Street that serves excellent pastries and opens at 7 AM on weekdays"
+
+ hf_inputs = tokenizer(test_text, padding=True, truncation=True, max_length=77, return_tensors="pt")
+ tt_prompt = ttnn.from_torch(
+ hf_inputs.input_ids,
+ dtype=ttnn.uint32,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ # === TT-DiT TextEmbeddings ====
+ config = CLIPConfig(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.hidden_size,
+ ff_dim=hf_model.config.intermediate_size,
+ num_heads=hf_model.config.num_attention_heads,
+ num_hidden_layers=hf_model.config.num_hidden_layers,
+ max_prompt_length=77,
+ layer_norm_eps=hf_model.config.layer_norm_eps,
+ attention_dropout=hf_model.config.attention_dropout,
+ hidden_act="quick_gelu",
+ )
+
+ tt_embedding = TextEmbeddings(config, encoder_submesh)
+ embeddings_state_dict = {}
+ for key, value in hf_model.state_dict().items():
+ if key.startswith("text_model.embeddings."):
+ new_key = key.replace("text_model.embeddings.", "")
+ embeddings_state_dict[new_key] = value
+ tt_embedding.load_state_dict(embeddings_state_dict)
+
+ # time TT model inference only
+ tt_start_time = time.time()
+ tt_embeddings_output = tt_embedding(tt_prompt)
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ # get HF outputs
+ with torch.no_grad():
+ hf_start_time = time.time()
+
+ hf_token_embeddings = hf_model.text_model.embeddings.token_embedding(hf_inputs.input_ids)
+ seq_len = hf_inputs.input_ids.shape[-1]
+ position_ids = torch.arange(seq_len).expand((1, -1))
+ hf_position_embeddings = hf_model.text_model.embeddings.position_embedding(position_ids)
+ hf_embeddings_output = hf_token_embeddings + hf_position_embeddings
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_embeddings_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_embeddings_output)[0])
+
+ logger.info(f"TT embeddings execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF embeddings execution time: {hf_execution_time:.4f} seconds")
+
+ assert hf_embeddings_output.shape == tt_embeddings_output_torch.shape
+
+ assert_quality(hf_embeddings_output, tt_embeddings_output_torch, pcc=expected_pcc)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/clip/test_clip_encoder_all_layers.py b/code/models/tt_dit/tests/encoders/clip/test_clip_encoder_all_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..dbc11637486791d51965c1eb77417d056f9f7a99
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/clip/test_clip_encoder_all_layers.py
@@ -0,0 +1,196 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer
+
+import ttnn
+from models.tt_dit.encoders.clip.model_clip import (
+ CLIPConfig,
+ CLIPStack,
+ TextEmbeddings,
+ create_4d_causal_attention_mask,
+)
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "clip_path, tokenizer_path, expected_pcc",
+ [
+ ("text_encoder", "tokenizer", 0.99),
+ ("text_encoder_2", "tokenizer_2", 0.985),
+ ],
+ ids=["encoder_1", "encoder_2"],
+)
+@pytest.mark.parametrize("mesh_device", [(2, 4)], ids=["t3k"], indirect=True)
+@pytest.mark.parametrize("submesh_shape", [(1, 4), (2, 2)], ids=["1x4", "2x2"])
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_clip_stack_all_layers(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ clip_path: str,
+ tokenizer_path: str,
+ expected_pcc: float,
+ topology: ttnn.Topology,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=encoder_submesh.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ hf_model = CLIPTextModelWithProjection.from_pretrained(
+ model_name_checkpoint, subfolder=clip_path, local_files_only=True
+ )
+ tokenizer = CLIPTokenizer.from_pretrained(model_name_checkpoint, subfolder=tokenizer_path, local_files_only=True)
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace CLIP Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.hidden_size}")
+ logger.info(f"intermediate_size: {hf_model.config.intermediate_size}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_attention_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_hidden_layers}")
+
+ test_text = "A coffee shop on Main Street that serves excellent pastries and opens at 7 AM on weekdays"
+
+ hf_inputs = tokenizer(test_text, padding=True, truncation=True, max_length=77, return_tensors="pt")
+ tt_prompt = ttnn.from_torch(
+ hf_inputs.input_ids,
+ dtype=ttnn.uint32,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ # === TT-DiT CLIPStack (all encoder layers) ====
+ config = CLIPConfig(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.hidden_size,
+ ff_dim=hf_model.config.intermediate_size,
+ num_heads=hf_model.config.num_attention_heads,
+ num_hidden_layers=hf_model.config.num_hidden_layers,
+ max_prompt_length=77,
+ layer_norm_eps=hf_model.config.layer_norm_eps,
+ attention_dropout=hf_model.config.attention_dropout,
+ hidden_act=hf_model.config.hidden_act,
+ )
+
+ logger.info(f"Testing CLIPStack with {config.num_hidden_layers} encoder layers")
+
+ tt_embedding = TextEmbeddings(config, encoder_submesh)
+ embeddings_state_dict = {}
+ for key, value in hf_model.state_dict().items():
+ if key.startswith("text_model.embeddings."):
+ new_key = key.replace("text_model.embeddings.", "")
+ embeddings_state_dict[new_key] = value
+ tt_embedding.load_torch_state_dict(embeddings_state_dict)
+
+ hidden_states = tt_embedding(tt_prompt)
+
+ causal_attention_mask = create_4d_causal_attention_mask(
+ tt_prompt.shape, encoder_submesh, dtype=hidden_states.get_dtype()
+ )
+
+ tt_clip_stack = CLIPStack(config, encoder_submesh, ccl_manager, parallel_config)
+
+ encoder_state_dict = {}
+ for key, value in hf_model.state_dict().items():
+ if key.startswith("text_model.encoder."):
+ new_key = key.replace("text_model.encoder.", "")
+ encoder_state_dict[new_key] = value
+ tt_clip_stack.load_torch_state_dict(encoder_state_dict)
+
+ tt_start_time = time.time()
+ tt_all_hidden_states = tt_clip_stack(hidden_states, causal_attention_mask, ccl_manager, parallel_config)
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ expected_outputs = config.num_hidden_layers + 1 # input + all encoder layers
+
+ # get HF reference outputs
+ with torch.no_grad():
+ hf_start_time = time.time()
+
+ # get HF embeddings
+ hf_token_embeddings = hf_model.text_model.embeddings.token_embedding(hf_inputs.input_ids)
+ seq_len = hf_inputs.input_ids.shape[-1]
+ position_ids = torch.arange(seq_len).expand((1, -1))
+ hf_position_embeddings = hf_model.text_model.embeddings.position_embedding(position_ids)
+ hf_hidden_states = hf_token_embeddings + hf_position_embeddings
+
+ # create causal attention mask for HF
+ batch_size, tgt_len = hf_inputs.input_ids.shape
+ hf_causal_mask = torch.full((tgt_len, tgt_len), float("-inf"))
+ mask_cond = torch.arange(hf_causal_mask.size(-1))
+ hf_causal_mask.masked_fill_(mask_cond < (mask_cond + 1).view(hf_causal_mask.size(-1), 1), 0)
+ hf_causal_mask = hf_causal_mask[None, None, :, :].expand(batch_size, 1, tgt_len, tgt_len)
+
+ # run through HF encoder layers manually to match CLIPStack behavior
+ hf_all_hidden_states = [hf_hidden_states] # input embeddings
+ current_hidden_states = hf_hidden_states
+
+ for layer_idx in range(hf_model.config.num_hidden_layers):
+ # transformers 5.x merged causal_attention_mask into attention_mask and returns
+ # the hidden states directly instead of a tuple.
+ layer_output = hf_model.text_model.encoder.layers[layer_idx](
+ current_hidden_states, hf_causal_mask # hidden_states, attention_mask
+ )
+ hf_all_hidden_states.append(layer_output)
+ current_hidden_states = layer_output
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ # convert mesh tensor to torch tensor for pcc (final layer output)
+ # since weights are replicated, can get the tensor from any single device
+ tt_final_layer_output = ttnn.to_torch(ttnn.get_device_tensors(tt_all_hidden_states[-1])[0])
+ hf_final_layer_output = hf_all_hidden_states[-1]
+
+ logger.info(f"TT CLIPStack (all {config.num_hidden_layers} layers) execution time: {tt_execution_time:.4f} seconds")
+ logger.info(
+ f"HF encoder (all {hf_model.config.num_hidden_layers} layers) execution time: {hf_execution_time:.4f} seconds"
+ )
+
+ assert hf_final_layer_output.shape == tt_final_layer_output.shape
+
+ assert_quality(hf_final_layer_output, tt_final_layer_output, pcc=expected_pcc)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/clip/test_clip_encoder_single_layer.py b/code/models/tt_dit/tests/encoders/clip/test_clip_encoder_single_layer.py
new file mode 100644
index 0000000000000000000000000000000000000000..d256ef499fc13fbd850cc2099a8f67a773303928
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/clip/test_clip_encoder_single_layer.py
@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer
+
+import ttnn
+from models.tt_dit.encoders.clip.model_clip import (
+ CLIPConfig,
+ CLIPEncoderLayer,
+ TextEmbeddings,
+ create_4d_causal_attention_mask,
+)
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "clip_path, tokenizer_path, expected_pcc",
+ [
+ ("text_encoder", "tokenizer", 0.99),
+ ("text_encoder_2", "tokenizer_2", 0.985),
+ ],
+ ids=["encoder_1", "encoder_2"],
+)
+@pytest.mark.parametrize("mesh_device", [(1, 1), (2, 4)], ids=["single_device", "t3k"], indirect=True)
+@pytest.mark.parametrize("submesh_shape", [(1, 1), (1, 4), (2, 2)], ids=["1x1", "1x4", "2x2"])
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_clip_encoder_layer(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ clip_path: str,
+ tokenizer_path: str,
+ expected_pcc: float,
+ topology: ttnn.Topology,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=encoder_submesh.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ hf_model = CLIPTextModelWithProjection.from_pretrained(
+ model_name_checkpoint, subfolder=clip_path, local_files_only=True
+ )
+ tokenizer = CLIPTokenizer.from_pretrained(model_name_checkpoint, subfolder=tokenizer_path, local_files_only=True)
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace CLIP Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.hidden_size}")
+ logger.info(f"intermediate_size: {hf_model.config.intermediate_size}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_attention_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_hidden_layers}")
+
+ test_text = "A coffee shop on Main Street that serves excellent pastries and opens at 7 AM on weekdays"
+
+ hf_inputs = tokenizer(test_text, padding=True, truncation=True, max_length=77, return_tensors="pt")
+ tt_prompt = ttnn.from_torch(
+ hf_inputs.input_ids,
+ dtype=ttnn.uint32,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ # === TT-DiT CLIPEncoderLayer ====
+ config = CLIPConfig(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.hidden_size,
+ ff_dim=hf_model.config.intermediate_size,
+ num_heads=hf_model.config.num_attention_heads,
+ num_hidden_layers=hf_model.config.num_hidden_layers,
+ max_prompt_length=77,
+ layer_norm_eps=hf_model.config.layer_norm_eps,
+ attention_dropout=hf_model.config.attention_dropout,
+ hidden_act="quick_gelu",
+ )
+
+ tt_embedding = TextEmbeddings(config, encoder_submesh)
+ embeddings_state_dict = {}
+ for key, value in hf_model.state_dict().items():
+ if key.startswith("text_model.embeddings."):
+ new_key = key.replace("text_model.embeddings.", "")
+ embeddings_state_dict[new_key] = value
+ tt_embedding.load_state_dict(embeddings_state_dict)
+
+ hidden_states = tt_embedding(tt_prompt)
+
+ causal_attention_mask = create_4d_causal_attention_mask(
+ tt_prompt.shape, encoder_submesh, dtype=hidden_states.get_dtype()
+ )
+
+ tt_encoder_layer = CLIPEncoderLayer(config, encoder_submesh, ccl_manager, parallel_config)
+
+ first_layer_state_dict = {}
+ for key, value in hf_model.state_dict().items():
+ if key.startswith("text_model.encoder.layers.0."):
+ new_key = key.replace("text_model.encoder.layers.0.", "")
+ first_layer_state_dict[new_key] = value
+ tt_encoder_layer.load_state_dict(first_layer_state_dict)
+
+ # time TT model inference only
+ tt_start_time = time.time()
+ tt_layer_output = tt_encoder_layer(hidden_states, causal_attention_mask, ccl_manager, parallel_config)
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ # get HF encoder layer output
+ with torch.no_grad():
+ hf_start_time = time.time()
+
+ # get HF embeddings
+ hf_token_embeddings = hf_model.text_model.embeddings.token_embedding(hf_inputs.input_ids)
+ seq_len = hf_inputs.input_ids.shape[-1]
+ position_ids = torch.arange(seq_len).expand((1, -1))
+ hf_position_embeddings = hf_model.text_model.embeddings.position_embedding(position_ids)
+ hf_hidden_states = hf_token_embeddings + hf_position_embeddings
+
+ # create causal attention mask for HF
+ batch_size, tgt_len = hf_inputs.input_ids.shape
+ hf_causal_mask = torch.full((tgt_len, tgt_len), float("-inf"))
+ mask_cond = torch.arange(hf_causal_mask.size(-1))
+ hf_causal_mask.masked_fill_(mask_cond < (mask_cond + 1).view(hf_causal_mask.size(-1), 1), 0)
+ hf_causal_mask = hf_causal_mask[None, None, :, :].expand(batch_size, 1, tgt_len, tgt_len)
+
+ # run the first encoder layer. transformers 5.x merged causal_attention_mask into
+ # attention_mask and returns the hidden states directly instead of a tuple.
+ hf_layer_output = hf_model.text_model.encoder.layers[0](
+ hf_hidden_states, # hidden_states
+ hf_causal_mask, # attention_mask
+ )
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_layer_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_layer_output)[0])
+
+ logger.info(f"TT encoder layer execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF encoder layer execution time: {hf_execution_time:.4f} seconds")
+
+ assert hf_layer_output.shape == tt_layer_output_torch.shape
+
+ assert_quality(hf_layer_output, tt_layer_output_torch, pcc=expected_pcc)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/clip/test_clip_full_no_projection.py b/code/models/tt_dit/tests/encoders/clip/test_clip_full_no_projection.py
new file mode 100644
index 0000000000000000000000000000000000000000..e6a2da3fd55a1f6ecd531b659147f47a3bf8db50
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/clip/test_clip_full_no_projection.py
@@ -0,0 +1,148 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers import CLIPTextModel, CLIPTokenizer
+
+import ttnn
+from models.tt_dit.encoders.clip.model_clip import CLIPConfig, CLIPEncoder
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "clip_path, tokenizer_path, expected_pcc",
+ [
+ ("text_encoder", "tokenizer", 0.99),
+ ("text_encoder_2", "tokenizer_2", 0.984),
+ ],
+ ids=["encoder_1", "encoder_2"],
+)
+@pytest.mark.parametrize("mesh_device", [(2, 4)], ids=["t3k"], indirect=True)
+@pytest.mark.parametrize("submesh_shape", [(1, 4), (2, 2)], ids=["1x4", "2x2"])
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_clip_encoder(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ clip_path: str,
+ tokenizer_path: str,
+ expected_pcc: float,
+ topology: ttnn.Topology,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=encoder_submesh.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ hf_model = CLIPTextModel.from_pretrained(model_name_checkpoint, subfolder=clip_path, local_files_only=True)
+ tokenizer = CLIPTokenizer.from_pretrained(model_name_checkpoint, subfolder=tokenizer_path, local_files_only=True)
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace CLIP Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.hidden_size}")
+ logger.info(f"intermediate_size: {hf_model.config.intermediate_size}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_attention_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_hidden_layers}")
+
+ # Test prompt. Cannot use randn tensor due to specific HF eos token id
+ test_text = "A coffee shop on Main Street that serves excellent pastries and opens at 7 AM on weekdays"
+
+ hf_inputs = tokenizer(test_text, padding=True, truncation=True, max_length=77, return_tensors="pt")
+ tt_prompt = ttnn.from_torch(
+ hf_inputs.input_ids,
+ dtype=ttnn.uint32,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ eos_token_id = hf_model.config.eos_token_id
+ logger.info(f"EOS token id: {eos_token_id}") # eos token id: 2
+
+ logger.info(f"Activation function: {hf_model.config.hidden_act}") # quick_gelu
+
+ # === TT-DiT CLIP ====
+ config = CLIPConfig(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.hidden_size,
+ ff_dim=hf_model.config.intermediate_size,
+ num_heads=hf_model.config.num_attention_heads,
+ num_hidden_layers=hf_model.config.num_hidden_layers,
+ max_prompt_length=77,
+ layer_norm_eps=hf_model.config.layer_norm_eps,
+ attention_dropout=hf_model.config.attention_dropout,
+ hidden_act=hf_model.config.hidden_act,
+ )
+
+ tt_clip = CLIPEncoder(config, encoder_submesh, ccl_manager, parallel_config, eos_token_id)
+ tt_clip.load_torch_state_dict(hf_model.state_dict())
+
+ # times TT model inference only
+ tt_start_time = time.time()
+ tt_sequence_output, tt_pooled_output = tt_clip(tt_prompt)
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ # get HF reference outputs
+ with torch.no_grad():
+ hf_start_time = time.time()
+ hf_output = hf_model(hf_inputs.input_ids, output_hidden_states=True)
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ hf_sequence_output = hf_output.hidden_states[-1]
+ hf_pooled_output = hf_output.pooler_output
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_sequence_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_sequence_output[-1])[0])
+ tt_pooled_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_pooled_output)[0])
+
+ logger.info(f"TT model execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF model execution time: {hf_execution_time:.4f} seconds")
+
+ assert hf_sequence_output.shape == tt_sequence_output_torch.shape
+ assert hf_pooled_output.shape == tt_pooled_output_torch.shape
+
+ assert_quality(hf_sequence_output, tt_sequence_output_torch, pcc=expected_pcc)
+ assert_quality(hf_pooled_output, tt_pooled_output_torch, pcc=expected_pcc)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/clip/test_clip_full_projection.py b/code/models/tt_dit/tests/encoders/clip/test_clip_full_projection.py
new file mode 100644
index 0000000000000000000000000000000000000000..14b1477d8f9c47dd665b0601d36f70f186a094f9
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/clip/test_clip_full_projection.py
@@ -0,0 +1,215 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers import CLIPTextConfig as HF_CLIPConfig
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer
+
+import ttnn
+from models.tt_dit.encoders.clip.model_clip import CLIPConfig, CLIPEncoder
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.tracing import Tracer
+
+
+def get_hf_config(model_name: str) -> HF_CLIPConfig:
+ # Clip config pulled from the model. Updated to use 2 layers for testing.
+ config_params = {
+ "text_encoder": {
+ "architectures": ["CLIPTextModelWithProjection"],
+ "bos_token_id": 0,
+ "eos_token_id": 2,
+ "hidden_act": "quick_gelu",
+ "hidden_size": 768,
+ "intermediate_size": 3072,
+ "num_attention_heads": 12,
+ "num_hidden_layers": 2,
+ "projection_dim": 768,
+ "vocab_size": 49408,
+ },
+ "text_encoder_2": {
+ "architectures": ["CLIPTextModelWithProjection"],
+ "bos_token_id": 0,
+ "eos_token_id": 2,
+ "hidden_act": "gelu",
+ "hidden_size": 1280,
+ "intermediate_size": 5120,
+ "num_attention_heads": 20,
+ "num_hidden_layers": 2,
+ "projection_dim": 1280,
+ "vocab_size": 49408,
+ },
+ }
+
+ return HF_CLIPConfig(**config_params[model_name])
+
+
+@pytest.mark.parametrize(
+ "dit_unit_test",
+ [{"1": True, "0": False}.get(os.environ.get("DIT_UNIT_TEST"), False)],
+)
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "clip_path, tokenizer_path, expected_pcc",
+ [
+ ("text_encoder", "tokenizer", 0.98),
+ ("text_encoder_2", "tokenizer_2", 0.98),
+ ],
+ ids=["encoder_1", "encoder_2"],
+)
+@pytest.mark.parametrize("mesh_device", [(2, 4), (1, 2)], ids=["t3k", "n300"], indirect=True)
+@pytest.mark.parametrize("submesh_shape", [(1, 4), (2, 2), (1, 1), (1, 2)], ids=["1x4", "2x2", "1x1", "1x2"])
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_clip_encoder(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: tuple[int, int],
+ model_name: str,
+ clip_path: str,
+ tokenizer_path: str,
+ expected_pcc: float,
+ topology: ttnn.Topology,
+ dit_unit_test: bool,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ # For N300 with parallel factor = 1, use factor=1 regardless of submesh shape
+ if tuple(mesh_device.shape) == (1, 2) and submesh_shape in [(1, 1), (1, 2)]:
+ parallel_factor = 1
+ else:
+ parallel_factor = encoder_submesh.shape[1]
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=parallel_factor, mesh_axis=1),
+ )
+ logger.info(f"Parallel factor: {parallel_factor}")
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ # HF model
+ if dit_unit_test:
+ hf_model = CLIPTextModelWithProjection(get_hf_config(clip_path))
+ else:
+ hf_model = CLIPTextModelWithProjection.from_pretrained(
+ model_name_checkpoint, subfolder=clip_path, local_files_only=True
+ )
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace CLIP Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.hidden_size}")
+ logger.info(f"intermediate_size: {hf_model.config.intermediate_size}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_attention_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_hidden_layers}")
+
+ tokenizer = CLIPTokenizer.from_pretrained(model_name_checkpoint, subfolder=tokenizer_path, local_files_only=True)
+
+ # Test prompt. Cannot use randn tensor due to specific HF eos token id
+ test_text = "A coffee shop on Main Street that serves excellent pastries and opens at 7 AM on weekdays"
+
+ hf_inputs = tokenizer(test_text, padding=True, truncation=True, max_length=77, return_tensors="pt")
+ tt_prompt = ttnn.from_torch(
+ hf_inputs.input_ids,
+ dtype=ttnn.uint32,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ eos_token_id = hf_model.config.eos_token_id
+ logger.info(f"EOS token id: {eos_token_id}") # eos token id: 2
+
+ logger.info(f"Activation function: {hf_model.config.hidden_act}") # quick_gelu
+
+ # === TT-DiT CLIP ====
+ config = CLIPConfig(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.hidden_size,
+ ff_dim=hf_model.config.intermediate_size,
+ num_heads=hf_model.config.num_attention_heads,
+ num_hidden_layers=hf_model.config.num_hidden_layers,
+ max_prompt_length=77,
+ layer_norm_eps=hf_model.config.layer_norm_eps,
+ attention_dropout=hf_model.config.attention_dropout,
+ hidden_act=hf_model.config.hidden_act,
+ projection_dim=hf_model.config.projection_dim,
+ )
+
+ tt_clip = CLIPEncoder(
+ config=config,
+ mesh_device=encoder_submesh,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ eos_token_id=eos_token_id,
+ )
+ tt_clip.load_torch_state_dict(hf_model.state_dict())
+ tracer = Tracer(tt_clip.forward, device=encoder_submesh, clone_prep_inputs=False)
+
+ # times TT model inference only
+ tt_start_time = time.time()
+ tt_sequence_output, tt_normalized_output = tracer(tt_prompt, skip_pooling=True)
+ tt_projected_output = tt_clip.pooled_output(
+ prompt_tokenized=tt_prompt,
+ normalized_final_state=tt_normalized_output,
+ )
+
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ # get HF reference outputs
+ with torch.no_grad():
+ hf_start_time = time.time()
+ hf_output = hf_model(hf_inputs.input_ids, output_hidden_states=True)
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ hf_sequence_output = hf_output.hidden_states[-1]
+ hf_projected_output = hf_output.text_embeds # projected/pooled output
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_sequence_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_sequence_output[-1])[0])
+ tt_projected_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_projected_output)[0])
+
+ logger.info(f"TT model execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF model execution time: {hf_execution_time:.4f} seconds")
+
+ assert hf_sequence_output.shape == tt_sequence_output_torch.shape
+ assert hf_projected_output.shape == tt_projected_output_torch.shape
+
+ assert_quality(hf_sequence_output, tt_sequence_output_torch, pcc=expected_pcc)
+ assert_quality(hf_projected_output, tt_projected_output_torch, pcc=expected_pcc)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/gemma/test_gemma_connector_blocks.py b/code/models/tt_dit/tests/encoders/gemma/test_gemma_connector_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..58e26dd88116f6dfc763dbb9b0b2eb1fe59a4024
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/gemma/test_gemma_connector_blocks.py
@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Isolated parity: device video connector vs diffusers ``LTX2ConnectorTransformer1d``
+on identical random features, asserting PCC. Isolating from the Gemma encoder keeps a
+regression here distinguishable from upstream bf16 drift.
+"""
+
+import glob
+import os
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from huggingface_hub.constants import HF_HUB_CACHE
+from loguru import logger
+from safetensors import safe_open
+
+import ttnn
+from models.tt_dit.encoders.gemma.embeddings_connector import EmbeddingsConnector
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.mochi import get_rot_transformation_mat
+from models.tt_dit.utils.tensor import bf16_tensor
+
+VIDEO_PREFIX = "model.diffusion_model.video_embeddings_connector."
+
+
+def _raw_video_connector_to_diffusers(raw: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Raw Lightricks video-connector weights → diffusers LTX2ConnectorTransformer1d keys
+ (transformer_1d_blocks → transformer_blocks, attn1.q_norm/k_norm → attn1.norm_q/norm_k)."""
+ out: dict[str, torch.Tensor] = {}
+ for k, v in raw.items():
+ sub = k[len(VIDEO_PREFIX) :]
+ sub = sub.replace("transformer_1d_blocks.", "transformer_blocks.")
+ sub = sub.replace(".attn1.q_norm.", ".attn1.norm_q.").replace(".attn1.k_norm.", ".attn1.norm_k.")
+ out[sub] = v
+ return out
+
+
+def _ltx_ckpt() -> str | None:
+ explicit = os.environ.get("LTX_CHECKPOINT")
+ if explicit and os.path.exists(explicit):
+ return explicit
+ # HF_HUB_CACHE resolves HF_HUB_CACHE / HF_HOME / the default. Hardcoding ~/.cache/huggingface
+ # here made this unfindable in CI, which points HF_HUB_CACHE at the shared weights mount.
+ cands = glob.glob(os.path.join(HF_HUB_CACHE, "models--Lightricks--LTX-2.3/snapshots/*/ltx-2.3-22b-dev.safetensors"))
+ return cands[0] if cands else None
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=["mesh_device"])
+@pytest.mark.parametrize("device_params", [{"l1_small_size": 8192}], indirect=["device_params"])
+def test_connector_blocks_isolated(*, mesh_device):
+ ckpt = _ltx_ckpt()
+ if not ckpt:
+ pytest.skip("LTX checkpoint not found")
+
+ pytest.importorskip("diffusers")
+ from diffusers.pipelines.ltx2.connectors import LTX2ConnectorTransformer1d
+
+ # --- device connector built directly (no pipeline): TP=1 on 1x1, weights from raw ckpt ---
+ enc_parallel = EncoderParallelConfig(tensor_parallel=ParallelFactor(factor=1, mesh_axis=1))
+ enc_ccl = CCLManager(mesh_device, num_links=1, topology=ttnn.Topology.Linear)
+ connector = EmbeddingsConnector(
+ output_dim=4096,
+ num_blocks=8,
+ num_heads=32,
+ mesh_device=mesh_device,
+ ccl_manager=enc_ccl,
+ parallel_config=enc_parallel,
+ )
+ raw_video = {}
+ with safe_open(ckpt, "pt") as f:
+ for k in f.keys():
+ if k.startswith(VIDEO_PREFIX):
+ raw_video[k] = f.get_tensor(k)
+ # The TT loader strips the prefix, renames blocks, and permutes Q/K split→interleaved.
+ connector.load_torch_state_dict({k[len(VIDEO_PREFIX) :]: v for k, v in raw_video.items()}, strict=False)
+
+ # --- reference: diffusers video connector (SPLIT rope = checkpoint's rotation, CPU fp32) ---
+ ref = (
+ LTX2ConnectorTransformer1d(
+ num_attention_heads=32,
+ attention_head_dim=128,
+ num_layers=8,
+ num_learnable_registers=128,
+ rope_base_seq_len=4096,
+ rope_theta=10000.0,
+ rope_double_precision=False,
+ rope_type="split",
+ gated_attention=True,
+ )
+ .float()
+ .eval()
+ )
+ inc = ref.load_state_dict(_raw_video_connector_to_diffusers(raw_video), strict=False)
+ logger.info(f"ref connector load: missing={len(inc.missing_keys)} unexpected={len(inc.unexpected_keys)}")
+
+ # --- identical random input WITH PADDING (left-pad) → register replacement active ---
+ # Mirrors the real encode path: only n_real real tokens, the rest become registers.
+ seq, dim, n_real = 256, 4096, 17
+ torch.manual_seed(0)
+ x = (torch.randn(1, seq, dim) * 0.5).bfloat16()
+ binary_mask = torch.zeros(1, seq, dtype=torch.long)
+ binary_mask[:, seq - n_real :] = 1 # left-pad: real tokens at the end
+
+ # reference: diffusers binarizes the additive mask at >= -9000.
+ add_mask = torch.where(binary_mask[:, None, None, :].bool(), 0.0, -1e9) # (1,1,1,seq)
+ with torch.no_grad():
+ ref_out = ref(x.float().clone(), add_mask)[0].float()
+
+ # device: connector forward (register replacement → on-device RoPE → blocks → norm).
+ tt_feat = ttnn.from_torch(x, device=mesh_device, layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16)
+ trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+ dev_out = connector(tt_feat, binary_mask, trans_mat=trans_mat)
+
+ logger.info(f"CONNECTOR (with registers) ref={tuple(ref_out.shape)} dev={tuple(dev_out.shape)}")
+ assert_quality(ref_out, dev_out, pcc=0.998)
diff --git a/code/models/tt_dit/tests/encoders/gemma/test_gemma_encoder_all_layers.py b/code/models/tt_dit/tests/encoders/gemma/test_gemma_encoder_all_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..d6e5ebb5778ceae522bb5d00484db087fcf9ff04
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/gemma/test_gemma_encoder_all_layers.py
@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Per-layer parity: TTNN GemmaEncoder vs HF reference, real-token PCC over the first
+``N_LAYERS`` decoder layers (the full 12B encoder doesn't fit one chip at tp=1).
+
+ pytest models/tt_dit/tests/encoders/gemma/test_gemma_encoder_all_layers.py -s
+"""
+
+import glob
+import os
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import gc
+
+import pytest
+import torch
+from huggingface_hub.constants import HF_HUB_CACHE
+from loguru import logger
+from safetensors.torch import load_file
+
+import ttnn
+from models.tt_dit.encoders.gemma.model_gemma import GemmaConfig, GemmaEncoder
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+
+N_LAYERS = int(os.environ.get("N_LAYERS", "6"))
+SEQ_LEN = 128 # must be >= SDPA q/k chunk size (128) in model_gemma to avoid a device hang
+# Per-state real-token PCC floor at the default N_LAYERS=6 (measured min ~0.9963 at the final
+# norm; shallower states ride ~0.9998). bf16 drifts with depth, so a much larger N_LAYERS may
+# need a lower floor.
+PCC_BAR = 0.995
+# Long enough to fill all SEQ_LEN slots with real tokens (no padding), so the reference
+# runs pure-causal and the device is_causal path matches exactly — no mask confound.
+PROMPT = (
+ "A plump orange tabby cat sits on a worn piano bench in a sunlit living room, "
+ "carefully playing the keys with its soft paws while dust motes drift through "
+ "the warm afternoon light. Outside the tall windows a gentle breeze stirs the "
+ "garden, and the faint melody echoes down the hallway past framed photographs, "
+ "old books stacked on the floor, and a sleeping dog curled on a faded rug near "
+ "the fireplace, utterly content in the quiet rhythm of the unhurried day. "
+ "Down the street a baker pulls warm loaves from the oven, their crust crackling "
+ "as steam rises into the cool morning air, and children chase a bright red ball "
+ "across the cobblestones while pigeons scatter and wheel above the rooftops. "
+ "Far beyond the town the mountains stand blue and patient under a wide pale sky, "
+ "their snowcaps catching the first gold of dawn as a river winds slowly through "
+ "the valley, carrying leaves and the reflections of clouds toward the distant sea."
+)
+
+
+def _gemma_path() -> str:
+ explicit = os.environ.get("GEMMA_PATH")
+ if explicit:
+ return explicit
+ # HF_HUB_CACHE resolves HF_HUB_CACHE / HF_HOME / the default. Hardcoding ~/.cache/huggingface
+ # here made this unfindable in CI, which points HF_HUB_CACHE at the shared weights mount.
+ cands = glob.glob(os.path.join(HF_HUB_CACHE, "models--google--gemma-3-12b-it-qat-q4_0-unquantized/snapshots/*/"))
+ if cands:
+ return cands[0].rstrip("/")
+ return "google/gemma-3-12b-it-qat-q4_0-unquantized"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=["mesh_device"])
+@pytest.mark.parametrize("device_params", [{"l1_small_size": 8192}], indirect=["device_params"])
+def test_gemma_layers_individually(*, mesh_device):
+ # _gemma_path falls back to the bare repo id, which from_pretrained resolves cache-then-hub.
+ # Gating an os.path.isdir on it made that fallback dead and skipped before any HF call.
+ gemma_path = _gemma_path()
+
+ from transformers import AutoModelForCausalLM, AutoTokenizer
+
+ tok = AutoTokenizer.from_pretrained(gemma_path)
+ tokens = tok(PROMPT, return_tensors="pt", padding="max_length", max_length=SEQ_LEN, truncation=True)
+ n_real = int(tokens.attention_mask.sum().item())
+ logger.info(f"tokens {tuple(tokens.input_ids.shape)}, real={n_real} (left-pad)")
+
+ # --- HF reference hidden states (first N_LAYERS only, for speed) ---
+ ref_model = AutoModelForCausalLM.from_pretrained(gemma_path, torch_dtype=torch.bfloat16).eval()
+ # Truncate the decoder stack to N_LAYERS so the CPU forward is ~8x cheaper.
+ lm = ref_model
+ for attr in ("model", "language_model", "layers"):
+ if hasattr(lm, attr):
+ if attr == "layers":
+ break
+ lm = getattr(lm, attr)
+ assert hasattr(lm, "layers"), "could not locate decoder layers on reference model"
+ lm.layers = lm.layers[:N_LAYERS]
+ if hasattr(lm, "config"):
+ lm.config.num_hidden_layers = N_LAYERS
+ with torch.no_grad():
+ out = ref_model(
+ input_ids=tokens.input_ids,
+ attention_mask=tokens.attention_mask,
+ output_hidden_states=True,
+ )
+ ref_hs = [h.float() for h in out.hidden_states] # embed + N_LAYERS
+ del ref_model, out
+ gc.collect()
+
+ # --- device encoder: first N_LAYERS only (tp=1) ---
+ weight_files = sorted(glob.glob(f"{gemma_path}/model-*.safetensors"))
+ full_sd = {}
+ for f in weight_files:
+ full_sd.update(load_file(f))
+
+ # keep embed_tokens, final norm, and first N_LAYERS decoder layers
+ keep = {}
+ for k, v in full_sd.items():
+ if k.startswith("language_model.model.layers."):
+ li = int(k.split("language_model.model.layers.")[1].split(".")[0])
+ if li >= N_LAYERS:
+ continue
+ keep[k] = v
+ del full_sd
+ gc.collect()
+
+ config = GemmaConfig(num_hidden_layers=N_LAYERS, max_position_embeddings=SEQ_LEN)
+ enc_parallel = EncoderParallelConfig(tensor_parallel=ParallelFactor(factor=1, mesh_axis=1))
+ enc_ccl = CCLManager(mesh_device, num_links=1, topology=ttnn.Topology.Linear)
+ encoder = GemmaEncoder(config, mesh_device, enc_ccl, enc_parallel)
+ inc = encoder.load_torch_state_dict(keep, strict=False)
+ logger.info(f"load: missing={len(inc.missing_keys)} unexpected={len(inc.unexpected_keys)}")
+ del keep
+ gc.collect()
+
+ tt_ids = ttnn.from_torch(tokens.input_ids, device=mesh_device, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT)
+ real_slice = slice(SEQ_LEN - n_real, SEQ_LEN) # left-pad → real tokens at the end
+
+ # Pure-causal path (no mask): all positions real, so the device is_causal path
+ # matches the reference exactly.
+ hs = encoder(tt_ids, attention_mask=None) # device: [embed, L0..L_{N-1}, final_norm]
+ # HF output_hidden_states is [embed, L0..L_{N-2}, final_norm(L_{N-1})] — its last entry is
+ # post-final-norm, not the raw last layer. Align the shared prefix directly and the final
+ # norm at [-1]; the device's raw last-layer state has no HF counterpart.
+ device_states = list(hs[: len(ref_hs) - 1]) + [hs[-1]]
+ for i, (ref, dev_tt) in enumerate(zip(ref_hs, device_states)):
+ dev = ttnn.to_torch(ttnn.get_device_tensors(dev_tt)[0]).float()
+ if dev.dim() == 4:
+ dev = dev[0]
+ assert torch.isfinite(dev).all(), f"state {i}: device output has NaN/Inf"
+ logger.info(" final_norm" if i == len(ref_hs) - 1 else f" L{i:02d}")
+ assert_quality(ref[:, real_slice, :], dev[:, real_slice, :], pcc=PCC_BAR)
+ for h in hs:
+ ttnn.deallocate(h)
diff --git a/code/models/tt_dit/tests/encoders/gemma/test_gemma_full.py b/code/models/tt_dit/tests/encoders/gemma/test_gemma_full.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c2b36be7629b87711cc326ed7d960e2bdc79c33
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/gemma/test_gemma_full.py
@@ -0,0 +1,294 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""End-to-end parity: on-device ``encode_prompts`` vs a torch reference built from HF
+Gemma-3 + diffusers ``LTX2TextConnectors`` (mirrors the diffusers ``LTX2Pipeline`` text
+path). PCC of the final video (4096) and audio (2048) context embeddings.
+
+``test_prof_gemma_ltx_devicetime`` (same file) is a Tracy per-op device-time profile of the
+encode; it skips unless GEMMA_PROF=1, so it never silently no-ops the default (asserting) suite.
+
+Do NOT run the 2x4 case under TT_METAL_WATCHER — the watcher overflows the active-eth
+fabric-router kernel-config buffer at device open.
+
+ pytest models/tt_dit/tests/encoders/gemma/test_gemma_full.py -s
+"""
+
+import glob
+import os
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from huggingface_hub.constants import HF_HUB_CACHE
+from loguru import logger
+from safetensors import safe_open
+
+import ttnn
+from models.tt_dit.encoders.gemma.model_gemma import GemmaEncoderLayer
+from models.tt_dit.pipelines.ltx.pipeline_ltx import LTXPipeline
+from models.tt_dit.utils.check import assert_quality
+
+
+def _connector_block_subkey(sub: str) -> str:
+ """Per-block raw Lightricks key → diffusers connector key: only the block-list name and
+ the QK-norm names differ (q_norm/k_norm → norm_q/norm_k); the rest is shared."""
+ sub = sub.replace("transformer_1d_blocks.", "transformer_blocks.")
+ sub = sub.replace(".attn1.q_norm.", ".attn1.norm_q.")
+ sub = sub.replace(".attn1.k_norm.", ".attn1.norm_k.")
+ return sub
+
+
+def _raw_connectors_to_diffusers(raw: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Raw Lightricks connector + aggregate_embed weights → diffusers LTX2TextConnectors keys.
+ aggregate_embed → text_proj_in; the connector blocks map 1:1 modulo _connector_block_subkey."""
+ VP = "model.diffusion_model.video_embeddings_connector."
+ AP = "model.diffusion_model.audio_embeddings_connector."
+ VA = "text_embedding_projection.video_aggregate_embed."
+ AA = "text_embedding_projection.audio_aggregate_embed."
+ out: dict[str, torch.Tensor] = {}
+ for k, v in raw.items():
+ if k.startswith(VA):
+ out["video_text_proj_in." + k[len(VA) :]] = v
+ elif k.startswith(AA):
+ out["audio_text_proj_in." + k[len(AA) :]] = v
+ elif k.startswith(VP):
+ out["video_connector." + _connector_block_subkey(k[len(VP) :])] = v
+ elif k.startswith(AP):
+ out["audio_connector." + _connector_block_subkey(k[len(AP) :])] = v
+ return out
+
+
+PROMPT = "A plump orange tabby cat sits on a piano bench playing keys with its paws."
+
+CONNECTOR_PREFIXES = (
+ "text_embedding_projection.video_aggregate_embed.",
+ "text_embedding_projection.audio_aggregate_embed.",
+ "model.diffusion_model.video_embeddings_connector.",
+ "model.diffusion_model.audio_embeddings_connector.",
+)
+
+
+def _gemma_path() -> str:
+ explicit = os.environ.get("GEMMA_PATH")
+ if explicit:
+ return explicit
+ # HF_HUB_CACHE resolves HF_HUB_CACHE / HF_HOME / the default. Hardcoding ~/.cache/huggingface
+ # here made this unfindable in CI, which points HF_HUB_CACHE at the shared weights mount.
+ cands = glob.glob(os.path.join(HF_HUB_CACHE, "models--google--gemma-3-12b-it-qat-q4_0-unquantized/snapshots/*/"))
+ return cands[0].rstrip("/") if cands else "google/gemma-3-12b-it-qat-q4_0-unquantized"
+
+
+def _ltx_ckpt() -> str | None:
+ explicit = os.environ.get("LTX_CHECKPOINT")
+ if explicit and os.path.exists(explicit):
+ return explicit
+ # HF_HUB_CACHE resolves HF_HUB_CACHE / HF_HOME / the default. Hardcoding ~/.cache/huggingface
+ # here made this unfindable in CI, which points HF_HUB_CACHE at the shared weights mount.
+ cands = glob.glob(os.path.join(HF_HUB_CACHE, "models--Lightricks--LTX-2.3/snapshots/*/ltx-2.3-22b-dev.safetensors"))
+ return cands[0] if cands else None
+
+
+def _encode_prompts_reference(
+ checkpoint_path: str, gemma_root: str, prompts: list[str]
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Torch reference: HF Gemma-3 text encoder + diffusers ``LTX2TextConnectors``, mirroring
+ the diffusers ``LTX2Pipeline`` text path (stacked per-layer hidden states → per-modality
+ feature extraction + connectors). Connector/aggregate weights come from the raw checkpoint;
+ the LTX-2 reference repo is no longer used. Returns (video_encoding, audio_encoding)."""
+ pytest.importorskip("diffusers")
+ from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
+ from transformers import AutoModelForCausalLM, AutoTokenizer
+
+ # LTX-2.3-22b connector config: per-modality projections, 8 blocks each, video 32x128=4096,
+ # audio 32x64=2048, gated attention, SPLIT rope (the raw checkpoint's rotation convention).
+ connectors = (
+ LTX2TextConnectors(
+ caption_channels=3840,
+ text_proj_in_factor=49,
+ video_connector_num_attention_heads=32,
+ video_connector_attention_head_dim=128,
+ video_connector_num_layers=8,
+ audio_connector_num_attention_heads=32,
+ audio_connector_attention_head_dim=64,
+ audio_connector_num_layers=8,
+ video_gated_attn=True,
+ audio_gated_attn=True,
+ per_modality_projections=True,
+ video_hidden_dim=4096,
+ audio_hidden_dim=2048,
+ rope_type="split",
+ rope_double_precision=False,
+ proj_bias=True,
+ )
+ .float()
+ .eval()
+ )
+ raw = {}
+ with safe_open(checkpoint_path, "pt") as f:
+ for k in f.keys():
+ if k.startswith(CONNECTOR_PREFIXES):
+ raw[k] = f.get_tensor(k)
+ inc = connectors.load_state_dict(_raw_connectors_to_diffusers(raw), strict=False)
+ logger.info(f"ref connectors load: missing={len(inc.missing_keys)} unexpected={len(inc.unexpected_keys)}")
+
+ tok = AutoTokenizer.from_pretrained(gemma_root)
+ tok.padding_side = "left" # Gemma-3 / LTX-2 use left padding
+ if tok.pad_token is None:
+ tok.pad_token = tok.eos_token
+ text_encoder = AutoModelForCausalLM.from_pretrained(gemma_root, torch_dtype=torch.bfloat16).eval()
+
+ ti = tok(
+ prompts,
+ padding="max_length",
+ max_length=1024,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ with torch.no_grad():
+ out = text_encoder(input_ids=ti.input_ids, attention_mask=ti.attention_mask, output_hidden_states=True)
+ # Pack the 49 per-layer hidden states D-major (B, seq, 3840, 49) → (B, seq, 188160).
+ hidden = torch.stack(out.hidden_states, dim=-1).flatten(2, 3).float()
+ video, audio, _ = connectors(hidden, ti.attention_mask)
+ return video.float(), audio.float()
+
+
+# 2x4 drives the encoder's TP all-gathers over CCL, which needs the 1D fabric up;
+# 1x1 is single-chip with no CCL, so it must run without fabric.
+@pytest.mark.parametrize(
+ ("mesh_device", "device_params"),
+ [
+ pytest.param((1, 1), {"l1_small_size": 8192}, id="1x1"),
+ pytest.param((1, 8), {"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, id="1x8"),
+ pytest.param((2, 4), {"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, id="2x4"),
+ pytest.param((4, 8), {"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING}, id="4x8"),
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+def test_gemma_encoder(*, mesh_device):
+ # _gemma_path falls back to the bare repo id, which from_pretrained resolves cache-then-hub.
+ # Gating an os.path.isdir on it made that fallback dead and skipped before any HF call.
+ gemma = _gemma_path()
+ ckpt = _ltx_ckpt()
+ if not ckpt:
+ pytest.skip("LTX checkpoint not found")
+
+ # Bare pipeline: checkpoint_name=None skips the heavy transformer/VAE load. dynamic_load=False
+ # keeps the encoder resident so the whole-encode trace (default when resident) is exercised.
+ pipe = LTXPipeline.create_pipeline(
+ mesh_device, checkpoint_name=None, gemma_path=gemma, mode="av", dynamic_load=False
+ )
+
+ # On-device Gemma encoder (full 48 layers). TP follows the T5 pattern (axis-1 width):
+ # TP=1 on 1x1, TP=4 on 2x4 — set inside the loader, no override needed.
+ pipe.gemma_encoder_pair.load_gemma_encoder(gemma)
+
+ # Load only the connector weights from the 46GB checkpoint.
+ conn_state = {}
+ with safe_open(ckpt, "pt") as f:
+ for k in f.keys():
+ if k.startswith(CONNECTOR_PREFIXES):
+ conn_state[k] = f.get_tensor(k)
+ logger.info(f"connector weights: {len(conn_state)} tensors")
+ pipe.gemma_encoder_pair.load_embeddings_connectors(conn_state, audio_num_blocks=8)
+
+ # Reference embeds (HF Gemma-3 + diffusers LTX2TextConnectors), local to this test.
+ v_ref, a_ref = _encode_prompts_reference(ckpt, gemma, [PROMPT])
+
+ # Device embeds. First call compiles + populates the program cache; the second is warm
+ # and is the one we time. use_cache=False to measure the real encode, not a cache load.
+ pipe.encode_prompts([PROMPT], use_cache=False)
+ t0 = time.perf_counter()
+ dev = pipe.encode_prompts([PROMPT], use_cache=False)
+ t_warm_ms = (time.perf_counter() - t0) * 1e3
+
+ v_dev = torch.as_tensor(dev[0][0]).float()
+ a_dev = torch.as_tensor(dev[0][1]).float()
+
+ logger.info(f"ENCODE warm wall-clock (mesh {tuple(mesh_device.shape)}): {t_warm_ms:.1f} ms")
+
+ # Device and reference both pad to GEMMA_SEQUENCE_LENGTH=1024 (the reference text
+ # encoder fixes its tokenizer to 1024 too), so the outputs are the same shape and
+ # compare directly — no alignment needed.
+ logger.info(f"VIDEO ref={tuple(v_ref.shape)} dev={tuple(v_dev.shape)}")
+ assert_quality(v_ref, v_dev, pcc=0.9995)
+ # fp32 connector weights (HiFi4) match the host reference, so audio holds ~0.9999 like video.
+ logger.info(f"AUDIO ref={tuple(a_ref.shape)} dev={tuple(a_dev.shape)}")
+ assert_quality(a_ref, a_dev, pcc=0.9995)
+
+
+def _walk(m):
+ yield m
+ for _, c in m.named_children():
+ yield from _walk(c)
+
+
+def _flush_after(mod, mesh_device):
+ """Flush the on-device profiler after each mod.forward so each Tracy zone window stays under
+ the 1000-zone buffer."""
+ orig = mod.forward
+
+ def timed(*a, _orig=orig, **k):
+ r = _orig(*a, **k)
+ ttnn.ReadDeviceProfiler(mesh_device)
+ return r
+
+ mod.forward = timed
+
+
+@pytest.mark.skipif(
+ os.environ.get("GEMMA_PROF") != "1",
+ reason="Tracy profiling harness — run: GEMMA_PROF=1 python -m tracy -p -r -m pytest "
+ "'::test_prof_gemma_ltx_devicetime[2x4]' -s",
+)
+@pytest.mark.parametrize(
+ "device_params", [{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True
+)
+@pytest.mark.parametrize("mesh_device", [(2, 4)], indirect=True)
+def test_prof_gemma_ltx_devicetime(mesh_device, device_params):
+ """Tracy per-op device-time profile of the full encode (tokenize -> 48 layers -> feature
+ extractor -> connectors). Per-layer ReadDeviceProfiler flush keeps each zone window under
+ Tracy's 1000-zone buffer; no PCC — pure timing. Sum DEVICE FW DURATION from the CSV vs the
+ printed GEMMA_ENCODE_HOST_WALL_MS — the gap is host dispatch."""
+ gemma = _gemma_path()
+ ckpt = _ltx_ckpt()
+ if not ckpt:
+ pytest.skip("LTX checkpoint not found")
+
+ pipe = LTXPipeline.create_pipeline(mesh_device, checkpoint_name=None, gemma_path=gemma, mode="av")
+ pipe.gemma_encoder_pair.load_gemma_encoder(gemma)
+
+ conn_state = {}
+ with safe_open(ckpt, "pt") as f:
+ for k in f.keys():
+ if k.startswith(CONNECTOR_PREFIXES):
+ conn_state[k] = f.get_tensor(k)
+ pipe.gemma_encoder_pair.load_embeddings_connectors(conn_state, audio_num_blocks=8)
+
+ # Flush after each decoder layer plus the feature extractor and each connector — keeps every
+ # window well under Tracy's 1000-zone buffer.
+ pair = pipe.gemma_encoder_pair
+ for mod in _walk(pair.gemma_encoder):
+ if isinstance(mod, GemmaEncoderLayer):
+ _flush_after(mod, mesh_device)
+ _flush_after(pair.feature_extractor, mesh_device)
+ _flush_after(pair.video_connector, mesh_device)
+ if pair.audio_connector is not None:
+ _flush_after(pair.audio_connector, mesh_device)
+
+ pipe.encode_prompts([PROMPT], use_cache=False) # warm program cache
+ ttnn.synchronize_device(mesh_device)
+ ttnn.ReadDeviceProfiler(mesh_device)
+
+ t0 = time.perf_counter()
+ pipe.encode_prompts([PROMPT], use_cache=False)
+ ttnn.synchronize_device(mesh_device)
+ host_wall = (time.perf_counter() - t0) * 1000
+ ttnn.ReadDeviceProfiler(mesh_device)
+ print(f"\nGEMMA_ENCODE_HOST_WALL_MS={host_wall:.2f}", flush=True)
diff --git a/code/models/tt_dit/tests/encoders/mistral3/test_mistral3.py b/code/models/tt_dit/tests/encoders/mistral3/test_mistral3.py
new file mode 100644
index 0000000000000000000000000000000000000000..c83d6296465c95c96b7dcebabbc31e69771fb3d9
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/mistral3/test_mistral3.py
@@ -0,0 +1,398 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import time
+
+import pytest
+import torch
+import transformers
+import transformers.generation.utils
+import transformers.models.mistral.modeling_mistral
+from loguru import logger
+
+import ttnn
+
+from ....blocks.rope import RopeConfig
+from ....encoders.mistral3.model_mistral3 import Mistral3Encoder
+from ....parallel.config import EncoderParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils import cache, tensor
+from ....utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "skip_layers"),
+ [
+ pytest.param((1, 4), 0, id="1x4"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "masked",
+ [
+ pytest.param(True, id="masked"),
+ # pytest.param(False, id="unmasked"),
+ ],
+)
+def test_generation(*, mesh_device: ttnn.MeshDevice, skip_layers: int, masked: bool) -> None:
+ torch.manual_seed(0)
+
+ tp_axis = 1
+ max_length = 40
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+ parallel_config = (
+ EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=mesh_device.shape[tp_axis], mesh_axis=tp_axis),
+ )
+ if tp_axis is not None
+ else None
+ )
+
+ tokenizer = transformers.LlamaTokenizerFast.from_pretrained("black-forest-labs/FLUX.2-dev", subfolder="tokenizer")
+ assert isinstance(tokenizer, transformers.LlamaTokenizerFast)
+
+ torch_model = transformers.Mistral3ForConditionalGeneration.from_pretrained(
+ "black-forest-labs/FLUX.2-dev", subfolder="text_encoder"
+ )
+ config = torch_model.model.language_model.config
+
+ num_layers = len(torch_model.model.language_model.layers)
+ del torch_model.model.language_model.layers[num_layers - skip_layers :]
+
+ generation_config = torch_model.generation_config
+ assert isinstance(generation_config, transformers.GenerationConfig)
+
+ model = Mistral3Encoder(
+ vocab_size=config.vocab_size,
+ head_size=config.head_dim,
+ embed_size=config.hidden_size,
+ ff_size=config.intermediate_size,
+ num_layers=config.num_hidden_layers - skip_layers,
+ num_heads=config.num_attention_heads,
+ num_kv_heads=config.num_key_value_heads,
+ norm_eps=config.rms_norm_eps,
+ attn_qkv_bias=False,
+ attn_out_bias=False,
+ rope_config=RopeConfig(theta=config.rope_parameters["rope_theta"]),
+ device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+
+ state_dict = torch_model.state_dict()
+ state_dict = Mistral3Encoder.convert_state(state_dict)
+ cache.load_model(
+ model,
+ get_torch_state_dict=lambda: state_dict,
+ model_name="flux2",
+ subfolder="text_encoder",
+ parallel_config=parallel_config,
+ mesh_shape=tuple(mesh_device.shape),
+ )
+
+ # This makes unmasked generation more similar in the two implementations, possibly because
+ # `transformers` masks out padding tokens.
+ tokenizer.pad_token_id = generation_config.bos_token_id
+
+ out = tokenizer.__call__(
+ ["Once upon a time", "Hello"],
+ padding="longest",
+ # padding side does not matter for our implementation but the
+ # transformers library complains if right padding is used
+ padding_side="left",
+ return_tensors="pt",
+ return_attention_mask=True,
+ )
+ tokens = out["input_ids"].to(torch_model.device)
+ mask = out["attention_mask"].to(torch_model.device) if masked else None
+
+ tt_tokens = tensor.from_torch(tokens, device=mesh_device, dtype=ttnn.uint32)
+ tt_mask = tensor.from_torch(mask, device=mesh_device) if mask is not None else None
+
+ generation_config.max_length = max_length
+ generation_config.repetition_penalty = None # repetition penalty is not implemented
+ generation_config.return_dict_in_generate = True
+ generation_config.output_logits = True
+
+ print("running ttnn model...")
+
+ start_time = time.time()
+ tt_out = model.generate(
+ tt_tokens,
+ mask=tt_mask,
+ eos_tokens=generation_config.eos_token_id,
+ max_length=generation_config.max_length,
+ top_k=generation_config.top_k if generation_config.do_sample else 1,
+ top_p=generation_config.top_p or 1,
+ temperature=generation_config.temperature,
+ )
+ tt_tokens_out = tensor.to_torch(tt_out.tokens)
+
+ print(f"generation took {time.time() - start_time:.2f} seconds")
+
+ for i in range(tt_tokens_out.size(0)):
+ print(tokenizer.decode(tt_tokens_out[i]))
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "skip_layers"),
+ [
+ pytest.param((1, 4), 0, id="1x4"),
+ pytest.param((1, 8), 0, id="1x8"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "masked",
+ [
+ pytest.param(True, id="masked"),
+ pytest.param(False, id="unmasked"),
+ ],
+)
+def test_guided_generation(*, mesh_device: ttnn.MeshDevice, skip_layers: int, masked: bool) -> None:
+ torch.manual_seed(0)
+
+ tp_axis = 1
+ max_length = 20
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+ parallel_config = (
+ EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=mesh_device.shape[tp_axis], mesh_axis=tp_axis),
+ )
+ if tp_axis is not None
+ else None
+ )
+
+ tokenizer = transformers.LlamaTokenizerFast.from_pretrained("black-forest-labs/FLUX.2-dev", subfolder="tokenizer")
+ assert isinstance(tokenizer, transformers.LlamaTokenizerFast)
+
+ torch_model = transformers.Mistral3ForConditionalGeneration.from_pretrained(
+ "black-forest-labs/FLUX.2-dev", subfolder="text_encoder"
+ )
+ config = torch_model.model.language_model.config
+
+ num_layers = len(torch_model.model.language_model.layers)
+ del torch_model.model.language_model.layers[num_layers - skip_layers :]
+
+ generation_config = torch_model.generation_config
+ assert isinstance(generation_config, transformers.GenerationConfig)
+
+ model = Mistral3Encoder(
+ vocab_size=config.vocab_size,
+ head_size=config.head_dim,
+ embed_size=config.hidden_size,
+ ff_size=config.intermediate_size,
+ num_layers=config.num_hidden_layers - skip_layers,
+ num_heads=config.num_attention_heads,
+ num_kv_heads=config.num_key_value_heads,
+ norm_eps=config.rms_norm_eps,
+ attn_qkv_bias=False,
+ attn_out_bias=False,
+ rope_config=RopeConfig(theta=config.rope_parameters["rope_theta"]),
+ device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+
+ state_dict = torch_model.state_dict()
+ state_dict = Mistral3Encoder.convert_state(state_dict)
+ cache.load_model(
+ model,
+ get_torch_state_dict=lambda: state_dict,
+ model_name="flux2",
+ subfolder="text_encoder",
+ parallel_config=parallel_config,
+ mesh_shape=tuple(mesh_device.shape),
+ )
+
+ # This makes unmasked generation more similar in the two implementations, possibly because
+ # `transformers` masks out padding tokens.
+ tokenizer.pad_token_id = generation_config.bos_token_id
+
+ out = tokenizer.__call__(
+ ["Once upon a time", "Hello"],
+ padding="longest",
+ # padding side does not matter for our implementation but the
+ # transformers library complains if right padding is used
+ padding_side="left",
+ return_tensors="pt",
+ return_attention_mask=True,
+ )
+ tokens = out["input_ids"].to(torch_model.device)
+ mask = out["attention_mask"].to(torch_model.device) if masked else None
+
+ tt_tokens = tensor.from_torch(tokens, device=mesh_device, dtype=ttnn.uint32)
+ tt_mask = tensor.from_torch(mask, device=mesh_device) if mask is not None else None
+
+ generation_config.max_length = max_length
+ generation_config.repetition_penalty = None # repetition penalty is not implemented
+ generation_config.return_dict_in_generate = True
+ generation_config.output_logits = True
+
+ print("running torch model...")
+ torch_mask_input = mask if mask is not None else torch.ones_like(tokens)
+ out_ref = torch_model.generate(tokens, attention_mask=torch_mask_input)
+ assert isinstance(out_ref, transformers.generation.utils.GenerateOutput)
+
+ tokens_out = out_ref.sequences
+ logits = torch.stack(out_ref.logits, dim=1)
+
+ # diffusers somtimes generates longer sequences than max_length, in particular when `max_length
+ # = 20` for whatever reason.
+ tokens_out = tokens_out[:, :max_length]
+ logits = logits[:, : max_length - tokens.size(1)]
+
+ print("running ttnn model...")
+ tt_out = model.generate(
+ tt_tokens,
+ guide=tokens_out,
+ mask=tt_mask,
+ eos_tokens=generation_config.eos_token_id,
+ max_length=generation_config.max_length,
+ top_k=generation_config.top_k if generation_config.do_sample else 1,
+ top_p=generation_config.top_p or 1,
+ temperature=generation_config.temperature,
+ return_logits=True,
+ )
+
+ tt_tokens_out = tensor.to_torch(tt_out.tokens)
+ tt_logits = tensor.to_torch(ttnn.stack(tt_out.logits, dim=1), mesh_axes=[..., tp_axis])
+
+ # To compare generated tokens, remove `guide` in the call to `model.generate`!
+ # for i in range(tt_tokens_out.size(0)):
+ # print(tokenizer.decode(tokens_out[i]))
+ # print(tokenizer.decode(tt_tokens_out[i]))
+
+ if mask is not None:
+ # Masked positions on the start of the sequence contain random values from computing softmax over all -inf
+ # so we remove them before comparison.
+ _, s, d = logits.shape
+ padded_mask = torch.nn.functional.pad(mask.bool(), [0, s - mask.size(1)], value=True)
+ logits = logits.masked_select(padded_mask.unsqueeze(-1)).view([-1, d])
+ tt_logits = tt_logits.masked_select(padded_mask.unsqueeze(-1)).view([-1, d])
+
+ assert_quality(logits, tt_logits, ccc=0.9980, relative_rmse=0.063)
+
+ assert tt_tokens_out.eq(tokens_out).all()
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "batch_size", "skip_layers"),
+ [
+ pytest.param((1, 1), 2, 32, id="1x1"),
+ pytest.param((1, 2), 2, 22, id="1x2"),
+ pytest.param((1, 4), 2, 0, id="1x4"),
+ pytest.param((1, 8), 2, 0, id="1x8"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "masked",
+ [
+ pytest.param(True, id="masked"),
+ pytest.param(False, id="unmasked"),
+ ],
+)
+def test_transformer(*, mesh_device: ttnn.MeshDevice, batch_size: int, skip_layers: int, masked: bool) -> None:
+ torch.manual_seed(0)
+
+ sequence_length = 512
+ tp_axis = 1
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+ parallel_config = (
+ EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=mesh_device.shape[tp_axis], mesh_axis=tp_axis),
+ )
+ if tp_axis is not None
+ else None
+ )
+
+ torch_model = transformers.Mistral3ForConditionalGeneration.from_pretrained(
+ "black-forest-labs/FLUX.2-dev", subfolder="text_encoder"
+ )
+ config = torch_model.model.language_model.config
+
+ num_layers = len(torch_model.model.language_model.layers)
+ del torch_model.model.language_model.layers[num_layers - skip_layers :]
+
+ model = Mistral3Encoder(
+ vocab_size=config.vocab_size,
+ head_size=config.head_dim,
+ embed_size=config.hidden_size,
+ ff_size=config.intermediate_size,
+ num_layers=config.num_hidden_layers - skip_layers,
+ num_heads=config.num_attention_heads,
+ num_kv_heads=config.num_key_value_heads,
+ norm_eps=config.rms_norm_eps,
+ attn_qkv_bias=False,
+ attn_out_bias=False,
+ rope_config=RopeConfig(theta=config.rope_parameters["rope_theta"]),
+ device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+
+ state_dict = torch_model.state_dict()
+ state_dict = Mistral3Encoder.convert_state(state_dict)
+ cache.load_model(
+ model,
+ get_torch_state_dict=lambda: state_dict,
+ model_name="flux2",
+ subfolder="text_encoder",
+ parallel_config=parallel_config,
+ mesh_shape=tuple(mesh_device.shape),
+ )
+
+ tokens = torch.randint(0, config.vocab_size, [batch_size, sequence_length])
+ lengths = torch.randint(sequence_length // 4, 3 * sequence_length // 4, [batch_size])
+ mask = torch.arange(sequence_length).flip([0]) < lengths.unsqueeze(1) if masked else None
+
+ tt_tokens = tensor.from_torch(tokens, device=mesh_device, dtype=ttnn.uint32)
+ tt_mask = tensor.from_torch(mask, device=mesh_device) if mask is not None else None
+
+ logger.info("running ttnn model...")
+ tt_hidden_states = model.forward(
+ tt_tokens,
+ mask=tt_mask,
+ skip_final_linear=True,
+ output_hidden_states=True,
+ )
+ tt_hidden_states_torch = [tensor.to_torch(t) for t in tt_hidden_states]
+
+ logger.info("running torch model...")
+ torch_mask_input = mask if mask is not None else torch.ones_like(tokens)
+ with torch.no_grad():
+ hidden_states = torch_model.forward(
+ tokens, attention_mask=torch_mask_input, output_hidden_states=True
+ ).hidden_states
+
+ if mask is not None:
+ # Masked positions on the start of the sequence contain undefined values from computing softmax over all -inf
+ # so we remove them before comparison.
+ _, _, d = hidden_states[0].shape
+ hidden_states = [t.masked_select(mask.unsqueeze(-1)).view([-1, d]) for t in hidden_states]
+ tt_hidden_states_torch = [t.masked_select(mask.unsqueeze(-1)).view([-1, d]) for t in tt_hidden_states_torch]
+
+ assert len(hidden_states) == len(tt_hidden_states_torch)
+
+ for x, tt_x in zip(hidden_states[-4:], tt_hidden_states_torch[-4:], strict=True):
+ assert_quality(x, tt_x, pcc=0.9979, relative_rmse=0.065)
diff --git a/code/models/tt_dit/tests/encoders/qwen25vl/test_qwen25vl.py b/code/models/tt_dit/tests/encoders/qwen25vl/test_qwen25vl.py
new file mode 100644
index 0000000000000000000000000000000000000000..8be2bc4da7df80aa85425ab6255d537628d2c3ea
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen25vl/test_qwen25vl.py
@@ -0,0 +1,354 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import diffusers.pipelines.qwenimage.pipeline_qwenimage
+import pytest
+import torch
+import transformers
+import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl
+from loguru import logger
+
+import ttnn
+
+from ....encoders.qwen25vl.encoder_pair import Qwen25VlTokenizerEncoderPair
+from ....encoders.qwen25vl.model_qwen25vl import (
+ Qwen25VlAttention,
+ Qwen25VlContext,
+ Qwen25VlTextEncoder,
+ create_rope_tensors,
+ prepare_attention_bias,
+)
+from ....parallel.config import EncoderParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.test import line_params_req_exact_devices
+
+
+def _text_config(config):
+ """transformers 5.x nests the language-model fields under `.text_config`; 4.x had them at the
+ top level of `Qwen2_5_VLConfig`. Returns `config` unchanged when already a text config."""
+ return getattr(config, "text_config", config)
+
+
+def _rope_params(config):
+ """transformers 5.x moved `rope_theta` into `rope_parameters` (aliased as `rope_scaling`)."""
+ params = getattr(config, "rope_parameters", None) or config.rope_scaling
+ rope_theta = params["rope_theta"] if "rope_theta" in params else config.rope_theta
+ return rope_theta, params["mrope_section"]
+
+
+def _real_tokens_only(x: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
+ """Drop padding rows before comparing against the transformers reference.
+
+ We and transformers assign different position numbers to *padding* tokens, and that
+ disagreement is the only thing separating the two outputs. For 5 real tokens and 3 padding:
+
+ real tokens padding
+ ours (create_rope_tensors) 0, 1, 2, 3, 4 1, 1, 1
+ theirs (auto-inferred) 0, 1, 2, 3, 4 5, 6, 7
+
+ Both tests here compare against the auto-inferred convention: test_qwen25vl_text_encoder
+ lets Qwen2_5_VLTextModel infer positions internally, and test_qwen25vl_attention builds the
+ same arange positions by hand because it drives a single attention module directly.
+
+ Real tokens get the same position number either way, so the outputs agree everywhere a caller
+ actually reads. Padding output is discarded downstream by every caller, so comparing it
+ measures nothing but the convention mismatch -- it dragged PCC to ~91.5% against a 95.2%
+ threshold while the real tokens were exact.
+
+ transformers 4.x used our convention, which is why this only surfaced with the 5.12.1 pin.
+ Whether `create_rope_tensors` should switch to theirs is a live question for the Qwen-Image
+ owners; it changes model code, not just this test, so it is deliberately not decided here.
+ """
+ if attention_mask is None:
+ return x
+ return x[attention_mask.bool()]
+
+
+@pytest.mark.parametrize(
+ "mesh_device",
+ [
+ pytest.param((1, 2), id="1x2"),
+ pytest.param((1, 8), id="1x8"),
+ ],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "masked",
+ [
+ pytest.param(True, id="masked"),
+ pytest.param(False, id="unmasked"),
+ ],
+)
+def test_qwen25vl_attention(*, mesh_device: ttnn.MeshDevice, masked: bool) -> None:
+ torch.manual_seed(0)
+
+ batch_size = 10
+ sequence_length = 512
+ tp_axis = 1
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+
+ parent_torch_model = transformers.Qwen2_5_VLForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen-Image", subfolder="text_encoder"
+ )
+ torch_model = parent_torch_model.model.language_model.layers[0].self_attn
+ assert isinstance(torch_model, transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLAttention)
+
+ attn_config = _text_config(torch_model.config)
+ rope_theta, mrope_section = _rope_params(attn_config)
+
+ model = Qwen25VlAttention(
+ hidden_size=attn_config.hidden_size,
+ num_heads=attn_config.num_attention_heads,
+ num_key_value_heads=attn_config.num_key_value_heads,
+ ctx=Qwen25VlContext(mesh_device, tp_axis, ccl_manager),
+ )
+ model.load_torch_state_dict(torch_model.state_dict())
+
+ sequence = torch.randn([batch_size, sequence_length, attn_config.hidden_size])
+ m = torch.randint(0, sequence_length + 1, [batch_size])
+ attention_mask = torch.arange(sequence_length) < m.unsqueeze(1) if masked else None
+ cos, sin = create_rope_tensors(
+ batch_size,
+ sequence_length,
+ attention_mask,
+ head_dim=attn_config.hidden_size // attn_config.num_attention_heads,
+ rope_theta=rope_theta,
+ mrope_section=mrope_section,
+ )
+
+ tt_sequence = tensor.from_torch(sequence, device=mesh_device)
+ tt_attention_mask = tensor.from_torch(attention_mask, device=mesh_device) if attention_mask is not None else None
+ tt_pos_embeds_cos = tensor.from_torch(cos, device=mesh_device)
+ tt_pos_embeds_sin = tensor.from_torch(sin, device=mesh_device)
+
+ tt_attention_bias = prepare_attention_bias(tt_attention_mask) if tt_attention_mask is not None else None
+
+ logger.info("running ttnn model...")
+ tt_out = model.forward(
+ tt_sequence,
+ attention_bias=tt_attention_bias,
+ pos_embeds=(tt_pos_embeds_cos, tt_pos_embeds_sin),
+ )
+ tt_out_torch = tensor.to_torch(tt_out)
+
+ logger.info("running torch model...")
+ # get_rope_index is a vision+text utility; its own docstring says pure-text callers should
+ # "rely on model's auto-inferred position ids". For a text-only sequence Qwen2_5_VLTextModel
+ # infers plain arange positions replicated across all three mrope axes, so build that directly
+ # instead of calling get_rope_index off-label (which additionally requires a mm_token_type_ids
+ # argument and an integer input_ids carrier it has no real use for here).
+ position_ids = torch.arange(sequence_length).view(1, 1, -1).expand(3, batch_size, -1)
+ # transformers 5.x loads at the checkpoint dtype (bf16) where 4.x loaded fp32, so the fp32
+ # `sequence` no longer matches the weights. Feed the reference what the device already sees.
+ ref_sequence = sequence.to(next(parent_torch_model.parameters()).dtype)
+ # transformers 5.x moved rotary_emb off the attention module up to Qwen2_5_VLTextModel.
+ position_embeddings = parent_torch_model.model.language_model.rotary_emb(ref_sequence, position_ids)
+ if attention_mask is not None:
+ causal_attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)
+ causal_attention_mask = causal_attention_mask.expand([-1, -1, sequence_length, -1])
+ causal_attention_mask = causal_attention_mask.tril()
+ causal_attention_mask = causal_attention_mask.bool()
+ else:
+ causal_attention_mask = None
+
+ with torch.no_grad():
+ out, _ = torch_model.forward(
+ ref_sequence,
+ attention_mask=causal_attention_mask,
+ position_embeddings=position_embeddings,
+ )
+
+ assert_quality(
+ _real_tokens_only(out, attention_mask),
+ _real_tokens_only(tt_out_torch, attention_mask),
+ pcc=0.988,
+ relative_rmse=0.15,
+ )
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "batch_size", "skip_layers"),
+ [
+ pytest.param((1, 2), 1, 0, id="1x2"),
+ pytest.param((1, 4), 1, 0, id="1x4"),
+ pytest.param((1, 8), 1, 0, id="1x8"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "masked",
+ [
+ pytest.param(True, id="masked"),
+ pytest.param(False, id="unmasked"),
+ ],
+)
+def test_qwen25vl_text_encoder(
+ *, mesh_device: ttnn.MeshDevice, batch_size: int, skip_layers: int, masked: bool
+) -> None:
+ torch.manual_seed(0)
+
+ sequence_length = 512
+ tp_axis = 1
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=mesh_device.shape[tp_axis], mesh_axis=tp_axis),
+ )
+
+ torch_model = transformers.Qwen2_5_VLForConditionalGeneration.from_pretrained(
+ "Qwen/Qwen-Image", subfolder="text_encoder"
+ )
+ torch_text_model = torch_model.model.language_model
+
+ mid = len(torch_text_model.layers) // 2
+ del torch_text_model.layers[mid - skip_layers // 2 : mid - (-skip_layers // 2)]
+
+ text_config = _text_config(torch_model.config)
+ rope_theta, mrope_section = _rope_params(text_config)
+
+ model = Qwen25VlTextEncoder(
+ vocab_size=text_config.vocab_size,
+ hidden_size=text_config.hidden_size,
+ intermediate_size=text_config.intermediate_size,
+ hidden_act=text_config.hidden_act,
+ num_hidden_layers=text_config.num_hidden_layers - skip_layers,
+ num_attention_heads=text_config.num_attention_heads,
+ num_key_value_heads=text_config.num_key_value_heads,
+ rms_norm_eps=text_config.rms_norm_eps,
+ rope_theta=rope_theta,
+ mrope_section=mrope_section,
+ device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+ model.load_torch_state_dict(torch_text_model.state_dict())
+
+ tokens = torch.randint(0, text_config.vocab_size, [batch_size, sequence_length])
+ m = torch.randint(0, sequence_length + 1, [batch_size])
+ attention_mask = torch.arange(sequence_length) < m.unsqueeze(1) if masked else None
+ cos, sin = model.create_rope_tensors(batch_size, sequence_length, attention_mask)
+
+ tt_tokens = tensor.from_torch(tokens, device=mesh_device, dtype=ttnn.uint32)
+ tt_attention_mask = tensor.from_torch(attention_mask, device=mesh_device) if attention_mask is not None else None
+ tt_pos_embeds_cos = tensor.from_torch(cos, device=mesh_device)
+ tt_pos_embeds_sin = tensor.from_torch(sin, device=mesh_device)
+
+ logger.info("running ttnn model...")
+ tt_hidden_states = model.forward(
+ tt_tokens,
+ attention_mask=tt_attention_mask,
+ pos_embeds=(tt_pos_embeds_cos, tt_pos_embeds_sin),
+ )
+ tt_prompt_embeds = tt_hidden_states[-1]
+ tt_prompt_embeds_torch = tensor.to_torch(tt_prompt_embeds)
+
+ logger.info("running torch model...")
+ with torch.no_grad():
+ out = torch_model.forward(
+ tokens,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ )
+ prompt_embeds = out.hidden_states[-1]
+
+ prompt_embeds = _real_tokens_only(prompt_embeds, attention_mask)
+ tt_prompt_embeds_torch = _real_tokens_only(tt_prompt_embeds_torch, attention_mask)
+ if masked:
+ assert_quality(prompt_embeds, tt_prompt_embeds_torch, pcc=0.994, relative_rmse=0.09)
+ else:
+ assert_quality(prompt_embeds, tt_prompt_embeds_torch, pcc=0.994, relative_rmse=0.09)
+
+
+@pytest.mark.parametrize(
+ "mesh_device , submesh_shape",
+ [[(2, 2), (1, 2)], [(2, 4), (1, 4)], [(4, 8), (1, 4)]],
+ ids=["2x2_1x2", "2x4_1x4", "4x8_1x4"],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "prompts",
+ [
+ [
+ "",
+ "Neon-lit cyberpunk alley, rain-soaked, cinematic wide shot",
+ ],
+ ],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{**line_params_req_exact_devices, "trace_region_size": 31000000}],
+ ids=["line"],
+ indirect=True,
+)
+def test_qwen25vl_encoder_pair(
+ *, mesh_device: ttnn.MeshDevice, submesh_shape: tuple[int, int], prompts: list[str]
+) -> None:
+ # There is a bug in the HF implementation where the prompt_embeds_mask is incorrectly repeated
+ # if num_images_per_prompt != 1.
+ # https://github.com/huggingface/diffusers/blob/v0.35.2/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py#L262
+ # is
+ # prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+ # but should be
+ # prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt)
+ num_images_per_prompt = 1
+ submesh_device = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ checkpoint = "Qwen/Qwen-Image"
+
+ torch_pipeline = diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.from_pretrained(checkpoint)
+
+ template = torch_pipeline.prompt_template_encode
+ start_idx = torch_pipeline.prompt_template_encode_start_idx
+ sequence_length = 512
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=submesh_shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(mesh_device=submesh_device, num_links=1, topology=ttnn.Topology.Linear)
+
+ tt_encoder_pair = Qwen25VlTokenizerEncoderPair(
+ checkpoint,
+ tokenizer_subfolder="tokenizer",
+ encoder_subfolder="text_encoder",
+ use_torch=False,
+ device=submesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+
+ logger.info("running torch model...")
+ with torch.no_grad():
+ embeds, mask = torch_pipeline.encode_prompt(
+ prompts,
+ num_images_per_prompt=num_images_per_prompt,
+ max_sequence_length=sequence_length,
+ )
+ embeds = torch.nn.functional.pad(embeds, [0, 0, 0, sequence_length - embeds.shape[1]], value=0)
+ mask = torch.nn.functional.pad(mask, [0, sequence_length - mask.shape[1]], value=0)
+
+ logger.info("running TT model...")
+ formatted_prompts = [template.format(e) for e in prompts]
+ tt_embeds, tt_mask = tt_encoder_pair.encode(
+ formatted_prompts,
+ num_images_per_prompt=num_images_per_prompt,
+ sequence_length=sequence_length + start_idx,
+ )
+ tt_embeds = tt_embeds[:, start_idx:]
+ tt_mask = tt_mask[:, start_idx:]
+ tt_embeds *= tt_mask.unsqueeze(-1)
+
+ assert torch.allclose(mask, tt_mask)
+ assert_quality(embeds, tt_embeds, pcc=0.988, relative_rmse=0.15)
diff --git a/code/models/tt_dit/tests/encoders/qwen3vl/common.py b/code/models/tt_dit/tests/encoders/qwen3vl/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..11be16895c62d3bc13e569631404882889bd1ef5
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen3vl/common.py
@@ -0,0 +1,134 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Shared helpers and constants for the Qwen3-VL encoder tests.
+
+import contextlib
+
+import pytest
+import torch
+import transformers
+
+import ttnn
+
+from ....encoders.qwen3vl.model_qwen3vl import Qwen3VlTextEncoder
+from ....parallel.config import EncoderParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils.tensor import bf16_tensor
+
+# FABRIC_1D on a 1x1 mesh has no ethernet partner and times out router init, so fabric is per-config, CCL-only.
+L1_SMALL = 32768
+FABRIC = {"fabric_config": ttnn.FabricConfig.FABRIC_1D, "l1_small_size": L1_SMALL}
+NO_FABRIC = {"l1_small_size": L1_SMALL}
+
+# MiniMax-H3's Qwen3-VL vision tower (also Qwen3-VL-32B's).
+HIDDEN_SIZE = 1152
+NUM_HEADS = 16
+HEAD_DIM = HIDDEN_SIZE // NUM_HEADS # 72 -- deliberately not tile-aligned
+INTERMEDIATE_SIZE = 4304
+SPATIAL_MERGE_SIZE = 2
+OUT_HIDDEN_SIZE = 5120
+NUM_POSITION_EMBEDDINGS = 2304
+NORM_EPS = 1e-6
+HIDDEN_ACT = "gelu_pytorch_tanh"
+
+
+def vision_config(depth, **overrides):
+ kwargs = dict(
+ depth=depth,
+ hidden_size=HIDDEN_SIZE,
+ num_heads=NUM_HEADS,
+ intermediate_size=INTERMEDIATE_SIZE,
+ in_channels=3,
+ patch_size=16,
+ temporal_patch_size=2,
+ spatial_merge_size=SPATIAL_MERGE_SIZE,
+ num_position_embeddings=NUM_POSITION_EMBEDDINGS,
+ out_hidden_size=OUT_HIDDEN_SIZE,
+ hidden_act=HIDDEN_ACT,
+ deepstack_visual_indexes=[0],
+ initializer_range=0.02,
+ )
+ kwargs.update(overrides)
+ return transformers.Qwen3VLVisionConfig(**kwargs)
+
+
+VISION_MESH = [
+ pytest.param((1, 1), (1, 1), None, None, 1, NO_FABRIC, id="single"),
+ pytest.param((8, 4), (8, 4), 0, None, 2, FABRIC, id="tp8_sp1"),
+ pytest.param((8, 4), (8, 4), 0, 1, 2, FABRIC, id="tp8_sp4"),
+]
+VISION_PARAMS = pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "tp_axis", "sp_axis", "num_links", "device_params"),
+ VISION_MESH,
+ indirect=["mesh_device", "device_params"],
+)
+
+
+def resolve_parallel(submesh, tp_axis, sp_axis, num_links):
+ if tp_axis is None and sp_axis is None:
+ return None, None
+ shape = tuple(submesh.shape)
+ cfg = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=shape[tp_axis] if tp_axis is not None else 1, mesh_axis=tp_axis),
+ sequence_parallel=(ParallelFactor(factor=shape[sp_axis], mesh_axis=sp_axis) if sp_axis is not None else None),
+ )
+ return cfg, CCLManager(submesh, num_links=num_links, topology=ttnn.Topology.Linear)
+
+
+def skip_if_sp_misaligned(total, submesh, sp_axis):
+ """Ring SDPA needs N_local_q % 32 == 0; misalignment is model geometry, not a port bug, so skip."""
+ if sp_axis is None:
+ return
+ sp = tuple(submesh.shape)[sp_axis]
+ if total % (sp * 32) != 0:
+ pytest.skip(f"{total} patches do not divide into {sp} tile-aligned shards (needs a multiple of {sp * 32})")
+
+
+def sp_shard(x, submesh, sp_axis):
+ if sp_axis is None:
+ return bf16_tensor(x, device=submesh)
+ return bf16_tensor(x, device=submesh, mesh_axis=sp_axis, shard_dim=0)
+
+
+@contextlib.contextmanager
+def capture_layer_outputs(lm, layers):
+ caps: dict[int, torch.Tensor] = {}
+ handles = [
+ lm.layers[i].register_forward_hook(
+ lambda m, i_, o, i=i: caps.__setitem__(i, (o[0] if isinstance(o, tuple) else o).detach())
+ )
+ for i in layers
+ ]
+ try:
+ yield caps
+ finally:
+ for h in handles:
+ h.remove()
+
+
+def hf_rope_params(cfg):
+ """transformers>=5 has no top-level cfg.rope_theta, so it must not be an (eagerly evaluated) dict.get default."""
+ rope_params = getattr(cfg, "rope_parameters", None) or cfg.rope_scaling
+ mrope_section = rope_params["mrope_section"]
+ rope_theta = rope_params["rope_theta"] if "rope_theta" in rope_params else cfg.rope_theta
+ return rope_theta, mrope_section
+
+
+def encoder_from_hf_config(cfg, **overrides):
+ rope_theta, mrope_section = hf_rope_params(cfg)
+ kwargs = dict(
+ vocab_size=cfg.vocab_size,
+ hidden_size=cfg.hidden_size,
+ intermediate_size=cfg.intermediate_size,
+ hidden_act="silu",
+ num_hidden_layers=cfg.num_hidden_layers,
+ num_attention_heads=cfg.num_attention_heads,
+ num_key_value_heads=cfg.num_key_value_heads,
+ rms_norm_eps=cfg.rms_norm_eps,
+ rope_theta=rope_theta,
+ mrope_section=mrope_section,
+ )
+ kwargs.update(overrides)
+ return Qwen3VlTextEncoder(**kwargs)
diff --git a/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl.py b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd9989a09f4c06455be76d5aeb0f6334997bbb8d
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl.py
@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Qwen3-VL-8B text encoder for Ideogram 4.0: qwen25vl port + per-head QK-RMSNorm,
+# no qkv bias, and the 13-layer feature tap. HF reference is built from config
+# (no 8B pull); "real" overlays the shipped Ideogram fp8 text_encoder weights.
+
+import os
+
+import pytest
+import torch
+import transformers
+from loguru import logger
+from safetensors.torch import load_file
+
+import ttnn
+
+from ....encoders.qwen3vl.model_qwen3vl import create_rope_tensors
+from ....parallel.config import EncoderParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....reference.ideogram4.constants import QWEN3_VL_ACTIVATION_LAYERS
+from ....reference.ideogram4.dequant import dequant_fp8_state_dict
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.tensor import bf16_tensor
+from .common import capture_layer_outputs, encoder_from_hf_config, hf_rope_params
+
+REPO = "Qwen/Qwen3-VL-8B-Instruct"
+FP8 = os.environ.get("IDEOGRAM4_WEIGHTS")
+_NEEDS_WEIGHTS = pytest.mark.skipif(not FP8, reason="IDEOGRAM4_WEIGHTS not set (gated fp8 checkpoint)")
+
+
+def _reference_lm(weights: str):
+ cfg = transformers.AutoConfig.from_pretrained(REPO)
+ hf = transformers.AutoModel.from_config(cfg).to(torch.bfloat16)
+ lm = hf.language_model if hasattr(hf, "language_model") else hf.model.language_model
+ if weights == "real":
+ sd = dequant_fp8_state_dict(load_file(f"{FP8}/text_encoder/model.safetensors"))
+ sd = {k[len("language_model.") :]: v for k, v in sd.items() if k.startswith("language_model.")}
+ incompat = lm.load_state_dict(sd, strict=False)
+ # empty missing/unexpected proves the weights landed; else both sides share the random init and PCC is vacuous
+ assert not incompat.missing_keys and not incompat.unexpected_keys, (
+ f"real Qwen3-VL load key mismatch: missing={incompat.missing_keys[:5]} "
+ f"unexpected={incompat.unexpected_keys[:5]}"
+ )
+ return lm.eval()
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "tp_axis"),
+ [
+ pytest.param((2, 4), (2, 4), 1, id="tp4_fsdp2"),
+ pytest.param((2, 4), (2, 4), 0, id="tp2_fsdp4"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "l1_small_size": 32768}], indirect=True
+)
+@pytest.mark.parametrize(
+ "weights",
+ [pytest.param("random", id="random"), pytest.param("real", id="real", marks=_NEEDS_WEIGHTS)],
+)
+@pytest.mark.parametrize("masked", [pytest.param(False, id="nomask"), pytest.param(True, id="masked")])
+@pytest.mark.parametrize("seq_len", [128])
+def test_qwen3vl_text_encoder(
+ *, mesh_device: ttnn.MeshDevice, submesh_shape, tp_axis, weights: str, masked: bool, seq_len: int
+) -> None:
+ torch.manual_seed(0)
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ tp_factor = tuple(submesh.shape)[tp_axis]
+
+ lm = _reference_lm(weights)
+ cfg = lm.config
+ head_dim = cfg.hidden_size // cfg.num_attention_heads
+ rope_theta, mrope_section = hf_rope_params(cfg)
+
+ ids = torch.randint(0, cfg.vocab_size, (1, seq_len))
+ with capture_layer_outputs(lm, QWEN3_VL_ACTIVATION_LAYERS) as caps:
+ with torch.no_grad():
+ lm(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False)
+ golden = [caps[i].float() for i in QWEN3_VL_ACTIVATION_LAYERS]
+
+ enc = encoder_from_hf_config(
+ cfg,
+ activation_layers=QWEN3_VL_ACTIVATION_LAYERS,
+ device=submesh,
+ parallel_config=EncoderParallelConfig(tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis)),
+ ccl_manager=CCLManager(submesh, num_links=1, topology=ttnn.Topology.Linear),
+ is_fsdp=True,
+ )
+ enc.load_torch_state_dict(lm.state_dict())
+
+ cos, sin = create_rope_tensors(1, seq_len, None, head_dim, rope_theta, mrope_section)
+ tt_ids = ttnn.from_torch(ids, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=submesh)
+ attn_mask = tensor.from_torch(torch.ones(1, seq_len, dtype=torch.bool), device=submesh) if masked else None
+ tt_caps = enc.forward(
+ tt_ids,
+ attention_mask=attn_mask,
+ pos_embeds=(bf16_tensor(cos, device=submesh), bf16_tensor(sin, device=submesh)),
+ )
+
+ pcc = 0.99 if weights == "real" else 0.98 # random weights over 36 layers accumulate more bf16 error
+ for layer_idx, g, tt_t in zip(QWEN3_VL_ACTIVATION_LAYERS, golden, tt_caps):
+ logger.info(f"qwen3vl [{weights}] TP={tp_factor} layer {layer_idx}:")
+ assert_quality(g, tensor.to_torch(tt_t, mesh_axes=[None, None, None]), pcc=pcc)
diff --git a/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_decoder_block.py b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_decoder_block.py
new file mode 100644
index 0000000000000000000000000000000000000000..e487baf3c3af87ae04800b022dfb081664a43e53
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_decoder_block.py
@@ -0,0 +1,191 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# One Qwen3-VL text decoder layer -- plus its attention and MLP in isolation -- at
+# the MiniMax-H3 conditioner dimensions, sized so a Tracy capture measures the
+# layer's arithmetic instead of a 32B whole-stack weight load.
+
+import pytest
+import torch
+import transformers
+from loguru import logger
+
+import ttnn
+
+from ....encoders.qwen3vl.model_qwen3vl import (
+ Qwen3VlAttention,
+ Qwen3VlContext,
+ Qwen3VlDecoderLayer,
+ Qwen3VlMlp,
+ create_rope_tensors,
+)
+from ....parallel.manager import CCLManager
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.tensor import bf16_tensor
+from .common import FABRIC, NO_FABRIC
+
+# MiniMax-H3's `text_encoder/config.json` (text_config).
+HIDDEN_SIZE = 5120
+INTERMEDIATE_SIZE = 25600
+NUM_HEADS = 64
+NUM_KV_HEADS = 8
+HEAD_DIM = 128 # explicit in the checkpoint; 5120 // 64 would give 80
+ROPE_THETA = 5_000_000.0
+MROPE_SECTION = [24, 20, 20] # sums to HEAD_DIM // 2
+NORM_EPS = 1e-6
+HIDDEN_ACT = "silu"
+
+VOCAB_SIZE = 256 # the layer never sees the embedding table; tiny keeps the fp32 reference small
+SEQ_LEN = 128 # == SEQ_BUCKET_SIZE, so the encoder's prompt bucketing does not pad
+
+_MESH = [
+ pytest.param((1, 1), (1, 1), None, 1, NO_FABRIC, id="single"),
+ pytest.param((4, 8), (4, 8), 1, 2, FABRIC, id="tp8_axis1"),
+ pytest.param((8, 4), (8, 4), 0, 2, FABRIC, id="tp8_axis0"),
+]
+_MESH_REPRESENTATIVE = [_MESH[1]] # tp8_axis1: the (4, 8) mesh the MiniMax-H3 pipeline actually deploys on
+
+_PARAMS = pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "tp_axis", "num_links", "device_params"),
+ _MESH,
+ indirect=["mesh_device", "device_params"],
+)
+_PARAMS_REPRESENTATIVE = pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "tp_axis", "num_links", "device_params"),
+ _MESH_REPRESENTATIVE,
+ indirect=["mesh_device", "device_params"],
+)
+
+
+def _config():
+ return transformers.Qwen3VLTextConfig(
+ vocab_size=VOCAB_SIZE,
+ hidden_size=HIDDEN_SIZE,
+ intermediate_size=INTERMEDIATE_SIZE,
+ num_hidden_layers=1,
+ num_attention_heads=NUM_HEADS,
+ num_key_value_heads=NUM_KV_HEADS,
+ head_dim=HEAD_DIM,
+ rms_norm_eps=NORM_EPS,
+ hidden_act=HIDDEN_ACT,
+ rope_theta=ROPE_THETA,
+ rope_scaling={"rope_type": "default", "mrope_section": MROPE_SECTION, "mrope_interleaved": True},
+ )
+
+
+@pytest.fixture(scope="module")
+def golden():
+ # fp32 reference: casting it to bf16 would degrade the rotary's inv_freq and poison the golden
+ torch.manual_seed(0)
+ cfg = _config()
+ assert sum(MROPE_SECTION) == HEAD_DIM // 2, f"mrope_section sums to {sum(MROPE_SECTION)}, want {HEAD_DIM // 2}"
+
+ lm = transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextModel._from_config(cfg).eval()
+ layer = lm.layers[0]
+
+ cap: dict[str, torch.Tensor | None] = {}
+
+ def grab(key):
+ def hook(_m, args, kwargs, out):
+ inp = args[0] if args else kwargs.get("hidden_states")
+ assert inp is not None, f"{key}: could not find hidden states in args/kwargs"
+ cap[f"{key}_in"] = inp.detach()
+ cap[f"{key}_out"] = (out[0] if isinstance(out, tuple) else out).detach()
+ if key == "layer":
+ cap["mask"] = kwargs.get("attention_mask")
+ cap["pos"] = kwargs.get("position_embeddings")
+
+ return hook
+
+ handles = [
+ layer.register_forward_hook(grab("layer"), with_kwargs=True),
+ layer.self_attn.register_forward_hook(grab("attn"), with_kwargs=True),
+ layer.mlp.register_forward_hook(grab("mlp"), with_kwargs=True),
+ ]
+ ids = torch.randint(0, VOCAB_SIZE, (1, SEQ_LEN))
+ with torch.no_grad():
+ lm(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False)
+ for h in handles:
+ h.remove()
+
+ # a future transformers materializing a mask here would silently break causal parity; fail loudly
+ assert cap["mask"] is None, (
+ f"reference passed an explicit attention_mask ({type(cap['mask'])}) to the decoder layer; "
+ "the is_causal=True path on our side no longer provably matches it"
+ )
+ return {"state": layer.state_dict(), **cap}
+
+
+def _ctx(mesh_device, submesh_shape, tp_axis, num_links):
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ ccl = CCLManager(submesh, num_links=num_links, topology=ttnn.Topology.Linear) if tp_axis is not None else None
+ return submesh, Qwen3VlContext(device=submesh, tp_axis=tp_axis, ccl_manager=ccl, fsdp_mesh_axis=None)
+
+
+def _rope(submesh):
+ # interleaved stays False (as in production): the layouts coincide while all three MRoPE axes share the position
+ cos, sin = create_rope_tensors(1, SEQ_LEN, None, HEAD_DIM, ROPE_THETA, MROPE_SECTION)
+ return bf16_tensor(cos, device=submesh), bf16_tensor(sin, device=submesh)
+
+
+@_PARAMS
+def test_decoder_block_on_device(golden, mesh_device, submesh_shape, tp_axis, num_links):
+ submesh, ctx = _ctx(mesh_device, submesh_shape, tp_axis, num_links)
+
+ block = Qwen3VlDecoderLayer(
+ hidden_size=HIDDEN_SIZE,
+ intermediate_size=INTERMEDIATE_SIZE,
+ hidden_act=HIDDEN_ACT,
+ num_attention_heads=NUM_HEADS,
+ num_key_value_heads=NUM_KV_HEADS,
+ head_dim=HEAD_DIM,
+ rms_norm_eps=NORM_EPS,
+ ctx=ctx,
+ )
+ block.load_torch_state_dict(golden["state"])
+
+ out = block.forward(
+ bf16_tensor(golden["layer_in"], device=submesh),
+ attention_bias=None,
+ pos_embeds=_rope(submesh),
+ )
+
+ tp_factor = tuple(submesh.shape)[tp_axis] if tp_axis is not None else 1
+ logger.info(f"qwen3vl decoder layer TP={tp_factor} (axis {tp_axis}):")
+ assert_quality(golden["layer_out"].float(), tensor.to_torch(out, mesh_axes=[None, None, None]), pcc=0.99)
+
+
+@_PARAMS_REPRESENTATIVE
+def test_decoder_attention_on_device(golden, mesh_device, submesh_shape, tp_axis, num_links):
+ submesh, ctx = _ctx(mesh_device, submesh_shape, tp_axis, num_links)
+
+ attn = Qwen3VlAttention(
+ hidden_size=HIDDEN_SIZE,
+ num_heads=NUM_HEADS,
+ num_key_value_heads=NUM_KV_HEADS,
+ head_dim=HEAD_DIM,
+ rms_norm_eps=NORM_EPS,
+ ctx=ctx,
+ )
+ prefix = "self_attn."
+ attn.load_torch_state_dict({k[len(prefix) :]: v for k, v in golden["state"].items() if k.startswith(prefix)})
+
+ out = attn.forward(
+ bf16_tensor(golden["attn_in"], device=submesh),
+ attention_bias=None,
+ pos_embeds=_rope(submesh),
+ )
+ assert_quality(golden["attn_out"].float(), tensor.to_torch(out, mesh_axes=[None, None, None]), pcc=0.99)
+
+
+@_PARAMS_REPRESENTATIVE
+def test_decoder_mlp_on_device(golden, mesh_device, submesh_shape, tp_axis, num_links):
+ submesh, ctx = _ctx(mesh_device, submesh_shape, tp_axis, num_links)
+
+ mlp = Qwen3VlMlp(hidden_size=HIDDEN_SIZE, intermediate_size=INTERMEDIATE_SIZE, hidden_act=HIDDEN_ACT, ctx=ctx)
+ mlp.load_torch_state_dict({k[len("mlp.") :]: v for k, v in golden["state"].items() if k.startswith("mlp.")})
+
+ out = mlp.forward(bf16_tensor(golden["mlp_in"], device=submesh))
+ assert_quality(golden["mlp_out"].float(), tensor.to_torch(out, mesh_axes=[None, None, None]), pcc=0.99)
diff --git a/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_mrope.py b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_mrope.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e33bbbc7893fd614c5cc7ecada4d8d50ea30d3e
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_mrope.py
@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Qwen3-VL M-RoPE: the multimodal position grid and the interleaved layout,
+# against the HF reference. Pure torch, no device. Text-only prompts make the
+# chunked and interleaved layouts coincide; vision runs make the layout observable.
+
+import pytest
+import torch
+import transformers
+
+from ....encoders.qwen3vl.model_qwen3vl import create_rope_tensors, mrope_position_ids, vision_position_ids
+
+# MiniMax-H3's conditioner: head_dim 128 (not hidden_size // num_heads), interleaved M-RoPE.
+HEAD_DIM = 128
+MROPE_SECTION = [24, 20, 20]
+ROPE_THETA = 5_000_000.0
+SPATIAL_MERGE_SIZE = 2
+
+
+@pytest.fixture(scope="module")
+def reference():
+ config = transformers.Qwen3VLConfig(
+ text_config={
+ "num_hidden_layers": 2,
+ "hidden_size": 64,
+ "intermediate_size": 128,
+ "num_attention_heads": 4,
+ "num_key_value_heads": 2,
+ "head_dim": 16,
+ },
+ vision_config={
+ "depth": 2,
+ "hidden_size": 32,
+ "num_heads": 2,
+ "out_hidden_size": 64,
+ "spatial_merge_size": SPATIAL_MERGE_SIZE,
+ },
+ )
+ return transformers.AutoModel.from_config(config).eval()
+
+
+_TIMESTAMP_TOKENS = 2 # frames are separated by timestamp text in the presentation, so each frame is its own run
+
+
+def _prompt(*runs):
+ """(mm_token_type_ids, image_grid_thw, video_grid_thw) from ("text", n) / ("image"|"video", (t, h, w)) runs."""
+ type_ids, images, videos = [], [], []
+ for kind, spec in runs:
+ t, h, w = spec if kind != "text" else (0, 0, 0)
+ if kind == "text":
+ type_ids += [0] * spec
+ elif kind == "image":
+ type_ids += [1] * (t * h * w // (SPATIAL_MERGE_SIZE**2))
+ images.append([t, h, w])
+ else:
+ videos.append([t, h, w])
+ for _ in range(t):
+ type_ids += [0] * _TIMESTAMP_TOKENS
+ type_ids += [2] * (h * w // (SPATIAL_MERGE_SIZE**2))
+ return (
+ torch.tensor([type_ids], dtype=torch.long),
+ torch.tensor(images, dtype=torch.long) if images else None,
+ torch.tensor(videos, dtype=torch.long) if videos else None,
+ )
+
+
+PROMPTS = {
+ "text_only": (("text", 24),),
+ "one_image": (("text", 6), ("image", (1, 4, 6)), ("text", 5)),
+ "two_images": (("text", 3), ("image", (1, 4, 4)), ("text", 2), ("image", (1, 6, 2)), ("text", 4)),
+ "image_first": (("image", (1, 2, 2)), ("text", 8)),
+ "one_video": (("text", 4), ("video", (2, 4, 4)), ("text", 3)),
+ "video_and_image": (("text", 2), ("video", (2, 2, 4)), ("text", 3), ("image", (1, 4, 4)), ("text", 2)),
+ "video_3_frames": (("text", 2), ("video", (3, 4, 2)), ("text", 2)),
+ # production fl2va working point; rope tables only hit rounding boundaries at production length
+ "keyframe_768x1344": (("text", 5), ("image", (1, 48, 84)), ("text", 41)),
+ "two_keyframes_768x1344": (
+ ("text", 5),
+ ("image", (1, 48, 84)),
+ ("text", 5),
+ ("image", (1, 48, 84)),
+ ("text", 41),
+ ),
+}
+
+
+@pytest.mark.parametrize("name", list(PROMPTS))
+def test_mrope_position_ids_matches_reference(reference, name):
+ type_ids, image_grid, video_grid = _prompt(*PROMPTS[name])
+ input_ids = torch.zeros_like(type_ids)
+
+ expected, _ = reference.get_rope_index(
+ input_ids,
+ mm_token_type_ids=type_ids,
+ image_grid_thw=None if image_grid is None else image_grid.clone(),
+ video_grid_thw=None if video_grid is None else video_grid.clone(),
+ )
+ actual = mrope_position_ids(
+ type_ids,
+ image_grid_thw=image_grid,
+ video_grid_thw=video_grid,
+ spatial_merge_size=SPATIAL_MERGE_SIZE,
+ )
+
+ assert actual.shape == expected.shape, f"{tuple(actual.shape)} != {tuple(expected.shape)}"
+ assert torch.equal(actual, expected), (
+ f"{name}: position grid differs at "
+ f"{(actual != expected).nonzero()[:5].tolist()}\n ours={actual[:, 0, :12].tolist()}\n"
+ f" ref ={expected[:, 0, :12].tolist()}"
+ )
+
+ grids = ([] if image_grid is None else list(image_grid)) + ([] if video_grid is None else list(video_grid))
+ for grid in grids:
+ for start in (0, 7):
+ block_expected = reference.get_vision_position_ids(start, grid, 1, SPATIAL_MERGE_SIZE, device=None)
+ block_actual = vision_position_ids(start, grid, spatial_merge_size=SPATIAL_MERGE_SIZE)
+ assert torch.equal(block_actual, block_expected), (
+ f"{name} grid {grid.tolist()} start {start}: vision block grid differs\n"
+ f" ours={block_actual[:, :12].tolist()}\n ref ={block_expected[:, :12].tolist()}"
+ )
+
+
+def test_text_only_layouts_are_identical():
+ seq = 64
+ chunked = create_rope_tensors(1, seq, None, HEAD_DIM, ROPE_THETA, MROPE_SECTION, interleaved=False)
+ interleaved = create_rope_tensors(1, seq, None, HEAD_DIM, ROPE_THETA, MROPE_SECTION, interleaved=True)
+ for a, b, which in zip(chunked, interleaved, ("cos", "sin")):
+ assert torch.equal(a, b), f"text-only {which} differs between layouts"
+
+
+@pytest.mark.parametrize("name", list(PROMPTS))
+def test_interleaved_rope_tensors_match_reference(name):
+ """Toleranced, not bitwise: theta**-x vs the reference's 1/theta**x inv_freq differ by one fp32 ulp."""
+ text_config = transformers.Qwen3VLTextConfig(
+ hidden_size=5120,
+ num_attention_heads=64,
+ head_dim=HEAD_DIM,
+ rope_parameters={
+ "rope_type": "default",
+ "rope_theta": ROPE_THETA,
+ "mrope_section": MROPE_SECTION,
+ "mrope_interleaved": True,
+ },
+ )
+ rotary = transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRotaryEmbedding(text_config)
+
+ type_ids, image_grid, video_grid = _prompt(*PROMPTS[name])
+ seq = type_ids.shape[1]
+ position_ids = mrope_position_ids(
+ type_ids, image_grid_thw=image_grid, video_grid_thw=video_grid, spatial_merge_size=SPATIAL_MERGE_SIZE
+ )
+
+ expected_cos, expected_sin = rotary(torch.zeros(1, seq, text_config.hidden_size), position_ids)
+ cos, sin = create_rope_tensors(
+ 1, seq, None, HEAD_DIM, ROPE_THETA, MROPE_SECTION, position_ids=position_ids, interleaved=True
+ )
+
+ for actual, expected, which in ((cos, expected_cos, "cos"), (sin, expected_sin, "sin")):
+ actual = actual.squeeze(1)
+ assert actual.shape == expected.shape, f"{which}: {tuple(actual.shape)} != {tuple(expected.shape)}"
+ max_diff = (actual - expected).abs().max().item()
+ assert max_diff < 1e-5, f"{name} {which}: max abs diff {max_diff:.2e} exceeds the 1-ulp inv_freq budget"
+
+
+def test_position_ids_default_is_the_shared_token_index():
+ seq = 32
+ implicit = create_rope_tensors(1, seq, None, HEAD_DIM, ROPE_THETA, MROPE_SECTION)
+ explicit = create_rope_tensors(
+ 1,
+ seq,
+ None,
+ HEAD_DIM,
+ ROPE_THETA,
+ MROPE_SECTION,
+ position_ids=torch.arange(seq).view(1, 1, -1).expand(3, 1, -1),
+ )
+ for a, b, which in zip(implicit, explicit, ("cos", "sin")):
+ assert torch.equal(a, b), f"{which} differs between the implicit and explicit text grid"
+
+
+def test_missing_grid_is_an_error(expect_error):
+ type_ids, image_grid, _ = _prompt(*PROMPTS["one_image"])
+ del image_grid
+ with expect_error(ValueError, "no matching grid"):
+ mrope_position_ids(type_ids, spatial_merge_size=SPATIAL_MERGE_SIZE)
+
+
+def test_adjacent_same_modality_blocks_are_unsupported(expect_error):
+ """The reference rejects this format too; do not "fix" ours to accept an input it cannot express."""
+ type_ids = torch.tensor([[0, 0] + [1] * 4 + [1] * 4 + [0, 0]], dtype=torch.long)
+ grids = torch.tensor([[1, 4, 4], [1, 4, 4]], dtype=torch.long)
+ with expect_error(RuntimeError, "must match the existing size"):
+ mrope_position_ids(type_ids, image_grid_thw=grids, spatial_merge_size=SPATIAL_MERGE_SIZE)
diff --git a/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_block.py b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_block.py
new file mode 100644
index 0000000000000000000000000000000000000000..6776134986fefa1a78a4a3594421d4d6d25504b9
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_block.py
@@ -0,0 +1,115 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Qwen3-VL vision transformer block on device against the HF reference; the
+# aspect sweep, capacity grids and perf mode live in test_qwen3vl_vision_tower.py.
+
+import pytest
+import torch
+import transformers
+
+import ttnn
+
+from ....encoders.qwen3vl.vision_qwen3vl import (
+ Qwen3VlVisionBlock,
+ resolve_vision_parallel,
+ vision_cu_seqlens,
+ vision_rope_tensors,
+)
+from ....utils import tensor
+from ....utils.check import assert_quality
+from .common import (
+ HEAD_DIM,
+ HIDDEN_ACT,
+ HIDDEN_SIZE,
+ INTERMEDIATE_SIZE,
+ NORM_EPS,
+ NUM_HEADS,
+ SPATIAL_MERGE_SIZE,
+ VISION_PARAMS,
+ resolve_parallel,
+ skip_if_sp_misaligned,
+ sp_shard,
+ vision_config,
+)
+
+PADDED_HEAD_DIM = 96 # head_dim 72 is not tile-aligned; q/k/v/o zero-pad to 96, scale must stay 72**-0.5
+
+GRIDS = {
+ "canvas_1to1": [[1, 48, 48]], # 2304 rows, 1 block -- cheapest real shape
+ "image_and_video": [[1, 128, 128], [3, 48, 48]], # 23296, 4 blocks, both sizing rules
+}
+
+
+@pytest.fixture(scope="module")
+def reference():
+ torch.manual_seed(0)
+ return transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionModel._from_config(vision_config(depth=1)).eval()
+
+
+def _golden_pos_embeds(grid):
+ from transformers.vision_utils import get_vision_position_ids
+
+ position_ids = get_vision_position_ids(grid, SPATIAL_MERGE_SIZE)
+ rot = transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionRotaryEmbedding(HEAD_DIM // 2)
+ freqs = rot(position_ids)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ return emb.cos(), emb.sin()
+
+
+def _parallel(submesh, tp_axis, sp_axis, num_links):
+ cfg, ccl = resolve_parallel(submesh, tp_axis, sp_axis, num_links)
+ if cfg is None:
+ return None
+ return resolve_vision_parallel(submesh, cfg, ccl)
+
+
+@VISION_PARAMS
+@pytest.mark.parametrize("name", list(GRIDS))
+def test_block_on_device(reference, mesh_device, submesh_shape, tp_axis, sp_axis, num_links, name):
+ """The whole pre-norm block; multi-block grids exercise the per-frame SDPA loop and the padded-head-dim scale."""
+ import math
+
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ grid = torch.tensor(GRIDS[name], dtype=torch.long)
+ total = sum(t * h * w for t, h, w in GRIDS[name])
+ skip_if_sp_misaligned(total, submesh, sp_axis)
+
+ assert HEAD_DIM == 72 and HEAD_DIM % 32 != 0, "the whole padding question presumes a misaligned 72"
+ assert math.ceil(HEAD_DIM / 32) * 32 == PADDED_HEAD_DIM
+ assert HEAD_DIM**-0.5 != pytest.approx(PADDED_HEAD_DIM**-0.5), "the two temperatures must differ"
+
+ cu_seqlens = vision_cu_seqlens(grid)
+ assert len(cu_seqlens) - 1 == int(grid[:, 0].sum()), f"expected one block per frame, got {cu_seqlens}"
+ assert cu_seqlens[0] == 0 and cu_seqlens[-1] == total, f"cu_seqlens must span [0, {total}]: {cu_seqlens}"
+
+ torch.manual_seed(0)
+ x = torch.randn(total, HIDDEN_SIZE)
+ cos, sin = _golden_pos_embeds(grid)
+ with torch.no_grad():
+ golden = reference.blocks[0](
+ x,
+ cu_seqlens=torch.tensor(list(cu_seqlens), dtype=torch.int32),
+ position_embeddings=(cos, sin),
+ ).float()
+
+ block = Qwen3VlVisionBlock(
+ hidden_size=HIDDEN_SIZE,
+ num_heads=NUM_HEADS,
+ intermediate_size=INTERMEDIATE_SIZE,
+ hidden_act=HIDDEN_ACT,
+ norm_eps=NORM_EPS,
+ mesh_device=submesh,
+ parallel=_parallel(submesh, tp_axis, sp_axis, num_links),
+ )
+ block.load_torch_state_dict(reference.blocks[0].state_dict())
+ tt_cos, tt_sin = vision_rope_tensors(grid, head_dim=HEAD_DIM, spatial_merge_size=SPATIAL_MERGE_SIZE)
+ out = block.forward(
+ sp_shard(x, submesh, sp_axis),
+ pos_embeds=(sp_shard(tt_cos, submesh, sp_axis), sp_shard(tt_sin, submesh, sp_axis)),
+ cu_seqlens=cu_seqlens,
+ )
+ actual = tensor.to_torch(out, mesh_axes=[sp_axis, None])
+
+ assert_quality(golden, actual, pcc=0.99)
diff --git a/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_injection.py b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_injection.py
new file mode 100644
index 0000000000000000000000000000000000000000..4aa7577394f5ee67da1b39d5c6af4fd569169435
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_injection.py
@@ -0,0 +1,242 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# How vision reaches the Qwen3-VL decoder: the tower's merged tokens REPLACE the
+# <|image_pad|> row embeddings before the block stack; its deepstack features are
+# ADDED to those same rows after each of the first len(deepstack) decoder layers.
+
+import pytest
+import torch
+import transformers
+from loguru import logger
+
+import ttnn
+
+from ....encoders.qwen3vl.model_qwen3vl import _scatter_rows, create_rope_tensors, vision_token_runs
+from ....parallel.config import EncoderParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils import tensor
+from ....utils.tensor import bf16_tensor
+from .common import encoder_from_hf_config
+
+IMAGE_TOKEN_ID = 151655 # <|image_pad|>
+HIDDEN = 128
+SEQ = 64
+
+
+def test_vision_token_runs_finds_one_run_per_image():
+ ids = torch.tensor([[1, 2, 3] + [IMAGE_TOKEN_ID] * 4 + [9, 9] + [IMAGE_TOKEN_ID] * 6 + [7]])
+ assert vision_token_runs(ids, IMAGE_TOKEN_ID) == [(3, 4), (9, 6)]
+
+
+def test_vision_token_runs_handles_the_edges():
+ assert vision_token_runs(torch.tensor([[IMAGE_TOKEN_ID] * 3 + [1, 2]]), IMAGE_TOKEN_ID) == [(0, 3)]
+ assert vision_token_runs(torch.tensor([[1, 2] + [IMAGE_TOKEN_ID] * 3]), IMAGE_TOKEN_ID) == [(2, 3)]
+ assert vision_token_runs(torch.tensor([[1, 2, 3]]), IMAGE_TOKEN_ID) == []
+
+
+def test_vision_token_runs_rejects_a_batch(expect_error):
+ """One request is one sequence; a batch would silently take only the first row."""
+ with expect_error(ValueError, "expected a single sequence"):
+ vision_token_runs(torch.zeros(2, 8, dtype=torch.long), IMAGE_TOKEN_ID)
+
+
+_MESH = [pytest.param((1, 1), (1, 1), id="single")]
+
+
+@pytest.mark.parametrize(("mesh_device", "submesh_shape"), _MESH, indirect=["mesh_device"])
+def test_scatter_rows_rejects_a_row_count_mismatch(mesh_device, submesh_shape, expect_error):
+ """Too few or too many value rows is a caller error, not a silent partial write."""
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ base = bf16_tensor(torch.zeros(1, SEQ, HIDDEN), device=submesh)
+ values = bf16_tensor(torch.zeros(1, 5, HIDDEN), device=submesh)
+ with expect_error(ValueError, "runs cover 16 rows but values has 5"):
+ _scatter_rows(base, values, [(8, 16)], add=False)
+
+
+def _reference_text_model(layers):
+ config = transformers.Qwen3VLTextConfig(
+ vocab_size=256,
+ hidden_size=HIDDEN,
+ intermediate_size=256,
+ num_hidden_layers=layers,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ head_dim=32,
+ rms_norm_eps=1e-6,
+ rope_parameters={
+ "rope_type": "default",
+ "rope_theta": 10000.0,
+ "mrope_section": [6, 5, 5],
+ "mrope_interleaved": True,
+ },
+ )
+ torch.manual_seed(0)
+ return transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextModel._from_config(config).eval()
+
+
+def _encoder(submesh, *, layers, activation_layers):
+ reference = _reference_text_model(layers)
+ enc = encoder_from_hf_config(
+ reference.config,
+ head_dim=32,
+ activation_layers=activation_layers,
+ device=submesh,
+ parallel_config=EncoderParallelConfig(tensor_parallel=ParallelFactor(factor=1, mesh_axis=0)),
+ ccl_manager=CCLManager(submesh, num_links=1, topology=ttnn.Topology.Linear),
+ )
+ enc.load_torch_state_dict(reference.state_dict())
+ return enc
+
+
+@pytest.mark.parametrize(("mesh_device", "submesh_shape"), _MESH, indirect=["mesh_device"])
+def test_text_only_path_is_unchanged(mesh_device, submesh_shape):
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ torch.manual_seed(0)
+ enc = _encoder(submesh, layers=4, activation_layers=(3,))
+ ids = ttnn.from_torch(
+ torch.randint(0, 256, (1, SEQ)), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=submesh
+ )
+ cos, sin = create_rope_tensors(1, SEQ, None, 32, 10000.0, [6, 5, 5])
+ pe = (bf16_tensor(cos, device=submesh), bf16_tensor(sin, device=submesh))
+
+ plain = enc.forward(ids, attention_mask=None, pos_embeds=pe)[0]
+ explicit_none = enc.forward(
+ ids, attention_mask=None, pos_embeds=pe, vision_embeds=None, vision_runs=None, deepstack_embeds=None
+ )[0]
+ a = tensor.to_torch(plain, mesh_axes=[None, None, None])
+ b = tensor.to_torch(explicit_none, mesh_axes=[None, None, None])
+ assert torch.equal(a, b), "explicitly passing None diverges from omitting the arguments"
+
+
+@pytest.mark.parametrize(("mesh_device", "submesh_shape"), _MESH, indirect=["mesh_device"])
+def test_vision_arguments_must_be_paired(mesh_device, submesh_shape, expect_error):
+ """Embeds without runs (or the reverse) is a caller error rather than a silent no-op."""
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ enc = _encoder(submesh, layers=2, activation_layers=(1,))
+ ids = ttnn.from_torch(
+ torch.zeros(1, SEQ, dtype=torch.long), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=submesh
+ )
+ cos, sin = create_rope_tensors(1, SEQ, None, 32, 10000.0, [6, 5, 5])
+ pe = (bf16_tensor(cos, device=submesh), bf16_tensor(sin, device=submesh))
+ embeds = bf16_tensor(torch.zeros(1, 16, HIDDEN), device=submesh)
+
+ with expect_error(ValueError, "must be passed together"):
+ enc.forward(ids, attention_mask=None, pos_embeds=pe, vision_embeds=embeds)
+ with expect_error(ValueError, "needs vision_runs"):
+ enc.forward(ids, attention_mask=None, pos_embeds=pe, deepstack_embeds=[embeds])
+
+
+@pytest.mark.parametrize(("mesh_device", "submesh_shape"), _MESH, indirect=["mesh_device"])
+def test_deepstack_is_applied_at_the_leading_layers(mesh_device, submesh_shape):
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ torch.manual_seed(0)
+ enc = _encoder(submesh, layers=4, activation_layers=(3,))
+ ids = ttnn.from_torch(
+ torch.randint(0, 256, (1, SEQ)), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=submesh
+ )
+ cos, sin = create_rope_tensors(1, SEQ, None, 32, 10000.0, [6, 5, 5])
+ pe = (bf16_tensor(cos, device=submesh), bf16_tensor(sin, device=submesh))
+ runs = [(8, 16)]
+ vision = bf16_tensor(torch.randn(1, 16, HIDDEN), device=submesh)
+ feat = [bf16_tensor(torch.randn(1, 16, HIDDEN), device=submesh) for _ in range(2)]
+
+ def run(deepstack):
+ out = enc.forward(
+ ids,
+ attention_mask=None,
+ pos_embeds=pe,
+ vision_embeds=vision,
+ vision_runs=runs,
+ deepstack_embeds=deepstack,
+ )[0]
+ return tensor.to_torch(out, mesh_axes=[None, None, None])
+
+ none, one, two = run(None), run(feat[:1]), run(feat[:2])
+ assert not torch.allclose(none, one, atol=1e-2), "a deepstack feature had no effect"
+ assert not torch.allclose(one, two, atol=1e-2), "the second deepstack feature had no effect"
+
+
+HIDDEN_REAL = 5120 # real conditioner width
+
+# production keyframe runs cross 31/63 tile boundaries; runs inside one 32-row tile never exercise the slicing hazard
+_TILE_RUNS = [
+ pytest.param(1054, [(5, 1008)], HIDDEN_REAL, id="production_keyframe_crosses_31"),
+ pytest.param(2067, [(5, 1008), (1018, 1008)], HIDDEN_REAL, id="production_two_keyframes"),
+ pytest.param(SEQ, [(0, 8)], HIDDEN, id="edge_run_at_row_0"),
+ pytest.param(SEQ, [(SEQ - 8, 8)], HIDDEN, id="edge_run_at_seq_end"),
+]
+
+
+@pytest.mark.parametrize(("mesh_device", "submesh_shape"), _MESH, indirect=["mesh_device"])
+@pytest.mark.parametrize(("seq", "runs", "hidden"), _TILE_RUNS)
+def test_scatter_rows_is_exact_across_tile_boundaries(mesh_device, submesh_shape, seq, runs, hidden):
+ """Replace is pure data movement, so torch.equal, not PCC (PCC would tolerate rows landing a tile off)."""
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ total = sum(length for _, length in runs)
+ assert all(
+ start % ttnn.TILE_SIZE or length % ttnn.TILE_SIZE for start, length in runs
+ ), "every run should be unaligned on at least one end, or this gates nothing"
+
+ torch.manual_seed(0)
+ # pre-rounded through bf16 so the comparison measures the scatter, not the host's fp32 -> bf16 cast
+ base = torch.randn(1, seq, hidden).bfloat16().float()
+ values = torch.randn(1, total, hidden).bfloat16().float()
+
+ golden = base.clone()
+ taken = 0
+ for start, length in runs:
+ golden[:, start : start + length, :] = values[:, taken : taken + length, :]
+ taken += length
+
+ out = _scatter_rows(bf16_tensor(base, device=submesh), bf16_tensor(values, device=submesh), runs, add=False)
+ actual = tensor.to_torch(out, mesh_axes=[None, None, None]).float()
+ assert actual.shape[-2:] == (seq, hidden), f"{tuple(actual.shape)} != (…, {seq}, {hidden})"
+
+ if not torch.equal(golden, actual):
+ wrong = (golden != actual).any(dim=-1).nonzero().flatten().tolist()
+ raise AssertionError(
+ f"{len(wrong)} of {seq} rows differ; first 10 at {wrong[:10]} "
+ f"(tile row-blocks {sorted({r // ttnn.TILE_SIZE for r in wrong[:10]})}), runs={runs}"
+ )
+
+
+@pytest.mark.parametrize(("mesh_device", "submesh_shape"), _MESH, indirect=["mesh_device"])
+@pytest.mark.parametrize(("seq", "runs", "hidden"), _TILE_RUNS)
+def test_scatter_rows_add_is_exact_across_tile_boundaries(mesh_device, submesh_shape, seq, runs, hidden):
+ """add=True: untouched rows bit-exact; written rows within 2 unbiased bf16 ulps (ttnn.add rounds differently)."""
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ total = sum(length for _, length in runs)
+
+ torch.manual_seed(0)
+ base = torch.randn(1, seq, hidden).bfloat16().float()
+ values = torch.randn(1, total, hidden).bfloat16().float()
+
+ golden = base.clone()
+ written = torch.zeros(seq, dtype=torch.bool)
+ taken = 0
+ for start, length in runs:
+ golden[:, start : start + length, :] += values[:, taken : taken + length, :]
+ written[start : start + length] = True
+ taken += length
+
+ out = _scatter_rows(bf16_tensor(base, device=submesh), bf16_tensor(values, device=submesh), runs, add=True)
+ actual = tensor.to_torch(out, mesh_axes=[None, None, None]).float()
+
+ assert torch.equal(golden[:, ~written], actual[:, ~written]), (
+ f"pass-through rows changed under add; "
+ f"{(golden[:, ~written] != actual[:, ~written]).any(dim=-1).sum().item()} of "
+ f"{int((~written).sum())} untouched rows differ"
+ )
+ expected, got = golden[:, written], actual[:, written]
+ # one bf16 ulp = 2**-7 of the binade; 2 ulps lets boundary entries fall either way
+ ulp = torch.ldexp(torch.ones_like(expected), torch.floor(torch.log2(expected.abs().clamp(min=1e-30))).int() - 7)
+ diff = (expected - got).abs()
+ worst = (diff / ulp).max().item()
+ assert worst <= 2.0, f"written rows are {worst:.2f} bf16 ulps out, past the 2-ulp floor"
+
+ # No systematic bias: round-to-nearest is unbiased, truncation is not.
+ bias = ((got - expected) / ulp).mean().item()
+ assert abs(bias) < 0.1, f"mean error {bias:+.3f} ulps suggests truncation rather than round-to-nearest"
+ logger.info(f"scatter add @ seq={seq}: worst {worst:.2f} ulps, mean bias {bias:+.3f} ulps")
diff --git a/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_tower.py b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_tower.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0333f5e6047a4b1f0268d7699f4685134350805
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/qwen3vl/test_qwen3vl_vision_tower.py
@@ -0,0 +1,158 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Qwen3-VL vision tower end to end on device against the HF reference: patch
+# embed, position embeddings, every block, the output and deepstack mergers, at
+# the released depth and only production grid shapes.
+
+import pytest
+import torch
+import transformers
+
+import ttnn
+
+from ....encoders.qwen3vl.vision_qwen3vl import Qwen3VlVisionModel, vision_cu_seqlens
+from ....utils import tensor
+from ....utils.check import assert_quality
+from .common import (
+ HIDDEN_ACT,
+ HIDDEN_SIZE,
+ INTERMEDIATE_SIZE,
+ NORM_EPS,
+ NUM_HEADS,
+ NUM_POSITION_EMBEDDINGS,
+ OUT_HIDDEN_SIZE,
+ SPATIAL_MERGE_SIZE,
+ VISION_PARAMS,
+ resolve_parallel,
+ skip_if_sp_misaligned,
+ sp_shard,
+ vision_config,
+)
+
+DEPTH = 27 # released tower: 27 blocks, deepstack taps at 8/16/24
+DEEPSTACK_INDEXES = [8, 16, 24]
+
+# Grids measured through the checkpoint's own image processor; orientation pairs catch h-vs-w swaps.
+# The 4032-patch canvases are 31.5 x 128 rows, so SP=4 skips there (model geometry, not a port bug).
+_CANVAS = {
+ "canvas_1to1": [[1, 48, 48]],
+ "canvas_4to3": [[1, 48, 64]],
+ "canvas_3to4": [[1, 64, 48]],
+ "canvas_16to9": [[1, 48, 84]],
+ "canvas_9to16": [[1, 84, 48]],
+ "canvas_4to1": [[1, 32, 126]], # area cap binds first, short edge forced to 32
+ "canvas_1to4": [[1, 126, 32]],
+}
+_REFERENCE = {
+ "ref_1to1": [[1, 128, 128]],
+ "ref_4to3": [[1, 128, 170]],
+ "ref_3to4": [[1, 170, 128]],
+ "ref_16to9": [[1, 128, 228]],
+ "ref_9to16": [[1, 228, 128]],
+ "ref_4to1": [[1, 128, 512]],
+ "ref_1to4": [[1, 512, 128]],
+}
+# an image is one block, a video one block PER FRAME; a tower ignoring boundaries passes single-block grids
+_MULTI = {
+ "two_refs": [[1, 128, 128], [1, 128, 170]],
+ "video_3_frames": [[3, 48, 48]],
+ "image_and_video": [[1, 128, 128], [3, 48, 48]],
+ "max_load": [[1, 128, 128]] * 9 + [[3, 48, 48]] * 3, # documented ceiling: 9 ref images + 3 ref videos
+}
+
+GRIDS = {**_CANVAS, **_REFERENCE, **_MULTI}
+
+
+@pytest.fixture(scope="module")
+def reference():
+ torch.manual_seed(0)
+ config = vision_config(depth=DEPTH, deepstack_visual_indexes=DEEPSTACK_INDEXES)
+ return transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionModel._from_config(config).eval()
+
+
+def _tower(reference, submesh, parallel_config=None, ccl_manager=None):
+ tower = Qwen3VlVisionModel(
+ hidden_size=HIDDEN_SIZE,
+ num_heads=NUM_HEADS,
+ depth=DEPTH,
+ intermediate_size=INTERMEDIATE_SIZE,
+ spatial_merge_size=SPATIAL_MERGE_SIZE,
+ num_position_embeddings=NUM_POSITION_EMBEDDINGS,
+ out_hidden_size=OUT_HIDDEN_SIZE,
+ hidden_act=HIDDEN_ACT,
+ norm_eps=NORM_EPS,
+ deepstack_visual_indexes=DEEPSTACK_INDEXES,
+ mesh_device=submesh,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+ tower.load_torch_state_dict(reference.state_dict())
+ return tower
+
+
+# perf mode skips the quadratic CPU golden (shapes/finiteness only) and says nothing about accuracy
+_PERF_GRIDS = ("max_load", "ref_4to1", "ref_1to4", "image_and_video")
+assert all(name in GRIDS for name in _PERF_GRIDS)
+_CASES = [pytest.param(name, True, id=f"check-{name}") for name in GRIDS] + [
+ pytest.param(name, False, id=f"perf-{name}") for name in _PERF_GRIDS
+]
+
+
+@VISION_PARAMS
+@pytest.mark.parametrize(("name", "check_pcc"), _CASES)
+def test_tower_on_device(reference, mesh_device, submesh_shape, tp_axis, sp_axis, num_links, name, check_pcc):
+ """The full tower through both outputs: merged tokens and one deepstack feature per tap."""
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ grid = torch.tensor(GRIDS[name], dtype=torch.long)
+ total = sum(t * h * w for t, h, w in GRIDS[name])
+ skip_if_sp_misaligned(total, submesh, sp_axis)
+ patch_dim = 3 * 2 * 16 * 16
+
+ torch.manual_seed(0)
+ patches = torch.randn(total, patch_dim)
+ golden_tokens, golden_deepstack = None, None
+ if check_pcc:
+ with torch.no_grad():
+ ref_out = reference(patches, grid_thw=grid, return_dict=True)
+ golden_tokens = ref_out.pooler_output.float()
+ golden_deepstack = [f.float() for f in ref_out.deepstack_features]
+ assert len(golden_deepstack) == len(DEEPSTACK_INDEXES), "reference did not emit one feature per index"
+
+ cu_seqlens = vision_cu_seqlens(grid)
+ assert len(cu_seqlens) - 1 == int(grid[:, 0].sum()), f"expected one block per frame, got {cu_seqlens}"
+ assert cu_seqlens[0] == 0 and cu_seqlens[-1] == total, f"cu_seqlens must span [0, {total}]: {cu_seqlens}"
+
+ tower = _tower(reference, submesh, *resolve_parallel(submesh, tp_axis, sp_axis, num_links))
+ cos, sin = tower.prepare_rope(grid)
+ pos = tower.prepare_pos_embeds(grid)
+ tokens, deepstack = tower.forward(
+ sp_shard(patches, submesh, sp_axis),
+ pos_embeds=sp_shard(pos, submesh, sp_axis),
+ rope=(sp_shard(cos, submesh, sp_axis), sp_shard(sin, submesh, sp_axis)),
+ cu_seqlens=cu_seqlens,
+ )
+
+ merged = total // SPATIAL_MERGE_SIZE**2
+ actual_tokens = tensor.to_torch(tokens, mesh_axes=[None, None])
+ assert actual_tokens.shape[-2:] == (merged, OUT_HIDDEN_SIZE), f"{tuple(actual_tokens.shape)}"
+ assert len(deepstack) == len(DEEPSTACK_INDEXES), f"expected {len(DEEPSTACK_INDEXES)} features"
+ features = [tensor.to_torch(f, mesh_axes=[None, None]) for f in deepstack]
+ for i, feature in enumerate(features):
+ assert feature.shape[-2:] == (merged, OUT_HIDDEN_SIZE), f"deepstack {i}: {tuple(feature.shape)}"
+
+ if check_pcc:
+ assert_quality(golden_tokens, actual_tokens, pcc=0.99)
+ for golden_feature, feature in zip(golden_deepstack, features):
+ assert_quality(golden_feature, feature, pcc=0.99)
+ else:
+ assert torch.isfinite(actual_tokens).all(), "merged tokens contain NaN or Inf"
+ for i, feature in enumerate(features):
+ assert torch.isfinite(feature).all(), f"deepstack {i} contains NaN or Inf"
+
+ # Routing by list index can yield the right count and shapes while tapping one layer twice.
+ for i in range(len(features) - 1):
+ assert not torch.allclose(
+ features[i], features[i + 1], atol=1e-2
+ ), f"deepstack features {i} and {i + 1} are identical; layer routing is wrong"
diff --git a/code/models/tt_dit/tests/encoders/t5/test_t5_embeddings.py b/code/models/tt_dit/tests/encoders/t5/test_t5_embeddings.py
new file mode 100644
index 0000000000000000000000000000000000000000..6efa05e1a6334adfe6b735628acbd29b25bb434b
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/t5/test_t5_embeddings.py
@@ -0,0 +1,150 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers.models.t5.modeling_t5 import T5EncoderModel
+
+import ttnn
+from models.tt_dit.encoders.t5.model_t5 import RelativePositionEmbeddings, T5Config
+from models.tt_dit.layers.embeddings import Embedding
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "mesh_device,submesh_shape", [[(2, 4), (1, 4)], [(4, 8), (1, 8)]], ids=["t3k", "glx"], indirect=["mesh_device"]
+)
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_t5_embeddings(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ topology: ttnn.Topology,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=encoder_submesh.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ hf_model = T5EncoderModel.from_pretrained(model_name_checkpoint, subfolder="text_encoder_3", local_files_only=True)
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace T5 Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.d_model}")
+ logger.info(f"intermediate_size: {hf_model.config.d_ff}")
+ logger.info(f"d_kv: {hf_model.config.d_kv}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_layers}")
+ logger.info(f"relative_attention_num_buckets: {hf_model.config.relative_attention_num_buckets}")
+ logger.info(f"relative_attention_max_distance: {hf_model.config.relative_attention_max_distance}")
+ logger.info(f"layer_norm_epsilon: {hf_model.config.layer_norm_epsilon}")
+
+ max_prompt_length = 256
+ torch.manual_seed(0)
+ tokens = torch.randint(hf_model.config.vocab_size, [1, max_prompt_length])
+
+ tt_prompt = ttnn.from_torch(
+ tokens,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ # === TT-DiT T5 ====
+ config = T5Config(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.d_model,
+ ff_dim=hf_model.config.d_ff,
+ kv_dim=hf_model.config.d_kv,
+ num_heads=hf_model.config.num_heads,
+ num_hidden_layers=hf_model.config.num_layers,
+ max_prompt_length=max_prompt_length,
+ layer_norm_eps=hf_model.config.layer_norm_epsilon,
+ relative_attention_num_buckets=hf_model.config.relative_attention_num_buckets,
+ relative_attention_max_distance=hf_model.config.relative_attention_max_distance,
+ )
+
+ state_dict = hf_model.state_dict()
+
+ tt_token_embed = Embedding(config.vocab_size, config.embed_dim, device=encoder_submesh)
+ tt_token_embed.load_torch_state_dict({"weight": state_dict["encoder.embed_tokens.weight"]})
+ tt_relative_position_embed = RelativePositionEmbeddings(config, encoder_submesh, ccl_manager, parallel_config)
+ tt_relative_position_embed.load_torch_state_dict(
+ {"weight": state_dict["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"]}
+ )
+
+ tt_start_time = time.time()
+ tt_embeddings_output = tt_token_embed(tt_prompt)
+ tt_position_bias = tt_relative_position_embed(tt_embeddings_output.shape[1])
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ with torch.no_grad():
+ hf_start_time = time.time()
+
+ hf_token_embeddings = hf_model.encoder.embed_tokens(tokens)
+
+ hf_position_bias = (
+ hf_model.encoder.block[0]
+ .layer[0]
+ .SelfAttention.compute_bias(
+ hf_token_embeddings.size(1), hf_token_embeddings.size(1), device=hf_token_embeddings.device
+ )
+ )
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_embeddings_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_embeddings_output)[0])
+ mesh_shape = list(encoder_submesh.shape)
+ mesh_shape[1 - parallel_config.tensor_parallel.mesh_axis] = 1
+
+ tt_position_bias_torch = ttnn.to_torch(
+ tt_position_bias,
+ mesh_composer=ttnn.create_mesh_composer(
+ encoder_submesh, ttnn.MeshComposerConfig([0, 1], ttnn.MeshShape(mesh_shape))
+ ), # [0,1] is the mesh dimensions to concatenate. Set replicated dimensions to 1.
+ )
+
+ logger.info(f"TT embeddings execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF embeddings execution time: {hf_execution_time:.4f} seconds")
+
+ assert_quality(hf_token_embeddings, tt_embeddings_output_torch, pcc=0.99)
+ assert_quality(hf_position_bias, tt_position_bias_torch, pcc=0.99)
diff --git a/code/models/tt_dit/tests/encoders/t5/test_t5_encoder_all_layers.py b/code/models/tt_dit/tests/encoders/t5/test_t5_encoder_all_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..fdd1616db4923dbf89ad6557907797cbaa21c136
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/t5/test_t5_encoder_all_layers.py
@@ -0,0 +1,197 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers.models.t5.modeling_t5 import T5EncoderModel
+
+import ttnn
+from models.tt_dit.encoders.t5.model_t5 import RelativePositionEmbeddings, T5Config, T5EncoderLayer
+from models.tt_dit.layers.embeddings import Embedding
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.substate import substate
+
+
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "mesh_device,submesh_shape", [[(2, 4), (1, 4)], [(4, 8), (1, 8)]], ids=["t3k", "glx"], indirect=["mesh_device"]
+)
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_t5_layers_individually(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ topology: ttnn.Topology,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=encoder_submesh.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ hf_model = T5EncoderModel.from_pretrained(model_name_checkpoint, subfolder="text_encoder_3", local_files_only=False)
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace T5 Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.d_model}")
+ logger.info(f"intermediate_size: {hf_model.config.d_ff}")
+ logger.info(f"d_kv: {hf_model.config.d_kv}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_layers}")
+ logger.info(f"relative_attention_num_buckets: {hf_model.config.relative_attention_num_buckets}")
+ logger.info(f"relative_attention_max_distance: {hf_model.config.relative_attention_max_distance}")
+ logger.info(f"layer_norm_epsilon: {hf_model.config.layer_norm_epsilon}")
+
+ max_prompt_length = 256
+ torch.manual_seed(0)
+ tokens = torch.randint(hf_model.config.vocab_size, [1, max_prompt_length])
+
+ tt_prompt = ttnn.from_torch(
+ tokens,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ # === TT-DiT T5 ====
+ config = T5Config(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.d_model,
+ ff_dim=hf_model.config.d_ff,
+ kv_dim=hf_model.config.d_kv,
+ num_heads=hf_model.config.num_heads,
+ num_hidden_layers=hf_model.config.num_layers,
+ max_prompt_length=max_prompt_length,
+ layer_norm_eps=hf_model.config.layer_norm_epsilon,
+ relative_attention_num_buckets=hf_model.config.relative_attention_num_buckets,
+ relative_attention_max_distance=hf_model.config.relative_attention_max_distance,
+ )
+
+ state_dict = hf_model.state_dict()
+
+ tt_token_embed = Embedding(config.vocab_size, config.embed_dim, device=encoder_submesh)
+ tt_token_embed.load_torch_state_dict({"weight": state_dict["encoder.embed_tokens.weight"]})
+ tt_relative_position_embed = RelativePositionEmbeddings(config, encoder_submesh, ccl_manager, parallel_config)
+ tt_relative_position_embed.load_torch_state_dict(
+ {"weight": state_dict["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"]}
+ )
+
+ # time TT model inference only
+ tt_start_time = time.time()
+ tt_embeddings_output = tt_token_embed(tt_prompt)
+ tt_position_bias = tt_relative_position_embed(tt_embeddings_output.shape[1])
+
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+ logger.info(f"TT embedding execution time: {tt_execution_time:.4f} seconds")
+
+ hf_start_time = time.time()
+ with torch.no_grad():
+ hf_token_embeddings = hf_model.encoder.embed_tokens(tokens)
+
+ # get position bias from first layer's attention (block 0 always has the bias)
+ hf_position_bias = (
+ hf_model.encoder.block[0] # Always use block[0] for position bias
+ .layer[0]
+ .SelfAttention.compute_bias(
+ hf_token_embeddings.size(1), hf_token_embeddings.size(1), device=hf_token_embeddings.device
+ )
+ )
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+ logger.info(f"HF embedding and position bias time: {hf_execution_time:.4f} seconds")
+
+ logger.info("\n=== Testing each layer individually ===")
+ pcc_values = []
+
+ tt_embeddings_output = tt_token_embed(tt_prompt)
+ tt_position_bias = tt_relative_position_embed(tt_embeddings_output.shape[1])
+ with torch.no_grad():
+ hf_token_embeddings = hf_model.encoder.embed_tokens(tokens)
+ hf_position_bias = (
+ hf_model.encoder.block[0] # Always use block[0] for position bias
+ .layer[0]
+ .SelfAttention.compute_bias(
+ hf_token_embeddings.size(1), hf_token_embeddings.size(1), device=hf_token_embeddings.device
+ )
+ )
+
+ for i in range(config.num_hidden_layers):
+ logger.info(f"\nTesting layer {i}")
+ tt_start_time = time.time()
+ layer = i
+
+ # Block 0 is the only one carrying relative_attention_bias in the HF checkpoint, so it
+ # must be built with the bias parameter or the strict state-dict load rejects the key.
+ # This mirrors T5Stack, which passes config.use_relative_position_bias[i].
+ tt_encoder_layer = T5EncoderLayer(
+ config, encoder_submesh, ccl_manager, parallel_config, config.use_relative_position_bias[i]
+ )
+ tt_encoder_layer.load_torch_state_dict(substate(hf_model.state_dict(), f"encoder.block.{layer}"))
+
+ tt_layer_output = tt_encoder_layer(tt_embeddings_output, tt_position_bias)
+
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+ logger.info(f"TT layer {i} execution time: {tt_execution_time:.4f} seconds")
+
+ hf_start_time = time.time()
+
+ hf_layer_output = hf_model.encoder.block[layer](
+ hf_token_embeddings, attention_mask=None, position_bias=hf_position_bias
+ )[0]
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+ logger.info(f"HF layer {i} execution time: {hf_execution_time:.4f} seconds")
+
+ tt_layer_output = ttnn.to_torch(ttnn.get_device_tensors(tt_layer_output)[0])
+
+ assert hf_layer_output.shape == tt_layer_output.shape
+
+ try:
+ assert_quality(hf_layer_output, tt_layer_output, pcc=0.947)
+ logger.info(f"Layer {i} passed PCC threshold")
+ except Exception as e:
+ # extract PCC value from error message
+ error_msg = str(e)
+ pcc = float(error_msg.split("%")[0].split("=")[1].strip()) / 100
+ logger.warning(f"Layer {i} failed: PCC = {pcc:.4f}")
+ pcc_values.append(pcc)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/t5/test_t5_encoder_single_layer.py b/code/models/tt_dit/tests/encoders/t5/test_t5_encoder_single_layer.py
new file mode 100644
index 0000000000000000000000000000000000000000..09813d51c4fa31e6507f23eda022e48c1749aebd
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/t5/test_t5_encoder_single_layer.py
@@ -0,0 +1,163 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers.models.t5.modeling_t5 import T5EncoderModel
+
+import ttnn
+from models.tt_dit.encoders.t5.model_t5 import RelativePositionEmbeddings, T5Config, T5EncoderLayer
+from models.tt_dit.layers.embeddings import Embedding
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.substate import substate
+
+
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "mesh_device,submesh_shape", [[(2, 4), (1, 4)], [(4, 8), (1, 8)]], ids=["t3k", "glx"], indirect=["mesh_device"]
+)
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_t5_single_layer(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ topology: ttnn.Topology,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=encoder_submesh.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ hf_model = T5EncoderModel.from_pretrained(model_name_checkpoint, subfolder="text_encoder_3", local_files_only=True)
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace T5 Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.d_model}")
+ logger.info(f"intermediate_size: {hf_model.config.d_ff}")
+ logger.info(f"d_kv: {hf_model.config.d_kv}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_layers}")
+ logger.info(f"relative_attention_num_buckets: {hf_model.config.relative_attention_num_buckets}")
+ logger.info(f"relative_attention_max_distance: {hf_model.config.relative_attention_max_distance}")
+ logger.info(f"layer_norm_epsilon: {hf_model.config.layer_norm_epsilon}")
+
+ max_prompt_length = 256
+ torch.manual_seed(0)
+ tokens = torch.randint(hf_model.config.vocab_size, [1, max_prompt_length])
+
+ tt_prompt = ttnn.from_torch(
+ tokens,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ # === TT-DiT T5 ====
+ config = T5Config(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.d_model,
+ ff_dim=hf_model.config.d_ff,
+ kv_dim=hf_model.config.d_kv,
+ num_heads=hf_model.config.num_heads,
+ num_hidden_layers=hf_model.config.num_layers,
+ max_prompt_length=max_prompt_length,
+ layer_norm_eps=hf_model.config.layer_norm_epsilon,
+ relative_attention_num_buckets=hf_model.config.relative_attention_num_buckets,
+ relative_attention_max_distance=hf_model.config.relative_attention_max_distance,
+ )
+
+ state_dict = hf_model.state_dict()
+
+ tt_token_embed = Embedding(config.vocab_size, config.embed_dim, device=encoder_submesh)
+ tt_token_embed.load_torch_state_dict({"weight": state_dict["encoder.embed_tokens.weight"]})
+ tt_relative_position_embed = RelativePositionEmbeddings(config, encoder_submesh, ccl_manager, parallel_config)
+ tt_relative_position_embed.load_torch_state_dict(
+ {"weight": state_dict["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"]}
+ )
+
+ layer = 23
+
+ tt_encoder_layer = T5EncoderLayer(config, encoder_submesh, ccl_manager, parallel_config)
+
+ tt_encoder_layer.load_torch_state_dict(
+ substate(hf_model.state_dict(), f"encoder.block.{layer}")
+ ) # first layer weights
+
+ tt_start_time = time.time()
+ tt_embeddings_output = tt_token_embed(tt_prompt)
+ tt_position_bias = tt_relative_position_embed(tt_embeddings_output.shape[1])
+ tt_layer_output = tt_encoder_layer(tt_embeddings_output, tt_position_bias)
+
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ with torch.no_grad():
+ hf_start_time = time.time()
+
+ hf_token_embeddings = hf_model.encoder.embed_tokens(tokens)
+
+ # get position bias from first layer's attention (block 0 always has the bias)
+ hf_position_bias = (
+ hf_model.encoder.block[0] # Always use block[0] for position bias
+ .layer[0]
+ .SelfAttention.compute_bias(
+ hf_token_embeddings.size(1), hf_token_embeddings.size(1), device=hf_token_embeddings.device
+ )
+ )
+
+ # run through encoder layer
+ hf_layer_output = hf_model.encoder.block[layer](
+ hf_token_embeddings, attention_mask=None, position_bias=hf_position_bias
+ )[0]
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_layer_output = ttnn.to_torch(ttnn.get_device_tensors(tt_layer_output)[0])
+
+ logger.info(f"TT layer execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF layer execution time: {hf_execution_time:.4f} seconds")
+
+ assert hf_layer_output.shape == tt_layer_output.shape
+
+ assert_quality(hf_layer_output, tt_layer_output, pcc=0.947)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/t5/test_t5_full.py b/code/models/tt_dit/tests/encoders/t5/test_t5_full.py
new file mode 100644
index 0000000000000000000000000000000000000000..1df8f0cee699f3304e8b3d72c90a07b3ddf3a073
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/t5/test_t5_full.py
@@ -0,0 +1,154 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parents[6]))
+
+import pytest
+import torch
+from loguru import logger
+from transformers.models.t5.configuration_t5 import T5Config as HF_T5Config
+from transformers.models.t5.modeling_t5 import T5EncoderModel
+
+import ttnn
+from models.tt_dit.encoders.t5.model_t5 import T5Config, T5Encoder
+from models.tt_dit.parallel.config import EncoderParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ "dit_unit_test",
+ [{"1": True, "0": False}.get(os.environ.get("DIT_UNIT_TEST"), False)],
+)
+@pytest.mark.parametrize(
+ "model_name",
+ [
+ "large",
+ ],
+)
+@pytest.mark.parametrize(
+ "mesh_device,submesh_shape", [[(2, 4), (1, 4)], [(4, 8), (1, 8)]], ids=["t3k", "glx"], indirect=["mesh_device"]
+)
+@pytest.mark.parametrize(
+ "device_params, topology",
+ [[{"l1_small_size": 8192, "fabric_config": ttnn.FabricConfig.FABRIC_1D}, ttnn.Topology.Linear]],
+ indirect=["device_params"],
+)
+def test_t5_encoder(
+ *,
+ mesh_device: ttnn.Device,
+ submesh_shape: ttnn.MeshShape,
+ model_name: str,
+ topology: ttnn.Topology,
+ dit_unit_test: bool,
+) -> None:
+ parent_mesh_shape = tuple(mesh_device.shape)
+ if any(x[0] < x[1] for x in zip(parent_mesh_shape, submesh_shape)):
+ pytest.skip("submesh shape is larger than parent mesh shape, skipping")
+ encoder_submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ print(f"Running on submesh {encoder_submesh.shape} of parent mesh {mesh_device.shape}")
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=encoder_submesh.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=encoder_submesh,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ model_name_checkpoint = f"stabilityai/stable-diffusion-3.5-{model_name}"
+
+ if dit_unit_test:
+ # Config for t5-xxl with minimal layers
+ hf_config = HF_T5Config(
+ d_model=4096,
+ d_ff=10240,
+ num_heads=64,
+ num_layers=2,
+ num_decoder_layers=2,
+ feed_forward_proj="gated-gelu",
+ output_past=True,
+ )
+ hf_model = T5EncoderModel(hf_config)
+ else:
+ hf_model = T5EncoderModel.from_pretrained(
+ model_name_checkpoint, subfolder="text_encoder_3", local_files_only=False
+ )
+ hf_model.eval()
+
+ logger.info("=== HuggingFace T5 Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.d_model}")
+ logger.info(f"intermediate_size: {hf_model.config.d_ff}")
+ logger.info(f"d_kv: {hf_model.config.d_kv}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_layers}")
+ logger.info(f"relative_attention_num_buckets: {hf_model.config.relative_attention_num_buckets}")
+ logger.info(f"relative_attention_max_distance: {hf_model.config.relative_attention_max_distance}")
+ logger.info(f"layer_norm_epsilon: {hf_model.config.layer_norm_epsilon}")
+
+ max_prompt_length = 256
+ torch.manual_seed(0)
+ tokens = torch.randint(hf_model.config.vocab_size, [1, max_prompt_length])
+
+ tt_prompt = ttnn.from_torch(
+ tokens,
+ layout=ttnn.TILE_LAYOUT,
+ device=encoder_submesh,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(encoder_submesh),
+ )
+
+ # === TT-DiT T5 ====
+ config = T5Config(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.d_model,
+ ff_dim=hf_model.config.d_ff,
+ kv_dim=hf_model.config.d_kv,
+ num_heads=hf_model.config.num_heads,
+ num_hidden_layers=hf_model.config.num_layers,
+ max_prompt_length=max_prompt_length,
+ layer_norm_eps=hf_model.config.layer_norm_epsilon,
+ relative_attention_num_buckets=hf_model.config.relative_attention_num_buckets,
+ relative_attention_max_distance=hf_model.config.relative_attention_max_distance,
+ )
+
+ tt_encoder = T5Encoder(config, encoder_submesh, ccl_manager, parallel_config)
+ tt_encoder.load_torch_state_dict(hf_model.state_dict())
+
+ # time TT model inference only
+ tt_start_time = time.time()
+ tt_output = tt_encoder(tt_prompt)
+
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ # get HF reference outputs
+ with torch.no_grad():
+ hf_start_time = time.time()
+
+ hf_outputs = hf_model(tokens).last_hidden_state
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_output[-1])[0])
+
+ logger.info(f"TT encoder execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF encoder execution time: {hf_execution_time:.4f} seconds")
+
+ assert hf_outputs.shape == tt_output_torch.shape
+
+ assert_quality(hf_outputs, tt_output_torch, pcc=0.99)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/encoders/umt5/test_umt5.py b/code/models/tt_dit/tests/encoders/umt5/test_umt5.py
new file mode 100644
index 0000000000000000000000000000000000000000..c5778d8921f8c6c3a8114c8f81372c9ef2ea032b
--- /dev/null
+++ b/code/models/tt_dit/tests/encoders/umt5/test_umt5.py
@@ -0,0 +1,304 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import time
+
+import pytest
+import torch
+from loguru import logger
+from transformers import AutoTokenizer
+from transformers.models.umt5.configuration_umt5 import UMT5Config as HF_UMT5Config
+from transformers.models.umt5.modeling_umt5 import UMT5EncoderModel
+
+import ttnn
+from models.perf.benchmarking_utils import BenchmarkProfiler
+
+from ....encoders.t5.model_t5 import RelativePositionEmbeddings as TT_UMT5RelativePositionEmbeddings
+from ....encoders.umt5.model_umt5 import UMT5Config as TT_UMT5Config
+from ....encoders.umt5.model_umt5 import UMT5Encoder as TT_UMT5Encoder
+from ....layers.embeddings import Embedding
+from ....parallel.config import EncoderParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils import cache
+from ....utils.check import assert_quality
+
+# Shared device configuration for UMT5 tests to reduce code duplication and ensure we are using consistent device configurations.
+
+
+def umt5_device_config(func):
+ """Decorator to apply standard UMT5 device/mesh parametrization to tests."""
+ func = pytest.mark.parametrize(
+ "dit_unit_test",
+ [{"1": True, "0": False}.get(os.environ.get("DIT_UNIT_TEST"), False)],
+ )(func)
+ func = pytest.mark.parametrize(
+ "mesh_device, num_links, max_execution_time",
+ [[(2, 4), 1, 0.13], [(4, 8), 4, 0.1], [(4, 8), 2, 0.1]],
+ ids=["t3k", "wh_glx", "bh_glx"],
+ indirect=["mesh_device", "num_links"],
+ )(func)
+ func = pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}],
+ indirect=True,
+ )(func)
+ return func
+
+
+@pytest.fixture
+def num_links(request):
+ return request.param
+
+
+@pytest.fixture
+def parallel_config_and_ccl_manager(mesh_device, num_links):
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=mesh_device.shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(
+ mesh_device=mesh_device,
+ num_links=num_links,
+ topology=ttnn.Topology.Linear,
+ )
+ return parallel_config, ccl_manager
+
+
+def get_random_weights_model() -> UMT5EncoderModel:
+ hf_config = HF_UMT5Config(
+ d_model=4096,
+ d_ff=10240,
+ num_heads=64,
+ num_layers=24,
+ num_decoder_layers=0,
+ relative_attention_max_distance=128,
+ relative_attention_num_buckets=32,
+ feed_forward_proj="gated-gelu",
+ output_past=True,
+ )
+ return UMT5EncoderModel(hf_config)
+
+
+@umt5_device_config
+def test_umt5_embeddings(
+ *,
+ mesh_device: ttnn.Device,
+ parallel_config_and_ccl_manager: tuple,
+ dit_unit_test: bool,
+ max_execution_time: float,
+) -> None:
+ torch.manual_seed(0)
+
+ parallel_config, ccl_manager = parallel_config_and_ccl_manager
+
+ if dit_unit_test:
+ hf_model = get_random_weights_model()
+ else:
+ model_name_checkpoint = f"Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+ hf_model = UMT5EncoderModel.from_pretrained(
+ model_name_checkpoint, subfolder="text_encoder", local_files_only=False
+ )
+
+ hf_model.eval()
+
+ logger.info("=== HuggingFace UMT5 Config ===")
+ logger.info(f"vocab_size: {hf_model.config.vocab_size}")
+ logger.info(f"hidden_size: {hf_model.config.d_model}")
+ logger.info(f"intermediate_size: {hf_model.config.d_ff}")
+ logger.info(f"d_kv: {hf_model.config.d_kv}")
+ logger.info(f"num_attention_heads: {hf_model.config.num_heads}")
+ logger.info(f"num_hidden_layers: {hf_model.config.num_layers}")
+ logger.info(f"relative_attention_num_buckets: {hf_model.config.relative_attention_num_buckets}")
+ logger.info(f"relative_attention_max_distance: {hf_model.config.relative_attention_max_distance}")
+ logger.info(f"layer_norm_epsilon: {hf_model.config.layer_norm_epsilon}")
+
+ max_prompt_length = 512
+ tokens = torch.randint(hf_model.config.vocab_size, [1, max_prompt_length])
+
+ tt_prompt = ttnn.from_torch(
+ tokens,
+ layout=ttnn.TILE_LAYOUT,
+ device=mesh_device,
+ mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+ )
+
+ # === TT-DiT UMT5 ====
+ config = TT_UMT5Config(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.d_model,
+ ff_dim=hf_model.config.d_ff,
+ kv_dim=hf_model.config.d_kv,
+ num_heads=hf_model.config.num_heads,
+ num_hidden_layers=hf_model.config.num_layers,
+ max_prompt_length=max_prompt_length,
+ layer_norm_eps=hf_model.config.layer_norm_epsilon,
+ relative_attention_num_buckets=hf_model.config.relative_attention_num_buckets,
+ relative_attention_max_distance=hf_model.config.relative_attention_max_distance,
+ )
+
+ state_dict = hf_model.state_dict()
+
+ tt_token_embed = Embedding(config.vocab_size, config.embed_dim, device=mesh_device)
+ tt_token_embed.load_torch_state_dict({"weight": state_dict["encoder.embed_tokens.weight"]})
+ tt_relative_position_embed = TT_UMT5RelativePositionEmbeddings(config, mesh_device, ccl_manager, parallel_config)
+ tt_relative_position_embed.load_torch_state_dict(
+ {"weight": state_dict["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"]}
+ )
+
+ tt_start_time = time.time()
+ tt_embeddings_output = tt_token_embed(tt_prompt)
+ tt_position_bias = tt_relative_position_embed(tt_embeddings_output.shape[1])
+ tt_end_time = time.time()
+ tt_execution_time = tt_end_time - tt_start_time
+
+ with torch.no_grad():
+ hf_start_time = time.time()
+
+ hf_token_embeddings = hf_model.encoder.embed_tokens(tokens)
+
+ hf_position_bias = (
+ hf_model.encoder.block[0]
+ .layer[0]
+ .SelfAttention.compute_bias(
+ hf_token_embeddings.size(1), hf_token_embeddings.size(1), device=hf_token_embeddings.device
+ )
+ )
+
+ hf_end_time = time.time()
+ hf_execution_time = hf_end_time - hf_start_time
+
+ # convert mesh tensor to torch tensor for pcc
+ # since weights are replicated, can get the tensor from any single device
+ tt_embeddings_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_embeddings_output)[0])
+ mesh_shape = list(mesh_device.shape)
+ mesh_shape[1 - parallel_config.tensor_parallel.mesh_axis] = 1
+
+ tt_position_bias_torch = ttnn.to_torch(
+ tt_position_bias,
+ mesh_composer=ttnn.create_mesh_composer(
+ mesh_device, ttnn.MeshComposerConfig([0, 1], ttnn.MeshShape(mesh_shape))
+ ), # [0,1] is the mesh dimensions to concatenate. Set replicated dimensions to 1.
+ )
+
+ logger.info(f"TT embeddings execution time: {tt_execution_time:.4f} seconds")
+ logger.info(f"HF embeddings execution time: {hf_execution_time:.4f} seconds")
+
+ assert_quality(hf_token_embeddings, tt_embeddings_output_torch, pcc=(1.0 - 1e-5), relative_rmse=0.005)
+ assert_quality(hf_position_bias, tt_position_bias_torch, pcc=(1.0 - 1e-5), relative_rmse=0.005)
+
+
+@umt5_device_config
+def test_umt5_encoder(
+ *,
+ mesh_device: ttnn.Device,
+ parallel_config_and_ccl_manager: tuple,
+ dit_unit_test: bool,
+ max_execution_time: float,
+) -> None:
+ torch.manual_seed(0)
+
+ parallel_config, ccl_manager = parallel_config_and_ccl_manager
+
+ model_name_checkpoint = f"Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+
+ if dit_unit_test:
+ hf_model = get_random_weights_model()
+ else:
+ hf_model = UMT5EncoderModel.from_pretrained(
+ model_name_checkpoint, subfolder="text_encoder", local_files_only=False
+ )
+ hf_model.eval()
+
+ max_prompt_length = 512
+
+ prompt = [
+ "A Roman general standing on a battlefield at dawn, torn red cape blowing in the wind, distant soldiers forming ranks, painterly brushwork in the style of Caravaggio, chiaroscuro lighting, epic composition",
+ "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+ ]
+ total_prompts = len(prompt)
+ num_devices = mesh_device.shape[1 - parallel_config.tensor_parallel.mesh_axis]
+ prompt += [" "] * ((num_devices - (total_prompts % num_devices)) % num_devices)
+ tokenizer = AutoTokenizer.from_pretrained(model_name_checkpoint, subfolder="tokenizer", trust_remote_code=True)
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_prompt_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_attention_mask=True,
+ return_tensors="pt",
+ )
+ text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+
+ dims = [None, None]
+ DP_axis = 1 - parallel_config.tensor_parallel.mesh_axis
+ dims[DP_axis] = 0
+ mesh_mapper = ttnn.ShardTensor2dMesh(mesh_device, mesh_shape=tuple(mesh_device.shape), dims=dims)
+
+ tt_prompt = ttnn.from_torch(
+ text_input_ids,
+ dtype=ttnn.uint32,
+ layout=ttnn.TILE_LAYOUT,
+ device=mesh_device,
+ mesh_mapper=mesh_mapper,
+ )
+
+ tt_mask = ttnn.from_torch(
+ mask,
+ dtype=ttnn.bfloat16,
+ layout=ttnn.TILE_LAYOUT,
+ device=mesh_device,
+ mesh_mapper=mesh_mapper,
+ )
+
+ # === TT-DiT UMT5 ====
+ config = TT_UMT5Config(
+ vocab_size=hf_model.config.vocab_size,
+ embed_dim=hf_model.config.d_model,
+ ff_dim=hf_model.config.d_ff,
+ kv_dim=hf_model.config.d_kv,
+ num_heads=hf_model.config.num_heads,
+ num_hidden_layers=hf_model.config.num_layers,
+ max_prompt_length=max_prompt_length,
+ layer_norm_eps=hf_model.config.layer_norm_epsilon,
+ relative_attention_num_buckets=hf_model.config.relative_attention_num_buckets,
+ relative_attention_max_distance=hf_model.config.relative_attention_max_distance,
+ )
+
+ tt_encoder = TT_UMT5Encoder(config, mesh_device, ccl_manager, parallel_config)
+ cache.load_model(
+ tt_encoder,
+ model_name=model_name_checkpoint,
+ subfolder="text_encoder",
+ parallel_config=parallel_config,
+ mesh_shape=tuple(mesh_device.shape),
+ mesh_device=mesh_device,
+ get_torch_state_dict=lambda: hf_model.state_dict(),
+ )
+
+ # warmup
+ tt_output = tt_encoder(tt_prompt, attention_mask=tt_mask)
+ ttnn.synchronize_device(mesh_device) # wait for all operations to complete
+ benchmark_profiler = BenchmarkProfiler()
+ num_runs = 10
+ for i in range(num_runs):
+ with benchmark_profiler("tt_umt5_encoder", i):
+ tt_output = tt_encoder(tt_prompt, attention_mask=tt_mask)[-1]
+ tt_output = ccl_manager.all_gather(tt_output, dim=0, mesh_axis=DP_axis, use_hyperparams=True)
+ # wait for all operations to complete to get correct dispatch + execution time
+ ttnn.synchronize_device(mesh_device)
+
+ # get HF reference outputs
+ with torch.no_grad():
+ hf_outputs = hf_model(text_input_ids, attention_mask=mask).last_hidden_state
+
+ # # convert mesh tensor to torch tensor for pcc
+ # # since weights are replicated, can get the tensor from any single device
+ tt_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_output)[0])
+ tt_execution_time = benchmark_profiler.get_duration_average("tt_umt5_encoder")
+ logger.info(f"TT encoder execution time benchmark: {tt_execution_time:.4f} seconds")
+ assert_quality(hf_outputs[:total_prompts], tt_output_torch[:total_prompts], pcc=0.99, relative_rmse=0.06)
+ # assert (
+ # tt_execution_time < max_execution_time
+ # ), f"TT Encoder execution time {tt_execution_time:.4f} seconds is greater than the max {max_execution_time} seconds"
diff --git a/code/models/tt_dit/tests/models/flux1/scaleout_configs/single_galaxy/mesh_graph_descriptor.textproto b/code/models/tt_dit/tests/models/flux1/scaleout_configs/single_galaxy/mesh_graph_descriptor.textproto
new file mode 100644
index 0000000000000000000000000000000000000000..93ee3b534ab562f8690393dc46523589b6dbba6a
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux1/scaleout_configs/single_galaxy/mesh_graph_descriptor.textproto
@@ -0,0 +1,17 @@
+# Single-rank single-galaxy mesh for the Flux multihost CI experiment.
+# Copied from models/demos/deepseek_v3_b1/scaleout_configs/rev_c/one_galaxy, which is proven to
+# load and map on the bh_sc1 (single BH galaxy) SKU. One connected local 4x2 BLACKHOLE block
+# (8 devices) owned by a single rank.
+
+mesh_descriptors {
+ name: "M0"
+ arch: BLACKHOLE
+ device_topology { dims: [ 4, 2 ] dim_types: [ RING, LINE ] }
+ host_topology { dims: [ 1, 1 ] }
+ channels {
+ count: 2
+ policy: RELAXED
+ }
+}
+
+top_level_instance { mesh { mesh_descriptor: "M0" mesh_id: 0 } }
diff --git a/code/models/tt_dit/tests/models/flux1/test_attention_flux1.py b/code/models/tt_dit/tests/models/flux1/test_attention_flux1.py
new file mode 100644
index 0000000000000000000000000000000000000000..080ec8d30ce7fc01195003ff21501187a2df2ec6
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux1/test_attention_flux1.py
@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import diffusers.models.attention_processor
+import pytest
+import torch
+
+import ttnn
+
+from ....blocks.attention import Attention
+from ....parallel.config import DiTParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.padding import PaddingConfig
+from ....utils.tensor import bf16_tensor
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis"),
+ [
+ pytest.param((1, 1), 0, 1, id="1x1sp0tp1"),
+ pytest.param((1, 2), 0, 1, id="1x2sp0tp1"),
+ pytest.param((1, 2), 1, 0, id="1x2sp1tp0"),
+ pytest.param((2, 1), 0, 1, id="2x1sp0tp1"),
+ pytest.param((2, 1), 1, 0, id="2x1sp1tp0"),
+ pytest.param((2, 2), 0, 1, id="2x2sp0tp1"),
+ pytest.param((2, 2), 1, 0, id="2x2sp1tp0"),
+ pytest.param((2, 4), 0, 1, id="2x4sp0tp1"),
+ pytest.param((2, 4), 1, 0, id="2x4sp1tp0"),
+ pytest.param((4, 4), 0, 1, id="4x4sp0tp1"),
+ pytest.param((4, 8), 0, 1, id="4x8sp0tp1"),
+ pytest.param((4, 8), 1, 0, id="4x8sp1tp0"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "spatial_seq_len", "prompt_seq_len"),
+ [
+ (1, 4096, 512),
+ # (1, 4096 + 512, 0),
+ ],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "require_exact_physical_num_devices": True,
+ }
+ ],
+ indirect=True,
+)
+def test_attention_flux1(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ batch_size: int,
+ spatial_seq_len: int,
+ prompt_seq_len: int,
+) -> None:
+ torch.manual_seed(0)
+
+ query_dim = 3072
+ added_kv_proj_dim = 3072
+ head_dim = 128
+ num_heads = 24
+ out_dim = 3072
+ context_pre_only = prompt_seq_len == 0
+ pre_only = prompt_seq_len == 0
+ batch_size = 1
+
+ joint_attention = added_kv_proj_dim != 0
+
+ sp_factor = mesh_device.shape[sp_axis]
+ tp_factor = mesh_device.shape[tp_axis]
+
+ torch_model = diffusers.models.attention_processor.Attention(
+ query_dim=query_dim,
+ added_kv_proj_dim=added_kv_proj_dim if added_kv_proj_dim != 0 else None,
+ dim_head=head_dim,
+ heads=num_heads,
+ out_dim=out_dim,
+ context_pre_only=context_pre_only,
+ pre_only=pre_only,
+ bias=True,
+ qk_norm="rms_norm",
+ eps=1e-6,
+ processor=diffusers.models.attention_processor.FluxAttnProcessor2_0(), # type: ignore[arg-type]
+ )
+ torch_model.eval()
+
+ ccl_manager = CCLManager(
+ mesh_device=mesh_device,
+ num_links=1,
+ topology=ttnn.Topology.Linear,
+ )
+
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=0, mesh_axis=0),
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+
+ if num_heads % tp_factor != 0:
+ padding_config = PaddingConfig.from_tensor_parallel_factor(num_heads, head_dim, tp_factor)
+ else:
+ padding_config = None
+
+ tt_model = Attention(
+ query_dim=query_dim,
+ head_dim=head_dim,
+ heads=num_heads,
+ out_dim=out_dim,
+ added_kv_proj_dim=added_kv_proj_dim,
+ context_pre_only=context_pre_only,
+ pre_only=pre_only,
+ eps=1e-6,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ padding_config=padding_config,
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ spatial = torch.randn((batch_size, spatial_seq_len, query_dim))
+ prompt = torch.randn((batch_size, prompt_seq_len, query_dim)) if joint_attention else None
+ rope_cos = torch.randn([spatial_seq_len + prompt_seq_len, 128])
+ rope_sin = torch.randn([spatial_seq_len + prompt_seq_len, 128])
+
+ spatial_padded = tt_model.pad_spatial_sequence(spatial, sp_factor=sp_factor)
+ spatial_rope_cos_padded = tt_model.pad_spatial_sequence(rope_cos[prompt_seq_len:], sp_factor=sp_factor)
+ spatial_rope_sin_padded = tt_model.pad_spatial_sequence(rope_sin[prompt_seq_len:], sp_factor=sp_factor)
+
+ tt_spatial = bf16_tensor(spatial_padded, device=mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ tt_prompt = bf16_tensor(prompt, device=mesh_device) if prompt is not None else None
+ tt_spatial_rope_cos = bf16_tensor(spatial_rope_cos_padded, device=mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ tt_spatial_rope_sin = bf16_tensor(spatial_rope_sin_padded, device=mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ tt_prompt_rope_cos = bf16_tensor(rope_cos[:prompt_seq_len], device=mesh_device) if joint_attention else None
+ tt_prompt_rope_sin = bf16_tensor(rope_sin[:prompt_seq_len], device=mesh_device) if joint_attention else None
+
+ with torch.no_grad():
+ torch_output = torch_model.forward(spatial, prompt, image_rotary_emb=(rope_cos, rope_sin))
+ torch_spatial, torch_prompt = torch_output if joint_attention else (torch_output, None)
+
+ tt_spatial_out, tt_prompt_out = tt_model.forward(
+ spatial=tt_spatial,
+ prompt=tt_prompt,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin) if joint_attention else None,
+ spatial_sequence_length=spatial_seq_len,
+ )
+
+ tt_spatial_torch = tensor.to_torch(tt_spatial_out, mesh_axes=[None, sp_axis, tp_axis])
+ assert_quality(torch_spatial, tt_spatial_torch, pcc=0.995, relative_rmse=0.13)
+
+ if torch_prompt is not None:
+ tt_prompt_torch = tensor.to_torch(tt_prompt_out, mesh_axes=[None, None, tp_axis])
+ assert_quality(torch_prompt, tt_prompt_torch, pcc=0.995, relative_rmse=0.13)
diff --git a/code/models/tt_dit/tests/models/flux1/test_performance_flux1.py b/code/models/tt_dit/tests/models/flux1/test_performance_flux1.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e646e508ea11da1aa5cb6fafb868ce5c730301b
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux1/test_performance_flux1.py
@@ -0,0 +1,525 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import statistics
+
+import pytest
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import is_blackhole
+from models.perf.benchmarking_utils import BenchmarkData, BenchmarkProfiler
+
+from ....parallel.config import DiTParallelConfig, EncoderParallelConfig, VAEParallelConfig
+from ....pipelines.events import profiler_event_callback
+from ....pipelines.flux1.pipeline_flux1 import Flux1Pipeline, Flux1PipelineConfig
+from ...dataset_eval.clip_encoder import CLIPEncoder
+
+# Test prompts - diverse set for comprehensive performance testing
+TEST_PROMPTS = [
+ """A neon-lit alley in a sprawling cyberpunk metropolis at night, rain-slick streets reflecting glowing holograms, dense atmosphere, flying cars in the sky, people in high-tech streetwear — ultra-detailed, cinematic lighting, 4K""",
+ """A colossal whale floating through a desert sky like a blimp, casting a long shadow over sand dunes, people in ancient robes watching in awe, golden hour lighting, dreamlike color palette — surrealism, concept art, Greg Rutkowski style""",
+ """A Roman general standing on a battlefield at dawn, torn red cape blowing in the wind, distant soldiers forming ranks, painterly brushwork in the style of Caravaggio, chiaroscuro lighting, epic composition""",
+ """A tiny, fluffy dragon curled up in a teacup, warm cozy lighting, big expressive eyes, intricate scale patterns, surrounded by books and potions — high detail, Studio Ghibli meets Pixar""",
+ """An epic, high-definition cinematic shot of a rustic snowy cabin glowing warmly at dusk, nestled in a serene winter landscape. Surrounded by gentle snow-covered pines and delicate falling snowflakes — captured in a rich, atmospheric, wide-angle scene with deep cinematic depth and warmth.""",
+]
+
+PERF_TEST_IMAGE_PARAMS = [
+ (1024, 1024, 3.5, 28),
+]
+
+PERF_TEST_MESH_PARAMS = [
+ pytest.param((1, 2), (1, 0), (2, 1), (2, 1), (2, 1), ttnn.Topology.Linear, 2, id="1x2sp0tp1"),
+ pytest.param((2, 2), (2, 0), (2, 1), (2, 1), (2, 1), ttnn.Topology.Linear, 2, id="2x2sp0tp1"),
+ pytest.param((2, 4), (2, 0), (4, 1), (4, 1), (4, 1), ttnn.Topology.Linear, 1, id="wh_2x4sp0tp1"),
+ pytest.param((2, 4), (2, 0), (4, 1), (4, 1), (4, 1), ttnn.Topology.Linear, 2, id="bh_2x4sp0tp1"),
+ pytest.param((4, 8), (4, 0), (8, 1), (4, 0), (4, 0), ttnn.Topology.Linear, 4, id="wh_4x8sp0tp1"),
+ pytest.param((4, 8), (4, 0), (8, 1), (4, 0), (4, 0), ttnn.Topology.Linear, 2, id="bh_4x8sp0tp1"),
+]
+
+PERF_TEST_DEVICE_PARAMS = [
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "l1_small_size": 32768,
+ "trace_region_size": 51000000,
+ "require_exact_physical_num_devices": True,
+ }
+]
+
+PERF_TEST_MODEL_VARIANTS = [
+ pytest.param("schnell", id="flux_schnell"),
+ pytest.param("dev", id="flux_dev"),
+]
+
+
+@pytest.mark.parametrize(
+ "model_variant",
+ PERF_TEST_MODEL_VARIANTS,
+)
+@pytest.mark.parametrize(
+ "image_w, image_h, guidance_scale, num_inference_steps",
+ PERF_TEST_IMAGE_PARAMS,
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp, tp, encoder_tp, vae_tp, topology, num_links",
+ PERF_TEST_MESH_PARAMS,
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ PERF_TEST_DEVICE_PARAMS,
+ indirect=True,
+)
+def test_flux1_pipeline_performance_speed(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ image_w,
+ image_h,
+ guidance_scale,
+ num_inference_steps,
+ sp,
+ tp,
+ encoder_tp,
+ vae_tp,
+ topology,
+ num_links,
+ model_variant,
+ model_location_generator,
+ is_ci_env,
+ galaxy_type,
+) -> None:
+ """Performance test for Flux.1 pipeline with detailed timing analysis."""
+
+ benchmark_profiler = BenchmarkProfiler()
+
+ # Skip 4U.
+ if galaxy_type == "4U":
+ # NOTE: Pipelines fail if a performance test is skipped without providing a benchmark output.
+ if is_ci_env:
+ with benchmark_profiler("run", iteration=0):
+ pass
+
+ benchmark_data = BenchmarkData()
+ benchmark_data.save_partial_run_json(
+ benchmark_profiler,
+ run_type="empty_run",
+ ml_model_name="empty_run",
+ )
+ pytest.skip("4U is not supported for this test")
+
+ logger.info(f" Image size: {image_w}x{image_h}")
+ logger.info(f" Guidance scale: {guidance_scale}")
+ logger.info(f" Inference steps: {num_inference_steps}")
+
+ pipeline = Flux1Pipeline(
+ device=mesh_device,
+ config=Flux1PipelineConfig.default(
+ mesh_shape=mesh_device.shape,
+ dit_parallel_config=DiTParallelConfig.from_tuples(cfg=(1, 0), sp=sp, tp=tp),
+ encoder_parallel_config=EncoderParallelConfig.from_tuple(encoder_tp),
+ vae_parallel_config=VAEParallelConfig.from_tuple(vae_tp),
+ topology=topology,
+ num_links=num_links,
+ width=image_w,
+ height=image_h,
+ checkpoint_name=model_location_generator(f"black-forest-labs/FLUX.1-{model_variant}"),
+ ),
+ )
+
+ # Warmup runs
+ for i in range(2):
+ logger.info("Running warmup iteration...")
+ with benchmark_profiler("run", iteration=0):
+ images = pipeline.run_single_prompt(
+ prompt=TEST_PROMPTS[0],
+ num_inference_steps=num_inference_steps,
+ )
+ images[0].save(f"flux1_{model_variant}_{image_w}_{image_h}_warmup_{i}.png")
+ logger.info(f"Warmup completed in {benchmark_profiler.get_duration('run', 0):.2f}s")
+
+ # Performance measurement runs
+ logger.info("Running performance measurement iterations...")
+ num_perf_runs = 4 # Use 4 different prompts for performance testing
+
+ # Optional Tracy profiling (if available)
+ profiler = None
+ try:
+ from tracy import Profiler
+
+ profiler = Profiler()
+ profiler.enable()
+ logger.info("Tracy profiling enabled")
+ except ImportError:
+ logger.info("Tracy profiler not available, continuing without profiling")
+
+ try:
+ for i in range(num_perf_runs):
+ logger.info(f"Performance run {i+1}/{num_perf_runs}...")
+
+ # Run pipeline with different prompt
+ prompt_idx = (i + 1) % len(TEST_PROMPTS)
+ with benchmark_profiler("run", iteration=i):
+ images = pipeline.run_single_prompt(
+ prompt=TEST_PROMPTS[prompt_idx],
+ num_inference_steps=num_inference_steps,
+ on_event=profiler_event_callback(benchmark_profiler, i),
+ )
+ images[0].save(f"flux1_{model_variant}_{image_w}_{image_h}_perf_run{i}.png")
+
+ # Collect timing data
+ logger.info(f" Run {i+1} completed in {benchmark_profiler.get_duration('run', i):.2f}s")
+
+ finally:
+ if profiler:
+ profiler.disable()
+ logger.info("Tracy profiling disabled")
+
+ # Calculate statistics
+ clip_times = [benchmark_profiler.get_duration("clip_encoding", i) for i in range(num_perf_runs)]
+ t5_times = [benchmark_profiler.get_duration("t5_encoding", i) for i in range(num_perf_runs)]
+ total_encoding_times = [benchmark_profiler.get_duration("encoder", i) for i in range(num_perf_runs)]
+ vae_times = [benchmark_profiler.get_duration("vae", i) for i in range(num_perf_runs)]
+ total_times = [benchmark_profiler.get_duration("run", i) for i in range(num_perf_runs)]
+
+ # Calculate per-step denoising times
+ all_denoising_steps = []
+ for i in range(num_perf_runs):
+ for j in range(num_inference_steps):
+ assert benchmark_profiler.contains_step(
+ f"denoising_step_{j}", i
+ ), f"All runs should have {num_inference_steps} denoising steps"
+ all_denoising_steps.append(benchmark_profiler.get_duration(f"denoising_step_{j}", i))
+
+ # Report results
+ sp_factor = sp[0] # pipeline.dit_parallel_config.sequence_parallel.factor
+ tp_factor = tp[0] # pipeline.dit_parallel_config.tensor_parallel.factor
+ encoder_tp_factor = encoder_tp[0] # pipeline.encoder_parallel_config.tensor_parallel.factor
+ vae_tp_factor = vae_tp[0] # pipeline.vae_parallel_config.tensor_parallel.factor
+
+ print("\n" + "=" * 80)
+ print(f"FLUX.1 {model_variant.upper()} PIPELINE PERFORMANCE RESULTS")
+ print("=" * 80)
+ print(f"Model: FLUX.1-{model_variant}")
+ print(f"Image Size: {image_w}x{image_h}")
+ print(f"Guidance Scale: {guidance_scale}")
+ print(f"Inference Steps: {num_inference_steps}")
+ print(f"Configuration: sp={sp_factor}, tp={tp_factor}, encoder_tp={encoder_tp_factor}, vae_tp={vae_tp_factor}")
+ print(f"Mesh Shape: {mesh_device.shape}")
+ print(f"Topology: {topology}")
+ print("-" * 80)
+
+ def print_stats(name, times):
+ if not times:
+ print(f"{name:25} | No data available")
+ return
+ mean_time = statistics.mean(times)
+ std_time = statistics.stdev(times) if len(times) > 1 else 0
+ min_time = min(times)
+ max_time = max(times)
+ print(
+ f"{name:25} | Mean: {mean_time:8.4f}s | Std: {std_time:8.4f}s | Min: {min_time:8.4f}s | Max: {max_time:8.4f}s"
+ )
+
+ print_stats("CLIP Encoding", clip_times)
+ print_stats("T5 Encoding", t5_times)
+ print_stats("Total Encoding", total_encoding_times)
+ print_stats("Denoising (per step)", all_denoising_steps)
+ print_stats("VAE Decoding", vae_times)
+ print_stats("Total Pipeline", total_times)
+
+ print("-" * 80)
+
+ # Additional metrics
+ if total_times and all_denoising_steps:
+ avg_total_time = statistics.mean(total_times)
+ avg_step_time = statistics.mean(all_denoising_steps)
+ total_denoising_time = avg_step_time * num_inference_steps
+
+ print(f"Average total denoising time: {total_denoising_time:.4f}s")
+ print(f"Denoising throughput: {num_inference_steps / total_denoising_time:.2f} steps/second")
+ print(f"Overall throughput: {1 / avg_total_time:.4f} images/second")
+
+ # Breakdown percentages
+ avg_encoding_time = statistics.mean(total_encoding_times)
+ avg_vae_time = statistics.mean(vae_times)
+
+ print(f"\nTime breakdown:")
+ print(f" Encoding: {avg_encoding_time/avg_total_time*100:.1f}%")
+ print(f" Denoising: {total_denoising_time/avg_total_time*100:.1f}%")
+ print(f" VAE: {avg_vae_time/avg_total_time*100:.1f}%")
+
+ print("=" * 80)
+
+ # Validate performance
+ measurements = {
+ "clip_encoding_time": statistics.mean(clip_times),
+ "t5_encoding_time": statistics.mean(t5_times),
+ "total_encoding_time": statistics.mean(total_encoding_times),
+ "denoising_steps_time": total_denoising_time,
+ "vae_decoding_time": statistics.mean(vae_times),
+ "total_time": statistics.mean(total_times),
+ }
+ if tuple(mesh_device.shape) == (2, 4) and is_blackhole():
+ expected_metrics = {
+ "clip_encoding_time": 0.05,
+ "t5_encoding_time": 0.25,
+ "total_encoding_time": 0.3,
+ "denoising_steps_time": 0.38 * num_inference_steps,
+ "vae_decoding_time": 1.2,
+ "total_time": 12.5,
+ }
+ elif tuple(mesh_device.shape) == (2, 4): # WH.
+ expected_metrics = {
+ "clip_encoding_time": 0.1,
+ "t5_encoding_time": 0.25,
+ "total_encoding_time": 0.3,
+ "denoising_steps_time": 0.75 * num_inference_steps,
+ "vae_decoding_time": 1.6,
+ "total_time": 23,
+ }
+ elif tuple(mesh_device.shape) == (4, 8):
+ expected_metrics = {
+ "clip_encoding_time": 0.1,
+ "t5_encoding_time": 0.21,
+ "total_encoding_time": 0.3,
+ "denoising_steps_time": 0.3 * num_inference_steps,
+ "vae_decoding_time": 1.4,
+ "total_time": 9.0,
+ }
+ elif tuple(mesh_device.shape) == (2, 2):
+ assert is_blackhole(), "2x2 is only supported for blackhole"
+ expected_metrics = {
+ "clip_encoding_time": 0.05,
+ "t5_encoding_time": 0.25,
+ "total_encoding_time": 0.3,
+ "denoising_steps_time": 0.55 * num_inference_steps,
+ "vae_decoding_time": 1.8,
+ "total_time": 17.5,
+ }
+ elif tuple(mesh_device.shape) == (1, 2):
+ assert is_blackhole(), "1x2 is only supported for blackhole"
+ expected_metrics = {
+ "clip_encoding_time": 0.1,
+ "t5_encoding_time": 0.25,
+ "total_encoding_time": 0.3,
+ "denoising_steps_time": 1.1 * num_inference_steps,
+ "vae_decoding_time": 1.6,
+ "total_time": 29.0,
+ }
+ else:
+ assert False, f"Unknown mesh device for performance comparison: {mesh_device}"
+
+ if is_ci_env:
+ # In CI, dump a performance report
+ benchmark_data = BenchmarkData()
+ for iteration in range(num_perf_runs):
+ for step_name, target in zip(
+ ["encoder", "denoising", "vae", "run"],
+ [
+ None,
+ expected_metrics["denoising_steps_time"],
+ expected_metrics["vae_decoding_time"],
+ expected_metrics["total_time"],
+ ],
+ ):
+ benchmark_data.add_measurement(
+ profiler=benchmark_profiler,
+ iteration=iteration,
+ step_name=step_name,
+ name=step_name,
+ value=benchmark_profiler.get_duration(step_name, iteration),
+ target=target,
+ )
+ device_name_map = {
+ (1, 2): "BH_P300",
+ (2, 2): "BH_QB",
+ (2, 4): "BH_LB" if is_blackhole() else "WH_T3K",
+ (4, 8): "BH_GLX" if is_blackhole() else "WH_GLX",
+ }
+ benchmark_data.save_partial_run_json(
+ benchmark_profiler,
+ run_type=device_name_map[tuple(mesh_device.shape)],
+ ml_model_name=f"Flux1{model_variant.capitalize()}",
+ batch_size=1,
+ config_params={
+ "width": image_w,
+ "height": image_h,
+ "num_frames": 1,
+ "num_steps": num_inference_steps,
+ "sp_factor": sp_factor,
+ "tp_factor": tp_factor,
+ "topology": str(topology),
+ "num_links": num_links,
+ "fsdp": False,
+ },
+ )
+
+ pass_perf_check = True
+ assert_msgs = []
+ for k in expected_metrics.keys():
+ if measurements[k] > expected_metrics[k]:
+ assert_msgs.append(
+ f"Warning: {k} is outside of the tolerance range. Expected: {expected_metrics[k]}, Actual: {measurements[k]}"
+ )
+ pass_perf_check = False
+
+ assert pass_perf_check, "\n".join(assert_msgs)
+
+ # Synchronize all devices
+ pipeline.synchronize_devices()
+
+ logger.info("Performance test completed successfully!")
+
+
+@pytest.mark.parametrize(
+ "model_variant",
+ PERF_TEST_MODEL_VARIANTS,
+)
+@pytest.mark.parametrize(
+ "image_w, image_h, guidance_scale, num_inference_steps",
+ PERF_TEST_IMAGE_PARAMS,
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp, tp, encoder_tp, vae_tp, topology, num_links",
+ PERF_TEST_MESH_PARAMS,
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ PERF_TEST_DEVICE_PARAMS,
+ indirect=True,
+)
+def test_flux1_pipeline_performance_accuracy(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ image_w,
+ image_h,
+ guidance_scale,
+ num_inference_steps,
+ sp,
+ tp,
+ encoder_tp,
+ vae_tp,
+ topology,
+ num_links,
+ model_variant,
+ model_location_generator,
+ is_ci_env,
+ galaxy_type,
+) -> None:
+ pipeline = Flux1Pipeline(
+ device=mesh_device,
+ config=Flux1PipelineConfig.default(
+ mesh_shape=mesh_device.shape,
+ dit_parallel_config=DiTParallelConfig.from_tuples(cfg=(1, 0), sp=sp, tp=tp),
+ encoder_parallel_config=EncoderParallelConfig.from_tuple(encoder_tp),
+ vae_parallel_config=VAEParallelConfig.from_tuple(vae_tp),
+ topology=topology,
+ num_links=num_links,
+ width=image_w,
+ height=image_h,
+ checkpoint_name=model_location_generator(f"black-forest-labs/FLUX.1-{model_variant}"),
+ ),
+ )
+
+ # Warmup runs
+ for i in range(2):
+ logger.info("Running warmup iteration...")
+ pipeline.run_single_prompt(
+ prompt=TEST_PROMPTS[0],
+ num_inference_steps=num_inference_steps,
+ )
+
+ output_images = []
+ for i in range(len(TEST_PROMPTS)):
+ images = pipeline.run_single_prompt(
+ prompt=TEST_PROMPTS[i],
+ num_inference_steps=num_inference_steps,
+ )
+ images[0].save(f"flux1_{model_variant}_{image_w}_{image_h}_perf_run{i}.png")
+ output_images.append(images[0])
+
+ clip = CLIPEncoder()
+ clip_scores = []
+
+ for prompt, image in zip(TEST_PROMPTS, output_images):
+ clip_scores.append(100 * clip.get_clip_score(prompt, image).item())
+
+ for prompt, score in zip(TEST_PROMPTS, clip_scores):
+ logger.info(f'Prompt "{prompt[:50]}". Score {score}')
+ # TODO: At some point this was above 30 for all prompts. I wasn't able to find out what caused
+ # the regression, so doing the hacky thing and reducing the value for now.
+ # Possible culprit (https://github.com/tenstorrent/tt-metal/actions/runs/27652185143) this
+ # regression occurred at the same time as this went down, but unclear if the two are related.
+ assert score > 28.0
+
+
+@pytest.mark.parametrize(
+ "model_variant",
+ PERF_TEST_MODEL_VARIANTS,
+)
+@pytest.mark.parametrize(
+ "image_w, image_h, guidance_scale, num_inference_steps",
+ PERF_TEST_IMAGE_PARAMS,
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp, tp, encoder_tp, vae_tp, topology, num_links",
+ PERF_TEST_MESH_PARAMS,
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ PERF_TEST_DEVICE_PARAMS,
+ indirect=True,
+)
+def test_flux1_pipeline_performance_determinism(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ image_w,
+ image_h,
+ guidance_scale,
+ num_inference_steps,
+ sp,
+ tp,
+ encoder_tp,
+ vae_tp,
+ topology,
+ num_links,
+ model_variant,
+ model_location_generator,
+ is_ci_env,
+ galaxy_type,
+) -> None:
+ pipeline = Flux1Pipeline(
+ device=mesh_device,
+ config=Flux1PipelineConfig.default(
+ mesh_shape=mesh_device.shape,
+ dit_parallel_config=DiTParallelConfig.from_tuples(cfg=(1, 0), sp=sp, tp=tp),
+ encoder_parallel_config=EncoderParallelConfig.from_tuple(encoder_tp),
+ vae_parallel_config=VAEParallelConfig.from_tuple(vae_tp),
+ topology=topology,
+ num_links=num_links,
+ width=image_w,
+ height=image_h,
+ checkpoint_name=model_location_generator(f"black-forest-labs/FLUX.1-{model_variant}"),
+ ),
+ )
+
+ import numpy as np
+
+ output_images = []
+ num_iters = 3
+ for i in range(num_iters):
+ images = pipeline.run_single_prompt(
+ prompt=TEST_PROMPTS[0],
+ num_inference_steps=num_inference_steps,
+ )
+ images[0].save(f"flux1_{model_variant}_{image_w}_{image_h}_perf_run{i}.png")
+ output_images.append(np.array(images[0]))
+
+ reference = output_images[0]
+ for i, img_array in enumerate(output_images[1:], start=1):
+ assert np.array_equal(
+ reference, img_array
+ ), f"Image from iteration {i} differs from iteration 0 (pixel-level mismatch)"
diff --git a/code/models/tt_dit/tests/models/flux1/test_pipeline_flux1.py b/code/models/tt_dit/tests/models/flux1/test_pipeline_flux1.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b67e7654c247638fc059ebafb0c558b7136062b
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux1/test_pipeline_flux1.py
@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+import os
+
+import pytest
+from loguru import logger
+
+import ttnn
+from models.perf.benchmarking_utils import BenchmarkProfiler
+
+from ....models.transformers.transformer_flux1 import Flux1Checkpoint
+from ....parallel.config import DiTParallelConfig, EncoderParallelConfig, VAEParallelConfig
+from ....pipelines.events import profiler_event_callback
+from ....pipelines.flux1.pipeline_flux1 import Flux1Pipeline, Flux1PipelineConfig
+
+
+@pytest.mark.parametrize(
+ "no_prompt",
+ [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), False)],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "l1_small_size": 32768,
+ "trace_region_size": 50000000,
+ "require_exact_physical_num_devices": True,
+ }
+ ],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ ("model_variant", "width", "height", "num_inference_steps"),
+ [
+ ("schnell", 1024, 1024, 4),
+ ("dev", 1024, 1024, 28),
+ ],
+ ids=["flux_schnell", "flux_dev"],
+)
+@pytest.mark.parametrize(
+ ("mesh_device", "sp", "tp", "encoder_tp", "vae_tp", "topology", "num_links", "mesh_test_id"),
+ [
+ pytest.param((1, 2), (1, 0), (2, 1), (2, 1), (2, 1), ttnn.Topology.Linear, 2, "1x2sp0tp1", id="1x2sp0tp1"),
+ pytest.param((2, 2), (2, 0), (2, 1), (2, 1), (2, 1), ttnn.Topology.Linear, 2, "2x2sp0tp1", id="2x2sp0tp1"),
+ pytest.param((2, 4), (2, 0), (4, 1), (4, 1), (4, 1), ttnn.Topology.Linear, 1, "2x4sp0tp1", id="2x4sp0tp1"),
+ pytest.param((4, 8), (4, 0), (8, 1), (4, 0), (4, 0), ttnn.Topology.Linear, 4, "4x8sp0tp1", id="4x8sp0tp1"),
+ pytest.param((4, 8), (8, 1), (4, 0), (4, 0), (4, 0), ttnn.Topology.Linear, 4, "4x8sp1tp0", id="4x8sp1tp0"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ ("enable_t5_text_encoder", "use_torch_t5_text_encoder", "use_torch_clip_text_encoder"),
+ [
+ # pytest.param(True, True, True, id="encoder_cpu"),
+ pytest.param(True, False, False, id="encoder_device"),
+ ],
+)
+@pytest.mark.parametrize(
+ "traced",
+ [
+ pytest.param(True, id="traced"),
+ pytest.param(False, id="not_traced"),
+ ],
+)
+@pytest.mark.parametrize(
+ "use_cache",
+ [
+ pytest.param(True, id="yes_use_cache"),
+ pytest.param(False, id="no_use_cache"),
+ ],
+)
+def test_flux1_pipeline(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ model_variant: str,
+ width: int,
+ height: int,
+ num_inference_steps: int,
+ sp: tuple[int, int],
+ tp: tuple[int, int],
+ encoder_tp: tuple[int, int],
+ vae_tp: tuple[int, int],
+ topology: ttnn.Topology,
+ num_links: int,
+ no_prompt: bool,
+ enable_t5_text_encoder: bool,
+ use_torch_t5_text_encoder: bool,
+ use_torch_clip_text_encoder: bool,
+ model_location_generator,
+ traced: bool,
+ mesh_test_id: str,
+ use_cache: bool,
+ is_ci_env: bool,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ # Setup CI environment
+ if is_ci_env and traced:
+ pytest.skip("Skipping traced test in CI environment. Use Performance test for detailed timing analysis.")
+
+ parallel_config = DiTParallelConfig.from_tuples(cfg=(1, 0), sp=sp, tp=tp)
+ encoder_parallel_config = EncoderParallelConfig.from_tuple(encoder_tp)
+ vae_parallel_config = VAEParallelConfig.from_tuple(vae_tp)
+
+ logger.info(f"Mesh device shape: {mesh_device.shape}")
+ logger.info(f"Parallel config: {parallel_config}")
+ logger.info(f"Encoder parallel config: {encoder_parallel_config}")
+ logger.info(f"VAE parallel config: {vae_parallel_config}")
+ logger.info(f"T5 enabled: {enable_t5_text_encoder}")
+ logger.info(f"HF_HUB_OFFLINE: {os.environ.get('HF_HUB_OFFLINE', '')}")
+
+ pipeline = Flux1Pipeline(
+ device=mesh_device,
+ config=Flux1PipelineConfig.default(
+ mesh_shape=mesh_device.shape,
+ dit_parallel_config=parallel_config,
+ encoder_parallel_config=encoder_parallel_config,
+ vae_parallel_config=vae_parallel_config,
+ enable_t5_text_encoder=enable_t5_text_encoder,
+ use_torch_t5_text_encoder=use_torch_t5_text_encoder,
+ use_torch_clip_text_encoder=use_torch_clip_text_encoder,
+ num_links=num_links,
+ topology=topology,
+ width=width,
+ height=height,
+ checkpoint_name=model_location_generator(f"black-forest-labs/FLUX.1-{model_variant}"),
+ ),
+ )
+
+ prompts = [
+ "A luxury sports car.",
+ # "Neon-lit cyberpunk alley, rain-soaked, cinematic wide shot",
+ # "Golden retriever astronaut drifting in sunlit space",
+ # "Minimalist Scandinavian kitchen, morning light, ultra clean",
+ # "Ancient desert temple at dawn, soft fog, wide angle",
+ # "Steampunk airship over Victorian city, dramatic clouds",
+ # "Macro dewdrops on fern, shallow depth of field",
+ # "Luxury wristwatch on marble, studio lighting, hyper-detail",
+ # "Stormy coastline lighthouse, crashing waves, long exposure",
+ # "Futuristic Tokyo street market, vibrant signage, motion blur",
+ ]
+
+ filename_prefix = f"flux_{model_variant}_{width}_{height}_{mesh_test_id}"
+ if enable_t5_text_encoder:
+ if use_torch_t5_text_encoder:
+ filename_prefix += "_t5cpu"
+ else:
+ filename_prefix += "_t5off"
+ if use_torch_clip_text_encoder:
+ filename_prefix += "_clipcpu"
+ if not traced:
+ filename_prefix += "_untraced"
+
+ def run(*, prompt: str, number: int, seed: int) -> None:
+ benchmark_profiler = BenchmarkProfiler()
+ with benchmark_profiler("run", iteration=0):
+ images = pipeline(
+ prompts=[prompt],
+ num_inference_steps=num_inference_steps,
+ seed=seed,
+ traced=traced,
+ vae_traced=False,
+ encoder_traced=False,
+ on_event=profiler_event_callback(benchmark_profiler, 0),
+ )
+
+ output_filename = f"{filename_prefix}_{number}.png"
+ images[0].save(output_filename)
+ logger.info(f"Image saved as {output_filename}")
+
+ logger.info(f"CLIP encoding time: {benchmark_profiler.get_duration('clip_encoding', 0):.2f}s")
+ logger.info(f"T5 encoding time: {benchmark_profiler.get_duration('t5_encoding', 0):.2f}s")
+ logger.info(f"Total encoding time: {benchmark_profiler.get_duration('encoder', 0):.2f}s")
+ logger.info(f"VAE decoding time: {benchmark_profiler.get_duration('vae', 0):.2f}s")
+ logger.info(f"Total pipeline time: {benchmark_profiler.get_duration('total', 0):.2f}s")
+ avg_step_time = benchmark_profiler.get_duration("denoising", 0) / num_inference_steps
+ logger.info(f"Average denoising step time: {avg_step_time:.2f}s")
+
+ if no_prompt:
+ for i, prompt in enumerate(prompts):
+ run(prompt=prompt, number=i, seed=0)
+ else:
+ prompt = prompts[0]
+ 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)
+
+
+def _log_hf_checkpoint_location(checkpoint_name) -> None:
+ """Log where HuggingFace will look for / resolve the checkpoint files."""
+ import huggingface_hub.constants as hf_const
+ from huggingface_hub import try_to_load_from_cache
+
+ logger.info(f"Resolved checkpoint_name (passed to from_pretrained): {checkpoint_name}")
+ logger.info(f"HF_HOME = {hf_const.HF_HOME}")
+ logger.info(f"HF_HUB_CACHE = {hf_const.HF_HUB_CACHE}")
+
+ # If checkpoint_name is a local path, files load from there directly.
+ if os.path.isdir(checkpoint_name):
+ logger.info(f"checkpoint_name is a local directory; loading from: {checkpoint_name}")
+ return
+
+ # Otherwise it's an HF repo id resolved out of the hub cache. Show the cached file paths.
+ for rel in ("transformer/config.json", "model_index.json"):
+ cached = try_to_load_from_cache(checkpoint_name, rel)
+ logger.info(f" {rel} -> {cached}")
+
+
+def test_flux1_dev_checkpoint_load(model_location_generator) -> None:
+ """Loads Flux1Checkpoint and verifies key metadata/state are available."""
+ checkpoint_name = model_location_generator("black-forest-labs/FLUX.1-dev")
+ _log_hf_checkpoint_location(checkpoint_name)
+ Flux1Checkpoint(checkpoint_name)
+
+
+def test_flux1_schnell_checkpoint_load(model_location_generator) -> None:
+ """Loads Flux1Checkpoint and verifies key metadata/state are available."""
+ checkpoint_name = model_location_generator("black-forest-labs/FLUX.1-schnell")
+ _log_hf_checkpoint_location(checkpoint_name)
+ Flux1Checkpoint(checkpoint_name)
diff --git a/code/models/tt_dit/tests/models/flux1/test_transformer_block_flux1.py b/code/models/tt_dit/tests/models/flux1/test_transformer_block_flux1.py
new file mode 100644
index 0000000000000000000000000000000000000000..8541b759132ec7cacb3faaeda8d88dce5f072eb0
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux1/test_transformer_block_flux1.py
@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from time import time
+
+import diffusers.models.transformers.transformer_flux
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+
+from ....blocks.attention import Attention
+from ....blocks.transformer_block import TransformerBlock
+from ....parallel.config import DiTParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.padding import PaddingConfig
+from ....utils.tensor import bf16_tensor, bf16_tensor_2dshard
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "sp_axis", "tp_axis", "num_links", "mesh_id"),
+ [
+ pytest.param((1, 2), (1, 2), 0, 1, 2, "1x2sp0tp1", id="1x2sp0tp1"),
+ pytest.param((2, 2), (2, 2), 0, 1, 2, "2x2sp0tp1", id="2x2sp0tp1"),
+ pytest.param((2, 4), (1, 4), 0, 1, 1, "1x4sp0tp1", id="1x4sp0tp1"),
+ pytest.param((2, 4), (2, 4), 0, 1, 1, "2x4sp0tp1", id="2x4sp0tp1"),
+ pytest.param((2, 4), (2, 4), 1, 0, 1, "2x4sp1tp0", id="2x4sp1tp0"),
+ pytest.param((4, 8), (4, 4), 0, 1, 4, "4x4sp0tp1", id="4x4sp0tp1"),
+ pytest.param((4, 8), (4, 8), 0, 1, 4, "4x8sp0tp1", id="4x8sp0tp1"),
+ pytest.param((4, 8), (4, 8), 1, 0, 4, "4x8sp1tp0", id="4x8sp1tp0"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "spatial_seq_len", "prompt_seq_len"),
+ [
+ (1, 4096, 512),
+ ],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "require_exact_physical_num_devices": True,
+ }
+ ],
+ indirect=True,
+)
+def test_transformer_block_flux1(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ submesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ batch_size: int,
+ spatial_seq_len: int,
+ prompt_seq_len: int,
+ mesh_id: str,
+ is_ci_env: bool,
+) -> None:
+ torch.manual_seed(0)
+
+ submesh_device = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ sp_factor = tuple(submesh_device.shape)[sp_axis]
+ tp_factor = tuple(submesh_device.shape)[tp_axis]
+
+ inner_dim = 3072
+ head_dim = 128
+ num_heads = 24
+
+ torch_model = diffusers.models.transformers.transformer_flux.FluxTransformerBlock(
+ dim=inner_dim, num_attention_heads=num_heads, attention_head_dim=head_dim
+ )
+ torch_model.eval()
+
+ ccl_manager = CCLManager(
+ mesh_device=submesh_device,
+ num_links=num_links,
+ topology=ttnn.Topology.Linear,
+ )
+
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=0, mesh_axis=0),
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+
+ if num_heads % tp_factor != 0:
+ padding_config = PaddingConfig.from_tensor_parallel_factor(num_heads, head_dim, tp_factor)
+ else:
+ padding_config = None
+
+ tt_model = TransformerBlock(
+ dim=inner_dim,
+ num_heads=num_heads,
+ head_dim=head_dim,
+ context_pre_only=False,
+ mesh_device=submesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ padding_config=padding_config,
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ spatial = torch.randn([batch_size, spatial_seq_len, inner_dim])
+ prompt = torch.randn([batch_size, prompt_seq_len, inner_dim])
+ time_embed = torch.randn([batch_size, inner_dim])
+ rope_cos = torch.randn([spatial_seq_len + prompt_seq_len, head_dim])
+ rope_sin = torch.randn([spatial_seq_len + prompt_seq_len, head_dim])
+
+ spatial_padded = Attention.pad_spatial_sequence(spatial, sp_factor=sp_factor)
+ spatial_rope_cos_padded = Attention.pad_spatial_sequence(rope_cos[prompt_seq_len:], sp_factor=sp_factor)
+ spatial_rope_sin_padded = Attention.pad_spatial_sequence(rope_sin[prompt_seq_len:], sp_factor=sp_factor)
+
+ tt_spatial = bf16_tensor_2dshard(spatial_padded, device=submesh_device, shard_mapping={sp_axis: 1, tp_axis: 2})
+ tt_prompt = bf16_tensor(prompt, device=submesh_device, mesh_axis=tp_axis, shard_dim=2)
+ tt_time_embed = bf16_tensor(time_embed.unsqueeze(1), device=submesh_device)
+ tt_spatial_rope_cos = bf16_tensor(spatial_rope_cos_padded, device=submesh_device, mesh_axis=sp_axis, shard_dim=0)
+ tt_spatial_rope_sin = bf16_tensor(spatial_rope_sin_padded, device=submesh_device, mesh_axis=sp_axis, shard_dim=0)
+ tt_prompt_rope_cos = bf16_tensor(rope_cos[:prompt_seq_len], device=submesh_device)
+ tt_prompt_rope_sin = bf16_tensor(rope_sin[:prompt_seq_len], device=submesh_device)
+
+ with torch.no_grad():
+ torch_prompt, torch_spatial = torch_model.forward(
+ spatial, prompt, temb=time_embed, image_rotary_emb=(rope_cos, rope_sin)
+ )
+
+ if not is_ci_env:
+ try:
+ from tracy import signpost
+
+ signpost("caching")
+ tt_spatial_out, tt_prompt_out = tt_model.forward(
+ tt_spatial,
+ tt_prompt,
+ tt_time_embed,
+ spatial_sequence_length=spatial_seq_len,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ )
+
+ signpost("performance")
+ except ImportError:
+ logger.info("Tracy profiler not available, continuing without profiling")
+
+ itr = 1
+ start = time()
+ for _ in range(itr):
+ tt_spatial_out, tt_prompt_out = tt_model.forward(
+ tt_spatial,
+ tt_prompt,
+ tt_time_embed,
+ spatial_sequence_length=spatial_seq_len,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ )
+
+ ttnn.synchronize_device(submesh_device)
+ logger.info(f"Time taken for {mesh_id}: {(time() - start) * 1000 / itr} ms")
+
+ tt_spatial_torch = tensor.to_torch(tt_spatial_out, mesh_axes=[None, sp_axis, tp_axis])
+ assert_quality(torch_spatial, tt_spatial_torch, pcc=0.99998, relative_rmse=0.006)
+
+ tt_prompt_torch = tensor.to_torch(tt_prompt_out, mesh_axes=[None, None, tp_axis])
+ assert_quality(torch_prompt, tt_prompt_torch, pcc=0.99998, relative_rmse=0.006)
diff --git a/code/models/tt_dit/tests/models/flux1/test_transformer_flux1.py b/code/models/tt_dit/tests/models/flux1/test_transformer_flux1.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f3f9f81762845851cb3b7fff4f7f371e389c2f4
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux1/test_transformer_flux1.py
@@ -0,0 +1,326 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# import diffusers.models.transformers.transformer_flux as reference
+from time import time
+
+import diffusers as reference
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+
+from ....models.transformers.transformer_flux1 import Flux1Checkpoint, Flux1SingleTransformerBlock
+from ....parallel.config import DiTParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils.check import assert_quality
+from ....utils.padding import PaddingConfig
+from ....utils.tensor import bf16_tensor, bf16_tensor_2dshard
+
+TEST_MESH_PARAMS = [
+ # BH Only
+ pytest.param((1, 2), (1, 2), 0, 1, 2, "1x2sp0tp1", id="1x2sp0tp1"),
+ pytest.param((2, 2), (2, 2), 0, 1, 2, "2x2sp0tp1", id="2x2sp0tp1"),
+ # WH Only
+ pytest.param((2, 4), (2, 4), 0, 1, 1, "2x4sp0tp1", id="2x4sp0tp1"),
+ pytest.param((2, 4), (2, 4), 1, 0, 1, "2x4sp1tp0", id="2x4sp1tp0"),
+ # BH and WH
+ pytest.param((4, 8), (4, 4), 0, 1, 4, "4x4sp0tp1", id="4x4sp0tp1"),
+ pytest.param((4, 8), (4, 8), 0, 1, 4, "4x8sp0tp1", id="4x8sp0tp1"),
+ pytest.param((4, 8), (4, 8), 1, 0, 4, "4x8sp1tp0", id="4x8sp1tp0"),
+]
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "sp_axis", "tp_axis", "num_links", "id"),
+ TEST_MESH_PARAMS,
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "spatial_seq_len", "prompt_seq_len"),
+ [
+ (1, 4096, 512),
+ ],
+)
+@pytest.mark.parametrize(
+ "model_variant",
+ ["schnell", "dev"],
+ ids=["flux_schnell", "flux_dev"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "require_exact_physical_num_devices": True,
+ }
+ ],
+ indirect=True,
+)
+def test_transformer_single_block(
+ mesh_device: ttnn.MeshDevice,
+ submesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ batch_size: int,
+ prompt_seq_len: int,
+ spatial_seq_len: int,
+ model_variant: str,
+ id: str,
+ model_location_generator,
+ is_ci_env: bool,
+) -> None:
+ submesh_device = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ sp_factor = tuple(submesh_device.shape)[sp_axis]
+ tp_factor = tuple(submesh_device.shape)[tp_axis]
+ model_name = model_location_generator(f"black-forest-labs/FLUX.1-{model_variant}", model_subdir="transformer")
+ parent_torch_model = reference.FluxTransformer2DModel.from_pretrained(
+ model_name, subfolder="transformer", torch_dtype=torch.bfloat16
+ )
+ torch_model = parent_torch_model.single_transformer_blocks[0]
+ assert isinstance(torch_model, reference.models.transformers.transformer_flux.FluxSingleTransformerBlock)
+ torch_model.eval()
+
+ inner_dim = torch_model.attn.inner_dim
+ num_heads = torch_model.attn.heads
+ head_dim = inner_dim // num_heads
+
+ ccl_manager = CCLManager(
+ mesh_device=submesh_device,
+ num_links=num_links,
+ topology=ttnn.Topology.Linear,
+ )
+
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=0, mesh_axis=0),
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+
+ if num_heads % tp_factor != 0:
+ padding_config = PaddingConfig.from_tensor_parallel_factor(num_heads, head_dim, tp_factor)
+ else:
+ padding_config = None
+
+ tt_model = Flux1SingleTransformerBlock(
+ dim=inner_dim,
+ num_heads=num_heads,
+ head_dim=head_dim,
+ mesh_device=submesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ padding_config=padding_config,
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ torch.manual_seed(0)
+ combined = torch.randn([batch_size, prompt_seq_len + spatial_seq_len, inner_dim]).to(torch.bfloat16)
+ time_embed = torch.randn([batch_size, inner_dim]).to(torch.bfloat16)
+ rope_cos = torch.randn([prompt_seq_len + spatial_seq_len, head_dim]).to(torch.bfloat16)
+ rope_sin = torch.randn([prompt_seq_len + spatial_seq_len, head_dim]).to(torch.bfloat16)
+
+ tt_spatial = bf16_tensor_2dshard(
+ combined[:, prompt_seq_len:], device=submesh_device, shard_mapping={sp_axis: 1, tp_axis: 2}
+ )
+ tt_prompt = bf16_tensor(combined[:, :prompt_seq_len], device=submesh_device, mesh_axis=tp_axis, shard_dim=2)
+ tt_time_embed = bf16_tensor(time_embed.unsqueeze(1), device=submesh_device)
+ tt_spatial_rope_cos = bf16_tensor(rope_cos[prompt_seq_len:], device=submesh_device, mesh_axis=sp_axis, shard_dim=0)
+ tt_spatial_rope_sin = bf16_tensor(rope_sin[prompt_seq_len:], device=submesh_device, mesh_axis=sp_axis, shard_dim=0)
+ tt_prompt_rope_cos = bf16_tensor(rope_cos[:prompt_seq_len], device=submesh_device)
+ tt_prompt_rope_sin = bf16_tensor(rope_sin[:prompt_seq_len], device=submesh_device)
+
+ with torch.no_grad():
+ torch_combined = torch_model.forward(
+ hidden_states=combined[:, prompt_seq_len:],
+ encoder_hidden_states=combined[:, :prompt_seq_len],
+ temb=time_embed,
+ image_rotary_emb=(rope_cos, rope_sin),
+ )
+
+ if not is_ci_env:
+ try:
+ from tracy import signpost
+
+ signpost("caching")
+ tt_spatial_out, tt_prompt_out = tt_model.forward(
+ spatial=tt_spatial,
+ prompt=tt_prompt,
+ time_embed=tt_time_embed,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ spatial_sequence_length=spatial_seq_len,
+ )
+
+ signpost("performance")
+ except ImportError:
+ logger.info("Tracy profiler not available, continuing without profiling")
+
+ itr = 1
+ start = time()
+ for _ in range(itr):
+ tt_spatial_out, tt_prompt_out = tt_model.forward(
+ spatial=tt_spatial,
+ prompt=tt_prompt,
+ time_embed=tt_time_embed,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ spatial_sequence_length=spatial_seq_len,
+ )
+ # ttnn.synchronize_device(submesh_device)
+ ttnn.synchronize_device(submesh_device)
+ logger.info(f"Time taken for {id}: {(time() - start)*1000/itr} ms")
+
+ shard_dims = [None, None]
+ shard_dims[sp_axis], shard_dims[tp_axis] = 1, 2
+ tt_spatial_torch = ttnn.to_torch(
+ tt_spatial_out,
+ mesh_composer=ttnn.create_mesh_composer(submesh_device, ttnn.MeshComposerConfig(shard_dims)),
+ )
+
+ shard_dims = [None, None]
+ shard_dims[sp_axis], shard_dims[tp_axis] = 0, 2
+ tt_prompt_torch = ttnn.to_torch(
+ tt_prompt_out,
+ mesh_composer=ttnn.create_mesh_composer(submesh_device, ttnn.MeshComposerConfig(shard_dims)),
+ )[:batch_size]
+
+ tt_combined_torch = torch.concat([tt_prompt_torch, tt_spatial_torch], dim=1)
+ torch_combined = torch.concat(torch_combined, dim=1)
+
+ assert_quality(torch_combined, tt_combined_torch, pcc=0.99959, relative_rmse=6.3)
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "sp_axis", "tp_axis", "num_links", "id"),
+ TEST_MESH_PARAMS,
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "spatial_seq_len", "prompt_seq_len"),
+ [
+ (1, 4096, 512),
+ ],
+)
+@pytest.mark.parametrize(
+ "model_variant",
+ ["schnell", "dev"],
+ ids=["flux_schnell", "flux_dev"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "trace_region_size": 31000000,
+ "require_exact_physical_num_devices": True,
+ },
+ ],
+ indirect=True,
+)
+def test_transformer(
+ mesh_device: ttnn.MeshDevice,
+ submesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ batch_size: int,
+ spatial_seq_len: int,
+ prompt_seq_len: int,
+ model_variant: str,
+ id: str,
+ model_location_generator,
+) -> None:
+ submesh_device = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+
+ sp_factor = tuple(submesh_device.shape)[sp_axis]
+ tp_factor = tuple(submesh_device.shape)[tp_axis]
+
+ model_name = model_location_generator(f"black-forest-labs/FLUX.1-{model_variant}", model_subdir="transformer")
+ torch_model = reference.FluxTransformer2DModel.from_pretrained(
+ model_name, subfolder="transformer", torch_dtype=torch.bfloat16
+ )
+ assert isinstance(torch_model, reference.FluxTransformer2DModel)
+ torch_model.eval()
+
+ in_channels = torch_model.config.in_channels
+ joint_attention_dim = torch_model.config.joint_attention_dim
+ pooled_projection_dim = torch_model.config.pooled_projection_dim
+ with_guidance_embeds = torch_model.config.guidance_embeds
+
+ ccl_manager = CCLManager(
+ mesh_device=submesh_device,
+ num_links=num_links,
+ topology=ttnn.Topology.Linear,
+ )
+
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=0, mesh_axis=0),
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+
+ checkpoint = Flux1Checkpoint(model_name)
+ tt_model = checkpoint.build(ccl_manager=ccl_manager, parallel_config=parallel_config)
+
+ torch.manual_seed(0)
+ spatial = torch.randn([batch_size, spatial_seq_len, in_channels]).to(torch.bfloat16)
+ prompt = torch.randn([batch_size, prompt_seq_len, joint_attention_dim]).to(torch.bfloat16)
+ pooled = torch.randn([batch_size, pooled_projection_dim]).to(torch.bfloat16)
+ timestep = torch.full([batch_size], fill_value=500).to(torch.bfloat16)
+ guidance = torch.full([batch_size], fill_value=3).to(torch.bfloat16) if with_guidance_embeds else None
+
+ # prepare for ROPE
+ text_ids = torch.zeros([prompt_seq_len, 3]).to(torch.bfloat16)
+ image_ids = torch.randint(1024 * 1024, [spatial_seq_len, 3]).to(torch.bfloat16)
+ ids = torch.cat((text_ids, image_ids), dim=0).to(torch.bfloat16)
+ rope_cos, rope_sin = torch_model.pos_embed.forward(ids)
+
+ tt_spatial = bf16_tensor(spatial, device=submesh_device, mesh_axis=sp_axis, shard_dim=1)
+ tt_prompt = bf16_tensor(prompt, device=submesh_device)
+ tt_pooled = bf16_tensor(pooled, device=submesh_device)
+ tt_timestep = ttnn.from_torch(
+ timestep.unsqueeze(-1), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=submesh_device
+ )
+ tt_guidance = bf16_tensor(guidance.unsqueeze(-1), device=submesh_device) if guidance is not None else None
+
+ tt_spatial_rope_cos = bf16_tensor(rope_cos[prompt_seq_len:], device=submesh_device, mesh_axis=sp_axis, shard_dim=0)
+ tt_spatial_rope_sin = bf16_tensor(rope_sin[prompt_seq_len:], device=submesh_device, mesh_axis=sp_axis, shard_dim=0)
+ tt_prompt_rope_cos = bf16_tensor(rope_cos[:prompt_seq_len], device=submesh_device)
+ tt_prompt_rope_sin = bf16_tensor(rope_sin[:prompt_seq_len], device=submesh_device)
+
+ logger.info("running torch model...")
+ with torch.no_grad():
+ torch_output = torch_model.forward(
+ hidden_states=spatial,
+ encoder_hidden_states=prompt,
+ pooled_projections=pooled,
+ timestep=timestep / 1000,
+ guidance=guidance,
+ img_ids=image_ids,
+ txt_ids=text_ids,
+ ).sample
+
+ logger.info("running TT model...")
+ tt_output = tt_model.forward(
+ spatial=tt_spatial,
+ prompt=tt_prompt,
+ pooled=tt_pooled,
+ timestep=tt_timestep,
+ guidance=tt_guidance,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ spatial_sequence_length=spatial_seq_len,
+ prompt_sequence_length=prompt_seq_len,
+ )
+
+ shard_dims = [None, None]
+ shard_dims[sp_axis], shard_dims[tp_axis] = 1, 0
+ tt_output_torch = ttnn.to_torch(
+ tt_output,
+ mesh_composer=ttnn.create_mesh_composer(submesh_device, ttnn.MeshComposerConfig(shard_dims)),
+ )[:batch_size]
+
+ assert_quality(torch_output, tt_output_torch, pcc=0.997, relative_rmse=8.1)
diff --git a/code/models/tt_dit/tests/models/flux1/test_vae.py b/code/models/tt_dit/tests/models/flux1/test_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..890eef8ed8b984ec991fbb3619b3ca7f43a40bf4
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux1/test_vae.py
@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import diffusers.models.autoencoders.autoencoder_kl_qwenimage as reference
+import pytest
+import torch
+
+import ttnn
+
+from ...models.vae.vae import VaeAttention, VaeContext, VaeNormDescRms, VaeRmsNorm
+from ...parallel.manager import CCLManager
+from ...utils import tensor
+from ...utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "tp_axis"),
+ [
+ pytest.param((1, 1), None, id="1x1"),
+ pytest.param((1, 2), None, id="1x2"),
+ # pytest.param((1, 2), 1, id="1x2tp"), # hangs, resolved by https://github.com/tenstorrent/tt-metal/pull/33254
+ pytest.param((1, 4), 1, id="1x4tp"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "l1_small_size": 65536, "trace_region_size": 20000000}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ ("batch_size", "height", "width", "channels"),
+ [
+ (2, 128, 128, 384),
+ ],
+)
+def test_vae_rms_norm(
+ *,
+ mesh_device: ttnn.Device,
+ batch_size: int,
+ height: int,
+ width: int,
+ channels: int,
+ tp_axis: int | None,
+) -> None:
+ torch.manual_seed(0)
+
+ torch_model = reference.QwenImageRMS_norm(channels)
+ torch_model.eval()
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+
+ ctx = VaeContext(tp_axis=tp_axis, device=mesh_device, ccl_manager=ccl_manager)
+ tt_model = VaeRmsNorm(channels, eps=1e-12, ctx=ctx)
+
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ inp = torch.randn(batch_size, channels, height, width)
+
+ tt_inp = tensor.from_torch(inp.permute(0, 2, 3, 1), mesh_axes=[..., tp_axis], device=mesh_device)
+
+ with torch.no_grad():
+ torch_output = torch_model.forward(inp)
+
+ tt_out = tt_model.forward(tt_inp)
+
+ tt_out_torch = tensor.to_torch(tt_out, mesh_axes=[..., tp_axis]).permute(0, 3, 1, 2)
+ assert_quality(torch_output, tt_out_torch, pcc=0.99999, relative_rmse=0.008)
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "tp_axis"),
+ [
+ pytest.param((1, 1), None, id="1x1"),
+ pytest.param((1, 2), None, id="1x2"),
+ # pytest.param((1, 2), 1, id="1x2tp"), # hangs, resolved by https://github.com/tenstorrent/tt-metal/pull/33254
+ pytest.param((1, 4), 1, id="1x4tp"),
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "l1_small_size": 65536, "trace_region_size": 20000000}],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ ("batch_size", "height", "width", "dim"),
+ [
+ (2, 128, 128, 384),
+ ],
+)
+def test_vae_qwenimage_attention(
+ *,
+ mesh_device: ttnn.Device,
+ batch_size: int,
+ height: int,
+ width: int,
+ dim: int,
+ tp_axis: int | None,
+) -> None:
+ torch.manual_seed(0)
+
+ torch_model = reference.QwenImageAttentionBlock(dim)
+ torch_model.eval()
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+
+ ctx = VaeContext(tp_axis=tp_axis, device=mesh_device, ccl_manager=ccl_manager)
+ tt_model = VaeAttention(num_channels=dim, norm=VaeNormDescRms(eps=1e-12), ctx=ctx)
+
+ state = torch_model.state_dict()
+ tt_model.load_torch_state_dict(
+ {
+ "to_out.0.weight": state["proj.weight"].squeeze(2, 3),
+ "to_out.0.bias": state["proj.bias"],
+ "to_q.weight": state["to_qkv.weight"].squeeze(2, 3).chunk(3)[0],
+ "to_k.weight": state["to_qkv.weight"].squeeze(2, 3).chunk(3)[1],
+ "to_v.weight": state["to_qkv.weight"].squeeze(2, 3).chunk(3)[2],
+ "to_q.bias": state["to_qkv.bias"].chunk(3)[0],
+ "to_k.bias": state["to_qkv.bias"].chunk(3)[1],
+ "to_v.bias": state["to_qkv.bias"].chunk(3)[2],
+ "norm.gamma": state["norm.gamma"],
+ }
+ )
+
+ inp = torch.randn(batch_size, dim, height, width)
+
+ tt_inp = tensor.from_torch(inp.permute(0, 2, 3, 1), mesh_axes=[..., tp_axis], device=mesh_device)
+
+ with torch.no_grad():
+ torch_output = torch_model.forward(inp.unsqueeze(2)).squeeze(2)
+
+ tt_out = tt_model.forward(tt_inp)
+
+ tt_out_torch = tensor.to_torch(tt_out, mesh_axes=[..., tp_axis]).permute(0, 3, 1, 2)
+ assert_quality(torch_output, tt_out_torch, pcc=0.99996, relative_rmse=0.009)
diff --git a/code/models/tt_dit/tests/models/flux2/device_params.py b/code/models/tt_dit/tests/models/flux2/device_params.py
new file mode 100644
index 0000000000000000000000000000000000000000..b41c716aadb0247a074d94510bd6d443d4ac830e
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux2/device_params.py
@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared ``device_params`` dicts for the Flux2 test suite.
+
+Every Flux2 test parametrizes the repo-global ``device_params`` fixture (indirect) with one of
+these. They build on the fabric/topology base dicts in ``models/tt_dit/utils/test.py`` and layer on
+the three knobs Flux2 needs: an L1_SMALL region for the VAE conv2d, a trace region for the
+transformer denoising loop, and ``require_exact_physical_num_devices`` so one CI command self-selects
+the right mesh per SKU. Centralizing them here keeps the L1_SMALL value and the req-exact overlay
+defined once instead of copy-pasted across the five test modules.
+"""
+
+from ....utils.test import line_params_8k, line_params_req_exact_devices, ring_params_8k
+
+# Applied to a row so it self-skips unless the mesh it asks for is exactly what the machine has.
+REQ_EXACT = {"require_exact_physical_num_devices": True}
+# The Flux2 VAE uses conv2d, which needs L1_SMALL buffers on every mesh.
+L1_SMALL_SIZE = 65536
+# The transformer traces the denoising loop.
+TRACE_REGION_SIZE = 31_000_000
+
+# Base fabric/topology params + L1_SMALL for the VAE conv2d. line_params_flux2 already carries
+# require_exact (via line_params_req_exact_devices).
+line_params_flux2 = {**line_params_req_exact_devices, "l1_small_size": L1_SMALL_SIZE}
+line_params_8k_flux2 = {**line_params_8k, "l1_small_size": L1_SMALL_SIZE}
+ring_params_8k_flux2 = {**ring_params_8k, "l1_small_size": L1_SMALL_SIZE}
+ring_params_8k_flux2_req_exact = {**ring_params_8k_flux2, **REQ_EXACT}
+
+# Transformer: base + trace region.
+line_params_flux2_transformer = {**line_params_flux2, "trace_region_size": TRACE_REGION_SIZE}
+
+# Performance: base + req-exact so each row self-selects its mesh.
+line_params_flux2_perf = {**line_params_flux2, **REQ_EXACT}
+line_params_8k_flux2_perf = {**line_params_8k_flux2, **REQ_EXACT}
+ring_params_8k_flux2_perf = {**ring_params_8k_flux2, **REQ_EXACT}
+
+# VAE single-device row: L1_SMALL only, no fabric and no require_exact.
+# Don't set require_exact so it runs on every machine rather than only on a literal one-chip host, which doesn't exist.
+
+# No fabric because FABRIC_1D on a 1-device submesh of a multi-chip box doesn't work: fabric gets initialized
+# on 1 device, but router sync still waits on a router
+# on another chip (that isn't configured), so the handshake stays at STARTED and times
+# out in fabric_firmware_initializer.cpp.
+
+# Even if the bug above were fixed, CCLManager.all_gather returns early when shape[mesh_axis] == 1, _all_gather_hw
+# when both factors are 1, vae_all_gather when the cluster axis is 1 -- so no CCL op is ever
+# issued and fabric is redundant.
+single_device_params_flux2 = {"l1_small_size": L1_SMALL_SIZE}
+
+# Prompt encoder: fabric + req-exact, no L1_SMALL (it has no conv2d). Same shape as
+# line_params_req_exact_devices; aliased here so all Flux2 device params have one home.
+prompt_encoder_params_flux2 = line_params_req_exact_devices
diff --git a/code/models/tt_dit/tests/models/flux2/run_flux2_tests.sh b/code/models/tt_dit/tests/models/flux2/run_flux2_tests.sh
new file mode 100644
index 0000000000000000000000000000000000000000..46d795b209f1c8e15687fbd1783a3b3965763264
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux2/run_flux2_tests.sh
@@ -0,0 +1,271 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Run the FLUX.2 unit + e2e test suites on a single host.
+#
+# The FLUX.2 CI legs in tests/pipeline_reorg/*.yaml cover only 1024x1024 -- one
+# quietbox config and one galaxy config -- so the legs below are derived from the
+# test files themselves to sweep wider. Two things worth knowing:
+#
+# * Every flux2 test sets require_exact_physical_num_devices, so a (2,2) or (1,8)
+# parametrization self-skips on a 32-chip galaxy. Legs still pin a mesh id where
+# a file offers more than one row matching the host. Use --mesh/-k to widen.
+# * black-forest-labs/FLUX.2-dev is a manually gated HF repo (CI is blocked on
+# it, see #54499) and is not staged under /mnt/models. You need an HF token
+# on its allowlist; see PREREQUISITES below.
+#
+# All resolutions run by default (1024 -> 4096 for perf); cap with --max-res.
+#
+# Usage:
+# ./run_flux2_tests.sh # everything, arch auto-detected
+# ./run_flux2_tests.sh --suite unit
+# ./run_flux2_tests.sh --max-res 2048 # skip the 4096 perf leg
+# ./run_flux2_tests.sh --list # collect only, run nothing
+# ./run_flux2_tests.sh -- -x # trailing args go to pytest
+
+set -uo pipefail
+
+SUITE=all
+ARCH=auto
+MESH=auto
+EXTRA_K=""
+MAX_RES=4096
+WITH_PROFILE=0
+LIST_ONLY=0
+DRY_RUN=0
+FAIL_FAST=0
+LOG_DIR=""
+PYTEST_EXTRA=()
+
+usage() {
+ sed -n '6,28p' "$0" | sed 's/^# \?//'
+ cat <<'EOF'
+
+Options:
+ --suite {unit,e2e,all} Which suite(s) to run (default: all)
+ --arch {bh,wh,auto} Pick the arch-specific parametrizations (default: auto)
+ --mesh ID Override the mesh-id filter for every leg (advanced;
+ the per-test ids differ, so this is usually wrong --
+ prefer -k)
+ -k EXPR Extra pytest -k expression, ANDed into every leg
+ --max-res N Skip perf/profile resolutions above N
+ (N in 1024 2048 4096; default 4096)
+ --with-profile Also run test_transformer_profile (72 heavy cases,
+ a profiling sweep, not a correctness test)
+ --list pytest --collect-only -q for each leg
+ --dry-run Print the commands without running them
+ --fail-fast Stop after the first failing leg
+ --log-dir DIR Where to write per-leg logs and generated images
+ (default: $TT_METAL_HOME/generated/flux2_tests/)
+ -h, --help This message
+
+PREREQUISITES
+ * Activate the venv: source $TT_METAL_HOME/python_env/bin/activate
+ * FLUX.2-dev is gated. Request access at
+ https://huggingface.co/black-forest-labs/FLUX.2-dev then either
+ `huggingface-cli login` or export HF_TOKEN=. Without it every leg
+ fails with a 403 GatedRepoError.
+ * First run downloads the full FLUX.2-dev checkpoint (tens of GB); it is not
+ pre-staged under /mnt/models (only FLUX.1-dev/schnell are).
+EOF
+}
+
+while [[ $# -gt 0 ]]; do
+ case "$1" in
+ --suite) SUITE="$2"; shift 2 ;;
+ --arch) ARCH="$2"; shift 2 ;;
+ --mesh) MESH="$2"; shift 2 ;;
+ -k) EXTRA_K="$2"; shift 2 ;;
+ --max-res) MAX_RES="$2"; shift 2 ;;
+ --with-profile) WITH_PROFILE=1; shift ;;
+ --list) LIST_ONLY=1; shift ;;
+ --dry-run) DRY_RUN=1; shift ;;
+ --fail-fast) FAIL_FAST=1; shift ;;
+ --log-dir) LOG_DIR="$2"; shift 2 ;;
+ -h|--help) usage; exit 0 ;;
+ --) shift; PYTEST_EXTRA=("$@"); break ;;
+ *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;;
+ esac
+done
+
+case "$SUITE" in unit|e2e|all) ;; *) echo "--suite must be unit|e2e|all" >&2; exit 2 ;; esac
+case "$ARCH" in bh|wh|auto) ;; *) echo "--arch must be bh|wh|auto" >&2; exit 2 ;; esac
+case "$MAX_RES" in 1024|2048|4096) ;; *) echo "--max-res must be 1024|2048|4096" >&2; exit 2 ;; esac
+
+# ---------------------------------------------------------------------------
+# Repo / arch discovery
+# ---------------------------------------------------------------------------
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+: "${TT_METAL_HOME:=$(cd "${SCRIPT_DIR}/../../../../.." && pwd)}"
+export TT_METAL_HOME
+export PYTHONPATH="${PYTHONPATH:-${TT_METAL_HOME}}"
+
+if [[ "$ARCH" == auto && -n "${ARCH_NAME:-}" ]]; then
+ case "$ARCH_NAME" in *blackhole*) ARCH=bh ;; *wormhole*) ARCH=wh ;; esac
+fi
+if [[ "$ARCH" == auto ]]; then
+ smi="$(tt-smi -ls 2>/dev/null)"
+ if grep -qi blackhole <<<"$smi"; then ARCH=bh
+ elif grep -qi wormhole <<<"$smi"; then ARCH=wh
+ else echo "could not detect arch, pass --arch bh|wh" >&2; exit 2
+ fi
+fi
+
+# Per-test mesh parametrization ids. The ids differ between files, hence one
+# variable per test rather than a single global mesh filter.
+VAE_ID="4x8" # test_vae_flux2.py also has 1x1, 1x8
+ENC_ID="4x8" # test_prompt_encoder.py only 4x8 exists
+PIPE_ID="4x8" # test_pipeline_flux2.py both 4x8 rows; the 4-link one self-skips on BH
+if [[ "$ARCH" == bh ]]; then
+ XF_ID="bh_4x8_ring" # test_transformer_flux2.py::test_transformer
+ PERF_IDS="bh_glx_linear or bh_glx_ring_sp0tp1_nofsdp or bh_glx_ring_sp1tp0 or bh_glx_ring_sp0tp1_fsdp"
+ PROFILE_IDS="bh_4x8_ring_nofsdp or bh_4x8_ring_sp1tp0_nofsdp or bh_4x8_ring_fsdp"
+else
+ XF_ID="wh_2x4_linear"
+ PERF_IDS="" # test_performance_flux2.py parametrizes BH meshes only
+ PROFILE_IDS="" # test_transformer_profile is BH-only too
+fi
+if [[ "$MESH" != auto ]]; then
+ VAE_ID="$MESH"; ENC_ID="$MESH"; XF_ID="$MESH"; PIPE_ID="$MESH"
+ PERF_IDS="$MESH"; PROFILE_IDS="$MESH"
+fi
+
+# Resolutions, capped by --max-res.
+ALL_RES=(1024 2048 4096)
+RES=()
+for r in "${ALL_RES[@]}"; do (( r <= MAX_RES )) && RES+=("$r"); done
+
+TS="$(date +%Y%m%d-%H%M%S)"
+: "${LOG_DIR:=${TT_METAL_HOME}/generated/flux2_tests/${TS}}"
+mkdir -p "$LOG_DIR"
+
+# ---------------------------------------------------------------------------
+# Environment (same conventions as the Flux.1 CI legs)
+# ---------------------------------------------------------------------------
+unset TT_METAL_WATCHER # blocked by #50886
+unset HF_HUB_OFFLINE # FLUX.2-dev is not pre-staged; the hub is the source
+export NO_PROMPT=1 # test_pipeline_flux2 blocks on input() without this
+export TT_DIT_CACHE_DIR="${TT_DIT_CACHE_DIR:-/tmp/TT_DIT_CACHE}"
+if [[ -z "${HF_TOKEN:-}" && ! -f "${HF_HOME:-$HOME/.cache/huggingface}/token" ]]; then
+ echo "WARNING: no HF credentials found (HF_TOKEN unset, no cached token)." >&2
+ echo " black-forest-labs/FLUX.2-dev is gated; expect 403 GatedRepoError." >&2
+fi
+
+FLUX2_DIR="models/tt_dit/tests/models/flux2"
+
+# ---------------------------------------------------------------------------
+# Leg runner
+# ---------------------------------------------------------------------------
+declare -a RESULTS=()
+FAILED=0
+
+# run [k_expr] [extra pytest args...]
+run() {
+ local name="$1" timeout="$2" target="$3" kexpr="${4:-}"
+ shift 4 2>/dev/null || shift $#
+
+ if [[ -n "$EXTRA_K" ]]; then
+ if [[ -n "$kexpr" ]]; then kexpr="($kexpr) and ($EXTRA_K)"; else kexpr="$EXTRA_K"; fi
+ fi
+
+ local -a cmd=(pytest --timeout "$timeout" "$target")
+ [[ -n "$kexpr" ]] && cmd+=(-k "$kexpr")
+ cmd+=("$@")
+ (( LIST_ONLY )) && cmd+=(--collect-only -q)
+ cmd+=("${PYTEST_EXTRA[@]+"${PYTEST_EXTRA[@]}"}")
+
+ local log="${LOG_DIR}/${name}.log"
+ echo
+ echo "=============================================================================="
+ echo "[$name] ${cmd[*]}"
+ echo " log: $log"
+ echo "=============================================================================="
+
+ if (( DRY_RUN )); then
+ RESULTS+=("DRYRUN $name")
+ return 0
+ fi
+
+ local start=$SECONDS
+ ( cd "$TT_METAL_HOME" && "${cmd[@]}" ) 2>&1 | tee "$log"
+ local rc=${PIPESTATUS[0]}
+ local dur=$(( SECONDS - start ))
+
+ # test_pipeline_flux2 saves flux2_*.png into the cwd; keep them with the logs.
+ shopt -s nullglob
+ for png in "$TT_METAL_HOME"/flux2_*.png; do mv -f "$png" "$LOG_DIR/"; done
+ shopt -u nullglob
+
+ if (( rc == 0 )); then
+ RESULTS+=("$(printf 'PASS %-34s %5ds' "$name" "$dur")")
+ else
+ RESULTS+=("$(printf 'FAIL(%d) %-34s %5ds %s' "$rc" "$name" "$dur" "$log")")
+ FAILED=1
+ if (( FAIL_FAST )); then summary; exit 1; fi
+ fi
+}
+
+skip_leg() {
+ RESULTS+=("$(printf 'SKIP %-34s %s' "$1" "$2")")
+ echo
+ echo "[skip] $1 -- $2"
+}
+
+summary() {
+ echo
+ echo "=============================================================================="
+ echo "FLUX.2 test summary (suite=$SUITE arch=$ARCH max-res=$MAX_RES)"
+ echo "logs/images: $LOG_DIR"
+ echo "=============================================================================="
+ printf '%s\n' "${RESULTS[@]+"${RESULTS[@]}"}"
+}
+
+want() { [[ "$1" == all || "$1" == "$2" ]]; }
+
+# ---------------------------------------------------------------------------
+# Unit legs -- PCC/correctness against the torch reference
+# ---------------------------------------------------------------------------
+if want "$SUITE" unit; then
+ # VAE decoder, pruned to 1 resnet per up_block by the test itself.
+ run flux2-unit-vae 3000 "$FLUX2_DIR/test_vae_flux2.py" "$VAE_ID"
+ # Mistral3 prompt encoder + the prompt-embedding upsampler.
+ run flux2-unit-encoder 3000 "$FLUX2_DIR/test_prompt_encoder.py" "$ENC_ID"
+ # Both all_blocks and single_blocks variants at 1024x1024.
+ run flux2-unit-transformer 10000 "$FLUX2_DIR/test_transformer_flux2.py::test_transformer" "$XF_ID"
+ # 12-step 1024x1024 generation; writes flux2_*.png next to the logs.
+ run flux2-unit-pipeline 10000 "$FLUX2_DIR/test_pipeline_flux2.py" "$PIPE_ID"
+fi
+
+# ---------------------------------------------------------------------------
+# e2e legs -- end-to-end throughput, one leg per resolution
+# ---------------------------------------------------------------------------
+if want "$SUITE" e2e; then
+ if [[ -z "$PERF_IDS" ]]; then
+ skip_leg flux2-e2e-perf "test_performance_flux2.py parametrizes Blackhole meshes only"
+ else
+ for r in "${RES[@]}"; do
+ run "flux2-e2e-perf-${r}" 10000 "$FLUX2_DIR/test_performance_flux2.py" \
+ "($PERF_IDS) and ${r}x${r}"
+ done
+ fi
+
+ if (( WITH_PROFILE )); then
+ if [[ -z "$PROFILE_IDS" ]]; then
+ skip_leg flux2-e2e-transformer-profile "test_transformer_profile is Blackhole-only"
+ else
+ for r in "${RES[@]}"; do
+ (( r > 4096 )) && continue # profile sweep tops out at 4096
+ run "flux2-e2e-xf-profile-${r}" 10000 \
+ "$FLUX2_DIR/test_transformer_flux2.py::test_transformer_profile" \
+ "($PROFILE_IDS) and ${r}"
+ done
+ fi
+ else
+ skip_leg flux2-e2e-transformer-profile "profiling sweep, opt in with --with-profile"
+ fi
+fi
+
+summary
+exit $FAILED
diff --git a/code/models/tt_dit/tests/models/flux2/test_performance_flux2.py b/code/models/tt_dit/tests/models/flux2/test_performance_flux2.py
new file mode 100644
index 0000000000000000000000000000000000000000..535b739257a17f03c1b834861d9a1d65986090a3
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux2/test_performance_flux2.py
@@ -0,0 +1,293 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import statistics
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import is_blackhole
+from models.perf.benchmarking_utils import BenchmarkData, BenchmarkProfiler
+
+from ....pipelines.flux2.pipeline_flux2 import Flux2Pipeline
+from .device_params import line_params_8k_flux2_perf, line_params_flux2_perf, ring_params_8k_flux2_perf
+
+NUM_INFERENCE_STEPS = 50
+NUM_PERF_RUNS = 3
+
+# Upper bounds in seconds, 1.5x the 1024x1024 / 50-step numbers
+# These are more like regression tripwires, not a-priori determined perf targets.
+#
+# The key is (mesh shape, width, topology, sp_axis, is_fsdp), which is what it takes to name one
+# parametrization uniquely.
+#
+# Configs absent from this table are reported but we don't asssert on any perf targets.
+PERF_THRESHOLDS = {
+ ((4, 8), 1024, ttnn.Topology.Ring, 0, False): {
+ "total_encoding_time": 0.13, # measured 0.0971
+ "denoising_steps_time": 0.20 * NUM_INFERENCE_STEPS, # measured 0.1476/step, 7.3997 total
+ "vae_decoding_time": 0.49, # measured 0.1797
+ "total_time": 9.8125, # measured 7.6956
+ },
+}
+
+# Profiler step name -> the PERF_THRESHOLDS key it is gated against.
+_STEP_TO_METRIC = {
+ "encoder": "total_encoding_time",
+ "denoising": "denoising_steps_time",
+ "vae": "vae_decoding_time",
+ "run": "total_time",
+}
+
+
+@pytest.mark.timeout(6000)
+@pytest.mark.parametrize(
+ "width, height",
+ [
+ (1024, 1024),
+ (2048, 2048),
+ (4096, 4096),
+ (8192, 8192),
+ ],
+ ids=[
+ "1024x1024",
+ "2048x2048",
+ "4096x4096",
+ "8192x8192",
+ ],
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, encoder_tp_axis, vae_tp_axis, topology, num_links, is_fsdp, dynamic_load, device_params",
+ [
+ # Without dynamic_load the encoder OOMs during weight conversion on 4 chips.
+ [(2, 2), 0, 1, 1, 1, ttnn.Topology.Linear, 2, True, True, line_params_flux2_perf],
+ [(2, 4), 0, 1, 1, 1, ttnn.Topology.Linear, 2, False, False, line_params_flux2_perf],
+ [(4, 8), 0, 1, 1, 0, ttnn.Topology.Linear, 2, False, False, line_params_8k_flux2_perf],
+ [(4, 8), 0, 1, 1, 0, ttnn.Topology.Ring, 2, False, False, ring_params_8k_flux2_perf],
+ [(4, 8), 1, 0, 1, 0, ttnn.Topology.Ring, 2, False, False, ring_params_8k_flux2_perf],
+ [(4, 8), 0, 1, 1, 0, ttnn.Topology.Ring, 2, True, False, ring_params_8k_flux2_perf],
+ ],
+ ids=[
+ "bh_qb",
+ "bh_lb",
+ "bh_glx_linear",
+ "bh_glx_ring_sp0tp1_nofsdp",
+ "bh_glx_ring_sp1tp0",
+ "bh_glx_ring_sp0tp1_fsdp",
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+def test_flux2_performance(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ encoder_tp_axis: int,
+ vae_tp_axis: int,
+ topology: ttnn.Topology,
+ num_links: int,
+ is_fsdp: bool,
+ dynamic_load: bool,
+ width: int,
+ height: int,
+ model_location_generator,
+ is_ci_env: bool,
+) -> None:
+ """Performance test for Flux2 pipeline across machine types and resolutions."""
+
+ benchmark_profiler = BenchmarkProfiler()
+
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ tp_factor = tuple(mesh_device.shape)[tp_axis]
+
+ logger.info(
+ f"Config: {width}x{height}, mesh={tuple(mesh_device.shape)}, "
+ f"sp={sp_factor}, tp={tp_factor}, topology={topology}"
+ )
+
+ # Pipeline creation includes internal warmup (2 steps at target resolution).
+ pipeline = Flux2Pipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=model_location_generator("black-forest-labs/FLUX.2-dev"),
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ encoder_tp_axis=encoder_tp_axis,
+ vae_tp_axis=vae_tp_axis,
+ vae_h_axis=1 - vae_tp_axis,
+ vae_w_axis=None, # 1-vae_tp_axis,
+ num_links=num_links,
+ topology=topology,
+ width=width,
+ height=height,
+ dynamic_load=dynamic_load,
+ is_fsdp=is_fsdp,
+ trace_warmup=True,
+ shard_prompt=True,
+ )
+
+ # Performance measurement runs
+ logger.info(f"Running {NUM_PERF_RUNS} timed iterations...")
+
+ prompts = [
+ "A photo of a cat sitting on a windowsill at sunset",
+ "A lifelike, warm portrait of a woman in her late 30s, shot on an 85mm lens, golden-hour lighting casting a soft glow, shallow depth of field, sharp focus on the eyes, highly detailed skin texture, photorealistic editorial style.",
+ "Create an image of an alien world with floating rock islands, massive glowing tropical plants, and twin moons in the night sky. Wide-angle shot, cinematic lighting, vivid colors, highly detailed sci-fi concept art.",
+ ]
+
+ results = []
+ for i in range(NUM_PERF_RUNS):
+ logger.info(f"Performance run {i + 1}/{NUM_PERF_RUNS}...")
+
+ ttnn.synchronize_device(mesh_device)
+ ttnn.distributed_context_barrier()
+
+ with benchmark_profiler("run", iteration=i):
+ with torch.no_grad():
+ images = pipeline(
+ prompts=[prompts[i % len(prompts)]],
+ num_inference_steps=NUM_INFERENCE_STEPS,
+ seed=42,
+ traced=True,
+ profiler=benchmark_profiler,
+ profiler_iteration=i,
+ )
+ results.append(images)
+ ttnn.synchronize_device(mesh_device)
+
+ logger.info(f" Run {i + 1} completed in {benchmark_profiler.get_duration('run', i):.2f}s")
+
+ # Collect timing data across all runs
+ encoder_times = [benchmark_profiler.get_duration("encoder", i) for i in range(NUM_PERF_RUNS)]
+ denoising_times = [benchmark_profiler.get_duration("denoising", i) for i in range(NUM_PERF_RUNS)]
+ vae_times = [benchmark_profiler.get_duration("vae", i) for i in range(NUM_PERF_RUNS)]
+ total_times = [benchmark_profiler.get_duration("run", i) for i in range(NUM_PERF_RUNS)]
+
+ all_denoising_steps = []
+ for i in range(NUM_PERF_RUNS):
+ for j in range(NUM_INFERENCE_STEPS):
+ if benchmark_profiler.contains_step(f"denoising_step_{j}", i):
+ all_denoising_steps.append(benchmark_profiler.get_duration(f"denoising_step_{j}", i))
+
+ # Save output image on rank 0 (from last run)
+ if not is_ci_env:
+ rank = int(ttnn.distributed_context_get_rank())
+ if rank == 0:
+ for i in range(NUM_PERF_RUNS):
+ mesh_tag = "x".join(str(s) for s in mesh_device.shape)
+ output_path = f"flux2_{mesh_tag}_{width}x{height}_{i}.png"
+ results[i][0].save(output_path)
+ logger.info(f"Image {i} saved as {output_path}")
+
+ # Report results
+ print("\n" + "=" * 80)
+ print("FLUX2 PERFORMANCE RESULTS")
+ print("=" * 80)
+ print(f"Image Size: {width}x{height}")
+ print(f"Inference Steps: {NUM_INFERENCE_STEPS}")
+ print(f"Performance Runs: {NUM_PERF_RUNS}")
+ print(f"DiT Configuration: sp={sp_factor}, tp={tp_factor}")
+ print(f"Mesh Shape: {tuple(mesh_device.shape)}")
+ print(f"Topology: {topology}")
+ print("-" * 80)
+
+ def print_stats(name, times):
+ if not times:
+ print(f"{name:25} | No data available")
+ return
+ mean_time = statistics.mean(times)
+ std_time = statistics.stdev(times) if len(times) > 1 else 0
+ min_time = min(times)
+ max_time = max(times)
+ print(
+ f"{name:25} | Mean: {mean_time:8.4f}s | Std: {std_time:8.4f}s | "
+ f"Min: {min_time:8.4f}s | Max: {max_time:8.4f}s"
+ )
+
+ print_stats("Encoding", encoder_times)
+ print_stats("Denoising (per step)", all_denoising_steps)
+ print_stats("Denoising (total)", denoising_times)
+ print_stats("VAE Decoding", vae_times)
+ print_stats("Total Pipeline", total_times)
+ print("-" * 80)
+
+ # Additional metrics
+ avg_total = statistics.mean(total_times)
+ avg_encoding = statistics.mean(encoder_times)
+ avg_denoising = statistics.mean(denoising_times)
+ avg_vae = statistics.mean(vae_times)
+
+ if all_denoising_steps:
+ print(f"Average total denoising time: {avg_denoising:.4f}s")
+ print(f"Denoising throughput: {NUM_INFERENCE_STEPS / avg_denoising:.2f} steps/second")
+
+ print(f"Overall throughput: {1 / avg_total:.4f} images/second")
+
+ print(f"\nTime breakdown:")
+ print(f" Encoding: {avg_encoding / avg_total * 100:5.1f}%")
+ print(f" Denoising: {avg_denoising / avg_total * 100:5.1f}%")
+ print(f" VAE: {avg_vae / avg_total * 100:5.1f}%")
+
+ print("=" * 80)
+
+ measurements = {
+ "total_encoding_time": avg_encoding,
+ "denoising_steps_time": avg_denoising,
+ "vae_decoding_time": avg_vae,
+ "total_time": avg_total,
+ }
+ perf_key = (tuple(mesh_device.shape), width, topology, sp_axis, is_fsdp)
+ expected_metrics = PERF_THRESHOLDS.get(perf_key)
+ if expected_metrics is None:
+ logger.warning(
+ f"No perf thresholds for {perf_key}: reporting timings only. "
+ "Add a PERF_THRESHOLDS entry to gate this config against regressions."
+ )
+
+ if is_ci_env:
+ device_name_map = {
+ (2, 2): "BH_QB",
+ (2, 4): "BH_LB" if is_blackhole() else "WH_T3K",
+ (4, 8): "BH_GLX" if is_blackhole() else "WH_GLX",
+ }
+
+ benchmark_data = BenchmarkData()
+ for i in range(NUM_PERF_RUNS):
+ for step_name in ["encoder", "denoising", "vae", "run"]:
+ value = benchmark_profiler.get_duration(step_name, i)
+ benchmark_data.add_measurement(
+ profiler=benchmark_profiler,
+ iteration=i,
+ step_name=step_name,
+ name=step_name,
+ value=value,
+ target=(expected_metrics or {}).get(_STEP_TO_METRIC[step_name], value),
+ )
+ benchmark_data.save_partial_run_json(
+ benchmark_profiler,
+ run_type=device_name_map.get(tuple(mesh_device.shape), "UNKNOWN"),
+ ml_model_name=os.path.basename(pipeline.checkpoint_name),
+ batch_size=1,
+ config_params={
+ "width": width,
+ "height": height,
+ "num_steps": NUM_INFERENCE_STEPS,
+ "num_perf_runs": NUM_PERF_RUNS,
+ "sp_factor": sp_factor,
+ "tp_factor": tp_factor,
+ "topology": str(topology),
+ },
+ )
+
+ if expected_metrics is not None:
+ regressions = [
+ f" {k}: {measurements[k]:.4f}s exceeds threshold {expected_metrics[k]:.4f}s"
+ for k in expected_metrics
+ if measurements[k] > expected_metrics[k]
+ ]
+ assert not regressions, f"FLUX.2 performance regression for {perf_key}:\n" + "\n".join(regressions)
+
+ logger.info("Performance test completed successfully!")
diff --git a/code/models/tt_dit/tests/models/flux2/test_pipeline_flux2.py b/code/models/tt_dit/tests/models/flux2/test_pipeline_flux2.py
new file mode 100644
index 0000000000000000000000000000000000000000..e19757a2a5b237ff04c830397bb05503651f7fa9
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux2/test_pipeline_flux2.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+import os
+
+import pytest
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import is_blackhole
+
+from ....pipelines.flux2.pipeline_flux2 import Flux2Pipeline
+from ....utils.test import skip_if_unsupported_num_links
+from .device_params import line_params_flux2
+
+
+@pytest.mark.parametrize(
+ "no_prompt",
+ [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), False)],
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [line_params_flux2],
+ indirect=True,
+)
+@pytest.mark.parametrize(("width", "height", "num_inference_steps"), [(1024, 1024, 12)])
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, encoder_tp_axis, vae_tp_axis, topology, num_links, is_fsdp, dynamic_load, traced",
+ [
+ # is_fsdp and dynamic_load are both mandatory for 2x2 -- the 32B transformer, the 24B encoder and the VAE
+ # cannot be co-resident on 4 chips.
+ [(2, 2), 0, 1, 1, 1, ttnn.Topology.Linear, 2, True, True, False],
+ [(1, 8), 0, 1, 1, 1, ttnn.Topology.Linear, 1, False, True, False],
+ [(4, 8), 0, 1, 1, 0, ttnn.Topology.Linear, 4, True, False, True],
+ [(4, 8), 0, 1, 1, 0, ttnn.Topology.Linear, 2, False, False, True],
+ ],
+ ids=[
+ "2x2sp0tp1vaetp1nl2_linear_is_fsdp1",
+ "1x8sp0tp1vaetp1nl1_linear_is_fsdp0",
+ "4x8sp0tp1vaetp0nl4_linear_is_fsdp1",
+ "4x8sp0tp1vaetp0nl2_linear_is_fsdp0",
+ ],
+ indirect=["mesh_device"],
+)
+def test_pipeline(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ width: int,
+ height: int,
+ num_inference_steps: int,
+ sp_axis: int,
+ tp_axis: int,
+ encoder_tp_axis: int,
+ vae_tp_axis: int,
+ topology: ttnn.Topology,
+ num_links: int,
+ no_prompt: bool,
+ traced: bool,
+ is_fsdp: bool,
+ dynamic_load: bool,
+ model_location_generator,
+ is_ci_env: bool,
+) -> None:
+ skip_if_unsupported_num_links(mesh_device, num_links)
+
+ pipeline = Flux2Pipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=model_location_generator("black-forest-labs/FLUX.2-dev"),
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ encoder_tp_axis=encoder_tp_axis,
+ vae_tp_axis=vae_tp_axis,
+ vae_h_axis=1 - vae_tp_axis,
+ vae_w_axis=None,
+ num_links=num_links,
+ topology=topology,
+ width=width,
+ height=height,
+ is_fsdp=is_fsdp,
+ dynamic_load=dynamic_load,
+ trace_warmup=traced,
+ shard_prompt=True,
+ )
+
+ prompts = [
+ "Neon-lit cyberpunk alley, rain-soaked, cinematic wide shot",
+ # "Golden retriever astronaut drifting in sunlit space",
+ # "Minimalist Scandinavian kitchen, morning light, ultra clean",
+ # "Ancient desert temple at dawn, soft fog, wide angle",
+ # "Steampunk airship over Victorian city, dramatic clouds",
+ # "Macro dewdrops on fern, shallow depth of field",
+ # "Luxury wristwatch on marble, studio lighting, hyper-detail",
+ # "Stormy coastline lighthouse, crashing waves, long exposure",
+ # "Futuristic Tokyo street market, vibrant signage, motion blur",
+ ]
+
+ arch = "bh" if is_blackhole() else "wh"
+ mesh_tag = "x".join(str(s) for s in mesh_device.shape)
+ filename_prefix = f"flux2_{arch}_{width}_{height}_{mesh_tag}"
+ if not traced:
+ filename_prefix += "_untraced"
+
+ def run(*, prompt: str, number: int, seed: int) -> None:
+ images = pipeline(
+ prompts=[prompt],
+ num_inference_steps=num_inference_steps,
+ seed=seed,
+ traced=traced,
+ )
+
+ output_filename = f"{filename_prefix}_{number}.png"
+ images[0].save(output_filename)
+ logger.info(f"Image saved as {output_filename}")
+
+ if no_prompt:
+ for i, prompt in enumerate(prompts):
+ run(prompt=prompt, number=i, seed=0)
+ else:
+ prompt = prompts[0]
+ 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)
diff --git a/code/models/tt_dit/tests/models/flux2/test_prompt_encoder.py b/code/models/tt_dit/tests/models/flux2/test_prompt_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c09da8c25b14fcca5b2e483155f4790cae7d553
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux2/test_prompt_encoder.py
@@ -0,0 +1,108 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+import pytest
+from loguru import logger
+
+import ttnn
+
+from ....parallel.config import EncoderParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....pipelines.flux2.prompt_encoder import PromptEncoder
+from ....utils.check import assert_quality
+from .device_params import prompt_encoder_params_flux2
+
+
+@pytest.mark.parametrize(
+ "mesh_device",
+ [
+ pytest.param((2, 2), id="2x2"),
+ pytest.param((4, 8), id="4x8"),
+ ],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [prompt_encoder_params_flux2],
+ indirect=True,
+)
+def test_encode(mesh_device: ttnn.MeshDevice) -> None:
+ tp_axis = 1
+ sequence_length = 64
+ checkpoint_name = "black-forest-labs/FLUX.2-dev"
+ prompts = ["A painting", "A futuristic city skyline at sunset during winter"]
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=mesh_device.shape[tp_axis], mesh_axis=tp_axis),
+ )
+
+ tt_encoder = PromptEncoder(
+ checkpoint_name=checkpoint_name,
+ use_torch_encoder=False,
+ device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+ tt_encoder.load_weights()
+
+ torch_encoder = PromptEncoder(
+ checkpoint_name=checkpoint_name,
+ use_torch_encoder=True,
+ device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+
+ logger.info("running Torch model...")
+ torch_embeds, torch_mask = torch_encoder.encode(prompts, num_images_per_prompt=2, sequence_length=sequence_length)
+
+ logger.info("running TT model...")
+ tt_embeds, tt_mask = tt_encoder.encode(prompts, num_images_per_prompt=2, sequence_length=sequence_length)
+
+ assert torch_mask.equal(tt_mask)
+
+ assert_quality(torch_embeds, tt_embeds, pcc=0.9997, relative_rmse=0.04)
+
+
+@pytest.mark.parametrize(
+ "mesh_device",
+ [
+ pytest.param((2, 2), id="2x2"),
+ pytest.param((4, 8), id="4x8"),
+ ],
+ indirect=True,
+)
+@pytest.mark.parametrize(
+ "device_params",
+ [prompt_encoder_params_flux2],
+ indirect=True,
+)
+def test_upsample(mesh_device: ttnn.MeshDevice) -> None:
+ tp_axis = 1
+ checkpoint_name = "black-forest-labs/FLUX.2-dev"
+ prompts = ["A futuristic city skyline at sunset"]
+ temperature = 0.15
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+
+ parallel_config = EncoderParallelConfig(
+ tensor_parallel=ParallelFactor(factor=mesh_device.shape[tp_axis], mesh_axis=tp_axis),
+ )
+
+ tt_encoder = PromptEncoder(
+ checkpoint_name=checkpoint_name,
+ use_torch_encoder=False,
+ device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+
+ tt_encoder.load_weights()
+
+ logger.info("running TT model...")
+ tt_output = tt_encoder.upsample(prompts, max_length=224, temperature=temperature)
+ print(tt_output)
diff --git a/code/models/tt_dit/tests/models/flux2/test_transformer_flux2.py b/code/models/tt_dit/tests/models/flux2/test_transformer_flux2.py
new file mode 100644
index 0000000000000000000000000000000000000000..81eac76333ef843c06e1aba3215fd7b6ad347394
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux2/test_transformer_flux2.py
@@ -0,0 +1,434 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Protocol
+
+import diffusers
+import diffusers.models.transformers.transformer_flux2
+import pytest
+import torch
+from loguru import logger
+from tracy import signpost
+
+import ttnn
+
+from ....models.transformers.transformer_flux2 import Flux2Transformer
+from ....parallel.config import DiTGParallelConfigNoCFG, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....utils import cache, tensor
+from ....utils.check import assert_quality
+from ....utils.padding import PaddingConfig
+from .device_params import line_params_flux2_transformer, ring_params_8k_flux2, ring_params_8k_flux2_req_exact
+
+
+class ModelLocationGenerator(Protocol):
+ def __call__(
+ self,
+ model_version: str,
+ *,
+ model_subdir: str = "",
+ download_if_ci_v2: bool = False,
+ ci_v2_timeout_in_s: int = 300,
+ endpoint_prefix: str = "",
+ download_dir_suffix: str = "",
+ ) -> str:
+ ...
+
+
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, topology, num_links, fsdp, shard_prompt, device_params",
+ [
+ [(2, 2), 0, 1, ttnn.Topology.Linear, 2, True, False, line_params_flux2_transformer],
+ [(1, 8), 0, 1, ttnn.Topology.Linear, 1, False, False, line_params_flux2_transformer],
+ [(2, 4), 0, 1, ttnn.Topology.Linear, 1, False, False, line_params_flux2_transformer],
+ # shard_prompt=True on 4x8. With it False the prompt is replicated across the sp axis
+ # while the spatial stream is sharded, and the joint attention PCC in the single blocks
+ # degrades as sp increases.
+ #
+ # In practice, for sp=1/2, PCC is not sufficiently degraded to matter, but it's terrible
+ # for sp=4 (way below PCC thresholds), which is why we _must_ set shard_prompt=True for this mesh.
+ # True is also what the pipeline and test_performance_flux2.py use.
+ [(4, 8), 0, 1, ttnn.Topology.Ring, 2, False, True, ring_params_8k_flux2_req_exact],
+ ],
+ ids=[
+ "bh_2x2_linear",
+ "1x8_linear",
+ "wh_2x4_linear",
+ "bh_4x8_ring",
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "height", "width", "prompt_seq_len"),
+ [
+ (1, 1024, 1024, 512),
+ ],
+)
+@pytest.mark.parametrize(
+ ("skip_layers", "skip_single_layers"),
+ [
+ pytest.param(0, 0, id="all_blocks"),
+ pytest.param(7, 47, id="single_blocks"),
+ ],
+)
+def test_transformer(
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ topology: ttnn.Topology,
+ num_links: int,
+ fsdp: bool,
+ shard_prompt: bool,
+ batch_size: int,
+ height: int,
+ width: int,
+ prompt_seq_len: int,
+ skip_layers: int,
+ skip_single_layers: int,
+ model_location_generator: ModelLocationGenerator,
+) -> None:
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ tp_factor = tuple(mesh_device.shape)[tp_axis]
+
+ logger.info(
+ f"test_transformer: mesh={tuple(mesh_device.shape)}, sp={sp_factor}, tp={tp_factor}, "
+ f"topology={topology}, num_links={num_links}, fsdp={fsdp}"
+ )
+
+ model_name = model_location_generator("black-forest-labs/FLUX.2-dev", model_subdir="transformer")
+ torch_model = diffusers.Flux2Transformer2DModel.from_pretrained(model_name, subfolder="transformer")
+ assert isinstance(torch_model, diffusers.Flux2Transformer2DModel)
+ torch_model.eval()
+
+ assert 0 <= skip_layers < len(torch_model.transformer_blocks)
+ assert 0 <= skip_single_layers < len(torch_model.single_transformer_blocks)
+
+ del torch_model.transformer_blocks[len(torch_model.transformer_blocks) - skip_layers :]
+ del torch_model.single_transformer_blocks[len(torch_model.single_transformer_blocks) - skip_single_layers :]
+
+ head_dim = torch_model.config.attention_head_dim
+ num_heads = torch_model.config.num_attention_heads
+ in_channels = torch_model.in_channels
+ joint_attention_dim = torch_model.config.joint_attention_dim
+
+ ccl_manager = CCLManager(
+ mesh_device=mesh_device,
+ num_links=num_links,
+ topology=topology,
+ )
+
+ parallel_config = DiTGParallelConfigNoCFG(
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+
+ if num_heads % tp_factor != 0:
+ padding_config = PaddingConfig.from_tensor_parallel_factor(num_heads, head_dim, tp_factor)
+ else:
+ padding_config = None
+
+ num_layers = torch_model.config.num_layers - skip_layers
+ num_single_layers = torch_model.config.num_single_layers - skip_single_layers
+
+ # HACK: The cache subfolder name modification is to ensure the cache contains all tensors for both the
+ # all_blocks and single_blocks test variants. Without this, we would get invalid caches for one or the other.
+ cache_subfolder = (
+ "transformer" if skip_layers == skip_single_layers == 0 else f"transformer_L{num_layers}_S{num_single_layers}"
+ )
+
+ tt_model = Flux2Transformer(
+ in_channels=in_channels,
+ num_layers=num_layers,
+ num_single_layers=num_single_layers,
+ attention_head_dim=head_dim,
+ num_attention_heads=num_heads,
+ joint_attention_dim=joint_attention_dim,
+ out_channels=torch_model.out_channels,
+ device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ padding_config=padding_config,
+ is_fsdp=fsdp,
+ shard_prompt=shard_prompt,
+ )
+
+ cache.load_model(
+ tt_model,
+ get_torch_state_dict=torch_model.state_dict,
+ model_name="FLUX.2-dev",
+ subfolder=cache_subfolder,
+ parallel_config=parallel_config,
+ mesh_shape=tuple(mesh_device.shape),
+ mesh_device=mesh_device,
+ is_fsdp=fsdp,
+ )
+
+ spatial_seq_len = height * width // 16**2
+
+ torch.manual_seed(0)
+ spatial = torch.randn([batch_size, spatial_seq_len, in_channels])
+ prompt = torch.randn([batch_size, prompt_seq_len, joint_attention_dim])
+ timestep = torch.full([batch_size], fill_value=500)
+ guidance = torch.full([batch_size], fill_value=3)
+
+ # prepare for ROPE
+ text_ids = _prepare_ids(text_sequence_length=prompt_seq_len)
+ image_ids = _prepare_ids(height=height // 16, width=width // 16)
+ prompt_rope_cos, prompt_rope_sin = torch_model.pos_embed.forward(text_ids)
+ spatial_rope_cos, spatial_rope_sin = torch_model.pos_embed.forward(image_ids)
+
+ tt_spatial = tensor.from_torch(spatial, device=mesh_device, mesh_axes=[None, sp_axis, None])
+ # The prompt and its RoPE follow shard_prompt, the same pairing test_transformer_profile uses.
+ prompt_mesh_axes = [None, sp_axis, None] if shard_prompt else None
+ prompt_rope_mesh_axes = [None, None, sp_axis, None] if shard_prompt else None
+ tt_prompt = tensor.from_torch(prompt, device=mesh_device, mesh_axes=prompt_mesh_axes)
+ tt_embedded_prompt = tt_model.context_embedder(tt_prompt)
+ tt_timestep = tensor.from_torch(timestep.unsqueeze(-1), dtype=ttnn.float32, device=mesh_device)
+ tt_guidance = tensor.from_torch(guidance.unsqueeze(-1), device=mesh_device)
+
+ tt_spatial_rope_cos = tensor.from_torch(
+ spatial_rope_cos.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=[None, None, sp_axis, None]
+ )
+ tt_spatial_rope_sin = tensor.from_torch(
+ spatial_rope_sin.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=[None, None, sp_axis, None]
+ )
+ tt_prompt_rope_cos = tensor.from_torch(
+ prompt_rope_cos.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=prompt_rope_mesh_axes
+ )
+ tt_prompt_rope_sin = tensor.from_torch(
+ prompt_rope_sin.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=prompt_rope_mesh_axes
+ )
+ tt_combined_rope = (
+ ttnn.concat([tt_spatial_rope_cos, tt_prompt_rope_cos], dim=2),
+ ttnn.concat([tt_spatial_rope_sin, tt_prompt_rope_sin], dim=2),
+ )
+
+ logger.info("running Torch model...")
+ with torch.no_grad():
+ torch_output = torch_model.forward(
+ hidden_states=spatial,
+ encoder_hidden_states=prompt,
+ timestep=timestep / 1000,
+ guidance=guidance,
+ img_ids=image_ids,
+ txt_ids=text_ids,
+ ).sample
+
+ logger.info("running TT model...")
+
+ signpost("t_start")
+ tt_output = tt_model.forward(
+ spatial=tt_spatial,
+ embedded_prompt=tt_embedded_prompt,
+ timestep=tt_timestep,
+ guidance=tt_guidance,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ combined_rope=tt_combined_rope,
+ spatial_sequence_length=spatial_seq_len,
+ prompt_sequence_length=prompt_seq_len,
+ )
+ signpost("t_end")
+ tt_output_torch = tensor.to_torch(tt_output, mesh_axes=[None, sp_axis, None])
+ assert_quality(torch_output, tt_output_torch, pcc=0.996, relative_rmse=0.09)
+
+
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, topology, num_links, fsdp, device_params",
+ [
+ [(4, 8), 0, 1, ttnn.Topology.Ring, 2, False, ring_params_8k_flux2],
+ [(4, 8), 1, 0, ttnn.Topology.Ring, 2, False, ring_params_8k_flux2],
+ [(4, 8), 0, 1, ttnn.Topology.Ring, 2, True, ring_params_8k_flux2],
+ ],
+ ids=[
+ "bh_4x8_ring_nofsdp",
+ "bh_4x8_ring_sp1tp0_nofsdp",
+ "bh_4x8_ring_fsdp",
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "height", "width", "prompt_seq_len"),
+ [
+ (1, 1024, 1024, 512),
+ (1, 2048, 2048, 512),
+ (1, 4096, 4096, 512),
+ ],
+ ids=[
+ "1024",
+ "2048",
+ "4096",
+ ],
+)
+@pytest.mark.parametrize("xformer_idx", [0, 1, 3, 4], ids=["x_single", "x_double", "x_both", "x_all"])
+@pytest.mark.parametrize("shard_prompt", [False, True], ids=["no_shard_prompt", "yes_shard_prompt"])
+def test_transformer_profile(
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ topology: ttnn.Topology,
+ num_links: int,
+ fsdp: bool,
+ batch_size: int,
+ height: int,
+ width: int,
+ prompt_seq_len: int,
+ xformer_idx: int,
+ shard_prompt: bool,
+) -> None:
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ tp_factor = tuple(mesh_device.shape)[tp_axis]
+
+ logger.info(
+ f"test_transformer: mesh={tuple(mesh_device.shape)}, sp={sp_factor}, tp={tp_factor}, "
+ f"topology={topology}, num_links={num_links}"
+ )
+
+ model_name = "black-forest-labs/FLUX.2-dev"
+ torch_model = diffusers.Flux2Transformer2DModel.from_pretrained(
+ model_name, subfolder="transformer", torch_dtype=torch.bfloat16
+ )
+ torch_model.eval()
+ logger.info(f"single transformer blocks: {len(torch_model.single_transformer_blocks)}")
+ logger.info(f"transformer blocks: {len(torch_model.transformer_blocks)}")
+ if xformer_idx != 4:
+ torch_model.transformer_blocks = torch.nn.ModuleList(
+ [torch_model.transformer_blocks[0]] if xformer_idx > 0 else []
+ )
+ torch_model.single_transformer_blocks = torch.nn.ModuleList(
+ [torch_model.single_transformer_blocks[0]] if xformer_idx != 1 else []
+ )
+
+ ccl_manager = CCLManager(
+ mesh_device=mesh_device,
+ num_links=num_links,
+ topology=topology,
+ )
+
+ parallel_config = DiTGParallelConfigNoCFG(
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+
+ pos_embed = torch_model.pos_embed
+ head_dim = 128
+ num_heads = 48
+ in_channels = 128
+ joint_attention_dim = 15360
+ num_channels_latents = 128 # after patchify
+
+ if num_heads % parallel_config.tensor_parallel.factor != 0:
+ padding_config = PaddingConfig.from_tensor_parallel_factor(
+ num_heads,
+ head_dim,
+ parallel_config.tensor_parallel.factor,
+ )
+ else:
+ padding_config = None
+
+ tt_model = Flux2Transformer(
+ in_channels=in_channels,
+ num_layers=torch_model.config.num_layers if xformer_idx == 4 else 1 if xformer_idx > 0 else 0,
+ num_single_layers=torch_model.config.num_single_layers if xformer_idx == 4 else 1 if xformer_idx != 1 else 0,
+ attention_head_dim=head_dim,
+ num_attention_heads=num_heads,
+ joint_attention_dim=joint_attention_dim,
+ out_channels=128,
+ device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ padding_config=padding_config,
+ is_fsdp=fsdp,
+ shard_prompt=shard_prompt,
+ )
+
+ cache.load_model(
+ tt_model,
+ get_torch_state_dict=torch_model.state_dict,
+ model_name=os.path.basename(model_name),
+ subfolder="transformer",
+ parallel_config=parallel_config,
+ mesh_shape=tuple(mesh_device.shape),
+ is_fsdp=fsdp,
+ )
+
+ spatial_seq_len = height * width // 16**2
+
+ torch.manual_seed(0)
+ spatial = torch.randn([batch_size, spatial_seq_len, num_channels_latents])
+ prompt = torch.randn([batch_size, prompt_seq_len, joint_attention_dim])
+ timestep = torch.full([batch_size], fill_value=500)
+ guidance = torch.full([batch_size], fill_value=3)
+
+ # prepare for ROPE
+ text_ids = _prepare_ids(text_sequence_length=prompt_seq_len)
+ image_ids = _prepare_ids(height=height // 16, width=width // 16)
+ prompt_rope_cos, prompt_rope_sin = pos_embed.forward(text_ids)
+ spatial_rope_cos, spatial_rope_sin = pos_embed.forward(image_ids)
+
+ prompt_mesh_axes = [None, sp_axis, None] if shard_prompt else None
+ prompt_rope_mesh_axes = [None, None, sp_axis, None] if shard_prompt else None
+
+ tt_spatial = tensor.from_torch(spatial, device=mesh_device, mesh_axes=[None, sp_axis, None])
+ tt_prompt = tensor.from_torch(prompt, device=mesh_device, mesh_axes=prompt_mesh_axes)
+ tt_embedded_prompt = tt_model.context_embedder(tt_prompt)
+ tt_timestep = tensor.from_torch(timestep.unsqueeze(-1), dtype=ttnn.float32, device=mesh_device)
+ tt_guidance = tensor.from_torch(guidance.unsqueeze(-1), device=mesh_device)
+
+ tt_spatial_rope_cos = tensor.from_torch(
+ spatial_rope_cos.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=[None, None, sp_axis, None]
+ )
+ tt_spatial_rope_sin = tensor.from_torch(
+ spatial_rope_sin.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=[None, None, sp_axis, None]
+ )
+ tt_prompt_rope_cos = tensor.from_torch(
+ prompt_rope_cos.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=prompt_rope_mesh_axes
+ )
+ tt_prompt_rope_sin = tensor.from_torch(
+ prompt_rope_sin.unsqueeze(0).unsqueeze(0), device=mesh_device, mesh_axes=prompt_rope_mesh_axes
+ )
+ tt_combined_rope = (
+ ttnn.concat([tt_spatial_rope_cos, tt_prompt_rope_cos], dim=2),
+ ttnn.concat([tt_spatial_rope_sin, tt_prompt_rope_sin], dim=2),
+ )
+
+ # warmup
+ tt_model.forward(
+ spatial=tt_spatial,
+ embedded_prompt=tt_embedded_prompt,
+ timestep=tt_timestep,
+ guidance=tt_guidance,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ combined_rope=tt_combined_rope,
+ spatial_sequence_length=spatial_seq_len,
+ prompt_sequence_length=prompt_seq_len,
+ compute_prompt_output=(xformer_idx in [0, 3]), # force computing eveything
+ )
+ ttnn.synchronize_device(mesh_device)
+ signpost("t_start")
+ tt_output = tt_model.forward(
+ spatial=tt_spatial,
+ embedded_prompt=tt_embedded_prompt,
+ timestep=tt_timestep,
+ guidance=tt_guidance,
+ spatial_rope=(tt_spatial_rope_cos, tt_spatial_rope_sin),
+ prompt_rope=(tt_prompt_rope_cos, tt_prompt_rope_sin),
+ combined_rope=tt_combined_rope,
+ spatial_sequence_length=spatial_seq_len,
+ prompt_sequence_length=prompt_seq_len,
+ compute_prompt_output=(xformer_idx in [0, 3]), # force computing eveything
+ )
+ ttnn.synchronize_device(mesh_device)
+ signpost("t_end")
+
+
+def _prepare_ids(*, height: int = 1, width: int = 1, text_sequence_length: int = 1) -> torch.Tensor:
+ t = torch.arange(1)
+ h = torch.arange(height)
+ w = torch.arange(width)
+ s = torch.arange(text_sequence_length)
+
+ return torch.cartesian_prod(t, h, w, s)
diff --git a/code/models/tt_dit/tests/models/flux2/test_vae_flux2.py b/code/models/tt_dit/tests/models/flux2/test_vae_flux2.py
new file mode 100644
index 0000000000000000000000000000000000000000..cbdb723288e638fb043d1105f570b55a3fa6b07e
--- /dev/null
+++ b/code/models/tt_dit/tests/models/flux2/test_vae_flux2.py
@@ -0,0 +1,163 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+import diffusers.models.autoencoders.autoencoder_kl_flux2 as reference
+import pytest
+import torch
+
+import ttnn
+
+from ....models.vae.vae_flux2 import Flux2VaeDecoder
+from ....parallel.config import Flux2VaeParallelConfig
+from ....parallel.manager import CCLManager
+from ....utils import tensor
+from ....utils.check import assert_quality
+from .device_params import line_params_flux2, single_device_params_flux2
+
+_LAYERS_PER_BLOCK = 1 # pruned from pretrained (2) to keep host run fast
+
+
+def _load_pruned_torch_model() -> reference.AutoencoderKLFlux2:
+ """Load the pretrained VAE and truncate each up_block to _LAYERS_PER_BLOCK resnets.
+
+ The mid_block is left intact — it already uses num_layers=1 (2 resnets, 1 attention)
+ matching TT's VaeMidBlock default, regardless of layers_per_block.
+ """
+ model = reference.AutoencoderKLFlux2.from_pretrained("black-forest-labs/FLUX.2-dev", subfolder="vae")
+ assert isinstance(model, reference.AutoencoderKLFlux2)
+ model.eval()
+ keep = _LAYERS_PER_BLOCK + 1
+ for up_block in model.decoder.up_blocks:
+ up_block.resnets = torch.nn.ModuleList(list(up_block.resnets)[:keep])
+ return model
+
+
+def _torch_decode_reference(torch_model: reference.AutoencoderKLFlux2, inp: torch.Tensor) -> torch.Tensor:
+ """Apply BN inv-normalize, unpatchify, then run the torch decoder.
+
+ inp: [B, C*p^2, H_t, W_t] patchified latent.
+ Returns [B, 3, H, W] (NCHW) for comparison with the TT output permuted to NCHW.
+ """
+ p = 2 # _PATCH_SIZE
+ z_channels = torch_model.config.latent_channels
+ bn_eps = torch_model.bn.eps
+ s = (torch_model.bn.running_var + bn_eps).sqrt() # [C*p^2]
+ m = torch_model.bn.running_mean # [C*p^2]
+ # Matches TT _inv_normalize: z * sqrt(var + eps) + mean
+ inp_norm = inp * s.view(1, -1, 1, 1) + m.view(1, -1, 1, 1)
+
+ # Unpatchify: [B, C*p^2, H_t, W_t] -> [B, C, H_lat, W_lat]
+ B, _, h_t, w_t = inp_norm.shape
+ z = inp_norm.reshape(B, z_channels, p, p, h_t, w_t)
+ z = z.permute(0, 1, 4, 2, 5, 3).contiguous()
+ z = z.reshape(B, z_channels, h_t * p, w_t * p)
+
+ with torch.no_grad():
+ return torch_model.decode(z).sample # [B, 3, H, W]
+
+
+def prep_data(
+ vae_parallel_config: Flux2VaeParallelConfig,
+ inp: torch.Tensor,
+ mesh_device: ttnn.Device,
+ tt_model: Flux2VaeDecoder,
+) -> ttnn.Tensor:
+ # inp: [B, C*p^2, H_t, W_t] patchified latent in torch.
+ # Convert to patchified token format [B, H_t*W_t, C*p^2].
+ inp_flat = inp.permute(0, 2, 3, 1).flatten(1, 2)
+
+ if vae_parallel_config.h_parallel is not None:
+ tt_latents = tensor.from_torch(
+ inp_flat,
+ device=mesh_device,
+ mesh_axes=(None, vae_parallel_config.h_parallel.mesh_axis, None),
+ )
+ else:
+ tt_latents = tensor.from_torch(inp_flat, device=mesh_device)
+
+ p = Flux2VaeDecoder._PATCH_SIZE
+ height_for_unp = inp.shape[2] * p
+ width_for_unp = inp.shape[3] * p
+
+ tt_latents = tt_model.preprocess_and_unpatchify(
+ tt_latents,
+ height=height_for_unp,
+ width=width_for_unp,
+ )
+
+ # W-sharding can only be applied after unpatchify because H and W are interleaved
+ # in the patchified token dimension.
+ if vae_parallel_config.w_parallel is not None:
+ tt_latents = ttnn.mesh_partition(
+ tt_latents,
+ dim=2,
+ cluster_axis=vae_parallel_config.w_parallel.mesh_axis,
+ memory_config=tt_latents.memory_config(),
+ )
+
+ return tt_latents
+
+
+@pytest.mark.parametrize(
+ "mesh_device, device_params",
+ [
+ [(1, 1), single_device_params_flux2],
+ [(2, 2), line_params_flux2],
+ [(1, 8), line_params_flux2],
+ [(4, 8), line_params_flux2],
+ ],
+ ids=["1x1", "2x2", "1x8", "4x8"],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "height", "width"),
+ [(1, 1024, 1024)],
+ ids=["1024"],
+)
+def test_vae_flux2_decoder(
+ *,
+ mesh_device: ttnn.Device,
+ batch_size: int,
+ height: int,
+ width: int,
+) -> None:
+ torch.manual_seed(0)
+
+ tp_axis = 0
+ h_axis = 1 - tp_axis
+ w_axis = None
+
+ torch_model = _load_pruned_torch_model()
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+ vae_parallel_config = Flux2VaeParallelConfig.from_axes(mesh_device, tp_axis=tp_axis, h_axis=h_axis, w_axis=w_axis)
+
+ z_channels = torch_model.config.latent_channels
+ patch_size = 2
+ vae_scale_factor = 8
+
+ tt_model = Flux2VaeDecoder(
+ out_channels=torch_model.config.out_channels,
+ block_out_channels=list(torch_model.config.block_out_channels),
+ layers_per_block=_LAYERS_PER_BLOCK,
+ z_channels=z_channels,
+ device=mesh_device,
+ parallel_config=vae_parallel_config,
+ ccl_manager=ccl_manager,
+ use_conv3d=False,
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ f = vae_scale_factor * patch_size
+ inp = torch.randn(batch_size, z_channels * patch_size**2, height // f, width // f)
+
+ torch_output = _torch_decode_reference(torch_model, inp) # [B, 3, H, W]
+
+ tt_inp = prep_data(vae_parallel_config, inp, mesh_device, tt_model=tt_model)
+ tt_out = tt_model.forward(tt_inp)
+ ttnn.synchronize_device(mesh_device)
+
+ tt_out_torch = tensor.to_torch(tt_out).permute(0, 3, 1, 2) # [B, 3, H, W]
+ assert_quality(torch_output, tt_out_torch, pcc=0.9978, relative_rmse=0.034)
diff --git a/code/models/tt_dit/tests/models/ideogram4/test_perf_ideogram4.py b/code/models/tt_dit/tests/models/ideogram4/test_perf_ideogram4.py
new file mode 100644
index 0000000000000000000000000000000000000000..1266a4b6a3a885997b5225133c919ab5d7026cf6
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ideogram4/test_perf_ideogram4.py
@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# =============================================================================
+# Ideogram 4.0 perf: end-to-end latency of the device-resident Ideogram4Pipeline.
+# Warms the program cache, then times encode / denoise / decode via the standard
+# tt_dit event -> BenchmarkProfiler path, emits BenchmarkData (in CI), and asserts
+# each section stays under a conservative ceiling so a gross regression fails the
+# test. Single-block device perf comes from the PCC test (test_transformer_ideogram4.py).
+# =============================================================================
+
+import os
+
+import pytest
+from loguru import logger
+
+from models.perf.benchmarking_utils import BenchmarkData, BenchmarkProfiler
+
+from ....pipelines.events import profiler_event_callback
+from ....pipelines.ideogram4.pipeline_ideogram4 import Ideogram4Pipeline
+from ....utils.test import line_params, ring_params
+
+# The pipeline needs the gated fp8 checkpoint (create_pipeline raises ValueError without it);
+# skip (don't error) when it isn't configured, matching the component tests.
+_NEEDS_WEIGHTS = pytest.mark.skipif(
+ not os.environ.get("IDEOGRAM4_WEIGHTS"), reason="IDEOGRAM4_WEIGHTS not set (gated fp8 checkpoint)"
+)
+
+PROMPT = "a watercolor painting of a red panda reading a book under a cherry tree, soft morning light"
+
+_LINE_PIPE = {**line_params, "l1_small_size": 32768, "trace_region_size": 60000000}
+_RING_PIPE = {**ring_params, "l1_small_size": 32768, "trace_region_size": 60000000}
+
+# Conservative upper-bound targets (seconds), program-cache warm, 512px V4_TURBO_12. Generous
+# placeholders that catch a gross (>~2x) regression without flaking; tighten once stable CI
+# baselines exist. Keyed by mesh shape.
+_PERF_TARGETS = {
+ (2, 4): {"encoder": 30.0, "denoising": 90.0, "vae": 30.0, "total": 130.0}, # BH LoudBox 2x4 SP4xTP2
+ (4, 8): {"encoder": 30.0, "denoising": 60.0, "vae": 30.0, "total": 100.0}, # BH Galaxy TP4xSP8
+}
+
+
+# Parallel config (tp/sp/num_links/topology) is DISCOVERED from the mesh shape via
+# Ideogram4Pipeline._PRESETS, so the perf test only picks a mesh shape + its device_params.
+@_NEEDS_WEIGHTS
+@pytest.mark.parametrize(
+ ("mesh_device", "device_params"),
+ [
+ pytest.param((2, 4), _LINE_PIPE, id="sp4tp2"), # BH loudbox 2x4: SP4 x TP2
+ pytest.param((4, 8), _RING_PIPE, id="bh_galaxy_sp8tp4"), # BH Galaxy: TP4 x SP8 (Ring)
+ ],
+ indirect=True,
+)
+@pytest.mark.parametrize(("height", "width", "preset"), [(512, 512, "V4_TURBO_12")], ids=["512px_turbo12"])
+def test_latency(*, mesh_device, is_ci_env, height, width, preset) -> None:
+ pipe = Ideogram4Pipeline.create_pipeline(mesh_device=mesh_device, height=height, width=width)
+
+ pipe(prompts=[PROMPT], preset=preset, seed=1) # warmup (compile + cache)
+ profiler = BenchmarkProfiler()
+ pipe(prompts=[PROMPT], preset=preset, seed=2, on_event=profiler_event_callback(profiler, 0)) # timed (warm)
+
+ steps = ("encoder", "denoising", "vae", "total")
+ measurements = {s: profiler.get_duration(s) for s in steps}
+ logger.info(f"LATENCY {height}px {preset}: " + " | ".join(f"{s} {measurements[s]:.2f}s" for s in steps))
+
+ targets = _PERF_TARGETS[tuple(mesh_device.shape)]
+ if is_ci_env:
+ benchmark_data = BenchmarkData()
+ for s in steps:
+ benchmark_data.add_measurement(
+ profiler=profiler, iteration=0, step_name=s, name=s, value=measurements[s], target=targets[s]
+ )
+ benchmark_data.save_partial_run_json(
+ profiler,
+ run_type="BH_LB" if tuple(mesh_device.shape) == (2, 4) else "BH_GLX",
+ ml_model_name="Ideogram4",
+ batch_size=1,
+ config_params={"width": width, "height": height, "preset": preset},
+ )
+
+ regressions = [f"{s}: {measurements[s]:.2f}s > {targets[s]:.1f}s" for s in steps if measurements[s] > targets[s]]
+ assert not regressions, "perf regression: " + "; ".join(regressions)
diff --git a/code/models/tt_dit/tests/models/ideogram4/test_pipeline_class_ideogram4.py b/code/models/tt_dit/tests/models/ideogram4/test_pipeline_class_ideogram4.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c071029b2640e62be4998e1938925e07389e17e
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ideogram4/test_pipeline_class_ideogram4.py
@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# =============================================================================
+# Drive the Ideogram4Pipeline class end-to-end on device (TP=4, all models
+# resident: encoder + cond/uncond transformers + VAE) with real weights, across
+# presets and resolutions. Saves each image and asserts it is non-degenerate
+# (correct shape, std > 1). A generation smoke test, not a PCC gate.
+# =============================================================================
+
+import os
+
+import numpy as np
+import pytest
+from loguru import logger
+
+from ....pipelines.ideogram4.pipeline_ideogram4 import Ideogram4Pipeline
+from ....utils.test import line_params, ring_params
+
+# The pipeline needs the gated fp8 checkpoint (create_pipeline raises ValueError without it);
+# skip (don't error) when it isn't configured, matching the component tests.
+_NEEDS_WEIGHTS = pytest.mark.skipif(
+ not os.environ.get("IDEOGRAM4_WEIGHTS"), reason="IDEOGRAM4_WEIGHTS not set (gated fp8 checkpoint)"
+)
+
+
+def _typography_prompt_json() -> str:
+ """Typography-heavy in-distribution JSON caption that showcases Ideogram 4's
+ best-in-class text rendering (crisp legible lettering + bbox layout + palette).
+
+ Ideogram 4 was trained exclusively on structured-JSON captions; quoting the exact
+ text strings and placing them with bbox is the reliable way to get them rendered.
+ """
+ import json
+
+ caption = {
+ "aspect_ratio": "1:1",
+ "high_level_description": (
+ "A bold retro travel poster with large, crisp typography. A condensed sans-serif "
+ 'headline reading "EXPLORE THE WILD" arcs across the top, a clean letter-spaced '
+ 'subtitle "NATIONAL PARKS · EST. 1872" sits beneath it, and a small banner at the '
+ 'bottom reads "ADVENTURE AWAITS". Set over a stylized sunrise mountain landscape '
+ "with pine silhouettes; flat vintage screen-print look, sharp legible lettering."
+ ),
+ "colour_palette": ["#F4A259", "#2E4600", "#1B3A4B", "#F2E8CF", "#BC4749"],
+ "compositional_deconstruction": {
+ "background": (
+ "Stylized sunrise sky in warm cream and orange bands, layered mountain ridges in "
+ "deep green and teal, a row of pine-tree silhouettes along the lower third, flat "
+ "screen-print texture with subtle grain."
+ ),
+ "elements": [
+ {
+ "type": "text",
+ "bbox": [70, 70, 290, 930],
+ "text": "EXPLORE THE WILD",
+ "desc": "Large bold condensed sans-serif headline in cream, all caps, gentle "
+ "upward arc, high contrast against the sky.",
+ },
+ {
+ "type": "text",
+ "bbox": [300, 180, 400, 820],
+ "text": "NATIONAL PARKS · EST. 1872",
+ "desc": "Smaller clean uppercase subtitle in warm orange, generously "
+ "letter-spaced, centered beneath the headline.",
+ },
+ {
+ "type": "text",
+ "bbox": [840, 250, 960, 750],
+ "text": "ADVENTURE AWAITS",
+ "desc": "Small bold uppercase banner text in cream on a deep-red rounded ribbon.",
+ },
+ {
+ "type": "obj",
+ "bbox": [410, 60, 880, 940],
+ "desc": "Layered mountain range at sunrise with a bold circular sun disc, deep "
+ "green and teal ridges, vintage poster shading.",
+ },
+ ],
+ },
+ }
+ return json.dumps(caption, ensure_ascii=False, separators=(",", ":"))
+
+
+# Typography showcase prompt (JSON is the known-good format; plain prose washes out
+# at >=1024px per the in-weights-moderation finding).
+PROMPT = _typography_prompt_json()
+
+
+# Trace region for the per-branch transformer-forward traces (cond + uncond).
+# 60MB headroom; bump if "trace region" OOM at higher resolution.
+# l1_small_size=32KB: only the VAE conv draws from L1_SMALL (proven sufficient in the VAE tests);
+# the denoiser matmuls use none, so 64KB over-reserved the shared L1 region and made the SP4xTP2
+# path's resident ring/all-gather buffers clash with the matmul circular buffers.
+_LINE_PIPE = {**line_params, "l1_small_size": 32768, "trace_region_size": 60000000}
+_RING_PIPE = {**ring_params, "l1_small_size": 32768, "trace_region_size": 60000000}
+
+
+# Pipeline mesh configs: the parallel config (tp/sp/num_links/topology) is DISCOVERED from the
+# mesh shape via Ideogram4Pipeline._PRESETS, so tests only pick a mesh shape + its device_params.
+_PIPE_CONFIGS = [
+ pytest.param((2, 4), _LINE_PIPE, id="sp4tp2"), # BH loudbox 2x4: SP4 x TP2
+ pytest.param((4, 8), _RING_PIPE, id="bh_galaxy_sp8tp4"), # BH Galaxy: TP4 x SP8 (Ring)
+]
+
+
+@_NEEDS_WEIGHTS
+@pytest.mark.parametrize(("mesh_device", "device_params"), _PIPE_CONFIGS, indirect=True)
+@pytest.mark.parametrize(
+ ("height", "width"), [(512, 512), (1024, 1024), (2048, 2048)], ids=["512px", "1024px", "2048px"]
+)
+@pytest.mark.parametrize("preset", ["V4_TURBO_12", "V4_QUALITY_48"])
+def test_pipeline_class(*, mesh_device, height, width, preset) -> None:
+ # Parallel config is discovered from the mesh shape (_PRESETS); no per-test tp/sp/links wiring.
+ pipe = Ideogram4Pipeline.create_pipeline(mesh_device=mesh_device, height=height, width=width)
+ img = pipe(prompts=[PROMPT], preset=preset, seed=1234, traced=False)[0]
+
+ out_dir = os.environ.get("IDEOGRAM4_OUT_DIR", "generated")
+ os.makedirs(out_dir, exist_ok=True)
+ out = os.path.join(out_dir, f"ideogram4_pipeline_class_{height}_{preset}.png")
+ img.save(out)
+ arr = np.asarray(img)
+ logger.info(f"pipeline class saved {out} shape={arr.shape} std={arr.std():.1f} preset={preset}")
+ assert arr.shape == (height, width, 3)
+ assert arr.std() > 1.0
+
+
+@_NEEDS_WEIGHTS
+@pytest.mark.parametrize(("mesh_device", "device_params"), _PIPE_CONFIGS, indirect=True)
+def test_2048_presets(*, mesh_device) -> None:
+ """UNTRACED 2048px generation across all 3 sampler presets, saved for manual review.
+
+ Builds the pipeline ONCE and loops the presets (avoids a per-preset model reload).
+ Untraced (traced=False) per the serving plan; 2048px doesn't benefit from tracing.
+ """
+ pipe = Ideogram4Pipeline.create_pipeline(mesh_device=mesh_device, height=2048, width=2048)
+ for preset in ("V4_TURBO_12", "V4_DEFAULT_20", "V4_QUALITY_48"):
+ img = pipe(prompts=[PROMPT], preset=preset, seed=1234, traced=False)[0]
+ out_dir = os.environ.get("IDEOGRAM4_OUT_DIR", "generated")
+ os.makedirs(out_dir, exist_ok=True)
+ out = os.path.join(out_dir, f"ideogram4_2048_{preset}.png")
+ img.save(out)
+ arr = np.asarray(img)
+ logger.info(f"2048px UNTRACED {preset}: saved {out} std={arr.std():.1f}")
+ assert arr.shape == (2048, 2048, 3)
+ assert arr.std() > 1.0
diff --git a/code/models/tt_dit/tests/models/ideogram4/test_pipeline_ideogram4.py b/code/models/tt_dit/tests/models/ideogram4/test_pipeline_ideogram4.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b00529782d8d3424eebd9266aa85414b9d82271
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ideogram4/test_pipeline_ideogram4.py
@@ -0,0 +1,121 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# =============================================================================
+# Ideogram 4.0 pipeline glue. The denoise loop (CFG blend + Euler) is covered by
+# the sampler + transformer tests and the text encode by the encoder test; this
+# verifies the pipeline-specific NEW device computation:
+# * the decode tail: per-channel latent denorm -> 2x2 unpatch -> VAE decode,
+# end-to-end on device vs a torch reference (latent_norm + diffusers
+# AutoencoderKL.decode);
+# * the asymmetric-CFG velocity blend v = gw*v_cond + (1-gw)*v_uncond.
+# =============================================================================
+
+import pytest
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from loguru import logger
+
+import ttnn
+
+from ....models.vae.vae_ideogram4 import Ideogram4VAEDecoder
+from ....parallel.config import ParallelFactor, VAEParallelConfig
+from ....parallel.manager import CCLManager
+from ....pipelines.ideogram4.pipeline_ideogram4 import (
+ Ideogram4DecodeStage,
+ cfg_blend,
+ interleave_layer_taps,
+ unpatchify_latent,
+)
+from ....reference.ideogram4.constants import QWEN3_VL_ACTIVATION_LAYERS
+from ....reference.ideogram4.latent_norm import get_latent_norm
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.tensor import bf16_tensor
+
+Z_CHANNELS = 32
+PATCH = 2
+
+
+def test_llm_tap_interleave_feature_major() -> None:
+ """Host-only: lock the 13-layer Qwen tap assembly to FEATURE-MAJOR (out[..., f*L + l] =
+ taps[l][..., f]). This is the vendor order the real llm_cond_proj weights expect; a layer-major
+ flip would silently scramble conditioning. Marks each tap l with the constant l so the index
+ mapping is checkable without weights/device."""
+ layers = len(QWEN3_VL_ACTIVATION_LAYERS) # 13
+ d, n = 4, 3 # tiny feature dim + seq for a fast host check
+ taps = [torch.full((1, n, d), float(l)) for l in range(layers)] # tap l is all-l
+ out = interleave_layer_taps(taps)
+
+ assert out.shape == (1, n, d * layers)
+ # feature-major: the value at flat index f*layers + l must be tap l's marker (l)
+ for f in range(d):
+ for l in range(layers):
+ assert out[0, 0, f * layers + l].item() == float(l), f"interleave not feature-major at f={f}, l={l}"
+
+
+def _build_akl(torch_dtype):
+ return (
+ AutoencoderKL(
+ in_channels=3,
+ out_channels=3,
+ latent_channels=Z_CHANNELS,
+ down_block_types=("DownEncoderBlock2D",) * 4,
+ up_block_types=("UpDecoderBlock2D",) * 4,
+ block_out_channels=(128, 256, 512, 512),
+ layers_per_block=2,
+ norm_num_groups=32,
+ )
+ .to(torch_dtype)
+ .eval()
+ )
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+@pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
+@pytest.mark.parametrize(("grid_h", "grid_w"), [(16, 16)], ids=["256px"])
+def test_pipeline_decode(*, mesh_device: ttnn.MeshDevice, grid_h: int, grid_w: int) -> None:
+ """Decode tail end-to-end (latent denorm + 2x2 unpatch + VAE) vs torch reference."""
+ torch.manual_seed(0)
+ torch_dtype = torch.bfloat16
+ tokens = grid_h * grid_w
+
+ akl = _build_akl(torch_dtype)
+ parallel_config = VAEParallelConfig(tensor_parallel=ParallelFactor(factor=1, mesh_axis=1))
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+ vae = Ideogram4VAEDecoder.from_torch(
+ akl, mesh_device=mesh_device, parallel_config=parallel_config, ccl_manager=ccl_manager
+ )
+ stage = Ideogram4DecodeStage(vae, mesh_device=mesh_device, patch=PATCH)
+
+ z = torch.randn(1, tokens, PATCH * PATCH * Z_CHANNELS, dtype=torch.float32) # patchified latent
+
+ # torch reference: denorm -> unpatch -> diffusers decode
+ shift, scale = get_latent_norm()
+ z_dn = (z * scale + shift).to(torch_dtype)
+ z_nchw = unpatchify_latent(z_dn, grid_h=grid_h, grid_w=grid_w, patch=PATCH)
+ with torch.no_grad():
+ ref = akl.decode(z_nchw).sample
+
+ out = stage.decode(bf16_tensor(z, device=mesh_device), grid_h=grid_h, grid_w=grid_w)
+
+ logger.info(f"ideogram4 pipeline decode: {tokens} tokens -> image {grid_h*PATCH*8}x{grid_w*PATCH*8}")
+ assert_quality(ref.float(), out.float(), pcc=0.99)
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+@pytest.mark.parametrize("device_params", [{}], indirect=True)
+def test_pipeline_cfg_blend(*, mesh_device: ttnn.MeshDevice) -> None:
+ """Asymmetric CFG: v = gw*v_cond + (1-gw)*v_uncond on device vs torch."""
+ torch.manual_seed(0)
+ gw = 7.0
+ v_cond = torch.randn(1, 4096, 128, dtype=torch.float32)
+ v_uncond = torch.randn(1, 4096, 128, dtype=torch.float32)
+ ref = gw * v_cond + (1.0 - gw) * v_uncond
+
+ out = cfg_blend(bf16_tensor(v_cond, device=mesh_device), bf16_tensor(v_uncond, device=mesh_device), gw)
+ out_torch = tensor.to_torch(out, mesh_axes=[None, None, None])
+
+ logger.info("ideogram4 pipeline CFG blend")
+ assert_quality(ref, out_torch, pcc=0.999)
diff --git a/code/models/tt_dit/tests/models/ideogram4/test_sampler_ideogram4.py b/code/models/tt_dit/tests/models/ideogram4/test_sampler_ideogram4.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b32f0a216a67b8724aa6df0a28e3935bad44dd2
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ideogram4/test_sampler_ideogram4.py
@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# =============================================================================
+# Ideogram 4.0 Euler flow-matching sampler. Two checks:
+# 1. host: the preset table + logit-normal schedule / guidance weights are well-formed
+# (num_steps/guidance/mu/std match spec; times in [0,1], monotonic; gw per schedule).
+# 2. device: a full Euler step loop z <- z + v*(s-t) matches a torch reference.
+# =============================================================================
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+
+from ....pipelines.ideogram4.sampler import PRESETS, Ideogram4Sampler
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.tensor import bf16_tensor
+
+# Expected preset spec (num_steps, guidance schedule, logit-normal mean/std) — the
+# official Ideogram 4.0 values, kept here as an independent regression against sampler.PRESETS.
+_EXPECTED = {
+ "V4_QUALITY_48": dict(num_steps=48, gw=(3.0,) * 3 + (7.0,) * 45, mu=0.0, std=1.5),
+ "V4_DEFAULT_20": dict(num_steps=20, gw=(3.0,) * 2 + (7.0,) * 18, mu=0.0, std=1.75),
+ "V4_TURBO_12": dict(num_steps=12, gw=(3.0,) * 1 + (7.0,) * 11, mu=0.5, std=1.75),
+}
+
+
+@pytest.mark.parametrize("preset", list(PRESETS))
+@pytest.mark.parametrize("resolution", [(1024, 1024), (2048, 768)], ids=["1024sq", "2048x768"])
+def test_sampler_schedule(*, preset: str, resolution: tuple[int, int]) -> None:
+ """Host-only: preset table matches spec and the logit-normal schedule is well-formed."""
+ h, w = resolution
+ params = PRESETS[preset]
+ exp = _EXPECTED[preset]
+
+ assert params.num_steps == exp["num_steps"]
+ assert tuple(float(g) for g in params.guidance_schedule) == exp["gw"]
+ assert (params.mu, params.std) == (exp["mu"], exp["std"])
+
+ samp = Ideogram4Sampler.from_preset(preset, height=h, width=w)
+ times = [samp.times_for_step(i) for i in range(params.num_steps)]
+
+ for i, (t_val, s_val) in enumerate(times):
+ assert 0.0 <= t_val <= 1.0, f"step {i}: t {t_val} out of [0,1]"
+ assert 0.0 <= s_val <= 1.0, f"step {i}: s {s_val} out of [0,1]"
+ assert samp.guidance_weight(i) == float(params.guidance_schedule[i])
+
+ # the schedule must sweep monotonically across the denoising loop
+ t_seq = [t for t, _ in times]
+ assert t_seq == sorted(t_seq) or t_seq == sorted(t_seq, reverse=True), f"non-monotonic schedule: {t_seq}"
+
+ logger.info(f"sampler schedule OK: {preset} @ {resolution}, {params.num_steps} steps")
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
+@pytest.mark.parametrize("device_params", [{}], indirect=True)
+@pytest.mark.parametrize("preset", ["V4_TURBO_12"])
+def test_sampler_euler_step_device(*, mesh_device: ttnn.MeshDevice, preset: str) -> None:
+ """Device: the full Euler step loop matches torch z <- z + v*(s-t)."""
+ torch.manual_seed(0)
+ h = w = 1024
+ samp = Ideogram4Sampler.from_preset(preset, height=h, width=w)
+
+ batch, tokens, ch = 1, 4096, 128
+ z = torch.randn(batch, tokens, ch, dtype=torch.float32)
+ tt_z = bf16_tensor(z, device=mesh_device)
+
+ for i in reversed(range(samp.num_steps)):
+ v = torch.randn(batch, tokens, ch, dtype=torch.float32) # stand-in velocity
+ t_val, s_val = samp.times_for_step(i)
+ z = z + v * (s_val - t_val) # torch reference
+ tt_z = samp.step(tt_z, bf16_tensor(v, device=mesh_device), i)
+
+ tt_z_torch = tensor.to_torch(tt_z, mesh_axes=[None, None, None])
+ logger.info(f"sampler device step OK: {preset}, {samp.num_steps} steps")
+ assert_quality(z, tt_z_torch, pcc=0.999)
diff --git a/code/models/tt_dit/tests/models/ideogram4/test_transformer_ideogram4.py b/code/models/tt_dit/tests/models/ideogram4/test_transformer_ideogram4.py
new file mode 100644
index 0000000000000000000000000000000000000000..41a0884c32a358cd1a9815dbb112f309da4f86a2
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ideogram4/test_transformer_ideogram4.py
@@ -0,0 +1,522 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# =============================================================================
+# Ideogram 4.0 single-stream transformer PCC vs the OFFICIAL reference. Two tests:
+# * test_transformer_block: one Ideogram4TransformerBlock (random init) on real
+# Ideogram dims across the TP/SP/segment sweep.
+# * test_transformer_model: the full transformer (top + blocks + bottom) under
+# parallelism, random or real weights.
+# Both load the reference state_dict into the tt module and assert PCC >= 0.99.
+#
+# Block coverage (2x4 BH mesh, FABRIC_1D, Linear): {B=1,2} x {1,2 segments} x
+# {1088, 4224 seq} x {tp1sp1, tp2, tp4(pad 20), sp2, sp2tp2, sp2tp4}.
+#
+# 1. rotate-half MRoPE (cos/sin = cat(freqs, freqs)). Precomputed on host here
+# from the reference Ideogram4MRoPE so the convention is identical.
+# 2. block-diagonal segment-id mask -> additive bias (0 / -inf) fed to SDPA;
+# with SP it is sharded on the query-row dim (K/V are all-gathered).
+# 3. tanh-gated 4-branch AdaLN (no shift terms).
+# 4. double-RMSNorm sandwich residual (norm2 on the attn/ff *output*).
+# 5. head_dim=256 / 18 heads — TP pads heads up to a multiple of tp_factor
+# (18 -> 20 for tp=4); SP sharding uses ring attention (or K/V gather + mask).
+# =============================================================================
+
+import math
+import os
+
+import pytest
+import torch
+from loguru import logger
+from safetensors.torch import load_file
+
+import ttnn
+
+from ....models.transformers.transformer_ideogram4 import (
+ Ideogram4Transformer,
+ Ideogram4TransformerBlock,
+ rope_halfsplit_to_interleaved,
+)
+from ....parallel.config import DiTParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....reference.ideogram4 import modeling_ideogram4
+from ....reference.ideogram4.constants import IMAGE_POSITION_OFFSET, LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR
+from ....reference.ideogram4.dequant import dequant_fp8_state_dict
+from ....utils import tensor
+from ....utils.check import assert_quality
+from ....utils.padding import PaddingConfig
+from ....utils.tensor import bf16_tensor
+from ....utils.test import line_params, ring_params
+
+# Real shipped fp8 checkpoint (ideogram-ai/ideogram-4-fp8), dequantized to bf16. Override via
+# env for CI. Used by the weights="real" parametrizations (weights="random" uses fresh init).
+FP8 = os.environ.get("IDEOGRAM4_WEIGHTS")
+# Real-weight cases need the gated fp8 checkpoint; skip (don't error) when it isn't configured.
+_NEEDS_WEIGHTS = pytest.mark.skipif(not FP8, reason="IDEOGRAM4_WEIGHTS not set (gated fp8 checkpoint)")
+
+
+def _real_transformer_sd(num_layers: int) -> dict:
+ """Dequantized real transformer weights, filtered to the first num_layers layers so the
+ (strict) tt loader accepts a truncated-depth model."""
+ sd = dequant_fp8_state_dict(load_file(f"{FP8}/transformer/diffusion_pytorch_model.safetensors"))
+ return {k: v for k, v in sd.items() if not (k.startswith("layers.") and int(k.split(".")[1]) >= num_layers)}
+
+
+# Ring SDPA (sequence parallel) only needs each per-device shard tile-aligned, i.e. the
+# padded sequence divisible by tile_size * sp_factor. Must match pipeline._sp_padded_len.
+_SP_TILE = 32
+
+
+def _sp_padded_len(seq_len: int, sp_factor: int) -> int:
+ if sp_factor == 1:
+ return seq_len
+ divisor = _SP_TILE * sp_factor
+ return math.ceil(seq_len / divisor) * divisor
+
+
+# Real Ideogram 4.0 denoiser dims (confirmed against reference Ideogram4Config).
+EMB_DIM = 4608
+NUM_HEADS = 18
+HEAD_DIM = EMB_DIM // NUM_HEADS # 256
+INTERMEDIATE_SIZE = 12288
+ADALN_DIM = 512
+ROPE_THETA = 5_000_000
+MROPE_SECTION = (24, 20, 20)
+NORM_EPS = 1e-5
+
+
+def _build_inputs(batch_size: int, text_len: int, image_len: int, torch_dtype, num_segments: int = 1):
+ """Construct a realistic single-stream batch: [text tokens | image tokens].
+
+ Returns torch tensors mirroring the reference block's forward signature plus
+ the derived cos/sin (B, L, head_dim) and a (B, 1, L, L) additive attn mask.
+
+ num_segments > 1 packs that many independent samples into the sequence (contiguous,
+ roughly equal spans) so the block-diagonal segment mask blocks cross-segment
+ attention — exercising the masked SDPA path.
+ """
+ seq_len = text_len + image_len
+
+ # hidden states (already-fused h = x + llm_features + indicator embed upstream).
+ x = torch.randn(batch_size, seq_len, EMB_DIM, dtype=torch_dtype)
+
+ # adaln_input: reference passes F.silu(adaln_proj(t_cond)); per-sample => (B,1,adaln_dim).
+ adaln_input = torch.randn(batch_size, 1, ADALN_DIM, dtype=torch_dtype)
+
+ # segment ids: contiguous spans => block-diagonal attention. num_segments==1 is the
+ # all-same-segment (full attention) case.
+ segment_ids = torch.zeros(batch_size, seq_len, dtype=torch.long)
+ if num_segments > 1:
+ bounds = torch.linspace(0, seq_len, num_segments + 1).round().long()
+ for s in range(num_segments):
+ segment_ids[:, bounds[s] : bounds[s + 1]] = s
+
+ # position_ids (B, L, 3) = (t, h, w). Text tokens: 1D positions on all 3 axes.
+ # Image tokens: a (h, w) grid offset so they never collide with text positions.
+ position_ids = torch.zeros(batch_size, seq_len, 3, dtype=torch.long)
+ # text: linear positions 0..text_len-1 on every axis.
+ text_pos = torch.arange(text_len)
+ position_ids[:, :text_len, 0] = text_pos
+ position_ids[:, :text_len, 1] = text_pos
+ position_ids[:, :text_len, 2] = text_pos
+ # image: square-ish grid.
+ grid_h = int(round(image_len**0.5))
+ while image_len % grid_h != 0:
+ grid_h -= 1
+ grid_w = image_len // grid_h
+ hh = torch.arange(grid_h).repeat_interleave(grid_w)
+ ww = torch.arange(grid_w).repeat(grid_h)
+ position_ids[:, text_len:, 0] = IMAGE_POSITION_OFFSET
+ position_ids[:, text_len:, 1] = IMAGE_POSITION_OFFSET + hh
+ position_ids[:, text_len:, 2] = IMAGE_POSITION_OFFSET + ww
+
+ # cos/sin from the reference MRoPE (ground-truth convention).
+ rope = modeling_ideogram4.Ideogram4MRoPE(head_dim=HEAD_DIM, base=ROPE_THETA, mrope_section=MROPE_SECTION)
+ cos, sin = rope(position_ids) # (B, L, head_dim), float32
+ cos = cos.to(torch_dtype)
+ sin = sin.to(torch_dtype)
+
+ # additive attention bias from segment ids: (B, 1, L, L), 0 attend / -inf block.
+ same_seg = segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1) # (B, L, L)
+ attn_bias = torch.zeros(batch_size, 1, seq_len, seq_len, dtype=torch.float32)
+ attn_bias.masked_fill_(~same_seg.unsqueeze(1), float("-inf"))
+
+ return x, adaln_input, segment_ids, cos, sin, attn_bias
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "sp_axis", "tp_axis", "num_links", "device_params", "topology"),
+ [
+ # factor == submesh axis size. TP shards heads (padded to a multiple of tp);
+ # SP shards the sequence (ring attention). Full mesh is the 2x4 BH loudbox.
+ # factor == submesh axis size; sp on axis 0 (size up to 2), tp on axis 1 (size up to 4).
+ pytest.param(
+ (2, 4), (1, 1), 0, 1, 1, line_params, ttnn.Topology.Linear, id="tp1sp1"
+ ), # regression: single device
+ pytest.param(
+ (2, 4), (1, 2), 0, 1, 1, line_params, ttnn.Topology.Linear, id="tp2"
+ ), # 18 heads % 2 == 0, no padding
+ pytest.param((2, 4), (1, 4), 0, 1, 1, line_params, ttnn.Topology.Linear, id="tp4_pad20"), # 18 -> 20 heads
+ pytest.param(
+ (2, 4), (2, 1), 0, 1, 1, line_params, ttnn.Topology.Linear, id="sp2"
+ ), # sequence parallel only (TP=1)
+ pytest.param((2, 4), (2, 2), 0, 1, 1, line_params, ttnn.Topology.Linear, id="sp2tp2"),
+ pytest.param(
+ (2, 4), (2, 4), 0, 1, 1, line_params, ttnn.Topology.Linear, id="sp2tp4_pad20"
+ ), # SP=2, TP=4 (pad 20)
+ pytest.param(
+ (2, 4), (2, 4), 1, 0, 2, line_params, ttnn.Topology.Linear, id="sp4tp2"
+ ), # SP=4 (axis 1), TP=2 (axis 0) — full 2x4 loudbox; num_links=2
+ # BH Galaxy 4x8, 2D torus Ring: SP=8 (axis 1), TP=4 (axis 0), 2 links/neighbor. 18 heads -> pad 20.
+ pytest.param((4, 8), (4, 8), 1, 0, 2, ring_params, ttnn.Topology.Ring, id="bh_galaxy_sp8tp4"),
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+ ("batch_size", "text_len", "image_len"),
+ [
+ pytest.param(1, 64, 1024, id="b1_text64_img1024"),
+ pytest.param(1, 128, 4096, id="b1_text128_img4096"), # larger sequence (4224)
+ ],
+)
+@pytest.mark.parametrize("num_segments", [1, 2], ids=["1seg", "2seg"])
+def test_transformer_block(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ submesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ batch_size: int,
+ text_len: int,
+ image_len: int,
+ num_segments: int,
+) -> None:
+ torch.manual_seed(0)
+ torch_dtype = torch.bfloat16
+
+ seq_len = text_len + image_len
+
+ submesh_device = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ sp_factor = tuple(submesh_device.shape)[sp_axis]
+ tp_factor = tuple(submesh_device.shape)[tp_axis]
+
+ # ---- reference (ground truth), random init ----
+ torch_block = modeling_ideogram4.Ideogram4TransformerBlock(
+ hidden_size=EMB_DIM,
+ intermediate_size=INTERMEDIATE_SIZE,
+ num_heads=NUM_HEADS,
+ norm_eps=NORM_EPS,
+ adanln_dim=ADALN_DIM,
+ ).to(dtype=torch_dtype)
+ torch_block.eval()
+
+ x, adaln_input, segment_ids, cos, sin, attn_bias = _build_inputs(
+ batch_size, text_len, image_len, torch_dtype, num_segments
+ )
+
+ with torch.no_grad():
+ torch_out = torch_block(
+ x,
+ segment_ids=segment_ids,
+ cos=cos,
+ sin=sin,
+ adaln_input=adaln_input,
+ )
+
+ # ---- tt block ----
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+ ccl_manager = CCLManager(submesh_device, num_links=num_links, topology=topology)
+
+ # Head padding for TP divisibility (18 heads -> next multiple of tp_factor).
+ padding_config = (
+ PaddingConfig.from_tensor_parallel_factor(NUM_HEADS, HEAD_DIM, tp_factor)
+ if NUM_HEADS % tp_factor != 0
+ else None
+ )
+
+ tt_block = Ideogram4TransformerBlock(
+ hidden_size=EMB_DIM,
+ intermediate_size=INTERMEDIATE_SIZE,
+ num_heads=NUM_HEADS,
+ norm_eps=NORM_EPS,
+ adaln_dim=ADALN_DIM,
+ mesh_device=submesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ padding_config=padding_config,
+ )
+ tt_block.load_torch_state_dict(torch_block.state_dict())
+
+ # Sequence-parallel sharding: pad the sequence to k_chunk*sp, shard x and the
+ # cos/sin rope tables on the sequence dim across the SP axis. The hidden dim stays
+ # replicated (TP only shards inside the matmuls), so a single-axis shard suffices.
+ padded_len = _sp_padded_len(seq_len, sp_factor)
+ cos4 = cos.unsqueeze(1) # (B, 1, L, head_dim) to broadcast over heads
+ sin4 = sin.unsqueeze(1)
+ # tt block uses interleaved RoPE (rotary_embedding_llama); permute cos/sin to match.
+ # The torch reference above keeps the original half-split tables (unchanged).
+ cos4, sin4 = rope_halfsplit_to_interleaved(cos4, sin4, HEAD_DIM)
+ # The residual stream is FRACTURED on hidden (TP axis) and sharded on sequence (SP).
+ # Shard x on both axes when tp>1: SP on the sequence dim, TP on the hidden dim.
+ if sp_factor > 1:
+ x = torch.nn.functional.pad(x, (0, 0, 0, padded_len - seq_len))
+ cos4 = torch.nn.functional.pad(cos4, (0, 0, 0, padded_len - seq_len))
+ sin4 = torch.nn.functional.pad(sin4, (0, 0, 0, padded_len - seq_len))
+ if tp_factor > 1:
+ from ....utils.tensor import bf16_tensor_2dshard
+
+ tt_x = bf16_tensor_2dshard(x, device=submesh_device, shard_mapping={sp_axis: 1, tp_axis: 2})
+ else:
+ tt_x = bf16_tensor(x, device=submesh_device, mesh_axis=sp_axis, shard_dim=1)
+ tt_cos = bf16_tensor(cos4, device=submesh_device, mesh_axis=sp_axis, shard_dim=2)
+ tt_sin = bf16_tensor(sin4, device=submesh_device, mesh_axis=sp_axis, shard_dim=2)
+ else:
+ if tp_factor > 1:
+ tt_x = bf16_tensor(x, device=submesh_device, mesh_axis=tp_axis, shard_dim=2)
+ else:
+ tt_x = bf16_tensor(x, device=submesh_device)
+ tt_cos = bf16_tensor(cos4, device=submesh_device)
+ tt_sin = bf16_tensor(sin4, device=submesh_device)
+ tt_adaln = bf16_tensor(adaln_input, device=submesh_device)
+
+ # Block-diagonal segment mask (additive 0 / -inf). num_segments==1 => full attention.
+ if num_segments == 1:
+ tt_attn_mask = None
+ elif sp_factor == 1:
+ # Replicated [B, 1, L, L] mask fed to plain SDPA.
+ assert torch.isneginf(attn_bias).any(), "multi-segment input should produce a non-trivial mask"
+ tt_attn_mask = bf16_tensor(attn_bias, device=submesh_device)
+ else:
+ # SP + mask: the block all-gathers K/V to the full (padded) sequence and runs
+ # masked SDPA with Q kept sequence-sharded, so the mask is [B, 1, padded_L, padded_L]
+ # sharded on the query-row dim. Pad segment ids with a sentinel id so real queries
+ # never attend to padded keys (and vice-versa).
+ seg_padded = torch.full((batch_size, padded_len), num_segments, dtype=torch.long)
+ seg_padded[:, :seq_len] = segment_ids
+ same_seg = seg_padded.unsqueeze(2) == seg_padded.unsqueeze(1) # (B, padded_L, padded_L)
+ bias = torch.zeros(batch_size, 1, padded_len, padded_len, dtype=torch.float32)
+ bias.masked_fill_(~same_seg.unsqueeze(1), float("-inf"))
+ tt_attn_mask = bf16_tensor(bias, device=submesh_device, mesh_axis=sp_axis, shard_dim=2)
+
+ tt_out = tt_block(
+ tt_x,
+ cos=tt_cos,
+ sin=tt_sin,
+ adaln_input=tt_adaln,
+ attn_mask=tt_attn_mask,
+ spatial_sequence_length=seq_len,
+ )
+
+ # Output is FRACTURED on hidden (TP), sharded on SP. mesh_axes maps tensor dim ->
+ # mesh axis: sequence (dim 1) on sp_axis, hidden (dim 2) on tp_axis.
+ out_mesh_axes = [None, sp_axis, tp_axis if tp_factor > 1 else None]
+ tt_out_torch = tensor.to_torch(tt_out, mesh_axes=out_mesh_axes)
+ tt_out_torch = tt_out_torch[:, :seq_len, :]
+
+ logger.info(
+ f"ideogram4 block: B={batch_size} seq={seq_len} tp={tp_factor} sp={sp_factor} "
+ f"padded_heads={tt_block.padded_heads} padded_len={padded_len}"
+ )
+ assert_quality(torch_out, tt_out_torch, pcc=0.99, relative_rmse=0.05)
+
+
+# =============================================================================
+# Full transformer (embeddings + input_proj + block stack + adaLN + final layer,
+# i.e. the top and bottom outside the decoder blocks) vs the OFFICIAL reference
+# Ideogram4Transformer, under TP / SP. 2 layers is enough to prove the wiring;
+# weights="random" is the parallelism sweep, weights="real" the shipped-weight
+# fidelity check. Asserts PCC >= 0.99 on the image-position velocity prediction.
+# =============================================================================
+
+
+def _build_model_inputs(config, batch_size, llm_len, image_len):
+ """[llm tokens | image tokens] packed single-stream batch for the full model."""
+ seq_len = llm_len + image_len
+ llm_features = torch.randn(batch_size, seq_len, config.llm_features_dim)
+ x = torch.randn(batch_size, seq_len, config.in_channels)
+ t = torch.rand(batch_size) # flow-matching time in [0, 1]
+ segment_ids = torch.zeros(batch_size, seq_len, dtype=torch.long) # single segment
+
+ indicator = torch.full((batch_size, seq_len), OUTPUT_IMAGE_INDICATOR, dtype=torch.long)
+ indicator[:, :llm_len] = LLM_TOKEN_INDICATOR
+
+ # position_ids (B, L, 3) = (t, h, w): llm tokens linear on all axes, image on an offset grid.
+ position_ids = torch.zeros(batch_size, seq_len, 3, dtype=torch.long)
+ lpos = torch.arange(llm_len)
+ position_ids[:, :llm_len, 0] = lpos
+ position_ids[:, :llm_len, 1] = lpos
+ position_ids[:, :llm_len, 2] = lpos
+ grid_h = int(round(image_len**0.5))
+ while image_len % grid_h != 0:
+ grid_h -= 1
+ grid_w = image_len // grid_h
+ hh = torch.arange(grid_h).repeat_interleave(grid_w)
+ ww = torch.arange(grid_w).repeat(grid_h)
+ position_ids[:, llm_len:, 0] = IMAGE_POSITION_OFFSET
+ position_ids[:, llm_len:, 1] = IMAGE_POSITION_OFFSET + hh
+ position_ids[:, llm_len:, 2] = IMAGE_POSITION_OFFSET + ww
+
+ return llm_features, x, t, position_ids, segment_ids, indicator
+
+
+def _seq_shard(t: torch.Tensor, seq_dim: int, padded_len: int, sp_factor: int, sp_axis: int, device, *, dtype):
+ """Pad a tensor on its sequence dim and shard it across the SP axis (or replicate)."""
+ if sp_factor > 1:
+ pad = padded_len - t.shape[seq_dim]
+ if pad:
+ spec = [0, 0] * (t.ndim - seq_dim - 1) + [0, pad]
+ t = torch.nn.functional.pad(t, spec)
+ mesh_axes = [None] * t.ndim
+ mesh_axes[seq_dim] = sp_axis
+ layout = ttnn.ROW_MAJOR_LAYOUT if dtype == ttnn.uint32 else ttnn.TILE_LAYOUT
+ return tensor.from_torch(t, device=device, dtype=dtype, layout=layout, mesh_axes=mesh_axes)
+ if dtype == ttnn.uint32:
+ return tensor.from_torch(t, device=device, dtype=dtype, layout=ttnn.ROW_MAJOR_LAYOUT)
+ return bf16_tensor(t, device=device)
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "submesh_shape", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "weights"),
+ [
+ # random-weight parallelism sweep (proves the full-model wiring under TP / SP).
+ pytest.param((1, 1), (1, 1), 0, 1, 1, {}, ttnn.Topology.Linear, "random", id="tp1sp1"),
+ pytest.param((2, 4), (1, 2), 0, 1, 1, line_params, ttnn.Topology.Linear, "random", id="tp2"),
+ pytest.param((2, 4), (1, 4), 0, 1, 1, line_params, ttnn.Topology.Linear, "random", id="tp4_pad20"),
+ pytest.param((2, 4), (2, 1), 0, 1, 1, line_params, ttnn.Topology.Linear, "random", id="sp2"),
+ pytest.param((2, 4), (2, 2), 0, 1, 1, line_params, ttnn.Topology.Linear, "random", id="sp2tp2"),
+ pytest.param((2, 4), (2, 4), 0, 1, 1, line_params, ttnn.Topology.Linear, "random", id="sp2tp4_pad20"),
+ # BH Galaxy 4x8, 2D torus Ring: SP=8 (axis 1), TP=4 (axis 0), 2 links/neighbor.
+ pytest.param((4, 8), (4, 8), 1, 0, 2, ring_params, ttnn.Topology.Ring, "random", id="bh_galaxy_sp8tp4"),
+ # real-weight fidelity: shipped checkpoint on single device and under production TP.
+ pytest.param((1, 1), (1, 1), 0, 1, 1, {}, ttnn.Topology.Linear, "real", id="tp1sp1_real", marks=_NEEDS_WEIGHTS),
+ pytest.param(
+ (2, 4), (1, 4), 0, 1, 1, line_params, ttnn.Topology.Linear, "real", id="tp4_real", marks=_NEEDS_WEIGHTS
+ ),
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize("num_layers", [2], ids=["layers2"])
+@pytest.mark.parametrize(
+ ("batch_size", "llm_len", "image_len"),
+ [pytest.param(1, 64, 256, id="llm64_img256")],
+)
+def test_transformer_model(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ submesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ weights: str,
+ num_layers: int,
+ batch_size: int,
+ llm_len: int,
+ image_len: int,
+) -> None:
+ torch.manual_seed(0)
+ torch_dtype = torch.bfloat16
+ seq_len = llm_len + image_len
+
+ submesh_device = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ sp_factor = tuple(submesh_device.shape)[sp_axis]
+ tp_factor = tuple(submesh_device.shape)[tp_axis]
+
+ config = modeling_ideogram4.Ideogram4Config(num_layers=num_layers)
+ torch_model = modeling_ideogram4.Ideogram4Transformer(config).to(dtype=torch_dtype)
+ if weights == "real":
+ state_dict = _real_transformer_sd(num_layers)
+ _load = torch_model.load_state_dict(state_dict, strict=False)
+ # strict=False only to tolerate non-persistent buffers (e.g. rotary_emb) — those are NOT in
+ # state_dict, so a correct real load has BOTH missing and unexpected empty. Enforce both, so a
+ # broken checkpoint key map fails loudly rather than silently dropping weights to random init.
+ assert not _load.unexpected_keys, f"unexpected checkpoint keys: {_load.unexpected_keys[:5]}"
+ assert not _load.missing_keys, f"missing checkpoint keys: {_load.missing_keys[:5]}"
+ else:
+ state_dict = torch_model.state_dict() # random init
+ torch_model.eval()
+
+ llm_features, x, t, position_ids, segment_ids, indicator = _build_model_inputs(
+ config, batch_size, llm_len, image_len
+ )
+ llm_features = llm_features.to(torch_dtype)
+ x = x.to(torch_dtype)
+
+ with torch.no_grad():
+ torch_out = torch_model(
+ llm_features=llm_features,
+ x=x,
+ t=t,
+ position_ids=position_ids,
+ segment_ids=segment_ids,
+ indicator=indicator,
+ )
+
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+ tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis),
+ sequence_parallel=ParallelFactor(factor=sp_factor, mesh_axis=sp_axis),
+ )
+ ccl_manager = CCLManager(submesh_device, num_links=num_links, topology=topology)
+ padding_config = (
+ PaddingConfig.from_tensor_parallel_factor(config.num_heads, config.emb_dim // config.num_heads, tp_factor)
+ if config.num_heads % tp_factor != 0
+ else None
+ )
+
+ tt_model = Ideogram4Transformer(
+ emb_dim=config.emb_dim,
+ num_layers=config.num_layers,
+ num_heads=config.num_heads,
+ intermediate_size=config.intermediate_size,
+ adaln_dim=config.adanln_dim,
+ in_channels=config.in_channels,
+ llm_features_dim=config.llm_features_dim,
+ norm_eps=config.norm_eps,
+ mesh_device=submesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ padding_config=padding_config,
+ )
+ tt_model.load_torch_state_dict(state_dict)
+
+ cos, sin = torch_model.rotary_emb(position_ids)
+ t_sin = Ideogram4Transformer.sinusoidal_embedding(t, config.emb_dim)
+ llm_token_mask = (indicator == LLM_TOKEN_INDICATOR).to(torch.float32).unsqueeze(-1)
+ output_image_mask = (indicator == OUTPUT_IMAGE_INDICATOR).to(torch.float32).unsqueeze(-1)
+ image_idx = (indicator == OUTPUT_IMAGE_INDICATOR).to(torch.int32)
+ padded_len = _sp_padded_len(seq_len, sp_factor)
+
+ dev = submesh_device
+ _cos_il, _sin_il = rope_halfsplit_to_interleaved(
+ cos.unsqueeze(1), sin.unsqueeze(1), config.emb_dim // config.num_heads
+ )
+ tt_out = tt_model(
+ x=_seq_shard(x, 1, padded_len, sp_factor, sp_axis, dev, dtype=ttnn.bfloat16),
+ llm_features=_seq_shard(llm_features, 1, padded_len, sp_factor, sp_axis, dev, dtype=ttnn.bfloat16),
+ t_sin=bf16_tensor(t_sin.unsqueeze(1), device=dev),
+ cos=_seq_shard(_cos_il, 2, padded_len, sp_factor, sp_axis, dev, dtype=ttnn.bfloat16),
+ sin=_seq_shard(_sin_il, 2, padded_len, sp_factor, sp_axis, dev, dtype=ttnn.bfloat16),
+ image_indicator_index=_seq_shard(image_idx, 1, padded_len, sp_factor, sp_axis, dev, dtype=ttnn.uint32),
+ llm_token_mask=_seq_shard(llm_token_mask, 1, padded_len, sp_factor, sp_axis, dev, dtype=ttnn.bfloat16),
+ output_image_mask=_seq_shard(output_image_mask, 1, padded_len, sp_factor, sp_axis, dev, dtype=ttnn.bfloat16),
+ attn_mask=None,
+ spatial_sequence_length=seq_len,
+ )
+ tt_out_torch = tensor.to_torch(tt_out, mesh_axes=[None, sp_axis, None])[:, :seq_len]
+
+ logger.info(
+ f"ideogram4 model [{weights}]: layers={num_layers} tp={tp_factor} sp={sp_factor} "
+ f"padded_heads={tt_model.layers[0].padded_heads} seq={seq_len} padded_len={padded_len}"
+ )
+ # Only OUTPUT_IMAGE positions carry a meaningful velocity (the reference zeros the
+ # image-latent input at llm positions; their output is unused conditioning noise).
+ image_mask = (indicator == OUTPUT_IMAGE_INDICATOR)[0] # (L,), same across batch
+ assert_quality(torch_out[:, image_mask], tt_out_torch[:, image_mask], pcc=0.99, relative_rmse=0.05)
diff --git a/code/models/tt_dit/tests/models/ideogram4/test_vae_ideogram4.py b/code/models/tt_dit/tests/models/ideogram4/test_vae_ideogram4.py
new file mode 100644
index 0000000000000000000000000000000000000000..85c96c3396f26b646917708ff930433ab873b148
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ideogram4/test_vae_ideogram4.py
@@ -0,0 +1,155 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+# =============================================================================
+# Ideogram 4.0 "Flux2 KL autoencoder" decode path. The reference autoencoder.py
+# ships a diffusers->Flux state-dict converter, so the model IS a diffusers
+# AutoencoderKL (block_out [128,256,512,512], latent_channels 32, mid attention).
+# We reuse vae_sd35.VAEDecoder + a post_quant_conv and verify the decode against a
+# diffusers AutoencoderKL with Ideogram's config. weights="random" is the wiring /
+# channel-parallel (GroupNorm shards channels) check across single device + TP;
+# weights="real" loads the shipped VAE checkpoint for a fidelity check.
+# =============================================================================
+
+import os
+
+import pytest
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from loguru import logger
+from safetensors.torch import load_file
+
+import ttnn
+
+from ....models.vae.vae_ideogram4 import Ideogram4VAEDecoder
+from ....parallel.config import ParallelFactor, VAEParallelConfig
+from ....parallel.manager import CCLManager
+from ....utils.check import assert_quality
+from ....utils.tensor import bf16_tensor
+from ....utils.test import line_params, ring_params
+
+Z_CHANNELS = 32
+FP8 = os.environ.get("IDEOGRAM4_WEIGHTS")
+# Real-weight cases need the gated fp8 checkpoint; skip (don't error) when it isn't configured.
+_NEEDS_WEIGHTS = pytest.mark.skipif(not FP8, reason="IDEOGRAM4_WEIGHTS not set (gated fp8 checkpoint)")
+_L1 = {"l1_small_size": 32768}
+
+
+def _make_akl(weights: str, torch_dtype) -> AutoencoderKL:
+ """diffusers AutoencoderKL with Ideogram's config; random init or the shipped VAE weights."""
+ akl = AutoencoderKL(
+ in_channels=3,
+ out_channels=3,
+ latent_channels=Z_CHANNELS,
+ down_block_types=("DownEncoderBlock2D",) * 4,
+ up_block_types=("UpDecoderBlock2D",) * 4,
+ block_out_channels=(128, 256, 512, 512),
+ layers_per_block=2,
+ norm_num_groups=32,
+ )
+ if weights == "real":
+ vae_sd = load_file(f"{FP8}/vae/diffusion_pytorch_model.safetensors") # diffusers layout, not fp8
+ # bn.* are checkpoint-only batchnorm running stats absent from diffusers AutoencoderKL.
+ incompat = akl.load_state_dict({k: v for k, v in vae_sd.items() if not k.startswith("bn.")}, strict=False)
+ # Prove the shipped weights actually landed; otherwise random init passes PCC without
+ # testing the real checkpoint. (Verified empty on the real load.)
+ assert not incompat.missing_keys and not incompat.unexpected_keys, (
+ f"real VAE load key mismatch: missing={incompat.missing_keys[:5]} "
+ f"unexpected={incompat.unexpected_keys[:5]}"
+ )
+ return akl.to(torch_dtype).eval()
+
+
+@pytest.mark.parametrize(
+ (
+ "mesh_device",
+ "submesh_shape",
+ "tp_axis",
+ "num_links",
+ "device_params",
+ "topology",
+ "latent_h",
+ "latent_w",
+ "weights",
+ ),
+ [
+ # single device (TP=1): small latents, random + real fidelity.
+ pytest.param((1, 1), (1, 1), 1, 1, _L1, ttnn.Topology.Linear, 32, 32, "random", id="tp1_256sq"),
+ pytest.param((1, 1), (1, 1), 1, 1, _L1, ttnn.Topology.Linear, 32, 48, "random", id="tp1_256x384"),
+ pytest.param(
+ (1, 1), (1, 1), 1, 1, _L1, ttnn.Topology.Linear, 32, 32, "real", id="tp1_256sq_real", marks=_NEEDS_WEIGHTS
+ ),
+ # TP=4 on the 2x4 loudbox (channel-parallel GroupNorm), 512px + 1024px.
+ pytest.param(
+ (2, 4), (1, 4), 1, 1, {**line_params, **_L1}, ttnn.Topology.Linear, 64, 64, "random", id="tp4_512px"
+ ),
+ pytest.param(
+ (2, 4), (1, 4), 1, 1, {**line_params, **_L1}, ttnn.Topology.Linear, 128, 128, "random", id="tp4_1024px"
+ ),
+ # SP4xTP2 denoiser: VAE TP=2, replicated on the size-4 axis.
+ pytest.param(
+ (2, 4), (2, 4), 0, 1, {**line_params, **_L1}, ttnn.Topology.Linear, 64, 64, "random", id="tp2_512px"
+ ),
+ pytest.param(
+ (2, 4), (2, 4), 0, 1, {**line_params, **_L1}, ttnn.Topology.Linear, 128, 128, "random", id="tp2_1024px"
+ ),
+ # BH Galaxy 4x8, 2D torus Ring: VAE TP=4 (axis 0), replicated on the size-8 axis, 2 links/neighbor.
+ pytest.param(
+ (4, 8), (4, 8), 0, 2, {**ring_params, **_L1}, ttnn.Topology.Ring, 64, 64, "random", id="bh_galaxy_tp4_512px"
+ ),
+ pytest.param(
+ (4, 8),
+ (4, 8),
+ 0,
+ 2,
+ {**ring_params, **_L1},
+ ttnn.Topology.Ring,
+ 128,
+ 128,
+ "random",
+ id="bh_galaxy_tp4_1024px",
+ ),
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+def test_vae_decoder(
+ *,
+ mesh_device: ttnn.MeshDevice,
+ submesh_shape: tuple[int, int],
+ tp_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ latent_h: int,
+ latent_w: int,
+ weights: str,
+) -> None:
+ """VAE decode, channel-parallel at TP across the mesh (GroupNorm shards channels)."""
+ torch.manual_seed(0)
+ torch_dtype = torch.bfloat16
+ submesh = mesh_device.create_submesh(ttnn.MeshShape(*submesh_shape))
+ tp_factor = tuple(submesh.shape)[tp_axis]
+
+ akl = _make_akl(weights, torch_dtype)
+
+ parallel_config = VAEParallelConfig(tensor_parallel=ParallelFactor(factor=tp_factor, mesh_axis=tp_axis))
+ ccl_manager = CCLManager(submesh, num_links=num_links, topology=topology)
+ tt_decoder = Ideogram4VAEDecoder.from_torch(
+ akl, mesh_device=submesh, parallel_config=parallel_config, ccl_manager=ccl_manager
+ )
+
+ z = torch.randn(1, Z_CHANNELS, latent_h, latent_w, dtype=torch_dtype)
+ with torch.no_grad():
+ ref = akl.decode(z).sample # post_quant_conv + decoder
+
+ # Feed replicated NHWC: post_quant_conv is replicated; the decoder's conv_in
+ # (out_mesh_axis=tp) shards channels from there. Output is replicated -> take device 0.
+ tt_z = bf16_tensor(z.permute(0, 2, 3, 1), device=submesh) # NCHW -> NHWC
+ tt_out = tt_decoder(tt_z)
+ tt_out_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0]).permute(0, 3, 1, 2) # NHWC -> NCHW
+
+ logger.info(
+ f"ideogram4 VAE decode [{weights}] TP={tp_factor}: latent {latent_h}x{latent_w} "
+ f"-> image {latent_h*8}x{latent_w*8}"
+ )
+ assert_quality(ref.float(), tt_out_torch.float(), pcc=0.99, relative_rmse=0.15)
diff --git a/code/models/tt_dit/tests/models/ltx/conftest.py b/code/models/tt_dit/tests/models/ltx/conftest.py
new file mode 100644
index 0000000000000000000000000000000000000000..56a5f3763f0684487f6c900fb25a44d2ee54345a
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/conftest.py
@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: (c) 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""LTX test hooks.
+
+The earliest pytest hook available to a conftest: pin the process to one hardware thread per physical
+core and ``os.execv`` re-exec BEFORE torch/ttnn thread pools and the tt-metal device threads spawn, so
+they inherit the pinned mask (taskset-equivalent placement, fully in-process). On by default so every
+LTX test run gets it (galaxy ring traced replay 6.2 s re-execed vs 6.6-6.7 s without, the same as a
+launch-time ``taskset``); ``LTX_PIN_PREIMPORT=0`` skips the re-exec, ``LTX_PIN_CORES=0`` disables all
+pinning. A no-op on hosts without SMT or whose mask is already one thread per core, so CI runners
+without SMT behave exactly as before.
+
+The re-exec fires before the mesh_device fixture opens the device -- re-execing after a device is open
+risks wedging it.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+
+@pytest.hookimpl(hookwrapper=True, tryfirst=True)
+def pytest_cmdline_main(config): # noqa: ARG001
+ """Runs before collection and before any fixture/device open -- the earliest safe re-exec point."""
+ if os.environ.get("LTX_PIN_PREIMPORT", "1") != "0":
+ # Imported lazily and locally: keeps this hook free of torch/ttnn.
+ from models.tt_dit.utils.host_affinity import reexec_pinned_before_torch
+
+ reexec_pinned_before_torch("LTX tests (pytest_cmdline_main)")
+ yield
diff --git a/code/models/tt_dit/tests/models/ltx/ltx_mesh_params.py b/code/models/tt_dit/tests/models/ltx/ltx_mesh_params.py
new file mode 100644
index 0000000000000000000000000000000000000000..4774cb5fc8e2d16170d8b3b174d87a2eff77c789
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/ltx_mesh_params.py
@@ -0,0 +1,292 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+
+import ttnn
+from models.tt_dit.utils.test import (
+ line_params_req_exact_devices,
+ ring_params_8k_req_exact_devices,
+ ring_params_req_exact_devices,
+)
+
+_line = line_params_req_exact_devices
+_ring = ring_params_req_exact_devices
+
+_1x1sp0tp1nl1_line_is_fsdp0 = pytest.param(
+ (1, 1), 0, 1, 1, {}, ttnn.Topology.Linear, False, id="1x1sp0tp1nl1_line_is_fsdp0"
+)
+_2x2sp0tp1nl2_line_is_fsdp1 = pytest.param(
+ (2, 2), 0, 1, 2, _line, ttnn.Topology.Linear, True, id="2x2sp0tp1nl2_line_is_fsdp1"
+)
+_2x4sp0tp1nl1_line_is_fsdp1 = pytest.param(
+ (2, 4), 0, 1, 1, _line, ttnn.Topology.Linear, True, id="2x4sp0tp1nl1_line_is_fsdp1"
+)
+_2x4sp1tp0nl1_line_is_fsdp1 = pytest.param(
+ (2, 4), 1, 0, 1, _line, ttnn.Topology.Linear, True, id="2x4sp1tp0nl1_line_is_fsdp1"
+)
+# is_fsdp=False mirrors production BH 2x4 (dynamic_load path, not per-layer FSDP gathers)
+_2x4sp1tp0nl2_line_is_fsdp0 = pytest.param(
+ (2, 4), 1, 0, 2, _line, ttnn.Topology.Linear, False, id="2x4sp1tp0nl2_line_is_fsdp0"
+)
+
+# WH has 4 links.
+_4x8sp1tp0nl4_ring_is_fsdp1 = pytest.param(
+ (4, 8), 1, 0, 4, _ring, ttnn.Topology.Ring, True, id="4x8sp1tp0nl4_ring_is_fsdp1"
+)
+_4x8sp1tp0nl2_line_is_fsdp0 = pytest.param(
+ (4, 8), 1, 0, 2, _line, ttnn.Topology.Linear, False, id="4x8sp1tp0nl2_line_is_fsdp0"
+)
+_4x8sp1tp0nl2_ring_is_fsdp0 = pytest.param(
+ (4, 8), 1, 0, 2, _ring, ttnn.Topology.Ring, False, id="4x8sp1tp0nl2_ring_is_fsdp0"
+)
+_4x32sp1tp0nl2_ring_is_fsdp0 = pytest.param(
+ (4, 32), 1, 0, 2, _ring, ttnn.Topology.Ring, False, id="4x32sp1tp0nl2_ring_is_fsdp0"
+)
+
+LTX_PIPELINE_MESH_PARAMS = [
+ _2x2sp0tp1nl2_line_is_fsdp1,
+ _2x4sp0tp1nl1_line_is_fsdp1,
+ _2x4sp1tp0nl2_line_is_fsdp0,
+ _4x8sp1tp0nl4_ring_is_fsdp1,
+ _4x8sp1tp0nl2_line_is_fsdp0,
+ _4x8sp1tp0nl2_ring_is_fsdp0,
+ _4x32sp1tp0nl2_ring_is_fsdp0,
+]
+
+
+def _with_dynamic_load(param, dynamic_load):
+ """Append dynamic_node arg to a 7-field mesh param, yielding the pipeline signature
+ (mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load).
+ Some tests use 7-field mesh_device params, some need an 8th dynamic_load. This function
+ allows us concatenate them together"""
+ v = param.values # (mesh_shape, sp, tp, num_links, device_params, topology, is_fsdp)
+ assert len(v) == 7, f"expected a 7-field mesh param, got {len(v)}: {param.id}"
+ return pytest.param(*v, dynamic_load, marks=param.marks, id=param.id)
+
+
+# dynamic_load=True only for the 2x4 (BH-like) configs — mirrors production, where 2x4 pages
+# weights in/out (dynamic_load path) to avoid init-time DRAM OOM rather than per-layer FSDP gathers.
+_PIPELINE_DL_TRUE = {"2x4sp0tp1nl1_line_is_fsdp1", "2x4sp1tp0nl2_line_is_fsdp0"}
+
+LTX_PIPELINE_MESH_PARAMS_DL = [_with_dynamic_load(p, p.id in _PIPELINE_DL_TRUE) for p in LTX_PIPELINE_MESH_PARAMS]
+
+
+def _override_base_device_params(base, device_params, *, id=None):
+ """Swap the device_params (element 4) of a shared building block, keeping its geometry.
+ Optional id override for variants that need a distinct test id (e.g. the i2v prefix)."""
+ v = base.values # (mesh_shape, sp, tp, num_links, device_params, topology, is_fsdp)
+ return pytest.param(*v[:4], device_params, *v[5:], marks=base.marks, id=id or base.id)
+
+
+# One-stage (Pro) AV needs the L1_SMALL pool for the audio-vocoder conv1d taps. The base pipeline
+# params leave l1_small_size=0, which OOMs the vocoder on the 4x8 galaxy ("L1_SMALL ... bank size is
+# 0 B" in bank_manager). Reserve 32 KB L1_SMALL on the BH 4x8 configs (untraced -> no trace region),
+# matching the distilled AV configs. Other geometries (deselected by the CI -k) pass through unchanged.
+_LTX_ONE_STAGE_L1SMALL_IDS = {"4x8sp1tp0nl2_ring_is_fsdp0", "4x8sp1tp0nl2_line_is_fsdp0"}
+LTX_ONE_STAGE_MESH_PARAMS_DL = [
+ _override_base_device_params(p, {**p.values[4], "l1_small_size": 32768}, id=p.id)
+ if p.id in _LTX_ONE_STAGE_L1SMALL_IDS
+ else p
+ for p in LTX_PIPELINE_MESH_PARAMS_DL
+]
+
+# Two-stages (Pro) drives the same audio vocoder, so it needs the same pool for the same reason.
+LTX_TWO_STAGES_MESH_PARAMS_DL = LTX_ONE_STAGE_MESH_PARAMS_DL
+
+
+# ---------------------------------------------------------------------------
+# Distilled AV pipeline mesh params. Same geometry as the pipeline configs, but the audio
+# decode chain needs bigger device pools on a few configs, and the WH 4x8 ring config runs
+# dynamic_load. Reuse the shared building blocks for geometry (single-sourced); override
+# device_params only where audio decode diverges, then append the distilled dynamic_load.
+# ---------------------------------------------------------------------------
+# Override device_params for the audio decode chain. Each dict augments the base line/ring
+# params only with the pool(s) that config actually needs:
+# l1_small_size: the audio vocoder's native conv1d/conv2d taps (depthwise audio filters) run an
+# UntilizeWithHalo gather whose sharding/config tensors allocate from the dedicated L1_SMALL
+# pool. It defaults to 0, which OOMs the vocoder in decode; 32 KB matches the audio component
+# tests.
+# worker_l1_size: the WH 4x8 ring config's RingAttention otherwise hits a kernel code-size error;
+# the larger worker L1 gives its command stream room.
+# trace_region_size: under LTX_TRACED=1 both stage traces' command streams (stage-1 + the
+# larger-sequence stage-2) live here; measured need is ~236 MB at 1080p (get_trace_buffers_size),
+# so 500 MB leaves headroom.
+_ring_worker_l1 = {"worker_l1_size": 1344544, **_ring}
+_line_trace = {**_line, "trace_region_size": 500_000_000, "l1_small_size": 32768}
+# fabric_router_config (8 KB payload): the strided all-gather packs up to 4 bf16 tiles per fabric packet
+_ring_trace = {**ring_params_8k_req_exact_devices, "trace_region_size": 500_000_000, "l1_small_size": 32768}
+
+LTX_DISTILLED_MESH_PARAMS_DL = [
+ _with_dynamic_load(_2x2sp0tp1nl2_line_is_fsdp1, False),
+ _with_dynamic_load(_2x4sp0tp1nl1_line_is_fsdp1, True),
+ # BH on 2x4 (production loudbox): L1_SMALL scratch for the vocoder conv taps + trace region for the traced
+ # decode -- the same device params the I2V and audio lists already give this id, so the distilled e2e runs
+ # the served (traced) path here too.
+ _with_dynamic_load(_override_base_device_params(_2x4sp1tp0nl2_line_is_fsdp0, _line_trace), True),
+ # WH (ring) on 4x8: bigger worker L1 for RingAttention.
+ _with_dynamic_load(_override_base_device_params(_4x8sp1tp0nl4_ring_is_fsdp1, _ring_worker_l1), True),
+ # BH (linear) on 4x8: L1_SMALL for the vocoder conv taps (without it every tap filter OOMs on the
+ # L1_SMALL pool and falls through to the MAC path, ~13x slower audio) + trace region for the traced
+ # decode -- the same device params the ring row and the I2V/Pro lists already give this id.
+ _with_dynamic_load(_override_base_device_params(_4x8sp1tp0nl2_line_is_fsdp0, _line_trace), False),
+ # BH (ring) on 4x8: trace region + L1_SMALL for the traced decode.
+ _with_dynamic_load(_override_base_device_params(_4x8sp1tp0nl2_ring_is_fsdp0, _ring_trace), False),
+ _with_dynamic_load(_4x32sp1tp0nl2_ring_is_fsdp0, False),
+]
+
+# I2V chained t2v->i2v; both configs are traced (trace region + L1_SMALL). 4x8 Galaxy (ring):
+# full-res 1088x1920 latent shards unevenly on the 4x8 mesh (s1 cond latent 17x30, full 34x60),
+# so the VAE encoder fold + even-shard padding must handle non-mesh-aligned dims here — the 2x4
+# loudbox shards evenly and never hits this, so it's kept as a distinct id.
+LTX_DISTILLED_I2V_MESH_PARAMS_DL = [
+ _with_dynamic_load(_override_base_device_params(_2x4sp1tp0nl2_line_is_fsdp0, _line_trace), True),
+ _with_dynamic_load(
+ _override_base_device_params(_4x8sp1tp0nl2_ring_is_fsdp0, _ring_trace, id="i2v_4x8sp1tp0nl2_ring_is_fsdp0"),
+ False,
+ ),
+]
+
+# Two-stages (Pro) I2V smoke reuses the same geometry + device params as the distilled i2v chain
+# (2x4 loudbox + 4x8 galaxy, traced + L1_SMALL for the audio vocoder), so the Pro-I2V leg can run on
+# the single 4x8 galaxy in addition to the 2x4 loudbox.
+LTX_TWO_STAGES_I2V_MESH_PARAMS_DL = LTX_DISTILLED_I2V_MESH_PARAMS_DL
+
+# Audio-decode-only profiling: both configs use the traced line params (trace region + L1_SMALL)
+# so LTX_TRACED=1 can capture the vocoder.
+LTX_DISTILLED_AUDIO_MESH_PARAMS_DL = [
+ _with_dynamic_load(_override_base_device_params(_2x4sp1tp0nl2_line_is_fsdp0, _line_trace), True),
+ _with_dynamic_load(_override_base_device_params(_4x8sp1tp0nl2_line_is_fsdp0, _line_trace), False),
+]
+
+# No 1x1 config: real-grid shapes require SP padding, and video self-attention only masks
+# padded keys via ring SDPA's logical_n (sp>1). Production never runs sp=1.
+LTX_TRANSFORMER_MESH_PARAMS = [
+ _2x4sp0tp1nl1_line_is_fsdp1,
+ _2x4sp1tp0nl2_line_is_fsdp0,
+ _4x8sp1tp0nl4_ring_is_fsdp1,
+ _4x8sp1tp0nl2_ring_is_fsdp0,
+ _4x8sp1tp0nl2_line_is_fsdp0,
+]
+
+LTX_ATTENTION_MESH_PARAMS = [
+ _1x1sp0tp1nl1_line_is_fsdp0,
+ _2x4sp0tp1nl1_line_is_fsdp1,
+ _2x4sp1tp0nl1_line_is_fsdp1,
+ _4x8sp1tp0nl4_ring_is_fsdp1,
+ _4x8sp1tp0nl2_line_is_fsdp0,
+]
+
+# Upsampler uses VaeHWParallelConfig (h_axis, w_axis) instead of DiT sp/tp axes.
+_2x4h0w1nl1_line = pytest.param(
+ (2, 4), 0, 1, 1, {**_line, "trace_region_size": 23887872}, ttnn.Topology.Linear, id="2x4h0w1nl1_line"
+)
+
+_4x8h0w1nl2_line = pytest.param(
+ (4, 8), 0, 1, 2, {**_line, "trace_region_size": 23887872}, ttnn.Topology.Linear, id="4x8h0w1nl2_line"
+)
+
+LTX_UPSAMPLER_MESH_PARAMS = [
+ _2x4h0w1nl1_line,
+ _4x8h0w1nl2_line,
+]
+
+# VAE decoder mesh params — multi-device configs for full LTXVideoDecoder parity.
+# Format: (mesh_shape, h_axis, w_axis, num_links, device_params, topology)
+_LTX_SKIP_SMALL_MESH_FABRIC = os.environ.get("LTX_VAE_FORCE_SMALL_MESH_FABRIC", "0") != "1"
+
+_line_fabric = {"fabric_config": ttnn.FabricConfig.FABRIC_1D, "require_exact_physical_num_devices": True}
+
+_1x1h0w1nl1_line_vae = pytest.param(
+ (1, 1),
+ 0,
+ 1,
+ 1,
+ _line_fabric,
+ ttnn.Topology.Linear,
+ marks=pytest.mark.skipif(
+ _LTX_SKIP_SMALL_MESH_FABRIC,
+ reason="Known flaky fabric handshake on 1x1 mesh; set LTX_VAE_FORCE_SMALL_MESH_FABRIC=1 to force-run.",
+ ),
+ id="1x1h0w1nl1_line",
+)
+
+_2x4h0w1nl1_line_vae = pytest.param((2, 4), 0, 1, 1, _line_fabric, ttnn.Topology.Linear, id="2x4h0w1nl1_line")
+
+_2x4h1w0nl1_line_vae = pytest.param((2, 4), 1, 0, 1, _line_fabric, ttnn.Topology.Linear, id="2x4h1w0nl1_line")
+
+_1x8h0w1nl1_line_vae = pytest.param((1, 8), 0, 1, 1, _line_fabric, ttnn.Topology.Linear, id="1x8h0w1nl1_line")
+
+_1x4h1w0nl1_line_vae = pytest.param(
+ (1, 4),
+ 1,
+ 0,
+ 1,
+ _line_fabric,
+ ttnn.Topology.Linear,
+ marks=pytest.mark.skipif(
+ _LTX_SKIP_SMALL_MESH_FABRIC,
+ reason="Known flaky fabric handshake on 1x4 mesh; set LTX_VAE_FORCE_SMALL_MESH_FABRIC=1 to force-run.",
+ ),
+ id="1x4h1w0nl1_line",
+)
+
+_4x8h0w1nl2_line_vae = pytest.param((4, 8), 0, 1, 2, _line_fabric, ttnn.Topology.Linear, id="4x8h0w1nl2_line")
+
+LTX_VAE_DECODER_MESH_PARAMS = [
+ _1x1h0w1nl1_line_vae,
+ _2x4h0w1nl1_line_vae,
+ _2x4h1w0nl1_line_vae,
+ _1x8h0w1nl1_line_vae,
+ _1x4h1w0nl1_line_vae,
+ _4x8h0w1nl2_line_vae,
+]
+
+# Multi-device only (2K decode doesn't fit on a single device).
+LTX_VAE_DECODER_MULTI_ONLY_MESH_PARAMS = [
+ _2x4h0w1nl1_line_vae,
+ _2x4h1w0nl1_line_vae,
+ _4x8h0w1nl2_line_vae,
+]
+
+# Encoder production mesh (mirrors the I2V pipeline layout).
+LTX_VAE_ENCODER_PROD_MESH_PARAMS = [
+ _4x8h0w1nl2_line_vae,
+]
+
+# ---------------------------------------------------------------------------
+# Audio mesh params — (mesh_device, mesh_shape, sp_axis, tp_axis, num_links,
+# dynamic_load, device_params, topology, is_fsdp)
+# Audio tests create a submesh from mesh_device using mesh_shape, hence the
+# duplicate shape field.
+# ---------------------------------------------------------------------------
+_line_audio = {**_line, "l1_small_size": 32768}
+_ring_audio = {**_ring, "l1_small_size": 32768}
+_ring_trace_audio = {**_ring_audio, "trace_region_size": 300_000_000}
+
+LTX_AUDIO_MESH_PARAMS_FULL = [
+ pytest.param(
+ (2, 2), (2, 2), 0, 1, 2, False, _line_audio, ttnn.Topology.Linear, True, id="2x2sp0tp1nl2_line_is_fsdp1"
+ ),
+ pytest.param(
+ (2, 4), (2, 4), 0, 1, 1, True, _line_audio, ttnn.Topology.Linear, True, id="2x4sp0tp1nl1_line_is_fsdp1"
+ ),
+ pytest.param(
+ (2, 4), (2, 4), 1, 0, 2, True, _line_audio, ttnn.Topology.Linear, False, id="2x4sp1tp0nl2_line_is_fsdp0"
+ ),
+ pytest.param(
+ (4, 8), (4, 8), 1, 0, 4, False, _ring_audio, ttnn.Topology.Ring, True, id="4x8sp1tp0nl4_ring_is_fsdp1"
+ ),
+ pytest.param(
+ (4, 8), (4, 8), 1, 0, 2, False, _line_audio, ttnn.Topology.Linear, False, id="4x8sp1tp0nl2_line_is_fsdp0"
+ ),
+ pytest.param(
+ (4, 8), (4, 8), 1, 0, 2, False, _ring_trace_audio, ttnn.Topology.Ring, False, id="4x8sp1tp0nl2_ring_is_fsdp0"
+ ),
+ pytest.param(
+ (4, 32), (4, 32), 1, 0, 2, False, _ring_audio, ttnn.Topology.Ring, False, id="4x32sp1tp0nl2_ring_is_fsdp0"
+ ),
+]
diff --git a/code/models/tt_dit/tests/models/ltx/scaleout_configs/single_galaxy/mesh_graph_descriptor.textproto b/code/models/tt_dit/tests/models/ltx/scaleout_configs/single_galaxy/mesh_graph_descriptor.textproto
new file mode 100644
index 0000000000000000000000000000000000000000..3c8df695db0cc55adecbf5a00849b2d86cd66486
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/scaleout_configs/single_galaxy/mesh_graph_descriptor.textproto
@@ -0,0 +1,20 @@
+# Single-rank single-galaxy mesh for the LTX-2.3 multihost CI (bh_sc1) legs.
+# LTX's 4x8sp1tp0nl2_ring config uses FabricConfig FABRIC_1D_RING, which requires the MGD to
+# provide TORUS_XY connectivity (BOTH dims wrap). A FabricConfig can only *restrict* topology
+# (torus->mesh), not add links, so a single-axis ring (TORUS_X, i.e. dim_types [LINE, RING]) fails
+# at mesh_graph init with "requires more connectivity than MGD provides". Both dims are therefore
+# RING, matching the proven tt_metal/fabric single_bh_galaxy_torus_xy descriptor (that one is [8,4];
+# here it is [4,8] because the LTX test requests a 4x8 mesh). One rank owns the whole 32-device
+# galaxy (confirmed on bh_sc1: a single rank sees all 32 chips).
+mesh_descriptors {
+ name: "M0"
+ arch: BLACKHOLE
+ device_topology { dims: [ 4, 8 ] dim_types: [ RING, RING ] }
+ host_topology { dims: [ 1, 1 ] }
+ channels {
+ count: 2
+ policy: RELAXED
+ }
+}
+
+top_level_instance { mesh { mesh_descriptor: "M0" mesh_id: 0 } }
diff --git a/code/models/tt_dit/tests/models/ltx/test_attention_ltx.py b/code/models/tt_dit/tests/models/ltx/test_attention_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..5c0597ae96ee9169f3c6a9496ff4ad82468d62dd
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_attention_ltx.py
@@ -0,0 +1,258 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""LTX-2 attention unit tests; mesh/fabric parametrization mirrors test_wan_attention."""
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.models.transformers.ltx.attention_ltx import LTXAttention
+from models.tt_dit.parallel.config import DiTParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.mochi import get_rot_transformation_mat
+from models.tt_dit.utils.tensor import bf16_tensor, bf16_tensor_2dshard
+from models.tt_dit.utils.test import skip_if_unsupported_num_links
+
+from .ltx_mesh_params import LTX_ATTENTION_MESH_PARAMS
+
+
+def _diffusers_qk_to_split(t: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor:
+ """diffusers (interleaved-rotation) Q/K → Lightricks SPLIT convention the TT loader expects."""
+ D = head_dim
+ D_half = D // 2
+ inv = torch.empty(D, dtype=torch.long)
+ inv[:D_half] = torch.arange(0, D, 2)
+ inv[D_half:] = torch.arange(1, D, 2)
+ rest = t.shape[1:]
+ return t.reshape(num_heads, D, *rest).index_select(1, inv).reshape(num_heads * D, *rest)
+
+
+def _convert_diffusers_attn_state(state: dict, num_heads: int, head_dim: int) -> dict:
+ """diffusers LTX2Attention state_dict → TT LTXAttention loader input (split Q/K convention)."""
+ out = dict(state)
+ for base in ("to_q", "to_k"):
+ for suffix in ("weight", "bias"):
+ k = f"{base}.{suffix}"
+ if k in out:
+ out[k] = _diffusers_qk_to_split(out[k], num_heads, head_dim)
+ for k in ("norm_q.weight", "norm_k.weight"):
+ if k in out:
+ out[k] = _diffusers_qk_to_split(out[k], num_heads, head_dim)
+ return out
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ LTX_ATTENTION_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_self_attention(
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+) -> None:
+ """
+ Test LTX-2 self-attention: compare TT LTXAttention vs PyTorch Attention.
+ """
+ from diffusers.models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoRotaryPosEmbed
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+
+ dim = 4096
+ num_heads = 32
+ head_dim = dim // num_heads
+ B = 1
+ seq_len = 256 # Small for fast test
+
+ # interleaved matches the (B, N, dim) cos/sin layout the device kernel expects
+ torch_model = LTX2Attention(
+ query_dim=dim, heads=num_heads, kv_heads=num_heads, dim_head=head_dim, norm_eps=1e-6, rope_type="interleaved"
+ )
+ torch_model.eval()
+ torch_state = torch_model.state_dict()
+
+ ccl_manager = CCLManager(mesh_device=mesh_device, num_links=num_links, topology=topology)
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+ sequence_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[sp_axis], mesh_axis=sp_axis),
+ tensor_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[tp_axis], mesh_axis=tp_axis),
+ )
+
+ tt_model = LTXAttention(
+ dim=dim,
+ num_heads=num_heads,
+ eps=1e-6,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ is_self=True,
+ )
+ tt_model.load_torch_state_dict(_convert_diffusers_attn_state(torch_state, num_heads, head_dim))
+
+ # seq_len=256 → N_local=32 for 4x8 (sp=8); exactly 1 tile, the minimum
+ torch.manual_seed(42)
+ x = torch.randn(B, seq_len, dim, dtype=torch.float32)
+
+ # Build RoPE (small grid for test) with the diffusers reference, like test_rope_ltx.
+ F, H, W = 4, 8, 8 # 4*8*8=256 = seq_len
+ t_ids = torch.arange(F)
+ h_ids = torch.arange(H)
+ w_ids = torch.arange(W)
+ grid_t, grid_h, grid_w = torch.meshgrid(t_ids, h_ids, w_ids, indexing="ij")
+ indices_grid = torch.stack([grid_t.flatten(), grid_h.flatten(), grid_w.flatten()], dim=0).float().unsqueeze(0)
+
+ rope = LTX2AudioVideoRotaryPosEmbed(
+ dim=dim,
+ base_num_frames=20,
+ base_height=2048,
+ base_width=2048,
+ theta=10000.0,
+ modality="video",
+ double_precision=False,
+ rope_type="interleaved",
+ num_attention_heads=num_heads,
+ )
+ # Interleaved cos/sin: (B, seq_len, dim)
+ cos_freq, sin_freq = rope(indices_grid)
+ # Reshape for the device kernel: (B, seq_len, num_heads, head_dim)
+ cos_apply = cos_freq.reshape(B, seq_len, num_heads, head_dim)
+ sin_apply = sin_freq.reshape(B, seq_len, num_heads, head_dim)
+
+ # diffusers applies interleaved RoPE to q/k on (B, T, dim) before the head split.
+ with torch.no_grad():
+ torch_out = torch_model(x, query_rotary_emb=(cos_freq, sin_freq))
+
+ logger.info(f"PyTorch output shape: {torch_out.shape}")
+
+ # Prepare TT tensors
+ # spatial: (1, B, N, D) fractured on SP/TP
+ spatial = x.unsqueeze(0) # (1, B, N, D)
+ tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3})
+
+ # rotary_embedding_llama expects per-head (B, H, N, head_dim)
+ rope_cos_bhnd = cos_apply.permute(0, 2, 1, 3) # (B, H, N, D)
+ rope_sin_bhnd = sin_apply.permute(0, 2, 1, 3)
+ tt_cos = bf16_tensor_2dshard(rope_cos_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+ tt_sin = bf16_tensor_2dshard(rope_sin_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ # TT forward
+ tt_out = tt_model(
+ spatial_1BND=tt_spatial,
+ N=seq_len,
+ rope_cos=tt_cos,
+ rope_sin=tt_sin,
+ trans_mat=tt_trans_mat,
+ )
+
+ # Gather and compare
+ concat_dims = [None, None]
+ concat_dims[sp_axis] = 2
+ concat_dims[tp_axis] = 3
+ tt_out_torch = ttnn.to_torch(
+ tt_out,
+ mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)),
+ ).squeeze(0)
+
+ # 8-way SP-ring all-gathers on 4x8 loosen the bound; <=8 devices hold the single-op bar.
+ pcc = 0.988 if mesh_device.get_num_devices() > 8 else 0.999
+ assert_quality(torch_out, tt_out_torch, pcc=pcc)
+ logger.info("PASSED: LTX self-attention matches PyTorch reference")
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ LTX_ATTENTION_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_cross_attention(
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+) -> None:
+ """
+ Test LTX-2 cross-attention: compare TT LTXAttention vs PyTorch Attention.
+ """
+ from diffusers.models.transformers.transformer_ltx2 import LTX2Attention
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+
+ dim = 4096
+ context_dim = 4096
+ num_heads = 32
+ head_dim = dim // num_heads
+ B = 1
+ seq_len = 256
+ prompt_len = 32
+
+ # diffusers reference cross-attention (cross_attention_dim set; no RoPE on text K/V)
+ torch_model = LTX2Attention(
+ query_dim=dim,
+ cross_attention_dim=context_dim,
+ heads=num_heads,
+ kv_heads=num_heads,
+ dim_head=head_dim,
+ norm_eps=1e-6,
+ )
+ torch_model.eval()
+ torch_state = torch_model.state_dict()
+
+ ccl_manager = CCLManager(mesh_device=mesh_device, num_links=num_links, topology=topology)
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+ sequence_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[sp_axis], mesh_axis=sp_axis),
+ tensor_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[tp_axis], mesh_axis=tp_axis),
+ )
+
+ tt_model = LTXAttention(
+ dim=dim,
+ num_heads=num_heads,
+ eps=1e-6,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ is_self=False,
+ context_dim=context_dim,
+ )
+ tt_model.load_torch_state_dict(torch_state)
+
+ torch.manual_seed(42)
+ x = torch.randn(B, seq_len, dim, dtype=torch.float32)
+ context = torch.randn(B, prompt_len, context_dim, dtype=torch.float32)
+
+ with torch.no_grad():
+ torch_out = torch_model(x, encoder_hidden_states=context)
+
+ logger.info(f"PyTorch cross-attn output shape: {torch_out.shape}")
+
+ spatial = x.unsqueeze(0)
+ prompt = context.unsqueeze(0)
+ tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3})
+ tt_prompt = bf16_tensor(prompt, device=mesh_device)
+
+ tt_out = tt_model(spatial_1BND=tt_spatial, N=seq_len, prompt_1BLP=tt_prompt)
+
+ concat_dims = [None, None]
+ concat_dims[sp_axis] = 2
+ concat_dims[tp_axis] = 3
+ tt_out_torch = ttnn.to_torch(
+ tt_out,
+ mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)),
+ ).squeeze(0)
+
+ # 8-way SP-ring all-gathers on 4x8 loosen the bound; <=8 devices hold the single-op bar.
+ pcc = 0.988 if mesh_device.get_num_devices() > 8 else 0.999
+ assert_quality(torch_out, tt_out_torch, pcc=pcc)
+ logger.info("PASSED: LTX cross-attention matches PyTorch reference")
diff --git a/code/models/tt_dit/tests/models/ltx/test_audio_ltx.py b/code/models/tt_dit/tests/models/ltx/test_audio_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..efa66a87662b4f384277f2e852d5402d05e8a639
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_audio_ltx.py
@@ -0,0 +1,1207 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Audio tests for LTX-2: per-component stage checks + post-denoise decode e2e.
+
+Component stage tests (random weights vs diffusers LTX-2 torch reference):
+- Stage A: ``audio_decoder`` (mel-VAE)
+- Stage B: ``vocoder`` (24 kHz waveform)
+- Stage C: ``vocoder_with_bwe`` (48 kHz BWE path)
+
+End-to-end decode (mel-VAE → vocoder → BWE, requires real checkpoints):
+- on-device decode runs on representative meshes
+- warm-path decode stays functional after cold compile/build
+- on-device decode quality remains close to host reference (PSNR)
+
+Denoise-path audio regressions live in sibling files.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import time
+from glob import glob
+
+import einops
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.models.audio_vae.audio_decoder_ltx import LTXAudioDecoderAdapter
+from models.tt_dit.models.audio_vae.bwe_ltx import MelSTFT, VocoderWithBWE
+from models.tt_dit.models.audio_vae.mel_decoder_ltx import AudioUpsample, MelDecoder, ResnetBlock
+from models.tt_dit.models.audio_vae.vocoder_ltx import AMPBlock1, Vocoder
+from models.tt_dit.parallel.config import AudioTCParallelConfig, DiTParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.pipelines.ltx.pipeline_ltx import LTXPipeline
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.test import skip_if_unsupported_num_links
+from models.tt_dit.utils.video import Audio
+
+from .ltx_mesh_params import LTX_AUDIO_MESH_PARAMS_FULL
+
+_WARM_ITERS = 3
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint discovery
+# ---------------------------------------------------------------------------
+# Default checkpoint filenames used by the e2e/perf decode tests. Any of these can
+# be overridden at runtime with LTX_CHECKPOINT, which is honored regardless of its
+# basename (see _resolve_checkpoint). The tuples below only drive *fallback*
+# discovery when LTX_CHECKPOINT is unset.
+_DEV_CHECKPOINT = "ltx-2.3-22b-dev.safetensors"
+_DISTILLED_CHECKPOINT = "ltx-2.3-22b-distilled-1.1.safetensors"
+_LOCAL_CHECKPOINT_DIR = "~/.cache/ltx-checkpoints"
+# HF cache snapshot glob for the LTX-2.3 repo; {name} is the checkpoint filename.
+_HF_SNAPSHOT_GLOB = "~/.cache/huggingface/hub/models--Lightricks--LTX-2.3/snapshots/*/{name}"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def _resolve_checkpoint(candidates: list[str], expected_filenames: tuple[str, ...]) -> str | None:
+ # An explicit LTX_CHECKPOINT override wins unconditionally as long as the file
+ # exists — that's the whole point of the env var. The expected_filenames tuple
+ # only gates the fallback discovery below, never the override.
+ explicit = os.environ.get("LTX_CHECKPOINT")
+ if explicit and os.path.exists(explicit):
+ return explicit
+ for path in candidates:
+ expanded = os.path.expanduser(path)
+ if os.path.exists(expanded):
+ return expanded
+ # Reuse default HF cache location if present; avoid network dependency.
+ for name in expected_filenames:
+ hits = glob(os.path.expanduser(_HF_SNAPSHOT_GLOB.format(name=name)))
+ if hits:
+ return hits[0]
+ return None
+
+
+def _build_pipeline(
+ mesh_device: ttnn.MeshDevice,
+ *,
+ sp_axis: int,
+ tp_axis: int,
+ checkpoint: str,
+ num_links: int,
+ topology: ttnn.Topology,
+) -> tuple[LTXPipeline, DiTParallelConfig]:
+ mesh_shape = tuple(mesh_device.shape)
+ parallel_config = DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+ sequence_parallel=ParallelFactor(factor=mesh_shape[sp_axis], mesh_axis=sp_axis),
+ tensor_parallel=ParallelFactor(factor=mesh_shape[tp_axis], mesh_axis=tp_axis),
+ )
+ ccl_manager = CCLManager(mesh_device, num_links=num_links, topology=topology)
+ pipeline = LTXPipeline(
+ mesh_device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ checkpoint_name=None,
+ )
+ pipeline.checkpoint_name = checkpoint
+ pipeline._audio_adapter = LTXAudioDecoderAdapter(
+ pipeline.checkpoint_name,
+ mesh_device=pipeline.mesh_device,
+ vae_ccl_manager=pipeline.vae_ccl_manager,
+ dit_parallel_config=pipeline.parallel_config,
+ traced=pipeline._traced,
+ )
+ return pipeline, parallel_config
+
+
+def _psnr(ref: torch.Tensor, test: torch.Tensor) -> float:
+ ref = ref.float()
+ test = test.float()
+ mse = torch.mean((ref - test) ** 2).item()
+ if mse == 0.0:
+ return float("inf")
+ peak = ref.abs().max().item()
+ if peak == 0.0:
+ return float("inf")
+ return 20.0 * math.log10(peak) - 10.0 * math.log10(mse)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+ (
+ "mesh_device",
+ "mesh_shape",
+ "sp_axis",
+ "tp_axis",
+ "num_links",
+ "dynamic_load",
+ "device_params",
+ "topology",
+ "is_fsdp",
+ ),
+ LTX_AUDIO_MESH_PARAMS_FULL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_audio_decode_perf(
+ mesh_device: ttnn.MeshDevice,
+ mesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ dynamic_load: bool,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+):
+ ckpt = _resolve_checkpoint(
+ [f"{_LOCAL_CHECKPOINT_DIR}/{_DEV_CHECKPOINT}"],
+ (_DEV_CHECKPOINT,),
+ )
+ if ckpt is None:
+ pytest.skip("checkpoint not found (set LTX_CHECKPOINT)")
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ _ = (dynamic_load, is_fsdp) # parity with fast_av config schema
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ torch.manual_seed(0)
+ pipeline, parallel_config = _build_pipeline(
+ mesh_device,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ checkpoint=ckpt,
+ num_links=num_links,
+ topology=topology,
+ )
+ mesh_shape = tuple(mesh_device.shape)
+
+ audio_latent = torch.randn(1, 256, 128, dtype=torch.float32)
+ num_frames, fps = 145, 24.0
+
+ t0 = time.time()
+ dev_audio = pipeline.decode_audio(audio_latent, num_frames, fps=fps)
+ t_cold = time.time() - t0
+ assert dev_audio is not None
+
+ warm_times = []
+ for _ in range(_WARM_ITERS):
+ t0 = time.time()
+ out = pipeline.decode_audio(audio_latent, num_frames, fps=fps)
+ warm_times.append(time.time() - t0)
+ assert out is not None
+
+ logger.info(
+ f"Audio decode perf (mesh {mesh_shape}, T-shard={parallel_config.sequence_parallel.factor}): "
+ f"cold={t_cold:.2f}s warm_min={min(warm_times):.2f}s warm_mean={sum(warm_times)/len(warm_times):.2f}s "
+ f"wave={tuple(dev_audio.waveform.shape)} sr={dev_audio.sampling_rate}"
+ )
+
+
+@pytest.mark.parametrize(
+ (
+ "mesh_device",
+ "mesh_shape",
+ "sp_axis",
+ "tp_axis",
+ "num_links",
+ "dynamic_load",
+ "device_params",
+ "topology",
+ "is_fsdp",
+ ),
+ LTX_AUDIO_MESH_PARAMS_FULL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_audio_decode_e2e_psnr(
+ mesh_device: ttnn.MeshDevice,
+ mesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ dynamic_load: bool,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+):
+ ckpt = _resolve_checkpoint(
+ [f"{_LOCAL_CHECKPOINT_DIR}/{_DISTILLED_CHECKPOINT}"],
+ (_DISTILLED_CHECKPOINT,),
+ )
+ if ckpt is None:
+ pytest.skip("distilled checkpoint not found (set LTX_CHECKPOINT)")
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ _ = (dynamic_load, is_fsdp) # parity with fast_av config schema
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ torch.manual_seed(0)
+ mesh_shape = tuple(mesh_device.shape)
+ pipeline, _ = _build_pipeline(
+ mesh_device,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ checkpoint=ckpt,
+ num_links=num_links,
+ topology=topology,
+ )
+
+ audio_latent = torch.randn(1, 256, 128, dtype=torch.float32)
+ num_frames, fps = 145, 24.0
+
+ ref_audio = _decode_audio_reference(ckpt, audio_latent, num_frames, fps=fps)
+ assert ref_audio is not None, "host reference audio decode returned None"
+
+ dev_audio = pipeline.decode_audio(audio_latent, num_frames, fps=fps)
+ assert dev_audio is not None, "on-device audio decode returned None"
+ assert dev_audio.sampling_rate == ref_audio.sampling_rate
+
+ n = min(ref_audio.waveform.shape[-1], dev_audio.waveform.shape[-1])
+ ref_w = ref_audio.waveform[..., :n].float()
+ dev_w = dev_audio.waveform[..., :n].float()
+ psnr = _psnr(ref_w, dev_w)
+
+ logger.info(
+ f"Audio decode PSNR (mesh {mesh_shape}): {psnr:.2f} dB "
+ f"(ref {tuple(ref_audio.waveform.shape)} -> dev {tuple(dev_audio.waveform.shape)})"
+ )
+ assert psnr >= 28.0, f"audio decode PSNR {psnr:.2f} dB < 28 dB"
+
+
+# ===========================================================================
+# Component stage tests (pipeline_ltx.py audio decode chain)
+#
+# Torch references use diffusers LTX-2 audio modules (same as transformer/rope tests).
+# Stage/component IDs match ``pipeline_ltx.py`` audio decode chain:
+# - Stage A: ``audio_decoder`` (mel-VAE)
+# - Stage B: ``vocoder`` (24 kHz waveform)
+# - Stage C: ``vocoder_with_bwe`` (48 kHz BWE path)
+# ===========================================================================
+_AUDIO_DECODER_CFG = dict(
+ ch=128,
+ out_ch=2,
+ ch_mult=(1, 2, 4),
+ num_res_blocks=2,
+ attn_resolutions=(),
+ mid_block_add_attention=False,
+ z_channels=8,
+ resolution=64,
+ mel_bins=64,
+ sample_rate=16000,
+ mel_hop_length=160,
+ is_causal=True,
+)
+
+_MAIN_VOCODER_CFG = dict(
+ resblock_kernel_sizes=[3, 7, 11],
+ upsample_rates=[5, 2, 2, 2, 2, 2],
+ upsample_kernel_sizes=[11, 4, 4, 4, 4, 4],
+ resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ upsample_initial_channel=1536,
+ resblock="AMP1",
+ output_sampling_rate=24000,
+ activation="snakebeta",
+ use_tanh_at_final=False,
+ apply_final_activation=True,
+ use_bias_at_final=False,
+)
+
+_BWE_VOCODER_CFG = dict(
+ resblock_kernel_sizes=[3, 7, 11],
+ upsample_rates=[6, 5, 2, 2, 2],
+ upsample_kernel_sizes=[12, 11, 4, 4, 4],
+ resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ upsample_initial_channel=512,
+ resblock="AMP1",
+ output_sampling_rate=48000,
+ activation="snakebeta",
+ use_tanh_at_final=False,
+ apply_final_activation=False,
+ use_bias_at_final=False,
+)
+
+_MEL_STFT_CFG = dict(
+ filter_length=512,
+ hop_length=80,
+ win_length=512,
+ n_mel_channels=64,
+)
+
+_INPUT_SR = 16000
+_OUTPUT_SR = 48000
+_HOP_LENGTH = 80
+
+
+# ---------------------------------------------------------------------------
+# Component helpers
+# ---------------------------------------------------------------------------
+def _vocoder_mel(t_frames: int = 120) -> torch.Tensor:
+ """Stereo mel input. t_frames=120 with T-shard factor 4 leaves t_pad=8 (128-aligned),
+ so the tile-align pad fits inside the last shard and the sharded path's per-op tail
+ materialization is exercised — the boundary the click came from."""
+ return torch.randn(1, 2, t_frames, 64, dtype=torch.float32) * 0.5
+
+
+def _assert_sharded_matches_unsharded(unsharded: torch.Tensor, sharded: torch.Tensor, *, name: str) -> None:
+ """The T-sharded vocoder must reproduce the unsharded device path up to the bf16 storage
+ floor. A t_pad boundary regression (the audible click) shows up as a localized tail spike
+ of order 1e-2 — far above the ~1e-4 floor — so a tight max|Δ| catches it where aggregate
+ PCC against the torch reference (with its own fp32-kernel floor) does not."""
+ max_abs = (unsharded - sharded).abs().max().item()
+ logger.info(f"{name}: sharded vs unsharded device max|Δ| = {max_abs:.3e}")
+ assert max_abs < 5e-3, f"{name}: sharded vs unsharded device max|Δ| {max_abs:.3e} — t_pad boundary drift"
+
+
+def _require_diffusers():
+ pytest.importorskip("diffusers")
+ from diffusers.models.autoencoders.autoencoder_kl_ltx2_audio import AutoencoderKLLTX2Audio
+ from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder, LTX2VocoderWithBWE
+
+ return AutoencoderKLLTX2Audio, LTX2Vocoder, LTX2VocoderWithBWE
+
+
+def _diffusers_vocoder_state_to_tt(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Invert diffusers LTX-2.3 vocoder key names to the TT/Lightricks layout."""
+ out: dict[str, torch.Tensor] = {}
+ for key, value in state_dict.items():
+ new_key = (
+ key.replace("conv_in.", "conv_pre.")
+ .replace("conv_out.", "conv_post.")
+ .replace("act_out.", "act_post.")
+ .replace("upsamplers.", "ups.")
+ .replace("resnets.", "resblocks.")
+ )
+ # TT derives downsample filter taps at init; these are not loadable Parameters.
+ if new_key.endswith("downsample.filter"):
+ continue
+ out[new_key] = value
+ return out
+
+
+def _diffusers_vocoder_with_bwe_state_to_tt(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ out: dict[str, torch.Tensor] = {}
+ for prefix in ("vocoder.", "bwe_generator."):
+ sub = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)}
+ for k, v in _diffusers_vocoder_state_to_tt(sub).items():
+ out[f"{prefix}{k}"] = v
+ for key, value in state_dict.items():
+ if key.startswith("mel_stft."):
+ out[key] = value
+ return out
+
+
+def _audio_decoder_state_from_diffusers(
+ vae,
+ *,
+ stats_std: torch.Tensor,
+ stats_mean: torch.Tensor,
+) -> dict[str, torch.Tensor]:
+ state = {k[len("decoder.") :]: v for k, v in vae.state_dict().items() if k.startswith("decoder.")}
+ state["per_channel_statistics.std-of-means"] = stats_std
+ state["per_channel_statistics.mean-of-means"] = stats_mean
+ return state
+
+
+def _decode_latent_ref(
+ vae,
+ latent: torch.Tensor,
+ *,
+ stats_std: torch.Tensor,
+ stats_mean: torch.Tensor,
+) -> torch.Tensor:
+ """Match the Lightricks MelDecoder / TT denormalize + decoder forward."""
+ b, channels, _t, mel_bins = latent.shape
+ patched = einops.rearrange(latent, "b c t f -> b t (c f)")
+ std = stats_std.to(dtype=patched.dtype, device=patched.device)
+ mean = stats_mean.to(dtype=patched.dtype, device=patched.device)
+ last_dim = patched.shape[-1]
+ if std.shape[0] != last_dim:
+ if std.shape[0] == channels and last_dim == channels * mel_bins:
+ std = std.repeat_interleave(mel_bins)
+ mean = mean.repeat_interleave(mel_bins)
+ else:
+ raise ValueError(f"per-channel stats {tuple(std.shape)} do not broadcast to last dim {last_dim}")
+ sample = einops.rearrange(patched * std + mean, "b t (c f) -> b c t f", c=channels, f=mel_bins)
+ return vae.decoder(sample)
+
+
+def _compute_mel_ref(vocoder_with_bwe, audio_bct: torch.Tensor) -> torch.Tensor:
+ batch, n_channels, _ = audio_bct.shape
+ flat = audio_bct.reshape(batch * n_channels, -1).unsqueeze(1)
+ log_mel, _, _, _ = vocoder_with_bwe.mel_stft(flat)
+ return log_mel.reshape(batch, n_channels, log_mel.shape[1], log_mel.shape[2])
+
+
+def _tt_vocoder_cfg(cfg: dict) -> dict:
+ # Torch reference config carries keys that TT constructor does not accept.
+ return {k: v for k, v in cfg.items() if k != "output_sampling_rate"}
+
+
+def _randomize_snake_alphas(module: torch.nn.Module) -> None:
+ with torch.no_grad():
+ for submodule in module.modules():
+ if hasattr(submodule, "alpha") and isinstance(submodule.alpha, torch.nn.Parameter):
+ submodule.alpha.data = torch.randn_like(submodule.alpha.data) * 0.1
+ if hasattr(submodule, "beta") and isinstance(submodule.beta, torch.nn.Parameter):
+ submodule.beta.data = torch.randn_like(submodule.beta.data) * 0.1
+
+
+def _build_torch_stage_b(seed: int = 42):
+ _, LTX2Vocoder, _ = _require_diffusers()
+
+ torch.manual_seed(seed)
+ model = LTX2Vocoder(
+ in_channels=128,
+ hidden_channels=_MAIN_VOCODER_CFG["upsample_initial_channel"],
+ out_channels=2,
+ upsample_kernel_sizes=_MAIN_VOCODER_CFG["upsample_kernel_sizes"],
+ upsample_factors=_MAIN_VOCODER_CFG["upsample_rates"],
+ resnet_kernel_sizes=_MAIN_VOCODER_CFG["resblock_kernel_sizes"],
+ resnet_dilations=_MAIN_VOCODER_CFG["resblock_dilation_sizes"],
+ act_fn=_MAIN_VOCODER_CFG["activation"],
+ antialias=True,
+ final_act_fn=None,
+ final_bias=_MAIN_VOCODER_CFG["use_bias_at_final"],
+ output_sampling_rate=_MAIN_VOCODER_CFG["output_sampling_rate"],
+ ).eval()
+ _randomize_snake_alphas(model)
+ return model
+
+
+def _audio_parallel_config(mesh_shape: tuple[int, int]) -> AudioTCParallelConfig | ParallelFactor | None:
+ t_axis = 0 if mesh_shape[0] >= mesh_shape[1] else 1
+ t_factor = mesh_shape[t_axis]
+ c_axis = 1 - t_axis
+ c_factor = mesh_shape[c_axis]
+ if t_factor > 1 and c_factor > 1:
+ return AudioTCParallelConfig(
+ time_parallel=ParallelFactor(factor=t_factor, mesh_axis=t_axis),
+ channel_parallel=ParallelFactor(factor=c_factor, mesh_axis=c_axis),
+ )
+ if t_factor > 1:
+ return ParallelFactor(factor=t_factor, mesh_axis=t_axis)
+ return None
+
+
+def _build_tt_stage_b(mesh_device: ttnn.MeshDevice, *, parallel_config, ccl_manager: CCLManager | None) -> Vocoder:
+ return Vocoder(
+ mesh_device=mesh_device,
+ dtype=ttnn.float32,
+ in_channels=128,
+ out_channels=2,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ **_tt_vocoder_cfg(_MAIN_VOCODER_CFG),
+ )
+
+
+def _build_torch_stage_c(seed: int = 42):
+ _, LTX2Vocoder, LTX2VocoderWithBWE = _require_diffusers()
+
+ torch.manual_seed(seed)
+ full = LTX2VocoderWithBWE(
+ in_channels=128,
+ hidden_channels=_MAIN_VOCODER_CFG["upsample_initial_channel"],
+ out_channels=2,
+ upsample_kernel_sizes=_MAIN_VOCODER_CFG["upsample_kernel_sizes"],
+ upsample_factors=_MAIN_VOCODER_CFG["upsample_rates"],
+ resnet_kernel_sizes=_MAIN_VOCODER_CFG["resblock_kernel_sizes"],
+ resnet_dilations=_MAIN_VOCODER_CFG["resblock_dilation_sizes"],
+ act_fn=_MAIN_VOCODER_CFG["activation"],
+ antialias=True,
+ final_act_fn=None,
+ final_bias=_MAIN_VOCODER_CFG["use_bias_at_final"],
+ bwe_in_channels=128,
+ bwe_hidden_channels=_BWE_VOCODER_CFG["upsample_initial_channel"],
+ bwe_out_channels=2,
+ bwe_upsample_kernel_sizes=_BWE_VOCODER_CFG["upsample_kernel_sizes"],
+ bwe_upsample_factors=_BWE_VOCODER_CFG["upsample_rates"],
+ bwe_resnet_kernel_sizes=_BWE_VOCODER_CFG["resblock_kernel_sizes"],
+ bwe_resnet_dilations=_BWE_VOCODER_CFG["resblock_dilation_sizes"],
+ bwe_act_fn=_BWE_VOCODER_CFG["activation"],
+ bwe_antialias=True,
+ bwe_final_act_fn=None,
+ bwe_final_bias=_BWE_VOCODER_CFG["use_bias_at_final"],
+ filter_length=_MEL_STFT_CFG["filter_length"],
+ hop_length=_MEL_STFT_CFG["hop_length"],
+ window_length=_MEL_STFT_CFG["win_length"],
+ num_mel_channels=_MEL_STFT_CFG["n_mel_channels"],
+ input_sampling_rate=_INPUT_SR,
+ output_sampling_rate=_OUTPUT_SR,
+ )
+ n_freqs = _MEL_STFT_CFG["filter_length"] // 2 + 1
+ with torch.no_grad():
+ full.mel_stft.stft_fn.forward_basis.copy_(torch.randn(n_freqs * 2, 1, _MEL_STFT_CFG["win_length"]) * 0.05)
+ full.mel_stft.stft_fn.inverse_basis.copy_(torch.randn(n_freqs * 2, 1, _MEL_STFT_CFG["win_length"]) * 0.05)
+ full.mel_stft.mel_basis.copy_(torch.randn(_MEL_STFT_CFG["n_mel_channels"], n_freqs).abs() * 0.1)
+ _randomize_snake_alphas(full)
+ return full.eval()
+
+
+def _build_torch_stage_c_real(checkpoint_name: str):
+ """diffusers ``LTX2VocoderWithBWE`` with the *real* checkpoint weights — the torch oracle
+ for the vocoder+BWE. The checkpoint uses the original BigVGAN key names; invert the
+ diffusers→TT rename (see ``_diffusers_vocoder_state_to_tt``) to load them."""
+ from safetensors import safe_open
+
+ full = _build_torch_stage_c(seed=0) # correct architecture; weights overwritten below
+
+ def ckpt_to_diffusers(s: str) -> str:
+ return (
+ s.replace("conv_pre.", "conv_in.")
+ .replace("conv_post.", "conv_out.")
+ .replace("act_post.", "act_out.")
+ .replace("ups.", "upsamplers.")
+ .replace("resblocks.", "resnets.")
+ .replace("downsample.lowpass.filter", "downsample.filter")
+ )
+
+ state: dict[str, torch.Tensor] = {}
+ with safe_open(checkpoint_name, framework="pt") as f:
+ for k in f.keys():
+ if k.startswith("vocoder."):
+ state[ckpt_to_diffusers(k[len("vocoder.") :])] = f.get_tensor(k).float()
+ res = full.load_state_dict(state, strict=False)
+ assert (
+ not res.missing_keys and not res.unexpected_keys
+ ), f"torch oracle weight load mismatch: {len(res.missing_keys)} missing, {len(res.unexpected_keys)} unexpected"
+ return full.eval()
+
+
+def _build_torch_audio_decoder_real(checkpoint_name: str):
+ """diffusers ``AutoencoderKLLTX2Audio`` (mel-VAE decoder) with the *real* checkpoint weights —
+ the host oracle for the mel-VAE stage. Architecture dims are read from the checkpoint's
+ embedded config, the same source ``LTXAudioDecoderAdapter`` uses, so this tracks the
+ real audio VAE rather than the toy component-test config.
+
+ Returns ``(vae, stats_std, stats_mean, z_channels)`` ready for ``_decode_latent_ref``.
+ """
+ import json
+
+ from safetensors import safe_open
+
+ AutoencoderKLLTX2Audio, _, _ = _require_diffusers()
+
+ state: dict[str, torch.Tensor] = {}
+ stats_std = stats_mean = None
+ with safe_open(checkpoint_name, framework="pt") as f:
+ config = json.loads(f.metadata()["config"])
+ ad = config["audio_vae"]["model"]["params"]
+ ddconfig = ad["ddconfig"]
+ stft_cfg = config["audio_vae"].get("preprocessing", {}).get("stft", {})
+ mel_cfg = config["audio_vae"].get("preprocessing", {}).get("mel", {})
+ mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels")
+ for k in f.keys():
+ if k.startswith("audio_vae.decoder."):
+ state[k[len("audio_vae.") :]] = f.get_tensor(k).float() # -> "decoder.*"
+ elif k == "audio_vae.per_channel_statistics.std-of-means":
+ stats_std = f.get_tensor(k).float()
+ elif k == "audio_vae.per_channel_statistics.mean-of-means":
+ stats_mean = f.get_tensor(k).float()
+
+ z_channels = ddconfig.get("z_channels", 8)
+ vae = AutoencoderKLLTX2Audio(
+ base_channels=ddconfig.get("ch", 128),
+ output_channels=ddconfig.get("out_ch", 2),
+ ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
+ num_res_blocks=ddconfig.get("num_res_blocks", 2),
+ attn_resolutions=tuple(ddconfig.get("attn_resolutions", ())),
+ resolution=ddconfig.get("resolution", 256),
+ latent_channels=z_channels,
+ norm_type="pixel",
+ causality_axis="height",
+ mid_block_add_attention=ddconfig.get("mid_block_add_attention", False),
+ sample_rate=ad.get("sampling_rate", 16000),
+ mel_hop_length=stft_cfg.get("hop_length", 160),
+ is_causal=stft_cfg.get("causal", True),
+ mel_bins=mel_bins,
+ ).eval()
+
+ # Decoder-only load: the checkpoint carries no encoder for this path, so encoder keys are
+ # legitimately missing — only decoder keys must be satisfied, with nothing left unexpected.
+ res = vae.load_state_dict(state, strict=False)
+ missing_decoder = [k for k in res.missing_keys if k.startswith("decoder.")]
+ assert not missing_decoder and not res.unexpected_keys, (
+ f"audio decoder oracle weight load mismatch: {len(missing_decoder)} decoder keys missing, "
+ f"{len(res.unexpected_keys)} unexpected"
+ )
+
+ if stats_std is None or stats_mean is None:
+ # Checkpoint without denormalize stats → identity (matches the component-test default).
+ stats_std = torch.ones(z_channels * mel_bins)
+ stats_mean = torch.zeros(z_channels * mel_bins)
+ return vae, stats_std, stats_mean, z_channels
+
+
+def _decode_audio_reference(checkpoint_name: str, audio_latent: torch.Tensor, num_frames: int, fps: float) -> Audio:
+ """Host (torch) reference for the full audio decode chain — mel-VAE → vocoder+BWE — using the
+ real checkpoint weights. Mirrors ``LTXPipeline.decode_audio`` on CPU so the e2e PSNR test can
+ compare the on-device decode against a torch oracle instead of skipping when the pipeline
+ revision doesn't expose its own reference decode.
+
+ Returns an ``Audio(waveform=(C, T), sampling_rate=...)`` matching the device output layout.
+ """
+ vae, stats_std, stats_mean, z = _build_torch_audio_decoder_real(checkpoint_name)
+ vocoder_with_bwe = _build_torch_stage_c_real(checkpoint_name)
+
+ # Unpatchify exactly as decode_audio does: (1, audio_N, 128) -> (1, z, audio_N, 128 // z).
+ audio_N = audio_latent.shape[1]
+ audio_spatial = audio_latent.reshape(1, audio_N, z, audio_latent.shape[2] // z).permute(0, 2, 1, 3).float()
+
+ with torch.no_grad():
+ mel = _decode_latent_ref(vae, audio_spatial, stats_std=stats_std, stats_mean=stats_mean)
+ waveform = vocoder_with_bwe(mel).squeeze(0).float()
+ sampling_rate = _OUTPUT_SR
+
+ # Trim to the clip duration, matching decode_audio.
+ target_samples = int((num_frames / fps) * sampling_rate)
+ if waveform.shape[-1] > target_samples:
+ waveform = waveform[..., :target_samples]
+ return Audio(waveform=waveform, sampling_rate=sampling_rate)
+
+
+def _build_tt_stage_c(
+ mesh_device: ttnn.MeshDevice, *, parallel_config, ccl_manager: CCLManager | None
+) -> VocoderWithBWE:
+ main_voc = Vocoder(
+ mesh_device=mesh_device,
+ dtype=ttnn.float32,
+ in_channels=128,
+ out_channels=2,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ **_tt_vocoder_cfg(_MAIN_VOCODER_CFG),
+ )
+ # Mirror the pipeline (LTXAudioDecoderAdapter): channel-TP the BWE generator only
+ # where the channel axis pays for the gather (factor > 2); single-axis on 2x4 (factor 2).
+ if isinstance(parallel_config, AudioTCParallelConfig):
+ bwe_pc = parallel_config if parallel_config.channel_parallel.factor > 2 else parallel_config.time_parallel
+ else:
+ bwe_pc = parallel_config
+ bwe_voc = Vocoder(
+ mesh_device=mesh_device,
+ dtype=ttnn.float32,
+ in_channels=128,
+ out_channels=2,
+ parallel_config=bwe_pc,
+ ccl_manager=ccl_manager,
+ **_tt_vocoder_cfg(_BWE_VOCODER_CFG),
+ )
+ mel_stft = MelSTFT(mesh_device=mesh_device, dtype=ttnn.float32, **_MEL_STFT_CFG)
+ return VocoderWithBWE(
+ vocoder=main_voc,
+ bwe_generator=bwe_voc,
+ mel_stft=mel_stft,
+ input_sampling_rate=_INPUT_SR,
+ output_sampling_rate=_OUTPUT_SR,
+ hop_length=_HOP_LENGTH,
+ mesh_device=mesh_device,
+ dtype=ttnn.float32,
+ )
+
+
+# ---------------------------------------------------------------------------
+# Component stage tests
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+ (
+ "mesh_device",
+ "mesh_shape",
+ "sp_axis",
+ "tp_axis",
+ "num_links",
+ "dynamic_load",
+ "device_params",
+ "topology",
+ "is_fsdp",
+ ),
+ LTX_AUDIO_MESH_PARAMS_FULL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_stage_a_audio_decoder(
+ mesh_device: ttnn.MeshDevice,
+ mesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ dynamic_load: bool,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+):
+ AutoencoderKLLTX2Audio, _, _ = _require_diffusers()
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ _ = (sp_axis, tp_axis, num_links, dynamic_load, topology, is_fsdp) # parity with fast_av config schema
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ torch.manual_seed(42)
+ ref_vae = AutoencoderKLLTX2Audio(
+ base_channels=_AUDIO_DECODER_CFG["ch"],
+ output_channels=_AUDIO_DECODER_CFG["out_ch"],
+ ch_mult=_AUDIO_DECODER_CFG["ch_mult"],
+ num_res_blocks=_AUDIO_DECODER_CFG["num_res_blocks"],
+ attn_resolutions=_AUDIO_DECODER_CFG["attn_resolutions"],
+ resolution=_AUDIO_DECODER_CFG["resolution"],
+ latent_channels=_AUDIO_DECODER_CFG["z_channels"],
+ norm_type="pixel",
+ causality_axis="height",
+ mid_block_add_attention=_AUDIO_DECODER_CFG["mid_block_add_attention"],
+ sample_rate=_AUDIO_DECODER_CFG["sample_rate"],
+ mel_hop_length=_AUDIO_DECODER_CFG["mel_hop_length"],
+ is_causal=_AUDIO_DECODER_CFG["is_causal"],
+ mel_bins=_AUDIO_DECODER_CFG["mel_bins"],
+ ).eval()
+
+ z_times_f = _AUDIO_DECODER_CFG["z_channels"] * _AUDIO_DECODER_CFG["mel_bins"]
+ stats_std = torch.ones(z_times_f)
+ stats_mean = torch.zeros(z_times_f)
+
+ tt_decoder = MelDecoder(mesh_device=mesh_device, **_AUDIO_DECODER_CFG)
+ tt_decoder.load_torch_state_dict(
+ _audio_decoder_state_from_diffusers(ref_vae, stats_std=stats_std, stats_mean=stats_mean)
+ )
+
+ latent = torch.randn(1, _AUDIO_DECODER_CFG["z_channels"], 64, _AUDIO_DECODER_CFG["mel_bins"], dtype=torch.float32)
+ with torch.no_grad():
+ torch_out = _decode_latent_ref(ref_vae, latent, stats_std=stats_std, stats_mean=stats_mean)
+ tt_out = tt_decoder(latent)
+
+ assert torch_out.shape == tt_out.shape
+ assert_quality(torch_out, tt_out, pcc=0.998, relative_rmse=0.05)
+
+
+@pytest.mark.parametrize(
+ (
+ "mesh_device",
+ "mesh_shape",
+ "sp_axis",
+ "tp_axis",
+ "num_links",
+ "dynamic_load",
+ "device_params",
+ "topology",
+ "is_fsdp",
+ ),
+ LTX_AUDIO_MESH_PARAMS_FULL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_stage_b_vocoder(
+ mesh_device: ttnn.MeshDevice,
+ mesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ dynamic_load: bool,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+):
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ _ = (sp_axis, tp_axis, dynamic_load, is_fsdp) # parity with fast_av config schema
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+ pc = _audio_parallel_config(mesh_shape)
+ ccl = CCLManager(mesh_device, num_links=num_links, topology=ttnn.Topology.Linear) if pc is not None else None
+
+ torch_voc = _build_torch_stage_b(seed=42)
+ tt_voc = _build_tt_stage_b(mesh_device, parallel_config=pc, ccl_manager=ccl)
+ tt_voc.load_torch_state_dict(_diffusers_vocoder_state_to_tt(torch_voc.state_dict()))
+
+ mel = _vocoder_mel()
+ with torch.no_grad():
+ ref_out = torch_voc(mel)
+ tt_out = tt_voc(mel)
+
+ assert tt_out.shape == ref_out.shape
+ assert_quality(ref_out, tt_out, pcc=0.99)
+
+ # When T-sharded, assert the sharded path reproduces the unsharded device path: the t_pad
+ # boundary click is localized to the tail and slips past aggregate PCC, but breaks this.
+ if pc is not None:
+ tt_voc_un = _build_tt_stage_b(mesh_device, parallel_config=None, ccl_manager=None)
+ tt_voc_un.load_torch_state_dict(_diffusers_vocoder_state_to_tt(torch_voc.state_dict()))
+ _assert_sharded_matches_unsharded(tt_voc_un(mel), tt_out, name="stage_b_vocoder")
+
+
+@pytest.mark.parametrize(
+ (
+ "mesh_device",
+ "mesh_shape",
+ "sp_axis",
+ "tp_axis",
+ "num_links",
+ "dynamic_load",
+ "device_params",
+ "topology",
+ "is_fsdp",
+ ),
+ LTX_AUDIO_MESH_PARAMS_FULL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_stage_c_vocoder_with_bwe(
+ mesh_device: ttnn.MeshDevice,
+ mesh_shape: tuple[int, int],
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ dynamic_load: bool,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+):
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ _ = (sp_axis, tp_axis, dynamic_load, is_fsdp) # parity with fast_av config schema
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+ pc = _audio_parallel_config(mesh_shape)
+ ccl = CCLManager(mesh_device, num_links=num_links, topology=ttnn.Topology.Linear) if pc is not None else None
+
+ torch_full = _build_torch_stage_c(seed=42)
+ tt_full = _build_tt_stage_c(mesh_device, parallel_config=pc, ccl_manager=ccl)
+
+ incompatible = tt_full.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict()))
+ assert not incompatible.missing_keys, f"missing keys: {incompatible.missing_keys}"
+
+ mel = _vocoder_mel()
+
+ # Stage B inside Stage C.
+ with torch.no_grad():
+ ref_main = torch_full.vocoder(mel)
+ tt_main = tt_full.vocoder(mel)
+ assert tt_main.shape == ref_main.shape
+ assert_quality(ref_main, tt_main, pcc=0.99)
+ if pc is not None:
+ tt_full_un = _build_tt_stage_c(mesh_device, parallel_config=None, ccl_manager=None)
+ tt_full_un.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict()))
+ _assert_sharded_matches_unsharded(tt_full_un.vocoder(mel), tt_main, name="stage_c_vocoder")
+
+ # Log-mel path (apples-to-apples same waveform).
+ remainder = ref_main.shape[-1] % _HOP_LENGTH
+ if remainder:
+ pad_right = _HOP_LENGTH - remainder
+ ref_main = torch.nn.functional.pad(ref_main, (0, pad_right))
+ with torch.no_grad():
+ ref_mel = _compute_mel_ref(torch_full, ref_main)
+ tt_mel = tt_full._compute_mel_device(ref_main)
+ assert tt_mel.shape == ref_mel.shape
+ assert_quality(ref_mel, tt_mel, pcc=0.99)
+
+ # BWE residual and skip-resample paths.
+ ref_mel_for_bwe = ref_mel.transpose(2, 3).contiguous()
+ with torch.no_grad():
+ ref_residual = torch_full.bwe_generator(ref_mel_for_bwe)
+ ref_skip = torch_full.resampler(ref_main)
+ tt_residual = tt_full.bwe_generator(ref_mel_for_bwe)
+ tt_skip = tt_full._resample_device(ref_main)
+ assert tt_residual.shape == ref_residual.shape
+ assert tt_skip.shape == ref_skip.shape
+ assert_quality(ref_residual, tt_residual, pcc=0.98)
+ assert_quality(ref_skip, tt_skip, pcc=0.99)
+
+ # Full Stage C forward: the muxed waveform is where the sharded tail click surfaced.
+ with torch.no_grad():
+ ref_out = torch_full(mel)
+ tt_out = tt_full(mel)
+ assert tt_out.shape == ref_out.shape
+ if pc is not None:
+ _assert_sharded_matches_unsharded(tt_full_un(mel), tt_out, name="stage_c_full")
+
+
+# Localized audio regression: aggregate PCC over a full clip masks a brief burst, so this
+# drives the real ~6s length and gates each 0.5s window vs the unsharded eager baseline.
+_STATIC_WIN_S = 0.5
+
+
+def _window_error_rows(ref: torch.Tensor, tt: torch.Tensor, *, sr: int, win_s: float):
+ """Per-window (start_s, max|Δ|, rmse/σ_ref) over the time axis of (B, C, T)."""
+ T = ref.shape[-1]
+ n = max(1, int(win_s * sr))
+ rows = []
+ for start in range(0, T, n):
+ r = ref[..., start : start + n]
+ t = tt[..., start : start + n]
+ max_abs = (r - t).abs().max().item()
+ rmse = ((r - t) ** 2).mean().sqrt().item()
+ sigma = ((r**2).mean().sqrt().item()) + 1e-9
+ rows.append((start / sr, max_abs, rmse / sigma))
+ return rows
+
+
+@pytest.mark.parametrize(
+ "device_params",
+ [
+ {
+ "l1_small_size": 32768,
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "trace_region_size": 300000000,
+ "require_exact_physical_num_devices": True,
+ }
+ ],
+ indirect=True,
+)
+@pytest.mark.parametrize("mesh_device", [(2, 4)], indirect=True)
+def test_stage_c_audio_decode(mesh_device, device_params):
+ """End-to-end audio synthesis (VocoderWithBWE): channel-TP main vocoder + conv1d depthwise,
+ gated against the unsharded eager baseline per 0.5s window so a localized burst (the demo
+ static) can't hide behind an aggregate PCC. Drives the real ~6s clip length (601 mel
+ frames)."""
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(2, 4))
+ pc = AudioTCParallelConfig(
+ time_parallel=ParallelFactor(factor=4, mesh_axis=1),
+ channel_parallel=ParallelFactor(factor=2, mesh_axis=0),
+ )
+ ccl = CCLManager(mesh_device, num_links=1, topology=ttnn.Topology.Linear)
+
+ torch_full = _build_torch_stage_c(seed=42)
+ tt_full = _build_tt_stage_c(mesh_device, parallel_config=pc, ccl_manager=ccl)
+ tt_full.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict()))
+ tt_un = _build_tt_stage_c(mesh_device, parallel_config=None, ccl_manager=None)
+ tt_un.load_torch_state_dict(_diffusers_vocoder_with_bwe_state_to_tt(torch_full.state_dict()))
+
+ torch.manual_seed(7)
+ _tf = int(os.environ.get("AUDIO_T_FRAMES", "601")) # 601 mel frames ≈ the 6s girl clip
+ mel = torch.randn(1, 2, _tf, 64, dtype=torch.float32) * 0.5
+ tt = tt_full(mel)
+ out_un = tt_un(mel)
+
+ sr = tt_full.output_sampling_rate
+ logger.info(f"[audio] output {tuple(tt.shape)} = {tt.shape[-1]/sr:.2f}s @ {sr}Hz")
+ su = _window_error_rows(out_un, tt, sr=sr, win_s=_STATIC_WIN_S)
+ for s, su_max, su_rmse in su:
+ logger.info(f"[audio] t={s:4.2f}s shard-vs-un max|Δ|={su_max:.3e} rmse/σ={su_rmse:.3e}")
+ worst_su_max = max(r[1] for r in su)
+ worst_su_rmse = max(r[2] for r in su)
+ print(
+ f"\nAUDIO_DECODE shard-vs-un worst_max|Δ|={worst_su_max:.3e} worst_rmse/σ={worst_su_rmse:.3e}",
+ flush=True,
+ )
+ # Regression guard: channel-TP audio must match the unsharded eager TT baseline. Catches
+ # broken sharding (the actual risk); torch-vs-TT on the full muxed BWE output is inherently
+ # loose with random weights and is checked per-component in test_stage_c_vocoder_with_bwe.
+ assert worst_su_max < 5e-3, f"audio window shard-vs-un max|Δ| {worst_su_max:.3e} — divergence"
+ assert worst_su_rmse < 1e-2, f"audio window shard-vs-un rmse/σ {worst_su_rmse:.3e} — divergence"
+
+
+# ===========================================================================
+# Tracy per-op device-time profile of the on-device audio decode chain.
+#
+# Goal: emit every op in the audio decode (mel-VAE -> vocoder -> BWE) with its
+# device duration, and have the summed device time land on the ~0.49s traced
+# steady-state audio decode (within a small delta).
+#
+# Why eager (traced=False): each op dispatches as its own program, so (a) the
+# on-device profiler can be flushed per-block to stay under the 125-scope/core
+# buffer, and (b) every op gets a named DEVICE FW / KERNEL DURATION. The device
+# kernel durations are identical to the traced replay — trace only removes host
+# dispatch — so the summed device time reproduces the traced ~0.49s while the
+# host wall-clock (printed) is far larger; that gap is pure host dispatch.
+#
+# Two entry points:
+# * test_audio_decode_profile — the workload (run under `python -m tracy`)
+# * test_audio_decode_perf_table — driver that runs the profiler on the above
+# and prints the per-op table + device total
+#
+# Manual one-liner (equivalent to the driver):
+# LTX_AUDIO_PROF=1 \
+# LTX_CHECKPOINT=~/.cache/ltx-checkpoints/ltx-2.3-22b-distilled-1.1.safetensors \
+# python -m tracy -p -r -v -o generated/profiler/ltx_audio_decode -m \
+# pytest 'models/tt_dit/tests/models/ltx/test_audio_ltx.py::test_audio_decode_profile' -s
+# then sum "DEVICE FW DURATION [ns]" from
+# generated/profiler/ltx_audio_decode/reports//ops_perf_results_.csv
+# ===========================================================================
+
+# The ~0.49s traced steady-state audio decode (4x8 ring, distilled, 145f@24fps),
+# expressed in ms; the summed device time should land within _PROF_TOL_MS of it.
+_PROF_TARGET_MS = 490.0
+_PROF_TOL_MS = 150.0
+
+# Compute-heavy leaf blocks; a ReadDeviceProfiler flush after each keeps every
+# Tracy zone window well under the per-core 125-scope device buffer.
+_PROF_FLUSH_TYPES = (ResnetBlock, AudioUpsample, AMPBlock1, MelSTFT)
+
+
+def _walk_tt(module):
+ """Depth-first walk over a tt_dit Module tree (root included)."""
+ yield module
+ for _, child in module.named_children():
+ yield from _walk_tt(child)
+
+
+def _flush_forward_after(module, mesh_device):
+ """Wrap ``module.forward`` so the on-device profiler is drained after every call,
+ keeping each Tracy zone window under the per-core scope buffer (see gemma harness)."""
+ orig = module.forward
+
+ def timed(*a, _orig=orig, **k):
+ r = _orig(*a, **k)
+ ttnn.ReadDeviceProfiler(mesh_device)
+ return r
+
+ module.forward = timed
+
+
+@pytest.mark.skipif(
+ os.environ.get("LTX_AUDIO_PROF") != "1",
+ reason="Tracy profiling workload — run via test_audio_decode_perf_table, or manually with "
+ "LTX_AUDIO_PROF=1 python -m tracy -p -r -o -m pytest '::test_audio_decode_profile' -s",
+)
+# Profiler-instrumented kernels compile fresh on the first profiling run and the eager decode
+# runs op-by-op, so this comfortably exceeds the default 300s per-test cap.
+@pytest.mark.timeout(3600)
+@pytest.mark.parametrize(
+ "device_params",
+ [
+ {
+ "l1_small_size": 32768,
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING,
+ "trace_region_size": 300_000_000,
+ "require_exact_physical_num_devices": True,
+ }
+ ],
+ indirect=True,
+)
+@pytest.mark.parametrize("mesh_device", [(4, 8)], indirect=True)
+def test_audio_decode_profile(mesh_device, device_params):
+ """Eager on-device audio decode wrapped in Tracy signposts, one profiled iteration.
+
+ Mirrors the 4x8 ring distilled pipeline's audio-decode config (T-shard=8, replicated
+ channel; the pipeline default with LTX_AUDIO_CHANNEL_TP off). Emits ``signpost("start")``
+ / ``signpost("stop")`` around a single decode so ``post_process_ops_log(has_signposts=True)``
+ slices out warmup/compile; per-block ReadDeviceProfiler flushes keep the device buffer live.
+ Pure timing — no PCC (correctness is covered by the sibling e2e/component tests)."""
+ ckpt = _resolve_checkpoint(
+ [f"{_LOCAL_CHECKPOINT_DIR}/{_DISTILLED_CHECKPOINT}"],
+ (_DISTILLED_CHECKPOINT,),
+ )
+ if ckpt is None:
+ pytest.skip("distilled checkpoint not found (set LTX_CHECKPOINT)")
+
+ from tracy import signpost
+
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(4, 8))
+
+ torch.manual_seed(0)
+ # traced=False (LTXPipeline default): decode_audio runs eager, one program per op.
+ pipeline, _ = _build_pipeline(
+ mesh_device,
+ sp_axis=1,
+ tp_axis=0,
+ checkpoint=ckpt,
+ num_links=2,
+ topology=ttnn.Topology.Ring,
+ )
+ assert not pipeline._traced, "profile expects the eager decode path (per-op programs)"
+
+ audio_latent = torch.randn(1, 256, 128, dtype=torch.float32)
+ num_frames, fps = 145, 24.0
+
+ # Warmup: compile kernels + one-time device weight prep, unprofiled. Two iters so the
+ # program cache and any lazy device state are fully warm before the measured pass.
+ for _ in range(2):
+ assert pipeline.decode_audio(audio_latent, num_frames, fps=fps) is not None
+ ttnn.synchronize_device(mesh_device)
+ ttnn.ReadDeviceProfiler(mesh_device) # drain warmup zones (they precede the "start" marker)
+
+ # Install per-block flushes only for the measured pass.
+ for root in (pipeline.tt_mel_decoder, pipeline.tt_vocoder_with_bwe):
+ for mod in _walk_tt(root):
+ if isinstance(mod, _PROF_FLUSH_TYPES):
+ _flush_forward_after(mod, mesh_device)
+
+ # Capture the device-op-id range around the single forward. GLOBAL CALL COUNT in the C++
+ # device digest (cpp_device_perf_report.csv) == this counter, so [op0, op1) slices the digest
+ # to exactly this forward — no dependence on the fragile Tracy -r ops report or signpost rows.
+ from ttnn import _ttnn
+
+ op0 = _ttnn.get_device_operation_id()
+ signpost("start")
+ t0 = time.perf_counter()
+ out = pipeline.decode_audio(audio_latent, num_frames, fps=fps)
+ ttnn.synchronize_device(mesh_device)
+ host_wall_ms = (time.perf_counter() - t0) * 1000.0
+ ttnn.ReadDeviceProfiler(mesh_device) # flush the tail ops while still inside the region
+ signpost("stop")
+ op1 = _ttnn.get_device_operation_id()
+
+ assert out is not None
+ logger.info(f"Audio decode profiled: host_wall={host_wall_ms:.1f}ms wave={tuple(out.waveform.shape)}")
+ print(f"\nAUDIO_DECODE_HOST_WALL_MS={host_wall_ms:.2f}", flush=True)
+ print(f"AUDIO_DECODE_OPID_RANGE={op0},{op1}", flush=True)
+
+
+@pytest.mark.skipif(
+ os.environ.get("LTX_AUDIO_PROF") != "1",
+ reason="Tracy driver — run: LTX_AUDIO_PROF=1 LTX_CHECKPOINT= pytest "
+ "'::test_audio_decode_perf_table' -s",
+)
+@pytest.mark.timeout(3600)
+def test_audio_decode_perf_table():
+ """Profile ``test_audio_decode_profile`` and print the per-op device-time table.
+
+ Spawns ``python -m tracy`` on the workload (opening the device itself; this driver does not
+ take the mesh_device fixture, matching test_ring_joint_sdpa_perf_table), then reads the
+ signpost-sliced ops log. Prints one row per op (device FW/kernel duration) plus the summed
+ device time, and asserts the total lands within ``_PROF_TOL_MS`` of the ~0.49s traced target.
+ """
+ import pandas as pd
+ from tracy.process_model_log import post_process_ops_log, run_device_profiler
+
+ subdir = "ltx_audio_decode"
+ test_id = "models/tt_dit/tests/models/ltx/test_audio_ltx.py::test_audio_decode_profile"
+ run_device_profiler(
+ f"-m 'pytest {test_id} -s'",
+ subdir,
+ check_test_return_code=True,
+ is_command_binary_exe=True,
+ )
+
+ df = post_process_ops_log(subdir, has_signposts=True) # DataFrame sliced between start/stop
+ fw_col, k_col = "DEVICE FW DURATION [ns]", "DEVICE KERNEL DURATION [ns]"
+
+ def _num(series):
+ return pd.to_numeric(series[series != "-"], errors="coerce").fillna(0.0)
+
+ fw_ns = _num(df[fw_col])
+ k_ns = _num(df[k_col])
+ total_fw_ms = fw_ns.sum() / 1e6
+ total_k_ms = k_ns.sum() / 1e6
+
+ # Per-op table, ops in dispatch order.
+ print("\n" + "=" * 78, flush=True)
+ print(f"{'#':>3} {'OP CODE':<34}{'FW [µs]':>12}{'KERNEL [µs]':>14}", flush=True)
+ print("-" * 78, flush=True)
+ codes = df["OP CODE"].tolist()
+ fw_list = _num(df[fw_col]).tolist()
+ k_list = _num(df[k_col]).tolist()
+ for i, (code, fw, kd) in enumerate(zip(codes, fw_list, k_list)):
+ print(f"{i:>3} {str(code):<34}{fw / 1e3:>12.2f}{kd / 1e3:>14.2f}", flush=True)
+ print("-" * 78, flush=True)
+
+ # Per-op-code aggregate (sorted by total FW time).
+ agg = (
+ pd.DataFrame({"code": codes, "fw_ns": fw_list, "k_ns": k_list})
+ .groupby("code")
+ .agg(count=("fw_ns", "size"), fw_us=("fw_ns", lambda s: s.sum() / 1e3), k_us=("k_ns", lambda s: s.sum() / 1e3))
+ .sort_values("fw_us", ascending=False)
+ )
+ print("\nBy op code (FW-descending):", flush=True)
+ print(f"{'OP CODE':<34}{'count':>7}{'FW [µs]':>12}{'KERNEL [µs]':>14}", flush=True)
+ for code, row in agg.iterrows():
+ print(f"{str(code):<34}{int(row['count']):>7}{row['fw_us']:>12.1f}{row['k_us']:>14.1f}", flush=True)
+ print("=" * 78, flush=True)
+
+ print(
+ f"\nAUDIO_DECODE_DEVICE_FW_MS={total_fw_ms:.2f} "
+ f"AUDIO_DECODE_DEVICE_KERNEL_MS={total_k_ms:.2f} "
+ f"(target ~{_PROF_TARGET_MS:.0f}ms, tol ±{_PROF_TOL_MS:.0f}ms, n_ops={len(df)})",
+ flush=True,
+ )
+ logger.info(f"Audio decode device time: FW={total_fw_ms:.1f}ms kernel={total_k_ms:.1f}ms over {len(df)} ops")
+ assert len(df) > 0, "no ops captured between signposts — check flushing / signpost naming"
+ assert abs(total_fw_ms - _PROF_TARGET_MS) <= _PROF_TOL_MS, (
+ f"summed device FW {total_fw_ms:.1f}ms is >{_PROF_TOL_MS:.0f}ms from the ~{_PROF_TARGET_MS:.0f}ms "
+ f"traced target — device-bound audio-decode time drifted"
+ )
diff --git a/code/models/tt_dit/tests/models/ltx/test_embeddings_ltx.py b/code/models/tt_dit/tests/models/ltx/test_embeddings_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c0cd747928443de4d5dfc65cc3f1844c8cd28e5
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_embeddings_ltx.py
@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.layers.embeddings import LTXAdaLayerNormSingle
+from models.tt_dit.utils.check import assert_quality
+
+
+@pytest.mark.parametrize(
+ "mesh_device",
+ [(1, 1)],
+ ids=["1x1"],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize("embedding_dim", [4096], ids=["dim4096"])
+def test_ltx_adaln_single(mesh_device: ttnn.MeshDevice, embedding_dim: int):
+ """
+ Test LTXAdaLayerNormSingle: compare TT output vs diffusers LTX2AdaLayerNormSingle.
+ """
+ from diffusers.models.transformers.transformer_ltx2 import LTX2AdaLayerNormSingle as TorchAdaLayerNormSingle
+
+ B = 2
+
+ # Create PyTorch reference model
+ torch_model = TorchAdaLayerNormSingle(
+ embedding_dim=embedding_dim, num_mod_params=6, use_additional_conditions=False
+ )
+ torch_model.eval()
+ torch_state = torch_model.state_dict()
+ logger.info(f"PyTorch AdaLayerNormSingle state keys: {list(torch_state.keys())}")
+
+ # Create TT model and load weights
+ tt_model = LTXAdaLayerNormSingle(
+ embedding_dim=embedding_dim,
+ embedding_coefficient=6,
+ mesh_device=mesh_device,
+ dtype=ttnn.float32,
+ )
+ tt_model.load_torch_state_dict(torch_state)
+
+ # Create timestep input
+ torch.manual_seed(42)
+ timestep = torch.rand(B) * 1000 # random timesteps in [0, 1000)
+
+ # PyTorch reference forward
+ with torch.no_grad():
+ torch_modulation, torch_embedded = torch_model(timestep)
+ logger.info(f"PyTorch modulation shape: {torch_modulation.shape}, embedded: {torch_embedded.shape}")
+
+ # TT forward
+ # Timesteps module expects (1, 1, B, 1) so B broadcasts with the 128-dim factor.
+ # Use float32 for timestep input — sinusoidal embeddings are precision-sensitive.
+ tt_timestep = ttnn.from_torch(
+ timestep.reshape(1, 1, B, 1),
+ device=mesh_device,
+ layout=ttnn.TILE_LAYOUT,
+ dtype=ttnn.float32,
+ )
+ tt_modulation, tt_embedded = tt_model(tt_timestep)
+
+ # Convert back and compare — output is (1, 1, B, dim) in ttnn
+ tt_modulation_torch = ttnn.to_torch(tt_modulation).squeeze(0).squeeze(0)
+ tt_embedded_torch = ttnn.to_torch(tt_embedded).squeeze(0).squeeze(0)
+
+ logger.info(f"TT modulation shape: {tt_modulation_torch.shape}, embedded: {tt_embedded_torch.shape}")
+
+ assert_quality(torch_modulation, tt_modulation_torch, pcc=0.999)
+ assert_quality(torch_embedded, tt_embedded_torch, pcc=0.999)
+ logger.info("PASSED: LTXAdaLayerNormSingle matches PyTorch reference")
diff --git a/code/models/tt_dit/tests/models/ltx/test_latent_upsampler.py b/code/models/tt_dit/tests/models/ltx/test_latent_upsampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e14605dd16b4e28748a69248e1c372cd64a9664
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_latent_upsampler.py
@@ -0,0 +1,256 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""PCC tests for ``LTXLatentUpsampler`` vs the diffusers ``LTX2LatentUpsamplerModel`` reference.
+
+All shapes match the production LTX-2 Fast pipeline (mid_channels=1024, num_frames=19,
+2x4 BH mesh, half-res latent 17x30 / 18x32)."""
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.models.upsampler.latent_upsampler_ltx import LTXLatentUpsampler, LTXUpsamplerResBlock
+from models.tt_dit.parallel.config import ParallelFactor, VaeHWParallelConfig
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.conv3d import ConvDims, conv_pad_height, conv_pad_width
+from models.tt_dit.utils.tensor import fast_device_to_host, typed_tensor_2dshard
+from models.tt_dit.utils.test import skip_if_unsupported_num_links
+
+from .ltx_mesh_params import LTX_UPSAMPLER_MESH_PARAMS
+
+
+def _parallel_kwargs(
+ mesh_device: ttnn.MeshDevice, h_axis: int, w_axis: int, num_links: int, topology: ttnn.Topology
+) -> dict:
+ mesh_shape = tuple(mesh_device.shape)
+ parallel_config = VaeHWParallelConfig(
+ height_parallel=ParallelFactor(factor=mesh_shape[h_axis], mesh_axis=h_axis),
+ width_parallel=ParallelFactor(factor=mesh_shape[w_axis], mesh_axis=w_axis),
+ )
+ ccl_manager = CCLManager(mesh_device, num_links=num_links, topology=topology)
+ return {"parallel_config": parallel_config, "ccl_manager": ccl_manager}
+
+
+@pytest.mark.parametrize(
+ "H, W",
+ [(17, 30), (18, 32)],
+ ids=["nondiv_17x30", "div_18x32"],
+)
+@pytest.mark.parametrize(
+ ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"),
+ LTX_UPSAMPLER_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_upsampler_resblock(
+ mesh_device: ttnn.MeshDevice, h_axis: int, w_axis: int, num_links: int, topology: ttnn.Topology, H: int, W: int
+):
+ """Single LTXUpsamplerResBlock at the production shape (channels=1024, 2x4 sharded,
+ T=19). ``nondiv_17x30`` exercises the sharded pad/mask/crop path; ``div_18x32`` is
+ the no-mask control."""
+ from diffusers.pipelines.ltx2.latent_upsampler import ResBlock as TorchResBlock
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ channels, T, B = 1024, 19, 1
+ torch.manual_seed(0xABBA)
+ pk = _parallel_kwargs(mesh_device, h_axis, w_axis, num_links, topology)
+ pc = pk["parallel_config"]
+ hf, wf = pc.height_parallel.factor, pc.width_parallel.factor
+ padded_h = ((H + hf - 1) // hf) * hf
+ padded_w = ((W + wf - 1) // wf) * wf
+
+ torch_block = TorchResBlock(channels=channels, dims=3)
+ torch_block.eval()
+
+ tt_block = LTXUpsamplerResBlock(
+ channels=channels,
+ gn_input_nhw=T * H * W,
+ mesh_device=mesh_device,
+ conv_dims=ConvDims(T=T + 2, H=padded_h // hf, W=padded_w // wf),
+ **pk,
+ )
+ tt_block.load_torch_state_dict(torch_block.state_dict())
+
+ x = torch.randn(B, channels, T, H, W, dtype=torch.float32)
+ with torch.no_grad():
+ torch_out = torch_block(x) # (B, C, T, H, W)
+
+ x_bthwc = x.permute(0, 2, 3, 4, 1).contiguous() # (B, T, H, W, C)
+ x_bthwc, logical_h = conv_pad_height(x_bthwc, hf)
+ x_bthwc, logical_w = conv_pad_width(x_bthwc, wf)
+ x_tt = typed_tensor_2dshard(
+ x_bthwc,
+ mesh_device,
+ shard_mapping={pc.height_parallel.mesh_axis: 2, pc.width_parallel.mesh_axis: 3},
+ layout=ttnn.ROW_MAJOR_LAYOUT,
+ dtype=ttnn.bfloat16,
+ )
+
+ out = tt_block(x_tt, logical_h, logical_w)
+
+ concat_dims = [None, None]
+ concat_dims[pc.height_parallel.mesh_axis] = 2
+ concat_dims[pc.width_parallel.mesh_axis] = 3
+ out_torch = fast_device_to_host(out, mesh_device, concat_dims, ccl_manager=pk["ccl_manager"])
+ out_torch = out_torch[:, :, :logical_h, :logical_w, :channels].permute(0, 4, 1, 2, 3)
+
+ assert out_torch.shape == torch_out.shape, f"shape mismatch: TT {out_torch.shape} vs torch {torch_out.shape}"
+ assert_quality(torch_out, out_torch, pcc=0.99)
+ logger.info(f"PASSED: LTXUpsamplerResBlock(1024 ch, {T}x{H}x{W}, 2x4) matches reference")
+
+
+@pytest.mark.parametrize(
+ "H, W",
+ [(17, 30), (18, 32)],
+ ids=["nondiv_17x30", "div_18x32"],
+)
+@pytest.mark.parametrize(
+ ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"),
+ LTX_UPSAMPLER_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_latent_upsampler(
+ mesh_device: ttnn.MeshDevice, h_axis: int, w_axis: int, num_links: int, topology: ttnn.Topology, H: int, W: int
+):
+ """Full ``LTXLatentUpsampler`` PCC vs reference, random weights at production shape
+ (in=128, mid=1024, T=19, 2x4). No HF checkpoint required."""
+ from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ in_c, mid_c, T, B = 128, 1024, 19, 1
+ torch.manual_seed(0xC0FFEE)
+
+ torch_model = LTX2LatentUpsamplerModel(
+ in_channels=in_c,
+ mid_channels=mid_c,
+ num_blocks_per_stage=4,
+ dims=3,
+ spatial_upsample=True,
+ temporal_upsample=False,
+ rational_spatial_scale=2.0,
+ use_rational_resampler=False,
+ )
+ torch_model.eval()
+
+ tt_model = LTXLatentUpsampler(
+ input_hw=(H, W),
+ in_channels=in_c,
+ mid_channels=mid_c,
+ num_blocks_per_stage=4,
+ spatial_upsample=True,
+ temporal_upsample=False,
+ spatial_scale=2.0,
+ rational_resampler=False,
+ mesh_device=mesh_device,
+ num_frames=T,
+ **_parallel_kwargs(mesh_device, h_axis, w_axis, num_links, topology),
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ latent = torch.randn(B, in_c, T, H, W, dtype=torch.float32)
+
+ with torch.no_grad():
+ torch_out = torch_model(latent)
+ tt_out = tt_model(latent)
+
+ assert tt_out.shape == torch_out.shape, f"shape mismatch: TT {tt_out.shape} vs torch {torch_out.shape}"
+ assert_quality(torch_out, tt_out, pcc=0.99)
+ logger.info(f"PASSED: LTXLatentUpsampler({in_c}->{mid_c}, {T}x{H}x{W} → {T}x{H * 2}x{W * 2}) matches reference")
+
+
+def _resolve_upsampler_checkpoint() -> str | None:
+ """Resolve the real spatial-upscaler safetensors (local HF cache, no download)."""
+ from huggingface_hub import hf_hub_download
+ from huggingface_hub.utils import LocalEntryNotFoundError
+
+ repo_id, filename = "Lightricks/LTX-2.3", "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
+ try:
+ return hf_hub_download(repo_id, filename, local_files_only=True)
+ except (LocalEntryNotFoundError, Exception): # noqa: BLE001
+ return None
+
+
+@pytest.mark.parametrize(
+ "T, H, W",
+ [(19, 17, 30), (19, 18, 32)],
+ ids=["nondiv_17x30", "div_18x32"],
+)
+@pytest.mark.parametrize(
+ ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"),
+ LTX_UPSAMPLER_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_latent_upsampler_real_checkpoint(
+ mesh_device: ttnn.MeshDevice,
+ h_axis: int,
+ w_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ T: int,
+ H: int,
+ W: int,
+):
+ """PCC vs reference with the real ``ltx-2.3-spatial-upscaler-x2`` weights (mid_channels=1024)
+ on the production 2x4 mesh. Input is a smooth, per-channel-scaled latent: white noise
+ is unrepresentative and lower-bounds bf16 conv PCC through the 8-block residual chain."""
+ import json
+
+ from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
+ from safetensors import safe_open
+ from safetensors.torch import load_file
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ ckpt = _resolve_upsampler_checkpoint()
+ if ckpt is None:
+ pytest.skip("real upsampler checkpoint not in local HF cache")
+
+ with safe_open(ckpt, framework="pt") as f:
+ cfg = json.loads(f.metadata()["config"])
+ in_c, mid_c = cfg["in_channels"], cfg["mid_channels"]
+ n_blocks = cfg["num_blocks_per_stage"]
+ B = 1
+ torch.manual_seed(0xC0FFEE)
+
+ torch_model = LTX2LatentUpsamplerModel(
+ in_channels=in_c,
+ mid_channels=mid_c,
+ num_blocks_per_stage=n_blocks,
+ dims=3,
+ spatial_upsample=True,
+ temporal_upsample=False,
+ rational_spatial_scale=2.0,
+ use_rational_resampler=False,
+ )
+ torch_model.load_state_dict(load_file(ckpt))
+ torch_model.eval()
+
+ tt_model = LTXLatentUpsampler(
+ input_hw=(H, W),
+ in_channels=in_c,
+ mid_channels=mid_c,
+ num_blocks_per_stage=n_blocks,
+ spatial_upsample=True,
+ temporal_upsample=False,
+ spatial_scale=2.0,
+ rational_resampler=False,
+ mesh_device=mesh_device,
+ num_frames=T,
+ **_parallel_kwargs(mesh_device, h_axis, w_axis, num_links, topology),
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ # Smooth, spatially-correlated latent with per-channel scale (mimics real VAE latents).
+ noise = torch.randn(B, in_c, T, H, W, dtype=torch.float32)
+ latent = torch.nn.functional.avg_pool3d(noise, kernel_size=(1, 3, 3), stride=1, padding=(0, 1, 1))
+ latent = latent * (torch.rand(in_c) * 4 + 0.5).view(1, in_c, 1, 1, 1)
+
+ with torch.no_grad():
+ torch_out = torch_model(latent)
+ tt_out = tt_model(latent)
+
+ assert tt_out.shape == torch_out.shape, f"shape mismatch: TT {tt_out.shape} vs torch {torch_out.shape}"
+ assert_quality(torch_out, tt_out, pcc=0.99)
+ logger.info(f"PASSED: real-checkpoint LTXLatentUpsampler({in_c}->{mid_c}, {n_blocks} blocks) matches reference")
diff --git a/code/models/tt_dit/tests/models/ltx/test_ltx_adapter_loader.py b/code/models/tt_dit/tests/models/ltx/test_ltx_adapter_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..7321c3dedc4bdd9eeca800b9e0725c1511f133fb
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_ltx_adapter_loader.py
@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Host-only unit tests for the LTX LoRA adapter loader's module resolution.
+
+No device, no checkpoint: these pin the loader's warn-and-skip contract for
+targets that do not exist on the transformer variant being loaded. A LoRA file
+authored against a deeper variant references block indices past this variant's
+depth; resolution must return "absent" so the caller skips, not raise IndexError.
+"""
+
+from types import SimpleNamespace
+
+from models.tt_dit.experimental.lora.ltx_adapter_loader import _get_attn, _resolve_singleton
+
+
+def _fake_transformer(num_blocks: int):
+ """A transformer stub with ``num_blocks`` blocks, each carrying an ``attn1``."""
+ blocks = [SimpleNamespace(attn1=SimpleNamespace(name=f"attn1_{i}")) for i in range(num_blocks)]
+ return SimpleNamespace(transformer_blocks=blocks)
+
+
+def test_get_attn_out_of_range_returns_none():
+ tf = _fake_transformer(2)
+ assert _get_attn(tf, 5, "attn1") is None
+ assert _get_attn(tf, -1, "attn1") is None
+
+
+def test_get_attn_in_range_resolves():
+ tf = _fake_transformer(2)
+ assert _get_attn(tf, 1, "attn1") is tf.transformer_blocks[1].attn1
+ # Missing attribute on an in-range block still degrades to None, not AttributeError.
+ assert _get_attn(tf, 1, "attn2") is None
+
+
+def test_resolve_singleton_out_of_range_returns_none_triple():
+ tf = _fake_transformer(2)
+ assert _resolve_singleton(tf, 9, "attn1", "to_out") == (None, None, None)
+ assert _resolve_singleton(tf, -1, "attn1", "ff1") == (None, None, None)
diff --git a/code/models/tt_dit/tests/models/ltx/test_ltx_lora_ondevice.py b/code/models/tt_dit/tests/models/ltx/test_ltx_lora_ondevice.py
new file mode 100644
index 0000000000000000000000000000000000000000..1099c141cf6ca78ae25fd1a601421aeeea50569e
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_ltx_lora_ondevice.py
@@ -0,0 +1,176 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""On-device fuse-mode LoRA swap for the LTX-2.3 pipeline.
+
+Unlike the host-fuse path (``fuse_loras_into`` + full weight reload), swapping here is a
+``bind_active`` on each LoRA Linear: ``weight.data += scale * A@B`` on device, no reload.
+
+Correctness anchor: the host-layout algebra (loader A/B == ``fuse_loras_into``) is already
+proven CPU-side. This test closes the remaining gap — that ``LoRAMixin._apply_delta`` applies
+that delta correctly on the *sharded device* weight — by comparing the on-device fused weight
+against a host recompute of the same registered A/B.
+
+The device cases need the real 22B checkpoint + a LoRA file + a mesh; they mirror the opt-in
+style of ``test_ltx_lora_runtime.py``. Set ``RUN_LORA_GEN=1`` for an end-to-end base/lora/unloaded
+generate A/B.
+"""
+
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_pcc
+from models.tt_dit.experimental.lora.ltx_adapter_loader import iter_lora_modules
+from models.tt_dit.pipelines.ltx.pipeline_ltx import LTXPipeline
+from models.tt_dit.utils import tensor
+from models.tt_dit.utils.test import line_params
+
+
+def _default_checkpoint() -> str:
+ explicit = os.environ.get("LTX_CHECKPOINT")
+ if explicit:
+ return explicit
+ local = os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-dev.safetensors")
+ if os.path.exists(local):
+ return local
+ return "Lightricks/LTX-2.3:ltx-2.3-22b-dev.safetensors"
+
+
+def _default_lora() -> str:
+ from models.tt_dit.utils.ltx_lora_asset import resolve_lora
+
+ return resolve_lora()
+
+
+def _read_weight(param) -> torch.Tensor:
+ """Reconstruct a sharded Parameter's full weight on host."""
+ return tensor.to_torch(param.data, mesh_axes=param.mesh_axes)
+
+
+def _first_qkv_lora_module(transformer):
+ """First self-attn to_qkv LoRA module (path, module) — the fused/permuted/interleaved
+ path, i.e. the one most sensitive to a layout bug."""
+ for path, mod in iter_lora_modules(transformer):
+ if path.endswith("attn1.to_qkv"):
+ return path, mod
+ return None, None
+
+
+@pytest.mark.timeout(3600)
+@pytest.mark.parametrize(
+ "mesh_device, mesh_shape, sp_axis, tp_axis, num_links, dynamic_load, device_params, topology, is_fsdp",
+ [
+ [(2, 2), (2, 2), 0, 1, 2, False, line_params, ttnn.Topology.Linear, True],
+ [(4, 8), (4, 8), 1, 0, 2, False, line_params, ttnn.Topology.Linear, False],
+ ],
+ ids=["2x2sp0tp1", "bh_4x8sp1tp0"],
+ indirect=["mesh_device", "device_params"],
+)
+def test_ondevice_lora_bind_matches_host_delta(
+ mesh_device,
+ mesh_shape,
+ sp_axis,
+ tp_axis,
+ num_links,
+ dynamic_load,
+ topology,
+ is_fsdp,
+):
+ """bind_active changes the weight to match base + host(scale*A@B); unbind restores the base
+ (within bf16 add/subtract drift); the transformer Module is reused throughout."""
+ ckpt = _default_checkpoint()
+ lora = _default_lora()
+ strength = float(os.environ.get("LORA_STRENGTH", "1.0"))
+
+ parent_mesh = mesh_device
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ pipeline = LTXPipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=ckpt,
+ gemma_path=None,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ dynamic_load=dynamic_load,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ mode="av",
+ num_frames=int(os.environ.get("NUM_FRAMES", "25")),
+ height=int(os.environ.get("HEIGHT", "256")),
+ width=int(os.environ.get("WIDTH", "256")),
+ lora_enabled=True,
+ )
+
+ if int(ttnn.distributed_context_get_rank()) != 0:
+ logger.info("Skipping assertions on non-zero rank")
+ return
+
+ transformer_id = id(pipeline.transformer)
+ assert pipeline._active_lora is None, "fresh pipeline must have no active LoRA"
+
+ path, module = _first_qkv_lora_module(pipeline.transformer)
+ assert module is not None, "no attn1.to_qkv LoRA module found — is lora_enabled wired?"
+ base_w = _read_weight(module.weight).clone()
+
+ # --- bind (on-device fuse) ---------------------------------------------
+ pipeline.load_lora_weights(lora, strength=strength)
+ assert id(pipeline.transformer) == transformer_id, "model object was rebuilt, not reused"
+ assert module.is_lora_active, f"{path} was not bound"
+ fused_w = _read_weight(module.weight)
+ assert not torch.allclose(fused_w, base_w), "bind did not change the weight"
+
+ # host recompute of the SAME registered delta; scale folded in exactly as _apply_delta does.
+ adapter = module.lora_bank[module.active_idx]
+ delta_oi = module.active_scale * (adapter.B.to(torch.float32) @ adapter.A.to(torch.float32)) # [out,in]
+ delta = delta_oi if delta_oi.shape == base_w.shape else delta_oi.T
+ assert delta.shape == base_w.shape, f"delta {tuple(delta.shape)} vs weight {tuple(base_w.shape)}"
+ expected = base_w.to(torch.float32) + delta
+
+ ok, pcc = comp_pcc(expected, fused_w.to(torch.float32), pcc=0.999)
+ logger.info(f"[{path}] on-device bind vs host delta: {pcc}")
+ assert ok, f"on-device fused weight disagrees with host delta recompute: {pcc}"
+
+ # --- unbind ------------------------------------------------------------
+ pipeline.unload_lora_weights()
+ assert pipeline._active_lora is None
+ assert id(pipeline.transformer) == transformer_id, "model object was rebuilt, not reused"
+ restored_w = _read_weight(module.weight)
+ # fuse-mode unbind subtracts in bf16 — close, not bit-exact (see lora.py drift note).
+ assert torch.allclose(
+ restored_w.to(torch.float32), base_w.to(torch.float32), atol=1e-2
+ ), "unbind did not restore the base weight within bf16 drift"
+ logger.info("On-device LoRA bind/unbind weight checks passed.")
+
+ if os.environ.get("RUN_LORA_GEN", "0") in ("1", "true", "True"):
+ steps = int(os.environ.get("STEPS", "6"))
+ num_frames = int(os.environ.get("NUM_FRAMES", "25"))
+ height = int(os.environ.get("HEIGHT", "256"))
+ width = int(os.environ.get("WIDTH", "256"))
+ seq = pipeline.gemma_encoder_pair.sequence_length
+ v_p = torch.zeros(1, seq, pipeline.gemma_encoder_pair.video_dim)
+ a_p = torch.zeros(1, seq, pipeline.gemma_encoder_pair.audio_dim)
+ for tag, spec in (("base", None), ("lora", lora), ("unloaded", None)):
+ if spec is not None:
+ pipeline.load_lora_weights(spec, strength=strength)
+ else:
+ pipeline.unload_lora_weights()
+ pipeline._prepare_transformer(0)
+ v_lat, a_lat = pipeline.call_av(
+ video_prompt_embeds=v_p,
+ audio_prompt_embeds=a_p,
+ neg_video_prompt_embeds=v_p,
+ neg_audio_prompt_embeds=a_p,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ num_inference_steps=steps,
+ seed=0,
+ ge_gamma=0.0,
+ )
+ logger.info(f"[{tag}] generate ran; video latent {tuple(v_lat.shape)}")
+ assert id(pipeline.transformer) == transformer_id, "model object was rebuilt during generation"
diff --git a/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py
new file mode 100644
index 0000000000000000000000000000000000000000..b2602d45116bbe31112bc5fa6d00f3568d845b10
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py
@@ -0,0 +1,757 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import itertools
+import math
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+
+import numpy as np
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.models.audio_vae.audio_decoder_ltx import LTXAudioDecoderAdapter
+from models.tt_dit.pipelines.ltx.pipeline_ltx_distilled import LTXDistilledPipeline
+from models.tt_dit.utils.ltx import (
+ DEFAULT_LTX_PROMPT,
+ STEADY_STATE_LTX_PROMPT,
+ STEADY_STATE_REPLAY_LTX_PROMPT,
+ default_ltx_checkpoint,
+ default_ltx_gemma,
+ print_ltx_timing_table,
+ traced_default,
+)
+from models.tt_dit.utils.patchifiers import AudioLatentShape, VideoPixelShape
+from models.tt_dit.utils.test import skip_if_unsupported_num_links
+from models.tt_dit.utils.vbench import assert_vbench_quality
+
+from .ltx_mesh_params import (
+ LTX_DISTILLED_AUDIO_MESH_PARAMS_DL,
+ LTX_DISTILLED_I2V_MESH_PARAMS_DL,
+ LTX_DISTILLED_MESH_PARAMS_DL,
+)
+
+
+def _ltx_checkpoint_cached(filename: str) -> bool:
+ """True if the LTX checkpoint is available -- either a real local file (LTX_CHECKPOINT / flat
+ ~/.cache/ltx-checkpoints) OR already in the HF hub cache. The old guard checked
+ os.path.exists(default_ltx_checkpoint(...)), but when the checkpoint resolves from the hub that
+ returns a "repo:file" string (never a path), so it always skipped in CI even with the weights
+ cached. Checking the hub cache too (local_files_only) lets the offline CI run actually execute."""
+ ref = default_ltx_checkpoint(filename)
+ if os.path.exists(ref):
+ return True
+ if ":" in ref: # "repo_id:filename" -> resolvable from the local hub cache without a network call
+ repo_id, fname = ref.split(":", 1)
+ from huggingface_hub import hf_hub_download
+ from huggingface_hub.utils import EntryNotFoundError, LocalEntryNotFoundError
+
+ try:
+ hf_hub_download(repo_id=repo_id, filename=fname, local_files_only=True)
+ return True
+ except (LocalEntryNotFoundError, EntryNotFoundError, FileNotFoundError):
+ return False
+ return False
+
+
+# Default-off: full AV gen needs the real LTX checkpoint + Gemma, so it skips without one.
+@pytest.mark.skipif(
+ not _ltx_checkpoint_cached("ltx-2.3-22b-distilled-1.1.safetensors"),
+ reason="needs the LTX checkpoint (set LTX_CHECKPOINT to a local .safetensors)",
+)
+@pytest.mark.parametrize(
+ "no_prompt",
+ [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), True)],
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load",
+ LTX_DISTILLED_MESH_PARAMS_DL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_pipeline_distilled(
+ mesh_device,
+ device_params,
+ sp_axis,
+ tp_axis,
+ num_links,
+ dynamic_load,
+ topology,
+ is_fsdp,
+ no_prompt,
+):
+ """LTX-2.3 distilled 2-stage AV pipeline."""
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ ckpt = default_ltx_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors")
+ gemma = default_ltx_gemma()
+
+ parent_mesh = mesh_device
+ mesh_shape = tuple(parent_mesh.shape)
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ # Stage-2 output target: 1080p @ 24fps, ~6s. Constraints from the pipeline:
+ # - height/width must be divisible by 64 (asserted at the top of generate(),
+ # so half-res stage-1 is divisible by 32 for the VAE).
+ # 1920 is already % 64; 1080 is not → round up to 1088 (next multiple of 64).
+ # If you need exactly 1080 lines, post-crop the 8 extra rows.
+ # - num_frames must satisfy (num_frames - 1) % 8 == 0 so the VAE decoder maps
+ # latent_frames → num_frames frames exactly. 144 (== 6s × 24fps) doesn't
+ # satisfy this; 145 does (19 latent frames → 145 decoded frames ≈ 6.04s).
+ num_frames = int(os.environ.get("NUM_FRAMES", "145"))
+ height = int(os.environ.get("HEIGHT", "1088"))
+ width = int(os.environ.get("WIDTH", "1920"))
+
+ run_warmup = os.environ.get("RUN_WARMUP", "0") in ("1", "true", "True")
+ # Traced by default wherever the mesh param reserves a trace region; LTX_TRACED=0/1 overrides
+ # (rule + rationale: utils.ltx.traced_default, unit-tested in tests/unit/test_ltx_traced_default.py).
+ traced = traced_default(device_params, os.environ.get("LTX_TRACED"))
+
+ # Conditioning image (I2V). Its mere presence drives image_conditioning: with a path the
+ # transformer builds the per-token video-timestep (I2V) modulation; without one pure T2V keeps
+ # the fast scalar-AdaLN path. Resolved before create_pipeline so the bool can gate the build.
+ image_path = os.environ.get("LTX_I2V_IMAGE")
+ images = None
+ if image_path:
+ if not os.path.exists(image_path):
+ pytest.skip(f"LTX_I2V_IMAGE set but file not found: {image_path}")
+ strength = float(os.environ.get("LTX_I2V_STRENGTH", "1.0"))
+ images = [(image_path, 0, strength)]
+
+ pipeline = LTXDistilledPipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=ckpt,
+ gemma_path=gemma,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ dynamic_load=dynamic_load,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ run_warmup=run_warmup,
+ traced=traced,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ image_conditioning=bool(image_path),
+ )
+
+ prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT)
+ # Gen #1 is the steady-state measurement, and on console every request encodes a prompt the
+ # cache has never seen — so it gets its own prompt to keep the encoder in the measured path.
+ # Under dynamic_load the encoder is coresident-excluded with the DiT: a second encode would
+ # evict the DiT and clobber the captured traces, so that path reuses the cached embedding.
+ steady_state_prompt = prompt if dynamic_load else os.environ.get("PROMPT_STEADY_STATE", STEADY_STATE_LTX_PROMPT)
+ replay_prompt = prompt if dynamic_load else STEADY_STATE_REPLAY_LTX_PROMPT
+
+ def run(*, prompt, number, seed):
+ output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_fast_{width}x{height}_{number}.mp4")
+ logger.info(f"Running LTX AV Fast: '{prompt[:80]}...'")
+ logger.info(f"Config: {height}x{width}, {num_frames} frames")
+ if images:
+ logger.info(f"I2V: conditioning image {images[0][0]} (strength={images[0][2]})")
+
+ if int(ttnn.distributed_context_get_rank()) != 0:
+ logger.info(f"Skipping generation on rank {ttnn.distributed_context_get_rank()}")
+ return
+
+ pipeline.generate(
+ prompt,
+ output_path=output_filename,
+ images=images,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ seed=seed,
+ )
+ logger.info(f"Saved video to: {output_filename}")
+ print_ltx_timing_table(
+ pipeline,
+ label="LTX DISTILLED",
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ mesh_shape=mesh_shape,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ topology=topology,
+ output_path=output_filename,
+ prompt=prompt,
+ )
+
+ vbench_thresholds_by_height = {
+ 1088: {
+ "subject_consistency": 0.92,
+ "background_consistency": 0.93,
+ "motion_smoothness": 0.955,
+ # Averaged over VBENCH_SEEDS clips: dynamic_degree is near-binary per clip and this content
+ # sits at ~0.8-1.0, so the floor requires ~4/5 seeds dynamic (1.0 would demand every seed).
+ "dynamic_degree": 0.8,
+ "imaging_quality": 0.645,
+ },
+ }
+
+ # RUN_VBENCH=0 skips the quality gate (e.g. perf-only iteration); defaults on so CI gates.
+ run_vbench = os.environ.get("RUN_VBENCH", "1") in ("1", "true", "True")
+ # RUN_CLIP=0 skips the CLIP prompt-alignment gate; defaults on (mirrors the wan2.2 test).
+ run_clip = os.environ.get("RUN_CLIP", "1") in ("1", "true", "True")
+
+ export_dir = os.environ.get("VBENCH_EXPORT_DIR")
+ if export_dir and not run_vbench:
+ raise ValueError("VBENCH_EXPORT_DIR requires RUN_VBENCH=1")
+ # Missing dependencies must fail an enabled gate, including in CI.
+ if run_vbench and not export_dir:
+ __import__("vbench")
+ if run_clip:
+ __import__("decord")
+
+ def check_output_with_vbench(prompt, number, seed=None):
+ if not run_vbench:
+ logger.info("RUN_VBENCH=0, skipping VBench quality gate")
+ return
+ if int(ttnn.distributed_context_get_rank()) != 0:
+ return
+ thresholds = vbench_thresholds_by_height.get(height)
+ if thresholds is None:
+ pytest.skip(f"no VBench thresholds calibrated for height {height}")
+
+ # Gate on the average of VBENCH_SEEDS traced replays (default 5): a single clip's VBench score
+ # is too noisy to gate on — dynamic_degree is near-binary per clip — so score the set together
+ # (VBench averages over a directory). The primary clip (`number`, `seed`) is already rendered;
+ # the rest are cheap warm replays. A single-clip caller (no seed) or OUTPUT_PATH (one pinned
+ # filename) keeps the one-clip path.
+ num_seeds = int(os.environ.get("VBENCH_SEEDS", "5"))
+ if export_dir and (seed is None or num_seeds != 5 or os.environ.get("OUTPUT_PATH")):
+ raise ValueError("CI VBench export requires five distinct traced seed clips")
+ if seed is None or num_seeds <= 1 or os.environ.get("OUTPUT_PATH"):
+ output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_fast_{width}x{height}_{number}.mp4")
+ assert_vbench_quality(output_filename, prompt=prompt, thresholds=thresholds)
+ return
+
+ cwd = os.path.abspath(os.getcwd())
+
+ with tempfile.TemporaryDirectory() as vbench_dir:
+
+ def _link(idx: int, mp4_number: int) -> None:
+ # The render always lands directly in CWD; resolve the path and confirm it stays there so
+ # the number threaded into the filename can't escape it, and give the link a plain integer
+ # name so nothing dynamic reaches the destination path either.
+ src = os.path.abspath(os.path.join(cwd, f"ltx_av_fast_{width}x{height}_{mp4_number}.mp4"))
+ if os.path.dirname(src) != cwd:
+ raise ValueError(f"render path escaped the working dir: {src!r}")
+ os.symlink(src, os.path.join(vbench_dir, f"seed_{idx}.mp4"))
+
+ _link(0, number)
+ for k in range(1, num_seeds):
+ run(prompt=prompt, number=1000 + k, seed=seed + k)
+ _link(k, 1000 + k)
+ logger.info(f"VBench gate averaged over {num_seeds} seeds (base seed {seed})")
+ if export_dir:
+ from models.tt_dit.utils.vbench_bundle import export_bundle
+
+ export_bundle(
+ vbench_dir,
+ export_dir,
+ prompt=prompt,
+ thresholds=thresholds,
+ temporal_width=int(os.environ.get("VBENCH_TEMPORAL_WIDTH", "0")),
+ )
+ else:
+ assert_vbench_quality(vbench_dir, prompt=prompt, thresholds=thresholds)
+
+ def check_output_with_clip(prompt, number, clip_threshold=None):
+ # Mirrors wan2.2's check_output_with_clip: sample ~8 evenly-spaced frames, score each
+ # against the prompt with CLIP, assert the mean clears a floor. LTX writes the video to
+ # disk (generate() returns only the path), so frames are read back from the mp4.
+ if not run_clip:
+ logger.info("RUN_CLIP=0, skipping CLIP score check")
+ return
+ if int(ttnn.distributed_context_get_rank()) != 0:
+ return
+ try:
+ from decord import VideoReader
+ from PIL import Image
+
+ from models.tt_dit.tests.dataset_eval.clip_encoder import CLIPEncoder
+ except ImportError as e:
+ pytest.skip(f"CLIP deps unavailable ({e})")
+
+ # LTX baseline mean ~31.3 (vs wan2.2's ~37): the LTX prompt exceeds CLIP's 77-token limit
+ # so only its head is scored, and frames are read back re-encoded (CRF=25). 28.0 leaves
+ # ~3pt headroom for run-to-run variance. Override with CLIP_THRESHOLD.
+ threshold = float(os.environ.get("CLIP_THRESHOLD", "28.0")) if clip_threshold is None else clip_threshold
+ output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_fast_{width}x{height}_{number}.mp4")
+
+ vr = VideoReader(output_filename)
+ total_frames = len(vr)
+ indices = np.linspace(0, total_frames - 1, min(8, total_frames), dtype=int)
+
+ clip_encoder = CLIPEncoder()
+ scores = []
+ for i in indices:
+ frame = vr[int(i)].asnumpy() # (H, W, 3) uint8 RGB
+ pil_img = Image.fromarray(frame.astype(np.uint8))
+ score = clip_encoder.get_clip_score(prompt, pil_img).item() * 100.0
+ scores.append(score)
+
+ clip_min = min(scores)
+ clip_max = max(scores)
+ clip_mean = sum(scores) / len(scores)
+ logger.info(f"CLIP scores: min={clip_min:.2f}, max={clip_max:.2f}, mean={clip_mean:.2f}")
+
+ assert clip_mean >= threshold, (
+ f"Mean CLIP score {clip_mean:.2f} is below threshold {threshold:.2f}. "
+ f"Per-frame scores: {[f'{s:.2f}' for s in scores]}"
+ )
+
+ if no_prompt:
+ seed = int(os.environ.get("SEED", "10"))
+ run(prompt=prompt, number=0, seed=seed)
+ # Traced: gen #0 captures the denoise/VAE/audio traces (lazily, on first step of each stage).
+ if traced:
+ # Gen #1 captures the encode trace (the pipeline opens that gate once its own captures
+ # are done), so gen #2 is the first pass where every trace replays — the one whose
+ # timings are the steady state.
+ logger.info("=== gen #1: encode trace capture ===")
+ run(prompt=steady_state_prompt, number=1, seed=seed)
+ logger.info("=== traced steady-state pass (gen #2, pure replay) ===")
+ run(prompt=replay_prompt, number=2, seed=seed)
+ # Gate gen #2: a corrupted full-replay pass is the failure this structure exists to
+ # catch, and gen #1 still has a capture in it.
+ check_output_with_clip(replay_prompt, 2)
+ check_output_with_vbench(replay_prompt, 2, seed=seed)
+ else:
+ check_output_with_clip(prompt, 0)
+ check_output_with_vbench(prompt, 0)
+ else:
+ for i in itertools.count():
+ new_prompt = input("Enter the input prompt, or q to exit: ")
+ if new_prompt:
+ prompt = new_prompt
+ if prompt[0] == "q":
+ break
+ run(prompt=prompt, number=i, seed=i)
+ check_output_with_clip(prompt, i)
+ check_output_with_vbench(prompt, i)
+
+ if traced:
+ pipeline.release_traces()
+
+
+@functools.lru_cache(maxsize=1)
+def _ffmpeg():
+ """An ffmpeg executable. Prefer the system binary; fall back to imageio-ffmpeg's
+ bundled static build (pip-installed on demand) so a host without ffmpeg still runs
+ the frame-decode checks. The static build is real ffmpeg, so decoded pixels — and
+ thus the calibrated seam thresholds — are unchanged."""
+ exe = shutil.which("ffmpeg")
+ if exe:
+ return exe
+ try:
+ import imageio_ffmpeg
+ except ImportError:
+ pip_install = [sys.executable, "-m", "pip", "install", "-q", "imageio-ffmpeg"]
+ if subprocess.run(pip_install).returncode != 0:
+ # uv-managed venvs ship without pip; bootstrap it from the stdlib, then retry.
+ subprocess.run([sys.executable, "-m", "ensurepip", "--upgrade"], check=True)
+ subprocess.run(pip_install, check=True)
+ import imageio_ffmpeg
+ return imageio_ffmpeg.get_ffmpeg_exe()
+
+
+def _ffmpeg_frames(path, n):
+ """Sample ``n`` luma frames evenly across ``path`` (via ffmpeg) as float ndarrays."""
+ import io
+
+ from PIL import Image
+
+ frames = []
+ for t in np.linspace(0.1, 5.5, n):
+ raw = subprocess.run(
+ [
+ _ffmpeg(),
+ "-v",
+ "error",
+ "-ss",
+ f"{t}",
+ "-i",
+ path,
+ "-vframes",
+ "1",
+ "-f",
+ "image2pipe",
+ "-vcodec",
+ "png",
+ "-",
+ ],
+ capture_output=True,
+ ).stdout
+ if raw:
+ frames.append(np.asarray(Image.open(io.BytesIO(raw)).convert("L")).astype(float))
+ return frames
+
+
+def _temporal_seam_score(path):
+ """Grid-seam strength at the 2x4 mesh boundaries (W/4,W/2,3W/4 vertical; H/2 horizontal).
+ Per-column/row TEMPORAL MEDIAN of the gradient isolates a static seam from moving content.
+ Each boundary is scored against its LOCAL neighbour gradient (a ring +/-12 px, skipping
+ +/-1 px for seam width) — not the global median, which a centred subject over a dark,
+ low-gradient background inflates (the subject straddles W/2, not a mesh line). A baked seam
+ is a thin gradient spike pinned to the boundary, so it stands well above its flat local
+ ring; clean content scores ~1."""
+ fs = _ffmpeg_frames(path, 16)
+ assert fs, f"no frames decoded from {path}"
+ h, w = fs[0].shape
+ gx = np.median(np.stack([np.abs(np.diff(f, axis=1)).mean(0) for f in fs]), 0)
+ gy = np.median(np.stack([np.abs(np.diff(f, axis=0)).mean(1) for f in fs]), 0)
+
+ def _ring_ratio(g, idx, skip=1, span=12):
+ lo, hi = max(0, idx - span), min(len(g), idx + span + 1)
+ ring = np.concatenate([g[lo : idx - skip], g[idx + skip + 1 : hi]])
+ return float(g[idx] / (np.median(ring) + 1e-6))
+
+ v = float(np.mean([_ring_ratio(gx, round(w * i / 4)) for i in (1, 2, 3)]))
+ hh = _ring_ratio(gy, round(h / 2))
+ return v, hh
+
+
+@pytest.mark.skipif(
+ not _ltx_checkpoint_cached("ltx-2.3-22b-distilled-1.1.safetensors"),
+ reason="needs the LTX checkpoint (set LTX_CHECKPOINT to a local .safetensors)",
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load",
+ LTX_DISTILLED_I2V_MESH_PARAMS_DL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_pipeline_distilled_i2v(mesh_device, sp_axis, tp_axis, num_links, dynamic_load, topology, is_fsdp, tmp_path):
+ """E2E chained t2v->i2v: generate a t2v clip, then i2v continuing from its LAST frame
+ (same DEFAULT_LTX_PROMPT), and splice both into one continuous ~12s clip. Asserts the
+ I2V output (a) reproduces the conditioning frame at frame-0 (conditioning works) and
+ (b) is free of the VAE 2x4 grid seam (guards the non-mesh-aligned i2v fix at 1088x1920,
+ whose s1 cond latent is the uneven 17x30 that used to seam). An external LTX_I2V_IMAGE
+ overrides the t2v pre-gen (single i2v, no splice)."""
+ import subprocess
+
+ from PIL import Image
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+
+ ckpt = default_ltx_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors")
+ parent_mesh = mesh_device
+ mesh_shape = tuple(parent_mesh.shape)
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ num_frames = int(os.environ.get("NUM_FRAMES", "145"))
+ height = int(os.environ.get("HEIGHT", "1088")) # 1088x1920 -> s1 latent 17x30 (uneven, the fixed case)
+ width = int(os.environ.get("WIDTH", "1920"))
+ traced = os.environ.get("LTX_TRACED", "0") in ("1", "true", "True")
+ seed = int(os.environ.get("SEED", "10"))
+ prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT)
+
+ pipeline = LTXDistilledPipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=ckpt,
+ gemma_path=default_ltx_gemma(),
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ dynamic_load=dynamic_load,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ run_warmup=traced, # trace capture needs warmup; eager doesn't
+ traced=traced,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ # This test conditions on a frame produced in-process (t2v last frame), so there's no
+ # conditioning-image path in hand at construction — force the I2V path on explicitly.
+ image_conditioning=True,
+ )
+
+ if int(ttnn.distributed_context_get_rank()) != 0:
+ return
+
+ def _gen(out, images):
+ pipeline.generate(
+ prompt, output_path=str(out), images=images, num_frames=num_frames, height=height, width=width, seed=seed
+ )
+
+ # Conditioning image: an explicit LTX_I2V_IMAGE conditions directly on it and SKIPS the t2v
+ # pre-gen (an external/user image, or a t2v last frame produced by a prior process). Else this
+ # generates a t2v clip and conditions on its LAST FRAME so the i2v continues the shot.
+ cond_override = os.environ.get("LTX_I2V_IMAGE")
+ strength = float(os.environ.get("LTX_I2V_STRENGTH", "1.0"))
+ t2v = None
+ if cond_override:
+ cond = cond_override
+ else:
+ t2v = tmp_path / "t2v.mp4"
+ _gen(t2v, None)
+ # -sseof seeks near the end and -update rewrites the same file per decoded frame, so it
+ # lands on the true final frame without buffering the whole clip into memory.
+ cond = tmp_path / "cond_lastframe.png"
+ subprocess.run(
+ [_ffmpeg(), "-v", "error", "-sseof", "-3", "-i", str(t2v), "-update", "1", "-y", str(cond)],
+ check=True,
+ )
+ if os.environ.get("LTX_T2V_OUT"):
+ shutil.copy(str(t2v), os.environ["LTX_T2V_OUT"])
+ if os.environ.get("LTX_COND_OUT"):
+ shutil.copy(str(cond), os.environ["LTX_COND_OUT"])
+ # On a 4x8 mesh the image encoder and DiT coresident-evict, so a t2v->i2v chain in one
+ # process reloads the encoder over t2v's captured traces and wedges. LTX_T2V_ONLY stops
+ # here so a fresh process runs i2v on LTX_COND_OUT (via LTX_I2V_IMAGE) and host-splices.
+ if os.environ.get("LTX_T2V_ONLY"):
+ print(
+ f"I2V_E2E t2v-only -> clip={os.environ.get('LTX_T2V_OUT')} cond={os.environ.get('LTX_COND_OUT')}",
+ flush=True,
+ )
+ if traced:
+ pipeline.release_traces()
+ return
+
+ # i2v conditioned on that frame
+ i2v = tmp_path / "i2v.mp4"
+ _gen(i2v, [(str(cond), 0, strength)])
+ if os.environ.get("LTX_I2V_OUT"):
+ shutil.copy(str(i2v), os.environ["LTX_I2V_OUT"])
+
+ # In-process splice of t2v + i2v into one continuous clip (each gen is 145f@24fps -> ~12s).
+ # Only reached where encoder+DiT+traces fit one process (not the 4x8 mesh — use LTX_T2V_ONLY
+ # plus a second i2v process there). An external LTX_I2V_IMAGE has no preceding t2v to splice.
+ if t2v is not None:
+ chained = os.environ.get("LTX_CHAINED_OUT") or str(tmp_path / "chained.mp4")
+ subprocess.run(
+ [
+ _ffmpeg(),
+ "-v",
+ "error",
+ "-i",
+ str(t2v),
+ "-i",
+ str(i2v),
+ "-filter_complex",
+ "[0:v:0][0:a:0][1:v:0][1:a:0]concat=n=2:v=1:a=1[v][a]",
+ "-map",
+ "[v]",
+ "-map",
+ "[a]",
+ "-y",
+ chained,
+ ],
+ check=True,
+ )
+ print(f"I2V_E2E chained t2v+i2v -> {chained}", flush=True)
+
+ # (a) conditioning works: i2v frame-0 reproduces the conditioning frame (VAE roundtrip + CRF,
+ # so not identity — but a far tighter correlation than an unconditioned gen of the same prompt).
+ def _luma0(path):
+ raw = subprocess.run(
+ [_ffmpeg(), "-v", "error", "-i", path, "-vframes", "1", "-f", "image2pipe", "-vcodec", "png", "-"],
+ capture_output=True,
+ ).stdout
+ return Image.open(__import__("io").BytesIO(raw)).convert("L")
+
+ def _vec(im):
+ return torch.from_numpy(np.asarray(im).astype("float32")).flatten()
+
+ f0_img = _luma0(str(i2v))
+ cond_img = _luma0(str(cond))
+ # Compare a matched field of view: an external LTX_I2V_IMAGE can differ in resolution/aspect
+ # from the generated frame (the pipeline center-crops the conditioning image to the target
+ # aspect before encoding). Mirror that — center-crop cond to the frame's aspect, then resize.
+ # No-op when cond is a same-size/same-aspect frame (the self-conditioned t2v-last-frame case).
+ fw, fh = f0_img.size
+ cw, ch = cond_img.size
+ aspect = fw / fh
+ if cw / ch > aspect:
+ nw = round(ch * aspect)
+ x0 = (cw - nw) // 2
+ cond_img = cond_img.crop((x0, 0, x0 + nw, ch))
+ elif cw / ch < aspect:
+ nh = round(cw / aspect)
+ y0 = (ch - nh) // 2
+ cond_img = cond_img.crop((0, y0, cw, y0 + nh))
+ if cond_img.size != f0_img.size:
+ cond_img = cond_img.resize(f0_img.size)
+
+ c, f0 = _vec(cond_img), _vec(f0_img)
+ pcc = torch.corrcoef(torch.stack([c, f0]))[0, 1].item()
+ print(f"\nI2V_E2E frame0-vs-cond PCC={pcc:.4f}", flush=True)
+
+ # (b) seam-free at the uneven 1088x1920 (the i2v grid-seam fix). Local-ring ratios sit at
+ # ~1.0 on clean content (measured 0.85-1.16 across all boundaries); a baked seam spikes
+ # well above its flat neighbourhood. Threshold 1.5 clears the clean band with margin.
+ v, hh = _temporal_seam_score(str(i2v))
+ print(f"I2V_E2E seam V={v:.2f} H={hh:.2f} (clean~1.0, seam>>1)", flush=True)
+
+ if traced:
+ pipeline.release_traces()
+
+ assert pcc > 0.85, f"i2v frame-0 does not reproduce the conditioning frame (PCC={pcc:.4f}) — conditioning broken"
+ assert v < 1.5 and hh < 1.5, f"i2v grid seam present (V={v:.2f}, H={hh:.2f}) at 1088x1920"
+
+
+@pytest.mark.skipif(
+ not _ltx_checkpoint_cached("ltx-2.3-22b-distilled-1.1.safetensors"),
+ reason="needs the LTX checkpoint (set LTX_CHECKPOINT to a local .safetensors)",
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load",
+ LTX_DISTILLED_AUDIO_MESH_PARAMS_DL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_audio_decode_girl(mesh_device, sp_axis, tp_axis, num_links, dynamic_load, topology, is_fsdp):
+ """Just the audio section of the AV pipeline: real checkpoint weights, real girl-clip
+ latent shape (num_frames -> als.frames), no transformer/gemma (decode_audio never
+ encodes prompts). Profiles cold (weight load + compile) vs warm (steady-state per-gen)
+ decode wall. LTX_TRACED=1 traces the main vocoder (cold/warm timing only — the eager
+ torch-oracle quality block is skipped under trace; run LTX_TRACED=0 for the stats)."""
+ skip_if_unsupported_num_links(mesh_device, num_links)
+
+ parent_mesh = mesh_device
+ mesh_shape = tuple(parent_mesh.shape)
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ num_frames = int(os.environ.get("NUM_FRAMES", "145")) # ~6.04s @ 24fps
+ height = int(os.environ.get("HEIGHT", "1088"))
+ width = int(os.environ.get("WIDTH", "1920"))
+ traced = os.environ.get("LTX_TRACED", "0") in ("1", "true", "True")
+
+ ckpt = default_ltx_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors")
+ # Audio-only build: create the pipeline WITHOUT a checkpoint so __init__ builds/primes
+ # nothing (skipping the ~10-min video warmup and the 22B transformer / VAE / upsampler
+ # prime that dominate a cold run — decode_audio only touches the audio decoder + vocoder).
+ # Then point it at the checkpoint and construct + prime just the audio decode chain;
+ # decode_audio compiles + captures its own trace lazily on its first call.
+ pipeline = LTXDistilledPipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=None,
+ gemma_path=default_ltx_gemma(), # built as a lazy shell only; never loaded here
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ dynamic_load=dynamic_load,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ run_warmup=False,
+ traced=traced,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ )
+ pipeline.checkpoint_name = LTXDistilledPipeline._resolve_checkpoint_file(ckpt)
+ pipeline._audio_adapter = LTXAudioDecoderAdapter(
+ pipeline.checkpoint_name,
+ mesh_device=pipeline.mesh_device,
+ vae_ccl_manager=pipeline.vae_ccl_manager,
+ dit_parallel_config=pipeline.parallel_config,
+ traced=pipeline._traced,
+ )
+ pipeline._prepare_audio_decoder()
+
+ vps = VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=24)
+ als = AudioLatentShape.from_video_pixel_shape(vps)
+ # Seeded synthetic latent scaled to a real girl-clip audio latent's statistics (mean≈0.2,
+ # std≈0.92). The PCC gate below is TT-vs-torch with both legs fed the same mel, so the input
+ # need not be "real" content — only AUDIO_OUT .wav listening would differ. AUDIO_LATENT
+ # (.npy or .pt) still overrides for anyone who wants to decode a real dumped latent.
+ _lat = os.environ.get("AUDIO_LATENT")
+ if _lat and os.path.exists(_lat):
+ latent = (torch.from_numpy(np.load(_lat)) if _lat.endswith(".npy") else torch.load(_lat)).float()
+ else:
+ torch.manual_seed(0)
+ latent = torch.randn(1, als.frames, pipeline.in_channels, dtype=torch.float32) * 0.92 + 0.2
+
+ t0 = time.perf_counter()
+ audio = pipeline.decode_audio(latent, num_frames, fps=24.0) # cold: weight load + compile (+ capture)
+ cold_ms = (time.perf_counter() - t0) * 1000
+
+ # One untimed decode to absorb a late first-replay cost (the BWE trace captures on the
+ # first post-cold call, not during the cold decode), so warm_ms is true steady state.
+ pipeline.decode_audio(latent, num_frames, fps=24.0)
+ # WARM_REPS shrinks the steady-state loop for fast dev iteration (the PCC oracle below is
+ # unaffected); default 5 keeps warm_ms a stable average for reported timings.
+ N = int(os.environ.get("WARM_REPS", "5"))
+ t0 = time.perf_counter()
+ for _ in range(N):
+ audio = pipeline.decode_audio(latent, num_frames, fps=24.0)
+ warm_ms = (time.perf_counter() - t0) * 1000 / N
+ wav = audio.waveform
+ sr = audio.sampling_rate
+ out_dir = os.environ.get("AUDIO_OUT")
+
+ def _save(w, name):
+ if not out_dir:
+ return
+ import soundfile as sf
+
+ os.makedirs(out_dir, exist_ok=True)
+ a = w[0] if w.dim() == 3 else w # (2, T) -> (T, 2)
+ path = os.path.join(out_dir, name)
+ sf.write(path, a.transpose(0, 1).cpu().numpy(), int(sr))
+ logger.info(f"wrote {path}")
+
+ _save(wav, f"girl_audio_conv1d{'_traced' if traced else ''}.wav")
+
+ # Torch-oracle quality gate vs the diffusers vocoder+BWE (real checkpoint, CPU) — the ground
+ # truth. The conv1d path runs under BOTH eager and traced (gating the traced vocoder/BWE
+ # output, which is otherwise unverified — exactly the gap that let a broken trace ship).
+ # Per-1s-interval error-to-signal (rmse/σ, dB) + overall PCC; the absolute level (~−18 dB
+ # conv1d) is the fp32/bf16 floor, gated only against a gross spike.
+ from models.tt_dit.tests.models.ltx.test_audio_ltx import _build_torch_stage_c_real
+
+ z = pipeline.tt_mel_decoder.z_channels
+ audio_spatial = latent.reshape(1, latent.shape[1], z, latent.shape[2] // z).permute(0, 2, 1, 3).float()
+ mel = pipeline.tt_mel_decoder(audio_spatial) # TT mel, fed to both TT and torch
+ w_conv = pipeline.tt_vocoder_with_bwe(mel).squeeze(0).float()
+ with torch.no_grad():
+ w_torch = _build_torch_stage_c_real(ckpt)(mel.float()).squeeze(0).float()
+ _save(w_conv, f"girl_audio_conv1d{'_traced' if traced else ''}.wav")
+ _save(w_torch, "girl_audio_torch.wav")
+
+ nn = min(w_torch.shape[-1], w_conv.shape[-1])
+
+ def _vs_torch(test, name):
+ rows = []
+ for start in range(0, nn, sr):
+ r, t = w_torch[..., start : start + sr], test[..., start : start + sr]
+ sig = r.pow(2).mean().sqrt().item() + 1e-9
+ rr = ((r - t) ** 2).mean().sqrt().item() / sig
+ rows.append(rr)
+ logger.info(
+ f"[oracle] t={start//sr}s {name}-vs-torch rmse/σ={rr:.3e} ({20*math.log10(rr+1e-12):+.1f} dB) "
+ f"max|Δ|={(r - t).abs().max().item():.3e}"
+ )
+ pcc = torch.corrcoef(torch.stack([w_torch[..., :nn].flatten(), test[..., :nn].flatten()]))[0, 1].item()
+ worst_, med_ = max(rows), sorted(rows)[len(rows) // 2]
+ print(
+ f"\nAUDIO_GIRL {name}-vs-torch traced={traced} worst_rmse/σ={worst_:.3e} "
+ f"({20*math.log10(worst_+1e-12):+.1f} dB) median={med_:.3e} ratio={worst_/(med_+1e-9):.2f} PCC={pcc:.5f}",
+ flush=True,
+ )
+ return worst_, med_, pcc
+
+ cw, cm, c_pcc = _vs_torch(w_conv, "conv1d")
+ # Gross localized spike only: an interval >4x the median AND past a clearly-audible floor
+ # (rmse/σ 0.5 = −6 dB). Also gate overall PCC vs torch — a broken (e.g. mis-traced) vocoder
+ # drops PCC far below the ~0.99 fp32/bf16 floor.
+ assert cw < 0.5 and cw < 4 * cm + 1e-6, f"conv1d-vs-torch localized spike: worst {cw:.3e} vs median {cm:.3e}"
+ assert c_pcc > 0.95, f"conv1d-vs-torch PCC {c_pcc:.4f} below floor (traced={traced}) — vocoder output is wrong"
+
+ pipeline.release_traces()
+ dur = wav.shape[-1] / sr
+ print(
+ f"\nAUDIO_GIRL traced={traced} "
+ f"latent_frames={als.frames} out={tuple(wav.shape)} {dur:.2f}s@{sr}Hz cold={cold_ms:.1f}ms warm={warm_ms:.1f}ms",
+ flush=True,
+ )
+ assert torch.isfinite(wav).all(), "decoded waveform has non-finite samples"
+ assert abs(dur - num_frames / 24.0) < 0.2, f"duration {dur:.2f}s != ~{num_frames/24.0:.2f}s"
diff --git a/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_one_stage.py b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_one_stage.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e3794a6424ec3d73354f5ed60ee4d6cde8df5b0
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_one_stage.py
@@ -0,0 +1,173 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+LTX-2 one-stage pipeline tests:
+- Euler step scheduler primitive (no device needed)
+- End-to-end one-stage AV generation (LTXOneStagePipeline) on a multi-chip mesh
+"""
+
+import itertools
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.pipelines.ltx.pipeline_ltx import euler_step
+from models.tt_dit.pipelines.ltx.pipeline_ltx_one_stage import LTXOneStagePipeline
+from models.tt_dit.utils.ltx import (
+ DEFAULT_LTX_PROMPT,
+ default_ltx_checkpoint,
+ default_ltx_gemma,
+ print_ltx_timing_table,
+)
+from models.tt_dit.utils.test import skip_if_unsupported_num_links
+
+from .ltx_mesh_params import LTX_ONE_STAGE_MESH_PARAMS_DL
+
+
+def test_euler_step():
+ """Test Euler step: x_{t+1} = x_t + velocity * dt."""
+ torch.manual_seed(42)
+ sample = torch.randn(1, 256, 128)
+ denoised = torch.randn(1, 256, 128)
+ sigma, sigma_next = 0.8, 0.6
+
+ result = euler_step(sample, denoised, sigma=sigma, sigma_next=sigma_next)
+
+ # Manual computation
+ velocity = (sample.float() - denoised.float()) / sigma
+ dt = sigma_next - sigma
+ expected = (sample.float() + velocity * dt).to(sample.dtype)
+
+ assert torch.allclose(
+ result, expected, atol=1e-6
+ ), f"Euler step mismatch: max diff {(result - expected).abs().max():.2e}"
+ logger.info("Euler step matches manual computation")
+
+
+# ---------------------------------------------------------------------------
+# End-to-end one-stage AV generation (requires device + checkpoint)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "no_prompt",
+ [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), True)],
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load",
+ LTX_ONE_STAGE_MESH_PARAMS_DL,
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(
+ "width, height",
+ [
+ (1920, 1088),
+ ],
+ ids=[
+ "resolution_1080p",
+ ],
+)
+def test_pipeline_one_stage(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ dynamic_load,
+ topology,
+ width,
+ height,
+ is_fsdp,
+ no_prompt,
+):
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ ckpt = default_ltx_checkpoint("ltx-2.3-22b-dev.safetensors")
+ gemma = default_ltx_gemma()
+ # ckpt / gemma always resolve (env var → local → HF repo string fallback). The pipeline's
+ # resolver downloads from HF if needed.
+
+ parent_mesh = mesh_device
+ mesh_shape = tuple(parent_mesh.shape)
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ num_frames = int(os.environ.get("NUM_FRAMES", "145"))
+ num_inference_steps = int(os.environ.get("NUM_STEPS", "30"))
+ width = int(os.environ.get("WIDTH", width))
+ height = int(os.environ.get("HEIGHT", height))
+
+ run_warmup = os.environ.get("RUN_WARMUP", "0") in ("1", "true", "True")
+
+ pipeline = LTXOneStagePipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=ckpt,
+ gemma_path=gemma,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ dynamic_load=dynamic_load,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ run_warmup=run_warmup,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ )
+
+ prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT)
+
+ def run(*, prompt, number, seed):
+ logger.info(f"Running inference with prompt: '{prompt}'")
+ logger.info(f"Parameters: {height}x{width}, {num_frames} frames, {num_inference_steps} steps")
+
+ output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_pro_{width}x{height}_{number}.mp4")
+
+ pipeline.generate(
+ prompt,
+ output_path=output_filename,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ num_inference_steps=num_inference_steps,
+ seed=seed,
+ )
+
+ if int(ttnn.distributed_context_get_rank()) == 0:
+ # Minimal artifact gate: the leg must not pass on missing/empty output. A calibrated
+ # VBench/CLIP quality check (like the distilled path) is the stronger follow-up.
+ assert os.path.exists(output_filename), f"pipeline produced no video at {output_filename}"
+ assert os.path.getsize(output_filename) > 0, f"pipeline produced an empty video at {output_filename}"
+ logger.info(f"Saved video to: {output_filename}")
+ print_ltx_timing_table(
+ pipeline,
+ label="LTX ONE-STAGE",
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ mesh_shape=mesh_shape,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ topology=topology,
+ output_path=output_filename,
+ prompt=prompt,
+ )
+ else:
+ logger.info(f"Skipping video export on rank {ttnn.distributed_context_get_rank()}")
+
+ if no_prompt:
+ run(prompt=prompt, number=0, seed=42)
+ else:
+ for i in itertools.count():
+ new_prompt = input("Enter the input prompt, or q to exit: ")
+ if new_prompt:
+ prompt = new_prompt
+ if prompt[0] == "q":
+ break
+ run(prompt=prompt, number=i, seed=i)
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py
new file mode 100644
index 0000000000000000000000000000000000000000..56eb790ef190f73eb0cfdd9d5125dfd3284a98ea
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py
@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+import os
+
+import pytest
+from loguru import logger
+
+import ttnn
+from models.tt_dit.pipelines.ltx.pipeline_ltx_two_stages import LTXTwoStagesPipeline
+from models.tt_dit.utils.ltx import (
+ DEFAULT_LTX_PROMPT,
+ default_ltx_checkpoint,
+ default_ltx_gemma,
+ print_ltx_timing_table,
+)
+from models.tt_dit.utils.test import skip_if_unsupported_num_links
+
+from .ltx_mesh_params import LTX_TWO_STAGES_I2V_MESH_PARAMS_DL, LTX_TWO_STAGES_MESH_PARAMS_DL
+
+
+def _default_distilled_lora() -> str:
+ """Resolve distilled stage-2 LoRA: env var > local file > HF (auto-downloaded here)."""
+ explicit = os.environ.get("DISTILLED_LORA_PATH")
+ if explicit:
+ return explicit
+ local = os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-distilled-lora-384-1.1.safetensors")
+ if os.path.exists(local):
+ return local
+ from huggingface_hub import hf_hub_download
+
+ logger.info(
+ "Resolving HuggingFace distilled LoRA Lightricks/LTX-2.3:ltx-2.3-22b-distilled-lora-384-1.1.safetensors"
+ )
+ return hf_hub_download(repo_id="Lightricks/LTX-2.3", filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors")
+
+
+@pytest.mark.parametrize(
+ "no_prompt",
+ [{"1": True, "0": False}.get(os.environ.get("NO_PROMPT"), True)],
+)
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load",
+ LTX_TWO_STAGES_MESH_PARAMS_DL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_pipeline_two_stages(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ dynamic_load,
+ topology,
+ is_fsdp,
+ no_prompt,
+):
+ """LTX-2.3 22B 2-stage AV pipeline: full-guidance s1 + distilled-LoRA s2 refine."""
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ ckpt = default_ltx_checkpoint("ltx-2.3-22b-dev.safetensors")
+ distilled_lora = _default_distilled_lora()
+ gemma = default_ltx_gemma()
+
+ parent_mesh = mesh_device
+ mesh_shape = tuple(parent_mesh.shape)
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ num_frames = int(os.environ.get("NUM_FRAMES", "145"))
+ height = int(os.environ.get("HEIGHT", "1088"))
+ width = int(os.environ.get("WIDTH", "1920"))
+ num_inference_steps = int(os.environ.get("NUM_INFERENCE_STEPS", "30"))
+
+ run_warmup = os.environ.get("RUN_WARMUP", "0") in ("1", "true", "True")
+ pipeline = LTXTwoStagesPipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=ckpt,
+ gemma_path=gemma,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ dynamic_load=dynamic_load,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ distilled_lora_path=distilled_lora,
+ run_warmup=run_warmup,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ )
+
+ prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT)
+
+ def run(*, prompt, number, seed):
+ output_filename = os.environ.get("OUTPUT_PATH", f"ltx_av_two_stages_{width}x{height}_{number}.mp4")
+ logger.info(f"Running LTX AV Two-Stages: '{prompt[:80]}...'")
+ logger.info(f"Config: {height}x{width}, {num_frames} frames, {num_inference_steps} stage-1 steps")
+
+ if int(ttnn.distributed_context_get_rank()) != 0:
+ logger.info(f"Skipping generation on rank {ttnn.distributed_context_get_rank()}")
+ return
+
+ pipeline.generate(
+ prompt,
+ output_path=output_filename,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ num_inference_steps=num_inference_steps,
+ seed=seed,
+ )
+ logger.info(f"Saved video to: {output_filename}")
+ print_ltx_timing_table(
+ pipeline,
+ label="LTX TWO-STAGES",
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ mesh_shape=mesh_shape,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ topology=topology,
+ output_path=output_filename,
+ prompt=prompt,
+ )
+
+ if no_prompt:
+ run(prompt=prompt, number=0, seed=int(os.environ.get("SEED", "10")))
+ else:
+ for i in itertools.count():
+ new_prompt = input("Enter the input prompt, or q to exit: ")
+ if new_prompt:
+ prompt = new_prompt
+ if prompt[0] == "q":
+ break
+ run(prompt=prompt, number=i, seed=i)
+
+
+# ---------------------------------------------------------------------------
+# I2V smoke test (image + text -> video)
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis, num_links, device_params, topology, is_fsdp, dynamic_load",
+ LTX_TWO_STAGES_I2V_MESH_PARAMS_DL,
+ indirect=["mesh_device", "device_params"],
+)
+def test_pipeline_two_stages_i2v_smoke(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ dynamic_load,
+ topology,
+ is_fsdp,
+):
+ """End-to-end I2V smoke: generate(images=[(img, 0, 1.0)]) produces a non-empty MP4 whose
+ first frame correlates with the conditioning image.
+
+ Gated behind LTX_I2V_IMAGE (path to a conditioning image); skips otherwise so it never runs
+ in the default CI sweep without an input image / checkpoint present.
+ """
+ image_path = os.environ.get("LTX_I2V_IMAGE")
+ if not image_path or not os.path.exists(image_path):
+ pytest.skip("set LTX_I2V_IMAGE to a conditioning image path to run the I2V smoke test")
+
+ skip_if_unsupported_num_links(mesh_device, num_links)
+
+ ckpt = default_ltx_checkpoint("ltx-2.3-22b-dev.safetensors")
+ distilled_lora = _default_distilled_lora()
+ gemma = default_ltx_gemma()
+
+ parent_mesh = mesh_device
+ mesh_shape = tuple(parent_mesh.shape)
+ mesh_device = parent_mesh.create_submesh(ttnn.MeshShape(*mesh_shape))
+
+ num_frames = int(os.environ.get("NUM_FRAMES", "25"))
+ height = int(os.environ.get("HEIGHT", "512"))
+ width = int(os.environ.get("WIDTH", "768"))
+ num_inference_steps = int(os.environ.get("NUM_INFERENCE_STEPS", "8"))
+ strength = float(os.environ.get("LTX_I2V_STRENGTH", "1.0"))
+
+ pipeline = LTXTwoStagesPipeline.create_pipeline(
+ mesh_device=mesh_device,
+ checkpoint_name=ckpt,
+ gemma_path=gemma,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ dynamic_load=dynamic_load,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ distilled_lora_path=distilled_lora,
+ run_warmup=False,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ )
+
+ if int(ttnn.distributed_context_get_rank()) != 0:
+ logger.info(f"Skipping generation on rank {ttnn.distributed_context_get_rank()}")
+ return
+
+ output_filename = os.environ.get("OUTPUT_PATH", f"ltx_i2v_two_stages_{width}x{height}.mp4")
+ prompt = os.environ.get("PROMPT", DEFAULT_LTX_PROMPT)
+ pipeline.generate(
+ prompt,
+ output_path=output_filename,
+ images=[(image_path, 0, strength)],
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ num_inference_steps=num_inference_steps,
+ seed=int(os.environ.get("SEED", "10")),
+ )
+
+ assert os.path.exists(output_filename), f"no MP4 produced at {output_filename}"
+ assert os.path.getsize(output_filename) > 0, "produced MP4 is empty"
+ logger.info(f"I2V smoke: saved {output_filename} ({os.path.getsize(output_filename)} bytes)")
+
+ # Optional first-frame correlation check (only if imageio + PIL are available).
+ try:
+ import imageio.v3 as iio
+ import numpy as np
+ from PIL import Image
+
+ frames = iio.imread(output_filename, index=None, plugin="pyav") # (T, H, W, 3), in-process
+ first = frames[0].astype(np.float32) / 255.0
+ cond = Image.open(image_path).convert("RGB").resize((width, height))
+ cond_np = np.asarray(cond, dtype=np.float32) / 255.0
+ a = first.reshape(-1) - first.mean()
+ b = cond_np.reshape(-1) - cond_np.mean()
+ pcc = float((a @ b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8))
+ logger.info(f"I2V smoke: first-frame vs conditioning-image PCC = {pcc:.3f}")
+ if strength >= 1.0:
+ assert pcc > 0.3, f"first frame does not resemble the conditioning image (PCC={pcc:.3f})"
+ except ImportError:
+ logger.info("imageio/PIL unavailable — skipping first-frame correlation check")
diff --git a/code/models/tt_dit/tests/models/ltx/test_rope_ltx.py b/code/models/tt_dit/tests/models/ltx/test_rope_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..2fa9718e8d7588efeca38973af0fcd1d1275d9f1
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_rope_ltx.py
@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.mochi import get_rot_transformation_mat
+from models.tt_dit.utils.tensor import bf16_tensor, bf16_tensor_2dshard
+
+
+def make_indices_grid(F: int, H: int, W: int) -> torch.Tensor:
+ """3D position indices grid (1, n_dims=3, seq_len) in the layout diffusers' RoPE expects."""
+ t_ids = torch.arange(F)
+ h_ids = torch.arange(H)
+ w_ids = torch.arange(W)
+ grid_t, grid_h, grid_w = torch.meshgrid(t_ids, h_ids, w_ids, indexing="ij")
+ indices_grid = torch.stack([grid_t.flatten(), grid_h.flatten(), grid_w.flatten()], dim=0).float()
+ return indices_grid.unsqueeze(0) # (1, n_dims=3, seq_len)
+
+
+@pytest.mark.parametrize(
+ "mesh_device, sp_axis, tp_axis",
+ [
+ [(1, 1), 0, 1],
+ [(1, 2), 0, 1],
+ [(2, 1), 0, 1],
+ [(2, 2), 0, 1],
+ [(2, 4), 0, 1],
+ [(2, 4), 1, 0],
+ ],
+ ids=[
+ "1x1sp0tp1",
+ "1x2sp0tp1",
+ "2x1sp0tp1",
+ "2x2sp0tp1",
+ "2x4sp0tp1",
+ "2x4sp1tp0",
+ ],
+ indirect=["mesh_device"],
+)
+@pytest.mark.parametrize(
+ "F, H, W",
+ [
+ (5, 16, 28), # Small: 2240 tokens
+ (5, 30, 52), # 480p-ish: 7800 tokens
+ (5, 34, 64), # 2K spatial grid (1088x2048 latent H×W): 10880 tokens
+ ],
+ ids=["small", "480p", "2k"],
+)
+def test_ltx_rope_interleaved(mesh_device: ttnn.MeshDevice, sp_axis: int, tp_axis: int, F: int, H: int, W: int):
+ """Test LTX-2 interleaved RoPE against the diffusers reference (mirrors test_wan_rotary_pos_embed)."""
+ # double_precision=False keeps the freq grid fp32 to match the device path.
+ from diffusers.models.transformers.transformer_ltx2 import (
+ LTX2AudioVideoRotaryPosEmbed,
+ apply_interleaved_rotary_emb,
+ )
+
+ dim = 4096
+ num_heads = 32
+ head_dim = dim // num_heads # 128
+ theta = 10000.0
+ max_pos = [20, 2048, 2048]
+ B = 1
+
+ indices_grid = make_indices_grid(F, H, W)
+ seq_len = F * H * W
+ logger.info(f"indices_grid shape: {indices_grid.shape}, seq_len: {seq_len}")
+
+ # Reference cos/sin on CPU; raw indices (ndim==3) normalize by base_{frames,height,width}.
+ rope = LTX2AudioVideoRotaryPosEmbed(
+ dim=dim,
+ base_num_frames=max_pos[0],
+ base_height=max_pos[1],
+ base_width=max_pos[2],
+ theta=theta,
+ modality="video",
+ double_precision=False,
+ rope_type="interleaved",
+ num_attention_heads=num_heads,
+ )
+ # Interleaved cos/sin: (B, seq_len, dim)
+ cos_freq, sin_freq = rope(indices_grid)
+ logger.info(f"cos_freq shape: {cos_freq.shape}, sin_freq shape: {sin_freq.shape}")
+
+ # Create random input on the flat (B, seq_len, dim) layout the diffusers apply expects.
+ torch.manual_seed(42)
+ input_flat = torch.randn(B, seq_len, dim, dtype=torch.float32)
+
+ # PyTorch reference: diffusers applies interleaved RoPE on (B, N, dim).
+ output_ref_flat = apply_interleaved_rotary_emb(input_flat, (cos_freq, sin_freq))
+
+ # Reshape into (B, seq_len, num_heads, head_dim) for the device kernel + comparison.
+ input_tensor = input_flat.reshape(B, seq_len, num_heads, head_dim)
+ output_ref = output_ref_flat.reshape(B, seq_len, num_heads, head_dim)
+ cos_for_apply = cos_freq.reshape(B, seq_len, num_heads, head_dim)
+ sin_for_apply = sin_freq.reshape(B, seq_len, num_heads, head_dim)
+ logger.info(f"Reference output shape: {output_ref.shape}")
+
+ # rotary_embedding_llama expects (B, num_heads, seq_len, head_dim)
+ input_bhnd = input_tensor.permute(0, 2, 1, 3) # (B, H, N, D)
+ cos_bhnd = cos_for_apply.permute(0, 2, 1, 3) # (B, H, N, D)
+ sin_bhnd = sin_for_apply.permute(0, 2, 1, 3) # (B, H, N, D)
+
+ tt_input = bf16_tensor_2dshard(input_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+ # cos/sin sharded same as input: SP axis on sequence, TP axis on heads
+ tt_cos = bf16_tensor_2dshard(cos_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+ tt_sin = bf16_tensor_2dshard(sin_bhnd, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ # Apply RoPE on device
+ tt_output = ttnn.experimental.rotary_embedding_llama(tt_input, tt_cos, tt_sin, tt_trans_mat)
+
+ # Gather output back to host
+ concat_dims = [None, None]
+ concat_dims[sp_axis] = 2
+ concat_dims[tp_axis] = 1
+ tt_output_torch = ttnn.to_torch(
+ tt_output,
+ mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)),
+ )
+ tt_output_torch = tt_output_torch.permute(0, 2, 1, 3) # Back to (B, N, H, D)
+
+ # Compare
+ assert_quality(output_ref, tt_output_torch, pcc=0.99)
+ logger.info("PASSED: LTX interleaved RoPE matches PyTorch reference")
diff --git a/code/models/tt_dit/tests/models/ltx/test_transformer_ltx.py b/code/models/tt_dit/tests/models/ltx/test_transformer_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..c811a9407c01211b4ad0ff53a002be861dfbaab6
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_transformer_ltx.py
@@ -0,0 +1,2113 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""LTX-2.3 transformer unit tests; structure mirrors test_transformer_wan.py.
+
+Torch reference is diffusers' LTX2VideoTransformer{Block,3DModel}. The TT runtime fuses RoPE
+into rotary_embedding_llama (INTERLEAVED + trans_mat) with a load-time Q/K permute, so every
+TT freq feed is INTERLEAVED. Modality axis selects video vs av; set LTX_SKIP_PCC=1 to assert
+shape/finiteness only.
+"""
+
+import os
+import time
+
+import pytest
+import torch
+from loguru import logger
+from safetensors.torch import load_file
+from tracy import signpost
+
+import ttnn
+from models.tt_dit.models.transformers.ltx.quant_config import LtxQuantProfile
+from models.tt_dit.models.transformers.ltx.rope_ltx import LTXRopeType, precompute_freqs_cis
+from models.tt_dit.models.transformers.ltx.transformer_ltx import (
+ LTXTransformerBlock,
+ LTXTransformerModel,
+ build_audio_masks,
+ build_video_pad_mask,
+)
+from models.tt_dit.parallel.config import DiTParallelConfig, ParallelFactor
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.mochi import get_rot_transformation_mat
+from models.tt_dit.utils.patchifiers import AudioLatentShape, VideoPixelShape
+from models.tt_dit.utils.tensor import bf16_tensor, bf16_tensor_2dshard
+from models.tt_dit.utils.test import (
+ line_params_req_exact_devices,
+ ring_params_4k,
+ ring_params_8k,
+ skip_if_unsupported_num_links,
+)
+from models.tt_dit.utils.tracing import Tracer
+
+from .ltx_mesh_params import LTX_TRANSFORMER_MESH_PARAMS, _4x8sp1tp0nl2_ring_is_fsdp0
+
+# ---------------------------------------------------------------------------
+# LTX-2.3-22B distilled transformer configuration
+# ---------------------------------------------------------------------------
+DIM = 4096
+NUM_HEADS = 32
+HEAD_DIM = DIM // NUM_HEADS # 128
+IN_CHANNELS = 128
+OUT_CHANNELS = 128
+CTX_DIM = 4096
+AUDIO_DIM = 2048
+AUDIO_HEAD_DIM = AUDIO_DIM // NUM_HEADS # 64
+AUDIO_IN_CHANNELS = 128
+AUDIO_CTX_DIM = 2048
+EPS = 1e-6
+PROMPT_LEN = 32
+
+# RoPE
+ROPE_THETA = 10000.0
+VIDEO_ROPE_MAX_POS = [20, 2048, 2048]
+AUDIO_ROPE_MAX_POS = [20]
+CROSS_PE_MAX_POS = [20]
+
+# AV-mode auxiliary
+AUDIO_N = 256 # sp-aligned for sp ∈ {1, 2, 4, 8}
+# 22B checkpoint variants (AV mode only): "fast" = distilled-1.1 (production), "dev" = base 22B.
+_CHECKPOINT_22B_VARIANTS = {
+ "fast": os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-distilled-1.1.safetensors"),
+ "dev": os.path.expanduser("~/.cache/ltx-checkpoints/ltx-2.3-22b-dev.safetensors"),
+}
+
+
+# Fallback: the same checkpoints as cached by huggingface_hub (Lightricks/LTX-2.3 repo).
+_CHECKPOINT_22B_HF_FILENAMES = {
+ "fast": "ltx-2.3-22b-distilled-1.1.safetensors",
+ "dev": "ltx-2.3-22b-dev.safetensors",
+}
+
+
+def _resolve_checkpoint_22b(variant: str) -> str:
+ """Path for the 22B variant: LTX_CHECKPOINT env, then local cache, then huggingface_hub snapshot."""
+ env = os.environ.get("LTX_CHECKPOINT")
+ if env:
+ return env
+ local = _CHECKPOINT_22B_VARIANTS[variant]
+ if os.path.exists(local):
+ return local
+ snap_root = os.path.expanduser("~/.cache/huggingface/hub/models--Lightricks--LTX-2.3/snapshots")
+ if os.path.isdir(snap_root):
+ for snap in sorted(os.listdir(snap_root)):
+ cand = os.path.join(snap_root, snap, _CHECKPOINT_22B_HF_FILENAMES[variant])
+ if os.path.exists(cand):
+ return cand
+ return local
+
+
+# Block-test reference uses small scaled random weights; default init blows up the adaln chains.
+WEIGHT_SEED = 1234
+INPUT_SEED = 42
+TIMESTEP_VAL = 0.01
+
+# Toggle PCC verification via env (mirrors Wan's `dit_unit_test`). Default ON.
+_RUN_PCC_DEFAULT = {"1": False, "0": True}.get(os.environ.get("LTX_SKIP_PCC"), True)
+
+
+# ---------------------------------------------------------------------------
+# Parametrize lists
+# ---------------------------------------------------------------------------
+
+# 1080p fast-pipeline real latent grids (F, H_lat, W_lat) — the exact (latent_frames, h//32,
+# w//32) production feeds RoPE. These are NOT tile/SP-aligned, so the test SP-pads the sequence
+# (round F·H·W up to TILE * sp_factor) exactly like LTXPipeline._sp_pad_len, then crops back.
+_LTX_TRANSFORMER_SHAPE_PARAMS = [
+ pytest.param(19, 17, 30, id="stage_1"), # F·H·W = 9690 (real); SP-padded on device
+ pytest.param(19, 34, 60, id="stage_2"), # F·H·W = 38760 (real); SP-padded on device
+]
+
+_LTX_TRANSFORMER_MODALITY_PARAMS = [
+ pytest.param(False, id="video"),
+ pytest.param(True, id="av"),
+]
+
+_LTX_TRANSFORMER_RUN_PCC_PARAMS = [pytest.param(_RUN_PCC_DEFAULT, id="pcc" if _RUN_PCC_DEFAULT else "nopcc")]
+
+# Checkpoint variant for AV mode ("fast" = production default); ignored in video mode.
+_LTX_TRANSFORMER_CKPT_PARAMS = [
+ pytest.param("fast", id="ckpt_fast"),
+ pytest.param("dev", id="ckpt_dev"),
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def _interleaved_to_bhnd(t: torch.Tensor, num_heads: int) -> torch.Tensor:
+ """(B, N, dim) INTERLEAVED freqs → (B, num_heads, N, head_dim) for rotary_embedding_llama."""
+ B, N, dim = t.shape
+ head_dim = dim // num_heads
+ return t.reshape(B, N, num_heads, head_dim).permute(0, 2, 1, 3).contiguous()
+
+
+def _sp_pad_len(n_real: int, sp_factor: int) -> int:
+ """Round a seq len up to ttnn.TILE_SIZE * sp_factor (mirrors LTXPipeline._sp_pad_len)."""
+ divisor = ttnn.TILE_SIZE * sp_factor
+ return ((n_real + divisor - 1) // divisor) * divisor
+
+
+def _pad_seq_dim(t: torch.Tensor, n_pad: int, dim: int = -2) -> torch.Tensor:
+ """Zero-pad a torch tensor's sequence dim up to n_pad (no-op if already ≥ n_pad)."""
+ n = t.shape[dim]
+ if n_pad <= n:
+ return t
+ shape = list(t.shape)
+ shape[dim] = n_pad
+ out = torch.zeros(shape, dtype=t.dtype)
+ idx = [slice(None)] * t.ndim
+ idx[dim] = slice(0, n)
+ out[tuple(idx)] = t
+ return out
+
+
+def _pad_rope_bhnd(cos_i: torch.Tensor, sin_i: torch.Tensor, pad_to: int | None):
+ """Right-pad BHND INTERLEAVED freqs on the seq dim (2) with identity rotation (cos=1, sin=0).
+
+ Same convention as rope_ltx.pad_video_rope_sp: padded slots are no-op rotations; SDPA
+ still masks them out via logical_n / padding masks.
+ """
+ if pad_to is None or pad_to <= cos_i.shape[2]:
+ return cos_i, sin_i
+ pad = pad_to - cos_i.shape[2]
+ cos_i = torch.nn.functional.pad(cos_i, (0, 0, 0, pad), value=1.0)
+ sin_i = torch.nn.functional.pad(sin_i, (0, 0, 0, pad), value=0.0)
+ return cos_i, sin_i
+
+
+def _diffusers_qk_to_split(t: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor:
+ """diffusers (interleaved-rotation) Q/K → Lightricks SPLIT convention the TT loader expects."""
+ D = head_dim
+ D_half = D // 2
+ inv = torch.empty(D, dtype=torch.long)
+ inv[:D_half] = torch.arange(0, D, 2)
+ inv[D_half:] = torch.arange(1, D, 2)
+ rest = t.shape[1:]
+ return t.reshape(num_heads, D, *rest).index_select(1, inv).reshape(num_heads * D, *rest)
+
+
+def _convert_diffusers_video_block_to_tt(state: dict, *, num_heads: int, head_dim: int) -> dict:
+ """diffusers LTX2VideoTransformerBlock state_dict → TT video block loader input (video subset, split Q/K)."""
+ keep_exact = {"scale_shift_table", "prompt_scale_shift_table"}
+ out = {k: v for k, v in state.items() if k in keep_exact or k.startswith(("attn1.", "attn2.", "ff."))}
+ for base in ("attn1", "attn2"):
+ for proj in ("to_q", "to_k"):
+ for suf in ("weight", "bias"):
+ kk = f"{base}.{proj}.{suf}"
+ if kk in out:
+ out[kk] = _diffusers_qk_to_split(out[kk], num_heads, head_dim)
+ for nk in ("norm_q.weight", "norm_k.weight"):
+ kk = f"{base}.{nk}"
+ if kk in out:
+ out[kk] = _diffusers_qk_to_split(out[kk], num_heads, head_dim)
+ return out
+
+
+def _make_diffusers_video_block():
+ """Build the diffusers LTX-2 block in the video config used as the PCC reference."""
+ from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformerBlock
+
+ return LTX2VideoTransformerBlock(
+ dim=DIM,
+ num_attention_heads=NUM_HEADS,
+ attention_head_dim=HEAD_DIM,
+ cross_attention_dim=CTX_DIM,
+ audio_dim=AUDIO_DIM,
+ audio_num_attention_heads=NUM_HEADS,
+ audio_attention_head_dim=AUDIO_HEAD_DIM,
+ audio_cross_attention_dim=AUDIO_CTX_DIM,
+ video_gated_attn=False,
+ video_cross_attn_adaln=True,
+ audio_gated_attn=False,
+ audio_cross_attn_adaln=False,
+ eps=EPS,
+ rope_type="interleaved",
+ )
+
+
+def _diffusers_video_block_ref(block, *, x, context, temb, prompt_temb, cos_i, sin_i):
+ """Run the diffusers block video-only (a2v/v2a disabled) and return the video output."""
+ B, N, _ = x.shape
+ a_N = 32
+ audio_hidden = torch.zeros(B, a_N, AUDIO_DIM)
+ audio_enc = torch.zeros(B, context.shape[1], AUDIO_CTX_DIM)
+ temb_audio = torch.zeros(B, 1, 6 * AUDIO_DIM) # audio_cross_attn_adaln=False → 6 mod params
+ temb_prompt_audio = torch.zeros(B, 1, 2 * AUDIO_DIM)
+ with torch.no_grad():
+ out_v, _ = block(
+ hidden_states=x,
+ audio_hidden_states=audio_hidden,
+ encoder_hidden_states=context,
+ audio_encoder_hidden_states=audio_enc,
+ temb=temb,
+ temb_audio=temb_audio,
+ temb_ca_scale_shift=None,
+ temb_ca_audio_scale_shift=None,
+ temb_ca_gate=None,
+ temb_ca_audio_gate=None,
+ temb_prompt=prompt_temb,
+ temb_prompt_audio=temb_prompt_audio,
+ video_rotary_emb=(cos_i, sin_i),
+ audio_rotary_emb=None,
+ use_a2v_cross_attention=False,
+ use_v2a_cross_attention=False,
+ )
+ return out_v
+
+
+def _make_diffusers_video_model(num_layers: int = 1):
+ """diffusers LTX2VideoTransformer3DModel configured to match the TT video model."""
+ from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel
+
+ return LTX2VideoTransformer3DModel(
+ in_channels=IN_CHANNELS,
+ out_channels=OUT_CHANNELS,
+ num_attention_heads=NUM_HEADS,
+ attention_head_dim=HEAD_DIM,
+ cross_attention_dim=CTX_DIM,
+ audio_in_channels=AUDIO_IN_CHANNELS,
+ audio_out_channels=AUDIO_IN_CHANNELS,
+ audio_num_attention_heads=NUM_HEADS,
+ audio_attention_head_dim=AUDIO_HEAD_DIM,
+ audio_cross_attention_dim=AUDIO_CTX_DIM,
+ num_layers=num_layers,
+ cross_attn_mod=True,
+ audio_cross_attn_mod=False,
+ norm_eps=EPS,
+ rope_theta=10000.0,
+ rope_double_precision=False,
+ use_prompt_embeddings=False,
+ rope_type="interleaved",
+ )
+
+
+def _convert_diffusers_video_model_to_tt(state: dict, *, num_heads: int, head_dim: int) -> dict:
+ """diffusers LTX2VideoTransformer3DModel state_dict → TT LTXTransformerModel loader input."""
+ drop_top = ("audio_", "av_cross_attn_", "caption_projection")
+ rename_top = {"proj_in.": "patchify_proj.", "time_embed.": "adaln_single.", "prompt_adaln.": "prompt_adaln_single."}
+ out: dict = {}
+ for k, v in state.items():
+ if k.startswith("transformer_blocks."):
+ _, _idx, rest = k.split(".", 2)
+ if rest in ("scale_shift_table", "prompt_scale_shift_table") or rest.startswith(
+ ("attn1.", "attn2.", "ff.")
+ ):
+ out[k] = v
+ continue
+ if k.startswith(drop_top):
+ continue
+ nk = k
+ for a, b in rename_top.items():
+ if nk.startswith(a):
+ nk = b + nk[len(a) :]
+ break
+ out[nk] = v
+ qk_suffixes = ("to_q.weight", "to_q.bias", "to_k.weight", "to_k.bias", "norm_q.weight", "norm_k.weight")
+ for k in list(out.keys()):
+ if (".attn1." in k or ".attn2." in k) and k.endswith(qk_suffixes):
+ out[k] = _diffusers_qk_to_split(out[k], num_heads, head_dim)
+ return out
+
+
+def _diffusers_video_model_ref(model, *, video_lat, video_prompt, sigma_val, F, H, W):
+ """Run the diffusers 3D model video-only (isolate_modalities=True) and return video output."""
+ B, N, _ = video_lat.shape
+ gt, gh, gw, _ = _video_grid(F, H, W)
+ video_coords = torch.stack([gt, gh, gw], dim=0).float().unsqueeze(0) # (1, 3, N)
+ a_N = 64
+ audio_lat = torch.zeros(B, a_N, AUDIO_IN_CHANNELS)
+ audio_prompt = torch.zeros(B, video_prompt.shape[1], AUDIO_CTX_DIM)
+ audio_coords = torch.arange(a_N).reshape(1, 1, a_N).float()
+ ts = torch.tensor([sigma_val * 1000.0])
+ with torch.no_grad():
+ out = model(
+ hidden_states=video_lat,
+ audio_hidden_states=audio_lat,
+ encoder_hidden_states=video_prompt,
+ audio_encoder_hidden_states=audio_prompt,
+ timestep=ts,
+ audio_timestep=ts,
+ sigma=ts,
+ audio_sigma=ts,
+ video_coords=video_coords,
+ audio_coords=audio_coords,
+ isolate_modalities=True,
+ return_dict=False,
+ )
+ return out[0] if isinstance(out, (tuple, list)) else out
+
+
+def _diffusers_video_model_ref_pertoken(model, *, video_lat, video_prompt, video_ts_real, sigma_scalar, F, H, W):
+ """Per-token oracle: same as ``_diffusers_video_model_ref`` but ``timestep`` is ``(1, N)``.
+
+ The diffusers LTX2 model flattens ``timestep`` through ``time_embed`` then reshapes to
+ ``(B, N, ...)`` modulation, so a per-token timestep yields genuine per-token modulation.
+ ``sigma`` (prompt/cross modulation) stays scalar, matching the TT scalar-timestep path.
+ """
+ B, N, _ = video_lat.shape
+ gt, gh, gw, _ = _video_grid(F, H, W)
+ video_coords = torch.stack([gt, gh, gw], dim=0).float().unsqueeze(0) # (1, 3, N)
+ a_N = 64
+ audio_lat = torch.zeros(B, a_N, AUDIO_IN_CHANNELS)
+ audio_prompt = torch.zeros(B, video_prompt.shape[1], AUDIO_CTX_DIM)
+ audio_coords = torch.arange(a_N).reshape(1, 1, a_N).float()
+ ts_pertoken = video_ts_real.reshape(1, -1) * 1000.0 # (1, N)
+ ts_scalar = torch.tensor([sigma_scalar * 1000.0])
+ with torch.no_grad():
+ out = model(
+ hidden_states=video_lat,
+ audio_hidden_states=audio_lat,
+ encoder_hidden_states=video_prompt,
+ audio_encoder_hidden_states=audio_prompt,
+ timestep=ts_pertoken,
+ audio_timestep=ts_scalar,
+ sigma=ts_scalar,
+ audio_sigma=ts_scalar,
+ video_coords=video_coords,
+ audio_coords=audio_coords,
+ isolate_modalities=True,
+ return_dict=False,
+ )
+ return out[0] if isinstance(out, (tuple, list)) else out
+
+
+def _lightricks_qk_to_interleaved(t: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor:
+ """Forward per-head SPLIT→INTERLEAVED Q/K permute (same as LTXAttention._permute_qk)."""
+ D = head_dim
+ D_half = D // 2
+ perm = torch.empty(D, dtype=torch.long)
+ perm[0::2] = torch.arange(D_half)
+ perm[1::2] = torch.arange(D_half, D)
+ rest = t.shape[1:]
+ return t.reshape(num_heads, D, *rest).index_select(1, perm).reshape(num_heads * D, *rest)
+
+
+# Per-attention head_dim for the AV block (video self/text-cross use 128; audio + cross-modal use 64).
+_AV_ATTN_HEAD_DIM = {
+ "attn1": HEAD_DIM,
+ "attn2": HEAD_DIM,
+ "audio_attn1": AUDIO_HEAD_DIM,
+ "audio_attn2": AUDIO_HEAD_DIM,
+ "audio_to_video_attn": AUDIO_HEAD_DIM,
+ "video_to_audio_attn": AUDIO_HEAD_DIM,
+}
+
+
+def _convert_lightricks_av_to_diffusers(state: dict, *, num_heads: int) -> dict:
+ """Raw Lightricks AV state_dict → diffusers LTX2VideoTransformer3DModel (AV) names."""
+ rename_top = {
+ "patchify_proj.": "proj_in.",
+ "adaln_single.": "time_embed.",
+ "prompt_adaln_single.": "prompt_adaln.",
+ "audio_patchify_proj.": "audio_proj_in.",
+ "audio_adaln_single.": "audio_time_embed.",
+ "audio_prompt_adaln_single.": "audio_prompt_adaln.",
+ "av_ca_video_scale_shift_adaln_single.": "av_cross_attn_video_scale_shift.",
+ "av_ca_audio_scale_shift_adaln_single.": "av_cross_attn_audio_scale_shift.",
+ "av_ca_a2v_gate_adaln_single.": "av_cross_attn_video_a2v_gate.",
+ "av_ca_v2a_gate_adaln_single.": "av_cross_attn_audio_v2a_gate.",
+ }
+ block_table_rename = {
+ "scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table",
+ "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table",
+ }
+ out: dict = {}
+ for k, v in state.items():
+ if k.startswith("transformer_blocks."):
+ prefix, idx, rest = k.split(".", 2)
+ if rest in block_table_rename:
+ rest = block_table_rename[rest]
+ rest = rest.replace("q_norm", "norm_q").replace("k_norm", "norm_k")
+ out[f"{prefix}.{idx}.{rest}"] = v
+ else:
+ nk = k
+ for a, b in rename_top.items():
+ if nk.startswith(a):
+ nk = b + nk[len(a) :]
+ break
+ out[nk] = v
+ qk_suffixes = ("to_q.weight", "to_q.bias", "to_k.weight", "to_k.bias", "norm_q.weight", "norm_k.weight")
+ for k in list(out.keys()):
+ if not k.startswith("transformer_blocks."):
+ continue
+ attn = k.split(".")[2]
+ if attn in _AV_ATTN_HEAD_DIM and k.endswith(qk_suffixes):
+ out[k] = _lightricks_qk_to_interleaved(out[k], num_heads, _AV_ATTN_HEAD_DIM[attn])
+ return out
+
+
+def _make_diffusers_av_model(num_layers: int = 1):
+ """diffusers ``LTX2VideoTransformer3DModel`` configured to match the TT AV model (gated
+ attention, 9-param video+audio modulation, cross-modal a2v/v2a, no caption projection)."""
+ from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel
+
+ return LTX2VideoTransformer3DModel(
+ in_channels=IN_CHANNELS,
+ out_channels=OUT_CHANNELS,
+ num_attention_heads=NUM_HEADS,
+ attention_head_dim=HEAD_DIM,
+ cross_attention_dim=CTX_DIM,
+ audio_in_channels=AUDIO_IN_CHANNELS,
+ audio_out_channels=AUDIO_IN_CHANNELS,
+ audio_num_attention_heads=NUM_HEADS,
+ audio_attention_head_dim=AUDIO_HEAD_DIM,
+ audio_cross_attention_dim=AUDIO_CTX_DIM,
+ num_layers=num_layers,
+ gated_attn=True,
+ audio_gated_attn=True,
+ cross_attn_mod=True,
+ audio_cross_attn_mod=True,
+ norm_eps=EPS,
+ rope_theta=ROPE_THETA,
+ rope_double_precision=False,
+ use_prompt_embeddings=False,
+ rope_type="interleaved",
+ # =1 matches Lightricks' av_ca_timestep_scale_multiplier (a2v/v2a gate adaln sees raw sigma).
+ cross_attn_timestep_scale_multiplier=1,
+ )
+
+
+def _diffusers_av_model_ref(model, *, video_lat, video_prompt, audio_lat, audio_prompt, sigma_val, F, H, W, audio_N):
+ """Run the diffusers AV model (cross-modal enabled) and return (video, audio) outputs."""
+ gt, gh, gw, _ = _video_grid(F, H, W)
+ video_coords = torch.stack([gt, gh, gw], dim=0).float().unsqueeze(0) # (1, 3, N)
+ audio_coords = torch.arange(audio_N).reshape(1, 1, audio_N).float()
+ ts = torch.tensor([sigma_val * 1000.0])
+ with torch.no_grad():
+ out = model(
+ hidden_states=video_lat,
+ audio_hidden_states=audio_lat,
+ encoder_hidden_states=video_prompt,
+ audio_encoder_hidden_states=audio_prompt,
+ timestep=ts,
+ audio_timestep=ts,
+ sigma=ts,
+ audio_sigma=ts,
+ video_coords=video_coords,
+ audio_coords=audio_coords,
+ isolate_modalities=False,
+ return_dict=False,
+ )
+ return out[0], out[1]
+
+
+def _make_parallel_config(mesh_device, sp_axis, tp_axis):
+ return DiTParallelConfig(
+ cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
+ sequence_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[sp_axis], mesh_axis=sp_axis),
+ tensor_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[tp_axis], mesh_axis=tp_axis),
+ )
+
+
+def _make_ccl_manager(mesh_device, num_links, topology):
+ return CCLManager(mesh_device=mesh_device, num_links=num_links, topology=topology)
+
+
+def _video_grid(F, H, W):
+ """Return (grid_t, grid_h, grid_w) flat tensors and stacked 3D positions (1, N, 3)."""
+ t_ids, h_ids, w_ids = torch.arange(F), torch.arange(H), torch.arange(W)
+ gt, gh, gw = torch.meshgrid(t_ids, h_ids, w_ids, indexing="ij")
+ gt_f, gh_f, gw_f = gt.flatten(), gh.flatten(), gw.flatten()
+ positions = torch.stack([gt_f, gh_f, gw_f], dim=-1).float().unsqueeze(0)
+ return gt_f, gh_f, gw_f, positions
+
+
+def _video_rope_freqs(F, H, W, rope_type=LTXRopeType.SPLIT):
+ _, _, _, positions = _video_grid(F, H, W)
+ return precompute_freqs_cis(
+ positions,
+ dim=DIM,
+ out_dtype=torch.float32,
+ theta=ROPE_THETA,
+ max_pos=VIDEO_ROPE_MAX_POS,
+ num_attention_heads=NUM_HEADS,
+ rope_type=rope_type,
+ )
+
+
+def _video_cross_pe_freqs(F, H, W, rope_type=LTXRopeType.SPLIT):
+ """T-only cross-PE for V↔A, dim matches audio (AUDIO_DIM)."""
+ _, _, _, positions = _video_grid(F, H, W)
+ return precompute_freqs_cis(
+ positions[..., 0:1],
+ dim=AUDIO_DIM,
+ out_dtype=torch.float32,
+ theta=ROPE_THETA,
+ max_pos=CROSS_PE_MAX_POS,
+ num_attention_heads=NUM_HEADS,
+ rope_type=rope_type,
+ )
+
+
+def _audio_rope_freqs(audio_N, rope_type=LTXRopeType.SPLIT):
+ a_pos = torch.arange(audio_N).float().unsqueeze(0).unsqueeze(-1)
+ return precompute_freqs_cis(
+ a_pos,
+ dim=AUDIO_DIM,
+ out_dtype=torch.float32,
+ theta=ROPE_THETA,
+ max_pos=AUDIO_ROPE_MAX_POS,
+ num_attention_heads=NUM_HEADS,
+ rope_type=rope_type,
+ )
+
+
+def _audio_cross_pe_freqs(audio_N, rope_type=LTXRopeType.SPLIT):
+ a_pos = torch.arange(audio_N).float().unsqueeze(0).unsqueeze(-1)
+ return precompute_freqs_cis(
+ a_pos,
+ dim=AUDIO_DIM,
+ out_dtype=torch.float32,
+ theta=ROPE_THETA,
+ max_pos=CROSS_PE_MAX_POS,
+ num_attention_heads=NUM_HEADS,
+ rope_type=rope_type,
+ )
+
+
+def _audio_seq_lens(F: int, sp_factor: int) -> tuple[int, int]:
+ """Fast-pipeline audio seq lengths from the latent frame count (audio_N_real, padded to 32*sp)."""
+ num_frames = (F - 1) * 8 + 1
+ vps = VideoPixelShape(batch=1, frames=num_frames, height=64, width=64, fps=24)
+ audio_N_real = AudioLatentShape.from_video_pixel_shape(vps).frames
+ audio_N = ((audio_N_real + 32 * sp_factor - 1) // (32 * sp_factor)) * (32 * sp_factor)
+ return audio_N, audio_N_real
+
+
+def _tt_rope(freqs_fn, *args, mesh_device, sp_axis, tp_axis, pad_to=None):
+ """Build INTERLEAVED freqs for the TT runtime: precompute → BHND reshape → SP-pad → 2D shard."""
+ cos_i, sin_i = freqs_fn(*args, rope_type=LTXRopeType.INTERLEAVED)
+ cos_i = _interleaved_to_bhnd(cos_i, NUM_HEADS)
+ sin_i = _interleaved_to_bhnd(sin_i, NUM_HEADS)
+ cos_i, sin_i = _pad_rope_bhnd(cos_i, sin_i, pad_to)
+ tt_cos = bf16_tensor_2dshard(cos_i, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+ tt_sin = bf16_tensor_2dshard(sin_i, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 1})
+ return tt_cos, tt_sin
+
+
+def _tt_rope_full(freqs_fn, *args, mesh_device, tp_axis, pad_to=None):
+ """INTERLEAVED freqs replicated full-seq (cross-attn K side): BHND reshape, SP-pad, no SP shard."""
+ cos_i, sin_i = freqs_fn(*args, rope_type=LTXRopeType.INTERLEAVED)
+ cos_i = _interleaved_to_bhnd(cos_i, NUM_HEADS)
+ sin_i = _interleaved_to_bhnd(sin_i, NUM_HEADS)
+ cos_i, sin_i = _pad_rope_bhnd(cos_i, sin_i, pad_to)
+ tt_cos = bf16_tensor(cos_i, device=mesh_device, mesh_axis=tp_axis, shard_dim=1)
+ tt_sin = bf16_tensor(sin_i, device=mesh_device, mesh_axis=tp_axis, shard_dim=1)
+ return tt_cos, tt_sin
+
+
+def _make_tt_block(*, mesh_device, ccl_manager, parallel_config, is_fsdp, has_audio, quant_config=None):
+ return LTXTransformerBlock(
+ video_dim=DIM,
+ video_ffn_dim=DIM * 4,
+ video_num_heads=NUM_HEADS,
+ video_cross_attention_dim=CTX_DIM,
+ audio_dim=AUDIO_DIM,
+ audio_ffn_dim=AUDIO_DIM * 4,
+ audio_num_heads=NUM_HEADS,
+ audio_cross_attention_dim=AUDIO_CTX_DIM,
+ eps=EPS,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ apply_gated_attention=has_audio,
+ cross_attention_adaln=True,
+ quant_config=quant_config,
+ )
+
+
+def _make_tt_model(
+ *, mesh_device, ccl_manager, parallel_config, is_fsdp, has_audio, num_layers, image_conditioning=False
+):
+ return LTXTransformerModel(
+ num_attention_heads=NUM_HEADS,
+ attention_head_dim=HEAD_DIM,
+ in_channels=IN_CHANNELS,
+ out_channels=OUT_CHANNELS,
+ num_layers=num_layers,
+ cross_attention_dim=CTX_DIM,
+ audio_num_attention_heads=NUM_HEADS,
+ audio_attention_head_dim=AUDIO_HEAD_DIM,
+ audio_in_channels=AUDIO_IN_CHANNELS,
+ audio_out_channels=AUDIO_IN_CHANNELS,
+ audio_cross_attention_dim=AUDIO_CTX_DIM,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ apply_gated_attention=has_audio,
+ cross_attention_adaln=True,
+ image_conditioning=image_conditioning,
+ )
+
+
+def _load_22b_state_dict(num_layers: int, checkpoint_path: str) -> dict | None:
+ """Load LTX-2.3-22B weights filtered to the first num_layers blocks; None if checkpoint missing."""
+ if not os.path.exists(checkpoint_path):
+ return None
+ raw = load_file(checkpoint_path)
+ prefix = "model.diffusion_model."
+ sd = {k[len(prefix) :]: v for k, v in raw.items() if k.startswith(prefix)}
+ return {k: v for k, v in sd.items() if not k.startswith("transformer_blocks.") or int(k.split(".")[1]) < num_layers}
+
+
+def _scale_init_(module: torch.nn.Module, seed: int = WEIGHT_SEED) -> None:
+ """Deterministic N(0, 0.1²) reinit in-place. See block test comment for why."""
+ torch.manual_seed(seed)
+ with torch.no_grad():
+ for p in module.parameters():
+ p.copy_(torch.randn(p.shape, dtype=p.dtype) * 0.1)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ LTX_TRANSFORMER_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(("F", "H", "W"), _LTX_TRANSFORMER_SHAPE_PARAMS)
+@pytest.mark.parametrize("has_audio", _LTX_TRANSFORMER_MODALITY_PARAMS)
+@pytest.mark.parametrize("run_pcc", _LTX_TRANSFORMER_RUN_PCC_PARAMS)
+@pytest.mark.parametrize("checkpoint_variant", _LTX_TRANSFORMER_CKPT_PARAMS)
+def test_ltx_transformer_block(
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+ F: int,
+ H: int,
+ W: int,
+ has_audio: bool,
+ run_pcc: bool,
+ checkpoint_variant: str,
+ reset_seeds,
+) -> None:
+ """Test LTXTransformerBlock: TT forward, with optional PCC vs the diffusers LTX-2 block."""
+ # Checkpoint variant only affects AV weight loading; skip the redundant copy in video mode.
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ if not has_audio and checkpoint_variant != "fast":
+ pytest.skip("checkpoint_variant only affects AV mode (video uses random scaled weights)")
+ checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant)
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ # Real latent grid → SP-padded sequence (mirrors LTXPipeline). video_N is the padded device
+ # length; video_N_real is the logical token count fed to SDPA's logical_n / padding masks.
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+ # Audio seq length mirrors the fast pipeline so op sizes / padding match production.
+ audio_N, audio_N_real = _audio_seq_lens(F, sp_factor)
+ assert video_N % (32 * sp_factor) == 0, f"video_N={video_N} not sp/tile-aligned for sp={sp_factor}"
+ assert audio_N % (32 * sp_factor) == 0, f"audio_N={audio_N} not sp/tile-aligned for sp={sp_factor}"
+
+ # AV mode lacks the LTXModel-level cross-PE plumbing, so it asserts shape/finiteness only.
+ do_pcc = run_pcc and not has_audio
+ torch_block = None
+ torch_out = None
+ if do_pcc:
+ # Reference is the diffusers LTX-2 block run video-only (audio + a2v/v2a disabled).
+ torch_block = _make_diffusers_video_block()
+ torch_block.eval()
+ _scale_init_(torch_block)
+
+ # TT block
+ ccl_manager = _make_ccl_manager(mesh_device, num_links, topology)
+ parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis)
+ tt_block = _make_tt_block(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ )
+
+ if do_pcc:
+ # diffusers→TT: keep video subset, convert attn Q/K to split. Clone to avoid aliasing.
+ conv = _convert_diffusers_video_block_to_tt(torch_block.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM)
+ tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()})
+ elif has_audio:
+ sd = _load_22b_state_dict(num_layers=1, checkpoint_path=checkpoint_22b)
+ if sd is None:
+ pytest.skip(f"22B checkpoint not found at {checkpoint_22b}")
+ block_sd = {
+ k[len("transformer_blocks.0.") :]: v for k, v in sd.items() if k.startswith("transformer_blocks.0.")
+ }
+ tt_block.load_torch_state_dict(block_sd, strict=False)
+ else:
+ # video mode, PCC disabled — harvest a shape-matching state dict from a throwaway block.
+ dummy = _make_diffusers_video_block()
+ _scale_init_(dummy)
+ conv = _convert_diffusers_video_block_to_tt(dummy.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM)
+ tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()})
+ del dummy
+
+ # Inputs (real-grid token count; the TT side gets a zero-padded copy below).
+ torch.manual_seed(INPUT_SEED)
+ x = torch.randn(1, video_N_real, DIM, dtype=torch.float32)
+ context = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+ temb = torch.randn(1, 1, 9 * DIM, dtype=torch.float32) # 9 adaln params
+ prompt_temb = torch.randn(1, 1, 2 * DIM, dtype=torch.float32) # 2 adaln params for prompt
+
+ # Torch forward (video-only branch): diffusers block consumes INTERLEAVED cos/sin, same
+ # convention as the TT runtime (tt_dit interleaved freqs ≈ diffusers rope, PCC ~1.0).
+ if do_pcc:
+ cos_int, sin_int = _video_rope_freqs(F, H, W, rope_type=LTXRopeType.INTERLEAVED)
+ torch_out = _diffusers_video_block_ref(
+ torch_block, x=x, context=context, temb=temb, prompt_temb=prompt_temb, cos_i=cos_int, sin_i=sin_int
+ )
+ logger.info(f"torch block output {tuple(torch_out.shape)}")
+
+ # TT video tensors. Runtime rope is INTERLEAVED + trans_mat. Sequence is SP-padded to video_N;
+ # rope pads with identity rotation; video_N (logical) goes to the block for SDPA masking.
+ spatial = _pad_seq_dim(x, video_N, dim=1).unsqueeze(0)
+ tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3})
+ tt_prompt = bf16_tensor(context.unsqueeze(0), device=mesh_device)
+ # Outer-param layout (coeff, B=1, 1, D) so the block's chunk(dim=0) is a free tile-aligned slice.
+ tt_temb = bf16_tensor(
+ temb.reshape(9, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3
+ )
+ tt_prompt_temb = bf16_tensor(prompt_temb.reshape(2, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device)
+ tt_cos, tt_sin = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ forward_kwargs = dict(
+ video_1BND=tt_spatial,
+ video_prompt=tt_prompt,
+ video_temb=tt_temb,
+ video_N=video_N_real,
+ video_rope_cos=tt_cos,
+ video_rope_sin=tt_sin,
+ trans_mat=tt_trans_mat,
+ video_prompt_temb=tt_prompt_temb,
+ )
+ if has_audio:
+ # Real tokens in [:audio_N_real], zeros in the padded tail (matches the padded audio latent).
+ a_x = torch.zeros(1, audio_N, AUDIO_DIM, dtype=torch.float32)
+ a_x[:, :audio_N_real, :] = torch.randn(1, audio_N_real, AUDIO_DIM, dtype=torch.float32)
+ a_ctx = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32)
+ a_temb = torch.randn(1, 1, 9 * AUDIO_DIM, dtype=torch.float32)
+ a_prompt_temb = torch.randn(1, 1, 2 * AUDIO_DIM, dtype=torch.float32)
+ av_ca_v = torch.randn(1, 1, 5 * DIM, dtype=torch.float32) # 4 scale-shift + 1 gate
+ av_ca_a = torch.randn(1, 1, 5 * AUDIO_DIM, dtype=torch.float32)
+
+ a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis)
+ vx_cos, vx_sin = _tt_rope(
+ _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ ax_cos, ax_sin = _tt_rope(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis
+ )
+ ax_cos_full, ax_sin_full = _tt_rope_full(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis
+ )
+
+ # Padding masks, same construction as the fast pipeline (audio padded, video aligned).
+ a_attn_mask, a_pad_sp, a_pad_full = build_audio_masks(
+ audio_N, audio_N_real, mesh_device=mesh_device, sp_axis=sp_axis
+ )
+ v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis)
+
+ forward_kwargs.update(
+ audio_1BND=bf16_tensor_2dshard(
+ a_x.unsqueeze(0), device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3}
+ ),
+ audio_prompt=bf16_tensor(a_ctx.unsqueeze(0), device=mesh_device),
+ audio_temb=bf16_tensor(
+ a_temb.reshape(9, AUDIO_DIM).unsqueeze(1).unsqueeze(1),
+ device=mesh_device,
+ mesh_axis=tp_axis,
+ shard_dim=3,
+ ),
+ audio_prompt_temb=bf16_tensor(
+ a_prompt_temb.reshape(2, AUDIO_DIM).unsqueeze(1).unsqueeze(1), device=mesh_device
+ ),
+ av_ca_temb=bf16_tensor(
+ av_ca_v.reshape(5, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3
+ ),
+ av_ca_audio_temb=bf16_tensor(
+ av_ca_a.reshape(5, AUDIO_DIM).unsqueeze(1).unsqueeze(1),
+ device=mesh_device,
+ mesh_axis=tp_axis,
+ shard_dim=3,
+ ),
+ audio_N=audio_N,
+ audio_rope_cos=a_cos,
+ audio_rope_sin=a_sin,
+ video_cross_pe_cos=vx_cos,
+ video_cross_pe_sin=vx_sin,
+ audio_cross_pe_cos=ax_cos,
+ audio_cross_pe_sin=ax_sin,
+ audio_cross_pe_cos_full=ax_cos_full,
+ audio_cross_pe_sin_full=ax_sin_full,
+ audio_attn_mask=a_attn_mask,
+ audio_padding_mask=a_pad_sp,
+ audio_padding_mask_full=a_pad_full,
+ video_padding_mask=v_pad_sp,
+ )
+
+ # Signpost-bracket one warm forward for tt-perf-report --start/end-signpost.
+ for _ in range(2):
+ tt_out = tt_block(**forward_kwargs)
+ ttnn.synchronize_device(mesh_device)
+ signpost("start")
+ tt_out = tt_block(**forward_kwargs)
+ ttnn.synchronize_device(mesh_device)
+ signpost("stop")
+
+ if has_audio:
+ tt_v, tt_a = tt_out
+ else:
+ tt_v = tt_out
+
+ # Gather video output
+ concat_dims = [None, None]
+ concat_dims[sp_axis] = 2
+ concat_dims[tp_axis] = 3
+ tt_v_torch = ttnn.to_torch(
+ tt_v,
+ mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)),
+ ).squeeze(0)
+
+ # Crop the SP-padding tail off the gathered video before checking/comparing real tokens.
+ tt_v_torch = tt_v_torch[:, :video_N_real, :]
+ assert tt_v_torch.shape == (1, video_N_real, DIM), f"video shape {tt_v_torch.shape}"
+ assert torch.isfinite(tt_v_torch).all(), "video output NaN/Inf"
+
+ if has_audio:
+ tt_a_torch = ttnn.to_torch(
+ tt_a,
+ mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)),
+ ).squeeze(0)[:, :audio_N_real, :]
+ assert tt_a_torch.shape == (1, audio_N_real, AUDIO_DIM), f"audio shape {tt_a_torch.shape}"
+ assert torch.isfinite(tt_a_torch).all(), "audio output NaN/Inf"
+
+ if do_pcc:
+ # 4x8 mesh has 8-way SP ring all-gathers — looser tolerance.
+ pcc = 0.988 if mesh_device.get_num_devices() > 8 else 0.999
+ rmse = 0.10 if mesh_device.get_num_devices() > 8 else 0.032
+ assert_quality(torch_out, tt_v_torch, pcc=pcc, relative_rmse=rmse)
+ logger.info(f"PASSED block PCC: video {tuple(tt_v_torch.shape)}")
+ else:
+ logger.info(f"PASSED block (no PCC): video {tuple(tt_v_torch.shape)}")
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ [_4x8sp1tp0nl2_ring_is_fsdp0],
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_transformer_block_bf8(
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+ is_fsdp: bool,
+ reset_seeds,
+) -> None:
+ """Shipped bf8 DiT-linear quant path (``all_bf8_lofi``) on the production 4x8 ring.
+
+ Covers two things a bf16-only run cannot: (1) the profile actually produced bf8 device weights with
+ the ``to_out`` carve-out kept bf16 — a mis-wired role or a silent no-quant fails the dtype asserts;
+ (2) the bf8 forward still tracks the fp32 diffusers oracle. Video-only, matching the block test's
+ PCC path; the gate-role carve-out only exists under gated (audio) attention, so it is out of scope.
+ """
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ F, H, W = 19, 17, 30 # stage_1 latent grid
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+
+ torch_block = _make_diffusers_video_block()
+ torch_block.eval()
+ _scale_init_(torch_block)
+
+ ccl_manager = _make_ccl_manager(mesh_device, num_links, topology)
+ parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis)
+ tt_block = _make_tt_block(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=False,
+ quant_config=LtxQuantProfile.all_bf8_lofi(),
+ )
+ conv = _convert_diffusers_video_block_to_tt(torch_block.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM)
+ tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()})
+
+ # Assert the loaded device-tensor dtypes the profile is supposed to map: projections bf8, the
+ # fused-addcmul to_out carve-out bf16. A wrong role or a silent no-quant (bf16 everywhere) fails here.
+ def wdtype(linear):
+ return linear.weight.data.dtype
+
+ assert wdtype(tt_block.attn1.to_qkv) == ttnn.bfloat8_b, "self-attn qkv should be bf8"
+ assert wdtype(tt_block.attn1.to_out) == ttnn.bfloat16, "self-attn to_out carve-out should stay bf16"
+ assert wdtype(tt_block.attn2.to_q) == ttnn.bfloat8_b, "cross-attn q should be bf8"
+ assert wdtype(tt_block.attn2.to_kv) == ttnn.bfloat8_b, "cross-attn kv should be bf8"
+ assert wdtype(tt_block.attn2.to_out) == ttnn.bfloat16, "cross-attn to_out carve-out should stay bf16"
+ assert wdtype(tt_block.ffn.ff1) == ttnn.bfloat8_b, "ff1 should be bf8"
+ assert wdtype(tt_block.ffn.ff2) == ttnn.bfloat8_b, "ff2 should be bf8"
+
+ # Activation casts ride LTX_QUANT_ACTIVATIONS (default on): the projections and the bf16-weight
+ # to_out both narrow their activation to bf8. PCC can't catch a missing cast (bf16 activations only
+ # raise it), so assert it directly.
+ assert tt_block.attn1.to_qkv.activation_dtype == ttnn.bfloat8_b, "self-attn qkv activation should be bf8"
+ assert tt_block.attn1.to_out.activation_dtype == ttnn.bfloat8_b, "to_out consumes the shared bf8 activation"
+ assert tt_block.attn2.to_q.activation_dtype == ttnn.bfloat8_b, "cross-attn q activation should be bf8"
+
+ torch.manual_seed(INPUT_SEED)
+ x = torch.randn(1, video_N_real, DIM, dtype=torch.float32)
+ context = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+ temb = torch.randn(1, 1, 9 * DIM, dtype=torch.float32)
+ prompt_temb = torch.randn(1, 1, 2 * DIM, dtype=torch.float32)
+ cos_int, sin_int = _video_rope_freqs(F, H, W, rope_type=LTXRopeType.INTERLEAVED)
+ torch_out = _diffusers_video_block_ref(
+ torch_block, x=x, context=context, temb=temb, prompt_temb=prompt_temb, cos_i=cos_int, sin_i=sin_int
+ )
+
+ spatial = _pad_seq_dim(x, video_N, dim=1).unsqueeze(0)
+ tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3})
+ tt_prompt = bf16_tensor(context.unsqueeze(0), device=mesh_device)
+ tt_temb = bf16_tensor(
+ temb.reshape(9, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3
+ )
+ tt_prompt_temb = bf16_tensor(prompt_temb.reshape(2, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device)
+ tt_cos, tt_sin = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+ tt_out = tt_block(
+ video_1BND=tt_spatial,
+ video_prompt=tt_prompt,
+ video_temb=tt_temb,
+ video_N=video_N_real,
+ video_rope_cos=tt_cos,
+ video_rope_sin=tt_sin,
+ trans_mat=tt_trans_mat,
+ video_prompt_temb=tt_prompt_temb,
+ )
+
+ concat_dims = [None, None]
+ concat_dims[sp_axis] = 2
+ concat_dims[tp_axis] = 3
+ tt_v = ttnn.to_torch(
+ tt_out,
+ mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)),
+ ).squeeze(0)[:, :video_N_real, :]
+ assert tt_v.shape == (1, video_N_real, DIM), f"video shape {tt_v.shape}"
+ assert torch.isfinite(tt_v).all(), "bf8 video output NaN/Inf"
+
+ # bf8 measures PCC 99.93% / RMSE 4.2% on BH 4x8. Floor PCC at the bf16 block test's ring-noise
+ # tolerance (0.988), not tighter — the RMSE gate (0.08 vs observed 0.042) is the real quant-
+ # regression catch; a broken quant path craters PCC or blows up RMSE and trips here regardless.
+ assert_quality(torch_out, tt_v, pcc=0.988, relative_rmse=0.08)
+ logger.info(f"PASSED bf8 block PCC: video {tuple(tt_v.shape)}")
+
+
+def _run_inner_step(
+ *,
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ has_audio,
+ run_pcc,
+ checkpoint_variant: str,
+ use_forward_alias: bool,
+):
+ """Shared body for the model forward tests; use_forward_alias picks tt_model(...) vs .forward(...)."""
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ # Checkpoint variant only affects AV weight loading; skip the redundant copy in video mode.
+ if not has_audio and checkpoint_variant != "fast":
+ pytest.skip("checkpoint_variant only affects AV mode (video uses random scaled weights)")
+ checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant)
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ # Real latent grid → SP-padded sequence (mirrors LTXPipeline). video_N_real is the logical
+ # token count (SDPA logical_n / pad masks); video_N is the padded on-device length.
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+ audio_N = AUDIO_N
+ assert video_N % (32 * sp_factor) == 0
+ if has_audio:
+ assert audio_N % (32 * sp_factor) == 0
+
+ # === Reference model ===
+ do_pcc = run_pcc
+ ref_video = None
+ ref_audio = None
+ state_dict = None
+
+ if has_audio:
+ # AV PCC path: 22B with strict=False.
+ state_dict = _load_22b_state_dict(num_layers=1, checkpoint_path=checkpoint_22b)
+ if state_dict is None:
+ pytest.skip(f"22B checkpoint not found at {checkpoint_22b}")
+ else:
+ # Video PCC path: random scaled weights (1 layer) from the diffusers 3D model.
+ torch_model = _make_diffusers_video_model(num_layers=1)
+ torch_model.eval()
+ _scale_init_(torch_model)
+ state_dict = _convert_diffusers_video_model_to_tt(
+ torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM
+ )
+ state_dict = {k: v.detach().clone() for k, v in state_dict.items()}
+
+ # === Inputs ===
+ # Real-grid latent for the torch reference; the TT side gets a zero-padded copy (video_N).
+ torch.manual_seed(INPUT_SEED)
+ video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32)
+ video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1)
+ video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+ sigma_val = TIMESTEP_VAL if not has_audio else 0.5 # AV path uses production sigma=0.5
+ timestep_torch = torch.tensor([sigma_val])
+
+ audio_lat = None
+ audio_prompt = None
+ if has_audio:
+ audio_lat = torch.randn(1, audio_N, AUDIO_IN_CHANNELS, dtype=torch.float32)
+ audio_prompt = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32)
+
+ # === Reference forward (before TT to avoid weight aliasing) ===
+ # The torch LTXModel computes rope internally from positions; only the TT side builds freqs.
+ if do_pcc and not has_audio:
+ ref_video = _diffusers_video_model_ref(
+ torch_model, video_lat=video_lat_real, video_prompt=video_prompt, sigma_val=sigma_val, F=F, H=H, W=W
+ )
+ logger.info(f"ref video {tuple(ref_video.shape)} range=[{ref_video.min():.3f}, {ref_video.max():.3f}]")
+ del torch_model
+
+ if do_pcc and has_audio:
+ # AV reference: load the raw 22B into the diffusers AV model and run the cross-modal forward.
+ ref_model = _make_diffusers_av_model(num_layers=1)
+ ref_model.load_state_dict(_convert_lightricks_av_to_diffusers(state_dict, num_heads=NUM_HEADS), strict=False)
+ ref_model.eval()
+ t0 = time.time()
+ ref_video, ref_audio = _diffusers_av_model_ref(
+ ref_model,
+ video_lat=video_lat_real,
+ video_prompt=video_prompt,
+ audio_lat=audio_lat,
+ audio_prompt=audio_prompt,
+ sigma_val=sigma_val,
+ F=F,
+ H=H,
+ W=W,
+ audio_N=audio_N,
+ )
+ logger.info(
+ f"AV ref forward: {time.time() - t0:.1f}s — video {tuple(ref_video.shape)} audio {tuple(ref_audio.shape)}"
+ )
+ del ref_model
+ import gc
+
+ gc.collect()
+
+ # === TT model ===
+ ccl_manager = _make_ccl_manager(mesh_device, num_links, topology)
+ parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis)
+ tt_model = _make_tt_model(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ num_layers=1,
+ )
+ t0 = time.time()
+ tt_model.load_torch_state_dict(state_dict, strict=not has_audio)
+ logger.info(f"state-dict load: {time.time() - t0:.1f}s")
+
+ # === Build TT-side RoPE / cross-PE / prompt tensors (INTERLEAVED + trans_mat) ===
+ tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device)
+ tt_vc, tt_vs = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ call_kwargs = dict(
+ video_1BNI_torch=video_lat.unsqueeze(0),
+ video_prompt_1BLP=tt_video_prompt,
+ video_rope_cos=tt_vc,
+ video_rope_sin=tt_vs,
+ video_N=video_N_real,
+ trans_mat=tt_trans_mat,
+ timestep_torch=timestep_torch,
+ )
+ if has_audio:
+ a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis)
+ vx_cos, vx_sin = _tt_rope(
+ _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ ax_cos, ax_sin = _tt_rope(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis
+ )
+ ax_cos_full, ax_sin_full = _tt_rope_full(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis
+ )
+ # Zero padded video tokens as V→A cross-attention keys (matches build_video_pad_mask).
+ v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis)
+ call_kwargs.update(
+ audio_1BNI_torch=audio_lat.unsqueeze(0),
+ audio_prompt_1BLP=bf16_tensor(audio_prompt.unsqueeze(0), device=mesh_device),
+ audio_rope_cos=a_cos,
+ audio_rope_sin=a_sin,
+ audio_N=audio_N,
+ video_cross_pe_cos=vx_cos,
+ video_cross_pe_sin=vx_sin,
+ audio_cross_pe_cos=ax_cos,
+ audio_cross_pe_sin=ax_sin,
+ audio_cross_pe_cos_full=ax_cos_full,
+ audio_cross_pe_sin_full=ax_sin_full,
+ video_padding_mask=v_pad_sp,
+ )
+
+ # === Forward ===
+ t0 = time.time()
+ result = tt_model(**call_kwargs) if use_forward_alias else tt_model.forward(**call_kwargs)
+ logger.info(f"TT forward: {time.time() - t0:.1f}s")
+
+ # Crop the SP-padding tail off the gathered video output before checking/comparing real tokens.
+ if has_audio:
+ tt_v_dev, tt_a_dev = result
+ tt_video = LTXTransformerModel.device_to_host(tt_v_dev).squeeze(0)[:, :video_N_real, :]
+ tt_audio = LTXTransformerModel.device_to_host(tt_a_dev).squeeze(0)
+ assert tt_video.shape == (1, video_N_real, OUT_CHANNELS), f"video shape {tt_video.shape}"
+ assert tt_audio.shape == (1, audio_N, AUDIO_IN_CHANNELS), f"audio shape {tt_audio.shape}"
+ assert torch.isfinite(tt_video).all() and torch.isfinite(tt_audio).all(), "NaN/Inf in TT output"
+ else:
+ tt_video = LTXTransformerModel.device_to_host(result).squeeze(0)[:, :video_N_real, :]
+ assert tt_video.shape == (1, video_N_real, OUT_CHANNELS), f"video shape {tt_video.shape}"
+ assert torch.isfinite(tt_video).all(), "NaN/Inf in TT output"
+
+ del tt_model
+
+ if do_pcc:
+ assert_quality(ref_video, tt_video, pcc=0.992, relative_rmse=0.15)
+ if has_audio:
+ assert_quality(ref_audio, tt_audio, pcc=0.992, relative_rmse=0.15)
+ logger.info("PASSED PCC")
+ else:
+ logger.info("PASSED (no PCC)")
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ LTX_TRANSFORMER_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(("F", "H", "W"), _LTX_TRANSFORMER_SHAPE_PARAMS)
+@pytest.mark.parametrize("has_audio", _LTX_TRANSFORMER_MODALITY_PARAMS)
+@pytest.mark.parametrize("run_pcc", _LTX_TRANSFORMER_RUN_PCC_PARAMS)
+@pytest.mark.parametrize("checkpoint_variant", _LTX_TRANSFORMER_CKPT_PARAMS)
+def test_ltx_transformer_model(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ has_audio,
+ run_pcc,
+ checkpoint_variant,
+ reset_seeds,
+) -> None:
+ """Full LTXTransformerModel forward via ``__call__``."""
+ _run_inner_step(
+ mesh_device=mesh_device,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ F=F,
+ H=H,
+ W=W,
+ has_audio=has_audio,
+ run_pcc=run_pcc,
+ checkpoint_variant=checkpoint_variant,
+ use_forward_alias=True,
+ )
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ LTX_TRANSFORMER_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(("F", "H", "W"), _LTX_TRANSFORMER_SHAPE_PARAMS)
+@pytest.mark.parametrize("has_audio", _LTX_TRANSFORMER_MODALITY_PARAMS)
+@pytest.mark.parametrize("run_pcc", _LTX_TRANSFORMER_RUN_PCC_PARAMS)
+@pytest.mark.parametrize("checkpoint_variant", _LTX_TRANSFORMER_CKPT_PARAMS)
+def test_ltx_transformer_inner_step(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ has_audio,
+ run_pcc,
+ checkpoint_variant,
+ reset_seeds,
+) -> None:
+ """LTXTransformerModel.forward — explicit call, denoising-loop path."""
+ _run_inner_step(
+ mesh_device=mesh_device,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ F=F,
+ H=H,
+ W=W,
+ has_audio=has_audio,
+ run_pcc=run_pcc,
+ checkpoint_variant=checkpoint_variant,
+ use_forward_alias=False,
+ )
+
+
+# ---------------------------------------------------------------------------
+# Per-token video timestep (I2V) — equivalence to the scalar path
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ [pytest.param((2, 4), 1, 0, 2, line_params_req_exact_devices, ttnn.Topology.Linear, False, id="2x4sp1tp0")],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 17, 30, id="stage_1")])
+def test_ltx_per_token_timestep_equivalence(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ reset_seeds,
+) -> None:
+ """A uniform per-token video timestep must reproduce the scalar-timestep path.
+
+ This is the weight-agnostic invariant behind I2V: with ``image_conditioning=True`` and a
+ per-token ``video_timestep = sigma`` for every token, the model output must match the
+ ``image_conditioning=False`` scalar path (T2V/AV stays bit-similar). The frame-0 pinning
+ behaviour itself is covered end-to-end by the pipeline; here we isolate the DiT plumbing.
+ """
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+
+ # Shared random-scaled weights (1 layer) from the diffusers 3D video model.
+ torch_model = _make_diffusers_video_model(num_layers=1)
+ torch_model.eval()
+ _scale_init_(torch_model)
+ state_dict = _convert_diffusers_video_model_to_tt(torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM)
+ state_dict = {k: v.detach().clone() for k, v in state_dict.items()}
+ del torch_model
+
+ torch.manual_seed(INPUT_SEED)
+ video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32)
+ video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1)
+ video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+ sigma_val = TIMESTEP_VAL
+ timestep_torch = torch.tensor([sigma_val])
+
+ ccl_manager = _make_ccl_manager(mesh_device, num_links, topology)
+ parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis)
+
+ # Shared TT-side RoPE / prompt tensors.
+ tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device)
+ tt_vc, tt_vs = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+ base_kwargs = dict(
+ video_1BNI_torch=video_lat.unsqueeze(0),
+ video_prompt_1BLP=tt_video_prompt,
+ video_rope_cos=tt_vc,
+ video_rope_sin=tt_vs,
+ video_N=video_N_real,
+ trans_mat=tt_trans_mat,
+ timestep_torch=timestep_torch,
+ )
+
+ # Scalar path (image_conditioning=False).
+ scalar_model = _make_tt_model(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=False,
+ num_layers=1,
+ image_conditioning=False,
+ )
+ scalar_model.load_torch_state_dict(state_dict, strict=True)
+ out_scalar = LTXTransformerModel.device_to_host(scalar_model.forward(**base_kwargs)).squeeze(0)[:, :video_N_real, :]
+ del scalar_model
+
+ # Per-token path (image_conditioning=True) with a uniform timestep = sigma over the padded grid.
+ pertoken_model = _make_tt_model(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=False,
+ num_layers=1,
+ image_conditioning=True,
+ )
+ pertoken_model.load_torch_state_dict(state_dict, strict=True)
+ video_timestep_torch = torch.full((video_N,), sigma_val, dtype=torch.float32)
+ out_pertoken = LTXTransformerModel.device_to_host(
+ pertoken_model.forward(video_timestep_torch=video_timestep_torch, **base_kwargs)
+ ).squeeze(0)[:, :video_N_real, :]
+ del pertoken_model
+
+ assert out_pertoken.shape == out_scalar.shape, f"{out_pertoken.shape} vs {out_scalar.shape}"
+ assert torch.isfinite(out_pertoken).all(), "NaN/Inf in per-token output"
+ # Same weights + uniform sigma: the per-token MLP collapses to the scalar broadcast (bf16-equal).
+ assert_quality(out_scalar, out_pertoken, pcc=0.999, relative_rmse=0.02)
+ logger.info("PASSED: uniform per-token timestep reproduces the scalar path")
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ [pytest.param((2, 4), 1, 0, 2, line_params_req_exact_devices, ttnn.Topology.Linear, False, id="2x4sp1tp0")],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 17, 30, id="stage_1")])
+def test_ltx_per_token_timestep_nonuniform(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ reset_seeds,
+) -> None:
+ """NON-uniform per-token timestep (the real I2V case) must match the diffusers per-token oracle.
+
+ This is the case the existing equivalence test does NOT cover: frame-0 tokens at sigma=0
+ (the pinned image anchor) and every other token at a high sigma. If the per-token gate does
+ not reach the fused attn/FFN epilogues, non-frame-0 tokens get the wrong (frame-0) gate and
+ the output diverges from the oracle — the DiT-level signature of "frame 0 ok, rest collapses".
+ """
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+ frame0_tokens = H * W # token order is f*H*W + h*W + w, so frame 0 == first H*W tokens
+ sigma_high = float(os.environ.get("LTX_TEST_SIGMA", "0.7"))
+
+ # Shared random-scaled weights (1 layer) from the diffusers 3D video model.
+ torch_model = _make_diffusers_video_model(num_layers=1)
+ torch_model.eval()
+ _scale_init_(torch_model)
+ state_dict = _convert_diffusers_video_model_to_tt(torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM)
+ state_dict = {k: v.detach().clone() for k, v in state_dict.items()}
+
+ torch.manual_seed(INPUT_SEED)
+ video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32)
+ video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1)
+ video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+
+ # Per-token timestep: frame-0 tokens pinned at sigma=0, everything else (incl. SP padding) at sigma_high.
+ video_ts_real = torch.full((video_N_real,), sigma_high, dtype=torch.float32)
+ video_ts_real[:frame0_tokens] = 0.0
+ video_timestep_torch = torch.full((video_N,), sigma_high, dtype=torch.float32)
+ video_timestep_torch[:frame0_tokens] = 0.0
+
+ # === Reference (per-token oracle) — before TT to avoid weight aliasing ===
+ ref_video = _diffusers_video_model_ref_pertoken(
+ torch_model,
+ video_lat=video_lat_real,
+ video_prompt=video_prompt,
+ video_ts_real=video_ts_real,
+ sigma_scalar=sigma_high,
+ F=F,
+ H=H,
+ W=W,
+ )
+ del torch_model
+
+ ccl_manager = _make_ccl_manager(mesh_device, num_links, topology)
+ parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis)
+ tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device)
+ tt_vc, tt_vs = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ model = _make_tt_model(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=False,
+ num_layers=1,
+ image_conditioning=True,
+ )
+ model.load_torch_state_dict(state_dict, strict=True)
+ out_tt = LTXTransformerModel.device_to_host(
+ model.forward(
+ video_1BNI_torch=video_lat.unsqueeze(0),
+ video_prompt_1BLP=tt_video_prompt,
+ video_rope_cos=tt_vc,
+ video_rope_sin=tt_vs,
+ video_N=video_N_real,
+ trans_mat=tt_trans_mat,
+ timestep_torch=torch.tensor([sigma_high]),
+ video_timestep_torch=video_timestep_torch,
+ )
+ ).squeeze(0)[:, :video_N_real, :]
+ del model
+
+ assert out_tt.shape == ref_video.shape, f"{out_tt.shape} vs {ref_video.shape}"
+ assert torch.isfinite(out_tt).all(), "NaN/Inf in per-token output"
+
+ # Localize frame-0 (pinned, sigma=0) vs the rest (sigma_high): a per-token gate that doesn't
+ # reach the fused epilogues makes the "rest" track the oracle far worse than frame 0.
+ def _pcc(a, b):
+ a32, b32 = a.float().flatten(), b.float().flatten()
+ return torch.corrcoef(torch.stack([a32, b32]))[0, 1].item()
+
+ pcc_all = _pcc(ref_video, out_tt)
+ pcc_f0 = _pcc(ref_video[:, :frame0_tokens, :], out_tt[:, :frame0_tokens, :])
+ pcc_rest = _pcc(ref_video[:, frame0_tokens:, :], out_tt[:, frame0_tokens:, :])
+ logger.info(f"non-uniform per-token (TT vs oracle): all={pcc_all:.5f} frame0={pcc_f0:.5f} rest={pcc_rest:.5f}")
+
+ assert_quality(ref_video, out_tt, pcc=0.99, relative_rmse=0.03)
+
+
+# --------------------------------------------------------------------------- Trace-based perf: VIDEO stage-2 ring
+def _build_video_trace_setup(
+ *,
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ has_audio=False,
+ checkpoint_variant="fast",
+ num_layers=1,
+):
+ """Build the TT model + device-only ``inner_step`` kwargs for the trace-perf test.
+
+ Mirrors the TT setup in ``_run_inner_step`` and reproduces ``LTXTransformerModel.forward``'s
+ host→device upload preamble (which ``inner_step`` normally receives), so the returned ttnn
+ kwargs can be captured/replayed directly under ttnn trace. Returns
+ ``(tt_model, inner_kwargs, video_N_real)``. Perf-only; no PCC.
+
+ Video mode uses random scaled weights. AV mode (``has_audio=True``) loads the 22B checkpoint
+ (``strict=False``) like the AV PCC path and adds the audio device tensors (audio latent,
+ audio/cross-modal rope, video pad mask) so a2v / audio shapes actually fire under the trace.
+ """
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+ assert video_N % (32 * sp_factor) == 0
+ audio_N = AUDIO_N
+ if has_audio:
+ assert audio_N % (32 * sp_factor) == 0
+
+ if has_audio:
+ # AV weights come from the 22B checkpoint (strict=False), same as the AV PCC path.
+ checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant)
+ state_dict = _load_22b_state_dict(num_layers=num_layers, checkpoint_path=checkpoint_22b)
+ if state_dict is None:
+ pytest.skip(f"22B checkpoint not found at {checkpoint_22b}")
+ else:
+ # Random scaled weights from the diffusers 3D video model — the video PCC path.
+ torch_model = _make_diffusers_video_model(num_layers=num_layers)
+ torch_model.eval()
+ _scale_init_(torch_model)
+ state_dict = _convert_diffusers_video_model_to_tt(
+ torch_model.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM
+ )
+ state_dict = {k: v.detach().clone() for k, v in state_dict.items()}
+ del torch_model
+
+ # Real-grid latent; the TT side gets a zero-padded copy (video_N).
+ torch.manual_seed(INPUT_SEED)
+ video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32)
+ video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1)
+ video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+ sigma_val = 0.5 if has_audio else TIMESTEP_VAL # AV path uses production sigma=0.5
+ timestep_torch = torch.tensor([sigma_val])
+
+ audio_lat = None
+ audio_prompt = None
+ if has_audio:
+ audio_lat = torch.randn(1, audio_N, AUDIO_IN_CHANNELS, dtype=torch.float32)
+ audio_prompt = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32)
+
+ ccl_manager = _make_ccl_manager(mesh_device, num_links, topology)
+ parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis)
+ tt_model = _make_tt_model(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ num_layers=num_layers,
+ )
+ t0 = time.time()
+ tt_model.load_torch_state_dict(state_dict, strict=not has_audio)
+ logger.info(f"state-dict load: {time.time() - t0:.1f}s")
+
+ # Stage-constant device tensors (INTERLEAVED rope + trans_mat + prompt).
+ tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device)
+ tt_vc, tt_vs = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ # Per-step device inputs — reproduce LTXTransformerModel.forward's upload for the video path.
+ video_1BNI_torch = video_lat.unsqueeze(0) # (1, B=1, video_N, IN_CHANNELS)
+ B_size = video_1BNI_torch.shape[1]
+ video_1BNI = bf16_tensor(video_1BNI_torch, device=mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ timestep = bf16_tensor(timestep_torch.reshape(1, 1, B_size, 1) * 1000.0, device=mesh_device)
+
+ # inner_step is keyword-only; these are exactly the device tensors it consumes.
+ inner_kwargs = dict(
+ video_1BNI=video_1BNI,
+ timestep=timestep,
+ video_prompt_1BLP=tt_video_prompt,
+ video_rope_cos=tt_vc,
+ video_rope_sin=tt_vs,
+ video_N=video_N_real,
+ trans_mat=tt_trans_mat,
+ )
+
+ if has_audio:
+ # Audio + cross-modal device tensors, mirroring the AV branch of _run_inner_step.
+ a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis)
+ vx_cos, vx_sin = _tt_rope(
+ _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ ax_cos, ax_sin = _tt_rope(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis
+ )
+ ax_cos_full, ax_sin_full = _tt_rope_full(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis
+ )
+ v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis)
+ audio_1BNI = bf16_tensor(audio_lat.unsqueeze(0), device=mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ inner_kwargs.update(
+ audio_1BNI=audio_1BNI,
+ audio_prompt_1BLP=bf16_tensor(audio_prompt.unsqueeze(0), device=mesh_device),
+ audio_rope_cos=a_cos,
+ audio_rope_sin=a_sin,
+ audio_N=audio_N,
+ video_cross_pe_cos=vx_cos,
+ video_cross_pe_sin=vx_sin,
+ audio_cross_pe_cos=ax_cos,
+ audio_cross_pe_sin=ax_sin,
+ audio_cross_pe_cos_full=ax_cos_full,
+ audio_cross_pe_sin_full=ax_sin_full,
+ video_padding_mask=v_pad_sp,
+ )
+ return tt_model, inner_kwargs, video_N_real
+
+
+def _build_trace_inner_kwargs(*, mesh_device, sp_axis, tp_axis, F, H, W, has_audio):
+ """Build ONLY the device inner_step kwargs for a given (F,H,W) — no model / CCLManager.
+
+ Split out so ONE model + CCLManager can be driven at multiple shapes (e.g. stage-1 then
+ stage-2) to reproduce the e2e's cross-stage device/semaphore-pool state. Mirrors the per-shape
+ tensor setup inside _build_video_trace_setup. Returns (inner_kwargs, video_N_real).
+ """
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+ assert video_N % (32 * sp_factor) == 0
+ audio_N = AUDIO_N
+ if has_audio:
+ assert audio_N % (32 * sp_factor) == 0
+
+ torch.manual_seed(INPUT_SEED)
+ video_lat_real = torch.randn(1, video_N_real, IN_CHANNELS, dtype=torch.float32)
+ video_lat = _pad_seq_dim(video_lat_real, video_N, dim=1)
+ video_prompt = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+ sigma_val = 0.5 if has_audio else TIMESTEP_VAL
+ timestep_torch = torch.tensor([sigma_val])
+
+ audio_lat = None
+ audio_prompt = None
+ if has_audio:
+ audio_lat = torch.randn(1, audio_N, AUDIO_IN_CHANNELS, dtype=torch.float32)
+ audio_prompt = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32)
+
+ tt_video_prompt = bf16_tensor(video_prompt.unsqueeze(0), device=mesh_device)
+ tt_vc, tt_vs = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ video_1BNI_torch = video_lat.unsqueeze(0)
+ B_size = video_1BNI_torch.shape[1]
+ video_1BNI = bf16_tensor(video_1BNI_torch, device=mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ timestep = bf16_tensor(timestep_torch.reshape(1, 1, B_size, 1) * 1000.0, device=mesh_device)
+
+ inner_kwargs = dict(
+ video_1BNI=video_1BNI,
+ timestep=timestep,
+ video_prompt_1BLP=tt_video_prompt,
+ video_rope_cos=tt_vc,
+ video_rope_sin=tt_vs,
+ video_N=video_N_real,
+ trans_mat=tt_trans_mat,
+ )
+ if has_audio:
+ a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis)
+ vx_cos, vx_sin = _tt_rope(
+ _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ ax_cos, ax_sin = _tt_rope(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis
+ )
+ ax_cos_full, ax_sin_full = _tt_rope_full(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis
+ )
+ v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis)
+ audio_1BNI = bf16_tensor(audio_lat.unsqueeze(0), device=mesh_device, mesh_axis=sp_axis, shard_dim=-2)
+ inner_kwargs.update(
+ audio_1BNI=audio_1BNI,
+ audio_prompt_1BLP=bf16_tensor(audio_prompt.unsqueeze(0), device=mesh_device),
+ audio_rope_cos=a_cos,
+ audio_rope_sin=a_sin,
+ audio_N=audio_N,
+ video_cross_pe_cos=vx_cos,
+ video_cross_pe_sin=vx_sin,
+ audio_cross_pe_cos=ax_cos,
+ audio_cross_pe_sin=ax_sin,
+ audio_cross_pe_cos_full=ax_cos_full,
+ audio_cross_pe_sin_full=ax_sin_full,
+ video_padding_mask=v_pad_sp,
+ )
+ return inner_kwargs, video_N_real
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ [
+ pytest.param((4, 8), 1, 0, 2, ring_params_4k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_4k"),
+ # 8k fabric payload variant (where stage-2 wedges at depth).
+ pytest.param((4, 8), 1, 0, 2, ring_params_8k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_8k"),
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize("has_audio", [pytest.param(False, id="video"), pytest.param(True, id="av")])
+@pytest.mark.parametrize("checkpoint_variant", [pytest.param("fast", id="ckpt_fast")])
+def test_ltx_transformer_crossstage_hang(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ has_audio,
+ checkpoint_variant,
+ reset_seeds,
+) -> None:
+ """FAST cross-stage repro of the e2e stage-2 deadlock: drive ONE model+CCLManager through
+ stage-1 then stage-2 (eager, no VAE/Gemma), so stage-2 reuses the shared strided-AGMM pool
+ that stage-1 left mid-rotation. PASSES -> not cross-stage; HANGS -> confirmed.
+ """
+ # Stage-2 shape; full depth (env-overridable) to stack many back-to-back strided ops per forward.
+ tt_model, kwargs_s2, vN2 = _build_video_trace_setup(
+ mesh_device=mesh_device,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ F=19,
+ H=34,
+ W=60,
+ has_audio=has_audio,
+ checkpoint_variant=checkpoint_variant,
+ num_layers=int(os.environ.get("CROSSSTAGE_NUM_LAYERS", "48")),
+ )
+ # Stage-1 shape (half spatial res, like the e2e's height//2,width//2) on the SAME model.
+ kwargs_s1, vN1 = _build_trace_inner_kwargs(
+ mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, F=19, H=17, W=30, has_audio=has_audio
+ )
+ logger.info(f"cross-stage: stage-1 vN={vN1} then stage-2 vN={vN2} (shared CCLManager)")
+
+ # CROSSSTAGE_SKIP_S1=1 -> run only stage-2, isolating the stage transition from full-depth big-M.
+ skip_s1 = os.environ.get("CROSSSTAGE_SKIP_S1", "0") != "0"
+ if not skip_s1:
+ # Stage 1 first (eager) — populate/rotate the shared strided-AGMM semaphore pool at stage-1 M.
+ for i in range(3):
+ _ = tt_model.inner_step(**kwargs_s1)
+ ttnn.synchronize_device(mesh_device)
+ logger.info(f"stage-1 forward {i} done")
+ logger.info("stage-1 complete; entering STAGE-2 (the e2e hang point) ...")
+ else:
+ logger.info("SKIP_S1: entering STAGE-2 directly (48 layers, no stage-1) ...")
+
+ # Stage 2 (eager) — reuses the same pool. This is where the e2e deadlocks.
+ for i in range(3):
+ _ = tt_model.inner_step(**kwargs_s2)
+ ttnn.synchronize_device(mesh_device)
+ logger.info(f"stage-2 forward {i} done")
+
+ logger.info("PASSED cross-stage: stage-1 -> stage-2 completed without deadlock")
+ del tt_model
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ # VIDEO stage-2 ring config only (8k fabric payload); do NOT add AV/other-mesh params here.
+ [pytest.param((4, 8), 1, 0, 2, ring_params_8k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_8k")],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 34, 60, id="stage_2")])
+@pytest.mark.parametrize("has_audio", [pytest.param(False, id="video"), pytest.param(True, id="av")])
+@pytest.mark.parametrize("checkpoint_variant", [pytest.param("fast", id="ckpt_fast")])
+def test_ltx_transformer_trace_perf(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ has_audio,
+ checkpoint_variant,
+ reset_seeds,
+) -> None:
+ """Trace-based perf for LTXTransformerModel.inner_step (VIDEO stage-2 ring).
+
+ Measures StridedAllGatherMinimalMatmulAsync device time and block wall-clock at steady state
+ under a ttnn trace (gap-free replay), which the eager op-by-op path can't expose. Perf-only:
+ asserts shape/finiteness, no PCC. inner_step is @traced_function(prep_run=False), so kernels
+ are precompiled with a couple of eager warm calls before trace capture.
+ """
+ tt_model, inner_kwargs, video_N_real = _build_video_trace_setup(
+ mesh_device=mesh_device,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ F=F,
+ H=H,
+ W=W,
+ has_audio=has_audio,
+ checkpoint_variant=checkpoint_variant,
+ )
+
+ # AV emits far more markers than video
+ n_warm = int(os.environ.get("LTX_PERF_N_WARM", "1" if has_audio else "2"))
+ n_replay = int(os.environ.get("LTX_PERF_N_REPLAY", "3" if has_audio else "5"))
+
+ # Warm / compile: eager (untraced) calls precompile kernels (prep_run=False on inner_step).
+ t0 = time.time()
+ for _ in range(n_warm):
+ _ = tt_model.inner_step(**inner_kwargs)
+ ttnn.synchronize_device(mesh_device)
+ logger.info(f"eager warm/compile ({n_warm}x): {time.time() - t0:.1f}s")
+
+ # Capture: first traced call captures the trace (and executes it once).
+ t0 = time.time()
+ _ = tt_model.inner_step(**inner_kwargs, traced=True)
+ ttnn.synchronize_device(mesh_device)
+ logger.info(f"trace capture: {time.time() - t0:.1f}s")
+
+ # Drain warm+capture markers so only the signposted replays occupy the profiler buffer.
+ ttnn.ReadDeviceProfiler(mesh_device)
+
+ # Steady-state replay under the signpost bracket for tt-perf-report --start/end-signpost.
+ signpost("start")
+ t0 = time.time()
+ result = None
+ for _ in range(n_replay):
+ result = tt_model.inner_step(**inner_kwargs, traced=True)
+ ttnn.synchronize_device(mesh_device)
+ replay_s = time.time() - t0
+ signpost("stop")
+ logger.info(f"trace replay ({n_replay}x): {replay_s:.3f}s → {replay_s / n_replay * 1e3:.2f} ms/step")
+
+ # Perf-only: crop SP-padding tail, assert shape + finiteness (no PCC)
+ video_out = result[0] if isinstance(result, tuple) else result
+ tt_video = LTXTransformerModel.device_to_host(video_out).squeeze(0)[:, :video_N_real, :]
+ assert tt_video.shape == (1, video_N_real, OUT_CHANNELS), f"video shape {tt_video.shape}"
+ assert torch.isfinite(tt_video).all(), "NaN/Inf in traced video output"
+ logger.info(f"PASSED trace-perf: video {tuple(tt_video.shape)} has_audio={has_audio}")
+
+ del tt_model
+
+
+# Trace-based perf for a SINGLE transformer block (isolates one block's program order).
+def _build_block_trace_setup(
+ *, mesh_device, sp_axis, tp_axis, num_links, topology, is_fsdp, F, H, W, has_audio, checkpoint_variant="fast"
+):
+ """Build an ``LTXTransformerBlock`` + device ``forward`` kwargs for the block trace-perf test.
+
+ Mirrors the non-PCC setup of ``test_ltx_transformer_block``: AV loads the 22B checkpoint
+ (``strict=False``) and builds the full video+audio+cross-modal device tensors; video harvests
+ random scaled weights from a throwaway diffusers block. Returns
+ ``(tt_block, forward_kwargs, video_N_real, audio_N_real)``. Perf-only; no PCC.
+ """
+ checkpoint_22b = _resolve_checkpoint_22b(checkpoint_variant)
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ # Real latent grid → SP-padded sequence (mirrors LTXPipeline)
+ video_N_real = F * H * W
+ video_N = _sp_pad_len(video_N_real, sp_factor)
+ audio_N, audio_N_real = _audio_seq_lens(F, sp_factor)
+ assert video_N % (32 * sp_factor) == 0, f"video_N={video_N} not sp/tile-aligned for sp={sp_factor}"
+ assert audio_N % (32 * sp_factor) == 0, f"audio_N={audio_N} not sp/tile-aligned for sp={sp_factor}"
+
+ ccl_manager = _make_ccl_manager(mesh_device, num_links, topology)
+ parallel_config = _make_parallel_config(mesh_device, sp_axis, tp_axis)
+ tt_block = _make_tt_block(
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ has_audio=has_audio,
+ )
+
+ if has_audio:
+ sd = _load_22b_state_dict(num_layers=1, checkpoint_path=checkpoint_22b)
+ if sd is None:
+ pytest.skip(f"22B checkpoint not found at {checkpoint_22b}")
+ block_sd = {
+ k[len("transformer_blocks.0.") :]: v for k, v in sd.items() if k.startswith("transformer_blocks.0.")
+ }
+ tt_block.load_torch_state_dict(block_sd, strict=False)
+ else:
+ # video mode — harvest a shape-matching state dict from a throwaway block (perf-only).
+ dummy = _make_diffusers_video_block()
+ _scale_init_(dummy)
+ conv = _convert_diffusers_video_block_to_tt(dummy.state_dict(), num_heads=NUM_HEADS, head_dim=HEAD_DIM)
+ tt_block.load_torch_state_dict({k: v.detach().clone() for k, v in conv.items()})
+ del dummy
+
+ # Inputs (real-grid token count; the TT side gets a zero-padded copy below).
+ torch.manual_seed(INPUT_SEED)
+ x = torch.randn(1, video_N_real, DIM, dtype=torch.float32)
+ context = torch.randn(1, PROMPT_LEN, CTX_DIM, dtype=torch.float32)
+ temb = torch.randn(1, 1, 9 * DIM, dtype=torch.float32) # 9 adaln params
+ prompt_temb = torch.randn(1, 1, 2 * DIM, dtype=torch.float32) # 2 adaln params for prompt
+
+ # TT video tensors
+ spatial = _pad_seq_dim(x, video_N, dim=1).unsqueeze(0)
+ tt_spatial = bf16_tensor_2dshard(spatial, device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3})
+ tt_prompt = bf16_tensor(context.unsqueeze(0), device=mesh_device)
+ tt_temb = bf16_tensor(
+ temb.reshape(9, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3
+ )
+ tt_prompt_temb = bf16_tensor(prompt_temb.reshape(2, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device)
+ tt_cos, tt_sin = _tt_rope(
+ _video_rope_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ tt_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=mesh_device)
+
+ forward_kwargs = dict(
+ video_1BND=tt_spatial,
+ video_prompt=tt_prompt,
+ video_temb=tt_temb,
+ video_N=video_N_real,
+ video_rope_cos=tt_cos,
+ video_rope_sin=tt_sin,
+ trans_mat=tt_trans_mat,
+ video_prompt_temb=tt_prompt_temb,
+ )
+ if has_audio:
+ # Real tokens in [:audio_N_real], zeros in the padded tail (matches the padded audio latent).
+ a_x = torch.zeros(1, audio_N, AUDIO_DIM, dtype=torch.float32)
+ a_x[:, :audio_N_real, :] = torch.randn(1, audio_N_real, AUDIO_DIM, dtype=torch.float32)
+ a_ctx = torch.randn(1, PROMPT_LEN, AUDIO_CTX_DIM, dtype=torch.float32)
+ a_temb = torch.randn(1, 1, 9 * AUDIO_DIM, dtype=torch.float32)
+ a_prompt_temb = torch.randn(1, 1, 2 * AUDIO_DIM, dtype=torch.float32)
+ av_ca_v = torch.randn(1, 1, 5 * DIM, dtype=torch.float32) # 4 scale-shift + 1 gate
+ av_ca_a = torch.randn(1, 1, 5 * AUDIO_DIM, dtype=torch.float32)
+
+ a_cos, a_sin = _tt_rope(_audio_rope_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis)
+ vx_cos, vx_sin = _tt_rope(
+ _video_cross_pe_freqs, F, H, W, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis, pad_to=video_N
+ )
+ ax_cos, ax_sin = _tt_rope(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, sp_axis=sp_axis, tp_axis=tp_axis
+ )
+ ax_cos_full, ax_sin_full = _tt_rope_full(
+ _audio_cross_pe_freqs, audio_N, mesh_device=mesh_device, tp_axis=tp_axis
+ )
+
+ # Padding masks, same construction as the fast pipeline (audio padded, video aligned).
+ a_attn_mask, a_pad_sp, a_pad_full = build_audio_masks(
+ audio_N, audio_N_real, mesh_device=mesh_device, sp_axis=sp_axis
+ )
+ v_pad_sp = build_video_pad_mask(video_N, video_N_real, mesh_device=mesh_device, sp_axis=sp_axis)
+
+ forward_kwargs.update(
+ audio_1BND=bf16_tensor_2dshard(
+ a_x.unsqueeze(0), device=mesh_device, shard_mapping={sp_axis: 2, tp_axis: 3}
+ ),
+ audio_prompt=bf16_tensor(a_ctx.unsqueeze(0), device=mesh_device),
+ audio_temb=bf16_tensor(
+ a_temb.reshape(9, AUDIO_DIM).unsqueeze(1).unsqueeze(1),
+ device=mesh_device,
+ mesh_axis=tp_axis,
+ shard_dim=3,
+ ),
+ audio_prompt_temb=bf16_tensor(
+ a_prompt_temb.reshape(2, AUDIO_DIM).unsqueeze(1).unsqueeze(1), device=mesh_device
+ ),
+ av_ca_temb=bf16_tensor(
+ av_ca_v.reshape(5, DIM).unsqueeze(1).unsqueeze(1), device=mesh_device, mesh_axis=tp_axis, shard_dim=3
+ ),
+ av_ca_audio_temb=bf16_tensor(
+ av_ca_a.reshape(5, AUDIO_DIM).unsqueeze(1).unsqueeze(1),
+ device=mesh_device,
+ mesh_axis=tp_axis,
+ shard_dim=3,
+ ),
+ audio_N=audio_N,
+ audio_rope_cos=a_cos,
+ audio_rope_sin=a_sin,
+ video_cross_pe_cos=vx_cos,
+ video_cross_pe_sin=vx_sin,
+ audio_cross_pe_cos=ax_cos,
+ audio_cross_pe_sin=ax_sin,
+ audio_cross_pe_cos_full=ax_cos_full,
+ audio_cross_pe_sin_full=ax_sin_full,
+ audio_attn_mask=a_attn_mask,
+ audio_padding_mask=a_pad_sp,
+ audio_padding_mask_full=a_pad_full,
+ video_padding_mask=v_pad_sp,
+ )
+ return tt_block, forward_kwargs, video_N_real, audio_N_real
+
+
+@pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ # Same ring stage-2 8k config as test_ltx_transformer_trace_perf
+ [pytest.param((4, 8), 1, 0, 2, ring_params_8k, ttnn.Topology.Ring, False, id="ring_bh_4x8sp1tp0_8k")],
+ indirect=["mesh_device", "device_params"],
+)
+@pytest.mark.parametrize(("F", "H", "W"), [pytest.param(19, 34, 60, id="stage_2")])
+@pytest.mark.parametrize("has_audio", [pytest.param(False, id="video"), pytest.param(True, id="av")])
+@pytest.mark.parametrize("checkpoint_variant", [pytest.param("fast", id="ckpt_fast")])
+def test_ltx_transformer_block_trace_perf(
+ mesh_device,
+ sp_axis,
+ tp_axis,
+ num_links,
+ topology,
+ is_fsdp,
+ F,
+ H,
+ W,
+ has_audio,
+ checkpoint_variant,
+ reset_seeds,
+) -> None:
+ """Trace-based perf for a SINGLE ``LTXTransformerBlock`` forward (AV stage-2 ring).
+
+ Isolates one block's program order — no model-level patch-embed / proj-out / tilize /
+ untilize / mesh-partition ops, and no multi-block scrambling — so tt-perf-report shows a clean
+ one-block, program-ordered view. Captures the block forward under a ttnn trace (via the same
+ ``Tracer`` machinery that backs ``inner_step``'s ``traced=True`` path) and replays it under the
+ same "start"/"stop" signpost bracket used by ``test_ltx_transformer_trace_perf``. Perf-only:
+ asserts shape/finiteness, no PCC.
+ """
+ tt_block, forward_kwargs, video_N_real, _audio_N_real = _build_block_trace_setup(
+ mesh_device=mesh_device,
+ sp_axis=sp_axis,
+ tp_axis=tp_axis,
+ num_links=num_links,
+ topology=topology,
+ is_fsdp=is_fsdp,
+ F=F,
+ H=H,
+ W=W,
+ has_audio=has_audio,
+ checkpoint_variant=checkpoint_variant,
+ )
+
+ # Trace the block forward with the same Tracer settings inner_step uses
+ tracer = Tracer(tt_block.forward, device=mesh_device, prep_run=False, clone_prep_inputs=False)
+
+ # AV runs many more ops than video (audio self/cross + a2v/v2a + audio ff/RS)
+ n_warm = 1 if has_audio else 2
+ n_replay = 3 if has_audio else 5
+
+ # Warm / compile: eager (untraced) calls precompile kernels before trace capture.
+ t0 = time.time()
+ for _ in range(n_warm):
+ _ = tt_block(**forward_kwargs)
+ ttnn.synchronize_device(mesh_device)
+ logger.info(f"eager warm/compile ({n_warm}x): {time.time() - t0:.1f}s")
+
+ # Capture: first traced call captures the trace (and executes it once).
+ t0 = time.time()
+ _ = tracer(**forward_kwargs, traced=True)
+ ttnn.synchronize_device(mesh_device)
+ logger.info(f"trace capture: {time.time() - t0:.1f}s")
+
+ # Drain warm+capture markers so ONLY the signposted replays occupy the profiler DRAM buffer
+ ttnn.ReadDeviceProfiler(mesh_device)
+
+ # Steady-state replay under the signpost bracket for tt-perf-report --start/end-signpost.
+ signpost("start")
+ t0 = time.time()
+ result = None
+ for _ in range(n_replay):
+ result = tracer(**forward_kwargs, traced=True)
+ ttnn.synchronize_device(mesh_device)
+ replay_s = time.time() - t0
+ signpost("stop")
+ logger.info(f"trace replay ({n_replay}x): {replay_s:.3f}s → {replay_s / n_replay * 1e3:.2f} ms/step")
+
+ # Perf-only: crop SP-padding tail, assert shape + finiteness (no PCC)
+ tt_v = result[0] if isinstance(result, tuple) else result
+ concat_dims = [None, None]
+ concat_dims[sp_axis] = 2
+ concat_dims[tp_axis] = 3
+ tt_v_torch = ttnn.to_torch(
+ tt_v,
+ mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=concat_dims, mesh_shape=tuple(mesh_device.shape)),
+ ).squeeze(0)[:, :video_N_real, :]
+ assert tt_v_torch.shape == (1, video_N_real, DIM), f"video shape {tt_v_torch.shape}"
+ assert torch.isfinite(tt_v_torch).all(), "NaN/Inf in traced video output"
+ logger.info(f"PASSED block trace-perf: video {tuple(tt_v_torch.shape)} has_audio={has_audio}")
+
+ tracer.release_trace()
+ del tt_block
diff --git a/code/models/tt_dit/tests/models/ltx/test_vae_ltx.py b/code/models/tt_dit/tests/models/ltx/test_vae_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..edbd3ba341923daf12e1da77ecdd034058384123
--- /dev/null
+++ b/code/models/tt_dit/tests/models/ltx/test_vae_ltx.py
@@ -0,0 +1,1151 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tests for LTX-2 Video VAE components.
+
+Torch references use diffusers LTX-2 VAE modules (mirrors test_vae_wan2_1.py / test_audio_ltx.py).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import pytest
+import torch
+import torch.nn as nn
+from loguru import logger
+from tracy import signpost
+
+import ttnn
+from models.tt_dit.models.vae.vae_ltx import (
+ LTXCausalConv3d,
+ LTXDepthToSpaceUpsample,
+ LTXResnetBlock3D,
+ LTXVideoDecoder,
+ LTXVideoEncoder,
+)
+from models.tt_dit.parallel.config import ParallelFactor, VaeHWParallelConfig
+from models.tt_dit.parallel.manager import CCLManager
+from models.tt_dit.utils.check import assert_quality
+from models.tt_dit.utils.conv3d import conv_pad_in_channels
+
+from .ltx_mesh_params import (
+ LTX_VAE_DECODER_MESH_PARAMS,
+ LTX_VAE_DECODER_MULTI_ONLY_MESH_PARAMS,
+ LTX_VAE_ENCODER_PROD_MESH_PARAMS,
+)
+
+
+def _require_diffusers_ltx_vae():
+ """Return diffusers LTX-2 VAE building blocks (preferred) or LTX-Video fallbacks."""
+ pytest.importorskip("diffusers")
+ try:
+ from diffusers.models.autoencoders.autoencoder_kl_ltx2 import (
+ LTX2VideoCausalConv3d,
+ LTX2VideoDownsampler3d,
+ LTX2VideoResnetBlock3d,
+ LTX2VideoUpsampler3d,
+ PerChannelRMSNorm,
+ )
+
+ return {
+ "causal_conv": LTX2VideoCausalConv3d,
+ "resnet": LTX2VideoResnetBlock3d,
+ "upsample": LTX2VideoUpsampler3d,
+ "downsample": LTX2VideoDownsampler3d,
+ "pixel_norm": PerChannelRMSNorm,
+ "ltx2": True,
+ }
+ except ImportError:
+ from diffusers.models.autoencoders.autoencoder_kl_ltx import (
+ LTXVideoCausalConv3d,
+ LTXVideoDownsampler3d,
+ LTXVideoResnetBlock3d,
+ LTXVideoUpsampler3d,
+ )
+ from diffusers.models.normalization import RMSNorm
+
+ return {
+ "causal_conv": LTXVideoCausalConv3d,
+ "resnet": LTXVideoResnetBlock3d,
+ "upsample": LTXVideoUpsampler3d,
+ "downsample": LTXVideoDownsampler3d,
+ "pixel_norm": RMSNorm,
+ "ltx2": False,
+ }
+
+
+def _diffusers_resnet_state_to_tt(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Remap diffusers LTX-2 resnet keys to the TT/Lightricks layout."""
+ out = dict(state)
+ if "conv_shortcut.weight" in out:
+ out["conv_shortcut.conv.weight"] = out.pop("conv_shortcut.weight")
+ if "conv_shortcut.bias" in out:
+ out["conv_shortcut.conv.bias"] = out.pop("conv_shortcut.bias")
+ return out
+
+
+def _run_diffusers_causal_conv(conv: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor:
+ if ltx2:
+ return conv(x, causal=causal)
+ return conv(x)
+
+
+def _run_diffusers_resnet(block: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor:
+ if ltx2:
+ return block(x, causal=causal)
+ return block(x)
+
+
+def _run_diffusers_upsample(block: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor:
+ if ltx2:
+ return block(x, causal=causal)
+ return block(x)
+
+
+def _run_diffusers_downsample(block: nn.Module, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor:
+ if ltx2:
+ return block(x, causal=causal)
+ return block(x)
+
+
+class _PerChannelStatistics(nn.Module):
+ """Minimal per-channel latent stats (Lightricks key names for TT weight load)."""
+
+ def __init__(self, latent_channels: int) -> None:
+ super().__init__()
+ self.register_buffer("mean-of-means", torch.zeros(latent_channels))
+ self.register_buffer("std-of-means", torch.ones(latent_channels))
+
+ def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
+ mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1)
+ std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1)
+ return x * std + mean
+
+ def normalize(self, x: torch.Tensor) -> torch.Tensor:
+ mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1)
+ std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1)
+ return (x - mean) / std
+
+
+def _diffusers_decoder_state_to_tt(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Remap diffusers LTX-2 decoder keys (incl. nested resnets) to TT layout."""
+ out: dict[str, torch.Tensor] = {}
+ for key, value in state.items():
+ if key.endswith("conv_shortcut.weight"):
+ out[key.replace("conv_shortcut.weight", "conv_shortcut.conv.weight")] = value
+ elif key.endswith("conv_shortcut.bias"):
+ out[key.replace("conv_shortcut.bias", "conv_shortcut.conv.bias")] = value
+ else:
+ out[key] = value
+ return out
+
+
+class _TorchUNetMidBlock3D(nn.Module):
+ def __init__(
+ self,
+ *,
+ in_channels: int,
+ num_layers: int,
+ resnet_cls: type,
+ ltx2: bool,
+ spatial_padding_mode: str = "zeros",
+ ) -> None:
+ super().__init__()
+ kwargs: dict[str, Any] = {"in_channels": in_channels, "out_channels": in_channels}
+ if ltx2 and spatial_padding_mode != "zeros":
+ kwargs["spatial_padding_mode"] = spatial_padding_mode
+ self.res_blocks = nn.ModuleList([resnet_cls(**kwargs) for _ in range(num_layers)])
+
+ def forward(self, x: torch.Tensor, *, causal: bool, ltx2: bool) -> torch.Tensor:
+ for block in self.res_blocks:
+ x = _run_diffusers_resnet(block, x, causal=causal, ltx2=ltx2)
+ return x
+
+
+def _unpatchify(x: torch.Tensor, *, patch_size_hw: int = 4, patch_size_t: int = 1) -> torch.Tensor:
+ batch_size, num_channels, num_frames, height, width = x.shape
+ p = patch_size_hw
+ p_t = patch_size_t
+ x = x.reshape(batch_size, -1, p_t, p, p, num_frames, height, width)
+ return x.permute(0, 1, 5, 2, 6, 4, 7, 3).flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+
+def _patchify(x: torch.Tensor, *, patch_size_hw: int = 4, patch_size_t: int = 1) -> torch.Tensor:
+ """Space-to-depth patchify (exact inverse of ``_unpatchify``): folds a
+ ``(p_t, p_hw, p_hw)`` patch into channels with order ``(c, p_t, p_width, p_height)``.
+ Mirrors the Lightricks/TT encoder patchify so both references agree.
+ """
+ batch_size, num_channels, num_frames, height, width = x.shape
+ p = patch_size_hw
+ p_t = patch_size_t
+ x = x.reshape(batch_size, num_channels, num_frames // p_t, p_t, height // p, p, width // p, p)
+ # (b, c, f', p_t, h', p_h, w', p_w) -> (b, c, p_t, p_w, p_h, f', h', w')
+ x = x.permute(0, 1, 3, 7, 5, 2, 4, 6)
+ return x.reshape(batch_size, num_channels * p_t * p * p, num_frames // p_t, height // p, width // p)
+
+
+class _TorchLTXVideoDecoder(nn.Module):
+ """Block-list LTX decoder built from diffusers VAE primitives (matches TT/Lightricks layout)."""
+
+ def __init__(
+ self,
+ *,
+ decoder_blocks: list[tuple[str, dict | int]],
+ in_channels: int = 128,
+ out_channels: int = 3,
+ patch_size: int = 4,
+ base_channels: int = 128,
+ causal: bool = False,
+ spatial_padding_mode: str = "zeros",
+ vae_mods: dict,
+ ) -> None:
+ super().__init__()
+ self.patch_size = patch_size
+ self.causal = causal
+ self.ltx2 = vae_mods["ltx2"]
+ causal_conv = vae_mods["causal_conv"]
+ resnet_cls = vae_mods["resnet"]
+ upsample_cls = vae_mods["upsample"]
+ pixel_norm = vae_mods["pixel_norm"]
+
+ out_channels_patched = out_channels * patch_size**2
+ feature_channels = base_channels * 8
+
+ conv_kwargs: dict[str, Any] = {}
+ if self.ltx2:
+ if spatial_padding_mode != "zeros":
+ conv_kwargs["spatial_padding_mode"] = spatial_padding_mode
+ else:
+ conv_kwargs["is_causal"] = causal
+
+ self.per_channel_statistics = _PerChannelStatistics(in_channels)
+ self.conv_in = causal_conv(
+ in_channels=in_channels,
+ out_channels=feature_channels,
+ kernel_size=3,
+ stride=1,
+ **conv_kwargs,
+ )
+
+ stride_map = {
+ "compress_all": (2, 2, 2),
+ "compress_space": (1, 2, 2),
+ "compress_time": (2, 1, 1),
+ }
+
+ self.up_blocks = nn.ModuleList()
+ ch = feature_channels
+ for block_name, block_params in reversed(decoder_blocks):
+ block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
+ if block_name == "res_x":
+ self.up_blocks.append(
+ _TorchUNetMidBlock3D(
+ in_channels=ch,
+ num_layers=block_config["num_layers"],
+ resnet_cls=resnet_cls,
+ ltx2=self.ltx2,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ elif block_name in stride_map:
+ multiplier = block_config.get("multiplier", 1)
+ upsample_kwargs: dict[str, Any] = {
+ "in_channels": ch,
+ "stride": stride_map[block_name],
+ "residual": block_config.get("residual", False),
+ "upscale_factor": multiplier,
+ }
+ self.up_blocks.append(upsample_cls(**upsample_kwargs))
+ ch = ch // multiplier
+ else:
+ raise ValueError(f"Unknown decoder block: {block_name}")
+
+ self.conv_norm_out = pixel_norm()
+ self.conv_act = nn.SiLU()
+ self.conv_out = causal_conv(
+ in_channels=ch,
+ out_channels=out_channels_patched,
+ kernel_size=3,
+ stride=1,
+ **conv_kwargs,
+ )
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ sample = self.per_channel_statistics.un_normalize(sample)
+ sample = _run_diffusers_causal_conv(self.conv_in, sample, causal=self.causal, ltx2=self.ltx2)
+ for up_block in self.up_blocks:
+ if isinstance(up_block, _TorchUNetMidBlock3D):
+ sample = up_block(sample, causal=self.causal, ltx2=self.ltx2)
+ else:
+ sample = _run_diffusers_upsample(up_block, sample, causal=self.causal, ltx2=self.ltx2)
+ if self.ltx2:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample.movedim(1, -1)).movedim(-1, 1)
+ sample = self.conv_act(sample)
+ sample = _run_diffusers_causal_conv(self.conv_out, sample, causal=self.causal, ltx2=self.ltx2)
+ return _unpatchify(sample, patch_size_hw=self.patch_size)
+
+
+def _vae_parallel_kwargs(mesh_device: ttnn.MeshDevice) -> dict:
+ """Build a VaeHWParallelConfig + CCLManager for a test mesh.
+
+ Maps mesh axis 0 → height shard, axis 1 → width shard. For a (1, 1) mesh
+ both factors are 1 and ``LTXCausalConv3d`` skips halo exchange.
+ """
+ mesh_shape = tuple(mesh_device.shape)
+ parallel_config = VaeHWParallelConfig(
+ height_parallel=ParallelFactor(factor=mesh_shape[0], mesh_axis=0),
+ width_parallel=ParallelFactor(factor=mesh_shape[1], mesh_axis=1),
+ )
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
+ return {"parallel_config": parallel_config, "ccl_manager": ccl_manager}
+
+
+# Single-chip: no fabric (avoids BH fabric router handshake when only one device is opened).
+# Stand-alone conv/res/upsample tests use raw (unsharded) tensors — 1x1 only.
+_LTX_VAE_MESH_DEVICE_PARAMS = [
+ ((1, 1), {}),
+]
+_LTX_VAE_MESH_DEVICE_IDS = ["1x1"]
+
+# (num_frames, height, width): H/W divisible by 64; (num_frames - 1) % 8 == 0 for VAE.
+_LTX_DECODER_SHAPE_PARAMS = [
+ pytest.param(145, 1088, 1920, id="145f_1088x1920"), # 1080p production (19 latent frames)
+]
+
+_LTX_DECODER_2K_SHAPE_PARAMS = [
+ pytest.param(9, 1088, 2048, id="9f_1088x2048"), # 2K (1088x2048, latent 34x64)
+]
+
+# Smoke resolution — only shape exercised on a 1x1 mesh in test_ltx_video_decoder.
+_LTX_DECODER_1X1_MAX_HW = (256, 256)
+
+
+def _skip_ltx_decoder_if_single_chip_too_large(mesh_device: ttnn.MeshDevice, height: int, width: int) -> None:
+ if tuple(mesh_device.shape) == (1, 1) and (
+ height > _LTX_DECODER_1X1_MAX_HW[0] or width > _LTX_DECODER_1X1_MAX_HW[1]
+ ):
+ pytest.skip(f"{height}x{width} LTX decode requires a multi-device mesh")
+
+
+def _unwrap_tt_output(tt_out):
+ if isinstance(tt_out, (tuple, list)):
+ # Newer TT path may return auxiliary tensors alongside the primary output.
+ return tt_out[0]
+ return tt_out
+
+
+# Full-decoder test — exercises the production decoder_blocks under sharding.
+_LTX_PROD_DECODER_BLOCKS = [
+ ("res_x", {"num_layers": 4}),
+ ("compress_space", {"multiplier": 2}),
+ ("res_x", {"num_layers": 6}),
+ ("compress_time", {"multiplier": 2}),
+ ("res_x", {"num_layers": 4}),
+ ("compress_all", {"multiplier": 1}),
+ ("res_x", {"num_layers": 2}),
+ ("compress_all", {"multiplier": 2}),
+ ("res_x", {"num_layers": 2}),
+]
+
+
+@pytest.mark.parametrize(
+ "in_c, out_c, kernel_size, stride, T, H, W",
+ [
+ (128, 128, 3, 1, 3, 16, 16), # small residual block conv
+ (128, 256, 3, 1, 3, 16, 16), # channel expansion
+ (48, 128, 3, 1, 5, 32, 32), # conv_in (after patchify)
+ (128, 1024, 3, 1, 21, 17, 15), # ltx_s0_conv_in @ 2x4 per-device (1080p path)
+ (1024, 1024, 3, 1, 21, 17, 15), # ltx_s0_res
+ (1024, 4096, 3, 1, 21, 17, 15), # ltx_s0_up (upsample projection)
+ (512, 512, 3, 1, 39, 34, 30), # ltx_s1_res
+ (128, 48, 3, 1, 147, 136, 120), # ltx_s4_out
+ ],
+ ids=[
+ "res_128_128",
+ "expand_128_256",
+ "conv_in_48_128",
+ "ltx_s0_conv_in",
+ "ltx_s0_res",
+ "ltx_s0_up",
+ "ltx_s1_res",
+ "ltx_s4_out",
+ ],
+)
+@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"])
+@pytest.mark.parametrize(
+ "mesh_device, device_params",
+ _LTX_VAE_MESH_DEVICE_PARAMS,
+ ids=_LTX_VAE_MESH_DEVICE_IDS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_causal_conv3d(
+ mesh_device: ttnn.MeshDevice,
+ in_c: int,
+ out_c: int,
+ kernel_size: int,
+ stride: int,
+ T: int,
+ H: int,
+ W: int,
+ mean: float,
+ std: float,
+):
+ """
+ Test LTXCausalConv3d against diffusers LTX causal conv reference.
+ """
+ vae_mods = _require_diffusers_ltx_vae()
+ TorchCausalConv3d = vae_mods["causal_conv"]
+ ltx2 = vae_mods["ltx2"]
+
+ B = 1
+ torch.manual_seed(42)
+
+ conv_kwargs: dict[str, Any] = {
+ "in_channels": in_c,
+ "out_channels": out_c,
+ "kernel_size": kernel_size,
+ "stride": stride if ltx2 else (stride, stride, stride),
+ }
+ if not ltx2:
+ conv_kwargs["is_causal"] = True
+
+ torch_model = TorchCausalConv3d(**conv_kwargs)
+ torch_model.eval()
+
+ # TT model
+ tt_model = LTXCausalConv3d(
+ in_channels=in_c,
+ out_channels=out_c,
+ kernel_size=kernel_size,
+ stride=stride,
+ mesh_device=mesh_device,
+ **_vae_parallel_kwargs(mesh_device),
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ # Input
+ x = torch.randn(B, in_c, T, H, W, dtype=torch.float32) * std + mean
+
+ # PyTorch forward (BCTHW format)
+ with torch.no_grad():
+ torch_out = _run_diffusers_causal_conv(torch_model, x, causal=True, ltx2=ltx2)
+
+ # TT forward (BTHWC format)
+ x_bthwc = x.permute(0, 2, 3, 4, 1) # (B, T, H, W, C)
+ x_bthwc = conv_pad_in_channels(x_bthwc) # Pad C to alignment
+ x_tt = ttnn.from_torch(x_bthwc, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=ttnn.bfloat16)
+
+ tt_out = _unwrap_tt_output(tt_model(x_tt))
+ tt_out_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0]) # (B, T_out, H_out, W_out, C_out)
+ tt_out_torch = tt_out_torch[:, :, :, :, :out_c] # Trim padded channels
+ tt_out_torch = tt_out_torch.permute(0, 4, 1, 2, 3) # Back to BCTHW
+
+ logger.info(f"PyTorch out: {torch_out.shape}, TT out: {tt_out_torch.shape}")
+ assert_quality(torch_out, tt_out_torch, pcc=0.999)
+ logger.info(f"PASSED: LTXCausalConv3d ({in_c}->{out_c}) matches reference")
+
+
+@pytest.mark.parametrize(
+ "in_c, out_c, T, H, W",
+ [
+ (128, 128, 3, 16, 16), # same channels (no shortcut)
+ (128, 256, 3, 16, 16), # channel expansion (with shortcut)
+ (1024, 1024, 21, 17, 15), # ltx_s0_res @ 2x4 per-device
+ (512, 512, 39, 34, 30), # ltx_s1_res
+ (256, 256, 147, 70, 66), # ltx_s3_res @ 2K per-device
+ (128, 128, 147, 138, 130), # ltx_s4_res @ 2K per-device
+ ],
+ ids=[
+ "same_channels",
+ "expand_channels",
+ "ltx_s0_res",
+ "ltx_s1_res",
+ "ltx_s3_res_2k",
+ "ltx_s4_res_2k",
+ ],
+)
+@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"])
+@pytest.mark.parametrize(
+ "mesh_device, device_params",
+ _LTX_VAE_MESH_DEVICE_PARAMS,
+ ids=_LTX_VAE_MESH_DEVICE_IDS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_resnet_block(
+ mesh_device: ttnn.MeshDevice, in_c: int, out_c: int, T: int, H: int, W: int, mean: float, std: float
+):
+ """
+ Test LTXResnetBlock3D against diffusers LTX ResnetBlock3D reference.
+ """
+ vae_mods = _require_diffusers_ltx_vae()
+ TorchResnetBlock3D = vae_mods["resnet"]
+ ltx2 = vae_mods["ltx2"]
+
+ B = 1
+ torch.manual_seed(42)
+
+ torch_model = TorchResnetBlock3D(in_channels=in_c, out_channels=out_c)
+ torch_model.eval()
+
+ # TT model
+ tt_model = LTXResnetBlock3D(
+ in_channels=in_c,
+ out_channels=out_c,
+ mesh_device=mesh_device,
+ **_vae_parallel_kwargs(mesh_device),
+ )
+ tt_model.load_torch_state_dict(_diffusers_resnet_state_to_tt(torch_model.state_dict()))
+
+ # Input
+ x = torch.randn(B, in_c, T, H, W, dtype=torch.float32) * std + mean
+
+ # PyTorch forward
+ with torch.no_grad():
+ torch_out = _run_diffusers_resnet(torch_model, x, causal=True, ltx2=ltx2)
+
+ # TT forward
+ x_bthwc = x.permute(0, 2, 3, 4, 1)
+ x_bthwc = conv_pad_in_channels(x_bthwc)
+ x_tt = ttnn.from_torch(x_bthwc, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=ttnn.bfloat16)
+
+ tt_out = _unwrap_tt_output(tt_model(x_tt))
+ tt_out_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0])
+ tt_out_torch = tt_out_torch[:, :, :, :, :out_c]
+ tt_out_torch = tt_out_torch.permute(0, 4, 1, 2, 3)
+
+ logger.info(f"PyTorch out: {torch_out.shape}, TT out: {tt_out_torch.shape}")
+ # Lower threshold for channel-expansion blocks (GroupNorm(1) vs layer_norm precision)
+ min_pcc = 0.995 if in_c != out_c else 0.999
+ assert_quality(torch_out, tt_out_torch, pcc=min_pcc)
+ logger.info(f"PASSED: LTXResnetBlock3D ({in_c}->{out_c}) matches reference")
+
+
+@pytest.mark.parametrize(
+ "in_c, stride, T, H, W",
+ [
+ (128, (2, 2, 2), 4, 8, 8), # full 3D upsample (compress_all)
+ (128, (1, 2, 2), 3, 8, 8), # spatial only (compress_space)
+ (128, (2, 1, 1), 4, 8, 8), # temporal only (compress_time)
+ (512, (2, 2, 2), 39, 34, 30), # compress_all @ ltx_s1 stage
+ (1024, (1, 2, 2), 21, 17, 15), # compress_space @ ltx_s0 stage
+ (256, (2, 1, 1), 75, 68, 60), # compress_time @ ltx_s2 stage
+ ],
+ ids=[
+ "upsample_all",
+ "upsample_space",
+ "upsample_time",
+ "upsample_all_s1",
+ "upsample_space_s0",
+ "upsample_time_s2",
+ ],
+)
+@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"])
+@pytest.mark.parametrize(
+ "mesh_device, device_params",
+ _LTX_VAE_MESH_DEVICE_PARAMS,
+ ids=_LTX_VAE_MESH_DEVICE_IDS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_depth_to_space_upsample(
+ mesh_device: ttnn.MeshDevice, in_c: int, stride: tuple, T: int, H: int, W: int, mean: float, std: float
+):
+ """Test LTXDepthToSpaceUpsample against diffusers LTX upsampler reference."""
+ vae_mods = _require_diffusers_ltx_vae()
+ TorchUpsample = vae_mods["upsample"]
+ ltx2 = vae_mods["ltx2"]
+
+ B = 1
+ torch.manual_seed(42)
+
+ torch_model = TorchUpsample(in_channels=in_c, stride=stride, upscale_factor=1)
+ torch_model.eval()
+
+ tt_model = LTXDepthToSpaceUpsample(
+ in_channels=in_c,
+ stride=stride,
+ mesh_device=mesh_device,
+ **_vae_parallel_kwargs(mesh_device),
+ )
+ tt_model.load_torch_state_dict(torch_model.state_dict())
+
+ x = torch.randn(B, in_c, T, H, W, dtype=torch.float32) * std + mean
+ with torch.no_grad():
+ torch_out = _run_diffusers_upsample(torch_model, x, causal=True, ltx2=ltx2)
+
+ x_bthwc = x.permute(0, 2, 3, 4, 1)
+ x_bthwc = conv_pad_in_channels(x_bthwc)
+ x_tt = ttnn.from_torch(x_bthwc, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=ttnn.bfloat16)
+
+ tt_out = _unwrap_tt_output(tt_model(x_tt))
+ tt_out_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0])
+ out_c = torch_out.shape[1]
+ tt_out_torch = tt_out_torch[:, :, :, :, :out_c].permute(0, 4, 1, 2, 3)
+
+ logger.info(f"Upsample stride={stride}: {x.shape} -> torch {torch_out.shape}, tt {tt_out_torch.shape}")
+ assert_quality(torch_out, tt_out_torch, pcc=0.999)
+ logger.info(f"PASSED: DepthToSpaceUpsample stride={stride}")
+
+
+def _run_ltx_decoder_parity(
+ mesh_device: ttnn.MeshDevice,
+ h_axis: int,
+ w_axis: int,
+ num_links: int,
+ num_frames: int,
+ height: int,
+ width: int,
+ *,
+ mean: float = 0.0,
+ std: float = 1.0,
+ pcc: float = 0.99,
+):
+ """Build torch + TT production decoders at ``(num_frames, height, width)`` and
+ assert numerical parity. Shared by the standard multi-resolution decoder test
+ and the 2K decoder test. The TT decoder is given num_frames/height/width so the
+ dim walker fires and ``_BLOCKINGS`` gets an exact-match lookup for every conv3d.
+ """
+ vae_mods = _require_diffusers_ltx_vae()
+ if not vae_mods["ltx2"]:
+ pytest.skip("Full decoder parity requires diffusers autoencoder_kl_ltx2 (runtime causal= support)")
+
+ B = 1
+ torch.manual_seed(42)
+ spatial_compression = 32
+ latent_frames = (num_frames - 1) // 8 + 1
+ latent_h = height // spatial_compression
+ latent_w = width // spatial_compression
+
+ # PyTorch reference — full production decoder shape (diffusers primitives).
+ torch_decoder = _TorchLTXVideoDecoder(
+ decoder_blocks=_LTX_PROD_DECODER_BLOCKS,
+ in_channels=128,
+ out_channels=3,
+ patch_size=4,
+ base_channels=128,
+ causal=False,
+ spatial_padding_mode="zeros",
+ vae_mods=vae_mods,
+ )
+ torch_decoder.eval()
+
+ # Per-channel stats set to identity so denormalization is a no-op (random
+ # weights produce garbage stats otherwise).
+ torch_state = torch_decoder.state_dict()
+ torch_state["per_channel_statistics.mean-of-means"] = torch.zeros(128)
+ torch_state["per_channel_statistics.std-of-means"] = torch.ones(128)
+ torch_decoder.load_state_dict(torch_state)
+
+ # Build VAE parallel config from h_axis / w_axis (mirrors test_wan_decoder).
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear, num_links=num_links)
+ parallel_config = VaeHWParallelConfig(
+ height_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[h_axis], mesh_axis=h_axis),
+ width_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[w_axis], mesh_axis=w_axis),
+ )
+
+ # TT decoder — pass num_frames/height/width so the dim walker fires and
+ # _BLOCKINGS gets an exact-match lookup for every conv3d.
+ tt_decoder = LTXVideoDecoder(
+ decoder_blocks=_LTX_PROD_DECODER_BLOCKS,
+ in_channels=128,
+ out_channels=3,
+ patch_size=4,
+ base_channels=128,
+ causal=False,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ mesh_device=mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+ tt_decoder.load_torch_state_dict(_diffusers_decoder_state_to_tt(torch_decoder.state_dict()))
+
+ # Latent shape matches what the pipeline would produce at this resolution.
+ latent = torch.randn(B, 128, latent_frames, latent_h, latent_w, dtype=torch.float32) * std + mean
+
+ with torch.no_grad():
+ torch_out = torch_decoder(latent)
+
+ tt_out = tt_decoder(latent)
+
+ logger.info(f"Decoder: {latent.shape} -> torch {torch_out.shape}, tt {tt_out.shape}")
+ assert_quality(torch_out, tt_out, pcc=pcc)
+ logger.info("PASSED: LTXVideoDecoder matches PyTorch reference")
+
+
+@pytest.mark.parametrize("num_frames, height, width", _LTX_DECODER_SHAPE_PARAMS)
+@pytest.mark.parametrize("mean, std", [(0, 1), (2, 3)], ids=["std1", "std3"])
+@pytest.mark.parametrize(
+ ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"),
+ LTX_VAE_DECODER_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_video_decoder(
+ mesh_device: ttnn.MeshDevice,
+ num_frames: int,
+ height: int,
+ width: int,
+ mean: float,
+ std: float,
+ h_axis: int,
+ w_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+):
+ """
+ Test full LTXVideoDecoder against PyTorch VideoDecoder reference using the
+ production LTX-2.3 22B decoder_blocks. Random weights — checks numerical
+ parity, not absolute output quality.
+
+ For 2x4 this validates the parallel halo-exchange path end-to-end, which
+ is otherwise only covered by the slow pipeline-level test.
+ """
+ _skip_ltx_decoder_if_single_chip_too_large(mesh_device, height, width)
+ _run_ltx_decoder_parity(mesh_device, h_axis, w_axis, num_links, num_frames, height, width, mean=mean, std=std)
+
+
+_NUM_FRAMES = 145
+_HEIGHT = 1088
+_WIDTH = 1920
+
+
+def _build_tt_decoder(mesh):
+ """Production decoder with random weights; torch module built for its
+ state_dict only (no host forward)."""
+ vae_mods = _require_diffusers_ltx_vae()
+ if not vae_mods["ltx2"]:
+ pytest.skip("LTX-2 decoder profiling requires diffusers autoencoder_kl_ltx2")
+
+ torch.manual_seed(42)
+ torch_decoder = _TorchLTXVideoDecoder(
+ decoder_blocks=_LTX_PROD_DECODER_BLOCKS,
+ in_channels=128,
+ out_channels=3,
+ patch_size=4,
+ base_channels=128,
+ causal=False,
+ spatial_padding_mode="zeros",
+ vae_mods=vae_mods,
+ )
+ torch_decoder.eval()
+ state = torch_decoder.state_dict()
+ # Identity per-channel stats so denorm is a no-op (random weights give garbage stats).
+ state["per_channel_statistics.mean-of-means"] = torch.zeros(128)
+ state["per_channel_statistics.std-of-means"] = torch.ones(128)
+ torch_decoder.load_state_dict(state)
+
+ # num_links=2 mirrors the BH production device config; NP_LINKS probes the W-transport bound.
+ ccl_manager = CCLManager(mesh, topology=ttnn.Topology.Linear, num_links=int(os.environ.get("NP_LINKS", "2")))
+ parallel_config = VaeHWParallelConfig(
+ height_parallel=ParallelFactor(factor=tuple(mesh.shape)[0], mesh_axis=0),
+ width_parallel=ParallelFactor(factor=tuple(mesh.shape)[1], mesh_axis=1),
+ )
+ tt_decoder = LTXVideoDecoder(
+ decoder_blocks=_LTX_PROD_DECODER_BLOCKS,
+ in_channels=128,
+ out_channels=3,
+ patch_size=4,
+ base_channels=128,
+ causal=False,
+ num_frames=_NUM_FRAMES,
+ height=_HEIGHT,
+ width=_WIDTH,
+ mesh_device=mesh,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+ tt_decoder.load_torch_state_dict(_diffusers_decoder_state_to_tt(torch_decoder.state_dict()))
+ return tt_decoder
+
+
+def _latent():
+ latent_frames = (_NUM_FRAMES - 1) // 8 + 1
+ return torch.randn(1, 128, latent_frames, _HEIGHT // 32, _WIDTH // 32, dtype=torch.float32)
+
+
+@pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(4, 8)], indirect=True)
+def test_prof_vae_ltx_devicetime(mesh_device, device_params):
+ mesh = mesh_device.create_submesh(ttnn.MeshShape(4, 8))
+ tt_decoder = _build_tt_decoder(mesh)
+ latent = _latent()
+
+ # One forward only: the decoder emits ~10x a transformer block's ops, so warm-ups
+ # would overflow the profiler DRAM buffer. Device kernel durations are warm-independent.
+ signpost("start")
+ _ = tt_decoder(latent)
+ ttnn.synchronize_device(mesh)
+ signpost("stop")
+ ttnn.ReadDeviceProfiler(mesh)
+
+
+@pytest.mark.parametrize("num_frames, height, width", _LTX_DECODER_2K_SHAPE_PARAMS)
+@pytest.mark.parametrize("mean, std", [(0, 1)])
+@pytest.mark.parametrize(
+ ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"),
+ LTX_VAE_DECODER_MULTI_ONLY_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_video_decoder_2k(
+ mesh_device: ttnn.MeshDevice,
+ num_frames: int,
+ height: int,
+ width: int,
+ mean: float,
+ std: float,
+ h_axis: int,
+ w_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+):
+ """Full LTXVideoDecoder at the 2K production resolution (2048x1080 -> 1088x2048).
+
+ Multi-device only (2x4 is the production layout; 4x8 also exercised). The two
+ 512->4096 / 1024->4096 upsample-projection convs are not in the 2K ``_BLOCKINGS``
+ table (they OOM a 2x4 mesh at full production T), so at 2K they fall back to the
+ generic blockings — correct, but not yet perf-tuned. See conv3d.py.
+ """
+ _run_ltx_decoder_parity(mesh_device, h_axis, w_axis, num_links, num_frames, height, width, mean=mean, std=std)
+
+
+# ---------------------------------------------------------------------------
+# Encoder parity (LTXVideoEncoder vs a diffusers-primitives reference encoder)
+# ---------------------------------------------------------------------------
+def _diffusers_encoder_state_to_tt(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Remap diffusers LTX-2 encoder keys (incl. nested resnets) to the TT/Lightricks layout."""
+ out: dict[str, torch.Tensor] = {}
+ for key, value in state.items():
+ if key.endswith("conv_shortcut.weight"):
+ out[key.replace("conv_shortcut.weight", "conv_shortcut.conv.weight")] = value
+ elif key.endswith("conv_shortcut.bias"):
+ out[key.replace("conv_shortcut.bias", "conv_shortcut.conv.bias")] = value
+ else:
+ out[key] = value
+ return out
+
+
+class _TorchLTXVideoEncoder(nn.Module):
+ """Block-list LTX encoder built from diffusers VAE primitives (matches TT/Lightricks layout).
+
+ Mirrors the reference ``VideoEncoder`` encode path: patchify -> conv_in -> down_blocks
+ (``res_x`` mid blocks + ``compress_*_res`` space-to-depth downsamplers) -> PixelNorm + SiLU
+ -> conv_out (latent means) -> per-channel normalize. Always causal.
+ """
+
+ _STRIDE_MAP = {
+ "compress_time_res": (2, 1, 1),
+ "compress_space_res": (1, 2, 2),
+ "compress_all_res": (2, 2, 2),
+ }
+
+ def __init__(
+ self,
+ *,
+ encoder_blocks: list[tuple[str, dict | int]],
+ in_channels: int = 3,
+ out_channels: int = 128,
+ patch_size: int = 4,
+ base_channels: int = 128,
+ causal: bool = True,
+ spatial_padding_mode: str = "zeros",
+ vae_mods: dict,
+ ) -> None:
+ super().__init__()
+ self.patch_size = patch_size
+ self.causal = causal
+ self.ltx2 = vae_mods["ltx2"]
+ causal_conv = vae_mods["causal_conv"]
+ resnet_cls = vae_mods["resnet"]
+ downsample_cls = vae_mods["downsample"]
+ pixel_norm = vae_mods["pixel_norm"]
+
+ in_channels_patched = in_channels * patch_size**2
+ feature_channels = base_channels
+
+ conv_kwargs: dict[str, Any] = {}
+ if self.ltx2:
+ if spatial_padding_mode != "zeros":
+ conv_kwargs["spatial_padding_mode"] = spatial_padding_mode
+ else:
+ conv_kwargs["is_causal"] = causal
+
+ self.conv_in = causal_conv(
+ in_channels=in_channels_patched,
+ out_channels=feature_channels,
+ kernel_size=3,
+ stride=1,
+ **conv_kwargs,
+ )
+
+ self.down_blocks = nn.ModuleList()
+ ch = feature_channels
+ for block_name, block_params in encoder_blocks:
+ block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
+ if block_name == "res_x":
+ self.down_blocks.append(
+ _TorchUNetMidBlock3D(
+ in_channels=ch,
+ num_layers=block_config["num_layers"],
+ resnet_cls=resnet_cls,
+ ltx2=self.ltx2,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ elif block_name in self._STRIDE_MAP:
+ multiplier = block_config.get("multiplier", 2)
+ new_ch = ch * multiplier
+ downsample_kwargs: dict[str, Any] = {
+ "in_channels": ch,
+ "out_channels": new_ch,
+ "stride": self._STRIDE_MAP[block_name],
+ }
+ if self.ltx2 and spatial_padding_mode != "zeros":
+ downsample_kwargs["spatial_padding_mode"] = spatial_padding_mode
+ self.down_blocks.append(downsample_cls(**downsample_kwargs))
+ ch = new_ch
+ else:
+ raise ValueError(f"Unknown encoder block: {block_name}")
+
+ self.conv_norm_out = pixel_norm()
+ self.conv_act = nn.SiLU()
+ self.conv_out = causal_conv(
+ in_channels=ch,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=1,
+ **conv_kwargs,
+ )
+ self.per_channel_statistics = _PerChannelStatistics(out_channels)
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ # Crop trailing frames so F == 1 + 8*k (mirrors the reference; no-op for the test shapes).
+ frames_count = sample.shape[2]
+ if (frames_count - 1) % 8 != 0:
+ sample = sample[:, :, : frames_count - ((frames_count - 1) % 8)]
+
+ sample = _patchify(sample, patch_size_hw=self.patch_size)
+ sample = _run_diffusers_causal_conv(self.conv_in, sample, causal=self.causal, ltx2=self.ltx2)
+ for down_block in self.down_blocks:
+ if isinstance(down_block, _TorchUNetMidBlock3D):
+ sample = down_block(sample, causal=self.causal, ltx2=self.ltx2)
+ else:
+ sample = _run_diffusers_downsample(down_block, sample, causal=self.causal, ltx2=self.ltx2)
+ if self.ltx2:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample.movedim(1, -1)).movedim(-1, 1)
+ sample = self.conv_act(sample)
+ sample = _run_diffusers_causal_conv(self.conv_out, sample, causal=self.causal, ltx2=self.ltx2)
+ return self.per_channel_statistics.normalize(sample)
+
+
+# Small encoder config: conv_in (48->128) -> res_x -> compress_space_res -> compress_time_res.
+# Avoids resnet channel-shortcut paths (res_x keeps channels; compress_*_res change channels via
+# space-to-depth), so the parity check isolates the new encoder primitives.
+_LTX_ENCODER_SMALL_BLOCKS = [
+ ("res_x", {"num_layers": 1}),
+ ("compress_space_res", {"multiplier": 2}),
+ ("compress_time_res", {"multiplier": 2}),
+]
+
+
+@pytest.mark.parametrize(
+ "num_frames, height, width",
+ [(9, 64, 64)],
+ ids=["9f_64x64"],
+)
+@pytest.mark.parametrize(
+ "mesh_device, device_params",
+ _LTX_VAE_MESH_DEVICE_PARAMS,
+ ids=_LTX_VAE_MESH_DEVICE_IDS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_video_encoder(mesh_device: ttnn.MeshDevice, num_frames: int, height: int, width: int):
+ """LTXVideoEncoder parity vs a diffusers-primitives reference encoder.
+
+ Random weights — checks the encode path (patchify -> conv_in -> down_blocks -> PixelNorm+SiLU
+ -> conv_out (means) -> per-channel normalize) matches the reference numerically.
+ """
+ vae_mods = _require_diffusers_ltx_vae()
+ if not vae_mods["ltx2"]:
+ pytest.skip("Encoder parity requires diffusers autoencoder_kl_ltx2 (runtime causal= support)")
+
+ torch.manual_seed(42)
+ latent_channels = 128
+
+ ref = _TorchLTXVideoEncoder(
+ encoder_blocks=_LTX_ENCODER_SMALL_BLOCKS,
+ in_channels=3,
+ out_channels=latent_channels,
+ patch_size=4,
+ base_channels=latent_channels,
+ causal=True,
+ spatial_padding_mode="zeros",
+ vae_mods=vae_mods,
+ )
+ # Set known per-channel stats so normalize() = (x - mean) / std is deterministic/well-conditioned.
+ with torch.no_grad():
+ ref.per_channel_statistics.get_buffer("mean-of-means").copy_(torch.randn(latent_channels))
+ ref.per_channel_statistics.get_buffer("std-of-means").copy_(torch.rand(latent_channels) + 0.5)
+ ref.eval()
+
+ x = torch.randn(1, 3, num_frames, height, width, dtype=torch.float32)
+ with torch.no_grad():
+ ref_out = ref(x.clone()) # (1, 128, F', H', W')
+
+ tt_model = LTXVideoEncoder(
+ encoder_blocks=_LTX_ENCODER_SMALL_BLOCKS,
+ in_channels=3,
+ out_channels=latent_channels,
+ patch_size=4,
+ mesh_device=mesh_device,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ **_vae_parallel_kwargs(mesh_device),
+ )
+ tt_model.load_torch_state_dict(_diffusers_encoder_state_to_tt(ref.state_dict()))
+
+ tt_out = _unwrap_tt_output(tt_model(x.clone())) # host (B, C, F', H', W')
+
+ logger.info(f"ref out {tuple(ref_out.shape)}, TT out {tuple(tt_out.shape)}")
+ assert tt_out.shape == ref_out.shape, f"shape mismatch: TT {tt_out.shape} vs ref {ref_out.shape}"
+ assert torch.isfinite(tt_out).all(), "NaN/Inf in TT encoder output"
+ assert_quality(ref_out, tt_out, pcc=0.99)
+ logger.info("PASSED: LTXVideoEncoder matches diffusers reference encoder")
+
+
+# ---------------------------------------------------------------------------
+# Production encoder parity under sharding (regression for the I2V VAE-encoder bug)
+# ---------------------------------------------------------------------------
+# Full LTX-2.3 22B encoder block list (from the checkpoint config metadata). The two
+# ``compress_all_res`` blocks are NOT exercised by ``test_ltx_video_encoder`` above, and
+# only matter under spatial sharding — they fold a (2,2,2) patch into channels per device.
+_LTX_PROD_ENCODER_BLOCKS = [
+ ("res_x", {"num_layers": 4}),
+ ("compress_space_res", {"multiplier": 2}),
+ ("res_x", {"num_layers": 6}),
+ ("compress_time_res", {"multiplier": 2}),
+ ("res_x", {"num_layers": 4}),
+ ("compress_all_res", {"multiplier": 2}),
+ ("res_x", {"num_layers": 2}),
+ ("compress_all_res", {"multiplier": 1}),
+ ("res_x", {"num_layers": 2}),
+]
+
+
+def _run_ltx_encoder_parity(
+ mesh_device: ttnn.MeshDevice,
+ h_axis: int,
+ w_axis: int,
+ num_links: int,
+ num_frames: int,
+ height: int,
+ width: int,
+ *,
+ pcc: float = 0.99,
+):
+ """Build torch (diffusers reference) + TT production encoders at ``(num_frames, height, width)``
+ and assert numerical parity. The TT encoder shards H/W across the mesh exactly like the I2V
+ pipeline; this is the path that fails when a per-device spatial dim is not divisible by the
+ downsample stride (patch groups would straddle a shard boundary)."""
+ vae_mods = _require_diffusers_ltx_vae()
+ if not vae_mods["ltx2"]:
+ pytest.skip("Encoder parity requires diffusers autoencoder_kl_ltx2 (runtime causal= support)")
+
+ torch.manual_seed(42)
+ latent_channels = 128
+
+ ref = _TorchLTXVideoEncoder(
+ encoder_blocks=_LTX_PROD_ENCODER_BLOCKS,
+ in_channels=3,
+ out_channels=latent_channels,
+ patch_size=4,
+ base_channels=latent_channels,
+ causal=True,
+ spatial_padding_mode="zeros",
+ vae_mods=vae_mods,
+ )
+ with torch.no_grad():
+ ref.per_channel_statistics.get_buffer("mean-of-means").copy_(torch.randn(latent_channels))
+ ref.per_channel_statistics.get_buffer("std-of-means").copy_(torch.rand(latent_channels) + 0.5)
+ ref.eval()
+
+ x = torch.randn(1, 3, num_frames, height, width, dtype=torch.float32)
+ with torch.no_grad():
+ ref_out = ref(x.clone()) # (1, 128, F', H', W')
+
+ ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear, num_links=num_links)
+ parallel_config = VaeHWParallelConfig(
+ height_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[h_axis], mesh_axis=h_axis),
+ width_parallel=ParallelFactor(factor=tuple(mesh_device.shape)[w_axis], mesh_axis=w_axis),
+ )
+
+ tt_model = LTXVideoEncoder(
+ encoder_blocks=_LTX_PROD_ENCODER_BLOCKS,
+ in_channels=3,
+ out_channels=latent_channels,
+ patch_size=4,
+ mesh_device=mesh_device,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+ tt_model.load_torch_state_dict(_diffusers_encoder_state_to_tt(ref.state_dict()))
+
+ tt_out = _unwrap_tt_output(tt_model(x.clone()))
+
+ logger.info(f"prod encoder: ref out {tuple(ref_out.shape)}, TT out {tuple(tt_out.shape)}")
+ assert tt_out.shape == ref_out.shape, f"shape mismatch: TT {tt_out.shape} vs ref {ref_out.shape}"
+ assert torch.isfinite(tt_out).all(), "NaN/Inf in TT encoder output"
+ assert_quality(ref_out, tt_out, pcc=pcc)
+ logger.info("PASSED: production LTXVideoEncoder matches diffusers reference encoder")
+
+
+# (num_frames, H, W). On a 4x8 mesh (H//4, W//8) the per-device spatial dim after patchify and
+# the early downsamples decides whether the last compress_all_res can fold locally:
+# - "clean": per-device dims stay even through every fold (fast local path).
+# - "boundary": per-device dim becomes odd at the last compress_all_res, so the space-to-depth
+# patch would straddle a shard boundary — the case that corrupts the I2V latent.
+_LTX_ENCODER_PROD_SHAPE_PARAMS = [
+ pytest.param(1, 128, 256, id="1f_128x256_clean"),
+ pytest.param(1, 192, 384, id="1f_192x384_boundary"),
+ # The two real I2V conditioning-image resolutions (single frame): stage-0 (544x960 -> latent
+ # 17x30) and stage-1 (1088x1920 -> latent 34x60). Both put an odd per-device spatial dim into
+ # the last compress_all_res on 4x8 — this is the exact case that corrupted the I2V latent.
+ pytest.param(1, 544, 960, id="1f_544x960_stage0"),
+ pytest.param(1, 1088, 1920, id="1f_1088x1920_stage1"),
+]
+
+
+@pytest.mark.parametrize("num_frames, height, width", _LTX_ENCODER_PROD_SHAPE_PARAMS)
+@pytest.mark.parametrize(
+ ("mesh_device", "h_axis", "w_axis", "num_links", "device_params", "topology"),
+ LTX_VAE_ENCODER_PROD_MESH_PARAMS,
+ indirect=["mesh_device", "device_params"],
+)
+def test_ltx_video_encoder_prod(
+ mesh_device: ttnn.MeshDevice,
+ num_frames: int,
+ height: int,
+ width: int,
+ h_axis: int,
+ w_axis: int,
+ num_links: int,
+ topology: ttnn.Topology,
+):
+ """Production-config LTXVideoEncoder parity vs the diffusers reference, sharded like the I2V pipeline.
+
+ The "boundary" resolution reproduces the I2V conditioning-latent corruption (device-vs-host
+ PCC ~0.3): a per-device spatial dim becomes odd at the second ``compress_all_res``, so the
+ per-device space-to-depth folds a patch that should straddle two shards.
+ """
+ _run_ltx_encoder_parity(mesh_device, h_axis, w_axis, num_links, num_frames, height, width)
diff --git a/code/models/tt_dit/tests/models/minimax_h3/.gitkeep b/code/models/tt_dit/tests/models/minimax_h3/.gitkeep
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/code/models/tt_dit/tests/models/minimax_h3/common.py b/code/models/tt_dit/tests/models/minimax_h3/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..5cddf6de79636766469a7cb90d718fbefc963e14
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/common.py
@@ -0,0 +1,330 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared helpers for the MiniMax-H3 bringup tests."""
+
+import json
+import math
+import os
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+import ttnn
+
+from ....utils.tensor import from_torch
+from ....utils.test import ring_params_8k_req_exact_devices, ring_params_req_exact_devices
+
+# Fixed VAE work units: encoder (17, 256, 256) tiles, decoder (7, 16, 16) latent chunks.
+TILE = 256
+CLIP_FRAMES = 17
+LATENT_TILE = 16
+DECODE_LATENT_FRAMES = 7
+
+
+def weights_subdir(subfolder: str) -> str | None:
+ base = os.environ.get("MINIMAX_H3_MODEL_PATH")
+ if not base:
+ return None
+ candidate = os.path.join(base, subfolder)
+ return candidate if os.path.isfile(os.path.join(candidate, "config.json")) else None
+
+
+def load_config(weights_dir: str) -> dict:
+ return {
+ k: v
+ for k, v in json.loads(open(os.path.join(weights_dir, "config.json")).read()).items()
+ if not k.startswith("_")
+ }
+
+
+def _reference_class(name: str):
+ pytest.importorskip("diffusers", reason="pinned diffusers reference not installed")
+ from diffusers.models.autoencoders import autoencoder_kl_minimax_h3 as ref
+
+ return getattr(ref, name)
+
+
+def random_encoder_state(config: dict) -> dict:
+ cls = _reference_class("MiniMaxH3VideoEncoder3d")
+ module = cls(
+ in_channels=3,
+ out_channels=2 * config["latent_channels"],
+ block_out_channels=tuple(config["block_out_channels"]),
+ layers_per_block=config["layers_per_block"],
+ spatial_downsample_factors=tuple(config["spatial_downsample_factors"]),
+ temporal_downsample_factors=tuple(config["temporal_downsample_factors"]),
+ norm_num_groups=config["norm_num_groups"],
+ norm_eps=config["norm_eps"],
+ spatial_padding_mode=config["spatial_padding_mode"],
+ )
+ return dict(module.state_dict())
+
+
+def random_decoder_state(config: dict, *, num_layers: int | None = None) -> dict:
+ cls = _reference_class("MiniMaxH3VideoViTDecoder3d")
+ module = cls(
+ in_channels=config["latent_channels"],
+ out_channels=config["out_channels"],
+ patch_size=16,
+ patch_size_t=4,
+ num_layers=config["decoder_num_layers"] if num_layers is None else num_layers,
+ num_attention_heads=config["decoder_num_attention_heads"],
+ attention_head_dim=config["decoder_attention_head_dim"],
+ num_register_tokens=config["decoder_num_register_tokens"],
+ ffn_mult=config["decoder_ffn_mult"],
+ rope_theta=config["decoder_rope_theta"],
+ rope_dim_ratio=config["decoder_rope_dim_ratio"],
+ norm_eps=config["decoder_norm_eps"],
+ )
+ return dict(module.state_dict())
+
+
+def psnr(reference: torch.Tensor, test: torch.Tensor) -> float:
+ """PSNR in dB, with the peak taken from the reference's own range."""
+ mse = torch.mean((reference.float() - test.float()) ** 2).item()
+ if mse == 0.0:
+ return float("inf")
+ peak = reference.abs().max().item()
+ return float("inf") if peak == 0.0 else 20.0 * math.log10(peak) - 10.0 * math.log10(mse)
+
+
+def create_fractal_image(width: int, height: int) -> Image.Image:
+ """Mandelbrot I2V seed: content the model would never generate, so a pipeline that ignores the keyframe cannot pass."""
+ c = np.linspace(-2.0, 1.0, width)[None, :] + 1j * np.linspace(-1.5, 1.5, height)[:, None]
+ z = np.zeros_like(c)
+ img = np.zeros(c.shape, dtype=np.uint8)
+ for i in range(32):
+ z = z * z + c
+ img[(img == 0) & (np.abs(z) > 2)] = 255 - 8 * i
+ return Image.fromarray(np.dstack((img, np.roll(img, width // 10, 1), np.roll(img, height // 10, 0))), "RGB")
+
+
+# The mesh shapes MiniMax-H3 is tuned for, named individually and composed into a list -- the shape
+# `ltx_mesh_params.py` uses. Both carry `require_exact_physical_num_devices`, so exactly one runs on
+# any given cluster and the other skips.
+#
+# The fabric must be FABRIC_1D_RING: `CCLManager` runs ring collectives, and on a line fabric one
+# cannot resolve a forwarding direction (`TT_FATAL fabric.cpp:174 forwarding_direction.has_value()`).
+# 4x32 additionally takes the 8 KB router payload, matching Wan's 4x32 rows, and a trace region for
+# the quad's `trace_denoise`; the region is only reserved, so 4x8 pays nothing but address space.
+_L1_SMALL = 65536
+_ring = {**ring_params_req_exact_devices, "l1_small_size": _L1_SMALL}
+_ring_8k = {**ring_params_8k_req_exact_devices, "l1_small_size": _L1_SMALL}
+_ring_8k_trace = {**ring_params_8k_req_exact_devices, "trace_region_size": 150_000_000, "l1_small_size": _L1_SMALL}
+
+MESH_4X8_RING = pytest.param((4, 8), _ring_8k, id="4x8")
+MESH_4X32_RING = pytest.param((4, 32), _ring_8k_trace, id="4x32")
+
+GALAXY_MESHES = [MESH_4X8_RING, MESH_4X32_RING]
+
+
+def randomize_norm_weights(module: torch.nn.Module, *, scale: float = 0.5) -> torch.nn.Module:
+ """Randomize every `nn.RMSNorm` affine in place, BEFORE taking `state_dict`: all-ones norms make norm-weight loading invisible to PCC."""
+ for submodule in module.modules():
+ if isinstance(submodule, torch.nn.RMSNorm) and submodule.weight is not None:
+ submodule.weight.data = 1.0 + scale * torch.randn_like(submodule.weight.data)
+ return module
+
+
+# ---- transformer test fixtures, shared by the correctness and perf tests so they cannot drift ----
+
+# The transformer tests take the axes explicitly: TP stays on axis 0 at factor 4 and SP absorbs the
+# rest, so 4x8 -> 4x32 moves only `sp_factor`, which every test body derives from `mesh_device.shape`.
+# `device_params` travels inside the tuple because the router payload differs per shape; crossing them
+# independently would pair a 4x8 mesh with the 4x32 router config.
+# The transformer tests take the axes explicitly: TP stays on axis 0 at factor 4 and SP absorbs the
+# rest, so 4x8 -> 4x32 moves only `sp_factor`, which every test body derives from `mesh_device.shape`.
+# `device_params` travels inside the tuple because the router payload differs per shape; crossing them
+# independently would pair a 4x8 mesh with the 4x32 router config.
+GALAXY_RING = pytest.mark.parametrize(
+ ("mesh_device", "sp_axis", "tp_axis", "num_links", "device_params", "topology", "is_fsdp"),
+ [
+ # 4x8 takes the 8 KB router payload like 4x32: sp_sim runs on it emulate the 4x32 machine,
+ # and the exp ring SDPA's fabric all-gather packs 4 tiles per packet only at 8 KB.
+ pytest.param((4, 8), 1, 0, 2, _ring_8k, ttnn.Topology.Ring, False, id="4x8sp1tp0nl2_ring_is_fsdp0"),
+ pytest.param((4, 32), 1, 0, 2, _ring_8k_trace, ttnn.Topology.Ring, False, id="4x32sp1tp0nl2_ring_is_fsdp0"),
+ ],
+ indirect=["mesh_device", "device_params"],
+)
+
+REAL_BLOCK_CONFIG = dict(
+ hidden_size=5376,
+ num_attention_heads=56,
+ attention_head_dim=128,
+ ffn_dim=14336,
+ time_embed_dim=2688,
+ norm_eps=1e-5,
+ qk_norm_eps=1e-5,
+)
+ROPE_FREQ_DIM = 16
+ROPE_THETA = 10000.0
+
+TT_BLOCK_CONFIG = dict(
+ hidden_size=REAL_BLOCK_CONFIG["hidden_size"],
+ num_heads=REAL_BLOCK_CONFIG["num_attention_heads"],
+ head_dim=REAL_BLOCK_CONFIG["attention_head_dim"],
+ ffn_dim=REAL_BLOCK_CONFIG["ffn_dim"],
+ time_embed_dim=REAL_BLOCK_CONFIG["time_embed_dim"],
+ norm_eps=REAL_BLOCK_CONFIG["norm_eps"],
+ qk_norm_eps=REAL_BLOCK_CONFIG["qk_norm_eps"],
+)
+
+# Reference token tags: 0 video, 1 text, 2 audio.
+TAG_VIDEO, TAG_TEXT, TAG_AUDIO = 0, 1, 2
+
+
+def packed_layout(
+ num_text: int,
+ num_audio: int,
+ num_video: int,
+ grid_hw: tuple[int, int] = (8, 8),
+ padded_len: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Build one packed layout `(position_ids, token_tags, timestep_indices)`: text, audio, video rows over two timesteps, zero-padded to `padded_len`."""
+ grid_h, grid_w = grid_hw
+ frame = grid_h * grid_w
+ assert num_video % frame == 0, "num_video must fill whole (h, w) frames"
+ grid_t = num_video // frame
+ assert grid_t >= 2, "need at least one conditioning frame and one target frame"
+
+ tags = torch.cat(
+ [
+ torch.full((num_text,), TAG_TEXT, dtype=torch.long),
+ torch.full((num_audio,), TAG_AUDIO, dtype=torch.long),
+ torch.full((num_video,), TAG_VIDEO, dtype=torch.long),
+ ]
+ )
+ timestep_indices = torch.cat(
+ [
+ torch.zeros(num_text, dtype=torch.long),
+ torch.ones(num_audio, dtype=torch.long),
+ torch.zeros(frame, dtype=torch.long),
+ torch.ones(num_video - frame, dtype=torch.long),
+ ]
+ )
+
+ vt, vh, vw = torch.meshgrid(torch.arange(grid_t), torch.arange(grid_h), torch.arange(grid_w), indexing="ij")
+ video_pos = torch.stack([vt.reshape(-1), vh.reshape(-1), vw.reshape(-1)], dim=-1)
+
+ def clock_pos(n: int) -> torch.Tensor:
+ return torch.stack([torch.arange(n), torch.zeros(n, dtype=torch.long), torch.zeros(n, dtype=torch.long)], -1)
+
+ position_ids = torch.cat([clock_pos(num_text), clock_pos(num_audio), video_pos], dim=0)
+
+ if padded_len is not None:
+ pad = padded_len - position_ids.shape[0]
+ if pad:
+ position_ids = torch.cat([position_ids, torch.zeros((pad, 3), dtype=position_ids.dtype)])
+ tags = torch.cat([tags, torch.zeros(pad, dtype=tags.dtype)])
+ timestep_indices = torch.cat([timestep_indices, torch.zeros(pad, dtype=timestep_indices.dtype)])
+ return position_ids, tags, timestep_indices
+
+
+def upload_rope(rope_cos: torch.Tensor, rope_sin: torch.Tensor, *, mesh_device, sp_axis: int):
+ """Upload prepared `(cos, sin)` tables: fp32, fractured on SP along the sequence, replicated on TP."""
+
+ def _upload(table: torch.Tensor) -> ttnn.Tensor:
+ return from_torch(
+ table.reshape(1, 1, *table.shape),
+ device=mesh_device,
+ dtype=ttnn.float32,
+ mesh_axes=[..., sp_axis, None],
+ )
+
+ return _upload(rope_cos), _upload(rope_sin)
+
+
+def build_visual_encoder(config: dict, mesh_device, num_frames: int, **overrides):
+ from ....models.vae.minimax_h3.encoder_minimax_h3 import MiniMaxH3Encoder3d
+
+ kwargs = dict(
+ num_frames=num_frames,
+ height=TILE,
+ width=TILE,
+ in_channels=3,
+ out_channels=2 * config["latent_channels"],
+ block_out_channels=tuple(config["block_out_channels"]),
+ layers_per_block=config["layers_per_block"],
+ spatial_downsample_factors=tuple(config["spatial_downsample_factors"]),
+ temporal_downsample_factors=tuple(config["temporal_downsample_factors"]),
+ mesh_device=mesh_device,
+ )
+ kwargs.update(overrides)
+ return MiniMaxH3Encoder3d(**kwargs)
+
+
+def build_visual_decoder(config: dict, mesh_device, **overrides):
+ from ....models.vae.minimax_h3.decoder_minimax_h3 import MiniMaxH3ViTDecoder3d
+
+ kwargs = dict(
+ num_frames=DECODE_LATENT_FRAMES,
+ height=LATENT_TILE,
+ width=LATENT_TILE,
+ in_channels=config["latent_channels"],
+ out_channels=config["out_channels"],
+ patch_size=16,
+ patch_size_t=4,
+ num_layers=config["decoder_num_layers"],
+ num_heads=config["decoder_num_attention_heads"],
+ head_dim=config["decoder_attention_head_dim"],
+ num_register_tokens=config["decoder_num_register_tokens"],
+ ffn_mult=config["decoder_ffn_mult"],
+ rope_theta=config["decoder_rope_theta"],
+ rope_dim_ratio=config["decoder_rope_dim_ratio"],
+ eps=config["decoder_norm_eps"],
+ mesh_device=mesh_device,
+ )
+ kwargs.update(overrides)
+ return MiniMaxH3ViTDecoder3d(**kwargs)
+
+
+def build_audio_decoder(config: dict, mesh_device, **overrides):
+ from ....models.audio_vae.minimax_h3.decoder_minimax_h3_audio import MiniMaxH3AudioDecoder
+
+ kwargs = dict(
+ latent_channels=config["latent_channels"],
+ latent_dim=config["latent_dim"],
+ decoder_dim=config["decoder_dim"],
+ decoder_rates=tuple(config["decoder_rates"]),
+ decoder_kernel_sizes=tuple(config["decoder_kernel_sizes"]),
+ resblock_kernel_sizes=tuple(config["resblock_kernel_sizes"]),
+ resblock_dilation_sizes=tuple(tuple(d) for d in config["resblock_dilation_sizes"]),
+ mesh_device=mesh_device,
+ )
+ kwargs.update(overrides)
+ return MiniMaxH3AudioDecoder(**kwargs)
+
+
+CONDITIONER_SUBFOLDER = "text_encoder"
+
+
+def conditioner_checkpoint_dir(patterns: list[str]) -> str:
+ """Resolve the conditioner directory from $MINIMAX_H3_MODEL_PATH; a missing checkpoint skips.
+
+ `patterns` are the shard patterns the test needs present (validation only, nothing is fetched).
+ """
+ import glob
+
+ root = os.environ.get("MINIMAX_H3_MODEL_PATH", "")
+ if not root or not os.path.isdir(root):
+ pytest.skip("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot")
+ missing = [pattern for pattern in patterns if not glob.glob(os.path.join(root, pattern))]
+ if missing:
+ pytest.skip(f"MiniMax-H3 conditioner checkpoint at {root} is missing {missing}")
+ return os.path.join(root, CONDITIONER_SUBFOLDER)
+
+
+def load_reference_conditioner(path: str):
+ """Load the conditioner with load-info checked: a silently partial load would pass parity while testing the fresh init."""
+ import transformers
+
+ hf, info = transformers.Qwen3VLForConditionalGeneration.from_pretrained(
+ path, dtype=torch.bfloat16, output_loading_info=True
+ )
+ bad = {k: sorted(info[k])[:5] for k in ("missing_keys", "unexpected_keys", "mismatched_keys") if info[k]}
+ assert not bad, f"conditioner load key mismatch: {bad}"
+ return hf
diff --git a/code/models/tt_dit/tests/models/minimax_h3/common_av.py b/code/models/tt_dit/tests/models/minimax_h3/common_av.py
new file mode 100644
index 0000000000000000000000000000000000000000..320a5496047ec0a01bf7741bdf166d15f21d8e0a
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/common_av.py
@@ -0,0 +1,563 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Reference-free A/V sanity checks plus the scaffolding shared by the MiniMax-H3 e2e gates.
+A/V sync is checked structurally (duration/ordering); envelope-vs-motion correlation is diagnostic only."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+
+from ....pipelines.minimax_h3.packing import MINIMAX_H3_FPS
+
+
+def check_audio_sanity(audio, *, sampling_rate, expected_seconds, tolerance_seconds=0.05):
+ """Guard against a soundtrack that is silent, clipped, constant, or the wrong length."""
+ if isinstance(audio, torch.Tensor):
+ audio = audio.detach().cpu().numpy()
+ audio = np.asarray(audio)
+ if audio.ndim == 3:
+ assert audio.shape[0] == 1, f"expected a single generation, got batch {audio.shape[0]}"
+ audio = audio[0]
+ assert audio.ndim == 2, f"expected (channels, samples), got shape {audio.shape}"
+
+ channels, samples = audio.shape
+ assert channels == 2, f"H3 generates stereo; got {channels} channel(s)"
+ assert np.isfinite(audio).all(), "soundtrack contains NaN/Inf"
+
+ seconds = samples / sampling_rate
+ assert abs(seconds - expected_seconds) <= tolerance_seconds, (
+ f"soundtrack is {seconds:.3f} s against {expected_seconds:.3f} s of video "
+ f"(off by {seconds - expected_seconds:+.3f} s, tolerance {tolerance_seconds:.3f} s)"
+ )
+
+ peak = float(np.abs(audio).max())
+ rms = float(np.sqrt((audio.astype(np.float64) ** 2).mean()))
+ assert peak > 1e-3, f"soundtrack is silent (peak {peak:.2e})"
+ assert rms > 1e-4, f"soundtrack is near-silent (rms {rms:.2e})"
+ for index in range(channels):
+ channel_std = float(audio[index].std())
+ assert channel_std > 1e-4, f"channel {index} is constant (std {channel_std:.2e})"
+ # Widespread clipping means the denormalization is wrong, not that the mix is loud.
+ clipped = float((np.abs(audio) >= 0.999).mean())
+ assert clipped < 0.01, f"{clipped:.1%} of samples are at full scale; suspect a scaling error"
+
+ logger.info(
+ f"Audio sanity OK: {channels}ch {seconds:.3f} s @ {sampling_rate} Hz, "
+ f"peak={peak:.3f}, rms={rms:.4f}, clipped={clipped:.3%}"
+ )
+
+
+def check_av_sync(frames, audio, *, sampling_rate, fps, tolerance_seconds=0.05):
+ """Durations must agree and stereo channels must differ -- identical channels mean
+ `unpack_audio_tokens` collapsed the channel-major layout. Envelope-motion lag is logged only."""
+ if isinstance(frames, torch.Tensor):
+ frames = frames.detach().cpu().numpy()
+ if isinstance(audio, torch.Tensor):
+ audio = audio.detach().cpu().numpy()
+ frames = np.asarray(frames)
+ audio = np.asarray(audio)
+ if audio.ndim == 3:
+ audio = audio[0]
+
+ num_frames = frames.shape[0]
+ video_seconds = num_frames / fps
+ audio_seconds = audio.shape[-1] / sampling_rate
+ assert num_frames > 1, f"need more than one frame to talk about sync, got {num_frames}"
+ assert audio.shape[-1] > sampling_rate // fps, "soundtrack is shorter than a single frame"
+ assert abs(video_seconds - audio_seconds) <= tolerance_seconds, (
+ f"video is {video_seconds:.3f} s but audio is {audio_seconds:.3f} s "
+ f"(off by {audio_seconds - video_seconds:+.3f} s)"
+ )
+
+ if audio.shape[0] == 2:
+ assert not np.allclose(audio[0], audio[1]), "stereo channels are identical; suspect audio row unpacking"
+
+ # Diagnostic only.
+ motion = np.abs(np.diff(frames.astype(np.float32), axis=0)).mean(axis=(1, 2, 3))
+ samples_per_frame = audio.shape[-1] / num_frames
+ envelope = np.array(
+ [
+ np.abs(audio[:, int(i * samples_per_frame) : int((i + 1) * samples_per_frame)]).mean()
+ for i in range(1, num_frames)
+ ]
+ )
+ lag_frames = 0
+ if motion.std() > 1e-8 and envelope.std() > 1e-8:
+ m = (motion - motion.mean()) / motion.std()
+ e = (envelope - envelope.mean()) / envelope.std()
+ correlation = np.correlate(m, e, mode="full") / len(m)
+ lag_frames = int(np.argmax(correlation)) - (len(m) - 1)
+ logger.info(
+ f"A/V envelope-motion best lag: {lag_frames:+d} frames ({lag_frames / fps:+.3f} s), "
+ f"peak r={correlation.max():.3f} (diagnostic, not asserted)"
+ )
+
+ logger.info(
+ f"A/V sync OK: video {video_seconds:.3f} s / {num_frames} frames @ {fps} fps, "
+ f"audio {audio_seconds:.3f} s @ {sampling_rate} Hz, delta {audio_seconds - video_seconds:+.4f} s"
+ )
+ return {"video_seconds": video_seconds, "audio_seconds": audio_seconds, "lag_frames": lag_frames}
+
+
+def check_spatial_seams(frames, *, vertical_boundaries, horizontal_boundaries, max_ratio=2.0):
+ """Gradient energy at the VAE's tile boundaries vs everywhere else (~1.0 = no seam; whole-frame means hide seams)."""
+ if isinstance(frames, torch.Tensor):
+ frames = frames.detach().cpu().numpy()
+ frames = np.asarray(frames).astype(np.float32)
+ if frames.ndim == 4:
+ frames = frames.mean(axis=-1)
+
+ def ratio(gradient, boundaries):
+ inside = np.array([b for b in boundaries if 1 <= b < len(gradient) - 1], dtype=int)
+ if not len(inside):
+ return float("nan")
+ # Exclude +-2 px around each boundary so a smeared seam cannot inflate its own baseline.
+ baseline = np.ones(len(gradient), dtype=bool)
+ for b in inside:
+ baseline[max(0, b - 2) : b + 3] = False
+ if not baseline.any() or gradient[baseline].mean() == 0:
+ return float("nan")
+ return float(gradient[inside].mean() / gradient[baseline].mean())
+
+ column_gradient = np.abs(np.diff(frames, axis=2)).mean(axis=(0, 1))
+ row_gradient = np.abs(np.diff(frames, axis=1)).mean(axis=(0, 2))
+ vertical = ratio(column_gradient, vertical_boundaries)
+ horizontal = ratio(row_gradient, horizontal_boundaries)
+
+ logger.info(
+ f"Spatial seam ratios (1.0 = no seam): vertical {vertical:.3f} at x={list(vertical_boundaries)}, "
+ f"horizontal {horizontal:.3f} at y={list(horizontal_boundaries)}"
+ )
+ for name, value in (("vertical", vertical), ("horizontal", horizontal)):
+ if np.isfinite(value):
+ assert value < max_ratio, (
+ f"{name} tile-boundary gradient is {value:.2f}x the surrounding image; "
+ "suspect the tile blend extent or per-tile normalization (artifact rubric: seams)"
+ )
+ return {"vertical": vertical, "horizontal": horizontal}
+
+
+def log_spectral_flatness(audio, *, sampling_rate, num_bands=64):
+ """Log-spectrum shape, logged only: white noise drives flatness toward 1.0, a stuck tone toward 0."""
+ if isinstance(audio, torch.Tensor):
+ audio = audio.detach().cpu().numpy()
+ audio = np.asarray(audio)
+ if audio.ndim == 3:
+ audio = audio[0]
+ mono = audio.mean(axis=0).astype(np.float64)
+
+ window = 2048
+ hop = window // 2
+ frames = [mono[i : i + window] for i in range(0, max(1, len(mono) - window), hop)]
+ if not frames:
+ return {"flatness": float("nan")}
+ spectrum = np.abs(np.fft.rfft(np.stack(frames) * np.hanning(window), axis=-1)) ** 2
+ power = spectrum.mean(axis=0)[1:]
+ flatness = float(np.exp(np.log(power + 1e-20).mean()) / (power.mean() + 1e-20))
+ band_edges = np.linspace(0, len(power), num_bands + 1).astype(int)
+ bands = np.array([power[a:b].mean() for a, b in zip(band_edges[:-1], band_edges[1:]) if b > a])
+ logger.info(
+ f"Audio log-spectrum: flatness={flatness:.4f}, "
+ f"band dB range=[{10 * np.log10(bands.min() + 1e-20):.1f}, {10 * np.log10(bands.max() + 1e-20):.1f}]"
+ )
+ return {"flatness": flatness}
+
+
+def _ffmpeg():
+ exe = shutil.which("ffmpeg")
+ if exe:
+ return exe
+ try:
+ import imageio_ffmpeg
+ except ImportError:
+ return None
+ return imageio_ffmpeg.get_ffmpeg_exe()
+
+
+def write_artifacts(frames, audio, sampling_rate, directory: Path, stem: str = "t2va"):
+ import wave
+
+ paths = {}
+ wav_path = directory / f"{stem}.wav"
+ interleaved = np.ascontiguousarray(audio[0].T if audio.ndim == 3 else audio.T)
+ pcm = np.clip(interleaved, -1.0, 1.0)
+ with wave.open(str(wav_path), "wb") as handle:
+ handle.setnchannels(pcm.shape[1])
+ handle.setsampwidth(2)
+ handle.setframerate(sampling_rate)
+ handle.writeframes((pcm * 32767.0).astype(" dict:
+ exe = shutil.which("ffprobe")
+ if exe is None:
+ return {}
+ import json as _json
+
+ result = subprocess.run(
+ [
+ exe,
+ "-v",
+ "error",
+ "-show_entries",
+ "stream=codec_type,duration,nb_frames,sample_rate,channels,width,height",
+ "-of",
+ "json",
+ str(path),
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ streams = {}
+ for stream in _json.loads(result.stdout).get("streams", []):
+ streams[stream.get("codec_type", "?")] = stream
+ return streams
+
+
+# Only reached when ffprobe is unavailable.
+_FALLBACK_HEIGHT = 768
+_FALLBACK_WIDTH = 1344
+
+
+def decoded_frames(path: Path, count: int, height: int = 0, width: int = 0) -> np.ndarray:
+ """Sample `count` frames evenly out of the written file, as uint8 luma.
+
+ `height`/`width` are the caller's known canvas. They matter: the frame count is derived by
+ dividing the raw byte stream by the frame size, and the only other source for that size is
+ `probe_streams`, which needs `ffprobe` -- absent wherever ffmpeg comes from imageio_ffmpeg.
+ Without them this fell back to a hardcoded 768x1344, which silently miscounts every canvas
+ with a different pixel total (4:3 read 94 frames for a correct 124-frame file). 16:9 and 21:9
+ both happen to be 1032192 px, so the constant looked fine until a sweep left that budget.
+ """
+ exe = _ffmpeg()
+ if exe is None:
+ return np.empty((0,))
+ result = subprocess.run(
+ [
+ exe,
+ "-v",
+ "error",
+ # Count what the file STORES. Muxing AAC writes an mp4 edit list (the codec's priming
+ # shifts the timeline), and a decoder honouring it drops frames at the edges: a 362-frame
+ # write reads back as 362 stored / 361 in-container / 359 after the edit list. The gate
+ # asks "were the frames written", so the edit list must not be applied here.
+ "-ignore_editlist",
+ "1",
+ "-i",
+ str(path),
+ # `select` with count<=1 is a no-op (`not(mod(n,1))` is always true) but combined with
+ # `-vsync 0` on an edit-list-shifted timeline it silently drops frames at the edges --
+ # 2 of them on a muxed 15 s clip. Only ask for the filter when actually subsampling.
+ *(["-vf", f"select='not(mod(n\\,{count}))'", "-vsync", "0"] if count > 1 else []),
+ "-f",
+ "rawvideo",
+ "-pix_fmt",
+ "gray",
+ "-",
+ ],
+ check=True,
+ capture_output=True,
+ )
+ probe = probe_streams(path).get("video", {})
+ height = int(probe.get("height") or height or _FALLBACK_HEIGHT)
+ width = int(probe.get("width") or width or _FALLBACK_WIDTH)
+ buffer = np.frombuffer(result.stdout, dtype=np.uint8)
+ usable = (buffer.size // (height * width)) * height * width
+ return buffer[:usable].reshape(-1, height, width)
+
+
+VBENCH_PYTHON = Path.home() / "vbench_env" / "bin" / "python"
+
+
+def run_vbench(video: Path, prompt: str, dimensions) -> dict[str, float]:
+ """Run VBench in its own interpreter (~/vbench_env): it pins numpy < 2 / transformers 4.33."""
+ if not VBENCH_PYTHON.is_file():
+ pytest.skip(
+ f"no VBench interpreter at {VBENCH_PYTHON}; create it with "
+ "`uv venv --python 3.10 ~/vbench_env && uv pip install --python ~/vbench_env/bin/python "
+ "vbench decord 'numpy==1.26.4' 'opencv-python-headless<4.11' 'setuptools<81'`"
+ )
+ interpreter = str(VBENCH_PYTHON)
+ runner = Path(__file__).with_name("tools") / "vbench_runner.py"
+ result = subprocess.run(
+ [interpreter, str(runner), str(video), ",".join(dimensions), "--prompt", prompt],
+ capture_output=True,
+ text=True,
+ timeout=5400,
+ )
+ marker = "VBENCH_JSON "
+ line = next((l for l in result.stdout.splitlines() if l.startswith(marker)), None)
+ if line is None:
+ raise AssertionError(
+ f"VBench produced no scores (exit {result.returncode}).\n"
+ f"stdout tail:\n{result.stdout[-2000:]}\nstderr tail:\n{result.stderr[-2000:]}"
+ )
+ import json as _json
+
+ return _json.loads(line[len(marker) :])
+
+
+def clip_prompt_alignment(frames: np.ndarray, prompt: str, num_frames: int = 8) -> dict[str, float]:
+ """Mean CLIP similarity between the prompt and evenly-spaced frames, x100."""
+ from PIL import Image
+
+ from ...dataset_eval.clip_encoder import CLIPEncoder
+
+ indices = np.linspace(0, frames.shape[0] - 1, num_frames).astype(int)
+ encoder = CLIPEncoder()
+ scores = [encoder.get_clip_score(prompt, Image.fromarray(frames[i])).item() * 100.0 for i in indices]
+ return {"mean": float(np.mean(scores)), "min": float(min(scores)), "max": float(max(scores))}
+
+
+def temporal_seam_score(frames: np.ndarray, period: int) -> float:
+ """Inter-frame delta at temporal chunk boundaries vs elsewhere (~1.0 = no seam)."""
+ deltas = np.abs(np.diff(frames.astype(np.float32), axis=0)).mean(axis=(1, 2))
+ if len(deltas) < 2 * period:
+ return float("nan")
+ index = np.arange(1, len(deltas) + 1)
+ at_boundary = deltas[index % period == 0]
+ elsewhere = deltas[index % period != 0]
+ if not len(at_boundary) or not elsewhere.mean():
+ return float("nan")
+ return float(at_boundary.mean() / elsewhere.mean())
+
+
+# ------------------------------------------------------------------ shared e2e gate scaffolding
+
+# Matched pair with the tier-6 bars (CLIP 37.37, imaging_quality 0.6896); imported by fl2va so it cannot drift.
+CALIBRATED_FOX_PROMPT = (
+ "A red fox trots across a snowy field at dawn, its breath visible in the cold air. "
+ "The low sun throws long blue shadows behind it, and loose snow lifts from each footfall."
+)
+
+
+def weights_dir(*required_subdirs: str) -> Path:
+ """The snapshot dir from MINIMAX_H3_MODEL_PATH; skips when it or a required partition is missing."""
+ root = os.environ.get("MINIMAX_H3_MODEL_PATH", "")
+ if not root or not Path(root).is_dir():
+ pytest.skip("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot")
+ directory = Path(root)
+ missing = [name for name in required_subdirs if not (directory / name).is_dir()]
+ if missing:
+ pytest.skip(f"MiniMax-H3 snapshot at {directory} is missing {missing}")
+ return directory
+
+
+def artifact_dir(name: str) -> Path:
+ """The gate's artifact directory `~/{name}`, created if absent."""
+ directory = Path.home() / name
+ directory.mkdir(parents=True, exist_ok=True)
+ return directory
+
+
+def run_warm_generation(pipeline, prompt: str, *, seed: int, **gen_kwargs):
+ """Warmup then the timed generation with identical kwargs; asserts padded-length agreement (programs are keyed on it)."""
+ pipeline.warmup(prompt=prompt, **gen_kwargs)
+ warm_padded_len = pipeline.last_padded_len
+
+ # The warmup is a full generation too, so its stage timings are the *cold* numbers -- program
+ # compilation, per-shape conv3d blocking and buffer allocation all land in them. Log them
+ # before the measured call overwrites `last_timings`. Not gated: cold cost depends on what the
+ # weight/JIT caches already held, so it is a diagnostic, not a bar.
+ cold_rows = list(pipeline.last_timings)
+ if cold_rows:
+ cold_total = sum(seconds for _, seconds in cold_rows)
+ logger.info("WARMUP (cold: program compile + buffer alloc included, not a perf target)")
+ for row_label, seconds in cold_rows:
+ share = 100 * seconds / cold_total if cold_total else 0.0
+ logger.info(f" {row_label:<18} {seconds:8.1f} s ({share:4.1f} %)")
+ logger.info(f" {'Total (warmup)':<18} {cold_total:8.1f} s")
+
+ ttnn.synchronize_device(pipeline.mesh_device)
+ if ttnn.using_distributed_env():
+ ttnn.distributed_context_barrier()
+
+ output = pipeline(prompt, seed=seed, **gen_kwargs)
+
+ assert pipeline.last_padded_len == warm_padded_len, (
+ f"warmup ran at padded_len {warm_padded_len} but the measured call ran at "
+ f"{pipeline.last_padded_len}; this number is not warm"
+ )
+ return output
+
+
+def log_timing_table(pipeline, label: str, num_forwards: int, video_seconds: float, expected_total_s=None, extra=""):
+ """The MEASUREMENT block; `expected_total_s`, when given, asserts the total. Returns the total."""
+ rows = pipeline.last_timings
+ total = sum(seconds for _, seconds in rows)
+ shape = tuple(pipeline.mesh_device.shape)
+ logger.info(
+ f"MEASUREMENT {label} fully warm | mesh {shape[0]}x{shape[1]} Blackhole, "
+ f"TP={pipeline.tp_factor} axis {pipeline.tp_axis} / SP={pipeline.sp_factor} axis {pipeline.sp_axis}, "
+ f"{pipeline.ccl_manager.topology}, {pipeline.ccl_manager.num_links} links{extra} "
+ f"| warm window: one full warmup generation at this shape, prepares and export excluded"
+ )
+ for row_label, seconds in rows:
+ logger.info(f" {row_label:<18} {seconds:8.1f} s ({100 * seconds / total:4.1f} %)")
+ logger.info(f" {'Total (compute)':<18} {total:8.1f} s")
+ denoise = dict(rows).get("Denoise")
+ if denoise:
+ logger.info(
+ f" per forward {denoise / num_forwards * 1000:8.1f} ms "
+ f"({num_forwards} forwards over {denoise:.1f} s)"
+ )
+ logger.info(f" realtime factor {total / video_seconds:8.1f} x (compute / video seconds)")
+ if expected_total_s is not None:
+ assert (
+ total < expected_total_s
+ ), f"fully-warm total {total:.1f} s exceeds the {expected_total_s:.0f} s floor bar"
+ return total
+
+
+def to_uint8_frames(output) -> np.ndarray:
+ """`(1, 3, F, H, W)` [0, 1] -> `(F, H, W, 3)` uint8; clamp + round first -- a bare astype wraps out-of-range values."""
+ video = output.video
+ assert video.ndim == 5 and video.shape[0] == 1, f"unexpected video shape {tuple(video.shape)}"
+ frames = video[0].permute(1, 2, 3, 0).clamp(0, 1).mul(255).round().to(torch.uint8)
+ return frames.cpu().numpy()
+
+
+def check_written_file(paths: dict, expected_frames: int, seam_period: int = 17, height: int = 0, width: int = 0):
+ """Gate the written *file* (streams, A/V skew, frame count, temporal seam); silently no-ops without an mp4."""
+ if "mp4" not in paths:
+ return
+ streams = probe_streams(paths["mp4"])
+ if streams:
+ logger.info(f"container streams: { {k: v.get('duration') for k, v in streams.items()} }")
+ assert "video" in streams and "audio" in streams, f"muxed file is missing a stream: {list(streams)}"
+ durations = {k: float(v["duration"]) for k, v in streams.items() if v.get("duration")}
+ if {"video", "audio"} <= set(durations):
+ skew = durations["audio"] - durations["video"]
+ # AAC pads to a frame boundary, so allow a little more than the tensor-level check.
+ assert abs(skew) < 0.15, f"muxed A/V skew {skew:+.3f} s"
+ logger.info(f"muxed A/V skew: {skew:+.4f} s")
+
+ decoded = decoded_frames(paths["mp4"], count=1, height=height, width=width)
+ if decoded.size:
+ assert (
+ decoded.shape[0] >= expected_frames - 1
+ ), f"the written mp4 decodes to {decoded.shape[0]} frames, expected ~{expected_frames}"
+ # The VAE's temporal chunk covers clip_length (17) pixel frames.
+ seam = temporal_seam_score(decoded, period=seam_period)
+ logger.info(f"temporal seam score at the {seam_period}-frame chunk period: {seam:.3f} (1.0 = no seam)")
+ if np.isfinite(seam):
+ assert seam < 3.0, (
+ f"inter-frame delta at chunk boundaries is {seam:.2f}x the delta elsewhere; "
+ "suspect temporal stitching (see the artifact rubric)"
+ )
+
+
+def gate_clip(frames: np.ndarray, prompt: str, threshold: float, label: str):
+ """CLIP prompt-alignment gate; skips only if `open_clip` is missing."""
+ pytest.importorskip("open_clip", reason="the CLIP gate needs open_clip, which is not installed")
+ alignment = clip_prompt_alignment(frames, prompt)
+ logger.info(
+ f"{label} CLIP prompt alignment: mean={alignment['mean']:.2f} "
+ f"min={alignment['min']:.2f} max={alignment['max']:.2f} (bar {threshold})"
+ )
+ assert alignment["mean"] >= threshold, (
+ f"{label} CLIP prompt alignment {alignment['mean']:.2f} is below the {threshold} bar; "
+ "the video does not match the prompt"
+ )
+ return alignment
+
+
+def gate_vbench(paths: dict, prompt: str, thresholds: dict, label: str, skip_without_mp4=False):
+ """VBench gate (skips only when its interpreter is missing); a requested dimension with no returned score FAILS, never silently passes."""
+ if "mp4" not in paths:
+ message = "the VBench gate needs the muxed mp4, which ffmpeg did not produce"
+ if skip_without_mp4:
+ pytest.skip(message)
+ raise AssertionError(message)
+ scores = run_vbench(paths["mp4"], prompt, tuple(thresholds))
+ for dimension, bar in thresholds.items():
+ assert dimension in scores, f"VBench returned no score for {dimension}"
+ logger.info(f"{label} VBench {dimension} = {scores[dimension]:.4f} (bar {bar})")
+ failures = [f"{d} = {scores[d]:.4f} < {bar:.4f}" for d, bar in thresholds.items() if scores[d] < bar]
+ assert not failures, "VBench below threshold: " + "; ".join(failures)
+ return scores
diff --git a/code/models/tt_dit/tests/models/minimax_h3/cpu_vs_device.py b/code/models/tt_dit/tests/models/minimax_h3/cpu_vs_device.py
new file mode 100644
index 0000000000000000000000000000000000000000..9dcc0b606e4adc5288787722d3477e5300a757a3
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/cpu_vs_device.py
@@ -0,0 +1,286 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""CPU reference vs device decode, fp32, scored on real clips.
+
+Per clip: resample to 32 kHz, take one production window (207 latents = 5.17 s), encode with the
+torch/diffusers reference, decode on CPU (ground truth) and on device, score device against CPU. WAVs
+are written so the difference can be heard. Batch defaults to 2 for comparability with numbers quoted
+before `CVD_BATCH` existed; the real per-clip cost is `CVD_BATCH=1`, since latency scales ~linearly
+with batch and one request is one clip. The decoder is built once, so times are steady-state.
+
+Usage -- single device, shipping defaults:
+
+ export MINIMAX_H3_MODEL_PATH=/path/to/MiniMax-H3-diffusers
+ python models/tt_dit/tests/models/minimax_h3/cpu_vs_device.py
+
+8-way T-shard on a 4x8 mesh with trace, all-fast levers (the sub-300 ms configuration):
+
+ env CVD_MESH=4x8 CVD_T_FACTOR=8 CVD_MESH_AXIS=1 CVD_TRACED=1 \
+ CVD_SPLIT_MODE=off python .../cpu_vs_device.py
+
+Env: CVD_MESH (default 1x1), CVD_T_FACTOR, CVD_MESH_AXIS, CVD_TRACED, CVD_BATCH (default 2, 1 for the
+real per-clip cost), CVD_SPLIT_MODE (off|weight|full),
+CVD_MAX_C_IN_BLOCK, CVD_OUT_DIR, CVD_BASELINE_PSNR.
+
+Accuracy is scored against the CPU reference, i.e. an absolute number; the T-parallel test only scores
+sharded against unsharded, which is a looser bar. Unset levers keep the constructor default (accurate
+mode), so `=0` is not the same as omitting. Read each run's `levers:` line rather than the filename.
+"""
+
+import json
+import os
+import time
+
+import librosa
+import numpy as np
+import soundfile as sf
+import torch
+
+import ttnn
+
+SR = 32000
+HOP = 800
+NUM_LATENT_FRAMES = 207
+# Under `generated/` rather than any developer's filesystem, so this runs on CI and a fresh clone.
+OUT_DIR = os.environ.get("CVD_OUT_DIR") or os.path.join(
+ os.environ.get("TT_METAL_HOME", os.getcwd()), "generated", "minimax_h3_audio", "clips"
+)
+# `MINIMAX_H3_MODEL_PATH` matches the test suite; the older var is still accepted.
+WEIGHTS = os.environ.get("MINIMAX_H3_MODEL_PATH") or os.environ.get("MINIMAX_H3_DIFFUSERS_DIR", "")
+if not WEIGHTS:
+ raise SystemExit("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot")
+
+# T_FACTOR must equal the length of the axis it shards: on a 4x8 mesh only (4, axis 0) and
+# (8, axis 1) work, anything else dies in `_partition_t` on a non-tile-aligned slice.
+MESH = tuple(int(v) for v in os.environ.get("CVD_MESH", "1x1").split("x"))
+T_FACTOR = int(os.environ.get("CVD_T_FACTOR", "1"))
+MESH_AXIS = int(os.environ.get("CVD_MESH_AXIS", "1"))
+TRACED = os.environ.get("CVD_TRACED", "0") == "1"
+# Default 2 keeps existing invocations/numbers unchanged; CVD_BATCH=1 is the honest per-clip cost.
+BATCH = int(os.environ.get("CVD_BATCH", "2"))
+
+
+# The precision levers select a different operator set, not just a different speed. Quoting one
+# row's latency with another's PSNR is what produced the retired "281.6 ms at 49.45 dB" claim.
+# `None` means "leave the constructor default alone".
+def _flag(name):
+ raw = os.environ.get(name)
+ return None if raw is None else raw not in ("0", "false", "False", "")
+
+
+SPLIT_MODE = os.environ.get("CVD_SPLIT_MODE")
+# conv3d's C_in_block cap. Moves accuracy on its own -- conv_pre's error falls as the block grows
+# (2.40e-03 at 32, 1.86e-03 at 128) -- but 256 buys 0.02 dB and 512 exceeds L1.
+MAX_C_IN_BLOCK = os.environ.get("CVD_MAX_C_IN_BLOCK")
+LEVERS = {
+ k: v
+ for k, v in (
+ ("split_mode", SPLIT_MODE),
+ ("max_c_in_block", int(MAX_C_IN_BLOCK) if MAX_C_IN_BLOCK else None),
+ )
+ if v is not None
+}
+
+IS_DEFAULT = MESH == (1, 1) and T_FACTOR == 1 and not TRACED and not LEVERS
+# Non-default runs get their config in the filename. Overwriting `{label}_2_device.wav` with a
+# differently-configured decode is how the stale `*_3_device_prefix.wav` confusion started.
+_lever_tag = "".join(f"_{k}-{v}" for k, v in (("sm", SPLIT_MODE), ("cin", MAX_C_IN_BLOCK)) if v is not None)
+TAG = "" if IS_DEFAULT else f"_{MESH[0]}x{MESH[1]}_f{T_FACTOR}ax{MESH_AXIS}{'_traced' if TRACED else ''}{_lever_tag}"
+
+# (label, librosa example key, seconds to skip -- past leading silence / into a busy passage)
+CLIPS = [
+ ("voice_libri1", "libri1", 0.5),
+ ("voice_libri2", "libri2", 0.5),
+ ("music_trumpet", "trumpet", 0.5),
+ ("music_brahms", "brahms", 8.0),
+]
+
+
+def load_clip(key: str, offset: float, num_samples: int, batch: int = BATCH) -> torch.Tensor:
+ """-> (batch, 1, num_samples) at SR, peak-normalised, mono duplicated across the batch slots."""
+ y, _ = librosa.load(librosa.ex(key), sr=SR, mono=True, offset=offset, duration=num_samples / SR + 1.0)
+ if len(y) < num_samples:
+ y = np.pad(y, (0, num_samples - len(y)))
+ y = y[:num_samples]
+ peak = float(np.abs(y).max()) or 1.0
+ mono = torch.from_numpy(0.85 * y / peak).float()
+ return mono.unsqueeze(0).unsqueeze(0).repeat(batch, 1, 1)
+
+
+def write_wav(path: str, wav: torch.Tensor) -> None:
+ data = wav.detach().float()[0, 0].numpy() # both batch slots are identical; write one
+ sf.write(path, np.clip(data, -1.0, 1.0), SR, subtype="FLOAT")
+
+
+def psnr(ref: torch.Tensor, test: torch.Tensor) -> float:
+ mse = torch.mean((ref.float() - test.float()) ** 2).item()
+ if mse == 0.0:
+ return float("inf")
+ return 20.0 * np.log10(ref.abs().max().item()) - 10.0 * np.log10(mse)
+
+
+def rel_rmse(ref: torch.Tensor, test: torch.Tensor) -> float:
+ return float((test.double() - ref.double()).pow(2).mean().sqrt() / ref.double().std())
+
+
+def log_spec_distance(a: torch.Tensor, b: torch.Tensor, n_fft: int = 1024, hop: int = 256) -> float:
+ window = torch.hann_window(n_fft)
+ spec = []
+ for sig in (a, b):
+ flat = sig.reshape(-1, sig.shape[-1]).float()
+ spec.append(
+ torch.log(
+ torch.stft(flat, n_fft=n_fft, hop_length=hop, window=window, return_complex=True).abs().clamp_min(1e-5)
+ )
+ )
+ return (spec[0] - spec[1]).abs().mean().item()
+
+
+def main():
+ os.makedirs(OUT_DIR, exist_ok=True)
+ from diffusers import AutoencoderKLMiniMaxH3Audio
+ from safetensors.torch import load_file
+
+ from models.tt_dit.models.audio_vae.minimax_h3.convert_minimax_h3_audio import convert_minimax_h3_audio_state_dict
+ from models.tt_dit.models.audio_vae.minimax_h3.decoder_minimax_h3_audio import MiniMaxH3AudioDecoder
+
+ audio_dir = os.path.join(WEIGHTS, "audio_vae")
+ with open(os.path.join(audio_dir, "config.json")) as fh:
+ config = {k: v for k, v in json.load(fh).items() if not k.startswith("_")}
+ reference = AutoencoderKLMiniMaxH3Audio(**config).eval()
+ reference.load_state_dict(load_file(os.path.join(audio_dir, "diffusion_pytorch_model.safetensors")))
+
+ num_samples = NUM_LATENT_FRAMES * HOP
+ sharded = T_FACTOR > 1
+ if MESH != (1, 1):
+ # open_mesh_device takes no fabric_config; conftest sets it separately, before opening.
+ ttnn.set_fabric_config(ttnn.FabricConfig.FABRIC_1D)
+ device = ttnn.open_mesh_device(
+ ttnn.MeshShape(*MESH),
+ l1_small_size=65536,
+ **({"trace_region_size": 450_000_000} if TRACED else {}),
+ )
+ rows = []
+ try:
+ print(
+ f"config: mesh {MESH[0]}x{MESH[1]} ({device.get_num_devices()} chips), t_factor={T_FACTOR} "
+ f"axis={MESH_AXIS}, traced={TRACED}",
+ flush=True,
+ )
+ parallel_config = None
+ ccl_manager = None
+ if sharded:
+ from models.tt_dit.parallel.config import ParallelFactor
+ from models.tt_dit.parallel.manager import CCLManager
+
+ parallel_config = ParallelFactor(factor=T_FACTOR, mesh_axis=MESH_AXIS)
+ ccl_manager = CCLManager(device, num_links=1, topology=ttnn.Topology.Linear)
+
+ if LEVERS:
+ print(f"levers: {LEVERS} (unset ones keep the constructor default)", flush=True)
+ decoder = MiniMaxH3AudioDecoder(
+ latent_channels=config["latent_channels"],
+ latent_dim=config["latent_dim"],
+ decoder_dim=config["decoder_dim"],
+ decoder_rates=tuple(config["decoder_rates"]),
+ decoder_kernel_sizes=tuple(config["decoder_kernel_sizes"]),
+ resblock_kernel_sizes=tuple(config["resblock_kernel_sizes"]),
+ resblock_dilation_sizes=tuple(tuple(d) for d in config["resblock_dilation_sizes"]),
+ mesh_device=device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ **LEVERS,
+ )
+ decoder.load_torch_state_dict(convert_minimax_h3_audio_state_dict(dict(reference.state_dict())), strict=False)
+
+ for label, key, offset in CLIPS:
+ src = load_clip(key, offset, num_samples)
+ write_wav(os.path.join(OUT_DIR, f"{label}_0_source.wav"), src)
+
+ with torch.no_grad():
+ latents = reference.encode(src).latent_dist.mode()[..., :NUM_LATENT_FRAMES]
+ t0 = time.perf_counter()
+ cpu_out = reference.decode(latents).sample
+ cpu_secs = time.perf_counter() - t0
+ write_wav(os.path.join(OUT_DIR, f"{label}_1_cpu.wav"), cpu_out)
+
+ # Warm with the same `traced` flag the timed runs use, so the trace is captured outside the
+ # measurement rather than inside the first one.
+ decoder(latents, traced=TRACED)
+ runs = []
+ for _ in range(3):
+ t0 = time.perf_counter()
+ dev_out = decoder(latents, traced=TRACED)
+ runs.append(time.perf_counter() - t0)
+ dev_secs = min(runs)
+ write_wav(os.path.join(OUT_DIR, f"{label}_2_device{TAG}.wav"), dev_out)
+
+ rows.append(
+ (
+ label,
+ cpu_secs,
+ dev_secs,
+ psnr(cpu_out, dev_out),
+ rel_rmse(cpu_out, dev_out),
+ log_spec_distance(cpu_out, dev_out),
+ )
+ )
+ print(f"done {label}: cpu {cpu_secs:.3f}s device {dev_secs:.3f}s psnr {rows[-1][3]:.2f} dB", flush=True)
+ finally:
+ if TRACED:
+ try:
+ decoder.release_trace()
+ except Exception: # nothing to release if the run died before capture
+ pass
+ ttnn.close_mesh_device(device)
+ if MESH != (1, 1):
+ # Leaving fabric enabled after close can wedge the next job's init.
+ ttnn.set_fabric_config(ttnn.FabricConfig.DISABLED)
+
+ print(f"\n=== fp32 decode, {num_samples / SR:.2f} s of 32 kHz audio per clip, batch {BATCH} ===")
+ print(f"config: mesh {MESH[0]}x{MESH[1]}, t_factor={T_FACTOR} axis={MESH_AXIS}, traced={TRACED}")
+ print(f"{'clip':<15} {'CPU s':>7} {'device s':>9} {'speedup':>8} {'PSNR dB':>9} {'rel_rmse':>11} {'log-spec':>9}")
+ print("-" * 74)
+ for label, c, d, p, r, ls in rows:
+ print(f"{label:<15} {c:>7.3f} {d:>9.3f} {c / d:>7.2f}x {p:>9.2f} {r:>11.3e} {ls:>9.4f}")
+ if rows:
+ print("-" * 74)
+ n = len(rows)
+ print(
+ f"{'mean':<15} {sum(r[1] for r in rows) / n:>7.3f} {sum(r[2] for r in rows) / n:>9.3f} "
+ f"{sum(r[1] for r in rows) / sum(r[2] for r in rows):>7.2f}x {sum(r[3] for r in rows) / n:>9.2f}"
+ )
+ mean_psnr = sum(r[3] for r in rows) / n
+ mean_secs = sum(r[2] for r in rows) / n
+ # The accuracy criterion is "no worse than the single-device path at the same levers", not a
+ # fixed number: the decoder's constructed defaults are accurate mode (split_mode='full')
+ # and score ~67 dB, where all-fast scores far lower. A bar that does not name
+ # its lever set is meaningless across that spread.
+ # So the baseline has to be *measured* on this branch -- run once without the CVD_* variables and
+ # pass the mean back in via CVD_BASELINE_PSNR. Unset, this reports and asserts nothing on
+ # accuracy, because a hard-coded bar from another configuration is worse than no bar.
+ baseline_env = os.environ.get("CVD_BASELINE_PSNR")
+ if baseline_env:
+ baseline = float(baseline_env)
+ drop = baseline - mean_psnr
+ accuracy = (
+ f"psnr {mean_psnr:.2f} dB vs {baseline:.2f} baseline "
+ f"({'no degradation' if drop <= 0.05 else f'DOWN {drop:.2f} dB'}): "
+ f"{'PASS' if drop <= 0.05 else 'FAIL'}"
+ )
+ else:
+ accuracy = (
+ f"psnr {mean_psnr:.2f} dB (no baseline set -- run single-device first, then pass CVD_BASELINE_PSNR)"
+ )
+ print(
+ f"\nACCEPTANCE {accuracy}"
+ f" | latency {mean_secs * 1e3:.1f} ms vs 300 ms: {'PASS' if mean_secs <= 0.300 else 'FAIL'}"
+ f" (222 ms stretch: {'PASS' if mean_secs <= 0.222 else 'FAIL'})"
+ )
+ print(f"\nWAVs in {OUT_DIR} (per clip: _0_source, _1_cpu, _2_device{TAG})")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/code/models/tt_dit/tests/models/minimax_h3/decode_bench.py b/code/models/tt_dit/tests/models/minimax_h3/decode_bench.py
new file mode 100644
index 0000000000000000000000000000000000000000..3775e9aa1f1da6bd01fe4fcf88db74091d4473dd
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/decode_bench.py
@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Median decode time over N runs, building the decoder once.
+
+Build once, warm once, then time N decodes and report median plus spread. Timing a single call per
+process is where an earlier script's run-to-run noise came from.
+
+Usage -- the label names the saved waveform:
+
+ export MINIMAX_H3_MODEL_PATH=/path/to/MiniMax-H3-diffusers
+ python models/tt_dit/tests/models/minimax_h3/decode_bench.py baseline
+
+Env: BENCH_N (default 5), BENCH_OUT (default `generated/minimax_h3_audio/bench`).
+
+For latency *with* PSNR against a CPU reference, and for mesh/sharding/trace configurations, use
+`cpu_vs_device.py` instead; this script is single-device and reports timing only.
+"""
+import json
+import os
+import re
+import statistics
+import sys
+import time
+
+import torch
+from safetensors.torch import load_file
+
+import ttnn
+from models.tt_dit.models.audio_vae.minimax_h3.convert_minimax_h3_audio import convert_minimax_h3_audio_state_dict
+from models.tt_dit.models.audio_vae.minimax_h3.decoder_minimax_h3_audio import MiniMaxH3AudioDecoder
+
+# Reaches a filename below, so keep it to one path-safe component.
+LABEL = re.sub(r"[^A-Za-z0-9._-]", "_", os.path.basename(sys.argv[1])) or "bench"
+N = int(os.environ.get("BENCH_N", "5"))
+OUT = os.environ.get("BENCH_OUT") or os.path.join(
+ os.environ.get("TT_METAL_HOME", os.getcwd()), "generated", "minimax_h3_audio", "bench"
+)
+# `MINIMAX_H3_MODEL_PATH` matches the test suite; the older var is still accepted.
+_root = os.environ.get("MINIMAX_H3_MODEL_PATH") or os.environ.get("MINIMAX_H3_DIFFUSERS_DIR", "")
+if not _root:
+ raise SystemExit("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot")
+wd = os.path.join(_root, "audio_vae")
+cfg = {k: v for k, v in json.load(open(os.path.join(wd, "config.json"))).items() if not k.startswith("_")}
+
+d = ttnn.open_mesh_device(ttnn.MeshShape(1, 1), l1_small_size=65536)
+try:
+ dec = MiniMaxH3AudioDecoder(
+ latent_channels=cfg["latent_channels"],
+ latent_dim=cfg["latent_dim"],
+ decoder_dim=cfg["decoder_dim"],
+ decoder_rates=tuple(cfg["decoder_rates"]),
+ decoder_kernel_sizes=tuple(cfg["decoder_kernel_sizes"]),
+ resblock_kernel_sizes=tuple(cfg["resblock_kernel_sizes"]),
+ resblock_dilation_sizes=tuple(tuple(x) for x in cfg["resblock_dilation_sizes"]),
+ mesh_device=d,
+ )
+ dec.load_torch_state_dict(
+ convert_minimax_h3_audio_state_dict(load_file(os.path.join(wd, "diffusion_pytorch_model.safetensors"))),
+ strict=False,
+ )
+ torch.manual_seed(2)
+ lat = torch.randn(2, cfg["latent_channels"], 207) * 0.1
+ out = dec(lat)
+ ttnn.synchronize_device(d)
+ s = []
+ for _ in range(N):
+ t0 = time.perf_counter()
+ out = dec(lat)
+ ttnn.synchronize_device(d)
+ s.append(time.perf_counter() - t0)
+ os.makedirs(OUT, exist_ok=True)
+ # Both halves come from outside the process (BENCH_OUT, argv[1]), so contain the result here
+ # rather than relying on the sanitising above.
+ _out_dir = os.path.realpath(OUT)
+ _dest = os.path.realpath(os.path.join(_out_dir, LABEL + ".pt"))
+ if os.path.commonpath([_out_dir, _dest]) != _out_dir:
+ raise SystemExit(f"refusing to write outside {_out_dir}: {_dest}")
+ torch.save({"wav": torch.as_tensor(out).float().cpu(), "samples": s}, _dest)
+ print(
+ f"RESULT {LABEL}: median {statistics.median(s):.4f}s min {min(s):.4f} max {max(s):.4f} n={N} "
+ f"[{', '.join(f'{v:.3f}' for v in s)}]",
+ flush=True,
+ )
+finally:
+ ttnn.close_mesh_device(d)
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_adaln_precompute_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_adaln_precompute_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb6b18ec6007c609e1a77c81f9b04df9090299c7
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_adaln_precompute_minimax_h3.py
@@ -0,0 +1,206 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Host-only gate for the MiniMax-H3 AdaLN precompute, on a small synthetic checkpoint.
+The design rests on ``adaln_proj`` being batch-stable while ``time_embedder`` is not. The full
+26 GB build was verified once against the real checkpoint (git history: tools/build_adaln_table.py)."""
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from ....pipelines.minimax_h3 import adaln_precompute as ap
+from ....pipelines.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG
+from ....pipelines.minimax_h3.scheduler import MiniMaxH3Scheduler
+
+HIDDEN = 64
+TIME_EMBED_DIM = 32
+FREQ_DIM = 16
+NUM_LAYERS = 3
+STEPS = 6
+
+
+def _schedules(steps=STEPS):
+ video = MiniMaxH3Scheduler(12.0)
+ audio = MiniMaxH3Scheduler(3.0)
+ video.set_timesteps(steps)
+ audio.set_timesteps(steps)
+ return video, audio
+
+
+def _step_timesteps(steps=STEPS):
+ video, audio = _schedules(steps)
+ return ap.request_step_timesteps(video.sigmas, audio.sigmas, MINIMAX_H3_KEYFRAME_NOISE_AUG)
+
+
+@pytest.fixture
+def synthetic_checkpoint(tmp_path):
+ torch.manual_seed(0)
+ tensors = {
+ "time_embedder.proj_in.weight": torch.randn(HIDDEN, FREQ_DIM),
+ "time_embedder.proj_in.bias": torch.randn(HIDDEN),
+ "time_embedder.proj_out.weight": torch.randn(TIME_EMBED_DIM, HIDDEN),
+ "time_embedder.proj_out.bias": torch.randn(TIME_EMBED_DIM),
+ "final_layer.adaln_proj.linear.weight": torch.randn(2 * HIDDEN, TIME_EMBED_DIM).bfloat16(),
+ "final_layer.adaln_proj.linear.bias": torch.randn(2 * HIDDEN).bfloat16(),
+ }
+ for layer in range(NUM_LAYERS):
+ out = ap.MINIMAX_H3_ADALN_PARAMS * HIDDEN * ap.MINIMAX_H3_MODALITY_NUM
+ tensors[f"blocks.{layer}.adaln_proj.linear.weight"] = torch.randn(out, TIME_EMBED_DIM).bfloat16()
+ tensors[f"blocks.{layer}.adaln_proj.linear.bias"] = torch.randn(out).bfloat16()
+ # Two shards, so the multi-shard key index is exercised.
+ keys = sorted(tensors)
+ save_file({k: tensors[k] for k in keys[: len(keys) // 2]}, tmp_path / "model-00001-of-00002.safetensors")
+ save_file({k: tensors[k] for k in keys[len(keys) // 2 :]}, tmp_path / "model-00002-of-00002.safetensors")
+ return tmp_path, tensors
+
+
+def _table(checkpoint, steps=STEPS):
+ return ap.precompute_adaln_table(
+ checkpoint,
+ _step_timesteps(steps),
+ num_layers=NUM_LAYERS,
+ hidden_size=HIDDEN,
+ freq_dim=FREQ_DIM,
+ )
+
+
+def test_frequency_embedding_is_cosine_then_sine():
+ diffusers_embeddings = pytest.importorskip("diffusers.models.embeddings")
+ timesteps = torch.tensor([0.0, 0.02, 0.5, 0.999, 1.0], dtype=torch.float32)
+ ours = ap.timestep_frequency_embedding(timesteps, 256)
+ theirs = diffusers_embeddings.get_timestep_embedding(timesteps, 256, flip_sin_to_cos=True, downscale_freq_shift=0)
+ assert torch.equal(ours, theirs)
+
+
+def test_silu_must_run_before_the_bfloat16_cast():
+ """bf16 activation shifts modulation ~7.8e-3, coherently across every block and step."""
+ torch.manual_seed(1)
+ weight = torch.randn(6 * HIDDEN * 3, TIME_EMBED_DIM).bfloat16()
+ bias = torch.randn(6 * HIDDEN * 3).bfloat16()
+ temb = torch.randn(4, TIME_EMBED_DIM, dtype=torch.float32)
+
+ correct = ap.project_block_adaln(temb, weight, bias, HIDDEN)
+ hoisted = torch.nn.functional.linear(torch.nn.functional.silu(temb.bfloat16()), weight, bias)
+ hoisted = torch.stack(hoisted.view(-1, 6 * HIDDEN).chunk(6, dim=-1), dim=1)
+
+ assert not torch.equal(correct, hoisted)
+ assert (correct.float() - hoisted.float()).abs().max() > 1e-3
+
+
+def test_adaln_projection_is_batch_stable():
+ torch.manual_seed(2)
+ weight = torch.randn(6 * HIDDEN * 3, TIME_EMBED_DIM).bfloat16()
+ bias = torch.randn(6 * HIDDEN * 3).bfloat16()
+ temb = torch.randn(16, TIME_EMBED_DIM, dtype=torch.float32)
+
+ batched = ap.project_block_adaln(temb, weight, bias, HIDDEN)
+ for row in (0, 5, 15):
+ alone = ap.project_block_adaln(temb[row : row + 1], weight, bias, HIDDEN)
+ assert torch.equal(batched[row * 3 : (row + 1) * 3], alone)
+
+
+def test_time_embedding_is_not_batch_stable(synthetic_checkpoint):
+ """Why ``temb`` is computed per step, asserted rather than silently depended on."""
+ _, tensors = synthetic_checkpoint
+ args = (
+ tensors["time_embedder.proj_in.weight"],
+ tensors["time_embedder.proj_in.bias"],
+ tensors["time_embedder.proj_out.weight"],
+ tensors["time_embedder.proj_out.bias"],
+ )
+ steps = _step_timesteps()
+ everything = torch.unique(torch.cat(steps), sorted=True)
+
+ batched = ap.time_embedding(everything, *args, freq_dim=FREQ_DIM)
+ per_step = ap.time_embedding(steps[0], *args, freq_dim=FREQ_DIM)
+ rows = [int((everything == value).nonzero()[0, 0]) for value in steps[0]]
+
+ assert torch.allclose(batched[rows], per_step, atol=1e-4)
+
+
+def test_step_timesteps_pin_conditioning_at_the_noise_aug_floor():
+ video, audio = _schedules()
+ steps = ap.request_step_timesteps(video.sigmas, audio.sigmas, MINIMAX_H3_KEYFRAME_NOISE_AUG)
+ for index, levels in enumerate(steps):
+ expected = max(float(video.timesteps[index]), MINIMAX_H3_KEYFRAME_NOISE_AUG)
+ assert pytest.approx(float(levels.max())) == expected
+
+
+def test_table_matches_per_step_recompute(synthetic_checkpoint):
+ checkpoint, tensors = synthetic_checkpoint
+ steps = _step_timesteps()
+ table = _table(checkpoint)
+
+ total_rows = sum(int(levels.numel()) for levels in steps)
+ assert len(steps) == _schedules()[0].num_inference_steps
+ assert table.num_layers == NUM_LAYERS
+ assert table.hidden_size == HIDDEN
+ assert table.num_steps == len(steps)
+ assert tuple(table.block_params.shape) == (
+ NUM_LAYERS,
+ total_rows * ap.MINIMAX_H3_MODALITY_NUM,
+ ap.MINIMAX_H3_ADALN_PARAMS,
+ HIDDEN,
+ )
+ assert table.block_params.dtype == torch.bfloat16
+ assert tuple(table.final_shift.shape) == (total_rows, HIDDEN)
+ assert int(table.step_offsets[0]) == 0
+ assert int(table.step_offsets[-1]) == total_rows
+
+ embed_args = (
+ tensors["time_embedder.proj_in.weight"],
+ tensors["time_embedder.proj_in.bias"],
+ tensors["time_embedder.proj_out.weight"],
+ tensors["time_embedder.proj_out.bias"],
+ )
+
+ for step, levels in enumerate(steps):
+ assert levels.numel() in (2, 3)
+ assert bool((levels[1:] > levels[:-1]).all())
+ assert torch.equal(table.step_timesteps(step), levels)
+
+ temb = ap.time_embedding(levels, *embed_args, freq_dim=FREQ_DIM)
+ rows = table.step_rows(step, torch.arange(levels.numel()))
+ low = int(rows[0]) * ap.MINIMAX_H3_MODALITY_NUM
+ high = (int(rows[-1]) + 1) * ap.MINIMAX_H3_MODALITY_NUM
+
+ for layer in range(NUM_LAYERS):
+ expected = ap.project_block_adaln(
+ temb,
+ tensors[f"blocks.{layer}.adaln_proj.linear.weight"],
+ tensors[f"blocks.{layer}.adaln_proj.linear.bias"],
+ HIDDEN,
+ )
+ assert torch.equal(table.block_params[layer, low:high], expected)
+
+ shift, scale = ap.project_final_adaln(
+ temb,
+ tensors["final_layer.adaln_proj.linear.weight"],
+ tensors["final_layer.adaln_proj.linear.bias"],
+ )
+ assert torch.equal(table.final_shift[rows[0] : rows[-1] + 1], shift)
+ assert torch.equal(table.final_scale[rows[0] : rows[-1] + 1], scale)
+
+
+def test_adaln_indices_address_timestep_and_modality(synthetic_checkpoint):
+ checkpoint, _ = synthetic_checkpoint
+ table = _table(checkpoint)
+ step = 3
+ offset = int(table.step_offsets[step])
+
+ timestep_indices = torch.tensor([0, 0, 1, 1])
+ token_tags = torch.tensor([0, 2, 1, -1])
+ indices = table.adaln_indices(step, timestep_indices, token_tags)
+
+ expected = torch.tensor(
+ [
+ (offset + 0) * 3 + 0,
+ (offset + 0) * 3 + 2,
+ (offset + 1) * 3 + 1,
+ # tag -1 padding rows clamp to the video slot; their output is discarded
+ (offset + 1) * 3 + 0,
+ ]
+ )
+ assert torch.equal(indices, expected)
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_audio_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_audio_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..57587da80848f42e3716e557c3958cecb0a47f03
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_audio_minimax_h3.py
@@ -0,0 +1,698 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Gate M8d: the MiniMax-H3 audio VAE (BigVGAN) at production shapes; stereo rides as batch 2.
+
+T-parallel decode works: ``test_audio_decode_t_parallel`` at the bottom of this file is resurrected from
+git history, as the previous version of this docstring asked for. The 8-way shard layout it called
+known-broken was `Vocoder.conv_pre` returning uninitialized memory when sharded (plus a halo bug in the
+since-retired tap-matmul path); both were fixed, and
+the test's PSNR assert is enforced for every factor rather than excused. ``test_neighbor_pad_t_minimax_h3.py``
+is still in git history if the halo itself ever needs its own gate again -- it was verified correct here
+(exact at pad 1/3/25) while hunting the conv_pre bug."""
+
+from __future__ import annotations
+
+import copy
+import os
+import time
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+
+from ....layers.audio_ops import Conv1dViaConv3d, depthwise_tap_filter
+from ....models.audio_vae.minimax_h3.blockings_minimax_h3_audio import register_h3_audio_blockings
+from ....models.audio_vae.minimax_h3.convert_minimax_h3_audio import (
+ assert_weight_norm_axes_consistent,
+ convert_minimax_h3_audio_state_dict,
+ fuse_attention_biases,
+ fuse_weight_norm,
+ remap_amp_activations,
+)
+from ....models.audio_vae.minimax_h3.decoder_minimax_h3_audio import MiniMaxH3AudioDecoder
+from ....models.audio_vae.vocoder_ltx import TILE_HEIGHT
+from ....parallel.config import ParallelFactor
+from ....parallel.manager import CCLManager
+from ....pipelines.minimax_h3.pipeline_minimax_h3 import resolve_mesh_preset
+from ....utils.check import assert_quality
+from .common import build_audio_decoder, load_config, psnr, weights_subdir
+
+# The vocoder needs extra L1 scratch, as the LTX audio tests do.
+SINGLE_DEVICE = [pytest.param((1, 1), {"l1_small_size": 65536}, id="single_device")]
+
+# Encode floor, calibrated before accurate mode became the constructed default (measured ~0.105 with
+# fp32 conv multipliers over ~130 convs); the accurate defaults only improve on it.
+AUDIO_RELATIVE_RMSE = 0.12
+
+CONV_PRE_SHAPE = (2048, 1024, 7) # conv_pre: the decoder's widest reduction, Cin 2048 x k 7
+CONV_PRE_LATENT_FRAMES = 207
+
+LATENT_CHANNELS = 32
+HOP_LENGTH = 800
+SAMPLING_RATE = 32000
+
+PRODUCTION_LATENT_FRAMES = [pytest.param(207, id="5s")] # 40 Hz latent; 10 s exercises no new path
+
+
+def _build_reference(load_weights: bool = True):
+ """The reference audio VAE, with weight norm still attached (the converter removes it)."""
+ weights_dir = weights_subdir("audio_vae")
+ if weights_dir is None:
+ pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH")
+ pytest.importorskip("diffusers", reason="pinned diffusers reference not installed")
+ from diffusers import AutoencoderKLMiniMaxH3Audio
+
+ config = load_config(weights_dir)
+ reference = AutoencoderKLMiniMaxH3Audio(**config).eval()
+ if load_weights:
+ from safetensors.torch import load_file
+
+ reference.load_state_dict(load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors")))
+ return reference, config
+
+
+def _golden_latents(reference, num_latent_frames: int):
+ """``(waveform, latents, expected_decode)``; latents come from the reference encoder, not ``randn``."""
+ waveform = torch.randn(2, 1, num_latent_frames * HOP_LENGTH) * 0.1
+ with torch.no_grad():
+ posterior = reference.encode(waveform).latent_dist
+ latents = posterior.mode()[..., :num_latent_frames]
+ expected = reference.decode(latents).sample
+ return waveform, latents, expected
+
+
+def _log_mel_distance(a: torch.Tensor, b: torch.Tensor, *, n_fft: int = 1024, hop: int = 256) -> float:
+ """Mean absolute log-magnitude-spectrogram difference; catches spectral drift an aggregate PCC hides."""
+ window = torch.hann_window(n_fft)
+ spectra = []
+ for signal in (a, b):
+ flat = signal.reshape(-1, signal.shape[-1]).float()
+ stft = torch.stft(flat, n_fft=n_fft, hop_length=hop, window=window, return_complex=True)
+ spectra.append(torch.log(stft.abs().clamp_min(1e-5)))
+ return (spectra[0] - spectra[1]).abs().mean().item()
+
+
+def _tt_decoder(config: dict, mesh_device) -> MiniMaxH3AudioDecoder:
+ return build_audio_decoder(config, mesh_device)
+
+
+@pytest.mark.parametrize("num_latent_frames", PRODUCTION_LATENT_FRAMES)
+@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"])
+def test_decode(mesh_device, num_latent_frames):
+ """The whole decode path against the reference, at a production duration, stereo.
+
+ Constructor defaults are accurate mode (split_mode='full'), so the bars are the
+ accurate-mode ones: measured 0.0045 rel RMSE / 99.9990% PCC / 67.5 dB PSNR.
+ """
+ reference, config = _build_reference()
+ torch.manual_seed(1)
+
+ _, latents, expected = _golden_latents(reference, num_latent_frames)
+ with torch.no_grad():
+ expected_proj = reference.dec_in_proj(latents)
+
+ tt_decoder = _tt_decoder(config, mesh_device)
+ # The precision levers are the constructed defaults; assert they landed where they matter.
+ assert tt_decoder.dec_in_proj.split_mode == "full", "split_mode='full' did not land on dec_in_proj"
+ assert tt_decoder.decoder.conv_post.split_mode == "full", "split_mode='full' did not land on conv_post"
+
+ tt_decoder.load_torch_state_dict(convert_minimax_h3_audio_state_dict(dict(reference.state_dict())), strict=False)
+
+ projected = tt_decoder._project_latents_device(latents)
+ assert projected.shape == expected_proj.shape, f"proj {tuple(projected.shape)} != {tuple(expected_proj.shape)}"
+ assert_quality(expected_proj, projected, pcc=0.999)
+
+ actual = tt_decoder(latents)
+
+ assert actual.shape == expected.shape, f"shape {tuple(actual.shape)} != reference {tuple(expected.shape)}"
+ assert_quality(expected, actual, pcc=0.9999, relative_rmse=0.006)
+
+ psnr_db = psnr(expected, actual)
+ mel_distance = _log_mel_distance(expected, actual)
+ assert psnr_db >= 60.0, f"decode PSNR {psnr_db:.2f} dB < 60 dB"
+ assert mel_distance <= 5.0, f"log-spectrogram distance {mel_distance:.3f} > 5.0"
+
+ left, right = actual[0, 0], actual[1, 0]
+ assert not torch.allclose(left, right, atol=1e-4), "stereo channels are identical -- a broadcast bug"
+
+
+@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"])
+def test_depthwise_chunked_conv1d_matches_torch(mesh_device):
+ """`depthwise_tap_filter`'s C-chunked `ttnn.conv1d` recovery, against a torch depthwise reference.
+
+ Covers the slicing, the per-chunk weight, and the concat in isolation, at the shape the decode
+ really chunks (T_pad=166, C=512, K=7, stride=1) -- the activation block is C*K wide there and
+ the slicer runs out of L1 at full C.
+ """
+ torch.manual_seed(0)
+ B, T_pad, C, K, stride = 1, 166, 512, 7, 1
+
+ x = torch.randn(B, T_pad, C, dtype=torch.float32)
+ taps = torch.randn(K, dtype=torch.float32)
+ x_dev = ttnn.from_torch(x, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=mesh_device)
+
+ cache: dict = {}
+ out = depthwise_tap_filter(
+ x_dev,
+ [float(t) for t in taps],
+ stride,
+ mesh_device=mesh_device,
+ dtype=ttnn.float32,
+ cache=cache,
+ )
+ actual = ttnn.to_torch(ttnn.get_device_tensors(out)[0]).float()
+
+ # Valid depthwise conv over T: one tap vector shared by every channel, groups=C.
+ weight = taps.view(1, 1, K).expand(C, 1, K).contiguous()
+ expected = torch.nn.functional.conv1d(x.transpose(1, 2), weight, stride=stride, groups=C).transpose(1, 2)
+
+ assert actual.shape == expected.shape, f"{actual.shape} != {expected.shape}"
+ psnr_db = psnr(expected, actual)
+ logger.info(f"chunked depthwise conv1d vs torch: {psnr_db:.2f} dB")
+ assert psnr_db > 60.0, f"chunked depthwise conv1d diverges from torch: {psnr_db:.2f} dB"
+
+ # Prove the chunked path ran: it keys its prepared weight on the chunk width where the unchunked
+ # path keys on C. Without this the test would pass on plain conv1d or a silent MAC fallback, i.e.
+ # assert torch equivalence while covering none of the code it exists for.
+ chunk_widths = [k[1] for k in cache if isinstance(k, tuple) and k and k[0] == "w" and k[1] < C]
+ assert chunk_widths, (
+ f"expected the C-chunked conv1d recovery to run at C={C}, K={K}, T_pad={T_pad}, but the weight "
+ f"cache has no sub-C chunk key: {sorted(k for k in cache if isinstance(k, tuple))}. Either "
+ "conv1d now fits at full C (then this shape no longer covers the chunked path and the test needs "
+ "a new one), or it fell back to MAC (then the recovery is broken)."
+ )
+ logger.info(f"chunked into widths {sorted(set(chunk_widths))} from C={C}")
+
+
+@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"])
+def test_conv_operand_split_improves_precision(mesh_device):
+ """Operand splitting really does beat the fp32 conv floor, at ``conv_pre``'s production shape."""
+ # Raw Conv1dViaConv3d does not register the H3 conv blockings (the H3 audio module constructors
+ # do); without this, every conv silently falls back to C_in_block=32 and the "off" bar fails.
+ register_h3_audio_blockings()
+ in_channels, out_channels, kernel = CONV_PRE_SHAPE
+ torch.manual_seed(0)
+ reference = torch.nn.Conv1d(in_channels, out_channels, kernel, padding=kernel // 2).float().eval()
+ x = torch.randn(2, in_channels, CONV_PRE_LATENT_FRAMES) * 0.1
+ with torch.no_grad():
+ golden = copy.deepcopy(reference).double()(x.double()).float()
+
+ x_device = None
+ errors = {}
+ for mode in ("off", "weight", "full"):
+ layer = Conv1dViaConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=kernel,
+ padding_mode="zeros",
+ bias=True,
+ mesh_device=mesh_device,
+ dtype=ttnn.float32,
+ split_mode=mode,
+ )
+ assert layer.split_mode == mode, f"layer resolved split_mode {layer.split_mode!r}, expected {mode!r}"
+ assert (layer.weight_lo is None) == (mode == "off"), f"weight residual allocation wrong for {mode!r}"
+ layer.load_torch_state_dict(
+ {
+ "weight": reference.weight.detach().contiguous(),
+ "bias": reference.bias.detach().contiguous(),
+ },
+ strict=False,
+ )
+ if x_device is None:
+ x_device = ttnn.from_torch(
+ x.transpose(1, 2).contiguous(), dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=mesh_device
+ )
+ actual = ttnn.to_torch(layer(x_device)).float().transpose(1, 2)
+ errors[mode] = float((actual.double() - golden.double()).pow(2).mean().sqrt() / golden.double().std())
+
+ logger.info("conv_pre rel_rmse by split mode: " + ", ".join(f"{k}={v:.3e}" for k, v in errors.items()))
+
+ # 2.1e-03 sits between production's 1.86e-03 and the unregistered C_in_block=32 fallback's 2.40e-03.
+ assert errors["off"] <= 2.1e-3, (
+ f"baseline conv_pre error {errors['off']:.3e} is above the production floor -- "
+ "C_in_block has probably fallen back to 32, meaning register_h3_audio_blockings() did not run"
+ )
+ # Measured 1.86e-03 / 1.25e-03 / 1.00e-03; ratios carry ~1.3x margin.
+ assert errors["weight"] <= 0.85 * errors["off"], f"weight-split did not improve: {errors}"
+ assert errors["full"] <= 0.70 * errors["off"], f"full split did not improve enough: {errors}"
+ assert errors["full"] < errors["weight"], f"full split should beat weight-only: {errors}"
+
+
+def _tt_encoder(config: dict, mesh_device):
+ from ....models.audio_vae.minimax_h3.encoder_minimax_h3_audio import MiniMaxH3AudioEncoder
+
+ return MiniMaxH3AudioEncoder(
+ encoder_dim=config["encoder_dim"],
+ encoder_rates=tuple(config["encoder_rates"]),
+ latent_dim=config["latent_dim"],
+ latent_channels=config["latent_channels"],
+ num_attention_heads=config["num_attention_heads"],
+ mesh_device=mesh_device,
+ )
+
+
+@pytest.mark.parametrize("num_latent_frames", PRODUCTION_LATENT_FRAMES)
+@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"])
+def test_encode(mesh_device, num_latent_frames):
+ """The whole encode path -- DAC trunk, ``pre_block``, posterior heads -- vs the reference."""
+ reference, config = _build_reference()
+ torch.manual_seed(2)
+
+ waveform = torch.randn(2, 1, num_latent_frames * HOP_LENGTH) * 0.1
+ with torch.no_grad():
+ posterior = reference.encode(waveform).latent_dist
+ expected_mean, expected_logs = posterior.mean, posterior.logs
+
+ tt_encoder = _tt_encoder(config, mesh_device)
+ tt_encoder.load_torch_state_dict(convert_minimax_h3_audio_state_dict(dict(reference.state_dict())), strict=False)
+ mean, logs = tt_encoder(waveform)
+
+ assert mean.shape == expected_mean.shape, f"mean shape {tuple(mean.shape)} != {tuple(expected_mean.shape)}"
+ assert mean.shape[2] == num_latent_frames, f"expected {num_latent_frames} latent frames, got {mean.shape[2]}"
+ assert_quality(expected_mean, mean, pcc=0.99, relative_rmse=AUDIO_RELATIVE_RMSE)
+ assert_quality(expected_logs, logs, pcc=0.98)
+
+
+@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"])
+def test_roundtrip(mesh_device):
+ """End to end: encode then decode on device, against the reference's own round trip."""
+ reference, config = _build_reference()
+ torch.manual_seed(3)
+ num_latent_frames = 207
+
+ waveform, _, expected = _golden_latents(reference, num_latent_frames)
+
+ tt_encoder = _tt_encoder(config, mesh_device)
+ tt_decoder = _tt_decoder(config, mesh_device)
+ converted = convert_minimax_h3_audio_state_dict(dict(reference.state_dict()))
+ tt_encoder.load_torch_state_dict(converted, strict=False)
+ tt_decoder.load_torch_state_dict(converted, strict=False)
+
+ mean, _ = tt_encoder(waveform)
+ actual = tt_decoder(mean)
+
+ assert actual.shape == expected.shape, f"shape {tuple(actual.shape)} != {tuple(expected.shape)}"
+ psnr_db = psnr(expected, actual)
+ assert psnr_db >= 28.0, f"round-trip PSNR {psnr_db:.2f} dB < 28 dB"
+ assert _log_mel_distance(expected, actual) <= 5.0, "round-trip spectrum drifted"
+
+
+# -------------------------------------------------------------------- weight-norm fusion and checkpoint conversion (host only)
+
+
+def test_fuse_weight_norm_matches_torch_conv_transpose1d():
+ """The load-bearing case: axis 0 of a ConvTranspose1d weight is ``in_channels``, not out."""
+ torch.manual_seed(1)
+ conv = torch.nn.ConvTranspose1d(32, 16, kernel_size=4, stride=2)
+ assert conv.weight.shape == (32, 16, 4), "ConvTranspose1d weight is (in, out, k)"
+ conv = torch.nn.utils.weight_norm(conv)
+ with torch.no_grad():
+ conv.weight_g.normal_(1.0, 0.2)
+ conv.weight_v.normal_()
+ assert conv.weight_g.shape == (32, 1, 1), "weight_g is per-in_channel, confirming dim=0"
+
+ weight_g = conv.weight_g.detach().clone()
+ weight_v = conv.weight_v.detach().clone()
+ fused = fuse_weight_norm(weight_g, weight_v)
+ torch.nn.utils.remove_weight_norm(conv)
+
+ relative = (fused - conv.weight).abs().max().item() / conv.weight.abs().max().item()
+ assert relative < 1e-6, f"ConvTranspose1d fusion differs from torch by {relative:.3e}"
+
+ wrong_norm = weight_v.transpose(0, 1).flatten(1).norm(dim=1).view(1, -1, 1)
+ wrong = weight_g * weight_v / wrong_norm
+ assert wrong.shape == fused.shape, "the wrong axis still type-checks -- hence this test"
+ assert not torch.allclose(wrong, fused, atol=1e-4), "the two axes agree, so this test proves nothing"
+
+
+def test_remap_amp_activations_interleaves_correctly():
+ """``activations.{0,2,4}`` -> ``acts1.{0,1,2}`` and ``{1,3,5}`` -> ``acts2.{0,1,2}``."""
+ state = {f"resblocks.0.activations.{i}.act.alpha": torch.tensor([float(i)]) for i in range(6)}
+ remapped = remap_amp_activations(state)
+
+ for i in range(3):
+ assert remapped[f"resblocks.0.acts1.{i}.act.alpha"].item() == 2 * i
+ assert remapped[f"resblocks.0.acts2.{i}.act.alpha"].item() == 2 * i + 1
+ assert not any("activations." in key for key in remapped), "an activations key survived"
+
+
+def test_fuse_attention_biases_rejects_a_nonzero_k_bias():
+ """A ``zero_k_bias`` that is not zero must fail loudly, not be dropped."""
+ state = {
+ "pre_block.attn.q_bias": torch.ones(8),
+ "pre_block.attn.v_bias": torch.full((8,), 2.0),
+ "pre_block.attn.zero_k_bias": torch.zeros(8),
+ }
+ fused = fuse_attention_biases(state)
+ bias = fused["pre_block.attn.qkv.bias"]
+ assert bias.shape == (24,)
+ assert torch.equal(bias[:8], torch.ones(8))
+ assert torch.equal(bias[8:16], torch.zeros(8))
+ assert torch.equal(bias[16:], torch.full((8,), 2.0))
+
+ state["pre_block.attn.zero_k_bias"] = torch.ones(8)
+ with pytest.raises(AssertionError, match="not all zero"): # allow-pytest.raises: guards a silent data-loss path
+ fuse_attention_biases(state)
+
+
+def test_real_checkpoint_fusion_matches_reference_module():
+ """Real checkpoint: axes hold, every pair fuses, fused weights equal torch's ``remove_weight_norm``."""
+ weights_dir = weights_subdir("audio_vae")
+ if weights_dir is None:
+ pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH")
+ pytest.importorskip("diffusers", reason="pinned diffusers reference not installed")
+ from diffusers import AutoencoderKLMiniMaxH3Audio
+ from safetensors.torch import load_file
+
+ config = load_config(weights_dir)
+ reference = AutoencoderKLMiniMaxH3Audio(**config)
+ state = load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors"))
+ reference.load_state_dict(state)
+
+ assert_weight_norm_axes_consistent(state)
+ num_pairs = len([k for k in state if k.endswith(".weight_g")])
+ assert num_pairs > 100, f"expected ~172 weight-normed convs, found {num_pairs}"
+
+ converted = convert_minimax_h3_audio_state_dict(state)
+ assert not [k for k in converted if k.endswith((".weight_g", ".weight_v"))], "a weight-norm pair survived"
+ assert not [k for k in converted if k.endswith(("q_bias", "v_bias", "zero_k_bias"))], "an attn bias survived"
+ assert not [k for k in converted if "activations." in k], "an activations key survived"
+ assert (
+ len(converted) == len(state) - num_pairs - 2
+ ), f"key count {len(converted)} does not match {len(state)} minus {num_pairs} g/v pairs and 2 folded biases"
+
+ for module in reference.modules():
+ if hasattr(module, "weight_g"):
+ torch.nn.utils.remove_weight_norm(module)
+ expected = dict(reference.state_dict())
+
+ checked = 0
+ worst_key, worst = None, 0.0
+ for key, value in expected.items():
+ if not key.endswith(".weight") or key not in converted:
+ continue
+ scale = max(value.abs().max().item(), 1e-12)
+ relative = (converted[key] - value).abs().max().item() / scale
+ if relative > worst:
+ worst_key, worst = key, relative
+ checked += 1
+ assert checked > 100, f"only compared {checked} fused weights"
+ assert worst < 1e-6, f"worst fused weight is {worst_key} at relative {worst:.3e}"
+
+
+# -------------------------------------------------------------------- traced audio decode
+
+# 450 MB trace region: the accurate-mode (default) decode graph needs 375463936 B.
+TRACED = [
+ pytest.param((1, 1), {"l1_small_size": 65536, "trace_region_size": 450_000_000}, id="single_device"),
+]
+NUM_LATENT_FRAMES = 207
+ITERS = 5
+
+
+def _best(fn) -> float:
+ fn()
+ best = float("inf")
+ for _ in range(ITERS):
+ t0 = time.perf_counter()
+ fn()
+ best = min(best, time.perf_counter() - t0)
+ return best
+
+
+@pytest.mark.parametrize(("mesh_device", "device_params"), TRACED, indirect=["mesh_device", "device_params"])
+def test_audio_decode_traced(mesh_device):
+ reference, config = _build_reference()
+ converted = convert_minimax_h3_audio_state_dict(dict(reference.state_dict()))
+
+ torch.manual_seed(2)
+ decoder = _tt_decoder(config, mesh_device)
+ decoder.load_torch_state_dict(converted, strict=False)
+
+ latents = torch.randn(2, config["latent_channels"], NUM_LATENT_FRAMES) * 0.1
+
+ plain = decoder(latents)
+ traced = decoder(latents, traced=True)
+ assert traced.shape == plain.shape, f"{traced.shape} != {plain.shape}"
+
+ # PSNR inf would also appear if traced=True silently fell through, so confirm a trace was captured.
+ tracers = type(decoder.decoder)._forward_device._tracers_keyed.get(decoder.decoder, {})
+ assert tracers, "traced=True captured no trace; the PSNR below would be meaningless"
+ assert plain.abs().max() > 1e-6, "decoder produced all-zero output; PSNR would be vacuous"
+
+ psnr_db = psnr(plain, traced)
+ logger.info(f"traced vs untraced PSNR: {psnr_db:.2f} dB ({len(tracers)} trace(s) captured)")
+
+ untraced_s = _best(lambda: decoder(latents))
+ traced_s = _best(lambda: decoder(latents, traced=True))
+ logger.info(
+ f"PERF audio_decode_5s untraced {untraced_s:.4f} s | traced {traced_s:.4f} s "
+ f"-> {untraced_s / traced_s:.2f}x"
+ )
+ decoder.release_trace()
+
+ assert psnr_db > 60.0, f"traced output diverges from untraced: PSNR {psnr_db:.2f} dB"
+
+
+MESH = [
+ pytest.param(
+ (4, 8),
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "require_exact_physical_num_devices": True,
+ "l1_small_size": 65536,
+ },
+ id="mesh4x8",
+ ),
+ # One Galaxy opened as a 32-wide line: the length of the quad Galaxy's inter-host axis, so the
+ # opt-in factor 32 is exercised on a single machine.
+ pytest.param(
+ (1, 32),
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D,
+ "require_exact_physical_num_devices": True,
+ "l1_small_size": 65536,
+ },
+ id="mesh1x32",
+ ),
+]
+# (t_factor, mesh_axis) per mesh shape. The factor must equal the length of the axis it shards: factor=2
+# or 4 on an 8-wide axis dies in `_partition_t` ("height begin index aligned to tiles"), because the
+# partition indexes by the device's coordinate along the axis and assumes it covers it.
+#
+# KNOWN_BROKEN must stay empty -- an entry silences the per-factor PSNR assert, the only guard against a
+# fast wrong answer. Short clips work via `Vocoder._upload_BCT`'s per-shard tile-floor.
+FACTORS_BY_MESH = {(4, 8): [(1, 1), (4, 0), (8, 1)], (1, 32): [(1, 1), (32, 1)]}
+KNOWN_BROKEN: set[tuple[int, int]] = set()
+# 207 is the production short clip. 192 (T mod 32 == 0) puts the last real row on the final row of its
+# shard with no pad buffer before the next shard: the layout where a fully-padded shard's garbage tail
+# reached real audio while the global PSNR still read 50 dB. Hence the separate tail gate below.
+T_PARALLEL_FRAMES = [207, 192]
+TAIL_ROWS = 32 # latent rows of the trailing suffix checked on its own
+
+
+def _shards_pad(num_latent_frames: int, factor: int) -> bool:
+ """True when `Vocoder._upload_BCT` pads T for this factor (per-shard rows are floored to a tile)."""
+ return max(-(-num_latent_frames // factor), TILE_HEIGHT) * factor > num_latent_frames
+
+
+def _build(mesh_device, config, converted, parallel_config, ccl_manager):
+ """The decoder at this file's shared defaults, plus a shard layout.
+
+ Goes through `build_audio_decoder` rather than constructing directly so the precision levers
+ (`split_mode`, `max_c_in_block`) stay at whatever the shipping default is -- a
+ sharded run must measure the same configuration the unsharded gates do, or a divergence could
+ be a lever difference rather than a sharding bug.
+ """
+ decoder = build_audio_decoder(
+ config,
+ mesh_device,
+ parallel_config=parallel_config,
+ ccl_manager=ccl_manager,
+ )
+ decoder.load_torch_state_dict(converted, strict=False)
+ return decoder
+
+
+def _localize_divergence(baseline, parallel, *, factor: int, logger) -> None:
+ """Say *where* a diverging shard layout diverges. "PSNR -10.2 dB" alone names no bug.
+
+ Three shapes of answer, each pointing somewhere different:
+
+ - concentrated at the internal shard boundaries -> the halo exchange is wrong, and each conv in the
+ stack needs `kernel_size - 1` samples of its neighbour that it is not getting;
+ - uniform across the whole signal -> the shard layout itself is wrong, not its edges;
+ - a prefix or suffix only -> the causal padding or the final trim is off by a shard.
+
+ Also reports the best cross-correlation lag: a diverging-but-highly-correlated output at a nonzero
+ lag is a *shift*, which is a trim bug rather than a numerics one, and PSNR cannot distinguish those.
+ """
+ import numpy as np
+
+ error = (parallel - baseline).abs()
+ total = baseline.shape[-1]
+ logger.info(
+ f" divergence localization, t_factor={factor}: overall mean {error.mean():.6f} "
+ f"max {error.max():.4f} against baseline absmax {baseline.abs().max():.4f}"
+ )
+ per_shard = []
+ for shard in range(factor):
+ lo, hi = shard * total // factor, (shard + 1) * total // factor
+ per_shard.append(float(error[..., lo:hi].mean()))
+ logger.info(f" per-shard mean error: {[f'{v:.6f}' for v in per_shard]}")
+
+ # Boundary vs interior. A halo bug puts the error in a narrow band at each internal boundary.
+ window = 128
+ boundary, interior = [], []
+ for shard in range(1, factor):
+ cut = shard * total // factor
+ boundary.append(float(error[..., max(0, cut - window) : cut + window].mean()))
+ mask = torch.ones(total, dtype=torch.bool)
+ for shard in range(1, factor):
+ cut = shard * total // factor
+ mask[max(0, cut - window) : cut + window] = False
+ interior = float(error[..., mask].mean())
+ logger.info(
+ f" boundary bands (+-{window}): {[f'{v:.6f}' for v in boundary]} interior {interior:.6f} "
+ f"ratio {max(boundary) / max(interior, 1e-12):.2f}"
+ )
+
+ a = baseline[0, 0].numpy()
+ c = parallel[0, 0].numpy()
+ best, best_lag = -2.0, None
+ for lag in range(-4096, 4097, 32):
+ x, y = (a[-lag:], c[: len(c) + lag]) if lag < 0 else (a[: len(a) - lag], c[lag:])
+ n = min(len(x), len(y))
+ if n < 2048:
+ continue
+ r = float(np.corrcoef(x[:n], y[:n])[0, 1])
+ if r > best:
+ best, best_lag = r, lag
+ logger.info(
+ f" correlation at lag 0: {float(np.corrcoef(a, c)[0, 1]):.4f}; "
+ f"best {best:.4f} at lag {best_lag} (nonzero lag => a shift, i.e. a trim bug)"
+ )
+
+
+# ~14 min: three decoder builds plus a decode per factor, against `pytest.ini`'s 300 s default.
+@pytest.mark.timeout(3600)
+@pytest.mark.parametrize("num_latent_frames", T_PARALLEL_FRAMES)
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESH, indirect=["mesh_device", "device_params"])
+def test_audio_decode_t_parallel(mesh_device, num_latent_frames):
+ weights_dir = weights_subdir("audio_vae")
+ if weights_dir is None:
+ pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH")
+ pytest.importorskip("diffusers", reason="pinned diffusers reference not installed")
+ from diffusers import AutoencoderKLMiniMaxH3Audio
+ from loguru import logger
+ from safetensors.torch import load_file
+
+ from ....models.audio_vae.minimax_h3.convert_minimax_h3_audio import convert_minimax_h3_audio_state_dict
+
+ config = load_config(weights_dir)
+ reference = AutoencoderKLMiniMaxH3Audio(**config).eval()
+ reference.load_state_dict(load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors")))
+ converted = convert_minimax_h3_audio_state_dict(dict(reference.state_dict()))
+
+ torch.manual_seed(2)
+ latents = torch.randn(2, config["latent_channels"], num_latent_frames) * 0.1
+ factors = FACTORS_BY_MESH[(mesh_device.shape[0], mesh_device.shape[1])]
+ # Only sharded factors that pad exercise the tail path; one that pads nothing (e.g. 4 at 192 latents,
+ # 48 rows/shard) decodes bit-identically to unsharded and would only add a cold compile.
+ factors = [(f, a) for f, a in factors if f == 1 or _shards_pad(num_latent_frames, f)]
+
+ baseline_out = None
+ baseline_s = None
+ results = []
+ # Run the pipeline's own CCL preset (Ring, two links on 4x8) so the gate exercises production settings;
+ # an unlisted mesh shape (1x32) has no preset and keeps the single-link line.
+ preset = resolve_mesh_preset(tuple(mesh_device.shape), required=False)
+ num_links = preset.get("num_links", 1)
+ topology = preset.get("topology", ttnn.Topology.Linear)
+ logger.info(f"CCL for the sharded decoders: num_links={num_links}, topology={topology}")
+
+ for factor, axis in factors:
+ pc = None if factor <= 1 else ParallelFactor(factor=factor, mesh_axis=axis)
+ ccl = None if pc is None else CCLManager(mesh_device, num_links=num_links, topology=topology)
+ try:
+ decoder = _build(mesh_device, config, converted, pc, ccl)
+ out = decoder(latents)
+ seconds = _best(lambda: decoder(latents))
+ except Exception as exc:
+ # Every FACTORS entry is a supported layout, so a raise is a regression, not a "result":
+ # recording it as unsupported let factor 4 break while factor 8 kept the test green, since
+ # the asserts below need only *some* factor to run. Optional layouts go in KNOWN_BROKEN.
+ logger.warning(f"t_factor={factor} axis={axis} FAILED: {str(exc)[:160]}")
+ results.append((factor, axis, None, None, None))
+ if (factor, axis) not in KNOWN_BROKEN:
+ raise
+ continue
+
+ if baseline_out is None:
+ baseline_out, baseline_s = out, seconds
+ psnr_db = tail_db = float("inf")
+ else:
+ assert out.shape == baseline_out.shape, f"factor {factor}: {out.shape} != {baseline_out.shape}"
+ # `psnr_db`, not `psnr`: binding the float to `psnr` shadows the imported helper, so the
+ # first factor's `float("inf")` makes the second factor's call a TypeError.
+ psnr_db = psnr(baseline_out, out)
+ # The global PSNR averages a corrupt tail away (50.6 dB overall hid a 33.8 dB last 4 rows).
+ n_tail = min(out.shape[-1], TAIL_ROWS * HOP_LENGTH)
+ tail_db = psnr(baseline_out[..., -n_tail:], out[..., -n_tail:])
+ results.append((factor, axis, seconds, psnr_db, tail_db))
+ logger.info(
+ f"PERF audio_decode t_factor={factor} axis={axis}: {seconds:.4f} s "
+ f"({baseline_s / seconds:.2f}x) PSNR vs 1-device {psnr_db:.1f} dB, last {TAIL_ROWS} rows {tail_db:.1f} dB"
+ )
+ if psnr_db < 40.0 and out is not baseline_out:
+ _localize_divergence(baseline_out.float(), out.float(), factor=factor, logger=logger)
+ del decoder
+
+ logger.info("=== audio decode T-parallel summary ===")
+ for factor, axis, seconds, psnr_db, tail_db in results:
+ if seconds is None:
+ logger.info(f" t_factor={factor:2d} axis={axis}: unsupported")
+ else:
+ logger.info(
+ f" t_factor={factor:2d} axis={axis}: {seconds:.4f} s {baseline_s / seconds:5.2f}x "
+ f"PSNR {psnr_db:6.1f} dB tail {tail_db:6.1f} dB"
+ )
+
+ # The baseline must have run, or there is nothing to compare against and every other factor was
+ # skipped for want of a reference. Without this the test PASSES when the whole stack is broken:
+ # observed once on a device left dirty by an unrelated crash, where all three factors raised
+ # TT_FATAL, each was swallowed as "unsupported", and the run reported green. A gate that cannot
+ # fail is worse than no gate.
+ baseline_ran = any(seconds is not None and factor == 1 for factor, _, seconds, _, _ in results)
+ assert baseline_ran, (
+ "the single-device baseline did not run, so nothing was compared. If this follows a crashed "
+ "run, reset the device (`tt-smi -glx_reset`) -- an allocator TT_FATAL here is usually a dirty "
+ "device, not a code failure"
+ )
+ # `baseline_ran` above catches "everything failed" but not the subtler case: if factor 1 raises,
+ # the loop's `baseline_out is None` branch promotes the *next* factor to baseline, and that factor
+ # then scores PSNR inf against itself and reads as correct. That happened for real -- a fusion
+ # crash killed factor 1, factor 4 became the reference, and a -10.1 dB configuration reported inf
+ # until the baseline was made to run. The first result must be the unsharded one.
+ assert results and results[0][0] == 1 and results[0][2] is not None, (
+ f"factor 1 must be the baseline, but the first result that ran was {results[0][:2] if results else None}; "
+ "a later factor has been promoted to baseline and is being compared against itself"
+ )
+ ran = [(f, a) for f, a, seconds, _, _ in results if seconds is not None and f != 1]
+ assert ran, "no parallel factor ran at all; the T-parallel path is entirely unavailable"
+
+ # Any factor that ran must agree with the single-device path; a fast wrong answer fails. The tail is
+ # gated on its own because the global number passed at 50.6 dB with the last rows at 33.8 dB.
+ for factor, axis, seconds, psnr_db, tail_db in results:
+ if seconds is None or (factor, axis) in KNOWN_BROKEN:
+ continue
+ assert psnr_db > 40.0, f"t_factor={factor} axis={axis} diverges from 1-device: PSNR {psnr_db:.1f} dB"
+ assert tail_db > 40.0, (
+ f"t_factor={factor} axis={axis}: last {TAIL_ROWS} latent rows diverge from 1-device: {tail_db:.1f} dB "
+ f"(global {psnr_db:.1f} dB) -- a shard's tail pad is being filled from the wrong row"
+ )
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_conditioning_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_conditioning_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..82efff2f880db12b41b49b63d8298dfd30e28376
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_conditioning_minimax_h3.py
@@ -0,0 +1,317 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Host-only gate for MiniMax-H3 FL2VA keyframe conditioning. Four silent contracts: ImageNet
+pixel statistics (not [-1, 1]), a sampled posterior under a request-independent seed, an fp16
+round trip on the sampled latent, and conditioning noise as the request generator's first draw."""
+
+import os
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+import transformers
+from PIL import Image
+
+from ....pipelines.minimax_h3 import conditioning as c
+from ....pipelines.minimax_h3.packing import (
+ MINIMAX_H3_KEYFRAME_NOISE_AUG,
+ MINIMAX_H3_TEXT_TAG,
+ MINIMAX_H3_VIDEO_TAG,
+ build_packed_sequence,
+ patchify_video_latents,
+ prepare_keyframe_image,
+ resolve_canvas_size,
+)
+from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline, draw_request_latents
+from ....pipelines.minimax_h3.scheduler import MiniMaxH3Scheduler
+
+LATENT_CHANNELS = 24
+LATENT_HEIGHT, LATENT_WIDTH = 34, 60
+PATCH = (1, 2, 2)
+CONDITION_SHAPES = ((1, LATENT_HEIGHT, LATENT_WIDTH),) * 2
+
+LATENTS_MEAN = tuple(0.1 * index for index in range(LATENT_CHANNELS))
+LATENTS_STD = tuple(1.0 + 0.05 * index for index in range(LATENT_CHANNELS))
+
+
+def _latents(seed=5, frames=1):
+ return torch.randn(
+ 1, LATENT_CHANNELS, frames, LATENT_HEIGHT, LATENT_WIDTH, generator=torch.Generator().manual_seed(seed)
+ )
+
+
+def _image(height=544, width=960, seed=3):
+ rng = np.random.default_rng(seed)
+ return Image.fromarray(rng.integers(0, 256, (height, width, 3), dtype=np.uint8))
+
+
+def test_pixels_use_imagenet_statistics():
+ assert c.MINIMAX_H3_PIXEL_MEAN == (0.485, 0.456, 0.406)
+ assert c.MINIMAX_H3_PIXEL_STD == (0.229, 0.224, 0.225)
+
+ image = _image()
+ pixels = c.normalize_keyframe_pixels(image)
+ assert pixels.shape == (1, 3, 1, image.size[1], image.size[0])
+ assert pixels.dtype == torch.float32
+
+ mean = torch.tensor(c.MINIMAX_H3_PIXEL_MEAN).view(1, -1, 1, 1, 1)
+ std = torch.tensor(c.MINIMAX_H3_PIXEL_STD).view(1, -1, 1, 1, 1)
+ raw = torch.from_numpy(np.array(image)).permute(2, 0, 1)[None, :, None].float().div(255.0)
+ assert torch.equal(pixels, (raw - mean) / std)
+ assert pixels.min() < -1.0
+
+
+def test_posterior_seed_is_independent_of_request_seed():
+ assert c.MINIMAX_H3_KEYFRAME_ENCODE_SEED == 42
+ moments = torch.randn(1, 2 * LATENT_CHANNELS, 1, 8, 8, generator=torch.Generator().manual_seed(1))
+
+ torch.manual_seed(999)
+ first = c.sample_posterior(moments)
+ torch.manual_seed(12345)
+ second = c.sample_posterior(moments)
+ assert torch.equal(first, second)
+
+ assert not torch.equal(first, c.sample_posterior(moments, seed=43))
+
+ # a sample, not the mean -- the mean would also be request-independent
+ mean = moments.chunk(2, dim=1)[0]
+ assert first.shape == mean.shape
+ assert not torch.equal(first, mean)
+
+
+def test_float16_round_trip_is_load_bearing():
+ """The fp16 round trip is load-bearing; asserted to change the rows so it cannot be optimized away."""
+ latents = _latents()
+ rows = c.keyframe_condition_rows(latents, LATENTS_MEAN, LATENTS_STD, PATCH)
+
+ expected_rows = (LATENT_HEIGHT // 2) * (LATENT_WIDTH // 2)
+ assert rows.shape == (expected_rows, LATENT_CHANNELS * 4)
+ assert rows.dtype == torch.float32
+
+ mean = torch.tensor(LATENTS_MEAN).view(1, LATENT_CHANNELS, 1, 1, 1)
+ std = torch.tensor(LATENTS_STD).view(1, LATENT_CHANNELS, 1, 1, 1)
+ exact = patchify_video_latents((latents - mean) / std, PATCH)
+ rounded = patchify_video_latents((latents.to(torch.float16).float() - mean) / std, PATCH)
+
+ assert torch.equal(rows, rounded)
+ assert not torch.equal(rows, exact)
+ assert (rows - exact).abs().max() > 1e-5
+
+
+def test_encode_keyframes_is_one_encode_per_image():
+ seen = []
+
+ def encode_clip(pixels):
+ seen.append(tuple(pixels.shape))
+ batch, _, frames, height, width = pixels.shape
+ return torch.zeros(batch, 2 * LATENT_CHANNELS, frames, height // 16, width // 16)
+
+ images = [_image(seed=1), _image(seed=2)]
+ rows = c.encode_keyframes(images, encode_clip, LATENTS_MEAN, LATENTS_STD, PATCH)
+
+ assert len(seen) == len(images)
+ assert all(shape[2] == 1 for shape in seen)
+ rows_per_frame = (544 // 16 // 2) * (960 // 16 // 2)
+ assert rows.shape == (len(images) * rows_per_frame, LATENT_CHANNELS * 4)
+
+
+def test_condition_noise_is_deterministic_per_seed():
+ seed = 42
+ first = c.keyframe_condition_noise(CONDITION_SHAPES, LATENT_CHANNELS, PATCH, torch.Generator().manual_seed(seed))
+ second = c.keyframe_condition_noise(CONDITION_SHAPES, LATENT_CHANNELS, PATCH, torch.Generator().manual_seed(seed))
+ assert torch.equal(first, second)
+
+ rows_per_frame = (LATENT_HEIGHT // 2) * (LATENT_WIDTH // 2)
+ assert first.shape == (len(CONDITION_SHAPES) * rows_per_frame, LATENT_CHANNELS * 4)
+
+
+def test_condition_noise_draws_per_condition_in_order():
+ """One concatenated draw would consume the same values, assign them differently, and leave the same generator state."""
+ generator = torch.Generator().manual_seed(0)
+ both = c.keyframe_condition_noise(CONDITION_SHAPES, LATENT_CHANNELS, PATCH, generator)
+
+ per_condition = torch.cat(
+ [
+ c.keyframe_condition_noise((shape,), LATENT_CHANNELS, PATCH, torch.Generator().manual_seed(0))
+ for shape in CONDITION_SHAPES[:1]
+ ]
+ )
+ assert torch.equal(both[: per_condition.shape[0]], per_condition)
+
+
+def test_condition_noise_is_the_first_draw_of_a_request():
+ ordered = torch.Generator().manual_seed(0)
+ c.keyframe_condition_noise(CONDITION_SHAPES, LATENT_CHANNELS, PATCH, ordered)
+ video_after_conditioning = torch.randn(3, 64, generator=ordered)
+
+ skipped = torch.Generator().manual_seed(0)
+ video_without_conditioning = torch.randn(3, 64, generator=skipped)
+
+ assert not torch.equal(video_after_conditioning, video_without_conditioning)
+
+
+def test_noise_augmentation_has_no_second_implementation():
+ """A local t*x0 + (1-t)*noise drifted 2.4e-7 against the scheduler (Python-double 1-t)."""
+ assert not hasattr(c, "noise_augment")
+ assert MINIMAX_H3_KEYFRAME_NOISE_AUG == 0.999
+
+ rows = torch.randn(20, 96, generator=torch.Generator().manual_seed(1))
+ noise = torch.randn(20, 96, generator=torch.Generator().manual_seed(2))
+ out = MiniMaxH3Scheduler(12.0).scale_noise(rows, MINIMAX_H3_KEYFRAME_NOISE_AUG, noise)
+
+ timestep = torch.tensor(MINIMAX_H3_KEYFRAME_NOISE_AUG, dtype=torch.float32)
+ assert torch.equal(out, timestep * rows + (1.0 - timestep) * noise)
+ drifted = MINIMAX_H3_KEYFRAME_NOISE_AUG * rows + (1.0 - MINIMAX_H3_KEYFRAME_NOISE_AUG) * noise
+ assert not torch.equal(out, drifted)
+
+
+PROD_LATENT_FRAMES = 37
+PROD_LATENT_H, PROD_LATENT_W = 48, 84
+PROD_AUDIO_LATENTS = 207
+PROD_AUDIO_CHANNELS = 32
+PROD_ROWS_PER_FRAME = (PROD_LATENT_H // 2) * (PROD_LATENT_W // 2) # 1008
+
+
+def _draw(num_keyframes: int, seed: int = 0):
+ return draw_request_latents(
+ torch.Generator().manual_seed(seed),
+ condition_latent_shapes=((1, PROD_LATENT_H, PROD_LATENT_W),) * num_keyframes,
+ latent_channels=LATENT_CHANNELS,
+ num_latent_frames=PROD_LATENT_FRAMES,
+ latent_height=PROD_LATENT_H,
+ latent_width=PROD_LATENT_W,
+ num_audio_latents=PROD_AUDIO_LATENTS,
+ audio_latent_channels=PROD_AUDIO_CHANNELS,
+ patch_size=PATCH,
+ )
+
+
+def test_t2va_draws_are_unchanged_by_the_keyframe_argument():
+ """The t2va no-regression proof: nothing may be drawn before the video draw."""
+ condition_noise, video_rows, audio_rows = _draw(0)
+ assert condition_noise is None, "t2va must draw no conditioning noise at all"
+
+ generator = torch.Generator().manual_seed(0)
+ expected_video = patchify_video_latents(
+ torch.randn(
+ (1, LATENT_CHANNELS, PROD_LATENT_FRAMES, PROD_LATENT_H, PROD_LATENT_W),
+ generator=generator,
+ dtype=torch.float32,
+ ),
+ PATCH,
+ )
+ expected_audio = torch.randn(
+ (PROD_AUDIO_LATENTS * 2, PROD_AUDIO_CHANNELS), generator=generator, dtype=torch.float32
+ )
+
+ assert torch.equal(video_rows, expected_video), "t2va video draw moved"
+ assert torch.equal(audio_rows, expected_audio), "t2va audio draw moved"
+
+
+@pytest.mark.parametrize("num_keyframes", [1, 2], ids=["first", "first_and_last"])
+def test_conditioning_noise_is_drawn_before_video_and_audio(num_keyframes):
+ condition_noise, video_rows, audio_rows = _draw(num_keyframes)
+ _, t2va_video, t2va_audio = _draw(0)
+
+ assert condition_noise is not None
+ assert condition_noise.shape == (num_keyframes * PROD_ROWS_PER_FRAME, LATENT_CHANNELS * PATCH[1] * PATCH[2])
+ expected_noise = c.keyframe_condition_noise(
+ ((1, PROD_LATENT_H, PROD_LATENT_W),) * num_keyframes,
+ LATENT_CHANNELS,
+ PATCH,
+ torch.Generator().manual_seed(0),
+ )
+ assert torch.equal(condition_noise, expected_noise), "conditioning noise is not the first draw"
+
+ # fl2va at a given seed does NOT reproduce t2va at that seed -- expected, not a regression
+ assert not torch.equal(video_rows, t2va_video), "the conditioning draw did not advance the stream"
+ assert not torch.equal(audio_rows, t2va_audio)
+
+
+def test_keyframe_rows_per_anchor_match_the_layout():
+ text_tags = torch.ones(39, dtype=torch.long)
+ for anchors in (("first",), ("last",), ("first", "last")):
+ layout = build_packed_sequence(
+ text_tags, PROD_LATENT_FRAMES, PROD_LATENT_H, PROD_LATENT_W, PROD_AUDIO_LATENTS, PATCH, anchors
+ )
+ assert layout.num_condition_video_rows == len(anchors) * PROD_ROWS_PER_FRAME
+ assert layout.num_condition_audio_rows == 0, "fl2va conditions video only; audio conditioning is ref2va"
+ expected = torch.arange(39, 39 + layout.num_condition_video_rows)
+ assert torch.equal(layout.video_indices[: layout.num_condition_video_rows], expected)
+
+ condition_noise, _, _ = _draw(len(anchors))
+ assert condition_noise.shape[0] == layout.num_condition_video_rows
+
+
+PROMPT = (
+ "A red fox trots across a snowy field at dawn, its breath visible in the cold air. "
+ "The low sun throws long blue shadows behind it, and loose snow lifts from each footfall."
+)
+
+
+def _weights_dir() -> str:
+ base = os.environ.get("MINIMAX_H3_MODEL_PATH")
+ if not base:
+ pytest.skip("set MINIMAX_H3_MODEL_PATH to a MiniMax-H3 diffusers snapshot")
+ return os.path.join(base, "text_encoder")
+
+
+def _snapshot_root() -> str:
+ return os.path.dirname(_weights_dir())
+
+
+def _presentation(prompt, keyframes):
+ """The real method called unbound: a reimplementation here would gate a second copy of the presentation."""
+ root = _snapshot_root()
+ stub = SimpleNamespace(
+ tokenizer=transformers.AutoTokenizer.from_pretrained(root, subfolder="tokenizer"),
+ image_processor=transformers.AutoImageProcessor.from_pretrained(root, subfolder="text_encoder"),
+ )
+ return MiniMaxH3Pipeline._build_presentation(stub, prompt, keyframes), stub
+
+
+@pytest.mark.parametrize("num_keyframes", [1, 2], ids=["first", "first_and_last"])
+def test_fl2va_presentation_matches_the_reference(num_keyframes):
+ """Token ids, H3 tags (whole vision block video) and Qwen's mm_token_type_ids (pads only) each fail silently alone."""
+ source = Image.fromarray(
+ (torch.rand(720, 1280, 3, generator=torch.Generator().manual_seed(0)) * 255).to(torch.uint8).numpy()
+ )
+ height, width = resolve_canvas_size(*source.size)
+ assert (height, width) == (768, 1344), f"expected the production canvas, got {width}x{height}"
+ keyframes = [prepare_keyframe_image(source, height, width, stretch=(index == 0)) for index in range(num_keyframes)]
+
+ (input_ids, tags, type_ids, pixel_values, grid_thw), stub = _presentation(PROMPT, keyframes)
+ tokenizer, processor = stub.tokenizer, stub.image_processor
+ merge = processor.merge_size**2
+
+ # the reference's builder: `encoders.py::MiniMaxH3TextEncoderStep.encode_prompt`
+ ref_ids, ref_tags = [], []
+ for index in range(num_keyframes):
+ num_image_tokens = int(grid_thw[index].prod()) // merge
+ label = tokenizer(f": ", add_special_tokens=False)["input_ids"]
+ block = (
+ [tokenizer.convert_tokens_to_ids("<|vision_start|>")]
+ + [tokenizer.convert_tokens_to_ids("<|image_pad|>")] * num_image_tokens
+ + [tokenizer.convert_tokens_to_ids("<|vision_end|>")]
+ )
+ ref_ids += label + block
+ ref_tags += [MINIMAX_H3_TEXT_TAG] * len(label) + [MINIMAX_H3_VIDEO_TAG] * len(block)
+ prompt_ids = tokenizer(PROMPT, add_special_tokens=False)["input_ids"]
+ ref_ids += prompt_ids
+ ref_tags += [MINIMAX_H3_TEXT_TAG] * len(prompt_ids)
+
+ assert input_ids.shape == (1, len(ref_ids)), f"{tuple(input_ids.shape)} vs {(1, len(ref_ids))}"
+ assert input_ids[0].tolist() == ref_ids, "token ids differ from the reference presentation"
+ assert tags.tolist() == ref_tags, "H3 row tags differ from the reference presentation"
+
+ full_processor = transformers.AutoProcessor.from_pretrained(_snapshot_root(), subfolder="processor")
+ expected_type_ids = torch.tensor(full_processor.create_mm_token_type_ids([input_ids[0].tolist()]))
+ assert torch.equal(type_ids, expected_type_ids), "mm_token_type_ids differ from the processor's own"
+ assert int(type_ids.sum()) == num_keyframes * 1008
+
+ assert int((tags == MINIMAX_H3_VIDEO_TAG).sum()) == num_keyframes * (1008 + 2)
+ assert grid_thw.shape[0] == num_keyframes
+ assert pixel_values.shape[0] == num_keyframes * 1008 * merge
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_packing_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_packing_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..e3c258bcd4eae26e1d5fc696fb99b541888fea10
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_packing_minimax_h3.py
@@ -0,0 +1,311 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Host-only parity gate for the MiniMax-H3 packed sequence: checkpoint contracts, exact
+assertions (drift is silent). Golden digests come from a run verified bit-exact against the
+minimax-h3 diffusers branch; when that branch is installed the layout is re-checked directly."""
+
+import hashlib
+
+import pytest
+import torch
+
+from ....pipelines.minimax_h3 import packing as p
+
+# (label, latent_height, latent_width, num_frames, keyframe_anchors)
+BRINGUP = ("bringup", 544 // 16, 960 // 16, 124, ("first", "last"))
+CANONICAL = ("canonical", 768 // 16, 1344 // 16, 192, ("first",))
+T2VA = ("t2va_768p_5s", 768 // 16, 1344 // 16, 124, ())
+
+# sequence_length, latent frames, audio latents, sha256[:16] of position_ids and token_tags
+GOLDEN = {
+ "bringup": (21301, 37, 207, "dcd77044930c7e5f", "320f230cf65b8e0b"),
+ "canonical": (60101, 57, 320, "3d2878587ddfdc7d", "bc71672198e60490"),
+}
+
+TEXT_LEN = 997
+VISION_BLOCK = slice(20, 70) # a keyframe's vision block: text rows tagged video
+
+
+def _text_tags():
+ tags = torch.ones(TEXT_LEN, dtype=torch.long)
+ tags[VISION_BLOCK] = p.MINIMAX_H3_VIDEO_TAG
+ return tags
+
+
+def _layout(latent_height, latent_width, num_frames, anchors):
+ return p.build_packed_sequence(
+ _text_tags(),
+ p.video_latent_num_frames(num_frames),
+ latent_height,
+ latent_width,
+ p.audio_latent_num_frames(num_frames),
+ (1, 2, 2),
+ anchors,
+ )
+
+
+def _digest(tensor):
+ return hashlib.sha256(tensor.contiguous().numpy().tobytes()).hexdigest()[:16]
+
+
+@pytest.mark.parametrize("num_frames,expected", [(1, 5), (5, 5), (6, 22), (22, 22), (120, 124), (192, 192)])
+def test_align_num_frames(num_frames, expected):
+ assert p.align_num_frames(num_frames) == expected
+
+
+@pytest.mark.parametrize("num_frames,latents", [(5, 2), (22, 7), (39, 12), (124, 37), (192, 57), (362, 107)])
+def test_video_latent_num_frames(num_frames, latents):
+ assert p.video_latent_num_frames(num_frames) == latents
+
+
+def test_video_latent_num_frames_rejects_unaligned(expect_error):
+ with expect_error(ValueError, "17n \\+ 5"):
+ p.video_latent_num_frames(120)
+
+
+@pytest.mark.parametrize("num_frames,latents", [(124, 207), (192, 320), (362, 603)])
+def test_audio_latent_num_frames(num_frames, latents):
+ assert p.audio_latent_num_frames(num_frames) == latents
+
+
+@pytest.mark.parametrize(
+ "aspect,canvas",
+ [
+ ((16, 9), (768, 1344)),
+ ((9, 16), (1344, 768)),
+ ((1, 1), (768, 768)),
+ ((4, 3), (768, 1024)),
+ ],
+)
+def test_resolve_canvas_size(aspect, canvas):
+ assert p.resolve_canvas_size(*aspect) == canvas
+ assert canvas[0] % p.MINIMAX_H3_CANVAS_MULTIPLE == 0
+ assert canvas[1] % p.MINIMAX_H3_CANVAS_MULTIPLE == 0
+
+
+@pytest.mark.parametrize(
+ "aspect,message",
+ [
+ ((5, 1), "aspect ratios from 1:4 to 4:1"),
+ ((1, 5), "aspect ratios from 1:4 to 4:1"),
+ ((0, 1), "must be positive"),
+ ((-16, 9), "must be positive"),
+ ],
+)
+def test_resolve_canvas_size_rejects_out_of_range(aspect, message, expect_error):
+ with expect_error(ValueError, message):
+ p.resolve_canvas_size(*aspect)
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL], ids=lambda c: c[0])
+def test_layout_matches_golden(case):
+ label, latent_height, latent_width, num_frames, anchors = case
+ sequence_length, latent_frames, audio_latents, position_sha, tag_sha = GOLDEN[label]
+
+ assert p.video_latent_num_frames(num_frames) == latent_frames
+ assert p.audio_latent_num_frames(num_frames) == audio_latents
+
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+ assert layout.sequence_length == sequence_length
+ assert layout.position_ids.dtype == torch.float64
+ assert _digest(layout.position_ids) == position_sha
+ assert _digest(layout.token_tags) == tag_sha
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL], ids=lambda c: c[0])
+def test_layout_structure(case):
+ _, latent_height, latent_width, num_frames, anchors = case
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+ rows_per_frame = (latent_height // 2) * (latent_width // 2)
+
+ assert layout.num_condition_video_rows == len(anchors) * rows_per_frame
+ assert layout.num_condition_audio_rows == 0
+
+ covered = torch.cat([layout.text_indices, layout.video_indices, layout.audio_indices]).sort().values
+ assert torch.equal(covered, torch.arange(layout.sequence_length))
+
+ assert torch.equal(layout.token_tags[layout.audio_indices[:1]], torch.tensor([p.MINIMAX_H3_AUDIO_TAG]))
+ assert (layout.token_tags[layout.video_indices] == p.MINIMAX_H3_VIDEO_TAG).all()
+ assert (layout.token_tags[VISION_BLOCK] == p.MINIMAX_H3_VIDEO_TAG).all()
+ assert (layout.token_tags[VISION_BLOCK.stop : TEXT_LEN] == p.MINIMAX_H3_TEXT_TAG).all()
+
+ audio_position_ids = layout.position_ids[layout.audio_indices]
+ assert (audio_position_ids[:, 1] == 0).all()
+ assert audio_position_ids[:, 2].unique().numel() == 2
+
+
+def test_keyframe_anchor_times():
+ """The `last` anchor overshoots the final frame by 5.0 on purpose -- looks like an off-by-one and is not."""
+ _, latent_height, latent_width, num_frames, _ = BRINGUP
+ latent_frames = p.video_latent_num_frames(num_frames)
+ layout = _layout(latent_height, latent_width, num_frames, ("first", "last"))
+ rows_per_frame = (latent_height // 2) * (latent_width // 2)
+
+ first_time = layout.position_ids[TEXT_LEN, 0]
+ last_time = layout.position_ids[TEXT_LEN + rows_per_frame, 0]
+ video_time = layout.position_ids[layout.video_indices[layout.num_condition_video_rows :], 0]
+
+ assert first_time == float(TEXT_LEN)
+ assert first_time == video_time.min()
+ assert float(last_time) == pytest.approx(
+ TEXT_LEN + p._temporal_position_span(latent_frames) - p._ROPE_FRAME_RESCALE
+ )
+ assert last_time > video_time.max()
+ assert float(last_time - video_time.max()) == pytest.approx(5.0)
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL], ids=lambda c: c[0])
+def test_row_timesteps_pin_condition_rows(case):
+ _, latent_height, latent_width, num_frames, anchors = case
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+ # expectations round through fp32: 0.7 is not representable
+ video_t, audio_t, cond_t = torch.tensor([0.7, 0.5, p.MINIMAX_H3_KEYFRAME_NOISE_AUG], dtype=torch.float32).tolist()
+ timesteps, indices = p.build_row_timesteps(layout, video_t, audio_t, cond_t, cond_t)
+
+ assert timesteps.tolist() == sorted({video_t, audio_t, cond_t})
+ resolved = timesteps[indices]
+ condition_rows = layout.video_indices[: layout.num_condition_video_rows]
+ target_rows = layout.video_indices[layout.num_condition_video_rows :]
+
+ assert (resolved[condition_rows] == cond_t).all()
+ assert (resolved[target_rows] == video_t).all()
+ assert (resolved[layout.audio_indices] == audio_t).all()
+ assert (resolved[layout.text_indices] == video_t).all()
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL], ids=lambda c: c[0])
+def test_adaln_index_ranges_round_trip(case):
+ _, latent_height, latent_width, num_frames, anchors = case
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+ _, indices = p.build_row_timesteps(layout, 0.7, 0.5, 0.999, 0.999)
+ table_rows = p.adaln_indices(layout.token_tags, indices)
+
+ runs = p.adaln_index_ranges(table_rows)
+ rebuilt = torch.empty_like(table_rows)
+ for start, stop, value in runs:
+ rebuilt[start:stop] = value
+ assert torch.equal(rebuilt, table_rows)
+ assert len(runs) <= 8
+
+
+def test_adaln_index_ranges_are_shard_local():
+ _, latent_height, latent_width, num_frames, anchors = CANONICAL
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+ _, indices = p.build_row_timesteps(layout, 0.7, 0.5, 0.999, 0.999)
+ table_rows = p.adaln_indices(layout.token_tags, indices)
+
+ sp_factor = 8
+ shard = layout.sequence_length // sp_factor
+ counts = [len(p.adaln_index_ranges(table_rows[i * shard : (i + 1) * shard])) for i in range(sp_factor)]
+
+ assert sum(count == 1 for count in counts) >= sp_factor - 2
+ assert max(counts) <= 8
+
+
+def test_adaln_index_ranges_empty():
+ assert p.adaln_index_ranges(torch.empty(0, dtype=torch.long)) == []
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL], ids=lambda c: c[0])
+def test_patchify_round_trip(case):
+ _, latent_height, latent_width, num_frames, _ = case
+ latent_frames = p.video_latent_num_frames(num_frames)
+ latents = torch.randn(1, 24, latent_frames, latent_height, latent_width)
+
+ rows = p.patchify_video_latents(latents, (1, 2, 2))
+ assert rows.shape == (latent_frames * (latent_height // 2) * (latent_width // 2), 24 * 4)
+ restored = p.unpatchify_video_tokens(rows, latent_frames, latent_height, latent_width, 24, (1, 2, 2))
+ assert torch.equal(restored, latents)
+
+
+def test_unpack_audio_tokens_is_channel_major():
+ num_audio_latents, channels = 207, p.MINIMAX_H3_AUDIO_CHANNELS
+ rows = torch.randn(num_audio_latents * channels, 32)
+ unpacked = p.unpack_audio_tokens(rows, num_audio_latents)
+
+ assert unpacked.shape == (channels, 32, num_audio_latents)
+ assert torch.equal(unpacked[0], rows[:num_audio_latents].transpose(0, 1))
+ assert torch.equal(unpacked[1], rows[num_audio_latents:].transpose(0, 1))
+
+
+ROPE_FREQ_DIM = 16
+ROPE_THETA = 10000.0
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL, T2VA], ids=lambda c: c[0])
+def test_rope_tables_match_reference(case):
+ """torch.equal, not PCC: merely-close is the regression."""
+ rope_module = pytest.importorskip(
+ "diffusers.models.transformers.transformer_minimax_h3",
+ reason="requires the minimax-h3 diffusers branch",
+ )
+ _, latent_height, latent_width, num_frames, anchors = case
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+
+ reference = rope_module.MiniMaxH3RotaryPosEmbed(rope_freq_dim=ROPE_FREQ_DIM, rope_theta=ROPE_THETA)
+ with torch.no_grad():
+ want_cos, want_sin = reference(layout.position_ids)
+ got_cos, got_sin = p.build_rope_tables(layout.position_ids, rope_freq_dim=ROPE_FREQ_DIM, rope_theta=ROPE_THETA)
+
+ assert torch.equal(got_cos, want_cos), "rope cos is not bit-exact against the reference"
+ assert torch.equal(got_sin, want_sin), "rope sin is not bit-exact against the reference"
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL, T2VA], ids=lambda c: c[0])
+def test_rope_tables_shape_and_dtype(case):
+ _, latent_height, latent_width, num_frames, anchors = case
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+ cos, sin = p.build_rope_tables(layout.position_ids, rope_freq_dim=ROPE_FREQ_DIM, rope_theta=ROPE_THETA)
+
+ assert cos.shape == sin.shape == (layout.sequence_length, 2 * 3 * ROPE_FREQ_DIM)
+ assert cos.shape[-1] == 96 < 128, "rotary width must stay narrower than head_dim"
+ assert cos.dtype == sin.dtype == torch.float32
+ assert torch.isfinite(cos).all() and torch.isfinite(sin).all()
+ half = cos.shape[-1] // 2
+ assert torch.equal(cos[:, :half], cos[:, half:])
+ assert torch.equal(sin[:, :half], sin[:, half:])
+
+
+def test_rope_tables_distinguish_the_three_axes():
+ """A collapse to 1-D rope would look like weak spatial coherence."""
+ _, latent_height, latent_width, num_frames, anchors = T2VA
+ layout = _layout(latent_height, latent_width, num_frames, anchors)
+ cos, _ = p.build_rope_tables(layout.position_ids, rope_freq_dim=ROPE_FREQ_DIM, rope_theta=ROPE_THETA)
+
+ row = int(layout.video_indices[-1])
+ t_block, h_block, w_block = (
+ cos[row, :ROPE_FREQ_DIM],
+ cos[row, ROPE_FREQ_DIM : 2 * ROPE_FREQ_DIM],
+ cos[row, 2 * ROPE_FREQ_DIM : 3 * ROPE_FREQ_DIM],
+ )
+ assert not torch.equal(t_block, h_block)
+ assert not torch.equal(h_block, w_block)
+
+
+@pytest.mark.parametrize("case", [BRINGUP, CANONICAL], ids=lambda c: c[0])
+def test_matches_diffusers_reference(case):
+ reference = pytest.importorskip(
+ "diffusers.modular_pipelines.minimax_h3.packing",
+ reason="requires the minimax-h3 diffusers branch",
+ )
+ _, latent_height, latent_width, num_frames, anchors = case
+ kwargs = dict(
+ text_token_tags=_text_tags(),
+ num_latent_frames=p.video_latent_num_frames(num_frames),
+ latent_height=latent_height,
+ latent_width=latent_width,
+ num_audio_latents=p.audio_latent_num_frames(num_frames),
+ patch_size=(1, 2, 2),
+ keyframe_anchors=anchors,
+ )
+ ours = p.build_packed_sequence(**kwargs)
+ theirs = reference.build_packed_sequence(**kwargs)
+
+ assert ours.sequence_length == theirs.sequence_length
+ assert torch.equal(ours.position_ids, theirs.position_ids)
+ assert torch.equal(ours.token_tags, theirs.token_tags)
+ assert torch.equal(ours.video_indices, theirs.video_indices)
+ assert torch.equal(ours.audio_indices, theirs.audio_indices)
+ assert ours.num_condition_video_rows == theirs.num_condition_video_rows
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_performance_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_performance_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd590a9f648ab9d108d74801389a98faf1e1a567
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_performance_minimax_h3.py
@@ -0,0 +1,222 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Device perf of one MiniMax-H3 transformer block at 768P; run under `scripts/run_safe_pytest.sh --profile`,
+then `tt-perf-report --start-signpost start --end-signpost stop`. Profile one duration at a time with
+`-k`: a multi-parameter profiled run yields a CSV containing only the first parameter's ops."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from diffusers.models.transformers.transformer_minimax_h3 import MINIMAX_H3_MODALITY_NUM, MiniMaxH3RotaryPosEmbed
+from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerBlock as TorchMiniMaxH3Block
+from loguru import logger
+from tracy import signpost
+
+import ttnn
+
+from ....models.transformers.minimax_h3.attention_minimax_h3 import prepare_rope_tables
+from ....models.transformers.minimax_h3.transformer_block_minimax_h3 import MiniMaxH3TransformerBlock
+from ....parallel.config import DiTParallelConfig, ParallelFactor
+from ....parallel.manager import CCLManager
+from ....pipelines.minimax_h3.packing import (
+ MINIMAX_H3_FPS,
+ align_num_frames,
+ audio_latent_num_frames,
+ resolve_canvas_size,
+ video_latent_num_frames,
+)
+from ....utils.tensor import bf16_tensor_2dshard, from_torch, local_device_to_torch
+from ....utils.test import skip_if_unsupported_num_links
+from .common import (
+ GALAXY_RING,
+ REAL_BLOCK_CONFIG,
+ ROPE_FREQ_DIM,
+ ROPE_THETA,
+ TT_BLOCK_CONFIG,
+ packed_layout,
+ upload_rope,
+)
+
+HIDDEN_SIZE = REAL_BLOCK_CONFIG["hidden_size"]
+HEAD_DIM = REAL_BLOCK_CONFIG["attention_head_dim"]
+TIME_EMBED_DIM = REAL_BLOCK_CONFIG["time_embed_dim"]
+
+PATCH_SIZE = (1, 2, 2)
+VAE_SPATIAL_DOWNSAMPLE = 16 # prod(spatial_downsample_factors) from the video VAE config
+NUM_TEXT_TOKENS = 512
+ASPECT = (16, 9)
+
+
+def _packed_sizes(duration_s: float) -> dict:
+ """Token counts for `duration_s` seconds of 768P video, derived from the pipeline's own packing helpers."""
+ height, width = resolve_canvas_size(*ASPECT)
+ tokens_per_latent_frame = (height // VAE_SPATIAL_DOWNSAMPLE // PATCH_SIZE[1]) * (
+ width // VAE_SPATIAL_DOWNSAMPLE // PATCH_SIZE[2]
+ )
+ num_frames = align_num_frames(int(duration_s * MINIMAX_H3_FPS))
+ latent_frames = video_latent_num_frames(num_frames)
+ num_audio = audio_latent_num_frames(num_frames)
+ num_video = latent_frames * tokens_per_latent_frame
+ return {
+ "height": height,
+ "width": width,
+ "num_frames": num_frames,
+ "latent_frames": latent_frames,
+ "grid_h": height // VAE_SPATIAL_DOWNSAMPLE // PATCH_SIZE[1],
+ "grid_w": width // VAE_SPATIAL_DOWNSAMPLE // PATCH_SIZE[2],
+ "num_video": num_video,
+ "num_audio": num_audio,
+ "num_text": NUM_TEXT_TOKENS,
+ "seq_len": NUM_TEXT_TOKENS + num_audio + num_video,
+ }
+
+
+@GALAXY_RING
+@pytest.mark.parametrize(
+ "duration_s",
+ [
+ pytest.param(5.0, id="5s_768p"),
+ pytest.param(10.0, id="10s_768p"),
+ pytest.param(15.0, id="15s_768p"),
+ ],
+)
+@pytest.mark.parametrize(
+ "sp_simulate",
+ [
+ pytest.param(1, id="sp_sim1"),
+ pytest.param(4, id="sp_sim4"),
+ ],
+)
+def test_minimax_h3_transformer_block_perf(
+ mesh_device: ttnn.MeshDevice,
+ sp_axis: int,
+ tp_axis: int,
+ num_links: int,
+ duration_s: float,
+ sp_simulate: int,
+ is_fsdp: bool,
+ topology: ttnn.Topology,
+ reset_seeds,
+) -> None:
+ skip_if_unsupported_num_links(mesh_device, num_links)
+ # Simulate a larger SP mesh (e.g. 4x32) on a smaller one (4x8) by shrinking the total sequence
+ # so each device carries a shard the larger mesh would produce. `sp_simulate` is that SP ratio.
+ SIM = sp_simulate
+
+ sp_factor = tuple(mesh_device.shape)[sp_axis]
+ tp_factor = tuple(mesh_device.shape)[tp_axis]
+
+ sizes = _packed_sizes(duration_s)
+ seq_len = sizes["seq_len"]
+ alignment = sp_factor * ttnn.TILE_SIZE * SIM
+ padded_len = ((seq_len + alignment - 1) // alignment) * alignment
+ padded_len = padded_len // SIM
+ logger.info(
+ f"{duration_s:g}s @ {sizes['height']}x{sizes['width']}: {sizes['num_frames']} frames -> "
+ f"{sizes['latent_frames']} latent frames x {sizes['grid_h']}x{sizes['grid_w']} patches = "
+ f"{sizes['num_video']} video + {sizes['num_audio']} audio + {sizes['num_text']} text "
+ f"= seq_len {seq_len} (padded {padded_len}, {padded_len // sp_factor} rows/device)"
+ )
+
+ num_timesteps = 2
+ # Simulate a 4x32 per-device shard on a 4x8 mesh: shrink the total sequence by SIM (32/8) so each
+ # 4x8 device carries a 4x32-sized shard. num_video must stay a whole number of (grid_h, grid_w)
+ # frames, so floor it to a frame boundary rather than dividing the raw token count.
+ frame = sizes["grid_h"] * sizes["grid_w"]
+ sim_num_video = (sizes["num_video"] // SIM // frame) * frame
+ sim_seq_len = sizes["num_text"] // SIM + sizes["num_audio"] // SIM + sim_num_video
+ position_ids, tags, timestep_indices = packed_layout(
+ sizes["num_text"] // SIM,
+ sizes["num_audio"] // SIM,
+ sim_num_video,
+ (sizes["grid_h"], sizes["grid_w"]),
+ padded_len,
+ )
+ adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + tags.clamp(min=0)
+
+ # Built only to source a correctly-keyed random state dict; its forward is never called.
+ torch_block = TorchMiniMaxH3Block(**REAL_BLOCK_CONFIG).to(torch.float32)
+
+ rope = MiniMaxH3RotaryPosEmbed(rope_freq_dim=ROPE_FREQ_DIM, rope_theta=ROPE_THETA)
+ with torch.no_grad():
+ rope_cos, rope_sin = rope(position_ids)
+ rope_cos, rope_sin = prepare_rope_tables(rope_cos, rope_sin, HEAD_DIM)
+
+ ccl_manager = CCLManager(mesh_device=mesh_device, num_links=num_links, topology=topology)
+ parallel_config = DiTParallelConfig(
+ tensor_parallel=ParallelFactor(mesh_axis=tp_axis, factor=tp_factor),
+ sequence_parallel=ParallelFactor(mesh_axis=sp_axis, factor=sp_factor),
+ cfg_parallel=None,
+ )
+
+ tt_block = MiniMaxH3TransformerBlock(
+ **TT_BLOCK_CONFIG,
+ rotary_dim=2 * 3 * ROPE_FREQ_DIM,
+ mesh_device=mesh_device,
+ ccl_manager=ccl_manager,
+ parallel_config=parallel_config,
+ is_fsdp=is_fsdp,
+ )
+ tt_block.load_torch_state_dict(torch_block.state_dict())
+ del torch_block
+
+ if SIM > 1:
+ # The exp-ring gate keys on sequence_parallel.factor == 32, a proxy for "the per-device
+ # shard is the 4x32 shape". SP simulation produces exactly that shard on a smaller mesh,
+ # but the model only sees the real mesh's factor — force the flag so the simulated run
+ # exercises what a real 4x32 would.
+ tt_block.attn.use_exp_ring_sdpa = tt_block.attn.use_ring and sp_factor * SIM == 32
+ if sp_factor * SIM == 32:
+ assert tt_block.attn.use_exp_ring_sdpa, "simulated 4x32 run must exercise the exp ring SDPA path"
+ assert tt_block.attn._exp_sdpa_program_config(padded_len // sp_factor) is not None, (
+ "exp ring SDPA has no valid program config for the simulated shard "
+ f"({padded_len // sp_factor} tokens/device); the run would silently measure the normal ring path"
+ )
+
+ tt_spatial = bf16_tensor_2dshard(
+ torch.randn(1, 1, padded_len, HIDDEN_SIZE),
+ device=mesh_device,
+ shard_mapping={sp_axis: 2, tp_axis: 3},
+ )
+ tt_temb = from_torch(torch.randn(1, 1, num_timesteps, TIME_EMBED_DIM), device=mesh_device, dtype=ttnn.float32)
+ tt_adaln = from_torch(
+ adaln_indices.to(torch.int32).reshape(1, 1, 1, padded_len),
+ device=mesh_device,
+ dtype=ttnn.int32,
+ layout=ttnn.Layout.ROW_MAJOR,
+ mesh_axes=[..., None, sp_axis],
+ )
+ tt_rope_cos, tt_rope_sin = upload_rope(rope_cos, rope_sin, mesh_device=mesh_device, sp_axis=sp_axis)
+
+ def run_block() -> ttnn.Tensor:
+ out = tt_block(
+ tt_spatial,
+ N=sim_seq_len, # simulated unpadded length: ring attention masks the pad tail via logical_n
+ temb=tt_temb,
+ adaln_indices=tt_adaln,
+ rope_cos=tt_rope_cos,
+ rope_sin=tt_rope_sin,
+ )
+ ttnn.synchronize_device(mesh_device)
+ return out
+
+ logger.info("iteration 1: compiling kernels and populating the program cache")
+ run_block()
+
+ logger.info("iteration 2: warm run (the profiled region)")
+ signpost("start")
+ tt_out = run_block()
+ signpost("stop")
+
+ assert tuple(tt_out.shape) == (
+ 1,
+ 1,
+ padded_len // sp_factor,
+ HIDDEN_SIZE // tp_factor,
+ ), f"unexpected output shape {tuple(tt_out.shape)}"
+ local = local_device_to_torch(tt_out).float()
+ assert torch.isfinite(local).all(), "block output contains NaN or Inf"
+ logger.info(f"output {tuple(tt_out.shape)}, local shard std={local.std().item():.4f}")
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae9baf07011bcdb474d62d5e1cba2dd1c5ce433b
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py
@@ -0,0 +1,327 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""Gate M8f: MiniMax-H3 VAE performance baselines, gated against ``EXPECTED_METRICS`` budgets
+(style of ``tests/models/wan2_2/test_performance_wan.py``). Tracy capture entry points live
+in ``tools/tracy_decode_harness.py``; quality gates live in the vae/audio test files."""
+
+import os
+import time
+
+import pytest
+import torch
+
+import ttnn
+
+from ....models.audio_vae.minimax_h3.convert_minimax_h3_audio import convert_minimax_h3_audio_state_dict
+from ....models.audio_vae.minimax_h3.encoder_minimax_h3_audio import MiniMaxH3AudioEncoder
+from ....models.vae.minimax_h3.vae_minimax_h3 import MiniMaxH3Vae, MiniMaxH3VaeConfig
+from ....parallel.manager import CCLManager
+from .common import (
+ CLIP_FRAMES,
+ DECODE_LATENT_FRAMES,
+ LATENT_TILE,
+ TILE,
+ build_audio_decoder,
+ build_visual_decoder,
+ build_visual_encoder,
+ load_config,
+ random_decoder_state,
+ random_encoder_state,
+ weights_subdir,
+)
+
+SINGLE_DEVICE = [pytest.param((1, 1), {"l1_small_size": 65536}, id="single_device")]
+
+HOP_LENGTH = 800
+
+# Work units per working point: 768P = 4x7 tiles, 1440P = 8x13; clip/chunk counts per the 17n+5 -> 5n+2 rule.
+WORK_UNITS = {
+ "768P_5s": {"tiles": 28, "encode_clips": 8, "decode_chunks": 7},
+ "768P_10s": {"tiles": 28, "encode_clips": 15, "decode_chunks": 14},
+ "768P_15s": {"tiles": 28, "encode_clips": 22, "decode_chunks": 21},
+ "1440P_5s": {"tiles": 104, "encode_clips": 8, "decode_chunks": 7},
+ "1440P_10s": {"tiles": 104, "encode_clips": 15, "decode_chunks": 14},
+ "1440P_15s": {"tiles": 104, "encode_clips": 22, "decode_chunks": 21},
+}
+
+AUDIO_LATENT_FRAMES_5S = 207 # ~5 s at 40 latent fps; budgets were calibrated at 207, not the exact 200
+
+# Seconds per measurement. Generous regression bars, not tuned targets; visual entries gate a 32-unit wave.
+# Audio decode runs the accurate-mode defaults (~13 s eager, ~3x the retired fast path); 60 s stays generous.
+EXPECTED_METRICS = {
+ "visual_encoder_clip_wave": 20.0,
+ "visual_encoder_keyframe_wave": 5.0,
+ "visual_decoder_wave": 20.0,
+ "audio_encode_5s": 60.0,
+ "audio_decode_5s": 60.0,
+}
+
+
+def _time_it(fn, *, warmup: int = 1, iterations: int = 2, mesh_device=None) -> float:
+ for _ in range(warmup):
+ fn()
+ if mesh_device is not None:
+ ttnn.synchronize_device(mesh_device)
+ best = float("inf")
+ for _ in range(iterations):
+ start = time.time()
+ fn()
+ if mesh_device is not None:
+ ttnn.synchronize_device(mesh_device)
+ best = min(best, time.time() - start)
+ return best
+
+
+def _report(measurements: dict[str, float]) -> None:
+ from loguru import logger
+
+ failures = []
+ for key, seconds in measurements.items():
+ budget = EXPECTED_METRICS[key]
+ logger.info(f"PERF {key}: {seconds:.3f} s (budget {budget:.1f} s)")
+ if seconds > budget:
+ failures.append(f"{key} took {seconds:.3f} s, budget {budget:.1f} s")
+ assert not failures, "\n".join(failures)
+
+
+MESH_4X8 = [
+ pytest.param(
+ (4, 8),
+ {"fabric_config": None, "require_exact_physical_num_devices": True, "l1_small_size": 65536},
+ id="mesh4x8",
+ )
+]
+
+
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESH_4X8, indirect=["mesh_device", "device_params"])
+def test_visual_data_parallel_throughput(mesh_device):
+ """Per-wave time with one work unit per device, gated, plus full-video projections."""
+ from loguru import logger
+
+ weights_dir = weights_subdir("vae")
+ if weights_dir is None:
+ pytest.skip("MiniMax-H3 vae not found; set MINIMAX_H3_MODEL_PATH")
+ config = load_config(weights_dir)
+ devices = mesh_device.get_num_devices()
+ torch.manual_seed(0)
+
+ measurements = {}
+
+ for label, num_frames, taps in (
+ ("visual_encoder_clip_wave", CLIP_FRAMES, 3),
+ ("visual_encoder_keyframe_wave", 1, 1),
+ ):
+ encoder = build_visual_encoder(config, mesh_device, num_frames, temporal_taps=taps)
+ encoder.load_torch_state_dict(random_encoder_state(config))
+ x = torch.randn(devices, num_frames, TILE, TILE, encoder.conv_in.in_channels)
+ x_device = ttnn.from_torch(
+ x,
+ dtype=ttnn.float32,
+ device=mesh_device,
+ layout=ttnn.ROW_MAJOR_LAYOUT,
+ mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=0),
+ )
+ seconds = _time_it(lambda: encoder(x_device), iterations=8, mesh_device=mesh_device)
+ logger.info(f"PERF {label} of {devices} units: {seconds:.3f} s ({seconds / devices:.4f} s/unit)")
+ measurements[label] = seconds
+
+ decoder = build_visual_decoder(config, mesh_device)
+ decoder.load_torch_state_dict(random_decoder_state(config))
+ tokens = torch.randn(devices, DECODE_LATENT_FRAMES * LATENT_TILE * LATENT_TILE, config["latent_channels"])
+ tokens_device = ttnn.from_torch(
+ tokens,
+ dtype=ttnn.bfloat16,
+ device=mesh_device,
+ layout=ttnn.TILE_LAYOUT,
+ mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=0),
+ )
+ decoder_wave = _time_it(lambda: decoder(tokens_device), iterations=8, mesh_device=mesh_device)
+ # 10.6 = TFLOP per invocation at 1797 tokens x 36 layers x dim 2048.
+ logger.info(
+ f"PERF visual_decoder_wave of {devices} units: {decoder_wave:.3f} s "
+ f"({decoder_wave / devices:.4f} s/unit, {10.6 / (decoder_wave / devices):.1f} TFLOP/s effective per unit)"
+ )
+ measurements["visual_decoder_wave"] = decoder_wave
+
+ encoder_wave = measurements["visual_encoder_clip_wave"]
+ for name, units in WORK_UNITS.items():
+ encode_units = units["tiles"] * units["encode_clips"]
+ decode_units = units["tiles"] * units["decode_chunks"]
+ encode_waves = -(-encode_units // devices)
+ decode_waves = -(-decode_units // devices)
+ encode_total = encode_waves * encoder_wave
+ decode_total = decode_waves * decoder_wave
+ logger.info(
+ f"PROJECTED {name}: encode {encode_units} units / {encode_waves} waves = {encode_total:.1f} s, "
+ f"decode {decode_units} units / {decode_waves} waves = {decode_total:.1f} s, "
+ f"total {encode_total + decode_total:.1f} s"
+ )
+
+ _report(measurements)
+
+
+@pytest.mark.parametrize(("mesh_device", "device_params"), SINGLE_DEVICE, indirect=["mesh_device", "device_params"])
+def test_audio_baselines(mesh_device):
+ """Audio encode/decode timing baselines at 5 s, against their budgets."""
+ weights_dir = weights_subdir("audio_vae")
+ if weights_dir is None:
+ pytest.skip("MiniMax-H3 audio_vae not found; set MINIMAX_H3_MODEL_PATH")
+ from safetensors.torch import load_file
+
+ config = load_config(weights_dir)
+ converted = convert_minimax_h3_audio_state_dict(
+ load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors"))
+ )
+
+ torch.manual_seed(2)
+ waveform = torch.randn(2, 1, AUDIO_LATENT_FRAMES_5S * HOP_LENGTH) * 0.1
+ latents = torch.randn(2, config["latent_channels"], AUDIO_LATENT_FRAMES_5S) * 0.1
+
+ encoder = MiniMaxH3AudioEncoder(
+ encoder_dim=config["encoder_dim"],
+ encoder_rates=tuple(config["encoder_rates"]),
+ latent_dim=config["latent_dim"],
+ latent_channels=config["latent_channels"],
+ num_attention_heads=config["num_attention_heads"],
+ mesh_device=mesh_device,
+ )
+ decoder = build_audio_decoder(config, mesh_device)
+ # `strict=False`: the converted dict carries both halves' tensors.
+ encoder.load_torch_state_dict(converted, strict=False)
+ decoder.load_torch_state_dict(converted, strict=False)
+
+ measurements = {
+ "audio_encode_5s": _time_it(lambda: encoder(waveform), iterations=1, mesh_device=mesh_device),
+ "audio_decode_5s": _time_it(lambda: decoder(latents), iterations=1, mesh_device=mesh_device),
+ }
+ _report(measurements)
+
+
+# ---- whole-stage `decode` timing, per output type -------------------------------------------
+#
+# `test_visual_data_parallel_throughput` times one decoder wave; this times the stage the pipeline
+# bills to "VAE decode" -- tiling, upload, readback, unpatchify, stitch and the temporal assembly --
+# and prints `_report_profile`'s breakdown, which is the point: a stage total alone cannot say which
+# of device / readback / stitch moved.
+#
+# RING, not the `fabric_config=None` the throughput test uses: `_decode_clip_device_stitched`
+# hardcodes `all_gather(topology=Ring)`, so `yuv420` cannot run on a line fabric. Built with a
+# `CCLManager` because `_read_wave_units` picks its reader on `ccl_manager is None` -- without one a
+# single-host run silently takes the `ConcatMeshToTensor` path and measures something production
+# never runs.
+#
+# export MINIMAX_H3_MODEL_PATH=/path/to/MiniMax-H3
+# pytest models/tt_dit/tests/models/minimax_h3/test_performance_vae_minimax_h3.py -k decode_stage -s
+
+MESH_4X8_RING = [
+ pytest.param(
+ (4, 8),
+ {
+ "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING,
+ "require_exact_physical_num_devices": True,
+ "l1_small_size": 65536,
+ },
+ id="4x8ring",
+ )
+]
+
+# 1344x768 at 16x spatial compression. Latent frame counts follow the 17n -> 5n rule less
+# `token_drop`, so 124 frames (5 s) is 37 and 362 (15 s) is 107.
+DECODE_STAGE_LATENT_HW = (48, 84)
+DECODE_STAGE_FRAMES = {5: 37, 15: 107}
+# 7 latent frames is one chunk, the cheapest warm-up that still exercises a whole wave.
+DECODE_STAGE_WARM_FRAMES = 7
+
+MINIMAX_H3_PIXEL_MEAN = (0.485, 0.456, 0.406)
+MINIMAX_H3_PIXEL_STD = (0.229, 0.224, 0.225)
+
+# Each mode is one lever, so a row of this table measures that lever alone.
+DECODE_STAGE_MODES = {
+ # Today's default: bf16 tiles in token space, blended and unpatchified on host.
+ "float": {},
+ # Same path, uint8 across PCIe. Halves the transfer.
+ "uint8": {"pixel_denorm": (MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD), "readback_uint8": True},
+ # Stitch, clamp and colour-convert on device; read one planar canvas at 1.5 bytes/pixel.
+ "yuv420": {"pixel_denorm": (MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD), "device_stitch": True},
+}
+
+
+def _decode_stage_state(weights_dir: str) -> dict[str, torch.Tensor]:
+ """Just the decoder-side tensors; over half the checkpoint is the encoder."""
+ import json
+
+ from safetensors.torch import load_file
+
+ index_path = os.path.join(weights_dir, "diffusion_pytorch_model.safetensors.index.json")
+ if not os.path.isfile(index_path):
+ return {
+ k: v
+ for k, v in load_file(os.path.join(weights_dir, "diffusion_pytorch_model.safetensors")).items()
+ if k.startswith(("decoder.", "post_quant_conv."))
+ }
+ weight_map = json.loads(open(index_path).read())["weight_map"]
+ wanted = {k: f for k, f in weight_map.items() if k.startswith(("decoder.", "post_quant_conv."))}
+ state: dict[str, torch.Tensor] = {}
+ for shard in sorted(set(wanted.values())):
+ loaded = load_file(os.path.join(weights_dir, shard))
+ state.update({k: loaded[k] for k in wanted if k in loaded})
+ return state
+
+
+@pytest.mark.timeout(3600)
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESH_4X8_RING, indirect=["mesh_device", "device_params"])
+@pytest.mark.parametrize("seconds", [5, 15], ids=["5s", "15s"])
+@pytest.mark.parametrize("mode", sorted(DECODE_STAGE_MODES), ids=sorted(DECODE_STAGE_MODES))
+def test_decode_stage(mesh_device, seconds, mode):
+ """Measurement driver, not a gate: reports the split, asserts nothing about wall time."""
+ from loguru import logger
+
+ weights_dir = weights_subdir("vae")
+ if weights_dir is None:
+ pytest.skip("MiniMax-H3 vae not found; set MINIMAX_H3_MODEL_PATH")
+
+ config = MiniMaxH3VaeConfig.from_pretrained(weights_dir)
+ ccl_manager = CCLManager(mesh_device, num_links=2, topology=ttnn.Topology.Ring)
+ vae = MiniMaxH3Vae(config, mesh_device=mesh_device, ccl_manager=ccl_manager, **DECODE_STAGE_MODES[mode])
+ output_type = "yuv420" if mode == "yuv420" else "float"
+
+ vae.load_decoder_state(_decode_stage_state(weights_dir))
+ # Build the per-shape decoder outside the timed region, as `_prepare_vae(decode_shape=...)` does:
+ # its weight upload is seconds and would otherwise land inside the first `decode`.
+ vae._decoder_for(config.tokens_chunk_size + config.token_overlap, 16, 16)
+
+ torch.manual_seed(0)
+ latent_h, latent_w = DECODE_STAGE_LATENT_HW
+ warm = torch.randn(1, config.latent_channels, DECODE_STAGE_WARM_FRAMES, latent_h, latent_w)
+ latents = torch.randn(1, config.latent_channels, DECODE_STAGE_FRAMES[seconds], latent_h, latent_w)
+
+ vae.decode(warm, output_type=output_type)
+
+ started = time.time()
+ video = vae.decode(latents, output_type=output_type)
+ elapsed = time.time() - started
+
+ profile = dict(vae.last_decode_profile)
+ logger.info(f"DECODE_STAGE mode={mode} seconds={seconds} total={elapsed:.3f}s video={tuple(video.shape)}")
+ logger.info(
+ f"DECODE_SPLIT mode={mode} seconds={seconds} "
+ + " ".join(
+ f"{k}={profile.get(k, 0.0):.3f}"
+ for k in (
+ "device",
+ "readback",
+ "stitch",
+ "unpatchify",
+ "blend",
+ "concat",
+ "tiling",
+ "upload",
+ "host_prep",
+ "residual",
+ )
+ )
+ + f" waves={int(profile.get('waves', 0))} units={int(profile.get('units', 0))}"
+ + f" readback_gb={profile.get('readback_mb', 0.0) / 1000:.2f}"
+ )
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_fl2va_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_fl2va_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d76c103278f92f54550cf70d147d515ab5b2f2f
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_fl2va_minimax_h3.py
@@ -0,0 +1,308 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""MiniMax-H3 `fl2va` e2e: prompt plus first/last keyframes in, video and audio out.
+`test_fl2va_follows_the_keyframe` (fractal keyframe) is the discriminating conditioning gate.
+Separate file so it runs in its own process: DiT programs + CCL buffers at several padded lengths is a memory risk."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from loguru import logger
+from PIL import Image
+
+from ....pipelines.minimax_h3.packing import align_num_frames, prepare_keyframe_image
+from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+from ..wan2_2.common import check_output_sanity
+from .common import GALAXY_MESHES, create_fractal_image
+from .common_av import (
+ CALIBRATED_FOX_PROMPT,
+ artifact_dir,
+ check_audio_sanity,
+ check_av_sync,
+ check_spatial_seams,
+ check_written_file,
+ gate_clip,
+ log_spectral_flatness,
+ log_timing_table,
+ run_warm_generation,
+ to_uint8_frames,
+ weights_dir,
+ write_artifacts,
+)
+
+HEIGHT, WIDTH = 768, 1344
+NUM_FRAMES = 124
+NUM_INFERENCE_STEPS = 50
+SEED = 0
+
+PROMPT = CALIBRATED_FOX_PROMPT # the tier-6 bars are calibrated against this exact prompt
+
+# Ring collectives require FABRIC_1D_RING.
+MESHES = GALAXY_MESHES
+
+ANCHOR_PCC_FLOOR = 0.95 # measured 0.9943-0.9971 across the three anchor cases
+
+CLIP_THRESHOLD = 33.0 # t2va's bar; measured to transfer (36.63-37.30 vs t2va's 37.37)
+
+EXPECTED_TOTAL_S = 400.0 # loose did-something-collapse bar, not a perf target
+
+
+def check_keyframe_anchor(frames, keyframe, *, index, stretch, width, height, pcc_floor=0.3):
+ """Decoded frame `index` must correlate with its keyframe; uses the pipeline's own
+ `prepare_keyframe_image` (stretch vs cover-crop) so the comparison canvas cannot drift."""
+ frame = frames[index]
+ if isinstance(frame, torch.Tensor):
+ frame = frame.cpu().numpy()
+ frame = np.asarray(frame).astype(np.float64)
+
+ prepared = prepare_keyframe_image(keyframe.convert("RGB"), height, width, stretch)
+ expected = np.asarray(prepared).astype(np.float64)
+ assert frame.shape == expected.shape, f"frame {index} shape {frame.shape} != keyframe {expected.shape}"
+
+ pcc = float(np.corrcoef(frame.ravel(), expected.ravel())[0, 1])
+ label = "first" if index == 0 else "last"
+ logger.info(f"fl2va {label}-keyframe anchor: decoded frame {index} vs keyframe PCC = {pcc:.4f}")
+ assert pcc > pcc_floor, (
+ f"decoded frame {index} barely correlates with the {label} keyframe (PCC={pcc:.3f}); "
+ "the fl2va conditioning path is likely broken"
+ )
+ return pcc
+
+
+def check_tile_boundary_gradient(frames, *, vertical_boundaries, horizontal_boundaries, max_ratio=3.0):
+ """One-pixel gradient at each tile boundary vs its neighbourhood -- catches seams the block-mean
+ check cannot see. Known-good measures 1.2-1.5x; the built-in control columns must sit near 1.0."""
+ frames = np.asarray(frames)
+ luma = frames.astype(np.float64).mean(-1) if frames.ndim == 4 else frames.astype(np.float64)
+ gx = np.abs(np.diff(luma, axis=2)).mean(axis=(0, 1))
+ gy = np.abs(np.diff(luma, axis=1)).mean(axis=(0, 2))
+
+ def ratio(profile, index):
+ near = np.concatenate([profile[index - 12 : index - 3], profile[index + 3 : index + 12]])
+ return float(profile[index - 1] / max(float(np.median(near)), 1e-9))
+
+ results = {}
+ for name, profile, boundaries in (("vertical", gx, vertical_boundaries), ("horizontal", gy, horizontal_boundaries)):
+ ratios = {int(b): ratio(profile, int(b)) for b in boundaries if 12 < int(b) < len(profile) - 12}
+ results[name] = ratios
+ if ratios:
+ logger.info(
+ f"{name} tile-boundary gradient ratios (1.0 = no seam): "
+ + ", ".join(f"x={b}:{r:.3f}" if name == "vertical" else f"y={b}:{r:.3f}" for b, r in ratios.items())
+ )
+
+ generator = np.random.default_rng(0)
+ candidates = generator.integers(30, len(gx) - 30, 24)
+ control = [c for c in candidates if all(abs(int(c) - int(b)) > 16 for b in vertical_boundaries)][:12]
+ control_ratios = [ratio(gx, int(c)) for c in control]
+ mean_control = float(np.mean(control_ratios))
+ logger.info(f"control non-boundary columns: mean ratio {mean_control:.3f}, max {max(control_ratios):.3f}")
+ assert mean_control < 1.15, (
+ f"control columns average {mean_control:.3f}; this statistic is tracking image structure rather "
+ "than tile boundaries, so its boundary numbers mean nothing"
+ )
+
+ worst = max((r, f"{n} {b}") for n, rs in results.items() for b, r in rs.items())
+ assert worst[0] < max_ratio, (
+ f"tile-boundary gradient at {worst[1]} is {worst[0]:.2f}x its neighbourhood (control "
+ f"{mean_control:.2f}); a visible seam. See the artifact rubric"
+ )
+ results["control"] = mean_control
+ return results
+
+
+def _gated_keyframe() -> Image.Image:
+ """Frame 0 of the calibrated t2va generation -- the only keyframe the tier-6 bars apply to.
+ As a conditioning signal it is confounded; `test_fl2va_follows_the_keyframe` rules that out."""
+ artifact = Path.home() / "h3_t2va_artifacts" / "t2va.mp4"
+ if not artifact.is_file():
+ pytest.skip(
+ f"no calibrated t2va artifact at {artifact}; run test_pipeline_minimax_h3.py first so the "
+ "fl2va keyframe comes from content the tier-6 thresholds were calibrated on"
+ )
+ frame = _first_frame(artifact)
+ assert frame.size == (WIDTH, HEIGHT), f"t2va artifact is {frame.size}, expected {(WIDTH, HEIGHT)}"
+ return frame
+
+
+def _first_frame(path: Path) -> Image.Image:
+ import imageio.v3 as iio
+
+ return Image.fromarray(np.asarray(iio.imread(path, index=0, plugin="pyav"))).convert("RGB")
+
+
+@pytest.mark.timeout(10800)
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESHES, indirect=["mesh_device", "device_params"])
+def test_fl2va_end_to_end(mesh_device, reset_seeds):
+ """The `first_and_last` case: both preparation paths and a two-run vision scatter; perf + quality."""
+ case = "first_and_last"
+ keyframe = _gated_keyframe()
+ image = keyframe
+ last_image = keyframe
+
+ pipeline = MiniMaxH3Pipeline.create_pipeline(mesh_device=mesh_device, weights_dir=weights_dir())
+
+ # Warmup must be fl2va-shaped (keyframes included): programs are keyed on padded length; the helper asserts it.
+ output = run_warm_generation(
+ pipeline,
+ PROMPT,
+ image=image,
+ last_image=last_image,
+ num_frames=NUM_FRAMES,
+ height=HEIGHT,
+ width=WIDTH,
+ num_inference_steps=NUM_INFERENCE_STEPS,
+ seed=SEED,
+ )
+
+ expected_frames = align_num_frames(NUM_FRAMES)
+ logger.info(
+ f"fl2va[{case}] padded_len={pipeline.last_padded_len} "
+ f"video={tuple(output.video.shape)} audio={tuple(output.audio.shape)}"
+ )
+
+ num_forwards = NUM_INFERENCE_STEPS - 1
+ video_seconds = expected_frames / output.fps
+ log_timing_table(
+ pipeline,
+ "fl2va",
+ num_forwards=num_forwards,
+ video_seconds=video_seconds,
+ expected_total_s=EXPECTED_TOTAL_S,
+ extra=(
+ f" | {WIDTH}x{HEIGHT}, {expected_frames} frames @ {output.fps} fps "
+ f"({video_seconds:.2f} s), {num_forwards} forwards, first+last anchors, "
+ f"padded_len {pipeline.last_padded_len}"
+ ),
+ )
+ assert output.video.shape[2] == expected_frames, f"{output.video.shape[2]} frames, expected {expected_frames}"
+
+ frames = to_uint8_frames(output)
+
+ check_output_sanity(frames, num_frames=expected_frames, height=HEIGHT, width=WIDTH)
+ check_audio_sanity(output.audio, sampling_rate=output.sampling_rate, expected_seconds=expected_frames / output.fps)
+ check_av_sync(frames, output.audio, sampling_rate=output.sampling_rate, fps=output.fps)
+ log_spectral_flatness(output.audio, sampling_rate=output.sampling_rate)
+
+ # Boundaries from the VAE's own tile grid (vertical seams come from the x starts).
+ ratio = pipeline.vae_config.spatial_compression_ratio
+ (y_starts, _, _), (x_starts, _, _) = pipeline.vae.decode_tile_grid(HEIGHT // ratio, WIDTH // ratio)
+ check_spatial_seams(frames, vertical_boundaries=x_starts[1:], horizontal_boundaries=y_starts[1:])
+ check_tile_boundary_gradient(frames, vertical_boundaries=x_starts[1:], horizontal_boundaries=y_starts[1:])
+
+ # `stretch` follows the pipeline's rule: the FIRST keyframe given is the geometry anchor.
+ if image is not None:
+ check_keyframe_anchor(
+ frames, image, index=0, stretch=True, width=WIDTH, height=HEIGHT, pcc_floor=ANCHOR_PCC_FLOOR
+ )
+ if last_image is not None:
+ check_keyframe_anchor(
+ frames,
+ last_image,
+ index=-1,
+ stretch=image is None,
+ width=WIDTH,
+ height=HEIGHT,
+ pcc_floor=ANCHOR_PCC_FLOOR,
+ )
+
+ artifacts = artifact_dir("h3_t2va_artifacts")
+ paths = write_artifacts(frames, output.audio.cpu().numpy(), output.sampling_rate, artifacts, stem=f"fl2va_{case}")
+ check_written_file(paths, expected_frames)
+ if "mp4" in paths:
+ for index in (0, 17, 62, expected_frames - 1):
+ Image.fromarray(frames[index]).save(artifacts / f"fl2va_{case}_frame_{index}.png")
+
+ gate_clip(frames, PROMPT, CLIP_THRESHOLD, f"fl2va[{case}]")
+
+ logger.info(
+ "REMINDER: read the artifact rubric against these frames -- seams and flicker are what every "
+ "whole-tensor metric averages away, and both are parallelism bugs"
+ )
+
+
+def test_lone_last_keyframe_is_stretched_not_cover_cropped():
+ """The pipeline keys `stretch` on list position, so a lone `last_image` (index 0) is stretched."""
+ source = np.zeros((512, 512, 3), dtype=np.uint8)
+ source[:64, :, 0] = 255
+ source[-64:, :, 2] = 255
+ source[:, :, 1] = 32 # non-zero interior so the crop is not comparing zeros to zeros
+ keyframe = Image.fromarray(source)
+
+ # The pipeline's own rule, replicated for the lone-last request (image=None).
+ image, last_image = None, keyframe
+ sources = [k for k in (image, last_image) if k is not None]
+ prepared = [prepare_keyframe_image(k, HEIGHT, WIDTH, stretch=(i == 0)) for i, k in enumerate(sources)]
+ assert len(prepared) == 1
+
+ stretched = prepare_keyframe_image(keyframe, HEIGHT, WIDTH, True)
+ cropped = prepare_keyframe_image(keyframe, HEIGHT, WIDTH, False)
+ got = np.asarray(prepared[0])
+
+ assert prepared[0].size == stretched.size == cropped.size == (WIDTH, HEIGHT)
+ assert not np.array_equal(np.asarray(stretched), np.asarray(cropped))
+
+ assert np.array_equal(got, np.asarray(stretched)), "a lone last_image must be stretched (it is the geometry anchor)"
+
+ assert np.asarray(stretched)[0, :, 0].mean() > 200, "stretch must keep the source's top edge"
+ assert np.asarray(stretched)[-1, :, 2].mean() > 200, "stretch must keep the source's bottom edge"
+ assert np.asarray(cropped)[0, :, 0].mean() < 50, "cover-crop of a 1:1 source must discard the top band"
+ assert np.asarray(cropped)[-1, :, 2].mean() < 50, "cover-crop of a 1:1 source must discard the bottom band"
+
+
+@pytest.mark.timeout(7200)
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESHES, indirect=["mesh_device", "device_params"])
+def test_fl2va_follows_the_keyframe(mesh_device, reset_seeds):
+ """Discriminating gate: a fractal keyframe the model would never produce must drive frame 0."""
+ fractal = create_fractal_image(WIDTH, HEIGHT)
+ pipeline = MiniMaxH3Pipeline.create_pipeline(mesh_device=mesh_device, weights_dir=weights_dir())
+ output = pipeline(
+ PROMPT,
+ image=fractal,
+ num_frames=NUM_FRAMES,
+ height=HEIGHT,
+ width=WIDTH,
+ num_inference_steps=NUM_INFERENCE_STEPS,
+ seed=SEED,
+ )
+ frames = to_uint8_frames(output)
+
+ def pcc(a, b):
+ a = np.asarray(a, dtype=np.float64).ravel()
+ b = np.asarray(b, dtype=np.float64).ravel()
+ return float(np.corrcoef(a, b)[0, 1])
+
+ prepared = np.asarray(prepare_keyframe_image(fractal, HEIGHT, WIDTH, True))
+ to_keyframe = pcc(frames[0], prepared)
+ to_t2va = pcc(frames[0], np.asarray(_gated_keyframe()))
+ tail = pcc(frames[-1], prepared)
+
+ logger.info(
+ f"fl2va keyframe-drives-generation: frame 0 vs fractal keyframe {to_keyframe:.4f}, "
+ f"frame 0 vs t2va's own frame 0 {to_t2va:.4f}, frame -1 vs fractal keyframe {tail:.4f}"
+ )
+ artifacts = artifact_dir("h3_t2va_artifacts")
+ for index in (0, 17, 62, align_num_frames(NUM_FRAMES) - 1):
+ Image.fromarray(frames[index]).save(artifacts / f"fl2va_fractal_frame_{index}.png")
+
+ # (1) frame 0 follows the supplied keyframe.
+ assert to_keyframe > ANCHOR_PCC_FLOOR, (
+ f"decoded frame 0 correlates only {to_keyframe:.3f} with the fractal keyframe it was "
+ "conditioned on; the keyframe is not reaching the DiT"
+ )
+ # (2) the margin rules out the t2va confound; a no-op pipeline inverts this comparison.
+ assert to_keyframe > to_t2va + 0.30, (
+ f"frame 0 resembles the fractal keyframe ({to_keyframe:.3f}) barely more than it resembles "
+ f"t2va's own frame 0 ({to_t2va:.3f}); conditioning may be a no-op"
+ )
+ # (3) it is still a video, not 124 copies of the keyframe.
+ assert tail < to_keyframe - 0.20, (
+ f"the last frame still correlates {tail:.3f} with the keyframe against frame 0's "
+ f"{to_keyframe:.3f}; the clip may be pinned throughout rather than anchored at one end"
+ )
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..170ab1802a5eba820d05af931bd16b7e2d5616af
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_minimax_h3.py
@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""End-to-end `t2va` perf + quality gate, swept over the published working points.
+
+Six aspect ratios (21:9 .. 9:16) x three durations (5 / 10 / 15 s), 50 steps. The canvas comes from
+`resolve_canvas_size` -- short edge 768 from 16:9 through 9:16, ~1 MPix for wider -- and the frame
+count from `align_num_frames`, so neither is tabulated here. Each case writes its own artifact stem,
+so a sweep does not overwrite itself.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+from loguru import logger
+
+import ttnn
+
+from ....pipelines.minimax_h3.packing import MINIMAX_H3_FPS, align_num_frames, resolve_canvas_size
+from ....pipelines.minimax_h3.pipeline_minimax_h3 import (
+ MiniMaxH3Pipeline,
+ _requested_audio_t_factor,
+ _resolve_audio_t_shard,
+)
+from ..wan2_2.common import check_output_sanity
+from .common import GALAXY_MESHES
+from .common_av import (
+ CALIBRATED_FOX_PROMPT,
+ artifact_dir,
+ check_audio_sanity,
+ check_av_sync,
+ check_spatial_seams,
+ check_written_file,
+ gate_clip,
+ gate_vbench,
+ log_spectral_flatness,
+ log_timing_table,
+ run_warm_generation,
+ to_uint8_frames,
+ weights_dir,
+ write_artifacts,
+)
+
+NUM_INFERENCE_STEPS = 50
+SEED = 0
+
+# The published working points: short edge 768 from 16:9 through 9:16, and ~1 MPix for anything
+# wider. `resolve_canvas_size` already implements exactly that (768 short edge, area capped at
+# 768*1344, both axes snapped to 32), so the canvas is derived rather than tabulated here --
+# 21:9 lands on 672x1536, which is the documented example.
+ASPECT_RATIOS = [(21, 9), (16, 9), (4, 3), (1, 1), (3, 4), (9, 16)]
+DURATIONS_S = [5, 10, 15]
+
+# Loose did-something-collapse bar, per second of video rather than absolute: the 400 s figure was
+# set for one 5.17 s clip, and 10 s / 15 s at 1 MPix cost proportionally more. 77 s of compute per
+# video second is the same allowance (400 / 5.167), i.e. still ~5x looser than the 13.4x realtime
+# factor measured at 16:9 / 5 s.
+EXPECTED_S_PER_VIDEO_SECOND = 77.0
+
+
+# calibrated 2026-08-04, fox prompt, seed 0 (single sample; margins are generous)
+VBENCH_THRESHOLDS = {
+ "subject_consistency": 0.95,
+ "background_consistency": 0.95,
+ "motion_smoothness": 0.97,
+ "dynamic_degree": 1.0,
+ "imaging_quality": 0.64,
+}
+CLIP_THRESHOLD = 33.0 # measured mean 37.05 (2026-08-04, fox prompt, seed 0)
+
+# tier-6 thresholds are calibrated against this exact prompt; swapping it invalidates both bars.
+PROMPT = CALIBRATED_FOX_PROMPT
+
+SWEEP = [
+ pytest.param(ratio, seconds, id=f"{ratio[0]}x{ratio[1]}_{seconds}s")
+ for seconds in DURATIONS_S
+ for ratio in ASPECT_RATIOS
+]
+
+
+# Pure (no-device) coverage for the audio T-shard resolver: axis selection, tie-break, 8->4->1 fallback.
+@pytest.mark.parametrize(
+ ("requested", "mesh_shape", "tp_axis", "sp_axis", "expected"),
+ [
+ (8, (4, 8), 0, 1, (8, 1)), # factor 8 -> size-8 SP axis
+ (8, (8, 4), 1, 0, (8, 0)), # reversed 8x4 -> size-8 axis 0
+ (4, (4, 8), 0, 1, (4, 0)), # factor 4 -> size-4 TP axis
+ (8, (4, 32), 0, 1, (4, 0)), # no size-8 axis -> fall back to 4
+ (8, (1, 1), 0, 1, (1, None)), # single device -> unsharded
+ (4, (4, 4), 0, 1, (4, 0)), # equal axes -> prefer TP
+ (1, (4, 8), 0, 1, (1, None)), # explicit unsharded
+ (2, (4, 8), 0, 1, (1, None)), # below the 4/8 chain -> unsharded
+ (32, (4, 32), 0, 1, (32, 1)), # opt-in 32 on a quad -> the 32-wide inter-host axis
+ (32, (4, 8), 0, 1, (8, 1)), # requesting 32 on a single galaxy -> falls to 8
+ ],
+)
+def test_resolve_audio_t_shard(requested, mesh_shape, tp_axis, sp_axis, expected):
+ assert _resolve_audio_t_shard(requested, mesh_shape, tp_axis, sp_axis) == expected
+
+
+def test_requested_audio_t_factor_precedence(monkeypatch, expect_error):
+ # explicit kwarg wins over the env var
+ monkeypatch.setenv("MINIMAX_H3_AUDIO_T_FACTOR", "32")
+ assert _requested_audio_t_factor(4) == (4, False)
+ # env var fills the default when the kwarg is None
+ assert _requested_audio_t_factor(None) == (32, True)
+ # no kwarg, no env -> default 8
+ monkeypatch.delenv("MINIMAX_H3_AUDIO_T_FACTOR", raising=False)
+ assert _requested_audio_t_factor(None) == (8, False)
+ # a non-integer env value is a clear error, not a silent fallback
+ monkeypatch.setenv("MINIMAX_H3_AUDIO_T_FACTOR", "thirty-two")
+ with expect_error(ValueError, "MINIMAX_H3_AUDIO_T_FACTOR"):
+ _requested_audio_t_factor(None)
+
+
+@pytest.mark.timeout(7200)
+@pytest.mark.parametrize(("aspect_ratio", "duration_s"), SWEEP)
+@pytest.mark.parametrize(("mesh_device", "device_params"), GALAXY_MESHES, indirect=["mesh_device", "device_params"])
+def test_t2va_end_to_end(mesh_device, reset_seeds, aspect_ratio, duration_s):
+ weights = weights_dir("transformer", "text_encoder", "vae", "audio_vae")
+ artifacts = artifact_dir("h3_t2va_artifacts")
+ prompt = PROMPT
+
+ HEIGHT, WIDTH = resolve_canvas_size(*aspect_ratio)
+ NUM_FRAMES = align_num_frames(round(duration_s * MINIMAX_H3_FPS))
+ # One artifact per working point, so a sweep does not overwrite itself.
+ stem = f"t2va_{aspect_ratio[0]}x{aspect_ratio[1]}_{WIDTH}x{HEIGHT}_{duration_s}s"
+ logger.info(f"working point: {aspect_ratio[0]}:{aspect_ratio[1]} -> {WIDTH}x{HEIGHT}, {NUM_FRAMES} frames, {stem}")
+
+ # A missing dependency must report SKIPPED before the long generation, never silently pass as green.
+ pytest.importorskip("open_clip", reason="the CLIP gate needs open_clip, which is not installed")
+
+ if not os.environ.get("TT_DIT_CACHE_DIR"):
+ logger.warning(
+ "TT_DIT_CACHE_DIR is unset, so every weight load reads safetensors. Prepares are excluded "
+ "from the total either way, but the run will take far longer than the reported compute."
+ )
+
+ pipeline = MiniMaxH3Pipeline.create_pipeline(mesh_device=mesh_device, weights_dir=weights)
+
+ output = run_warm_generation(
+ pipeline,
+ prompt,
+ num_frames=NUM_FRAMES,
+ height=HEIGHT,
+ width=WIDTH,
+ num_inference_steps=NUM_INFERENCE_STEPS,
+ seed=SEED,
+ )
+
+ expected_frames = align_num_frames(NUM_FRAMES)
+ logger.info(
+ f"generated {output.num_frames} frames ({output.video_seconds:.3f} s) and "
+ f"{output.audio_seconds:.3f} s of audio at {output.sampling_rate} Hz"
+ )
+
+ num_forwards = NUM_INFERENCE_STEPS - 1
+ video_seconds = expected_frames / MINIMAX_H3_FPS
+ log_timing_table(
+ pipeline,
+ "t2va",
+ num_forwards=num_forwards,
+ video_seconds=video_seconds,
+ expected_total_s=EXPECTED_S_PER_VIDEO_SECOND * video_seconds,
+ extra=(
+ f" | {aspect_ratio[0]}:{aspect_ratio[1]} {WIDTH}x{HEIGHT}, {expected_frames} frames "
+ f"@ {MINIMAX_H3_FPS} fps ({video_seconds:.2f} s), {num_forwards} forwards, "
+ f"padded_len {pipeline.last_padded_len}"
+ ),
+ )
+
+ frames = to_uint8_frames(output)
+
+ check_output_sanity(frames, num_frames=expected_frames, height=HEIGHT, width=WIDTH)
+ check_audio_sanity(
+ output.audio,
+ sampling_rate=output.sampling_rate,
+ expected_seconds=expected_frames / MINIMAX_H3_FPS,
+ )
+ sync = check_av_sync(
+ frames,
+ output.audio,
+ sampling_rate=output.sampling_rate,
+ fps=MINIMAX_H3_FPS,
+ )
+ log_spectral_flatness(output.audio, sampling_rate=output.sampling_rate)
+
+ # Boundaries from the VAE's own tile grid; re-deriving them with a wrong overlap gates non-boundaries.
+ ratio = pipeline.vae_config.spatial_compression_ratio
+ (y_starts, _, _), (x_starts, _, _) = pipeline.vae.decode_tile_grid(HEIGHT // ratio, WIDTH // ratio)
+ check_spatial_seams(frames, vertical_boundaries=x_starts[1:], horizontal_boundaries=y_starts[1:])
+
+ # Rank 0 alone writes and scores, and its outcome is held rather than raised: every other rank
+ # goes straight to the barrier below, so a gate failure, a VBench `pytest.skip` or an artifact I/O
+ # error raised here would leave them blocked and skip the trace release. `pytest.skip.Exception`
+ # is named alongside `Exception` because `Skipped` does not derive from it.
+ outcome: BaseException | None = None
+ is_distributed = ttnn.using_distributed_env()
+ if not is_distributed or int(ttnn.distributed_context_get_rank()) == 0:
+ try:
+ paths = write_artifacts(frames, output.audio.cpu().numpy(), output.sampling_rate, artifacts, stem=stem)
+ check_written_file(paths, expected_frames, height=HEIGHT, width=WIDTH)
+
+ # CLIP_THRESHOLD was measured at 16:9 / 5 s. Applying it across the sweep is an extrapolation,
+ # but a generous one: the calibrated point measures ~37 against a bar of 33, and the score is a
+ # prompt-alignment number rather than a resolution-dependent one. A prompt change would
+ # invalidate it outright -- recalibrate before swapping PROMPT.
+ gate_clip(frames, prompt, CLIP_THRESHOLD, stem)
+ # RUN_VBENCH=0 drops the VBench gate. It needs its own interpreter (~/vbench_env, pinned to
+ # numpy<2 / transformers 4.33), and without one `run_vbench` skips -- which marks the whole
+ # test SKIPPED *after* the full generation, hiding the perf and A/V results behind a
+ # non-result. Off means "not measured": everything above still gates.
+ if os.environ.get("RUN_VBENCH", "1") not in ("0", "false", "False"):
+ gate_vbench(paths, prompt, VBENCH_THRESHOLDS, stem)
+ else:
+ logger.warning("RUN_VBENCH=0, so the VBench gate did not run; generative quality is UNMEASURED")
+
+ logger.info(f"artifacts in {artifacts}: {sorted(p.name for p in artifacts.iterdir())}")
+ except (Exception, pytest.skip.Exception) as exc:
+ outcome = exc
+ if is_distributed:
+ ttnn.distributed_context_barrier()
+ logger.info(
+ "REMINDER: read the artifact rubric against these frames -- the seam and flicker scores above "
+ "are statistics, and only looking at the output catches what they average away"
+ )
+ # Every rank, and before the fixture closes the mesh: a captured trace holds device buffers
+ # for the request. Only the 4x32 preset traces, so this is a no-op at 4x8.
+ pipeline.release_traces()
+ # Now that every rank has resynchronised and released its trace, rank 0 reports. mpirun surfaces
+ # a single failing rank as a failing job, so the run fails rather than hanging.
+ if outcome is not None:
+ raise outcome
+ assert sync["video_seconds"] > 0
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_ref2va_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_ref2va_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..76ae9238e92b699b1520a38c5f0966c2ea352d6f
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_pipeline_ref2va_minimax_h3.py
@@ -0,0 +1,274 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""MiniMax-H3 ``ref2va`` e2e on the ``mixed`` case (largest shape, all three reference modalities).
+``test_ref2va_conditioning_is_not_a_no_op`` is the discriminating gate, judged against a measured run-to-run floor.
+Separate file so it runs in its own process: ref2va padded lengths are 1.2-3.0x t2va's, a memory risk."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from loguru import logger
+from PIL import Image
+
+from ....pipelines.minimax_h3.packing_ref2va import MiniMaxH3Reference, reference_from_video_file
+from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+from .common import GALAXY_MESHES, create_fractal_image
+from .common_av import (
+ artifact_dir,
+ check_audio_sanity,
+ check_av_sync,
+ check_spatial_seams,
+ gate_clip,
+ gate_vbench,
+ log_timing_table,
+ run_warm_generation,
+ to_uint8_frames,
+ weights_dir,
+ write_artifacts,
+)
+
+WIDTH, HEIGHT, NUM_FRAMES, STEPS, SEED = 1344, 768, 124, 50, 0
+FPS = 24
+
+PROMPT = "a slow push-in through a quiet room as afternoon light moves across the floor"
+
+# Bars set below the minimum measured across the three ref2va cases; t2va's bars do NOT transfer.
+REF2VA_CLIP_THRESHOLD = 25.0 # min measured 29.05
+REF2VA_VBENCH_THRESHOLDS = {
+ "subject_consistency": 0.90,
+ "background_consistency": 0.89,
+ "motion_smoothness": 0.97, # t2va's, and ref2va measures *better* than t2va on it
+ "dynamic_degree": 1.0, # binary over one video: 1.0 or the clip is frozen
+ "imaging_quality": 0.44,
+}
+
+# 16384, not the other gates' 65536: the video VAE's taps=3 encoder (only ref2va reaches it) clashes with L1 above it.
+_L1_SMALL = 16384
+
+# Ring collectives require FABRIC_1D_RING; only `l1_small_size` differs from the other gates.
+# The shared shapes, but this suite runs a smaller L1_SMALL pool than they carry and reports it in
+# the measurement line, so the size is overridden rather than inherited.
+MESHES = [
+ pytest.param(shape, {**params, "l1_small_size": _L1_SMALL}, id=param.id)
+ for param in GALAXY_MESHES
+ for shape, params in [param.values]
+]
+
+
+REFERENCE_MEDIA = Path.home() / "h3_fl2va_artifacts" / "fl2va_first.mp4"
+
+
+def reference_video() -> Path:
+ if not REFERENCE_MEDIA.is_file():
+ pytest.skip(f"no reference video at {REFERENCE_MEDIA}; place a clip with a soundtrack there")
+ return REFERENCE_MEDIA
+
+
+def ref2va_references(case: str) -> list[MiniMaxH3Reference]:
+ """The reference set per e2e case (a fractal: content nothing in the prompt could produce)."""
+ if case == "one_image":
+ return [MiniMaxH3Reference(image=create_fractal_image(1024, 1024))]
+ if case == "video_with_sound":
+ return [reference_from_video_file(reference_video())]
+ if case == "mixed":
+ # image, then SILENT video, then bare audio: stresses split_condition_blocks' per-modality row cursors.
+ sounded = reference_from_video_file(reference_video())
+ return [
+ MiniMaxH3Reference(image=create_fractal_image(1024, 1024)),
+ reference_from_video_file(reference_video(), with_audio=False),
+ MiniMaxH3Reference(audio=sounded.audio, sample_rate=sounded.sample_rate),
+ ]
+ raise ValueError(case)
+
+
+def _real_frame_image() -> Image.Image:
+ from ....pipelines.minimax_h3.packing_ref2va import decode_reference_video
+
+ frames, _, _ = decode_reference_video(reference_video())
+ return Image.fromarray(frames[0])
+
+
+def _inverted(image: Image.Image) -> Image.Image:
+ """Colour-inverted copy: same size, so the packed layout and noise stream are unchanged."""
+ return Image.fromarray(255 - np.asarray(image.convert("RGB")))
+
+
+def _clip_resemblance(output, image: Image.Image, num_frames: int = 8) -> float:
+ from ...dataset_eval.clip_encoder import CLIPEncoder
+
+ encoder = CLIPEncoder()
+ frames = to_uint8_frames(output)
+ indices = np.linspace(0, frames.shape[0] - 1, num_frames).round().astype(int)
+
+ with torch.no_grad():
+ reference_features = encoder.model.encode_image(encoder.preprocess(image).unsqueeze(0)).float()
+ reference_features /= reference_features.norm(dim=-1, keepdim=True)
+ scores = []
+ for index in indices:
+ frame = encoder.preprocess(Image.fromarray(frames[index])).unsqueeze(0)
+ features = encoder.model.encode_image(frame).float()
+ features /= features.norm(dim=-1, keepdim=True)
+ scores.append(float((features @ reference_features.T).squeeze()))
+ return float(np.mean(scores))
+
+
+def _colour_distance(output, image: Image.Image) -> float:
+ """Euclidean distance between mean RGBs, over [0, 1]. Logged only, never asserted."""
+ output_mean = to_uint8_frames(output).reshape(-1, 3).mean(axis=0) / 255.0
+ reference_mean = np.asarray(image.convert("RGB")).reshape(-1, 3).mean(axis=0) / 255.0
+ return float(np.linalg.norm(output_mean - reference_mean))
+
+
+def _divergence(a: torch.Tensor, b: torch.Tensor) -> float:
+ return float((a.float() - b.float()).abs().mean())
+
+
+def _write(output, stem: str) -> dict:
+ directory = artifact_dir("h3_ref2va_artifacts")
+ frames = to_uint8_frames(output)
+ for index in (0, 17, NUM_FRAMES // 2, NUM_FRAMES - 1):
+ Image.fromarray(frames[index]).save(directory / f"{stem}_frame_{index}.png")
+ paths = write_artifacts(frames, output.audio.cpu().numpy(), output.sampling_rate, directory, stem=stem)
+ logger.info(f"{stem}: artifacts in {directory} ({sorted(paths)})")
+ return paths
+
+
+def _record_quality(frames: np.ndarray, paths: dict, case: str) -> None:
+ """CLIP prompt alignment and the five VBench dimensions, against the ref2va-calibrated bars."""
+ gate_clip(frames, PROMPT, REF2VA_CLIP_THRESHOLD, f"QUALITY ref2va[{case}]")
+ # A missing muxed mp4 skips rather than fails here, unlike t2va.
+ gate_vbench(paths, PROMPT, REF2VA_VBENCH_THRESHOLDS, f"QUALITY ref2va[{case}]", skip_without_mp4=True)
+
+
+def _pipeline(mesh_device) -> MiniMaxH3Pipeline:
+ """Pipeline bound to the `transformer_ref` partition (62 GB, fixed at construction)."""
+ return MiniMaxH3Pipeline.create_pipeline(
+ mesh_device=mesh_device, weights_dir=weights_dir("transformer_ref"), task="ref2va"
+ )
+
+
+CASE = "mixed"
+EXPECTED_PADDED_LEN = (
+ 89856 # measured e2e (a host estimate says 90112); asserted so the case cannot drift off the probed shape
+)
+
+
+@pytest.mark.timeout(10800)
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESHES, indirect=["mesh_device", "device_params"])
+def test_ref2va_end_to_end(mesh_device, reset_seeds):
+ """Full ``mixed`` generation: geometry agreement across every reference path, plus warm latency."""
+ case = CASE
+ references = ref2va_references(case)
+ pipeline = _pipeline(mesh_device)
+
+ # Warmup must use the SAME references: `padded_len` depends on them, and the helper asserts agreement.
+ output = run_warm_generation(
+ pipeline,
+ PROMPT,
+ references=references,
+ num_frames=NUM_FRAMES,
+ height=HEIGHT,
+ width=WIDTH,
+ num_inference_steps=STEPS,
+ seed=SEED,
+ )
+
+ assert output.video.shape == (1, 3, NUM_FRAMES, HEIGHT, WIDTH), tuple(output.video.shape)
+ assert output.video.min() >= 0.0 and output.video.max() <= 1.0, "decoded video must be in [0, 1]"
+ assert torch.isfinite(output.video).all() and torch.isfinite(output.audio).all()
+ assert (
+ pipeline.last_padded_len == EXPECTED_PADDED_LEN
+ ), f"{case} ran at padded_len {pipeline.last_padded_len}, not the probed {EXPECTED_PADDED_LEN}"
+
+ frames = to_uint8_frames(output)
+ # Artifacts before the checks, so a failing check still leaves frames to inspect.
+ paths = _write(output, f"ref2va_{case}")
+
+ num_forwards = STEPS - 1
+ log_timing_table(
+ pipeline,
+ f"ref2va[{case}]",
+ num_forwards=num_forwards,
+ video_seconds=NUM_FRAMES / FPS,
+ extra=(
+ f", l1_small_size {_L1_SMALL} | {WIDTH}x{HEIGHT}, {NUM_FRAMES} frames @ {FPS} fps "
+ f"({NUM_FRAMES / FPS:.2f} s), {num_forwards} forwards, padded_len {pipeline.last_padded_len}"
+ ),
+ )
+
+ check_audio_sanity(
+ output.audio, sampling_rate=output.sampling_rate, expected_seconds=NUM_FRAMES / FPS, tolerance_seconds=0.05
+ )
+ check_av_sync(frames, output.audio, sampling_rate=output.sampling_rate, fps=FPS)
+ # Horizontal bar is 3.0: `video_with_sound` measures 2.29x there from scene content, not a seam.
+ check_spatial_seams(frames, vertical_boundaries=(448, 896), horizontal_boundaries=(), max_ratio=2.0)
+ check_spatial_seams(frames, vertical_boundaries=(), horizontal_boundaries=(384,), max_ratio=3.0)
+
+ _record_quality(frames, paths, case)
+
+
+@pytest.mark.timeout(9000)
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESHES, indirect=["mesh_device", "device_params"])
+def test_ref2va_conditioning_is_not_a_no_op(mesh_device, reset_seeds):
+ """Discriminator: a same-size inverted reference keeps the noise bit-identical, so any output
+ divergence above the measured run-to-run floor is attributable to reference content alone."""
+ pipeline = _pipeline(mesh_device)
+ normal = _real_frame_image()
+ inverted = _inverted(normal)
+
+ def generate(image):
+ return pipeline(
+ PROMPT,
+ references=[MiniMaxH3Reference(image=image)],
+ num_frames=NUM_FRAMES,
+ height=HEIGHT,
+ width=WIDTH,
+ num_inference_steps=STEPS,
+ seed=SEED,
+ )
+
+ a = generate(normal)
+ a_repeat = generate(normal)
+ b = generate(inverted)
+
+ floor = _divergence(a.video, a_repeat.video)
+ signal = _divergence(a.video, b.video)
+ to_normal = (_clip_resemblance(a, normal), _clip_resemblance(b, normal))
+ to_inverted = (_clip_resemblance(a, inverted), _clip_resemblance(b, inverted))
+ colour = (_colour_distance(a, normal), _colour_distance(b, inverted))
+ crossed = (_colour_distance(a, inverted), _colour_distance(b, normal))
+
+ logger.info(
+ f"ref2va discriminator: run-to-run floor {floor:.6f}, reference-swap signal {signal:.6f} "
+ f"(ratio {signal / max(floor, 1e-9):.1f}x) | CLIP to normal A={to_normal[0]:.4f} "
+ f"B={to_normal[1]:.4f} | CLIP to inverted A={to_inverted[0]:.4f} B={to_inverted[1]:.4f} | "
+ f"colour distance to own reference A={colour[0]:.4f} B={colour[1]:.4f}, "
+ f"to the other A={crossed[0]:.4f} B={crossed[1]:.4f}"
+ )
+ for stem, output in (("discriminate_normal", a), ("discriminate_inverted", b)):
+ _write(output, stem)
+
+ assert floor < 0.02, f"the same request twice diverged by {floor:.6f}; the pipeline is not reproducible"
+
+ # The null hypothesis (reference ignored) scores signal == floor exactly, the noise being bit-identical.
+ assert signal > 10 * floor, (
+ f"swapping the reference moved the output by {signal:.6f} against a run-to-run floor of "
+ f"{floor:.6f}: the reference is not conditioning the output"
+ )
+ assert signal > 0.01, (
+ f"swapping the reference moved the output by only {signal:.6f} mean absolute pixel value; "
+ "the effect is present but too small to call conditioning"
+ )
+
+ # Direction is logged, not asserted: no known instrument here can fail on a correct pipeline only.
+ logger.info(
+ f"ref2va direction (recorded, not asserted): CLIP own-vs-other "
+ f"A {to_normal[0]:.4f} vs {to_inverted[0]:.4f}, B {to_inverted[1]:.4f} vs {to_normal[1]:.4f}; "
+ f"colour own-vs-other A {colour[0]:.4f} vs {crossed[0]:.4f}, B {colour[1]:.4f} vs {crossed[1]:.4f}"
+ )
diff --git a/code/models/tt_dit/tests/models/minimax_h3/test_references_minimax_h3.py b/code/models/tt_dit/tests/models/minimax_h3/test_references_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..19a218444123dcd10851a48fcf2eb66402c74b71
--- /dev/null
+++ b/code/models/tt_dit/tests/models/minimax_h3/test_references_minimax_h3.py
@@ -0,0 +1,809 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""MiniMax-H3 ``ref2va`` reference conditioning, gated against the reference implementation's
+own path (``diffusers.modular_pipelines.minimax_h3`` at PR #14355): host gates are bit-exact
+(drift is silent), the device VAE encode is PCC on real media. importorskips that branch."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from loguru import logger
+
+from ....pipelines.minimax_h3 import packing as p
+from ....pipelines.minimax_h3 import packing_ref2va as rp
+from ....pipelines.minimax_h3 import references as R
+from ....utils.check import assert_quality
+from ....utils.test import ring_params_req_exact_devices
+
+# Base `packing` too: the rotary-clock gate pins both `_temporal_position_span` implementations.
+reference_before_encoder, reference_packing, reference_encoders, reference_base_packing = (
+ pytest.importorskip(module, reason="requires the minimax-h3 diffusers branch")
+ for module in (
+ "diffusers.modular_pipelines.minimax_h3.before_encoder",
+ "diffusers.modular_pipelines.minimax_h3.packing_ref2va",
+ "diffusers.modular_pipelines.minimax_h3.encoders",
+ "diffusers.modular_pipelines.minimax_h3.packing",
+ )
+)
+
+AUDIO_RATE = 32000
+TARGET_FRAMES = 124
+DURATION = TARGET_FRAMES / p.MINIMAX_H3_FPS # 5.1667 s
+
+TARGET_HEIGHT, TARGET_WIDTH = 768, 1344
+VAE_RATIO = 16
+AUDIO_HOP = 800
+PATCH_SIZE = (1, 2, 2)
+
+
+def _components():
+ return SimpleNamespace(audio_sampling_rate=AUDIO_RATE)
+
+
+def _rng():
+ return np.random.default_rng(0)
+
+
+def _image(width: int, height: int) -> np.ndarray:
+ """Structured content, not flat: a resize of a constant field cannot show an error."""
+ rng = _rng()
+ return (rng.random((height, width, 3)) * 255).astype(np.uint8)
+
+
+def _video(width: int, height: int, num_frames: int) -> np.ndarray:
+ rng = _rng()
+ return (rng.random((num_frames, height, width, 3)) * 255).astype(np.uint8)
+
+
+def _waveform(seconds: float, sample_rate: int = AUDIO_RATE, channels: int = 2) -> torch.Tensor:
+ torch.manual_seed(0)
+ return torch.randn(channels, int(seconds * sample_rate)) * 0.1
+
+
+def _pair(specs: list[dict]):
+ ours = [rp.MiniMaxH3Reference(**spec) for spec in specs]
+ theirs = [reference_packing.MiniMaxH3Reference(**spec) for spec in specs]
+ return ours, theirs
+
+
+SPECS = {
+ "image": dict(image=_image(1024, 1024)),
+ "image_16to9": dict(image=_image(1920, 1080)),
+ "video_24fps_at_canvas": dict(video=_video(1344, 768, TARGET_FRAMES), fps=24.0),
+ "video_30fps_off_canvas": dict(video=_video(1920, 1080, 60), fps=30.0),
+ "video_with_sound": dict(video=_video(1344, 768, TARGET_FRAMES), fps=24.0, audio=_waveform(DURATION)),
+ "audio": dict(audio=_waveform(DURATION)),
+ "audio_44k_mono": dict(audio=_waveform(DURATION, 44100, 1), sample_rate=44100),
+}
+
+
+def test_reference_limits_match_reference_and_are_enforced(expect_error):
+ assert rp.MINIMAX_H3_MAX_REFERENCE_IMAGES == reference_packing.MINIMAX_H3_MAX_REFERENCE_IMAGES == 9
+ assert rp.MINIMAX_H3_MAX_REFERENCE_VIDEOS == reference_packing.MINIMAX_H3_MAX_REFERENCE_VIDEOS == 3
+ assert rp.MINIMAX_H3_MAX_REFERENCE_AUDIOS == reference_packing.MINIMAX_H3_MAX_REFERENCE_AUDIOS == 3
+ assert rp.MINIMAX_H3_MAX_REFERENCES == reference_packing.MINIMAX_H3_MAX_REFERENCES == 12
+ assert rp.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE == reference_packing.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE == 2048
+
+ image = rp.MiniMaxH3Reference(image=_image(256, 256))
+ video = rp.MiniMaxH3Reference(video=_video(256, 256, 10), fps=24.0)
+ audio = rp.MiniMaxH3Reference(audio=_waveform(DURATION))
+
+ assert R.check_references([image, video, audio]) == ["image", "video", "audio"]
+
+ with expect_error(ValueError, "at least one reference"):
+ R.check_references([])
+ with expect_error(ValueError, "at most 9 image"):
+ R.check_references([image] * 10)
+ with expect_error(ValueError, "at most 3 video"):
+ R.check_references([video] * 4)
+ with expect_error(ValueError, "at most 3 audio"):
+ R.check_references([image, audio, audio, audio, audio])
+ with expect_error(ValueError, "at most 12 references"):
+ R.check_references([image] * 9 + [video] * 3 + [audio])
+ with expect_error(ValueError, "paired with at least one image or video"):
+ R.check_references([audio])
+
+
+@pytest.mark.parametrize("name", list(SPECS), ids=list(SPECS))
+def test_prepare_references_matches_reference(name):
+ spec = SPECS[name]
+ # audio may not be alone; pair with an image
+ specs = [spec] if "audio" not in name else [SPECS["image"], spec]
+ ours, theirs = _pair(specs)
+
+ got, got_frames = R.prepare_references(ours, TARGET_FRAMES, AUDIO_RATE)
+ want, want_frames = reference_before_encoder.MiniMaxH3Ref2VASetupStep.prepare_references(
+ _components(), theirs, TARGET_FRAMES
+ )
+
+ assert got_frames == want_frames == TARGET_FRAMES
+ assert len(got) == len(want)
+ for index, (a, b) in enumerate(zip(got, want)):
+ assert a.kind == b.kind, index
+ assert a.has_audio == b.has_audio, index
+ if a.kind == "image":
+ assert np.array_equal(np.asarray(a.image), np.asarray(b.image)), f"prepared image {index} differs"
+ assert a.image.size == b.image.size
+ if a.kind == "video":
+ assert np.array_equal(a.frames, b.frames), f"prepared frames {index} differ"
+ if a.has_audio:
+ assert torch.equal(a.waveform, b.waveform), f"prepared waveform {index} differs"
+ assert a.waveform.shape[0] == p.MINIMAX_H3_AUDIO_CHANNELS
+
+
+def test_prepare_references_uses_each_references_own_resolution():
+ references = [
+ rp.MiniMaxH3Reference(image=_image(1024, 1024)),
+ rp.MiniMaxH3Reference(video=_video(1920, 1080, 30), fps=24.0),
+ ]
+ prepared, _ = R.prepare_references(references, TARGET_FRAMES, AUDIO_RATE)
+
+ assert prepared[0].image.size == (2048, 2048)
+ assert min(prepared[0].image.size) == rp.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE
+ assert prepared[1].frames.shape[1:3] == (768, 1344)
+ assert min(prepared[1].frames.shape[1:3]) == p.MINIMAX_H3_SHORT_EDGE
+ assert min(prepared[0].image.size) > min(prepared[1].frames.shape[1:3])
+
+
+def test_prepare_references_truncates_a_video_and_its_soundtrack_to_the_target():
+ long_frames = _video(1344, 768, TARGET_FRAMES + 60)
+ references = [rp.MiniMaxH3Reference(video=long_frames, fps=24.0, audio=_waveform(10.0))]
+ prepared, num_frames = R.prepare_references(references, TARGET_FRAMES, AUDIO_RATE)
+
+ assert num_frames == TARGET_FRAMES
+ assert prepared[0].frames.shape[0] == TARGET_FRAMES
+ assert prepared[0].waveform.shape[1] == int(DURATION * AUDIO_RATE)
+
+
+def test_num_frames_may_be_derived_from_a_single_audio_bearing_reference():
+ references = [rp.MiniMaxH3Reference(image=_image(512, 512)), rp.MiniMaxH3Reference(audio=_waveform(6.0))]
+ _, num_frames = R.prepare_references(references, None, AUDIO_RATE)
+ assert num_frames == p.align_num_frames(round(6.0 * p.MINIMAX_H3_FPS))
+ assert num_frames % p.MINIMAX_H3_FRAMES_PER_CHUNK == p.MINIMAX_H3_LATENTS_PER_CHUNK
+ theirs = [
+ reference_packing.MiniMaxH3Reference(image=_image(512, 512)),
+ reference_packing.MiniMaxH3Reference(audio=_waveform(6.0)),
+ ]
+ _, want = reference_before_encoder.MiniMaxH3Ref2VASetupStep.prepare_references(_components(), theirs, None)
+ assert num_frames == want
+
+
+def test_num_frames_is_ambiguous_with_two_soundtracks(expect_error):
+ references = [
+ rp.MiniMaxH3Reference(video=_video(768, 768, 24), fps=24.0, audio=_waveform(6.0)),
+ rp.MiniMaxH3Reference(audio=_waveform(7.0)),
+ ]
+ with expect_error(ValueError, "exactly one of them carries audio"):
+ R.prepare_references(references, None, AUDIO_RATE)
+
+
+@pytest.mark.parametrize("seconds", [4.0, 16.0])
+def test_a_derived_duration_outside_the_models_range_is_rejected(seconds, expect_error):
+ references = [rp.MiniMaxH3Reference(image=_image(512, 512)), rp.MiniMaxH3Reference(audio=_waveform(seconds))]
+ with expect_error(ValueError, "seconds"):
+ R.prepare_references(references, None, AUDIO_RATE)
+
+
+@pytest.mark.parametrize("seconds", [DURATION, 5.0, 8.0, 15.0])
+def test_pad_waveform_to_hop_reproduces_the_reference_latent_count(seconds):
+ waveform = _waveform(seconds)
+ padded = R.pad_waveform_to_hop(waveform)
+
+ assert padded.shape[-1] % R.MINIMAX_H3_AUDIO_HOP == 0
+ assert padded.shape[-1] - waveform.shape[-1] < R.MINIMAX_H3_AUDIO_HOP
+ assert torch.equal(padded[:, : waveform.shape[-1]], waveform), "padding must not disturb the samples"
+ assert (padded[:, waveform.shape[-1] :] == 0).all(), "the pad is zeros, as F.pad's default is"
+ assert padded.shape[-1] // R.MINIMAX_H3_AUDIO_HOP == int(np.ceil(waveform.shape[-1] / R.MINIMAX_H3_AUDIO_HOP))
+
+
+def test_production_soundtrack_is_not_a_whole_number_of_hops():
+ """The case that makes the padding load-bearing rather than defensive."""
+ samples = int(DURATION * AUDIO_RATE)
+ assert samples == 165333
+ assert samples % R.MINIMAX_H3_AUDIO_HOP != 0
+ assert R.pad_waveform_to_hop(_waveform(DURATION)).shape[-1] // R.MINIMAX_H3_AUDIO_HOP == 207
+
+
+def test_span_matches_both_reference_summation_orders():
+ for n in range(1, 61):
+ assert p._temporal_position_span(n) == reference_base_packing._temporal_position_span(
+ n
+ ), f"pairwise span disagrees with the reference at n={n}"
+ assert rp._temporal_position_span(n) == reference_packing._temporal_position_span(
+ n
+ ), f"sequential span disagrees with the reference at n={n}"
+
+
+def test_the_two_span_orders_actually_differ():
+ assert p._temporal_position_span(15) == rp._temporal_position_span(15), "expected agreement below 16"
+ assert p._temporal_position_span(16) != rp._temporal_position_span(16), "expected divergence from 16 latent frames"
+
+ production = p.video_latent_num_frames(TARGET_FRAMES)
+ assert production == 37
+ pairwise, sequential = p._temporal_position_span(production), rp._temporal_position_span(production)
+ assert pairwise != sequential, "the two orders must differ at the production frame count"
+ assert rp._temporal_position_span(16) > p._temporal_position_span(16)
+ assert sequential < pairwise
+
+
+def _target():
+ return (
+ p.video_latent_num_frames(TARGET_FRAMES),
+ TARGET_HEIGHT // VAE_RATIO,
+ TARGET_WIDTH // VAE_RATIO,
+ p.audio_latent_num_frames(TARGET_FRAMES),
+ )
+
+
+def _image_geometry(source_width: int, source_height: int):
+ height, width = rp.resolve_reference_image_size(source_width, source_height)
+ return dict(num_latent_frames=1, latent_height=height // VAE_RATIO, latent_width=width // VAE_RATIO)
+
+
+def _video_geometry(source_width: int, source_height: int, num_frames: int = TARGET_FRAMES):
+ frames = np.zeros((num_frames, source_height, source_width, 3), dtype=np.uint8)
+ frames = rp.prepare_reference_frames(frames, TARGET_FRAMES)
+ trimmed = rp.trim_reference_num_frames(frames.shape[0])
+ return dict(
+ num_latent_frames=p.video_latent_num_frames(trimmed),
+ latent_height=frames.shape[1] // VAE_RATIO,
+ latent_width=frames.shape[2] // VAE_RATIO,
+ )
+
+
+def _audio_latents(seconds: float) -> int:
+ return int(np.ceil(int(seconds * AUDIO_RATE) / AUDIO_HOP))
+
+
+def _spec_to_prepared(spec, dataclass):
+ kind = spec["kind"]
+ geometry = {}
+ if kind == "image":
+ geometry = _image_geometry(*spec["source"])
+ elif kind == "video":
+ geometry = _video_geometry(*spec["source"], spec.get("num_frames", TARGET_FRAMES))
+ audio_seconds = spec.get("audio_seconds")
+ return dataclass(
+ kind=kind,
+ has_audio=audio_seconds is not None,
+ num_audio_latents=0 if audio_seconds is None else _audio_latents(audio_seconds),
+ **geometry,
+ )
+
+
+def _both(spec_list):
+ ours = [_spec_to_prepared(spec, rp.MiniMaxH3PreparedReference) for spec in spec_list]
+ theirs = [_spec_to_prepared(spec, reference_packing.MiniMaxH3PreparedReference) for spec in spec_list]
+ return ours, theirs
+
+
+IMAGE_1TO1 = dict(kind="image", source=(1024, 1024))
+IMAGE_16TO9 = dict(kind="image", source=(1920, 1080))
+VIDEO_SOUND = dict(kind="video", source=(1344, 768), audio_seconds=DURATION)
+# 1:1 sounded video: its soundtrack rows must pin to its OWN width grid, not the target's.
+VIDEO_SOUND_1TO1 = dict(kind="video", source=(768, 768), audio_seconds=DURATION)
+VIDEO_SILENT = dict(kind="video", source=(1344, 768))
+AUDIO_ONLY = dict(kind="audio", audio_seconds=DURATION)
+
+# Awkward order on purpose: exercises the label counters, soundtrack-before-video order and the
+# shared rotary clock; the reversed copy is the host order-sensitivity gate (order is semantic).
+_NINE_MIXED = [
+ VIDEO_SOUND,
+ AUDIO_ONLY,
+ VIDEO_SILENT,
+ IMAGE_1TO1,
+ AUDIO_ONLY,
+ IMAGE_16TO9,
+ VIDEO_SOUND,
+ IMAGE_1TO1,
+ AUDIO_ONLY,
+]
+
+CASES = {
+ "one_image": [IMAGE_1TO1],
+ "one_audio": [IMAGE_1TO1, AUDIO_ONLY], # audio may never be the only reference
+ "video_with_sound": [VIDEO_SOUND],
+ "video_with_sound_1to1": [VIDEO_SOUND_1TO1],
+ "video_without_sound": [VIDEO_SILENT],
+ "nine_mixed_awkward": _NINE_MIXED,
+ "nine_mixed_reversed": list(reversed(_NINE_MIXED)),
+}
+
+TEXT_LEN = 1024
+# A reference's vision block lands inside the text rows tagged as VIDEO, not text.
+VISION_BLOCK = slice(30, 120)
+
+
+def _text_tags():
+ tags = torch.ones(TEXT_LEN, dtype=torch.long)
+ tags[VISION_BLOCK] = p.MINIMAX_H3_VIDEO_TAG
+ return tags
+
+
+@pytest.mark.parametrize("case", list(CASES), ids=list(CASES))
+def test_layout_matches_reference(case):
+ ours, theirs = _both(CASES[case])
+ num_latent_frames, latent_height, latent_width, num_audio_latents = _target()
+ args = (num_latent_frames, latent_height, latent_width, num_audio_latents, PATCH_SIZE)
+
+ got = rp.build_ref2va_packed_sequence(_text_tags(), ours, *args)
+ want = reference_packing.build_ref2va_packed_sequence(_text_tags(), theirs, *args)
+
+ assert got.sequence_length == want.sequence_length
+ assert got.position_ids.dtype == torch.float64
+ assert torch.equal(got.position_ids, want.position_ids), "the fp64 rotary grid is not bit-exact"
+ assert torch.equal(got.token_tags, want.token_tags)
+ assert torch.equal(got.video_indices, want.video_indices)
+ assert torch.equal(got.audio_indices, want.audio_indices)
+ assert torch.equal(got.text_indices, want.text_indices)
+ assert got.num_condition_video_rows == want.num_condition_video_rows
+ assert got.num_condition_audio_rows == want.num_condition_audio_rows
+
+
+# `(width, height)` in, `(height, width)` out; both orientations listed so a transpose cannot pass.
+@pytest.mark.parametrize(
+ "source,expected",
+ [
+ ((1024, 1024), (2048, 2048)),
+ ((1920, 1080), (2048, 3648)),
+ ((1080, 1920), (3648, 2048)),
+ ((4096, 1024), (2048, 8192)),
+ ],
+)
+def test_resolve_reference_image_size(source, expected):
+ assert rp.resolve_reference_image_size(*source) == expected
+ assert rp.resolve_reference_image_size(*source) == reference_packing.resolve_reference_image_size(*source)
+ height, width = expected
+ assert height % p.MINIMAX_H3_CANVAS_MULTIPLE == 0 and width % p.MINIMAX_H3_CANVAS_MULTIPLE == 0
+ assert min(height, width) == rp.MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE
+
+
+def test_resolve_reference_image_size_rejects_out_of_range(expect_error):
+ with expect_error(ValueError, "within 1:4 and 4:1"):
+ rp.resolve_reference_image_size(5000, 1000)
+ with expect_error(ValueError, "positive size"):
+ rp.resolve_reference_image_size(0, 100)
+
+
+@pytest.mark.parametrize("fps", [24.0, 30.0, 25.0, 12.0, 60.0, 23.976])
+def test_resample_reference_frames_matches_reference(fps):
+ frames = np.arange(50 * 2 * 2 * 3, dtype=np.uint8).reshape(50, 2, 2, 3)
+ got = rp.resample_reference_frames(frames, fps)
+ want = reference_packing.resample_reference_frames(frames, fps)
+ assert got.shape == want.shape
+ assert np.array_equal(got, want)
+ if fps == 24.0:
+ assert got is frames
+
+
+def test_prepare_reference_frames_uses_its_own_aspect_canvas():
+ frames = np.zeros((10, 480, 480, 3), dtype=np.uint8)
+ prepared = rp.prepare_reference_frames(frames, TARGET_FRAMES)
+ assert prepared.shape[1:3] == p.resolve_canvas_size(480, 480) == (768, 768)
+ assert np.array_equal(prepared, reference_packing.prepare_reference_frames(frames, TARGET_FRAMES))
+
+ # `shares_memory`, not `is`: the frame-count cap slices first, so a view comes back.
+ at_canvas = np.zeros((10, 768, 1344, 3), dtype=np.uint8)
+ passthrough = rp.prepare_reference_frames(at_canvas, TARGET_FRAMES)
+ assert np.shares_memory(passthrough, at_canvas)
+
+
+@pytest.mark.parametrize("num_frames", [1, 5, 12, 24, 25, 124, 192])
+def test_sample_reference_video_frames_matches_reference(num_frames):
+ frames = np.zeros((num_frames, 4, 4, 3), dtype=np.uint8)
+ got_frames, got_ts = rp.sample_reference_video_frames(frames)
+ want_frames, want_ts = reference_packing.sample_reference_video_frames(frames)
+ assert len(got_frames) == len(want_frames)
+ assert got_ts == want_ts
+ assert len(got_ts) == -(-len(got_frames) // rp.MINIMAX_H3_QWEN_TEMPORAL_PATCH)
+
+
+def test_block_timestamps_round_half_to_even():
+ frames = np.zeros((124, 4, 4, 3), dtype=np.uint8)
+ _, timestamps = rp.sample_reference_video_frames(frames)
+ assert timestamps[0] == 0.25
+ assert f"{timestamps[0]:.1f}" == "0.2"
+
+
+@pytest.mark.parametrize("num_frames", [1, 5, 22, 39, 124, 125, 130])
+def test_trim_reference_num_frames_matches_reference(num_frames):
+ got = rp.trim_reference_num_frames(num_frames)
+ assert got == reference_packing.trim_reference_num_frames(num_frames)
+ assert got % p.MINIMAX_H3_FRAMES_PER_CHUNK == p.MINIMAX_H3_LATENTS_PER_CHUNK
+
+
+def test_trim_reference_num_frames_has_a_22_frame_floor():
+ """1..23 frames all map to 22 -- not a plain round-down; a "fixed" floor of 5 silently changes the latent count."""
+ for num_frames in (1, 5, 10, 22, 23):
+ assert rp.trim_reference_num_frames(num_frames) == 22
+ assert rp.trim_reference_num_frames(num_frames) == reference_packing.trim_reference_num_frames(num_frames)
+ assert rp.trim_reference_num_frames(39) == 39
+ assert rp.trim_reference_num_frames(55) == 39
+ assert rp.trim_reference_num_frames(125) == 124
+
+
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("sample_rate", [32000, 44100])
+def test_prepare_reference_waveform_matches_reference(channels, sample_rate):
+ torch.manual_seed(0)
+ waveform = torch.randn(channels, int(6.0 * sample_rate))
+ got = rp.prepare_reference_waveform(waveform, sample_rate, AUDIO_RATE, DURATION)
+ want = reference_packing.prepare_reference_waveform(waveform, sample_rate, AUDIO_RATE, DURATION)
+ assert got.shape == want.shape
+ assert torch.equal(got, want)
+ assert got.shape[0] == p.MINIMAX_H3_AUDIO_CHANNELS
+ if sample_rate == AUDIO_RATE:
+ assert got.shape[1] == int(DURATION * sample_rate)
+
+
+def test_reference_media_to_uint8_matches_reference():
+ torch.manual_seed(0)
+ tensor = torch.rand(3, 8, 6)
+ array = (np.random.default_rng(0).random((8, 6, 3)) * 255).astype(np.uint8)
+ for media in (tensor, array, [tensor, tensor]):
+ got = rp.reference_media_to_uint8(media)
+ want = reference_packing.reference_media_to_uint8(media)
+ assert got.dtype == np.uint8
+ assert np.array_equal(got, want)
+
+
+def test_reference_requires_exactly_one_medium(expect_error):
+ """A video may carry audio; nothing else may carry two media."""
+ image = np.zeros((4, 4, 3), dtype=np.uint8)
+ video = np.zeros((2, 4, 4, 3), dtype=np.uint8)
+ waveform = torch.zeros(2, 800)
+
+ assert rp.MiniMaxH3Reference(image=image).kind == "image"
+ assert rp.MiniMaxH3Reference(video=video).kind == "video"
+ assert rp.MiniMaxH3Reference(audio=waveform).kind == "audio"
+ sounded = rp.MiniMaxH3Reference(video=video, audio=waveform)
+ assert sounded.kind == "video" and sounded.has_audio
+
+ with expect_error(ValueError, "exactly one"):
+ rp.MiniMaxH3Reference(image=image, video=video)
+ with expect_error(ValueError, "exactly one"):
+ rp.MiniMaxH3Reference(image=image, audio=waveform)
+ with expect_error(ValueError, "exactly one"):
+ rp.MiniMaxH3Reference()
+
+
+def test_reference_refuses_a_path(expect_error):
+ with expect_error(ValueError, "never opens media files"):
+ rp.MiniMaxH3Reference(image="subject.png")
+
+
+def test_reference_defaults_fps_to_the_models_own():
+ video = np.zeros((2, 4, 4, 3), dtype=np.uint8)
+ assert rp.MiniMaxH3Reference(video=video).fps == float(p.MINIMAX_H3_FPS)
+
+
+def _prepared(kind, *, frames=1, height=48, width=84, audio_latents=0):
+ return rp.MiniMaxH3PreparedReference(
+ kind=kind,
+ has_audio=audio_latents > 0,
+ num_latent_frames=frames,
+ latent_height=height if kind != "audio" else 0,
+ latent_width=width if kind != "audio" else 0,
+ num_audio_latents=audio_latents,
+ )
+
+
+def test_split_condition_blocks_is_in_packed_order():
+ """Per-modality order would keep every row count right and desynchronize the whole region."""
+ references = [
+ _prepared("video", frames=37, audio_latents=207),
+ _prepared("audio", audio_latents=100),
+ _prepared("image"),
+ ]
+ video_rows = torch.arange(sum(r.num_video_rows for r in references if r.kind != "audio")).float()[:, None]
+ audio_rows = torch.arange(sum(r.num_audio_rows for r in references)).float()[:, None] + 10000
+
+ blocks = R.split_condition_blocks(references, video_rows, audio_rows)
+ assert [kind for _, kind in blocks] == ["audio", "video", "audio", "video"]
+ assert [block.shape[0] for block, _ in blocks] == [
+ references[0].num_audio_rows,
+ references[0].num_video_rows,
+ references[1].num_audio_rows,
+ references[2].num_video_rows,
+ ]
+
+ assert torch.equal(blocks[0][0], audio_rows[: references[0].num_audio_rows])
+ assert torch.equal(blocks[1][0], video_rows[: references[0].num_video_rows])
+ assert torch.equal(blocks[2][0], audio_rows[references[0].num_audio_rows :])
+ assert torch.equal(blocks[3][0], video_rows[references[0].num_video_rows :])
+
+
+def test_split_condition_blocks_row_count_matches_the_layout():
+ references = [
+ _prepared("image", height=128, width=128),
+ _prepared("video", frames=37, audio_latents=207),
+ _prepared("audio", audio_latents=207),
+ ]
+ layout = rp.build_ref2va_packed_sequence(torch.ones(64, dtype=torch.long), references, 37, 48, 84, 207, (1, 2, 2))
+ video_rows = torch.zeros(layout.num_condition_video_rows, 96)
+ audio_rows = torch.zeros(layout.num_condition_audio_rows, 32)
+
+ blocks = R.split_condition_blocks(references, video_rows, audio_rows)
+ assert sum(block.shape[0] for block, _ in blocks) == (
+ layout.num_condition_video_rows + layout.num_condition_audio_rows
+ )
+
+
+def test_split_condition_blocks_rejects_a_row_count_mismatch(expect_error):
+ references = [_prepared("image")]
+ with expect_error(ValueError, "consumed"):
+ R.split_condition_blocks(references, torch.zeros(references[0].num_video_rows + 7, 96), None)
+
+
+def test_reference_condition_shapes_skips_audio_references():
+ references = [_prepared("image"), _prepared("audio", audio_latents=50), _prepared("video", frames=37)]
+ shapes = R.reference_condition_shapes(references)
+ assert shapes == ((1, 48, 84), (37, 48, 84))
+
+
+def test_normalize_reference_pixels_uses_imagenet_statistics():
+ """Not [-1, 1]: H3 conditions on VLM-style normalized pixels."""
+ frames = np.full((1, 32, 32, 3), 128, dtype=np.uint8)
+ pixels = R.normalize_reference_pixels(frames)
+ assert pixels.shape == (1, 3, 1, 32, 32)
+ expected = [(128 / 255.0 - m) / s for m, s in zip(R.MINIMAX_H3_PIXEL_MEAN, R.MINIMAX_H3_PIXEL_STD)]
+ for channel, value in enumerate(expected):
+ assert torch.allclose(pixels[0, channel], torch.full((1, 32, 32), value), atol=1e-6)
+ assert R.normalize_reference_pixels(np.zeros((5, 8, 8, 3), dtype=np.uint8)).shape == (1, 3, 5, 8, 8)
+
+
+REFERENCE_MEDIA = Path.home() / "h3_fl2va_artifacts" / "fl2va_first.mp4"
+
+
+def _real_media() -> Path:
+ if not REFERENCE_MEDIA.is_file():
+ pytest.skip(f"no reference media at {REFERENCE_MEDIA}; place a video with a soundtrack there")
+ return REFERENCE_MEDIA
+
+
+def _weights_dir() -> Path:
+ directory = Path(os.environ.get("MINIMAX_H3_MODEL_PATH", ""))
+ if not directory.is_dir():
+ pytest.skip("set MINIMAX_H3_MODEL_PATH to a diffusers snapshot of the checkpoint")
+ return directory
+
+
+_L1_SMALL = 16384 # 65536 and 32768 clash with the taps=3 encoder CBs
+
+# The e2e ring fabric params: a LINE config happens to pass here but production never runs it.
+MESH_4X8 = [
+ pytest.param(
+ (4, 8),
+ {**ring_params_req_exact_devices, "l1_small_size": _L1_SMALL},
+ id="4x8",
+ )
+]
+
+REFERENCE_PCC = 0.99 # floor set by the encoder's ttnn.group_norm calls (no fp32 path)
+# 17*1 + 5: smallest multi-chunk count (7 latent frames); the reference video VAE encode is CPU-hours at 124.
+VIDEO_FRAMES = 22
+
+
+def _reference_components(weights: Path, device: str = "cpu"):
+ from diffusers.models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio
+
+ vae = AutoencoderKLMiniMaxH3.from_pretrained(str(weights), subfolder="vae").to(device).eval()
+ audio_vae = AutoencoderKLMiniMaxH3Audio.from_pretrained(str(weights), subfolder="audio_vae").to(device).eval()
+ return SimpleNamespace(
+ vae=vae,
+ audio_vae=audio_vae,
+ patch_size=(1, 2, 2),
+ audio_latent_channels=audio_vae.config.latent_channels,
+ _execution_device=torch.device(device),
+ )
+
+
+def _real_frame() -> np.ndarray:
+ """Not randn: the fp16 round trip is only exercised on natural image statistics."""
+ frames, _, _ = rp.decode_reference_video(_real_media())
+ return frames[0]
+
+
+def _reference_case(case: str):
+ """Audio pairs with an image (audio-only is forbidden); the pairing is inert for the comparison."""
+ media = _real_media()
+ frames, fps, soundtrack = rp.decode_reference_video(media)
+ waveform, sample_rate = soundtrack
+
+ if case == "image":
+ specs = [dict(image=_real_frame())]
+ elif case == "video":
+ specs = [dict(video=frames[:VIDEO_FRAMES], fps=fps)]
+ elif case == "audio":
+ specs = [dict(image=_real_frame()), dict(audio=waveform, sample_rate=sample_rate)]
+ else:
+ raise ValueError(case)
+
+ ours, _ = R.prepare_references([rp.MiniMaxH3Reference(**s) for s in specs], VIDEO_FRAMES, AUDIO_RATE)
+ theirs, _ = reference_before_encoder.MiniMaxH3Ref2VASetupStep.prepare_references(
+ _components(), [reference_packing.MiniMaxH3Reference(**s) for s in specs], VIDEO_FRAMES
+ )
+ return ours, theirs
+
+
+@pytest.mark.timeout(7200)
+@pytest.mark.parametrize(("mesh_device", "device_params"), MESH_4X8, indirect=["mesh_device", "device_params"])
+@pytest.mark.parametrize("case", ["image", "audio", "video"])
+def test_encode_references_matches_reference(mesh_device, case, reset_seeds):
+ """Per modality (the three encode paths fail independently); geometry asserted too, since the layout is built from it."""
+ from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+ weights = _weights_dir()
+ ours, theirs = _reference_case(case)
+
+ pipeline = MiniMaxH3Pipeline.create_pipeline(mesh_device=mesh_device, weights_dir=weights, task="ref2va")
+ needs_video = any(reference.kind == "video" for reference in ours)
+ vae = pipeline._prepare_vae()
+ vae = pipeline._prepare_vae(encode_shape=(1, vae.tile_size, vae.tile_size), encode_taps=1)
+ if needs_video:
+ vae = pipeline._prepare_vae(
+ encode_shape=(pipeline.vae_config.clip_length, vae.tile_size, vae.tile_size), encode_taps=3
+ )
+ audio_encoder = pipeline._prepare_audio_encoder() if any(r.has_audio for r in ours) else None
+
+ got_video, got_audio = R.encode_references(
+ ours,
+ encode_clip=vae.encode_clip,
+ encode_video=vae.encode if needs_video else None,
+ encode_audio=(lambda waveform: audio_encoder(waveform)[0]) if audio_encoder else None,
+ latents_mean=pipeline.vae_config.latents_mean,
+ latents_std=pipeline.vae_config.latents_std,
+ audio_latents_mean=pipeline.audio_config["latents_mean"],
+ audio_latents_std=pipeline.audio_config["latents_std"],
+ patch_size=pipeline.patch_size,
+ audio_latent_channels=pipeline.audio_config["latent_channels"],
+ )
+
+ components = _reference_components(weights)
+ with torch.no_grad():
+ want_video, want_audio = reference_encoders.MiniMaxH3Ref2VAReferenceEncoderStep.encode_references(
+ components, theirs, device=torch.device("cpu")
+ )
+
+ for index, (a, b) in enumerate(zip(ours, theirs)):
+ assert (a.num_latent_frames, a.latent_height, a.latent_width) == (
+ b.num_latent_frames,
+ b.latent_height,
+ b.latent_width,
+ ), f"reference {index} resolved to different visual latent geometry"
+ assert a.num_audio_latents == b.num_audio_latents, f"reference {index} resolved a different audio length"
+ logger.info(
+ f"[{case}] resolved geometry: "
+ + ", ".join(
+ f"{r.kind}({r.num_latent_frames}x{r.latent_height}x{r.latent_width}, audio {r.num_audio_latents})"
+ for r in ours
+ )
+ )
+
+ assert got_video.shape == want_video.shape, f"{tuple(got_video.shape)} != {tuple(want_video.shape)}"
+ logger.info(f"[{case}] video condition rows {tuple(got_video.shape)}")
+ assert_quality(want_video, got_video, pcc=REFERENCE_PCC)
+
+ if got_audio is not None:
+ assert got_audio.shape == want_audio.shape, f"{tuple(got_audio.shape)} != {tuple(want_audio.shape)}"
+ logger.info(f"[{case}] audio condition rows {tuple(got_audio.shape)}")
+ # The soundtrack takes the posterior MEAN, so it carries no sampling noise at all.
+ assert_quality(want_audio, got_audio, pcc=REFERENCE_PCC)
+ else:
+ assert want_audio is None, "the reference produced audio rows where we produced none"
+
+
+def _tokenizer():
+ weights = os.environ.get("MINIMAX_H3_MODEL_PATH")
+ if not weights or not (Path(weights) / "tokenizer").is_dir():
+ pytest.skip("needs MINIMAX_H3_MODEL_PATH pointing at a diffusers snapshot with tokenizer/")
+ from transformers import AutoTokenizer
+
+ return AutoTokenizer.from_pretrained(str(weights), subfolder="tokenizer")
+
+
+@pytest.mark.parametrize("case", list(CASES), ids=list(CASES))
+def test_presentation_matches_reference(case):
+ """The tags are the load-bearing half: H3 tags the WHOLE vision block video, sentinels included."""
+ tokenizer = _tokenizer()
+ ours, theirs = _both(CASES[case])
+ image_counts = [4096 for reference in ours if reference.kind == "image"]
+ video_counts = [1008 for reference in ours if reference.kind == "video"]
+ for prepared in (ours, theirs):
+ for reference in prepared:
+ if reference.kind == "video":
+ reference.block_timestamps = [0.25, 1.25, 2.25, 3.25, 4.25, 5.0]
+
+ prompt = "a red fox stepping through wet grass at dawn"
+ got_ids, got_tags = rp.build_ref2va_presentation(tokenizer, prompt, ours, image_counts, video_counts)
+ want_ids, want_tags = reference_packing.build_ref2va_presentation(
+ tokenizer, prompt, theirs, image_counts, video_counts
+ )
+ assert got_ids == want_ids
+ assert got_tags == want_tags
+ assert len(got_tags) == len(got_ids)
+
+
+class _ProcessorStub:
+ def __init__(self, weights: Path):
+ from transformers import AutoImageProcessor, AutoTokenizer, AutoVideoProcessor
+
+ self.tokenizer = AutoTokenizer.from_pretrained(str(weights), subfolder="tokenizer")
+ self.image_processor = AutoImageProcessor.from_pretrained(str(weights), subfolder="text_encoder")
+ self.video_processor = AutoVideoProcessor.from_pretrained(str(weights), subfolder="text_encoder")
+
+
+def _build(stub, prompt, references):
+ from ....pipelines.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+ return MiniMaxH3Pipeline._build_ref2va_presentation(stub, prompt, references)
+
+
+def test_presentation_orders_vision_patches_by_reference_not_by_batch():
+ """`_scatter_rows` consumes tower rows in run order, so processor-batch order would swap blocks."""
+ stub = _ProcessorStub(_weights_dir())
+ video = rp.reference_from_video_file(_real_media(), with_audio=False)
+ image = rp.MiniMaxH3Reference(image=_image(1024, 1024))
+
+ prepared, _ = R.prepare_references([video, image], TARGET_FRAMES, AUDIO_RATE)
+ input_ids, tags, type_ids, pixel_values, grid_thw, kinds = _build(stub, "a prompt", prepared)
+
+ assert kinds == ["video", "image"], f"vision entries are in {kinds}, not presentation order"
+ assert int(grid_thw[0][0]) == len(prepared[0].block_timestamps) > 1
+ assert tuple(grid_thw[1].tolist()) == (1, 128, 128)
+ assert pixel_values.shape[0] == sum(int(grid.prod()) for grid in grid_thw)
+
+ reversed_prepared, _ = R.prepare_references([image, video], TARGET_FRAMES, AUDIO_RATE)
+ _, _, _, _, reversed_grid, reversed_kinds = _build(stub, "a prompt", reversed_prepared)
+ assert reversed_kinds == ["image", "video"]
+ assert tuple(reversed_grid[0].tolist()) == (1, 128, 128)
+
+
+def test_presentation_vision_runs_line_up_with_the_towers_rows():
+ from ....encoders.qwen3vl.model_qwen3vl import vision_token_runs
+
+ stub = _ProcessorStub(_weights_dir())
+ references = [
+ rp.MiniMaxH3Reference(image=_image(1024, 1024)),
+ rp.reference_from_video_file(_real_media(), with_audio=False),
+ rp.MiniMaxH3Reference(audio=_waveform(DURATION)),
+ ]
+ prepared, _ = R.prepare_references(references, TARGET_FRAMES, AUDIO_RATE)
+ input_ids, tags, type_ids, pixel_values, grid_thw, kinds = _build(stub, "a prompt", prepared)
+
+ pad_ids = [stub.tokenizer.convert_tokens_to_ids(t) for t in ("<|image_pad|>", "<|video_pad|>")]
+ runs = vision_token_runs(input_ids, pad_ids)
+ merge = stub.image_processor.merge_size**2
+
+ assert len(runs) == sum(int(grid[0]) for grid in grid_thw)
+ assert sum(length for _, length in runs) == sum(int(grid.prod()) for grid in grid_thw) // merge
+ # Runs are sorted and disjoint, which `_scatter_rows` requires outright.
+ assert all(a[0] + a[1] <= b[0] for a, b in zip(runs, runs[1:]))
+ assert stub.tokenizer.decode(input_ids[0]).count("