KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

문제의 배경 — KV 캐시라는 병목

트랜스포머 기반 언어 모델에서 각 입력 토큰은 레이어마다 키-값 (KV) 벡터 쌍을 생성하고, 이를 자동회귀 생성 내내 메모리에 보관한다. KV 캐시의 형상은 (2, L, H, T, D)로, L은 레이어 수, H는 헤드 수, T는 시퀀스 길이, D는 키/값 차원이다.

"For example, in bfloat16 precision, the KV cache for a vanilla transformer like Llama1-65B (L=80, H=64, D=128) requires 335 GB of memory at T=128k."
예컨대 bfloat16 정밀도에서 Llama1-65B의 KV 캐시는 시퀀스 길이 128k 시 무려 335 GB에 달한다.

시퀀스 길이가 수만~수십만 토큰으로 길어지면서 KV 캐시는 GPU 최대 메모리 점유와 TTFT(Time To First Token)를 급격히 증가시키고 디코딩 처리량을 크게 낮춘다.

기존 접근법과 그 한계

KV 캐시 크기를 줄이는 시도는 크게 네 축(H, D, L, T)을 따라 이루어져 왔다.

H축 (헤드 수) GQA (Grouped Query Attention)는 여러 쿼리가 키-값을 공유하게 해 4~16× 압축을 달성한다 (Llama3: 4×, Qwen3-235B: 16×). 그러나 이는 사전학습 때부터 모델 구조를 바꿔야 하는 아키텍처 수준의 선택이다.

D축 (차원) DeepSeek V2의 MLA (Multi-head Latent Attention)는 키-값에 저차원 분해를 적용해 4H/9 압축을 얻는다. 역시 사전학습 구조 변경이 필요하다.

L축 (레이어 수) 슬라이딩 윈도우 어텐션이나 SSM과 어텐션 레이어를 교차 배치하면 2~8× 압축이 가능하다 (Gemma3: 6×, Jamba: 8×). 마찬가지로 사전학습 구조 선택이다.

T축 (시퀀스 길이) — 실질 공백 가장 직관적인 방향인 T축 압축, 즉 "덜 중요한 토큰의 KV 쌍을 완전히 제거"하는 접근은 광범위하게 채택된 아키텍처적 해법이 아직 없다. DSA 같은 희소 어텐션은 메모리 자체는 줄이지 못하고, 20여 가지 프루닝 방법들은 vLLM·SGLang·TRT-LLM 같은 주요 추론 엔진에 아직 통합되지 못했다.

왜 T축 프루닝이 아직 채택되지 않았는가?

논문은 실용적인 KV 캐시 프루닝이 충족해야 할 4가지 기준을 명시한다:

Criterion 1 빠르고 경량 (Fast & lightweight) — 프루닝 오버헤드가 무시 가능해야 한다.

Criterion 2 단계 무관 (Phase-agnostic) — 프리필링(긴 입력)과 디코딩(긴 출력, 추론 태스크) 모두에서 작동해야 한다.

Criterion 3 최적화 호환 (Optimization-friendly) — FlashAttention2, PagedAttention 같은 커널과 호환되어야 한다.

Criterion 4 충실 (Faithful) — 어떤 태스크에서도 정확도 저하를 최소화해야 한다.

현존 최고 성능 메서드인 KVzip은 Criterion 4를 탁월히 달성하지만, Criterion 1과 2를 충족하지 못해 실제 추론 엔진 통합이 어렵다. 이것이 KVzap이 해결하려는 핵심 갭이다.

이 논문의 선택 — 핵심 아이디어와 트레이드오프

핵심 통찰: 오라클을 근사하라

KVzip은 이미 훌륭하다. 4× 압축에서도 정확도 손실이 미미하고, KVpress 리더보드에서 SOTA다. 문제는 비용이다 — 이중 포워드 패스(입력의 2배 길이 프롬프트)를 필요로 하고, 디코딩 중에는 적용 자체가 불가능하다.

핵심 아이디어 "KVzip 점수를 직접 계산하는 대신, 히든 스테이트(hidden state)에서 그 점수를 예측하는 경량 대리 모델을 학습시키자."

히든 스테이트 h ∈ ℝ^Dh는 어텐션 연산 이전에 이미 존재하며, 행렬 곱 1~2번으로 점수를 뽑을 수 있다. 비용은 전체 레이어 연산의 ~1% 이하.

세 가지 핵심 선택과 근거

선택 1: 입력 적응적 임계값 (Thresholding)

KVzip이 "상위 k%만 남긴다"는 고정 비율 전략을 쓰는 반면, KVzap은 임계값 τ 기반 프루닝을 사용한다. 점수가 τ 미만인 KV 쌍만 버린다. 같은 τ라도 반복이 많은 합성 데이터(RULER)에서는 압축률이 높고, 밀도 높은 실세계 데이터(LongBench)에서는 압축률이 낮게 자동 조정된다. 입력의 정보 밀도에 따라 캐시 크기가 동적으로 결정되는 것이다.

