CompKV는 보상 잔차를 기준으로 KV 블록을 고르는 스파스 어텐션이다.

무엇인가

긴 문맥 추론에서 LLM은 KV 캐시 메모리 트래픽에 병목을 겪는다. 자기회귀 디코딩에서 KV 캐시는 문맥 길이에 따라 커지고, dense attention은 매 스텝 전체 캐시를 읽는다. 쿼리 인지 스파스 어텐션은 현재 쿼리에 대해 선택된 소수 토큰만 읽어 트래픽을 줄이고, 최근 방법들은 토큰을 블록으로 나눈 뒤 선택되지 않은 블록의 어텐션 기여를 요약 통계로 보상(compensation)한다. 문제는 선택과 재구성이 분리되어 있다는 점이다. 기존 방법은 예측된 관련도나 어텐션 질량으로 블록 순위를 매긴 뒤에야 보상을 고려한다. 질량이 큰 블록은 재구성이 쉬울 수 있고, 질량은 작아도 블록 내 로짓 변동이 큰 블록은 더 큰 오차를 남길 수 있다. 논문은 이 불일치를 지적하며, 생략되었을 때 보상 오차를 가장 크게 남기는 블록을 우선해야 한다고 주장한다.

어떻게 동작하나

이론은 Mean compensation, 즉 선택되지 않은 블록의 모든 로짓을 블록 평균으로 대체하는 모델 위에서 전개된다. 압축 요약이 정확한 log-partition 항등식을 주어 선택과 보상의 상호작용을 해석적으로 다룰 수 있기 때문이다. GQA로 하나의 KV 헤드를 공유하는 G개 쿼리 헤드를 두고, 선택 손실을 보상된 분포와 원래 분포 사이 KL divergence의 합으로 정의한 뒤 K개 블록 예산 아래 최소화한다. KL을 전개하면 각 블록의 기여가 블록 어텐션 질량 p_g,b와 블록 내 로짓 분산 σ²_g,b의 곱으로 근사된다. 블록 평균 로짓 주변 2차 테일러 전개로 p_g,b(1 - |b|e^{z̄}/Z) = p_g,b(1/2 σ² + O(σ³))를 얻고, 결국 Σ_b Σ_g p_g,b σ²_g,b를 최대화하는 K개 블록을 고르는 것이 2차 목표를 최소화한다는 결론에 도달한다. 로짓이 거의 일정한 블록은 질량이 커도 손실에 거의 기여하지 않는다. 부록에서는 Mean compensation을 목표에서 제거하면 단일 쿼리 헤드에 대한 표준 Top-K 어텐션 질량 선택 규칙이 복원된다는 것도 증명한다.

무엇과 다른가

실전 구현에서 정확한 점수를 계산하려면 모든 후보 블록의 키를 읽어야 하므로, CompKV는 블록별로 저장한 평균 키와 그룹화된 키 분산으로 점수를 추정한다. 키 좌표 d개를 r개 그룹으로 나누고 각 그룹의 평균 좌표 분산을 담은 벡터 D̄_b를 저장한다. 로짓 분산 추정은 σ̂²_g,b = (1/d) Σ_t (Σ_{i∈G_t} q²_{g,i}) (D̄_b)_t 이고, 어텐션 질량 추정은 p̂_g,b = |b| exp(z̄_g,b + ½σ̂²_g,b) / Σ_c |c| exp(z̄_g,c + ½σ̂²_g,c) 이다. 최종 점수는 S_b = Σ_g p̂_g,b σ̂²_g,b. 블록 요약은 {|b|, k̄_b, D̄_b, v̄_b}로 블록당 2d + r + 1개 스칼라다. 선택 시 첫 sink 블록과 최근 두 블록은 필수 집합 F로 항상 포함되고 예산에 포함되며, 나머지는 점수로 Top-K를 뽑는다. 선택 블록은 원본 KV를, 나머지는 평균 요약을 쓰고 두 기여를 공동 정규화해 출력을 만든다.

어떻게 쓰나

구현 스케줄도 논문의 기여다. CompKV는 전체 BF16 KV 캐시와 FP32 블록 요약을 pinned CPU 메모리에 둔다. 매 스텝 융합 GPU 커널이 현재 블록 통계를 갱신하고 새 K/V를 패킹한 뒤 D2H 전송 한 번과 CPU writeback이 이어진다. 메인 CUDA 스트림은 선택 영역 전송과 블록 통계·정규화 Top-K 스코어링 두 커널을 실행하고, 보조 스트림은 평균값을 독립적으로 전송한다. 선택 후 CPU는 선택된 K/V를 연속 pinned 버퍼로 모아 H2D 전송 한 번과 FlashInfer 정확 어텐션을 수행하고, 보조 스트림은 평균값과 선택 마스크가 준비되는 대로 Mean compensation을 계산해 CPU gather, 선택 KV 전송, 정확 어텐션과 겹친다. 메인 스트림은 공유 정규화로 두 기여를 합치기 직전에 보상을 기다린다.

전제와 한계

정확도 실험은 Llama-3.1-8B-Instruct, Qwen3-8B, Qwen3-32B 세 모델에서 Quest, InfLLM, Quest+RESA(λ=0.25), Full attention과 비교했다. RULER는 32K, LongBench-Pro는 최대 128K, 블록 크기 16, 기본 r=4 설정이다. CompKV는 RULER 32K·512토큰 예산에서 세 모델 모두 최고 AVG를, LongBench-Pro에서도 최고 평균을 기록했고 대부분의 태스크 수준 비교에서 최고 또는 공동 최고였다. 특히 MK3에서 격차가 두드러졌다. 다만 Llama-3.1에서 Full attention조차 CWE에서 9.9%에 그쳤고 모든 스파스 방법이 낮은 점수에 머물렀다. 128~1024 토큰 예산 스윕에서도 CompKV가 대부분의 예산·태스크·모델에서 최고를 유지했다.

효율은 H100 80GB, 배치 크기 1에서 CPU 오프로드 단일 레이어 파이프라인으로 측정했다. 32K/64K/128K 문맥과 512/1024/2048 토큰 예산의 9개 조합 모두에서 CompKV(r=4)가 최저 평균 지연을 기록했고 Full attention 대비 최대 6.85배 속도를 냈다. Ablation에서 Mean compensation 추가는 AVG를 0.48~2.47점 올렸고, 외부 분산 인자는 추가로 0.35~1.19점, 2차 질량 보정은 1차 대비 1.96~3.39점을 더했다. 분산 그룹 세분화 실험에서 r=128이 두 예산 모두 최고 AVG였고 기본값 r=4는 256·512 예산에서 각각 1.87점, 1.08점 뒤졌지만 분산 메타데이터는 1/32만 쓴다.

개발자에게 이 논문의 실질적 의미는 CPU 오프로드 KV 캐시를 쓰는 긴 문맥 서빙에서 블록 선택 기준만 바꿔 정확도를 올릴 수 있다는 점이다. 블록당 2d+r+1 스칼라 요약만 유지하면 되고, r로 정확도와 메타데이터 비용을 조절할 수 있다. 확인해야 할 전제는 두 가지다. 분석이 Mean compensation과 작은 분산 영역을 가정한 2차 테일러 근사에 기대고 있고, CWE처럼 dense 모델조차 낮은 점수를 내는 태스크는 스파스 방법으로 해결되지 않는다. 논문은 보상 방식 자체가 Mean compensation에 국한되지 않는다고 밝히지만, 해석적으로 다룰 수 있는 근거는 Mean compensation의 log-partition 항등식에서 나온다.