OxyGen: Unified KV Cache Management for Vision-Language-Action Models under Multi-Task Parallelism

본 논문은 비전 - 언어 - 행동 모델의 다중 태스크 병렬 실행 시 발생하는 비효율성을 해결하기 위해 KV 캐시를 공유 자원으로 관리하는 'OxyGen'을 제안하여, 기존 대비 최대 3.7 배의 속도 향상과 동시에 고품질 액션 생성을 가능하게 합니다.

핵심

🤖 OxyGen: 로봇의 '멀티태스킹'을 위한 숨겨진 엔진

이 논문은 **로봇이 여러 일을 동시에 할 때, 어떻게 하면 더 빠르고 효율적으로 움직일 수 있을까?**라는 질문에 답합니다.

지금까지의 로봇 AI는 "한 번에 한 가지 일만 제대로" 하도록 설계되어 있었습니다. 예를 들어, 로봇이 컵을 집어 올리는 동작을 하는 동안, 동시에 사용자에게 "지금 컵을 들고 있어요"라고 말하거나 주변 환경을 기억해야 한다면, 기존 시스템은 이 일을 순서대로 처리하거나 별도의 뇌를 따로 써야 했습니다. 그 결과 로봇은 느려지고, 배터리도 빨리 닳았죠.

이 논문은 OxyGen이라는 새로운 시스템을 제안하며, "로봇의 기억 (KV Cache) 을 공유해서 모든 일을 동시에 빠르게 처리하자"고 말합니다.

---

🧠 비유로 이해하는 OxyGen

이 시스템을 이해하기 위해 **요리사 (로봇)**와 **메모장 (KV Cache)**을 예로 들어볼게요.

1. 기존 방식 (고립된 실행): "혼자서 모든 걸 다시 하는 요리사"

기존 로봇 AI 는 다음과 같이 행동합니다.

• 상황: 요리사가 손님이 "파스타를 만들어줘" (언어 작업) 라고 하고, 동시에 "소스를 저어줘" (동작 작업) 라고 합니다.
• 기존 방식: 요리사는 먼저 파스타 레시피를 다시 읽고 메모장에 적고, 소스를 저으라는 지시를 다시 읽고 메모장에 적습니다.
• 문제점: 같은 재료를 보고 같은 메모를 두 번이나 하는 셈입니다. (중복 계산) 게다가 요리사가 메모장을 쓰느라 소스를 저을 시간이 늦어지거나, 반대로 소스를 저으느라 메모를 못 쓰는 경쟁이 발생합니다.

2. OxyGen 방식 (통합된 KV 캐시 관리): "공유 메모장을 쓰는 팀"

OxyGen 은 이 문제를 해결하기 위해 공유 메모장 (Unified KV Cache) 개념을 도입합니다.

• 상황: 같은 재료를 보고 같은 상황을 인식했습니다.
• OxyGen 방식:
  1. 한 번만 읽기 (Cross-task KV Sharing): 요리사는 재료를 한 번만 보고 메모장에 핵심 정보를 적습니다. 이 메모장은 '파스타 만들기'와 '소스 저어주기' 두 작업이 함께 공유합니다.
  2. 동시 작업 (Cross-frame Continuous Batching): 파스타를 만드는 동안, 요리사는 동시에 다른 손님의 주문 (이전 프레임의 언어 작업) 도 처리합니다. 마치 여러 주문을 한 번에 묶어서 (Batching) 효율적으로 처리하듯, 로봇은 여러 작업을 끊기지 않고 연속적으로 처리합니다.

---

🚀 OxyGen 이 가져온 놀라운 변화

이 논문의 실험 결과 (π0.5 모델 사용) 는 다음과 같은 성과를 보여줍니다.

1. 속도 3.7 배 증가: 로봇이 컵을 집는 속도 (Action Frequency) 가 초당 19 회에서 70 회로 빨라졌습니다. 이는 로봇이 훨씬 더 부드럽고 빠르게 움직일 수 있음을 의미합니다.
2. 동시성 확보: 로봇이 동시에 "이제 컵을 들었어요"라고 말하며 (언어 생성), 손으로 컵을 집는 동작을 멈추지 않고 수행할 수 있게 되었습니다.
3. 에너지 절약: 불필요한 계산을 줄여 전기를 최대 78% 까지 절약했습니다. (배터리가 오래 가는 로봇!)

