본문으로 건너뛰기

[SGLang] Model Runner: 포워드 패스 실행 엔진의 핵심

들어가며

ModelRunner는 SGLang 추론 엔진의 심장부이다. 모델 로딩, 메모리 풀 초기화, 포워드 패스 실행, CUDA Graph 관리를 모두 책임진다. TP Worker가 배치를 전달하면 ModelRunner가 실제 GPU 연산을 수행한다.

이 글에서는 python/sglang/srt/model_executor/model_runner.py를 중심으로 ModelRunner의 설계를 분석한다.

전체 구조

ModelRunner는 모델, 메모리, 그래프를 통합 관리하는 허브이다.

┌─────────────────────────────────────────────────────┐
│                    ModelRunner                       │
│                                                     │
│  ┌──────────┐  ┌──────────────┐  ┌───────────────┐ │
│  │  Model    │  │ Memory Pool  │  │  Graph Runner │ │
│  │ (nn.Module)│  │ (KV Cache)   │  │ (CUDA Graph)  │ │
│  └──────────┘  └──────────────┘  └───────────────┘ │
│                                                     │
│  ┌──────────┐  ┌──────────────┐  ┌───────────────┐ │
│  │ Attention │  │   Sampler    │  │  LoRA Manager │ │
│  │ Backend   │  │              │  │  (Optional)   │ │
│  └──────────┘  └──────────────┘  └───────────────┘ │
└─────────────────────────────────────────────────────┘

초기화 과정

ModelRunner의 __init__은 다양한 설정을 파싱한 후 initialize() 메서드로 실질적 초기화를 위임한다.

class ModelRunner(ModelRunnerKVCacheMixin):
    def __init__(self, model_config, mem_fraction_static, gpu_id,
                 tp_rank, tp_size, server_args, ...):
        self.device = server_args.device
        self.tp_rank = tp_rank
        self.model_config = model_config
        self.is_generation = model_config.is_generation

        # 분산 환경 초기화
        pre_model_load_memory = self.init_torch_distributed()

        # Forward 스트림 (overlap schedule용)
        self.forward_stream = torch.get_device_module(self.device).Stream()

        # 모델 초기화
        self.initialize(pre_model_load_memory)

initialize() 메서드는 아래 순서로 진행된다.

def initialize(self, pre_model_load_memory):
    # 1. 모델 가중치 로딩
    self.sampler = create_sampler()
    self.load_model()

    # 2. KV 캐시 메모리 풀 할당
    self.configure_kv_cache_dtype()
    self.init_memory_pool(pre_model_load_memory)

    # 3. Attention 백엔드 초기화
    self.init_attention_backend()

    # 4. CUDA Graph 캡처
    self.init_device_graphs()

    # 5. Piecewise CUDA Graph 초기화
    self.init_piecewise_cuda_graphs()

모델 로딩

load_model()은 LoadConfig를 기반으로 적절한 로더를 선택하고 모델을 GPU에 올린다.

def load_model(self):
    self.load_config = LoadConfig(
        load_format=self.server_args.load_format,
        download_dir=self.server_args.download_dir,
        ...
    )

    self.loader = get_model_loader(
        load_config=self.load_config,
        model_config=self.model_config,
    )
    self.model = self.loader.load_model(
        model_config=self.model_config,
        device_config=DeviceConfig(self.device, self.gpu_id),
    )

Memory Saver 어댑터로 감싸서 가중치 백업도 지원한다.

with self.memory_saver_adapter.region(
    GPU_MEMORY_TYPE_WEIGHTS,
    enable_cpu_backup=enable_cpu_backup,
):
    self.model = self.loader.load_model(...)

Forward 패스: 4가지 모드

ModelRunner의 forward() 메서드는 _forward_raw()를 호출하며, 이 메서드는 ForwardMode에 따라 4가지 경로로 분기한다.

def forward(self, forward_batch, ...):
    self.forward_pass_id += 1
    output = self._forward_raw(forward_batch, ...)
    return output

def _forward_raw(self, forward_batch, ...):
    # CUDA Graph 실행 가능 여부 판단
    can_run_graph = bool(
        forward_batch.forward_mode.is_cuda_graph()
        and self.graph_runner
        and self.graph_runner.can_run(forward_batch)
    )

    if can_run_graph:
        ret = self.graph_runner.replay(forward_batch)
    elif forward_batch.forward_mode.is_decode():
        ret = self.forward_decode(forward_batch)
    elif forward_batch.forward_mode.is_extend():
        ret = self.forward_extend(forward_batch)
    elif forward_batch.forward_mode.is_idle():
        ret = self.forward_idle(forward_batch)

