본문으로 건너뛰기

[vLLM] Preemption & Async Scheduling: 선점과 비동기 스케줄링

들어가며

LLM 서빙 스케줄러의 핵심 과제는 한정된 GPU 메모리에서 최대 처리량을 달성하면서, 새 요청과 진행 중인 요청 사이의 리소스 경합을 관리하는 것이다. vLLM v1의 스케줄러(vllm/v1/core/sched/scheduler.py)는 선점(preemption), 비동기 KV 전송, 멀티모달 인코더 예산 관리를 통합적으로 처리한다.

핵심 구조/코드 분석

스케줄러 초기화

class Scheduler(SchedulerInterface):
    def __init__(self, vllm_config, kv_cache_config, structured_output_manager,
                 block_size, mm_registry=MULTIMODAL_REGISTRY, ...):
        self.max_num_running_reqs = self.scheduler_config.max_num_seqs
        self.max_num_scheduled_tokens = (
            self.scheduler_config.max_num_scheduled_tokens
            or self.scheduler_config.max_num_batched_tokens
        )

        # 요청 큐
        self.requests: dict[str, Request] = {}
        self.waiting = create_request_queue(self.policy)  # 대기 큐
        self.skipped_waiting = create_request_queue(self.policy)  # 스킵된 대기 요청
        self.running: list[Request] = []  # 실행 중

        # 완료된 요청 추적
        self.finished_req_ids: set[str] = set()

세 가지 주요 상태를 관리한다: waiting(대기), running(실행 중), finished(완료). skipped_waiting은 비동기 의존성이나 제약 조건으로 인해 현재 스텝에서 스킵된 요청을 임시로 보관한다.

스케줄링 정책

self.policy = SchedulingPolicy(self.scheduler_config.policy)
self.waiting = create_request_queue(self.policy)

create_request_queue는 정책에 따라 우선순위 큐를 생성한다. FCFS(First Come First Served), 우선순위 기반 등 다양한 정책을 지원한다.

KV 커넥터 통합

# P/D(Prefill/Decode) 분리 아키텍처를 위한 KV 커넥터
self.connector = None
if self.vllm_config.kv_transfer_config is not None:
    self.connector = KVConnectorFactory.create_connector(
        config=self.vllm_config,
        role=KVConnectorRole.SCHEDULER,
        kv_cache_config=self.kv_cache_config,
    )

# 비동기 KV 수신 추적
self.finished_recving_kv_req_ids: set[str] = set()
self.failed_recving_kv_req_ids: set[str] = set()

Prefill-Decode 분리 아키텍처에서, KV 캐시를 원격으로 전송/수신하는 KV 커넥터가 스케줄러에 통합된다. 비동기 KV 로딩 실패 시 recompute_kv_load_failures 정책에 따라 재계산하거나 에러를 반환한다.

멀티모달 인코더 예산

supports_mm_inputs = mm_registry.supports_multimodal_inputs(vllm_config.model_config)
mm_budget = MultiModalBudget(vllm_config, mm_registry) if supports_mm_inputs else None

self.max_num_encoder_input_tokens = mm_budget.encoder_compute_budget if mm_budget else 0
encoder_cache_size = mm_budget.encoder_cache_size if mm_budget else 0
self.encoder_cache_manager = (
    EncoderDecoderCacheManager(cache_size=encoder_cache_size)
    if self.is_encoder_decoder
    else EncoderCacheManager(cache_size=encoder_cache_size)
)

멀티모달 모델의 인코더(예: Vision Encoder)도 GPU 리소스를 사용하므로, 스케줄러가 인코더 토큰 예산과 캐시를 관리한다. 인코더-디코더 모델과 디코더 전용 멀티모달 모델에 대해 별도의 캐시 매니저를 사용한다.

KV 캐시 매니저와 블록 풀

self.kv_cache_manager = KVCacheManager(
    kv_cache_config=kv_cache_config,
    max_model_len=self.max_model_len,
    enable_caching=self.cache_config.enable_prefix_caching,
    use_eagle=self.use_eagle,
    enable_kv_cache_events=self.enable_kv_cache_events,
    dcp_world_size=self.dcp_world_size,
    pcp_world_size=self.pcp_world_size,
    hash_block_size=self.block_size,
)

# KV 커넥터에 GPU 블록 풀 바인딩
if self.connector is not None and hasattr(self.connector, "bind_gpu_block_pool"):
    self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool)

KV 캐시 매니저는 블록 풀 기반으로 메모리를 관리하며, EAGLE 투기적 디코딩, prefix caching, decode context parallelism(DCP), prefill context parallelism(PCP) 등의 고급 기능을 모두 지원한다.

투기적 디코딩 통합

speculative_config = vllm_config.speculative_config
self.use_eagle = False
self.num_spec_tokens = self.num_lookahead_tokens = 0
if speculative_config:
    self.num_spec_tokens = speculative_config.num_speculative_tokens
    if speculative_config.use_eagle():
        self.use_eagle = True
        self.num_lookahead_tokens = self.num_spec_tokens

투기적 디코딩이 활성화되면, 스케줄러는 num_lookahead_tokens만큼의 추가 토큰 슬롯을 고려하여 블록을 할당한다.

왜 이 설계인가

  1. 스킵된 대기 큐 분리: 비동기 KV 수신이나 구조화된 출력 준비 등으로 현재 스텝에서 스케줄링할 수 없는 요청을 별도 큐에 보관한다. 이렇게 하면 메인 대기 큐의 순회가 빠르고, 스킵된 요청이 다음 스텝에서 재시도된다.

  2. 블록 풀의 KV 커넥터 바인딩: KV 커넥터가 블록 풀에 직접 접근하면, 원격 KV 수신 시 블록을 즉시 할당하고 선점(preemption)과 같은 메모리 관리 이벤트에 반응할 수 있다.

  3. DCP/PCP 월드 사이즈: Decode Context Parallelism과 Prefill Context Parallelism은 긴 시퀀스를 여러 워커에 분산하여 처리한다. 스케줄러가 이 월드 사이즈를 알아야 블록 할당 시 올바른 크기 배수를 사용할 수 있다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글