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

CacheManager: allocation, release and admission control

At any moment every page of the KV pool is in one of three states: free, owned by one request, or in the prefix cache (possibly shared by several requests). CacheManager moves pages between those states and answers the question the scheduler cares about most: if we admit one more request, does its KV still fit? mini-sglang's answer is conservative, reserving for the worst case, which is why it never needs preemption.

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

1. vLLM preempts running requests when KV runs short. Why does mini-sglang not need to? What does that cost?
2. Why does available_size count the evictable part of the prefix cache?
3. Why does prefix matching stop at the prompt's second-to-last token?
4. With a page size of 4, a request has cached_len=6 and device_len=9. How many pages does this step allocate?

Answers (try it yourself first, then expand)

1. Admission control reserves KV for the worst case (the rest of the prompt plus max_tokens), so an admitted request is guaranteed to run to completion and preemption is never needed. The cost is that most of the reservation goes unused, fewer requests run at once, and concurrency is lower.
2. Any node in the prefix cache that is not locked can be evicted at any moment to make room. It occupies memory but counts as available to a new request; leaving it out underestimates the available space and rejects requests for nothing.
3. Every step has to compute at least one token: even with the whole prompt cached, the hidden state at the last position is still needed for the logits. So matching stops at the second-to-last token, keeping cached_len < device_len.
4. It allocates the pages [ceil(6/4), ceil(9/4)) = [2, 3), so one page: positions 6 and 7 are still in the already allocated page 1 (covering positions 4-7), and only position 8 needs a new page 2.

Files you will write: scheduler/cache.py; and fill in PrefillAdder._try_allocate_one in scheduler/prefill.py.

