统一缓存管理让多任务机器人模型更快更稳运行
OxyGen: Unified KV Cache Management for VLA Inference under Multi-Task Parallelism
- 把键值缓存当作共享资源,跨任务复用观察数据
- 实测语言生成超200词/秒,动作频率达70赫兹
- 适合边缘设备上多任务并行的机器人推理场景
具身智能体需并行执行多种任务(如操作、对话、记忆构建),基于共享观测在不同时间约束下运行。当前混合变换器架构的视觉-语言-动作模型支持异构输出,但现有推理系统因冗余计算和资源竞争,难以实现高效设备端多任务并行。我们发现孤立的键值缓存管理是根本原因。为此提出统一键值缓存管理:将KV缓存作为跨任务、跨时间的共享资源,实现两项优化——跨任务缓存共享消除重复预填充,跨帧连续批处理解耦可变长度语言解码与固定速率动作生成。我们在π_{0.5}模型上实现该设计,分别在NVIDIA GeForce RTX 4090与Jetson AGX Thor平台评估。相比孤立执行,最高提速3.7倍,同时达成超过200词/秒的语言吞吐量和70赫兹的动作频率,且动作质量无损;进一步在搭载Jetson AGX Thor的真实人形机器人上验证了性能提升。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。