[vllm] vLLM의 동적 추측 디코딩(Dynamic Speculative Decoding) 도입

2026년 6월 14일수정: 2026년 6월 14일

PR 링크: vllm-project/vllm#32374 상태: Merged | 변경: +586 / -7

들어가며

추측 디코딩(Speculative Decoding, SD)은 작은 드래프트 모델을 사용하여 여러 토큰을 미리 생성하고, 타겟 모델이 이를 검증함으로써 추론 속도를 높이는 강력한 기법입니다. 하지만 기존 SD 방식은 고정된 수의 추측 토큰(K)을 사용합니다. 배치 크기(Batch Size, BS)가 커지면, 검증 과정에서 필요한 연산량이 BS * K로 증가하여 오히려 전체 추론 지연 시간(TPOT)이 늘어나는 병목 현상이 발생합니다. 본 PR은 이러한 문제를 해결하기 위해 Dynamic Speculative Decoding (DSD)을 도입하여, 현재 배치 크기에 따라 최적의 K 값을 실시간으로 조정함으로써 고부하 상황에서도 효율적인 추론 성능을 유지하도록 개선했습니다.

코드 분석

1. 스케줄러의 동적 K 결정 (`vllm/v1/core/sched/scheduler.py`)

기존에는 고정된 K 값을 사용했으나, 이제는 num_speculative_tokens_per_batch_size 설정을 통해 배치 크기 범위별로 K를 다르게 할당합니다.

Before:

# 기존에는 설정된 고정 K 값을 사용
self.num_spec_tokens = num_speculative_tokens

After:

# 배치 크기에 따라 동적으로 K를 조회하는 lookup 테이블 활용
self.dynamic_sd_lookup = build_dynamic_sd_schedule_lookup(
    num_speculative_tokens_per_batch_size=speculative_config.num_speculative_tokens_per_batch_size,
    vllm_max_batch_size=max_batch_size,
    vllm_num_speculative_tokens=num_speculative_tokens,
)

2. 모델 러너의 드래프트 토큰 처리 (`vllm/v1/worker/gpu_model_runner.py`)

모델 러너는 이제 스케줄러가 전달한 최적의 K 값에 맞춰 드래프트 토큰을 처리합니다. 특히, 고정된 패딩 대신 가변적인 K를 처리하기 위해 prev_num_spec_tokens를 추적합니다.

After (핵심 변경):

# 이전 단계의 K를 저장하여 다음 단계에서 정확한 스트라이드로 인덱싱
self.prev_num_spec_tokens = current_k
# ...
# 드래프트 토큰을 CPU로 복사할 때 최적화된 K 사용
draft_token_ids = draft_token_ids[:, :self.prev_num_spec_tokens]

왜 이게 좋은가

이 최적화는 '연산 효율성'과 '처리량' 사이의 균형을 맞춥니다.

성능 향상: 벤치마크 결과, 배치 크기가 256인 고부하 상황에서 기존 SD는 Vanilla 모델보다 느려지는 현상이 발생했으나, DSD는 K를 0으로 낮추어 Vanilla 수준의 성능을 유지하거나 더 나은 TPOT를 보여줍니다.
유연성: 사용자는 num_speculative_tokens_per_batch_size를 통해 하드웨어와 모델 특성에 맞는 최적의 K 스케줄을 직접 정의할 수 있습니다.
일반적 교훈: 고정된 하이퍼파라미터는 특정 환경에서는 최적일 수 있으나, 가변적인 워크로드(예: RL rollout, 트래픽 변동)에서는 시스템 상태(배치 크기 등)에 따라 파라미터를 동적으로 조정하는 것이 필수적입니다.

리뷰어 피드백 반영

초기 설계는 오프라인 프로파일링 파이프라인을 포함하여 복잡도가 높았으나, 리뷰어들의 피드백을 수용하여 '설정 기반의 단순한 범위 매핑' 방식으로 변경되었습니다. 이로 인해 코드 유지보수성이 크게 향상되었으며, LucasWilkinson과 benchislett의 제안에 따라 스케줄러가 직접 K를 결정하고 모델 러너에 전달하는 구조로 리팩토링되어 설계의 깔끔함과 확장성을 확보했습니다.