专题二 · 外传　·　L100　·　主线 17 分钟（§三 可跳，展开另加 2 分钟）

算力强，显存弱 —— 我们把这样一颗芯片用成了什么样

TPU v6e 与扩散模型。前十分钟一条线加两张芯片图，判你的活配不配； 后八分钟摊开我们真跑过的东西。

开场这一讲要回答的两件事

① 配不配 —— 一条线，跑之前就能判。② 怎么落地 —— 在哪儿切、三段放哪台机器。后一半是我们趟出来的。

⛔ 先说清这一讲不做什么 —— 不比 benchmark，一个快慢数都没有

专题二立过一条线：算力 ÷ 显存带宽 —— 「每搬一个字节，这台机器本来能算多少次」。 那一讲量出 B200 和 TPU v7 几乎一模一样；这一讲多量一颗，它立刻破了那个「都一样」。

屋脊点是结构量，它只说「这台机器的胃口有多大」，不说「这顿饭多久吃完」。

⚠️ 这不是回避：我们手上唯一一组同模型双平台实测，本身就不支持「v6e 更快」。 与其挑一组好看的数，不如把边界说死 —— 这一讲讲结构匹配。

扩散模型本身的原理在专题十一；每个数怎么推出来的、以及两颗芯片的 显微镜级拆解在 L200。

一同一把尺子量四颗芯片 —— 再把其中两颗拆开

先把这颗芯片的四个官方数摆出来 ——  整节课要用的原料就这四个。

918TFLOPS BF16 / 芯片
1,836TOPS INT8 / 芯片
32 GBHBM2e / 芯片
1,638GB/s HBM 带宽

TPU v6e（研发代号 Trillium）是 Google 第六代张量处理单元， 单 Pod 规模 256 芯片，针对 Transformer、文生图与卷积网络的 训练、微调与在线服务做了系统级优化。

⭐ 这一讲的头两个数就是从上面这排里除出来的： 918 ÷ 1.638 ＝ 560 —— 下面那张图上 v6e 那一根，就是这么来的。 另外两个数（32 GB、以及 1,638 本身）在 §三 还会各用一次。

规格出处：Cloud TPU v6e 官方文档 docs.cloud.google.com/tpu/docs/v6e —— 本课自己从 JAX tpu_info 读到的是 bf16 920 / int8 1,840 TOPS， 与官方差 0.3% 以内（取整口径不同，不是矛盾）。

同一把尺子，量四颗芯片 —— 只有 v6e 落在别处 尺子还是专题二那一把：算力 ÷ 显存带宽 ＝ 每搬一个字节，这台机器本来能算多少次。数越大，越挑「算得多、搬得少」的活。 三颗旗舰：几乎重合 v6e：跑出去将近一倍 295 ～ 313：三颗挤在这一线 NVIDIA H100 SXM 989.5 ÷ 3.35 TB/s 数据表（稠密 ＝ 1,979 ÷ 2） 295 NVIDIA B200 2500 ÷ 8.0 TB/s 专题二 §1 已算过 312 TPU v7（Ironwood） 2307 ÷ 7.38 TB/s 专题二 §1 已算过 313 TPU v6e（Trillium） 918 ÷ 1.638 TB/s 官方 v6e 规格表 560 0 100 200 300 400 500 600 屋脊点（FLOP / byte）—— 每搬一个字节，本来能算多少次 v6e 的胃口是 H100 的 1.9 倍 「算力强、显存弱」不是形容词 —— 它就是这个 560 把 v6e 跟 H100 逐项摆开：算力 918 对 989.5，是它的 93%，基本打平；而 带宽 1,638 对 3,350，只有 49%，容量 32 GB 对 80 GB，只有 40%。 分子几乎没变、分母砍掉一半 —— 屋脊点就翻了一倍。这就是那句「算力强、显存弱」唯一精确的说法。 所以这一讲只回答一个问题：什么样的活凑得够 560 屋脊点是结构量，不是性能量。它只说「这台机器的胃口有多大」，一个字都没说谁跑得快 —— 本讲从头到尾不比 benchmark。 接下来两张图把两颗芯片拆开，看这个 560 是从什么样的硅片布局里长出来的；然后再看扩散模型的计算长什么样，为什么它恰好是那种「凑得够」的活。
图 X-1　三颗旗舰挤在 295–313，v6e 是 560。「算力强、显存弱」不是形容词 —— 它就是这个数。

这个 560 从什么样的布局里长出来？先看 v6e 这一颗 —— H100 的对照放在下面那一折里，需要时再点开。

