SAGA: Workflow-Atomic Scheduling for AI Agent Inference on GPU Clusters

"GPU 클러스터에서의 AI 에이전트 추론을 위한 워크플로우-원자적 스케줄링 SAGA" 논문에 대한 상세한 기술 요약입니다.

1. 문제 제기

현재 GPU 클러스터 스케줄러 (예: vLLM, SGLang) 는 단일 회차 추론에 최적화되어 있어, 각 대규모 언어 모델 (LLM) 요청을 독립적인 단위로 취급합니다. 그러나 현대의 AI 에이전트는 반복적인 "사고 - 행동 - 관찰" 루프를 통해 작동하며, 각 작업당 수십에서 수백 개의 체인화된 LLM 호출을 외부 도구 호출 (코드 실행, 웹 브라우징, 데이터베이스 쿼리 등) 과 교차하여 실행합니다.

이는 기존 스케줄링 시스템과 근본적인 불일치를 초래하여 다음과 같은 세 가지 치명적인 비효율성을 야기합니다:

1. KV 캐시 재생성 오버헤드: 기존 시스템은 도구 호출 중 발생하는 유휴 시간 (수 밀리초에서 수 분까지 지속 가능) 동안 키 - 값 (KV) 캐시를 제거합니다. 에이전트가 resumed 될 때 캐시를 재생성해야 하므로, 실행 시간의 38% 가 낭비되고 지연 시간이 3~8 배 증가합니다.
2. 분편화된 메모리 활용: 독립적인 스케줄링으로 인해 메모리 지역성이 저하됩니다. 캐시가 분편화되고 조기 제거되므로 잠재적인 높은 재사용 가능성에도 불구하고 GPU 메모리 활용도는 평균 약 42% 에 불과합니다.
3. 작업 수준 공정성 부재: 기존 스케줄러는 요청당 지연 시간이나 처리량을 최적화합니다. 다중 테넌트 환경에서 이는 짧은 요청이 자원을 독점함으로써 길고 복잡한 워크플로우를 가진 에이전트의 '기아 (starvation)'를 초래합니다.

2. 방법론: SAGA 아키텍처

SAGA 는 요청 수준에서 프로그램 수준 (워크플로우-원자적) 스케줄링으로의 패러다임 전환을 제안합니다. 이는 개별 요청의 시퀀스가 아닌 전체 에이전트 워크플로우를 단일 스케줄링 단위로 취급합니다. 시스템은 세 가지 주요 레이어로 구성됩니다:

A. 에이전트 실행 그래프 (AEGs)

SAGA 는 에이전트 워크플로우를 노드가 LLM 추론 단계를, 엣지가 (도구 호출 포함) 의존성을 나타내는 방향 그래프로 모델링합니다.

• 구성: AEG 는 프레임워크 힌트 (예: LangChain 콜백) 에서 파생되거나, 힌트가 없는 경우 패턴 매칭을 통해 추론됩니다.
• 활용: 이 그래프를 통해 스케줄러는 미래 단계를 예측하고, 전이 확률을 계산하며, 도구 호출 경계 간 캐시 재사용 가능성을 예측할 수 있습니다.

B. 세 가지 핵심 메커니즘

1. 워크플로우 인지 KV 캐시 관리 (WA-LRU):

   • 표준 LRU 대신 SAGA 는 워크플로우 인지 LRU (WA-LRU) 제거 정책을 사용합니다.
   • 제거 점수는 다음을 기반으로 계산됩니다: (1) 최근성, (2) 예측 재사용 확률 (AEG 에서 유도됨), (3) 캐시 크기.
   • 도구 호출 인지 TTL: 시스템은 호출되는 특정 도구와 현재 메모리 압력에 따라 캐시된 KV 상태의 생존 시간 (TTL) 을 동적으로 조정합니다. 이를 통해 긴 도구 대기 시간 동안의 조기 제거를 방지합니다.
   • 예측적 프리페칭: 다음 단계가 예측 가능할 경우, SAGA 는 도구 실행 단계 중 필요한 KV 캐시를 프리페칭하여 I/O 와 연산을 중첩합니다.

2. 세션 애너티 배치 및 작업 도용:

   • 애너티 라우팅: 동일한 에이전트 세션에 속한 요청은 캐시 히트를 극대화하기 위해 동일한 GPU 워커로 라우팅됩니다.
   • 작업 도용: 부하 불균형 (지체자) 을 방지하기 위해 글로벌 조정자가 무작위 작업 도용을 수행합니다. 다른 워커가 과부하 상태인 동안 한 워커가 유휴 상태라면 작업을 도용합니다. 핵심적으로, 시스템은 국소성을 유지하기 위해 KV 캐시 상태를 이전합니다 (Llumnix 와 같은 메커니즘 사용). 이를 통해 도용된 작업이 캐시된 컨텍스트를 잃지 않도록 보장합니다.