💡 핵심 요약

• 문제: 로봇이 여러 일을 동시에 할 때, 같은 정보를 반복해서 계산하고 자원을 다투느라 느려졌습니다.
• 해결책 (OxyGen): **공유 메모장 (Unified KV Cache)**을 만들어, 같은 정보는 한 번만 읽고 모든 작업이 이를 공유하게 했습니다.
• 결과: 로봇은 더 빠르고, 더 부드럽게, 더 적은 전기로 여러 일을 동시에 처리할 수 있게 되었습니다.

마치 한 명의 요리사가 여러 개의 주문을 동시에 처리하되, 재료를 한 번만 다듬고 모든 요리에 활용하는 것처럼, OxyGen 은 로봇 AI 의 효율성을 혁신적으로 높여줍니다. 이제 로봇은 인간처럼 "말하면서 일하고, 일하면서 기억하는" 자연스러운 존재가 될 수 있는 발판을 마련했습니다.

기술 요약

1. 문제 정의 (Problem)

최근 embodied AI(로봇 등) 는 단일 입력 (관측치) 을 바탕으로 동시에 여러 작업을 수행해야 하는 요구가 증가하고 있습니다. 예를 들어, 물체를 조작 (Action) 하면서 사용자와 대화 (Language) 하거나 환경을 기억 (Memory) 하는 등의 다중 작업 병렬 실행 (Multi-Task Parallelism) 이 필요합니다.

• 현재의 한계: 최근 등장한 Mixture-of-Transformers (MoT) 기반의 비전 - 언어 - 행동 (VLA) 모델 (예: $\pi0.5$) 은 하나의 모델로 이러한 이종 출력을 생성할 수 있는 아키텍처를 제공합니다. 그러나 기존 추론 시스템은 각 작업을 독립적으로 처리하는 '분리된 실행 (Isolated Execution)' 방식을 사용합니다.
• 근본 원인:
  1. 중복 계산 (Redundant Computation): 동일한 관측치 (Observation) 를 기반으로 여러 작업 (행동 생성, 언어 생성 등) 이 수행될 때, 각 작업마다 VLM 백본을 다시 실행하여 동일한 KV Cache 를 생성합니다.
  2. 자원 경쟁 (Resource Contention): 제한된 하드웨어 (예: 로봇 내 단일 GPU) 에서 각 작업이 KV Cache 를 독점적으로 사용하려 하거나, 서로 다른 시간 제약 (행동은 고정 주파수, 언어는 유연한 주파수) 으로 인해 병목 현상이 발생합니다.
• 결과: 이는 불필요한 지연을 초래하여 로봇 제어 주파수 (Action Frequency) 와 언어 처리량 (Throughput) 을 동시에 저하시킵니다.

2. 방법론 (Methodology)

저자들은 OxyGen을 제안하며, 이는 MoT VLA 추론을 위한 효율적인 시스템입니다. 핵심 아이디어는 KV Cache 를 작업 간, 시간 간에 공유되는 '1 등 자원 (First-class Shared Resource)'으로 통합 관리하는 것입니다.

핵심 최적화 기법

1. 작업 간 KV 공유 (Cross-Task KV Sharing):

   • 동일한 프레임 (관측치) 에서 시작하는 여러 작업 (행동, 언어 등) 에 대해 관측치 인코딩 (Prefill) 을 한 번만 수행합니다.
   • 생성된 단일 KV Cache 를 모든 전문가 (Expert) 가 공유하여 중복 계산을 제거합니다.

2. 프레임 간 연속 배치 (Cross-Frame Continuous Batching):

   • 기존 시스템은 프레임 단위로 작업을 끊어 처리하지만, OxyGen 은 언어 생성과 같은 스트리밍 작업을 프레임 경계 없이 연속적으로 배치합니다.
   • 비동기 처리: 행동 생성은 하드 데드라인 (고정 제어 주기) 을 맞추기 위해 프레임 내에서 완료되지만, 언어 생성은 여러 프레임에 걸쳐 배치되어 처리됩니다.
   • 상태 관리: 각 요청의 진행 상태 (생성된 토큰, KV Cache 확장 등) 를 관리하여 프레임이 넘어가도 재계산 없이 추론을 이어나갈 수 있습니다.

