트랜스포머 — 어텐션을 손으로 계산한다 · KV 캐시가 줄이는 계산량 · 이론

토큰 하나를 더 만들 때 무엇을 다시 계산하는가

한 줄 요약

자기회귀 생성은 토큰을 한 번에 하나씩 만든다. 캐시가 없으면 토큰 하나를 더 만들 때마다 앞 문맥 전체의 k·v 를 처음부터 다시 만든다. KV 캐시는 그 다시 만들기를 없앤다 — 수식은 그대로 두고 계산 횟수만 바꾼다.

왜 이게 필요했나

어텐션 자체는 이미 손으로 계산해 봤을 것이다. 점수를 내고, 소프트맥스를 씌우고, 값의 가중 평균을 낸다. 세 줄이다.

문제는 그 세 줄을 생성 반복 안에서 부를 때 생긴다. 길이 100짜리 답을 만들려면 모델을 100번 부른다. 한 번 부를 때마다 모델은 "지금까지의 문맥 전체" 를 입력으로 받는다. 그리고 그 문맥의 모든 자리에 대해 q·k·v 를 만든다.

그런데 두 번째 부름에서 만드는 첫 토큰의 k·v 는 첫 번째 부름에서 만든 것과 완전히 같은 값이다. 세 번째 부름에서도 같다. 백 번째 부름에서도 같다. 같은 값을 백 번 만들고 아흔아홉 번 버린다.

이 낭비는 코드를 읽어서는 잘 보이지 않는다. 어텐션 함수는 잘못된 것이 하나도 없기 때문이다. 보이게 하려면 횟수를 세어야 한다.

왜 시간이 아니라 곱셈을 세는가

"캐시를 켰더니 빨라졌다" 는 측정으로는 아무것도 배울 수 없다. 빨라진 정도는 기계, 부하, 배치 크기, 메모리 대역폭에 따라 달라지고, 같은 기계에서도 다시 재면 다른 수가 나온다.

곱셈 횟수는 다르다. 같은 코드를 같은 입력으로 돌리면 늘 같은 수다. 그리고 그 수는 길이에 대해 어떻게 늘어나는지를 그대로 보여 준다. 그래서 이 실습은 시간을 재지 않는다. 대신 곱셈을 하는 함수를 하나 만들어 그 안에서 센다.

def mul(a, b):    global _MULS    _MULS += 1    return a * b

정직한 방법이다. 곱셈이 일어나는 자리를 전부 이 함수로 통과시키면, 누가 어디서 몇 번 곱했는지 셀 수 있다. 계산량을 짐작하지 않고 센다.

세어 보면 나오는 모양

한 층 한 헤드만 두고, 모델 차원을 d, 지금 문맥 길이를 n 이라고 하자. q·k·v 를 만드는 가중치는 d 곱하기 d 짜리 행렬 셋이다.

캐시 없이 한 토큰:

앞의 모든 자리에 대해 q·k·v 를 만든다 — 3 곱하기 n 곱하기 d 곱하기 d 번
마지막 질의가 n 개의 키를 본다 — n 곱하기 d 번
그 가중치로 n 개의 값을 섞는다 — n 곱하기 d 번

캐시로 한 토큰:

새 토큰 하나에 대해서만 q·k·v 를 만든다 — 3 곱하기 d 곱하기 d 번
질의가 n 개의 키를 본다 — n 곱하기 d 번
그 가중치로 n 개의 값을 섞는다 — n 곱하기 d 번

차이는 첫 줄 하나뿐이다. 캐시 없는 쪽에서 투영 비용이 문맥 길이에 비례해 늘어나고, 캐시 쪽에서는 그 비용이 길이와 무관하게 고정된다. 어텐션 자체(뒤의 두 줄)는 두 쪽 다 길이에 비례해 늘어난다 — 캐시는 그것을 없애 주지 않는다. 캐시가 없애는 것은 다시 계산하기이지 어텐션이 아니다.

토큰을 다섯 번 만드는 동안 각 자리의 키·값을 몇 번 만드는지 센 그림. 캐시가 없으면 매번 앞 자리 전부를 다시 만들어 1, 2, 3, 4, 5 번 해서 모두 15번이고, 캐시를 쓰면 새로 붙은 자리 하나만 만들어 매번 1번, 모두 5번이다

이 실습에서는 그 수를 손으로 세어 길이별로 표를 만든다. 남이 적어 둔 배수를 옮겨 적는 것이 아니라 자기가 센 수를 쓴다.

캐시가 성립하는 이유는 인과 마스크다

왜 앞 자리의 k·v 를 그대로 재사용해도 되는가? 새 토큰이 붙었는데 앞 자리의 값이 안 바뀐다고 어떻게 확신하는가?

