JAX/XLA 原语分析
MaxText 是配置驱动的引擎：模型代码在 trace 期被组装成 jaxpr，最终落成一张 XLA 程序，在 TPU pod 上以 SPMD 方式执行。读懂这台引擎，本质上就是读懂它用了哪些 JAX 原语、每个原语在追踪期与编译/运行期分别做了什么。
本系列不是 API 手册，而是**「原语在真实场景下的行为分析」**。对每个原语，结合 MaxText 中实际调用它的场景（训练主循环 / 分区 / Decoder 堆叠 / 注意力 / MoE / 流水线 / 上下文并行 / 量化），讲清三件事：
- JAX 追踪期语义 —— trace 出什么、引入什么副作用；
- XLA 落地形态 —— 编译成什么 HLO（collective、While、Conditional、Dot、Mosaic kernel）、如何参与 GSPMD 分区、是否触发重计算/offload；
- MaxText 为什么这样用 —— 性能（MFU）、显存、还是数值对齐的动机。
通用背景（全系列共识）
后续各场景页都建立在下面这套共识之上。
jit→ trace → jaxpr → HLO：一次jit调用把 Python 函数 trace 成 jaxpr 并编译成一张 XLA 可执行程序；Python 控制流只在 trace 期执行一次，运行期跑的是编译后的 HLO。- SPMD / GSPMD：单程序、多设备。张量带 sharding 标注，XLA 的 GSPMD partitioner 自动传播 sharding 并在需要处插入集合通信（all-gather / reduce-scatter / all-reduce）。
- 两套分区世界：
- GSPMD 自动分区 —— 用
with_sharding_constraint/ 逻辑轴约束给张量贴 sharding，编译器决定何时通信； shard_map手动分区 —— 切换到 per-shard 视角，你显式写集合通信，通信直接成为 HLO。
- GSPMD 自动分区 —— 用
- 逻辑轴 → 物理轴：MaxText 用 Flax 逻辑轴名（如
embed/mlp/heads）标注激活，再由logical_axis_rules映射到 mesh 的物理轴，把「模型代码」与「并行策略」解耦。 - Pallas / Mosaic：用 Python 写 TPU kernel，经 Mosaic 编译成自定义 HLO，绕开 XLA 的默认 codegen（splash attention、megablox GMM、KDA 等）。
- remat / checkpoint policy：用重计算换显存。policy 决定哪些中间值在前向存下、哪些反向重算、哪些 offload 到 host。
custom_vjp：手写反向，绕开自动微分——用于 Pallas kernel 反向、数值稳定、节省显存。
原语速查表
按类别速查；「页」列指向对应场景深入页。
并行 / 分区
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
shard_map | 在子函数内切换到 per-shard 手动分区视角，axis 名进入 manual mesh | 不走 GSPMD 自动通信，内部显式 collective 直接成 HLO | MoE / 注意力 / 流水线（04,05,06） |
with_sharding_constraint | 给中间张量贴物理 sharding（NamedSharding） | 约束 GSPMD 传播，必要处插入 reshard | 分区（02） |
nn.with_logical_constraint（Flax） | 给激活贴逻辑轴名 | 经 logical_axis_rules 解析为物理 sharding 后再约束 GSPMD | 激活分区主力（02） |
PartitionSpec / NamedSharding | 描述张量在 mesh 上如何切分 | 喂给 GSPMD partitioner | 全局（02） |
logical_axis_rules | 逻辑轴名 → mesh 物理轴名的映射规则 | 决定上述约束的物理含义 | configs/types.py、utils/sharding.py（02） |
pjit | 带 in/out sharding 的 jit | GSPMD 分区编译 | 少量显式使用（02） |
编译变换
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
jit | trace 函数→jaxpr，缓存编译产物 | 一张 XLA 可执行程序 | 全局（01） |
eval_shape | 只 trace 形状/dtype，不做实际计算 | 不产生执行 HLO，仅抽象推导 | 初始化 / AOT（01） |
ShapeDtypeStruct | 抽象张量占位（shape + dtype + sharding） | AOT lower 的输入规格 | 01 |
lower() + compile()（AOT） | 提前 trace 并编译，不执行 | 把编译从 step 0 提前，产出可复用可执行体 | trainers/pre_train/train_compile.py（01） |
make_jaxpr | 导出 jaxpr | 不执行 | 调试用（01 附注） |
控制流
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
lax.scan | 循环体 trace 一次，带 carry | 编译成 While HLO，循环体不展开（缩小程序体积与编译时间） | 层堆叠 / microbatch（03,06） |
lax.cond | 两个分支都 trace | Conditional HLO，运行期择一执行 | MoE 路由分支等（05） |
lax.fori_loop | 固定/动态次数循环 | While HLO 或展开 | 流水线（06） |
lax.while_loop | 条件循环 | While HLO | 流水线（06） |
lax.select | 逐元素择一 | Select HLO（无分支，两路都算） | 注意力 mask 等（04） |
集合通信（多在 shard_map / 流水线内）
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
lax.all_gather | 沿 axis 收集各 shard | AllGather HLO | FSDP 权重收集 / CP 环（05,06,07） |
lax.ppermute | 按 perm 在设备间点对点轮转 | CollectivePermute HLO | 流水线 ring / CP 环（06,07） |
lax.psum | 跨 axis 求和 | AllReduce HLO | MoE / 梯度（05） |
lax.psum_scatter | 求和并 scatter（reduce-scatter） | ReduceScatter HLO | megablox GMM（05） |
lax.axis_index | 取当前设备在 axis 上的下标 | 设备 id 相关 HLO | 路由 / mask（05,06） |
自定义梯度
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
custom_vjp | 指定前向 + 自定义反向 | 反向不由 autodiff 生成，接入你给的 HLO/kernel | Pallas 反向 / 注意力 / token 排序（04,05,07） |
stop_gradient | 反向梯度截断为 0 | 反向图剪枝 | 全局 |
value_and_grad | 同时求值与反向 | 前向 + 反向 HLO | 训练 loss（01） |
算子内核（Pallas / 底层）
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
pallas_call | 调用 Pallas kernel（grid / BlockSpec） | Mosaic 编出自定义 TPU HLO | splash / megablox / KDA（04,05,07） |
dot_general | 通用张量收缩 | Dot HLO（matmul），量化在此插入 | 量化 / GMM（05,08） |
显存 / remat
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
remat（jax.checkpoint） | 标记重计算边界 | 反向时重算而非物化激活 | 03 |
nn.remat（Flax） | 作用于带参 module 的 remat | 同上，作用于 module | 03 |
jax.checkpoint_policies.* | 选择 存 / 算 / offload 的命名集合（save_only_these_names、save_and_offload_only_these_names） | 控制 remat 物化策略与 host offload | 03 |
Flax lifted 变换
| 原语 | JAX 追踪期语义 | XLA 落地形态 | 主要场景（页） |
|---|---|---|---|
nn.scan | 在带参 module 上 scan（层堆叠/共享） | 底层 lax.scan → While HLO | Decoder 堆叠（03） |
nn.vmap | module 上批量化 | 新增向量化维度 | 03 / 05（按用到处） |
阅读导航
| 页 | 主题 | 串起的原语 |
|---|---|---|
| 01 | 训练主循环与 AOT 编译 | jit / value_and_grad / lower+compile / eval_shape / ShapeDtypeStruct |
| 02 | 分区：逻辑轴 → 物理 mesh | with_logical_constraint + logical_axis_rules / with_sharding_constraint / PartitionSpec·NamedSharding / pjit |
| 03 | Decoder 堆叠与重计算 / Offload | nn.scan / nn.remat / remat_policy / checkpoint_policies |
| 04 | 注意力 | shard_map / pallas_call(splash) / custom_vjp / cond·select |
| 05 | MoE / 专家并行 | shard_map / all_gather·psum·psum_scatter·axis_index / custom_vjp / megablox GMM pallas_call / dot_general |
| 06 | 流水线并行 | ppermute / scan / all_gather / fori_loop·while_loop |
| 07 | 上下文并行 / KDA 线性注意力 | ring all_gather·ppermute / GLA scan kernel / custom_vjp |
| 08 | 量化 | dot_general / AQT·Qwix 拦截 / sub-channel |