Efficient, VRAM-Constrained xLM Inference on Clients

1. Probleemstelling

Het artikel behandelt de kritieke uitdaging van het draaien van hoog-accurate, grootschalige AI-modellen (gezamenlijk aangeduid als xLM's, waaronder Dense LLM's, Mixture-of-Experts (MoE) LLM's en Vision-Language Models (VLM's)) op client-systemen (desktops en laptops) met beperkt Video RAM (VRAM).

• De Beperking: Client-GPU's hebben vaak aanzienlijk minder VRAM dan de totale geheugenbehoefte van moderne hoog-accurate modellen (bijvoorbeeld een 77GB-model op een 16GB-GPU).
• De Beperking van Bestaande Oplossingen:
  • Kwantiseren/Pruning: Deze technieken verkleinen het model, maar ten koste van de nauwkeurigheid, wat onacceptabel is voor toepassingen met hoge fideliteit.
  • Statische Offloading: Bestaande hybride planning (bijvoorbeeld in llama.cpp) vertrouwt vaak op handmatige configuratie (bijvoorbeeld het specificeren van het aantal GPU-lagen) en faalt om dynamisch aan te passen aan variërende systeemomstandigheden (CPU-draden, PCIe-bandbreedte) of inferentiefasen (prefill versus decode).
  • Gebrek aan Universaliteit: Geen enkel bestaand product behandelt interactieve modi, batchverwerking, hoge-resolutie VLM's en zowel dense- als MoE-architecturen tegelijkertijd efficiënt onder strikte VRAM-budgetten.

2. Methodologie

De auteurs stellen een nieuwe, benchmark-profielgeleide hybride planningstechniek voor genaamd Pipelined Sharding, gecombineerd met een reeks visie-specifieke optimalisaties genaamd VLMOpt.

A. Pipelined Sharding (Kerninnovatie)

Dit is een volledig automatische planningstechniek die het inferentiegraf op het sub-lageniveau verbreekt (bijvoorbeeld het splitsen van Attention- en FFN-blokken) in plaats van alleen op het lageniveau. Het werkt in drie fasen:

1. Install-fase (Profilering):

   • Een benchmark-suite draait op de specifieke client-hardware om kernelprestaties (FLOPS) te profileren voor verschillende tensorgroottes, kwantisatieniveaus en CPU-draadtellingen.
   • Cruciaal wordt geprofieleerd onder gelijktijdige belasting (simulatie van PCIe-verkeer en geheugencontroller-concurrentie) om nauwkeurige timing-schattingen te genereren voor hybride CPU-GPU-uitvoering.
   • Dit creëert een kleine profieldatabase (~170KB) die wordt gebruikt voor kostenschatting.

2. Planningsfase (Schema-Generatie):

   • Het systeem definieert een reeks "Token Tiers" (bijvoorbeeld 1, 4, 16, 64 tokens per iteratie) die verschillende inferentielasten vertegenwoordigen.
   • Voor elke tier genereert het drie verschillende planningsplannen op basis van het huidige VRAM-budget en systeemomstandigheden:
     • GPU-Only: Alle niet-gepinde sub-lagen draaien op de GPU; gewichten worden gestreamd via PCIe. Optimaal voor hoge rekenlasten (grote batches/contexts) waarbij PCIe-latentie wordt verborgen door rekenwerk.
     • Static: Hoge-prioriteit sub-lagen (Attention, KV Cache) zijn gepind naar VRAM; lage-prioriteit ones (FFN) draaien op de CPU. Minimaliseert PCIe-verkeer maar vertrouwt op CPU-rekenkracht.
     • Dynamic: Een hybride aanpak waarbij gewichten voor GPU-sub-lagen worden gestreamd vanuit sysRAM terwijl de CPU eerdere sub-lagen uitvoert. Hierdoor overlappen PCIe-overdrachten zich met rekenwerk.
   • Een Roofline-gebaseerde timing-schatter gebruikt de profieldatabase om de latentie van elk plan te voorspellen en selecteert het optimale plan voor de specifieke token-tier.

3. Inferentiefase (Uitvoering):

   • Tijdens inferentie bepaalt het systeem het aantal "nieuwe tokens" voor de huidige batch.
   • Het zoekt het vooraf berekende optimale schema voor die token-tier op en voert dit uit.
   • Hierdoor kan het systeem naadloos schakelen tussen interactieve modus (Batch Size = 1, lage tokens) en gebudgette modus (Batch Size > 1, hoge tokens).

B. VLMOpt (Visie-Taal Optimalisatie)

