Inside JAX：XLA 编译器如何将 Python 函数编译为 GPU PTX 代码

当你调用 `jax.jit(func)(x)`，Python 解释器里的几行代码經歷了一场从抽象语法树到硅片指令的完整旅行。本文拆解 JAX/XLA 的核心编译机制，带你理解 JAX 如何在不牺牲 Python 表达力的前提下，达到手写 CUDA kernel 80% 以上的性能。

一、JIT 的起点：Tracing 而非编译

JAX 的 jit 不是传统意义上的 JIT 编译器。它的工作方式是 Tracing：运行你的 Python 函数一次，但用 Tracer 对象代替真实张量。Tracer 并不做数学运算，只 记录 遇到的每一个 primitive 操作。

import jax
import jax.numpy as jnp

def predict(w, b, x):
    return jnp.maximum(0, w @ x + b)  # 简单 ReLU 全连接

# 首次调用会触发 tracing，生成 jaxpr（JAX 的 IR）
jax.make_jaxpr(predict)(jnp.ones((4, 8)), jnp.ones(8), jnp.ones(8))

输出一个类似 lisp 的 jaxpr（JAX Expression），它就是你的函数在 JAX 内部的表示：

{ lambda ; a:f32[4,8] b:f32[8] c:f32[8]. let
    d:f32[4] = dot[in_dtype=float32] a c
    e:f32[4] = add d b
    f:f32[4] = max 0.0 e
  in (f,) }

关键洞察：tracing 产生的 jaxpr 是 shape-typed 但 value-agnostic 的。同一个 trace 可以处理不同数值、相同 shape 的所有输入。这意味着 jit 缓存了编译结果——相同 shape 的第二次调用会复用缓存的 XLA 可执行文件，而零开销。

当 Python 控制流（if/for）依赖于输入值时，tracing 时它已经确定了，jit 会展开为单一执行路径。如果需要值相关的条件分支，必须显式使用 jax.lax.cond、scan 或 while_loop 这些 traced primitive。

二、XLA 编译管线：从 jaxpr 到 HLO

jaxpr 随后被翻译为 XLA HLO（High Level Operations）——XLA 编译器的核心中间表示。HLO 本身分为两层：硬件无关层（StableHLO）和硬件优化层。

核心编译 passes 包括：

• Fusion Pass：将多个逐元素操作合并为单个 kernel，避免写回 HBM。例如 relu(batch_norm(x)) 被融合为从 shared memory 读一次、写一次的单个 kernel。
• Layout Assignment：为每个 tensor 选择内存布局（major/minor dimension 排列），最大化 coalesced memory access。
• Tiling & Vectorization：针对具体 GPU SM 数量、shared memory 大小选择分块策略。

算子融合是 JAX/XLA 性能的核心杠杆。考虑三层 MLP 的前向传播，如果不做融合，每一层的矩阵乘、加偏置、ReLU 各自是一个 kernel，每次 kernel 启动都有约 5μs 的开销，加上中间结果写回显存的开销。XLA 将整层融合后，计算在寄存器/shared memory 中完成流水线化，性能提升可达 3-5x。

三、自动并行化：从 pjit 到 GSPMD

JAX 的并行列演进经历了三代 API：pmap → xmap → pjit。当前 jax.jit + NamedSharding 注解是推荐范式。

核心概念是 逻辑分片轴到物理设备网格的映射。声明式分片让编译器自动插入通信原语：

from jax.sharding import Mesh, PartitionSpec as P
import numpy as np

# 4 个 GPU 排成 (data=2, model=2) 的网格
devices = np.array(jax.devices()).reshape(2, 2)
mesh = Mesh(devices, axis_names=('data', 'model'))

# GPT 注意力头的张量并行
# query, key, value 沿 model 轴分片，每个 device 计算部分 head
with mesh:
    q = jax.jit(
        compute_attention,
        out_shardings=PartitionSpec('data', 'model', None)
    )(q_sharded, k_sharded, v_sharded)

分片传播（sharding propagation）规则决定了哪些操作触发 all-reduce、all-gather。矩阵乘在不同维度分片时，XLA 自动决定是 local matmul + all-reduce 还是 all-gather + local matmul，哪个更便宜。这种 GSPMD（Generalized SPMD）分区方式意味着如果一个 tensor 是分片的而另一个不是，编译器会自动在分片的那个上做 collective op。

一个实战陷阱：不恰当的分片会导致大量隐式通信。例如当 weight 按 model 维度分片、activation 按 data 维度分片时，每个 matmul 都触发一次全局 all-reduce。常见的优化是让权重和输入在同一维度分片，输出在另一个维度分片，实现 通信与计算重叠。

四、自动微分的实现：VJP 与 operator overloading

