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

Core data structures: Req, Batch, Context

Only a few things flow through the whole system: a request Req, a batch of them Batch, and the global Context that holds "which batch is being computed right now". They all live in one 136-line file, yet they decide how every later module is written. This chapter writes them, with the focus on the length fields on Req, because the scheduler's entire bookkeeping is built on them.

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

  1. A request has a prompt of 6 tokens, its first 2 tokens hit the prefix cache, and max_tokens=3. How many tokens does the prefill step compute? What is device_len after prefill?
  2. Why is Req declared @dataclass(eq=False)?
  3. What is the difference between a Batch's reqs and its padded_reqs?
  4. How does each model layer get hold of the current batch's positions?
Answers (try it yourself first, then expand)
  1. Four (extend_len = device_len − cached_len = 6 − 2). device_len is 6 for that step; after the forward pass complete_one() moves cached_len to 6 and device_len to 7, so the next decode step computes position 7.
  2. So that Req compares and hashes by identity. The scheduler puts requests into a set, and two different requests can have every field equal (two identical prompts arriving at once), which field-wise comparison would collapse into one.
  3. reqs are the real requests; padded_reqs also includes the dummy requests used to pad the batch size when CUDA Graph is on.
  4. From the global Context: the scheduler fills in tensors like positions while preparing the batch, and each layer reads what it needs.

Files you will write: minisgl/core.py.

Files so far: 2 (2 new, 0 modified, 0 brought in early)
minisgl/
├── __init__.py  ← new in this chapter
├── core.py  ← new in this chapter
This step's main: examples/ch01_req.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

SamplingParams

minisgl/core.py · SamplingParams
@dataclass
class SamplingParams:
    temperature: float = 0.0
    top_k: int = -1
    top_p: float = 1.0
    ignore_eos: bool = False
    max_tokens: int = 1024

    @property
    def is_greedy(self) -> bool:
        return (self.temperature <= 0.0 or self.top_k == 1) and self.top_p == 1.0
is_greedy lets the sampler take a fast path: when the whole batch is greedy it goes straight to argmax without a softmax. Note that top_k == 1 counts as greedy but top_p < 1 does not, because even at temperature 0 it still goes through the top-p filtering (the result is still deterministic, just a different code path).

Req: the three lengths of a request

A request's state inside the engine is fully described by three lengths:

at admission p0 p1 p2 p3 p4 p5 t0 t1 t2 cached_len = 2 extend_len = 4 device_len = 6 after prefill p0 p1 p2 p3 p4 p5 t0 t1 t2 cached_len = 6 extend_len = 1 device_len = 7 after decode 1 p0 p1 p2 p3 p4 p5 t0 t1 t2 cached_len = 7 extend_len = 1 device_len = 8 max_device_len = 9 (prompt 6 + max_tokens 3) already cached to compute now
the same request's three lengths at admission, after prefill, and after one decode step
  • cached_len: how many leading tokens already have their KV in the cache (computed earlier, or hit in the prefix cache);
  • device_len: how many tokens will have KV in the cache after this forward pass;
  • max_device_len = prompt length + max_tokens: the longest it can grow.

Everything the scheduler needs follows from these:

  • extend_len = device_len - cached_len: how many tokens go into the model this step. During prefill it is the part of the prompt that missed the cache; during decode it is 1;
  • remain_len = max_device_len - device_len: how many more tokens it may generate;
  • can_decode = remain_len > 0: whether another decode step is allowed.

After each forward pass complete_one() advances the state: this step's tokens are now cached (cached_len = device_len) and the next step computes one more position (device_len += 1). The sampled token is appended by the scheduler to the CPU-side input_ids with append_host.

minisgl/core.py · Req
@dataclass(eq=False)  # eq=False: compare and hash by identity so a Req can go into a set
class Req:
    input_ids: torch.Tensor  # a one-dimensional int32 tensor on the CPU: the prompt plus what has been generated
    table_idx: int  # the row it occupies in the page table / token pool
    cached_len: int  # the first cached_len tokens already have their KV in the cache
    output_len: int  # how many more tokens it may generate (max_tokens)
    uid: int
    sampling_params: SamplingParams
    cache_handle: BaseCacheHandle  # the stretch locked in the prefix cache

    def __post_init__(self) -> None:
        assert self.input_ids.is_cpu
        # device_len: how many tokens will have KV in the cache after this forward pass
        self.device_len = len(self.input_ids)
        self.max_device_len = len(self.input_ids) + self.output_len
        assert 0 <= self.cached_len < self.device_len <= self.max_device_len

    @property
    def remain_len(self) -> int:
        return self.max_device_len - self.device_len

    @property
    def extend_len(self) -> int:
        """How many tokens go into the model this step: the part of the prompt that missed the cache during prefill, 1 during decode."""
        return self.device_len - self.cached_len

    def complete_one(self) -> None:
        """A forward pass is done: this step's tokens are now cached, and the next step computes one more."""
        self.cached_len = self.device_len
        self.device_len += 1

    def append_host(self, next_token: torch.Tensor) -> None:
        self.input_ids = torch.cat([self.input_ids, next_token])

    @property
    def can_decode(self) -> bool:
        return self.remain_len > 0

    def __repr__(self) -> str:
        return (
            f"Req(uid={self.uid}, table_idx={self.table_idx}, cached_len={self.cached_len}, "
            f"device_len={self.device_len}, max_device_len={self.max_device_len})"
        )
