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

이 논문은 분산 LLM 학습 및 추론 시 계산과 통신을 공동으로 최적화하여 기존 접근법보다 최대 1.57 배의 지연 시간 감소를 달성하는 자동화된 CUDA 커널 생성 에이전트 프레임워크인 CUCo 를 제안합니다.

핵심

CUCo: GPU 를 위한 '자동화 요리사'와 '마스터 셰프'의 협업

이 논문은 인공지능 (AI) 모델을 훈련하거나 실행할 때, 수천 개의 그래픽 카드 (GPU) 가 함께 일하는 과정에서 발생하는 '지루하고 복잡한 요리 과정'을 자동화하는 새로운 시스템을 소개합니다.

이 시스템을 CUCo라고 부릅니다. CUCo 는 사람이 직접 손으로 코드를 짜는 대신, **AI 에이전트 (지능형 로봇)**가 스스로 최적의 코드를 찾아내게 해줍니다.

이해하기 쉽게 거대한 식당에 비유해서 설명해 드릴게요.

---

1. 문제 상황: "요리사 (계산) 와 배달부 (통신) 가 따로 놀고 있어요"

AI 모델을 학습시킬 때 GPU 는 두 가지 일을 합니다.

1. 계산 (Computation): 데이터를 처리하고 계산하는 일 (요리사가 음식을 만드는 일).
2. 통신 (Communication): 여러 GPU 간에 필요한 데이터를 주고받는 일 (배달부가 재료를 다른 주방으로 옮기는 일).

기존 방식의 문제점:

• 비효율적인 분리: 예전에는 요리사가 요리를 다 끝내면, 배달부가 와서 재료를 옮기고, 다시 요리사가 다음 요리를 시작했습니다. 요리사가 기다리는 동안 배달부가 오기를 기다리는 **'빈 시간 (대기 시간)'**이 너무 많았습니다.
• 수작업의 고통: 이 두 일을 동시에 시키려면 (요리하면서 재료를 옮기기), 매우 정교하고 복잡한 지시를 사람이 직접 내려야 했습니다. 하지만 이 과정은 매우 어렵고 실수가 잦아, 최고의 성능을 내기 힘들었습니다.

2. CUCo 의 해결책: "요리와 배달을 동시에 하는 '자동화 주방'"

CUCo 는 이 문제를 해결하기 위해 **두 명의 AI 로봇 (에이전트)**을 투입합니다. 마치 식당에 신입 요리사와 베테랑 셰프가 팀을 이루는 것과 같습니다.

1 단계: 신입 요리사 (Fast-Path Agent) - "일단 안전하게!"

• 역할: 가장 먼저 틀리지 않고 작동하는 기본 레시피를 만듭니다.
• 방식: "요리와 배달을 동시에 하려고 너무 욕심내지 말고, 일단 순서대로 하되 두 일을 하나의 큰 냄비 (커널) 에 넣어서 시작해라"라고 지시합니다.
• 효과: 성능이 최고는 아니더라도, 절대 실패하지 않는 안정적인 기본기를 확보합니다. 이 단계가 없으면 AI 가 엉뚱한 코드를 만들어서 식당이 멈춰버릴 수 있습니다.

2 단계: 베테랑 셰프 (Slow-Path Agent) - "이제 맛을 살려보자!"

• 역상: 신입 요리사가 만든 기본 레시피를 바탕으로, **최고의 맛 (최고의 성능)**을 찾아냅니다.
• 방식 (진화 과정):
  • "배달이 오기 전에 요리사가 미리 준비할까?", "재료를 조금씩 나누어 옮기면 더 빠를까?", "어떤 순서로 하면 가장 효율적일까?" 등을 수천 번 실험해 봅니다.
  • 실패한 레시피는 버리고, 성공한 레시피는 변형해서 더 좋은 버전을 만듭니다.
  • 이 과정을 반복하며 **가장 빠르고 효율적인 '초고속 레시피'**를 찾아냅니다.

