1. 从一次推理延迟抖动说起为什么要折腾 JAX 版 Pi05 与 RTC先说结论这套东西的核心目标只有一个——把 Pi05 这类策略模型在 JAX 上的推理链路压到实时可控让 RTCReal-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首帧延迟单块平均延迟连续性评分81622ms35ms0.9216838ms41ms0.9632465ms48ms0.98642110ms62ms0.99连续性评分是我自己定义的一个指标用相邻块边界处的动作二阶差分来衡量越接近 1 越平滑。从数据看L16 是个甜点首帧延迟 38ms 可接受连续性 0.96 已经足够好单块延迟也稳定。但这里有个隐藏约束L 必须能整除模型训练时的动作序列长度否则边界处的上下文会对不齐连续性会突然掉到 0.8 以下。我一开始试了 L12结果边界处动作明显跳变查了半天才发现是训练时序列长度是 6412 不能整除导致的。4. 实操过程从 40ms 抖动到 60ms 内 P99 的完整记录4.1 环境准备与依赖锁定环境这块我踩过最大的坑是 JAX 版本和 CUDA 版本的匹配。JAX 对 CUDA 版本很敏感版本不匹配时不会报错而是静默回退到 CPU性能直接掉一个数量级。我的配置是# 关键依赖版本 jax0.4.30 jaxlib0.4.30 cuda12.1 cudnn8.9安装时用官方推荐的 pip 源不要混用 conda 和 pip否则容易出现库冲突。装完之后一定要验证设备import jax print(jax.devices()) # 应该输出 [CudaDevice(id0)]如果是 [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 块时推理器已经在生成第 i1 块了。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 实测数据与对比优化前后的对比数据如下指标优化前优化后首帧延迟120ms38ms单块平均延迟40ms36msP99 延迟120ms58ms连续性评分0.880.96显存占用4.2GB3.8GBP99 从 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 降到 bfloat16JAX 对 bfloat16 的支持很好精度损失在策略模型上几乎看不出来但显存能省 40%。# 用 bfloat16 存储 KV Cache kv_cache kv_cache.astype(jnp.bfloat16)注意bfloat16 在累加操作上精度损失较大如果你的模型有大量累加比如 attention 的 softmax建议只在存储时用 bfloat16计算时转回 float32。5.4 常见问题速查表现象可能原因排查方法延迟突然翻倍计算回退到 CPUjax.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。这个技巧在块长度较小、调用频率高的场景下特别有效。
