Gemma 4 MTP: 드래프터가 초안을 쓰고 타깃이 검증하는 추론 가속 구조

왜 디코딩이 느린가

LLM 추론에서 디코딩은 본질적으로 순차적이다. 토큰 하나를 생성하려면 전체 모델이 한 번의 포워드 패스를 완료해야 한다. GPU는 병렬 처리에 최적화됐지만, 토큰 하나씩 나오는 자기회귀(autoregressive) 디코딩 루프는 GPU를 충분히 포화시키지 못한다. 결과적으로 배치 크기가 작을수록, 그리고 모델이 클수록 디코딩 처리량 대비 GPU 활용률이 낮아진다.

2026년 5월 5일 Google은 이 문제에 대한 실용적인 답변을 Gemma 4에 탑재해 공개했다. Multi-Token Prediction(MTP) 드래프터 모델이다. 드래프터가 여러 토큰을 미리 예측하면 타깃 모델이 이를 한꺼번에 검증한다. 결과물의 품질은 동일하면서도 초당 토큰 수는 최대 3배 높아진다.

투기적 디코딩(Speculative Decoding)의 기본 구조

투기적 디코딩은 두 모델의 비대칭성을 이용한다.

드래프터(draft model): 작고 빠르다. 여러 토큰을 연속으로 예측한다.
타깃(target model): 크고 정확하다. 드래프터의 예측 전체를 병렬로 검증한다.

타깃 모델은 드래프터가 제안한 k개의 토큰을 한 번의 포워드 패스로 동시에 처리할 수 있다. 드래프터의 예측이 맞으면 k개를 한꺼번에 확정한다. 틀린 위치가 있으면 그 위치까지만 확정하고 드래프터를 재시작한다.

중요한 점은 결과의 무손실 보장이다. 타깃 모델이 거부한 토큰은 타깃이 직접 교정 토큰을 생성하므로, 드래프터 없이 생성한 결과와 분포적으로 동일하다. 속도만 빨라지고 품질 저하는 없다.

Gemma 4 MTP 투기적 디코딩 동작 흐름

Gemma 4 MTP 드래프터 모델 구성

Google은 Gemma 4의 네 가지 사이즈 각각에 대해 별도의 드래프터 체크포인트를 Hugging Face에 공개했다. 모델 ID는 google/gemma-4-{variant}-it-assistant 패턴을 따른다.

타깃 모델	드래프터 모델 ID	드래프터 hidden dim
Gemma 4 E2B	`google/gemma-4-E2B-it-assistant`	256
Gemma 4 E4B	`google/gemma-4-E4B-it-assistant`	256
Gemma 4 26B-A4B	`google/gemma-4-26B-A4B-it-assistant`	1024
Gemma 4 31B	`google/gemma-4-31B-it-assistant`	1024

"assistant" 접미사가 드래프터 체크포인트의 표시다. 드래프터는 타깃 모델보다 훨씬 작고, 타깃의 KV 캐시에 cross-attention으로 연결되도록 훈련됐다.

드래프터 내부 아키텍처

드래프터의 MTP 헤드는 다음 요소로 구성된다.

별도 embedder: 드래프터 전용 토큰 임베딩 행렬. 타깃 모델과 공유하지 않는다.
4-layer Transformer block: 슬라이딩 윈도 어텐션 3층 + 글로벌 어텐션 1층.
Cross-attention to target KVs: 각 드래프터 레이어가 타깃 모델의 KV 캐시를 직접 참조한다. 이를 통해 드래프터가 타깃의 문맥을 유지한 채 다음 토큰을 예측한다.

포워드 패스 입력은 "이전 스텝의 타깃 마지막 레이어 활성화 + 드래프터 임베딩"이다. 드래프터는 이를 기반으로 k개의 후보 토큰을 순차적으로 생성하고, 타깃이 단일 포워드 패스로 k+1개 위치를 동시에 검증한다.

Hugging Face Transformers로 사용하기

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 타깃 모델
target = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
# MTP 드래프터
drafter = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it-assistant",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-E4B-it")

inputs = tokenizer("데이터베이스 인덱스 설계 원칙은?", return_tensors="pt").to("cuda")

