LU-KV — KV 캐시에서 “지금 중요한 토큰”이 아니라 “앞으로 쓸모 있을 토큰”을 남겨라
📡

LU-KV — KV 캐시에서 “지금 중요한 토큰”이 아니라 “앞으로 쓸모 있을 토큰”을 남겨라

Tags
KVcache
Attention
LLM Reasoning
NLP
AI
Published
August 25, 2026
Author

1. KV 캐시 축출의 2단 구조

긴 컨텍스트 LLM 추론의 병목은 KV 캐시다. 시퀀스가 길어질수록 캐시 메모리가 선형으로 불어나고 attention 계산의 제곱 복잡도와 맞물려 처리량을 갉아먹는다. 그래서 덜 중요한 토큰의 K/V를 버리는 KV 캐시 축출(eviction)이 필수 기법이 됐다.
기존 연구는 보통 두 축으로 나뉜다.
  1. Head 내부 스코어링(intra-head scoring). 각 head 안에서 어떤 토큰을 남길지 점수를 매긴다. StreamingLLM(attention sink), H2O(누적 attention), SnapKV(관찰 창 기반), KeyDiff(Key 벡터 기하), CriticalKV(출력 섭동), KVZip(컨텍스트 재구성) 등이 여기에 속한다.
  1. Head 간 예산 배분(cross-head budget allocation). 전체 메모리 예산을 layer·head에 어떻게 나눌지 정한다. Uniform(균등), PyramidKV(층이 깊을수록 감축하는 피라미드 구조), HeadKV·CAKE(태스크 prior 기반 규칙), AdaKV(실시간 통계 기반 동적 배분) 등.
논문의 표적은 두 번째 축이다. 스코어링 지표는 정교해졌지만 배분 전략은 여전히 점수가 큰 곳에 예산을 더 준다는 전제를 말없이 깔고 있다.

2. 휴리스틱 점수는 head마다 신뢰도가 다르다

저자들이 짚는 근본 결함은 가정 하나에 있다. 휴리스틱 점수의 크기(magnitude)가 중요도의 일관된 대리 지표라는 가정이다. 이 가정은 head 간 예측 충실도의 이질성을 무시한다. 어떤 head는 순간 기여를 우선하고 어떤 head는 오래 갈 효용을 포착한다. 점수가 높다고 예산을 몰아주면, 점수가 미래 중요도와 어긋나는 head에서는 투자 대비 수익이 급감한다.
이를 정량화하기 위해 Oracle Importance를 정의한다. head (ℓ,h)의 캐시 위치 j가 미래 디코딩 창(K_max 스텝) 안에서 출력 벡터에 기여할 수 있는 최대치다:
이 기준으로 정렬한 이상적 지표(Oracle Metric π*)의 선택 집합을 실제 휴리스틱 지표 π의 것과 견주면 남긴 토큰이 셋으로 갈린다. Hit(맞게 보존), Miss(중요한데 버림), False Positive(안 중요한데 남김). 축출 손실을 전개하면 다음 분해가 나온다:
앞 항은 압축률이 정해지면 고정되는 이론상 하한이다. 뒤 항이 바로 휴리스틱이 미래를 못 맞혀서 추가로 잃는 손실이다.
Figure 1
Figure 1
Figure 1
Figure 1. Mistral-7B-v0.3에서 압축률을 바꿔 가며 측정한 Oracle Importance 리콜. 빨간 Oracle 곡선과 1.0 사이가 고정 손실, SnapKV·KeyDiff 곡선과 Oracle 곡선 사이의 세로 간격이 Optimality Gap 손실이다. 압축이 세질수록 휴리스틱 지표의 리콜이 Oracle에서 크게 벌어진다.

3. LU-KV의 세 가지 구성 요소

3.1 예산 배분을 전역 조합최적화로 풀기

목표는 전체 예산 B_total 제약 아래 모든 layer·head의 축출 손실 합을 최소화하는 head별 예산 {b_ℓ,h}를 찾는 것이다:
문제는 π가 유도한 순서 탓에 손실 수열이 예산에 대해 비볼록이라는 데 있다(증명은 부록 A.1. π 순서에 Oracle 중요도의 역전이 하나라도 있으면 이차 차분이 음수가 된다는 게 핵심이다). 정확해는 DP로 구할 수 있지만 프로파일링 규모에서는 비용이 크다.

