Skip to content
mini-sglang from Scratch Stage 3 · Serving: make it a service

Messages, serialization and ZMQ

Everything in the first two parts ran in one process. To become an online service, the API server, the tokenizer and the scheduler processes have to pass messages. mini-sglang's approach is plain: each kind of message is a dataclass, a 70-line generic serializer turns it into a dictionary of primitives, msgpack turns that into bytes, and a few kinds of ZMQ socket carry them. This chapter writes all of it.

Self-test: if you can answer these, skip the chapter

  1. What kinds of messages are there, and which process sends each to which?
  2. How does a tensor (the prompt's tokens) go into msgpack?
  3. What is the difference between ZMQ's PUSH/PULL and PUB/SUB? Where does the scheduler use each?
  4. Why does rank 0 forward the raw bytes to the other ranks rather than the decoded message?
Answers (try it yourself first, then expand)
  1. Three groups, by receiver: to the tokenizer / detokenizer (TokenizeMsg from the API server, DetokenizeMsg carrying a new token from the scheduler, and AbortMsg); to the scheduler (UserMsg from the tokenizer, AbortBackendMsg, ExitMsg); and to the API server (UserReply from the detokenizer). Each group also has a Batch...Msg that packs several together, and each is deserialized by class name inside its own module.
  2. A one-dimensional tensor is stored as raw bytes (along with its dtype and so on) in a msgpack bytes field, and the receiver rebuilds the tensor from them.
  3. PUSH / PULL is a point-to-point pipe (one receiver takes each message); PUB / SUB is a broadcast (every subscriber gets a copy). The scheduler PULLs requests from upstream and PUSHes results out, and rank 0 PUBs what it receives to the other ranks.
  4. Forwarding the raw bytes guarantees every rank sees exactly the same message, saves a decode-then-encode round trip, and avoids any difference re-encoding might introduce.

Files you will write: message/utils.py, message/backend.py, message/tokenizer.py, message/frontend.py, message/__init__.py, utils/mp.py.

Files so far: 63 (6 new, 0 modified, 0 brought in early)
minisgl/
├── __init__.py
├── core.py
├── env.py
├── attention/  (4 files, unchanged)
├── distributed/  (3 files, unchanged)
├── engine/  (5 files, unchanged)
├── kernel/  (2 files, unchanged)
├── kvcache/  (5 files, unchanged)
├── layers/  (8 files, unchanged)
├── llm/  (2 files, unchanged)
├── message/
│   ├── __init__.py  ← new in this chapter
│   ├── backend.py  ← new in this chapter
│   ├── utils.py  ← new in this chapter
│   ├── tokenizer.py  ← new in this chapter
│   ├── frontend.py  ← new in this chapter
├── models/  (7 files, unchanged)
├── moe/  (2 files, unchanged)
├── scheduler/  (9 files, unchanged)
├── utils/
│   ├── __init__.py
│   ├── misc.py
│   ├── hf.py
│   ├── torch_utils.py
│   ├── registry.py
│   ├── logger.py
│   ├── device.py
│   ├── mp.py  ← new in this chapter
This step's main: examples/ch12_message.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

Three groups of messages

They split into three groups by receiver, each with a base class:

Base class Receiver Messages
BaseTokenizerMsg tokenizer / detokenizer TokenizeMsg (text or a conversation plus sampling parameters), DetokenizeMsg (one new token plus whether it finished), AbortMsg
BaseBackendMsg the scheduler UserMsg (a token tensor plus sampling parameters), AbortBackendMsg, ExitMsg
BaseFrontendMsg the API server UserReply (the delta text plus whether it finished)

Each group also has a Batch...Msg that packs several messages into one and cuts the message count.

minisgl/message/backend.py
"""Messages sent to the scheduler (the backend)."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List

import torch
from minisgl.core import SamplingParams

from .utils import deserialize_type, serialize_type


@dataclass
class BaseBackendMsg:
    def encoder(self) -> Dict:
        return serialize_type(self)

    @staticmethod
    def decoder(json: Dict) -> BaseBackendMsg:
        return deserialize_type(globals(), json)


@dataclass
class BatchBackendMsg(BaseBackendMsg):
    data: List[BaseBackendMsg]


@dataclass
class ExitMsg(BaseBackendMsg):
    pass


@dataclass
class UserMsg(BaseBackendMsg):
    uid: int
    input_ids: torch.Tensor  # a one-dimensional int32 tensor on the CPU
    sampling_params: SamplingParams


@dataclass
class AbortBackendMsg(BaseBackendMsg):
    uid: int
Note the globals() in decoder: deserialization looks the class up by name in the global namespace of the module where the message is defined. That is why the three groups live in three files with their own decoder, since each receiver only knows the group addressed to it.

Generic serialization

minisgl/message/utils.py
"""Message serialization: dataclass to a dict of primitives and back, with msgpack turning it into bytes.

Each object serializes to {"__type__": class name, fields...}; a one-dimensional tensor becomes raw
bytes plus a dtype. Deserialization looks the class up by name in a given table (usually the
globals() of the module defining the message) and restores the fields recursively.
"""

from __future__ import annotations

from typing import Any, Dict, Type

import numpy as np
import torch


def _serialize_any(value: Any) -> Any:
    if isinstance(value, dict):
        return {k: _serialize_any(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return type(value)(_serialize_any(v) for v in value)
    if isinstance(value, (int, float, str, type(None), bool, bytes)):
        return value
    return serialize_type(value)


def serialize_type(obj: Any) -> Dict:
    if isinstance(obj, torch.Tensor):
        assert obj.dim() == 1, "only 1D tensors are supported"
        return {"__type__": "Tensor", "buffer": obj.numpy().tobytes(), "dtype": str(obj.dtype)}
    result: Dict[str, Any] = {"__type__": type(obj).__name__}
    for k, v in obj.__dict__.items():
        result[k] = _serialize_any(v)
    return result


def _deserialize_any(cls_map: Dict[str, Type], data: Any) -> Any:
    if isinstance(data, dict):
        if "__type__" in data:
            return deserialize_type(cls_map, data)
        return {k: _deserialize_any(cls_map, v) for k, v in data.items()}
    if isinstance(data, (list, tuple)):
        return type(data)(_deserialize_any(cls_map, d) for d in data)
    return data


def deserialize_type(cls_map: Dict[str, Type], data: Dict) -> Any:
    type_name = data["__type__"]
    if type_name == "Tensor":
        np_dtype = getattr(np, data["dtype"].replace("torch.", ""))
        return torch.from_numpy(np.frombuffer(data["buffer"], dtype=np_dtype).copy())
    kwargs = {k: _deserialize_any(cls_map, v) for k, v in data.items() if k != "__type__"}
    return cls_map[type_name](**kwargs)
The rules are simple: any object becomes {"__type__": class name, field: value...} with its fields handled recursively; primitives stay as they are; and a one-dimensional tensor becomes raw bytes plus a dtype string. Deserialization goes the other way: a __type__ means construct an object of that class.

examples/ch12_message.py
"""Chapter 12: one message from object to bytes and back; ZMQ's PUSH/PULL and PUB/SUB."""

import threading

import msgpack
import torch
from minisgl.core import SamplingParams
from minisgl.message import BaseBackendMsg, BatchBackendMsg, UserMsg
from minisgl.utils import ZmqPubQueue, ZmqPullQueue, ZmqPushQueue, ZmqSubQueue

msg = BatchBackendMsg(data=[
    UserMsg(uid=7, input_ids=torch.tensor([9707, 11, 1879], dtype=torch.int32),
            sampling_params=SamplingParams(temperature=0.6, max_tokens=64)),
])
d = msg.encoder()
print("the dict encoder() produces:")
print(" ", {k: v for k, v in d.items() if k != "data"}, "data[0] =")
inner = dict(d["data"][0])
print("   ", inner)
raw = msgpack.packb(d, use_bin_type=True)
print(f"{len(raw)} bytes after msgpack")
back = BaseBackendMsg.decoder(msgpack.unpackb(raw, raw=False))
print("decoded back:", back.data[0].uid, back.data[0].input_ids, back.data[0].sampling_params)

# PUSH/PULL: a one-to-one queue; PUB/SUB: a one-to-many broadcast
addr = "ipc:///tmp/minisgl_example_pushpull"
pull = ZmqPullQueue(addr, create=True, decoder=BaseBackendMsg.decoder)
push = ZmqPushQueue(addr, create=False, encoder=BaseBackendMsg.encoder)
threading.Thread(target=lambda: [push.put(UserMsg(uid=i, input_ids=torch.tensor([i], dtype=torch.int32),
                                                  sampling_params=SamplingParams())) for i in range(3)]).start()
print("PULL received uid:", [pull.get().uid for _ in range(3)], " queue empty:", pull.empty())

paddr = "ipc:///tmp/minisgl_example_pubsub"
pub = ZmqPubQueue(paddr, create=True, encoder=BaseBackendMsg.encoder)
subs = [ZmqSubQueue(paddr, create=False, decoder=BaseBackendMsg.decoder) for _ in range(2)]
pub.wait_for_subscribers(2)  # anything published before a subscription takes effect is dropped, so wait for both subscribers first
pub.put_raw(raw)  # rank 0 forwards the raw bytes verbatim, with no re-encoding
print("both SUBs received it:", [s.get().data[0].uid for s in subs])
for q in [pull, push, pub, *subs]:
    q.stop()
What it prints
the dict encoder() produces:
  {'__type__': 'BatchBackendMsg'} data[0] =
    {'__type__': 'UserMsg', 'uid': 7, 'input_ids': {'__type__': 'Tensor', 'buffer': b'\xeb%\x00\x00\x0b\x00\x00\x00W\x07\x00\x00', 'dtype': 'torch.int32'}, 'sampling_params': {'__type__': 'SamplingParams', 'temperature': 0.6, 'top_k': -1, 'top_p': 1.0, 'ignore_eos': False, 'max_tokens': 64}}
229 bytes after msgpack
decoded back: 7 tensor([9707,   11, 1879], dtype=torch.int32) SamplingParams(temperature=0.6, top_k=-1, top_p=1.0, ignore_eos=False, max_tokens=64)
PULL received uid: [0, 1, 2]  queue empty: True
both SUBs received it: [7, 7]
A request message carrying 3 tokens is a little over 200 bytes, most of it field names. Repeating the field names in every message is what generic serialization costs; what it buys is that a new kind of message needs only a dataclass and no encoding or decoding code at all. The prompt's tokens are stored as raw bytes (4 bytes each as int32), so even an 8000-token prompt is only 32 KB.

ZMQ queues

minisgl/utils/mp.py · ZmqPushQueue
class ZmqPushQueue(Generic[T]):
    def __init__(self, addr: str, create: bool, encoder: Callable[[T], Dict]):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PUSH)
        self.socket.bind(addr) if create else self.socket.connect(addr)
        self.encoder = encoder

    def put(self, obj: T) -> None:
        self.socket.send(msgpack.packb(self.encoder(obj), use_bin_type=True), copy=False)

    def stop(self) -> None:
        self.socket.close(linger=0)
        self.context.term()
