KV 캐시가 터지지 않으려면 — 압축 기법들이 실제로 하는 일

컨텍스트 길이가 늘어날수록 KV 캐시가 메모리를 잡아먹는다는 사실은 이제 어느 정도 알려져 있다. 그런데 그 문제를 해결하겠다고 나온 기법들이 실제로 어떤 원리로 동작하는지는 의외로 피상적으로만 알려진 경우가 많다. PagedAttention이 “페이지 테이블에서 착안했다”는 설명은 들어봤어도, 정확히 어느 지점에서 어떤 이득이 생기는지, GQA가 MHA와 구체적으로 무엇이 다른지, H2O가 어떤 토큰을 어떤 기준으로 버리는지는 실제 논문이나 구현 코드를 들여다본 적 없으면 모호하게 남기 쉽다.

왜 압축이 필요한지부터 — 숫자로 다시 한번

먼저 문제 규모를 수식으로 정확히 잡아두자. KV 캐시가 차지하는 바이트 수는 다음 곱으로 결정된다.

KV Cache 크기 = 2 × L × H_kv × D_head × S × B × dtype_bytes

여기서 2는 Key와 Value 두 벌, L은 레이어 수, H_kv는 KV 헤드 수, D_head는 헤드 차원, S는 시퀀스 길이, B는 배치 크기, dtype_bytes는 데이터 타입(FP16 = 2바이트, INT8 = 1바이트)이다.

Llama 3 8B를 예로 들면 L=32, H_kv=8(GQA 적용 후), D_head=128, dtype=FP16(2바이트)이다. 배치 크기 1, 시퀀스 4K(4,096)에서는 2 × 32 × 8 × 128 × 4,096 × 1 × 2 = 약 536MB다. 이 시퀀스를 32K로 늘리면 8배인 약 4.3GB, 128K라면 다시 4배인 약 17GB가 된다. 모델 가중치 자체가 FP16 기준 약 16GB이니, 128K 컨텍스트에서는 가중치보다 캐시가 더 많은 메모리를 차지하는 역전 현상이 생긴다. API 서버처럼 배치 크기가 커지면 이 수치는 그 배수만큼 다시 곱해진다.

이 글에서 다루는 기법들은 크게 세 범주로 나뉜다. 첫째는 애초에 캐시를 작게 설계하는 아키텍처 수준의 접근(MQA, GQA), 둘째는 이미 생성된 캐시를 물리 메모리에서 효율적으로 배치하는 시스템 수준 접근(PagedAttention), 셋째는 캐시를 동적으로 선별·압축해 유지하는 런타임 접근(StreamingLLM, H2O, KV 양자화)이다. 각각 해결하는 문제가 다르고 트레이드오프도 다르다.

헤드를 줄이면 캐시가 줄어든다 — MQA와 GQA의 원리

트랜스포머(Transformer)의 원래 멀티헤드 어텐션(Multi-Head Attention, MHA)은 Query, Key, Value 각각을 H개의 헤드로 병렬 처리한다. 헤드가 많을수록 모델이 여러 “관점”에서 어텐션을 볼 수 있어 표현력은 올라가지만, KV 캐시 크기도 헤드 수에 비례해 늘어난다.

2019년 Google이 제안한 Multi-Query Attention(MQA)는 단순하지만 효과적인 아이디어다. Query 헤드는 H개 그대로 유지하되, Key와 Value는 딱 1개 헤드만 쓰고 모든 Query 헤드가 이 단일 KV 쌍을 공유한다. KV 캐시 크기가 이론상 H분의 1로 줄어드는 셈이다. Llama 계열에서 쓰는 H=32 기준이라면 캐시가 32분의 1이 된다. 문제는 품질 저하다. Key-Value를 공유하면 각 헤드가 서로 다른 정보를 포착하는 특성이 희석돼 긴 문서나 복잡한 추론 작업에서 성능이 눈에 띄게 떨어진다는 게 경험적으로 확인됐다.

Grouped Query Attention(GQA)은 이 두 극단 사이의 절충이다. 2023년 Google Research 논문에서 제안됐으며, Query 헤드를 G개 그룹으로 나누고 각 그룹 내의 Query들이 하나의 KV 쌍을 공유한다. G=1이면 MQA, G=H이면 원래 MHA가 된다. Llama 3 8B는 Query 헤드 32개, KV 헤드 8개를 쓰는데, 이는 G=4인 GQA 설정이다. 메모리는 MHA 대비 4분의 1이면서, MQA처럼 단 하나의 KV를 공유하는 것보다 훨씬 나은 품질을 보인다. Llama 3, Mistral 7B, Gemma 2, Falcon 등 현재 주요 오픈소스 모델들이 대부분 GQA를 채택한 이유다.