一颗 TPU v6e 拆开看 —— 1 颗芯片 ＝ 1 个核 ＝ 1 个 device 记住三件事就够了：算力集中在 2 个大方阵里、片上有一整块 128 MiB 的暂存、通往片外的那根管子只有 1,638 GB/s。 算：矩阵与向量单元 存：片上暂存 片外主存 HBM SparseCore（本讲不展开） 一颗 v6e 芯片 对软件 ＝ 1 个 device 不用除以 2 918 TFLOP/s bf16 TensorCore × 1 v7 是 2 个 —— 这一代只有 1 个 稠密算力全在这里面 MXU 0 256 × 256 MXU 1 256 × 256 矩阵乘全部在这两块里发生 一整块 256×256 的方阵，不切成小片 向量单元 VPU 8 × 128 标量单元 发指令 · 发搬运 归一化、激活、softmax 走这儿 VMEM　128 MiB 片上暂存（不是缓存） SMEM　1 MiB 标量 / 描述符 v7 是「两个核各 64 MiB」，v6e 是「一个核独占 128 MiB」—— 总量一样，但 v6e 不用把大张量切两半 SparseCore × 2 两个职责：① 稀疏 / 大表　② 集合通信卸载 只要跨卡就用得上它（all-gather / reduce-scatter）—— 见图 X-7 HBM　32 GB 片外主存 —— 模型和中间结果都住这儿 通往片上的带宽只有 1,638 GB/s —— H100 是 3,350，只有它的 49% 1,638 GB/s—— 整颗芯片的数据，都从这儿过 对外怎么连 官方规格表 ICI 端口 4 个 v7 是 6 个 ICI 带宽 800 GB/s 每 chip 双向合计 拓扑 2D 环面 v7 是 3D 一个 Pod 256 颗 v7 是 9,216 颗 每台主机 8 颗 DRAM 1,536 GiB 这一张只要记住三条 ① 1 颗 ＝ 1 个核 ＝ 1 个 device 专题二那个「除以 2」的坑，这一代没有 ② 片上一整块 128 MiB 暂存 大张量不用切成两半，编译器好排 ③ 片外那根管子只有 1,638 GB/s 这就是 560 的分母 —— 全讲的症结 把这张布局压成一句话：算的地方很大，进料的门很窄 两块 256×256 的方阵摆在那儿，每一拍能吞下的乘加数是很大的；可所有数据进出片外，只能挤 1,638 GB/s 这一根管子。 这就是 560 的来历：分子（算力）没少，分母（带宽）被砍掉一半，于是「每搬一个字节得算多少次才不亏」这个门槛，就抬到了别人的两倍。
图 X-2　一颗 TPU v6e：算力集中在两个 256×256 的方阵里，片上一整块 128 MiB 暂存（全由编译器安排，没有硬件缓存兜底），通往片外那道门 1,638 GB/s。
⭐ 对照 · 一颗 H100 拆开看（图 X-3，默认折起）—— 同样的画法，算力摊成 528 个小单元，片上多一整层硬件自动管的 L2，门宽一倍
一颗 NVIDIA H100 拆开看 —— 同样的画法，相反的选择 跟上一张比着看：算力摊成 528 个小单元、片上多了一整层硬件自动管的缓存、通往片外的门宽一倍。 算：SM 与 Tensor Core 存：软件管的 L1 / 共享 存：硬件管的 L2 —— v6e 没有这一层 片外主存 HBM3 一颗 H100 SXM 8 GPC · 66 TPC · 每 TPC 2 个 SM 989.5 TFLOP/s bf16 稠密 SM × 132 算力摊在 132 个小核里 —— v6e 是 1 个大核 每格 ＝ 1 个 SM 每个 SM 里：128 个 CUDA Core ＋ 4 个第四代 Tensor Core 全片合计 16,896 个 CUDA Core、528 个 Tensor Core 一个 SM 里的存 L1 ＋ 共享内存　256 KB 其中共享可配到 228 KB 寄存器堆　256 KB L2 缓存　50 MB 全片共享，硬件自动管 —— 程序管不着它留什么、赶走什么 v6e 完全没有这一层—— JAX 源码里 v6e 的 CMEM 直接写着 0。片上放什么，那边全部由编译期决定。 HBM3　80 GB 容量是 v6e 的 2.5 倍（80 对 32） 通往片上的带宽 3,350 GB/s —— v6e 是 1,638，这道门宽了一倍 3,350 GB/s 这道门的宽度跟上一张按同一比例画 三个相反的选择 这一讲的全部内容 算力怎么摆 v6e：2 块大方阵 H100：528 个小单元 一个是「一次吞一大块」，一个是「同时应付很多小块」 片上谁做主 v6e：编译期写死 H100：多一层硬件缓存 v6e 的 CMEM ＝ 0，那一层它压根没有 门有多宽 v6e：1,638 GB/s H100：3,350 GB/s 这一条直接决定了那个 560 对 295 一个反直觉的账 H100 全部片上内存 132 × 256 KB ＋ 50 MB ≈ 84 MB v6e 光 VMEM 一项 128 MiB ≈ 134 MB 门窄，但屋里的台面更大 —— 这两件事是配套的 两颗芯片在同一件事上做了相反的选择 —— 而它们各自都是自洽的 H100 摊成 528 个小单元、再压一层硬件缓存，是为了应付「我不知道你要跑什么」：形状随时会变、访问随时可能落空，那就多留后手。 v6e 压成 2 块大方阵、片上全交给编译器，是为了吃透「我早就知道你要跑什么」：形状固定、访问可提前排，后手就是浪费，不如把面积全给算力和台面。
图 X-3　一颗 H100：同样的画法。算力摊成 528 个小单元，片上多了一整层硬件自动管的 L2，门宽一倍。⭐ 两张图里那道门是按带宽等比画的，可以直接叠着比。
⭐⭐ 两颗都是自洽的：一个在防「我不知道你要跑什么」，一个在吃透「我早就知道你要跑什么」

同样三个问题，两边给了相反的答案：

