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

Stage 2 milestone: the scheduler takes over

At the end of stage 1 you still fill the batch's fields yourself: which requests enter this step, which pages hold their KV, where their positions start. The five chapters of stage 2 turn those lines into a scheduler — continuous batching, admission control, page allocation and release, the radix cache, chunked prefill, overlap scheduling — leaving a single offline interface, LLM.generate. This page runs the end state first: three sets of numbers, each for something the skeleton could not do.

What you hold at the end of this stage

  • LLM.generate(prompts): any number of requests at once, of any lengths; the ones that finish leave the batch, new ones join at any step;
  • A shared system prompt has its KV computed once and reused by later requests;
  • A long prompt is cut into fixed-size chunks that other requests can slip between;
  • The CPU's preparation of the next step hides behind the computation (overlap scheduling; this page turns it off so the numbers reflect the scheduler alone).

Run it first

cd minisgl && python examples/stage2_milestone.py

examples/stage2_milestone.py
"""Stage 2 milestone: the scheduler takes over the lines that stage 1 filled by hand.

Three things: the offline interface swallows 8 requests at once (continuous batching, admission,
page allocation all automatic); 4 questions on one system prompt, where the radix cache cuts prefill
by three quarters; a long prompt cut into several steps by chunked prefill.
"""

import time

import torch
from minisgl.core import SamplingParams
from minisgl.env import ENV
from minisgl.llm import LLM

ENV.DISABLE_OVERLAP_SCHEDULING.value = True                   # overlap scheduling off for now, so the numbers reflect the scheduler alone


def make(**kw):
    return LLM("models/Qwen3-0.6B", dtype=torch.float32, max_running_req=8, num_page_override=2048,
               max_seq_len_override=512, **kw)


def count_prefill(llm):
    """Wrap the engine's forward to count how many tokens prefill actually computed, and in how many steps"""
    computed = []
    orig = llm._forward
    llm._forward = lambda fi: (computed.append(sum(r.extend_len for r in fi.batch.reqs))
                               if fi.batch.is_prefill else None, orig(fi))[1]
    return computed


# ------------------------------------------------------------------ 1. 8 requests, one generate
prompts = ["The capital of France is", "List three prime numbers:", "def fibonacci(n):", "Once upon a time,",
           "1 + 1 =", "The quick brown fox", "import torch\n", "Roses are red,"]
llm = make()
t = time.perf_counter()
outs = llm.generate(prompts, SamplingParams(max_tokens=16, ignore_eos=True))
dt = time.perf_counter() - t
print(f"8 requests x 16 tokens: {dt:.2f} s, {8 * 16 / dt:.1f} tokens/s in total (stage 1's hand-assembled batch did 68.9, step 0 did 26.8)")
print(f"  the first one: {outs[0]['text']!r}")
llm.shutdown()

# ------------------------------------------------------------------ 2. one system prompt: the radix cache
system = "You are a helpful assistant. Answer concisely and accurately. " * 4
questions = ["What is 2 + 2?", "Name a primary color.", "What is the capital of Japan?", "Spell cat."]
for cache_type in ("naive", "radix"):
    llm = make(cache_type=cache_type)
    computed = count_prefill(llm)
    llm.generate([system + questions[0]], SamplingParams(max_tokens=4))          # the first request goes in first
    t = time.perf_counter()
    llm.generate([system + q for q in questions[1:]], SamplingParams(max_tokens=4))
    dt = time.perf_counter() - t
    total = sum(len(llm.tokenizer(system + q).input_ids) for q in questions)
    print(f"{cache_type:<5}: 4 prompts of {total} tokens in total, prefill actually computed {sum(computed)}, the last 3 requests took {dt * 1000:.0f} ms")
    llm.shutdown()

# ------------------------------------------------------------------ 3. a long prompt: chunked prefill
long_prompt = "Explain, step by step, why the sky is blue. " * 25
llm = make(max_extend_tokens=96)                              # at most 96 prefill tokens per step
computed = count_prefill(llm)
llm.generate([long_prompt], SamplingParams(max_tokens=2))
print(f"chunked prefill: a {len(llm.tokenizer(long_prompt).input_ids)}-token prompt took {len(computed)} steps, "
      f"{computed} per step; other requests can slip in between them")