선택 2: 어텐션 이후 프루닝

DMS처럼 어텐션 이전에 토큰을 제거하면 희소 프리필링이 가능하지만 구현이 복잡하다. KVzap은 어텐션 연산 후에 프루닝을 적용해 기존 FlashAttention2, PagedAttention 커널과 완벽히 호환된다. 속도 향상보다 통합 용이성을 우선한 실용주의적 선택이다.

선택 3: 슬라이딩 윈도우 유지 (w=128)

히든 스테이트는 위치 정보를 명시적으로 인코딩하지 않는다. 따라서 최근 128개 토큰을 항상 보존하는 슬라이딩 윈도우(StreamingLLM 방식)를 추가한다. 이를 제거하면(w=0) LCC 정확도가 28%로 붕괴하고, w=128에서 62%로 회복한다. w=512로 늘려도 추가 이득은 없다.

무엇을 포기했는가

트레이드오프 KVzap은 훈련이 필요하다 — KVzip처럼 학습 없이 바로 적용할 수 없다. 모델별로 서로게이트를 학습해야 하며, 1250 토큰 미만 데이터에서 훈련된 모델이 훨씬 긴 시퀀스에 일반화해야 한다는 분포 이동(distribution shift) 위험이 있다. 또한 비균일 캐시 길이(헤드마다 다른 수의 KV 쌍 보존)는 PagedAttention 같은 특수 커널 없이는 실제 메모리·속도 이득으로 이어지기 어렵다.

방법론

전체 아키텍처: KVzap 파이프라인

각 트랜스포머 레이어에서 KVzap은 다음 흐름으로 동작한다. 각 블록을 클릭하면 상세 설명이 표시된다.

3.1 KVzip — 문맥 재구성으로 중요 토큰 식별

KVzip은 "copy-and-paste" 사전학습 태스크를 활용해 KV 중요도를 측정한다. 입력 컨텍스트 user: <prompt>가 주어지면, 다음과 같은 연장 프롬프트를 구성한다:

user: <prompt>
Repeat the previous context exactly.
assistant: <prompt>

그런 다음 각 헤드에서 위치 i의 KV 쌍 점수를 아래와 같이 정의한다:

\[ s_i = \max_{j \in \langle\text{prompt}\rangle} a_{ji} \quad \text{(Eq. 1 — KVzip 점수)} \]

변수	의미	비고
`s_i`	위치 i의 KV 쌍 중요도 점수	값이 낮을수록 제거 후보
`a_{ji}`	반복 <prompt>의 위치 j에서 원본 위치 i로의 어텐션 가중치	softmax 출력, 0~1
`j ∈ <prompt>`	반복 구간 내 쿼리 위치	GQA 사용 시 같은 그룹 헤드 모두 포함

직관적 해설

모델이 원문을 "그대로 반복"할 때, 어떤 토큰을 많이 참조했는지 보면 그 토큰이 얼마나 중요한지 알 수 있다. 반복 과정에서 전혀 참조되지 않은 토큰은 정보를 담고 있지 않으므로 버려도 된다는 발상이다. 이는 인간이 긴 글을 읽고 요약할 때 중요한 단어에 집중하는 것과 유사하다.

수학적 유도 및 한계

이 점수는 어텐션 가중치 자체를 중요도로 사용한다. 그러나 어텐션 가중치 a_{ji}는 현재 쿼리 j와 키 i의 유사도를 나타낼 뿐, 그 KV 쌍이 실제로 출력 히든 스테이트에 얼마나 기여하는지는 반영하지 않는다. 값 벡터 v_i의 크기(||W_O v_i||)가 매우 작다면, 아무리 어텐션 가중치가 높아도 실질적 기여는 미미할 수 있다. KVzip+는 이 점을 보정한다.

3.2 KVzip+ — 값 벡터 크기로 정규화

먼저 트랜스포머 어텐션에서 디코딩 스텝 j의 히든 스테이트 업데이트 공식을 살펴보자:

\[ h_j^{out} = h_j + \sum_{i \leq j} a_{ji} W_O v_i \quad \text{(Eq. 2 — 히든 스테이트 업데이트)} \]

변수	의미	비고
`h_j`	스텝 j의 입력 히든 스테이트	잔차 스트림
`h_j^out`	스텝 j의 출력 히든 스테이트
`W_O`	어텐션 출력 투영 행렬	학습된 파라미터
`v_i`	위치 i의 값(value) 벡터
`a_{ji} W_O v_i`	토큰 i의 잔차 스트림 기여분	이 항이 실질 기여량

직관적 해설

잔차 연결(residual connection)에 의해 출력은 입력에 "어텐션 기여분"을 더한 형태다. 토큰 i의 실질 기여는 a_{ji} × W_O v_i이며, 이 크기가 |h_j|에 비해 작으면 없애도 무방하다. KVzip+는 바로 이 기여의 크기를 정규화 항으로 활용한다.