Following the request from the self-test through:

examples/ch01_req.py
"""Chapter 1: following a request's length fields through prefill and decode."""

import torch
from minisgl.core import Req, SamplingParams

req = Req(input_ids=torch.arange(6, dtype=torch.int32), table_idx=0, cached_len=2, output_len=3,
          uid=0, sampling_params=SamplingParams(max_tokens=3), cache_handle=None)


def show(stage: str) -> None:
    print(f"{stage:<18} cached_len={req.cached_len}  device_len={req.device_len}  "
          f"extend_len={req.extend_len}  remain_len={req.remain_len}  can_decode={req.can_decode}")


show("at admission (2 hit)")
for step in ("after prefill", "after decode 1", "after decode 2"):
    req.complete_one()
    req.append_host(torch.tensor([100], dtype=torch.int32))  # pretend token 100 was sampled
    show(step)
print("final input_ids:", req.input_ids.tolist())
What it prints
at admission (2 hit) cached_len=2  device_len=6  extend_len=4  remain_len=3  can_decode=True
after prefill      cached_len=6  device_len=7  extend_len=1  remain_len=2  can_decode=True
after decode 1     cached_len=7  device_len=8  extend_len=1  remain_len=1  can_decode=True
after decode 2     cached_len=8  device_len=9  extend_len=1  remain_len=0  can_decode=False
final input_ids: [0, 1, 2, 3, 4, 5, 100, 100, 100]
A few details worth noticing:

  • The prefill step has extend_len = 4: the 2 tokens that hit do not need computing again.
  • The KV of the third generated token (the one sampled at decode 2) is never computed: device_len = 9 while cached_len = 8. The last token only has to be emitted, never fed back in.
  • input_ids lives on the CPU (there is an assertion in __post_init__). The GPU keeps its own copy in the token_pool of chapter 4; the CPU copy is used for prefix-cache insertion and for deciding when a request is finished.
  • @dataclass(eq=False) makes Req compare and hash by identity. The scheduler puts requests into a set, and two different requests can have every field equal, for instance two identical prompts arriving at once.

Batch: who fills which field

minisgl/core.py · Batch
@dataclass
class Batch:
    reqs: List[Req]
    phase: Literal["prefill", "decode"]
    # the fields below are filled by the scheduler
    input_ids: torch.Tensor = field(init=False)
    positions: torch.Tensor = field(init=False)
    out_loc: torch.Tensor = field(init=False)  # where in the KV pool this step's new tokens write their KV
    padded_reqs: List[Req] = field(init=False)  # the request list after padding for CUDA Graph
    # the fields below are filled by the attention backend
    attn_metadata: BaseAttnMetadata = field(init=False)

    @property
    def is_prefill(self) -> bool:
        return self.phase == "prefill"

    @property
    def is_decode(self) -> bool:
        return self.phase == "decode"

    @property
    def size(self) -> int:
        return len(self.reqs)

    @property
    def padded_size(self) -> int:
        return len(self.padded_reqs)
Batch only takes reqs and phase in its constructor; the other fields are filled by different modules at different points in scheduling:

Field Filled by What it holds
padded_reqs GraphRunner.pad_batch (chapter 18) to use CUDA Graph, a decode batch is padded with dummy requests up to a captured batch size; without graphs it equals reqs
positions the scheduler's _prepare_batch this step's position for each token, [cached_len, device_len), concatenated into one dimension
out_loc the scheduler's _prepare_batch where in the KV pool each of this step's tokens writes its KV
input_ids the scheduler's _forward this step's input tokens, gathered by index from the token_pool on the GPU
attn_metadata the attention backend's prepare_metadata the variable-length and paging information the attention kernel needs

Every tensor packs the batch's requests end to end into one dimension ([total tokens]), with no padding. This is standard in modern inference engines: prefill and decode requests differ wildly in length, and padding would waste a lot of compute. The boundaries between requests are described by cu_seqlens in the attention metadata (chapter 5).

Context: the global "current batch"

