KV キャッシュは再計算しなくてよいものを覚えておくこと

한국어 원문으로 표시합니다.

한 줄 요약

생성은 토큰 하나를 뽑을 때마다 모델을 한 번 돌리는 일이다. 캐시가 없으면 매번 지금까지의 전체를 다시 넣어 계산하고, KV 캐시가 있으면 새 토큰 하나만 넣고 앞 토큰들의 K·V 는 꺼내 쓴다. 결과는 한 토큰도 다르지 않고 계산만 준다. 뽑는 방법(욕심쟁이·온도·top-p)은 그와 별개로 '어떤 토큰을 고를지' 를 정한다. 이 모듈에서는 MiniMind 의 generate 로 두 가지를 모두 숫자로 잰다.

왜 이게 필요했나

트랜스포머 코스에서 KV 캐시의 원리와 샘플링 공식을 손으로 계산했다. 여기서는 실제로 학습한 모델과 실제 생성 루프에서 그것이 어떻게 보이는지를 본다. 서빙 비용은 대부분 생성 단계에서 나오고, 그 비용의 모양은 캐시가 정한다. 그리고 같은 모델이 욕심쟁이로는 늘 같은 답을, 온도를 올리면 매번 다른 답을 내는데, 그 차이가 '창의성' 인지 '헛소리' 인지는 작은 모델에서 특히 극명하게 드러난다.

어떻게 동작하나

MiniMind 의 generate 는 이렇게 돈다.

for _ in range(max_new_tokens):
    past_len = past_key_values[0][0].shape[1] if past_key_values else 0
    outputs = self.forward(input_ids[:, past_len:], past_key_values=past_key_values, use_cache=use_cache)
    logits = outputs.logits[:, -1, :] / temperature
    ... top_k · top_p 로 자르기 ...
    next_token = torch.multinomial(softmax(logits), 1) if do_sample else argmax(logits)
    input_ids = torch.cat([input_ids, next_token], -1)
    past_key_values = outputs.past_key_values if use_cache else None

KV 캐시. 첫 걸음에는 프롬프트 P 토큰을 한꺼번에 넣는다(prefill). 캐시가 있으면 그 뒤로는 걸음마다 한 토큰만 넣으므로 N 개를 만드는 동안 forward 에 들어간 토큰은 P + (N − 1) 개다. 캐시가 없으면 걸음마다 P, P+1, … 를 다시 넣어 P·N + N(N−1)/2 개다. 프롬프트 183토큰으로 64토큰을 만들면 246 대 13,728 — 56배다. 인과 마스크 덕분에 앞 토큰의 K·V 는 뒤 토큰이 생겨도 바뀌지 않으므로, 꺼내 써도 결과가 같다.

캐시의 크기. 캐시는 층마다 K·V 이고, 모양은 (배치, 길이, KV 헤드, head_dim) 이다. MiniMind 는 repeat_kv 로 복사하기 전의 K·V 를 캐시에 넣으므로 GQA 로 줄인 만큼 그대로 준다. 토큰 100개면 2 × 4층 × 100 × 2헤드 × 32 × 4바이트 = 204,800 바이트다. 서빙에서는 이 숫자가 동시 사용자 수를 정한다.

온도. 로짓을 온도로 나눈 뒤 소프트맥스한다. 1 보다 작으면 분포가 뾰족해져 가장 그럴듯한 토큰으로 몰리고, 크면 평평해져 드문 토큰도 뽑힌다. 온도 0 에 가까우면 욕심쟁이와 같다.

top-p(nucleus). 확률이 큰 순서로 정렬해 누적 확률이 p 를 넘는 지점까지만 남기고 나머지를 −∞ 로 지운다. MiniMind 는 마스크를 한 칸 밀어 p 를 처음 넘기는 토큰까지 살리고, 맨 위 토큰은 언제나 살린다. 분포가 뾰족할 때는 몇 개만, 평평할 때는 많이 남는다 — 고정 개수를 남기는 top-k 와 다른 점이다. MiniMind 의 generate 기본값은 온도 0.85·top_p 0.85·top_k 50 이다.

현장에서 만나는 모습

'캐시를 켰는데 답이 달라졌다' 는 버그 신고는 대개 캐시가 아니라 샘플링 쪽 문제다 — 난수를 고정하지 않았거나, 한쪽만 do_sample 이다. 욕심쟁이로 두 방식을 견줘 한 토큰도 다르지 않은지 먼저 확인하면 원인을 반으로 줄인다. 반대로 긴 문서를 요약할 때 첫 토큰까지의 시간(TTFT)이 긴 것은 캐시로 줄지 않는다 — prefill 은 원래 프롬프트 전체를 계산해야 한다. 캐시가 줄이는 것은 그 뒤 토큰 사이의 시간이다.

MiniMind 원본과 이 코스가 다른 점

MiniMind 의 eval_llm.py 와 웹 데모는 온도 0.85·top_p 0.95 로 뽑고, 반복 벌점(repetition_penalty)을 줄 수 있다. 이 코스는 실험을 깨끗하게 하려고 top_k 를 끄고(top_k=0), 한 번에 한 가지만 바꾼다. 또 MiniMind 의 generate 는 배치마다 끝난 줄을 따로 기억해(finished) 모두 끝날 때 멈추고, 끝난 줄에는 eos 를 계속 채운다. 30개를 한 번에 뽑을 때 짧은 답 뒤에 eos 가 이어지는 이유다. 실제 서빙 엔진(vLLM 등)은 여기에 요청마다 다른 길이의 캐시를 한 배치에 묶는 장치를 더한다 — 원리는 같고, 관리가 복잡해질 뿐이다.

다음 실습에서 할 것

기준 사전학습 모델로 캐시가 있을 때와 없을 때 같은 토큰이 나오는지 보고, forward 에 들어간 토큰 수를 세어 식과 맞추고, 시간을 잰다. 캐시의 실제 바이트를 식과 맞추고, SFT 모델로 온도를 바꿔 30번씩 뽑아 서로 다른 답의 수를 세고, top-p 가 남기는 후보 수를 센다.