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

Il paper presenta CUCo, un framework agenziale senza addestramento che genera automaticamente kernel CUDA ad alte prestazioni ottimizzando congiuntamente calcolo e comunicazione, riducendo la latenza end-to-end fino a 1,57 volte rispetto alle soluzioni esistenti.

L'essenziale

Immagina di dover organizzare una grande festa di gruppo (l'addestramento di un'intelligenza artificiale) dove centinaia di cuochi (le GPU) devono preparare un pasto enorme insieme.

Fino a poco tempo fa, c'era un capocuoco umano (la CPU) che stava in mezzo a loro. Il suo compito era:

1. Dire ai cuochi: "Tagliate le verdure!" (Calcolo).
2. Fermare tutto, correre in cucina, prendere i piatti pronti e dirli agli altri cuochi (Comunicazione).
3. Dire di nuovo: "Ora cuocete la pasta!" (Calcolo).

Il problema? Il capocuoco è lento. Mentre lui corre avanti e indietro a coordinare lo scambio dei piatti, i cuochi restano fermi a guardare il muro. È un'enorme perdita di tempo.

Ora, i cuochi hanno imparato a parlare direttamente tra loro (grazie a nuove tecnologie chiamate NVSHMEM e NCCL). Potrebbero scambiarsi i piatti mentre continuano a cucinare, senza fermarsi. Ma c'è un ostacolo: scrivere le istruzioni per farlo è un incubo. È come chiedere a un cuoco di scrivere un manuale di fisica quantistica mentre sta tagliando cipolle. Se sbagliano una virgola, la festa va in tilt e il cibo viene buttato.

Cosa fa CUCo?

CUCo è come un assistente robotico super-intelligente (un "agente") che entra in cucina e dice: "Lasciate fare a me!". Il suo obiettivo è scrivere automaticamente le istruzioni perfette per far sì che i cuochi scambino i piatti mentre continuano a cucinare, senza che il capocuoco umano debba intervenire.

Ecco come funziona, diviso in tre parti magiche:

1. La Mappa del Tesoro (Design Space)

Prima di scrivere qualsiasi cosa, CUCo non indovina a caso. Ha una mappa dettagliata che dice: "Ehi, puoi scambiare i piatti in questi modi specifici, con queste regole di sicurezza". Questo evita che il robot inventi cose impossibili che farebbero esplodere la cucina.

2. Il "Fai da Te" Veloce (Fast-Path Agent)

Immagina di dover costruire un ponte.

• L'approccio vecchio: Provare a costruire un ponte di vetro subito. Se sbagli, crolla tutto e devi ricominciare da capo.
• L'approccio CUCo (Fast-Path): Prima, il robot costruisce un ponte di legno robusto e sicuro. Non è il più veloce del mondo, ma funziona. Non crolla mai. Questo è il "punto di partenza sicuro". Il robot dice: "Ok, ora abbiamo un ponte che regge. Non dobbiamo più preoccuparci che crolli, possiamo concentrarci sulla velocità".

3. Il "Raffinatore" Lento (Slow-Path Agent)

Una volta che il ponte di legno è sicuro, arriva il Raffinatore. Questo è un robot che ama l'evoluzione.

• Prende il ponte di legno e dice: "Proviamo a spostare una trave qui... no, troppo lento. Proviamo a cambiare il materiale lì... meglio!".
• Usa un po' di intelligenza artificiale per fare migliaia di piccoli esperimenti (mutazioni).
• Se un esperimento funziona, lo tiene. Se fa crollare il ponte, lo butta via e impara dall'errore.
• Dopo molti tentativi, trasforma quel ponte di legno sicuro in un ponte in acciaio super-veloce che permette ai cuochi di scambiare i piatti alla massima velocità possibile.

Perché è così importante?

Il risultato è che CUCo riesce a creare queste istruzioni complesse senza bisogno che un umano le scriva a mano.

• Risultato: I cuochi (le GPU) lavorano molto più velocemente perché non devono più aspettare il capocuoco umano.
• Vantaggio: Il sistema è stato testato su quattro tipi di "feste" diverse (diversi modelli di intelligenza artificiale) e ha dimostrato di essere fino al 57% più veloce rispetto ai metodi tradizionali.

