Skip to content
mini-sglang from Scratch Stage 2 · Scheduling it well: many requests at once

Overlap scheduling: hiding the CPU cost

One decode step of a small model takes the GPU a few milliseconds; what the scheduler has to do on the CPU each step, namely receive messages, assemble the batch, prepare metadata, handle the previous step's results and send messages, takes a few milliseconds too. Taking turns leaves the GPU waiting on the CPU almost half the time. Overlap scheduling, from NanoFlow and introduced in SGLang 0.4, lets the CPU handle step N's results while the GPU computes step N+1. mini-sglang's implementation adds barely a dozen lines, but it changes when you know what, which is where the four problems this reimplementation found come from.

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

  1. When step N+1 is launched the CPU does not yet know what step N sampled. Where does step N+1's input come from?
  2. A request samples EOS at step N. Can it appear in step N+1?
  3. Why do the scheduler and the engine need two different CUDA streams?
  4. Under overlap scheduling, when step N's results are handled, which step does req.device_len reflect?
Answers (try it yourself first, then expand)
  1. From the token pool on the GPU: step N's sample is written straight into the token pool on the GPU, and step N+1's forward pass reads its input from there, so the CPU never has to know what it is.
  2. It can: when step N+1 is launched the CPU has not handled step N's results and does not know it has finished, so it is scheduled one more time; the extra token is discarded when the results are handled.
  3. So that the scheduler's work preparing the next step (copying metadata, assembling the batch) runs in parallel with the engine's compute, with events establishing the dependencies; on one stream they would be serialized again.
  4. After step N+1: the request state on the CPU was advanced at launch time, so it runs one step ahead of the results being handled.

Files you will write: overlap_loop and run_forever in scheduler/scheduler.py, plus the overlap-related parts of _process_last_data, _free_req_resources and _process_one_msg.