utils/mp.py wraps ZMQ sockets into five kinds of queue, all thin layers over "encoder plus msgpack plus socket":

  • PUSH / PULL: a one-to-one (or many-to-one) pipe where messages arrive in order and are not dropped. tokenizer to scheduler, scheduler to detokenizer and detokenizer to API server all use it;
  • the async versions: the API server runs in an asyncio event loop and uses zmq.asyncio sockets, so await on send and receive does not block the loop;
  • PUB / SUB: a one-to-many broadcast, used by scheduler rank 0 to forward messages to the other ranks.

Whichever side passes create=True does the bind (claiming the address) and the other does the connect. The address is ipc:///tmp/minisgl_N.pid=the launching process's PID, so several services on one machine do not collide.

PUB/SUB has a classic trap, the slow joiner: a subscription only takes effect asynchronously after the connection is established, and anything published before that is dropped outright. If rank 0's first broadcast beats rank 1's subscription, rank 1 waits forever. Upstream dodges it by startup timing (after the scheduler initializes it still waits for the tokenizer and the API server, which leaves plenty of room), but that is a timing gap, not a guarantee. We switch the publishing side to XPUB, which is told when a subscriber subscribes, so wait_for_subscribers(n) returns only once all n have arrived:

minisgl/utils/mp.py · ZmqPubQueue
class ZmqPubQueue(Generic[T]):
    """The broadcast queue. XPUB rather than PUB: XPUB is told when a subscriber subscribes, so it can
    wait for all of them before sending.

    A classic PUB/SUB trap (the "slow joiner"): a subscription only takes effect asynchronously after
    the connection is established, and anything published before that is dropped outright.
    """

    def __init__(self, addr: str, create: bool, encoder: Callable[[T], Dict]):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.XPUB)
        self.socket.setsockopt(zmq.XPUB_VERBOSE, 1)  # every subscriber's notification is wanted, with no deduplication
        self.socket.bind(addr) if create else self.socket.connect(addr)
        self.encoder = encoder

    def wait_for_subscribers(self, n: int) -> None:
        for _ in range(n):
            self.socket.recv()  # a subscription notification: b"\x01" plus the topic

    def put_raw(self, raw: bytes) -> None:
        self.socket.send(raw, copy=False)

    def put(self, obj: T) -> None:
        self.socket.send(msgpack.packb(self.encoder(obj), use_bin_type=True), copy=False)

    def stop(self) -> None:
        self.socket.close(linger=0)
        self.context.term()
