软件栈
PJRT
也称为 PJRT、PJRT buffer、PJRT 插件
先用大白话
PJRT 是位于 JAX、TorchTPU 等框架与 TPU 硬件之间的运行时接口，负责管理设备内存缓冲区和执行。
技术定义
PJRT 是基于 XLA 的框架用来分配设备缓冲区、启动编译后可执行文件和传输数据的插件式运行时 API，PJRTBuffer 对象封装了张量在设备上的原始硬件内存描述符。
工程细节
TPU 上的 JAX 数组和 TorchTPU 张量都由 PJRT 缓冲区支撑。能够提取原生 PJRTBuffer 硬件描述符的库可以在设备之间直接搬运数据，不经过框架中转，TPU-Sync 正是这样为分离式服务实现零拷贝 KV cache 传输的。由于 JAX 和 TorchTPU 都使用 PJRT，一个传输库可以同时服务两套栈。
为什么重要
运行时层是分离式服务和 offload 管道所在的位置。零拷贝传输省去了每个 KV block 一次额外的 HBM 往返，当 prefill 池以高请求速率向 decode 池发送大量 cache 时尤为重要。同样的描述符很可能也是 TPU 上 Mooncake Store 支持的底层基础。
如何在 InferenceX 中解读
TPU InferenceX 预览描述 TPU-Sync（原名 TPU-raiden）通过提取原生 PJRTBuffer 描述符实现零拷贝传输，并支持 DRAM offload。基于它的 TPUv7 分离式服务结果将在后续文章发布。