본문으로 건너뛰기

[SGLang] Multi-Tokenizer: 다중 모델 토크나이저 동시 관리

들어가며

LLM serving에서 HTTP 서버의 tokenization 처리가 병목이 되는 경우가 있다. 특히 대량의 동시 요청을 처리할 때, 단일 프로세스의 TokenizerManager가 모든 토큰화와 응답 처리를 담당하면 CPU 사용률이 포화된다. SGLang은 이 문제를 multi-http-worker 모드로 해결한다. 여러 TokenizerWorker 프로세스를 실행하고, 각각이 독립적으로 토큰화와 응답 처리를 수행한다.

이 글에서는 python/sglang/srt/managers/multi_tokenizer_mixin.py를 중심으로 multi-worker 아키텍처의 설계를 분석한다.

전체 구조

single-worker 모드와 multi-worker 모드의 아키텍처 차이를 비교해보자.

=== Single Worker Mode (tokenizer_worker_num=1) ===

  Client ──▶ HTTP Server ──▶ TokenizerManager ──▶ Scheduler
                                   ▲                  │
                                   └── Detokenizer ◀──┘

=== Multi Worker Mode (tokenizer_worker_num=N) ===

  Client ──▶ HTTP Server (worker 1) ──▶ TokenizerWorker 1 ─┐
  Client ──▶ HTTP Server (worker 2) ──▶ TokenizerWorker 2 ─┤
  Client ──▶ HTTP Server (worker N) ──▶ TokenizerWorker N ─┤
                                                            │
                                            ┌───────────────┤
                                            ▼               │
                                   MultiTokenizerRouter     │
                                       │       ▲            │
                                       ▼       │            │
                                   Scheduler    │            │
                                       │       │            │
                                       ▼       │            │
                                  Detokenizer   │            │
                                       │       │            │
                                       ▼       │            │
                          MultiHttpWorkerDetokenizerMixin    │
                               (SocketMapping)               │
                                       │                     │
                                       └─── 결과를 각 Worker ─┘
                                            에게 개별 전송

핵심 차이는 세 가지다. (1) TokenizerWorker가 TokenizerManager를 상속하되, 요청에 자신의 IPC 이름을 부착한다. (2) MultiTokenizerRouter가 Worker들의 요청을 Scheduler로 라우팅한다. (3) DetokenizerManager가 SocketMapping을 통해 결과를 올바른 Worker에게 돌려보낸다.

핵심 코드 분석

SocketMapping: 동적 소켓 관리

multi-worker 모드에서 DetokenizerManager는 여러 TokenizerWorker에게 결과를 보내야 한다. SocketMapping이 이 동적 라우팅을 담당한다.

class SocketMapping:
    def __init__(self):
        self._zmq_context = zmq.Context()
        self._mapping: Dict[str, zmq.Socket] = {}

    def _register_ipc_mapping(self, ipc_name: str, is_tokenizer: bool):
        type_str = "tokenizer" if is_tokenizer else "detokenizer"
        if ipc_name in self._mapping:
            logger.warning(f"{type_str} already registered {ipc_name=}, skipping...")
            return
        logger.info(f"Registering {type_str} {ipc_name=} in SocketMapping...")
        socket = get_zmq_socket(self._zmq_context, zmq.PUSH, ipc_name, False)
        self._mapping[ipc_name] = socket

    def send_output(self, ipc_name: str, output: Any):
        if ipc_name is None:
            logger.warning(f"IPC name is None, output type={type(output)}, skipping...")
            return
        if ipc_name not in self._mapping:
            self._register_ipc_mapping(ipc_name, is_tokenizer=False)
        self._mapping[ipc_name].send_pyobj(output)

IPC 이름을 키로 ZMQ PUSH 소켓을 관리하는 딕셔너리다. 처음 보는 IPC 이름이 등장하면 자동으로 소켓을 생성하여 등록한다. 이 lazy initialization 전략으로 Worker의 수나 IPC 이름을 미리 알 필요가 없다.

MultiHttpWorkerDetokenizerMixin: Detokenizer 확장

이 mixin은 DetokenizerManager에 multi-worker 지원을 추가한다. single-worker 모드의 event_loop를 대체하는 multi_http_worker_event_loop를 제공한다.

class MultiHttpWorkerDetokenizerMixin:
    """Mixin class for DetokenizerManager"""

    def maybe_clear_socket_mapping(self: DetokenizerManager):
        if hasattr(self, "socket_mapping"):
            self.socket_mapping.clear_all_sockets()

    def multi_http_worker_event_loop(self: DetokenizerManager):
        """The event loop that handles requests, for multi multi-http-worker mode"""
        self.socket_mapping = SocketMapping()
        while True:
            recv_obj = self.recv_from_scheduler.recv_pyobj()
            output = self._request_dispatcher(recv_obj)
            if output is None:
                continue

            assert isinstance(
                recv_obj, BaseBatchReq
            ), "for multi-http-worker, recv_obj must be BaseBatchReq"

            for i, ipc_name in enumerate(recv_obj.http_worker_ipcs):
                new_output = _handle_output_by_index(output, i)
                self.socket_mapping.send_output(ipc_name, new_output)

single-worker 모드에서는 디토큰화된 배치 전체를 하나의 TokenizerManager에 보내면 되지만, multi-worker 모드에서는 배치 내 각 요청을 원래 요청을 보낸 Worker에게 개별적으로 보내야 한다. recv_obj.http_worker_ipcs[i]가 각 요청이 어느 Worker에서 왔는지를 추적한다.

프로세스 시작 시 어떤 루프를 사용할지는 run_detokenizer_process에서 결정된다.

def run_detokenizer_process(
    server_args: ServerArgs,
    port_args: PortArgs,
    detokenizer_manager_class=DetokenizerManager,
):
    kill_itself_when_parent_died()
    setproctitle.setproctitle("sglang::detokenizer")
    configure_logger(server_args)

    manager = detokenizer_manager_class(server_args, port_args)
    if server_args.tokenizer_worker_num == 1:
        manager.event_loop()
    else:
        manager.multi_http_worker_event_loop()

_handle_output_by_index: 배치 분할

배치 출력에서 개별 요청의 출력을 추출하는 유틸리티 함수다. BatchTokenIDOutput, BatchStrOutput, BatchEmbeddingOutput 세 가지 타입을 모두 처리한다.

def _handle_output_by_index(output, i):
    if isinstance(output, BatchTokenIDOutput):
        new_output = BatchTokenIDOutput(
            rids=[output.rids[i]],
            spec_verify_ct=_extract_field_by_index(output, "spec_verify_ct", i),
            spec_accepted_tokens=_extract_field_by_index(
                output, "spec_accepted_tokens", i
            ),
            finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
            decoded_texts=_extract_field_by_index(output, "decoded_texts", i),
            decode_ids=_extract_field_by_index(output, "decode_ids", i),
            # ... 20+ more fields
        )
    elif isinstance(output, BatchEmbeddingOutput):
        # ...
    elif isinstance(output, BatchStrOutput):
        # ...

각 필드를 개별적으로 추출해야 하므로 코드가 길지만, _extract_field_by_index 헬퍼로 중복을 줄였다.

def _extract_field_by_index(
    output: Any, field_name: str, index: int, check_length: bool = True
) -> Any:
    field = getattr(output, field_name, None)
    if field is None:
        return None

    if isinstance(field, dict):
        new_field = {}
        for k, v in field.items():
            if len(v) <= index:
                new_field[k] = None
            new_field[k] = v[index]
        return new_field

    if check_length:
        if len(field) <= index:
            return None
    return [field[index]]

check_length=False는 logprobs 같은 필드에 사용된다. 이 필드들은 요청에 따라 존재하지 않을 수 있으므로 길이 검사를 건너뛴다.

TokenizerWorker: Worker 프로세스

TokenizerWorker는 TokenizerManager를 상속하여 multi-worker 모드에서 독립적으로 동작하는 Worker 프로세스다.

class TokenizerWorker(TokenizerManager):
    """Tokenizer Worker in multi-http-worker mode"""

    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
    ):
        setproctitle.setproctitle(f"sglang::tokenizer_worker:{os.getpid()}")
        disaggregation_mode = server_args.disaggregation_mode
        server_args.disaggregation_mode = "null"
        super().__init__(server_args, port_args)

        self.worker_id = os.getpid()
        self.tokenizer_ipc_name = port_args.tokenizer_ipc_name

    def _attach_multi_http_worker_info(self, req: Union[BaseReq, BaseBatchReq]):
        if isinstance(req, BaseReq):
            req.http_worker_ipc = self.tokenizer_ipc_name
        elif isinstance(req, BaseBatchReq):
            req.http_worker_ipcs = [self.tokenizer_ipc_name] * len(req.rids)

_attach_multi_http_worker_info가 핵심이다. 각 요청에 자신의 IPC 이름(tokenizer_ipc_name)을 부착하여, DetokenizerManager가 결과를 돌려보낼 때 이 IPC 이름을 사용하도록 한다. TokenizerManager.generate_request에서 tokenizer_worker_num > 1이면 이 메서드를 호출하는 구조다.

MultiTokenizerRouter: 요청 라우팅

여러 TokenizerWorker와 Scheduler 사이에서 요청과 결과를 라우팅하는 Router다.

class MultiTokenizerRouter:
    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
    ):
        context = zmq.asyncio.Context(3)
        self.recv_from_detokenizer = get_zmq_socket(
            context, zmq.PULL, port_args.tokenizer_ipc_name, True
        )
        self.send_to_scheduler = get_zmq_socket(
            context, zmq.PUSH, port_args.scheduler_input_ipc_name, True
        )
        self.receive_from_worker = get_zmq_socket(
            context, zmq.PULL, port_args.tokenizer_worker_ipc_name, True
        )
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()