llm.shutdown()
What it prints
8 requests x 16 tokens: 1.29 s, 99.5 tokens/s in total (stage 1's hand-assembled batch did 68.9, step 0 did 26.8)
  the first one: ' Paris. The capital of Italy is Rome. The capital of Spain is Madrid.'
naive: 4 prompts of 231 tokens in total, prefill actually computed 231, the last 3 requests took 345 ms
radix: 4 prompts of 231 tokens in total, prefill actually computed 73, the last 3 requests took 190 ms
chunked prefill: a 302-token prompt took 4 steps, [96, 96, 96, 14] per step; other requests can slip in between them

Reading the numbers

8 requests in 1.29 seconds, 99.5 tokens/s. Step 0 did 26.8; stage 1's hand-assembled batch of 4 did 68.9. Two things stack: the batch is larger (8 requests per step, one read of the weights yields 8 tokens), and nobody fills fields by hand any more — at every step inside generate, the scheduler does what the few lines in stage 1's step() did: admission, page allocation, computing positions and out_loc, reclaiming the pages of finished requests. In a real service requests arrive one after another with different lengths, and the point of continuous batching is that the batch's membership keeps changing while the computation never stops.

A shared system prompt: prefill drops from 231 tokens to 73, and the last three requests from 345 ms to 190 ms. The four questions share a 158-token system prompt. With the naive cache every request recomputes it; with the radix cache, once the first request has computed the prefix it stays in the tree, and the next three compute only their own question. That is where the severalfold throughput difference in multi-turn and shared-system-prompt workloads comes from — not one of the 158 × 3 skipped tokens was wasted.

A 302-token prompt in 4 steps: 96, 96, 96, 14. Without chunked prefill a long prompt takes over the engine for a whole step and every decoding request waits for it; with chunks, each step computes at most max_extend_tokens tokens, decode requests slip in between, and a slightly longer time to first token buys steady per-token latency for everyone else.

What the hand-written lines became

What you did by hand in stage 1's step() Stage 2 Where
decide which requests enter this step's batch admission control, the continuous-batching main loop The scheduler skeleton and the offline interface
page_table[row, :256] = arange(...) — "allocating" KV by hand CacheManager: allocate on demand, reclaim on finish, wait when full CacheManager
cached_len=0 — every request from scratch the radix cache: match prefixes, reuse KV, reference counts, LRU eviction Radix Cache
the whole prompt in one prefill step chunked prefill: at most max_extend_tokens tokens per step Chunked prefill
prepare the next step after this one finishes overlap scheduling: prepare the next step while the GPU is still on this one Overlap scheduling

How to read these five chapters

Read the scheduler skeleton first — it defines the main loop, and the other four chapters add to that loop. At Radix Cache, flip this page's cache_type back and forth and watch the "prefill actually computed" number; at chunked prefill, change max_extend_tokens and watch the step count and tokens per step. After the last chapter, overlap scheduling, delete this page's first line ENV.DISABLE_OVERLAP_SCHEDULING.value = True and rerun: on a CPU the gain is small (there is no GPU computing at the other end), but on a real card it is the last stretch of throughput.

Checkpoint: do these before stage 3

  • Draw a request's state changes from entering the queue to releasing its KV pages, marking each scheduler decision;
  • Explain why the radix tree splits nodes and what reference counting protects against;
  • Say what chunked prefill gives up and what it buys;
  • Set this page's max_running_req to 2 and explain why throughput drops, and to what.

Summary

  • Stage 2 turns the hand-filled batch lines of stage 1 into a scheduler, leaving only LLM.generate.
  • Throughput goes from 68.9 to 99.5 tokens/s: a larger batch whose membership changes while computation never stops.
  • The radix cache computes a shared prefix once: 231 prefill tokens become 73.
  • Chunked prefill stops long prompts from monopolising the engine: 302 tokens in 4 steps, with other requests in between.

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