KV 캐시 양자화 — 가중치보다 이쪽이 먼저 넘친다

지난 글에서 양자화의 스케일을 정하는 보정을 다뤘다. 거기서 다룬 대상은 가중치와 활성값이었다. 그런데 LLM을 실제로 서빙해 보면 가중치를 아무리 줄여도 메모리가 모자란 상황을 만나게 된다. 문맥을 길게 열어 두고 동시 요청을 늘리는 순간이다.

범인은 KV 캐시다. 트랜스포머가 토큰을 하나씩 만들어 갈 때, 앞서 지나간 토큰들의 키와 값 벡터를 매번 다시 계산하지 않으려고 저장해 두는 것이 KV 캐시다. 저장하지 않으면 토큰 하나를 만들 때마다 전체 시퀀스를 다시 통과시켜야 하니 캐시는 선택이 아니다. 문제는 이게 요청마다, 토큰마다 쌓인다는 것이다.

얼마나 자라는지 먼저 세어 본다

캐시 크기는 정확히 계산할 수 있다. 층 수 $L$, KV 헤드 수 $H_{kv}$, 헤드 차원 $d_h$, 요소 하나의 바이트 수 $p$, 동시 요청 수 $B$, 시퀀스 길이 $S$일 때 이렇다.

$$\text{bytes} = 2 \cdot L \cdot H_{kv} \cdot d_h \cdot p \cdot B \cdot S$$

앞의 2는 K와 V 둘이라는 뜻이다. 여기서 눈여겨볼 것은 $B$와 $S$가 곱으로 들어간다는 점이다. 가중치는 이 식에 없다. 요청이 늘어도 문맥이 길어져도 가중치는 그대로다.

8B급 모델을 예로 넣어 보자. 층 32개, KV 헤드 8개(GQA), 헤드 차원 128, FP16이면 $p = 2$다. 토큰 하나가 차지하는 크기는 이렇다.

$$2 \times 32 \times 8 \times 128 \times 2 = 131{,}072 \ \text{바이트} = 128 \ \text{KiB}$$

토큰 하나에 128 KiB다. 문맥 32k를 열면 요청 하나당 4 GiB, 동시 요청 16건이면 64 GiB가 된다.

FP16 가중치 16 GiB를 더하면 80 GiB짜리 GPU가 정확히 꽉 찬다. 여기서 요청을 더 받으려면 문맥을 줄이거나, GPU를 더 붙이거나, 캐시를 압축해야 한다. KV 캐시를 INT8로 바꾸면 그 자리에서 32 GiB가 비고 동시 요청을 두 배로 받을 수 있다.

가중치 양자화와 성격이 다르다

같은 양자화라도 KV 캐시 쪽은 조건이 꽤 다르다. 이 차이를 모르고 가중치 양자화의 감각을 그대로 가져오면 어긋난다.

	가중치 양자화	KV 캐시 양자화
언제 하나	배포 전에 한 번	토큰이 생길 때마다 실시간으로
보정 데이터	필요하다	대개 필요 없다. 그 자리에서 잰다
비용	한 번 치르고 끝	매 토큰 붙는다. 무거우면 이득이 사라진다
오차의 성격	고정된 오차가 계속 남는다	앞선 토큰의 오차가 뒤 토큰에 누적된다
되돌리기	원본 파일로 교체	설정 하나로 껐다 켰다 할 수 있다

가장 조심할 것이 누적이다. 가중치 오차는 매 스텝 같은 크기로 들어오지만, KV 캐시의 오차는 다르다. 3번 토큰을 만들 때 생긴 오차가 캐시에 저장되고, 100번 토큰을 만들 때 그 캐시를 다시 읽는다. 긴 생성일수록 앞쪽에서 생긴 오차를 오래 끌고 간다. 짧은 답변에서는 멀쩡한데 긴 답변의 뒷부분이 흐트러지는 증상이 여기서 나온다.

반대로 유리한 점도 있다. 압축에 드는 비용이 그대로 이득으로 돌아온다. 디코딩 단계는 계산보다 메모리 대역폭이 병목인 구간이라, 매 토큰 읽어야 하는 캐시가 절반이 되면 읽는 시간도 절반이 된다. 메모리를 아끼려고 한 일이 속도까지 같이 올리는 흔치 않은 경우다. 물론 역양자화 커널이 그 이득을 까먹지 않을 만큼 가벼워야 한다.

