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

Each language version is independently generated for its own context, not a direct translation.

¡Claro que sí! Imagina que tienes un robot doméstico muy avanzado, como un mayordomo futurista. Este robot no solo tiene que mover sus brazos para recoger un vaso (una tarea de acción), sino que al mismo tiempo tiene que hablarte, explicarte lo que está haciendo (una tarea de lenguaje), y recordar dónde dejó las llaves para más tarde (una tarea de memoria).

El problema es que, hasta ahora, los sistemas que controlan a estos robots eran como un cocinero solitario en una cocina muy pequeña: si el robot necesita cocinar (acción) y hablar (lenguaje) al mismo tiempo, el cocinero tiene que hacer todo el trabajo dos veces por separado. Primero prepara los ingredientes para cocinar, luego limpia la mesa, y luego vuelve a preparar los mismos ingredientes para hablar. ¡Es un desperdicio de tiempo y energía!

Aquí es donde entra OxyGen, el nuevo sistema presentado en este paper. Vamos a explicarlo con una analogía sencilla:

1. El Problema: La "Cocina Desordenada"

Imagina que el cerebro del robot (llamado Modelo VLA) tiene que procesar lo que ve con sus cámaras (la "observación").

Antes (Sistema Antiguo): El robot ve una manzana. Para mover la mano, vuelve a "leer" la imagen de la manzana y la guarda en su memoria temporal. Luego, para hablar sobre la manzana, vuelve a "leer" la misma imagen y la guarda de nuevo.
Resultado: El robot pierde tiempo haciendo el mismo trabajo dos veces y se atasca porque su "mesa de trabajo" (la memoria de la tarjeta gráfica) se llena rápido. El robot se mueve lento y habla con tartamudeos.

2. La Solución: OxyGen y el "Cuaderno de Notas Compartido"

OxyGen introduce una idea brillante: ¿Por qué no compartir el cuaderno de notas?

En lugar de que cada tarea (moverse, hablar, recordar) tenga su propio cuaderno, OxyGen crea un Cuaderno de Notas Único y Compartido (esto es el Unified KV Cache).

La Analogía del Cuaderno: Cuando el robot ve la manzana, escribe la información una sola vez en el cuaderno compartido.
- La tarea de "mover la mano" lee el cuaderno y actúa.
- La tarea de "hablar" lee el mismo cuaderno y habla.
- Beneficio: ¡No hay que volver a escribir la información! Se ahorra mucho tiempo.

3. Dos Trucos Mágicos de OxyGen

OxyGen no solo comparte el cuaderno, sino que organiza el trabajo de dos formas geniales:

A. Compartir el Cuaderno entre Tareas (Cross-Task KV Sharing)

Es como si en una reunión de equipo, en lugar de que cada persona tenga que leer el informe completo desde el principio, todos tengan el mismo informe abierto frente a ellos.

En el robot: Si el robot está moviendo un objeto y al mismo tiempo narrando lo que hace, ambos usan la misma información visual guardada en el cuaderno. No hay que "recalcular" la imagen una y otra vez.

B. La "Cinta de Ensamblaje" Continua (Cross-Frame Continuous Batching)

Imagina una fábrica de juguetes.

El problema anterior: La fábrica producía un juguete, lo empaquetaba, lo enviaba, y luego empezaba a hacer el siguiente. Si el juguete era complejo, la línea se detenía.
La solución de OxyGen: OxyGen convierte la fábrica en una cinta de montaje continua.
- Mientras el robot termina de mover su mano (una tarea rápida que debe hacerse en milisegundos), la parte del cerebro que "habla" no espera. Sigue trabajando en su historia, agregando una palabra cada vez que puede, sin detener la línea de producción.
- OxyGen agrupa todas las tareas pendientes (hablar de la manzana, hablar del vaso, recordar la llave) y las procesa todas juntas en un solo "impulso" de energía.

4. ¿Qué logran con esto? (Los Resultados)

Gracias a este sistema, el robot se vuelve increíblemente eficiente:

