기술 요약: KVzap - 빠르고, 적응적이며, 충실한 KV 캐시 프루닝(Pruning)

1. 문제 정의

트랜스포머 기반 언어 모델의 컨텍스트 길이가 수만 또는 수십만 토큰으로 확장됨에 따라, Key-Value (KV) 캐시는 효율적인 대규모 언어 모델(LLM) 추론의 지배적인 병목 구간이 되었습니다. 중간 어텐션 상태를 저장하는 KV 캐시는 상당한 GPU 메모리를 소비하며(예: 128k 시퀀스 길이에서 65B 모델의 경우 335 GB), 첫 번째 토큰 생성 시간(time-to-first-token)을 증가시키고 디코딩 처리량(throughput)을 감소시킵니다.

Grouped Query Attention, Multi-head Latent Attention과 같은 구조적 변경이나 희소 어텐션(sparse attention) 메커니즘이 헤드($H$), 차원($D$), 레이어($L$) 축을 따라 압축을 해결해 왔지만, 시퀀스 길이 축($T$)은 널리 채택된 방법들로 인해 여전히 미해결 상태로 남아 있습니다. 기존의 임의적인(ad-hoc) 프루닝 방법들은 실제 운영 환경 배포를 위한 네 가지 핵심 기준을 충족하지 못하는 경우가 많습니다:

1. 빠르고 가벼움: 무시할 수 있는 수준의 오버헤드.
2. 단계 불가지론(Phase-agnostic): 프리필링(prefilling, 긴 컨텍스트)과 디코딩(reasoning) 모두에 적용 가능.
3. 최적화 친화적: FlashAttention2 및 PagedAttention과 같은 커널과 호환 가능.
4. 충실함(Faithful): 최소한의 정확도 저하.

현재의 최첨단 방식인 KVzip은 높은 압축률을 달-성하지만, 확장된 프롬프트에 대해 계산 비용이 많이 드는 이중 순전파(double forward pass)를 요구하므로 온라인 추론 및 디코딩 작업에는 부적합합니다.

2. 방법론

저자들은 KVzip의 고비용 문제를 해결하기 위해 프리필링과 디코딩 과정 모두에서 작동하는, KVzip의 빠르고 입력 적응적인 근사치인 KVzap를 소개합니다. 이 방법론은 세 가지 핵심 구성 요소로 이루어져 있습니다.

2.1. 향상된 스코어링 정책 (KVzip+)

저자들은 먼저 KVzip의 스코어링 메커니즘을 개선했습니다. KVzip은 프롬프트가 반복될 때 토큰이 받는 최대 어텐션 가중치를 측정하여 "컨텍스트 재구성(context reconstruction)"을 점수화합니다. 저자들은 Devoto 등(2025)에서 영감을 얻은 정규화 항을 통합하여 이를 강화했는데, 이는 토큰의 값 벡터(value vector)가 잔차 스트림(residual stream)에 기여하는 정도를 고려합니다.
KVzip+ 점수($s^+_i$)는 다음과 같이 정의됩니다:
$$s^+_i = \max_{j \in \text{<prompt>}} \frac{a_{ji}}{\|W_O v_i\| / \|h_j\|}$$
여기서 $a_{ji}$는 어텐션 가중치, $W_O$는 출력 투영 행렬, $v_i$는 값 벡터, $h_j$는 은닉 상태(hidden state)입니다.

2.2. 대리 모델링 (Surrogate Modeling)

KVzip의 과도한 비용이 드는 이중 순전파를 피하기 위해, KVzap는 모델의 입력 은닉 상태($h$)로부터 로그 스코어($\log(s^+)$)를 직접 예측하는 경량 대리 모델(선형 레이어 또는 2계층 MLP)을 학습시킵니다.

• 입력: 은닉 상태 $h_t \in \mathbb{R}^{D_h}$.
• 출력: 각 KV 헤드에 대한 점수 $\mathbb{R}^H$.
• 학습: 모델은 Nemotron-Pretraining-Dataset에서 샘플링된 120만 개의 (은닉 상태, 로그 스코어) 쌍을 사용하여 학습됩니다. MLP 변형은 GELU 활성화 함수를 사용하는 $D_h/8$의 은닉 폭을 가집니다.
• 성능: 대리 모델은 KVzip+ 오라클(oracle)과 평균 $R^2$ 상관관계 0.60–0.80을 달성하여, 고비용의 점수를 은닉 상태로부터 정확하게 근사할 수 있음을 입증했습니다.

2.3. 적응형 프루닝 및 퇴거 (Adaptive Pruning and Eviction)

