LLM Research2026.10.089 min read

MHA · MQA · GQA · KV Cache: 생성할 때 같은 계산을 매번 다시 하지 않는 법

LLM이 글자를 한 토큰씩 생성할 때, 이전 토큰의 Key/Value를 매번 다시 계산하면 얼마나 낭비인지 직접 확인한다. 이를 막는 KV Cache를 내 GPT에 붙이고, 캐시 메모리를 줄이기 위해 Key/Value 헤드를 공유하는 MQA·GQA 구조로 MHA를 바꿔 본다. 캐시 유무에 따른 생성 속도와 GQA의 메모리 감소량을 직접 측정한다.

By 송민성

오늘의 목표

LLM이 답변을 생성할 때 왜 느려지고 메모리를 많이 쓰는지 이해하고, 이를 줄이는 두 가지 구조(KV Cache, Key/Value 헤드 공유)를 내 GPT 코드에 직접 넣어 본다. 오늘 끝나면 "캐시가 왜 필요한지", "GQA가 왜 캐시를 줄이는지"를 숫자로 설명할 수 있어야 한다.

개념과 직관

1. 생성은 한 토큰씩 반복된다

GPT는 문장을 한 번에 쓰지 않는다. "오늘 날씨는"을 넣으면 다음 토큰 하나를 예측하고, 그 토큰을 입력 끝에 붙여 다시 모델에 넣는다. 이 반복을 자기회귀 생성(autoregressive generation, 방금 만든 출력을 다음 입력에 이어 붙이며 한 토큰씩 만드는 방식)이라고 한다.

문제는 매 단계마다 앞의 모든 토큰을 처음부터 다시 계산한다는 점이다. 길이 100짜리 문장이면 101번째 토큰을 만들 때 앞의 100개를 또 계산한다. 이미 계산한 결과가 바뀌지 않는데도 말이다.

2. Attention에서 무엇이 재사용되나

Self-Attention에서는 각 토큰이 세 벡터로 변환된다.

Query(Q, "나는 무엇을 찾고 있나"라는 질문 벡터)
Key(K, "나는 이런 내용이다"라는 색인 벡터)
Value(V, 실제로 가져다 쓸 내용 벡터)

새 토큰의 Q가 과거 토큰들의 K와 비교되어 가중치를 만들고, 그 가중치로 과거 V를 섞는다. 핵심 관찰은 이것이다.

GPT는 미래 토큰을 보지 않는 causal mask를 쓴다. 그래서 과거 토큰의 K, V는 새 토큰이 추가되어도 바뀌지 않는다.
새 토큰 단계에서 필요한 Q는 "새 토큰의 것" 하나뿐이다.

비유하면 도서관 사서다. 책 100권의 색인 카드(K)와 내용 요약(V)을 이미 만들어 두었다면, 101번째 책이 들어와도 앞 100권의 카드를 다시 쓰지 않는다. 새 카드 한 장만 만들어 서랍에 추가하면 된다. 이 서랍이 KV Cache(과거 토큰의 K, V를 저장해 두고 재사용하는 메모리)다.

3. 수식 최소한으로

캐시가 없을 때 길이 T까지 생성하는 총 토큰 계산량은 대략 1 + 2 + ... + T = T(T+1)/2 에 비례한다(K, V 투영 계산 기준). 캐시가 있으면 매 단계 새 토큰 1개만 투영하므로 T에 비례한다. 단, Attention 점수 계산 자체(새 Q와 과거 K 전체의 비교)는 캐시가 있어도 매 단계 과거 길이만큼 필요하다. 줄어드는 것은 "과거 토큰을 다시 통과시키는 낭비"다.

4. KV Cache의 메모리 크기

캐시는 공짜가 아니다. 크기는 이렇게 계산한다.

KV 캐시 바이트 = 2 × 레이어 수 × 배치 크기 × 시퀀스 길이 × KV 헤드 수 × 헤드 차원 × 원소당 바이트

맨 앞의 2는 K와 V 두 개를 저장하기 때문이다. 시퀀스 길이와 배치가 커질수록 선형으로 늘어나, 긴 문맥이나 동시 요청이 많은 서비스에서는 모델 가중치보다 캐시가 더 큰 병목이 되기도 한다.

5. MHA → MQA → GQA: KV 헤드를 줄이는 아이디어

MHA(Multi-Head Attention, 헤드마다 자기만의 Q, K, V를 갖는 원래 방식): Q 헤드 수 = KV 헤드 수. 캐시가 가장 크다.
MQA(Multi-Query Attention, 모든 Q 헤드가 K, V 한 세트를 공유): KV 헤드 = 1. 캐시가 가장 작지만, 표현력이 줄어 품질이 떨어질 수 있다.
GQA(Grouped-Query Attention, Q 헤드를 그룹으로 묶어 그룹마다 K, V 한 세트를 공유): KV 헤드 수가 그 중간. 품질과 메모리를 절충한다.

예: Q 헤드 8개일 때

text