同一个问题TPU v6eH100这个选择在防什么
算力怎么摆集中成 2 个 256×256 方阵摊成 528 个小单元 摊开是为了不管来什么形状都有人能接；集中是为了大矩阵上把利用率吃满
片上谁做主128 MiB 全归编译器，CMEM ＝ 0 50 MB L2 由硬件自动管 硬件自动管是为了应付负载未知；交给编译器是为了吃透负载已知
门开多宽1,638 GB/s3.35 TB/s 门宽一倍，代价是屋脊点低一半 —— 它挡的是「强度不够」的活

⭐ H100 那一列全部指向「负载未知」：摊开、缓存兜底、门开大 —— 三个都是为不确定性买的保险。 v6e 那一列全部指向「负载已知」：集中、编译期写死、门只开够用 —— 三个都是把保险费省下来换算力密度。
—— 所以问题从来不是谁更强，是你手上的活属于哪一种。

⚠️ 那个 CMEM ＝ 0 是有代价的：编译器排错了就没有后手。 形状动态、访存模式运行时才知道的负载，在这套设计上会很难受 —— 而扩散恰好把这个前提喂得满满的（§四会用到这一条）。

⚠️ 还有一条：v6e 是 1 颗 ＝ 1 个核 ＝ 1 个 device。 专题二反复强调的 v7「容量除以 2」那个坑，这一代没有。

这条线怎么用，以及它为什么跑之前就能画

把你要跑的那段计算也算出一个强度（每搬一个字节实际算了多少次）， 两个数一比 —— 落线右边＝算力受限，落左边＝带宽受限。

三个比值：跟 H100 比，算力 93%、带宽 49%、容量 40%。 分子基本没动，分母砍一半，商自然涨到两倍。

⭐ 分子分母都取官方规格，不取任何实测值 —— 所以这条线在跑任何东西之前就能画出来。 客户还在犹豫要不要给你机器的时候，你已经能下判断了。

560 怎么除出来的、H100 那个 989.5 为什么不是 1979、我们拿什么复现过这条公式 —— 见 L200 §一。

二那你的活落在哪一边

拿 Wan2.1 的官方配置当场算，不引用结论。

为什么扩散是计算密集型 —— 拿 Wan2.1 的真配置当场算一遍 结论先说：一段 720P、81 帧的视频 ＝ 75,600 个 token，而注意力是全局的—— 于是算力的七成花在注意力上，那是纯矩阵乘。 ① 视频怎么变成 token ② 算力花在哪 ③ 为什么这配 v6e ① 一段视频 ＝ 多少个 token Wan2.1 官方配置 像素 你要的那段视频 720 × 1280 × 81 帧 ↓ VAE stride (4, 8, 8) 时间 4 倍、空间 8×8 下采样 latent (81−1)/4+1 ＝ 21 21 × 90 × 160 ↓ patch (1, 2, 2) 空间 2×2 合成一个 token token 21 × 45 × 80 N ＝ 75,600 个 token ② 每层每步的算力花在哪 d ＝ 5120 · ffn ＝ 13824 注意力本身 4·N²·d 72% FFN 4·N·d·ffn 13% 自注意力投影 8·N·d² 10% 交叉注意力（文本 512） ≈ 5% 每层合计 1.63e14 FLOP　→　40 层一步 6.52e15 50 步一段视频 3.26e17 FLOP 注意力占七成，而它是 N² 的纯矩阵乘 —— config 里 window_size ＝ (−1,−1)，不开窗口，全局都算 ③ 于是它正好配 v6e 强度 ≈ 每份权重服务多少 token 75,600 ／ 门槛 560 → 高出 135 倍，那根窄管子完全不是瓶颈 而且七成算力是 N² 的注意力——  纯矩阵乘，MXU 的主场， 正是那两块 256×256 最擅长的形状 剩下要比的只有算力： 918 对 989.5 ＝ 93%，基本打平 一颗 v6e 的纯算力下界约 5.9 分钟 （3.26e17 ÷ 918 TFLOP/s，100% 利用率，实际做不到） 把这一整套压成三句话 ① 一段视频被压成七万五千个 token，而它们共用同一份权重 —— 每层那 702 MB 的权重搬一次，服务 75,600 个位置。 ② 注意力是全局的，所以算力随 token 数平方涨。七万五千个 token 两两都算，光这一项就占了七成 FLOP —— 这就是「计算密集」四个字的全部来历。 这张图刻意留白的地方 FLOP 只数了矩阵乘，没数 norm、激活、RoPE、softmax 那些向量运算。它给的是量级和占比，不是精确账。 那个 5.9 分钟是理论下界，不是性能预测 —— 它假设 100% 利用率，而实际做不到。 本讲不给性能数，这一条也不例外，它只是让「3.26e17」这个数有个体感。
图 X-9　720P、81 帧经 VAE 与 patch 化后是 75,600 个 token；config 里 window_size = (−1,−1)，不开窗口、全局注意力，于是七成算力压在 N² 的注意力上。

强度 ≈ 75,600，是 560 的 135 倍。放回轴上：

