본문으로 건너뛰기

[vLLM] Speculative Decoding: 드래프트 모델로 LLM 디코딩을 가속하는 원리

들어가며

LLM 디코딩은 본질적으로 순차적이다. 한 토큰을 생성해야 다음 토큰을 생성할 수 있다. 그런데 대부분의 디코딩 스텝에서 GPU 연산 자원은 크게 남는다(memory-bound). Speculative Decoding은 이 남는 자원을 활용한다. 작은 드래프트 모델이 K개 토큰을 빠르게 예측하고, 큰 타겟 모델이 이를 한 번에 검증하여 맞는 만큼 수락(accept)한다. 수학적으로 출력 분포가 원래 모델과 동일함이 보장된다.

공식 문서

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의 기반이 되는 SpecDecodeBaseProposereagle.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개 후보를 생성한다. 수락률이 높아지는 대신 연산량이 늘어나므로, 모델과 워크로드에 맞게 트리 구조를 조절한다.

왜 이 설계인가

  1. 무손실 가속: 거부 샘플링(rejection sampling) 알고리즘 덕분에 출력 분포가 타겟 모델과 수학적으로 동일하다. 품질 저하 없이 2-3배 속도 향상이 가능하다.

  2. 다양한 Proposer 지원: 독립 드래프트 모델, EAGLE, Medusa, N-gram 등 다양한 드래프팅 방식을 동일한 인터페이스로 지원한다.

  3. 트리 드래프팅: 단일 시퀀스 대신 트리 구조로 여러 후보를 탐색하여 수락률을 높인다. GPU의 남는 병렬 처리 능력을 효과적으로 활용한다.

  4. 스케줄러 레벨 통합: 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 의 다른글