본문으로 건너뛰기

[vLLM] Warmup: 커널 JIT 사전 컴파일

들어가며

JIT(Just-In-Time) 컴파일 커널은 처음 호출될 때 컴파일되면서 수 초의 지연이 발생한다. LLM 서빙에서 첫 요청의 지연 시간이 비정상적으로 길어지는 원인이다. vLLM은 vllm/model_executor/warmup/에서 서빙 시작 전에 주요 커널을 미리 실행하여 이 문제를 해결한다.

핵심 구조/코드 분석

kernel_warmup 진입점

def kernel_warmup(worker: "Worker"):
    # 1. Deep GEMM 워밍업
    do_deep_gemm_warmup = (
        envs.VLLM_USE_DEEP_GEMM
        and is_deep_gemm_supported()
        and envs.VLLM_DEEP_GEMM_WARMUP != "skip"
    )
    if do_deep_gemm_warmup:
        model = worker.get_model()
        max_tokens = worker.scheduler_config.max_num_batched_tokens
        deep_gemm_warmup(model, max_tokens)

    # 2. FlashInfer 오토튜닝
    enable_flashinfer_autotune = worker.vllm_config.kernel_config.enable_flashinfer_autotune
    if enable_flashinfer_autotune is not False:
        if has_flashinfer() and current_platform.has_device_capability(90):
            flashinfer_autotune(worker.model_runner)

    # 3. FlashInfer 어텐션 워밍업
    if (not worker.model_runner.is_pooling_model
        and worker.model_runner.attn_groups
        and all(_is_flashinfer_backend(group.backend)
                for groups in worker.model_runner.attn_groups
                for group in groups)):
        worker.model_runner._dummy_run(num_tokens=16, skip_eplb=True,
            is_profile=True, force_attention=True, create_mixed_batch=True)

세 단계로 구성된다: Deep GEMM 워밍업, FlashInfer 오토튜닝, FlashInfer 어텐션 워밍업. 각 단계는 조건부로 실행되며, skip 환경변수로 개별 비활성화가 가능하다.

FlashInfer 오토튜닝

def flashinfer_autotune(runner: "GPUModelRunner") -> None:
    """FlashInfer 오토튜닝: 여러 구현 중 최적을 벤치마크로 선택"""
    import vllm.utils.flashinfer as fi_utils

    with torch.inference_mode(), fi_utils.autotune():
        fi_utils._is_fi_autotuning = True
        runner._dummy_run(
            runner.scheduler_config.max_num_batched_tokens,
            skip_eplb=True,
            is_profile=True,
        )
        fi_utils._is_fi_autotuning = False

FlashInfer는 동일한 연산에 대해 여러 구현을 가지고 있다. 오토튜닝은 최대 토큰 수로 각 구현을 벤치마크하여 가장 빠른 것을 선택한다. 결과는 투명하게 캐시되어 이후 호출에서 자동으로 사용된다.

핵심은 max_num_batched_tokens으로 실행한다는 점이다. FlashInfer의 오토튜닝에서 토큰 수 m으로 튜닝하면 m 이하의 모든 토큰 수에 대한 최적 구현이 결정되므로, 최대값 한 번만 실행하면 된다.

Mixed Batch 어텐션 워밍업

worker.model_runner._dummy_run(
    num_tokens=16,
    skip_eplb=True,
    is_profile=True,
    force_attention=True,
    create_mixed_batch=True,  # prefill + decode 혼합 배치
)

실제 서빙에서는 prefill과 decode 요청이 같은 배치에 섞인다. create_mixed_batch=True로 이 혼합 시나리오를 미리 실행하여, prefill 커널과 decode 커널 모두 워밍업한다.

Deep GEMM 워밍업

from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup

deep_gemm_warmup(model, max_tokens)

Deep GEMM은 DeepSeek 계열 모델의 MoE 레이어에서 사용하는 고성능 행렬곱 커널이다. JIT 컴파일 시간이 길어서 별도 워밍업이 필요하다.

Hopper/Blackwell 제한

if has_flashinfer() and current_platform.has_device_capability(90):
    flashinfer_autotune(worker.model_runner)

FlashInfer 오토튜닝은 SM 9.0(Hopper) 이상에서만 실행된다. 이전 아키텍처에서는 구현 선택지가 적어 오토튜닝의 이점이 크지 않기 때문이다.

왜 이 설계인가

  1. 서빙 시작 전 완료: 워밍업을 서버 시작 시점에 수행하여, 첫 번째 실제 요청부터 일관된 지연 시간을 보장한다. 프로덕션 환경에서 cold start 문제를 방지한다.

  2. skip_eplb 플래그: 워밍업 중에는 Expert Parallel Load Balancing 메트릭을 기록하지 않는다. 더미 데이터로 실행하므로 실제 워크로드를 반영하지 않기 때문이다.

  3. 오토튜닝 중 일부 커널 스킵: _is_fi_autotuning 플래그로 오토튜닝과 호환되지 않는 커널(예: NVFP4 MoE)을 자동 스킵한다. 오토튜닝이 이런 커널을 호출하면 크래시할 수 있기 때문이다.

  4. 조건부 워밍업: 모든 모델이 모든 워밍업을 필요로 하지 않는다. Deep GEMM을 사용하지 않는 모델, FlashInfer를 사용하지 않는 백엔드, 풀링 모델 등은 각각의 워밍업을 건너뛴다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글