본문으로 건너뛰기

[SGLang] TP Worker: GPU별 텐서 병렬 워커의 설계

들어가며

SGLang에서 LLM 추론의 실제 GPU 연산은 TP Worker가 담당한다. Scheduler가 배치를 구성하면 TP Worker는 해당 배치를 받아 Model Runner에게 포워드 패스를 위임하고, 샘플링 결과를 반환한다. 텐서 병렬(Tensor Parallelism)이 활성화되면 GPU마다 하나의 TP Worker가 독립 실행된다.

이 글에서는 python/sglang/srt/managers/tp_worker.py를 중심으로 TP Worker의 설계를 분석한다.

전체 구조

TP Worker는 Scheduler와 Model Runner 사이의 중간 계층이다.

┌──────────────┐    ModelWorkerBatch    ┌──────────────────┐
│  Scheduler   │ ────────────────────→  │   TpModelWorker  │
│              │                        │                  │
│ (CPU 스케줄링)│ ←────────────────────  │ (GPU별 1개 인스턴스)│
└──────────────┘   GenerationBatchResult└──────┬───────────┘
                                               │
                                               │ ForwardBatch
                                               ▼
                                        ┌──────────────────┐
                                        │   ModelRunner     │
                                        │ (포워드 패스 실행) │
                                        └──────────────────┘

BaseTpWorker: 추상 기반 클래스

모든 TP Worker의 공통 인터페이스를 정의한다. 핵심은 forward_batch_generation 추상 메서드와 model_runner 프로퍼티이다.

class BaseTpWorker(ABC):
    @abstractmethod
    def forward_batch_generation(self, forward_batch: ForwardBatch):
        pass

    @property
    @abstractmethod
    def model_runner(self) -> "ModelRunner":
        pass

    @property
    def sliding_window_size(self) -> Optional[int]:
        return self.model_runner.sliding_window_size

기반 클래스는 가중치 업데이트, LoRA 로딩 같은 공통 작업도 위임 패턴으로 처리한다. 모든 요청은 self.model_runner로 전달된다.

def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
    success, message = self.model_runner.update_weights_from_disk(
        recv_req.model_path,
        recv_req.load_format,
        recapture_cuda_graph=recv_req.recapture_cuda_graph,
    )
    return success, message

TpModelWorker 초기화

TpModelWorker.__init__은 GPU 워커의 전체 생애주기를 설정한다. 핵심 단계를 순서대로 살펴보자.

class TpModelWorker(BaseTpWorker):
    def __init__(self, server_args, gpu_id, tp_rank, moe_ep_rank,
                 pp_rank, attn_cp_rank, moe_dp_rank, dp_rank, nccl_port, ...):
        self.tp_rank = tp_rank
        self.gpu_id = gpu_id

        self._init_model_config()   # 1. 모델 설정
        self._init_model_runner()   # 2. ModelRunner 생성
        self._init_dllm_algorithm() # 3. dLLM 알고리즘 초기화

1단계: 모델 설정 - ModelConfig.from_server_args로 HuggingFace 설정을 파싱한다. Draft Worker인 경우 별도 경로를 사용한다.

2단계: ModelRunner 생성 - GPU별 ModelRunner를 초기화한다. 여기서 모델 로딩, 메모리 풀 할당, CUDA Graph 캡처까지 수행된다.

def _init_model_runner(self):
    self._model_runner = ModelRunner(
        model_config=self.model_config,
        gpu_id=self.gpu_id,
        tp_rank=self.tp_rank,
        tp_size=self.tp_size,
        ...
    )

3단계: 리소스 제한 설정 - 초기화 후 메모리 풀 크기에서 최대 요청 길이를 계산한다.

self.max_total_num_tokens = self.model_runner.max_total_num_tokens
self.max_req_len = min(
    self.model_config.context_len - 1,
    self.model_runner.max_token_pool_size - 1,
)
self.max_req_input_len = self.max_req_len - 5

4단계: 랜덤 시드 동기화 - TP Worker 간 동일한 샘플링 결과를 보장하기 위해 랜덤 시드를 broadcast한다.

self.random_seed = broadcast_pyobj(
    [server_args.random_seed],
    self.tp_size * self.pp_rank + tp_rank,
    self.world_group.cpu_group,
    src=self.world_group.ranks[0],
)[0]
set_random_seed(self.random_seed)

Forward 패스: Generation 분기

forward_batch_generation은 TP Worker의 핵심 메서드이다. 세 가지 분기로 나뉜다.

def forward_batch_generation(self, model_worker_batch, forward_batch=None,
                             pp_proxy_tensors=None, is_verify=False, ...):
    # 1. ModelWorkerBatch → ForwardBatch 변환
    if model_worker_batch is not None:
        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)

    # 2. dLLM 분기
    if self.is_dllm():
        return self._forward_batch_generation_dllm(forward_batch)

    # 3. Pipeline Parallel 분기
    if self.pp_group.is_last_rank:
        # 마지막 PP 스테이지: logits 생성 + 샘플링
        out = self.model_runner.forward(forward_batch, ...)
        logits_output = out.logits_output
        batch_result.next_token_ids = self.model_runner.sample(
            logits_output, forward_batch
        )
        return batch_result
    else:
        # 중간 PP 스테이지: hidden states만 전달
        out = self.model_runner.forward(forward_batch, ...)
        return GenerationBatchResult(
            pp_hidden_states_proxy_tensors=out.logits_output
        )

마지막 PP 스테이지에서는 overlap schedule이 활성화된 경우 delay_sample_func을 통해 샘플링을 지연시킨다. 이로써 GPU 연산과 CPU 스케줄링이 겹칠 수 있다.

if self.enable_overlap and not self.enable_spec and ...:
    def sample_batch_func():
        batch_result.next_token_ids = self.model_runner.sample(
            logits_output, forward_batch
        )
        return batch_result
    batch_result.delay_sample_func = sample_batch_func

Prefill-Only 처리

Prefill 전용 요청(임베딩 생성 등)은 샘플링을 건너뛴다. 더미 토큰 ID를 반환하되, logprob이 필요한 경우 별도로 계산한다.

if model_worker_batch.is_prefill_only:
    batch_result.next_token_ids = torch.zeros(
        len(model_worker_batch.seq_lens),
        dtype=torch.long,
        device=model_worker_batch.input_ids.device,
    )
    if model_worker_batch.return_logprob and logits_output.next_token_logits is not None:
        self.model_runner.compute_logprobs_only(logits_output, model_worker_batch)

설계 근거: 왜 Worker 계층이 필요한가

Model Runner만으로도 포워드 패스는 실행할 수 있다. TP Worker가 별도 계층으로 존재하는 이유는 다음과 같다.

관심사 TP Worker Model Runner
배치 변환 ModelWorkerBatch → ForwardBatch 텐서 연산만 수행
PP 분기 마지막/중간 스테이지 분기 PP 무관
샘플링 지연 샘플링 결정 샘플링 실행
가중치 관리 요청 라우팅 실제 로딩/업데이트
토크나이저 보유 미보유

이 분리 덕분에 Model Runner는 순수한 텐서 연산에만 집중하고, TP Worker가 시스템 레벨의 조율을 담당한다.

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글