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

Yoga Sri Varshan Varadharajan; Bodun Hu; Saurabh Agarwal; Aditya Akella

cs.DCMar 2, 2026

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

Authors: Yoga Sri Varshan Varadharajan, Bodun Hu, Saurabh Agarwal, Aditya Akella

Organizations: UT Austin

Abstract

Computation and communication in distributed LLM training and inference are traditionally optimized in isolation; expert-crafted systems such as DeepEP, FLUX, and TokenWeave show the potential of co-design but require deep systems expertise and hardware-specific tuning; CUCo is an agentic framework that automates compute-communication co-design of CUDA kernels by combining a structured design-space formalization with a correctness-first fast-path agent for reliable baselines and an evolution-driven slow-path agent for high-performance strategies, achieving up to 1.57x speedup across four multi-GPU workloads and discovering a two-stream overlap strategy on a DeepSeek-V3 MoE layer that hides dispatch behind local compute at an LLM inference cost under $10 per workload.

Figures & tables

System	Multi-GPU	Comm Support	No Training	Device-Init
CUDAForge Zhang et al. (2025)	✗	✗	✓	✗
STARK Dong et al. (2025)	✗	✗	✓	✗
Kevin-32B Baronio et al. (2025)	✗	✗	✗	✗
KernelFalcon Wang et al. (2025)	✗	✗	✓	✗
CUCo	✓	✓	✓	✓

Table 1: Comparison of agentic CUDA kernel generation systems.

Concrete API Dimensions	Allowed Values
$\mathcal{B}$ : Backend	GIN \| LSA \| Hybrid
$\mathcal{M}$ : Sync Mechanism	Barrier \| Signal \| SignalShadow \| Counter
$\mathcal{I}$ : Issuer	Thread \| Warp \| WarpSpan \| Tile⟨N⟩ \| CTA
$\mathcal{O}$ : Ordering	Relaxed \| Acquire \| Release \| AcqRel
$\mathcal{K}$ : Contexts	$1$ \| $2$ \| $N$
Intent-Based Dimensions	Allowed Values

Table 2: Optimization directive dimensions. Concrete dimensions correspond to fixed NCCL API identifiers. Intent-based dimensions express high-level objectives that the agent interprets during code generation.

System	$\mathcal{B}$	$\mathcal{M}$	$\mathcal{P}$	$\mathcal{S}$	$\mathcal{I}$	$\mathcal{G}$	$\mathcal{O}$	$\mathcal{K}$
DeepEP (NVL)	LSA	Barrier	Deferred	Local	CTA	Per-peer	Release	–
DeepEP (IB)	GIN	Signal	Deferred	World	CTA	Per-peer	Acquire	1
FLUX	LSA	Barrier	Tile-fused	Local	Warp	Per-tile	AcqRel	–

Table 3: Expert-crafted systems mapped onto the design space

\mathcal{C}

by semantic equivalence. These systems use custom transport implementations, not NCCL’s device API. DeepEP’s IB path uses custom RDMA one-sided puts (GIN semantics); its NVLink path uses direct peer memory stores (LSA semantics).

Workload	Dominant bottleneck	Discovered strategy	Key design choices
Flash Attention	Pipeline bubbles between ring exchange rounds	Persistent cooperative kernel with per-tile pipelining; comm block polls tile readiness via atomics	$\mathcal{P}$ : tile-fused; $\mathcal{M}$ : counter; $\mathcal{I}$ : CTA (1 comm + 80 compute)
MoE Dispatch/ Combine	Sequential comm–compute phases; padding overhead on imbalanced ranks	Self/remote compute split with two-stream dispatch overlap; variable-size per-peer transfers	$\mathcal{B}$ : GIN (inter-node); $\mathcal{P}$ : stream-split; $\mathcal{G}$ : per-peer
KV-Cache Transfer	Compute-to-send idle gap from host-launched sends	GPU-triggered chained send: K GEMM $\to$ send K $\to$ V GEMM $\to$ send V with signal	$\mathcal{I}$ : device; $\mathcal{O}$ : release; $\mathcal{M}$ : signal
GEMM + AllGather	Post-compute idle before collective	Immediate device-initiated broadcast fused after GEMM completion	$\mathcal{P}$ : deferred; $\mathcal{B}$ : LSA (intra) / GIN (inter)

