第 12 讲 · TPU 与 JAX · TPU and JAX: From Python Code to Silicon

对应 Maps to：CMU 11-868 / 11-968 Large Language Model Systems（Spring 2026，Lei Li），Lecture 12 · Mon 2/23 · Guest Lecture by Srinath Mandalapu（Google）: TPU and JAX（slides：llmsys-12-Introduction_to_JAX_XLA_TPU.pdf，107 页，标题 Decoding the JAX AI Stack: From Python Code to Silicon）。 阅读材料：JAX 文档（jax.readthedocs.io）中的 “Thinking in JAX”、“JAX internals: the jaxpr language”、“Distributed arrays and automatic parallelization”、“shard_map”；How to Scale Your Model（jax-ml.github.io/scaling-book）的 TPU 与分片两章；XLA 与 StableHLO 的文档。 依据说明：本讲按嘉宾 slides 的四部分组织：JAX AI 技术栈（JAX、Flax NNX、Optax、Orbax、Grain；用 JAX 训练 GPT-2 的配置、网格、jit 训练步、检查点、XProf）、TPU 架构（Ironwood 的规格与 3D 环面、张量核心的 HBM / VMEM、向量寄存器与打包、VPU、两种归约、XLU、MXU 脉动阵列与 3N − 1、指令流）、XLA 编译（jit 的生命周期、追踪与 jaxpr、StableHLO、HLO、布局与分块、编译器的各个 pass、注意力的三个融合块与重算、存储空间与异步拷贝、LLO 与 VLIW 指令包、MXU 上的分块、exp 的改写、实测结果）、分片（Mesh、PartitionSpec、GSPMD 与 shard_map、SPMD）。迷你 JAX（追踪、jaxpr、grad、vmap、jit 缓存）、迷你 XLA（化简、CSE、DCE、融合与重算、代价分析）、逐拍的脉动阵列模拟、填充与布局的计算、bf16 累加实验、分片与集合通信的模拟，都是本讲义自己设计和实现的，不是 JAX / XLA 的真实代码；本机没有安装 JAX，所有数字来自这些迷你实现或 slides 的规格。 配套代码 Code：code/lecture12/（6 个脚本 + minijax.py、xla.py、tpu.py、sharding.py；都在 1 秒左右） 预计阅读时间：3–4 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 说出 JAX 的核心思想——可组合的函数变换（jit、grad、vmap），以及它们共同的基础：把 Python 函数追踪成 jaxpr 这种中间表示。
2. 解释 jit 的缓存（按形状、dtype、分片）、副作用只在追踪时发生、数据相关的 if 为什么不能出现在 jit 里，以及两种改法。
3. 写出从 Python 到 TPU 机器码的编译链：jaxpr → StableHLO → HLO → 优化后的 HLO → LLO → VLIW 指令包，说出每一层的作用。
4. 用”内核数、读写字节、超越函数次数”分析 XLA 的融合结果；解释注意力为什么被编译成三个融合块，“重算”省了什么、花了什么。
5. 描述 TPU 张量核心的结构：HBM 与 VMEM、8 × 128 的向量寄存器、VPU、XLU、256 × 256 的 MXU；算出脉动阵列的周期数 B + 2N − 1 和利用率，解释为什么 TPU 需要大批量、偏爱 128 的倍数。
6. 用 Mesh + PartitionSpec 描述并行；说出给定切法时编译器要插入的集合通信，以及 jit（GSPMD）和 shard_map 的区别。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  背景：另一条技术路线；JAX AI 技术栈；用 JAX 训练 GPT-2 的样子
§2  JAX：追踪与 jaxpr；jit 的缓存与"怪现象"；grad 与 vmap 是解释 jaxpr 的另一种方式   ← 代码 01、02
§3  TPU：Ironwood 与 pod；张量核心里的 HBM / VMEM、VPU、XLU、MXU                    ← 代码 04、05
§4  XLA：从 jaxpr 到 VLIW 指令包；融合、重算、布局、异步拷贝；注意力的三个融合块      ← 代码 03、05
§5  分片：Mesh、PartitionSpec、GSPMD 与 shard_map；GPT-2 的两种网格                  ← 代码 06
§6 小结   §7 常见误区   §8 练习   §9 延伸阅读   §10 下一讲预告

1. 背景：另一条路线 · Background

1.1 为什么要看 TPU 和 JAX

前 11 讲走的是”GPU + CUDA + PyTorch”这条路：框架一个运算一个运算地急切执行，性能靠人写融合内核（第 10–11 讲的 LightSeq）。这一讲换一条路线，它在三个层次上都不同：

两条路线在不断互相借鉴：PyTorch 有了 torch.compile（先捕获图再编译）、torch.func（vmap、grad）；JAX 有了 Pallas（手写内核，第 13 讲）。理解两边，才能看懂”框架—编译器—硬件”之间的分工。

1.2 JAX AI 技术栈

slides 第 5–9 页：