把活放回那两条线上 —— 扩散模型落在很右边，而且是甩开两个数量级 一条规则就够：强度 ≈ 同一份权重被多少个「位置」共用。权重从 HBM 搬上来是固定开销，服务的位置越多，这笔搬运摊得越薄。 喂不饱：卡在搬运上 喂得饱 扩散：甩开两个数量级 H100　295 v6e　560 LLM decode　batch 1 一次只产 1 个 token 1 LLM decode　batch 64 64 条请求共用一次权重搬运 64 LLM prefill　8K 上下文 整段一次过 8,192 文生图　1024 × 1024 latent 128×128 ＝ 16,384 个位置 16,384 文生视频 再多一到两个数量级（只标量级） 10 万 ～ 100 万 1 10 100 1,000 10,000 100,000 1,000,000 算术强度（FLOP / byte，对数轴）—— 越往右，越是「算得多、搬得少」 ← 这一边：卡在搬运上，算力再强也闲着 这一边：喂得饱，此时比的是算力 → 所以「v6e 适合扩散」的精确说法是这个 不是它跑扩散更快。是扩散把它的短板挡在了瓶颈之外 —— 强度上万，离 560 差着一个半数量级，那根窄管子根本没成为瓶颈。 于是只剩下没被挡住的那一项在起作用：算力 918 对 989.5，是 H100 的 93%。短板不参与，长板打平 —— 这就是「合适」的全部含义。 同一条道理，原样适用于 prefill 重、decode 轻的任务 prefill 是整段一次过，几千上万个位置共用一次权重搬运 —— 它跟扩散落在轴上的同一边。 所以选型判据可以压成一句：先看你的活在这根轴上落在哪儿，再看那一边比的是什么。落右边比算力，v6e 打平；落左边比带宽，v6e 吃亏。
图 X-4　一条规则就够：强度 ≈ 同一份权重被多少个「位置」共用。小 batch 的 decode 落在最左边，prefill 和扩散甩到右边两个数量级。
⭐⭐ 于是「v6e 适合扩散」的精确说法是：短板不参与，长板打平

不是它跑扩散更快。是扩散把它的短板挡在了瓶颈之外 ——  强度上万，那根窄带宽根本没参与；只剩没被挡住的那一项在起作用： 算力 918 对 989.5，是 H100 的 93%。

⭐ 顺带说我们怎么确认自己没算错：同一份 config 把参数量数出来， 每层 351.3 M × 40 层 ＝ 14.05 B，正好对上官方标称的 14B。 参数量能对上，说明我们对「这一层里有哪些矩阵」的理解是对的。

⛔ 反过来同样成立，这条决定别推错场景：小 batch 的 decode 在轴最左边， 那里比的全是带宽，而 v6e 只有 H100 的 49% —— 那种活它最吃亏。

⚠️ 还有一条：落在算力侧不等于算力就吃满了。 我们在 Wan2.2 上抓 profile，Splash Attention 那个算子的 MFU 只有 37% （roofline 15.974 ms ÷ 实测 43.93 ms），整模型优化后是 34%（基线 12%）—— 卡在 head_dim 128 对 MXU 的 256，方阵有一半是空的。细节在 L200 §三点五。

三一颗装得下吗⏱ 主线可跳

一句话带过就够：权重决定装不装得进，激活决定跑不跑得动 —— 而扩散是后者说了算（我们撞 OOM 时 XLA 报的是只剩 13.10G，而 VAE 一步要 19.00G）。展开的账在下面，主线不讲。

⏱ 整节展开（约 2 分钟）：那根权重体积轴、实测配置，以及我们自己撞 OOM 的那笔实测账
装得下吗 —— 要算的是权重 ＋ 峰值激活，而扩散是激活说了算 下面那根轴只画权重—— 而扩散模型真正吃显存的是运行时的激活，它比权重大、而且随分辨率和帧数涨。图下半部分是我们自己撞 OOM 时的实测账。 扩散模型 对照：大语言模型 一颗 v6e ＝ 32 GB 实测：单颗 实测：8 卡一台主机 一颗 v6e　32 GB 一台主机 8 颗　256 GB S3Diff（SD-Turbo） 单步超分 4× 6.6 GB 单颗 8 卡实测反而更慢 SDXL 文生图 7 GB v6e-1　（也测过 v6e-4 / v6e-8 数据并行） HunyuanVideo-1.5 文生视频 16.6 GB v6e-8 FLUX.1 [dev] 文生图 24 GB 未记录实测配置 Wan2.1-T2V-14B 文生视频 权重贴着 32 GB 的边 28 GB v6e-8（dp=1, tp=8） Wan2.2-T2V-A14B MoE：总 27B / 每步激活 14B 54 GB v6e-16（dp=2, sp=1, tp=8） Qwen3.5-397B 对照：LLM 794 GB —　我们没在 v6e 上跑过 DeepSeek-V3 对照：LLM 1342 GB —　我们没在 v6e 上跑过 10 GB 100 GB 1,000 GB bf16 权重体积（对数轴）——  只算权重，不含激活与中间结果 同一颗 v6e 上的真实预算（CogVideoX VAE，实测） 多少 出处 / 说明 芯片标称 32 GB 官方规格表上的 HBM 容量 那次 OOM 时实际可用 13.10 GB XLA 报错原文：There are 13.10G free —— 其余被权重与运行时占着 VAE 解码一步要多少 19.00 GB XLA 报错原文：Attempting to reserve 19.00G ——  19 > 13.1，OOM 对照：这个模型的权重 10 GB CogVideoX-5B，bf16 ——  激活是权重的近两倍 改实现之后 < 13 GB 逐帧解码 ＋ 共享缓存 ——  解法不是换更大的卡，是改实现
图 X-6　上半是权重体积与我们真跑过的配置；下半是同一颗芯片上的真实预算 —— 那是我们自己撞 OOM 时的实测账。⛔ 权重决定装不装得进，激活决定跑不跑得动 —— 而扩散是后者说了算。
⛔⛔ 「28 GB < 32 GB 所以能跑」这句话分子分母都错

