
에이전트에서 팽창하는 KV 캐시를 prefix 캐시를 파괴하지 않고 쳐내는 IntentKV
요약
에이전트 실행 시 발생하는 KV 캐시 팽창 문제를 해결하기 위해, prefix 캐시를 파괴하지 않고 메모리를 관리하는 IntentKV 기술을 소개합니다. 기존의 KV 축소 방식이 가진 위치 정보 왜곡 문제를 '파수꾼 슬롯(Sentinel slot)' 교체 방식으로 해결하여 prefix 캐시 히트율을 획기적으로 높였습니다.
핵심 포인트
- 에이전트의 긴 컨텍스트 사용 시 KV 캐시가 메모리 병목의 주원인이 됨
- 기존 KV 축소 기법은 토큰 재정렬로 인해 prefix 캐시 재사용을 방해함
- IntentKV는 물리적 슬롯을 유지하며 가중치를 0으로 만드는 방식으로 문제 해결
- prefix 캐시 히트율을 기존 0~3%에서 20.7%까지 회복 가능
도구 호출(Tool calling)을 반복하는 에이전트를 길게 실행하다 보면, 가장 먼저 한계에 부딪히는 것은 모델의 연산 그 자체가 아니라 대개 메모리 쪽이다. 검색 도구가 반환한 10만 건 규모의 문서, 통째로 붙여넣은 웹 페이지의 HTML, 과거의 추론 로그. 이것들이 컨텍스트(Context)에 계속 추가되며 다시는 빠지지 않는다. 이렇게 쌓인 이력을 유지하는 것이 KV 캐시 (KV Cache) 이며, 에이전트의 실운용에서는 이곳이 병목(Bottleneck)이 된다.
KV 캐시는 과거의 토큰마다 계산한 Key/Value(주의 집중 메커니즘(Attention mechanism)의 중간 결과)를 저장해 두는 영역이다. 이것 덕분에 새로운 토큰을 생성할 때마다 전체 이력을 재계산하지 않아도 된다. 다만 크기는 토큰 수에 비례하여 팽창한다. 단발성 채팅이라면 문제가 되지 않지만, 수십 턴에 걸쳐 도구 출력을 받아들이는 에이전트에서는 파라미터 계산보다 KV의 읽기/쓰기 대역폭이 먼저 막히게 된다.
6월에 arXiv에 발표된 「IntentKV」는 이 문제에 정면으로 대응하면서, 기존의 KV 축소 기법들이 간과했던 부작용을 짚어내고 있다.
KV가 팽창한다면 오래된 토큰을 솎아내면 된다는 발상의 기법은 이미 여러 개가 있다. 말미의 주의 분포(Attention distribution)에서 중요 토큰을 뽑아내는 SnapKV, 영향력이 큰 '헤비 히터(Heavy hitter)'를 남기는 H2O, 도입부의 싱크(Sink) + 최근 윈도우(Window)만 남기는 StreamingLLM 등이 대표적이다.
문제는 깎아낸 뒤에 KV 행을 다시 채우는(Compaction) 과정에서 토큰의 위치가 어긋나 버린다는 점에 있다. 위치가 어긋나면 RoPE(회전 위치 임베딩 (Rotary Positional Embedding))의 위상도, 캐시 상의 슬롯 식별자도 변한다. 여기서 망가지는 것이 바로 prefix 캐시 (prefix cache) 이다. vLLM의 automatic prefix caching을 대표로 하는 이 메커니즘은 여러 요청에서 앞부분의 문맥이 일치할 때, 그 부분의 KV 계산을 통째로 재사용한다. 에이전트는 매 턴 거의 동일한 긴 전제(시스템 프롬프트 + 지금까지의 이력)를 다시 보내기 때문에, prefix 캐시의 히트율(Hit rate)이 그대로 응답 속도와 비용에 직결된다.
from vllm import LLM
# 앞부분이 일치하는 요청 간에 KV 계산을 재사용함
llm = LLM(model="Qwen/Qwen2.5-14B", enable_prefix_caching=True)
즉, 기존의 KV 축소는 메모리를 절약하는 이면에, 재정렬(Compaction)로 인해 prefix 캐시의 일치를 깨뜨려 재계산이라는 형태로 절약분을 다시 지불하고 있었다. 논문의 측정에 따르면, compaction 계열의 prefix 히트율은 0~3%까지 떨어진다. 개인적으로는 이 부분이 가장 핵심적인 포인트라고 생각한다. KV 축소와 캐시 재사용은 별개의 최적화로 다뤄지기 쉽지만, 실운용에서는 명확한 트레이드오프(Trade-off) 관계에 있다.
IntentKV의 첫 번째 아이디어는 허탈할 정도로 소박하다. 버리고 싶은 토큰의 KV 행을 삭제하고 채우는 것이 아니라, 해당 물리 슬롯을 하나의 '파수꾼(Sentinel) 슬롯'으로 교체하는 것이다. 파수꾼의 Key에는 극단적으로 작은 값(K = -10⁴)을 넣어두기 때문에, softmax를 통과하면 주의 가중치가 0으로 뭉개져 사실상 해당 토큰은 존재하지 않는 것과 다름없는 취급을 받게 된다.
살아남은 토큰은 원래의 논리적 위치와 물리적 슬롯에 그대로 머물러 있기 때문에, RoPE의 위상도 슬롯 식별자도 변하지 않는다. 결과적으로 prefix 캐시의 트리 구조(Radix tree)는 다음 턴에도 계속 매칭된다. 이 교체만으로 8k 토큰 예산에서의 prefix 히트율이 20.7%까지 회복된다. 0~3%와의 차이는 곧 에이전트의 1턴당 재계산량의 차이로 이어진다.
두 번째는 어떤 토큰을 남길지에 대한 판단을 턴을 넘나들며 영리하게 만드는 부분이다. IntentKV는 세션 단위로 '의도 메모리(Intent memory)'를 가지며, 턴마다의 쿼리(Query)를 감쇠시키며 더해 나간다.
M_t = e^(-λ)·M_(t-1) + Enc(q_t) (λ=0.5)
직전의 프롬프트 하나만으로 토큰의 중요도를 측정하는 것이 아니라, 에이전트가 이 세션에서 무엇을 계속 추적하고 있는지라는 '의도의 이력'에 비추어 남길 토큰을 선택한다. 나아가 동결된 모델 위에 제로 초기화(Zero-initialized)된 잔차 헤드(Residual head)를 얹어 학습 보정을 가한다. 제로 초기화이므로 학습이 효과를 발휘하기 전까지는 규칙 기반(Rule-based)의 점수가 하한선으로서 보장되는 구조다. 베이스 모델 자체에는 일절 손을 대지 않는다.
핵심적인 효과는 동일한 KV 예산으로 제한했을 때의 정답률이다. 에이전트용 심층 검색 벤치마크(약 10만 문서를 대상으로 하는 830문항의 BrowseComp-Plus, 다단계 QA인 FRAMES)에서 KV 예산을 8k 토큰으로 고정하고 비교한 결과가 이것이다(Qwen2.5-14B).
| 방법 | 8k 예산에서의 정답률 |
|---|---|
| IntentKV | 18.55% |
| ... |
절대값이 낮은 것은 벤치마크 자체가 어렵기 때문이며, 주목해야 할 점은 상대적 차이다. 동일한 메모리 예산에서 타 방법론의 두 배 이상을 기록하고 있다. Qwen2.5-8B에서는 8k 예산의 IntentKV가 14.10%로, 16k 예산의 풀 캐시 (Full Cache, 15.06%)에 거의 육박했다. 절반의 예산으로 정확도를 거의 떨어뜨리지 않는다는 주장은 설득력이 있다. 메모리 측 수치도 놀라운데, 긴 쿼리 100건의 최악의 케이스(Worst Case)에서 피크 토큰(Peak Token)이 92.3k $\rightarrow$ 20.5k, KV 읽기(KV Read)가 411M $\rightarrow$ 31M (약 92.6% 감소)까지 떨어졌다.
그렇다고 해서 당장 내일부터 사용할 수 있는 것은 아니다. 논문 자체에서도 한계를 인정하고 있다. KV 예산 $C$는 쿼리마다 동적으로 조정되지 않고 전체적으로 고정되며, 잔차 헤드 (Residual Head)의 학습은 ToolBench의 멀티 턴 궤적 (Multi-turn Trajectory)에 국한되었고, 검증은 8B~14B 모델까지 이루어졌다. 코드 공개 또한 현 시점에서는 확인되지 않는다. 따라서 지금 당장 도입해야 한다는 이야기가 아니라, 방향성을 제시한 연구로서 읽는 것이 옳다.
다만, 센티넬 슬롯 (Sentinel Slot)을 통해 무효화한다는 발상은 단순하며, KV 축소와 prefix 캐시를 하나의 문제로 설계하는 관점은 추론 스택 (Inference Stack) 측에서 독립적으로 수용할 여지가 있다. 실무에서 vLLM 등으로 에이전트를 구동하는 사람들에게 당면한 교훈은, enable_prefix_caching을 확실히 활성화해 두는 것, 그리고 적극적인 KV 압축을 도입할 때는 prefix 히트율 (Prefix Hit Rate)이 조용히 떨어지고 있지는 않은지 함께 측정하는 것이다. 한쪽만 보고 있으면, 절약했다고 생각한 비용이 재계산 (Re-computation)으로 인해 상쇄될 수 있다. IntentKV는 바로 그 함정에 이름을 붙인 논문이기도 하다.
AI 자동 생성 콘텐츠
본 콘텐츠는 Qiita AI의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기