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

The scheduler's IO and multi-rank synchronization

The scheduler's main loop calls only two IO methods, receive_msg(blocking) and send_result(reply). Behind them is SchedulerIOMixin: on one GPU it talks ZMQ directly, and under tensor parallelism only rank 0 connects to the tokenizer and every rank has to see exactly the same messages each step. This chapter implements it and explains why "exactly the same" matters so much.

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

1. Under tensor parallelism, why does every rank run a full scheduler instead of rank 0 scheduling and the others following orders?
2. The number of messages rank 0 receives varies per step. How do the other ranks know how many to take?
3. Who sends the results to the detokenizer?
4. When is receive_msg(blocking=True) called? What happens before it blocks?

Answers (try it yourself first, then expand)

1. Each rank's KV pool and page table live on its own GPU, and the scheduling decisions determine the KV layout directly. Running a full scheduler on every rank means that, given the same input messages, deterministic scheduling produces the same batch and the same KV layout everywhere, with no need to send the scheduling result (a lot of metadata) to the other ranks each step.
2. rank 0 uses one broadcast to tell the other ranks how many messages this step has, and they take that many from PUB / SUB.
3. Only rank 0: every rank samples the same token anyway, so one copy is all that is needed.
4. Only when there is no work at all (no running requests, nothing pending) does it block for a new message, so it does not spin; before blocking it finishes handling every in-flight result and runs the idle-time checks.

Files you will write: scheduler/io.py.

