CUCo: An Agentic Framework for Compute and Communication Co-design

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

Stel je voor dat je een enorm groot puzzelstuk moet maken, maar je hebt twee verschillende teams nodig: Team Rekenen (de GPU's die de zware wiskunde doen) en Team Boodschappen (de netwerkverbindingen die data tussen de GPU's vervoeren).

In het verleden werkten deze twee teams als twee totaal gescheiden afdelingen in een fabriek.

Team Rekenen deed zijn werk.
Dan stopte alles.
Team Boodschappen kwam langs, haalde de data op, bracht het naar de volgende GPU en gaf een seintje.
Pas toen mocht Team Rekenen weer verder.

Dit "stop-en-start" ritme is heel inefficiënt. Het is alsof je een auto rijdt, maar elke 100 meter moet stoppen om de brandstof te tanken, terwijl de motor uit staat. De auto (de GPU) staat vaak stil terwijl de tankwagen (het netwerk) langskomt.

Het probleem: De "Handmatige" Oplossing is te moeilijk

Om dit op te lossen, willen programmeurs deze twee teams laten samensmelten. Ze willen dat de GPU terwijl hij rekent, ook al data opstuurt. Dit heet "Compute-Communication Co-design".

Het probleem is echter dat dit schrijven van computercode (CUDA-kernels) extreem moeilijk is. Het is als proberen een heel complex orkest te dirigeren waarbij je tegelijkertijd de muziek moet componeren én de instrumenten moet stemmen. Als je één foutje maakt, stopt de hele machine. Tot nu toe moesten mensen dit handmatig doen, wat jaren kost en vol fouten zit.

De Oplossing: CUCo (De Slimme Agent)

De auteurs van dit paper hebben CUCo bedacht. Je kunt CUCo zien als een super-slimme, geautomatiseerde architect die dit hele proces voor je regelt, zonder dat je zelf de code hoeft te schrijven.

CUCo werkt in twee fasen, net als het leren van een nieuwe vaardigheid:

1. De "Snelle Weg" (De Veilige Starter)

Stel je voor dat je een nieuwe taal leert. Je begint niet met poëzie schrijven, maar met simpele zinnen die zeker correct zijn.

De Snelle Agent van CUCo kijkt naar de oude, trage code.
Hij maakt een nieuwe versie die misschien niet supersnel is, maar wel 100% werkt.
Hij zorgt ervoor dat de data op het juiste moment wordt verstuurd en dat niemand in de war raakt.
Analogie: Het is alsof je eerst een veilige fietsroute bouwt waar niemand kan vallen, voordat je begint met racen.

2. De "Langzame Weg" (De Optimale Rennende)

Nu hebben we een werkend fundament. De Langzame Agent gaat nu op zoek naar de snelste manier om dit te doen.

Deze agent is als een evolutionaire bioloog. Hij maakt duizenden variaties van de code.
Hij probeert dingen zoals: "Wat als we de data in kleinere stukjes sturen?", "Wat als we de rekenopdrachten en het versturen door elkaar husselen?"
Hij test elke variant, kijkt welke het snelst is, en laat de langzamere versies "sterven".
Na veel iteraties vindt hij de perfecte, razendsnelle versie die precies past bij jouw specifieke computerhardware.

Waarom is dit zo speciaal?

Vroeger dachten computersystemen dat rekenen en communiceren twee aparte dingen waren. CUCo leert de computer om ze als één geheel te zien.

Vroeger: De CPU (het hoofd) moest constant tussen de GPU's (de spierballen) en het netwerk (de zenuwen) heen en weer lopen om instructies te geven. Dat kostte tijd.
Nu met CUCo: De GPU's praten direct met elkaar en met het netwerk, terwijl ze tegelijkertijd rekenen. De CPU hoeft niet meer tussenbeide te komen.

Het Resultaat

In de tests bleek dat CUCo de taken tot 1,57 keer sneller kon uitvoeren dan de oude methoden.