인과 마스크 때문이다. [Attention Is All You Need](https://arxiv.org/abs/1706.03762) 의 디코더는 각 자리가 자기 앞만 보게 되어 있다. 3번 자리의 k 와 v 는 3번 자리의 입력만으로 만들어지고, 4번 토큰이 뒤에 붙는다고 해서 3번의 값이 달라질 이유가 없다.

만약 마스크가 없어 모든 자리가 서로를 본다면 캐시는 성립하지 않는다. 뒤에 뭔가 붙을 때마다 앞의 표현이 바뀌기 때문이다. 그래서 KV 캐시는 디코더 전용 구조의 성질이지, 어디에나 쓸 수 있는 최적화가 아니다.

여기서 중요한 결론이 하나 나온다. 캐시는 근사가 아니다. 답이 달라지지 않는다. 같은 값을 다시 만들지 않을 뿐이다. 그러니 캐시를 켜고 껐을 때 출력이 달라진다면 그건 캐시의 성질이 아니라 구현의 결함이다.

캐시가 먹는 메모리

아끼는 것이 있으면 치르는 것도 있다. 캐시는 메모리를 먹는다. 원소 수는 곱으로 적힌다.

원소 수 = 2 (K 와 V) × 층 수 × 헤드 수 × 문맥 길이 × 헤드 차원바이트  = 원소 수 × 자료형 한 원소의 바이트 수

여기서 눈여겨볼 것은 문맥 길이가 곱해진다는 점이다. 길이가 두 배면 캐시도 두 배다. 이것이 긴 문맥이 비싼 이유 중 하나이고, 한 대의 기계에서 동시에 처리할 수 있는 요청 수를 정하는 것도 대개 이 표다.

이 실습에서는 층·헤드·헤드 차원을 여러분이 정한 작은 값으로 두고 그 값으로만 계산한다. 실제 모델의 GB 숫자는 쓰지 않는다 — 여기서 재지 않았기 때문이다. 재지 않은 숫자를 옮겨 적는 순간 그 글은 근거를 잃는다. 자료형을 바꾸면 마지막 항만 바뀐다는 것, 길이가 곱해진다는 것, 이 두 가지는 자기가 만든 표에서 직접 보인다.

현장에서 만나는 모습

첫째, 긴 프롬프트를 매번 다시 보내면서 캐시를 못 쓴다. 대화를 이어 갈 때 앞 내용을 통째로 다시 보내면 서버 쪽에서 캐시를 이어 쓸 근거가 없다. [Hugging Face 의 Cache strategies 문서](https://huggingface.co/docs/transformers/en/kv_cache)가 캐시 객체를 호출 사이에 직접 들고 다니는 방법을 따로 설명하는 이유가 이것이다.

둘째, 동시 요청 수의 상한이 계산이 아니라 메모리에서 온다. 캐시는 요청마다 따로 잡히고 길이에 비례해 자란다. 그래서 긴 대화 몇 건이 짧은 요청 수십 건보다 더 많은 자리를 차지한다.

셋째, 캐시를 낮은 정밀도로 들고 있으면 출력이 미세하게 달라진다. 캐시 자체는 근사가 아니지만, 캐시에 덜 정확하게 적는 것은 근사다. 이 실습에서는 캐시에 넣기 전에 반올림해 보고 그 차이가 허용 오차를 넘는지 직접 확인한다.

넷째, 첫 토큰과 그다음 토큰의 성격이 다르다. 프롬프트를 통째로 읽는 첫 계산은 캐시가 비어 있어 아낄 것이 없다. 캐시가 이득을 내는 것은 두 번째 토큰부터다. [Hugging Face 의 Text generation 문서](https://huggingface.co/docs/transformers/en/llm_tutorial)도 생성을 이 두 국면으로 나누어 설명한다.

다섯째, 캐시를 쓰면서 배치를 섞으면 자리가 어긋난다. 캐시의 줄 순서가 곧 토큰 순서다. 요청을 묶었다 풀었다 하면서 줄을 잘못 이어 붙이면 오류가 나지 않고 말이 이상해질 뿐이다.

실무에서 진짜 중요한 것

시간 대신 횟수를 세라. 시간은 환경을 재고 횟수는 알고리즘을 잰다. 무엇이 왜 비싼지는 횟수에서만 보인다.
캐시를 켜고 끈 출력이 같은지 시험으로 고정하라. 한 줄짜리 시험인데, 이게 깨지면 캐시 구현이 조용히 틀린 것이다. 견줄 때는 등호가 아니라 허용 오차를 쓴다.
아끼는 쪽과 여전히 자라는 쪽을 구분하라. 투영은 고정이 되지만 어텐션은 길이에 비례해 계속 자란다. 이걸 섞어 말하면 긴 문맥의 비용을 잘못 잡는다.
메모리 계산을 곱셈 다섯 개로 적어 두라. 층·헤드·길이·헤드 차원·자료형. 어느 항이 바뀌면 무엇이 바뀌는지가 한눈에 보인다.
재지 않은 숫자를 옮겨 적지 마라. "몇 배 빨라진다" 는 말은 그 사람의 기계에서 나온 것이다. 자기 설정으로 자기가 세면 옮겨 적을 일이 없다.

다음 실습에서 할 것

/root/work/tf-kv/kv.py 를 한 단계씩 키운다. 표준 라이브러리만 쓴다 — 이 파드의 시스템 파이썬에는 numpy 가 없고(/opt/onnx-lab/bin/python 안에만 있다) torch 도 transformers 도 없다. 그래서 여기 나오는 숫자는 전부 여러분이 만든 코드가 센 것이다.

곱셈을 세는 곱셈 함수와 내적에서 시작해, 토큰 하나를 투영하는 함수, 질의 하나가 쌓인 K·V 를 보는 어텐션, 캐시 없이 도는 판, 캐시로 도는 판을 차례로 만든다. 그다음 두 판을 같은 입력으로 돌려 출력이 같은 값인지 허용 오차로 확인하고, 길이마다 곱셈 횟수를 세어 표로 만든다.

마지막 두 단계가 이 실습의 요점이다. 캐시에 반올림해서 넣어 보면 두 방식의 출력이 더 이상 허용 오차 안에 들어오지 않는다 — 캐시는 근사가 아니지만 캐시에 덜 정확하게 적는 것은 근사라는 사실이 숫자로 나온다. 그리고 여러분이 정한 층·헤드·헤드 차원·자료형으로 캐시가 먹는 원소 수와 바이트를 계산해, 곱셈 횟수 표와 함께 기록으로 남긴다. 채점기는 여러분의 모듈을 실제로 불러 매번 다른 크기의 입력으로 함수를 두드려 보고, 값과 곱셈 횟수를 따로 계산해 대조한다.