수학적 유도

a_{ji} W_O v_i 는 잔차 스트림 h_j에 대한 토큰 i의 벡터 기여분이다. 이 기여분의 L2 노름은 a_{ji} × ||W_O v_i||이다 (a_{ji} ≥ 0). 이를 ||h_j||로 정규화하면 상대적 기여 비율을 얻는다. KVzip+는 이 정규화 항을 Eq. 1에 곱한다.

\[ s_i^+ = \max_{j \in \langle\text{prompt}\rangle} a_{ji} \cdot \frac{\|W_O v_i\|}{\|h_j\|} \quad \text{(Eq. 3 — KVzip+ 점수)} \]

변수	의미	비고
`s_i^+`	정규화된 KVzip+ 중요도 점수	Eq. 1의 개선판
`\|\|W_O v_i\|\|`	값 벡터의 출력 투영 후 L2 노름	값의 실질 크기
`\|\|h_j\|\|`	히든 스테이트의 L2 노름	정규화 기준

직관적 해설

KVzip 점수에 "이 토큰의 값이 실제로 얼마나 영향력이 있는가"를 곱한 것이다. 어텐션 가중치가 높더라도 값 벡터가 작으면(||W_O v_i||가 작으면) 점수가 낮아져 버려질 수 있다. 반대로 어텐션 가중치가 낮더라도 값이 크면 보존될 수 있다. 이는 Devoto et al. (2025)의 Expected Attention 분석에서 영감받은 것이다.

수학적 유도

Eq. 2에서 토큰 i의 기여 크기는 a_{ji} × ||W_O v_i||이다. 이를 기준(히든 스테이트 크기 ||h_j||)으로 나누면 상대적 기여율이 된다. max 연산을 통해 반복 구간 내에서 가장 많이 참조될 때의 최대 상대 기여를 점수로 삼는다. 실험에서 KVzip+는 KVzip과 동등하거나 더 나은 성능을 보여 정규화의 유효성이 입증되었다.

3.3 KVzap — 경량 서로게이트로 근사

KVzip+는 여전히 이중 포워드 패스가 필요하다. KVzap은 이를 근본적으로 우회한다: 레이어별로 가벼운 선형 레이어 또는 2-layer MLP를 훈련해, 히든 스테이트 h에서 직접 log(s⁺)를 예측하게 한다.

"The model acts independently at each sequence position t: it maps h_t ∈ ℝ^{D_h} to scores in ℝ^H, where D_h is the hidden dimension and H is the number of KV heads."
모델은 각 시퀀스 위치 t에서 독립적으로 작동한다: h_t ∈ ℝ^{D_h}를 H개의 점수 ℝ^H로 매핑한다.

log-공간을 사용하는 이유: softmax 어텐션의 지수 함수적 특성에 맞추기 위함이다(log(s⁺) 예측 → s⁺ = exp(예측값)).

임계값 기반 프루닝과 슬라이딩 윈도우

KVzap은 점수 기반 이진 결정으로 프루닝한다:

\[ \text{보존} \iff \text{score}_{t,h} \geq \tau \quad \text{또는} \quad t \geq T - w \]

변수	의미	비고
`score_{t,h}`	토큰 t, 헤드 h의 예측 점수	log(s⁺) 공간
`τ`	프루닝 임계값	Qwen3: -4, Llama: -7 등 모델별로 다름
`T`	전체 시퀀스 길이
`w`	슬라이딩 윈도우 크기	기본값 128

직관적 해설

점수가 τ 이상이면 KV 쌍을 보존하고, 미만이면 제거한다. 단, 최근 w개 토큰은 항상 보존한다. 임계값이 높을수록 더 많이 제거(높은 압축률), 낮을수록 덜 제거(낮은 압축률). 같은 τ라도 정보가 조밀한 입력에서는 더 많은 토큰이 높은 점수를 받아 압축률이 낮아지고, 반복이 많은 단순 입력에서는 압축률이 높아진다. 이것이 입력 적응적 특성이다.

슬라이딩 윈도우의 필요성

히든 스테이트 h는 토큰의 절대 위치 정보를 명시적으로 담지 않는다. 따라서 최근 토큰들이 낮은 점수를 받아 제거될 위험이 있다. w=0일 때 LCC 정확도가 28%로 붕괴하는 것이 이 문제를 보여준다. w=128 슬라이딩 윈도우로 62%를 회복하며, w=512로 늘려도 추가 이득이 없다.

τ (임계값) 변화에 따른 압축률 시뮬레이션

임계값 τ = -4 (높을수록 더 많이 제거)

가상의 점수 분포에 대한 시뮬레이션. 실제 압축률은 입력에 따라 달라진다.

Algorithm 1 — KVzap 프루닝 단계별 시연

스텝 1 / 4

