TPU 显微镜 —— 一颗 v7 从封装拆到 lane

0读之前：三件必须先说清的事

口径、来源等级、以及这份文档跟原来那份 embedding 专题的关系。

1先看全景：一颗 v7 chip 里有些什么

封装 → 两个 die → TensorCore ＋ SparseCore → HBM → 六个对外出口。

这张图最值得停下来的一处，是它和 B200 正好相反的那个选择。

两家都在同一个物理事实面前：一颗 die 做不到想要的算力，只能上双 die 封装； 而双 die 之间的通路必然比 die 内部慢。区别只在于把这条缝交给谁。

NVIDIA 用 NV-HBI 这条一致性总线把两个 die 缝起来，对软件装成一个 GPU。 好处是你什么都不用改；代价是缝还在 —— L2 跨 die 访问从约 21 TB/s 掉到 16.8 TB/s， 而这件事不会出现在任何 API 里，只会出现在你的性能曲线里。

Google 选了另一头：如实暴露成两个 device。代价是你必须自己决定怎么切分， 好处是缝在明处 —— 一旦你写下了分片策略，跨不跨 die 就是你自己写的， 不会有一条看不见的慢路径在背后偷走性能。

这不是谁对谁错，是两种成本转移。但它解释了后面七张图里几乎所有的设计差异： 当你决定把并行性交给软件之后，硬件里那些专门用来「自动藏起复杂度」的部件 —— 调度器、乱序、大寄存器堆、多级缓存 —— 就都可以省掉。省下的面积去了哪里， 就是下一张图的内容。

2把一个 TensorCore 拆开 —— 省掉的部件比留下的更说明问题

2 个 MXU、1 个 VPU、标量单元、VMEM、SMEM，一个不落地画出来。

数一遍留下的：2 个 MXU、1 个 VPU、2 个跨 lane 单元、一个标量单元、 64 MiB VMEM、1 MiB SMEM。就这些。

再数一遍没有的：没有 warp 调度器，没有几十份并发上下文，没有 256 KB 那种规模的寄存器堆， 没有会自动填充的 L1，没有记分板和乱序发射。这五样在 GPU 的 SM 里加起来占了相当可观的一片硅， 而它们的共同职能只有一个 —— 在运行时动态地藏住延迟和分支。

省下的面积去了哪里，图上给了一个直接的对照：一颗 TPU v7 chip 上的 VMEM 是 128 MiB （两个核各 64 MiB），而一整颗 B200 的共享内存合计约 34 MB。 差别不只在容量 —— 更在于那 128 MiB 的每一个字节都由编译器显式安排， 不是硬件按访问历史猜着填的。§6 会把这件事讲透。

还有一处容易被略过：标量单元在这颗芯片上比你以为的重要得多。 它不参与矩阵乘，但 DMA 的描述符是它写出来的 —— 也就是说「什么时候搬、搬多少、搬到哪」 这件事，是标量指令流里的一条指令，不是硬件后台自动发生的行为。 一颗以矩阵乘为业的芯片，把「安排搬运」这件事放在了标量单元上， 这个安排在 §6 和 §7 会各出现一次。

3并行层级：lane / sublane 分别对应硬件上的什么

和 GPU 的 thread / warp / block / grid 逐层对照 —— 层数少得多，少掉的那几层正是关键。

七个问题里，TPU 有两个答不出来：「共享一块暂存的是哪一组」和 「运行时谁决定接下来跑哪个」。

这两行不是随便哪两行。GPU 的 thread block 是一个运行时概念 —— 有几个 block、落在哪个 SM 上，启动的时候才知道；而 warp 调度器则是 延迟出现之后才发挥作用的东西。它们俩合起来，就是 GPU 藏延迟的全部机制。 TPU 把这两层都拿掉了，于是也就把「运行时还能补救」这条路一起拿掉了。

顺带解释了一件常被问到的事：为什么在 TPU 上「跑一个不规则的算法」这么别扭。 不是编译器不肯，是硬件的层级里根本没有一个能承载「不规则」的单位 —— 最小的可指名单位是「向量里的一格」，而它没有程序计数器，也不能走自己的分支。 你能表达的最小的「不同」，是让数据的形状不同，不是让某一格的行为不同。

4MXU：256×256 到底意味着什么

脉动阵列里数是怎么流的，以及那个「每 cell 每周期几个乘加」的常数 —— 这一节还附带一个我自己犯过的错。

先说结论，因为它反直觉：K 比 256 大一点不亏，K 比 256 小才是真亏。

阵列的收缩边是物理的 256 行。K ＝ 1,024 就是走 4 趟，累加器一直不落地， 效率跟一趟算完几乎一样；而 K ＝ 128 时，有一半的行在空转， 并且 —— 这才是关键 —— 没有办法把别的活塞进那一半。 回到 §3：TPU 没有「换一个跑」这一层，所以空着就是空着。

插一段：我在这个常数上错了两次，而第二次错得更值得讲