Table 4: What CUCo ’s agent discovered: per-workload dominant bottleneck, the co-design strategy the evolutionary search converged on, and the key design-space choices (

\mathcal{C}

) that realize it.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 14: NCCL device-initiated API: All-to-All CUDA kernel using the GPU-Initiated Networking (GIN) and Load/Store Accessible (LSA) API. GIN issues RDMA-style puts to remote peers (lines 31, 36); LSA accesses intra-node peer memory via direct loads and stores (lines 46–47). Fine-grained synchronization primitives (lines 12–15) and gin.waitSignal (line 53) enable compute–communication fusion that is out of reach for host-driven APIs.

Field	Content
Purpose	Device-side peer memory access and barrier synchronization within local GPU teams via direct load/store.
Core API	ncclLsaBarrierSession (provides synchronization), ncclGetLsaPointer (team-indexed access), ncclGetLocalPointer (local access), …
Sync Scope	ncclTeamLsa (intra-node local team), ncclTeamWorld (all ranks); tighter scope reduces overhead but limits overlap potential.
Invariants	Barrier index must be unique per CTA ( blockIdx.x ), …
Canonical Pattern	Instantiate ncclLsaBarrierSession with blockIdx.x ; perform memory operations; call sync(ncclCoopCta(), memory_order_release) ; access peers via ncclGetLsaPointer .

Appendix

Table 6: Sample knowledge base structure as presented within the LSA interface.

Workload	Fast-Path Iter.	Slow-Path Gen.	Wall Time	LLM Cost
Flash Attention	3	20	${\sim}$ 1.8 hr	$10.96
MoE (GIN)	2	18	${\sim}$ 1.5 hr	$10.20
KV-Cache	2	18	${\sim}$ 1.5 hr	$9.74
GEMM + AllGather	2	10	${\sim}$ 51 min	$4.98

Appendix

Table 7: Search cost per workload. All runs use Claude Sonnet 4.5.

GPU Kernel Optimization

GPU kernel optimization improves the speed and memory efficiency of computations by tailoring code to GPU hardware and workloads.

See all 209 papers

AgentCompile: An LLM-Guided Compiler for Direct CUDA Inference
Jun 4, 2026Xuanzhe Li, Ziyan Weng, Zhiyu Zhu +1High-Performance Computing GPU Kernel Generation
City University of Hong Kong (Dongguan) · City University of Hong Kong
Coco: An Agentic Copilot for the Hardware--Software Co-Design Lifecycle
Oct 1, 2026Samuel Kushnir, Kavya Sreedhar, Yeshwanth Reddy Pogula +8Hardware-Software Co-Design LLM Agents
Google DeepMind · Google · MIT
CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe
Apr 1, 2026Tara Saba, Zhiyang Chen, Jikai Jason Li +3GPU Kernel Generation GPU Kernel Optimization
Department of Computer Science, University of Toronto · Standard Kernel

Jun 4, 2026J/K move · Enter open · S save

AgentCompile: An LLM-Guided Compiler for Direct CUDA Inference

Xuanzhe Li, Ziyan Weng, Zhiyu Zhu, Junhui Hou

City University of Hong Kong (Dongguan) · City University of Hong Kong

High-Performance Computing GPU Kernel Generation GPU Acceleration Efficient Transformer Inference Large Language Model-Guided Optimization LLM Inference Acceleration+1

Open paper AskPDF arXiv

Transformer inference increasingly depends on specialized compiler and runtime support, but real model graphs still require semantic decisions about which regions are worth specializing and which CUDA implementation families are plausible. We present AgentCompile, an LLM-guided CUDA inference compiler that uses LLM outputs only as advisory search metadata. Given compiler-derived region summaries and bounded candidate spaces, the LLM proposes semantic labels, candidate priorities, parameter hints, and risk annotations; the compiler materializes CUDA candidates through templates, checks interface and hardware constraints, validates candidates empirically, selects implementations by measured latency, and falls back when specialization is unsupported or unprofitable. In end-to-end autoregressive generation, AgentCompile averages 5.66x, 4.05x, and 4.26x speedup over PyTorch eager on Qwen3-1.7B, Qwen3-4B, and Llama-3.2-1B-Instruct, respectively, across five representative workloads. We will open-source the project.

Abstract

GPU Kernel Optimization

Explore similar work

AgentCompile: An LLM-Guided Compiler for Direct CUDA Inference

Coco: An Agentic Copilot for the Hardware--Software Co-Design Lifecycle

CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe