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

CUCo es un marco de trabajo impulsado por agentes que automatiza la generación de kernels CUDA de alto rendimiento mediante la co-optimización conjunta de la computación y la comunicación, logrando reducir la latencia de extremo a extremo hasta en un 1,57 veces en comparación con los enfoques existentes.

Lo esencial

Imagina que tienes un equipo de cocineros de élite (las GPUs) trabajando en una receta gigante para una cena masiva (entrenar una Inteligencia Artificial).

Hasta ahora, había un problema grave en la cocina:

1. Los cocineros eran muy rápidos, pero tenían que esperar a que un jefe de cocina humano (la CPU) les dijera cuándo pasar los ingredientes de una mesa a otra.
2. Mientras el jefe de cocina corría a dar instrucciones, los cocineros se quedaban parados, sin hacer nada, perdiendo tiempo valioso.
3. Además, si querían mezclar el acto de cocinar (computación) con el acto de pasar los platos (comunicación), tenían que hacerlo por pasos separados, como si cocinaran un plato, lo dejaran en la mesa, esperaran a que alguien lo recogiera, y luego cocinaran el siguiente. Era un proceso lento y propenso a errores.

¿Qué es CUCo?

CUCo es como un nuevo sistema de gestión inteligente que elimina al jefe de cocina humano de la ecuación y le da a los cocineros una tablet mágica (un agente de IA) que les permite:

• Cocinar y pasar los platos al mismo tiempo.
• Coordinarse entre ellos sin esperar órdenes externas.
• Encontrar la forma más rápida y eficiente de hacerlo, probando miles de combinaciones en segundos.

¿Cómo funciona? (La analogía de los dos agentes)

El sistema CUCo no intenta adivinar la solución perfecta de inmediato. Usa un enfoque de dos pasos, como si tuviera dos tipos de ayudantes:

1. El Agente "Rápido" (El que asegura que no se queme la casa)

• Su misión: La seguridad y la lógica.
• Cómo actúa: Este agente es conservador. Su único objetivo es crear un código que funcione y no tenga errores. No le importa si es el más rápido del mundo, solo que no se detenga ni se rompa.
• La analogía: Es como un arquitecto que dibuja los planos básicos de un puente. Asegura que el puente no se caiga, que los pilares estén en su sitio y que el tráfico pueda circular. Es la "base sólida" sobre la que se construirá todo lo demás.

2. El Agente "Lento" (El que busca la velocidad de la luz)

• Su misión: La optimización y la velocidad.
• Cómo actúa: Este agente toma el plano seguro del primer agente y empieza a hacer "experimentos locos". Es como un chef que prueba mil variaciones de la misma receta: "¿Qué pasa si cocino el plato mientras lo paso? ¿Y si cambio el orden de los ingredientes? ¿Y si uso un camino más corto?".
• La analogía: Imagina un equipo de ingenieros que toma el puente seguro y empieza a probar materiales más ligeros, formas más aerodinámicas y nuevos motores. Usan un sistema de evolución: prueban una idea, si funciona mejor, la guardan; si falla, la descartan y aprenden de su error. Repiten esto miles de veces hasta encontrar la versión perfecta para ese puente específico.

¿Por qué es tan importante?

En el mundo de la Inteligencia Artificial moderna, los modelos son tan grandes que necesitan usar cientos de tarjetas gráficas a la vez.

• Antes: El tiempo se perdía esperando a que las tarjetas se "hablaran" entre sí. Era como tener una conversación telefónica donde cada vez que una persona hablaba, tenía que colgar, esperar a que la otra contestara, y luego volver a marcar.
• Con CUCo: Las tarjetas gráficas se hablan directamente mientras trabajan. Es como si todos estuvieran en una misma sala gritándose las instrucciones mientras cocinan.

Los Resultados (El "Milagro")

El papel muestra que CUCo ha logrado:

• Reducir el tiempo de espera: En algunos casos, ha hecho que las tareas sean 1.57 veces más rápidas.
• Ahorro de tiempo: En una tarea específica (Flash Attention), ahorró más de 139 milisegundos, lo cual parece poco, pero en el mundo de los superordenadores, es una eternidad.
• Adaptabilidad: Funciona bien tanto si las tarjetas están en la misma máquina (como en una cocina pequeña) o en edificios diferentes conectados por fibra óptica (como cocinas en diferentes ciudades).

En resumen

CUCo es un arquitecto y un ingeniero de carreras trabajando en equipo.

1. El primero asegura que el coche no se rompa.
2. El segundo lo modifica hasta que se convierte en un Fórmula 1.

Antes, los humanos tenían que hacer ambos trabajos manualmente, lo cual era lento, difícil y propenso a errores. CUCo automatiza este proceso, permitiendo que la Inteligencia Artificial crezca más rápido porque sus "motores" (las GPUs) ya no pierden tiempo esperando órdenes, sino que trabajan y se comunican en perfecta armonía.

Resumen técnico

Resumen Técnico: CUCo

1. El Problema

El entrenamiento e inferencia de modelos de lenguaje grandes (LLM) a escala requiere un uso eficiente de clústeres de GPU distribuidos. Históricamente, la optimización de estos sistemas se ha centrado casi exclusivamente en los kernels de computación, dejando los kernels de comunicación (como las operaciones colectivas de NCCL) como un cuello de botella significativo.

Aunque las APIs recientes (como NVSHMEM y las APIs de dispositivo de NCCL) permiten que las GPU inicien directamente la comunicación, integrando la lógica distribuida dentro del propio kernel, esto ha creado un nuevo desafío:

• Complejidad de Diseño: Escribir kernels que orquesten conjuntamente la computación y la comunicación es extremadamente difícil, propenso a errores y laborioso. El espacio de diseño es vasto y contra-intuitivo (ej. granularidad de tiles, sincronización, solapamiento).
• Falta de Herramientas: Los compiladores ML actuales asumen que la comunicación ocurre fuera del kernel (lado del host), lo que impide la fusión óptima.
• Heterogeneidad: La configuración óptima depende fuertemente de la topología de red específica (InfiniBand vs. RoCE) y del hardware, haciendo que el ajuste manual sea inviable.

2. Metodología: El Marco CUCo

CUCo es un marco de trabajo agente-driven (impulsado por agentes de IA) que automatiza la generación de kernels CUDA de alto rendimiento que fusionan computación y comunicación. No requiere entrenamiento previo (training-free) y se basa en tres componentes interconectados:

A. Especificación del Espacio de Diseño (Design Space Specification)

Para evitar que los agentes "alucinen" estrategias inválidas, CUCo define un espacio de búsqueda estructurado y restringido mediante directivas explícitas. Este espacio se compone de cinco dimensiones:

1. Backend (B): Elección entre GIN (GPU-Initiated Networking, RDMA) o LSA (Load/Store Accessible, memoria local).
2. Colocación (P): Agresividad del solapamiento vs. complejidad de ordenación.
3. Alcance de Sincronización (S): Qué tan estrecha es la sincronización.
4. Granularidad del Emisor (I): Nivel de hilo, warp o bloque (CTA) que inicia la comunicación.
5. Tamaño de Chunk (G): Granularidad de la transferencia de datos.

B. Agente de Ruta Rápida (Fast-Path Agent)

Su objetivo es la corrección, no el rendimiento máximo.

• Función: Convierte rápidamente un programa basado en host (CPU) en un kernel de dispositivo iniciado (device-initiated).
• Estrategia: Utiliza un bucle de retroalimentación LLM-Juez para garantizar que el kernel compile y sea funcionalmente correcto.
• Enfoque Conservador: Genera un kernel con fases globales explícitas y barreras de sincronización estrictas para evitar condiciones de carrera o deadlocks distribuidos.
• Resultado: Produce una "semilla" (seed) correcta que sirve como punto de partida seguro para la optimización.