Velocidad: El robot se mueve 3.7 veces más rápido que antes.
Fluidez: Puede hablar a una velocidad normal (como si estuvieras leyendo un libro rápido) y al mismo tiempo controlar sus brazos a una velocidad de videojuego (70 veces por segundo), sin que ninguno de los dos se trabe.
Calidad: Lo mejor es que, al ser más rápido, el robot no comete errores. Sigue siendo tan bueno moviendo objetos como antes, pero ahora es mucho más ágil.

En Resumen

OxyGen es como contratar a un director de orquesta para el cerebro del robot. Antes, cada músico (tarea) tocaba su instrumento solo, a veces chocando con los otros o repitiendo la misma nota. Ahora, el director asegura que todos lean la misma partitura (el cuaderno compartido) y toquen al unísono, permitiendo que la música (la acción del robot) sea rápida, fluida y perfecta.

Esto es crucial para que los robots del futuro puedan vivir en nuestras casas, ayudándonos con tareas complejas mientras nos charlan, sin quedarse "pensando" demasiado tiempo.

Each language version is independently generated for its own context, not a direct translation.

Resumen Técnico: OxyGen

1. El Problema: Ineficiencia en la Ejecución Multi-Tarea de Agentes Embodiment

Los agentes de IA corporales (Embodied AI) requieren cada vez más la ejecución paralela de múltiples tareas simultáneas (ej. manipulación robótica, conversación y construcción de memoria) a partir de las mismas observaciones visuales, pero con restricciones de tiempo heterogéneas.

Arquitectura Base: Los recientes Modelos Visión-Lenguaje-Acción (VLA) basados en Mixture-of-Transformers (MoT), como $\pi0.5$ , permiten arquitecturalmente generar múltiples modalidades (acción y lenguaje) desde un solo modelo.
Cuello de Botella Actual: Los sistemas de inferencia existentes fallan al desplegar estos modelos en dispositivos con recursos limitados (como una sola GPU en robots). Sufren de dos ineficiencias principales debido a una gestión de caché aislada:
1. Cómputo Redundante: Cuando múltiples tareas comparten la misma observación de entrada, el sistema codifica (prefill) la observación repetidamente para cada tarea, generando entradas de caché KV idénticas pero duplicadas.
2. Contención de Recursos: Las tareas compiten por los recursos de hardware limitados. Las tareas de acción tienen plazos estrictos (ej. 50-70 Hz para control suave), mientras que las tareas de lenguaje son más flexibles. La ejecución aislada bloquea las tareas críticas de acción mientras se procesan las de lenguaje, reduciendo la frecuencia de control.

2. Metodología: OxyGen y la Gestión Unificada de Caché KV

El artículo propone OxyGen, un sistema de inferencia que introduce un nuevo paradigma: la Gestión Unificada de Caché KV. La idea central es tratar la caché de llaves y valores (KV) como un recurso compartido de primera clase a través de tareas y en el tiempo, en lugar de gestionarla de forma aislada.

OxyGen implementa dos optimizaciones clave habilitadas por este paradigma:

A. Compartición de KV entre Tareas (Cross-Task KV Sharing):
- En lugar de codificar la observación visual para cada tarea por separado, el sistema realiza una sola fase de prefill por cuadro (frame).
- La caché KV resultante ( $K_t$ ) se comparte inmediatamente entre el experto de acción y el experto de lenguaje dentro del mismo cuadro. Esto elimina el cómputo redundante de la fase de prefill.
B. Agrupación Continua entre Cuadros (Cross-Frame Continuous Batching):
- Desacopla el flujo de inferencia del bucle de control por cuadro tradicional.
- Tareas de Acción: Se completan dentro del cuadro actual para cumplir con el plazo estricto (deadline duro).
- Tareas de Lenguaje: Se gestionan como solicitudes de transmisión (streaming) que se agrupan (batching) a través de múltiples cuadros. El sistema mantiene un estado de generación resumible para cada solicitud de lenguaje, permitiendo decodificar tokens de múltiples solicitudes pendientes en paralelo en un solo paso de inferencia.
- Esto permite que el decodificador de lenguaje aproveche la capacidad de procesamiento paralelo de la GPU sin bloquear la generación de acciones en tiempo real.