고정된 예산(예: 정확히 50%의 토큰 유지)을 강제하는 KVzip과 달리, KVzap는 임계값 기반 프루닝을 사용합니다:

• 임계값 설정: 예측된 점수가 고정된 임계값 $\tau$ 미만인 토큰은 버려집니다. 이를 통해 압축률이 입력 프롬프트의 정보 밀도에 따라 동적으로 적응할 수 있습니다.
• 슬라이딩 윈도우: 로컬 컨텍스트를 보존하기 위해, StreamingLLM 방식을 따라 가장 최근의 $w=128$개 토큰은 항상 유지합니다.
• 구현: 프루닝은 어텐션 연산 후에 수행됩니다. 이 과정은 계산 효율적이며, 매우 낮은 오버헤드(MLP의 경우 1.1%, 선형 모델의 경우 0.02% 이내)를 가집니다.

3. 주요 기여

1. KVzip+: 은닉 상태 정규화를 통합하여 개선된 KV 캐시 중요도 스코어링 지표를 제시했습니다.
2. 대리 근사(Surrogate Approximation): KVzip+ 점수를 은닉 상태로 학습된 경량 모델로 효과적으로 근사할 수 있음을 입증하여, 이중 순전파의 필요성을 제거했습니다.
3. KVzap: 이러한 대리 모델을 임계값 기반 퇴거 및 슬라이딩 윈도우와 함께 적용하여, 네 가지 실무적 기준(속도, 단계 불가지론, 최적화 친화성, 충실함)을 모두 충족하면서 최첨단 성능을 달성하는 새로운 프루닝 기술을 제안했습니다.

4. 실험 결과

Qwen3-8B, Llama-3.1-8B-Instruct, Qwen3-32B 모델을 대상으로 RULER(긴 컨텍스트 검색/집계), LongBench(다중 작업 긴 컨텍스트), AIME25(수학적 추론)의 세 가지 벤치마크에서 실험을 진행했습니다.

• 압축 대 정확도: KVzap는 무시할 수 있는 수준의 정확도 손실과 함께 2~4배의 KV 캐시 압축을 달성합니다.
  • RULER 4k에서 KVzap는 KVzip+ 오라클의 성능을 능가하거나 대등한 수준을 보였으며, Expected Attention 및 Duo Attention을 포함한 15개의 다른 프루닝 방법보다 우수한 성능을 보였습니다.
  • LongBench에서 KVzap는 2~3배 압축 시에도 거의 완벽한 정확도를 유지합니다.
  • AIME25에서 KVzap-MLP는 KV 캐시의 50% 이상을 버릴 때도 추론 정확도를 보존합니다.
• 적응성: 임계값 기반 접근 방식은 작업 난이도에 따라 압축률을 자동으로 조정합니다. 예를 들어, 밀도가 높은 실제 데이터인 LongBench에 비해 합성 데이터인 RULER에서 더 높은 압축률을 달성합니다.
• 효율성: 이 방법은 계산 및 메모리 오버헤드를 거의 추가하지 않습니다. 디코딩 중에 발생하는 추가 FLOPs는 메모리 대역폭 제한으로 인해 유휴 상태가 될 GPU 사이클을 활용합니다.

5. 의의 및 주장

본 논문은 KVzap가 학술적 프루닝 연구와 실제 추론 엔진 사이의 간극을 메운다고 주장합니다. 이 연구의 의의는 다음과 같습니다:

• 실용성: 속도를 희생하거나 구조적 변경을 요구하지 않으면서 최첨단 압축을 달성하는 최초의 방법으로서, 주요 추론 엔진(vLLM, SGLang, TRT-LLM)과 호환됩니다.
• 충실함: LLM이 KV 캐시를 완전히 활용하지 않는다는 점과, 베이스 모델을 엔드 투 엔드로 재학습하지 않고도 은닉 상태를 통해 사용되지 않는 쌍(pair)을 식별할 수 있음을 입증했습니다.
• 확장성: 8B에서 32B 파라미터에 이르는 모델 전반에 걸쳐 검증되었으며, 미래의 희소 어텐션 아키텍처와 호환되도록 설계되었습니다.

저자들은 KVzap가 엔드 투 엔드로 학습된 목적 함수가 아닌 사후(post-hoc) 추가 방식이라는 점을 한계로 언급하며, 압축률을 최적화된 커널(예: PagedAttention)을 통해 실제 실행 시간(wall-clock time) 단축으로 전환하기 위한 추가적인 엔지니어링이 필요하다고 인정했습니다. 그러나 KVzap의 단순성, 높은 압축률, 견고한 성능의 조합이 실제 운영 환경 배포의 유력한 후보라고 주장합니다.