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

OxyGen introduceert een unificatie van KV-cache-beheer voor Vision-Language-Action-modellen die multi-task parallelisme mogelijk maakt door gedeelde observaties en decoupling van taal- en actiegeneratie, wat resulteert in een tot 3,7 keer hogere snelheid zonder kwaliteitsverlies.

De kern

OxyGen: De Slimme Regisseur voor Robot-Geesten

Stel je voor dat je een robot hebt die niet alleen een taak uitvoert, maar er meerdere tegelijk doet. Hij moet een kopje koffie zetten (beweging), tegelijkertijd aan je vertellen wat hij doet (spraak), en onthouden waar hij de suiker heeft neergezet (geheugen).

In de wereld van kunstmatige intelligentie (AI) is dit een enorme uitdaging. De "hersenen" van deze robot zijn een enorm model dat alles tegelijk moet regelen. Het probleem? De huidige systemen zijn als een inefficiënte chef-kok in een drukke keuken.

Het Probleem: De Verkeerde Keuken

Stel je voor dat de chef-kok (de AI) twee recepten tegelijk moet maken: een soep (beweging) en een taart (spraak).

• Huidige systemen (Isolated Execution): De chef begint met het snijden van groenten voor de soep. Daarna wast hij het mes, pakt een nieuw bord, en begint opnieuw met het snijden van exact dezelfde groenten voor de taart. Hij doet hetzelfde werk twee keer! Bovendien staat hij in de weg van zichzelf; terwijl hij de soep maakt, kan hij de taart niet maken, en andersom. Dit kost veel tijd en energie.
• Het resultaat: De robot wordt traag. Hij stottert bij het bewegen en praat heel langzaam.

De Oplossing: OxyGen

De onderzoekers van Tsinghua University hebben OxyGen bedacht. Je kunt OxyGen zien als een slimme regisseur die de keuken volledig herinricht.

OxyGen introduceert twee slimme trucs:

1. De "Gedeelde Notitie" (Cross-Task KV Sharing)

In plaats van dat de chef twee keer dezelfde groenten snijdt, zegt OxyGen: "Wacht, we hebben die groenten al gesneden voor de soep! Laten we diezelfde gesneden groenten ook gebruiken voor de taart."

• In robot-taal: De robot kijkt één keer naar de wereld (de camera-beelden). In plaats van die beelden twee keer te verwerken (een keer voor bewegen, een keer voor praten), slaat hij de "herinnering" van die beelden één keer op. Die herinnering wordt gedeeld.
• Het effect: De robot hoeft geen tijd te verspillen aan het opnieuw bekijken van wat hij al ziet.

2. De "Trein van Sprekers" (Cross-Frame Continuous Batching)

Stel je voor dat de chef de taart moet maken, maar dat het een langzaam proces is. De soep moet echter elke seconde op het vuur worden geroerd (dat is de robotbeweging).

• Huidige systemen: De chef wacht tot de taart helemaal klaar is voordat hij weer naar de soep kijkt. Of hij doet alles in kleine stukjes, maar dan moet hij constant stoppen en starten, wat traag is.
• OxyGen: OxyGen werkt als een trein. De chef maakt de soep (beweging) op exact vast tijdstip. Maar voor de taart (spraak) pakt hij niet één taart, maar een hele trein van taarten die tegelijkertijd in de oven gaan.
  • Terwijl de robot zijn hand beweegt (de trein rijdt), werkt hij aan alle lopende gesprekken tegelijk.
  • Het is alsof je een bus vol mensen (verschillende zinnen) in één keer laat instappen, in plaats van dat ze één voor één moeten wachten.

Wat levert dit op?

Door deze twee trucs te combineren, krijgt de robot een enorme boost:

• Snelheid: De robot wordt tot 3,7 keer sneller.
• Gladheid: Hij kan tegelijkertijd heel snel bewegen (70 keer per seconde, net als een menselijke reflex) en razendsnel praten (meer dan 200 woorden per seconde).
• Geen kwaliteitsverlies: De robot maakt geen fouten door deze snelheid; hij is gewoon slimmer in het organiseren van zijn tijd.

Samenvatting in één zin

OxyGen is als een super-efficiënte manager die zorgt dat een robot niet twee keer hetzelfde werk doet en alle gesprekken tegelijk laat plaatsvinden, zodat de robot soepel, snel en slim kan handelen in onze echte wereld.

Technische samenvatting

Probleemstelling

