[vllm] vLLM의 첫 추론 지연 문제 해결: forward_native 샘플러 커널 웜업 최적화

PR 링크: vllm-project/vllm#41375 상태: Merged | 변경: +None / -None

들어가며

vLLM은 고성능 추론 엔진으로서 다양한 최적화 기법을 적용하고 있습니다. 최근 vLLM v1 엔진에 FlashInfer 기반의 top-k/top-p 샘플러가 기본값으로 도입되면서, 특정 환경에서 첫 추론 시 약 1초에 달하는 JIT(Just-In-Time) 컴파일 지연(Cold-start) 문제가 발생했습니다. 이는 특히 asyncio.sleep과 같은 짧은 타임아웃을 사용하는 테스트 코드에서 타임아웃 오류를 유발하는 원인이 되었습니다. 본 글에서는 이 문제를 해결하기 위해 forward_native 샘플러 커널을 어떻게 엔진 초기화 단계에서 미리 웜업(Warmup)했는지 분석합니다.

코드 분석

핵심 변경 사항은 vllm/v1/worker/gpu_model_runner.py 파일 내 _dummy_sampler_run 함수에 있습니다. 기존에는 더미 샘플러 실행 시 generators를 빈 딕셔너리로 설정하여 호출했는데, 이로 인해 forward_native Triton 커널이 실행되지 않고 JIT 컴파일이 지연되는 문제가 있었습니다.

Before

# 기존 코드: generators가 비어있어 forward_native 경로가 웜업되지 않음
sampler_output = self.sampler(
    logits=logits, sampling_metadata=dummy_metadata
)

After

# 개선 코드: 명시적으로 시드(seed)가 포함된 generator를 전달하여 forward_native 경로를 강제 웜업
if self.sampler.logprobs_mode not in (
    "processed_logits",
    "processed_logprobs",
):
    self.sampler(
        logits=logits,
        sampling_metadata=replace(
            dummy_metadata,
            generators={
                0: torch.Generator(device=self.device).manual_seed(0)
            },
        ),
    )

개선된 코드는 logprobs_mode가 특정 모드(이미 forward_native가 바인딩된 상태)가 아닐 경우, torch.Generator를 포함한 메타데이터를 생성하여 샘플러를 한 번 더 호출합니다. 이를 통해 런타임에 첫 요청이 들어왔을 때 발생하는 Triton 커널 컴파일 비용을 엔진 초기화 시점으로 앞당겼습니다.

왜 이게 좋은가

이 최적화의 핵심은 '런타임 지연을 초기화 비용으로 전가(Shift)'하는 것입니다.

1. 성능 수치: B200 하드웨어 환경에서 test_multi_abort 테스트가 기존에는 2/2 실패했으나, 본 수정 이후 2/2 성공으로 전환되었습니다. 이는 1초 이상의 JIT 컴파일 지연이 제거되었음을 의미합니다.
2. 일반적 교훈:
   • Cold-start 관리: Triton과 같은 JIT 기반 커널을 사용할 때는 실제 추론이 시작되기 전, 더미 데이터를 사용하여 주요 실행 경로(Hot path)를 미리 실행하는 것이 필수적입니다.
   • 조건부 실행: 모든 경우에 웜업을 수행하는 것이 아니라, logprobs_mode와 같이 이미 최적화된 경로를 사용하는 경우에는 중복 호출을 방지하여 메모리 오버헤드를 최소화했습니다.

리뷰 과정에서 rejection sampler에 대한 웜업 필요성도 논의되었으나, 해당 샘플러는 입력에 따른 분기 처리가 없어 추가 웜업이 불필요하다는 점이 확인되었습니다. 이는 불필요한 웜업 로직 추가를 방지하여 코드 복잡도를 낮추는 좋은 의사결정이었습니다.

참고 자료

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

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