OxyGen

OxyGen: Unified KV Cache Management for VLA Inference under Multi-Task Parallelism (Arxiv 2026)

Paper

---

Motivation

• 痛点：在多任务并发（动作控制、语言交互、环境记忆）的具身智能场景中，Mixture-of-Transformers架构 VLA 模型在端侧设备上的推理延迟高、显存浪费严重。
• 现有系统的缺陷：
  • 重复 Prefill 计算：各任务独立维护上下文，导致同一帧摄像机输入的视觉图像被 Action 专家和 Language 专家重复编码多次，产生巨大的无谓算力开销。
  • 资源争用与控制阻塞：变长且耗时的文本自回归生成与高频定拍的动作去噪生成共享底层计算流时相互抢占，长文本解码直接破坏动作控制的硬实时性，造成丢帧与机械臂抖动。

• 洞察：机器人在同一物理时刻捕获的环境视觉观测天然具有跨任务共享性。视觉编码产生的 KV 缓存生命周期虽然在各任务中不同（动作仅需当前帧，语言需跨多帧自回归），但在物理显存与计算流上完全可以进行统一管理与跨帧解耦调度。

---

Contribution

1. Cross-Task KV Sharing：在每帧控制周期内对共享视觉观测仅执行单次 Prefill 编码，通过多视图映射零拷贝分发给动作专家与语言专家，消除冗余计算。
2. Cross-Frame Continuous Batching：解耦变长语言自回归生成与固定频率动作控制时钟。将跨多帧的文本 Token 解码流水线化，穿插在高优先级动作去噪周期的计算空隙中执行。

---

Limitation

。

---

Architecture

在统一显存池与异步调度引擎之上构建，无需修改模型权重与损失函数。

• 统一视觉 Prefill 与全局 KV Cache 管理：
  • 每一帧传感器图像输入后，视觉 Backbone 仅前向传播一次，生成全局视觉 Key-Value 缓存对： $$K_{vis} = W_K \cdot \text{ViT}(I_t), \quad V_{vis} = W_V \cdot \text{ViT}(I_t)$$
  • 多视图访问控制：Action 模块通过去噪 Cross-Attention 直接寻址当前帧的 $(K_{vis}, V_{vis})$；Language 模块通过虚拟引用计数跨时间步持有该 KV Block。
• 异构注意力映射：
  • 支持 Action 专家的全连接去噪注意力与 Language 专家的下三角因果注意力在同一显存分页拓扑中并发寻址，消除内存数据拷贝。
• 文本流水线优化：
  • Action Slot：严格对齐机器人控制时序周期。
  • Language Batching Slot：文本自回归请求被切分为单 Token 执行粒度，跨帧流式填充入动作计算后的剩余算力窗口。

---

Experiments

硬件平台：NVIDIA RTX 4090 与 NVIDIA Jetson AGX Thor。

测试任务：LIBERO 模拟基准、DROID / ALOHA 数据集策略以及 Humanoid 双臂操作与人机交互任务。

指标：相比 openpi 原生基线与传统串行推理，实现最高 3.7× 的端到端推理加速。