Dynamic Speculative Decoding

Dynamic Speculative Decoding (동적 추측 디코딩)

Dynamic Speculative Decoding(DSD)은 동시성(concurrency, 배치 크기)에 따라 추측 토큰 수 K를 동적으로 조절해서, 배치가 커져도 추측 디코딩의 이점을 계속 유지하게 해주는 기능이에요. 고정된 K로는 배치가 커질 때 오히려 검증 연산이 늘어나 디코드 속도(TPOT)가 나빠질 수 있는데, DSD가 K를 최적값으로 튜닝해 이를 해결해요.

출처: 문서

본문

왜 Dynamic SD가 필요한가?

SD 방법은 디코딩 중 각 시퀀스에 대해 K개의 토큰을 검증해야 해요. 배치 크기(BS)가 커지면 유효 배치 크기가 BS*K가 되어 검증 중 연산 요구량이 늘어나요. 이 BS*K가 임계 배치 크기를 넘어서면 SD는 디코드 속도(TPOT)에 부정적인 영향을 줘요. DSD는 K를 최적값으로 튜닝해서 SD의 이점을 계속 누리게 해줘요.

사용 사례

변동 동시성 워크로드: 같은 배포에서 동시성이 변하는 경우. 동시성이 커지면 K가 감소해요.
RL 롤아웃 중: 높은 BS로 시작하지만, 매우 적은 수의 긴 꼬리(long tail) 요청이 많은 토큰을 생성해 현재 롤아웃의 진행을 지연시키면서 BS가 작아지는 경우. 여기서는 롤아웃이 끝날 무렵 K가 올라가요.

--speculative-config 스키마

Dynamic SD를 사용하려면 SD 방법의 config에 num_speculative_tokens_per_batch_size를 추가하면 돼요. 이 값은 리스트의 리스트이며, 각 항목은 [start_bs, end_bs, optimal_K] 형태예요. 즉 동시성이 [start_bs, end_bs] 범위에 있을 때 optimal_K개의 드래프트 토큰을 사용한다는 뜻이에요. 예를 들어:

--speculative-config '{
    "method": "eagle",
    "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
    "num_speculative_tokens": 3,
    "num_speculative_tokens_per_batch_size": [
      [1, 64, 3],
      [65, 128, 1],
      [129, 512, 0]
    ]
  }'

이것은 다음을 의미해요:

동시성이 [1, 64] 범위이면 K=3 사용
동시성이 [65, 128] 범위이면 K=1 사용
동시성이 [129, 512] 범위이면 K=0 사용, 즉 드래프트 토큰을 생성하지 않음

온라인 예시

Dynamic SD Eagle Drafter

VLLM_USE_V2_MODEL_RUNNER=0 \
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --speculative-config '{
    "method": "eagle",
    "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
    "num_speculative_tokens": 3,
    "num_speculative_tokens_per_batch_size": [
      [1, 64, 3],
      [65, 128, 1],
      [129, 512, 0]
    ]
  }'

Dynamic SD Eagle3 Drafter

VLLM_USE_V2_MODEL_RUNNER=0 \
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --speculative-config '{
    "method": "eagle3",
    "model": "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
    "num_speculative_tokens": 3,
    "num_speculative_tokens_per_batch_size": [
      [1, 16, 5],
      [17, 32, 4],
      [33, 64, 3],
      [65, 128, 1],
      [129, 512, 0]
    ]
  }'

제한 사항

Eagle, Eagle-3, DFlash에서 테스트되었어요. 다른 SD 방법은 바로 동작할 수도, 아닐 수도 있어요.
전체 Cudagraph는 Model Runner V2에서만 동작해요. MRv1은 이 기능과 함께 piece-wise cuda graph만 지원해요.
데이터 병렬화(--data-parallel-size > 1)와는 호환되지 않아요. 각 DP rank가 독립적으로 스케줄링하므로 rank마다 다른 K 값을 골라 DP 컬렉티브 불일치와 데드락이 발생할 수 있어요. DP가 활성화되면 vLLM은 자동으로 num_speculative_tokens_per_batch_size를 비활성화하고 정적 num_speculative_tokens 값으로 폴백해요.

Dynamic Speculative Decoding