XPUB_VERBOSE makes every subscriber's notification arrive (by default repeated subscriptions to the same topic are collapsed into one).

Why the raw bytes are forwarded

After rank 0 receives a message from the tokenizer it forwards it verbatim to the other ranks. ZmqPullQueue offers both get_raw() and decode(): rank 0 takes the raw bytes, put_raws them to the other ranks, and only then decodes them itself. That saves a decode-then-encode round trip and, more importantly, guarantees that every rank receives a byte-for-byte identical message, since all ranks must make exactly the same scheduling decisions (chapter 14).

The official implementation

Our queues use close(linger=0) in stop(); upstream uses the default close(), which can hang on exit waiting for messages that have not been sent.

Difference from upstream: XPUB

The official ZmqPubQueue uses a plain PUB socket and does not wait for subscribers. The demo program of chapter 14 (the main process starts two ranks and sends a message immediately) hangs every single time on this book's machine with the official version: rank 0 receives the first message, broadcasts it, then waits for rank 1 in the collective that counts the broadcasts, while rank 1 never receives the broadcast that was dropped. With XPUB and waiting for subscribers it does not. tests/test_ch12_message.py has a test for it.

Tests

tests/test_ch12_message.py · test_zmq_push_pull_over_ipc
def test_zmq_push_pull_over_ipc(ipc_dir):
    addr = f"ipc://{ipc_dir}/q"
    pull = ZmqPullQueue(addr, create=True, decoder=BaseTokenizerMsg.decoder)
    push = ZmqPushQueue(addr, create=False, encoder=BaseTokenizerMsg.encoder)
    msgs = [TokenizeMsg(uid=1, text=[{"role": "user", "content": "你好"}], sampling_params=SamplingParams()),
            DetokenizeMsg(uid=1, next_token=42, finished=True)]
    threading.Thread(target=lambda: [push.put(m) for m in msgs]).start()
    got = [pull.get(), pull.get()]
    assert got == msgs and pull.empty()
    push.stop()
    pull.stop()