In sintesi

CUCo è un architetto robotico che:

1. Disegna un piano sicuro (per non rompere nulla).
2. Costruisce una versione di base che funziona (per avere un punto di partenza).
3. Sperimenta migliaia di varianti per trovare la versione più veloce e efficiente possibile.

Grazie a CUCo, le intelligenze artificiali future potranno imparare e ragionare molto più velocemente, perché i "cuochi" del computer non dovranno più fermarsi ad aspettare che qualcuno porti loro gli ingredienti.

Sommario tecnico

1. Il Problema

Nell'era dell'addestramento e dell'inferenza di Large Language Models (LLM) su larga scala, l'efficienza delle GPU è fondamentale. Sebbene i kernel di calcolo (computation) e quelli di comunicazione siano stati ottimizzati separatamente, la loro co-progettazione (co-design) rimane una sfida enorme.

• Limitazioni attuali: Storicamente, l'orchestrazione tra calcolo e comunicazione è gestita dalla CPU (host-side), che introduce overhead di sincronizzazione e impedisce un sovrapposizione (overlap) stretta. Anche con le nuove API lato dispositivo (come NVSHMEM e NCCL Device-Side), scrivere kernel CUDA che fondono calcolo e comunicazione è estremamente complesso, soggetto a errori e richiede una conoscenza profonda dell'hardware e della topologia di rete.
• Complessità dello spazio di ricerca: Le decisioni di progettazione (dove iniziare la comunicazione, quando sincronizzare, dimensioni dei blocchi, granularità) sono controintuitive e altamente dipendenti dal contesto hardware. I compilatori ML esistenti trattano la comunicazione come un effetto collaterale lato host, rendendo difficile integrarla direttamente nei kernel.
• Inadeguatezza degli approcci esistenti: I sistemi di generazione di kernel basati su agenti (LLM) attuali falliscono in ambienti distribuiti perché mancano di supporto per le comunicazioni collettive multi-GPU, richiedono spesso training specifici e tendono a generare codice errato saltando la fase di decomposizione logica.

2. Metodologia: Il Framework CUCo

CUCo è un framework agentic (guidato da agenti AI) che automatizza la generazione di kernel CUDA ad alte prestazioni, fondendo calcolo e comunicazione senza necessità di addestramento (training-free). L'architettura si basa su tre componenti principali:

A. Specifica dello Spazio di Progettazione (Design Space Specification)

Per guidare l'agente, CUCo definisce uno spazio di ottimizzazione strutturato basato su cinque dimensioni:

1. Backend (B): Scelta tra GIN (GPU-Initiated Networking, RDMA) o LSA (Load/Store Accessible, memoria peer intra-nodo).
2. Issuer (I): Granularità dell'emissione (Thread, Warp, o CTA).
3. Placement (P): Aggressività dell'overlap rispetto alla complessità dell'ordinamento.
4. Sync Scope (S): Ampiezza della sincronizzazione.
5. Granularity (G): Dimensione dei chunk di trasferimento.
   Questo spazio vincola l'agente a scelte semanticamente valide, riducendo le "allucinazioni" di codice.

B. Agente Fast-Path (Priorità alla Correttezza)

L'obiettivo è produrre rapidamente un kernel funzionante e corretto, anche se non ottimizzato.

• Flusso: Analizza il codice host esistente, trasforma le chiamate NCCL host-driven in primitive device-initiated e annota le regioni modificabili.
• Strategia: Utilizza una pipeline a tre fasi (Analisi, Trasformazione Host-to-Device, Annotazione) con un ciclo di feedback "LLM-Judge" per garantire che il kernel compili e sia corretto prima di passare alla fase successiva.
• Risultato: Genera un baseline sicuro che serve come "seme" per l'ottimizzazione successiva.

C. Agente Slow-Path (Priorità alle Prestazioni)

Questo agente esegue una ricerca evolutiva per ottimizzare le prestazioni partendo dal seme del Fast-Path.

