[vLLM] Speculative Decoding: 드래프트 모델로 LLM 디코딩을 가속하는 원리
들어가며
LLM 디코딩은 본질적으로 순차적이다. 한 토큰을 생성해야 다음 토큰을 생성할 수 있다. 그런데 대부분의 디코딩 스텝에서 GPU 연산 자원은 크게 남는다(memory-bound). Speculative Decoding은 이 남는 자원을 활용한다. 작은 드래프트 모델이 K개 토큰을 빠르게 예측하고, 큰 타겟 모델이 이를 한 번에 검증하여 맞는 만큼 수락(accept)한다. 수학적으로 출력 분포가 원래 모델과 동일함이 보장된다.
- 논문: Fast Inference from Transformers via Speculative Decoding (arxiv 2211.17192)
- 공식 문서: https://docs.vllm.ai
공식 문서
vLLM 공식 문서: Speculative Decoding
핵심 구조/코드 분석
Spec Decode 모듈 구조
vLLM의 speculative decoding 구현은 vllm/v1/spec_decode/ 디렉토리에 있다:
spec_decode/
├── eagle.py # EAGLE 기반 proposer (기본 클래스 포함)
├── draft_model.py # 독립 드래프트 모델 proposer
├── medusa.py # Medusa 헤드 기반 proposer
├── ngram_proposer.py # N-gram 기반 proposer (모델 불필요)
├── metadata.py # Spec decode 메타데이터
├── metrics.py # 수락률 등 메트릭
└── utils.py # 유틸리티 함수
SpecDecodeBaseProposer: 공통 기반 클래스
모든 speculative decoding proposer의 기반이 되는 SpecDecodeBaseProposer가 eagle.py에 정의되어 있다:
class SpecDecodeBaseProposer:
def __init__(self, vllm_config, device, pass_hidden_states_to_model, ...):
self.num_speculative_tokens = (
self.speculative_config.num_speculative_tokens
)
self.hidden_size = self.draft_model_config.get_hidden_size()
# 순차 드래프팅 vs 병렬 드래프팅
self.parallel_drafting = self.speculative_config.parallel_drafting
self.extra_slots_per_request = (
1 if not self.parallel_drafting
else self.num_speculative_tokens
)
# 트리 구조 드래프팅
spec_token_tree = self.speculative_config.speculative_token_tree
self.tree_choices = ast.literal_eval(spec_token_tree)
num_speculative_tokens는 한 번에 예측할 드래프트 토큰 수이다. parallel_drafting이 활성화되면 모든 드래프트 토큰을 한 번의 forward로 생성한다.
DraftModelProposer: 독립 드래프트 모델
가장 직관적인 방식은 별도의 작은 모델을 드래프트로 사용하는 것이다:
class DraftModelProposer(SpecDecodeBaseProposer):
def __init__(self, vllm_config, device, runner=None):
super().__init__(
vllm_config=vllm_config,
device=device,
pass_hidden_states_to_model=False, # 타겟 모델과 독립
runner=runner,
)
self._raise_if_vocab_size_mismatch()
self._raise_if_draft_tp_mismatch()
def _raise_if_draft_tp_mismatch(self):
tgt_tp = self.speculative_config.target_parallel_config \
.tensor_parallel_size
draft_tp = self.speculative_config.draft_parallel_config \
.tensor_parallel_size
if draft_tp != tgt_tp:
raise ValueError(
"Currently, draft_tensor_parallel_size and "
"tensor_parallel_size must be the same."
)
드래프트 모델과 타겟 모델의 vocab 크기가 같아야 하고, 현재는 TP 크기도 동일해야 한다는 제약이 있다.
스케줄러와의 연동
스케줄러는 speculative decoding을 인식하고 추가 슬롯을 할당한다:
# scheduler.py
speculative_config = vllm_config.speculative_config
self.num_spec_tokens = self.num_lookahead_tokens = 0
if speculative_config:
self.num_spec_tokens = speculative_config.num_speculative_tokens
if speculative_config.use_eagle():
self.use_eagle = True
self.num_lookahead_tokens = self.num_spec_tokens
# 블록 할당 시 lookahead 토큰 반영
new_blocks = self.kv_cache_manager.allocate_slots(
request, num_new_tokens,
num_lookahead_tokens=self.num_lookahead_tokens,
)
num_lookahead_tokens만큼 추가 KV 블록을 미리 할당하여, 드래프트 토큰이 수락되었을 때 바로 사용할 수 있게 한다.
트리 구조 드래프팅
단순히 1개 시퀀스를 드래프팅하는 대신, 트리 구조로 여러 가지를 동시에 탐색할 수 있다:
spec_token_tree = self.speculative_config.speculative_token_tree
self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree)
tree_depth = len(self.tree_choices[-1])
# 레벨별 드래프트 수 계산
num_drafts_per_level = [0] * tree_depth
for node in self.tree_choices:
num_drafts_per_level[len(node) - 1] += 1
예를 들어 tree_choices = [(0,), (0,0), (0,1), (1,)]이면 첫 번째 레벨에서 2개, 두 번째 레벨에서 2개 후보를 생성한다. 수락률이 높아지는 대신 연산량이 늘어나므로, 모델과 워크로드에 맞게 트리 구조를 조절한다.
왜 이 설계인가
-
무손실 가속: 거부 샘플링(rejection sampling) 알고리즘 덕분에 출력 분포가 타겟 모델과 수학적으로 동일하다. 품질 저하 없이 2-3배 속도 향상이 가능하다.
-
다양한 Proposer 지원: 독립 드래프트 모델, EAGLE, Medusa, N-gram 등 다양한 드래프팅 방식을 동일한 인터페이스로 지원한다.
-
트리 드래프팅: 단일 시퀀스 대신 트리 구조로 여러 후보를 탐색하여 수락률을 높인다. GPU의 남는 병렬 처리 능력을 효과적으로 활용한다.
-
스케줄러 레벨 통합: spec_token_ids가 스케줄러에서 일급 시민으로 관리되어, KV 캐시 할당과 토큰 예산이 자연스럽게 연동된다.
Speculative decoding은 "디코딩이 memory-bound"라는 LLM 서빙의 근본적 특성을 활용한 가속 기법이며, vLLM은 다양한 proposer를 플러그인 형태로 지원하는 유연한 구조를 갖추고 있다.
논문 핵심 내용
Speculative Decoding 논문(Leviathan et al., 2022)의 가장 중요한 기여는 두 가지다. 첫째, 작은 드래프트 모델의 예측을 큰 타겟 모델이 검증하는 방식으로 출력 분포의 동일성을 수학적으로 보장하면서 디코딩을 가속할 수 있음을 증명했다. 둘째, 이 방식이 실제로 유의미한 벽시계 속도 향상을 가져온다는 것을 실험으로 보여줬다.
핵심 수치:
| 항목 | 결과 |
|---|---|
| T5-XXL 모델 가속 | 표준 T5X 구현 대비 2-3배 속도 향상 |
| 출력 품질 | 원본 모델과 분포 동일 (수학적 증명) |
| 재학습 필요 | 불필요 (기존 모델 그대로 사용) |
| 아키텍처 수정 | 불필요 |
이 논문의 핵심 통찰은 LLM의 autoregressive 디코딩이 memory-bound라는 점이다. 한 토큰을 생성할 때 GPU 연산 자원의 대부분이 놀고 있다. Speculative decoding은 이 남는 연산 자원을 활용하여, 드래프트 모델이 예측한 K개 토큰을 타겟 모델이 한 번의 forward pass로 병렬 검증한다. 검증은 modified rejection sampling으로 수행되어 수락된 토큰의 분포가 원래 모델의 분포와 정확히 일치한다. 품질 저하가 전혀 없는 순수한 가속이라는 점이 이 기법의 가장 강력한 장점이다.
관련 포스트
vLLM 의 다른글
- 이전글 [vLLM] Multi-head Latent Attention: KV 캐시를 압축하는 DeepSeek의 어텐션
- 현재글 : [vLLM] Speculative Decoding: 드래프트 모델로 LLM 디코딩을 가속하는 원리
- 다음글 [vLLM] EAGLE: 은닉 상태 기반 드래프트로 Speculative Decoding을 강화하다
댓글