CONF-KV: Confidence-Aware KV Cache Eviction with Mixed-Precision Storage for Long-Horizon LLM

기술 요약: CONF-KV

문제 제기
장기적 (Long-horizon) 대규모 언어 모델 (LLM) 추론은 키 - 밸류 (KV) 캐시가 GPU 메모리의 주요 소비자가 되어 토큰당 어텐션 지연 시간이 시퀀스 길이에 비례하여 선형적으로 증가하는 치명적인 병목 현상에 직면해 있습니다. 슬라이딩 윈도우나 과거 어텐션 기반 방법 (예: H2O, SCISSORHANDS) 과 같은 기존 제거 정책들은 어떤 토큰을 유지할지 결정하기 위해 정적인 최근성이나 과거 신호에 의존합니다. 이러한 접근 방식은 매 디코딩 단계에서 이용 가능한 신호인 모델의 현재 불확실성을 활용하지 못하는 경우가 많습니다. 결과적으로 정적 정책은 어려운 생성 구간에서는 컨텍스트를 조기에 제거하거나, 모델이 확신할 때는 불필요한 토큰을 유지하여 메모리 사용량과 생성 품질 간의 비최적 균형을 초래합니다.

방법론: CONF-KV
저자들은 모델의 현재 다음 토큰 분포에 기반하여 캐시 예산을 동적으로 조정하는 신뢰도 인식형 KV 캐시 관리자 CONF-KV를 제안합니다. 이 시스템은 다음과 같은 메커니즘을 통해 작동합니다:

1. 신뢰도 추정: 각 디코딩 단계에서 로짓 (logits) 을 스칼라 신뢰도 점수 ($c$) 로 변환합니다. 이 점수는 정규화된 엔트로피, 상위 두 토큰 간의 로그 확률 마진, 그리고 상위 토큰의 확률의 가중치 조합입니다.
   $$c = w_H(1 - \hat{H}) + w_m\sigma(m) + w_p p_{(1)}$$
   여기서 $\hat{H}$는 정규화된 엔트로피, $m$은 로그 마진이며, 가중치는 $(0.4, 0.3, 0.3)$으로 설정됩니다.

2. 적응형 예산 선택: 임계값 $\tau$가 현재 단계의 캐시 예산 ($N$) 을 결정합니다.

   • $c \ge \tau$인 경우 (높은 신뢰도), 공격적인 제거를 허용하는 엄격한 예산 ($N_{high}$) 이 선택됩니다.
   • $c < \tau$인 경우 (낮은 신뢰도/불확실성), 더 많은 컨텍스트를 유지하기 위해 느슨한 예산 ($N_{low}$) 이 선택됩니다.
     이로 인해 불확실한 구간에서는 캐시가 확장되고 확신 있는 구간에서는 축소되는 "톱니바퀴 (sawtooth)" 형태의 메모리 프로파일이 생성됩니다.

3. 토큰 순위 매기기 및 제거: 선택된 예산 내에서 토큰들은 지수 이동 평균 (EMA) 어텐션 질량과 최근성을 결합한 복합 점수에 따라 순위가 매겨집니다. 하드 보호 윈도우 ($P$) 는 지역적 일관성을 유지하기 위해 가장 최근 토큰이 절대 제거되지 않도록 보장합니다.

4. 시스템 통합:

   • 혼합 정밀도 저장: 가장 최근 $W$개의 토큰은 FP16 으로 저장되고, 유지되는 이전 토큰들은 헤드 및 채널 단위로 INT8 로 양자화됩니다. 이는 더 낮은 비트 양자화에 비해 퍼플렉시티 저하를 최소화하면서 메모리 발자국을 줄입니다.
   • 피라미드형 예산 (CONF-KV-L): 선택적 변형으로, 더 깊은 레이어가 종종 정보를 더 적은 토큰에 집중한다는 관찰에 기반하여 레이어 간에 예산을 비균일하게 할당합니다.
   • 구현: 관리자는 어텐션 커널을 위한 밀집 레이아웃을 유지하기 위해 연속적인 압축 (contiguous compaction) 을 사용하고, 간접 테이블을 피하며, 블록 단위 어텐션 읽기에 디양자화를 융합합니다.

주요 기여