minisgl/core.py · Context
@dataclass
class Context:
    page_size: int
    # note: page_table always stores as if page_size were 1, each element a token's location in the KV pool
    page_table: torch.Tensor = field(init=False)
    attn_backend: BaseAttnBackend = field(init=False)
    moe_backend: BaseMoeBackend = field(init=False)
    kv_cache: BaseKVCachePool = field(init=False)
    _batch: Batch | None = field(default=None, init=False)

    @property
    def batch(self) -> Batch:
        assert self._batch is not None, "No active batch in context"
        return self._batch

    @contextmanager
    def forward_batch(self, batch: Batch) -> Iterator[None]:
        assert self._batch is None, "Nested forward_batch is not allowed"
        try:
            self._batch = batch
            yield
        finally:
            self._batch = None
Context holds the things there is exactly one of per process: the page table, the attention backend, the MoE backend, the KV pool, and "the batch being computed right now". forward_batch is a context manager that sets the current batch only inside the with block, clears it on exit, and refuses to nest.

minisgl/core.py · get_global_ctx
def get_global_ctx() -> Context:
    assert _GLOBAL_CTX is not None, "Global context is not set"
    return _GLOBAL_CTX
So the model code can be written as model.forward(), with no arguments. The attention layer reads get_global_ctx().batch.positions when it needs positions, and the LM head reads batch.is_prefill when it needs to know whether this is prefill. The upside is that no arguments are threaded through the model and every layer's signature stays clean; the price is an implicit dependency, since calling the model outside a forward_batch context fails an assertion outright.

The official implementation

The whole of core.py. Our version has exactly the same fields and methods plus one extra reset_global_ctx() for tests: the official set_global_ctx refuses to be called twice, and tests create engines repeatedly in one process, so they need to clear it first.

Tests

tests/test_ch01_core.py · test_req_lengths_through_prefill_and_decode
def test_req_lengths_through_prefill_and_decode():
    req = make_req(prompt_len=5, max_tokens=3, cached_len=2)  # the first 2 tokens hit the cache
    assert (req.cached_len, req.device_len, req.extend_len, req.remain_len) == (2, 5, 3, 3)
    req.complete_one()  # prefill done: 5 tokens are cached and the next step computes position 6
    assert (req.cached_len, req.device_len, req.extend_len, req.remain_len) == (5, 6, 1, 2)
    req.complete_one()
    req.complete_one()
    assert req.remain_len == 0 and not req.can_decode
The other two tests in tests/test_ch01_core.py check the edge cases of is_greedy, and that forward_batch cannot nest and that the current batch is unreachable after it exits.

How to explain it

On request state: mini-sglang's Req describes its state with just three lengths, cached_len (already in the cache), device_len (in the cache after this step) and max_device_len (the ceiling), and this step computes extend_len = device_len - cached_len, of which prefill, chunked prefill and decode are all special cases. With a prompt of 6 and a 2-token prefix hit, the prefill step computes 4 tokens. The tensors in a Batch are one-dimensional arrays with the requests packed end to end and no padding (CUDA Graph pads the batch size through padded_reqs), and every model layer takes the current batch's positions and attention metadata from the global Context. Req uses eq=False so that it compares and hashes by identity, which is what lets it go into a set or serve as a dictionary key.

Exercises

  1. A request has a prompt of 100 tokens, max_tokens=50, and its first 64 tokens hit the cache. Write down cached_len, device_len and remain_len at admission, after prefill, and after the 10th decode step.
  2. If Req used the default eq=True, what would go wrong with the scheduler's running_reqs: Set[Req]? (Hint: a dataclass with eq=True is unhashable.)
  3. Why does Req.__post_init__ require cached_len < device_len rather than <=? When could that be violated?
Answers
  1. At admission (64, 100, 50); after prefill (100, 101, 49); after the 10th decode step (110, 111, 39).
  2. With eq=True (the default) and no frozen=True, @dataclass sets __hash__ to None, so Req cannot go into a set; and even if it could be hashed, two different requests with identical contents would be treated as one. eq=False keeps object's identity comparison and hashing.
  3. Every step has to compute at least one token: even when the whole prompt is cached, the hidden state at the last position is still needed for the logits. So the scheduler deliberately leaves the last token unmatched when matching a prefix (the input_len - 1 in match_req, chapter 8), which keeps cached_len < device_len.

Summary

  • Req describes its state with three lengths: cached_len (already cached), device_len (cached after this step) and max_device_len (the ceiling); extend_len, remain_len and can_decode all derive from them.
  • complete_one() advances the state after each forward pass, and the scheduler appends the new token to the CPU-side input_ids with append_host.
  • The tensors in a Batch are one-dimensional arrays with the requests packed end to end and no padding; the fields are filled by the scheduler, the attention backend and the GraphRunner at different points.
  • The model reads the current batch from the global Context, so forward() needs no arguments.

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