활성화 체크포인팅과 하이브리드 캐싱으로 효율적인 LLM 추론

초록

본 논문은 호스트 메모리를 활용한 대형 언어 모델(LLM) 추론에서 GPU-CPU 간 데이터 전송 병목을 완화하기 위해 활성화 체크포인팅 기반의 활성화 캐시와 KV‑활성화 하이브리드 캐시 방식을 제안한다. 활성화 캐시는 중간 레이어의 활성값을 저장해 KV 캐시 재계산 시 투영·FFN 연산을 건너뛰게 하며, KV와 활성화 캐시 비율을 최적화하는 알고리즘을 통해 통신량과 재계산 비용을 균형 있게 조절한다. 구현 결과 FlexGen 대비 2.19배, 순수 활성화 캐시 대비 1.35배의 토큰 처리량 향상을 달성하였다.

상세 분석

본 연구는 대규모 트랜스포머 디코더 기반 LLM이 GPU 메모리 한계에 부딪히는 상황에서, 호스트 메모리와의 PCIe 전송이 GPU 연산을 크게 지연시킨다는 근본적인 문제를 지적한다. 기존의 FlexGen은 가중치와 KV 캐시를 호스트에 저장하고 필요 시 스트리밍하지만, KV 캐시가 토큰 수에 비례해 급증하면서 전송량이 GPU 활용도를 7% 수준으로 떨어뜨린다. 이를 해결하기 위해 저자는 두 가지 핵심 아이디어를 도입한다. 첫째, 활성화 체크포인팅을 활용한 활성화 캐시(Activation Cache)이다. 각 레이어의 중간 활성값(A_i)을 호스트에 저장하고, KV 캐시 재구성 시 이 활성값을 이용해 Q·K·V를 바로 생성한다. 이렇게 하면 토큰 ID 기반의 전체 재계산에 비해 투영(projection)과 피드포워드 네트워크(FFN) 연산을 생략할 수 있어 연산량이 절반 이하로 감소한다. 둘째, KV‑Activation 하이브리드 캐시 메커니즘이다. 호스트 메모리에는 KV 캐시와 활성화 캐시를 혼합 저장하고, GPU 메모리에는 오직 활성화 캐시만 유지한다. 시스템은 주어진 GPU 메모리 용량과 목표 처리량을 고려해 KV와 활성화 캐시의 비율을 최적화하는 탐색 알고리즘을 제공한다. 이 비율은 배치 크기, 시퀀스 길이, 모델 깊이 등에 따라 동적으로 조정되며, 최적화 과정은 전송 지연과 재계산 지연을 최소화하도록 설계되었다. 구현은 vLLM 0.4.3을 기반으로 하며, FlexGen과 동일한 호스트‑GPU 오프로드 파이프라인에 위의 두 메커니즘을 삽입한다. 실험은 OPT‑30B, OPT‑66B 등 네 가지 모델에 대해 수행했으며, KV‑Activation 하이브리드 캐시가 KV 전용 오프로드 대비 2.19배, 활성화 캐시 전용 대비 1.35배의 토큰/초 향상을 보였다. 특히 배치 크기가 128 이상으로 증가할 때 발생하던 PCIe 전송 포화 현상이 크게 완화되어, GPU 활용도가 7%에서 15% 이상으로 상승한다. 이 결과는 호스트 메모리를 활용하면서도 정확도를 손상시키지 않고, 비용 효율적인 단일 GPU LLM 서비스가 가능함을 입증한다.