mini-sglang from Scratch Stage 1 · Computing it right: from skeleton to real structure

Stage 1 milestone: the same output, the real structure¶

The six chapters of stage 1 replace every naive part of step 0's 126 lines with a real module: Req / Batch / Context, an op layer that does not depend on nn.Module, a model with streaming weight loading, a paged KV pool, pluggable attention backends, and an Engine with a sampler. This page runs the end state first — the same prompt, the same 16 tokens, plus one thing the skeleton could not do: four requests in one batch, computed in one step. Come back to these numbers after the six chapters.

What you hold at the end of this stage

An Engine: give it a Batch with its fields filled in and it returns each request's next token;
No scheduler yet — you fill the batch's fields yourself (exactly the lines stage 2 automates);
A single request is as fast as step 0; several requests computed together give three times step 0's throughput;
Every step's output still matches Hugging Face token by token.

Run it first¶

cd minisgl && python examples/stage1_milestone.py

examples/stage1_milestone.py

"""Stage 1 milestone: the same prompt and the same 16 tokens, but the skeleton has become the real structure.

Two comparisons with step 0: tokens/s for a single request (where the real structure gains nothing),
and what the skeleton could not do at all: four requests in one batch, computed in one step.
"""

import time

import torch
from minisgl.core import Batch, Req, SamplingParams
from minisgl.distributed import DistributedInfo
from minisgl.engine import Engine, EngineConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

path = "models/Qwen3-0.6B"
engine = Engine(EngineConfig(model_path=path, tp_info=DistributedInfo(0, 1), dtype=torch.float32,
                             max_running_req=4, num_page_override=1024, max_seq_len_override=256))
tok = AutoTokenizer.from_pretrained(path)
print(f"the real structure: {engine.device}, attention backend {engine.config.attention_backend}, "
      f"KV pool {engine.num_pages} pages, {torch.get_num_threads()} threads (set by the Engine itself)")


def make_reqs(prompts, n):
    reqs = []
    for row, p in enumerate(prompts):
        engine.page_table[row, :256] = torch.arange(row * 256, row * 256 + 256)   # "allocate" the KV locations by hand
        reqs.append(Req(input_ids=torch.tensor(tok(p).input_ids, dtype=torch.int32), table_idx=row,
                        cached_len=0, output_len=n, uid=row, sampling_params=SamplingParams(max_tokens=n),
                        cache_handle=None))
    return reqs


def step(reqs, phase):
    """No scheduler: fill the batch's fields yourself, forward + sample, then write the new token back to each request."""
    batch = Batch(reqs=reqs, phase=phase)
    engine.graph_runner.pad_batch(batch)
    batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).int()
    batch.input_ids = torch.cat([r.input_ids[r.cached_len:r.device_len] for r in reqs])
    batch.out_loc = torch.cat([engine.page_table[r.table_idx, r.cached_len:r.device_len] for r in reqs])
    engine.attn_backend.prepare_metadata(batch)
    out = engine.forward_batch(batch, engine.sampler.prepare(batch))
    for i, r in enumerate(reqs):
        r.append_host(out.next_tokens_cpu[i:i + 1])
    return out.next_tokens_cpu.tolist()


def generate(prompts, n=16):
    reqs = make_reqs(prompts, n)
    outs = [[] for _ in reqs]
    t0 = time.perf_counter()
    for i, t in enumerate(step(reqs, "prefill")):
        outs[i].append(t)
    t1 = time.perf_counter()
    for _ in range(n - 1):
        for i, t in enumerate(step(reqs, "decode")):
            outs[i].append(t)
    t2 = time.perf_counter()
    return outs, t1 - t0, (t2 - t1) / (n - 1)


# ------------------------------------------------------------------ 1. the same request as step 0
prompt = "The capital of France is"
[out], t_prefill, t_step = generate([prompt])
print(f"single request: {tok.decode(out)!r}")
print(f"  prefill {t_prefill * 1000:.0f} ms, decode {t_step * 1000:.0f} ms per step ({1 / t_step:.1f} tokens/s)")

# ------------------------------------------------------------------ 2. what the skeleton could not do: 4 requests in one batch
prompts = [prompt, "List three prime numbers:", "def fibonacci(n):", "Once upon a time,"]
outs, t_prefill, t_step4 = generate(prompts)
print(f"4 requests in one batch: prefill {t_prefill * 1000:.0f} ms, decode {t_step4 * 1000:.0f} ms per step"
      f" (4 tokens per step, {4 / t_step4:.1f} tokens/s in total; a step is only {(t_step4 / t_step - 1) * 100:.0f}% slower than a single request)")
written = sum(len(tok(p).input_ids) + 15 for p in prompts)               # per request: the prompt + the 15 new tokens fed back in
print(f"KV pool: {written} / {engine.num_pages} pages written; each token writes only its own page, no copies at all")