Om hoge-resolutie VLM's te ondersteunen (specifiek het Cosmos-Reason1-model), hebben de auteurs drie optimalisaties geïmplementeerd binnen llama.cpp:

1. Vision Tensor CPU Offloading: Visiemodel-gewichten worden gepind naar CPU-geheugen en alleen naar de GPU gestreamd wanneer nodig, wat de VRAM-footprint verkleint.
2. Flash Attention voor Visie: Implementeert Flash Attention in de visie-encoder om de $O(N^2)$ geheugencomplexiteit van attentiescores voor hoge-resolutie afbeeldingen te verminderen.
3. VRAM Overlap Avoidance: Herstructureert de uitvoering om ervoor te zorgen dat visie-encoding voltooid is en VRAM vrijmaakt voordat taaldecoding begint, zodat het piekgebruik van VRAM het maximum van de twee delen is in plaats van hun som.

3. Belangrijkste Bijdragen

1. Pipelined Sharding Algorithm: Een nieuwe, automatische hybride planner die zich aanpast aan systeemomstandigheden (CPU-draden, PCIe-bandbreedte) en inferentiestatussen (context versus decode) zonder handmatige afstelling.
2. VLMOpt Implementatie: Een suite van optimalisaties die inferentie van hoge-resolutie VLM's op client-hardware mogelijk maakt met drastisch verminderde VRAM-eisen.
3. Open-Source C++ Implementatie: De oplossing is geïmplementeerd in llama.cpp (een C++-first framework), waardoor het toegankelijk is voor game-ontwikkelaars en fysieke AI-ingenieurs.
4. Uitgebreide Evaluatie: Rigoureuze testen op Dense LLM's, MoE LLM's en VLM's op diverse client-hardware (laptops tot high-end desktops).

4. Resultaten

De evaluatie toont significante prestatieverbeteringen aan ten opzichte van agressieve baselines (handmatig afgestelde llama.cpp en vLLM):

• Interactieve Modus (Batch Size = 1):
  • TTFT (Time-to-First-Token): Verbeterd met tot 6,7×.
  • TPS (Tokens Per Second): Verbeterd met tot 30× (bijvoorbeeld voor het 77GB Qwen235B-model).
  • Leefbaarheid: Maakt het mogelijk dat enorme modellen (bijvoorbeeld 77GB) interactief draaien op GPU's met zo weinig als 2GB VRAM (bereikend >5 TPS, de drempel voor menselijke leessnelheid).
• Gebudgette Modus (Batch Size > 1):
  • Doorvoerverbeteringen tot 8,2× ten opzichte van baselines.
  • De planner verplaatst het werk effectief naar de GPU naarmate de batchgrootte toeneemt, waardoor CPU-bottlenecks worden vermeden.
• VLM Prestaties (Cosmos-Reason1):
  • VRAM Vermindering: Verminderde VRAM-vraag met 10× (van ~20GB naar ~2GB), waardoor hoge-resolutie (1440p) inferentie mogelijk wordt op client-systemen waar dit eerder onmogelijk was.
  • Latentie: Bereikte tot 1,78× versnelling in End-to-End Latentie (E2EL) ten opzichte van baselines bij equivalente VRAM-budgetten.
• Gelijktijdige Uitvoering: In scenario's waarbij een game (Cyberpunk 2077) en een LLM gelijktijdig draaien, laat de oplossing toe dat de game meer assets in VRAM behoudt, waardoor zowel game-FPS als LLM-TPS verbetert door resource-concurrentie te vermijden.

5. Betekenis

Dit werk overbrugt de kloof tussen de hoge nauwkeurigheidsvereisten van moderne AI en de hardwarebeperkingen van consument-clientapparaten.

• Verliesvrije Inferentie: Het bewijst dat hoog-accurate modellen kunnen draaien op client-hardware zonder modelfideliteit te offeren via kwantisatie of pruning.
• Democratisering: Door hoge-resolutie VLM's en grote MoE-modellen op standaard laptops en desktops mogelijk te maken, stelt het ontwikkelaars in staat geavanceerde "Physical AI" en in-game AI-assistenten te bouwen zonder afhankelijk te zijn van dure cloud-infrastructuur.
• Aanpasbaarheid: De token-tier- en profielgedreven aanpak biedt een robuust kader voor toekomstige hardware-evolutie, waarbij inferentieprestaties schalen met beschikbare middelen (CPU-draden, PCIe-bandbreedte) in plaats van statisch gebonden te zijn aan een vaste configuratie.

Kortom, het artikel presenteert een productieklaar, geautomatiseerde oplossing die de volgende generatie client-side AI haalbaar, efficiënt en verliesvrij maakt.