① 分母错了。32 GB 是标称，不是预算。真跑起来权重和运行时先占掉一大块 ——  我们那次 OOM 时 XLA 报的是：There are 13.10G free，只剩 13.1 GB。

② 分子也错了。要放进去的不只是权重，还有峰值激活。 而扩散这一族激活比权重大：CogVideoX-5B 权重 10 GB， 而它的 VAE 解码一步要 19 GB（Attempting to reserve 19.00G）——  19 > 13.1，于是 OOM。

⭐⭐ 而且激活随分辨率与帧数涨，权重一个字节不涨 ——  Wan2.1 的 480P 跑得动、720P OOM，用的是同一份权重。 权重决定装不装得进，激活决定跑不跑得动。

⭐ 解法也不是换更大的卡：逐帧解码 ＋ 共享缓存，把那 19 GB 压到 < 13 GB —— 改实现。

另外两处「按体积猜会猜错」的：

③ S3Diff 只有 6.6 GB，按体积猜「一颗绰绰有余」—— 对，但那不是重点。 真正的发现是：我们把它摊到 8 卡做张量并行，实测反而更慢 （5.46 秒 对 5.28 秒），而预热长了 15 倍。模型太小，通信开销盖过了收益。

④ Wan2.1 的 28 GB 按体积猜「贴边能塞进一颗」——  而实测从来没人这么跑：它是在 v6e-8 上 dp=1、tp=8 摊开跑的。

⭐ 前者看体积就能答，后者只能实测 —— 而我们十个模型每一个都测过。

同一台 v6e-8 的两种用法（都实测过）：SDXL 7 GB 一颗装得下 → 开 8 路各生成各的，2.40 张/秒；Wan2.1 → TP 摊到 8 颗。

实测分布：小的单颗，主力清一色 8 卡 —— 没有一个需要跨主机

单颗：S3Diff · SDXL（延迟最优）· Real-ESRGAN
8 卡：HunyuanVideo-1.5 / Wan2.1 / CogVideoX 在 v6e-8 · Flux.2 在 v4-8 · Wan2.2 I2V 在 v6e-16

⭐ 这正好解释 X-2 里那三条看着像减配的规格 （4 个 ICI 口、二维环面、Pod 只有 256）： v6e 不打「一个模型摊在几千颗上」那场仗 —— 不打，就不用付那个成本。

这张图的数是从哪儿来的（出处 ＋ 三条边界）

右列「实测配置」全部取自各模型 README 的测试环境段： SDXL v6e-1/4/8 · HunyuanVideo-1.5 / Wan2.1 / CogVideoX 在 v6e-8 · Flux.2 在 v4-8 · Wan2.2 I2V 在 v6e-16（分片配置见该模型优化指南第三章）· S3Diff 与 Real-ESRGAN 单颗。

「8 卡反而更慢」出自 S3Diff README 的 Why Not Multi-Chip 段； 权重体积按 bf16 每参数 2 字节换算，参数量出自各家官方模型卡。

⛔ 那根轴只算权重，不含激活与编译缓存 ——  它能回答「装不装得下」，回答不了「该用几颗」。
⚠️ FLUX.1 我们没有记录实测配置，图上如实留空。

四仓库里为什么同一个模型有两份例子

下面这一半是我们自己趟出来的。仓库里每个扩散模型都有两份例子：一体化脚本，和一个 *_staged/ 目录 —— 后者存在的理由不是更快，是它把切口露在外面。

拿 Wan 2.2 图生视频当例子 —— 十个模型都是这套布局： github.com/yangwhale/gpu-tpu-pedia/tree/main/tpu/Wan2.2

tpu/Wan2.2/
├── generate_i2v_torchax.py                   ← ① 一体化：一个进程从头跑到尾
├── generate_diffusers_i2v_torchax_staged/    ← ② 三阶段
│   ├── stage1_encoder.py                        文本 ＋ 首帧 → embedding
│   ├── stage2_transformer.py                    五十步去噪 → latent
│   ├── stage3_vae_decoder.py                    latent → 成片
│   ├── utils.py                                 落盘 / 读盘的 helper 全在这
│   └── stage_outputs/                        ← ⭐ 切口就在这个目录里
│       ├── stage1_embeddings.safetensors     ← 段与段之间唯一交接的，就这三个文件
│       ├── stage2_latents.safetensors
│       ├── generation_config.json
│       └── output_video.mp4                     成片
├── docs/wan_tpu_optimization_guide.md        约 1,970 行迁移与优化指南
└── README.md

⚠️ 下面这张图上的字节数取自 Wan2.1 那一份同名目录 （tpu/Wan2.1/generate_diffusers_torchax_staged/stage_outputs/）—— 两个模型目录结构相同，数不同。

