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

TL;DR

OxyGen introduces a unified KV cache management system for multi-task VLA inference, significantly improving efficiency and throughput on edge devices by sharing resources across tasks and decoupling decoding processes.

Contribution

It presents a novel inference design that treats KV cache as a shared resource, enabling cross-task sharing and cross-frame batching for efficient multi-task VLA inference.

Findings

Achieves up to 3.7× speedup over isolated execution.

Delivers over 200 tokens/s language throughput.

Maintains 70 Hz action frequency without degrading quality.

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 because of 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 the 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}$, a popular MoT VLA, and evaluate it on both NVIDIA GeForce RTX 4090 and Jetson AGX Thor, two representative platforms for on-device VLA inference. OxyGen achieves up to 3.7$\times$ speedup over isolated execution, delivering over 200 tokens/s language throughput and 70 Hz action frequency simultaneously without degrading action quality, and we further validate the gains on a real humanoid robot with on-board Jetson AGX Thor.