루나에게 말 걸기: 같은 시작 문장에 왜 매번 다른 이야기가 나올까

KV 캐시: 앞에서 계산한 키와 값을 적어 두기

루나에게 200 토큰짜리 이야기를 쓰게 하면 모델을 200번 돌려야 한다. 스텝마다 지금까지의 글 전체를 넣으므로, 100번째 스텝에서는 토큰 100개를, 200번째 스텝에서는 200개를 처음부터 다시 계산한다. 그런데 어텐션은 뒤 토큰을 가려 두었으니, 앞 토큰들의 계산 결과는 뒤에 무엇이 붙든 바뀌지 않는다. 바뀌지 않는 것을 스텝마다 다시 계산할 필요가 있을까?

우리 GPT 의 어텐션은 캐시가 있으면 앞에서 적어 둔 키와 값 뒤에 새 토큰의 것만 이어 붙인다.

# sources/common/gpt.py 의 CausalSelfAttention.forward (발췌)
q, k, v = self.c_attn(x).split(C, dim=2)       # 새로 들어온 토큰의 쿼리·키·값
...
if cache is not None:  # KV 캐시: 앞에서 계산한 키·값 뒤에 이어 붙인다
    if cache.get("k") is not None:
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v

캐시를 쓰면 스텝마다 모델에 들어가는 토큰은 새 토큰 하나뿐이고, 그 토큰의 쿼리가 적어 둔 키 전체를 읽는다. 계산량으로 보면 200 토큰을 쓸 때 스텝마다 평균 100개 정도 하던 계산이 1개로 줄어드니 수십 배 빨라져야 할 것 같다. 그리디 선택으로 캐시 없이와 캐시 있이 같은 글을 쓰게 하고 시간을 쟀다(두 글이 같은지도 확인했다).

GPU 에서 글 한 줄(배치 1)을 쓸 때 캐시가 거의 소용없었던 이유는 스텝 하나에 걸린 시간으로 나눠 보면 보인다. 캐시 없이 500 스텝에 0.75초면 스텝당 1.5 ms 이고, 캐시를 써도 0.61초 ÷ 500 = 1.2 ms 다. 계산할 토큰을 수백 개에서 1개로 줄였는데 스텝당 시간은 0.3 ms 만 줄었으니, 스텝마다 드는 시간의 대부분은 토큰 수와 상관없는 비용이다. 파이토치는 층마다 행렬곱, 정규화, 어텐션을 GPU 에 따로따로 맡기는데, 맡길 때마다 드는 고정 비용이 8층을 지나는 동안 쌓인다. RTX 4090 에게 2,760만 파라미터 모델의 토큰 수백 개는 작은 일이라, 토큰 1개를 계산하든 수백 개를 계산하든 맡기는 비용이 시간을 정한다. 이 설명은 스텝당 시간에서 나온 우리의 짐작이고, 고정 비용을 따로 잰 것은 아니다.

배치를 64로 늘리면 스텝마다 다시 계산할 토큰이 64배가 되어 GPU 도 일다운 일을 하게 된다. 캐시 없이는 10.90초 ÷ 500 = 스텝당 21.8 ms 였고 캐시를 쓰면 2.7 ms 로, 8.0배 빨라졌다. CPU 는 계산 자체가 느려서 배치 1 에서도 다시 계산하는 양이 그대로 시간이 되고, 3.6배 빨라졌다. 캐시는 다시 계산하는 일을 없애 주므로, 그 계산이 시간의 대부분을 차지할 때만 효과가 크다. 같은 날 앞서 돌린 실행에서는 GPU 배치 1 에서 캐시가 오히려 더 느린 적도 있었는데, 이 머신은 다른 작업도 함께 돌고 있어서 속도 숫자는 실행마다 흔들린다.

그럼 우리 반 서른 명이 동시에 루나한테 말을 걸면 그때 캐시가 제값을 하겠네. 한 명씩이면 별 차이 없고.

서른 명이 루나랑 얘기할 일은 없어. 너 혼자 서른 명만큼 떠들 거면 몰라도.

문제 7. KV 캐시는 메모리를 얼마나 먹나

