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

O CUCo é um framework agênico sem treinamento que automatiza a geração de kernels CUDA de alto desempenho, otimizando conjuntamente computação e comunicação para reduzir a latência em treinamento e inferência de LLMs distribuídos.

O essencial

Imagine que você está organizando uma grande festa de dança em um prédio com vários andares (os GPUs). O objetivo é que todos os dançarinos se movam perfeitamente sincronizados para criar um espetáculo incrível (o treinamento de uma Inteligência Artificial).

Neste cenário, existem dois tipos de trabalho principais:

1. Dançar: Fazer os passos de dança complexos (o Cálculo ou Computation).
2. Passar mensagens: Correr para o próximo andar para entregar uma nota ou um objeto para outro dançarino (a Comunicação ou Communication).

O Problema: O "Gerente" Exausto

Antigamente, existia um Gerente Humano (o processador da CPU) que ficava no centro de controle.

• Ele dizia: "Ok, dançarinos do 1º andar, parem de dançar! Agora, corram para o 2º andar e entreguem a nota."
• Depois, ele esperava todos chegarem e dizia: "Agora, voltem a dançar!"

Isso era muito lento. O Gerente ficava cansado, e os dançarinos ficavam parados esperando ordens. Além disso, escrever as regras para essa dança era tão difícil e cheio de erros que só especialistas conseguiam fazer, e mesmo assim, cometiam muitos tropeços.

A Solução: CUCo (O Maestro Robô)

Os autores deste artigo criaram o CUCo, que é como um Maestro Robô Inteligente que não precisa de um Gerente Humano. Ele ensina os próprios dançarinos (os GPUs) a se organizarem sozinhos.

O CUCo funciona em duas etapas, como se fosse um estágio de treinamento:

1. O Treinador Rápido (Fast-Path Agent)

Imagine um treinador que chega e diz: "Esqueçam a velocidade por um momento. O mais importante é que ninguém bata cabeça e que a música pare se alguém errar."

• Ele pega o código antigo (cheio de ordens do Gerente Humano) e o transforma em um código onde os dançarinos podem falar diretamente entre si.
• Ele garante que a dança funcione perfeitamente, mesmo que seja um pouco lenta. É a base segura.

2. O Treinador de Elite (Slow-Path Agent)

Agora que temos uma dança que funciona, entra o Treinador de Elite. Ele é como um cientista louco que gosta de experimentar coisas novas.

• Ele usa uma técnica de "evolução": cria centenas de versões levemente diferentes da dança.
• Ele testa cada uma: "E se o dançarino do 3º andar entregar a nota enquanto o do 1º ainda está dançando? E se eles trocarem de lugar mais rápido?"
• Ele descarta as versões que dão errado e mantém as que são mais rápidas.
• Com o tempo, ele encontra a dança perfeita, onde a entrega de mensagens acontece enquanto a dança acontece, sem ninguém parar.

O Grande Truque: A "Fusão"

O segredo do CUCo é a Fusão.

• Antes: Dança -> Parada -> Entrega -> Parada -> Dança.
• Com CUCo: Dança e Entrega acontecem ao mesmo tempo, como se fosse um único movimento contínuo.

O CUCo usa uma "Inteligência Artificial" (um modelo de linguagem) para escrever o código dessa dança, mas com regras muito estritas para garantir que não haja erros. Ele não chuta; ele explora um espaço de possibilidades de forma inteligente.

Os Resultados

Quando eles testaram isso em tarefas reais (como treinar modelos de IA gigantes):

• A dança ficou até 1,57 vezes mais rápida.
• Isso significa que o tempo de espera do Gerente Humano foi eliminado.
• Os dançarinos (GPUs) trabalharam o tempo todo, sem ficar ociosos.

Resumo em uma Analogia Final

Pense em uma cozinha de restaurante:

• Método Antigo: O Chef (CPU) grita para o cozinheiro: "Pare de cortar cebola! Corra até a geladeira, pegue o tomate, volte e corte." O cozinheiro perde tempo correndo e esperando ordens.
• Método CUCo: O Chef ensina o cozinheiro a ter uma esteira rolante. Enquanto ele corta a cebola, a esteira traz o tomate automaticamente para a mão dele. Ele nunca para de trabalhar.

O CUCo é o sistema que aprende a desenhar essa esteira rolante perfeita para cada tipo de cozinha, garantindo que a comida (os dados) chegue mais rápido e a cozinha funcione no limite da eficiência.

Resumo técnico

Título: CUCo: Um Framework Agente para Co-design de Computação e Comunicação

1. O Problema

O treinamento e inferência de Grandes Modelos de Linguagem (LLMs) em larga escala dependem de clusters distribuídos de GPUs. Para maximizar a eficiência, é crucial otimizar não apenas os kernels de computação, mas também a sobreposição (overlap) entre computação e comunicação.

• Limitações Atuais: Historicamente, os kernels de computação (ex: cuBLAS, cuDNN) e comunicação (ex: NCCL) evoluíram de forma isolada. A orquestração era feita pelo host (CPU), introduzindo sobrecarga de sincronização e impedindo um entrelaçamento (interleaving) fino entre as duas fases.
• O Novo Paradigma: Interfaces recentes (como NVSHMEM e APIs de dispositivo do NCCL) permitem que os kernels CUDA invoquem primitivas de comunicação diretamente, eliminando a intervenção da CPU.
• O Desafio: Projetar manualmente kernels que gerenciem conjuntamente computação e comunicação é extremamente difícil, propenso a erros e laborioso. O espaço de design é vasto e complexo, envolvendo decisões sobre granularidade de tiles, sincronização, escolha de primitivas (GIN vs. LSA) e topologia de rede. Compiladores ML modernos não suportam bem essa fusão, tratando a comunicação como um efeito colateral do lado do host. Além disso, a heterogeneidade de hardware (topologias de rede, largura de banda) torna o ajuste manual inviável.

2. Metodologia: O Framework CUCo

O CUCo é um framework agentic (acionado por agentes de IA) que gera automaticamente kernels CUDA de alto desempenho, fundindo computação e comunicação. Ele opera sem treinamento prévio (training-free) e utiliza uma abordagem de dois caminhos (fast-path e slow-path).

A. Especificação do Espaço de Design (Design Space Specification)

Para evitar que o agente "alucine" estratégias inválidas, o CUCo define um espaço de design estruturado e limitado por cinco dimensões:

1. Backend (B): Escolha entre GIN (GPU-Initiated Networking, para RDMA entre nós) ou LSA (Load/Store Accessible, para memória intra-nó).
2. Issuer (I): Granularidade de emissão (Thread, Warp ou CTA).
3. Placement (P): Aggressividade da sobreposição (fusão vs. ordenação).
4. Sync Scope (S): Escopo de sincronização (global vs. local).
5. Granularity (G): Tamanho dos blocos de transferência.

O agente deve emitir uma Diretiva de Otimização explícita antes de gerar o código, garantindo que as decisões sejam fundamentadas na semântica da API e no contexto de hardware.

B. Agente de Caminho Rápido (Fast-Path Agent)

• Objetivo: Garantir correção funcional.
• Funcionamento: Transforma um programa base (orquestrado pelo host) em uma implementação iniciada pelo dispositivo (device-initiated).
• Estratégia: Prioriza a segurança sobre o desempenho. Usa barreiras globais explícitas e fases bem definidas para evitar deadlocks e estados inconsistentes em sistemas distribuídos.
• Saída: Gera um "seed" (semente) correto, mas não otimizado, que serve como base para a próxima fase.

C. Agente de Caminho Lento (Slow-Path Agent)

• Objetivo: Maximizar o desempenho.
• Mecanismo: Utiliza uma busca evolutiva guiada por LLM (Large Language Model).
• População Multi-Ilha: Mantém várias populações independentes ("ilhas") para evitar convergência prematura para ótimos locais.
• Operador de Mutação LLM: Em vez de edições aleatórias, o LLM propõe mutações semânticas (patches de diff, reescritas completas ou cruzamento de código) baseadas em:
  • Feedback de avaliação em cascata (compilação, execução, benchmark).
  • Um banco de dados histórico de candidatos (sucessos e falhas).
  • Resumo de padrões de gerações anteriores.