Files so far: 61 (1 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/  (4 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  ← modified in this chapter
│   ├── scheduler.py
│   ├── cache.py  ← new in this chapter
│   ├── io.py
├── utils/  (8 files, unchanged)

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

The free-page list

minisgl/scheduler/cache.py · CacheManager.__init__

def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, type: str):
    device = page_table.device
    self.free_slots = torch.arange(num_pages, dtype=torch.int32, device=device) * page_size
    self.prefix_cache = create_prefix_cache(device=device, type=type)
    self.device = device
    self.num_pages = num_pages
    self.page_table = page_table
    self.page_size = page_size

free_slots holds the location of each free page's first token (with a page size of 4 that is [0, 4, 8, …]), not the page number. That matches chapter 4's "the page table stores token locations": an allocated page only needs the offset within it added to give a token location, and freeing takes every page_size-th token location to get back to the page's start.

minisgl/scheduler/cache.py · CacheManager._allocate

def _allocate(self, needed_pages: int) -> torch.Tensor:
    if needed_pages > (free_pages := len(self.free_slots)):  # not enough free pages: evict from the prefix cache
        evicted = self.prefix_cache.evict((needed_pages - free_pages) * self.page_size)
        self.free_slots = torch.cat([self.free_slots, evicted[:: self.page_size]])
        assert len(self.free_slots) >= needed_pages, "Eviction did not free enough space."
    allocated = self.free_slots[:needed_pages]
    self.free_slots = self.free_slots[needed_pages:]
    return allocated

Allocating is slicing from the front of free_slots. When there are not enough free pages it first has the prefix cache evict some (chapter 9), adds what that frees, and then allocates.

Allocating pages for this step's tokens

minisgl/scheduler/cache.py · CacheManager.allocate_paged

def allocate_paged(self, reqs: List[Req]) -> None:
    """Allocate KV pages for this step's tokens and write them into the page table. Pages already allocated are not allocated again."""
    needed_pages = 0
    info: List[Tuple[int, int, int]] = []
    for req in reqs:
        first_page = div_ceil(req.cached_len, self.page_size)
        last_page = div_ceil(req.device_len, self.page_size)
        if last_page > first_page:
            needed_pages += last_page - first_page
            info.append((req.table_idx, first_page, last_page))
    if needed_pages > 0:
        allocated = self._page_to_token(self._allocate(needed_pages))
        _write_page_table(self.page_table, allocated, info, self.page_size)

For each request the pages already allocated are [0, ceil(cached_len / page_size)) and the pages needed after this step are [0, ceil(device_len / page_size)), and the difference is what gets allocated. With a page size of 4, cached_len=6 and device_len=9: there are 2 pages already (covering positions 0-7), 3 are needed, so 1 is allocated. During decode most steps need no new page, with one allocation every page_size steps.

Every request's pages are allocated in one go and written into the page table in one go:

minisgl/scheduler/cache.py · _write_page_table

def _write_page_table(page_table: torch.Tensor, allocated: torch.Tensor,
                      info: List[Tuple[int, int, int]], page_size: int) -> None:
    """Write every newly allocated location into the page table at once: compute the (row, column) indices on the CPU first, then one indexed assignment."""
    n = len(allocated)
    use_pin = pin(page_table.device)
    rows = torch.empty(n, dtype=torch.int64, pin_memory=use_pin)
    cols = torch.empty(n, dtype=torch.int64, pin_memory=use_pin)
    offset = 0
    for table_idx, first_page, last_page in info:
        first_pos, last_pos = first_page * page_size, last_page * page_size
        length = last_pos - first_pos
        rows[offset : offset + length].fill_(table_idx)
        torch.arange(first_pos, last_pos, out=cols[offset : offset + length])
        offset += length
    assert offset == n
    device = page_table.device
    page_table[rows.to(device, non_blocking=True), cols.to(device, non_blocking=True)] = allocated

All the (row, column) indices are computed on the CPU first and then written with a single indexed assignment. On a GPU that is one asynchronous copy and one kernel, rather than one operation per request per token.

Freeing, and freeing lazily

minisgl/scheduler/cache.py · CacheManager.lazy_free_region

@contextmanager
def lazy_free_region(self) -> Iterator[None]:
    """Frees inside this region are collected and concatenated into free_slots once at the end (saving many torch.cat calls)."""
    lazy_list: List[torch.Tensor] = []

    def lazy_free(indices: torch.Tensor) -> None:
        if len(indices) > 0:
            lazy_list.append(indices[:: self.page_size])

    try:
        self._free = lazy_free  # type: ignore[method-assign]
        yield
    finally:
        del self._free
        self.free_slots = torch.cat([self.free_slots] + lazy_list)

When a batch's results are handled, many requests may finish at once and each has pages to free. One torch.cat on free_slots per request would be N copies. Inside lazy_free_region, _free only records what to free, and the whole lot is concatenated once on exit. The scheduler's _process_last_data is wrapped entirely in that region.

Admission control

The second half of the animation at the start of chapter 7 shows the rules below: reserve for the worst case, and never preempt.

Whether a request can be admitted is decided by PrefillAdder._try_allocate_one:

minisgl/scheduler/prefill.py · PrefillAdder._try_allocate_one

def _try_allocate_one(self, req: PendingReq) -> Tuple[BaseCacheHandle, int] | None:
    if self.table_manager.available_size == 0:
        return None
    handle = self.cache_manager.match_req(req).cuda_handle
    cached_len = handle.cached_len
    # admission control: the KV this request needs at worst (the rest of the input plus the maximum output) has to fit
    estimated_len = req.input_len - cached_len + req.output_len
    if estimated_len + self.reserved_size > self.cache_manager.available_size:
        return None
    self.cache_manager.lock(handle)
    # once locked, what was hit is no longer evictable, so the available space may shrink and needs rechecking
    if estimated_len + self.reserved_size > self.cache_manager.available_size:
        self.cache_manager.unlock(handle)
        return None
    table_idx = self.table_manager.allocate()
    if cached_len > 0:  # the prefix that hit: its tokens and KV locations go straight into this row
        pin = self.cache_manager.device.type == "cuda"
        src = req.input_ids[:cached_len].pin_memory() if pin else req.input_ids[:cached_len]
        self.table_manager.token_pool[table_idx, :cached_len].copy_(src, non_blocking=True)
        self.table_manager.page_table[table_idx, :cached_len].copy_(handle.get_matched_indices())
    return handle, table_idx

The criterion is the worst case: the KV this request still needs is the part of the prompt that missed the cache plus max_tokens. Together with reserved_size, what is already promised to others, it must not exceed available_size (the free pages plus the evictable part of the prefix cache).

reserved_size comes from two places:

• the space the decoding requests may still need before this batch was assembled, namely DecodeManager.inflight_tokens, the sum of each request's remain_len plus one extra page of slack per request (the last page may be only partly used);
• the requests already admitted into this batch (reserved_size += remain_len + output_len in _add_one_req).

The check runs twice with lock(handle) in between: before the lock the prefix that was hit is evictable in the prefix cache and counts towards available_size; after the lock it is protected and available_size shrinks accordingly. Failing the first check can return immediately and save a lock-then-unlock; the second check is the accurate one.

Because every request reserves its worst case at admission, a running request never stalls for want of KV, so mini-sglang needs neither preemption nor recomputation. The price is being conservative: a request with a large max_tokens that hits EOS early leaves most of its reservation unused, which lowers how many requests can run at once. vLLM takes the opposite line, admitting optimistically and, when KV runs short, preempting the most recently arrived request, freeing its KV and recomputing later.

Here is an example where the KV pool holds only 64 tokens:

examples/ch08_cache_manager.py

"""Chapter 8: how admission control queues requests when the KV pool holds only 64 tokens."""

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

ENV.DISABLE_OVERLAP_SCHEDULING.value = True
llm = LLM("models/Qwen3-0.6B", dtype=torch.float32, max_running_req=8, num_page_override=64,
          max_seq_len_override=128, cache_type="naive")
prompts = ["The capital of France is", "List three prime numbers:", "def fibonacci(n):",
           "Once upon a time, in a small village by the sea, there lived", "1 + 1 ="]
print("prompt lengths:", [len(llm.tokenizer(p).input_ids) for p in prompts], " max_tokens = 8")
orig = llm._forward


def traced(fi):
    b = fi.batch
    out = orig(fi)
    cm = llm.cache_manager
    print(f"  {b.phase:<7} uids={[r.uid for r in b.reqs]!s:<16} waiting={len(llm.prefill_manager.pending_list)} "
          f"free pages={len(cm.free_slots):<3} reserved for running={llm.decode_manager.inflight_tokens}")
    return out


llm._forward = traced
out = llm.generate(prompts, SamplingParams(max_tokens=8, ignore_eos=True))
llm.cache_manager.check_integrity()
print("all done, free pages:", len(llm.cache_manager.free_slots), "/ 64, integrity check passed")
llm.shutdown()

What it prints

prompt lengths: [5, 5, 4, 15, 5]  max_tokens = 8
  prefill uids=[0, 1, 2, 3]     waiting=1 free pages=35  reserved for running=28
  decode  uids=[0, 1, 2, 3]     waiting=1 free pages=31  reserved for running=24
  decode  uids=[0, 1, 2, 3]     waiting=1 free pages=27  reserved for running=20
  decode  uids=[0, 1, 2, 3]     waiting=1 free pages=23  reserved for running=16
  decode  uids=[0, 1, 2, 3]     waiting=1 free pages=19  reserved for running=12
  decode  uids=[0, 1, 2, 3]     waiting=1 free pages=15  reserved for running=8
  decode  uids=[0, 1, 2, 3]     waiting=1 free pages=11  reserved for running=4
  decode  uids=[0, 1, 2, 3]     waiting=1 free pages=7   reserved for running=0
  prefill uids=[4]              waiting=0 free pages=59  reserved for running=7
  decode  uids=[4]              waiting=0 free pages=58  reserved for running=6
  decode  uids=[4]              waiting=0 free pages=57  reserved for running=5
  decode  uids=[4]              waiting=0 free pages=56  reserved for running=4
  decode  uids=[4]              waiting=0 free pages=55  reserved for running=3
  decode  uids=[4]              waiting=0 free pages=54  reserved for running=2
  decode  uids=[4]              waiting=0 free pages=53  reserved for running=1
  decode  uids=[4]              waiting=0 free pages=52  reserved for running=0
all done, free pages: 64 / 64, integrity check passed

The first four requests need 13 + 13 + 12 + 23 = 61 tokens in the worst case, which fits; the fifth needs another 13, which is over 64, so it waits. Notice that the pool has free pages the whole time until all four finish (7 pages at the end), which is what conservative reservation costs. Once the first four are done the space comes back and the fifth is admitted. When everything has finished all 64 pages are back in the free list and the integrity check passes.

Handing the KV to the prefix cache

When and how a request's KV goes to the prefix cache is the most delicate piece of logic in CacheManager:

minisgl/scheduler/cache.py · CacheManager.cache_req

def cache_req(self, req: Req, *, finished: bool) -> None:
    """Hand the KV this request has computed to the prefix cache.

    [0, old.cached_len)              already cached before prefill (this request locks it)
    [old.cached_len, cached_len)     computed by this request, but another request inserted the
                                     same prefix meanwhile, so the cache already has a copy and
                                     this one is redundant: free it
    [cached_len, new.cached_len)     newly inserted into the cache, ownership passing to it
    [new.cached_len, req.cached_len) the tail too short for a page: freed when the request ends,
                                     kept otherwise
    """
    insert_ids = req.input_ids[: req.cached_len]
    page_indices = self.page_table[req.table_idx, : req.cached_len]
    old_handle = req.cache_handle
    cached_len, new_handle = self.prefix_cache.insert_prefix(insert_ids, page_indices)
    self.unlock(old_handle)  # unlock only after everything depending on the old handle is done
    self._free(page_indices[old_handle.cached_len : cached_len])
    if finished:
        self._free(page_indices[new_handle.cached_len :])
    else:
        req.cache_handle = new_handle
        self.lock(new_handle)

It is called at two moments: right after prefill (finished=False, the prompt's KV enters the cache while the request keeps running) and when the request finishes (finished=True). The comments split a request's KV into four regions:

cache_req handles a request's KV in four regions

Where does the middle region, "computed here but already in the cache", come from? Two requests with the same prefix arrive together and prefill together, so neither hits the cache; the first to finish inserts the prefix, and when the second tries to insert it finds it already there. Its own copy is redundant and has to be freed, or it leaks. With the naive cache insert_prefix keeps nothing, so a request's pages are all freed when it ends. Once chapter 9 implements the radix cache, all four regions really show up.

The integrity check

minisgl/scheduler/cache.py · CacheManager.check_integrity

def check_integrity(self) -> None:
    self.prefix_cache.check_integrity()
    cache_pages = self.prefix_cache.size_info.total_size // self.page_size
    if len(self.free_slots) + cache_pages != self.num_pages:
        raise RuntimeError(
            f"CacheManager integrity check failed: free_pages({len(self.free_slots)}) + "
            f"cache_pages({cache_pages}) != num_pages({self.num_pages})"
        )
    if self.page_size > 1:
        assert torch.all(self.free_slots % self.page_size == 0)

The scheduler runs a check when it goes idle (run_when_idle): at that moment there are no running requests, so every page is either free or in the cache, and the two must add up to the total. Anything forgotten (or freed twice) shows up right here.

The official implementation

• scheduler/cache.py · CacheManager (lines 15-124)
• admission control: scheduler/prefill.py · PrefillAdder._try_allocate_one (lines 39-63)
• the reservation: scheduler/decode.py · DecodeManager.inflight_tokens (lines 27-30)

The official repository's tests/core/test_cache_allocate.py specifically tests that with a page size above 1 the pages obtained by "evict, then allocate" are page-aligned and do not overlap.

Tests

tests/test_ch08_cache_manager.py · test_allocate_paged_writes_page_aligned_token_positions

def test_allocate_paged_writes_page_aligned_token_positions():
    cm = _cm(num_pages=8, page_size=4, cache="naive")
    handle = cm.prefix_cache.match_prefix(torch.arange(1)).cuda_handle
    req = _req(table_idx=2, prompt_len=6, cached_len=0, handle=handle)
    cm.allocate_paged([req])  # 6 tokens need 2 pages
    row = cm.page_table[2, :8].tolist()
    assert row == [0, 1, 2, 3, 4, 5, 6, 7] and len(cm.free_slots) == 6
    req.complete_one()  # device_len 6 -> 7: still inside page 2, no new page needed
    cm.allocate_paged([req])
    assert len(cm.free_slots) == 6
    req.complete_one()
    req.complete_one()  # device_len 9: a third page is needed
    cm.allocate_paged([req])
    assert len(cm.free_slots) == 5 and cm.page_table[2, 8].item() % 4 == 0

tests/test_ch08_cache_manager.py · test_integrity_check_detects_leak

def test_integrity_check_detects_leak():
    cm = _cm(num_pages=8, page_size=1, cache="naive")
    cm.check_integrity()
    cm._allocate(3)  # allocated yet handed to nobody: a leak
    with pytest.raises(RuntimeError, match="integrity"):
        cm.check_integrity()

How to explain it

On KV management: the free-page list stores a start location per page, each step allocates only the pages [ceil(cached_len / page size), ceil(device_len / page size)), and decode allocates once every "page size" steps; when pages run short the prefix cache evicts, which is why the available space has to count the cache's evictable part; and frees are batched (lazy free). Admission control reserves for the worst case (the rest of the prompt plus max_tokens), so preemption is never needed, at the cost of lower concurrency than vLLM, which admits optimistically and preempts and recomputes when short. Prefix matching stops at the prompt's second-to-last token so that at least one token is computed and there are logits. When a request ends its KV goes to the prefix cache region by region, every page has exactly one owner, and the idle-time integrity check catches leaks.

Exercises

1. inflight_tokens reserves an extra page_size - 1 tokens per running request. Give an example of what goes wrong with a page size of 16 if it does not.
2. Implement an "optimistic" admission policy: reserve only half of max_tokens, and when KV runs short preempt the most recently admitted request (free its resources and put it back at the head of the waiting queue). What has to change? When the preempted request prefills again, how does the prefix cache help?
3. If cache_req forgot to free the [old.cached_len, cached_len) region, where would it be noticed?

Answers

1. A request's KV is allocated in pages, while remain_len is counted in tokens and the next allocation is a whole page. A request with device_len=17 and remain_len=1 still needs a whole new page of 16 tokens. Add up several such requests and a token-based estimate underestimates the pages actually needed, so an allocation can fail.
2. Check the free pages before allocate_paged in _prepare_batch and, when short, pick a request from the decode set: table_manager.free, cache_manager.cache_req(finished=True), and wrap it back into a PendingReq at the head of the queue. The tokens it has already generated have to be spliced into input_ids. Because cache_req handed its KV to the radix cache, most of the prefix hits directly when it prefills again (unless it was evicted), so recomputation is cheap. This is exactly what production SGLang does.
3. Those pages would be neither in the free list nor owned by the prefix cache (which holds another copy of the same content), so nobody would hold them once the request ended. The next check_integrity while the scheduler is idle would find that free pages plus cached pages do not add up to the total and raise. Chapter 9's end-to-end test exercises this path with two requests sharing a prefix prefilling at once.

Summary

• free_slots stores a start location per page; allocation slices from the front, eviction from the prefix cache covers a shortfall, and frees are batched inside lazy_free_region.
• Each step allocates only the pages [ceil(cached_len/ps), ceil(device_len/ps)), so decode allocates once every page_size steps.
• Admission control reserves for the worst case (the rest of the prompt plus max_tokens), so there is never any preemption, at the cost of lower concurrency.
• cache_req handles a request's KV in four regions, so every page has exactly one owner, and the idle-time integrity check catches any leak.

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.

• 192. cache_req：把请求的 KV 交给前缀缓存 (medium)

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