PJRT
Also known as PJRT, PJRT buffer, PJRT plugin
In plain English
PJRT is the runtime interface that sits between a framework like JAX or TorchTPU and the TPU hardware, managing device memory buffers and execution.
Technical definition
PJRT is the plugin runtime API that XLA-based frameworks use to allocate device buffers, launch compiled executables, and transfer data, with PJRTBuffer objects wrapping the raw hardware memory descriptors for a tensor on a device.
Engineering details
JAX arrays and TorchTPU tensors on TPU are backed by PJRT buffers. A library that can extract the native PJRTBuffer hardware descriptor can move data directly between devices without staging through the framework, which is how TPU-Sync performs zero-copy KV cache transfers for disaggregated serving. Because both JAX and TorchTPU speak PJRT, one transfer library serves both stacks.
Why it matters
The runtime layer is where disaggregation and offload plumbing lives. Zero-copy transfers avoid an extra HBM round trip per KV block, which matters when a prefill pool ships large caches to decode pools at high request rates. The same descriptors are the likely substrate for Mooncake Store support on TPU.
How to read it in InferenceX
The TPU InferenceX preview describes TPU-Sync, formerly TPU-raiden, as extracting native PJRTBuffer descriptors for zero-copy transfers and supporting DRAM offload. Disaggregated TPUv7 results built on it are promised in a follow-up article.