• Ricerca Evolutiva Multi-Isola: Utilizza popolazioni indipendenti (isole) per evitare la convergenza prematura su soluzioni subottimali.
• Mutazione Guidata da LLM: Invece di mutazioni casuali, l'LLM propone modifiche semantiche (patch diff, riscritture complete, incroci) basandosi su:
  • Contesto del codice genitore.
  • Ispirazioni da un archivio di soluzioni di successo (MAP-Elites).
  • Raccomandazioni meta-riassunte dai trend storici.
• Valutazione a Cascata: I candidati vengono valutati in tre livelli: compilazione, verifica di correttezza numerica (eseguita su cluster reale) e benchmark delle prestazioni. Solo i candidati validi procedono.
• Strategia Explore-Exploit: La ricerca inizia con una fase di esplorazione (ricerca di architetture strutturalmente diverse) per poi passare a una fase di sfruttamento (raffinamento fine dei parametri).

3. Contributi Chiave

1. Primo Framework Agentic per Co-design: CUCo è il primo sistema che automatizza la fusione di calcolo e comunicazione in kernel CUDA distribuiti, superando le limitazioni dei compilatori tradizionali.
2. Separazione Correttezza/Prestazioni: L'architettura a due percorsi (Fast-Path per la correttezza, Slow-Path per le prestazioni) risolve il problema della complessità, permettendo all'agente di non bloccarsi in cicli di generazione infiniti per correggere errori di base.
3. Spazio di Ricerca Vincolato e Semantico: L'uso di direttive esplicite e di un contesto condizionato all'hardware (API NCCL specifiche, topologia di rete) garantisce che le soluzioni generate siano valide e portabili.
4. Open Source: Il framework è implementato in 18.000 righe di Python e sarà rilasciato come open source per accelerare la ricerca nella co-progettazione.

4. Risultati Sperimentali

CUCo è stato valutato su quattro carichi di lavoro rappresentativi su cluster multi-GPU (NVIDIA A100, NVLink e RoCE):

• Flash Attention con Parallelismo del Contesto: Riduzione della latenza fino al 11.3% (speedup di 1.13x). Il guadagno deriva dall'eliminazione dei tempi di inattività della CPU e dal pipelining per-tile.
• DeepSeek-V3 MoE (Dispatch/Combine): Riduzione della latenza fino al 18.0% in scenari con squilibrio di routing. CUCo elimina l'overhead di padding e sovrappone il calcolo degli esperti con il trasferimento di rete.
• Trasferimento KV Cache: Riduzione del 7.6% - 13.6% eliminando i gap tra calcolo e invio dati.
• GEMM + AllGather: Riduzione della latenza fino al 36.2% per matrici grandi, grazie alla fusione completa.
• Risultato Complessivo: CUCo ha dimostrato riduzioni della latenza end-to-end tra il 5.3% e il 26.2% rispetto alle baseline host-driven, con un speedup massimo di 1.57x.
• Ablation Study: Gli esperimenti dimostrano che saltare il Fast-Path porta a un fallimento iniziale nella generazione di kernel corretti e a una convergenza più lenta, confermando la necessità dell'approccio a due fasi.

5. Significato e Impatto

CUCo segna un cambio di paradigma nella programmazione di sistemi distribuiti per il Machine Learning.

• Democratizzazione dell'ottimizzazione: Rende accessibile la creazione di kernel di fusione complessi, che in precedenza richiedevano mesi di lavoro manuale da parte di esperti di CUDA.
• Sfruttamento delle nuove API: Abilita pienamente le potenzialità delle nuove API di comunicazione lato dispositivo (NVSHMEM, NCCL Device-Side) che finora erano difficili da utilizzare correttamente.
• Futuro della Co-progettazione: Dimostra che l'uso di agenti AI strutturati può esplorare spazi di progettazione inaccessibili agli umani, aprendo la strada a kernel sempre più efficienti man mano che le architetture hardware e le topologie di rete evolvono.

In sintesi, CUCo risolve il collo di bottiglia della comunicazione nei cluster GPU moderni automatizzando la complessa arte di fondere calcolo e comunicazione, garantendo correttezza e massimizzando le prestazioni.