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

The engine and the sampler

Engine is the compute engine on one TP rank: it initializes the model, the KV pool, the page table, the attention backend, the sampler and CUDA Graph in a fixed order, and then exposes exactly one method, forward_batch, which takes a prepared batch and returns each request's next token. This chapter finishes it and drives it by hand through prefill and decode with no scheduler at all, matching Hugging Face token by token. That completes the "computing it right" part.

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

1. How does the engine decide how many pages the KV pool gets? Why does it measure free memory before loading the model?
2. Why does forward_batch return the tokens both on the GPU and on the CPU?
3. A batch holds some greedy requests and some with temperature=0.7, top_p=0.9. How does the sampler handle them in one call?
4. When top-k and top-p are both set, which comes first?

Answers (try it yourself first, then expand)

1. Pages = (memory_ratio × free memory before loading the model − what the model takes) ÷ bytes per page, with the minimum taken across ranks under TP. The early measurement is what reveals how much the model itself takes: the difference between free memory before and after loading.
2. The GPU copy is written straight back into the token pool as the next step's input, never touching the CPU; the CPU copy (an async transfer) lets the scheduler decide whether a request is finished and hand it to detokenization.
3. The per-request parameters are laid out as tensors and handled in one batch: when everything is greedy it goes straight to argmax, and in a mixed batch the greedy requests get a tiny temperature (equivalent to argmax) and travel the sampling path with the others.
4. top-k first, then top-p.

Files you will write: engine/config.py, engine/engine.py, engine/sample.py, utils/device.py (the device abstraction).

Files so far: 46 (5 new, 0 modified, 3 brought in early)

minisgl/
├── __init__.py
├── core.py
├── attention/  (4 files, unchanged)
├── distributed/  (3 files, unchanged)
├── engine/
│   ├── __init__.py  ← new in this chapter
│   ├── config.py  ← new in this chapter
│   ├── engine.py  ← new in this chapter
│   ├── sample.py  ← new in this chapter
│   ├── graph.py  ← brought in early (explained in chapter 18)
├── kernel/  (2 files, unchanged)
├── kvcache/  (4 files, unchanged)
├── layers/  (8 files, unchanged)
├── models/  (7 files, unchanged)
├── moe/
│   ├── __init__.py  ← brought in early (explained in chapter 20)
│   ├── base.py  ← brought in early (explained in chapter 20)
├── scheduler/  (2 files, unchanged)
├── utils/
│   ├── __init__.py
│   ├── misc.py
│   ├── hf.py
│   ├── torch_utils.py
│   ├── registry.py
│   ├── logger.py
│   ├── device.py  ← new in this chapter

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

Configuration

minisgl/engine/config.py · EngineConfig

@dataclass(frozen=True)
class EngineConfig:
    model_path: str
    tp_info: DistributedInfo
    dtype: torch.dtype
    device: str = "auto"  # auto: cuda:<rank> with a GPU, cpu otherwise (upstream supports only CUDA)
    max_running_req: int = 256
    attention_backend: str = "auto"
    moe_backend: str = "auto"
    cuda_graph_bs: List[int] | None = None
    cuda_graph_max_bs: int | None = None
    page_size: int = 1
    memory_ratio: float = 0.9
    cpu_kv_cache_bytes: int = 2 << 30  # the KV pool is capped at 2 GiB on a CPU by default
    distributed_timeout: float = 60.0
    distributed_port: int = 2333
    use_dummy_weight: bool = False
    max_seq_len_override: int | None = None
    num_page_override: int | None = None

    @cached_property
    def hf_config(self):
        return cached_load_hf_config(self.model_path)

    @cached_property
    def model_config(self) -> ModelConfig:
        from minisgl.models import ModelConfig

        return ModelConfig.from_hf(self.hf_config)

    @property
    def max_seq_len(self) -> int:
        if self.max_seq_len_override is not None:
            return self.max_seq_len_override
        return self.model_config.rotary_config.max_position

    @property
    def max_forward_len(self) -> int:
        return self.max_seq_len

    @property
    def distributed_addr(self) -> str:
        return f"tcp://127.0.0.1:{self.distributed_port}"

Most fields map to command-line arguments. device and cpu_kv_cache_bytes are ours; model_config is a cached_property parsed from config.json on first access (chapter 3).

The order of initialization

minisgl/engine/engine.py · Engine.__init__

