https://henryhmko.github.io/posts/tpu/tpu.html

令人印象深刻的几点：

  1. TPU 是软硬件的协同优化，硬件上是 systolic array + 流水线，软件上是 XLA 编译器和 AOT。
  2. TPU 相比 GPU，有较小的 HBM 大小（比如 32GB）（不小了吧，这可是 TPUv4）和较大的缓存（VMEM）。从 TPUv4 的结构来看，TPU 的 TensorCore 也是包含向量（各种非矩阵的 elementwise）和标量计算（比如地址解析、控制流判断）能力的。GPU 上这些能力不属于 TensorCore，而在 TPU 中属于。个人感觉 TPU TensorCore 可能相当于 GPU 的 SM。
  3. TPU 有 tray、rack、pod 和 slice 的概念。
    1. Tray 是包含 4 个芯片的芯片组
    2. Rack 是 4x4x4 的芯片组（也叫 cube）
    3. Pod 是按照最大连接数连接的芯片组（TPUv4 可以包含 4096 个芯片，TPUv7 "Ironwood” 甚至可以达到 9216 个芯片）
    4. Slice 泛指一组芯片，没有固定的大小
  4. TPU 有惊人的可扩展性，通过 ICI: Inter-Core Interconnect 和 OCS: Optical Circuit Switching 连接。前者用来 tray 内连接 4 个芯片，后者用来连接 tray。OCS 使得 TPU 拓扑结构可以灵活变化。

不过从这个图来看，TPU 的显存其实并不小，甚至比 GPU 更大。

https://blog.google/products/google-cloud/ironwood-tpu-age-of-inference/

image.png