hf = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32).eval()
for p, o in zip(prompts, outs):
    ids = tok(p, return_tensors="pt").input_ids
    with torch.no_grad():
        ref = hf.generate(ids, max_new_tokens=16, do_sample=False)[0, ids.shape[1]:].tolist()
    print(f"  {p!r:30} -> {tok.decode(o)[:32]!r:36} matches HF: {o == ref}")
print(f"run alone or inside the 4-request batch, the first request's 16 tokens are identical: {outs[0] == out}")
engine.shutdown()

What it prints

the real structure: cpu, attention backend torch, KV pool 1024 pages, 16 threads (set by the Engine itself)
single request: ' Paris. The capital of Italy is Rome. The capital of Spain is Madrid.'
  prefill 56 ms, decode 36 ms per step (28.1 tokens/s)
4 requests in one batch: prefill 91 ms, decode 58 ms per step (4 tokens per step, 68.9 tokens/s in total; a step is only 63% slower than a single request)
KV pool: 79 / 1024 pages written; each token writes only its own page, no copies at all
  'The capital of France is'     -> ' Paris. The capital of Italy is '   matches HF: True
  'List three prime numbers:'    -> ' 2, 3, 5, 7, 11,'                   matches HF: True
  'def fibonacci(n):'            -> '\n    if n == 0:\n        return 0' matches HF: True
  'Once upon a time,'            -> ' there were 3000 people in a tow'   matches HF: True
run alone or inside the 4-request batch, the first request's 16 tokens are identical: True

Reading the numbers¶

A single request: 28.1 tokens/s, versus 26.8 in step 0. The real structure gains nothing on a single request — the difference is run-to-run noise. No surprise: paged KV writes are scattered by position, the attention backend has to slice out each request's own KV, and the Req / Batch bookkeeping is all Python. What those costs buy is invisible in the single-request number; they exist for the next line.

Four requests in one batch: 58 ms per step, only 63% slower than the 36 ms of a single request, 68.9 tokens/s in total. This is what step 0's skeleton could not do. Decode is memory bound: one step reads the 2.4 GB of fp32 weights from memory once, and whether the matrix multiply's "batch" has 1 row or 4 makes almost no difference to that read. So gathering requests into a batch adds only the small extra compute per step, while the weights are read once. That is the whole argument for continuous batching, and where an inference system's throughput comes from.

79 KV pages written, no copies at all. Step 0 copied the entire KV with torch.cat on every step; now each token writes only its own page, and the page table records where it is. Run alone or inside the batch, the first request's 16 tokens are identical — batching changes no request's result, something every later chapter's tests confirm again and again.

What replaced each part of the skeleton¶

Step 0	Stage 1	Why
one `ids` tensor, one `pos` tensor	`Req`, `Batch`, `Context`	requests have their own lengths and state; the model reads the current batch from the global `Context`, so `forward()` takes no arguments
four functions plus a weight `dict`	The op layer: `BaseOP` and the `Layer`s	each `Layer` owns its weights' shapes, loading and sharding — the only way tensor parallelism can split them later
one `safe_open` reading everything	Models and weight loading	streaming loading on demand; one decoder file covers Qwen3, Qwen2.5 and Llama 3
`list` + `torch.cat`	The KV pool, the page table and the token pool	pre-allocated, written by page, zero copies; many requests share one pool
`einsum` + an explicit mask	Attention backends and the reference implementation	metadata for variable-length batches, paged KV and prefix reuse; FlashInfer / FlashAttention on a GPU
`argmax`	The Engine and the sampler	temperature, top-k and top-p per request, sampled in one batched call; `Engine` initialises everything above in the right order

How to read these six chapters¶

In order: read the self-test, write your own version, compare with this book's and the official implementation, then run the chapter's tests. After The Engine and the sampler, come back and rerun this page's script: the few lines in step() that "fill the batch fields by hand" each correspond to one duty of stage 2's scheduler — where positions and out_loc come from (page allocation), why cached_len changes (prefix caching), who decides which requests enter this step's batch (admission control). Stage 2 is those lines, automated.

Checkpoint: do these before stage 2

Without looking, draw Req's three length fields before and after prefill;
Explain what one row of the page table is and how out_loc is computed;
Rerun this page's script with page_size=4 and get the same output (chapter 4's exercise);
Say why four requests per step cost far less than four times one.

Summary¶

Stage 1 replaces the skeleton's six naive parts with six real modules whose interfaces match the official ones.
A single request is no faster than the skeleton; several requests in one batch are the point of the real structure — decode reads the weights once, however many rows it computes.
KV is written by page with zero copies; batching changes no request's result.
The lines filled by hand in step() are what stage 2's scheduler does.

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