ipc_dir, not pytest's tmp_path

A Unix domain socket's path lives in sockaddr_un.sun_path, which has a hard length limit: 108 bytes on Linux, only 104 on macOS. On macOS pytest's tmp_path looks like /private/var/folders/9k/3m_.../T/pytest-of-you/pytest-12/test_zmq_push_pull_over_ip0/ — the directory alone is over a hundred bytes, and building a socket path on top of it is rejected outright by ZMQ (ipc path ... is longer than 103 characters). So ipc_dir in conftest.py creates a short directory under /tmp instead. Real inference engines face the same limit: this book's ipc:///tmp/minisgl_0.pid=... is that short on purpose.

How to explain it

On inter-process communication: three groups of messages, from the API server to the tokenizer (the request), from the tokenizer to the scheduler (the tokenized request), and from the scheduler through the detokenizer back to the API server (the results). Generic serialization turns an object into {"__type__": class name, ...} and encodes it with msgpack, with one-dimensional tensors stored as raw bytes, which is faster and safer than pickle. ZMQ's PUSH/PULL is a point-to-point pipe and PUB/SUB is a broadcast (subscribers have to connect first or early messages are lost); rank 0 broadcasts the raw bytes of a request verbatim to the other ranks so that every rank sees exactly the same message. To explain why NCCL is not used for control messages: they are small and irregular, and NCCL only suits tensors.

Exercises

  1. Serialization only supports one-dimensional tensors. Which lines change to send a two-dimensional [batch, seq] tensor?
  2. Generic serialization carries the field names in every message. Design a more compact format (storing the fields as a list in order, say). What does it cost?
  3. Why does the API server use the async queues while the scheduler and the tokenizer use the synchronous ones?
Answers
  1. Store shape as well when serializing (list(obj.shape)) and np.frombuffer(...).reshape(shape) when deserializing; drop the dim() == 1 assertion.
  2. Both sides have to agree on the field order, so adding or removing a field means changing both at once and processes at different versions cannot talk; readability and debugging suffer too. msgpack is already compact and this is rarely the bottleneck.
  3. The API server handles hundreds or thousands of HTTP connections in one event loop, where any blocking call stalls all of them; the scheduler and the tokenizer are single-threaded loops that already alternate between receiving and processing, so synchronous calls are simpler.

Summary

  • Three groups of messages go to the tokenizer, the scheduler and the API server, each deserialized by class name inside its own module.
  • Generic serialization: object to {"__type__": class name, ...} to msgpack bytes, with one-dimensional tensors stored as raw bytes.
  • PUSH/PULL makes the point-to-point pipes and PUB/SUB the broadcast from rank 0 to the other ranks; subscribers must connect first.
  • rank 0 forwards the raw bytes verbatim so that every rank sees the same message.

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

Found a mistake or a better way to say it? Edit this page on GitHubReport an issue