那条峰值公式 —— 262,144 FLOP / 周期 / MXU × 2.2 GHz × 4 个 MXU ＝ 2,307 TFLOP/s —— 是 Google 工程博客逐字给出的。但在找到这条公式之前，我先后写下过两个版本， 两个都是错的，而第二个错误比第一个危险得多。

v2 那次的错法值得单独说。我用的方法本身是对的：找一代四个变量全公开的老芯片， 拿它反算公式，验证通过了再套到新芯片上。问题在于 v6e 那次反算里， 我算的是「4 个 MXU × 256×256 × 每 cell 1 次 @ 1.75 GHz ＝ 917.5 TFLOP/s ✓」—— 数字完美对上了官方的 918。

但官方文档白纸黑字写着「每个 TensorCore 有 2 个 MXU」，而 v6e 是 1 个 TensorCore/chip， 所以是 2 个 MXU，不是 4 个。正确的算式是 「2 个 MXU × 262,144 FLOP/周期 × 1.75 GHz ＝ 918 ✓」—— 同样对得上。

还有一件事图上写了、这里再强调一遍：「每 cell 2 次乘加」在硅上到底怎么实现， 公开资料答不了。是每个 cell 里真放了两个乘法器，还是 256×256 只是一个逻辑视图、 底下另有排布？我不知道，所以本文不猜。这一条会出现在 §10。

5SparseCore：第二种核，以及它到底在干什么

原来那份 embedding 专题的核心结论，压缩成一节 —— 包括那个最反直觉的一条。

关于 SparseCore，流传最广的一句话是「TPU 上有个专门查 embedding 表的核」。 这句话不算错，但它会让你在看自己的 profile 的时候彻底看错。

先把结构说清楚：一颗 chip 上有 4 个 SparseCore（每个 device 2 个）， 每个由 1 个标量子核 ＋ 16 个向量子核组成，lane 宽 16，私有 SRAM 512 KiB。 它没有 MXU，做不了矩阵乘。它能做的是 cumsum、sort_key_val、 fetch_and_add、addupdate_scatter 这一类 不规则访存 ＋ 归约的活 —— 这些在开源的 Pallas SparseCore 接口里能直接看到。

那生产里的 embedding 到底怎么跑的？

MaxText 里有一个开关叫 use_iota_embed。打开时，「查表」被写成 「把 token id 展成 one-hot，再和整张词表做一次矩阵乘」 —— 一条 dot，跑在 MXU 上。关掉才是真的 gather。

上游的默认值是关的（configs/base.yml 里写 false）， 但仓库里 34 份配置显式打开、0 份显式关闭 —— 连给 GPU 用的那几份也是打开的。 所以判断某一次跑走了哪条路，只能看那一次实际生效的配置，不能看模型的 yml。

至于 SparseCore 为什么没被派去干这件事：看的是同一批里「重复取同一行」的程度。 推荐系统那边少量热行被反复命中，重复率极高，那是 SparseCore 的主场； 而语言模型的取行重复度不够高 —— 注意不是「没有重复」： 一条十几万 token 的序列里不同的词只有几千个，绝大部分命中最高频那一两千行。 重复确实有，只是跟搜广推差着量级，落在最不划算的那一端， 编译器于是根本没把它派过去。

所以这一节真正的收获不是「SparseCore 是什么」，是「不要用宣传口径去解释自己的 profile」。 那颗核确实存在、确实为稀疏访问而生，但在你手上这个任务里它可能在干完全不同的事 （比如集合通信卸载），甚至什么都没干。看 trace，不要看宣传页。

6一个数走完全程：HBM → VMEM → MXU → 累加器

每一站的容量、谁负责搬、以及「暂存」和「缓存」到底差在哪。

几乎每份 TPU 材料都会列容量和带宽。但真正决定你写代码时会撞上什么的， 是图上那一整行 「谁决定搬」。

把这张小表念一遍就明白：TPU 上不存在「缓存没命中」这种性能问题。 不是因为它命中率高，是因为根本没有「命中」这个事件 —— 数据要么按班表到了，要么没到。 没到就是编译器排错了，而排错了没有兜底。

再说一处很多人第一次看会愣一下的：搬运不是后台自动发生的。 HBM ↔ VMEM 的每一次搬运都由 DMA 完成，而 DMA 的描述符是标量单元写出来的。 也就是说「什么时候搬、搬多少、搬到哪」跟乘加指令一样，占着同一条指令流里的槽位。 这解释了 §2 里那个看起来很怪的设计：为什么标量单元在一颗以矩阵乘为业的芯片上还这么重要 —— 它不算数，它安排搬运。

带宽落差是这条通路的真正约束：片上暂存的读写带宽比 HBM 高约一个数量级（具体数值官方未公开，这里只给量级）。所以「尽量让数在暂存里多待一会儿」不是风格建议 —— 它是绝大多数 TPU 调优工作的实际内容。