MHA: Q 8개 / KV 8개  → 캐시 1배
GQA: Q 8개 / KV 2개  → 캐시 1/4  (Q 4개가 KV 1세트 공유)
MQA: Q 8개 / KV 1개  → 캐시 1/8

교실 비유: MHA는 학생 8명이 각자 참고서를 갖는 것, GQA는 4명씩 한 권을 같이 보는 것, MQA는 반 전체가 한 권을 보는 것이다. 질문(Q)은 각자 다르지만 참고서(K, V)는 공유한다. 참고서 보관 공간이 줄어드는 대신 선택지가 조금 줄어든다.

코드로 직접 해보기

전제: 앞선 Day에서 만든 GPT의 Attention을 단순화해 쓴다. 아래는 따라 칠 수 있게 독립 실행되는 최소 버전이다.

1단계: GQA Attention (MHA는 n_kv_head = n_head인 특수 경우)

python

import torch, torch.nn as nn, torch.nn.functional as F

class GQAttention(nn.Module):
    def __init__(self, d_model=256, n_head=8, n_kv_head=2):
        super().__init__()
        assert n_head % n_kv_head == 0
        self.n_head, self.n_kv_head = n_head, n_kv_head
        self.hd = d_model // n_head
        self.q = nn.Linear(d_model, n_head * self.hd, bias=False)
        self.k = nn.Linear(d_model, n_kv_head * self.hd, bias=False)
        self.v = nn.Linear(d_model, n_kv_head * self.hd, bias=False)
        self.o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, cache=None):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.n_head, self.hd).transpose(1, 2)
        k = self.k(x).view(B, T, self.n_kv_head, self.hd).transpose(1, 2)
        v = self.v(x).view(B, T, self.n_kv_head, self.hd).transpose(1, 2)

        if cache is not None:                      # 과거 K, V 뒤에 이어 붙이기
            k = torch.cat([cache[0], k], dim=2)
            v = torch.cat([cache[1], v], dim=2)
        new_cache = (k, v)                         # 공유 전(KV 헤드 개수 그대로) 저장

        rep = self.n_head // self.n_kv_head        # KV를 Q 헤드 수만큼 복제해서 맞춤
        k = k.repeat_interleave(rep, dim=1)
        v = v.repeat_interleave(rep, dim=1)

        S = k.shape[2]
        # 캐시가 있으면 Q는 T(보통 1)개, K는 S개: 마스크는 오른쪽 아래 정렬
        mask = torch.ones(T, S, dtype=torch.bool, device=x.device).tril(diagonal=S - T)
        att = (q @ k.transpose(-2, -1)) / self.hd ** 0.5
        att = att.masked_fill(~mask, float("-inf")).softmax(dim=-1)
        out = (att @ v).transpose(1, 2).reshape(B, T, -1)
        return self.o(out), new_cache

포인트는 두 가지다. 캐시에는 복제 전 K, V(KV 헤드 수 기준)를 저장해야 메모리가 줄어든다. 그리고 tril(diagonal=S-T)는 "새 토큰은 자기와 과거만 본다"를 캐시 상황에서도 유지한다.

2단계: 정확성 확인 (캐시 있음 = 없음)

python

torch.manual_seed(0)
attn = GQAttention().eval()
x = torch.randn(1, 10, 256)

full, _ = attn(x)                       # 한 번에 전체 계산
cache, outs = None, []
for t in range(10):                     # 한 토큰씩 + 캐시
    o, cache = attn(x[:, t:t+1], cache)
    outs.append(o)
print(torch.allclose(full, torch.cat(outs, dim=1), atol=1e-5))  # True여야 함

True가 나오지 않으면 속도 실험으로 넘어가지 말고 먼저 고친다. 틀린 캐시로 빠르게 돌리는 건 의미가 없다.

3단계: 내 GPT에 붙이기

내 GPT의 Block이 forward(x)만 받았다면 다음처럼 바꾼다.

Block, GPT의 forward가 caches 리스트(레이어당 하나)를 받고 돌려주게 한다.
위치 정보를 넣는 방식이 있다면, 캐시 사용 시 새 토큰의 위치를 과거 길이부터 시작하도록 해야 한다. 학습형 위치 임베딩이라면 pos = torch.arange(past_len, past_len + T)를 쓴다.
생성 루프: 첫 단계에서 프롬프트 전체를 한 번에 넣어 캐시를 만들고(이를 prefill이라 부른다), 이후에는 토큰 1개씩 넣는다.

python

@torch.no_grad()
def generate(model, idx, max_new, use_cache=True):
    caches = None
    for _ in range(max_new):
        inp = idx if (not use_cache or caches is None) else idx[:, -1:]
        logits, caches = model(inp, caches if use_cache else None)
        nxt = logits[:, -1].argmax(-1, keepdim=True)
        idx = torch.cat([idx, nxt], dim=1)
    return idx

model의 반환값이 (logits, caches)가 되도록 내 GPT를 맞춘다. 학습 코드는 logits, _ = model(x)로 고치면 된다.

오늘의 실험

실험 1: KV Cache 유무에 따른 생성 속도 (토큰/초)

