← Back to feed
Papers·2주 전

CMU, KV 캐시 예산을 모델 불확실성으로 조절하는 CONF-KV — 32K Needle 91.4%, 메모리 2.8배 절감

CMU, KV 캐시 예산을 모델 불확실성으로 조절하는 CONF-KV — 32K Needle 91.4%, 메모리 2.8배 절감

CMU 팀이 LLM 디코딩 시 모델의 현재 불확실성(다음 토큰 분포)을 활용해 KV 캐시 예산을 동적으로 조절하는 CONF-KV를 제안했습니다. 모델이 불확실할 때는 더 많은 컨텍스트를 유지하고, 확신이 있을 때는 과감히 가지치기합니다. 4K 생성 길이에서 full KV 대비 perplexity 1.5~2.1pt 손실에 그치면서 고정 512 슬라이딩 윈도우 수준의 메모리를 사용했고, 32K Needle-in-a-Haystack에서 91.4% 검색 정확도(슬라이딩 윈도우 53.8%, H2O 80.6%)를 기록했습니다. VisualWebArena 75개 태스크에서는 full-KV 성공률의 95.3%를 유지하면서 최대 메모리는 2.8배 절감했습니다.

CMU 연구진이 LLM 디코딩 시 모델의 불확실성 신호를 활용해 KV 캐시 예산을 동적으로 조절하는 방법을 공개했습니다.

핵심 결론

  • 성능4K 생성에서 full KV 대비 perplexity 1.5~2.1pt 손실, 32K Needle 검색 정확도 91.4%.
  • 메모리VisualWebArena에서 full-KV 성공률 95.3% 유지, 최대 메모리 2.8배 절감.

방법

  • 불확실성 기반 예산다음 토큰 분포의 엔트로피를 스칼라 신뢰도 점수로 변환, 불확실할수록 더 많은 캐시를 유지합니다.
  • 블록 내 정책각 예산 내에서는 누적 어텐션 가중치와 최신성의 합성 점수로 토큰을 순위화하며, 최근 윈도우는 보호합니다.
  • 구현블록별 online-softmax 어텐션, FP16/INT8 혼합 저장, 피라미드형 레이어별 예산 변형을 함께 사용합니다.

한계·조건

  • 벤치 범위실험은 최대 32K 컨텍스트, 4K 생성 길이에 국한 — 더 긴 컨텍스트에서의 추세는 추가 검증 필요.
  • 코드현재 논문 내 abstract와 figure만 공개, 코드는 추후 공개 예정.

편집자 한 줄

불확실성이라는 디코딩 중에도 계산되는 신호를 캐시 관리에 직접 연결한 점이 깔끔합니다. 다만 피라미드 예산 변형이 추가 오버헤드를 만들지는 확인이 필요해 보입니다.

  • #kv-cache
  • #inference
  • #llm
  • #carnegie-mellon
Carnegie Mellon University
원문 보기 →

Comments

— 첫 댓글을 남겨보세요 —