def __init__(self, config: EngineConfig):
    set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size)
    self.config = config = _adjust_config(config)
    self.device = _pick_device(config)
    if is_cuda(self.device):
        torch.cuda.set_device(self.device)
    else:
        _set_cpu_threads(config.tp_info.size)
    torch.manual_seed(42)
    self.stream = create_stream(self.device)
    set_stream(self.stream)
    self.dtype = config.dtype
    self.ctx = Context(config.page_size)
    set_global_ctx(self.ctx)

    self.tp_cpu_group = self._init_communication(config)
    init_free_memory = self._sync_get_memory()[1]
    logger.info_rank0(f"Free memory before loading model: {_gib(init_free_memory)}")

    # ======================= the model: an empty shell on the meta device, then the weights swapped in
    set_rope_device(self.device)
    with torch.device("meta"), torch_dtype(config.dtype):
        self.model = create_model(config.model_config)
    self.model.load_state_dict(self._load_weight_state_dict(config))

    # ======================= the KV pool: it gets all the memory left
    self.num_pages = self._determine_num_pages(init_free_memory, config)
    num_tokens = self.num_pages * config.page_size
    self.ctx.kv_cache = self.kv_cache = create_kvcache_pool(
        model_config=config.model_config,
        num_pages=self.num_pages + 1,  # the extra page belongs to the dummy request, whose writes land here
        page_size=config.page_size,
        device=self.device,
        dtype=self.dtype,
    )

    # ======================= the page table: one row per request, holding each token's KV location
    self.max_seq_len = min(config.max_seq_len, num_tokens)
    aligned_max_seq_len = _align_up_32(self.max_seq_len)
    self.ctx.page_table = self.page_table = torch.zeros(
        (config.max_running_req + 1, aligned_max_seq_len),  # the last row belongs to the dummy request
        dtype=torch.int32, device=self.device,
    )

    # ======================= the attention backend, the MoE backend, the sampler
    self.ctx.attn_backend = self.attn_backend = create_attention_backend(
        config.attention_backend, config.model_config)
    if config.model_config.is_moe:
        self.ctx.moe_backend = self.moe_backend = create_moe_backend(config.moe_backend)
    self.sampler = Sampler(self.device, config.model_config.vocab_size)

    # ======================= the dummy request and CUDA Graph
    self.dummy_req = Req(
        input_ids=torch.tensor([0], dtype=torch.int32),
        table_idx=config.max_running_req,
        cached_len=0,
        output_len=1,
        uid=-1,
        sampling_params=None,  # type: ignore[arg-type]
        cache_handle=None,  # type: ignore[arg-type]
    )
    self.page_table[self.dummy_req.table_idx].fill_(num_tokens)  # points at the dummy page
    self.graph_runner = GraphRunner(
        stream=self.stream,
        device=self.device,
        model=self.model,
        attn_backend=self.attn_backend,
        cuda_graph_bs=config.cuda_graph_bs,
        cuda_graph_max_bs=config.cuda_graph_max_bs,
        free_memory=init_free_memory,
        max_seq_len=aligned_max_seq_len,
        vocab_size=config.model_config.vocab_size,
        dummy_req=self.dummy_req,
    )

The order is deliberate:

1. TP info and the device. Every layer built afterwards reads the TP info to decide how it shards.
2. The stream and the global Context. On CUDA all of the engine's compute runs on its own stream (overlap scheduling in chapter 11 needs one stream for the scheduler and one for the engine).
3. Communication. With TP>1 the process group is created (chapter 16).
4. Measure free memory, then load the model. The difference before and after is what the model takes.
5. The KV pool. The remaining memory gives the page count (next section), plus one extra page for the dummy request.
6. The page table, whose last row belongs to the dummy request and points at the dummy page (fill_(num_tokens), the dummy page's first location).
7. The attention backend and the MoE backend. The attention backend reads the KV pool as it is constructed (get_global_ctx().kv_cache), so it has to come after the pool.
8. The sampler and CUDA Graph. Capturing a graph really runs the model, so it has to be last.

Planning the memory

minisgl/engine/engine.py · Engine._determine_num_pages

def _determine_num_pages(self, old_free_memory: int, config: EngineConfig) -> int:
    cache_per_page = (
        2  # K and V
        * config.model_config.head_dim
        * div_even(config.model_config.num_kv_heads, config.tp_info.size, allow_replicate=True)
        * config.page_size
        * self.dtype.itemsize
        * config.model_config.num_layers
    )
    num_pages = config.num_page_override
    if num_pages is None:
        new_free_memory = self._sync_get_memory()[1]
        model_memory = old_free_memory - new_free_memory
        available = int(config.memory_ratio * old_free_memory) - model_memory
        if not is_cuda(self.device):
            available = min(available, config.cpu_kv_cache_bytes)
        num_pages = available // cache_per_page
    assert num_pages > 1, "Not enough memory for KV cache"
    logger.info_rank0(f"KV cache: {num_pages * config.page_size} tokens, "
                      f"{_gib(num_pages * cache_per_page)}")
    return num_pages

Bytes per page is 2 (K and V) × head_dim × this rank's KV heads × page_size × dtype bytes × layers. Usable memory works out as:

usable = memory_ratio × free memory before loading the model − what the model takes
what the model takes = free before loading − free after loading

memory_ratio defaults to 0.9, leaving 10% for activations, CUDA Graph and various workspaces. Under tensor parallelism the minimum across ranks is taken (_sync_get_memory does one all-reduce) so that every rank's KV pool is the same size, because the scheduler makes identical decisions on every rank only if the resources are identical.

On a CPU "free memory" is the whole machine's, and working proportionally would give a KV pool of tens of GB, so we add a cap, cpu_kv_cache_bytes (2 GiB by default); tests usually set the page count outright with num_page_override.

forward_batch

minisgl/engine/engine.py · Engine.forward_batch

def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
    with self.ctx.forward_batch(batch):
        if self.graph_runner.can_use_cuda_graph(batch):
            logits = self.graph_runner.replay(batch)
        else:
            logits = self.model.forward()
    for req in batch.reqs:
        req.complete_one()
    next_tokens_gpu = self.sampler.sample(logits[: batch.size], args).to(torch.int32)
    next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True)
    copy_done_event = create_event(self.device)
    copy_done_event.record(self.stream)
    return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)

