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

Das Paper stellt CUCo vor, ein agentenbasiertes Framework, das automatisch hochoptimierten CUDA-Code für die gemeinsame Gestaltung von Berechnung und Kommunikation generiert und so die End-to-End-Latenz im Vergleich zu bestehenden Ansätzen um bis zu 1,57-fach reduziert.

Das Wesentliche

Stell dir vor, du leitest ein riesiges Orchester, das eine komplexe Symphonie spielt. In diesem Orchester gibt es zwei Gruppen: die Musiker (die die eigentliche Musik berechnen, also die Berechnung) und die Boten (die Notenblätter zwischen den Musikern hin und her tragen, also die Kommunikation).

In der Welt der modernen KI (wie bei großen Sprachmodellen) müssen diese Gruppen extrem schnell zusammenarbeiten. Das Problem bisher war: Die Musiker und die Boten haben nicht direkt miteinander gesprochen. Stattdessen stand ein Dirigent am Pult (der Computer-Prozessor, die CPU) und rief: „Musiker, haltet inne! Boten, bringt die Noten! Musik, weiter! Boten, haltet inne!"

Dieser Dirigent war langsam. Er musste ständig hin und her rennen, die Gruppen synchronisieren und wartete oft darauf, dass die Boten fertig waren, bevor er die Musiker wieder anwies. Das Orchester stand oft still, während die Boten noch unterwegs waren. Das ist ineffizient und kostet wertvolle Zeit.

Was ist CUCo?

CUCo ist wie ein genialer neuer Dirigent, der eine revolutionäre Idee hat: „Lasst die Musiker und Boten direkt miteinander reden, ohne dass ich mich ständig einmischen muss!"

