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.

Imaginez un robot domestique très intelligent, un peu comme un majordome futuriste. Ce robot doit faire plusieurs choses en même temps : il doit manipuler des objets (comme mettre une tasse sur une table), parler avec vous (vous expliquer ce qu'il fait), et se souvenir de l'environnement pour plus tard.

Le problème, c'est que les robots actuels sont un peu comme des gens qui essaient de faire trois choses à la fois en fermant les yeux et en les rouvrant tour à tour. Ils font une tâche, puis une autre, puis une autre. C'est lent, et le robot semble brouillon ou lent à réagir.

Voici comment le papier OxyGen propose de régler ce problème, expliqué simplement :

1. Le Problème : Le "Gaspillage" de Mémoire

Pour comprendre OxyGen, il faut d'abord comprendre comment le cerveau du robot (son modèle d'intelligence artificielle) fonctionne.
Quand le robot regarde une image (une "observation"), il crée une sorte de mémoire temporaire (appelée KV Cache dans le jargon technique) pour se souvenir de ce qu'il a vu.

L'ancienne méthode (Isolée) : Imaginez que le robot doit à la fois mettre la tasse sur la table ET vous raconter l'histoire. Avec l'ancien système, le robot crée une mémoire pour la tâche "mettre la tasse", puis il efface tout et recrée exactement la même mémoire pour la tâche "raconter l'histoire".
- Analogie : C'est comme si vous deviez cuisiner un gâteau et écrire une lettre. Au lieu d'utiliser les mêmes ingrédients (farine, œufs) pour les deux, vous achetez deux fois la farine, deux fois les œufs, et vous faites deux cuissons séparées. C'est un gaspillage énorme de temps et d'énergie !

2. La Solution : OxyGen, le "Chef d'Orchestre"

OxyGen change la règle du jeu. Au lieu de traiter chaque tâche séparément, il gère la mémoire du robot comme une ressource partagée.

Voici les deux super-pouvoirs d'OxyGen :

A. Le Partage de Mémoire (Cross-task KV Sharing)

L'idée : Si le robot regarde la même image pour deux tâches différentes, il ne crée la mémoire qu'une seule fois. Il la partage ensuite avec les deux tâches.
L'analogie : C'est comme si le chef cuisinier préparait une seule grande assiette de pâte à gâteau (la mémoire de l'image) et la donnait à deux personnes : l'une pour faire le gâteau (action), l'autre pour écrire la recette (langage). Plus besoin de refaire la pâte !
Résultat : Le robot gagne un temps précieux car il ne perd pas de temps à "re-regarder" la même image deux fois.

B. La File d'Attente Continue (Cross-frame Continuous Batching)

Le défi : Les tâches ont des rythmes différents.
- Action : Le robot doit bouger très vite (70 fois par seconde !). C'est une urgence.
- Langage : Le robot peut parler plus lentement, mot par mot, sur plusieurs secondes.
L'ancienne méthode : Le robot attendait que la tâche "parler" soit finie avant de faire la tâche "bouger", ou vice-versa.
La méthode OxyGen : Imaginez un train de métro.
- Les tâches "Action" sont des passagers pressés qui montent et descendent à chaque arrêt (très rapide).
- Les tâches "Langage" sont des passagers qui voyagent sur plusieurs arrêts.
- Au lieu d'avoir un train pour chaque passager, OxyGen met tout le monde dans le même train. Le train (le processeur) avance à la vitesse du passager le plus pressé (l'action), mais il transporte aussi les autres passagers (le langage) en même temps, sans les ralentir.
Résultat : Le robot peut bouger très vite (comme un humain agile) tout en parlant couramment, sans que l'un ne ralentisse l'autre.

3. Les Résultats Concrets

Grâce à cette gestion intelligente de la mémoire partagée :

Le robot devient 3,7 fois plus rapide.
Il peut parler à une vitesse impressionnante (plus de 200 mots par seconde) tout en bougeant ses mains 70 fois par seconde.
Il consomme moins d'énergie car il ne fait pas de calculs inutiles.

En Résumé

OxyGen est comme un chef d'orchestre génial qui dit à son robot : "Arrête de faire les choses une par une en recommençant tout depuis le début ! Regarde, nous avons déjà vu cette image, utilisons cette mémoire pour tout le monde en même temps."

Cela permet aux robots de devenir plus fluides, plus rapides et plus capables de gérer plusieurs tâches complexes en même temps, exactement comme un humain le ferait dans la vie réelle.

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

1. Problématique

Les agents d'IA incarnée (Embodied AI) évoluent vers des systèmes capables d'exécuter simultanément plusieurs tâches hétérogènes (manipulation, conversation, construction de mémoire) à partir d'observations partagées, mais avec des contraintes temporelles différentes.