Embodied AI-agenten (zoals robots) moeten steeds vaker meerdere taken gelijktijdig uitvoeren op basis van dezelfde waarnemingen, zoals het manipuleren van objecten, converseren met gebruikers en het opbouwen van een geheugen. Recent werk in Vision-Language-Action Models (VLA's), specifiek die gebaseerd op de Mixture-of-Transformers (MoT) architectuur (zoals $\pi_0.5$), biedt de architecturale mogelijkheid om deze heterogene outputs (acties en taal) binnen één model te genereren.

Echter, bestaande inferentiesystemen falen in het efficiënt realiseren van deze multi-task parallelisme voor on-device implementatie. De auteurs identificeren twee hoofdoorzaken voor inefficiëntie:

1. Redundante Berekening: Bestaande systemen behandelen elke taak (bijv. actiegeneratie en taalgeneratie) als een volledig gescheiden doorloop van het model. Hierdoor wordt de gedeelde observatie (beeld + instructie) meerdere keren gecodeerd, wat leidt tot identieke maar onnodig berekende KV-cache (Key-Value cache) entries.
2. Ressourcenconcurrentie: Zelfs als de KV-cache gedeeld zou worden, concurreren taken om beperkte hardwarebronnen (meestal één GPU op een robot). Omdat taken verschillende tijdsbeperkingen hebben (acties moeten binnen een vaste controlecyclus klaar zijn, terwijl taaldecoding over meerdere frames kan duren), blokkeren taken elkaar, wat leidt tot een lage doorvoer en trage reactietijden.

De kern van het probleem is dat bestaande systemen de KV-cache van elke taak geïsoleerd beheren, waardoor kansen voor delen en gecoördineerde planning verloren gaan.

Methodologie: OxyGen

Om dit op te lossen, stellen de auteurs OxyGen voor, een inferentiesysteem dat geünificeerde KV-cache management introduceert. De kernidee is om de KV-cache te abstraheren als een eersteklasse, gedeelde resource die over taken en tijd wordt beheerd.

Dit paradigma maakt twee cruciale optimalisaties mogelijk:

1. Cross-Task KV Sharing (Gedeelde Cache per Frame):

   • Wanneer meerdere taken werken op dezelfde waarneming (bijv. binnen één frame $t$), wordt de observatie slechts één keer gecodeerd door de gemeenschappelijke VLM-ruggegraat.
   • De gegenereerde KV-cache ($K_t$) wordt vervolgens gedeeld tussen de actie-expert (voor robotbesturing) en de taal-expert (voor tekstgeneratie). Dit elimineert de redundante "prefill"-stap voor elke taak.

2. Cross-Frame Continuous Batching (Doorlopende Batchverwerking):

   • Dit ontkoppelt de taaldecoding van de vaste controlecyclus van de robot.
   • Actietaken moeten binnen een frame worden voltooid (harde deadline).
   • Taal taken worden continu gebatcht over meerdere frames (zachte deadline). Het systeem beheert een set van actieve taalverzoeken die over frames worden doorgevoerd.
   • Door meerdere taalverzoeken (van verschillende frames) tegelijkertijd te batchen, wordt de hardware-efficiëntie van de GPU gemaximaliseerd zonder de harde deadline voor acties te schaden.

Technische Implementatie:
OxyGen introduceert een Unified KV Cache Manager die de staat van elk lopend verzoek bijhoudt. Deze manager ondersteunt:

• Resumable Generation: Taalgeneratie kan worden onderbroken en hervat zonder herberekening, omdat de staat (KV-cache + gegenereerde tokens) wordt opgeslagen.
• Batching/Unbatching: De manager kan individuele verzoeken samenvoegen tot een gebatchte staat voor parallelle decoding en deze later weer splitsen.

Belangrijkste Bijdragen

1. Probleemformulering: De auteurs formuleren multi-task parallelisme als een specifiek inferentiescenario voor MoT-VLA's en identificeren geïsoleerde KV-cache-beheer als de fundamentele oorzaak van inefficiëntie.
2. Nieuw Paradigma: Ze introduceren "Unified KV Cache Management", een aanpak die de cache behandelt als een gedeelde resource, wat leidt tot optimalisaties zoals cross-task sharing en cross-frame batching.
3. Implementatie en Validatie: Ze hebben OxyGen geïmplementeerd voor $\pi_0.5$ (de populairste MoT VLA) en uitgebreid gevalideerd op robotische benchmarks, waarbij ze aantonen dat deze aanpak de snelheid aanzienlijk verhoogt zonder de kwaliteit van de acties te verminderen.

Resultaten

Het systeem is geëvalueerd op een enkele NVIDIA RTX 4090 GPU (een representatief platform voor on-device inferentie) over drie benchmarks: LIBERO, DROID en ALOHA.

• Snelheidswinst: OxyGen bereikt een snelheidswinst van maximaal 3.7x ten opzichte van geïsoleerde uitvoering (zowel sequentieel als naief parallel).
• Doorvoer: Het systeem bereikt gelijktijdig:
  • >200 tokens/s voor taaldecoding.
  • 70 Hz voor actiegeneratie (robotbesturing).
• Ablatie-studies:
  • Cross-task KV sharing alleen levert al een snelheidswinst van 1.4x op door redundante prefill te elimineren.
  • Cross-frame continuous batching is cruciaal voor langere decoding-taken en zorgt ervoor dat de actie-frequentie stabiel blijft (rond 60-70 Hz) zelfs bij toenemende taalcomplexiteit, terwijl baselines hierin sterk dalen (tot 19 Hz).
• Kwaliteit: De taaksuccesratio op de LIBERO-benchmarks blijft gelijk aan de originele $\pi_0.5$ implementatie, wat aantoont dat de optimalisaties de kwaliteit van de robotacties niet beïnvloeden.
• Efficiëntie: OxyGen verlaagt het energieverbruik per verzoek met tot 78% en het piekgeheugengebruik is slechts 15% hoger dan de baseline, terwijl naief parallelisme (MPS) het geheugengebruik verdubbelt.

Betekenis

OxyGen is een doorbraak voor de implementatie van geavanceerde Embodied AI-agenten. Het bewijst dat het mogelijk is om complexe, multi-task robotica (manipulatie + communicatie + planning) efficiënt uit te voeren op beperkte hardware door slimme resource-management in plaats van alleen te vertrouwen op grotere modellen of meer hardware. De methode lost het fundamentele conflict op tussen harde real-time eisen voor robotbesturing en de variabele, langdurige aard van taalgeneratie, en maakt zo vloeiende, mensachtige interactie in dynamische omgevingen haalbaar.