세 개의 ZMQ 소켓을 운영한다. Worker들로부터 토큰화된 요청을 받는 PULL 소켓, Scheduler에 전달하는 PUSH 소켓, DetokenizerManager로부터 특수 요청(non-batch)을 받아 해당 Worker에 라우팅하는 PULL 소켓이다.

두 개의 비동기 루프가 동시에 실행된다.

async def router_worker_obj(self):
    while True:
        recv_obj = await self.receive_from_worker.recv_pyobj()
        await self.send_to_scheduler.send_pyobj(recv_obj)

async def handle_loop(self):
    self.socket_mapping = SocketMapping()
    while True:
        recv_obj = await self.recv_from_detokenizer.recv_pyobj()
        await self._distribute_result_to_workers(recv_obj)

router_worker_obj는 Worker → Scheduler 방향의 단순 전달이다. handle_loop는 DetokenizerManager → Worker 방향으로, SocketMapping을 통해 결과를 올바른 Worker에게 분배한다.

SenderWrapper: 투명한 IPC 이름 부착

multi-worker 모드에서 TokenizerManager의 send_to_scheduler를 감싸는 래퍼다.

class SenderWrapper:
    def __init__(self, port_args: PortArgs, send_to_scheduler: zmq.Socket):
        self.port_args = port_args
        self.send_to_scheduler = send_to_scheduler

    def send_pyobj(self, obj):
        if isinstance(obj, BaseReq):
            obj.http_worker_ipc = self.port_args.tokenizer_ipc_name
        self.send_to_scheduler.send_pyobj(obj)

send_pyobj를 호출할 때마다 자동으로 http_worker_ipc를 부착한다. TokenizerManager의 기존 코드를 수정하지 않고도 multi-worker 지원을 추가할 수 있는 깔끔한 설계다.

Shared Memory: 프로세스 간 설정 공유

multi-worker 모드에서 각 Worker 프로세스는 서버 설정을 공유해야 한다. pickle과 shared memory를 사용하여 이를 구현한다.

def write_to_shared_memory(obj, name: str) -> shared_memory.SharedMemory:
    serialized = pickle.dumps(obj)
    size = len(serialized)
    try:
        shm = shared_memory.SharedMemory(name=name)
        if shm.size < size:
            shm.close()
            shm.unlink()
            shm = shared_memory.SharedMemory(create=True, size=size, name=name)
    except FileNotFoundError:
        shm = shared_memory.SharedMemory(create=True, size=size, name=name)

    shm.buf[:size] = serialized
    return shm

def read_from_shared_memory(name: str) -> Any:
    shm = shared_memory.SharedMemory(name=name)
    data = pickle.loads(bytes(shm.buf))
    shm.close()
    return data

메인 프로세스가 port_args, server_args, scheduler_info를 shared memory에 기록하면, 각 Worker 프로세스가 이를 읽어 초기화에 사용한다. shared memory 이름에 PID를 포함시켜(f"multi_tokenizer_args_{current_pid}") 여러 SGLang 인스턴스가 충돌하지 않도록 한다.

왜 이 설계인가

Mixin 패턴: MultiHttpWorkerDetokenizerMixin은 DetokenizerManager의 기본 구현을 변경하지 않고 multi-worker 기능을 추가한다. single-worker와 multi-worker 코드가 같은 클래스에 공존하면서도 서로 간섭하지 않는다. event_loop와 multi_http_worker_event_loop는 tokenizer_worker_num 설정에 따라 선택적으로 호출된다.

http_worker_ipc 기반 라우팅: 각 요청에 출발지 Worker의 IPC 이름을 부착하는 방식은, 중앙 라우팅 테이블 없이도 결과를 올바른 Worker에게 돌려보낼 수 있게 한다. 요청 자체가 라우팅 정보를 운반하므로, DetokenizerManager나 Router가 상태를 유지할 필요가 없다.

Lazy Socket Creation: SocketMapping은 처음 보는 IPC 이름에 대해 소켓을 동적으로 생성한다. Worker가 추가되거나 재시작되더라도 별도의 등록 절차 없이 자연스럽게 연결된다. 분산 시스템에서 흔히 사용되는 service discovery 패턴의 단순한 구현이다.

상속을 통한 확장: TokenizerWorker는 TokenizerManager를 상속하여 모든 토큰화 로직을 재사용한다. 오버라이드하는 것은 _attach_multi_http_worker_info 같은 라우팅 관련 메서드뿐이다. 이는 TokenizerManager의 버그 수정이나 기능 추가가 자동으로 모든 Worker에 반영됨을 의미한다.

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글