OxyGen: Unified KV Cache Management for VLA Inference under Multi-Task Parallelism
Xiangyu Li ⋅ Huaizhi Tang ⋅ Xin Ding ⋅ Weijun Wang ⋅ Ting Cao ⋅ Yunxin Liu
Abstract
Embodied AI agents increasingly require parallel execution of multiple tasks, such as manipulation, conversation, and memory construction, from shared observations under distinct time constraints. Recent Mixture-of-Transformers (MoT) Vision-Language-Action Models (VLAs) architecturally support such heterogeneous outputs, yet existing inference systems fail to achieve efficient multi-task parallelism for on-device deployment due to redundant computation and resource contention. We identify isolated KV cache management as the root cause. To address this, we propose *unified KV cache management*, an inference design that treats KV cache as a first-class shared resource across tasks and over time. This abstraction enables two key optimizations: *cross-task KV sharing* eliminates redundant prefill of shared observations, while *cross-frame continuous batching* decouples variable-length language decoding from fixed-rate action generation across control cycles. We implement this design for $\pi_{0.5}$, the most popular MoT VLA, and evaluate on both NVIDIA GeForce RTX 4090 and Jetson AGX Thor, two representative platforms for on-device VLA inference. \nickname achieves up to 3.7$\times$ speedup over isolated execution, delivering over 200 tokens/s language throughput and 70 Hz action frequency simultaneously without action quality degradation, and we further validate the gains on a real humanoid robot with on-board Jetson AGX Thor.
Chat is not available.
Successful Page Load