Bij het trainen van grote AI-modellen (zoals de modellen die dit gesprek mogelijk maken) betekent dit dat je minder tijd en minder geld kwijt bent.
Het is alsof je van een oude, trage trein overstapt op een hogesnelheidstrein die onderweg ook nog eens de lading kan lossen en laden zonder te stoppen.

Kortom: CUCo is een slimme robot die de moeilijke, saaie en foutgevoelige taak van het programmeren van deze complexe samenwerking overneemt, zodat AI-systemen sneller en efficiënter kunnen werken.

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

1. Het Probleem

De schaalvergroting van machine learning-modellen vereist distributieve training en inferentie op grote GPU-clusters. Hoewel er al veel geoptimaliseerd is op het gebied van berekening (compute) en communicatie, blijven deze twee componenten vaak gescheiden.

Huidige beperkingen: Traditioneel worden communicatiekernels (zoals NCCL) door de host (CPU) aangestuurd, wat synchronisatie-overhead introduceert en het overlappen van berekening en communicatie beperkt.
De uitdaging: Nieuwe interfaces (zoals NVSHMEM en NCCL Device-Side APIs) maken het mogelijk om communicatie direct vanuit de GPU-kernel te initiëren. Dit biedt enorme potentie voor prestaties, maar het handmatig schrijven van deze "gefuseerde" kernels is extreem complex, foutgevoelig en vereist kennis van een enorm ontwerpruimte (tiling, synchronisatiegranulariteit, netwerktopologie).
Moeilijkheid voor bestaande tools: Bestaande ML-compilers en LLM-agenten zijn niet ontworpen voor deze gecoördineerde zoekruimte. Ze behandelen communicatie vaak als een host-zijde bijwerking of missen de nodige semantische kennis over multi-GPU-communicatie, wat leidt tot incorrecte of inefficiënte code.

2. Methodologie: Het CUCo Framework

CUCo is een agent-gestuurde, trainingsvrije workflow die automatisch hoog-presterende CUDA-kernels genereert die berekening en communicatie co-ontwerpen. Het framework bestaat uit drie kerncomponenten:

A. Gestructureerde Ontwerpruimte Specificatie

Om de agent te voorkomen dat hij "hallucineert" of onmogelijke strategieën bedenkt, definieert CUCo een beperkte ontwerpruimte ( $C$ ) gebaseerd op vijf dimensies:

Backend (B): GIN (GPU-Initiated Networking) of LSA (Load/Store Accessible).
Issuer (I): Welke thread-groep de communicatie start (Thread, Warp, CTA).
Placement (P): Hoe agressief communicatie en berekening worden over elkaar gelegd.
Sync Scope (S): De reikwijdte van synchronisatie (lokaal vs. globaal).
Granularity (G): De grootte van de data-chunks voor overdracht.
De agent moet eerst een expliciete "optimization directive" kiezen binnen deze ruimte voordat code wordt gegenereerd, wat de zoekruimte grondig beperkt tot geldige strategieën.

B. De Fast-Path Agent (Correctheid Eerst)

Deze agent is verantwoordelijk voor het snel genereren van een functioneel correcte basisimplementatie.

Doel: Zorgen dat de kernel compiles en correct werkt over meerdere GPUs, zonder direct te focussen op maximale prestaties.
Aanpak: Het transformeert bestaande host-gedreven NCCL-code naar device-initiated code. Het gebruikt een conservatieve strategie (bijv. globale synchronisatie, grote blokken) om race conditions en deadlocks te voorkomen.
Output: Een betrouwbare "seed" die dient als startpunt voor de volgende fase.

C. De Slow-Path Agent (Prestatie Optimalisatie)

Deze agent voert een evolutionaire zoektocht uit om de seed te optimaliseren voor specifieke werklasten en hardware.