Files so far: 62 (0 new, 1 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/  (4 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  ← modified in this chapter
│   ├── cache.py
│   ├── io.py
├── utils/  (8 files, unchanged)
This step's main: examples/ch11_overlap.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

Animation: overlap scheduling and the four problems it brings (about 2 minutes, Chinese narration and subtitles)

The two loops

minisgl/scheduler/scheduler.py · Scheduler.overlap_loop
def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None:
    """The overlap loop: launch step N's compute first, then handle step N-1's results on the CPU.

    While the GPU runs step N, the CPU handles step N-1's results and schedules step N+1, so
    the CPU cost is hidden. This works because step N's input token never has to reach the CPU:
    step N-1's sample writes it into token_pool on the GPU directly (see _forward).
    """
    self._defer_slot_free = True
    if last_data is None:  # no batch in flight: the slots whose release was deferred can safely be reused
        self._release_deferred_slots()
    blocking = not (last_data is not None or self.prefill_manager.runnable
                    or self.decode_manager.runnable)
    for msg in self.receive_msg(blocking=blocking):
        self._process_one_msg(msg)
    forward_input = self._schedule_next_batch()
    ongoing = None
    if forward_input is not None:
        with self.engine_stream_ctx:
            self.engine.stream.wait_stream(self.stream)  # wait for this step's metadata to be ready
            ongoing = (forward_input, self._forward(forward_input))
    self._process_last_data(last_data)
    return ongoing
Against normal_loop there is one difference: launch this step first, then handle the previous one. overlap_loop takes the previous step's data and returns this step's, and run_forever strings them together:

minisgl/scheduler/scheduler.py · Scheduler.run_forever
@torch.inference_mode()
def run_forever(self) -> NoReturn:
    if ENV.DISABLE_OVERLAP_SCHEDULING:
        with self.engine_stream_ctx:
            self.engine.stream.wait_stream(self.stream)
            while True:
                self.normal_loop()
    else:
        data = None
        while True:
            data = self.overlap_loop(data)
With one request and 3 tokens, the order of events (the first half of the example):

normal loop:   launch 1  handle 1  launch 2  handle 2  launch 3  handle 3
overlap loop:  launch 1  launch 2  handle 1  launch 3  handle 2  handle 3

On a GPU, "launch" only queues kernels on a stream and returns immediately; "handle" has to wait for the result to be copied back to the CPU. In the overlap loop, by the time step 1 is handled the GPU is already computing step 2, and the CPU's cost hides behind the GPU's compute:

normal loop: the CPU and the GPU take turns CPU GPU sched 1 compute 1 handle 1 sched 2 compute 2 handle 2 sched 3 compute 3 handle 3 sched 4 compute 4 handle 4 GPU idle overlap loop: launch step N+1, then handle step N CPU GPU sched 1 compute 1 sched 2 compute 2 sched 3 handle 1 compute 3 sched 4 handle 2 compute 4 sched 5 handle 3 compute 5 sched 6 handle 4 compute 6 This needs the previous sample to write the next input into the token pool on the GPU, and the request state to advance at launch (complete_one). The price: while step N is handled the CPU state is one step ahead, so a request that hits EOS at step N gets scheduled once more.
the CPU / GPU timeline of the normal loop and the overlap loop

Three prerequisites

All of this rests on three designs laid down in earlier chapters.

First, the next step's input never passes through the CPU. Step N+1's decode input is the token step N sampled, which the CPU does not know when step N+1 is launched. But _forward wrote the sample straight into the token pool's next position on the GPU, and step N+1 reads its input from there (the write_tuple and input_tuple of chapter 7). A GPU executes a stream in order, so the read always comes after the write.

Second, the state scheduling needs is advanced early on the CPU. Engine.forward_batch calls complete_one() on each request right after launching, so by the time step N+1 is scheduled the request's cached_len and device_len are already the "after step N" values, which is enough to compute positions and allocate KV pages correctly. The only unknown is step N's sample, that is, whether the request hit EOS at step N. So a request that hits EOS at step N gets scheduled once more.

Third, two streams. The scheduler prepares metadata on its own stream (copying positions, indices and so on asynchronously from pinned memory to the GPU) while the engine computes on another. Before launching, self.engine.stream.wait_stream(self.stream) makes the compute wait for the metadata copies; meanwhile the next step's preparation on the scheduler's stream need not wait for the current compute. ForwardInput keeps every tensor this step uses, indices included, alive until the results are handled, so none of them is freed and its memory reused while the GPU is still reading it.

The four problems that follow

"The CPU state is one step ahead" and "a request may be scheduled one step too many" make handling results subtle. The official implementation handles one case (avoiding a double free), and this reimplementation found four more problems, all reproducible in the example and the tests:

examples/ch11_overlap.py
"""Chapter 11: the order of events under overlap scheduling, and the two problems upstream's finished check causes."""

from typing import List, Set

import torch
from minisgl.core import Req, SamplingParams
from minisgl.env import ENV
from minisgl.llm import LLM
from minisgl.message import DetokenizeMsg
from minisgl.scheduler.prefill import ChunkedReq

PATH = "models/Qwen3-0.6B"
KW = dict(dtype=torch.float32, max_running_req=8, num_page_override=512, max_seq_len_override=128)


def trace_events(overlap: bool) -> None:
    ENV.DISABLE_OVERLAP_SCHEDULING.value = not overlap
    llm = LLM(PATH, **KW)
    step = {"n": 0}
    orig_forward, orig_process = llm._forward, llm._process_last_data

    def forward(fi):
        step["n"] += 1
        fi.batch.step = step["n"]
        print(f"    launch step {step['n']} ({fi.batch.phase})")
        return orig_forward(fi)

    def process(data):
        if data is not None:
            print(f"    handle step {data[0].batch.step}'s results")
        return orig_process(data)

    llm._forward, llm._process_last_data = forward, process
    llm.generate(["The capital of France is"], SamplingParams(max_tokens=3, ignore_eos=True))
    llm.shutdown()


print("normal loop:")
trace_events(overlap=False)
print("overlap loop:")
trace_events(overlap=True)


class UpstreamLogicLLM(LLM):
    """Swap _process_last_data for upstream's check (everything else unchanged), to reproduce the problem."""

    def _process_last_data(self, last_data) -> None:
        if last_data is None:
            return
        batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1]
        copy_done.synchronize()
        reply: List[DetokenizeMsg] = []
        new_finished: Set[Req] = set()
        with self.cache_manager.lazy_free_region():
            for i, req in enumerate(batch.reqs):
                if isinstance(req, ChunkedReq):
                    continue
                req.append_host(next_tokens_cpu[i].unsqueeze(0))
                token = int(next_tokens_cpu[i].item())
                finished = not req.can_decode  # upstream: judged by can_decode
                if not req.sampling_params.ignore_eos:
                    finished |= token == self.eos_token_id
                reply.append(DetokenizeMsg(uid=req.uid, next_token=token, finished=finished))
                if finished and req not in self.finished_reqs:  # upstream: it only guards against a double free
                    self.decode_manager.remove_req(req)
                    self._free_req_resources(req)
                    new_finished.add(req)
                elif batch.is_prefill:
                    self.cache_manager.cache_req(req, finished=False)
        self.finished_reqs = new_finished
        self.send_result(reply)


def replies_of(cls, prompt: str, params: SamplingParams, eos: int | None = None):
    ENV.DISABLE_OVERLAP_SCHEDULING.value = False
    llm = cls(PATH, **KW)
    if eos is not None:
        llm.eos_token_id = eos
    replies = []
    orig = llm.offline_send_result
    llm.send_result = lambda reply: (replies.extend(reply), orig(reply))
    llm.generate([prompt], params)
    llm.shutdown()
    return [(m.next_token, m.finished) for m in replies]


print("\nProblem one: max_tokens=3, the (token, finished) sent to the detokenizer:")
for cls in (UpstreamLogicLLM, LLM):
    print(f"  {cls.__name__:<17}", replies_of(cls, "The capital of France is",
                                          SamplingParams(max_tokens=3, ignore_eos=True)))
print("Problem two: eos set to the 3rd output token (576), max_tokens=8:")
for cls in (UpstreamLogicLLM, LLM):
    print(f"  {cls.__name__:<17}", replies_of(cls, "The capital of France is",
                                          SamplingParams(max_tokens=8), eos=576))
What it prints
normal loop:
    launch step 1 (prefill)
    handle step 1's results
    launch step 2 (decode)
    handle step 2's results
    launch step 3 (decode)
    handle step 3's results
overlap loop:
    launch step 1 (prefill)
    launch step 2 (decode)
    handle step 1's results
    launch step 3 (decode)
    handle step 2's results
    handle step 3's results

Problem one: max_tokens=3, the (token, finished) sent to the detokenizer:
  UpstreamLogicLLM  [(12095, False), (13, True), (576, True)]
  LLM               [(12095, False), (13, False), (576, True)]
Problem two: eos set to the 3rd output token (576), max_tokens=8:
  UpstreamLogicLLM  [(12095, False), (13, False), (576, True), (6722, False)]
  LLM               [(12095, False), (13, False), (576, True)]

Problem one: the finished flag arrives one token early

Upstream uses finished = not req.can_decode to tell whether max_tokens has been reached. But when step N is handled, step N+1 has already been launched and device_len has already moved on, so can_decode reflects the state "after step N+1". A request with max_tokens=3 is therefore marked finished at its second token (the (13, True) above), and the third token comes with another finished message.

The offline interface happens to be unaffected, since it records the token of every message, but an online frontend ends its response at the first finished=True, so the client only ever receives max_tokens - 1 tokens. The fix is to use a quantity independent of the scheduling rhythm, namely how many tokens the CPU has actually received:

finished = len(req.input_ids) >= req.max_device_len

Problem two: one stale message after EOS

A request that hits EOS at step N is scheduled once more (prerequisite two). Upstream records the requests that finished last step in finished_reqs and skips releasing their resources when handling step N+1, but still builds a reply message for them, the (6722, False) above. The frontend has already deleted that request and ignores the message; the detokenizer, however, creates a fresh decoding state for it that never gets an end, which leaks a little memory. The fix is to skip these requests entirely at the top of the loop.

minisgl/scheduler/scheduler.py · Scheduler._process_last_data
def _process_last_data(self, last_data: ForwardData | None) -> None:
    if last_data is None:
        return
    batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1]
    copy_done.synchronize()  # wait for the sample to reach the CPU, which also means this batch and every earlier one are done
    self._release_deferred_slots()
    reply: List[DetokenizeMsg] = []
    new_finished: Set[Req] = set()
    with self.cache_manager.lazy_free_region():
        for i, req in enumerate(batch.reqs):
            if isinstance(req, ChunkedReq):  # a chunked request produces no output this step
                continue
            if req in self.finished_reqs:
                # under overlap scheduling the request finished last step (it hit EOS) but that was not handled yet, so this step had already
                # it was already included. This step's result is redundant and is dropped (upstream sends one extra message, which the frontend happens to ignore).
                continue
            next_token = next_tokens_cpu[i]
            req.append_host(next_token.unsqueeze(0))
            token = int(next_token.item())
            # req.can_decode cannot decide whether max_tokens was reached: under overlap scheduling, by the time step N's results are handled
            # step N+1 has already been launched and req's device_len has moved on, so it would mark finished one token early.
            # judge by the number of tokens the CPU has received, which is independent of the scheduling rhythm.
            finished = len(req.input_ids) >= req.max_device_len
            if not req.sampling_params.ignore_eos:
                finished |= token == self.eos_token_id
            reply.append(DetokenizeMsg(uid=req.uid, next_token=token, finished=finished))
            if finished:
                self.decode_manager.remove_req(req)
                self._free_req_resources(req)
                new_finished.add(req)
            elif batch.is_prefill:  # prefill has just finished: hand the prompt's KV to the prefix cache right away
                self.cache_manager.cache_req(req, finished=False)
    self.finished_reqs = new_finished
    self.send_result(reply)