• JAX：NumPy 风格的数组 API（jax.numpy）加上可组合的函数变换 jit、grad、vmap；通过 XLA 编译，同一份代码可以跑在 CPU、GPU、TPU 上；从一个设备扩展到整个 TPU pod，用户只需描述数据怎么切，并行和通信交给 XLA。
• Flax NNX：神经网络库，用普通的 Python 类（nnx.Module，__init__ 里定义层、__call__ 里写计算），模块本身是 JAX 的 pytree，可以直接被 jit / grad / vmap。
• Optax：优化器和梯度处理；nnx.Optimizer 把 Optax 的状态和模型参数绑在一起。
• Orbax（检查点）和 Grain（数据加载）。

1.3 用 JAX 训练 GPT-2 的样子

slides 第 10–20 页的例子：GPT-2 式的模型，24 个 Transformer 块、序列 1024、d = 1024、16 个头、前馈 4096、全局批 32、dropout 0.1、学习率 5e-4、权重衰减 0.1，结构是 pre-LN、因果注意力、GELU 的 MLP，输出层与 token 嵌入共享权重（第 6–7 讲）。几个关键写法：

• 网格：Mesh(create_device_mesh((8, 1)), ('batch', 'model'))——8 个设备排成 8 × 1，两个维度分别叫 batch 和 model。换成 (4, 2) 就是”数据并行 4 路 × 张量并行 2 路”（§5）。
• 参数的切法：nnx.with_partitioning(init, P(None, 'model')) 在创建参数时就标好”第 1 维沿 model 切”。
• 一个 jit 的训练步：@nnx.jit def train_step(model, optimizer, metrics, batch) 里调用 nnx.value_and_grad(loss_fn) 和 optimizer.update(grads)——前向、反向、优化器更新被编译成一个 XLA 图；因为参数带着分片标注，jit 自动插入梯度的 all-reduce。
• 数据：jax.device_put(batch, NamedSharding(mesh, P('batch', None))) 把一批 32 个样本分给 8 个设备，每个 4 个。
• 检查点：Orbax 异步地把分片的状态写到磁盘；恢复时先用 nnx.eval_shape 得到”不分配显存的蓝图”，再按蓝图加载。
• XProf（slides 第 20 页）：在时间线上能看到一个 Transformer 块的前向、反向，以及最后的 all-reduce。

后面几节把这几行代码背后的东西一层层拆开。

2. JAX：追踪与函数变换 · Tracing and transformations

2.1 追踪：用抽象值把函数执行一遍

slides 第 53–55 页。jax.make_jaxpr(f)(x) 并不真的计算 f：它用抽象值（tracer：只有形状和类型、没有数据）把 Python 函数执行一遍，每遇到一个基本运算就记下来，得到一个 jaxpr。01 号脚本第 1 部分（minijax.py 的迷你实现）：

def f(x, w):  h = tanh(x @ w);  return sum(h * h, axis=-1)

{ lambda ; a:f32[4,3] b:f32[3,2]. let
    c:f32[4,2] = dot a b
    d:f32[4,2] = tanh c
    e:f32[4,2] = mul d d
    f:f32[4] = reduce_sum[axis=-1 keepdims=False] e
  in (f,) }

追踪时函数里的 x 是 Traced<f32[4,3]>。jaxpr 只有”变量 = 基本运算(变量)“的方程，每个变量标明类型和形状——没有 Python 的循环、对象、函数调用。这正是编译器和函数变换需要的形式：

• 为什么要一个框架自己的 IR（slides 第 55 页）：它是 grad、vmap 等变换的稳定目标。每个基本运算（如 add_p）登记几样东西：怎样降到 StableHLO、在 grad 下怎样求导、在 vmap 下怎样批处理。
• 副作用只在追踪时执行一次（slides 第 53 页）：print、往列表里追加之类的 Python 副作用，只在追踪时发生，不进入编译后的图。

2.2 jit 的缓存与几个”怪现象”

slides 第 49 页：jax.jit(f) 返回一个 PjitFunction。每次调用时，它把参数的形状、dtype、分片做哈希，在编译缓存里查找：命中就直接执行编译好的程序；没命中才追踪、编译、存入缓存。01 号脚本第 2 部分：5 次调用（同形状不同值、新形状、新 dtype、再回到第一种）只追踪了 3 次，函数里的副作用也只记录了 3 次。由此有几个实践上的推论：

1. 在 jit 里 print 只会打印一次，而且打印出来的是 tracer。（要看运行时的值，用 jax.debug.print。）
2. 形状一直变的输入会让 jit 反复编译：不定长的句子要填充到少数几档长度。
3. 数据相关的 if 不能写：追踪时 x > 0 的结果也是 tracer，没有真假。01 号脚本第 3 部分里 if x > 0 抛出 ConcretizationError。两种改法：用 where 把两边都算出来、逐元素选择（把控制流变成数据流）；或者当条件只取决于”配置”（字符串、层数）时，把它声明成静态参数——它进入缓存的键，每个取值各编译一份。真 JAX 还有 lax.cond、lax.scan 等控制流原语，把分支和循环本身记进 jaxpr。
4. 异步执行（slides 第 49 页）：jit 的调用把程序发给 TPU 就返回一个 jax.Array（相当于 future），数据不会自动拷回 CPU，直到你显式地取值（如 .block_until_ready()）。

