JAX jit 怎样减少动态形状导致的重复编译

如果同一个 jax.jit 函数第一次调用很慢，后续请求又因为批次大小或序列长度变化不断变慢，通常不是“缓存失效”这么简单：JAX 会根据参数的类型信息生成并缓存编译结果，输入的 shape、dtype 或静态参数变化，都可能让缓存键发生变化。处理这类问题的重点，是把无限多种输入形状收敛到有限集合，同时把函数定义和静态参数保持稳定。

官方文档：https://docs.jax.dev/en/latest/jit-compilation.html

先打开编译日志确认是不是重复编译，再决定用 padding、分桶或静态参数整理。不要一上来调用 jax.clear_caches()；清缓存只能改变当前进程的缓存状态，不能解决输入形状不断变化的问题。

先判断：到底是什么让 jit 重新编译

JAX 的第一次 jit 调用会追踪函数并生成编译产物，后续调用只有在缓存键匹配时才能复用。最常见的差异有三类：

• 输入类型不同：相同逻辑但 shape 或 dtype 不同，都会产生新的专门化版本。例如 [8, 128] 和 [8, 256] 不是同一个数组类型。
• 静态参数不同：被标记为 static 的参数值会参与编译。如果把大量变化的业务值标成静态参数，编译数量会随取值增长。
• 函数对象不稳定：在循环或请求处理函数内部不断创建临时函数，再对它调用 jax.jit，会让缓存难以复用。

因此，先记录“哪些输入变化对应一次新编译”，比盲目增加缓存目录更重要。若变化来自真实的动态 shape，应该在输入侧收敛；若变化来自函数或静态参数，就应先整理代码边界。

打开编译日志，把缓存未命中变成线索

排查阶段可以开启较轻量的编译日志，也可以开启更详细的缓存未命中解释。下面示例只展示配置方式，日志本身由 JAX 输出。

import jax
import jax.numpy as jnp

# 只在排查阶段打开；它会让每次编译都留下明显日志。
jax.config.update("jax_log_compiles", True)
# 需要知道缓存键差异时再打开更详细的原因说明。
jax.config.update("jax_explain_cache_misses", True)

@jax.jit
def reduce_batch(x):
    # 让函数只依赖数组输入，避免把业务配置误标成静态参数。
    return jnp.mean(x, axis=-1)

# 这两个 shape 不同，排查时应观察是否出现两次编译日志。
first = reduce_batch(jnp.ones((8, 128), dtype=jnp.float32))
second = reduce_batch(jnp.ones((8, 256), dtype=jnp.float32))
print(first.shape, second.shape)  # 输出只是示例，不代表本文已在本机运行

日志中如果明确显示不同输入类型或 shape 导致 cache miss，就进入输入治理；如果提示函数对象或静态参数不同，就不要继续增加 bucket，而要把 jax.jit 提到稳定的模块级函数上。排查结束后可以关闭这些高噪声配置，避免把诊断日志混入正常服务。

处理动态输入：padding 和分桶比无限编译更稳

可变序列长度、最后一个不完整 batch 是重复编译的典型来源。实用做法不是为每个长度编译一个版本，而是预先设计少量 bucket，例如 128、256、512；输入落入某个 bucket 后补齐，计算时用 mask 忽略补齐部分。

import jax.numpy as jnp

def pad_to_bucket(tokens, bucket_size, pad_value=0):
    # 只把长度补到预设 bucket，超过最大 bucket 时交给上层拒绝或截断策略。
    length = tokens.shape[-1]
    if length > bucket_size:
        raise ValueError("输入长度超过当前 bucket")
    padding = bucket_size - length
    padded = jnp.pad(tokens, ((0, 0), (0, padding)), constant_values=pad_value)
    mask = jnp.arange(bucket_size) 

