ProphetKV: 질의 기반 선택적 재계산으로 효율적인 KV 캐시 재사용

초록

ProphetKV는 검색 기반 생성(RAG)에서 프리필 단계의 고비용을 줄이기 위해, 사용자 질의와의 의미적 연관성을 기준으로 토큰을 선택하고, 두 단계의 재계산 파이프라인으로 레이어별 어텐션 정보를 융합한다. 20% 토큰만 재계산해도 전체 프리필 정확도의 96%‑101%를 유지하며, 기존 방법 대비 RULER와 LongBench에서 각각 8.8‑24.9%·18.6‑50.9%의 정확도 향상을 달성한다.

상세 분석

본 논문은 장문 컨텍스트를 필요로 하는 Retrieval‑Augmented Generation(RAG) 시스템에서 프리필 단계가 전체 추론 지연의 주된 병목임을 지적한다. 기존의 KV‑Cache 재사용 기법은 두 가지 한계를 가진다. 첫째, 프리픽스 매칭 기반 방법은 문서 순서가 바뀌면 재사용이 불가능하고, 두 번째는 위치‑독립(Position‑Independent) 방식으로 사전 계산된 KV‑Cache를 단순히 연결할 경우 문서 간 교차 어텐션이 손실돼 정확도가 급격히 떨어진다. 이를 보완하기 위해 EPIC, CacheBlend, KVShare와 같은 부분 재계산(partial recomputation) 기법이 제안되었지만, 이들 모두 “전역 어텐션 복원”을 목표로 하여 전역적으로 중요한 토큰을 기준으로 선택한다. 결과적으로, 질의와 무관한 전역적으로 활발한 토큰이 제한된 재계산 예산을 차지해, 실제 질의 응답에 핵심적인 토큰이 배제되는 ‘crowding‑out effect’를 초래한다.

ProphetKV는 이러한 문제를 두 가지 핵심 아이디어로 해결한다. ① 질의‑구동 토큰 선택: 질의 토큰이 생성 단계에서 집중하는 어텐션 가중치를 역추적해, 질의와 의미적으로 연관된 컨텍스트 토큰을 정량화한다. 이때 어텐션 스코어의 상위 비율을 선택 기준으로 삼아, 질의가 ‘예언자(prophet)’ 역할을 하도록 설계한다. ② 이중‑단계 재계산 파이프라인: 첫 단계에서는 질의‑기반 중요도 점수를 산출하고, 두 번째 단계에서는 레이어별 어텐션 메트릭(예: 각 레이어의 Q‑K 유사도, 헤드별 가중치)을 융합해 최종 재계산 토큰 집합을 결정한다. 레이어 융합 알고리즘은 낮은 레이어에서 잡히지 않는 고차 의미 토큰까지 포괄하도록 설계돼, 기존 방법이 놓치기 쉬운 깊은 레이어의 핵심 토큰을 보존한다.

실험에서는 Llama‑3‑8B‑Inst, Qwen2.5‑14B‑Inst, Qwen3‑14B‑Thor 등 다양한 모델에 ProphetKV를 적용했으며, 재계산 비율을 20%로 제한했음에도 전체 프리필 정확도의 96%‑101%를 유지했다. 특히 RULER와 LongBench 벤치마크에서 기존 최첨단 방법 대비 8.8‑24.9%·18.6‑50.9%의 정확도 향상을 기록했다. 토큰 선택 정밀도는 ‘오버랩 비율(overlap ratio)’ 지표로 평가했을 때, ProphetKV가 질의‑어텐션 토큰과의 일치율이 0.85‑0.95 수준으로 가장 높았다. 이는 질의‑기반 선택이 실제 디코딩 단계에서 중요한 토큰을 효과적으로 포착한다는 것을 의미한다.

또한, ProphetKV는 훈련‑프리(free) 방식으로 구현돼 추가 모델 파인튜닝이 필요 없으며, 기존 KV‑Cache 재사용 파이프라인에 플러그인 형태로 쉽게 통합할 수 있다. 재계산 비용은 선택된 토큰 수에 선형적으로 비례하므로, 실시간 서비스 환경에서도 예측 가능한 지연 감소 효과를 제공한다.

요약하면, ProphetKV는 “전역 어텐션 복원”이라는 비현실적인 목표를 포기하고, 질의‑중심의 토큰 선택과 레이어‑통합 재계산이라는 두 축을 통해 제한된 연산 예산을 가장 가치 있는 교차 어텐션 복원에 집중한다. 이는 장문 RAG 시스템에서 프리필 지연을 크게 낮추면서도, 기존 전면 재계산 수준의 정확도를 유지할 수 있는 실용적인 솔루션으로 평가된다.