우리 GPT 는 8층, 폭 512 이고, 캐시는 층마다 토큰 하나당 키 512개와 값 512개를 32비트 실수(4바이트)로 적어 둔다. (가) 토큰 하나당 캐시는 몇 바이트인가? (나) 500 토큰 글 한 줄의 캐시는? 배치 64 이면? (다) 캐시 크기는 글 길이, 층 수, 폭 가운데 무엇에 비례하는가? 문맥 길이가 이 모델의 512 토큰보다 훨씬 긴 모델에서 무엇이 문제가 될까?

(가)는 곱하기만 하면 돼. 층 8 × 키와 값 2 × 512 = 8,192 개고, 4바이트씩이면 32,768 바이트. 32 KiB.

(나)는 500 을 곱해서 16,384,000 바이트, 16 MB 쯤. 배치 64 면 64배라 1 GB 를 조금 넘어.

배치 64 에서 8.0배 빨라진 대가가 GPU 메모리 1 GB 정도가었던 거예요. (다)는요?

셋 다에 비례해요. 글 길이를 두 배로 하면 두 배, 층이나 폭을 두 배로 해도 두 배. 큰 모델이 긴 글을 여러 줄 동시에 쓰면 곱이 금방 커지겠다.

그러면 모델 파라미터보다 캐시가 더 커질 수도 있겠네. 이 모델은 파라미터가 2,760만 개, 4바이트씩 110 MB 쯤인데, 배치 64 의 캐시가 벌써 그 열 배 가까이니까.

그래서 큰 모델들은 캐시를 줄이는 구조를 따로 연구해요. 김민준 학생이 떠올린 디퓨전 강의의 그 절에 나온 연구들이 그 이야기였어요.

빨라지는 만큼 메모리를 내주는 거네요. 배치를 무작정 키울 수도 없고.

문제 8. 배치 1 에서 캐시가 소용없었던 이유를 숫자로 따지기

GPU 배치 1, 500 토큰에서 캐시 없이 0.75초, 캐시 있이 0.61초였다. 스텝당 시간을 「토큰 수와 상관없는 고정 비용 a」와 「토큰 하나당 계산 비용 b × 한 스텝에 계산하는 토큰 수」의 합으로 놓자. 시작 토큰은 1개이고 캐시 없이는 스텝마다 지금까지의 토큰을 모두 계산한다(평균 약 250개). (가) 두 측정으로 a 와 b 를 짐작하라. (나) 짐작한 a, b 로 보면 스텝당 시간의 몇 % 가 고정 비용인가? (다) 이 두 항짜리 계산으로 배치 64 의 8.0배를 설명할 수 있는가? 무엇이 빠졌나?

캐시 있이는 스텝마다 토큰 하나니까 a + b = 0.61 ÷ 500 = 1.22 ms. 캐시 없이는 a + 250b = 1.5 ms. 빼면 249b = 0.28 ms 니까 b 는 0.0011 ms 쯤이고, a 는 거의 1.22 ms 야.

그럼 캐시 있이는 스텝당 시간의 99% 넘게가 고정 비용이고, 캐시 없이도 1.22 ÷ 1.5 니까 80% 정도가 고정 비용이야. 캐시가 줄여 주는 건 나머지 20% 뿐이니 1.2배가 나올 수밖에 없지.

숫자 둘로 미지수 둘을 풀었으니 맞는지 검산할 데이터는 없어요. (다)에서 검산해 볼까요?

일단 배치 64 면 b 가 64배가 된다고 쳐 볼게요. 그러면 캐시 없이 a + 250 × 64b = 1.22 + 17.6 = 18.8 ms 인데 실제는 21.8 ms. 캐시 있이 a + 64b = 1.29 ms 인데 실제는 2.7 ms 예요. 캐시 있이가 두 배나 어긋나요.

그럼 두 항짜리 계산에서 빠진 게 있다는 거지. 캐시를 쓰면 새 토큰 하나가 적어 둔 키 전체를 읽어야 하잖아. 글이 길어질수록 읽을 키가 늘어나는 비용은 a 에도 b 에도 없어.

좋은 데를 짚었어요. 캐시 읽기도 일이고, 배치 64 면 매 스텝 1 GB 가까운 캐시의 일부를 메모리에서 읽어 와요. 「토큰 수와 상관없는 비용 + 토큰 수에 비례하는 비용」은 배치 1 에서 캐시가 왜 소용없었는지를 설명하는 데까지만 쓰고, 그 너머는 직접 재서 확인해야 해요.