Problem three: the request slot is reused too early

request X samples EOS at step N launch N+1 (X included) handle N: X ends, slot freed schedule N+2: can the slot be reused? handle N+1: X's result is stale GPU: running step N+1 (it will write its sample into X's token-pool row) step N+2 upstream: X's slot goes to a new request at once → the scheduler stream and the engine stream write the same row with no ordering this book: the slot is returned only at "handle N+1", once the in-flight batch is done, and X's stale result is dropped whole the same logic also fixes: finished one token early, stale messages, and pages freed twice when a request aborts mid-prefill
what happens after a request hits EOS at step N

A request hits EOS at step N, and when step N's results are handled its resources are released, including the row it holds in the page table / token pool. But step N+1 is running on the GPU right then and still contains that request: it will write its sample into that row of the token pool on the engine stream. If the next scheduling step hands the row to a new request immediately, the write on the scheduler stream (the new request's prompt) has no ordering guarantee against the old write on the engine stream, and the new request's input can be overwritten.

The KV pages have no such problem, because the batch that writes the new request's KV is queued on the engine stream after step N+1. The request slot is different because it is written by the scheduler stream. The fix: in overlap mode a released request slot is only really returned at "the next time results are handled", by which point copy_done.synchronize() has guaranteed the in-flight batch is done. With no batch in flight (last_data is None) it is returned at once, or an interactive setup with only one request slot would stall.