仓库里为什么同一个模型有两份例子 —— 三阶段那份是拿来「看得见中间」的 一体化跑得快，但你看不见中间那份 latent 长什么样。三阶段把切口露出来：两个 safetensors ＋ 一份 config —— 下面这些数全是直接解文件头得到的，不是算的。 一体化：一个进程 三阶段：三个进程 ＋ 落盘交接 可逐字节核对 ① 一体化 generate_torchax.py 出片 / benchmark 一个进程，从 prompt 直接到成片。 看不见中间态：文本 embedding、 latent 全在内存里，跑完就没了。 它的用处是快 —— 验证一次 改动、量一次端到端，用这个。 但一出问题（视频全黑、出 NaN）， 你没有任何中间产物可看。 调试时它帮不上忙 ② 三阶段 —— 切口露在外面 generate_diffusers_torchax_staged/ 调试 / 部署 / 教学 stage1 ⏷　stage1_embeddings.safetensors 7,406,544 B prompt_embeds F32 [1, 226, 4096] negative_prompt_embeds　F32 [1, 226, 4096] metadata 的 dtype_info 记着：原始是 bfloat16 stage2 ⏷　stage2_latents.safetensors 19,353,872 B latents　F32 [1, 16, 21, 90, 160] ＝ 批 1 · 通道 16 · 帧 21 · 高 90 · 宽 160 跨机时唯一要搬的就是这一份 stage3 ⏷　generation_config.json 817 B height 720 · width 1280 · num_frames 81 steps 50 · guidance 5.0 · seed 2025 · model_id 没有它，前两个文件无法自解释 三段之间只认这三个文件、不认进程—— 所以它们天然可以跑在三台机器上 拿到一份 latent，三步自检 怎么做 看什么 ① shape 对不对 由分辨率直接推：帧 (81−1)/4+1 ＝ 21 · 高 720/8 ＝ 90 · 宽 1280/8 ＝ 160 · 通道 16 期望 [1, 16, 21, 90, 160]。对不上就别往下跑 —— 后面只会得到全黑或 NaN ② dtype 在哪看 读 safetensors 头的 dtype 字段，再读 metadata 里的 dtype_info 两者可能不一样：Wan 的 embedding 盘上是 F32，而 dtype_info 写着原始是 bfloat16 ③ 字节数对不对 元素数 × 每元素字节 ＋ 头 ＝ 文件大小 16×21×90×160×4 ＋ 272 ＝ 19,353,872，跟文件一个字节不差
图 X-10　三阶段之间只交接三个文件：两个 safetensors ＋ 一份 config。⭐ 图上的字节数、shape、dtype 全是直接解 safetensors 文件头得到的，这三个文件仓库里就有，可自行复核。
⭐⭐ 拿到一份 latent，三步自检：shape、dtype、字节数 —— 而「看文件大小」一步都不算

① shape 对不对：由分辨率直接推 —— 帧 (81−1)/4+1 ＝ 21、高 720/8 ＝ 90、 宽 1280/8 ＝ 160、通道 16 → 期望 [1, 16, 21, 90, 160]。 对不上就别往下跑，后面只会得到全黑或 NaN。

② dtype 在哪看：读 safetensors 头的 dtype 字段，再读 metadata 里的 dtype_info。⭐ 两者可能不一样 —— Wan 的 embedding 盘上是 F32， 而 dtype_info 写着原始是 bfloat16（保存时转的，加载时按这条恢复）。

③ 字节数对不对：16×21×90×160×4 ＋ 272 字节头 ＝ 19,353,872， 跟文件一个字节不差。

⛔ 为什么「看文件大小」不算验证：Wan2.1 的 latents 是 F32、19,353,872 B； CogVideoX 的 latents 形状一模一样，但 BF16、9,677,064 B ——  大小差一倍，形状相同。只能读头。

⭐ 顺带说切口有多便宜：跨机要搬的就是那 19 MB 出头，走 100 Gbps 约 1.5 毫秒， 而被切开的那一段本身要算 229 秒 —— 「能不能切」这个问题在这里根本不成立。

⭐⭐ 三阶段那份例子的真正用途：它是一个「把切口露出来」的装置

教学上：让人去 stage_outputs/ 把那两个文件打开看一眼， 比讲十分钟结构管用。

调试上：全黑 / NaN / 动作快进这些问题，在一体化脚本里无从下手， 三阶段能逐段定位。

部署上：既然只认文件不认进程，三段就能跑在三台机器上 ——  下一节讲各放哪台。

⛔ 这张图我推翻重写过一次 —— 注解抵消不了图形本身的断言（那个错值得讲给学员听）

初版画的是「数据体积对数轴」，最粗的一根柱子标着 457 GB —— 注意力分数矩阵，旁边我还很小心地注了一句「从不落地」。

⛔ 那是错的：那个矩阵从来没有被物化出来过。 Flash / Splash Attention 是分块算的，算完即弃， HBM 里根本不存在这么一块。

⭐ 而「从不落地」那句注解并不能救它： 把一个不存在的量画成「管子最粗处」， 整张图的比例尺就锚在了虚构上 ——  读者记住的是柱子，不是柱子旁边那行小字。 注解抵消不了图形本身的断言。

这张图的数是从哪儿来的（含一处 config 写的 ≠ 实际跑的）

三个文件的字节数、张量形状、dtype 与 metadata ——  直接解 safetensors 文件头得到；文件就在本仓库 tpu/Wan2.1/generate_diffusers_torchax_staged/stage_outputs/ 下， 可自行复核。CogVideoX 那条对照取自它同名目录下的 stage2_latents.safetensors。

⚠️ 官方 wan_t2v_14B.py 写 text_len ＝ 512， 而 diffusers 这条路实际存下来的文本 embedding 是 226 ——  不影响 X-9 的结论（交叉注意力只占约 5% 算力）， 但「config 写的」与「实际跑的」在这里确实不是一个数。