Architecture MoT (Mixture-of-Transformers) : Les modèles récents Vision-Language-Action (VLA) comme $\pi0.5$ adoptent une architecture MoT qui permet à un seul modèle de générer à la fois des actions (contrôle robotique) et du langage.
Le Goulot d'Étranglement : Malgré cette capacité architecturale, les systèmes d'inférence actuels échouent à réaliser un parallélisme efficace pour le déploiement sur appareil (on-device).
Cause Racine : Les systèmes existants gèrent le cache KV (Key-Value) de manière isolée pour chaque tâche. Cela entraîne deux inefficacités majeures :
1. Calcul Redondant : Les observations partagées sont encodées à plusieurs reprises pour chaque tâche, générant des entrées de cache KV identiques et inutiles.
2. Contention des Ressources : Les tâches (ex: génération d'actions à haute fréquence vs génération de texte asynchrone) se disputent les ressources matérielles limitées (ex: un seul GPU), bloquant les unes les autres.
Conséquence : Les systèmes actuels ne parviennent pas à dépasser les performances d'une inférence multi-modèle naïve, limitant la fréquence de contrôle et le débit de langage.

2. Méthodologie : OxyGen

Les auteurs proposent OxyGen, un système d'inférence qui introduit un nouveau paradigme : la gestion unifiée du cache KV. L'idée centrale est de traiter le cache KV comme une ressource partagée de premier ordre, gérée à la fois entre les tâches et dans le temps.

Cette approche repose sur deux optimisations clés :

A. Partage de KV Inter-tâches (Cross-Task KV Sharing)

Principe : Au lieu d'encoder l'observation visuelle et linguistique séparément pour la tâche d'action et la tâche de langage, le système effectue une seule phase de "prefill" (remplissage initial).
Mécanisme : Le cache KV résultant ( $K_t$ ) est partagé entre les experts spécifiques aux modalités (expert d'action et expert de langage) au sein d'une même trame temporelle.
Gain : Élimination complète du calcul redondant de préfill pour les observations partagées.

B. Regroupement Continu Inter-trames (Cross-Frame Continuous Batching)

Principe : Découpler le flux d'inférence du langage du cycle de contrôle rigide par trame.
Mécanisme :
- Les tâches temps réel (actions) doivent être complétées dans la trame courante (deadline stricte).
- Les tâches de streaming (langage) sont regroupées (batchées) sur plusieurs trames. Le système maintient un état de génération incrémental pour chaque requête de langage en cours.
- À chaque trame, le système effectue un décodage par lots (batched decoding) sur toutes les requêtes de langage actives (nouvelles et en cours) simultanément.
Gain : Cela permet d'amortir le coût de décodage sur plusieurs requêtes, augmentant l'utilisation du matériel (GPU) et le débit de tokens sans sacrifier la fréquence des actions.

3. Contributions Clés

Formalisation du problème : Identification de la gestion isolée du cache KV comme cause fondamentale de l'inefficacité dans les inférences multi-tâches parallèles sur les modèles MoT-VLA.
Nouveau Paradigme : Proposition de la "gestion unifiée du cache KV", abstraite comme une ressource partagée, permettant le partage inter-tâches et le regroupement temporel.
Implémentation et Validation : Développement de OxyGen pour le modèle $\pi0.5$ (le modèle MoT VLA le plus populaire) et évaluation sur des configurations robotiques réalistes, démontrant des gains significatifs sans dégradation de la qualité.

4. Résultats Expérimentaux

Les évaluations ont été menées sur un GPU NVIDIA RTX 4090 (plateforme représentative pour l'inférence embarquée) avec trois benchmarks : LIBERO, DROID et ALOHA.

Accélération Globale : OxyGen atteint un accélération allant jusqu'à 3,7x par rapport à l'exécution isolée (séquentielle ou parallèle naïve).
Performance Simultanée :
- Débit de Langage : Plus de 200 tokens/seconde.
- Fréquence d'Action : Jusqu'à 70 Hz (contre ~19 Hz pour la baseline).
Étude d'Ablation :
- Le partage de KV inter-tâches apporte un gain initial de ~1,4x en éliminant le préfill redondant.
- Le regroupement inter-trames est crucial pour les séquences longues, maintenant une fréquence d'action stable (~60-70 Hz) même lorsque le nombre d'étapes de décodage augmente, là où la baseline chute drastiquement.
Qualité des Actions : La réussite des tâches sur LIBERO reste statistiquement identique à celle du modèle de base ( $\pi0.5$ ), confirmant qu'aucune dégradation de la qualité de contrôle n'est introduite.
Efficacité Énergétique : Réduction de la consommation d'énergie par requête jusqu'à 78 % grâce à une meilleure utilisation de la mémoire et moins d'accès redondants aux poids du modèle.

5. Signification et Impact

Ce travail est significatif car il comble le fossé entre la capacité architecturale des modèles VLA modernes (MoT) et les contraintes pratiques de leur déploiement sur des robots physiques.

Viabilité du Déploiement : OxyGen rend possible l'exécution fluide et réactive d'agents robotiques complexes sur du matériel embarqué standard, là où les systèmes précédents échouaient à cause de la contention des ressources.
Paradigme de Gestion : Il redéfinit la gestion des ressources d'inférence pour les agents incarnés, passant d'une approche "tâche par tâche" à une approche "ressource unifiée", essentielle pour les agents capables de multitâche dynamique.
Futur de l'IA Incarnée : En permettant la coexistence de contrôles à haute fréquence et de raisonnement linguistique complexe, OxyGen ouvre la voie à des robots plus autonomes, capables de converser, de planifier et d'agir en temps réel sans latence perceptible.