K와 V를 같은 방식으로 다루면 안 된다

KV 캐시 양자화에서 가장 실용적인 발견은 이것이다. K와 V는 값이 퍼진 모양이 다르다.

K 쪽에는 특정 채널만 계속 크게 나오는 패턴이 있다. 어떤 채널은 토큰이 무엇이든 절댓값이 크고, 나머지 채널은 계속 작다. 이 상태에서 토큰 하나의 모든 채널을 묶어 스케일을 잡으면 — 즉 토큰 축으로 묶으면 — 그 큰 채널이 범위를 혼자 정해 버리고 나머지 채널이 몇 칸 안에 뭉개진다. 지난 글에서 본 이상치 문제와 정확히 같은 모양이다.

그래서 K는 채널 축으로 묶는다. 채널마다 스케일을 따로 두면 큰 채널은 큰 채널끼리, 작은 채널은 작은 채널끼리 범위를 나눠 갖는다.

V 쪽은 그런 편향이 약해서 토큰 축으로 묶어도 손해가 적다. 그리고 토큰 축이 구현하기 훨씬 쉽다. 새 토큰이 도착하면 그 토큰의 값만 보고 바로 스케일을 정해 저장하면 끝이기 때문이다.

채널 축 묶기가 까다로운 이유가 여기 있다. 채널 하나의 스케일을 정하려면 그 채널에 속한 토큰들을 모아 봐야 하는데, 토큰은 하나씩 도착한다. 그래서 실제 구현은 토큰을 일정 개수 모아 한 덩어리로 양자화하고, 아직 덩어리를 못 채운 최근 토큰들은 원본 정밀도로 들고 있는다. 이 원본 구간을 잔여 창(residual window)이라 부른다. 2비트 KV 캐시를 다룬 KIVI가 이 조합 — K는 채널 축, V는 토큰 축, 최근 토큰은 원본 — 을 정리해 알려진 방식이다.

최근 토큰을 원본으로 남기는 것은 구현 편의만은 아니다. 방금 만든 토큰일수록 다음 토큰에 미치는 영향이 크다. 어텐션 가중치가 대체로 최근 쪽에 몰리기 때문에, 오래된 토큰을 거칠게 저장하는 것보다 최근 토큰을 거칠게 저장하는 쪽이 훨씬 크게 다친다.

어디까지 줄여도 되는가

실무에서 쓰이는 대역은 대체로 이렇게 갈린다.

• FP8 또는 INT8 — 거의 안전한 구간이다. 대부분의 과제에서 품질 차이를 재기 어렵고, 메모리는 절반이 된다. 서빙 엔진들이 설정 한 줄로 지원하는 것도 이 구간이다
• INT4 — 이득이 크지만 방식을 가려야 한다. K를 채널 축으로 묶고 그룹 크기를 작게 잡고 잔여 창을 두는 조합이면 쓸 만하고, 토큰 축으로 통째로 묶으면 무너진다
• 2비트 — 연구 영역이다. 긴 생성이나 정확한 검색이 필요한 과제에서는 깨진다

vLLM 같은 엔진에서는 대개 플래그 하나로 켠다.

# KV 캐시만 FP8로 두고 가중치는 그대로 둔다
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --kv-cache-dtype fp8 \
  --max-model-len 32768 \
  --gpu-memory-utilization 0.92

여기서 가중치 양자화와 KV 캐시 양자화는 독립된 선택이라는 점이 중요하다. 가중치는 FP16으로 두고 캐시만 FP8로 둘 수 있고, 그 반대도 된다. 문맥이 짧고 배치가 작은 서비스라면 가중치만 줄이는 것이 맞고, 문맥이 길거나 동시 요청이 많다면 캐시 쪽이 먼저다. 어느 쪽이 메모리를 더 많이 쓰는지 실제로 세어 보고 정한다.

캐시를 줄이는 다른 방법들과의 관계

KV 캐시를 줄이는 길이 양자화만 있는 것은 아니고, 서로 곱해서 쓸 수 있는 것과 그렇지 않은 것이 섞여 있다.

