OxyGen：多任务并行下视觉-语言-动作模型的统一KV缓存管理 OxyGen: Unified KV Cache Management for Vision-Language-Action Models under Multi-Task Parallelism
通过统一KV缓存管理实现MoT VLA模型多任务并行推理，达3.7倍加速
前置知识
KV缓存
KV缓存是Transformer架构中注意力机制的核心优化技术，存储每个token的Key(K)和Value(V)向量。在自回归生成过程中，模型可以重用之前计算过的KV对，而不需要重新计算，显著减少重复计算开销。对于长序列生成，KV缓存能节省大量计算资源，是现代LLM服务系统的基础优化手段。
本文完全围绕KV缓存管理展开，理解KV缓存的工作原理和存储机制是理解跨任务KV共享和跨帧连续批处理两个核心优化的基础。
Mixture-of-Transformers (MoT)
MoT是一种模型架构，通过共享通用的Vision-Language Model (VLM)骨干网络，同时路由不同输出模态到专门的专家网络。与传统的VLA模型只能输出动作不同，MoT VLA可以同时生成动作和语言文本，如π0.5模型。这种架构在多任务场景下比运行多个独立模型更高效，因为它共享了大部分计算。
本文针对MoT VLA架构设计，理解MoT的架构特点是理解为什么KV缓存可以在不同专家间共享的前提。
连续批处理
连续批处理是LLM服务中的优化技术，通过将多个请求分组到单个批次中处理，利用硬件的并行计算能力。不同于静态批处理需要等待所有请求完成，连续批处理允许请求在不同时间加入和离开批次，实时处理完成的部分。这能显著提高吞吐量，特别是在有多个并发请求的场景下。
本文的跨帧连续批处理是连续批处理的扩展，但增加了机器人控制的硬性时间约束。
研究动机
现有系统在多任务并行推理时存在根本性低效问题。以一个自主家用机器人为例，它在操作物体时需要同时记录环境变化、向用户叙述进展、更新计划。现有的Mixture-of-Transformers (MoT) Vision-Language-Action Models虽然架构上支持多模态输出，但推理系统仍然采用孤立执行范式，即每个任务通过独立的模型前向传播执行，即使任务共享相同的观测输入。这导致两个具体问题：冗余计算使共享观测被重复编码，为每个任务产生相同的KV缓存条目，造成1.4倍慢；资源竞争导致不同任务竞争有限的GPU资源并互相阻塞，造成2.6倍慢。在NVIDIA GeForce RTX 4090上的测量显示，基线系统的动作频率从49.9 Hz下降到19.1 Hz，语言吞吐量仅为57.2 tokens/s，无法满足平滑机器人控制所需的50Hz最低频率要求。
本文的目标是本文的目标是设计一个高效的推理系统，使MoT VLA能够在机器人平台上实现高效的多任务并行推理。具体来说，系统需要同时满足两个看似矛盾的要求：动作生成必须满足硬性的每帧截止时间以维持最小控制频率f_min（如50Hz），语言生成需要在满足动作截止时间的前提下最大化token吞吐量。给定动作视界H、平均批大小B、每帧解码步数k和端到端推理延迟T，应用层目标转化为模型层形式：最大化语言吞吐量tau = Bk/T，约束条件是动作频率f = H/T >= f_min。作者希望在不降低动作质量的前提下，同时在RTX 4090和Jetson AGX Thor这两个代表性的机器人GPU平台上实现显著的加速。
与已有工作不同的是，本文的独特切入点是识别出孤立KV缓存管理是多任务并行低效的根本原因。虽然KV缓存管理在LLM服务中已被充分研究，如vLLM的prefix caching和SGLang的prefix sharing，但两个MoT特定的假设使这些技术直接失效：异构专家KV语义假设被打破——prefix sharing假设同构的自回归读者以相同方式扩展缓存，但在MoT VLA中，动作专家将prefill缓存作为跨S个去噪步的只读上下文，而语言专家为每个解码token追加新的KV条目，因此天真地共享可变缓存会将语言侧的追加混入动作专家的只读缓存视图；帧界限的语言解码假设被打破——连续批处理会不间断运行每个请求直到完成，但物理控制循环需要语言解码在每帧内按校准的步预算推进并在帧间恢复而不重新计算，现有服务系统不支持这种模式。本文将KV缓存抽象为跨任务和时间的头等共享资源，通过统一KV缓存管理同时解决这两个问题。
核心方法
OxyGen的核心思想是将KV缓存作为跨任务和时间的统一资源来管理。从直觉上看，多个任务（如动作生成和语言生成）都从相同的观测输入推导，因此由VLM骨干网络产生的KV缓存K_t是自然可以共享的。现有系统孤立地管理每个任务的KV缓存，错失了共享和协调调度的机会。OxyGen通过一个统一的KV缓存管理器M来实现两个关键优化：跨任务KV共享在帧t内将单个K_t分发给异构的动作和语言专家，消除冗余prefill；跨帧连续批处理将来自不同帧的进行中语言请求分组到单个解码批次中，通过校准的每帧步预算推进每个请求，使动作满足硬截止时间的同时语言吞吐量随活跃批大小扩展。这两个优化组合在一起，在不降低动作质量的前提下显著提高了推理速度。
核心创新点是将KV缓存从每个任务的私有资源转变为跨任务和时间的头等共享资源。与现有系统为每个任务执行独立的模型前向传播不同，OxyGen识别到VLM骨干网络产生的KV缓存是模态无关的：它编码了观测信息，可以被多个专家消费而不需要承诺特定的输出模态。这启发了两个关键优化。跨任务KV共享通过编码共享观测一次，然后在帧内向每个专家分发per-expert KV视图，消除了孤立执行中不同任务为同一观测重复编码prefill的冗余。同时，每个专家保持自己的访问语义：动作专家将prefill缓存作为只读上下文，语言专家为每个解码token追加新的KV条目。跨帧连续批处理将跨帧的进行中语言请求分组为单个解码批次，通过校准的每帧步预算k推进每个请求，使动作生成在每帧内满足硬性截止时间，而语言吞吐量随活跃批大小扩展。这种方法与现有LLM服务技术的本质区别在于：它针对的是资源受限的端到端多任务场景，而不是云规模多租户LLM推理；它最大化的是有硬性动作截止时间约束下的语言吞吐量，而不是给定充足计算和高并发下的聚合请求吞吐量。
方法步骤详情
方法步骤包含两个核心组件：统一的KV缓存管理器和每帧执行流程。KV缓存管理器维护每个进行中请求的生成状态，表示为sigma_t = (K_t, y_t, delta_t)，其中K_t是请求启动于帧t的KV缓存（初始时是prefill缓存，然后随着生成进展扩展解码token的KV），y_t是存储生成token的token缓冲区，delta_t是终止标志（当发出EOS或达到最大长度N时设为1）。管理器暴露标准CRUD操作在per-request状态上，并维护活跃请求集合R = {r_t1, ..., r_tm}。对于并行解码，管理器将活跃状态沿batch维度堆叠成batched state，使VLM在单个前向传播中推进所有m = |R|个请求。每帧执行流程分为两个阶段：首先，系统对新观测o_t运行一次prefill以获得K_t，管理器将这个缓存分发给动作专家（它去噪A_t跨S步）和新鲜初始化的语言状态sigma_t（存储在请求IDr_t下），消除了孤立执行所需的冗余prefill；其次，r_t加入活跃集合R，管理器检索并将所有m个活跃状态batch成进行k步联合自回归解码，完成的请求被驱逐，其余的带着增量tokeny_ti持久化以供下一帧使用。每帧预算k从目标帧率和硬件离线校准，使典型活跃批大小的batched解码在动作去噪后剩余的时间内完成，保持f_min的同时允许语言吞吐量随活跃批大小扩展。
技术新颖性
技术新颖性体现在三个方面。首先，这是第一个针对MoT VLA多任务并行推理的系统性优化工作。虽然VLA优化存在于模型级（压缩、token剪枝、层跳过、动作token重用、KV缓存剪枝、图优化）和应用级（异步流水线），但OxyGen作为调度层针对多任务推理，很大程度上与这些优化正交，可以在部署中结合。其次，跨任务KV共享处理了异构专家KV语义的挑战，通过提供per-expert KV视图而不是天真地共享可变缓存，避免了将语言侧追加混入动作专家的只读缓存视图的问题。第三，跨帧连续批处理实现了帧界限的语言解码，这是现有服务系统不支持的。它通过维护每个请求的resumable state和校准的每帧步预算k，使语言解码在帧间恢复而不重新计算，同时满足动作的硬截止时间。实验结果显示，在GeForce RTX 4090上，当k=1且N=30时，动作频率从基线的19.1 Hz提升到70.5 Hz（3.7倍），语言吞吐量从57.2 tokens/s提升到212.9 tokens/s（3.7倍）；在Jetson AGX Thor上，动作频率从基线的7.1 Hz提升到27 Hz（3.8倍），语言吞吐量从21.3 tokens/s提升到79.8 tokens/s（3.7倍）。
实验结果
OxyGen在多个配置和平台上取得了显著的性能提升。在LIBERO、DROID和ALOHA配置上，OxyGen一致地加速了并行多任务推理，最高达3.7倍。具体来说，在GeForce RTX 4090上，当每帧解码步数k=1且总解码步数N=30时，动作频率从基线的19.1 Hz提升到70.5 Hz（3.7倍），语言吞吐量从57.2 tokens/s提升到212.9 tokens/s（3.7倍）；当k=5且N=30时，动作频率从基线的33.3 Hz提升到66.7 Hz（2.0倍），语言吞吐量从96.6 tokens/s提升到193.2 tokens/s（2.0倍）；当k=10且N=30时，动作频率从基线的49.9 Hz提升到61.7 Hz（1.2倍），语言吞吐量从166.4 tokens/s提升到201.7 tokens/s（1.2倍）。在Jetson AGX Thor上，当k=1且N=30时，动作频率从基线的7.1 Hz提升到27 Hz（3.8倍），语言吞吐量从21.3 tokens/s提升到79.8 tokens/s（3.7倍）；当k=5且N=30时，动作频率从基线的12.5 Hz提升到25 Hz（2.0倍），语言吞吐量从38.5 tokens/s提升到77.0 tokens/s（2.0倍）；当k=10且N=30时，动作频率从基线的18.2 Hz提升到22.7 Hz（1.2倍），语言吞吐量从57.5 tokens/s提升到68.1 tokens/s（1.2倍）。消融研究显示，跨任务KV共享提供了初始加速（在RTX 4090上对于短解码步数且无batching时为1.4倍），通过消除不同任务对同一观测的冗余prefill；跨帧连续批处理对于长解码步数（≥10）场景至关重要，它保持几乎恒定的动作频率（在RTX 4090上约60 Hz，在Jetson AGX Thor上约27 Hz），而其他无batching的设置随解码步数增加显著下降（如基线在RTX 4090上从49.9 Hz下降到19.1 Hz，下降2.6倍）。在不同工作负载泛化实验中，包括均匀到达（每帧请求数变化）、泊松到达（变化强度lambda）和随机长度请求（不同长短解码长度比），OxyGen在所有工作负载下都比基线保持更高的动作频率。在真实机器人部署中，在配备Jetson AGX Thor的人形机器人上，基线上设备推理（1030 ms）远超333 ms动作执行窗口并阻塞下一个控制循环。OxyGen将总推理减少到393 ms；更重要的是，动作关键的prefill+denoise阶段耗时198 ms，适合执行窗口，而剩余195 ms语言生成在动作产生后运行，很大程度上隐藏在动作执行后面。OxyGen在不降低动作质量的前提下优化了VLA推理流水线：使用pi0.5-LIBERO checkpoint，它在LIBERO-{Spatial, Long, Goal, 10}上复现了openpi报告的成功率在±0.8%范围内。在相同硬件上，OxyGen也资源高效：它增加适度的内存开销（15%）同时减少平均功耗高达47%和每请求能量高达78%，归功于batched解码期间对VLM权重的较少内存访问。相比之下，naive并行化几乎使峰值内存翻倍甚至每请求能量增加3%，因为冗余计算和内存带宽竞争。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LIBERO多任务推理 | 动作频率 (Hz) | 70.5 (k=1, N=30) | 19.1 | 3.7倍 |
| LIBERO多任务推理 | 语言吞吐量 (tokens/s) | 212.9 (k=1, N=30) | 57.2 | 3.7倍 |
| Jetson AGX Thor推理 | 动作频率 (Hz) | 27 (k=1, N=30) | 7.1 | 3.8倍 |
| Jetson AGX Thor推理 | 语言吞吐量 (tokens/s) | 79.8 (k=1, N=30) | 21.3 | 3.7倍 |
| 真实机器人部署 | 语言生成延迟 (ms) | 195 | 822.3 | 4.2倍减少 |
| 能量效率 | 每请求能量 (mJ) | 减少78% | 基线 | 节能显著 |
局限与改进
作者承认的局限性包括：实证评估专注于一个现有的MoT VLA（pi0.5）和两个专家（动作和语言）；对其他MoT骨干网络和更丰富的多专家配置的泛化在第3.2节中讨论但留给未来工作。此外，我观察到：性能提升依赖于特定的工作负载特征。当语言任务较长（N大）且每帧解码步数较小时（k小），平均批大小B = N/k较大，batching效果更好；相反，如果语言任务很短或每帧解码步数较大，加速效果会减弱。另一个限制是系统假设请求可以跨帧挂起和恢复，这在某些对延迟极度敏感的应用中可能不适用。此外，虽然OxyGen与模型级优化正交，但如何在压缩、剪枝等模型修改下保持兼容性需要进一步验证。最后，虽然实验在两个代表性平台上进行，但更广泛的硬件环境（如更低的端设备）下的性能表现仍需探索。
独立分析的弱点
独立分析的弱点包括：第一，系统依赖离线校准的每帧解码步数k，这在动态工作负载下可能需要自适应调整。当请求到达模式频繁变化时，固定的k可能导致资源利用不足或违反动作截止时间。改进方向可以是设计自适应的k选择机制，基于实时负载和剩余帧预算动态调整。第二，KV缓存管理器的内存开销随活跃请求数线性增长，在极端场景下可能成为瓶颈。可以通过智能的请求驱逐策略（如LRU或优先级队列）来限制活跃集合大小，在性能和内存之间取得平衡。第三，系统假设所有语言请求使用相同的语言骨干网络，当多个语言任务使用不同的LoRA适配器时，per-expert KV视图的管理会更复杂。可以探索与cross-model KV transfer和multi-LoRA serving技术的集成。第四，在真实机器人部署中，观测质量的变化（如遮挡、光照变化）可能影响KV缓存的重用价值，可以引入KV缓存质量评估机制来决定是否需要重新编码观测。
未来方向
作者提出的未来工作方向包括：对其他MoT骨干网络和更丰富的多专家配置的泛化。基于成果可延伸的方向包括：扩展支持更多专家和更复杂的任务组合，如同时处理动作生成、语言记忆构建、视觉问答等；集成模型压缩和异步动作流水线，探索OxyGen与这些优化的协同效应；设计更智能的请求调度策略，不仅基于时间约束，还考虑任务优先级、资源可用性等；探索动态KV缓存管理，根据请求的相似性和重用价值决定共享粒度；在更多硬件平台上评估，包括更低端的端设备和云计算环境；研究KV缓存的持久化和跨会话重用，以支持长期运行的机器人任务；开发形式化的性能模型，帮助系统设计者在不同硬件和工作负载下预测和优化性能。
复现评估
论文在开源和复现方面做得很好。作者在GitHub上开源了代码：https://github.com/air-embodied-brain/OxyGen，提供了完整的实现。实验使用的模型pi0.5是一个广泛使用的MoT VLA（GitHub上超过1万星），其官方推理框架openpi也是开源的。论文详细描述了实验设置，包括硬件平台（NVIDIA GeForce RTX 4090和Jetson AGX Thor）、配置（LIBERO、DROID、ALOHA）、评估指标和基线实现。论文还提供了消融研究、工作负载泛化实验和真实机器人部署的详细结果。虽然在论文正文中没有提供超参数的完整列表，但开源代码应该包含了这些细节。实验使用了多个配置和数据集，增强了结果的普适性。总体来说，复现难度中等，主要挑战可能是需要特定的硬件（如Jetson AGX Thor）和pi0.5模型的下载和设置。开源代码、详细描述和多个实验设置使研究者能够验证结果并在其基础上开展工作。
论文图表
Figure 1左侧展示了部署MoT VLA模型进行并行多任务推理的示例：从每帧观测开始，VLA在每帧内生成机器人动作，同时跨多帧连续生成基于语言的记忆。右侧比较了两种MoT VLA推理范式：现有系统孤立地管理KV缓存，通过冗余计算和资源竞争减慢推理；我们的方法采用统一KV缓存管理，通过跨任务KV共享和跨帧连续批处理实现高达3.7倍的加速。
这张图对理解论文至关重要，因为它直观地展示了问题场景（多任务并行推理）和两种方法的对比。左侧展示了实际应用场景（家用机器人同时执行操作、记忆和对话），右侧的时间线图清楚地显示了孤立执行的资源竞争问题和统一KV缓存管理的优化效果，包括具体的加速数字（3.7倍）。这是论文motivation和method的核心视觉化。