안녕하세요, 에디터 배니입니다.

지난해 DeepSeek에 이어 한 번 더 AI 연구가 주식 시장에 크게 영향을 미쳤습니다. 이번에는 Google의 TurboQuant가 그 주인공인데요. TurboQuant의 핵심을 한 줄로 요약하자면, KV Cache를 손실 없이 압축하여 메모리 사용을 6분의 1 수준으로 절감할 수 있다는 것입니다. 메모리 사용이 줄어든다고 하니 불쑥 ‘고성능 메모리 수요가 줄어들지 않을까’하는 우려가 시장에 반영된 것으로 보입니다. 이제 적은 메모리를 사용하고도 기존 모델을 동작시킬 수 있기 때문에 고사양 GPU가 필요하지 않으니까요.

그러나 꼭 그렇지만은 않다는 전문가 의견이 잇따르고 있습니다. 오히려 메모리 수요가 늘 것이라고 예측하기도 합니다. 그간 기술 개발에 제한이 되었던 부분이 메모리 효율의 증가로 해소되기 때문입니다. (이런 현상을 ‘제본스의 역설’이라고도 부릅니다.) 또한, KV Cache는 모델 추론(Inference)에서 비용을 절감하는 것이기 때문에 학습용 GPU와는 크게 관련이 없기도 합니다.

결국, TurboQuant가 어떻게 작동하는지 이해하기 위해서는 KV Cache가 무엇인지 알아야 합니다. KV Cache가 무엇이길래 이런 난리가 일어난 걸까요? 이번주 뉴스레터에서는 KV Cache가 무엇인지 그림과 함께 소개합니다.

KV Cache를 이해하기 위해서는 Attention을 알아야 한다.

KV Cache에서 K는 Key를, V는 Value를 의미합니다. 그리고 이 Key와 Value는 Attention 연산에서 활용됩니다. 이 Attention 연산, 더 정확히 Self-Attention 연산은 Transformer 블록의 핵심이고요. Transformer는 LLM 모델의 핵심 구조입니다. 이 Self-Attention 구조만 알아도 사실상 LLM의 핵심은 파악한 셈입니다.

Decoder-only Transformer 구조의 현대 LLM은 주어진 토큰을 보고 다음 토큰을 예측합니다. 다음 토큰을 예측한다니, 감이 잘 오지 않을 수도 있는데요. <나는 어제 치킨을 먹>까지 토큰이 주어진 상황에서 다음 토큰을 예측하라고 하면, <었>을 예측할 확률이 높습니다. 그리고 <나는 어제 치킨을 먹었>까지 토큰이 주어진다면 다음은 <다>를 예측할 확률이 높겠죠. 정말 LLM이 지능을 가지고 똑똑해서 예측을 잘했다기보다는 모델이 학습한 데이터셋에서 ‘먹’이 나오면 다음은 ‘었’이, ‘먹었’이 나오면 ‘다’가 나오는 패턴이 많았기 때문에 확률적으로 예측한 것뿐입니다.

다음 토큰을 예측한다는 의미를 이해했다면, 어떻게 계산되길래 다음 토큰을 잘 예측하는지 알아야 합니다. 저는 종종 AI 모델을 ‘거대한 계산기’라고 부르는데요. 결국, 모델 훈련 과정에서 가중치를 최적화시키고, 추론 과정에서 달라지는 입력값에 따라 이미 결정된 가중치와 연산만 하면 되기 때문입니다. 말은 쉽지만, 실제로 매우 복잡한 훈련과 연산 과정을 거칩니다. 사실상 AI 연구란 가중치를 빠르게, 정확히 최적화시키기 위해서 세계의 수많은 학자들이 연구하고 있는 셈인데요. 그중 대다수의 상황에서 준수한 성능을 보이는 하나의 연산 방식을 찾았으니, 그것이 바로 Attention 메커니즘입니다.

거두절미 하고, 가장 기본적인 Attention 연산의 수식을 알아보도록 하겠습니다. 그 유명한 <Attention Is All You Need> (Vaswani et al., 2017)에서 Attention 연산을 아래와 같이 정의합니다.

Attention Is All You Need에서 정의한 Scaled Dot-Product Attention(SDPA) 출처: <Attention Is All You Need> (Vaswani et al., 2017)

이런 수학 수식이 낯설다면 어려워보일 수도 있지만, Softmax 함수와 스케일을 조절하기 위한 상수 d를 제외하면 꽤 단순한 행렬 연산이 됩니다. Q, K, V는 각각 입력 토큰에 따라 계산된 행렬입니다. 아래는 입력 토큰별로 계산된 임베딩 행렬의 예시입니다. 입력 토큰은 총 4개(나는 / 어제 / 치킨을 / 먹)이고, 차원은 8입니다. 4 × 8 행렬이 생성됐습니다. (여기서는 쉬운 예를 들기 위해 차원을 8로 설정했습니다.)