3. CUCo 가 사용하는 '비밀 무기': 설계 도면 (Design Space)

AI 로봇이 막연하게 코드를 짜면 엉망이 됩니다. 그래서 CUCo 는 AI 에게 **정해진 설계 도면 (규칙)**을 줍니다.

• "배달은 이 문 (API) 으로만 하고, 재료는 이 크기 (Chunk) 로 나누고, 신호는 이렇게 보내라"라고 구체적인 규칙을 정해줍니다.
• 이렇게 하면 AI 가 엉뚱한 상상을 하거나 (할루시네이션), 위험한 실험을 하는 것을 막아줍니다. 마치 레고 블록을 주면, AI 는 그 블록으로만 가장 멋진 성을 짓는다는 뜻입니다.

4. 결과: 얼마나 빨라졌나요?

CUCo 를 테스트한 결과, 기존 사람이 직접 만든 방식보다 최대 1.57 배 더 빨라졌습니다.

• 비유: 기존에는 요리사가 100 분 걸리던 일을, CUCo 가 만든 자동화 주방에서는 약 64 분 만에 끝냈습니다.
• 특히, 데이터 양이 작거나 GPU 가 많을수록 (요리할 게 많고 배달이 복잡할수록) 그 차이가 더 커졌습니다.

5. 요약: 왜 이것이 중요한가요?

앞으로 AI 모델은 더 거대해지고, 더 많은 GPU 가 필요해질 것입니다. 사람이 일일이 최적의 코드를 짜는 것은 도저히 불가능해집니다.

CUCo는 AI 가 스스로 "어떻게 하면 가장 효율적으로 일할까?"를 고민하고 코드를 만들어내는 시스템입니다.

• 사람의 역할: "이 요리를 만들어줘"라고 주문만 하면 됩니다.
• AI 의 역할: 안전하게 시작하고, 실험을 반복하며 최고의 성능을 찾아냅니다.

이 기술은 앞으로 거대한 AI 시대에 컴퓨터 자원을 아끼고, 더 빠른 AI 서비스를 가능하게 하는 핵심 열쇠가 될 것입니다.

기술 요약

CUCo: 컴퓨팅 및 통신 공동 설계 (Co-design) 를 위한 에이전트 기반 프레임워크 기술 요약

이 논문은 대규모 분산 LLM(대규모 언어 모델) 학습 및 추론에서 GPU 활용도를 극대화하기 위해 **컴퓨팅 (Computation) 과 통신 (Communication) 을 병합하여 최적화하는 자동화된 에이전트 기반 프레임워크인 'CUCo'**를 제안합니다.

1. 문제 정의 (Problem)

대규모 분산 ML 환경에서는 GPU 간 통신 오버헤드가 전체 실행 시간의 상당 부분을 차지합니다. 기존 접근 방식과 CUCo 가 해결하려는 핵심 문제는 다음과 같습니다.

• 수동 커널 개발의 비효율성: 컴퓨팅과 통신을 동시에 활용하는 고성능 CUDA 커널을 수동으로 작성하는 것은 매우 노동 집약적이며 오류가 발생하기 쉽습니다.
• 통신 커널의 소외: 기존 커널 최적화 연구는 주로 계산 (Computation) 에 집중하여, 통신 커널은 거의 개선되지 않았습니다.
• 호스트 (Host) 기반 오케스트레이션의 한계: 과거에는 CPU 가 계산과 비동기 통신 간의 전환을 관리했습니다. 이로 인해 동기화 오버헤드가 발생하고 두 작업을 밀접하게 교차 (Interleave) 시키는 것이 제한되었습니다.
• 디바이스 기반 통신의 복잡성: NVSHMEM 및 NCCL Device-side API 와 같은 최신 기술은 GPU 커널 내부에서 직접 통신을 수행할 수 있게 했으나, 이를 효과적으로 설계하는 것은 매우 어렵습니다.
  • 설계 공간의 복잡성: 통신 시작 시점, 동기화 타이밍, 타일링 (Tiling) 경계, 청크 크기 등 수많은 변수가 상호작용하며 최적 구성은 하드웨어와 작업 부하에 따라 달라집니다.
  • 기존 도구의 부재: ML 컴파일러는 통신을 커널 외부의 호스트 측 부작용으로 가정하므로, 커널 내부로 통합할 경우 메모리 접근 패턴이 깨지거나 스케줄링이 실패합니다.
  • 하드웨어 이질성: 네트워크 대역폭, 지연 시간, 토폴로지 (InfiniBand vs RoCE 등) 의 차이로 인해 수동 튜닝은 불가능에 가깝습니다.