2.3 grad：在 jaxpr 上反向走一遍

grad 是”换一种方式解释 jaxpr”：先追踪 f 得到前向的 jaxpr，然后从输出往回走，对每个方程套用它的 VJP 规则（第 5 讲的反向模式自动微分）。规则本身调用的又是同样的基本运算，所以 grad(f) 也是一个可以被追踪的函数。02 号脚本第 1 部分：f(x) = Σ tanh(x)·x 的 grad 被追踪成 9 个方程，前 3 个是前向，后面是反向（tanh 的规则 ct·(1 − out²)、mul 的规则 (ct·y, ct·x)、reduce_sum 的规则把 ct 广播回去）。正因如此，grad(grad(f)) 也能算二阶导数，jit(grad(f)) 能把前向和反向一起编译。

第 2 部分把两层 MLP 的交叉熵与 PyTorch autograd 对比：损失和三个梯度都一致到 2.2e-16。两者的算法相同，区别在于什么时候决定计算图：PyTorch 每次前向现场记一张图（动态图），JAX 先追踪成 jaxpr 再变换（静态图，第 5 讲 §6），后者让编译器看得到”前向 + 反向 + 更新”的全部。

2.4 vmap：自动批处理

vmap 同样是解释 jaxpr：每个基本运算登记一条批处理规则。02 号脚本第 3 部分：只为一个样本写的 predict(W, x)（x: f32[4]），vmap 之后的 jaxpr 里 dot 的输入变成了 f32[6,4]——逐元素运算靠广播，归约的 axis 往后挪一位（−1 不用改），矩阵乘用批维度广播。结果是”一次处理一批”的程序，而不是 6 次循环，正是 GPU / TPU 想要的形状。

2.5 组合：jit(vmap(grad(…)))

02 号脚本第 4 部分：一行 jit(vmap(grad(loss))) 得到 16 个样本各自的梯度（差分隐私训练 DP-SGD 要先把每个样本的梯度裁剪到固定范数再平均），与 Python 循环逐位相同，平均值就是整批的梯度。三层变换自动生成了 48 个方程，其中不少是用不上的（例如对常量求导）；编译器清理后剩 31 个，再融合成 8 个内核。“变换生成冗余、编译器负责清理”是 JAX 的分工：变换的规则可以写得简单直接，性能交给 XLA。

3. TPU 硬件 · TPU architecture

3.1 Ironwood 与 pod

slides 第 22–26 页（05 号脚本第 1 部分）：

每个 Ironwood 芯片有两个张量核心（TensorCore）和四个 SparseCore（处理嵌入查表这类稀疏访问）。立方体之间用光路交换（OCS）连接成 pod（256 片）或 superpod（9,216 片）：某个立方体坏了可以用光路绕过、换上备用的；也可以按需切出 4 × 4 × 4、8 × 8 × 8 这样的逻辑”切片”。

3.2 张量核心：存储层次

slides 第 27–29 页：数据在 HBM（片外，几十到上百 GiB）和 VMEM（片上的向量存储，约 0.1 GiB，紧挨着计算单元）之间搬运；还有 SMEM（标量存储）放循环变量、地址这类标量。计算单元只从 VMEM / 寄存器取数，所以每个运算都是”HBM → VMEM → 计算 → VMEM → HBM”的流水线，TPU 靠流水线把搬运和计算重叠起来。

按 slides 第 92 页每个张量核心的规格（峰值 1028.75 TFLOP/s、HBM 3433 GiB/s、VMEM 读 27180 GiB/s），roofline 的平衡点：从 HBM 读约 279 FLOP/字节，从 VMEM 读约 35。slides 第 28 页的说法是 VMEM 带宽高得多，强度 10–20 的算法也能接近峰值（数字取决于芯片和数据类型）。结论一样：能留在 VMEM 里反复用的数据，就别回 HBM——这是第 13 讲 Pallas 分块的出发点。

3.3 向量寄存器与 VPU

• 向量寄存器（slides 第 30–31 页）：64 个，每个是 8 个 sublane × 128 个 lane 的 32 位字。bf16 时两个元素压进一个字，等效 16 行。
• VPU（slides 第 32 页）：8 × 128 = 1024 个计算位置按同一条指令同步执行（2D SIMD），每个位置 4 个 ALU，一个周期可以同时做 4 条不同的向量指令（如 vadd、vsub、vmul、vmax）。峰值 1024 × 4 × 2 × 2.2 GHz = 18.0 TFLOP/s。
• 两种归约（slides 第 33 页）：sublane 归约在一个 lane 内把 8 个 sublane 加起来；cross-lane 归约在一个 sublane 内跨 128 个 lane 加起来。softmax 的行最大值、行求和就是这两种归约的组合。
• XLU（slides 第 34 页）：跨 lane 的数据移动要靠专门的跨通道单元：归约（二叉树折叠）、转置、置换、广播（128 × 128 的交叉开关）。Ironwood 每个张量核心有两个 XLU。