위 행렬에 Q, K, V를 얻기 위해 가중치를 한 번씩 더 계산합니다. 결과적으로 차원의 크기는 달라지지 않으면서, Q, K, V 각각은 조금씩 다른 값을 얻게 됩니다. 아래는 헤드가 1개일 때 예시입니다. 편의상 Multi-head Attention이 아닌, Single-head Attention이라고 생각하겠습니다.

이제 Query와 Key의 Transpose를 연산해봅시다. 4 × 8 행렬과 8 × 4 행렬을 곱하여 4 × 4의 행렬을 얻게 됩니다. 즉, (입력 토큰 수) × (입력 토큰 수)만큼의 크기가 되죠.

여기까지가 QK^T 연산과정입니다. 여기에 이제 d_k로 나누고, Softmax 연산 이후에 다시 V를 계산해주면 Attention Score를 얻을 수 있습니다. 동일하게 토큰화된 문장간의 연관성을 계산한 것이기 때문에, 자기 자신과의 Attention 값을 계산한다고 하여 Self-Attention이라고 합니다. 아무리 Q와 K의 값이 다르더라도 일반적으로는 대칭적인 구조로 나타나게 됩니다. 실제로는 위 연산 결과를 도식화하면 다음과 같습니다. 위 수식과 연결이 되시나요?

Q와 K^T를 계산한 결과

실제로는 다음에 나올 토큰을 예측하기 위해, Softmax 이전에 미래의 토큰을 보지 못하도록 상삼각행렬에 마스크를 씌우고, Softmax 함수를 통과시킨 후 V와 계산합니다. V도 4 × 8의 행렬이었던 것 기억하시나요? 최종적으로 얻게 되는 행렬은 4 × 4 행렬과 4 × 8 행렬을 곱한 4 × 8 행렬입니다.

다시 입력 차원과 동일한 크기의 행렬을 얻게 됐습니다. 이와 같은 연산을 반복하고 최종적으로 이 행렬을 LM Head와 곱하여 가장 확률이 높은 토큰 하나를 선택합니다. LM Head 행렬은 8 × (전체 단어 수) 입니다. 즉, 최종적으로 (입력 토큰 수) × (전체 단어 수)만큼이 계산되는 것입니다. 기술 용어로 설명하자면 (Context Length) × (Vocab Size)로 볼 수 있습니다.

‘먹’의 경우 ‘었’이 가장 높게 나왔습니다. LLM의 목표는 다음 토큰을 생성하는 것이라고 했기 때문에 다른 예측값들은 필요하지 않습니다. 이런 경우 다음 토큰으로 ‘었’을 생성하고, Context Length가 1 늘어난 채로 다시 처음부터 동일한 계산을 하게 됩니다.

그래서 KV Cache는 왜 필요할까?

새로운 토큰이 입력된 뒤에도 앞서 설명한 Attention 연산을 반복하면서 다음 토큰을 하나씩 생성해나갈 수 있습니다. 그러나 사실 더이상 앞 부분의 맥락끼리는 다시 계산할 필요가 없습니다. 우리에게 필요한 건, 생성된 다음 토큰(여기서는 ‘었’)이 입력됐을 때, 그 다음 토큰이 무엇일지만 예측된 결과입니다. 이때 KV Cache가 등장합니다.

K와 V가 생성됐을 때, 각각의 값을 미리 캐싱(Caching)해뒀다고 생각해봅시다. 이 경우, 모든 토큰을 다시 생성할 필요 없이 새롭게 생성된 하나의 토큰만 계산하면 됩니다. 위 예시에서는 ‘었’이 새롭게 생성됐었죠. 첫 번째 과정에서는 (나는 / 어제 / 치킨을 / 먹)까지 모두 입력해야 했지만, 이번에는 ‘었’만 입력해도 된다는 의미입니다. 그러면 ‘었’과 관련된 Q, K, V 값을 모두 얻을 수 있습니다.

그리고 캐싱된 KV 값을 불러와 새로 생성된 ‘었’의 K, V와 연결하고, ‘었’의 Q와 계산하면 앞서 계산한 것처럼 다음 토큰을 계산할 수 있습니다. 다음 토큰인 ‘다’를 얻었다면 마찬가지로 반복하면 됩니다.

처음에 한 번만 전체 토큰을 계산한 뒤로는 KV를 캐싱하여 새로 생성된 토큰에 대해서만 연산을 하면 됩니다. 이때 처음 한 번 전체 과정을 연산하는 것을 Prefill이라고 하고, 이후 과정을 Decoding이라고 합니다. 모델의 추론 과정은 이렇게 두 가지 단계로 구분할 수 있습니다. KV를 캐싱하고 빠르게 메모리를 불러와 사용한다면 추론 속도를 훨씬 빠르게 개선할 수 있습니다.

TurboQuant는 무엇을 해결했나?