3.2 Convex-hull 완화 + 탐욕 해법

저자들은 각 head의 이산 손실 수열에서 한계 감소량을 계산한다. 그다음 PAVA(Pool Adjacent Violators Algorithm) 기반 등위회귀로 비음·비증가 수열에 사영한다. 이를 누적 차감으로 재구성하면 볼록·비증가 대리 손실이 나오고 i번째 토큰을 보존할 때의 유효 한계 이득 g(i)가 단조 비증가가 된다. 그러면 지금 한계 이득이 가장 큰 head에 예산을 1단위씩 붓는 전역 탐욕 전략이 완화 문제의 전역 최적을 달성한다.
Figure 2
Figure 2
Figure 2
Figure 2. (a) convex-hull 완화 기반 탐욕 해법이 DP 최적해와 정확히 겹친다. (b) 같은 전역 압축률 σ 아래에서 head 수준 최적 압축률 r_ℓ,h가 hotpotqa·qmsum·trec 등 서로 다른 다운스트림 태스크에서 일관된 곡선을 그린다. 이 안정성이 다음 단계인 오프라인 프로파일링의 근거다.

3.3 오프라인 프로파일링으로 미래를 미리 구워 둔다

Oracle Importance는 미래 디코딩 결과가 필요하므로 실시간 추론 중엔 계산할 수 없다. 저자들이 붙잡은 건 Figure 2b의 관찰이다. head별 최적 로컬 압축률의 프로파일이 태스크가 바뀌어도 좀처럼 흔들리지 않는다. 그래서 다음 3단계 오프라인 프로토콜로 프로파일을 미리 만든다.
  1. 컨텍스트 생성. 긴 입력을 하나 만든다. 4,000토큰 정도이고 서사가 일관되게 이어지는 글이다. 기본 프로파일은 AI가 생성한 중국어 소설 발췌이며 모든 평가 벤치마크와 겹치지 않는다.
  1. Oracle 계산. 서로 다른 정보 구간을 겨냥한 질문 30개를 만든 다음, 질문마다 full-attention 디코딩으로 실제 Oracle Importance를 계산한다.
  1. 프로파일 집계. 전역 압축률 ρ의 촘촘한 그리드마다 완화 문제를 풀어 질문별 최적 로컬 비율을 얻는다. 이걸 질문 전체로 평균 내면 정적 프로파일 Φ(ρ)가 나온다.
추론 시에는 목표 압축률로 프로파일을 lookup해 비율을 정수 예산으로 바꾼 뒤 각 head가 독립적으로 휴리스틱 지표 π로 top-b 토큰만 남기면 끝이다. 온라인 최적화가 전혀 없으니 추가 비용도 없다시피 하다.
구현 세부(부록 E): 스코어링에는 투영된 Value 벡터의 L2 노름을 쓰고 layer 내 정규화를 적용한다. attention sink 4개 보존, 최근 토큰 창은 KeyDiff 1 / SnapKV 32, head당 최대 압축 99%(최소 1% 보존) 제약을 둔다.

4. 80% 압축에서 벌어지는 일

평가는 NVIDIA KVPress 프레임워크 위에서 question-agnostic 프로토콜로 진행됐다. 질문이 도착하기 전에 문서만 보고 캐시를 압축하는 방식이다. 프롬프트를 미리 캐싱해 두고 나중에 어떤 질문이 올지 모르는 실제 서빙 환경과 조건이 같다. 질문을 미리 보는 query-aware 설정보다 훨씬 가혹한 시험대다. 모델은 Llama-3.1-8B-Instruct, Mistral-7B-Instruct-v0.3, Qwen2.5-32B-Instruct.

LongBench (80% 압축, 표 1)

  • Mistral-7B + KeyDiff(π2): AdaKV 40.54 → LU-KV 46.21 (Full-KV 47.30). 압축으로 잃은 갭의 84%를 회복했다.
  • Mistral-7B + SnapKV(π1): Uniform 40.67 / Pyramid 40.94 / AdaKV 41.89 → LU-KV 45.79.
  • Qwen2.5-32B + KeyDiff: AdaKV 42.32 → LU-KV 48.26 (Full-KV 48.51). 32B 모델에선 Full-KV 코앞까지 간다.