각 모드의 역할을 정리하면 다음과 같다.

┌─────────────────────────────────────────────────────┐
│  ForwardMode 분기                                    │
│                                                     │
│  DECODE ─→ forward_decode()                         │
│            - 토큰 1개 생성                            │
│            - CUDA Graph 실행 가능                     │
│                                                     │
│  EXTEND ─→ forward_extend()                         │
│            - Prefill (시스템 프롬프트 처리)             │
│            - Piecewise CUDA Graph 가능               │
│                                                     │
│  IDLE ───→ forward_idle()                           │
│            - DP Attention 패딩용                      │
│                                                     │
│  SPLIT_PREFILL ─→ forward_split_prefill()           │
│                   - 레이어별 분할 실행                  │
└─────────────────────────────────────────────────────┘

forward_decode: Decode 모드

Decode 모드는 가장 단순하다. Attention 백엔드 메타데이터를 초기화하고 모델의 forward를 호출한다.

def forward_decode(self, forward_batch, skip_attn_backend_init=False, ...):
    if not skip_attn_backend_init:
        self.attn_backend.init_forward_metadata(forward_batch)
    return self.model.forward(
        forward_batch.input_ids,
        forward_batch.positions,
        forward_batch,
    )

forward_extend: Prefill 모드

Prefill 모드는 Piecewise CUDA Graph가 활성화된 경우 그래프 재생을 우선 시도한다.

def forward_extend(self, forward_batch, ...):
    can_run_graph = (
        self.piecewise_cuda_graph_runner is not None
        and self.piecewise_cuda_graph_runner.can_run(forward_batch)
    )

    if can_run_graph:
        return (
            self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs),
            can_run_graph,
        )

    # 일반 실행
    self.attn_backend.init_forward_metadata(forward_batch)
    return (
        self.model.forward(forward_batch.input_ids, forward_batch.positions, ...),
        can_run_graph,
    )

입력 임베딩 오버라이드도 이 단계에서 처리한다.

if forward_batch.replace_embeds is not None and forward_batch.replace_positions is not None:
    if "input_embeds" not in kwargs:
        embed_layer = self.model.get_input_embeddings()
        kwargs["input_embeds"] = embed_layer(forward_batch.input_ids)
    kwargs["input_embeds"][forward_batch.replace_positions] = (
        forward_batch.replace_embeds.to(kwargs["input_embeds"].dtype)
    )

CUDA Graph 초기화

init_device_graphs()에서 CudaGraphRunner를 생성한다. 이 과정에서 다양한 배치 크기에 대해 CUDA Graph를 사전 캡처한다.

def init_device_graphs(self):
    if self.device == "cuda":
        self.init_cublas()
        self.init_attention_backend()
        self.kernel_warmup()
        self.init_device_graphs()

Graph Runner는 Decode 모드용 CudaGraphRunner와 Prefill 모드용 PiecewiseCudaGraphRunner 두 가지이다.

Expert 분산 모니터링

Forward 패스가 끝난 후 MoE 전문가 분포를 기록한다.

def forward(self, forward_batch, ...):
    with get_global_expert_distribution_recorder().with_forward_pass(
        self.forward_pass_id, forward_batch,
    ) as recorder_outputs:
        output = self._forward_raw(forward_batch, ...)
    output.expert_distribution_metrics = recorder_outputs.get("metrics")

이 데이터는 EPLB(Expert Parallel Load Balancing)에서 전문가 재배치 결정에 활용된다.

설계 근거: ModelRunner의 책임 범위

ModelRunner는 "무엇을 실행할 것인가"와 "어떻게 실행할 것인가"를 모두 결정한다.

관심사 구현
모델 로딩 load_model() + get_model_loader()
메모리 관리 init_memory_pool()
Forward 모드 분기 _forward_raw()
CUDA Graph CudaGraphRunner / PiecewiseCudaGraphRunner
샘플링 sample() via Sampler
가중치 업데이트 update_weights_from_*()

TP Worker와의 역할 분담이 명확하다. TP Worker는 배치 변환과 PP 분기를, ModelRunner는 텐서 연산과 메모리를 담당한다.

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글