# assistant_model 인자 하나로 투기적 디코딩 활성화
outputs = target.generate(
    **inputs,
    assistant_model=drafter,
    max_new_tokens=512,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

assistant_model 인자를 generate()에 넘기면 투기적 디코딩이 자동으로 활성화된다. 코드 변경은 드래프터를 로드하고 인자 하나를 추가하는 것뿐이다. 결과 품질은 드래프터 없이 생성한 것과 동일하다.

vLLM과 Ollama 지원

vLLM에서는 --speculative-model 플래그로 MTP 드래프터를 지정한다.

vllm serve google/gemma-4-E4B-it \
  --speculative-model google/gemma-4-E4B-it-assistant \
  --num-speculative-tokens 5 \
  --speculative-draft-tensor-parallel-size 1

--num-speculative-tokens는 드래프터가 한 번에 생성하는 후보 토큰 수다. 일반적으로 4~6이 균형점이다. 드래프터는 텐서 병렬 크기를 타깃보다 작게 설정할 수 있다(--speculative-draft-tensor-parallel-size 1).

Ollama는 v0.23.1부터 DRAFT Modelfile 지시자를 지원한다.

FROM gemma4:e4b
DRAFT gemma4:e4b-assistant

실제 속도 향상: 무엇이 결정하는가

속도 향상 크기는 수락률(acceptance rate)이 결정한다. 드래프터가 예측한 토큰을 타깃이 얼마나 자주 수락하는지에 달려 있다.

수락률이 높아지는 상황:

반복적인 패턴이 많은 출력 (코드, 목록, 표)
도메인이 좁은 프롬프트 (전문 용어, 정형화된 텍스트)
문맥이 길고 드래프터가 패턴을 학습한 경우

수락률이 낮아지는 상황:

창의적 글쓰기, 다양한 어휘 사용
프롬프트와 크게 다른 도메인
temperature가 높은 샘플링

실측값: 코드 생성이나 구조화된 출력에서 1.8~3x, 자유로운 대화체에서 1.3~1.8x를 기대할 수 있다. 어느 경우에도 출력 품질은 변하지 않는다.

메모리 요구사항과 트레이드오프

드래프터도 VRAM을 차지한다. E2B/E4B용 드래프터는 hidden dim이 256으로 매우 작아 메모리 부담이 미미하다. 31B용은 dim 1024로 더 크지만 여전히 타깃의 수십 분의 일 수준이다.

타깃 모델과 드래프터를 동일한 GPU에 올릴 수 있는지 확인해야 한다. 모델이 두 개 올라가므로 KV 캐시에 할당 가능한 VRAM이 줄어든다. 배치 크기가 크고 긴 컨텍스트를 처리하는 서빙 환경이라면 메모리 예산 계획이 필요하다.

운영 도입 체크리스트

[ ] 타깃-드래프터 버전 쌍 확인: google/gemma-4-{variant}-it ↔ google/gemma-4-{variant}-it-assistant
[ ] 사용 중인 프레임워크 버전 확인: Transformers assistant_model, vLLM --speculative-model, Ollama DRAFT
[ ] 드래프터와 타깃을 같은 dtype(bfloat16 권장)으로 로드
[ ] VRAM 예산: 타깃 + 드래프터 + KV 캐시 합계 계산
[ ] num_speculative_tokens 실측 튜닝: 4에서 시작해 6까지 테스트
[ ] 출력 샘플 수십 개로 품질 동일성 확인 (무손실 보장이지만 구현별 검증)
[ ] throughput 벤치마크: 단일 요청 vs. 배치 요청 분리 측정

Open question

Ollama의 MLX 경로(Apple Silicon)에서 드래프터-타깃 크로스 어텐션을 어떻게 처리하는지 공식 문서에 명시되지 않았다. macOS에서의 드래프터 동작은 커뮤니티 실측이 필요하다.

References

https://ai.google.dev/gemma/docs/mtp/overview
https://ai.google.dev/gemma/docs/mtp/mtp
https://blog.google/innovation-and-ai/technology/developers-tools/multi-token-prediction-gemma-4/
https://www.infoq.com/news/2026/05/gemma4-multi-token-prediction/
https://huggingface.co/google/gemma-4-E4B-it-assistant
https://huggingface.co/google/gemma-4-31B-it-assistant
https://arxiv.org/pdf/2607.02770