진짜 격차가 드러나는 RULER-16K (80% 압축, 표 2)

  • Mistral-7B + SnapKV: Uniform 29.53, AdaKV 37.48에 그치는데 LU-KV는 69.98.
  • 특히 multikey-3 태스크에서 Uniform 1.00 → LU-KV 67.40. 희소하지만 없으면 안 되는 정보를 지키는 능력의 차이다.
  • Qwen2.5-32B + SnapKV: AdaKV 44.04 → LU-KV 80.47. KeyDiff에서도 80.81 → 88.88.
검색형 스트레스 테스트일수록 어느 head에 얼마를 주느냐가 성패를 뚜렷하게 가른다.

손실 분포로 본 기제 (Figure 3·4)

Figure 3
Figure 3
Figure 3
Figure 3. layer별 축출 손실 합계. PyramidKV(분홍)는 깊은 층(27~32)에서 손실이 폭등한다. 정보가 얕은 층으로 몰린다는 피라미드 가설이 이 구간의 높은 의미 밀도와 어긋나기 때문이다. AdaKV(파랑 점선)는 층 간 재배분을 못 해 Uniform과 거의 같은 궤적을 그린다. LU-KV(주황 실선)만 전 층에 걸쳐 손실이 낮고 평탄하다.
Figure 4
Figure 4
Figure 4
Figure 4. head 단위 손실 히트맵. 베이스라인들은 특정 head에서 “loss burst”(짙은 빨강 블록)를 겪는다. 점수 크기만 믿고 배분한 결과다. LU-KV는 전체 손실 3.87e-4로 AdaKV(4.98e-4)보다 낮으며 버스트 자체가 사라진다.
PyramidKV는 layer 축만, AdaKV는 (메모리 제약 때문에 사실상 층 내부에서만) head 축만 만진다. LU-KV는 오프라인 정적 프로파일 덕분에 런타임 메모리 스파이크 없이 층·head 양축을 진짜 전역 최적화한다.

효율 (Figure 5)

Figure 5
Figure 5
Figure 5
Figure 5. Llama-3.1-8B에서 잰 피크 메모리와 TTFT. LU-KV는 SnapKV 베이스라인과 같은 수준의 지연시간을 유지하면서 120k 컨텍스트에서 피크 메모리를 43.6GB → 31.5GB(80% 압축)로 줄인다. lookup 테이블 방식이라 배분 자체의 오버헤드가 없다는 주장 그대로다.

Ablation: Optimality Gap 항의 값어치 (표 3)

Gap 항을 빼고 Oracle Metric Loss만으로 배분하면 Mistral π2 기준 46.21 → 41.98로 떨어진다. Qwen2.5-32B π2에서도 48.26 → 43.18. 저자들은 휴리스틱 지표가 head마다 얼마나 못 미더운가를 최적화 목적에 명시적으로 넣는 것이 성능의 핵심이라고 봤다. 실험이 그 이론적 주장을 뒷받침한다.

5. 부록에서 건진 것들

압축률 민감도 (표 9, Llama-3.1-8B, Full 49.29)

방법
σ=50%
σ=80%
σ=90%
PyramidKV (SnapKV)
46.65
42.36
38.02
AdaKV (SnapKV)
47.91
43.98
40.20
LU-KV (SnapKV)
49.38
47.70
44.69
압축이 셀수록 격차가 벌어진다. 90%에서도 Full 대비 -4.6점에 그친다.

프로파일링 데이터 전이성 (표 10)

기본 프로파일(중국어 AI 소설 4k토큰 + 질문 30개)을 영어 SAT 독해 지문 150토큰짜리로 갈아 끼워도 LongBench 성능이 거의 같다(σ=0.5에서 48.47 vs 49.06, σ=0.8에서 47.16 vs 46.67). 교정 데이터는 언어도 길이도 도메인도 전혀 다르다. 그런데 거기서 뽑은 head 프로파일이 그대로 통한다. 이 프로파일이 잡아낸 것은 데이터가 아니라 모델 아키텍처에 내재한 head의 역할 분담으로 보인다. 이 논문에서 가장 놀라운 결과다.

