본문으로 건너뛰기

[vLLM] GPUModelRunner: GPU 모델 포워드 패스

들어가며

GPUModelRunner는 vLLM v1에서 실제 GPU 모델 실행을 담당하는 핵심 클래스이다. 스케줄러가 결정한 배치를 받아 입력 텐서를 준비하고, attention 메타데이터를 구성하며, 모델 포워드 패스를 실행하고, 샘플링까지 수행한다. 텍스트/멀티모달, 생성/임베딩, 퍼블릭/프라이빗 모든 모델이 이 하나의 러너를 공유한다.

소스 경로: vllm/v1/worker/gpu/model_runner.py

공식 문서

vLLM 공식 문서: Model Runner V2

핵심 구조/코드 분석

클래스 초기화

class GPUModelRunner(LoRAModelRunnerMixin):
    def __init__(self, vllm_config: VllmConfig, device: torch.device):
        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
        self.max_num_reqs = self.scheduler_config.max_num_seqs

        # Pipeline Parallelism
        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
        self.is_first_pp_rank = get_pp_group().is_first_rank
        self.is_last_pp_rank = get_pp_group().is_last_rank

        # Decode Context Parallelism
        self.dcp_size = self.parallel_config.decode_context_parallel_size
        self.use_dcp = self.dcp_size > 1

        # Speculative Decoding
        if self.speculative_config is not None:
            self.num_speculative_steps = self.speculative_config.num_speculative_tokens
            if self.is_last_pp_rank:
                self.speculator = init_speculator(self.vllm_config, self.device)

        # CUDA Graphs
        self.cudagraph_manager = ModelCudaGraphManager(
            self.vllm_config, self.device,
            self.compilation_config.cudagraph_mode,
            decode_query_len=self.decode_query_len,
        )

초기화에서 PP, DP, DCP(Decode Context Parallelism), Speculative Decoding, CUDA Graph 등 모든 병렬화 전략과 최적화가 설정된다.

execute_model() - 핵심 실행 함수

def execute_model(self, scheduler_output, intermediate_tensors=None, 
                  dummy_run=False, ...) -> ModelRunnerOutput | None:
    if not dummy_run:
        self.finish_requests(scheduler_output)
        self.add_requests(scheduler_output)
        self.update_requests(scheduler_output)
        self.block_tables.apply_staged_writes()
        if scheduler_output.total_num_scheduled_tokens == 0:
            return self.kv_connector.no_forward(scheduler_output)

    # CUDA 그래프 디스패치 결정
    batch_desc = self.cudagraph_manager.dispatch(
        num_reqs, num_toks, uniform_tok_count
    )

    # 입력 준비
    input_batch = self.prepare_inputs(scheduler_output, batch_desc)
    block_tables, slot_mappings = self.prepare_attn(input_batch)

    # 모델 포워드 패스
    hidden_states = self.model(
        input_ids=input_batch.token_ids,
        positions=input_batch.positions,
        intermediate_tensors=intermediate_tensors,
        ...
    )

실행 흐름은 다음과 같다:

  1. 요청 상태 업데이트 (추가/삭제/변경)
  2. CUDA 그래프 모드 결정 (FULL, PIECEWISE, NONE)
  3. 입력 텐서 준비 (token_ids, positions, block_tables)
  4. Attention 메타데이터 구성
  5. 모델 포워드 패스 실행

CUDA 그래프 캡처

@torch.inference_mode()
def capture_model(self) -> int:
    start_free_gpu_memory = torch.cuda.mem_get_info()[0]
    self.cudagraph_manager.capture(
        self.model, self.model_state, self.input_buffers,
        self.intermediate_tensors, self.block_tables,
        self.attn_groups, self.kv_cache_config, ...
    )
    if self.speculator is not None:
        self.speculator.capture_model()
    cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory

모델과 speculator 모두에 대해 CUDA 그래프를 캡처한다. 이 과정은 보통 5~20초 소요되며, 이후 디코드 단계에서의 커널 런칭 오버헤드를 제거한다.

KV 캐시 초기화

def initialize_kv_cache(self, kv_cache_config) -> None:
    self.block_tables = BlockTables(
        block_sizes=block_sizes,
        max_num_reqs=self.max_num_reqs,
        max_model_len=block_table_max_model_len,
        cp_size=self.dcp_size, cp_rank=self.dcp_rank,
        ...
    )
    self.attn_backends, self.attn_groups = init_attn_backend(
        self.kv_cache_config, self.vllm_config, self.device
    )
    kv_caches_dict = init_kv_cache(
        self.kv_caches, ..., self.kv_cache_config, self.attn_backends, ...
    )
    self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

KV 캐시는 블록 단위로 관리되며, Context Parallelism을 사용할 경우 블록 테이블이 DCP 랭크에 따라 분할된다. KV Connector가 설정되면 원격 KV 캐시 전송도 가능하다.

왜 이 설계인가

  1. 단일 러너 원칙: 소스 코드 주석에도 명시되어 있듯이, 이 파일은 모든 모델이 공유하는 코드만 포함한다. 모델별 로직은 별도 파일에 위치하여 안정성을 유지한다.

  2. Staged Write: block_tables.apply_staged_writes()는 블록 테이블 변경을 한 번에 적용한다. 이는 prefill과 decode가 혼합된 배치에서 일관성을 보장한다.

  3. DP 동기화: 데이터 병렬 처리 시 sync_cudagraph_and_dp_padding()으로 모든 DP 랭크의 배치 크기를 맞추어 CUDA 그래프 재사용을 극대화한다.

참고

댓글

관련 포스트

vLLM 의 다른글