본문 바로가기
Inference/Optimization

KV Cache 란

by AteN 2024. 7. 2.

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-bound

2.2 동작을 단계로

① prefill
   프롬프트 토큰들의 K, V 를 계산해 캐시에 저장
   │
   ▼
② decode 한 스텝
   새 토큰의 Q 를 만들어 캐시된 모든 K, V 와 어텐션
   다음 토큰을 예측
   │
   ▼
③ 그 토큰의 K, V 를 캐시에 추가
   │
   └──► ②로 되돌아감 (자기회귀)

2.3 왜 성격이 갈리나

이 구분이 최적화의 출발점입니다.

prefill
   토큰이 수백에서 수만 개다
   한 번에 큰 행렬곱을 돌린다
   연산기가 바쁘다 ──► compute-bound

decode
   토큰이 1개다
   연산량이 작은데 모델 가중치 전체를 읽어야 한다
   연산기는 놀고 메모리만 왕복한다 ──► memory-bound

2.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 참고자료

'Inference > Optimization' 카테고리의 다른 글

FlashAttention  (0) 2024.08.27
[Opt] prefill vs decode  (0) 2024.03.11

댓글