고정: 모델, 프롬프트 길이(예: 16토큰), 디바이스, torch.no_grad()
바꾸는 값: use_cache ∈ {False, True}, 생성 길이 ∈ {64, 256, 512}
관찰: 토큰/초 = 생성 토큰 수 / 걸린 시간

python

import time
def bench(model, idx, n, use_cache):
    if idx.is_cuda: torch.cuda.synchronize()
    t0 = time.perf_counter()
    generate(model, idx, n, use_cache)
    if idx.is_cuda: torch.cuda.synchronize()
    return n / (time.perf_counter() - t0)

GPU에서는 synchronize를 빼면 시간이 실제보다 짧게 측정된다. 첫 실행은 준비 비용이 섞이므로 버리고 2~3번 평균을 낸다. 예상 경향: 생성이 길수록 캐시 없는 쪽이 눈에 띄게 느려진다. 실제 숫자는 환경마다 다르니 지어내지 말고 본인 결과를 적는다.

실험 2: GQA로 KV 캐시 메모리 감소량

바꾸는 값: n_kv_head ∈ {8(MHA), 4, 2, 1(MQA)}, n_head=8 고정
측정: 생성 후 캐시 텐서의 실제 크기

python

def cache_bytes(caches):
    return sum(k.numel() * k.element_size() + v.numel() * v.element_size()
               for k, v in caches)

이론값(위의 공식)과 측정값이 같은지 비교하고, MHA 대비 비율이 1, 1/2, 1/4, 1/8로 나오는지 확인한다. 가능하다면 같은 설정으로 짧게 학습시켜 검증 손실(validation loss)도 함께 적어, 메모리를 줄인 대가로 품질이 얼마나 변하는지 본다(작은 모델·짧은 학습의 차이는 노이즈일 수 있으니 단정하지 않는다).

기록표

n_kv_head	캐시 바이트(측정)	이론값	MHA 대비	토큰/초	검증 손실
8			1
4
2
1

왜 이걸 배우나 (LLM / Post-training 연결)

현대 오픈 LLM 상당수가 GQA 계열의 KV 헤드 공유를 쓰는 이유가 바로 오늘의 캐시 메모리 문제다. 추론 서버의 동시 처리량과 최대 문맥 길이는 KV Cache 크기에 직접 좌우된다. Post-training에서도 RLHF 류의 학습 중 모델이 답변을 수없이 샘플링(생성)해야 하므로 생성 속도가 곧 실험 속도가 되고, 평가 단계의 대량 생성도 같은 구조 위에서 돌아간다. LLM Research Engineer라면 모델 설정 파일의 head 수와 KV head 수를 보고 캐시 비용을 바로 계산할 수 있어야 한다.

초보자가 막히는 지점

캐시 ON/OFF 결과가 다르게 나옴: 대부분 위치 정보 오류다. 캐시 사용 시 새 토큰의 위치 인덱스가 0부터 시작하면 안 된다. 과거 길이만큼 이동해야 한다. 또 마스크를 T×T로 만들면 T=1일 때 과거 전체를 못 보거나 형상이 맞지 않는다.
GQA인데 메모리가 안 줄어듦: repeat_interleave로 복제한 K, V를 캐시에 저장하고 있을 가능성이 크다. 복제 전 텐서를 저장한다.
`view` 에러(shape mismatch): Q는 n_head, K/V는 n_kv_head로 reshape하는지 확인한다. n_head % n_kv_head == 0이어야 그룹이 나뉜다.
"캐시를 쓰면 모든 게 빨라진다"는 오해: 프롬프트를 처음 처리하는 prefill 비용은 그대로다. 그리고 캐시가 길어지면 매 단계 메모리 읽기가 늘어난다. 작은 모델·짧은 길이에서는 차이가 작게 나올 수 있다.
GPU 시간 측정 오류: 위에서 말한 synchronize 누락, 첫 실행 포함이 대표적이다.
`torch.cat`으로 캐시를 이어 붙이는 방식은 매번 새 메모리를 할당하므로 학습용 구현으로는 충분하지만 실서비스 방식은 아니다. 오늘은 원리 이해가 목적이다.

알고리즘 (병행 1시간)

LeetCode Medium 그래프 2문제를 푼다. 추천 유형은 아래 중 고른다(번호는 직접 검색해서 확인).

섬의 개수(Number of Islands): 격자 DFS/BFS
코스 스케줄(Course Schedule): 방향 그래프의 사이클 판정(위상 정렬)

풀이마다 다음을 기록한다.

접근: 그래프를 어떻게 표현했나(인접 리스트/격자), DFS인지 BFS인지와 그 이유, 방문 처리를 어디서 했나
시간·공간 복잡도: 정점 수 V, 간선 수 E 기준으로 직접 유도해서 적는다(예: 인접 리스트 순회는 O(V+E)가 되는 이유)
틀린 이유: 첫 제출이 실패했다면 원인(방문 표시 누락, 경계 조건, 재귀 깊이, 방향 착각 등)을 한 줄로 남기고, 다음에 확인할 체크 항목으로 바꿔 적는다