C. Agente de Ruta Lenta (Slow-Path Agent)

Su objetivo es la optimización de rendimiento mediante una búsqueda evolutiva.

• Estrategia: Utiliza un algoritmo evolutivo de "islas múltiples" (multi-island) para explorar el espacio de diseño definido por las directivas.
• Mutación Guiada por LLM: En lugar de mutaciones aleatorias, un LLM propone cambios semánticos (parches diff, reescrituras completas o cruce de código) dentro de bloques de código marcados como mutables.
• Fases de Búsqueda:
  1. Exploración: Alta temperatura, busca arquitecturas estructuralmente diversas (ej. pipelines divididos, fusión agresiva).
  2. Explotación: Baja temperatura, refina las mejores arquitecturas encontradas ajustando parámetros.
• Evaluación en Cascada: Los candidatos se evalúan en tres niveles: compilación, verificación de corrección numérica y benchmarking de rendimiento. Los fallos se almacenan para evitar repetir errores.

3. Contribuciones Clave

1. Marco Agente para Co-diseño: Primer sistema que automatiza la fusión de computación y comunicación en kernels CUDA utilizando agentes de IA, superando las limitaciones de los compiladores tradicionales.
2. Directivas de Optimización Estructuradas: Introduce una interfaz declarativa que obliga al agente a razonar sobre decisiones críticas (backend, sincronización, granularidad) antes de generar código, reduciendo errores semánticos.
3. Arquitectura de Dos Fases (Fast/Slow Path): Demuestra que separar la búsqueda de "corrección" (ruta rápida) de la búsqueda de "rendimiento" (ruta lenta) mejora drásticamente la eficiencia de la muestra y la estabilidad del proceso.
4. Open Source: El marco está implementado en 18k líneas de Python y será de código abierto para la comunidad.

4. Resultados y Evaluación

CUCo se evaluó en cuatro cargas de trabajo representativas de GPU múltiples:

1. Flash Attention con Parallelismo de Contexto: (4 GPUs, NVLink).
2. Capa MoE de DeepSeek-V3: (2 GPUs, enlaces RoCE inter-nodo).
3. Transferencia de KV-Cache: (2 GPUs, NVLink).
4. GEMM + AllGather: (4 GPUs, mix intra/inter-nodo).

Hallazgos Principales:

• Reducción de Latencia: CUCo logró reducciones de latencia de extremo a extremo entre 5.3% y 26.2% en comparación con las implementaciones basadas en host (NCCL estándar).
• Aceleración: Se observó una aceleración de hasta 1.57x en el caso de Flash Attention.
• Mecanismos de Mejora:
  • Eliminación de la sobrecarga de la orquestación del host (llamadas API NCCL, sincronización de eventos).
  • Eliminación de "burbujas" en la tubería (pipeline bubbles) al permitir que la computación se solape con la comunicación a nivel de tile.
  • Uso eficiente de la memoria y la red mediante el solapamiento fino.
• Estudios de Ablación:
  • Sin el agente de ruta rápida, el agente evolutivo tarda muchas generaciones más en encontrar un kernel correcto y converge a un rendimiento inferior.
  • La estrategia de dos fases (exploración seguida de explotación) es superior a la búsqueda de una sola fase, logrando mejores resultados con menos generaciones.

5. Significancia

CUCo representa un cambio de paradigma en el desarrollo de sistemas distribuidos para IA. Demuestra que la complejidad de la co-optimización de computación y comunicación, que anteriormente requería ingeniería manual experta y propensa a errores, puede ser automatizada mediante agentes de IA estructurados.

A medida que la comunicación iniciada por dispositivo se vuelve central en el aprendizaje distribuido, CUCo ofrece una vía para descubrir oportunidades de optimización que son intratables para los humanos, permitiendo que los clústeres de GPU operen más cerca de su límite teórico de eficiencia. Esto es crucial para escalar modelos de IA cada vez más grandes en hardware heterogéneo.