- run the model inside the forward_batch context (or replay the CUDA Graph) to get [batch_size, vocab] logits; - advance every request's state with complete_one(). Note that this happens before sampling and completes immediately on the CPU, so from here on the scheduler can treat those requests as having one more token and schedule the next step without waiting for the GPU (chapter 11); - sample to get next_tokens_gpu, then copy it to the CPU non_blocking and record an event. A caller that needs the CPU values first calls copy_done_event.synchronize().

The two copies serve different purposes: the GPU one is written back into the token pool by the scheduler as the next step's input and never leaves the GPU; the CPU one decides whether a request is finished and goes to the detokenizer.

The device abstraction

Upstream uses torch.cuda.Stream(), torch.cuda.Event() and pin_memory=True directly here. We replace them with a few small functions:

minisgl/utils/device.py · create_stream

def create_stream(device: torch.device) -> Any:
    return torch.cuda.Stream(device=device) if is_cuda(device) else NullStream()

minisgl/utils/device.py · NullEvent

class NullEvent:
    """An "event" on a CPU: recording and waiting are no-ops (the work is done by the time record runs)."""

    def record(self, stream: Any = None) -> None:
        pass

    def synchronize(self) -> None:
        pass

    def query(self) -> bool:
        return True

On a CPU every computation finishes synchronously: by the time record runs the result is long since computed, so synchronize has nothing to wait for. Pinned memory only means something on CUDA, where it lets host-to-device copies run asynchronously, so pin(device) returns False on a CPU.

The sampler

minisgl/engine/sample.py · Sampler

@dataclass
class Sampler:
    device: torch.device
    vocab_size: int

    def prepare(self, batch: Batch) -> BatchSamplingArgs:
        """Prepare the whole batch's sampling-parameter tensors on the scheduler side (the CPU)."""
        params = [r.sampling_params for r in batch.reqs]
        if all(p.is_greedy for p in params):
            return BatchSamplingArgs(temperatures=None)
        MIN_T = MIN_P = 1e-6
        # the greedy requests in a mixed batch: a tiny temperature collapses the softmax to one-hot, which equals argmax
        ts = [max(0.0 if p.is_greedy else p.temperature, MIN_T) for p in params]
        top_ks = [p.top_k if p.top_k >= 1 else self.vocab_size for p in params]
        top_ps = [min(max(p.top_p, MIN_P), 1.0) for p in params]
        top_k = top_p = None
        if any(k != self.vocab_size for k in top_ks):
            top_k = make_device_tensor(top_ks, torch.int32, self.device)
        if any(p < 1.0 for p in top_ps):
            top_p = make_device_tensor(top_ps, torch.float32, self.device)
        return BatchSamplingArgs(make_device_tensor(ts, torch.float32, self.device), top_k, top_p)

    def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
        if args.temperatures is None:
            return torch.argmax(logits, dim=-1)
        impl = sample_torch
        if logits.is_cuda:
            try:
                import flashinfer.sampling  # noqa: F401

                impl = sample_flashinfer
            except ImportError:
                pass
        return impl(logits.float(), args.temperatures, args.top_k, args.top_p)

Sampling takes two steps:

• prepare(batch) is called on the scheduler side, before the forward pass, and lays each request's sampling parameters out as tensors copied asynchronously to the device. When the whole batch is greedy it returns temperatures=None and sample goes straight to argmax, skipping the softmax.
• sample(logits, args) is called after the forward pass.

Requests in one batch may each have their own parameters. The greedy requests in a mixed batch are given a tiny temperature (1e-6), which after the softmax collapses to one-hot, so the sample is the argmax and the whole batch travels one code path. A top-k of -1 (no limit) becomes the vocabulary size, and the tensors for top-k (or top-p) are only built when at least one request in the batch actually sets them.

The PyTorch reference implementation:

minisgl/engine/sample.py · sample_torch

def sample_torch(logits: torch.Tensor, temperatures: torch.Tensor, top_k: torch.Tensor | None,
                 top_p: torch.Tensor | None) -> torch.Tensor:
    """Each request may have its own temperature, top-k and top-p. Like FlashInfer, top-k comes first and top-p second."""
    probs = torch.softmax(logits / temperatures.unsqueeze(1), dim=-1)
    if top_k is not None or top_p is not None:
        sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True)
        if top_k is not None:
            rank = torch.arange(probs.shape[1], device=probs.device).unsqueeze(0)
            sorted_probs = sorted_probs.masked_fill(rank >= top_k.unsqueeze(1), 0.0)
        if top_p is not None:
            sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
            before = sorted_probs.cumsum(dim=-1) - sorted_probs  # the sum of the probabilities ahead of it
            sorted_probs = sorted_probs.masked_fill(before >= top_p.unsqueeze(1), 0.0)
        probs = torch.zeros_like(probs).scatter_(1, sorted_idx, sorted_probs)
    return torch.multinomial(probs, num_samples=1).squeeze(1)

Like FlashInfer's default, top-k comes first and top-p applies to what is left. top-p keeps the tokens whose preceding probabilities sum to less than p, so the most likely token is always kept. On a GPU it switches to FlashInfer's sampling kernel (sample_flashinfer), which uses rejection sampling to avoid a sort and gives the same distribution.

Driving the engine by hand

With no scheduler we play the scheduler ourselves: "allocate" a stretch of KV locations for each request by writing the page table directly, fill in the batch's four fields each step, call forward_batch, and append the new token to the request.

examples/ch06_engine.py

"""Chapter 6: no scheduler, driving the Engine through prefill and decode by hand against HF generate."""

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))
print(f"device {engine.device}, attention backend {engine.config.attention_backend}, KV pool {engine.num_pages} pages, "
      f"page table {tuple(engine.page_table.shape)}")

tok = AutoTokenizer.from_pretrained(path)
prompts = ["The capital of France is", "def fibonacci(n):"]
reqs = []
for row, p in enumerate(prompts):
    ids = tok(p).input_ids
    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(ids, dtype=torch.int32), table_idx=row, cached_len=0,
                    output_len=8, uid=row, sampling_params=SamplingParams(max_tokens=8),
                    cache_handle=None))

outs = [[] for _ in reqs]
phase = "prefill"
for step in range(8):
    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))  # forward and sample, advancing the Req state
    for i, r in enumerate(reqs):
        r.append_host(out.next_tokens_cpu[i:i + 1])
        outs[i].append(int(out.next_tokens_cpu[i]))
    if step == 0:
        print("prefill input tokens:", len(batch.input_ids), " then per decode step:", len(reqs))
    phase = "decode"

hf = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32)
for p, o in zip(prompts, outs):
    ids = tok(p, return_tensors="pt").input_ids
    ref = hf.generate(ids, max_new_tokens=8, do_sample=False)[0, ids.shape[1]:].tolist()
    print(f"{p!r}: {tok.decode(o)!r}  matches HF: {o == ref}")
engine.shutdown()

What it prints

device cpu, attention backend torch, KV pool 1024 pages, page table (5, 256)
prefill input tokens: 9  then per decode step: 2
'The capital of France is': ' Paris. The capital of Italy is Rome'  matches HF: True
'def fibonacci(n):': '\n    if n == 0:\n'  matches HF: True

The first step is prefill: two requests' 9 tokens are packed into one dimension and fed to the model; after that each step decodes 1 token per request. All 8 tokens match Hugging Face's generate exactly. What the scheduler of the next part has to do is automate the "allocate KV locations", "fill the batch fields" and "append the token" of this hand-written code, and let requests join and leave at any time.

The official implementation