05 号脚本第 3 部分算了 8 × 128 分块的填充浪费：f32[1024, 128] 100% 有效；头维度 64 的 f32[1024, 64] 只有 50%；行数很少的 bf16[4, 1024]（解码时批 4）只有 25%；一行的 f32[1, 4096] 12.5%。这就是 TPU 上的模型偏爱 128 的倍数（头维度 128、d_model 是 128 的倍数）的原因。

3.4 MXU：脉动阵列

slides 第 35–45 页：MXU 是 256 × 256 = 65,536 个乘加单元组成的二维网格，乘法用 bf16、累加用 FP32。数据流是权重驻留（weight-stationary）的：权重 W 预先装进网格，每个 PE 存一个 W[i, j]；输入向量的各个分量错开一拍从顶上进入、每拍往下走一行；部分和从每一行的左边以 0 开始、每拍往右走一列，经过 PE(i, j) 时加上 W[i, j]·x[j]；最右边出来的就是 y[i]。数据在相邻的 PE 之间用短导线直接传递，不回存储器，所以省能量；每个输入在一行里被复用 256 次。

04 号脚本第 1 部分逐拍模拟了 slides 的 3 × 3 例子：

一共 8 = 3N − 1 个周期（slides 第 45 页的”填充 N 拍、排空”）。一般地，流过 B 个向量需要 B + 2N − 1 个周期，PE 利用率 B / (B + 2N − 1)——04 号脚本第 2 部分在 N = 16 时的逐拍模拟与公式完全一致。对 256 × 256 的 MXU：

填充和排空的 2N − 1 拍是固定成本，只有足够多的向量才能摊薄。这和第 9 讲”批 1 的解码受限于访存”是同一件事的两面，也说明了为什么推理服务要把很多请求拼成一批（第 22 讲）。

精度（04 号脚本第 4 部分）：64 个长度 K 的点积，bf16 输入 + FP32 累加的相对误差在 K = 16 到 4096 之间稳定在 2.5e-03 到 2.9e-03；累加也用 bf16 时从 6.2e-03 涨到 8.8e-02（K = 4096）。和越加越大，每加一项都要舍入到和的量级——所以 MXU 的累加器一定是 FP32。

3.5 指令流

slides 第 46 页：核心的定序器（sequencer）执行全部标量指令；向量指令排进队列，由 VPU 的 128 个 lane 同步执行；MXU 和 XLU 的数据都从 VPU 送进去。指令的组织方式是 VLIW（超长指令字，§4.6）：调度完全由编译器决定，硬件不做动态调度。

4. XLA：从 jaxpr 到机器码 · The XLA compiler

4.1 编译链

slides 第 49–51 页：

Python 函数 ──追踪──▶ jaxpr ──降级──▶ StableHLO ──▶ HLO ──优化──▶ 优化后的 HLO
    ──TPU 后端──▶ LLO（低层算子）──调度──▶ VLIW 指令包 ──▶ 可执行文件（加入缓存）

对应的 API：jax.make_jaxpr(f)(…)、jax.jit(f).lower(…)（.as_text() 是 StableHLO、.compiler_ir('hlo') 是 HLO）、.compile()、.cost_analysis()。

• StableHLO（slides 第 56 页）：与硬件无关、有规范和序列化格式、保证向后兼容的中间表示。它解决 N × M 问题：N 个框架（JAX、TensorFlow、PyTorch）× M 种硬件，各自都只需对接 StableHLO。
• HLO（slides 第 58 页）：XLA 的输入语言。有向无环图、没有顺序和递归，大多数指令是纯函数；所有数组的形状在编译时都已知，所以显存用量能在编译时完全确定；操作码很少（逐元素、Transpose / Dot 这类数组运算、AllGather / ReduceScatter 这类集合通信）。

4.2 布局与分块

slides 第 60 页：HLO 里的 f32[1024,512]{1,0:T(8,128)} 由三部分组成：元素类型和维度、minor-to-major（先列出变化最快的维度，{1,0} 就是行优先）、分块（T(8,128)：先一块一块地排，块内再排）。05 号脚本第 4 部分把 slides 的例子 f32[3,5]{1,0:T(2,2)} 展开：

        列0 列1 列2 列3 列4
行 0      0   1   4   5   8
行 1      2   3   6   7  10
行 2     12  13  16  17  20

3 × 5 的数组被补成 4 × 6、占 24 个位置，其中 9 个是填充。TPU 上的分块是 T(8,128)：一块正好是一个向量寄存器。

4.3 编译器的 pass

slides 第 61–62 页列了 XLA 的几类变换：代数化简（例如把除以常数变成乘以倒数）、常量折叠、死代码消除；reshape 移动（把布局变化推到图的边缘，让数学运算挨在一起好融合）；布局与分块；融合；缓冲区分配与拷贝插入；调度与延迟隐藏；存储空间分配（决定张量最终放在 VMEM 还是别处）。