• Fases de Busca:
  • Exploração: Alta temperatura, focada em diversidade estrutural (ex: novos padrões de pipeline).
  • Exploração (Exploitation): Baixa temperatura, focada em refinar os melhores candidatos encontrados.

3. Principais Contribuições

1. Framework Agente para Co-design: Primeira abordagem que automatiza a geração de kernels fundidos de computação e comunicação, superando as limitações de compiladores e ferramentas manuais.
2. Arquitetura Híbrida (Fast/Slow Path): A separação entre a geração de um baseline correto (Fast-Path) e a otimização evolutiva (Slow-Path) resolve o dilema entre correção e desempenho, permitindo que o agente explore um espaço de design complexo sem falhar catastróficamente.
3. Especificação Declarativa: A introdução de diretrizes explícitas (Backend, Granularidade, etc.) ancora o raciocínio do LLM na realidade das APIs de comunicação (NCCL Device-Side), reduzindo erros de sintaxe e semântica.
4. Busca Evolutiva Guiada por Contexto: Uso de um banco de dados compartilhado e feedback de avaliação em cascata para guiar o LLM, permitindo que o sistema aprenda com falhas passadas e padrões de sucesso.

4. Resultados

O CUCo foi avaliado em quatro cargas de trabalho representativas de multi-GPU:

1. Flash Attention com Parallelismo de Contexto: (4 GPUs, intra-nó).
2. DeepSeek-V3 MoE (Mixture of Experts): (2 GPUs, inter-nó, roteamento esparsos).
3. Transferência de KV Cache: (2 GPUs, intra-nó).
4. GEMM + AllGather: (4 GPUs, intra e inter-nó).

Desempenho:

• O CUCo reduziu consistentemente a latência de ponta a ponta em todas as cargas de trabalho.
• Aceleração Máxima: Até 1.57x (redução de latência) em comparação com as bases lineares baseadas em host (NCCL tradicional).
• Reduções Específicas:
  • Flash Attention: ~11.3% de redução de latência (eliminação de bubbles de pipeline e sobrecarga de API do host).
  • DeepSeek MoE: Até 18% de redução em cenários de desbalanceamento de tokens.
  • GEMM + AllGather: Até 36.2% de redução em matrizes grandes.

Estudos de Ablação:

• A remoção do Fast-Path resultou em uma busca evolutiva muito mais lenta e com muitas falhas iniciais (o agente demorou 5 gerações para produzir um kernel correto, contra 2 iterações do Fast-Path).
• A estratégia de duas fases (Exploração + Exploração) superou significativamente a busca puramente exploratória ou puramente exploradora, alcançando melhores resultados com menos gerações.

5. Significado e Impacto

O trabalho CUCo demonstra que a otimização manual de kernels para sistemas distribuídos de ML está se tornando um gargalo insustentável devido à complexidade da integração entre computação e comunicação.

• Mudança de Paradigma: O CUCo valida que a geração automática de kernels via agentes de IA é viável e superior para cenários de co-design, especialmente com o advento de APIs de comunicação iniciadas pelo dispositivo.
• Acesso ao Conhecimento: Ao abrir o código-fonte e o framework, o trabalho permite que a comunidade de pesquisa explore otimizações que antes eram inacessíveis devido à complexidade de implementação manual.
• Escalabilidade: A abordagem é adaptável a diferentes topologias de hardware e cargas de trabalho, oferecendo uma solução robusta para a próxima geração de clusters de treinamento de IA.

Em resumo, o CUCo preenche a lacuna crítica entre a capacidade teórica de fusão de kernels (disponível em novas APIs de hardware) e a dificuldade prática de implementá-los, automatizando a descoberta de estratégias de otimização que superam o estado da arte atual.