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

TL;DR

CUCo is an agent-driven framework that automatically generates CUDA kernels to optimize both computation and communication in GPU workloads, significantly reducing latency in large-scale distributed training.

Contribution

CUCo introduces a novel, training-free, agent-based approach for co-optimizing computation and communication kernels in CUDA, addressing a gap in existing kernel optimization methods.

Findings

Outperforms state-of-the-art baselines in GPU kernel efficiency.

Reduces end-to-end latency by up to 1.57 times.

Automates kernel generation without training overhead.

Abstract

Custom CUDA kernel development is essential for maximizing GPU utilization in large-scale distributed LLM training and inference, yet manually writing kernels that jointly leverage both computation and communication remains a labor-intensive and error-prone process. Prior work on kernel optimization has focused almost exclusively on computation, leaving communication kernels largely untouched even though they constitute a significant share of total execution time. We introduce CUCo, a training-free agent-driven workflow that automatically generates high-performance CUDA kernels that jointly orchestrate computation and communication. By co-optimizing these traditionally disjoint components, CUCo unlocks new optimization opportunities unavailable to existing approaches, outperforming state-of-the-art baselines and reducing end-to-end latency by up to $1.57\times$.