긴 추론에서 KV 캐시 줄이기
LLM은 다음 토큰을 만들 때 앞에서 계산한 키와 값을 KV 캐시에 저장해 다시 씁니다. 풀이가 길어질수록 이 캐시도 커집니다. 이전 토큰을 지우면 메모리는 아낄 수 있지만 풀이에 필요한 정보까지 사라져 압축 전과 다른 답을 낼 수 있습니다.
CASK에서는 답을 만드는 데 필요한 상태를 남기고 반복된 중간 추론은 묶어 저장합니다. 실험에서는 캐시를 줄인 뒤 원본 모델의 추론을 얼마나 유지하는지 확인했습니다.
중요도 점수의 한계
TriAttention처럼 토큰에 중요도 점수를 매겨 일부를 남기는 접근을 살폈습니다. 점수 계산을 바꿔도 실제로 남는 토큰은 거의 같을 때가 있었습니다. 점수만 더 정교하게 매겨서는 추론에 필요한 정보를 지키기 어려웠습니다.
그래서 생성 중 쌓이는 상태를 두 부류로 나눴습니다. 답을 결정하거나 다음 풀이를 이어가는 데 쓰이는 상태는 protected core로 남깁니다. 반복 탐색이나 재확인에서 나온 비슷한 상태는 mergeable scratch로 구분해 병합합니다.
병합할 때는 각 상태의 기여도를 반영해 대표 키와 값을 계산합니다. 이후 생성할 토큰과 얼마나 관련 있는지 고려한 거리로 병합 대상을 고릅니다. 합치는 과정에서 잃는 정보도 따로 점검합니다. core는 이 병합 대상에 넣지 않습니다.
입력과 생성 구간 분리
입력 문서가 길면 답을 만들기도 전에 프롬프트가 캐시 예산을 대부분 차지합니다. 생성 구간을 압축하는 규칙이 있어도 적용할 여유가 없습니다. 그래서 압축을 두 단계로 나눴습니다.
- 입력 구간에서는 일부 캐시를 제거해 생성에 쓸 공간을 확보합니다.
- 생성 구간에서는 core를 보존하고 scratch만 선택적으로 합칩니다.
실험 방법과 결과
실험에는 Qwen3-8B를 사용했습니다. H100 PCIe에서 bfloat16·SDPA 설정으로 실행했고 AIME24와 AIME25에서 각각 6개 예제를 비교했습니다.
캐시를 압축하지 않은 모델의 출력을 먼저 얻은 뒤, 같은 토큰을 압축 모델에 하나씩 넣었습니다. 이때 다음 토큰 예측이 얼마나 일치하는지, 원본 토큰에 얼마나 높은 확률을 주는지, 어디서 처음 예측이 어긋나는지 기록했습니다.
| 데이터 | 방법 | 다음 토큰 일치율 | 평균 NLL |
|---|---|---|---|
| AIME24 | TriAttention | 88.2% | 0.383 |
| AIME24 | CASK | 90.7% | 0.268 |
| AIME25 | TriAttention | 89.1% | 0.357 |
| AIME25 | CASK | 90.3% | 0.313 |
KV 예산을 384로 맞췄을 때 두 데이터셋 모두 CASK의 토큰 일치율이 높고 평균 NLL이 낮았습니다. CASK의 일치율은 예산을 512로 늘린 TriAttention의 AIME24 89.4%, AIME25 89.9%보다도 높았습니다.
긴 입력에서 확인한 한계
LongBench의 multi_news와 vcsum에서는 생성 구간의 병합이 실제로 실행됐고 원본 모델의 출력도 더 잘 따라갔습니다. qasper의 개선은 주로 입력 구간을 정리한 데서 나왔습니다.
qmsum과 gov_report에서는 프롬프트만으로 예산이 소진되는 문제가 남았습니다. 2wikimqa도 예산 384에서는 CASK의 토큰 일치율이 더 낮았습니다.
논문과 구현
CASK: Core-Aware Selective KV Compression for Reasoning Traces
Buseong Kim · Heejun Gwon
arXiv에 공개했습니다. arXiv:2604.10900.
공개 저장소에는 CASK와 TriAttention 구현, 실험 실행 스크립트, 측정 결과와 논문 도표를 함께 담았습니다.