def compress(hidden_states, keys, values, kvzap_model, threshold, window=128):
    scores = kvzap_model(hidden_states)              # ← 현재 스텝
    scores[..., -window:] = float("inf")
    indices = torch.where(scores >= threshold)
    return keys[indices], values[indices]

스텝 1: KVzap 모델이 히든 스테이트에서 각 헤드의 중요도 점수를 예측한다. 입력: (T, Dh), 출력: (T, H). 이 연산이 전체 레이어 비용의 ~1% 이하다.

구현 세부사항 및 훈련

훈련 데이터

훈련 데이터로 Nemotron-Pretraining-Dataset-sample을 사용한다. 27k 프롬프트를 Common Crawl, 다국어, 수학, 코드 등 9개 서브셋으로 구성. 길이를 750~1,250 토큰으로 제한해 시퀀스 길이가 어텐션 가중치에 미치는 영향을 최소화한다. 서브셋당 최대 500개 훈련 프롬프트, 5개 검증 프롬프트를 선택해 총 약 2.4k 프롬프트. 프롬프트당 무작위로 500개 토큰을 샘플링해 KV 헤드당 1.2M 훈련 쌍(h, log(s⁺))을 생성하고, 23k는 검증에 사용한다.

평가 훈련 프롬프트가 750~1,250 토큰으로 제한되어 있어, 추론 시 훨씬 긴 시퀀스(수만 토큰)에 대한 분포 이동이 존재한다. 논문 자체도 이를 잠재적 한계로 인정한다. 실험에서 실제로 문제가 발생하지는 않았지만, 긴 시퀀스에서 히든 스테이트의 통계적 특성이 달라질 수 있다.

훈련 방법

KVzap-Linear는 scikit-learn으로 훈련하고, KVzap-MLP는 skorch(PyTorch 래퍼)로 훈련한다. 입력 차원은 모델 히든 크기(Dh=4096 또는 5120)이고 출력 차원은 KV 헤드 수(H=8)다.

설정	Qwen3-8B	Llama-3.1-8B-Instruct	Qwen3-32B
히든 차원 D_h	4096	4096	5120
KV 헤드 수 H	8	8	8
쿼리 헤드 수 H_Q	32	32	64
MLP 숨은층 너비	512 (=4096/8)	512	640 (=5120/8)
MLP 파라미터 수	76M	-	210M
Linear 파라미터 수	-	1.1M	-
최적 임계값 τ	-4	-7	-4

컴퓨팅 오버헤드 (Appendix B)

트랜스포머 레이어의 선형 투영 비용(어텐션 행렬 W_Q, W_K, W_V, W_O 및 FFN)만 고려해 보수적 상한을 계산하면:

\[ C = C_{\text{attn}} + C_{\text{ffn}} = 4D_h(H_Q D + HD) + 6D_h D_{\text{int}} \] \[ C_{\text{KVzap-MLP}} = \frac{D_h}{4}(D_h + H), \quad C_{\text{KVzap-Linear}} = 2D_h H \]

변수	의미	비고
`D_h`	히든 차원	4096 또는 5120
`H_Q`	쿼리 헤드 수	32 또는 64
`H`	KV 헤드 수	8
`D`	헤드 차원	128
`D_int`	FFN 중간 차원	12288~25600

직관적 해설

KVzap-MLP의 FLOPs는 Dh×(Dh/8)×2 + (Dh/8)×H×2 ≈ Dh²/4이다. 전체 레이어 비용 C에서 이 값이 차지하는 비율이 최대 1.09%에 불과하다. KVzap-Linear는 단일 행렬 곱 Dh×H로 0.02% 수준이다.

[Appendix 부연] 모델별 오버헤드

Qwen3-8B (H_Q=32, H=8, D=128, D_h=4096, D_int=12288): MLP 1.09%, Linear 0.02%. Llama-3.1-8B-Instruct (D_int=14336): MLP 0.96%, Linear 0.02%. Qwen3-32B (H_Q=64, D_h=5120, D_int=25600): MLP 0.67%, Linear 0.01%.

결과 — 수치 비교 & 주요 발견

모든 실험은 Qwen3-8B, Llama-3.1-8B-Instruct, Qwen3-32B에서 수행되며 NVIDIA/kvpress로 완전 재현 가능하다. 훈련 모델은 NVIDIA/KVzap 컬렉션에서 공개한다.

2~4×

KV 캐시 압축률

정확도 손실 무시 가능

3.5×

Qwen3-8B 평균 압축

RULER+LongBench+AIME 평균

#1

KVpress 리더보드

15개 방법 대비 SOTA

<1.1%

연산 오버헤드

레이어 선형 투영 대비

4.1 서로게이트 훈련 품질 (R² 점수)

검증 세트에서 H×L개 KV 헤드에 대한 평균 제곱 피어슨 상관계수(R²)를 측정한다. R²가 1에 가까울수록 KVzip+ 점수를 더 정확히 예측한다.

