OxyGen: Unified KV Cache Management for VLA Inference under Multi-Task Parallelism
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
- A Modern System Recipe for Situated Embodied Human-Robot Conversation with Real-Time Multimodal LLMs and Tool-Calling
- Hi Robot: Open-Ended Instruction Following with Hierarchical Vision-Language-Action Models
- Mixture-of-Transformers: A Sparse and Scalable Architecture for Multi-Modal Foundation Models
- OpenVLA: An Open-Source Vision-Language-Action Model
- $\pi_0$: A Vision-Language-Action Flow Model for General Robot Control
- $\pi_{0.5}$: a Vision-Language-Action Model with Open-World Generalization
- GR00T N1: An Open Foundation Model for Generalist Humanoid Robots
- Igniting VLMs toward the Embodied Space
- CogACT: A Foundational Vision-Language-Action Model for Synergizing Cognition and Action in Robotic Manipulation
- Galaxea Open-World Dataset and G0 Dual-System VLA Model
- DexVLA: Vision-Language Model with Plug-In Diffusion Expert for General Robot Control
- UniVLA: Learning to Act Anywhere with Task-centric Latent Actions
- RynnVLA-002: A Unified Vision-Language-Action and World Model
- FAST: Efficient Action Tokenization for Vision-Language-Action Models
- How Fast Can I Run My VLA? Demystifying VLA Inference Performance with VLA-Perf
- Running VLAs at Real-time Speed
- BitVLA: 1-bit Vision-Language-Action Models for Robotics Manipulation
- SQAP-VLA: A Synergistic Quantization-Aware Pruning Framework for High-Performance Vision-Language-Action Models
- EfficientVLA: Training-Free Acceleration and Compression for Vision-Language-Action Models
- Quantization-Aware Imitation-Learning for Resource-Efficient Robotic Control
Related papers
- FMT x: An Efficient and Asymptotically Optimal Extension of the Fast Marching Tree for Dynamic Replanning
- MPCFormer: A physics-informed data-driven approach for explainable socially-aware autonomous driving
- RoboLab: A High-Fidelity Simulation Benchmark for Analysis of Task Generalist Policies
- ARC: A Reasoning Recipe for Robot Foundation Models
- HRDexDB: A 4D Dexterous Grasping Dataset Across Human and Multiple Robot Embodiments
- APT: Action Expert Pretraining Improves Instruction Generalization of Vision-Language-Action Policies