OxyGen: Unified KV Cache Management for VLA Inference under Multi-Task Parallelism

arXiv:2603.14371 · cs.RO, cs.AI · Submitted 2026-03-15 · Read on arXiv

cs.RO, cs.AI

Submitted: 2026-03-15

Updated: 2026-10-14

Comments: NeurIPS 2026

Code: https://github.com/air-embodied-brain/OxyGen

License: http://creativecommons.org/licenses/by/4.0/

The gist: Embodied AI agents increasingly require parallel execution of multiple tasks, such as manipulation, conversation, and memory construction, from shared observations under distinct time constraints.

Terminology

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) support such heterogeneous outputs, yet existing inference systems incur redundant computation and resource contention when serving concurrent tasks on-device. We identify isolated KV cache management as the root cause. To address this, we propose unified KV cache management, an inference design that treats the KV cache as a first-class shared resource across tasks and control cycles. This abstraction enables two key optimizations: cross-task KV sharing eliminates redundant prefill of shared observations, while cross-frame continuous batching batches unfinished language requests from different frames and advances them within a per-frame decoding budget. We implement this design in OxyGen and evaluate π 0.5 on NVIDIA GeForce RTX 4090 and Jetson AGX Thor. OxyGen achieves up to 3.7 times speedup over isolated execution, delivering over 200 tokens/s language throughput and 70 Hz action frequency simultaneously on RTX 4090. We also validate the design on Xiaomi-Robotics-0 and StarVLA across different backbones and expert configurations. We further demonstrate the inference gains on a humanoid robot with on-board Jetson AGX Thor.

Sources

Related papers