2. 방법론 (Methodology)

CUCo 는 학습이 필요 없는 (Training-free) 에이전트 기반 워크플로우로, 세 가지 핵심 구성 요소로 이루어져 있습니다.

2.1 구조화된 설계 공간 명세 (Structured Design Space Specification)

에이전트가 합리적인 추론을 할 수 있도록 통신 설계 결정을 위한 명시적인 인터페이스를 제공합니다.

• 최적화 지시문 (Optimization Directive): 에이전트가 코드 생성 전에 반드시 emit 해야 하는 명세입니다.
  • 구체적 차원 (Concrete): 백엔드 (GIN/LSA), 발행자 (Thread/Warp/CTA) 등 NCCL API 의 고정된 식별자.
  • 의도 기반 차원 (Intent-based): 배치 (Placement), 동기화 범위 (Sync Scope), 세분도 (Granularity) 등 에이전트가 작업 부하와 하드웨어에 따라 해석해야 하는 목표.
• 맥락 (Context): 에이전트에게 NCCL Device API 문서, 하드웨어 정보 (GPU 아키텍처, 토폴로지), 그리고 검증된 참조 코드를 제공하여 환각 (Hallucination) 을 방지하고 올바른 API 사용법을 유도합니다.

2.2 패스트-패스 에이전트 (Fast-Path Agent)

정확성 (Correctness) 을 최우선으로 하여 작동 가능한 디바이스 기반 커널을 빠르게 생성합니다.

• 작동 방식: 호스트 기반 NCCL 코드를 분석하여 통신 의존성 그래프를 추출한 후, 이를 디바이스 기반 통신 (Device-initiated) 으로 변환합니다.
• 전략: 계산과 통신을 엄격한 장벽 (Barrier) 으로 구분된 단계로 구성하여, 분산 환경에서의 교착 상태 (Deadlock) 를 방지하고 검증 가능한 기준선 (Baseline) 을 제공합니다.
• 역할: 느린 경로 에이전트 (Slow-Path) 를 위한 '씨앗 (Seed)'을 제공하며, 초기 생성 단계에서 발생하는 오류를 최소화합니다.

2.3 슬로우-패스 에이전트 (Slow-Path Agent)

성능 (Performance) 을 최적화하기 위해 설계 공간을 탐색하는 진화적 (Evolutionary) 검색 에이전트입니다.

• 멀티-아일랜드 진화 (Multi-Island Evolutionary Search): 여러 독립적인 개체군 (Island) 을 유지하여 조기 수렴을 방지하고 다양한 최적화 전략을 탐색합니다.
• LLM 기반 돌연변이 (LLM Mutation): 무작위 편집 대신, LLM 을 활용하여 의미 있는 코드 수정 (Diff, 전체 재작성, 교차) 을 수행합니다.
  • 탐색 (Explore) 단계: 초기에는 구조적으로 다양한 설계 (멀티 스트림 오버랩, 파이프라이닝 등) 를 탐색합니다.
  • 활용 (Exploit) 단계: 이후 가장 유망한 아키텍처를 미세 조정합니다.
• 캐스케이드 평가 (Cascade Evaluation):
  1. 컴파일: 문법 오류 확인.
  2. 실행 및 검증: 수치적 정확성 확인.
  3. 벤치마킹: 성능 측정 (지연 시간, 처리량).