모델	KVzap-Linear (R²)	KVzap-MLP (R²)
Qwen3-8B	0.671	0.711
Llama-3.1-8B-Instruct	0.743	0.772
Qwen3-32B	0.629	0.668

R²가 0.60~0.80 범위에 있다는 것은 완벽한 예측은 아니지만 충분히 유용한 근사임을 의미한다. 모든 모델에서 MLP가 Linear를 능가한다. 한편 Llama-3.1-8B-Instruct에서는 R²가 더 낮음에도 불구하고 KVzap-Linear가 실제 벤치마크에서 KVzip+ 오라클을 능가하는 흥미로운 결과가 나온다(R²가 낮아도 프루닝 결정의 순위 정확도는 높을 수 있기 때문).

4.2 RULER 벤치마크 결과

RULER (Hsieh et al., 2024)는 검색, 다단계 추적, 집계, QA 등 4가지 범주, 13개 서브셋, 4k~128k 시퀀스 길이를 다루는 장문 맥락 벤치마크다.

"On RULER 4k, KVzap achieves state-of-the-art results for both Qwen3-8B and Llama-3.1-8B-Instruct, significantly outperforming 15 concurrent KV cache pruning methods."
RULER 4k에서 KVzap은 Qwen3-8B와 Llama-3.1-8B-Instruct 모두에서 SOTA를 달성했으며, 15개의 동시대 KV 캐시 프루닝 방법들을 크게 능가한다.

※ StreamingLLM, H₂O, Expected Attention, Duo Attention, Compactor 수치는 논문 Figure 1 및 Figure 2의 개략적 수치다. KVzip, KVzip+, KVzap 수치는 논문 Table 2에서 가져옴.

발견 1 KVzip+가 KVzip과 동등하거나 더 좋다 — 정규화 항의 유효성을 입증한다.

발견 2 KVzap은 3~4× 압축까지 완벽한 정확도를 유지한다.

발견 3 Qwen 모델에서는 MLP가 Linear를 능가하지만, Llama-3.1-8B-Instruct에서는 KVzap-Linear가 심지어 자신이 근사하는 KVzip+ 오라클보다 뛰어난 성능을 보인다.

4.3 모델별 최적 설정 및 종합 성능 (Table 2)

항목	Qwen3-8B (MLP)	Llama-3.1-8B (Linear)	Qwen3-32B (MLP)
임계값 τ	-4	-7	-4
파라미터 수	76M	1.1M	210M
RULER 4k	95.32 → 95.09	95.69 → 95.55	95.65 → 95.95
RULER 16k	92.99 → 92.78	93.42 → 93.29	95.19 → 94.96
LongBench	46.74 → 46.49	45.25 → 44.65	50.56 → 50.40
AIME25 pass@4	0.77 → 0.77	미평가	0.83 → 0.87
평균 압축률	0.72 (3.5×)	0.67 (3.0×)	0.63 (2.7×)

화살표(→)는 전체 KV 캐시에서 압축 후로의 변화를 나타낸다. Qwen3-32B에서는 RULER 4k 점수가 오히려 올라가는 경우도 있다(95.65 → 95.95) — 이는 노이즈 범위 내이지만, 적어도 압축이 성능을 저해하지 않음을 보여준다.

4.4 AIME25 — 추론 태스크에서의 성능

AIME25는 2025년 AMC(미국 초청 수학 시험)의 올림피아드급 문제 30개로 구성된다. 문항당 4번의 롤아웃(32k 토큰 생성 한도), pass@1과 pass@4를 평가한다. 디코딩 중 수만 토큰을 생성하는 태스크이므로 KVzip(디코딩 미지원)으로는 불가능한 평가다.

"KVzap-MLP preserves reasoning accuracy even at compression ratios exceeding 2×."
KVzap-MLP는 2× 이상의 압축 비율에서도 추론 정확도를 유지한다.

보충 [Appendix 데이터] AIME25 상세 결과(4 롤아웃 각각의 정답 수, n=30): Full KV(Qwen3-8B)는 18, 20, 21, 21이고, KVzap-MLP τ=-4는 16, 17, 19, 20으로 근접하다. τ=-3에서는 0, 0, 0, 0으로 붕괴하는데 이는 96% 과압축 때문이다.

4.5 LongBench 결과 및 주의 사항

LongBench (Bai et al., 2024)는 단일 문서 QA, 다중 문서 QA, 요약, 퓨샷, 합성, 코드 완성 등 21개 영어/중국어 서브셋을 포함한다. KVzap은 2~3× 압축에서 거의 완벽한 정확도를 유지한다.

한 가지 흥미로운 관찰: Expected Attention이 낮은 압축률에서 전체 KV 기준치를 능가하는 것처럼 보이는데, 이는 TREC 서브셋 이상치(outlier)가 주된 원인이다. TREC를 제외하면 KVzap이 일관되게 우위다. 이는 퓨샷 학습 태스크(TREC)에서 예시를 너무 많이 주면(over-prompting) 오히려 정확도가 낮아지는 현상과 관련될 수 있다.