01 号脚本第 4 部分用迷你 XLA（xla.py）在一个故意写了冗余的 GELU 层上跑了一遍：追踪得到 25 个方程；除以常数 → 乘倒数 2 处、加 0 去掉 1 处、公共子表达式 9 处、死代码 2 处，剩 13 个；再融合成 1 个内核（矩阵乘法把加偏置和整个 GELU 吸进输出端），读写从 1953 KiB 降到 225 KiB。注意编译器不会把 x·0 化简成 0：浮点里 inf·0 = NaN。

4.4 融合：注意力的三个融合块

slides 第 48–73 页拿最简单的注意力做例子（单头、不缩放，q、k、v: f32[1024, 512]，mask: bool[1024, 1024]）：

def attention(mask, q, k, v):
    s = where(mask, q @ k.T, -1e38)
    return softmax(s) @ v

XLA 把它编译成三个融合块。03 号脚本第 1 部分用迷你 XLA 的几条规则得到了同样的结构：

用到的规则只有几条：

• 输出融合（kOutput）：dot 新开一块，后面的逐元素运算和归约被吸进来——MXU 一边吐出分数，VPU 一边加掩码、找行最大值，原始分数矩阵从不落地。这一块输出两个结果（“多输出融合”）：掩码后的分数和每行的最大值。
• 输入融合：归约把产生它输入的逐元素运算吸进来——exp 的结果在寄存器里就被求和了。（slides 的 HLO 里这一块标为 kLoop。）
• 生产者融合进 dot：最后的归一化只为第二个 dot 服务，被并进 dot 的那一块，算好的概率直接喂进 MXU。

合计：

FLOPs 2.15 G 和超越函数 2.10 M 都与 slides 第 52 页的 cost_analysis 一致——后者正好是 1024 × 1024 个 exp 算了两遍：重算让 exp 的次数翻倍。（字节数和 slides 的 40.9 MB 不能直接比：XLA 的统计还包括布局拷贝、VMEM 读写等。）

4.5 重算：数学便宜，带宽贵

slides 第 68 页的论证：“存下 exp 要写 4 MB、再读 4 MB，这 8 MB 的带宽税比在 VPU 上重算 exp(x − max) 贵得多。” 03 号脚本第 2 部分算了一下账：重算时第 3 块要改读掩码后的分数（也是 4 MiB），所以净省的是一次写：4 MiB，不是 8 MiB。即便如此仍然划算：按张量核心的规格，省下的读写约 1.14 μs，多算的 1.05 M 个 exp（每个约 3 条向量指令）约 0.35 μs。而且重算省下的是存储容量——序列越长，越重要。这就是第 21 讲 FlashAttention 的核心思想：注意力矩阵一次都不落地，反向时再重算。

MLP 子层（03 号脚本第 3 部分，LayerNorm + 线性 + GELU + 线性 + 残差，512 × 256）：25 个方程融合成 4 个内核——LayerNorm 的两次归约各一块（均值；减均值、平方、求方差、加 ε 开方）；归一化、乘 γ 加 β 被吸进第一个 dot 的输入端，加偏置和整个 GELU 被吸进它的输出端；第二个 dot 吸进加偏置和残差。读写从 63.5 MiB 降到 8.5 MiB。这和第 10 讲按 LightSeq 手写的融合几乎一样，只是没有人去写 CUDA。自动融合的边界也很清楚：LayerNorm 的两段归约照样是两个内核——把方差改写成 E[x²] − E[x]² 需要代数上的洞察（和数值上的判断，第 10 讲），这几条融合规则不会自己发现。

一整步训练（03 号脚本第 4 部分）：前向、反向、SGD 更新是一个 jaxpr（追踪 1 次，之后 29 步都命中缓存），43 个运算融合成 12 个内核；参数更新 W − 0.1·∂W 被融进了算 ∂W 的那个 dot 的输出端——第 11 讲 LightSeq2 手写训练器想要的效果，这里由编译器自动得到。

4.6 存储空间、异步拷贝与 VLIW

• 存储空间（slides 第 71 页）：最终的 HLO 给张量标注 S(0)（HBM）、S(1)（VMEM）、S(2)（同步令牌，硬件信号量）。
• 异步拷贝（slides 第 72 页）：copy-start 启动一次把 V 从 HBM 搬到 VMEM 的 DMA，计算单元不等待，继续算第 2 个融合块；copy-done 是轻量的栅栏，保证第 3 块开始前 V 已经到位。搬运和计算重叠，延迟被藏起来。
• LLO 与 VLIW 指令包（slides 第 74–78 页）：每个融合块被降到 LLO（直接对应 MXU、VPU、XLU 的低层算子），再打包成 512 位的 VLIW 指令包，一个包里的操作同步发射：标量与谓词槽（跳转、WAIT）、4 个向量 ALU 槽、矩阵的 push / pop 槽、2 个向量 load 和 1 个 store 槽、6 个立即数槽。编译器用固定的硬件延迟精确排程，典型地让 2–3 个 load 和一次矩阵运算重叠，喂饱 65,536 个乘加单元。
• MXU 上的分块（slides 第 80–83 页）：Q·Kᵀ 的 1024 × 1024 输出切成 4 × 4 个 256 × 256 的块，两个 MXU 各负责一半的行；K = 512 要累加两遍；每装一块权重要推 64 次 8 × 128 的块（vmatpush），左矩阵的块一拍一个地流过（vmatmul），结果从结果寄存器弹出（vpop）。04 号脚本第 3 部分按”装权重 + 流过 M 行 + 填充排空”的简单模型，M = 1024 时利用率上限约 57.2%。
• exp 的改写（slides 第 85 页）：硬件只做 2 的幂（vpow2），e^x 被改写成 2^(x·log₂e)：先 vmul.f32 1.442695，再 vpow2。05 号脚本第 5 部分：x 到 −30 时最大相对误差 1.8e-06（直接 exp 是 6.0e-08），对 softmax 完全够用。slides 第 86 页的指令包里，一个周期同时在处理 6 个不同数据块的 softmax 的不同阶段。