CUCo ist ein intelligentes System (ein „Agent"), das automatisch neue, hochleistungsfähige Anweisungen (sogenannte CUDA-Kernel) für die Grafikkarten (GPUs) schreibt. Diese Anweisungen erlauben es den Musikern und Boten, Hand in Hand zu arbeiten.

Wie funktioniert das? (Die drei Schritte)

Das CUCo-System arbeitet in zwei Phasen, ähnlich wie ein Architekt, der zuerst ein stabiles Haus baut und es dann perfektioniert:

1. Der schnelle Weg (Fast-Path): „Bauen wir erst einmal ein sicheres Haus."

   • Zuerst sorgt CUCo dafür, dass alles funktioniert. Es nimmt die alten, langsamen Anweisungen und wandelt sie so um, dass die Boten direkt von den Musikern aus losgeschickt werden können.
   • Die Analogie: Stell dir vor, der Dirigent sagt: „Okay, wir lassen die Boten jetzt direkt zu den Musikern laufen, aber wir machen es ganz vorsichtig. Alle warten auf ein Signal, bevor sie weitermachen." Das ist vielleicht nicht die schnellste Methode, aber es funktioniert garantiert und niemand stürzt ein Haus ein. Es ist die sichere Basis.

2. Der langsame Weg (Slow-Path): „Jetzt machen wir es zum Rennwagen."

   • Sobald das sichere Fundament steht, beginnt die eigentliche Magie. CUCo nutzt eine Art „Evolution" (wie in der Natur), um Tausende von Variationen dieser Anweisungen zu testen.
   • Die Analogie: Der Dirigent probiert jetzt aus: „Was passiert, wenn die Boten nur die Hälfte der Noten tragen und dann sofort zurückkommen, während die Musiker schon die nächste Zeile spielen?" Oder: „Was, wenn die Boten und Musiker gleichzeitig an verschiedenen Tischen arbeiten?"
   • Das System testet diese Ideen, misst die Geschwindigkeit und behält nur die besten. Es lernt aus Fehlern und verbessert sich immer weiter, bis es die absolut schnellste Kombination findet.

Warum ist das so wichtig?

Bisher war es extrem schwer, solche Anweisungen von Hand zu schreiben. Es ist wie zu versuchen, ein komplexes Uhrwerk zu bauen, indem man blind nach den Zahnrädern greift. Ein kleiner Fehler führt dazu, dass die Uhr stehen bleibt.

CUCo löst dieses Problem, indem es:

• Kein menschliches Genie braucht: Es schreibt den Code selbst.
• Lernen aus Fehlern: Wenn eine Idee nicht funktioniert, merkt es sich das und probiert etwas anderes.
• Anpassungsfähigkeit: Es passt sich an verschiedene Computer-Hardware und Netzwerke an (wie ein Schauspieler, der sich an jede Bühne anpasst).

Das Ergebnis

In Tests hat CUCo gezeigt, dass es die Kommunikation zwischen den Grafikkarten so stark beschleunigen kann, dass die gesamte Aufgabe bis zu 1,57-mal schneller erledigt wird. Das ist, als würde man ein Orchester, das bisher 10 Minuten für ein Stück brauchte, plötzlich in 6 Minuten fertig spielen lassen – und das ohne dass die Musik schlechter klingt.

Zusammenfassend: CUCo ist ein intelligenter Assistent, der die starre Trennung zwischen „Rechnen" und „Datenübertragung" aufhebt. Er verwandelt ein chaotisches Orchester mit einem langsamen Dirigenten in ein hochpräzises, selbstorganisierendes Ensemble, das die volle Geschwindigkeit der modernen Hardware ausnutzt.

Technische Zusammenfassung

1. Problemstellung

Die effiziente Nutzung von GPUs für das Training und die Inferenz großer verteilter LLMs (Large Language Models) erfordert eine enge Verzahnung von Berechnung (Compute) und Kommunikation.

• Herausforderung: Das manuelle Schreiben von CUDA-Kernen, die Berechnung und Kommunikation gemeinsam optimieren (Co-Design), ist extrem arbeitsintensiv und fehleranfällig.
• Limitationen bestehender Ansätze: Bisherige Optimierungen konzentrierten sich fast ausschließlich auf die Berechnung, während Kommunikationskerne oft separat und ineffizient behandelt wurden. Traditionelle Modelle erzwangen eine Host-seitige Orchestrierung (CPU steuert den Wechsel zwischen Berechnung und asynchroner Kommunikation), was Synchronisations-Overhead erzeugt und eine enge Überlappung verhindert.
• Komplexität: Der Designraum für solche Kerne ist enorm. Entscheidungen über Tiling-Grenzen, Synchronisationszeitpunkte und Chunk-Größen interagieren auf nicht-triviale Weise und hängen stark von der spezifischen Hardware-Architektur und dem Netzwerk-Topologie ab. Herkömmliche Compiler können diese Aufgaben nicht lösen, da sie Kommunikation oft als Host-seitigen Nebeneffekt betrachten.

2. Methodik: Das CUCo-Framework

CUCo ist ein agentengetriebener, trainingsfreier Workflow, der automatisch hochperformante CUDA-Kerne generiert, die Berechnung und Kommunikation gemeinsam orchestrieren. Das Framework besteht aus drei wesentlichen Komponenten:

A. Strukturierte Designraum-Spezifikation

Anstatt dem Agenten freien Spielraum zu lassen, definiert CUCo einen eingeschränkten, aber flexiblen Suchraum durch eine optimierungsgerichtete Schnittstelle (Optimization Directive). Diese umfasst fünf Dimensionen:

1. Backend (B): Wahl zwischen GIN (GPU-Initiated Networking, RDMA) oder LSA (Load/Store Accessible, intra-node).
2. Issuer (I): Granularität des Auslösers (Thread, Warp, CTA).
3. Placement (P): Aggressivität der Überlappung vs. Ordnungskomplexität.
4. Sync Scope (S): Enge vs. weite Synchronisationsbereiche.
5. Granularity (G): Chunk-Größe für Datenübertragungen.
   Der Agent muss diese Parameter explizit angeben, bevor Code generiert wird, was Halluzinationen reduziert und die Semantik sicherstellt.

B. Fast-Path Agent (Korrektur-First)

Dieser Agent hat das primäre Ziel, Korrektheit zu gewährleisten, nicht maximale Performance.

• Aufgabe: Er wandelt einen bestehenden, host-basierten NCCL-Code in einen device-initiierten Kernel um.
• Prozess:
  1. Analyse der Kommunikationsabhängigkeiten.
  2. Transformation von Host-Collectives zu Device-Primitiven (GIN/LSA) unter strengen Konsistenzregeln.
  3. Markierung von Codebereichen („Evolve-Blocks"), die später für Optimierungen verändert werden dürfen.
• Ergebnis: Ein funktionsfähiger, sicherer Basis-Kernel, der als Startpunkt („Seed") für die nächste Phase dient. Dies verhindert, dass der Suchprozess in der komplexen Syntax von Device-APIs stecken bleibt.

C. Slow-Path Agent (Performance-First Evolution)

Dieser Agent führt eine evolutionäre Suche durch, um den Seed-Kernel zu optimieren.

• Strategie: Multi-Island-Evolutionary Search (mehrere unabhängige Populationen), um vorzeitige Konvergenz zu vermeiden.
• LLM-Mutation: Anstatt zufälliger Änderungen nutzt ein LLM kontextbewusste Mutationen (Diff-Patches, vollständige Neuschreibungen, Crossover), gesteuert durch Feedback aus vorherigen Läufen.
• Zwei-Phasen-Ansatz:
  1. Exploration: Hohe Temperatur, um strukturell diverse Strategien (z. B. Pipelining, Fusion) zu entdecken.
  2. Exploitation: Geringe Temperatur, um die vielversprechendsten Architekturen durch Feinabstimmung zu optimieren.
• Bewertung (Cascade Evaluation): Kandidaten werden in drei Stufen geprüft: Kompilierung -> Laufzeitkorrektheit (mpirun) -> Performance-Benchmarking. Fehler werden als negatives Feedback in die Datenbank aufgenommen, um zukünftige Mutationen zu leiten.

3. Schlüsselbeiträge

• Erster agenter Framework für Compute-Communication Co-Design: CUCo automatisiert die Erstellung von Kerneln, die Kommunikation direkt im GPU-Kernel ausführen (Device-Side), anstatt sie vom Host zu steuern.
• Hybrider Agenten-Ansatz: Die Trennung in einen „Fast-Path" (für Korrektheit und Dekomposition) und einen „Slow-Path" (für Performance-Optimierung) löst das Problem, dass reine Evolution auf rohem Host-Code oft scheitert, weil sie die komplexen Device-API-Semantiken nicht versteht.
• Strukturierte Suche: Durch die Definition eines expliziten Designraums (Directives) und die Nutzung von Hardware-Kontext (Topologie, Bandbreite) werden Halluzinationen minimiert und die Suche auf physikalisch sinnvolle Konfigurationen gelenkt.
• Open Source: Das Framework wird vollständig quelloffen bereitgestellt, um die Forschung voranzutreiben.

4. Ergebnisse

CUCo wurde auf vier repräsentativen Workloads evaluiert (Flash Attention mit Context Parallelism, DeepSeek-V3 MoE, KV-Cache Transfer, GEMM + AllGather) auf einem Cluster mit NVIDIA A100 GPUs (NVLink und RoCE).

• Performance-Gewinn: CUCo reduziert die End-to-End-Latenz um 5,3 % bis 26,2 % im Vergleich zu herkömmlichen Host-basierten NCCL-Implementierungen.
• Maximaler Speedup: Bis zu 1,57-fach schneller als die Baseline.
• Mechanismen der Verbesserung:
  • Beseitigung von Host-Orchestrierungs-Overhead (CPU-Wartezeiten, Kernel-Launch-Lücken).
  • Ermöglichung von feingranularer Pipelining (z. B. Berechnung des nächsten Tiles, während das vorherige Tile noch übertragen wird), was bei Host-Steuerung aufgrund von SM-Sättigung oft unmöglich ist.
  • Reduktion von Synchronisations-Overhead durch device-seitige Barrieren und Signale.
• Ablationsstudien:
  • Ohne den Fast-Path-Agenten benötigt die Evolution deutlich mehr Generationen, um korrekte Kernel zu finden, und erreicht ein niedrigeres Leistungsoptimum.
  • Der Zwei-Phasen-Ansatz (Exploration gefolgt von Exploitation) ist signifikant effizienter und erreicht bessere Ergebnisse als reine Exploitation-Strategien.

5. Bedeutung und Fazit

CUCo markiert einen Paradigmenwechsel in der Entwicklung verteilter ML-Systeme. Es zeigt, dass die manuelle Optimierung von Compute-Communication-Kernen für heterogene Hardware und Netzwerktopologien nicht mehr skalierbar ist. Durch den Einsatz von agentic AI kann CUCo komplexe, hardwarebewusste Optimierungen automatisch entdecken, die für menschliche Entwickler zu schwer zu finden sind. Mit dem Aufkommen von Device-Initiated Communication (NVSHMEM, NCCL Device APIs) wird CUCo zu einem essenziellen Werkzeug, um das volle Potenzial verteilter GPU-Cluster für zukünftige LLMs auszuschöpfen.