五拆完怎么摆 —— 三段吃的不是同一种资源

切得动只是可行性。收益在于：三段吃的不是同一种资源。

谁吃算力、谁吃显存 —— 于是三段该摆在不同的地方 上一张说明切得动（腰只有 19.4 MB）；这一张说明为什么值得切 —— 三段的胃口根本不是一回事，塞进同一台机器必然有人吃不饱、有人撑着。 ① 文本编码 ② DiT 去噪（吃算力） ③ VAE 解码（吃显存） 问的是同一件事 ① 文本编码 ② DiT 去噪 ③ VAE 解码 旁注 它在干什么 把 prompt 编成向量 五十步去噪，全部算力在这 把 latent 展开成像素 吃的是哪种资源 几乎不吃／读一遍权重 吃算力 吃显存峰值 三种胃口完全不同 —— 这才是该分开的真正理由 正式运行占多少时间 3 秒 ／ 1.3% 230 秒 ／ 98.3% 1 秒 ／ 0.4% Wan2.1 720P，v6e-8 实测 编译（预热）要多久 — 110 秒 80 秒 → 之后每次 1 秒 VAE 是 80 倍差 —— 这个形状天生该做成常驻服务 峰值内存卡在哪 权重本身 权重 28 GB ＋ 激活 774 MB 展开成 448 MB 像素 当年 maxdiffusion 那版就是在第三段 OOM 于是摆在哪 CPU 或 TPU —— 先量再定 TPU 多芯片，TP／CP 分片 TPU 单芯片，常驻 第一格别写死，见下方 Flux.2 反例 方式一：端到端一体化 generate_torchax.py · 一台 v6e-8，一个进程 验证 / 演示 / benchmark ① 文本编码 T5 权重驻留 算完就闲着，但内存一直占着 ② DiT 去噪 28 GB 权重 ＋ 激活 全程只有这一段真的在忙 ③ VAE 解码 峰值展开 448 MB 跑 1 秒，编译 80 秒每次重付 三种胃口共用一份资源—— 按最馋的那一段配机器，另外两段的钱就白付了 方式二：三段分开部署 stage1 / stage2 / stage3 生产 / 多机 / 异构 ① CPU 节点（或小 TPU） 算一次，换 50 个 seed 反复用 前提是它在 CPU 上够快 —— 见下方反例 7.4 MB embedding ⏷ ② TPU DiT 节点 · v6e-8 8 颗做 TP：40 头 ÷ 8 ＝ 5 头每颗 算力全压这儿，扩容也只扩这一层 19.4 MB latent ⏷ ③ TPU VAE 服务 · 1 颗常驻 编译好放着，谁要解码谁来调 80 秒编译摊到成千上万次调用 ≈ 0 三段各配各的机器—— 而把它们串起来的成本，是那两根绿线上的几十兆 「text encoder 放 CPU」不是普适结论 —— Flux.2 就是反例 Wan2.1 用 T5，那一段 3 秒，放哪都行。但 Flux.2 用 Mistral3，放 CPU 要 30 秒 —— 而它的 DiT 在 TPU 上跑完 50 步只要 13.5 秒。 也就是说 Flux.2 那个「轻量」的第一段，反而是全程最慢的一段。照着「编码器放 CPU」的顺口结论摆，等于把整条链的瓶颈从 TPU 挪到了 CPU 上。 「一个节点能放几路」—— 先看权重塞不塞得下，再看它够不够忙 一台 v6e-8 是 8 颗 × 32 GB ＝ 256 GB。Wan2.1 的 28 GB 权重单颗只剩 4 GB 余量，所以走 TP 摊到 8 颗，每颗 3.5 GB，宽裕得多。 反过来，SDXL 只有 7 GB，一颗就装得下 —— 于是 v6e-8 上的正确用法不是切一个模型，而是开 8 路各生成各的（我们实测 2.40 img/s）。
图 X-11　上半是三段的资源画像，下半是两种部署拓扑。⭐ 左边一体化那栏的问题不是「慢」，是按最馋的那一段配机器，另外两段的钱就白付了。
⭐ 三段的胃口各一条，外加一个当场推翻顺口结论的反例

① 文本编码：几乎不吃算力，占全程 1.3%。 ② DiT 去噪：吃算力，占 98.3%。 ③ VAE 解码：吃显存峰值，把 480 万个数展开成 2.24 亿个，只占 0.4%。

⭐ VAE 还有个特别的形状：预热 80 秒、之后每次 1 秒 —— 80 倍差。 这种东西天生该做成常驻服务：编译好放着，那 80 秒摊到上万次调用等于零。

⛔ 反例：看完这张图最容易得出「文本编码那么轻，放 CPU 就行」—— 在 Wan2.1 上成立（T5，3 秒）。但 Flux.2 立刻推翻它： 它用 Mistral3，放 CPU 要 30 秒，而它的 DiT 在 TPU 上跑完 50 步只要 13.5 秒 ——  那个「轻量」的第一段，反而是全程最慢的一段。

⭐ 判据不是「哪一段天生该放哪」，而是先量一量它在目标硬件上要多久。 分段的价值恰恰在这里：拆开之后，每一段的账才第一次能单独算清楚。

六这套说法，我们自己验过