4.7 实测结果

slides 第 92 页（Ironwood 每个张量核心：峰值 1028.75 TFLOP/s）：

第 2 块的 FLOP 利用率只有 0.32%，因为它只用 VPU——05 号脚本第 2 部分：MXU 的峰值是 VPU 的 57 倍。和第 10 讲 GPU 上”非 GEMM 运算 FLOPs 不到 1%、时间却占一半”是同一个现象，解决办法也一样：把它们融合到矩阵乘法的前后，藏在 MXU 的时间里。

5. 分片：Mesh、PartitionSpec、GSPMD 与 shard_map · Sharding

5.1 描述”怎么切”

slides 第 94–97 页：

• Mesh：给设备网格的每一维起名字，例如 2 × 4 的网格叫 (‘X’, ‘Y’)，或 slides 的 (‘batch’, ‘model’)。
• PartitionSpec：数组的每一维写一个网格维度名（沿它切）或 None（复制）。记号 A[I_X, J_Y] 表示 I 维沿 X 切、J 维沿 Y 切；全局形状 [I, J] 在每个设备上的局部形状是 [I / |X|, J / |Y|]。

06 号脚本第 1 部分：Mesh(batch = 4, model = 2) 上的 f32[8, 6]，P(‘batch’, None) 每设备 (2, 6)、全体共 2 份（model 维上复制）；P(None, ‘model’) 每设备 (8, 3)、共 4 份；P(‘batch’, ‘model’) 每设备 (2, 3)、只有 1 份；P(None, None) 每设备一整份、共 8 份。同一份模型代码配不同的网格，只是切法不同。

5.2 GSPMD：编译器自动插通信

slides 第 98、100–103 页：用 jit 时，用户只写全局的程序和输入 / 输出的切法，XLA 的 GSPMD 自动切分计算、插入需要的集合通信。规则只有一条：每个设备要算自己那一块输出，就必须在本地拿到所需的全部输入。06 号脚本第 2 部分（4 个设备，C = A @ B）：

五种切法的结果都与 A @ B 一致。slides 的例子（A: f32[8192, 8192]、B: f32[8192, 1024]，都按行切）：设备 k 有 A 的第 k 块行，却只有 B 的第 k 块行，而算一行输出需要 B 的全部行——编译器插入 all-gather B，每设备发 24 MiB；如果改成沿 K 切再 all-reduce，要 48 MiB。slides 第 103 页还看到编译器为了通信效率，临时把分块从 T(8,128) 改成 T(16,128)，通信完再用 bitcast 改回来对齐 MXU。

5.3 shard_map：自己写局部程序和通信

slides 第 99–104 页：shard_map 里用户写的是每个设备的局部程序（看到的是局部形状 2048，而不是全局的 8192），通信用 jax.lax.all_gather、jax.lax.psum 显式写出。06 号脚本第 3 部分：手写的”all_gather(rhs_block) 再 lhs_block @ rhs_full”与自动版本在每个设备上逐位相同。两者都是 SPMD（单程序多数据，slides 第 105 页）：所有设备运行同一段程序，只是数据不同。自动的写法省心，但编译器的选择不一定最好；手动的写法能精确控制通信的时机——例如把通信和计算重叠（第 15 讲的 DDP、第 18 讲的 ZeRO）。

5.4 GPT-2 的两种网格

06 号脚本第 4 部分按 slides 的 GPT-2 配置（355 M 参数，其中注意力和 MLP 的矩阵 302 M 可以沿 model 维切）估算：

(8, 1) 是纯数据并行：每台设备存全部参数，每步把梯度在 8 台之间 all-reduce。(4, 2) 把注意力的头、MLP 的前馈维度在 model 维上一分为二（slides：每台设备只存 8 个头），每台设备的训练状态从 5.29 降到 3.04 GiB，能训练更大的模型；代价是每层都要在 2 台设备之间交换激活，而且在前向、反向的关键路径上。两种通信性质不同：数据并行传梯度（和参数量成正比，可以和反向重叠），张量并行传激活（和批 × 序列成正比，每层都要等）——所以张量并行通常只放在高带宽的链路上（同一个立方体的 ICI、同一台机器的 NVLink）。第 14–16 讲会展开。

6. 小结 · Summary