Flujo de Ejecución (Algoritmo 1):

Prefill: Se procesa la nueva observación $o_t$ una sola vez para generar $K_t$ .
Generación de Acción: Se consume $K_t$ inmediatamente para generar el comando de acción $A_t$ .
Gestión de Estado: Se inicializa un estado de generación para la nueva solicitud de lenguaje y se almacena en el gestor unificado.
Decodificación Agrupada: Se recuperan los estados de todas las solicitudes de lenguaje activas (nuevas y pendientes), se agrupan en un lote y se decodifican $k$ tokens en paralelo.
Actualización: Se actualizan los estados de las solicitudes activas y se eliminan las finalizadas.

3. Contribuciones Clave

Formulación del Problema: Identifican la gestión aislada de la caché KV como la causa raíz de la ineficiencia en la inferencia multi-tarea de VLAs y definen la "paralelización multi-tarea" como un escenario objetivo crítico.
Nuevo Paradigma de Inferencia: Proponen la gestión unificada de caché KV, que permite la reutilización de cómputo (compartición) y la coordinación de ejecución (agrupación continua) entre tareas con plazos asimétricos.
Implementación y Evaluación: Implementan OxyGen en $\pi0.5$ (el modelo MoT VLA más popular) sobre el sistema openpi, demostrando mejoras significativas en plataformas robóticas reales sin degradar la calidad de la acción.

4. Resultados Experimentales

Los experimentos se realizaron en una GPU NVIDIA RTX 4090 utilizando tres benchmarks: LIBERO, DROID y ALOHA.

Aceleración: OxyGen logra una aceleración de hasta 3.7x en comparación con la ejecución aislada (secuencial o paralela naive).
Rendimiento Simultáneo: Logra simultáneamente:
- Frecuencia de Acción: > 70 Hz (crítico para control robótico suave).
- Rendimiento de Lenguaje: > 200 tokens/segundo.
Análisis de Componentes (Ablación):
- La compartición de KV entre tareas aporta un speedup inicial de ~1.4x al eliminar el prefill redundante.
- La agrupación continua entre cuadros es crucial para mantener la frecuencia de acción estable (~60-70 Hz) incluso cuando los pasos de decodificación de lenguaje son largos, evitando que la frecuencia de acción caiga a ~19 Hz (como ocurre en la línea base).
Calidad y Eficiencia:
- Calidad: La tasa de éxito en tareas (LIBERO) se mantiene idéntica a la del sistema original (sin degradación).
- Energía: Reduce el consumo de energía por solicitud hasta en un 78% debido a la menor contención de ancho de banda de memoria y la menor potencia promedio (hasta un 47% menos).
- Memoria: Añade una sobrecarga de memoria modesta (~15%) en comparación con la duplicación de memoria en la ejecución paralela naive.

5. Significado e Impacto

OxyGen demuestra que la gestión eficiente de recursos de inferencia es tan crítica como el diseño del modelo para la viabilidad de los agentes corporales avanzados.

Viabilidad en el Borde (Edge): Permite ejecutar modelos VLA complejos y multi-tarea en hardware de robot estándar (una sola GPU) sin sacrificar la fluidez del control ni la capacidad de diálogo.
Escalabilidad: La capacidad de manejar plazos asimétricos (duras para acción, blandas para lenguaje) mediante la gestión unificada de caché abre la puerta a agentes más inteligentes que pueden "pensar" (planificar/lenguaje) y "actuar" simultáneamente de manera eficiente.
Paradigma Futuro: Establece un nuevo estándar para la inferencia de modelos multimodales, moviéndose más allá de la simple ejecución secuencial o paralela naive hacia una orquestación inteligente de recursos compartidos.