[vllm] DFlash - Block Diffusion 기반 Speculative Decoding

PR 링크: vllm-project/vllm#36847 상태: Merged | 변경: +1,577/-107

들어가며

DFlash는 Block Diffusion 기반의 speculative decoding 기법으로, 기존 Eagle이나 NGram 방식보다 더 긴 acceptance length를 달성한다. 이 PR은 DFlash를 vLLM의 spec decode 프레임워크에 통합하여, Qwen3 계열 모델에서 최대 6.54의 평균 acceptance length를 기록한다.

핵심 코드 분석

DFlash Draft Model 등록

# tests/models/registry.py
"DFlashDraftModel": _HfExamplesInfo(
    "Qwen/Qwen3.5-4B",
    speculative_model="z-lab/Qwen3.5-4B-DFlash",
    use_original_num_layers=True,  # DFlash는 >1 layer 사용
    max_model_len=8192,
    max_num_seqs=32,
),

DFlash draft 모델은 기존 spec decode 모델과 달리 여러 레이어를 사용한다. use_original_num_layers=True로 설정하여 전체 레이어를 유지한다.

E2E 검증: 논문 대비 Acceptance Rate

# 논문(arxiv:2602.06036) Table 1 기준
expected_acceptance_lengths = {
    "mt-bench": 4.24,
    "humaneval": 6.50,
    "gsm8k": 6.54 * 0.95,  # subset이므로 5% tolerance
}

# 설정: Temp=0, MaxTokens=2048 (논문 재현 조건)
spec_llm.generate(
    [prompt_text],
    SamplingParams(temperature=0, max_tokens=2048),
    use_tqdm=False,
)

논문의 벤치마크 결과를 직접 테스트에 포함시켜, 구현의 정확성을 검증한다. 90% 이상의 acceptance length를 달성해야 테스트를 통과한다.

점진적 메트릭 계산

# Before: 누적 메트릭만 계산
n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value

# After: 이전 메트릭과의 차분으로 구간별 acceptance rate 계산
if prev_metrics is not None:
    n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value
    n_accepted_toks -= prev_name2metric["vllm:spec_decode_num_accepted_tokens"].value

왜 이게 좋은가

1. 높은 Acceptance Length: GSM8k에서 6.54, HumanEval에서 6.50으로 기존 방식 대비 2-3배 향상
2. 다양한 워크로드 지원: 수학(GSM8k), 코딩(HumanEval), 대화(MT-Bench) 모두에서 효과적
3. CUDA Graph 호환: enforce_eager=False로 CUDA Graph와 함께 동작하여 추가 최적화 가능
4. 배치 처리 지원: max_num_seqs=128으로 높은 동시성에서도 동작

정리

DFlash는 speculative decoding의 새로운 패러다임을 제시한다. Block Diffusion을 통해 기존 autoregressive draft 모델의 한계를 넘어서며, vLLM 통합으로 실제 프로덕션 환경에서 바로 사용할 수 있게 되었다.

참고 자료

• PR #36847
• DFlash Paper (arXiv:2602.06036)

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.