한 단계 더 나아간 방식이 DeepSeek V2부터 도입한 Multi-head Latent Attention(MLA)이다. MLA는 KV 헤드 수를 줄이는 대신, Key-Value 자체를 저차원 잠재 벡터(latent vector)로 압축해 캐시한다. 실제 어텐션 계산 시에는 이 잠재 벡터를 다시 원래 차원으로 복원(up-projection)한다. 압축 비율이 GQA보다 공격적이어서 DeepSeek V2의 경우 KV 캐시가 표준 MHA 대비 약 93.3% 감소했다고 보고됐다. 추론 시 복원 연산이 추가되지만, 메모리 대역폭 절감이 이 오버헤드를 상쇄한다.

가상 메모리에서 착안한 — PagedAttention이 단편화를 없애는 방식

GQA나 MLA가 모델 설계 단계에서 캐시를 작게 만드는 방식이라면, PagedAttention은 이미 생성된 캐시를 GPU 메모리에 어떻게 배치할 것인가의 문제다.

기존 추론 프레임워크에서 KV 캐시는 요청별로 연속된 메모리 블록을 미리 예약(pre-allocation)했다. 최대 시퀀스 길이만큼 공간을 잡아두는 방식인데, 여기서 두 가지 낭비가 생긴다. 첫째는 내부 단편화(internal fragmentation)다. 실제 생성 길이가 최대치에 못 미치면 예약만 해놓고 쓰지 않는 공간이 생긴다. 둘째는 외부 단편화(external fragmentation)다. 여러 요청이 서로 다른 길이로 완료되고 나면 메모리 여기저기에 작은 빈 공간이 흩어지는데, 각각은 새 요청에 쓰기 부족해도 합치면 충분할 수 있는 상황이 된다. 연구에 따르면 이런 단편화로 인한 낭비가 실제 KV 캐시 메모리의 60~80%에 달하기도 했다.

2023년 UC Berkeley의 vLLM 프로젝트가 제안한 PagedAttention은 운영체제(OS)의 가상 메모리 페이지 테이블 개념을 그대로 가져온다. KV 캐시를 고정 크기의 “KV 블록”(논문 설정에서는 16토큰 단위)으로 나누고, 각 블록이 GPU 메모리 어느 위치에 있는지를 블록 테이블(block table)로 관리한다. 시퀀스가 반드시 연속된 물리 메모리를 점유하지 않아도 되고, 블록 테이블이 논리적 시퀀스 순서와 물리적 메모리 위치 사이의 매핑을 담당한다. 덕분에 요청이 끝나면 해당 블록들이 즉시 반환돼 재사용된다. 단편화가 원리적으로 사라지는 것이다.

여기에 Copy-on-Write(쓰기 시 복사) 방식의 블록 공유가 추가된다. 병렬 샘플링(parallel sampling)이나 빔 서치(beam search)처럼 동일한 프리픽스(prefix)를 공유하는 여러 출력을 생성할 때, 프리픽스에 해당하는 KV 블록을 복사하지 않고 참조만 한다. 실제로 다른 내용이 생성되는 시점에 비로소 해당 블록을 복사해 독립된 공간을 만든다. 프리픽스 공유만으로도 배치 요청에서 메모리를 수십 퍼센트 절약할 수 있고, 시스템 처리량(throughput)이 기존 HuggingFace Transformers 대비 최대 24배 증가했다는 게 vLLM 논문의 보고다.

토큰을 선별해 버린다 — H2O와 StreamingLLM의 다른 전략

위의 방식들이 캐시를 작게 만들거나 효율적으로 배치하는 데 집중했다면, H2O와 StreamingLLM은 캐시에 담긴 토큰 자체를 동적으로 선별해 중요한 것만 남기는 접근이다.

H2O(Heavy Hitter Oracle)는 2023년 Meta AI 연구팀이 제안한 KV 캐시 eviction(퇴출) 알고리즘이다. 핵심 관찰은 어텐션 점수가 특정 토큰에 집중(heavy hitter)되는 현상이 실험적으로 일관되게 나타난다는 것이다. 전체 토큰 중 소수가 전체 어텐션 점수의 대부분을 차지한다. 이 관찰을 바탕으로 H2O는 각 토큰의 “누적 어텐션 점수”를 중요도 지표로 삼아 KV 캐시를 동적으로 관리한다.