4.6 에블레이션 (Ablation)

임계값 vs 고정 비율

KVzap은 top-k(고정 비율 선택) 대신 임계값 기반 프루닝을 사용한다. Figure 5는 같은 τ값에서 프롬프트별로 압축률이 최대 20%까지 변동함을 보여준다. 실험에서 임계값 방식이 AdaKV(헤드별) 및 레이어별 top-k 방식 모두를 능가한다.

※ 위 차트의 Top-k 및 AdaKV 수치는 논문 Figure 5 right 기반 개략 수치다.

한계점 & 트레이드오프

한계 1: 범위 및 일반화 현재 결과는 최대 32B 모델까지만 검증되었다. GLM 4.7(130B), Qwen3-235B-A22B, DeepSeek V3.2(희소 어텐션 구조) 같은 대형 모델에 대한 추가 검증이 필요하다. 또한 평가가 RULER, LongBench, AIME25에 집중되어 있어 에이전틱 태스크나 단기 지식 QA에서의 성능은 미검증이다.

LLM scaling sparse attention

한계 2: 훈련 필요성 (post-hoc 추가) KVzap은 훈련이 필요하다. 훈련이 필요 없는 KVzip과 달리 모델별로 서로게이트를 학습해야 한다. 이는 단기 실용성은 높이지만, 장기적으로 엔드-투-엔드 학습(DMS처럼 출력 분포를 직접 보존하도록 훈련)이 더 나은 성능을 낼 가능성을 배제하지 못한다. 논문은 깊은 학습에서 "멀티 토큰 예측이 투기적 디코딩(Medusa)을 대체했듯" 엔드-투-엔드 프루닝이 결국 우세할 수 있다고 인정한다.

end-to-end pruning dynamic memory sparsification

한계 3: 실제 속도/메모리 절감 미실현 현재 논문은 KV 캐시 압축률(제거된 비율)을 측정하지만, 이것이 실제 GPU 메모리 절감이나 추론 속도 향상으로 이어지는지 측정하지 않았다. 헤드별로 보존되는 KV 쌍 수가 다른 비균일 캐시 구조를 효율적으로 처리하려면 PagedAttention 같은 가변 길이 블록 커널이 필요하고, 이는 복잡한 엔지니어링을 요구한다. DMS, Compactor, AdaKV 등 선행 연구가 가능성은 보였지만, 커널 최적화는 여전히 비자명하다.

PagedAttention FlashAttention sparse

한계 4: 훈련 분포 이동 서로게이트 모델은 750~1,250 토큰 길이의 프롬프트에서 훈련되었지만, 추론 시에는 수만 토큰을 다룬다. 히든 스테이트의 통계적 특성이 달라질 수 있는 분포 이동 문제가 잠재적으로 존재한다. 실험에서 실제 문제는 발생하지 않았지만, 논문 자체도 "potential train-test distribution shift"를 인정한다.

length generalization OOD generalization

얻은 것 vs 잃은 것 요약

얻은 것	잃은 것/포기한 것
빠른 추론 (오버헤드 <1.1%)	모델별 훈련 필요
프리필링 + 디코딩 모두 지원	실제 메모리/속도 이득 미측정
입력 적응적 압축	비균일 캐시 구조 → 특수 커널 필요
FlashAttention 호환	짧은 프롬프트 훈련 → 분포 이동
SOTA 정확도 (KVpress 리더보드 1위)	대형 모델(>32B) 미검증

영향력 & 후속 연구

누구에게 도움이 되는가

KVzap은 두 가지 주요 시나리오에서 즉각적 가치를 가진다. 첫째, 수만~수십만 토큰을 처리하는 장문 컨텍스트 추론(RAG, 문서 요약, 법률/의료 문서 분석)에서 GPU 메모리 절감과 처리량 증가를 실현할 수 있다. 둘째, 수천 토큰을 생성하는 추론(reasoning) 태스크(수학 올림피아드, 코드 생성, 복잡한 계획 수립)에서 디코딩 중에도 캐시를 압축할 수 있는 최초의 고성능 방법 중 하나다.

평가 KVzap이 정말 중요한 이유는 단순한 성능 개선이 아니라, 학술 연구와 실제 추론 엔진 사이의 간극을 메울 가능성에 있다. vLLM, SGLang, TRT-LLM에 이런 방법이 통합된 적이 없는 이유는 속도/정확도 트레이드오프 때문이었는데, KVzap은 그 기준을 충족하는 첫 번째 후보다.

저자가 제안한 후속 과제

논문은 다음 방향을 명시적으로 제안한다: (1) 더 나은 데이터 선택과 하이퍼파라미터 튜닝을 통한 R² 개선, (2) 어텐션 이전에 프루닝을 적용해 프리필링 자체를 가속하는 방향, (3) PagedAttention 커널 최적화를 통한 실제 메모리 및 속도 절감 달성, (4) 대형 모델 및 희소 어텐션 아키텍처로의 확장.

