OxyGen: Unified KV Cache Management for Vision-Language-Action Models under Multi-Task Parallelism

Imagine you have a highly intelligent robot assistant named "Robo." Robo is trying to do three things at once:

Move its arm to pick up a cup (Action).
Talk to you about what it's doing (Language).
Remember where it put the cup for later (Memory).

In the past, if you asked Robo to do all these things, it was like hiring three different people to do the same job, but they were all standing in a tiny kitchen trying to use the same stove. They kept bumping into each other, waiting in line, and re-cooking the same ingredients over and over again. The result? Robo was slow, clumsy, and often dropped the cup because it couldn't move its arm fast enough while it was busy talking.

This paper introduces OxyGen, a new way to manage Robo's brain so it can multitask smoothly. Here is how it works, using simple analogies:

The Problem: The "Isolated Kitchen"

Current robot systems treat every task as if it's happening in a separate, isolated kitchen.

Redundant Cooking: If Robo needs to look at a picture of a cup to move its arm and to talk about the cup, the old system makes the robot "cook" (process) that picture twice. It's a waste of time and energy.
The Traffic Jam: Even if they share the stove, the "Talking" task might take a long time to finish a sentence. While it's talking, the "Moving" task has to wait, even though moving the arm needs to happen instantly (like a reflex). This causes the robot to freeze or move jerkily.

The Solution: OxyGen (The Unified Manager)

OxyGen acts like a super-efficient head chef who manages the entire kitchen as one big, shared space. It uses two main tricks:

1. The "Shared Recipe Book" (Cross-Task KV Sharing)

In AI, the "KV Cache" is like a scratchpad or a working memory where the robot writes down what it has already seen and understood.

Old Way: Every time the robot starts a new task (move, talk, remember), it throws away the scratchpad and starts writing from scratch, even if it's looking at the same cup.
OxyGen Way: The head chef says, "Hey, we already looked at the cup! Let's just keep that page in the scratchpad."
The Result: The robot doesn't waste time re-reading the cup. It instantly shares that understanding between the arm-moving task and the talking task. It's like sharing a single Google Doc instead of emailing three different Word files back and forth.

2. The "Assembly Line" (Cross-Frame Continuous Batching)

Robots have different deadlines. Moving an arm needs to happen now (every 1/60th of a second). Talking can happen a bit slower, over several seconds.

Old Way: The robot stops everything to finish a whole sentence before moving its arm again. This makes the arm stop and start (jerky motion).
OxyGen Way: Imagine an assembly line. While the robot is moving its arm for this second, it is also quietly working on the next few words of the sentence in the background. It batches (groups) the talking tasks together so the computer chip works on them all at once, efficiently.
The Result: The arm moves smoothly and quickly (like a human), while the robot is simultaneously churning out a long, detailed story without slowing down.

The Real-World Impact

The researchers tested this on a powerful robot brain (called $\pi0.5$ ) running on a standard gaming graphics card (RTX 4090).

Speed: The robot became 3.7 times faster.
Smoothness: It could move its arm at 70 times per second (super smooth) while still talking at a very fast speed (200 words per second).
Quality: Crucially, it didn't make mistakes. The robot still picked up the cup correctly; it just did it much faster and without getting "tired" from the traffic jams.

Summary

Think of OxyGen as the difference between a chaotic kitchen where three cooks fight over one stove, versus a well-organized kitchen with one head chef who shares ingredients and runs an assembly line. It allows our robot friends to finally do what humans do best: walk and chew gum at the same time, but with the speed and precision of a machine.

1. Problem Statement

Embodied AI agents increasingly require multi-task parallelism, where a single agent must simultaneously execute diverse tasks (e.g., robot manipulation, conversational narration, and memory construction) based on shared visual observations. While recent Mixture-of-Transformers (MoT) Vision-Language-Action (VLA) models (like $\pi^0.5$ ) architecturally support generating heterogeneous outputs (actions and language) from a single model, existing inference systems fail to leverage this efficiently on resource-constrained devices (e.g., single GPUs).

The paper identifies isolated KV cache management as the root cause of inefficiency in current systems. Existing approaches treat each task independently, leading to two critical bottlenecks:

Redundant Computation: Shared observations are encoded repeatedly for each task, generating identical Key-Value (KV) cache entries multiple times.
Resource Contention: Tasks compete for limited hardware resources (GPU memory/bandwidth) without coordination. Crucially, tasks have asymmetric deadlines: action generation requires hard, real-time deadlines (e.g., 50–70 Hz) for smooth control, while language generation has soft deadlines and can span multiple frames. Current systems cannot decouple these constraints, leading to significant slowdowns (up to 2.6 $\times$ ) and reduced throughput.

2. Methodology: OxyGen

The authors propose OxyGen, an inference system built on a novel paradigm called Unified KV Cache Management. This approach treats the KV cache not as a task-specific artifact, but as a first-class, shared resource across tasks and time.

Core Components:

Unified KV Cache Manager ( $\mathcal{M}$ ): A central scheduler that maintains the state of all in-flight requests. It tracks generation states ( $\sigma_t$ ) consisting of the KV cache, token buffers, and termination flags.
Optimization 1: Cross-Task KV Sharing:
- Instead of re-encoding the visual observation for both the action expert and the language expert, the system performs a single prefill pass to generate the shared KV cache ( $K_t$ ).
- This $K_t$ is then duplicated and consumed by both experts within the same frame, eliminating redundant prefill computation.
Optimization 2: Cross-Frame Continuous Batching:
- The system decouples language decoding from the rigid per-frame control loop.
- Action tasks are processed immediately within the current frame to meet hard deadlines.
- Language tasks are treated as streaming requests. The manager aggregates multiple active language requests from different frames into a single batched state ( $\hat{\sigma}$ ).
- The VLM performs autoregressive decoding on this batch in parallel, amortizing the decoding cost across multiple requests and maximizing GPU utilization.

Execution Flow (Per Frame):

Prefill: Encode the new observation $o_t$ once to produce shared $K_t$ .
Action Generation: Immediately use $K_t$ to generate the action chunk $A_t$ (via diffusion/flow matching).
State Management: Initialize a new language request state $\sigma_t$ using $K_t$ and store it in the manager.
Batched Decoding: Retrieve all active language requests (including the new one), batch them, and perform $k$ steps of parallel decoding.
Update: Unbatch the results, update the manager with new tokens, and evict completed requests.

3. Key Contributions

Problem Formulation: The paper formally defines multi-task parallelism with asymmetric deadlines as a critical inference scenario for MoT-VLAs and identifies isolated KV cache management as the primary bottleneck.
Unified KV Cache Paradigm: Introduces a novel inference paradigm that abstracts KV cache as a shared resource, enabling cross-task sharing (eliminating redundant prefill) and cross-frame continuous batching (decoupling variable-length language decoding from fixed-rate action generation).
System Implementation & Evaluation: Implemented OxyGen for $\pi^0.5$ (the most popular MoT VLA) and validated it on representative robotic benchmarks (LIBERO, DROID, ALOHA) using an NVIDIA RTX 4090.

4. Experimental Results

Evaluated on a single NVIDIA RTX 4090 GPU, OxyGen demonstrates significant performance gains over both sequential and naive parallel baselines (using CUDA MPS):

Speedup: Achieves up to 3.7 $\times$ speedup in action frequency and language throughput compared to isolated execution.
Throughput: Simultaneously achieves >200 tokens/s for language decoding and 70 Hz for action generation without degrading control quality.
Ablation Studies:
- Cross-task KV sharing alone provides a ~1.4 $\times$ speedup by removing redundant prefill.
- Cross-frame continuous batching is crucial for long decoding tasks, maintaining high action frequency (approx. 60 Hz) even as language decoding steps increase, whereas baselines drop to ~19 Hz.
Quality & Efficiency:
- Action Quality: Task success rates on LIBERO benchmarks match the original $\pi^0.5$ performance (within statistical noise), confirming no degradation in control quality.
- Energy Efficiency: Reduces energy per request by up to 78% and average power consumption by 47% compared to baselines, due to reduced memory access and better hardware utilization.
- Memory: Adds only ~15% memory overhead compared to the baseline, whereas naive parallelization nearly doubles memory usage.

5. Significance

OxyGen addresses a critical gap in deploying advanced Embodied AI agents. By rethinking KV cache management as a unified resource, it enables MoT-VLAs to truly realize their architectural potential for simultaneous multi-task execution. This work demonstrates that efficient on-device inference for complex, multi-modal agents is achievable without sacrificing control frequency or language fluency, paving the way for more capable, real-time, and energy-efficient home and industrial robots.