Files so far: 67 (1 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/  (5 files, unchanged)
├── models/  (7 files, unchanged)
├── moe/  (2 files, unchanged)
├── scheduler/
│   ├── table.py
│   ├── __init__.py
│   ├── config.py
│   ├── utils.py
│   ├── decode.py
│   ├── prefill.py
│   ├── scheduler.py
│   ├── cache.py
│   ├── io.py  ← new in this chapter
├── tokenizer/  (4 files, unchanged)
├── utils/  (8 files, unchanged)

This step's main: examples/ch14_scheduler_io.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

Why every rank schedules

Under tensor parallelism each rank holds part of the model and part of the KV cache (split by KV head, chapter 16). Every step, all ranks must compute the same batch: the same requests, in the same order, writing the same tokens to the same locations in the KV pool. Otherwise an all-reduce adds different requests' partial sums together.

One option is to have rank 0 schedule and broadcast each step's batch (the request list, positions, page allocations and so on) to the others. mini-sglang's way uses less communication: every rank runs a full scheduler, and as long as they see exactly the same input messages, deterministic scheduling logic reaches exactly the same decisions on every rank, keeping the KV pool, the page table and the radix cache in agreement throughout. The only thing to broadcast is the raw messages, a few hundred bytes each.

That requires the scheduling logic to be deterministic, which is why DecodeManager sorts by uid when assembling a batch in chapter 7, to remove any difference in set iteration order between processes.

Receiving

minisgl/scheduler/io.py · SchedulerIOMixin.__init__

def __init__(self, config: SchedulerConfig, tp_cpu_group: Any):
    tp_info = config.tp_info
    self.tp_cpu_group = tp_cpu_group
    if config.offline_mode:  # offline mode: the LLM subclass implements the IO
        self.receive_msg = self.offline_receive_msg
        self.send_result = self.offline_send_result
        return
    if tp_info.is_primary():
        self._recv_from_tokenizer = ZmqPullQueue(
            config.zmq_backend_addr, create=True, decoder=BaseBackendMsg.decoder)
        self._send_into_tokenizer = ZmqPushQueue(
            config.zmq_detokenizer_addr, create=config.backend_create_detokenizer_link,
            encoder=BaseTokenizerMsg.encoder)
    recv, send = self._recv_msg_single_rank, self._reply_tokenizer_rank0
    if tp_info.size > 1:
        if tp_info.is_primary():
            recv = self._recv_msg_multi_rank0
            self._send_into_ranks = ZmqPubQueue(
                config.zmq_scheduler_broadcast_addr, create=True, encoder=BaseBackendMsg.encoder)
            # wait for the other ranks' subscriptions to take effect: otherwise the first broadcast can be dropped and that rank waits forever
            self._send_into_ranks.wait_for_subscribers(tp_info.size - 1)
        else:
            recv, send = self._recv_msg_multi_rank1, self._reply_tokenizer_rank1
            self._recv_from_rank0 = ZmqSubQueue(
                config.zmq_scheduler_broadcast_addr, create=False, decoder=BaseBackendMsg.decoder)
    self.receive_msg = recv
    self.send_result = send

The constructor picks the receive and send methods by whether this is rank 0 and whether there are other ranks. Offline mode swaps in the two methods LLM provides (chapter 7).

The single-GPU version:

minisgl/scheduler/io.py · SchedulerIOMixin._recv_msg_single_rank

def _recv_msg_single_rank(self, blocking: bool = False) -> List[BaseBackendMsg]:
    msgs: List[BaseBackendMsg] = []
    if blocking:  # nothing to do: block for the first message
        self.run_when_idle()
        msgs.append(self._recv_from_tokenizer.get())
    while not self._recv_from_tokenizer.empty():  # then take everything that has already arrived
        msgs.append(self._recv_from_tokenizer.get())
    return msgs

With nothing to do (blocking=True) it first calls run_when_idle (one memory integrity check, chapter 8), then blocks for the first message; after that it takes everything that has already arrived without waiting further.

Rank 0 with several GPUs:

minisgl/scheduler/io.py · SchedulerIOMixin._recv_msg_multi_rank0

def _recv_msg_multi_rank0(self, blocking: bool = False) -> List[BaseBackendMsg]:
    msgs: List[BaseBackendMsg] = []
    if blocking:
        self.run_when_idle()
        raw = self._recv_from_tokenizer.get_raw()
        self._send_into_ranks.put_raw(raw)
        msgs.append(self._recv_from_tokenizer.decode(raw))
    raws: List[bytes] = []
    while not self._recv_from_tokenizer.empty():
        raws.append(self._recv_from_tokenizer.get_raw())
    count = torch.tensor(len(raws))
    dist.broadcast(count, src=0, group=self.tp_cpu_group)  # tell the other ranks how many to take this step
    for raw in raws:
        self._send_into_ranks.put_raw(raw)
        msgs.append(self._recv_from_tokenizer.decode(raw))
    return msgs

1. the first message received while blocking is forwarded verbatim right away; 2. everything that has already arrived is taken (only rank 0 knows how many that is); 3. one broadcast tells the other ranks how many more there are this step; 4. those messages are forwarded verbatim and decoded locally.

The other ranks:

minisgl/scheduler/io.py · SchedulerIOMixin._recv_msg_multi_rank1

def _recv_msg_multi_rank1(self, blocking: bool = False) -> List[BaseBackendMsg]:
    msgs: List[BaseBackendMsg] = []
    if blocking:
        self.run_when_idle()
        msgs.append(self._recv_from_rank0.get())
    count = torch.tensor(-1)
    dist.broadcast(count, src=0, group=self.tp_cpu_group)
    for _ in range(int(count)):
        msgs.append(self._recv_from_rank0.get())
    return msgs

While blocking they first take one message from PUB/SUB (matching rank 0's first), then join the broadcast to learn the count and take that many. The blocking argument always agrees on both sides, because it depends on the scheduler's state and every rank's state is identical.

The count is broadcast over the CPU gloo process group (tp_cpu_group). On a GPU the tensors go over NCCL and the control information over a separate gloo group, out of each other's way.

What the two ranks received

Without loading a model, using only the scheduler's IO: the main process plays the tokenizer and sends 3 messages in a row, and two processes play the two ranks.

examples/ch14_scheduler_io.py

"""Chapter 14: how two "scheduler ranks" receive exactly the same sequence of messages.

No model is loaded, only the scheduler's IO: the main process plays the tokenizer and sends 3
messages, rank 0 receives them over ZMQ, broadcasts "how many this step" and forwards the raw
bytes, and rank 1 takes that many from PUB/SUB.
"""

import multiprocessing as mp
import time

import torch
import torch.distributed as dist
from minisgl.core import SamplingParams
from minisgl.distributed import DistributedInfo
from minisgl.message import BaseBackendMsg, UserMsg
from minisgl.scheduler.config import SchedulerConfig
from minisgl.scheduler.io import SchedulerIOMixin
from minisgl.utils import ZmqPushQueue


class IOOnly(SchedulerIOMixin):
    def run_when_idle(self) -> None:
        pass


def rank_main(rank: int, suffix: str, out: mp.Queue) -> None:
    dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29555", rank=rank, world_size=2)
    config = SchedulerConfig(model_path="unused", tp_info=DistributedInfo(rank, 2),
                             dtype=torch.float32, _unique_suffix=suffix)
    io = IOOnly(config, dist.group.WORLD)
    dist.barrier()  # matching sync_all_ranks: start only once every subscriber has connected
    out.put(("ready", rank))
    got = []
    while len(got) < 3:
        got += [m.uid for m in io.receive_msg(blocking=not got)]
    out.put((rank, got))
    dist.barrier()


if __name__ == "__main__":
    mp.set_start_method("spawn")
    suffix = f".example{int(time.time())}"
    out: mp.Queue = mp.Queue()
    procs = [mp.Process(target=rank_main, args=(r, suffix, out)) for r in range(2)]
    for p in procs:
        p.start()
    for _ in range(2):
        out.get()
    config = SchedulerConfig(model_path="unused", tp_info=DistributedInfo(0, 2), dtype=torch.float32,
                             _unique_suffix=suffix)
    push = ZmqPushQueue(config.zmq_backend_addr, create=False, encoder=BaseBackendMsg.encoder)
    for uid in range(3):
        push.put(UserMsg(uid=uid, input_ids=torch.tensor([1, 2], dtype=torch.int32),
                         sampling_params=SamplingParams()))
    results = dict(out.get() for _ in range(2))
    for rank in (0, 1):
        print(f"the uid sequence rank {rank} received: {results[rank]}")
    for p in procs:
        p.join()
    push.stop()

What it prints

the uid sequence rank 0 received: [0, 1, 2]
the uid sequence rank 1 received: [0, 1, 2]

rank 0 blocks and receives message 1 on the first step, then takes whatever else had arrived by then; rank 1 learns the count from the broadcast and receives the same messages from PUB/SUB. However the messages arrive and in whatever batches, both ranks see the same sequence. (This demo hung the first time it ran, because of PUB/SUB's slow-joiner problem; see "Difference from upstream: XPUB" in the last chapter.)

Sending results

minisgl/scheduler/io.py · SchedulerIOMixin._reply_tokenizer_rank0

def _reply_tokenizer_rank0(self, reply: List[DetokenizeMsg]) -> None:
    if len(reply) == 1:
        self._send_into_tokenizer.put(reply[0])
    elif len(reply) > 1:
        self._send_into_tokenizer.put(BatchTokenizerMsg(data=reply))  # type: ignore[arg-type]

Only rank 0 sends results to the detokenizer, and send_result does nothing on the other ranks. Every rank samples the same thing anyway: after the all-gather every rank has the complete LM head logits, so greedy decoding agrees by construction, and with random sampling every rank uses the same seed (torch.manual_seed(42) in Engine) and the same call sequence, so the results agree too.

The official implementation

• scheduler/io.py · SchedulerIOMixin (lines 15-133)
• the most recent official commit (#113) fixes exactly the case where the decode batches on different ranks ordered their requests differently, which violates the "exactly the same" principle head-on.

Tests

Multi-rank IO is covered by the end-to-end tensor-parallel tests of chapter 16: TP=2 and TP=4 services produce output identical to single-GPU Hugging Face.

How to explain it

On multi-rank synchronization: under tensor parallelism every rank runs a full scheduler, and as long as all ranks receive exactly the same messages, deterministic scheduling assembles the same batch and allocates the same KV locations, so rank 0 never has to send its scheduling results down each step. The mechanism is that rank 0 receives messages, forwards the raw bytes verbatim, and uses one broadcast to tell the others how many there are this step, which they then take from PUB/SUB; only rank 0 sends results to the detokenizer, since every rank sampled the same thing. It blocks for a message only when there is nothing to run, and finishes the in-flight batch before blocking. The prerequisite for determinism is that scheduling never depends on anything that differs between ranks, such as the clock or a set's iteration order, which is why batches sort by uid.

Exercises

1. With random sampling, what happens if the ranks' random number generators are in different states? Only rank 0 sends results, so it looks fine. Is it?
2. What goes wrong if the count broadcast is removed and rank 1 decides for itself with empty() whether more messages remain?
3. Why does rank 0 "forward first, decode second" rather than "decode first, re-encode and forward"?

Answers

1. The ranks sample different tokens, and each rank writes its own token pool, so the next step's inputs differ, an all-reduce adds data from different sequences together, and every rank's computation is wrong (rank 0's output merely still looks like some sequence). So every rank must use the same seed and make the same number of random calls in the same order.
2. PUB/SUB messages reach rank 1 with a delay, so when rank 1 decides "no more messages", rank 0 may already have forwarded several that are still in flight. The two sides then handle different sets of messages this step and their scheduling decisions diverge.
3. It saves an encode, and it guarantees that what is forwarded is byte-for-byte identical.

Summary

• Every rank runs a full scheduler, and given the same input messages, deterministic scheduling produces the same batch and the same KV layout.
• rank 0 forwards messages verbatim and uses one broadcast to tell the other ranks how many there are this step; they take that many from PUB/SUB.
• Only rank 0 replies to the detokenizer, since every rank sampled the same thing.

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