1. JAX 的核心是可组合的函数变换。追踪用抽象值把 Python 函数执行一遍得到 jaxpr；jit、grad、vmap 都是”处理 jaxpr”的不同方式，所以能任意组合。
2. jit 按（形状, dtype, 分片）缓存；副作用只在追踪时发生；数据相关的 if 要改成 where 或静态参数；执行是异步的。
3. 编译链：jaxpr → StableHLO（与框架、硬件都无关的契约）→ HLO（静态形状的数据流图）→ 优化（化简、融合、布局、调度、存储分配）→ LLO → VLIW 指令包。
4. XLA 的融合规则（输出融合、输入融合、生产者融合、收尾）自动得到第 10 讲手写的大部分融合；重算用便宜的计算换昂贵的读写（注意力的 exp 算两遍，省下一次 4 MiB 的写）。
5. TPU 张量核心：HBM ↔ VMEM 的流水线；8 × 128 的向量寄存器（最后两维最好是 8、128 的倍数）；VPU 18 TFLOP/s；XLU 负责跨 lane；256 × 256 的 MXU 是权重驻留的脉动阵列，bf16 乘、FP32 加，流过 B 个向量要 B + 2N − 1 拍——只有大批量才喂得饱它。
6. 并行用 Mesh + PartitionSpec 描述；jit（GSPMD）自动插集合通信，shard_map 手写局部程序和通信；数据并行传梯度、张量并行传激活。

7. 常见误区 · Common pitfalls

• “jit 里的 print 会每次都打印。” 只在追踪时打印一次，打印的是 tracer。
• “jit 只看函数，不看参数。” 缓存的键是参数的形状、dtype、分片（和静态参数的值）；形状一变就重新编译。
• “jit 里可以随便写 if。” 条件依赖数据时不行：要么 where，要么静态参数，要么 lax.cond。
• “grad 是另一套机制，和 jit 无关。” grad 生成的仍是基本运算组成的 jaxpr，可以被 jit 编译、被 vmap 批处理、再被 grad。
• “编译器能化简一切。” x·0 → 0、E[x²] − E[x]² 这类改写会改变浮点结果，编译器默认不做。
• “重算总是浪费。” 当读写比计算贵时（大部分逐元素运算），重算反而更快、更省显存。
• “TPU 的利用率低说明代码写得差。” 批很小时（解码）MXU 的利用率天生就低：填充和排空的 2N − 1 拍摊不开。
• “MXU 全程用 bf16。” 只有乘法输入是 bf16，累加是 FP32；否则长点积的误差会大一个数量级以上。
• “头维度 64 和 128 在 TPU 上一样。” 64 的那一维要填充到 128，向量寄存器一半是空的。
• “shard_map 和 jit 是两种不同的并行。” 都是 SPMD；区别只在于通信是编译器插的还是用户写的。

8. 练习 · Exercises

1. 写出 f(x) = exp(x) * x + 1.0（x: f32[3]）的 jaxpr，再写出 grad(lambda x: sum(f(x))) 里反向部分的方程。
2. 一个 jit 函数先后用 f32[8]、f32[8]、f32[16]、f64[8]、f32[16] 调用，一共追踪几次？如果函数里有 print("hi")，会打印几次？
3. 把 def g(x, n): return x * 2 if n > 3 else x 改写成能被 jit 的两种形式（n 是整数配置）。
4. 用脉动阵列的公式算：N = 128 的阵列流过 B = 128 个向量要几个周期？利用率多少？B 要多大利用率才超过 90%？
5. 一个 256 × 256 的 MXU、2.2 GHz、每个乘加算 2 FLOP，峰值是多少 TFLOP/s？Ironwood 每个张量核心 1028.75 TFLOP/s，相当于几个这样的 MXU？
6. bf16[24, 100] 的数组按 (16, 128) 的块填充后是多少？有效比例多少？
7. 写出 f32[4,4]{1,0:T(2,2)} 里元素 (2, 1) 和 (1, 3) 的内存位置。
8. 注意力里如果不融合 select 和 reduce_max 到第一个 dot 里，读写会多多少（按 03 号脚本的口径，f32[1024,1024]）？
9. 为什么重算 exp 时第 3 个融合块要读入”掩码后的分数”而不是原始分数？如果第 1 块不输出掩码后的分数、只输出最大值，第 3 块要读什么？
10. 4 个设备上 A: f32[4096, 4096] 按列切、B: f32[4096, 4096] 按行切，C 要复制在每个设备上。插什么通信？每设备发多少字节？如果 C 只要按行切呢？
11. GPT-2（355 M 参数，fp32 梯度）在 8 个设备上做数据并行，每步梯度 all-reduce 每设备发多少？在 ICI 1200 GB/s 下至少要多久？
12. 为什么张量并行的激活通信”在关键路径上”，而数据并行的梯度通信可以和计算重叠？

参考答案