minisgl/scheduler/scheduler.py · Scheduler._free_req_resources
def _free_req_resources(self, req: Req) -> None:
    # the request slot (a row of the page table / token pool): under overlap scheduling the GPU may still hold a batch containing this request
    # batch is running (it will write its sample into this token_pool row). Handing the row to a new request now would let the scheduler stream's
    # writes to this row race the old batch on the engine stream. So it is deferred to the next time results are handled, once the in-flight batch is known to be done.
    # KV pages need no deferral: the batch writing the new request's KV queues behind the old one on the engine stream.
    if self._defer_slot_free:
        self._deferred_slots.append(req.table_idx)
    else:
        self.table_manager.free(req.table_idx)
    self.cache_manager.cache_req(req, finished=True)

Problem four: an abort while prefill is in flight

A disconnecting client triggers an abort. If the abort arrives just after this request's prefill was launched but before its results were handled, upstream's sequence is: the abort releases the request (cache_req(finished=True), so the prompt's KV enters the radix cache); then the prefill results are handled, the request has not finished, and cache_req(finished=False) is called once more, on a request that has already been released. The second insertion finds the prefix already in the cache and frees "the duplicate copy", which happens to be the cache's own copy: the same pages end up both in the free list and in the cache tree. They may then be handed to a new request and overwrite the cached KV, and the integrity check while the scheduler is idle raises as well (under the radix cache it reproduces as free_pages(1024) + cache_pages(5) != num_pages(1024)).

The fix is one line: add the request to finished_reqs on abort, and problem two's logic discards the in-flight batch's results for it wholesale.

Differences from upstream

All four of these fix the official overlap scheduling, and each has a test: put the official version back and the test fails. Reproducing them on a CPU relies only on the ordering "step N+1 is already launched when step N is handled", not on real GPU concurrency; problem three's test checks the invariant directly, that a request slot being allocated does not belong to any request in an in-flight batch.

The official implementation

What it means on a CPU

There is no asynchronous execution on a CPU: NullStream and NullEvent are no-ops and the computation is already done by the time it is "launched". So overlap scheduling is no faster on a CPU, but the scheduling logic is identical, and the four problems above still appear and are still caught by the tests. The performance gain has to be measured on a GPU; the official README suggests an ablation with MINISGL_DISABLE_OVERLAP_SCHEDULING=1. The smaller the model and the shorter each step on the GPU, the larger the CPU's share and the more overlap wins. Chapter 21 shows how to run that ablation on a GPU.

Tests

tests/test_ch11_overlap.py · test_overlap_equals_normal_and_hf
@pytest.mark.parametrize("cache,chunk,page_size", [("radix", 8192, 1), ("naive", 5, 1), ("radix", 7, 4)])
def test_overlap_equals_normal_and_hf(cache, chunk, page_size):
    results = {}
    for overlap in (True, False):
        ENV.DISABLE_OVERLAP_SCHEDULING.value = not overlap
        llm = build_llm(QWEN3, cache_type=cache, max_extend_tokens=chunk, page_size=page_size,
                        num_page_override=1024 // page_size)
        results[overlap] = [o["token_ids"] for o in llm.generate(PROMPTS, greedy(10))]
        llm.cache_manager.check_integrity()
        llm.shutdown()
    assert results[True] == results[False] == hf_greedy(QWEN3, PROMPTS, 10)
Both loops across three configurations (radix / naive, chunking, page size 4) produce output identical to Hugging Face. Four more tests cover the four problems above:

tests/test_ch11_overlap.py · test_abort_while_prefill_is_in_flight
def test_abort_while_prefill_is_in_flight():
    """An abort arriving after prefill was launched but before its results were handled: the request is released and the in-flight batch's results must be discarded."""
    from minisgl.message import AbortBackendMsg

    ENV.DISABLE_OVERLAP_SCHEDULING.value = False
    llm = build_llm(QWEN3, cache_type="radix")
    calls = {"n": 0}
    orig = llm.offline_receive_msg

    def receive(blocking=False):
        calls["n"] += 1
        if calls["n"] == 2:  # abort it at the start of the second step, while the prefill is "on the GPU"
            return [AbortBackendMsg(uid=0)]
        return orig(blocking)

    llm.receive_msg = receive
    replies = []
    llm.send_result = lambda reply: replies.extend(reply)
    llm.generate(PROMPTS[:1], greedy(8))
    assert replies == []  # an aborted request produces no output
    llm.cache_manager.check_integrity()
    assert llm.table_manager.available_size == 8
    llm.shutdown()

How to explain it

On overlap scheduling: launch step N+1 first, then handle step N's results on the CPU, so the CPU cost of scheduling, batching and detokenization hides behind the GPU's compute. There are three prerequisites: step N+1's input token is written on the GPU by step N's sample (the CPU never needs the value); the request state is advanced at launch time; and the scheduler and the engine use two streams. The price is that the CPU state runs a step ahead of the GPU: a request that sampled EOS at step N may already be in step N+1, so its extra result has to be discarded; the finished check has to use the number of tokens actually received; and a request slot can only be reused once the in-flight batch is done. The gain is largest with small models and small batches, where a GPU step is very short.

Exercises

  1. In overlap_loop, what do you get if you swap "launch this step" and "handle the previous one" back while keeping the way last_data is passed along?
  2. The extra step in problem two wastes one computation. Could you know whether step N hit EOS before scheduling step N+1? At what cost?
  3. Why does problem three affect only the request slot, while KV pages can be reused at once? If a new request hits the prefix cache, the scheduler stream writes the hit locations into the page table; does that raise a similar problem?
Answers
  1. It becomes equivalent to the normal loop (launch, then wait for the result immediately), losing the overlap but staying correct.
  2. You could synchronize on step N's result before launching step N+1, but that is exactly the wait overlap scheduling exists to avoid. The other view is to accept the waste: a request computes at most one extra token, so the effect on throughput is proportional to how many requests finish per step and is usually negligible.
  3. The KV pool is only written by kernels on the engine stream, so a new request's KV writes queue behind the old batch; the token-pool row and page-table row of a request slot, however, are written by the scheduler stream (the new request's prompt, the hit locations, the newly allocated pages). On a prefix hit the write goes to the new request's own row, which is fine as long as that row was not just released while an in-flight batch still uses it, which is precisely what deferring the request slot's release guarantees.

Summary

  • Overlap scheduling launches this step before handling the previous one, hiding the CPU cost behind the GPU's compute.
  • Prerequisites: the input is written on the GPU by the previous step's sample; the request state is advanced at launch time; the scheduler and the engine use two streams.
  • The price: the CPU state runs a step ahead, and a request may be scheduled one step too many after EOS.
  • This book fixes four problems in the official version: the finished check uses the number of tokens received; stale results are discarded; a request slot is reused only once the in-flight batch is done; and an aborted request joins finished_reqs.

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