Q&A — 연구자의 고민과 독자의 질문

Q1. 왜 히든 스테이트(h)에서 예측하는가? 키-값(k, v)에서 예측하면 안 되는가?

논문 Appendix A에서 명시적으로 실험한 결과, 키-값(k, v)에서 직접 점수를 예측하면 히든 스테이트 h를 사용할 때보다 R²가 항상 낮았다. 히든 스테이트는 레이어의 입력으로서 해당 위치의 토큰에 대한 풍부한 맥락 표현을 담고 있다. 반면 키-값 벡터는 어텐션 특화 투영을 거친 이후의 값으로, 중요도 예측에 덜 정보적일 수 있다.

평가 또한 h는 어텐션 연산 이전에 이미 존재하여 병렬 계산이 가능한 반면, k, v는 어텐션 투영 후에야 얻을 수 있어 약간의 지연이 생긴다.

Q2. DMS와 어떻게 다른가? 결국 같은 접근 아닌가?

두 방법 모두 히든 스테이트에서 중요도를 학습한다는 점에서 유사하지만, 핵심적 차이가 있다:

프루닝 시점: DMS는 어텐션 이전에 토큰을 제거(희소 프리필링), KVzap은 어텐션 이후에 제거(기존 커널 호환)
훈련 목표: DMS는 KL 발산 + 희소성 패널티로 엔드-투-엔드 학습, KVzap은 KVzip+ 점수를 근사하도록 학습(대리 모델 방식)
공개 체크포인트: DMS는 공개 체크포인트가 없어 벤치마크 비교 불가

"We were unable to include DMS in our benchmark comparison due to the lack of publicly available checkpoints."

Q3. Llama에서 KVzap-Linear가 KVzip+ 오라클보다 더 좋은 이유는?

"Surprisingly, KVzap-Linear excels on Llama-3.1-8B-Instruct despite lower R² than KVzap-MLP and even outperforms the KVzip+ oracle it approximates."

평가 논문은 이를 명확히 설명하지 않지만, 가능한 해석은 다음과 같다: R²는 점수의 수치적 정확도를 측정하지만, 프루닝에서 중요한 것은 순위(ranking) 정확도다. KVzap-Linear의 예측이 KVzip+ 점수의 순위 패턴을 더 잘 포착할 수 있다. 또한 KVzip+ 자체가 완벽한 기준점이 아닐 수 있다 — 서로게이트가 때로는 더 나은 프루닝 결정을 학습할 수 있다.

Q4. 이 방법을 실제 추론 엔진(vLLM 등)에 통합할 수 있는가?

원칙적으로는 가능하다. KVzap은 어텐션 연산 후에 적용되며 FlashAttention2와 호환된다. 그러나 핵심 도전은 비균일 캐시 길이다 — 헤드마다 보존된 KV 쌍 수가 다르기 때문에 고정 크기 행렬로 배치 처리하기 어렵다. PagedAttention 같은 가변 길이 블록 시스템이 필요하다. DMS, Compactor, AdaKV가 이 경로를 보였지만, 커널 수준 최적화는 아직 미완성이다.

논문이 "production deployment의 유력 후보"라고 말하는 이유는 설계의 단순함(히든 스테이트 + 행렬 곱)과 높은 압축률-정확도 트레이드오프 때문이다. 실제 통합은 후속 엔지니어링 작업의 과제다.

Q5. 양자화(quantization)와 함께 사용할 수 있는가?

가능하며 논문도 이를 언급한다. KIVI(INT4 KV), ZipCache(INT2 KV) 같은 양자화는 KV 쌍을 낮은 정밀도로 저장해 메모리를 줄이는데, 이는 토큰 프루닝과 직교(orthogonal)하다 — 즉 두 방법을 조합하면 더 큰 압축이 가능하다. "이 접근법들은 추가 압축을 위해 프루닝과 결합될 수 있다"고 논문은 명시한다.

Q6. 코드와 모델은 어디서 구할 수 있는가?

코드: github.com/NVIDIA/kvpress (20개 이상의 KV 캐시 프루닝 방법 구현 포함)
훈련된 모델: NVIDIA/KVzap HuggingFace 컬렉션
평가 로그: KVzap predictions (논문에 링크)

원본 Figures — 논문 원문 그대로

Figure 1 — KVpress 리더보드 (논문 §1, p.1)

Figure 1: KVpress Leaderboard — Figure 1: KVpress 리더보드 — Qwen3-8B(좌)와 Llama-3.1-8B-Instruct(우). x축: KV 캐시 압축 비율, y축: RULER 4k 정확도.

주목할 포인트: KVzap은 KVzip(자신이 근사하는 오라클)과 거의 동일한 곡선을 그리며 15개 방법 모두를 능가한다. Expected Attention은 낮은 압축률에서 경쟁력 있지만 높은 압축률에서 급격히 하락한다.

