MatKV 플래시 저장소로 연산을 교환한 LLM 추론

초록

MatKV는 RAG 작업에서 문서의 키‑밸류(KV) 벡터를 사전에 계산해 고속 플래시 SSD에 저장하고, 추론 시 GPU 대신 이를 불러와 사용함으로써 KV 계산 비용을 절반 이하로 줄이고 전력 소모와 지연 시간을 크게 감소시킨다.

상세 요약

본 논문은 대규모 언어 모델(LLM) 기반 생성 AI에서 최근 두드러지는 현상, 즉 추론 비용이 학습 비용을 앞지르고 RAG(Retrieval Augmented Generation) 방식이 보편화되는 추세를 출발점으로 삼는다. RAG에서는 외부 문서들을 검색해 프리픽스 단계에서 입력 텍스트와 함께 키‑밸류(KV) 쌍을 생성해야 하는데, 이 과정은 트랜스포머의 self‑attention 연산 특성상 GPU 메모리 대역폭과 연산량을 크게 요구한다. 특히 긴 입력을 다루는 경우, 프리픽스 단계가 전체 추론 파이프라인의 병목이 되며 전력 소모도 급증한다.

MatKV는 이러한 병목을 “연산을 저장소로 교환”한다는 관점에서 해결한다. 구체적으로, RAG에 사용되는 문서 집합(예: 위키피디아 문단, 기업 데이터베이스 등)을 사전에 토크나이즈하고, 해당 토큰 시퀀스에 대해 트랜스포머 레이어별 KV 벡터를 계산한다. 계산된 KV는 압축 및 정렬 과정을 거쳐 NVMe 기반 플래시 SSD에 순차적으로 기록된다. 이때 SSD는 비용 대비 높은 읽기 대역폭과 낮은 전력 소비를 제공하므로, GPU가 KV를 재계산하는 것보다 훨씬 효율적이다.

추론 시에는 질문을 토크나이즈하고, 검색 엔진을 통해 관련 문서를 식별한다. 식별된 문서에 대응하는 KV 블록을 SSD에서 직접 스트리밍하여 GPU 메모리로 전송한다. GPU는 이 KV를 바로 사용해 디코딩(생성) 단계만 수행한다. 논문은 두 가지 추가 최적화를 제시한다. 첫째, GPU가 현재 토큰을 디코딩하는 동안 SSD에서 다음 인스턴스의 KV를 비동기적으로 로드함으로써 I/O 대기 시간을 최소화한다. 둘째, 디코딩 연산은 KV 계산에 비해 GPU 성능에 덜 민감하므로, 저가형 GPU(예: RTX 3060)라도 KV가 미리 로드된 상황에서는 전체 지연 시간에 큰 영향을 주지 않는다.

실험은 Hugging Face Transformers 라이브러리를 기반으로 LLaMA‑2‑7B, Falcon‑40B 등 최신 모델을 사용했으며, NVIDIA A100, RTX 4090 등 고성능 GPU와 Samsung PM983 NVMe SSD를 조합했다. 결과는 KV 사전 계산·저장 방식이 전체 추론 시간과 전력 소비를 평균 48 %·52 % 감소시켰으며, 질문‑응답 정확도는 0.3 % 이하의 미세한 감소만을 보였다. 이는 KV 재사용이 모델 출력 품질에 큰 영향을 미치지 않음을 입증한다.

또한, 메모리 사용량 측면에서 사전 저장된 KV는 압축률 2.3×을 달성해 GPU VRAM 부담을 크게 경감시켰다. 이는 더 큰 컨텍스트 길이를 다루거나 멀티‑모델 파이프라인을 구성할 때 유용하다. 논문은 향후 KV 압축 알고리즘 고도화, SSD 캐시 계층 설계, 그리고 분산 환경에서 KV 공유 메커니즘을 연구할 방향을 제시한다.

전반적으로 MatKV는 연산‑저장소 트레이드오프를 통해 LLM 추론 비용을 구조적으로 낮추는 실용적 솔루션이며, 특히 클라우드·엣지 환경에서 비용·전력 효율성을 동시에 개선하고자 하는 기업·연구기관에 큰 의미를 가진다.