다른 지표들 위에 얹기 (표 12·13)

LU-KV는 배분 레이어라서 스코어링 지표를 바꿔 끼울 수 있다. KVZip σ=0.9에서 40.84 → 44.79, CAKE σ=0.9에서 39.93 → 45.36, EA σ=0.9에서 36.58 → 46.28. 압축을 세게 걸수록 이득이 커진다. 같은 오프라인 교정 계열인 HeadKV와의 직접 비교(표 11)에서도 SnapKV 기준 45.27 → 46.37, EA 기준 45.34 → 47.12로 앞선다.

부록 그림 (Figure 6~9)

Figure 6
Figure 6
Figure 6
Figure 6. SnapKV·KeyDiff·EA 세 지표 모두에서 layer 손실 곡선과 head 히트맵의 패턴이 재현된다. Optimality Gap이 특정 지표의 문제가 아니라 휴리스틱 일반의 문제라는 보편성 검증.
Figure 7
Figure 7
Figure 7
Figure 8
Figure 8
Figure 8
Figure 9
Figure 9
Figure 9
Figure 7~9. Mistral-7B(KeyDiff)의 32개 layer × 8개 head 전체에 대한 최적 배분 프로파일. 검정 점선(본 방법)이 주황 점선(MCKP DP 최적해)과 겹치고 배경의 색 곡선들(각 다운스트림 태스크)이 대체로 같은 모양을 그린다. head마다 곡선의 형태가 확연히 다르다. 이 논문이 출발한 자리가 바로 여기다.

6. 비판적으로 읽기

강점 - 배분 문제를 손실 분해(고정 하한 + Optimality Gap)로 정식화한 것이 깔끔하다. “지표를 더 잘 만들자”가 아니라 “지표의 못 미더움 자체를 배분에 반영하자”는 관점 전환. - 지표-불가지론(metric-agnostic) 설계라 기존 방법들과 경쟁이 아니라 결합 관계다. KVZip·CAKE·EA에 얹은 실험이 그 증거다. - 오프라인 프로파일 + 런타임 lookup 구조는 배포 관점에서 온라인 동적 방법(AdaKV의 메모리 스파이크 문제)보다 현실적이다.
한계와 열린 질문 - LongBench의 Code 태스크에서는 LU-KV-π1이 57.69로 오히려 베이스라인(Uniform 60.98, AdaKV 62.05)보다 낮다. 서사형 교정 데이터에서 뽑은 프로파일이 코드처럼 구조가 다른 도메인에는 덜 맞을 수 있다는 신호다. 표 10의 전이성 결과와 겹쳐 보면 어디까지 전이되는지가 아직 불분명하다. - 프로파일은 모델별로 새로 구워야 한다. 교정 비용 자체는 작지만(문서 1개 + 질문 30개의 full-attention 디코딩) 모델 업데이트마다 반복해야 하고 K_max(미래 창 길이) 등 프로파일링 하이퍼파라미터의 민감도 분석은 본문에 없다. - 저자들도 Impact Statement에서 이 위험을 인정한다. 프로파일링 설정과 다른 도메인·컨텍스트 길이·디코딩 세팅에서 압축을 세게 걸면 드물지만 중요한 정보가 날아갈 수 있다. 배포 전에 목표 워크로드로 검증해야 한다. - 비교 대상이 모두 축출(eviction) 계열이다. 양자화·저랭크 근사 등 다른 압축 축과의 결합이나 비교는 다루지 않는다. - Oracle Importance 자체가 max 기반 정의라 여러 스텝에 걸쳐 조금씩 기여하는 토큰의 누적 효용은 과소평가될 수 있다.

7. 직접 써보기 (설치와 실행)

공개 코드는 NVIDIA KVPress를 포크해 LU-KV 배분을 얹은 구조다.
git clone https://github.com/baidu-baige/LU-KV
cd LU-KV
pip install -e .
pip install transformers==4.53.0 flash-attn datasets accelerate \
    fastparquet rouge nltk jieba fuzzywuzzy bert-score fire seaborn scikit-learn

