[vllm] vLLM의 Dynamic Speculative Decoding을 위한 Full CUDA Graph 최적화

PR 링크: vllm-project/vllm#45953 상태: Merged | 변경: +436 / -36

들어가며

vLLM의 Speculative Decoding은 모델의 추론 속도를 높이는 핵심 기능 중 하나입니다. 특히 Dynamic Speculative Decoding(DSD)은 배치 크기나 상황에 따라 동적으로 추론 토큰 수를 조절하여 효율성을 높입니다. 하지만 기존에는 MRv1(Model Runner V1) 환경에서만 제한적으로 작동하거나, Full CUDA Graph(FCG)와의 호환성 문제로 인해 Piecewise CUDA Graph를 강제해야 하는 제약이 있었습니다. 이번 PR은 DSD를 MRv2와 FCG 환경에서 완벽하게 지원하도록 개선하여, 추론 지연 시간(TPOT)을 최적화하고 유연한 그래프 캡처 전략을 도입했습니다.

코드 분석

1. CUDA Graph 캡처 전략의 유연화 (vllm/config/compilation.py)

기존에는 고정된 크기의 그래프만 캡처했으나, DSD에서는 스케줄링에 따라 최적의 draft 토큰 수 k가 변할 수 있습니다. 이를 해결하기 위해 필요한 모든 가능한 shape을 캡처하도록 로직을 수정했습니다.

Before:

# 기존에는 고정된 크기만 고려
compilation_config.cudagraph_capture_sizes = list(range(1, 17))

After:

# DSD 설정에 따라 동적으로 캡처 크기 조정
compilation_config.resolve_cudagraph_mode_and_sizes(
    AttentionCGSupport.ALWAYS,
    "FakeAttentionBackend",
    uniform_decode_query_len=4,
    use_v2_model_runner=use_v2_model_runner,
    tensor_parallel_size=1,
)

2. MRv2 기반의 Dispatch 로직 개선 (vllm/v1/worker/gpu/cudagraph_utils.py)

CudaGraphManager가 DSD 스케줄에 따라 적절한 그래프를 선택하도록 로직을 개선했습니다. 이제는 num_speculative_tokens_per_batch_size에 정의된 모든 가능한 shape에 대해 FULL 그래프를 캡처하여, 런타임 시 오버헤드 없이 즉시 그래프를 전환할 수 있습니다.

# DSD가 활성화된 경우, 모든 가능한 k에 대해 FULL 그래프 후보 생성
if use_dynamic_sd:
    if num_spec_per_batch_size is None:
        num_spec_per_batch_size = [
            (qlen, qlen, qlen - 1) for qlen in range(1, max_decode_query_len + 1)
        ]

왜 이게 좋은가

이번 최적화의 핵심은 '필요한 그래프만 정확히 캡처'하고 '런타임 오버헤드를 제거'하는 것입니다.

1. 성능 향상: 벤치마크 결과, Eagle3 모델 기준 DSD와 FCG를 결합했을 때 TPOT(Time Per Output Token)가 2.91ms로, 기존 Piecewise 방식(2.97ms) 대비 향상된 성능을 보였습니다.
2. 유연성: Prefill과 Decode 단계의 그래프 캡처 크기를 분리하여, 각 단계에 최적화된 메모리 사용이 가능해졌습니다.
3. 교훈: CUDA Graph를 사용할 때, 단순히 최대 크기만 캡처하는 것이 아니라, 동적 스케줄링이 발생하는 경우 해당 스케줄의 모든 가능한 상태를 미리 캡처(Warm-up)해두는 것이 런타임 성능 저하를 막는 핵심임을 확인했습니다.

리뷰어 피드백 반영

리뷰어들은 특히 '불필요한 메모리 낭비'를 우려했습니다. 이에 따라 모든 가능한 k를 무조건 캡처하는 대신, num_speculative_tokens_per_batch_size에 명시된 설정값만을 기반으로 캡처하도록 로직을 정교화했습니다. 또한, MRv1과 MRv2 간의 호환성을 위해 use_v2_model_runner 플래그를 명시적으로 분리하여 안정성을 확보했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.vllm.ai/en/latest/models/speculative_decoding.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.