구체적으로는 미리 정해둔 캐시 버짓(budget)에서 Heavy Hitter 토큰과 최근 토큰(Recent Token)을 각각 일정 비율로 유지한다. 최근 토큰을 무조건 남기는 이유는, 어텐션 메커니즘 특성상 직전 토큰들과의 관계가 항상 중요하기 때문이다. 캐시가 버짓을 초과하면 누적 어텐션 점수가 가장 낮은 토큰을 추방(evict)한다. 이론적으로 캐시를 원하는 크기로 고정할 수 있어 시퀀스가 아무리 길어져도 메모리가 선형으로 증가하지 않는다. 실험 결과 전체 KV 캐시의 약 20%만 유지해도 OPT, LLaMA, GPT-NeoX 등에서 full cache 대비 성능 저하가 미미하다고 보고됐다.

StreamingLLM은 MIT CSAIL이 2023년 제안한 방식으로, 문제 설정이 조금 다르다. 무한하게 이어지는 대화나 스트리밍 텍스트를 처리할 때, 슬라이딩 윈도우(sliding window)로 최근 N개 토큰만 KV 캐시에 유지하는 것이 자연스러운 발상이다. 그런데 단순 슬라이딩 윈도우를 적용하면 초기 토큰들이 캐시에서 밀려나는 순간 성능이 급격히 떨어지는 현상이 관찰됐다. 왜일까.

연구팀이 발견한 것은 “어텐션 싱크(attention sink)” 현상이다. 첫 번째 토큰, 즉 위치 0의 토큰은 실제 의미적으로 중요하지 않아도 전체 어텐션 점수에서 불균형적으로 큰 비중을 차지한다. 소프트맥스(Softmax) 기반 어텐션은 점수를 반드시 어딘가에 분배해야 하는데, 모델이 학습 과정에서 “어디에도 집중할 필요 없을 때”의 여분 확률을 초기 토큰으로 몰아주는 패턴을 학습한 것으로 해석된다. 이 초기 토큰들이 캐시에서 빠지면 어텐션 계산의 균형이 깨지면서 출력이 난무하게 된다.

StreamingLLM의 해법은 간단하다. 슬라이딩 윈도우를 그대로 쓰되, 초기 몇 개(보통 4개) 토큰의 KV 캐시를 항상 “고정(anchor)”해 두는 것이다. 최근 토큰들은 윈도우 크기에 따라 교체되지만, 이 초기 토큰들은 절대 퇴출되지 않는다. 덕분에 이론상 무한 길이의 스트리밍 텍스트를 거의 일정한 메모리에서 처리할 수 있다. 단, 윈도우를 벗어난 중간 토큰들의 정보가 유실된다는 한계는 있다. “긴 컨텍스트의 완전한 이해”가 목표가 아니라 “스트리밍 상황에서 일관된 출력 유지”가 목표인 시나리오에 적합하다.

서버 메모리 최적화 — 데이터센터 인프라

이미지 출처: Unsplash

비트를 줄인다 — KV 캐시 양자화

모델 가중치 양자화(quantization)는 이미 로컬 LLM 사용자에게 익숙한 개념이다. FP16(16비트 부동소수점) 가중치를 INT8(8비트 정수)이나 INT4(4비트)로 줄여 메모리를 절약하는 방식이다. 같은 원리를 KV 캐시에 적용하면 어떻게 될까.

KV 캐시 양자화는 캐시에 저장되는 Key와 Value 텐서를 더 낮은 비트폭으로 저장하는 방식이다. FP16 대신 INT8을 쓰면 캐시 크기가 절반, INT4라면 4분의 1이 된다. 언뜻 단순해 보이지만, 가중치 양자화와는 다른 도전이 있다. 가중치는 학습 후 고정값이라 통계적 분포가 안정적이다. 반면 KV 캐시는 매 추론마다 입력에 따라 달라지는 동적인 값이고, 특히 Key 텐서는 이상치(outlier)가 많이 등장하는 것으로 알려져 있다. 이상치를 낮은 비트폭으로 표현하면 정보 손실이 커진다.

이 문제를 해결하기 위해 2023년 이후 여러 기법이 제안됐다. KIVI(Key-Value cache Quantization without Information Loss)는 Key와 Value를 서로 다른 전략으로 양자화한다. 이상치가 많은 Key는 채널(channel) 단위로 양자화해 이상치를 각 채널 범위 내에 가두고, Value는 토큰 단위로 처리한다. KVQuant는 더 세밀하게 각 KV 텐서의 분포를 분석해 비균일(non-uniform) 양자화 격자를 적용한다. FlexGen 같은 오프로딩 시스템에서는 INT4 KV 캐시를 CPU 메모리나 디스크로 스왑하는 방식으로, 더 적은 GPU 메모리로 훨씬 긴 컨텍스트를 처리할 수 있게 한다.