• initialization: engine/engine.py · Engine.__init__ (lines 30-110)
• memory planning: engine/engine.py · Engine._determine_num_pages (lines 148-168)
• the forward pass: engine/engine.py · Engine.forward_batch (lines 191-206)
• sampling: engine/sample.py · Sampler (lines 48-75) (FlashInfer only)

Differences from upstream

• Device: upstream supports only CUDA; we support CPU as well through utils/device.py.
• Communication: upstream initializes a gloo process group even at TP=1; we only do so when TP>1, which lets tests create engines repeatedly in one process (shutdown clears the global state).
• CPU thread count: on a CPU we set PyTorch's thread count to the physical core count divided by the TP size, or decode's small matmuls get dragged down tens of times by thread synchronization (see the architecture chapter).

Tests

tests/test_ch06_engine.py has three tests: the prefill logits differ from HF by less than 1e-4 at the maximum; the hand-written loop's greedy output matches HF token by token; and the sampler's top-k, top-p and mixed-greedy behaviour and probabilities (with top-k=2 the top token is drawn about e⁵ / (e⁵ + e⁴) ≈ 0.731 of the time).

tests/test_ch06_engine.py · test_sampler_greedy_and_filters

def test_sampler_greedy_and_filters():
    sampler = Sampler(torch.device("cpu"), vocab_size=6)
    assert sampler.prepare(_batch([SamplingParams(), SamplingParams()])).temperatures is None
    logits = torch.tensor([[5.0, 4.0, 3.0, 0.0, 0.0, 0.0]] * 3)
    args = sampler.prepare(_batch([SamplingParams(temperature=1.0, top_k=2),
                                   SamplingParams(temperature=1.0, top_p=0.5),
                                   SamplingParams()]))  # the third request is greedy
    torch.manual_seed(0)
    draws = torch.stack([sampler.sample(logits, args) for _ in range(2000)])
    assert set(draws[:, 0].tolist()) == {0, 1}  # top-k=2: only the top two are possible
    assert set(draws[:, 1].tolist()) == {0}  # the top probability is about 0.67 > 0.5, so top-p keeps only it
    assert set(draws[:, 2].tolist()) == {0}  # the greedy requests in a mixed batch
    p0 = (draws[:, 0] == 0).float().mean().item()  # renormalized after top-k: e^5/(e^5+e^4) = 0.731
    assert abs(p0 - 0.731) < 0.03

How to explain it

On the engine: the initialization order is TP info, then the stream and the context, then communication, then loading the model, then the KV pool, the page table, the attention backend, the sampler and CUDA Graph. The KV pool's size is the memory ratio times the free memory before loading, minus what the model takes, converted into pages, with the minimum across ranks under TP (and the measurement has to happen before loading to know what the model takes). forward_batch runs the forward pass, advances the request state and samples, returning the tokens on the GPU (written straight back into the token pool) and an asynchronous copy on the CPU (for the scheduler to decide what has finished). The sampler handles requests with different parameters in one call: argmax when everything is greedy, a tiny temperature for the greedy ones in a mixed batch, and top-k before top-p.

Exercises

1. Running Qwen3-0.6B (bf16) on a 24 GB card, with the model taking about 1.2 GB and memory_ratio=0.9, roughly how many tokens does the KV pool hold?
2. Add min_p to the sampler: keep the tokens whose probability is at least "the maximum probability × min_p". Should it come before or after top-k and top-p?
3. Does it affect correctness if forward_batch calls complete_one() after sampling instead? What about overlap scheduling in chapter 11?

Answers

1. About 0.9 × 24 − 1.2 ≈ 20.4 GB usable, at 112 KiB per token, so roughly 180,000 tokens (CUDA Graph and friends come off the top in practice).
2. min_p's threshold depends on the maximum probability and usually applies to the full distribution after temperature scaling: compute its mask before top-k and top-p, then intersect it with theirs.
3. Correctness is unaffected, since complete_one only changes counters on the CPU. Overlap scheduling is not fundamentally affected either, because the scheduler assembles the next step only after forward_batch returns; what matters is that it does not depend on the sampled values, so the CPU never waits for the GPU.

Summary

• Engine initializes in this order: TP info, stream and context, communication, model loading, KV pool, page table, backends, sampler, CUDA Graph.
• The KV pool's size is memory_ratio × free memory before loading − what the model takes, converted to pages by the bytes per page, with the minimum taken across ranks under TP.
• forward_batch runs the forward pass, advances the request state and samples, returning the tokens on the GPU and an in-flight asynchronous copy on the CPU.
• The sampler supports per-request parameters: argmax when everything is greedy, a tiny temperature for greedy requests in a mixed batch, and top-k before top-p.

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.

• 190. 批量采样：每个请求一组参数 (medium)

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