Draft Models

Draft Models (드래프트 모델)

드래프트 모델(draft model) 방식은 작고 빠른 별도의 모델이 여러 개의 토큰을 먼저 제안(speculate)하고, 큰 타깃 모델이 이를 검증하는 추측 디코딩 방법이에요. 아래 코드는 vLLM을 오프라인 모드로 구성해 드래프트 모델로 추측 디코딩을 수행하고, 한 번에 5개 토큰을 추측하는 예시예요.

출처: 문서

본문

오프라인 모드 예시

from vllm import LLM, SamplingParams

prompts = ["The future of AI is"]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="Qwen/Qwen3-8B",
    tensor_parallel_size=1,
    speculative_config={
        "model": "Qwen/Qwen3-0.6B",
        "num_speculative_tokens": 5,
        "method": "draft_model",
    },
)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

온라인(서버) 모드 예시

동일한 구성을 온라인 모드로 실행하려면 서버 코드는 다음과 같아요:

vllm serve Qwen/Qwen3-4B-Thinking-2507 \
    --host 0.0.0.0 \
    --port 8000 \
    --seed 42 \
    -tp 1 \
    --max-model-len 2048 \
    --gpu-memory-utilization 0.8 \
    --speculative-config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'

클라이언트로 완성(completion)을 요청하는 코드는 동일하게 유지돼요:

from openai import OpenAI

# Modify OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"

client = OpenAI(
    # defaults to os.environ.get("OPENAI_API_KEY")
    api_key=openai_api_key,
    base_url=openai_api_base,
)

models = client.models.list()
model = models.data[0].id

# Completion API
stream = False
completion = client.completions.create(
    model=model,
    prompt="The future of AI is",
    echo=False,
    n=1,
    stream=stream,
)

print("Completion results:")
if stream:
    for c in completion:
        print(c)
else:
    print(completion)

교차 어휘(heterogeneous vocab)를 사용하는 드래프트 모델

기본적으로 vLLM은 드래프트와 타깃 모델이 동일한 어휘(vocabulary)를 공유할 것을 요구해요. use_heterogeneous_vocab: true를 설정하면 Token-Level Intersection (TLI) 알고리즘이 활성화되어, 다른 tokenizer를 가진 다른 모델 계열의 드래프트 모델을 사용할 수 있게 돼요.

현재 use_heterogeneous_vocab는 draft_sample_method='greedy'(기본값)를 요구해요. 확률적 드래프트 샘플링은 아직 미지원이며 향후 릴리스에서 추가될 예정이에요.

from vllm import LLM, SamplingParams

llm = LLM(
    model="Qwen/Qwen3-8B",
    speculative_config={
        "method": "draft_model",
        "model": "HuggingFaceTB/SmolLM2-135M-Instruct",
        "num_speculative_tokens": 3,
        "use_heterogeneous_vocab": True,
    },
    gpu_memory_utilization=0.5,
)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

경고: 추측 디코딩과 관련된 모든 설정은 --speculative-config를 사용해 지정하세요. 기존에 --speculative-model로 모델을 지정하고 --num-speculative-tokens 같은 파라미터를 따로 추가하던 방식은 deprecated 되었어요. 지원 키와 예시는 --speculative-config 스키마를 참고하세요.

더 알아보기 (Learn more)