정밀도 손실 관점에서 보면, INT8 KV 캐시는 대부분의 벤치마크에서 FP16 대비 1% 내외의 성능 저하를 보인다. INT4는 태스크에 따라 2~5% 수준의 저하가 나타나며, 특히 복잡한 추론이나 긴 문맥 이해가 필요한 작업에서는 체감이 될 수 있다.

기법별 효과 정리

주요 기법들의 메모리 절감 효과를 수치로 정리하면 다음과 같다.

기법 절감 방식 대략적 KV 캐시 절감 주요 트레이드오프
MQA KV 헤드 1개로 통일 ~1/H (H=헤드 수) 품질 저하, 특히 복잡 추론
GQA (G그룹) KV 헤드를 G개로 ~G/H 수준으로 절감 품질 거의 유지 (실용적 절충)
MLA (DeepSeek) 잠재 벡터 압축 저장 ~93% 절감 보고 복원 연산 오버헤드
PagedAttention 블록 단편화 제거 낭비 60~80% 제거 추가 블록 테이블 관리 비용
H2O 하위 어텐션 토큰 eviction 80% 토큰 제거 가능 유실 토큰 정보 복원 불가
StreamingLLM 윈도우+어텐션 싱크 고정 고정 메모리 유지 윈도우 밖 정보 손실
INT8 양자화 KV 텐서 비트 절감 50% (FP16→INT8) 정밀도 ~1% 손실
INT4 양자화 KV 텐서 비트 절감 75% (FP16→INT4) 정밀도 2~5% 손실

시나리오별로 무엇을 선택할 것인가

이 기법들은 서로 배타적이지 않다. 실제 프로덕션 시스템에서는 여러 기법이 중첩 적용된다. vLLM은 PagedAttention 위에 GQA, FP8/INT8 KV 양자화를 함께 지원한다. 다만 시나리오에 따라 우선순위는 달라진다.

API 서버나 클라우드 서빙처럼 높은 동시성(concurrency)과 배치 크기가 중요한 환경에서는 PagedAttention이 사실상 필수다. 단편화 문제가 워낙 커서, 이를 해결하는 것만으로도 같은 GPU에서 처리할 수 있는 동시 요청 수가 몇 배 늘어난다. GQA가 적용된 모델을 쓰는 것도 서버 환경에서는 기본 조건에 가깝다. 여기에 INT8 KV 양자화를 추가하면 품질 손실을 거의 감수하지 않고도 throughput을 추가로 높일 수 있다.

로컬 실행 환경, 특히 VRAM이 제한된 소비자용 GPU나 통합 메모리를 쓰는 맥(Mac) 환경에서는 선택지가 좁아진다. GQA가 이미 적용된 모델을 선택하는 것이 출발점이다. INT4 KV 양자화 지원 여부가 모델 런타임 선택 기준이 될 수 있다. llama.cpp는 캐시 타입을 --cache-type-k, --cache-type-v 인자로 조정할 수 있고, f16/q8_0/q4_0 등을 지원한다. 매우 긴 문서 처리가 목적이 아니라 긴 대화 유지가 목적이라면, StreamingLLM의 방식을 런타임 수준에서 구현한 옵션(일부 추론 프레임워크에서 sliding window KV cache로 제공)을 쓰는 것도 선택지다.

요약 또는 RAG(Retrieval-Augmented Generation) 파이프라인처럼 긴 문서 전체를 완전히 이해해야 하는 작업에서는 H2O나 StreamingLLM처럼 토큰을 선별·유실하는 방식은 피해야 한다. 이런 경우에는 GQA + PagedAttention + INT8 양자화 조합으로 메모리 효율을 높이되, 모든 토큰의 KV를 보존하는 방식을 유지해야 한다.

현재 연구의 방향은 이 기법들을 더 정교하게 조합하는 쪽으로 향하고 있다. 어떤 토큰을 버릴지를 사전에 고정된 규칙이 아니라 레이어별·태스크별로 학습 가능한 방식으로 결정하려는 시도(Adaptive KV compression)나, 캐시 중요도를 토큰 레벨이 아니라 어텐션 헤드 레벨로 따로 추적하는 연구들이 이어지고 있다. MLA처럼 압축과 표현력을 동시에 확보하려는 방향도 계속 발전하고 있어, 몇 년 안에 “KV 캐시 크기는 시퀀스 길이에 비례한다”는 전제가 더 이상 당연하지 않은 시대가 올 가능성이 크다. 다만 지금 당장 모델을 배포하거나 로컬에서 돌리는 사람에게는 위에서 정리한 기법들의 조합이 현실적인 선택지이고, 어느 기법이 자신의 시나리오에 맞는지 이해하는 것이 막연히 “최신 모델을 쓰면 되겠지” 하는 것보다 훨씬 실질적인 차이를 만든다.


출처

댓글 남기기