IceCache: 긴 시퀀스 LLM을 위한 메모리 효율적인 KV 캐시 관리

LLM이 처리하는 문맥이 길어질수록 KV 캐시가 차지하는 메모리가 기하급수적으로 늘어나는 문제를 해결합니다. 의미론적 토큰 클러스터링을 통해 필요한 정보만 CPU에서 GPU로 빠르게 불러옴으로써, 저사양 하드웨어에서도 긴 문서 요약이나 복잡한 추론을 가능하게 합니다.

Transformer 모델의 Self-Attention은 이전의 모든 토큰 정보를 KV 캐시에 저장해야 하므로 시퀀스가 길어질수록 메모리 사용량이 선형적으로 증가한다. 기존 방식은 메모리 부족을 해결하기 위해 덜 중요한 토큰을 삭제하거나 단순히 생성 순서대로 CPU에 옮겼는데, 이는 나중에 중요한 토큰이 다시 필요할 때 정확도가 떨어지거나 전송 효율이 낮아지는 한계가 있었다.

IceCache는 토큰의 Embedding 벡터가 가진 의미적 유사성에 주목한다. 비슷한 의미를 가진 토큰들은 특정 쿼리에 대해 동시에 중요할 가능성이 높다는 점을 이용해, 이들을 하나의 메모리 페이지(PagedAttention 단위)로 묶어 관리한다. 이는 도서관에서 관련 서적을 같은 서가에 꽂아두어 한 번에 찾아오기 쉽게 만드는 것과 같은 원리이다.

결과적으로 쿼리가 들어왔을 때 관련 없는 토큰이 섞인 페이지를 여러 번 불러오는 대신, 핵심 토큰들이 밀집된 소수의 페이지 만을 GPU로 로드한다. 이를 통해 전체 KV 캐시의 극히 일부(예: 256개 토큰 분량)만 GPU에 유지하면서도 모델의 원래 성능을 거의 그대로 유지할 수 있게 된다.

IceCache의 워크플로우는 인덱싱, 페이지 선택, 벌크 로딩의 세 단계로 구성된다. 먼저 프리필 단계에서 각 토큰의 Key Embedding을 기반으로 계층적 트리 구조인 DCI-tree를 구축한다. 이때 임베딩 벡터 kj를 [kj/c, sqrt(1-||kj||²/c²)] 형태로 변환하여 내적(Inner Product) 연산을 유클리드 거리 기반의 근사 최근접 이웃(ANN) 검색으로 치환할 수 있도록 한다.

디코딩 단계에서는 현재 쿼리 q에 대해 M-DCI 알고리즘을 실행한다. 상위 레벨 노드에서 시작해 쿼리와 가장 유사한 k개의 토큰이 포함된 하위 페이지들을 독립적인 Attention Head별로 검색한다. [쿼리 벡터 입력 → 트리 검색 연산 → 상위 k개 페이지 식별] 과정을 통해 연산에 필요한 최소한의 데이터셋을 결정한다.

마지막으로 선택된 비연속적인 메모리 페이지들을 벌크 로딩(Bulk Loading) 기법으로 처리한다. CPU 백로드 버퍼에 흩어진 페이지들을 모아 하나의 연속된 덩어리로 만든 뒤 PCIe를 통해 GPU로 한 번에 전송한다. [비연속 페이지 수집 → 연속 버퍼 생성 → PCIe 전송 → GPU 캐시 테이블 배포] 순서로 진행하여 메모리 대역폭 활용도를 극대화한다.

LongBench 벤치마크 실험 결과, IceCache는 256개의 토큰 예산만으로도 전체 KV 캐시를 사용한 모델 정확도의 99%를 유지했다. 특히 Llama-3.1-8B 모델에서 64개의 토큰만 사용했을 때 47.8%의 정확도를 기록하며, 4배 더 많은 예산을 사용한 기존 SOTA 모델인 PQCache(47.3%)를 능가하는 효율성을 보였다.

지연 시간 측면에서도 우수한 성능을 입증했다. 36k 시퀀스 길이에서 두 번째 토큰 생성 시간(TT2T)은 7.7초를 기록했으며, 페이지 재사용(Reuse) 기법을 적용할 경우 5.9초까지 단축되어 OmniKV와 대등한 수준의 속도를 확보했다. 이는 정확도를 희생하지 않으면서도 검색 기반 방식의 고질적인 지연 문제를 해결했음을 의미한다.

Passkey Retrieval 테스트에서는 10k에서 100k에 이르는 모든 문맥 길이와 삽입 위치에 관계없이 100%의 회수율을 달성했다. 또한 RULER 벤치마크를 통해 300k 이상의 초장기 문맥에서도 Full KV 모델 대비 디코딩 지연 시간 증가 폭이 현저히 낮음을 확인했다.

IceCache 아키텍처의 핵심은 PagedAttention의 물리적 페이지 구조와 M-DCI(Multi-level Dynamic Continuous Indexing) 검색 알고리즘의 결합이다. 기존의 Quest나 ArkVale이 토큰 생성 순서대로 페이지를 구성하여 의미적 응집도가 낮았던 것과 달리, IceCache는 DCI-tree를 통해 의미론적으로 유사한 토큰들을 동일한 물리적 페이지에 강제로 할당한다.

수학적으로는 내적 최대화 문제를 ANN 검색으로 변환하기 위해 임베딩 정규화 기법을 사용한다. 모든 Key 벡터를 단위 구면 내부로 투영하고 추가 차원을 부여함으로써, 쿼리와의 내적이 클수록 변환된 공간에서의 유클리드 거리가 가까워지도록 설계했다. 이는 검색 정확도를 보장하면서도 연산 복잡도를 획기적으로 낮추는 기반이 된다.

구현 측면에서는 CPU와 GPU의 하드웨어 특성을 고려한 비동기 파이프라이닝이 돋보인다. GPU가 현재 레이어의 프리필 연산을 수행하는 동안, PCIe를 통해 이미 전송된 이전 레이어의 KV 데이터를 CPU가 인덱싱하도록 설계하여 전체 추론 파이프라인의 병목 현상을 제거했다. 또한 GQA(Group-Query Attention) 구조를 지원하여 여러 헤드가 페이지를 공유할 수 있도록 최적화되었다.

논문은 인덱싱 과정이 CPU 집약적인 작업이므로 CPU 스레드 수가 부족한 환경에서는 지연 시간이 발생할 수 있음을 언급한다. 또한 매우 짧은 시퀀스에서는 오프로딩과 검색에 따른 오버헤드가 이득보다 클 수 있다는 점이 한계로 작용할 수 있다.