[arXiv’ 24] Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache 논문 리뷰
Paper: Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Arxiv, 2024)
대규모 언어 모델(LLM) 추론 서비스에서는 요청마다 지원하는 최대 입력 문맥 길이가 크게 달라진다.
예를 들어,
- OpenAI의 ChatGPT는 128K 토큰,
- 구글 Gemini는 1000K 토큰,
- LongRoPE 연구는 2000K 토큰까지 지원한다.
그러나 LLM은 AR(Autoregressive) 방식으로 순차적으로 토큰을 생성하기 때문에, 응답이 완료되기 전까지 필요한 메모리와 연산 자원을 정확히 예측하기 어렵다. 특히, 긴 문맥의 경우 KV-Cache(Key-Value 캐시)에 저장해야 하는 중간 생성 벡터 수가 기하급수적으로 증가한다.
KV Cache란? LLM의 생성단계중 Key, Value를 따로 저장해 두었다가 새로운 토큰 생성을 생성하기 위한 추론에 재사용하는 캐시 메모리
예를 들어, LLaMA-7B 모델에서는
- 1천 토큰 처리 시 약 15GB 메모리가 필요하지만,
- 100만 토큰 처리 시 500GB 이상의 메모리가 요구되어 단일 GPU 메모리 한계를 훨씬 초과하게 된다.
반면 짧은 문맥은 낮은 메모리 수요와 동시에 배치 처리로 GPU 연산 자원을 최대한 활용해야 하므로, 고정된 자원 할당 방식은 두 상황을 모두 만족시키기 어렵다.
1. 문제 정의
- 문맥 길이의 다양성:
요청에 따라 적은(수십-수백) 토큰에서 최대 2000K 토큰까지 극단적인 차이가 있음. 이런 경우 토큰이 점점 많아지면서 Attention 연산의 경우 필요로 하는 메모리 양이 Non-Attention 연산 대비해서 많이 증가하는 것을 관찰할 수 있음.
- 자원 불균형:
긴 문맥은 KV Cache 저장으로 인해 메모리 집약적이며, 짧은 문맥은 배치 처리를 통한 연산 집약적 특성을 가짐.
- 정적 자원 할당의 한계:
한 인스턴스(요청 데이터를 병렬로 처리하기 위해 서버에 배치된 모델 복제본)에 긴 요청이 몰리면 GPU 메모리가 포화되어 배치 크기를 줄여야 하고, 짧은 요청만 있는 경우 남는 메모리를 다른 인스턴스와 공유하지 못해 전체 클러스터 활용률이 낮아짐.
2. Infinite-LLM의 핵심 아이디어
Infinite-LLM은 LLM의 어텐션 연산과 KV Cache를 기존 모델 내부에서 분리하여, 클러스터 전체 GPU 메모리를 하나의 풀(pool)처럼 활용하는 시스템이다.
Infinite-LLM은 다음과 같은 핵심 아이디어들로 문제를 해결했다.
- 어텐션 연산을 나눠서 수행 (DistAttention) → 어텐션 계산을 작은 서브블록 단위로 쪼개 다수 GPU에 동적으로 분산 처리하여, 긴 문맥 처리 시 통신 부하를 줄임.
- 분산 KV Cache 관리 → 한 요청의 KV Cache를 단일 인스턴스가 아닌, 클러스터 내 여러 GPU에 분산 저장하고 필요에 따라 이동(migrate)시킴.
- 동적 메모리 대여/회수 (Debtor와 Creditor) → 메모리 부족 인스턴스(Debtors)는 다른 인스턴스(Creditors)로부터 KV Cache 블록을 빌려 처리량을 유지하며, 작업 완료 또는 메모리 여유 시 이를 회수.
- 중앙 및 지역 자원 관리자 (gManager와 rManager) → 전체 클러스터 상태와 KVCache 분포를 실시간 모니터링하여, KVCache 배치 및 이동을 동적으로 조율하는 관리 체계를 도입.
3. 핵심 기술(Idea)설명
DistAttention: 분산 어텐션 메커니즘
- 기존 문제점:
- 전통적인 트랜스포머 어텐션은 모든 쿼리(Q)와 키밸류(KV) 데이터에 대해 소프트맥스(Softmax) 연산을 수행해야 하므로, 긴 시퀀스의 경우 다른 인스턴스에서 계산된 KV값을 전송해야 하므로, GB~TB 단위의 데이터 전송이 필요해 병목이 발생함.
- 서브블록 분할: 긴 시퀀스를 여러 조각(chunks)으로 나누어, 각 GPU가 해당 블록에 대해 부분 어텐션(마이크로 어텐션)을 독립 계산.
- 어텐션 중간 정보 교환: 각 조각에서 로컬 최대값($m_j$)과 합계($e_j$), 그리고 부분 결과 벡터($MA_j$)를 산출한 후, 주 인스턴스가 이 요약 정보를 합산해 전체 소프트맥스 정규화를 수행함.
- 통신 오버헤드 감소: 쿼리 벡터와 최소한의 정보만 교환하므로, 기존 방식보다 수십 배 빠른 통신 성능을 구현.
- 정확도 유지: 수학적으로 기존 어텐션과 동등한 결과를 보장하므로 모델 정확도가 손상되지 않음.
각 인스턴스에서 계산된 어텐션 결과가 어떻게 합쳤을 때 값이 정확하게 보정되는지에 대한 자세한 내용은 Online Softmax에 대한 논문을 참조해주세요.
🔗link: (arXiv.org) Online normalizer calculation for softmax
분산 KVCache 관리: 클러스터 메모리 풀 활용
기존 문제점: 한 인스턴스에 요청의 모든 KV Cache를 저장하면, 문맥 길이가 극도로 길어질 때 메모리 부족 문제가 발생하거나, 오래된 KV 캐시를 삭제해야 하는 문제가 있음.
- 분산 저장: KV Cache 블록을 여러 GPU 인스턴스에 균등하게 분산시켜 단일 GPU의 메모리 한계를 극복.
- 동적 이동: 만약 한 인스턴스의 긴 대화 문맥으로 인해 KV Cache가 과도하게 늘어나면, 클러스터 내 여유 메모리를 가진 다른 인스턴스(Creditors)로부터 일부 KV Cache 블록을 빌려 저장함.
- 오버랩 최적화: KV Cache 이동을 모델 계산과 겹쳐 수행하여, 예를 들어 32토큰 분량 이동 시 처리량 저하가 8.6% 나타나지만 16토큰 단위로 이동하면 오버헤드를 숨김.
- 정확성 보장: DistAttention의 설계 덕분에, KV Cache가 분산되어 있더라도 최종 어텐션 결과는 기존과 동일함.
동적 메모리 대여 및 회수 (Debtor와 Creditor)
- 개념 및 필요성:
- Debtor (채무자): 메모리 부족 상황에 처한 인스턴스는, 긴 문맥 처리 시 배치 크기 유지와 토큰 생성 속도 향상을 위해 다른 인스턴스(채권자)로부터 KV Cache 블록을 빌림.
- Creditor (채권자): 여유 메모리를 보유한 인스턴스는, 필요 시 자신의 자원을 채무자에게 제공하되, 자신의 작업 부하가 증가하면 즉시 회수할 수 있음.
- 얼마만큼의 KV Cache를 대여할지에 대한 수학적 모델링:
- 한 인스턴스는 동시에 채무자와 채권자 역할을 겸하지 않으며, 채무자는 요청 완료나 메모리 여유 발생 시 빌린 KV Cache를 회수. (모델링을 통해 최적 퍼포먼스를 예측 후 최적값 Offloading)
- 이 메커니즘은 마치 원장(ledger)처럼 메모리 대여 관계를 관리하며, 실험 결과 약 200개의 KV Cache 블록 이동이 전체 처리량 향상에 유의미함을 보임.
글로벌 및 로컬 자원 관리자 (gManager와 rManager)
- gManager (글로벌 관리자):
- 클러스터 내 모든 GPU 인스턴스의 메모리 사용량, 요청 진행 상황, 그리고 KV Cache 분포(전역 배치)를 추적함.
- 배치 크기(β)와 메모리 사용률(𝑈) 등의 기준을 바탕으로, 채무자와 채권자 인스턴스를 선정하고, 그리디 알고리즘을 통해 최적의 KV Cache 이동 계획을 수립함.
- rManager (지역 관리자):
- 각 인스턴스에 상주하면서 로컬 메모리 사용량과 요청 상태를 모니터링하며, 주기적으로 gManager에 상태를 보고함.
- gManager의 명령에 따라 KVCache 블록 이동 등 로컬 작업을 수행하며, 대상 인스턴스의 상태 변화에 따라 유연하게 대응.
- 운영 프로토콜:
- FCFS(First-Come-First-Serve) 방식을 사용하여 공간 할당을 관리하며, 만약 대상 인스턴스의 메모리 여유가 부족하면 해당 작업을 건너뛰고 다음 주기 상태를 반영하여 계획을 재수립함.
4. 성능 개선 및 실험 결과
- 실험 환경: 32개의 A100 GPU로 구성된 클러스터.
- Context Length 범위: ~2000K 토큰에 이르는 초장문 Context Length 까지 평가.
- 주요 결과:
- 최신 LLM 서비스 시스템(vLLM 등) 대비 엔드투엔드 처리 성능이 1.35배에서 최대 3.4배 향상됨.
- 긴 문맥 요청 처리 시, 단일 인스턴스의 메모리 한계를 극복하고 클러스터 전체 자원을 활용함으로써 1.4배~3.4배 높은 처리량을 달성.
- 짧은 요청의 경우에도 여러 인스턴스 간 남는 메모리를 모아 배치 크기를 늘려 1.3배 이상의 개선을 보임.
- DistAttention은 LLaMA2-13B 모델의 4GPU 분산 실험에서, 기존 헤드 분할 방식 대비 최대 25% 빠르고, RingAttention 대비 7.7배~19.8배 빠른 성능을 입증함.
- KV Cache 이동의 오버헤드 또한 소규모 블록 단위 이동과 통신-연산 오버랩 최적화를 통해 거의 무시할 수 있는 수준으로 감소됨.
5. 결론
Infinite-LLM은 기존 정적 자원 할당 방식의 한계를 극복하고,
- 어텐션 연산과 KV Cache 관리의 분산 처리를 통해 클러스터 전체 GPU 메모리를 하나의 거대한 풀처럼 활용
- 동적 메모리 대여/회수(Debtor와 Creditor) 및 중앙-분산 관리(gManager와 rManager) 체계를 도입해 자원 불균형 문제를 해결
- 긴 문맥 처리에서도 높은 처리량과 낮은 지연시간을 유지할 수 있었음.