추측 디코딩과 KV 캐시 최적화의 자리

추측 디코딩은 서빙 최적화야

작은 초안 모델이 후보 토큰 K개를 빠르게 만들면 큰 검증 모델이 한꺼번에 확인하고, 자신이 냈을 접두 부분만 받아들여. 결과 분포는 일반 디코딩과 수학적으로 같아. 병렬 연산을 조금 더 쓰는 대신 순차 호출을 줄여 지연 시간을 낮추는 방식이야.

두 모델의 가중치는 그대로야

추측 디코딩은 완성된 모델 둘을 함께 실행하는 런타임 기법이야. 초안 토큰을 언제 받아들이고 버릴지는 오케스트레이션 코드가 결정하며 어느 모델의 백본도 바뀌지 않아.

GQA, MQA, MLA는 KV 캐시를 줄여

이들은 어텐션 층 안의 실제 변화지만 새 계열이라기보다 트랜스포머 안에서 이뤄진 정제에 가까워.

MHA: 각 어텐션 헤드가 자기 K와 V 투영을 가져.
MQA: 모든 헤드가 K와 V 투영 하나를 공유해 KV 캐시를 가장 공격적으로 줄여.
GQA: 헤드를 그룹으로 묶고 그룹 안에서 K와 V를 공유해 MHA와 MQA 사이를 절충해. Llama 3, Gemma, Qwen이 써.
MLA: DeepSeek가 K와 V를 저차원 잠재 공간으로 압축했다가 필요할 때 복원해 캐시를 줄이는 방식이야.

서빙에는 크지만 분류표에는 작은 변화

KV 캐시가 줄면 긴 컨텍스트와 다중 사용자 추론에서 큰 이득을 얻어. 그래도 Query와 Key의 유사도를 softmax로 가중해 Value에 적용하는 기본 패턴은 남아 있어. 같은 아키텍처 계열 안의 중요한 공학 선택이야.

기본 계산이 남았는지 물어

새 어텐션 변형을 봤다면 Q×K→softmax×V라는 기본 패턴이 유지되는지 확인해. 유지된다면 트랜스포머 어텐션의 정제이고, 사라진다면 Mamba나 RWKV처럼 정말 다른 설계를 보고 있을 수 있어.

Code

추측 디코딩 골격·python

def speculative_decode(prompt, draft_model, verifier_model, k=4):
    output = []
    while not done(output):
        draft_tokens = draft_model.generate(prompt + output, max_new=k)
        verifier_logprobs = verifier_model.score(prompt + output, draft_tokens)
        accepted = []
        for tok, lp in zip(draft_tokens, verifier_logprobs):
            if accept(tok, lp):
                accepted.append(tok)
            else:
                accepted.append(verifier_model.sample_one(prompt + output + accepted))
                break
        output.extend(accepted)
    return output
# Output is identical in distribution to verifier_model.generate alone.

MHA, GQA, MQA의 KV 헤드 수·python

# MHA: num_kv_heads == num_q_heads (e.g., 32 query heads, 32 KV heads)
# GQA: num_kv_heads <  num_q_heads (e.g., 32 query heads,  8 KV heads in Llama 3)
# MQA: num_kv_heads == 1            (e.g., 32 query heads,  1 KV head)

추측 디코딩은 서빙 최적화야

두 모델의 가중치는 그대로야

GQA, MQA, MLA는 KV 캐시를 줄여

서빙에는 크지만 분류표에는 작은 변화

기본 계산이 남았는지 물어

Code

External links

Exercise

Progress

댓글 0