分桶数量要和业务分布一起决定：桶太多会重新引入编译组合爆炸，桶太少则会浪费更多 padding 计算。线上常见的判断方式是统计长度分布，优先覆盖大多数请求，再为超长输入保留明确的拒绝、截断或异步路径。无论选择哪种策略，目标都是让 jit 看到有限的 shape 集合。

稳定函数定义，并减少静态参数的变化

把装饰后的函数放在模块级，避免在请求循环里重复创建等价函数。对于确实需要控制编译分支的配置，可以使用少量、可枚举的静态值；不要把用户 ID、随机业务开关或每次变化的字符串作为静态参数。

import functools
import jax
import jax.numpy as jnp

def _scale(x, factor):
    # factor 只有在确实改变计算图时才适合成为静态参数。
    return x * factor

# 在模块加载时创建稳定的 jit 函数，服务请求只传数组数据。
scale = jax.jit(_scale, static_argnames=("factor",))

def apply_scale(x, factor):
    # 业务层把 factor 限制为少量已知档位，避免无限生成编译版本。
    if factor not in (1.0, 2.0, 4.0):
        raise ValueError("factor 必须是预设档位")
    return scale(x, factor)

如果 factor 并不会改变控制流或算子形状，也可以重新评估是否需要把它标成 static。静态参数的优势是能参与专门化，代价是每个不同值都可能对应新的编译产物。函数身份也同样重要：不要在循环内创建新的 lambda、functools.partial 或装饰器结果。

跨进程运行时，再考虑持久化编译缓存

当输入 shape 已经收敛、函数定义也稳定，但服务重启后仍要为同一批计算重新编译，可以配置 JAX 的持久化编译缓存。它把编译产物写到磁盘或对象存储，主要解决跨进程、跨重启的重复等待；它不能替代 shape 治理，也不会让互不相同的 shape 自动变成同一个程序。

import jax

# 生产环境应选择权限受控、生命周期清楚的缓存目录。
jax.config.update("jax_compilation_cache_dir", "/var/cache/my-jax")
# 排查阶段可降低写入门槛；上线时按官方配置和容量策略调整。
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)

官方文档特别提醒，持久化缓存目录应被视为受信任内容：不要让不受信任的用户拥有写权限。另一个边界是版本、后端、设备拓扑和编译选项都会影响缓存键，所以升级 JAX、切换硬件或调整 XLA 选项后，出现新的编译并不一定是配置失效。

出现异常时的回退路径

如果新分桶策略导致结果差异，先保留原始变长输入路径作为回退，再逐个 bucket 对比数值和延迟。不要直接清空生产缓存来“修复”问题；更稳妥的顺序是：

1. 关闭详细编译日志，确认问题不是诊断开关造成的噪声或额外开销。
2. 临时减少 bucket 数量或回到单一固定形状，确认计算结果和服务稳定性。
3. 检查 padding 的 mask 是否参与了归一化、损失或聚合，避免补齐值改变业务结果。
4. 确认缓存目录权限、容量和版本边界，再逐步恢复持久化缓存。

把重复编译控制在可解释范围内

这类问题的最终目标不是让 JAX 永远只编译一次，而是让编译次数与少量、可解释的输入档位对应。先用日志找出 cache miss 的差异，再用固定形状、分桶、稳定函数定义和有限静态参数减少组合；跨重启仍有成本时，最后再接入持久化编译缓存。这样即使业务输入继续变化，也能把编译等待从不可控的请求级事件，变成可监控、可回退的工程参数。

相关问题

• 为什么只改变 dtype 也可能触发新的 JAX 编译？因为 dtype 属于输入类型信息，通常会进入专门化和缓存匹配。
• padding 会不会让推理更慢？会增加一部分无效计算，但通常能用有限编译次数换取更稳定的尾延迟，需要按长度分布权衡。
• 持久化缓存能否解决每个长度都编译的问题？不能；它主要减少跨进程重复编译，动态 shape 仍然会产生不同缓存键。
• 什么时候应该用静态参数？只有当参数确实改变控制流或计算图，并且取值集合很小时才适合。