1. 从一次推理延迟抖动说起：为什么要折腾 JAX 版 Pi05 与 RTC
先说结论：这套东西的核心目标只有一个——把 Pi05 这类策略模型在 JAX 上的推理链路压到实时可控，让 RTC（Real-Time Chunking，实时分块）真正跑出“边生成边执行”的效果，而不是名义上实时、实际上一卡一卡。如果你正在做机器人策略部署、具身智能推理加速，或者单纯对 JAX 生态下的低延迟推理感兴趣，这篇记录应该能帮你少走几天弯路。
Pi05 是当前具身智能领域讨论度很高的一类策略模型架构，它的特点是参数量不算夸张，但对推理时序极其敏感——因为下游是真实的机械臂或移动底盘，动作块（action chunk）晚出来几十毫秒，整条轨迹的平滑度就会肉眼可见地变差。RTC 这套实时分块机制，本质上是把一次长序列的动作预测拆成多个短块，边预测边下发执行，用“流水线”换“低首帧延迟”。听起来很美好，但真把它落到 JAX 上，坑一个接一个：JIT 编译的首次开销、分块边界的对齐、KV Cache 的复用、设备间同步的抖动，每一个都能让实时性崩掉。
我这次优化的起点，是一个能跑通但延迟不稳的版本：平均单块推理 40ms 上下，但 P99 能飙到 120ms 以上，抖动大到下游控制器直接报警。目标是把 P99 压到 60ms 以内，同时保证动作块之间的连续性不出现跳变。下面我把整个拆解、选型、实操和踩坑过程完整写出来，涉及参数的地方我都会给出计算依据，涉及取舍的地方我都会说清楚为什么这么选。
2. 整体设计思路：为什么是 JAX + RTC 这个组合
2.1 Pi05 推理链路的瓶颈到底在哪
在动手优化之前，必须先定位瓶颈，否则就是瞎调参。我把 Pi05 的推理链路拆成四段：输入预处理（图像编码 + 状态拼接）、主干网络前向、动作头解码、后处理（动作裁剪与平滑）。用 JAX 的jax.profiler抓了一次 trace，结果很典型：
- 输入预处理占了约 18% 的时间，主要是图像 resize 和归一化在 CPU 上做，然后才搬到设备；
- 主干前向占 55% 左右，这部分是纯计算，理论上最该被 JIT 吃掉；
- 动作头解码占 15%，但因为 RTC 分块，这部分被反复调用，累积开销不小；
- 后处理占 12%，看着不多，但它是纯 Python 循环，没法向量化，延迟完全不可控。
关键发现是：首次调用和后续调用的延迟差异巨大。第一次前向因为 JIT 编译，花了将近 2 秒；后续稳定在 40ms 左右。这说明 JIT 本身没问题，问题在于编译时机和缓存策略——如果每次输入 shape 变了就重新编译，那实时性根本无从谈起。
2.2 RTC 分块机制与 JAX 的天然契合点
RTC 的核心思想是把长度为 T 的完整动作序列切成 K 个长度为 L 的块（T = K × L），每次只预测当前块，同时用上一块的尾部作为当前块的上下文前缀，保证块与块之间的连续性。这个机制和 JAX 的几个特性天然契合：
第一，JAX 的jit对固定 shape 的函数编译一次就能反复用，而 RTC 的每个块 shape 是固定的（都是 L），这正好避开了动态 shape 触发重编译的坑。第二，JAX 的vmap可以很方便地把多个块的推理并行化，如果你有足够的显存，甚至可以一次预测多个候选块然后择优。第三，JAX 的donate机制（jax.jit(..., donate_argnums=...)）能在块之间复用缓冲区，减少内存分配带来的抖动。
但契合归契合，真做起来有两个硬约束：一是块长度 L 必须和模型的训练时上下文窗口对齐，否则连续性会崩；二是 KV Cache 的跨块复用需要手动管理，JAX 不会自动帮你做。
2.3 方案选型：为什么不用 PyTorch 而坚持 JAX
这里必须解释一下选型逻辑，因为很多人第一反应是“PyTorch 生态更成熟，为什么不用”。我实际对比过两版实现，结论是：在固定 shape、高频调用的推理场景下，JAX 的 XLA 编译优势是实打实的。
PyTorch 的 eager 模式单块推理约 55ms，用torch.compile后能到 45ms 左右，但编译本身有开销，且对动态控制流支持一般。JAX 版本在 JIT 之后稳定在 38-42ms，而且因为 XLA 会把整个计算图融合优化，算子间的 kernel launch 开销被大幅削减。在 RTC 这种“小块高频”的场景下，kernel launch 开销占比很高，JAX 的优势就体现出来了。
当然代价也有：JAX 的调试体验确实不如 PyTorch 直观，报错信息经常让人摸不着头脑，而且动态 shape 支持较弱。但既然 RTC 的块 shape 是固定的，这个代价可以接受。
3. 核心细节拆解：JIT 编译、KV Cache 与分块对齐
3.1 JIT 编译策略：怎么避免“编译抖动”
JIT 编译是 JAX 性能的命门，也是抖动的最大来源。我的做法是预热编译 + 静态 shape 锁定。
具体来说，在服务启动阶段，先用一个 dummy 输入跑一遍完整的推理链路，触发所有函数的编译。这个 dummy 输入的 shape 必须和实际推理时完全一致，包括 batch size、序列长度、图像分辨率。任何一处不一致，都会导致运行时重新编译，延迟瞬间飙到秒级。
import jax import jax.numpy as jnp # 预热：用真实 shape 的 dummy 输入触发编译 dummy_obs = jnp.zeros((1, 3, 224, 224)) dummy_state = jnp.zeros((1, 32)) dummy_prefix = jnp.zeros((1, L, action_dim)) # 这一步会触发编译，耗时约 1-2 秒，但只发生一次 _ = jitted_infer(dummy_obs, dummy_state, dummy_prefix)这里有个细节：jax.jit默认会对输入 shape 和 dtype 做缓存，但如果你用了static_argnums，那被标记为 static 的参数值变了也会重编译。所以我把所有会影响计算图的参数（比如块长度 L、是否使用 KV Cache）都设成 static，并且在启动时就固定下来，运行时不改。
注意：预热用的 dummy 输入一定要覆盖所有可能的分支。比如如果你的模型有“是否使用历史动作”的条件分支，那两种情况的输入都要预热一遍，否则第一次走到未预热的分支时还是会卡。
3.2 KV Cache 的跨块复用：省下的不只是时间
RTC 分块推理时，每个块都需要前面所有块的上下文。最朴素的做法是每个块都把完整历史重新喂进去，但这样计算量随块数线性增长，到后面几块延迟会明显上升。正确做法是维护一个 KV Cache，每个块只计算新增部分的 KV，然后拼接。
在 JAX 里实现 KV Cache 有个坑：JAX 的数组是不可变的，你不能像 PyTorch 那样原地更新 cache。我的做法是用jax.lax.dynamic_update_slice来做原地更新的语义，同时配合donate_argnums把旧 cache 的内存捐出去，避免额外分配。
def update_kv_cache(cache, new_kv, pos): # cache shape: (num_layers, 2, max_len, num_heads, head_dim) # new_kv shape: (num_layers, 2, L, num_heads, head_dim) return jax.lax.dynamic_update_slice(cache, new_kv, (0, 0, pos, 0, 0))实测下来，用了 KV Cache 复用之后，第 8 块的推理延迟从 52ms 降到了 41ms，而且块数越多，收益越明显。这是因为省下的不只是计算量，还有内存带宽——重新计算历史 KV 需要反复读写大数组，而复用只需要读一次。
3.3 分块对齐：L 取多少才不崩连续性
块长度 L 的选择是个权衡：L 太小，块数多，每块都要走一遍完整的前向，总延迟高；L 太大，首帧延迟高，实时性差。我做了几组对比实验：
| 块长度 L | 块数 K | 首帧延迟 | 单块平均延迟 | 连续性评分 |
|---|---|---|---|---|
| 8 | 16 | 22ms | 35ms | 0.92 |
| 16 | 8 | 38ms | 41ms | 0.96 |
| 32 | 4 | 65ms | 48ms | 0.98 |
| 64 | 2 | 110ms | 62ms | 0.99 |
连续性评分是我自己定义的一个指标，用相邻块边界处的动作二阶差分来衡量，越接近 1 越平滑。从数据看，L=16 是个甜点：首帧延迟 38ms 可接受，连续性 0.96 已经足够好，单块延迟也稳定。
但这里有个隐藏约束：L 必须能整除模型训练时的动作序列长度，否则边界处的上下文会对不齐，连续性会突然掉到 0.8 以下。我一开始试了 L=12，结果边界处动作明显跳变，查了半天才发现是训练时序列长度是 64，12 不能整除导致的。
4. 实操过程：从 40ms 抖动到 60ms 内 P99 的完整记录
4.1 环境准备与依赖锁定
环境这块我踩过最大的坑是 JAX 版本和 CUDA 版本的匹配。JAX 对 CUDA 版本很敏感，版本不匹配时不会报错，而是静默回退到 CPU，性能直接掉一个数量级。我的配置是：
# 关键依赖版本 jax==0.4.30 jaxlib==0.4.30 cuda==12.1 cudnn==8.9安装时用官方推荐的 pip 源，不要混用 conda 和 pip，否则容易出现库冲突。装完之后一定要验证设备：
import jax print(jax.devices()) # 应该输出 [CudaDevice(id=0)]，如果是 [CpuDevice] 就说明没装对提示：如果你用的是多卡环境，记得用
jax.local_devices()确认每张卡都被识别到。我遇到过一张卡因为驱动问题没被识别，结果所有计算都挤到另一张卡上，延迟翻倍。
4.2 推理函数的 JIT 化改造
原始版本的推理函数是纯 Python 写的，里面有不少if-else和循环，这些在 JAX 里都是性能杀手。改造的核心原则是：把所有控制流改成jax.lax.cond和jax.lax.scan，把所有循环改成向量化操作。
@jax.jit def infer_block(params, obs, state, kv_cache, pos): # 图像编码 img_feat = image_encoder(params['img'], obs) # 状态拼接 token = jnp.concatenate([img_feat, state], axis=-1) # 主干前向，带 KV Cache new_kv, hidden = backbone(params['backbone'], token, kv_cache, pos) # 动作头解码 action = action_head(params['head'], hidden) # 后处理：裁剪到合法范围 action = jnp.clip(action, -1.0, 1.0) return action, new_kv这里的关键改动是把后处理的clip也放进 JIT 里。原本clip是在 Python 里用 numpy 做的，每次都要把设备上的数组搬到 CPU 再搬回去，光这一来一回就是 5-8ms。放进 JIT 后，这部分开销直接归零。
4.3 RTC 调度器的实现与参数调优
RTC 调度器负责管理块的生成和执行节奏。我的实现是一个双缓冲队列：一个缓冲区放正在执行的块，另一个放正在生成的块，两者交替。这样当执行器在跑第 i 块时，推理器已经在生成第 i+1 块了。
class RTCScheduler: def __init__(self, infer_fn, block_len, num_blocks): self.infer_fn = infer_fn self.block_len = block_len self.num_blocks = num_blocks self.kv_cache = init_kv_cache() self.pos = 0 def step(self, obs, state): action, self.kv_cache = self.infer_fn(obs, state, self.kv_cache, self.pos) self.pos += self.block_len return action参数调优上，我重点调了两个：预生成块数和执行器提前量。预生成块数设为 2 时，延迟最稳；设为 3 时，虽然理论上更平滑，但显存占用增加，反而因为内存压力导致抖动。执行器提前量设为 1 个块的时间（约 40ms），这样执行器永远不会等推理器。
4.4 实测数据与对比
优化前后的对比数据如下：
| 指标 | 优化前 | 优化后 |
|---|---|---|
| 首帧延迟 | 120ms | 38ms |
| 单块平均延迟 | 40ms | 36ms |
| P99 延迟 | 120ms | 58ms |
| 连续性评分 | 0.88 | 0.96 |
| 显存占用 | 4.2GB | 3.8GB |
P99 从 120ms 压到 58ms，主要贡献来自三块：JIT 预热消除了首次编译抖动（贡献约 30ms）、KV Cache 复用降低了后续块延迟（贡献约 15ms）、后处理搬进 JIT 消除了 CPU-GPU 往返（贡献约 8ms）。剩下的抖动来自设备同步和调度器本身的 Python 开销，这部分很难完全消除，但 58ms 已经满足下游控制器的要求。
5. 常见问题与排查技巧实录
5.1 延迟突然翻倍：先查设备再查代码
遇到延迟突然翻倍，第一件事不是看代码，而是确认计算到底跑在哪个设备上。JAX 有个很坑的行为：如果某个操作不支持当前设备，它会静默回退到 CPU，而且不报错。我遇到过一次，因为某个jnp函数在 CUDA 上没实现，整个前向被拆成两半，一半在 GPU 一半在 CPU，延迟直接翻三倍。
排查方法很简单：
with jax.profiler.trace("/tmp/jax-trace"): result = infer_block(...) # 然后用 TensorBoard 打开 trace，看每个算子的设备5.2 连续性崩了：九成是分块边界没对齐
动作块之间出现跳变，最常见的原因是分块边界和训练时的序列边界没对齐。检查方法：把相邻块边界处的动作打印出来，看二阶差分是否突变。如果是，先确认 L 能否整除训练序列长度，再确认 KV Cache 的 pos 是否正确递增。
还有一个隐蔽原因：位置编码的偏移。如果你的模型用了 RoPE 之类的相对位置编码，每个块的起始位置必须正确传入，否则位置编码会错位，导致动作漂移。我在这上面卡了一整天，最后发现是 pos 在块之间没有正确累加。
5.3 显存不够：donate 和 dtype 两手抓
显存不够时，优先用donate_argnums把不再需要的数组捐出去。比如 KV Cache 更新后，旧的 cache 就没用了，直接 donate 掉能省下一半显存。其次是把不必要的 float32 降到 bfloat16，JAX 对 bfloat16 的支持很好，精度损失在策略模型上几乎看不出来，但显存能省 40%。
# 用 bfloat16 存储 KV Cache kv_cache = kv_cache.astype(jnp.bfloat16)注意：bfloat16 在累加操作上精度损失较大，如果你的模型有大量累加（比如 attention 的 softmax），建议只在存储时用 bfloat16，计算时转回 float32。
5.4 常见问题速查表
| 现象 | 可能原因 | 排查方法 |
|---|---|---|
| 延迟突然翻倍 | 计算回退到 CPU | jax.profiler看算子设备 |
| 动作块边界跳变 | 分块边界未对齐 | 打印边界处二阶差分 |
| 显存溢出 | KV Cache 未复用或未 donate | 检查donate_argnums配置 |
| 首次推理极慢 | JIT 未预热 | 启动时用 dummy 输入预热 |
| 多卡负载不均 | 设备未全部识别 | jax.local_devices()确认 |
6. 几个我踩过的坑和最后的小技巧
第一个坑是过度依赖jax.jit的自动缓存。JAX 的 JIT 缓存是基于输入 shape 和 dtype 的，但如果你在函数里用了 Python 的全局变量，那缓存不会感知到变量变化，会导致结果错误。我的做法是把所有外部依赖都显式作为参数传入，绝不在 JIT 函数里引用全局状态。
第二个坑是调度器的 Python 开销。RTC 调度器本身是 Python 写的，每块都要走一遍 Python 逻辑，这部分开销在 40ms 的块延迟里占了 3-5ms。后来我把调度器的核心逻辑也用 JAX 重写了一遍，虽然可读性差了，但延迟降到了 1ms 以内。
最后分享一个小技巧：如果你发现 P99 延迟总是比平均值高很多，但找不到明显原因，试试把推理函数的donate_argnums打开，同时把XLA_FLAGS里的--xla_gpu_autotune_level调到 4。前者减少内存分配抖动，后者让 XLA 花更多时间做算子调优，两者结合能把 P99 再压 5-8ms。这个技巧在块长度较小、调用频率高的场景下特别有效。