Figure 2 — RULER 4k 확대 뷰 (논문 §4.4, p.5-6)

Figure 2: RULER 4k results — Figure 2: RULER 4k 결과 (확대 뷰, y축 범위 [90, 100]). Qwen3-8B(좌), Llama-3.1-8B-Instruct(중), Qwen3-32B(우).

주목할 포인트: y축이 [90, 100]으로 확대된 뷰다. KVzap은 3~4× 압축까지 완벽한 정확도를 유지한다. Llama에서 KVzap-Linear가 KVzip+ 오라클보다 뛰어난 것이 이 그래프에서 명확히 보인다.

Figure 3 — LongBench 결과 (논문 §4.5, p.6)

Figure 3: LongBench results — Figure 3: LongBench 결과. Qwen3-8B(좌), Llama-3.1-8B-Instruct(중), Qwen3-32B(우).

주목할 포인트: Expected Attention이 Qwen3-8B에서 기준치를 넘어서 보이는데, 이는 TREC 서브셋 이상치 때문이다. TREC 제외 시 KVzap이 일관되게 우위다(Figure 12 참조). LongBench는 실세계 데이터로 합성 데이터(RULER)보다 정보 밀도가 높아 같은 τ에서 더 낮은 압축률이 나온다.

Figure 4 — AIME25 추론 성능 (논문 §4.6, p.7)

Figure 4: AIME25 reasoning — Figure 4: AIME25 추론 성능. pass@1(실선)과 pass@4(점선). Qwen3-8B(좌), Qwen3-32B(우).

주목할 포인트: 32k 토큰을 생성하는 추론 태스크에서도 KVzap-MLP는 2× 이상 압축에서 정확도를 유지한다. 이는 KVzip이 아예 지원하지 못하는 시나리오다. τ=-3에서 극단적 과압축(96%)이 발생해 점수가 0으로 붕괴하는 것도 주목할만하다.

Figure 5 — 입력 적응적 압축 및 에블레이션 (논문 §4.8, p.7)

Figure 5: Adaptive compression — Figure 5: 압축률 분포(좌)와 대안 방법 비교(우). Qwen3-8B, KVzap-MLP.

주목할 포인트: 왼쪽 그래프는 같은 τ=-4에서 RULER(합성, 반복적), LongBench(실세계), AIME25(추론)에서 압축률이 최대 20%까지 차이남을 보인다. 이것이 "입력 적응적" 특성이다. 오른쪽 그래프는 임계값 방식이 top-k(고정 비율) 방식과 AdaKV 모두를 능가함을 보인다.

KVzap: Fast, Adaptive, and Faithful KV Cache Pruning

문제의 배경 — KV 캐시라는 병목

기존 접근법과 그 한계

왜 T축 프루닝이 아직 채택되지 않았는가?

이 논문의 선택 — 핵심 아이디어와 트레이드오프

핵심 통찰: 오라클을 근사하라

세 가지 핵심 선택과 근거

선택 1: 입력 적응적 임계값 (Thresholding)

선택 2: 어텐션 이후 프루닝

선택 3: 슬라이딩 윈도우 유지 (w=128)

무엇을 포기했는가

방법론

전체 아키텍처: KVzap 파이프라인

3.1 KVzip — 문맥 재구성으로 중요 토큰 식별

3.2 KVzip+ — 값 벡터 크기로 정규화

3.3 KVzap — 경량 서로게이트로 근사

임계값 기반 프루닝과 슬라이딩 윈도우

τ (임계값) 변화에 따른 압축률 시뮬레이션

Algorithm 1 — KVzap 프루닝 단계별 시연

구현 세부사항 및 훈련

훈련 데이터

훈련 방법

컴퓨팅 오버헤드 (Appendix B)

결과 — 수치 비교 & 주요 발견

4.1 서로게이트 훈련 품질 (R² 점수)

4.2 RULER 벤치마크 결과

4.3 모델별 최적 설정 및 종합 성능 (Table 2)

4.4 AIME25 — 추론 태스크에서의 성능

4.5 LongBench 결과 및 주의 사항

4.6 에블레이션 (Ablation)

임계값 vs 고정 비율

한계점 & 트레이드오프

얻은 것 vs 잃은 것 요약

영향력 & 후속 연구

누구에게 도움이 되는가

저자가 제안한 후속 과제

관련 분야 더 탐색하기

Q&A — 연구자의 고민과 독자의 질문

원본 Figures — 논문 원문 그대로

Figure 1 — KVpress 리더보드 (논문 §1, p.1)

Figure 2 — RULER 4k 확대 뷰 (논문 §4.4, p.5-6)

Figure 3 — LongBench 결과 (논문 §4.5, p.6)

Figure 4 — AIME25 추론 성능 (논문 §4.6, p.7)

Figure 5 — 입력 적응적 압축 및 에블레이션 (논문 §4.8, p.7)