• 신뢰도 인식형 정책: 역사적 신호에만 의존하는 것이 아니라 현재 출력 분포를 사용하여 단계별 캐시 예산 결정을 주도하는 새로운 제거 전략.
• 시스템 설계: 적응형 제거와 블록 단위 온라인 소프트맥스 어텐션, 혼합 FP16/INT8 저장, 선택적 피라미드형 레이어 예산을 결합한 완전한 구현.
• 기작적 검증: 최근 컨텍스트를 제거할 때 유발되는 KL 시프트와 신뢰도 점수가 강력하게 반비례 (피어슨 $r = -0.77$) 한다는 것을 보여주는 테스트로, 낮은 신뢰도가 더 많은 유지 컨텍스트 필요성을 나타낸다는 가정을 검증함.
• 종합적 평가: 퍼플렉시티, 장기 컨텍스트 검색, 웹 에이전트 작업을 포함한 다양한 워크로드와 GPT-2, Qwen-14B/32B, gpt-oss-20b 등 네 가지 모델 계열에 걸친 광범위한 벤치마크.

실험 결과

• 메모리 - 품질 트레이드오프: 2048 토큰 생성 길이의 GPT-2 에서 CONF-KV+INT8 은 고정 512 토큰 슬라이딩 윈도우 (~38.7 MB) 와 동일한 메모리 발자국을 가지지만, 슬라이딩 윈도우 대비 퍼플렉시티를 3.11 포인트 감소시킵니다. CONF-KV-L(피라미드형 변형) 은 메모리를 34.2 MB 로 추가로 줄이면서 퍼플렉시티를 3.89 포인트 개선하여, 슬라이딩 윈도우와 전체 KV 간의 품질 격차의 74% 를 해소합니다.
• 장기 컨텍스트 검색: 최대 32K 토큰까지의 "Haystack 속 바늘 찾기 (Needle-in-a-Haystack, NIAH)" 테스트에서 CONF-KV 는 91.4% 의 검색 정확도를 달성하여 슬라이딩 윈도우 (53.8%) 와 H2O(80.6%) 를 크게 능가합니다.
• 에이전트 작업: 75 개의 VisualWebArena 작업에서 CONF-KV 는 전체 KV 성공률의 95.3% 를 유지하면서 피크 메모리를 2.8 배 감소시킵니다.
• 처리량 및 지연 시간: CONF-KV 는 GPT-2 와 Qwen-32B 에서 전체 KV 대비 단계별 지연 시간을 1.8 배 감소시킵니다. 처리량은 배치 크기에 따라 확장되며, 배치 크기 8 에서 CONF-KV 는 전체 KV 대비 2.06 배의 처리량을 달성합니다.
• 애블레이션: 신뢰도 신호를 분리하여 분석한 결과, 동일한 비율의 무작위 제거는 성능을 36.54 PPL 로 저하시키는 반면, 전체 CONF-KV 는 30.92 PPL 을 달성하여 신뢰도 게이트와 순위 매기기 메커니즘 모두 필수적임을 확인했습니다.

의의 및 주장
본 논문은 CONF-KV 가 현재의 불확실성이 메모리 - 품질 파레토 프론티어를 개선하는 데 유용한 시스템 메타데이터임을 입증한다고 주장합니다. 이 방법은 미래 지향적이기 때문에 제거된 토큰을 복구할 수는 없지만 어려운 구간에서의 조기 제거를 방지한다는 점에서 독특합니다. 이 접근법은 **학습 불필요 (training-free)**로 설계되었으며 모델 가중치를 변경하지 않고, 헤드 단위 할당, 레이어 단위 할당, 정밀도 선택과 같은 기존 기술과 직교하여 이를 조합할 수 있습니다.

저자들은 이 방법이 제거 임계값을 초과하지 않는 짧은 컨텍스트에서는 "작동 없음 (no-op)"이 된다는 점과, 연속적인 압축이 수집 (gather) 비용을 발생시킨다는 점 (비록 어텐션 절감 효과에 비해 작지만) 을 제한 사항으로 지적합니다. 또한 고온 (high-temperature) 샘플링 하에서는 신뢰도 신호의 정보량이 감소한다는 점도 인정합니다. 마지막으로, 논문은 낮은 추론 비용이 효율성을 높이는 동시에 바람직하지 않은 장기 자동화의 장벽을 낮출 수 있다는 이중 사용 고려 사항을 강조하며, 기반 모델에 대한 안전 장치가 여전히 필요함을 시사합니다.