1. 前向：b = exp a; c = mul b a; d = add c 1.0，再 e = reduce_sum d。反向（ct = 1 广播成 f32[3] 记作 g）：add 的规则把 g 传给 c；mul 的规则得到对 b 的 g·a 和对 a 的 g·b；exp 的规则把对 b 的余切乘以 out：(g·a)·b；最后对 a 的两份相加：g·b + g·a·b，即 exp(x)(1 + x)。
2. 键是（形状, dtype）：f32[8]、f32[16]、f64[8] 三种 → 追踪 3 次，打印 3 次。
3. 一：把 n 声明为静态参数，jit(g, static_argnums=(1,))（每个 n 值编译一份，n > 3 的那些共享同一个分支的代码但缓存键不同）。二：where(n > 3, x * 2, x)，n 作为普通参数传入（两边都算，按条件选）。
4. B + 2N − 1 = 128 + 255 = 383 个周期，利用率 128/383 ≈ 33.4%。B / (B + 255) > 0.9 ⇔ B > 2295，即 B ≥ 2296。
5. 65,536 × 2 × 2.2e9 ≈ 288 TFLOP/s。1028.75 / 288 ≈ 3.6 个——每个张量核心有不止一个 MXU（slides 第 80 页提到 MXU0、MXU1；具体个数和频率以官方规格为准）。
6. 行 24 → 32（16 的倍数），列 100 → 128：[32, 128]，有效比例 2400 / 4096 ≈ 58.6%。
7. 4 × 4 切成 2 × 2 的块，每行 2 块。(2, 1)：块 (1, 0) 是第 2 块，块内位置 (0, 1) → 2·4 + 0·2 + 1 = 9。(1, 3)：块 (0, 1) 是第 1 块，块内 (1, 1) → 1·4 + 1·2 + 1 = 7。
8. 不融合时 dot 要写出原始分数（4 MiB），select 读它和掩码（4 + 1 MiB）再写出（4 MiB），reduce_max 再读一遍（4 MiB）——比融合版多出 4（写）+ 4（读）+ 4（读）= 12 MiB 左右，外加两次内核发射。
9. 因为 exp 要对”加过掩码的分数减最大值”求：被掩掉的位置必须是 −1e38，exp 后才是 0。如果第 1 块不输出掩码后的分数，第 3 块就得读原始分数和掩码（4 + 1 MiB）再重新做 select——甚至可以连原始分数都不读、重新算 Q·Kᵀ（FlashAttention 的选择）。
10. 每个设备算出一份完整大小的部分和 [4096, 4096]，要 all-reduce：2 × (3/4) × 64 MiB = 96 MiB。C 按行切时用 reduce-scatter：(3/4) × 64 MiB = 48 MiB。
11. 355 M × 4 字节 ≈ 1.42 GB，all-reduce 每设备发 2 × 7/8 × 1.42 ≈ 2.48 GB（06 号脚本：2.31 GiB），÷ 1200 GB/s ≈ 2.1 ms（这里把双向带宽全算上，是下限）。
12. 张量并行：下一层的计算要用这一层 all-reduce 之后的激活，不通信完就没法继续——通信在数据依赖链上。数据并行：某一层的梯度一算出来就可以开始 all-reduce，同时反向继续算前面的层；只有到参数更新时才需要全部梯度（第 15 讲的 DDP 正是利用这一点）。

9. 延伸阅读 · Further reading

• JAX 文档：“Thinking in JAX”、“JAX internals: the jaxpr language”、“Distributed arrays and automatic parallelization”、“shard_map”。
• Austin et al. How to Scale Your Model（jax-ml.github.io/scaling-book）：TPU、roofline、分片记号 A[I_X, J_Y] 与集合通信的系统讲解（slides 第 96–97 页的图来自这里）。
• Frostig, Johnson & Leary (2018) Compiling machine learning programs via high-level tracing（SysML）：JAX 的原始论文。
• Xu et al. (2021) GSPMD: General and Scalable Parallelization for ML Computation Graphs：jit 自动分片的编译器。
• Jouppi et al. (2017) In-Datacenter Performance Analysis of a Tensor Processing Unit（ISCA）：第一代 TPU 与脉动阵列；以及后续 TPU v4（OCS、3D 环面）的论文。
• Kung (1982) Why Systolic Architectures?：脉动阵列的经典文章。
• OpenXLA / StableHLO 的规范与文档；Autodidax（JAX 文档里”从零写一个 JAX”的教程）——本讲的 minijax 就是按这个思路简化的。

10. 下一讲预告 · Next lecture

这一讲看到 XLA 能自动融合出很好的结果——但它也有做不到的时候：分块大小、什么数据留在 VMEM、怎样把搬运和计算重叠，这些都由编译器的启发式决定；遇到稀疏的注意力掩码，它更不知道哪些块可以整块跳过。第 13 讲（同一位嘉宾）讲 Pallas：JAX 里手写 TPU 内核的方式。我们会看到 BlockSpec 怎样描述”每一步把哪一块搬进 VMEM”、双缓冲的流水线怎样让 DMA 和计算重叠、在线 softmax 怎样让 Flash Attention 不需要整行的分数，以及 Splash Attention 怎样利用块稀疏的掩码（因果、局部窗口、文档边界）整块跳过计算。