3. 에이전트 공정 공유 (AFS):

   • 요청 수 대신 작업 완료 시간을 기반으로 한 새로운 공정성 지표입니다.
   • 라이아푸노프 드리프트 분석을 사용하여 긴급하고 장기 실행되는 작업을 가진 테넌트가 비례하는 자원을 받도록 보장하여 기아를 방지합니다. 시스템은 각 에이전트 작업의 "긴급성" (잔여 작업 대 마감기한) 에 따라 스케줄링 우선순위를 동적으로 조정합니다.

3. 주요 기여

• 워크플로우-원자적 추상화: 전체 에이전트 프로그램을 1 급 스케줄링 단위로 취급하고 워크플로우 구조 (DAG) 를 스케줄러에 명시적으로 노출하는 최초의 시스템입니다.
• 이론적 경계:
  • 생산 트래픽에서 WA-LRU가 벨라디의 최적 오프라인 정책 (이론적으로 가능한 최상의 캐시 제거) 대비 1.31 배 이내의 경쟁 비율을 달성함을 증명했습니다.
  • **에이전트 공정 공유 (AFS)**에 대해 라이아푸노프 드리프트를 사용한 공식적 증명을 제공하여, 제한된 수요 이질성 하에서 작업 완료 시간의 유계된 편차를 보장합니다.
• 통합 시스템 설계: 분산 스케줄링, 예측적 캐싱, 공정성 보장을 통합된 프레임워크로 결합하여 기존 작업 (예: vLLM 의 접두어 캐싱, KVFlow 의 제거, Llumnix 의 이전) 이 남긴 간극을 해소했습니다.

4. 실험 결과

64-GPU 클러스터(8 노드 × 8 NVIDIA A100) 에서 SWE-bench(코딩 에이전트) 및 WebArena(브라우저 에이전트) 워크로드를 사용하여 평가했습니다.

• 지연 시간 감소:
  • SAGA 는 최첨단 vLLM v0.15.1(자동 접두어 캐싱, APC 포함) 대비 **작업 완료 시간 (TCT)**을 1.64 배(기하 평균) 감소시켰습니다.
  • 일반 vLLM 대비 개선 폭은 3.01 배입니다.
  • 구체적으로 SWE-bench 에서 TCT 는 vLLM 의 약 612 초에서 SAGA 의 203 초로 감소했습니다.
• 메모리 효율성:
  • GPU 메모리 활용도는 vLLM 의 약 42% 에서 SAGA 의 **71.3%**로 증가했습니다.
  • KV 캐시 재생성 시간이 총 시간의 38% 에서 **8%**로 감소했습니다.
• 공정성:
  • 다중 테넌트 간섭 하에서 99.2% SLO 달성률(마감기한을 준수한 작업) 을 기록했으며, vLLM 의 67.3% 대비 우위를 보였습니다.
  • 경량 테넌트의 '긴 꼬리' 지연을 제거하여 무거운 워크로드와 가벼운 워크로드 전반에 걸쳐 일관된 성능을 보장했습니다.
• 트레이드오프:
  • SAGA 는 지연 시간과 작업 완료를 최적화하므로, 처리량 최적화 배치 스케줄링 대비 최대 처리량이 약 30% 감소합니다. 저자들은 지연 시간에 민감한 상호작용형 에이전트 배포의 경우 이는 수용 가능한 트레이드오프라고 주장합니다.

5. 의의

• 격차 해소: SAGA 는 복합 AI 시스템의 비효율성이 모델 자체의 문제가 아니라 요청 수준 추상화의 산물임을 입증했습니다. 워크플로우 구조를 노출함으로써 스케줄러는 오프라인 최적 효율에 근접할 수 있습니다.
• 이론의 경험적 검증: 이 논문은 에이전트 워크플로우의 온라인 캐시 관리에 대한 최초의 경험적 경쟁 비율 경계 (1.31 배) 를 제공하여, 워크플로우 인식이 최적의 오프라인 정책을 거의 달성하기에 충분함을 검증했습니다.
• 생산 준비성: 이 시스템은 vLLM 의 확장 기능으로 구현되었으며, 에이전트 자체의 코드 변경 없이 주요 에이전트 프레임워크 (LangChain, AutoGen) 를 지원합니다.
• 미래 방향: 이는 "복합 AI 시스템"에 대한 새로운 기준을 설정하며, 향후 인프라가 정적 요청 제공에서 동적이고 상태가 있는 워크플로우 오케스트레이션으로 진화해야 함을 시사합니다.

요약하자면, SAGA는 개별 요청에서 전체 워크플로우로 스케줄링 단위를 전환하고, 예측적 캐싱과 공정성 인지 자원 할당을 활용하여 GPU 클러스터에서 거의 최적의 성능을 달성함으로써 AI 에이전트의 치명적인 지연 시간 및 효율성 병목 현상을 해결합니다.