最后一站顺手给一个能自己验的推导：累加器是 128 个、形状 (8, 256)、32 bit， 所以 128 × 8 × 256 × 4 B ＝ 1,048,576 B，正好 1 MiB / MXU。 这类能对上整数的推导值得多做几次 —— 它是检查自己有没有把口径搞错的最便宜的办法。

7延迟被藏到哪里去了：VLIW 与编译期排班

GPU 靠切换 warp 藏延迟，TPU 靠编译器提前排好 —— 代价是排错了没有兜底。

VLIW 的意思很朴素：一拍要干的所有事，打成一包。包里给标量、向量、矩阵、 杂项各留了固定的槽，谁往哪个槽里填、在第几拍填 —— 全部写死在指令流里。 没有乱序，没有记分板，没有「运行时再看」。

看甘特图的时候要盯两处。第一处是第 1–6 拍的向量槽：那里塞的是 上一块的后处理，和这次搬运没有依赖关系，所以能被提前挪过来盖住 DMA 的延迟。 这就是编译器「藏延迟」的全部手段 —— 找独立的活。 第二处是第 14 拍之后：独立的活用完了，所有槽都空着， 而且不会有任何东西自动顶上来，因为没有第二个线程可切。

这个取舍在日常里长成三件事，而这三件事其实是同一件：

1. 形状一不规则，性能就掉得很难看。batch 小、序列不齐、专家路由不均 —— 这几种情况的共同点都是「后一步紧跟着前一步」，空拍没东西填。
2. 编译很慢，而且慢得有道理。v7 上一次 XLA 编译常见 10–17 分钟。 它不是在「翻译」，是在替硬件把整张班表排完；GPU 那边这件事是每次运行时由调度器现做的， 成本摊在了跑的时候。
3. 但跑起来非常稳。同一个形状重复跑，步时几乎没有抖动 —— 因为根本没有运行时决策可抖。 可预测性是这个设计买来的东西，不是附带效果：容量规划和性能回归检测在 TPU 上都因此简单不少。

8从一颗到一个 pod —— 全文的落点

ICI、3D 环面、cube、slice：为什么说 scale-out 是这颗芯片的第一性设计。

每颗 v7 有 6 条 ICI 物理链路，对应三维的正负方向，接成 3D 环面。 「环面」而不是「网格」的差别是实打实的：4×4×4 如果只是网格，最远要走 3+3+3 ＝ 9 跳； 接成环面之后是 2+2+2 ＝ 6 跳。规模越大差得越多，而这直接决定了 all-reduce 的最坏时延。

但环面不是白拿的。它要求切片在物理上必须是连续的一块立方体 —— 所以在 TPU 上你申请的不是「64 颗芯片」，是「一个 4×4×4」。 形状本身是调度的一部分，这一点在 §3 那张表里就已经埋下了伏笔。

两份文档合起来的那一句话

《GPU 显微镜》的最后一张图讲的是：一颗 B200 里有 592 个 Tensor Core， 而一颗 TPU v7 里只有 4 个 MXU —— 同样一块矩阵乘的活，两边被切成的份数差 128 倍。 GPU 的协调主要发生在芯片内部。

这张图讲的是另一半：TPU 一颗芯片里只有两个核要协调， 但不换协议能一路连到 9,216 颗，而 GPU 在 72 颗上就得换。 TPU 的协调主要发生在芯片之间。

两个「128 倍」出现在完全不同的位置，这不是巧合，是同一个设计取向的两个侧面： 一边把复杂度收进一颗芯片里由硬件消化，另一边把复杂度摊到芯片之间由软件面对。 八张图从头到尾说的都是这同一件事 —— 不是谁更强，是把同一份复杂度放在了不同的地方。

9来源等级总表

逐条列出每个数字来自哪里 —— 这样你可以只信你愿意信的那几行。

四级的划分标准写在 §0。这里只补一句：「第三方」不等于「不可信」， 但它确实被降级了 —— 因为本文实测到其中至少一条是错的（那本流传很广的 TPU 教科书式博客 把 256×256 说成 131,072 FLOP/周期，比官方的 262,144 少一半，本文不采用）。

官方 —— 可直接引用

第三方 —— 一律标灰

本文推导 —— 每条都给推导链

本文实测 —— 我们自己在 v7 上跑出来的

下面这几条不是从任何文档里查来的，是在 v7 上跑开源模型量出来的。 所以它们既不属于「官方」也不属于「第三方」—— 判断它们可不可信， 唯一的依据是这份材料自己说了实验是怎么做的。

10还没查实的

查不到就写查不到 —— 这一节是这份文档最该被信任的部分。

下面每一条都是我试着查了、没查到的。它们没有出现在前面任何一张图的数字里， 出现的地方一律画成灰色虚线或直接写「查不到」。

一、公开资料里找不到的

二、官方资料自己对不上的

这类比「查不到」更麻烦 —— 它会让你以为自己查到了。