Transformers 로컬 LLM 2배 가속 - KV 캐시·SpecDecode

서빙 환경이 아니라 로컬에서 LLM을 돌릴 때 가장 답답한 순간은 “한 토큰씩 찍히는 속도”입니다. GPU를 꽂아도 생각보다 느린 경우가 많고, 특히 긴 프롬프트를 반복해서 넣는 워크플로(에이전트/툴콜/리트라이)가 겹치면 지연이 눈덩이처럼 커집니다.

이 글은 Hugging Face Transformers 기준으로 추론 속도를 2배 수준까지 끌어올리는 대표적인 두 축을 다룹니다.

• KV 캐시(Key-Value cache): 디코딩 단계에서 이미 계산한 어텐션의 K/V를 재사용해, 매 토큰마다 과거 토큰을 다시 계산하지 않게 함
• Speculative Decoding(SpecDecode): 작은 드래프터 모델이 여러 토큰을 “미리 제안”하고, 큰 타깃 모델이 검증해 한 번에 여러 토큰을 확정

추가로, 실제로 속도가 안 나오는 원인(프롬프트 재처리, 캐시 미사용, 설정 충돌)을 체크리스트로 정리합니다. 메모리/VRAM 이슈가 같이 터지는 경우도 많아, 최적화 관점은 이미지 모델에서의 VRAM 최적화 경험과도 닮았습니다. 필요하면 Stable Diffusion VRAM OOM - xFormers·SDPA로 해결도 함께 참고하면 좋습니다.

1) 왜 로컬 LLM이 느린가: 디코딩의 비용 구조

LLM 추론은 크게 두 단계입니다.

1. Prefill(프롬프트 처리): 입력 프롬프트 전체 길이 L에 대해 한 번에 forward
2. Decode(생성): 토큰을 1개 생성할 때마다 forward를 반복(생성 길이 T만큼)

문제는 decode에서, 매 스텝마다 “지금까지의 모든 토큰”을 다시 어텐션에 넣으면 비용이 커진다는 점입니다. 이때 필요한 게 KV 캐시입니다.

• KV 캐시 없음: decode 스텝마다 과거 토큰의 K/V를 다시 계산
• KV 캐시 있음: 과거 토큰의 K/V를 저장하고, 새 토큰 1개에 대한 K/V만 추가

결과적으로 decode 단계의 연산량이 크게 줄어듭니다.

2) KV 캐시: 원리와 체감 포인트

KV 캐시가 실제로 줄여주는 것

Self-attention에서 각 레이어는 K/V를 만들고, 이후 쿼리와 곱해 가중합을 만듭니다. 생성 중에는 과거 토큰의 K/V가 바뀌지 않으므로, 이를 저장해두면 다음 토큰 생성 시 재사용할 수 있습니다.

• 장점: 생성 토큰당 연산량 감소, 토큰/초(token/s) 증가
• 단점: 메모리 사용량 증가(레이어 수, 헤드 수, 시퀀스 길이에 비례)

Transformers에서 KV 캐시 켜는 법(기본)

최근 Transformers는 기본적으로 generate에서 캐시를 사용하도록 설계되어 있지만, 설정/모델/코드 경로에 따라 꺼져 있거나 무효화되기도 합니다. 가장 안전한 패턴은 명시적으로 use_cache=True를 주는 것입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-3.1-8B-Instruct"  # 예시

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)
model.generation_config.use_cache = True

prompt = "로컬 LLM 속도를 높이는 방법을 3가지로 요약해줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=False,
        use_cache=True,
    )

print(tokenizer.decode(out[0], skip_special_tokens=True))

캐시가 “안 먹는” 흔한 함정 5가지

1. use_cache=False로 꺼짐: 모델 config나 generation config에서 비활성화
2. 훈련용 설정이 남아 있음: gradient_checkpointing 같은 옵션이 켜져 있으면 캐시와 충돌하는 경우가 있음
3. 프롬프트를 매번 새로 토크나이즈하고, 매번 prefill을 반복: 캐시는 “한 번의 생성 루프” 안에서 효과가 크며, 대화 히스토리를 매번 다시 넣으면 prefill 비용이 누적
4. 너무 긴 컨텍스트: 캐시는 메모리를 먹기 때문에 VRAM이 부족하면 오히려 느려지거나 OOM
5. 백엔드 커널 비최적: PyTorch/드라이버/FlashAttention(SDPA) 조합에 따라 토큰/초가 크게 달라짐