시스템 아키텍처

• 통합 KV Cache 관리자 (Unified KV Cache Manager): 각 요청의 상태 ($\sigma_t$) 를 저장하고, 활성 요청들을 배치 (Batch) 하여 병렬 디코딩을 수행합니다.
• 실행 흐름:
  1. 새로운 관측치에 대해 한 번만 Prefill 수행 ($K_t$ 생성).
  2. $K_t$ 를 복사하여 행동 전문가 (Action Expert) 가 즉시 행동 생성.
  3. $K_t$ 를 언어 생성 상태에 저장하고, 이전 프레임에서 이어진 다른 언어 요청들과 함께 배치 디코딩 수행.

3. 주요 기여 (Key Contributions)

1. 문제 공식화: MoT-VLA 에서 다중 작업 병렬 실행을 위한 추론 시나리오를 정의하고, 기존 시스템의 비효율성이 '분리된 KV Cache 관리'에 기인함을 규명했습니다.
2. 새로운 추론 패러다임 제안: KV Cache 를 공유 자원으로 간주하여 작업 간 공유와 프레임 간 연속 배치를 가능하게 하는 통합 관리 방식을 제안했습니다.
3. 실제 구현 및 검증: 가장 인기 있는 MoT VLA 모델인 $\pi0.5$와 오픈소스 프레임워크인 openpi 위에 OxyGen 을 구현했습니다.

4. 실험 결과 (Results)

NVIDIA RTX 4090 GPU 에서 LIBERO, DROID, ALOHA 등 3 가지 벤치마크를 통해 평가했습니다.

• 성능 향상:
  • 기존 분리 실행 (Baseline) 대비 최대 3.7 배의 속도 향상을 달성했습니다.
  • 동시 달성: 언어 처리량 200 토큰/초 이상과 로봇 제어 주파수 70 Hz를 동시에 유지했습니다.
  • 구체적 개선:
    • 짧은 디코딩 단계에서는 작업 간 KV 공유로 인해 1.4 배 향상.
    • 긴 디코딩 단계에서는 프레임 간 배치로 인해 2.6 배 이상 추가 향상 (전체 3.7 배).
• 품질 보장: 행동 생성의 품질 저하 없이, LIBERO 벤치마크에서 기존 모델의 작업 성공률 (Task Success Rate) 을 통계적 오차 범위 내에서 재현했습니다.
• 에너지 효율성:
  • 메모리 오버헤드는 약 15% 증가에 그쳤으나, 에너지 소비는 최대 78% 감소했습니다.
  • 단순 병렬 실행 (Naive Parallel) 은 메모리 사용량을 2 배로 늘리고 에너지를 증가시켰으나, OxyGen 은 이를 효율적으로 해결했습니다.

5. 의의 및 결론 (Significance)

• 온디바이스 배포 가능성: 제한된 하드웨어 자원 (단일 GPU) 을 가진 로봇에서도 복잡한 다중 작업 (조작 + 대화 + 기억) 을 실시간으로 수행할 수 있는 기반을 마련했습니다.
• 아키텍처와 시스템의 시너지: MoT 아키텍처가 가진 잠재력을 실제 추론 시스템 수준에서 최대한 끌어올리는 방법을 제시했습니다.
• 미래 방향: 비동기 데드라인을 가진 이종 작업들을 효율적으로 스케줄링하는 새로운 추론 패러다임을 제시하여, 더 복잡하고 자연스러운 embodied AI 에이전트 개발에 기여합니다.

요약하자면, OxyGen은 로봇이 여러 일을 동시에 할 때 발생하는 계산 중복과 자원 경합을 해결하기 위해 KV Cache 를 통합 관리함으로써, 속도와 효율성을 극대화하면서도 작업 품질을 유지하는 획기적인 추론 시스템입니다.