• 피드백 루프: 평가 결과와 실패 사례를 데이터베이스에 저장하고, 이를 다음 세대의 돌연변이 맥락으로 활용하여 학습을 가속화합니다.

3. 주요 기여 (Key Contributions)

1. CUCo 프레임워크: 컴퓨팅과 통신을 통합하여 최적화하는 최초의 에이전트 기반 자동화 프레임워크를 제안했습니다.
2. 학습 불필요 (No Training): 특정 도메인 데이터로 LLM 을 미세 조정 (Fine-tuning) 할 필요 없이, 구조화된 맥락과 에이전트 워크플로우만으로 고성능 커널을 생성합니다.
3. 정확성과 성능의 균형: '정확성 우선'의 패스트-패스와 '성능 최적화'의 슬로우-패스를 분리함으로써, 복잡한 분산 커널의 개발 난이도를 획기적으로 낮췄습니다.
4. 오픈소스 및 재현성: 18K 줄의 Python 코드로 구현되었으며, 연구 커뮤니티의 발전을 위해 오픈소스로 공개될 예정입니다.

4. 실험 결과 (Results)

CUCo 는 4 가지 대표적인 멀티 GPU 워크로드 (Flash Attention with Context Parallelism, DeepSeek-V3 MoE, KV Cache Transfer, GEMM + AllGather) 에서 평가되었습니다.

• 성능 향상: 호스트 기반 NCCL 베이스라인 대비 **최대 1.57 배 (약 36.2% 지연 시간 감소)**의 속도 향상을 달성했습니다.
• 구체적 개선 사례 (Flash Attention):
  • 파이프라인 버블 제거: 통신과 계산 간의 유휴 시간을 제거하여 37.7ms 절감.
  • 호스트 API 오버헤드 제거: 768 개의 NCCL API 호출 그룹을 디바이스 기반 put/flush 로 대체하여 37.6ms 절감.
  • 타일 단위 파이프라이닝: 계산 블록이 통신이 완료될 때까지 기다리지 않고 즉시 다음 타일을 처리하도록 하여 63.7ms 절감.
• 에이전트 효율성:
  • 패스트-패스 에이전트가 없으면, 진화 에이전트가 올바른 디바이스 통신 모델을 발견하는 데 실패하거나 많은 세대를 소모하여 성능이 낮아졌습니다 (최대 2.5% 성능 저하).
  • '탐색 - 활용 (Explore-Exploit)' 2 단계 전략이 단일 단계 전략보다 더 빠르고 높은 성능을 달성했습니다.

5. 의의 및 결론 (Significance)

CUCo 는 분산 머신러닝 시스템의 패러다임 전환을 보여줍니다.

• 수동 엔지니어링의 대체: 컴퓨팅과 통신을 함께 설계하는 것은 인간에게 너무 복잡하여 수동으로 최적화하기 어렵습니다. CUCo 는 이러한 복잡성을 자동화하여 최적의 커널을 찾아냅니다.
• 디바이스 기반 통신의 실용화: NVSHMEM 및 NCCL Device API 와 같은 최신 기술의 잠재력을 실제로 활용하여, 네트워크 지연을 계산 작업 뒤에 숨기는 (Hiding latency) 것을 가능하게 합니다.
• 미래 지향성: 하드웨어 토폴로지와 워크로드에 따라 적응적으로 최적화되는 에이전트 기반 접근법은 차세대 분산 ML 시스템의 핵심 기술로 자리 잡을 것으로 기대됩니다.

요약하자면, CUCo 는 에이전트 기반의 자동화를 통해 **컴퓨팅과 통신의 공동 설계 (Co-design)**를 실현함으로써, 대규모 AI 모델 학습 및 추론의 지연 시간을 획기적으로 단축하고 GPU 활용 효율을 극대화하는 획기적인 솔루션입니다.