방법	무엇을 줄이나	양자화와 함께 쓰나
GQA·MQA	식의 $H_{kv}$를 줄인다	그렇다. 이미 대부분의 최신 모델에 들어 있다
MLA(잠재 압축)	K·V를 저차원으로 압축해 저장한다	그렇다. 모델 구조 자체의 선택이라 나중에 바꿀 수 없다
슬라이딩 윈도우	오래된 토큰을 아예 버려 $S$를 묶는다	그렇다. 다만 버린 토큰은 되돌릴 수 없다
토큰 축출	덜 중요한 토큰만 골라 버린다	그렇다. 무엇이 덜 중요한지 판단이 들어간다
프리픽스 캐싱	같은 앞부분을 요청끼리 공유한다	그렇다. 요청이 겹칠 때만 효과가 있다
CPU·디스크 이관	안 쓰는 캐시를 밖으로 내린다	그렇다. 다시 올릴 때 지연이 붙는다

성격을 나눠 보면 명확하다. 양자화는 정보를 흐리게 만들고, 축출과 슬라이딩 윈도우는 정보를 버린다. 흐린 정보는 어느 토큰에나 조금씩 남아 있지만, 버린 토큰은 없다. 문서 안의 특정 문장을 정확히 다시 꺼내야 하는 과제라면 버리는 쪽이 훨씬 위험하다.

무엇을 재야 도입 판단이 되는가

캐시 양자화를 켤지 정하려면 두 종류의 숫자가 필요하다.

메모리 쪽에서 보는 것은 동시 요청 수의 상한이다. 캐시를 반으로 줄이면 같은 GPU에서 몇 건까지 받을 수 있는지, 그래서 초당 처리량이 얼마나 올라가는지를 본다. 여기가 도입의 실제 이유다. 지연 시간 개선은 덤이다.

품질 쪽에서 보는 것은 세 가지다.

1. 긴 생성. 2,000토큰짜리 답변을 만들게 하고 뒷부분이 흐트러지지 않는지 본다. 짧은 답변만 보면 오차 누적을 못 잡는다
2. 긴 문맥에서의 정확한 인용. 문서 안의 특정 숫자나 이름을 찾아 그대로 옮기게 한다. 캐시가 흐려지면 여기서 먼저 틀린다
3. 여러 번 주고받는 대화. 앞선 차례에서 정한 것을 뒤에서 지키는지 본다

그리고 이 셋은 전부 원본 정밀도와 나란히 놓고 비교한다. 캐시 양자화는 껐다 켜는 것이 설정 한 줄이라, 이 비교가 다른 양자화보다 쉽다는 것이 장점이다.

도입 순서는 이렇게 잡는 것이 무난하다. FP8·INT8로 먼저 켜고 위 세 가지를 재 본다. 차이가 없으면 그대로 두고, 메모리가 여전히 모자라면 INT4로 내려가되 K의 축 처리와 잔여 창을 지원하는 구현인지 확인한다. 지원하지 않는 구현에서 INT4를 켜면 어느 순간 답변이 무너지는데, 원인이 캐시라는 것을 알아채기까지 오래 걸린다.

정리

• 긴 문맥·많은 동시 요청에서 GPU 메모리를 실제로 잡아먹는 것은 가중치가 아니라 KV 캐시다. 캐시는 $B \times S$에 비례해 자란다
• 8B급 모델에서 토큰 하나가 128 KiB다. 32k 문맥 16건이면 64 GiB로, 80 GiB GPU가 가중치와 함께 꽉 찬다
• 가중치 양자화와 달리 실시간이고 오차가 누적된다. 긴 생성의 뒷부분에서 증상이 먼저 나온다
• 대신 디코딩은 메모리 대역폭 병목이라 캐시를 줄이면 속도도 같이 올라간다
• K는 채널 축, V는 토큰 축으로 묶는다. K에는 특정 채널만 계속 큰 편향이 있어 토큰 축으로 묶으면 나머지가 뭉개진다
• 최근 토큰은 원본 정밀도로 남긴다. 어텐션이 최근 쪽에 몰려서, 오래된 토큰보다 최근 토큰을 거칠게 저장하는 쪽이 훨씬 크게 다친다
• FP8·INT8은 안전한 구간이고, INT4는 축 처리와 잔여 창이 갖춰진 구현에서만 쓴다
• 양자화는 정보를 흐리게 하고 축출·슬라이딩 윈도우는 정보를 버린다. 정확한 인용이 필요한 과제라면 버리는 쪽이 더 위험하다
• 검증은 긴 생성·긴 문맥 인용·여러 차례 대화 셋을 원본과 나란히 놓고 본다

---

읽어주셔서 감사합니다. 😊