Source-linked AI summary
OxyGen: Unified KV Cache Management for VLA Inference under Multi-Task Parallelism
Xiangyu Li, Huaizhi Tang, Xin Ding, Weijun Wang, Ting Cao, Yunxin Liu
TL;DR
On-device MoT VLA systems struggle with parallel tasks because isolated KV cache management causes redundant computation and resource contention. OxyGen unifies KV cache management across tasks and time, using cross-task sharing and cross-frame batching, and reports up to 3.7× speedup with concurrent language and action performance on robotic platforms.
Problem
Existing MoT VLA inference systems use isolated KV caches, causing redundant shared-observation computation and resource contention during multi-task execution with different time constraints.
Method
OxyGen treats KV cache as a unified shared resource, encoding shared observations once and batching language decoding continuously across control frames while preserving action deadlines.
Results
3.7× speedup is achieved for π0.5 across RTX 4090 and Jetson AGX Thor, with over 200 tokens/s language throughput and sustained action frequencies reported on both platforms.
Takeaways & Limitations
Unified KV cache management supports efficient multi-task MoT VLA inference on robotic hardware and is validated on a humanoid robot with on-board Jetson AGX Thor.
Takeaways & Limitations
The evaluation covers one MoT VLA, π0.5, with two experts, leaving generalization to other backbones and richer expert configurations for future work.
Abstract
from arXiv · showhide
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 $π_{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.
1. Introduction
OxyGen addresses inefficient multi-task parallelism in MoT VLAs by unifying KV cache management across tasks and time. Its optimizations reduce redundant computation and coordinate language and action execution on robotic hardware.
- Motivation: Multi-task parallelism runs temporally independent tasks concurrently from shared observations, each with its own time constraints.Examples include manipulating while memorizing environmental changes, narrating progress, and planning ahead.
- Motivation: MoT VLAs support language and action outputs through modality-specific experts, but this capability does not automatically improve inference speed over naive multi-model inference.Earlier conventional VLAs are restricted to action outputs and require multiple models for heterogeneous tasks.
- Problem: 1.4× slowdown comes from repeatedly encoding shared observations, while 2.6× slowdown comes from tasks competing for limited hardware despite different time constraints.Action denoising must finish within each control cycle, whereas language decoding may span multiple frames.
- Approach: OxyGen treats KV cache as a shared resource across tasks and time through cross-task KV sharing and cross-frame continuous batching.Shared observations are encoded once and exposed through per-expert views, while language requests advance by calibrated budgets across frames.
- Results: 3.7× speedup is achieved for π0.5 across LIBERO, DROID, and ALOHA configurations on RTX 4090 and Jetson AGX Thor.The system sustains over 200 tokens/s language throughput at 60 Hz action frequency on RTX 4090 and 27 Hz on Jetson AGX Thor, and is validated on a humanoid robot.
- Contributions: OxyGen is presented as an efficient MoT VLA inference system that demonstrates unified KV cache management on two representative robotic platforms.The contribution summary identifies unified KV cache management as the basis for its reported acceleration.
2. Related Work
Prior work optimizes VLAs, LLM KV reuse, or aggregate cloud throughput, but on-device MoT VLA inference must coordinate language throughput with a hard per-frame action deadline. OxyGen targets this distinct scheduling problem rather than applying cloud-serving techniques directly.
- MoT VLAs: MoT VLAs share a VLM backbone while routing action and language outputs to separate experts, yet existing systems still use independent forward passes.This yields no acceleration over naive multi-model inference despite architectural multitasking.
- VLA inference optimizations: VLA optimizations span model-level techniques such as compression, token pruning, and KV pruning, and application-level techniques such as asynchronous pipelines.OxyGen manages KV cache at the model level without modifying operators.
- KV cache reuse: LLM serving systems study prefix, nonprefix, and cross-model KV reuse for cloud-scale inference aimed at maximizing aggregate request throughput.These systems assume abundant compute and high concurrency, unlike on-device robotic inference.
- On-device setting: On-device MoT VLA inference must maximize language throughput subject to a hard per-frame action deadline tied to the physical control loop.Directly applying prefix caching is insufficient because uncoordinated language decoding can block action generation and violate the frame budget.
3. Method
OxyGen manages KV caches as a unified resource across tasks and frames for MoT VLA inference. It combines shared-observation prefill with resumable, cross-frame language decoding while preserving action deadlines.
- Unified KV cache management: OxyGen treats the KV cache from the shared VLM backbone as a shared resource for computation reuse and execution coordination.The cache encodes a common observation and can be consumed by multiple modality-specific experts.
- MoT VLA inference: At each frame, modality-agnostic prefill produces K_t, which conditions both action-chunk generation and autoregressive language decoding.The action expert generates A_t through S denoising steps, while the language expert generates y_t token by token.
- Cross-task KV sharing: Cross-task KV sharing fans one prefilled K_t out to action and language experts, eliminating repeated encoding of the shared observation.The experts retain separate access semantics while consuming the same within-frame prefill cache.
- Cross-frame continuous batching: Cross-frame continuous batching groups in-flight language requests from different frames, decoupling language decoding from the per-frame control loop.Resumable request states preserve each request’s KV cache, generated-token buffer, and termination flag, allowing active requests to advance in one batched forward pass.
- Per-frame execution: The system calibrates k, the number of language decoding steps per frame, so batched decoding fits after action denoising while preserving the minimum action frequency.Cross-task sharing provides the initial speedup, while continuous batching further reduces latency as decoding length increases.
- Generalization: OxyGen’s manager interface extends to more than two experts and remains compatible with speculative or parallel decoding schemes.These extensions rely on treating experts as opaque consumers of K_t and retaining resumable generation state.
4. Experiments
OxyGen is evaluated against sequential and parallel isolated execution across robotic configurations, hardware platforms, workloads, and ablations. It improves the action-frequency/language-throughput tradeoff through unified KV management while preserving action quality and enabling on-device deployment.
- Experimental Setup: OxyGen is evaluated with π0.5 on LIBERO, DROID, and ALOHA configurations using GeForce RTX 4090 and Jetson AGX Thor.The evaluation measures action frequency, language throughput, and average batch size.
- Baselines: The baselines comprise sequential isolated execution and parallel isolated execution, implemented with CUDA Multi-Process Service on one GPU.The parallel baseline runs each task in an individual process, while sequential execution runs tasks independently within each frame.
- End-to-End Results: 1.2–3.7× speedup is achieved on both action frequency and language throughput, extending the Pareto frontier instead of trading one axis for the other.Larger N and smaller k increase speedup by producing larger average batches, which better utilize hardware parallelism.
- End-to-End Results: Naive parallelization provides modest improvement, showing that concurrent execution alone does not remove redundant computation.The result supports the need for coordinated KV-cache management rather than simple process-level parallelism.
- Ablation Study: Cross-task KV sharing provides a 1.4× initial speedup on RTX 4090, while cross-frame continuous batching maintains around 60 Hz on RTX 4090 and 27 Hz on Jetson AGX Thor for decoding steps ≥10.The gap from the single-frame oracle remains modest on both platforms, indicating scheduling overhead.
5. Conclusion
OxyGen presents unified KV cache management for efficient MoT VLA inference under multi-task parallelism. Its reported scope includes speed improvements, lower deployment cost, and a limitation to one two-expert VLA.
- OxyGen treats the KV cache as a shared resource across tasks and over time, enabling cross-task KV sharing and cross-frame continuous batching.
- Up to 3.7× speedup is achieved on NVIDIA GeForce RTX 4090 and Jetson AGX Thor for multi-task MoT VLA inference.
- The empirical evaluation focuses on π0.5 with two experts, leaving generalization to other MoT backbones and richer multi-expert configurations for future work.
- OxyGen lowers the compute and energy cost of deploying multi-task embodied agents on commodity robotic hardware.
A. Method Details
The method details define persistent per-request language states, active-request batching, and per-frame execution around shared KV caches. The execution flow combines new-request prefill, action generation, and batched language decoding.
- Per-frame execution: The action expert generates an action chunk through iterative denoising, with every denoising step conditioned on the shared KV cache.
- State persistence: The manager exposes Store, Retrieve, Update, and Remove operations for persistent per-request generation states.
- Batched decoding: Active requests are maintained in R, and newly created requests immediately join batched decoding with unfinished requests from earlier frames.
- Batched decoding: The batched state lets the VLM decode all active requests in parallel in one forward pass, amortizing decoding cost across requests.
- Per-frame execution: At each frame, the system runs prefill once, sends the shared KV cache to action and language states, then batches active states for joint decoding.
- Tradeoff analysis: Figure 9 sweeps total decoding steps N and per-frame decoding steps k to compare the action-frequency versus language-throughput tradeoff.
B. Additional Experimental Results
The additional experimental-results section provides supplementary end-to-end speed, action-quality, memory, and energy measurements omitted from the main text.
- The supplementary experiments report additional end-to-end speed results.
- The section includes per-suite action-quality numbers.
- The section also reports memory and energy measurements.
B.1. Additional End-to-End Results
Additional end-to-end results extend the main LIBERO-focused analysis with ALOHA and DROID tradeoff curves, decoding-step heatmaps, and underlying latency–throughput curves.
- The additional results report ALOHA and DROID action-frequency versus language-throughput tradeoff curves.
- A complementary speedup heatmap varies the total number of decoding steps.
- Raw latency–throughput curves provide the underlying measurements for the tradeoff plots.
B.1.1. End-to-End Tradeoff on ALOHA and DROID
Across ALOHA and DROID configurations, OxyGen pushes the action-frequency/language-throughput Pareto frontier outward, showing that the tradeoff conclusions generalize beyond LIBERO. Additional sweeps indicate that speedup grows with total decoding steps while configuration and action-denoising steps have smaller effects.
- Cross-configuration tradeoff: OxyGen pushes the action-frequency/language-throughput Pareto frontier outward across LIBERO, ALOHA, and DROID configurations.This indicates the reported tradeoff is not specific to LIBERO.
- Decoding-step sweep: Speedup increases with larger total decoding steps N, while configuration choice and action-denoising steps S have smaller impact.The sweep fixes decoding steps per frame at k = 5.
- Decoding-step sweep: Figure 10 complements the per-frame sweep by varying total decoding steps N at fixed decoding steps per frame k = 5.This experiment probes how total decoding length affects action-frequency speedup.
- Underlying measurements: Raw action-frequency and language-throughput curves provide the underlying measurements for the tradeoff plots in Figures 4 and 9.These curves are reported across the evaluated settings.
B.2. Additional Action-Quality Results
OxyGen preserves the task-success performance reported by openpi while accelerating inference.
- Task success: OxyGen reproduces openpi-reported task success rates within statistical noise.The comparison is reported for the LIBERO test suites.
- Action quality: Task success rate is used to support the claim that OxyGen does not degrade action quality while accelerating inference.The evaluation uses the officially released π0.5-LIBERO checkpoint.
- Evaluation: Table 1 reports per-suite task success rates for the LIBERO test suites.These measurements underpin the action-quality comparison with openpi.
B.3. Additional Memory and Energy Measurements
Additional measurements compare OxyGen’s memory and energy costs, showing modest memory overhead alongside substantial energy savings.
- Resource costs: Table 2 compares OxyGen’s memory and energy costs.The measurements support the headline numbers reported for the system.
- Memory: OxyGen adds modest memory overhead.This is reported as part of the memory-cost comparison.
- Energy: OxyGen achieves substantial energy savings.The savings are reported alongside the memory-overhead measurement.