특히 3번은 에이전트/툴콜에서 자주 발생합니다. “대화 이력 전체를 매번 다시 넣는 구조”를 유지하되, 가능한 경우 프롬프트를 짧게 유지하거나 요약을 끼워 넣는 전략이 필요합니다. 에이전트가 무한 루프나 리트라이를 반복하는 상황이면 속도 최적화가 더 중요해지므로, LangChain 에이전트 무한 루프 끊는 실전 디버깅도 같이 보면 실전에서 도움이 됩니다.

3) KV 캐시를 더 잘 쓰는 실전 패턴

(1) 스트리밍 출력으로 “체감 지연” 줄이기

토큰/초가 동일해도, 첫 토큰이 늦게 나오면 느리게 느껴집니다. 스트리밍은 UX를 크게 개선합니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

model_id = "mistralai/Mistral-7B-Instruct-v0.3"  # 예시

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)

streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

inputs = tokenizer("한 문단으로 KV 캐시를 설명해줘.", return_tensors="pt").to(model.device)

with torch.inference_mode():
    _ = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=False,
        use_cache=True,
        streamer=streamer,
    )

(2) 배치/동시성에서는 캐시가 병목이 될 수 있음

로컬에서 여러 요청을 동시에 처리하면, KV 캐시가 요청별로 쌓이며 VRAM이 빠르게 찹니다. 이때는 다음 중 하나를 선택해야 합니다.

• 동시성 낮추고 단일 요청 토큰/초를 높이기
• 컨텍스트 길이 제한을 강하게 걸기
• 더 공격적으로는 양자화(4bit/8bit) 또는 작은 모델 + SpecDecode로 전환

4) Speculative Decoding: “큰 모델”을 덜 돌리는 방법

KV 캐시는 “같은 큰 모델을 더 효율적으로” 쓰는 방법이라면, SpecDecode는 “큰 모델 호출 횟수 자체를 줄이는” 방법입니다.

동작 개요

1. 작은 드래프터(drafter) 모델이 앞으로 나올 토큰을 k개 정도 빠르게 제안
2. 큰 타깃(target) 모델이 그 제안을 검증
3. 맞으면 한 번에 여러 토큰을 확정, 틀리면 일부만 확정하고 다시 진행

핵심은 드래프터가 충분히 정확하면, 타깃 모델의 비싼 forward를 더 적게 수행해도 동일한 품질을 얻는다는 점입니다.

언제 효과가 큰가

• 큰 모델이 무겁고(예: 30B 이상), 작은 모델이 상대적으로 매우 빠를 때
• 출력이 비교적 “예측 가능”할 때(요약/정리/형식화)
• 샘플링을 과하게 쓰지 않을 때(높은 temperature는 일치율을 떨어뜨릴 수 있음)

반대로, 창의적 생성(높은 temperature)이나 도메인 특화(드래프터가 잘 모르는 분포)에서는 이득이 줄 수 있습니다.

5) Transformers에서 SpecDecode 적용 예시

Transformers는 버전에 따라 speculative decoding 지원 방식이 달라질 수 있습니다. 가장 중요한 포인트는 다음입니다.

• 타깃 모델과 드래프터 모델을 각각 로드
• 생성 시 speculative 관련 인자를 켜기

아래 코드는 “구조를 이해하기 위한 예시”입니다. 실제 사용 시에는 설치된 Transformers 버전의 generate 시그니처에 맞춰 파라미터 이름을 확인하세요.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

target_id = "meta-llama/Llama-3.1-8B-Instruct"   # 큰 모델
rafter_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # 작은 모델(예시)

tokenizer = AutoTokenizer.from_pretrained(target_id, use_fast=True)

target = AutoModelForCausalLM.from_pretrained(
    target_id,
    torch_dtype=torch.float16,
    device_map="auto",
)

drafter = AutoModelForCausalLM.from_pretrained(
    rafter_id,
    torch_dtype=torch.float16,
    device_map="auto",
)

prompt = "다음 내용을 5줄 bullet로 요약해줘: KV 캐시와 SpecDecode의 차이"
inputs = tokenizer(prompt, return_tensors="pt").to(target.device)