JAX 的自动微分基于 Vector-Jacobian Product（VJP） 模式。每个 primitive 注册一个 transpose rule，描述前向操作的雅可比矩阵转置如何作用于余切向量。

@jax.custom_vjp
def logsumexp(x):
    """数值稳定的 log(sum(exp(x)))"""
    a = jnp.max(x)
    return a + jnp.log(jnp.sum(jnp.exp(x - a)))

def logsumexp_fwd(x):
    y = logsumexp(x)
    return y, (x, y)

def logsumexp_bwd(res, g):
    x, y = res
    # 自定义反向：避免 exp(x-y) 再次溢出
    return (g * jnp.exp(x - y),)

logsumexp.defvjp(logsumexp_fwd, logsumexp_bwd)

custom_vjp 的价值在于你可以对反向传播做数值稳定性优化。例如上面标准实现的 grad(logsumexp) 虽然数学正确，但 exp(x-y) 在前向已做过一次，反向又做一次，如果 x-y 很大就会溢出。自定义 VJP 复用前向中间结果，彻底消除重复计算和溢出风险。

JAX 内部每个数学运算都对应一个 primitive，primitive 同时注册了 impl_rule（实际计算）和 jvp_rule（前向微分）。grad(f) 在 tracing 阶段构建前向计算图后，再对图中的每个节点应用转置规则生成反向图。这意味着 JAX 的 autodiff 支持任意前向代码中的控制流、高阶函数甚至自定义 kernel——只要转置规则正确。

五、性能调优实战

5.1 消除重编译

最常见的性能杀手：shape 变化触发重编译。jit 对输入 shape 是特化的，不同 shape 意味着不同的 tile size 选择、不同的 shared memory 配置。

# 反模式：dynamic shape 导致每次重编译
for batch in dataset:
    # batch_size 不同 → 每次 tracing → 编译
    params = train_step(params, batch)

# 正解：填充或固定 micro-batch size
dataset = dataset.batch(MICRO_BATCH, drop_remainder=True)

另一个源是 Python 标量被 trace 为 constant。jit 对 Python int/float 参数的数值变化会触发重编译。解决方式是传入 jnp 数组替代标量，或使用 static_argnums 显式标记。

5.2 操作融合模式

XLA 的融合器使用启发式规则避免寄存器溢出。当融合后的 kernel 需要太多寄存器的中间值超过 SM 的 256 寄存器文件时，会拆分为多个 kernel。可以用 jax.jit 的 jax_tracebackloc 或 JAX_DUMP_FLOPS 查看。

反直觉的优化：有时拆分大 kernel 反而更快——允许 HBM→L2→Shared Memory 的流水级重叠。通过 XLA_FLAGS=--xla_dump_to 可以导出 HLO 图进行分析。

5.3 设备内存管理与 donatable

在长训练循环中 JAX 会积累中间 buffer。jax.jit 的 donate_argnums 参数允许调用者放弃输入 buffer 的所有权，XLA 在编译期内复用其内存：

# 训练步骤中 params 每次都被覆盖，可以 donate
@partial(jax.jit, donate_argnums=(0,))
def train_step(params, batch):
    grads = jax.grad(loss_fn)(params, batch)
    return jax.tree_map(lambda p, g: p - lr * g, params, grads)

这个提示让 XLA 的 buffer assignment 器直接将梯度结果写回 params 的原内存位置，峰值内存使用降低约 30%。

六、与手写 CUDA 的性能差距

在实际 benchmark 中（GPT-2 small，8x A100），JAX 实现的训练 throughput 约为手写 CUDA FasterTransformer 的 82-88%。差距主要来自：

1. XLA 的通用 kernel 无法做到针对特定 block size 的极致寄存器分配

2. CUDA graph 的 launch overhead 在 JAX 中略多（XLA runtime 层）

3. 自定义 attention flash 需要在 JAX 中绕道 custom_call（虽然 FlashAttention-3 已官方支持）

JAX 80%+ 的性能但 30% 的开发时间，这就是它在 2024-2026 年成为 AI 研究首选框架的原因。从 Flax/Haiku 到现在的原生 Equinox、NNX 等框架，JAX 的生态已经从"学术研究工具"演进为大规模生产训练的标准基础设施（包括 DeepMind 的 Gemini 预训练、Mistral 的模型微调管线等）。

结语

JAX 的设计哲学是 "可组合的变换"：jit、grad、vmap、pm 是四个可交换顺序的函数变换算子。理解 XLA 编译链路不是为了取代 JAX 的抽象，而是为了在关键时刻——merging kernel 不足、内存爆炸、通信模式出问题时——能够定位根因并给出正确的 surgeon-level 修复。当 Python 变成世界上最慢的语言时，是编译器在背后提速。