Multi-Island Evolutie: Het onderhoudt meerdere onafhankelijke populaties ("eilanden") om vroegtijdige convergentie naar lokale optima te voorkomen.
LLM-gestuurde Mutatie: In plaats van willekeurige wijzigingen, gebruikt de agent een LLM om semantisch gefundeerde wijzigingen aan te brengen in gemarkeerde codeblokken ("Evolve-Blocks"). De LLM krijgt context over de huidige prestaties, fouten, en historische patronen.
Twee-fase strategie:
1. Explore: Hoog-temperatuur mutaties om fundamenteel verschillende architecturen te ontdekken (bijv. gesplitste put/wait pipelines).
2. Exploit: Laag-temperatuur fijne afstelling van de beste kandidaten.
Cascaderende Evaluatie: Kandidaten worden gefilterd via drie niveaus: compilatie, numerieke correctheid, en uiteindelijk benchmarking. Fouten worden opgeslagen in een database om de agent te leren van eerdere mislukkingen.

3. Belangrijkste Bijdragen

Eerste Agentic Framework voor Compute-Communication Co-design: CUCo is het eerste systeem dat automatisch kernels genereert die communicatie en berekening volledig fuseren op GPU-niveau, zonder dat de gebruiker specifieke optimalisatie-instructies hoeft te geven.
Trainingsvrij en Hardware-bewust: Het systeem vereist geen fine-tuning op specifieke kernel-corpora. Het gebruikt dynamisch gegenereerde context (API-documentatie, hardware-specificaties) om de agent te "grounden" in de realiteit van de target-hardware.
De "Fast-Path / Slow-Path" Architectuur: Een innovatieve aanpak die het probleem decomposeert in een correctheidsfase en een prestatiefase, wat de efficiëntie van de zoektocht drastisch verhoogt in vergelijking met directe evolutie op ruwe code.
Open Source: De auteurs kondigen aan CUCo volledig open-source te maken om de gemeenschap te ondersteunen.

4. Resultaten

CUCo werd geëvalueerd op vier representatieve werklasten (Flash Attention met Context Parallelism, DeepSeek-V3 MoE, KV-Cache Transfer, en GEMM + AllGather) op clusters met NVIDIA A100 GPU's (binnen- en tussen-nodes).

Prestatieverbetering: CUCo reduceerde de end-to-end latentie met 5,3% tot 26,2% ten opzichte van traditionele host-gedreven baselines.
Snelheidswinst: In de beste scenario's (Flash Attention) werd een snelheidswinst van 1,57x bereikt.
Case Study (Flash Attention): De grootste winst kwam voort uit het elimineren van "pipeline bubbles". Waar de host-baseline moest wachten op alle communicatie voordat de volgende berekening kon starten, liet de CUCo-kernel berekening en communicatie per "tile" parallel lopen. Dit elimineerde 37,7 ms aan idle-tijd en 37,6 ms aan host-API-overhead.
Ablatie Studies:
- Zonder de Fast-Path agent moest de Slow-Path agent zelf de complexe NCCL-device-API's ontdekken, wat leidde tot 25% meer mislukte generaties en een lagere eindscore.
- Een "twee-fase" zoekstrategie (eerst exploreren, dan exploiteren) presteerde aanzienlijk beter dan een puur exploitatie-gedreven aanpak, wat aantoont dat diversiteit in de zoekruimte cruciaal is.

5. Betekenis en Conclusie

CUCo markeert een fundamentele verschuiving in hoe gedistribueerde machine learning-kernels worden ontworpen. Het bewijst dat agente zoektochten complexe, handmatig ondoenlijke optimalisaties kunnen vinden die de grenzen van huidige compiler-technieken overstijgen.

Door de drempel voor het schrijven van geavanceerde, device-initiated communicatiekernels te verlagen, stelt CUCo onderzoekers en ingenieurs in staat om de volledige potentie van moderne GPU-clusters en netwerkinfrastructuur (zoals NVLink en RoCE) te benutten. Dit is essentieel voor de volgende generatie grote taalmodellen (LLMs) waar communicatie-overhead een steeds groter bottleneck wordt.