with torch.inference_mode():
    out = target.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=False,
        use_cache=True,
        # 아래 옵션들은 버전에 따라 이름이 다를 수 있음
        # speculative_draft_model=drafter,
        # num_speculative_tokens=5,
    )

print(tokenizer.decode(out[0], skip_special_tokens=True))

버전 차이로 speculative 옵션이 바로 동작하지 않는다면, 대안은 다음과 같습니다.

• Transformers 최신 버전으로 업데이트
• text-generation-inference 같은 서빙 엔진의 SpecDecode 지원 확인
• vLLM 등 다른 런타임의 speculative 기능 사용(로컬에서도 가능)

6) “2배”를 만들기 위한 조합 전략

KV 캐시와 SpecDecode는 경쟁 관계가 아니라 서로 보완합니다.

• KV 캐시: 디코딩 루프에서 토큰당 비용을 낮춤
• SpecDecode: 큰 모델의 디코딩 스텝 수 자체를 줄임

실전에서 2배를 목표로 할 때는 아래 순서가 성공 확률이 높습니다.

1. KV 캐시가 확실히 켜졌는지 확인(use_cache=True 명시)
2. SDPA/FlashAttention 경로가 활성화되어 있는지 확인(환경/버전 영향 큼)
3. 프롬프트가 길다면 prefill 비용을 줄이는 설계(요약/컨텍스트 슬라이딩)
4. 여전히 느리면 SpecDecode로 큰 모델 호출 빈도 줄이기

이 과정에서 VRAM이 빠듯해지면, 캐시가 늘어날수록 OOM 가능성이 커집니다. 이미지 생성에서 xFormers/타일링으로 VRAM을 아끼는 것처럼, LLM도 “빠르게 하려다 메모리로 막히는” 패턴이 자주 나옵니다. 비슷한 문제 해결 감각은 Stable Diffusion VRAM OOM - xFormers·VAE 타일링에서 얻을 수 있습니다.

7) 속도 측정: token/s를 직접 찍어보기

최적화는 측정 없이는 감으로 하게 됩니다. 아래처럼 간단히 토큰/초를 측정해두면, 설정 변경의 효과가 명확해집니다.

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "mistralai/Mistral-7B-Instruct-v0.3"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)
model.generation_config.use_cache = True

prompt = "로컬 LLM 추론을 빠르게 하는 방법을 단계별로 정리해줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

max_new = 256

start = time.perf_counter()
with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=max_new,
        do_sample=False,
        use_cache=True,
    )
end = time.perf_counter()

# 생성된 새 토큰 수를 대략 계산(정확히 하려면 입력 길이 차감)
input_len = inputs["input_ids"].shape[1]
out_len = out.shape[1]
new_tokens = out_len - input_len

tps = new_tokens / (end - start)
print(f"new_tokens={new_tokens}, seconds={end-start:.3f}, token_per_sec={tps:.2f}")

측정 시 주의점:

• 첫 실행은 커널 컴파일/캐시 워밍업으로 느릴 수 있어 2~3회 반복 후 평균을 보세요.
• max_new_tokens를 너무 작게 잡으면 오차가 커집니다.
• 스트리밍 출력은 I/O 때문에 token/s 측정에는 불리할 수 있습니다(UX용과 벤치용을 분리).

8) 마무리: 빠르게 만들되, 실패 모드도 함께 관리하기

로컬 LLM 최적화는 “한 방 세팅”보다, 병목을 분해해서 하나씩 제거하는 쪽이 확실합니다.

• KV 캐시는 기본 중의 기본이고, 꺼져 있으면 손해가 큽니다.
• SpecDecode는 큰 모델을 덜 돌리는 방식이라, 조건이 맞으면 체감 가속이 큽니다.
• 다만 둘 다 메모리/환경/버전 영향을 강하게 받으므로, 반드시 token/s로 측정하면서 진행하세요.

다음 글을 확장으로 쓴다면, 로컬 런타임별(vLLM, llama.cpp, TGI)로 KV 캐시 구현 차이와 SpecDecode 지원 현황, 그리고 “긴 컨텍스트에서 캐시 메모리 폭증”을 어떻게 다루는지까지 비교해보는 것도 좋습니다.