前面是应该怎样，这一页是实际怎样 —— 十个模型，五个月。

五个月，十个模型 —— 而三代方法的分界，写在文件名里 每一行的起止和提交数都是 git 提交历史直接数出来的，不是凭印象写的 —— 印象里最早做的是 Wan，而历史说是 HunyuanVideo，早了整整一周。 视频生成 图像生成 超分 / 修复 交付形态 第三代 · torchax ＋ diffusers-tpu 第二代 · 手写 Flax / 纯 JAX 移植方法 第一代 · 官方 maxdiffusion —— 起点，无独立区间 2025 年 12 月 2026 年 1 月 2 月 3 月 4 月 HunyuanVideo-1.5 文生视频 · 8.3B 73 次提交 第一个吃螃蟹的，提交数全场最多 Wan 2.1 T2V 文生视频 · 14B 44 次提交 720P 从 OOM 到 229 秒 CogVideoX 1.5 文生视频 · 5B 49 次提交 720P 106 秒 Wan 2.2 I2V 图生视频 · 27B/14B 26 次提交 MoE，两颗起 Flux.2 文生图 · — 11 次提交 一上来就是三代打法，不走弯路 ComfyUI on TPU 图形界面 · 4 个节点 28 次提交 给不写代码的人用 SDXL 文生图 · 3.5B 9 次提交 补上最经典的 UNet，单颗 3 秒 S3Diff 单步超分 · 3.3B 11 次提交 编译后 5.5 倍 Flux.1 文生图 · 12B 2 次提交 只交一体化版 Real-ESRGAN 超分 · 8.8M 3 次提交 纯卷积网 —— 验证非 Transformer 路径 ▲ 三阶段首次出现 ▲ flax 后缀集体删除 —— 第三代从这天算起 ▲ 优化完全指南成稿 读这张图的三条线索 第一代没有独立的时间段：maxdiffusion 是起点不是阶段 —— 12 月 3 日第一天就已经在写 Flax 版了。图上只给它一个起点标记，不编一段区间出来。 ① 前两周极密：12 月 3 日到 14 日，四个视频模型全部开工，同时试完了三条移植路线 —— 代价最大的探索集中在最前面。 提交数掉下去，才是这套方法真正立住的证据 HunyuanVideo 73 次、CogVideoX 49 次、Wan2.1 44 次；而 SDXL 9 次、Real-ESRGAN 3 次、Flux.1 2 次。 差别不在模型难度 —— SDXL 和 Flux.1 都不比 CogVideoX 简单。差别在于前面几个是在「发明方法」，后面几个是在「套用方法」。
图 X-12　每一行的起止与提交数都是 git 提交历史直接数出来的，不是凭印象写的。⭐ 12 月 10 日那道红线：那天手写 Flax 的文件被删除，第三代从那一刻算起。
⭐⭐ 真正的落点不是「我们做了十个」，是提交数从 73 掉到 2

HunyuanVideo 73 次 → CogVideoX 49 → Wan2.1 44 → SDXL 9 → Real-ESRGAN 3 → Flux.1 2。

差别不在模型难度 —— SDXL 和 Flux.1 都不比 CogVideoX 简单。 差别在于前面几个是在发明方法，后面几个是在套用方法。
—— 所以落点是：接第十一个模型的成本，已经不是前十个的量级了。

⭐ 这十个盖了五种架构，不是凑数： DiT（Wan2.1 / CogVideoX / HunyuanVideo）· MMDiT（Flux.1 / Flux.2）· UNet（SDXL / S3Diff）· MoE（Wan2.2 I2V）· 纯卷积（Real-ESRGAN，8.8 M）。 最后那个是唯一非 Transformer、非扩散的 —— 放它进来就是看这套框架 在完全不同的架构上还成不成立。答案是成立的。

⚠️ 提交数只反映改动次数，不等于工作量或难度 —— 这里只用它做「发明 vs 套用」的量级判断。

收尾带走三句

① 先看你的活落在轴的哪一边 —— 跑之前就能判。
② 切口只有 19 MB 出头，所以三段能拆开、各配各的资源。
③ 十个模型五个月，提交数从 73 掉到 2。

想往下挖：这一讲的每一节在 L200 的哪儿（含本讲没讲的 SparseCore 与环面）
本讲L200 精讲版那边多给了什么
§一 一条线 ＋ 两颗芯片拆开§一 ＋ §二点一、二点二 560 怎么除出来的、989.5 为什么不是 1979、拿 v5p 复现公式；两颗芯片的逐项对照表
§二 你的活落在哪一边§三 ＋ §四 FLOP 全表、window_size 的反事实、锚点当场抓到的一个错、MFU 37% 的原因
§三 一颗装得下吗（主线可跳）§二点三 那三条「减配」为什么是配套的
§四 两份例子 ／ 三步自检§五点一 切口成本的几笔账、三个落盘产物的完整口径
§五 拆完怎么摆§五点二 ～ §五点四 资源画像逐项、一个节点摆几路
§六 我们真的跑过§六 三代移植路线的逐项对照、十个模型各自验了哪一条断言
本讲没讲的§二点四 SparseCore 的两个职责（它不只给推荐系统用）、二维环面与 Pod 为什么只有 256 颗

⚠️ 本讲没有端到端性能数：讲的是结构匹配，不是快慢 —— 后者要实测，而实测要连口径一起给。
想看扩散模型本身：专题十一。