# LongBench 16개 태스크 평가
cd evaluation
bash lukv_longbench.sh
# RULER 평가
bash lukv_ruler.sh
# 베이스라인 재현
bash baseline_longbench.sh

# 오프라인 프로파일(곡선 데이터) 직접 생성
bash curve_data/generate_curve.sh
저장소 구조는 evaluation/(벤치마크 스크립트), kvpress/(핵심 구현), notebooks/(예시), tests/로 나뉜다. 새 모델에 적용하려면 generate_curve로 그 모델의 프로파일부터 만들어야 한다.

8. 해외 리서처들은 이 문제를 어떻게 다루고 있나

이 논문을 실전 감각으로 읽으려면 주변 생태계를 같이 봐야 한다.
NVIDIA가 만든 KVPress는 30개 이상의 압축 방법을 HuggingFace transformers에 끼우는 훅(press) 형태로 통일했다. RULER·InfiniteBench·Loogle 벤치마크 CLI와 공개 리더보드도 같이 제공한다. 이제 이 라이브러리가 공용 실험대나 다름없다. 새 방법을 press 클래스 하나로 구현해 같은 조건에서 비교하는 것이 요즘 이 분야 논문의 표준 워크플로다. LU-KV도 정확히 이 위에서 평가됐다. 사용감은 이 정도로 단순하다:
from transformers import pipeline
from kvpress import ExpectedAttentionPress

pipe = pipeline("kv-press-text-generation",
                model="meta-llama/Llama-3.1-8B-Instruct", device="cuda")
press = ExpectedAttentionPress(compression_ratio=0.5)
answer = pipe(context, question=question, press=press)["answer"]
배분 연구는 프로덕션으로 넘어가는 중이다. LU-KV가 주요 비교 대상으로 삼은 AdaKV(NeurIPS 2025)는 Cloudflare 팀이 vLLM에 통합했다. 연구 코드가 아니라 서빙 스택에 head-wise 배분이 들어가기 시작했다.
LU-KV와 같은 시기에 나온 LKV(arXiv 2605.06676)는 head별 예산과 토큰 선택을 아예 end-to-end로 학습하면서 “성능의 주 동인은 토큰 선택이 아니라 예산 배분”이라고 주장한다. LAVa(2509.09754)는 layer별 동적 배분을, ThinKV(2510.01290)·ReasonAlloc(2606.11164)은 reasoning 모델의 디코딩 시점 배분으로 문제를 확장했다. 배분이 주 동인이라는 합의가 이렇게 형성되는 중이다. LU-KV의 오프라인 전역 최적화 노선과 LKV의 학습 기반 노선이 같은 결론(휴리스틱 배분의 한계)을 다른 방법으로 공격하는 형국이다. 후속 비교가 기다려진다.
question-agnostic 평가가 새 표준이 되고 있다. 프롬프트를 미리 캐싱하고 질문은 나중에 오는 서빙 시나리오(prompt caching)가 보편화되면서 질문을 보고 압축하는 기존 query-aware 평가는 점점 현실성이 떨어진다는 인식이 퍼지는 중이다. KVZip이 이 프로토콜을 밀었고 LU-KV도 전 실험을 이 조건으로 돌렸다.

9. 우리 연구(ARS) 관점의 메모

  • 이 논문의 head별 “예측 충실도 이질성”은 우리가 attention head 역할 분화를 다루는 방식과 문제 구조가 같다. 특히 Figure 7~9의 head별 배분 곡선은 head 기능 분류의 외부 근거로 인용할 수 있다.
  • 손실 분해(고정 하한 + 정책 기인 갭) 프레임은 우리 within-question 분석에서 지표 고유 한계와 선택 정책 한계를 갈라 보는 템플릿으로 차용할 만하다.
  • 전이성 실험(중국어 소설 → 영어 SAT)은 프로파일이 데이터가 아니라 아키텍처를 학습한다는 주장의 증거 설계로 참고할 가치가 있다. 우리 되튐 사이트 위치=아키텍처 vs 집중도=학습 발견과 논증 구조가 유사하다.

참고 링크

  • 관련 최신 연구: LKV (arXiv 2605.06676) · LAVa (2509.09754) · ThinKV (2510.01290) · ReasonAlloc (2606.11164)