KV 캐시 완전 정리, 속도의 토대이자 메모리 병목
요약
- KV 캐시는 자기회귀 생성에서 이전 토큰들의 Key와 Value를 저장하고 재사용해, 매 토큰마다 전체 시퀀스를 다시 계산하는 낭비를 없앱니다.
- 두 얼굴을 가집니다. LLM 추론 속도의 토대이면서, 동시에 긴 컨텍스트에서 메모리 병목의 주범입니다.
- KV 캐시로 얻는 속도에는 대가가 있습니다. 시간과 메모리의 트레이드오프입니다. 계산을 아끼는 대신 메모리를 많이 먹습니다.
- 생성은 성격이 전혀 다른 두 단계로 나뉩니다. prefill은 연산에 묶이고(compute-bound), decode는 메모리 대역폭에 묶입니다(memory-bound). 이 구분이 최적화 기법 선택을 좌우합니다.
- 캐시 크기는 길이와 배치에 선형으로 커집니다. 긴 컨텍스트나 대량 동시요청에서는 KV 캐시가 모델 가중치보다 큰 메모리를 차지하기도 합니다.
- 그래서 추론 최적화 기법 상당수가 결국 KV 캐시 축소로 수렴합니다. GQA와 MLA, KV 양자화, PagedAttention, 프롬프트 캐싱이 전부 이 하나를 겨냥합니다.
1장. KV 캐시가 무엇인가
1.1 없으면 무슨 일이 벌어지나
자기회귀 생성은 토큰을 하나씩 만듭니다. 그런데 어텐션은 지금까지의 모든 토큰을 봐야 합니다.
캐시가 없으면
토큰 1 생성 토큰 1개의 K, V 계산
토큰 2 생성 토큰 2개의 K, V 계산 ← 1번을 다시 계산
토큰 3 생성 토큰 3개의 K, V 계산 ← 1, 2번을 다시 계산
...
토큰 n 생성 토큰 n개의 K, V 계산
누적 비용이 길이의 제곱으로 폭증한다같은 값을 계속 다시 만드는 게 문제입니다. 앞 토큰들의 K와 V는 뒤 토큰이 생성돼도 바뀌지 않기 때문입니다.
1.2 그래서 저장한다
이미 계산한 각 토큰의 K, V 벡터를 메모리에 저장해 두고
다음 토큰을 만들 때 재계산 없이 재사용한다
누적 비용이 제곱에서 선형으로 내려간다1.3 왜 Q는 캐시하지 않나
자주 나오는 질문이라 짚어 둡니다.
K, V 과거 토큰의 것이 계속 재사용된다
한 번 계산하면 끝까지 유효하다 ──► 캐시한다
Q 매 스텝 "지금 생성 중인 토큰"의 것만 필요하다
다음 스텝에서는 쓸 일이 없다 ──► 캐시할 이유가 없다어텐션의 비대칭성입니다. 새 토큰이 과거를 보는 구조라, 과거 쪽만 남기면 됩니다.
1.4 메모 비유
책을 한 줄씩 쓰는데
매번 처음부터 다시 읽지 않는다
지금까지 읽은 핵심(K, V)을 메모해 두고
새 줄만 참고해 이어 쓴다2장. prefill과 decode
2.1 두 단계로 나뉜다
prefill
프롬프트 전체를 한 번에 처리하며 K, V 캐시를 채운다
토큰이 많아 병렬 처리가 된다
compute-bound
decode
토큰을 하나씩 생성한다
새 토큰의 Q 만 계산해 캐시된 K, V 와 어텐션하고
새 K, V 를 캐시에 추가한다
memory-bound2.2 동작을 단계로
① prefill
프롬프트 토큰들의 K, V 를 계산해 캐시에 저장
│
▼
② decode 한 스텝
새 토큰의 Q 를 만들어 캐시된 모든 K, V 와 어텐션
다음 토큰을 예측
│
▼
③ 그 토큰의 K, V 를 캐시에 추가
│
└──► ②로 되돌아감 (자기회귀)2.3 왜 성격이 갈리나
이 구분이 최적화의 출발점입니다.
prefill
토큰이 수백에서 수만 개다
한 번에 큰 행렬곱을 돌린다
연산기가 바쁘다 ──► compute-bound
decode
토큰이 1개다
연산량이 작은데 모델 가중치 전체를 읽어야 한다
연산기는 놀고 메모리만 왕복한다 ──► memory-bound2.4 그래서 최적화 대상이 다르다
prefill 가속 연산을 줄이거나 잘 배치하는 기법
flash attention 이 여기서 크게 먹힌다
decode 가속 메모리 이동을 줄이는 기법
speculative decoding, 양자화, 배칭"이 최적화가 왜 안 먹히지"의 절반이 여기서 설명됩니다. decode 병목인데 연산 최적화를 붙이면 효과가 없습니다.
3장. 캐시 크기
3.1 계산식
KV 캐시 크기
2 (K와 V) x 층 수 x KV헤드 수 x head_dim
x 시퀀스 길이 x 배치 크기 x dtype 바이트3.2 무엇이 곱으로 들어가나
2 K 와 V 둘 다 저장한다
층 수 모든 층이 각자 캐시를 갖는다
KV헤드 수 여기를 줄이는 게 MQA 와 GQA
head_dim 헤드 하나의 차원
시퀀스 길이 여기에 선형. 컨텍스트를 늘리면 그대로 늘어난다
배치 크기 여기에도 선형. 동시 요청 수
dtype fp16이면 2바이트, fp8이면 1바이트곱해지는 항이 일곱 개라 하나만 키워도 체감이 큽니다.
3.3 실무에서 무엇을 의미하나
길이와 배치에 선형으로 커진다
긴 컨텍스트나 대량 동시요청에서는
KV 캐시가 모델 가중치보다 큰 메모리를 차지하기도 한다
긴 컨텍스트(수만에서 수십만 토큰)와 큰 배치에서는
수 GB 에서 수십 GB 까지 커져
GPU 메모리의 주 소비처가 된다이게 추론 비용과 최대 배치의 한계를 만듭니다. "왜 배치를 더 못 키우나"의 답이 대개 KV 캐시입니다.
3.4 그래서 컨텍스트와 배치는 함께 잡아야 한다
컨텍스트를 2배로 늘리면 KV 캐시도 2배가 된다
같은 VRAM 이면 배치를 절반으로 줄여야 한다
즉 긴 컨텍스트와 높은 처리량은 같은 예산을 놓고 다툰다4장. 캐시를 줄이는 다섯 갈래
4.1 KV 헤드 공유
MQA (Multi-Query Attention)
Q 헤드는 여럿인데 K, V 헤드는 하나만 둔다
캐시가 헤드 수만큼 줄어든다
대신 표현력 손실이 있다
GQA (Grouped-Query Attention)
Q 헤드를 그룹으로 묶고 그룹당 K, V 헤드 하나
MQA 와 MHA 사이의 절충
지금 대부분의 오픈 모델이 이걸 쓴다3.1 식의 "KV헤드 수" 항을 직접 줄이는 접근입니다.
4.2 저랭크 압축
MLA (Multi-head Latent Attention)
K, V 를 저차원 잠재 벡터로 압축해 그것만 저장한다
쓸 때 다시 펼친다
캐시에 들어가는 건 작은 잠재 벡터 하나뿐이다4.3 KV 양자화
K, V 를 fp16 대신 8비트나 4비트로 저장한다
정밀도를 내주고 메모리를 얻는다
K 가 V 보다 양자화에 민감하다는 게 알려져 있어
비대칭 조합을 쓰기도 한다4.4 메모리 관리
PagedAttention
운영체제의 가상 메모리 페이징에서 아이디어를 빌린다
KV 캐시를 고정 크기 블록으로 쪼개 관리해
단편화를 제거한다
"얼마나 쓸지 모르니 최대치로 미리 잡아두는" 낭비가 사라진다캐시를 줄이는 게 아니라 낭비를 없애는 접근이라 성격이 다릅니다.
4.5 재사용
프롬프트 캐싱
공통 prefix 의 K, V 를 캐시해 재사용한다
같은 시스템 프롬프트로 여러 요청이 들어오면
그 부분의 prefill 을 건너뛴다4.6 다섯 갈래를 한눈에
줄인다 MQA/GQA 헤드 수를 줄여서
MLA 저차원으로 압축해서
KV 양자화 비트 수를 줄여서
안 낭비한다 PagedAttention 단편화 제거
다시 쓴다 프롬프트 캐싱 공통 prefix 재사용셋의 성격이 다르므로 함께 쓸 수 있습니다. GQA 모델에 KV 양자화를 걸고 PagedAttention 위에서 프롬프트 캐싱을 하는 게 실제 서빙 구성입니다.
5장. 실무
5.1 어디서 켜고 끄나
transformers
use_cache=True 가 기본이다
끌 이유는 거의 없다
메모리가 극단적으로 부족한 학습 시나리오에서만
끄는 경우가 있다
고속 서빙
vLLM 의 PagedAttention
HF TGI
둘 다 KV 캐시 관리가 핵심 기능이다5.2 캐시를 줄이려면
① 모델 선택 단계
GQA 나 MLA 를 쓰는 모델을 고른다
이게 가장 큰 레버다. 나중에 못 바꾼다
② 배포 설정 단계
KV 양자화를 켠다
컨텍스트 길이를 필요한 만큼만 잡는다
③ 애플리케이션 단계
시스템 프롬프트를 고정해 프롬프트 캐싱이 먹게 한다순서가 중요합니다. 1번을 잘못 고르면 2, 3번으로 메우는 데 한계가 있습니다.
5.3 KV 양자화를 어디까지 내릴까
8비트 (q8_0 급)
품질 손실이 사실상 무시할 수준이라는 보고가 많다
메모리는 절반
4비트
품질 저하가 눈에 띈다는 후기가 다수
최후 수단으로 본다
특히 K 를 4비트로 내리는 건 위험하다
K 가 어텐션 점수 계산에 직접 들어가
오차가 softmax 를 거치며 증폭되기 때문이다정량 수치는 모델과 태스크에 따라 크게 달라지므로 자기 환경에서 재봐야 합니다.
6장. 정리
6.1 발전 흐름
| 시기 | 발전 |
|---|---|
| 2017 이후 | 자기회귀 디코딩에 KV 캐시가 기본 내장 |
| 2019 | MQA. KV 헤드를 하나로 |
| 2023-05 | GQA. MQA와 MHA 사이의 절충 |
| 2023-09 | PagedAttention. 메모리 관리로 단편화 제거 |
| 2024-05 | MLA. 저랭크 압축으로 캐시 축소 |
| 2024 이후 | KV 양자화 실용화 |
흐름이 명확합니다. 처음엔 캐시가 "당연한 최적화"였다가, 컨텍스트가 길어지면서 캐시 자체가 병목이 되어 이를 줄이는 연구가 계속 나왔습니다.
6.2 핵심 셋
① KV 캐시는 과거 토큰의 K, V 를 저장하고 재사용해
제곱 비용을 선형으로 낮춘다
② prefill 은 compute-bound, decode 는 memory-bound 다
성격이 달라 최적화 대상도 다르다
③ 캐시는 길이와 배치에 선형으로 커져
긴 컨텍스트에서 모델 가중치보다 큰 메모리를 먹기도 한다
그래서 GQA/MLA, 양자화, PagedAttention, 프롬프트 캐싱이
전부 이 하나를 겨냥한다6.3 용어 정리
| 용어 | 한 줄 뜻 |
|---|---|
| KV 캐시 | 이전 토큰의 K와 V를 저장해 재계산을 피하는 기법 |
| prefill / decode | 프롬프트 일괄 처리 / 토큰 하나씩 생성하는 두 단계 |
| compute-bound / memory-bound | 연산에 묶임 / 메모리 대역폭에 묶임 |
| MQA / GQA | KV 헤드를 하나로 / 그룹당 하나로 줄이는 어텐션 변형 |
| MLA | K와 V를 저차원 잠재 벡터로 압축해 저장하는 방식 |
| KV 양자화 | K와 V를 8비트나 4비트로 저장해 캐시 메모리를 줄임 |
| PagedAttention | KV 캐시를 블록 단위로 관리해 단편화를 제거하는 기법 |
| 프롬프트 캐싱 | 공통 prefix의 K와 V를 재사용해 prefill을 건너뛰는 것 |
| dtype | 수치 정밀도. 캐시 크기에 직접 곱해진다 |
6.4 참고자료
- Vaswani et al., "Attention Is All You Need" (NeurIPS 2017, arXiv:1706.03762)
- Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need" (MQA, arXiv:1911.02150)
- Ainslie et al., "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" (EMNLP 2023, arXiv:2305.13245)
- Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention" (SOSP 2023, arXiv:2309.06180)
- DeepSeek-AI, "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model" (MLA, arXiv:2405.04434)
'Inference > Optimization' 카테고리의 다른 글
| FlashAttention (0) | 2024.08.27 |
|---|---|
| [Opt] prefill vs decode (0) | 2024.03.11 |
댓글