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

Chunked prefill

Prefilling a 30,000-token prompt in one go can blow out memory with its activations, especially the MLP's intermediate results and the hidden state before the logits, and it also holds the GPU for a long time while everything else waits. Chunked prefill, from Sarathi-Serve, cuts a long prompt into chunks and computes one per step. mini-sglang's implementation is very compact: a ChunkedReq subclass for "a request that cannot decode", plus a few lines of budget logic in PrefillAdder.

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

1. Is chunked prefill mathematically equivalent to prefilling all at once? Why?
2. A request is cut into 4 chunks. When chunk 2 computes attention, where does the KV of the first 16 tokens come from?
3. Should the logits a chunk produces mid-way be sampled?
4. In mini-sglang, what happens to other decoding requests while one request is being prefilled in chunks?

Answers (try it yourself first, then expand)

1. Equivalent: in causal attention each token depends only on the tokens before it, and a later chunk reads the KV earlier chunks already wrote through the page table, so the result is the same as computing it all at once (numerically there is only the tiny difference of floating-point ordering).
2. From the KV pool: when the earlier chunks finished, their KV was written into the pool and the page table recorded where, so chunk 2's attention backend reads it back by the page table. This is the same situation as a prefix-cache hit.
3. No: only when the last chunk finishes do the logits at the last position sample the first output token; a middle chunk's logits are thrown away.
4. They pause: mini-sglang does only prefill or only decode per step, so decode waits during a chunked prefill. Production SGLang solves this by putting chunked prefill and decode in the same batch.

Files you will write: fill in ChunkedReq, PrefillAdder._add_one_req and PrefillAdder.try_add_one in scheduler/prefill.py.

Files so far: 62 (0 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/  (5 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
│   ├── io.py
├── utils/  (8 files, unchanged)

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

Why it is equivalent

Attention is causal: position j depends only on tokens at positions ≤ j. Cut the prompt into [0,16), [16,32) and so on, and once the first chunk is done its KV is already in the KV pool; the second chunk's queries read the first 32 positions' KV through the page table (16 from the previous chunk and 16 just written by this one), which is exactly what positions 16 to 31 would have seen in a single prefill. The MLP and RMSNorm are per-position anyway. So chunked prefill is mathematically equivalent to prefilling at once, only in a different order.

This is the same situation as the "prefix hit" of chapter 5: to the attention backend there is no difference at all between "the first 16 tokens are in the cache and we compute 16 this step" and "the first 16 tokens hit the prefix cache". So chunked prefill needs no change to the attention backend.

ChunkedReq

minisgl/scheduler/prefill.py · ChunkedReq

class ChunkedReq(Req):
    """A request only partly prefilled: it neither samples this step nor joins decode."""

    def append_host(self, next_token: torch.Tensor) -> None:
        raise NotImplementedError("ChunkedReq should not be sampled")

    @property
    def can_decode(self) -> bool:
        return False

A request in the middle of chunking is a subclass of Req that changes two things: it cannot be sampled (the logits the LM head computes this step belong to a middle position of the prompt and mean nothing) and it cannot join the decode set.

The budget

minisgl/scheduler/prefill.py · PrefillAdder._add_one_req

def _add_one_req(self, pending_req: PendingReq, cache_handle: BaseCacheHandle,
                 table_idx: int, cached_len: int) -> Req:
    remain_len = pending_req.input_len - cached_len
    chunk_size = min(self.token_budget, remain_len)
    is_chunked = chunk_size < remain_len
    self.token_budget -= chunk_size
    self.reserved_size += remain_len + pending_req.output_len
    # write only this chunk's tokens into the token pool; the scheduler allocates the KV pages together in _prepare_batch
    s = slice(cached_len, cached_len + chunk_size)
    pin = self.cache_manager.device.type == "cuda"
    src = pending_req.input_ids[s].pin_memory() if pin else pending_req.input_ids[s]
    self.table_manager.token_pool[table_idx, s].copy_(src, non_blocking=True)
    cls = ChunkedReq if is_chunked else Req
    return cls(
        input_ids=pending_req.input_ids[: cached_len + chunk_size],
        table_idx=table_idx,
        cached_len=cached_len,
        output_len=pending_req.output_len,
        uid=pending_req.uid,
        cache_handle=cache_handle,
        sampling_params=pending_req.sampling_params,
    )

Each prefill batch has a token budget, token_budget (that is max_extend_tokens, --max-prefill-length on the command line, 8192 by default). When a request is admitted, this chunk's size is the smaller of "budget left" and "request left"; if the chunk does not cover the rest of the prompt, a ChunkedReq is created, otherwise an ordinary Req.

reserved_size reserves for the request's entire remaining prompt plus max_tokens, not just this chunk, because the moment a request is first admitted all the space it will ever need is committed and no later chunk can get stuck for want of room.

minisgl/scheduler/prefill.py · PrefillAdder.try_add_one

def try_add_one(self, pending_req: PendingReq) -> Req | None:
    if self.token_budget <= 0:
        return None
    if chunked := pending_req.chunked_req:  # chunking is under way: the resources are already allocated, so carry on with the next chunk
        return self._add_one_req(pending_req, chunked.cache_handle, chunked.table_idx,
                                 chunked.cached_len)
    if resource := self._try_allocate_one(pending_req):
        handle, table_idx = resource
        return self._add_one_req(pending_req, handle, table_idx, handle.cached_len)
    return None

PendingReq.chunked_req remembers how far the chunking has got. When the request comes round again on the next step, it reuses the request slot and cache handle it already has and carries on from the previous cached_len, with no prefix matching or admission check.

Back in PrefillManager.schedule_next_batch from the last chapter, an unfinished chunked request goes back to the head of the waiting queue and continues first next step:

self.pending_list = chunked_list + self.pending_list[len(reqs):]

The scheduler skips ChunkedReq when handling results (the isinstance check at the start of _process_last_data in chapter 7) and does not hand it to the prefix cache; its column is −1 when writing back to the token pool, since there is no sample to use as the next input.

An example

the prefill steps of three requests with max_extend_tokens=16

examples/ch10_chunked.py

"""Chapter 10: how a long prompt is chunked at max_extend_tokens=16, and how it interleaves with other requests."""

import torch
from minisgl.core import SamplingParams
from minisgl.env import ENV
from minisgl.llm import LLM
from minisgl.scheduler.prefill import ChunkedReq

ENV.DISABLE_OVERLAP_SCHEDULING.value = True
llm = LLM("models/Qwen3-0.6B", dtype=torch.float32, max_running_req=8, num_page_override=1024,
          max_seq_len_override=256, cache_type="radix", max_extend_tokens=16)
long_prompt = "Once upon a time, in a small village by the sea, there lived an old fisherman. " * 3
prompts = [long_prompt, "1 + 1 =", "The capital of France is"]
print("prompt lengths:", [len(llm.tokenizer(p).input_ids) for p in prompts])
orig = llm._forward


def traced(fi):
    b = fi.batch
    parts = []
    for r in b.reqs:
        tag = "chunk" if isinstance(r, ChunkedReq) else "full"
        parts.append(f"uid{r.uid}[{r.cached_len}:{r.device_len}]{tag}" if b.is_prefill else f"uid{r.uid}")
    print(f"  {b.phase:<7} {len(b.positions):>2} tokens: {', '.join(parts)}")
    return orig(fi)


llm._forward = traced
out = llm.generate(prompts, SamplingParams(max_tokens=3, ignore_eos=True))
llm.shutdown()
from transformers import AutoModelForCausalLM  # noqa: E402

hf = AutoModelForCausalLM.from_pretrained("models/Qwen3-0.6B", dtype=torch.float32)
refs = []
for p in prompts:
    ids = llm.tokenizer(p, return_tensors="pt").input_ids
    refs.append(hf.generate(ids, max_new_tokens=3, min_new_tokens=3, do_sample=False)[0, ids.shape[1]:].tolist())
print("matches HF (prefilled whole):", [o["token_ids"] for o in out] == refs)

What it prints

prompt lengths: [61, 5, 5]
  prefill 16 tokens: uid0[0:16]chunk
  prefill 16 tokens: uid0[16:32]chunk
  prefill 16 tokens: uid0[32:48]chunk
  prefill 16 tokens: uid0[48:61]full, uid1[0:3]chunk
  prefill  7 tokens: uid1[3:5]full, uid2[0:5]full
  decode   3 tokens: uid0, uid1, uid2
  decode   3 tokens: uid0, uid1, uid2
matches HF (prefilled whole): True

The 61-token prompt is cut into 16 + 16 + 16 + 13; the fourth step has 3 tokens of budget left, which go to uid1 behind it, making that a chunked request too; the fifth step finishes uid1's remaining 2 tokens and does uid2 in one go. Decode only starts once all three have prefilled. The final output matches Hugging Face prefilling the whole prompt at once.

The cost

Look at the trace above: there is no decode at all while the long prompt is being prefilled in chunks. That follows directly from "prefill first": as long as the waiting queue is non-empty the scheduler prefills, and the chunked request is always at the head. If other requests were decoding at the time, their next token would wait for every chunk to finish, and TPOT spikes.

What production SGLang and vLLM do is put the prefill chunk and the decoding requests in one batch: the decoding requests compute their one token as usual each step and the remaining budget goes to the prefill chunk. That way a long prompt never blocks decode, at the cost of two shapes in one batch, so CUDA Graph only covers pure decode steps and the scheduling logic gets more complex. mini-sglang chose "one kind per step" for simplicity.

The official implementation

• scheduler/prefill.py · ChunkedReq (lines 23-29)
• scheduler/prefill.py · PrefillAdder._add_one_req (lines 65-90)
• The official docs/features.md warns that a chunk that is too small, say 128, noticeably hurts performance: every chunk runs the whole model and rereads the weights, and the smaller the chunk the less of the GPU is used.

Tests

tests/test_ch10_chunked.py · test_chunked_prefill_is_exact_and_bounded

def test_chunked_prefill_is_exact_and_bounded():
    ENV.DISABLE_OVERLAP_SCHEDULING.value = True
    llm = build_llm(QWEN3, cache_type="radix", max_extend_tokens=6)
    sizes = []
    orig = llm._forward

    def spy(fi):
        if fi.batch.is_prefill:
            sizes.append(sum(r.extend_len for r in fi.batch.reqs))
        return orig(fi)

    llm._forward = spy
    out = llm.generate(PROMPTS, greedy(8))
    assert [o["token_ids"] for o in out] == hf_greedy(QWEN3, PROMPTS, 8)
    total = sum(len(llm.tokenizer(p).input_ids) for p in PROMPTS)
    # every prefill batch uses exactly its 6-token budget (except the last), and each prompt token is computed once
    assert sum(sizes) == total and all(s == 6 for s in sizes[:-1]) and sizes[-1] <= 6
    llm.cache_manager.check_integrity()
    llm.shutdown()

Every prefill batch uses exactly its 6-token budget (except the last), every prompt token is computed exactly once, and the output matches HF. Chapter 11's tests verify it again in the combination of chunking, overlap scheduling and a page size of 4.

How to explain it

On chunked prefill: causal attention makes chunking mathematically equivalent to prefilling at once, because a later chunk reads the KV earlier chunks already wrote through the page table, which is the same situation as a prefix-cache hit; middle chunks do not sample (their logits are discarded) and only the last chunk samples and joins decode. Each prefill batch is bounded by max_extend_tokens, and an unfinished chunked request returns to the head of the queue to continue first. The point is to bound the length and the memory peak of a single step; but mini-sglang does one kind per step, so decode still pauses during chunking, while production SGLang and vLLM mix chunked prefill with decode in one batch, letting decode advance every step, just more slowly.

Exercises

1. Implement a "mixed batch": in _schedule_next_batch, when the decode set is non-empty, put all the decoding requests into the batch first (one token each) and give the rest of the budget to prefill. What should Batch.phase be? What do the attention backend and the LM head need?
2. When a chunked request is aborted (the client disconnects), PrefillManager.abort_req returns its chunked_req and the scheduler then calls _free_req_resources. Why does aborting a non-chunked request that is still waiting need no release?
3. What happens if max_extend_tokens is 1?

Answers

1. You can keep phase="prefill" (there are requests with more than one query) and treat the decoding requests as "prefills" with extend_len = 1: the attention backend's metadata already supports any cu_seqlens_q, and the LM head takes each request's last position through get_last_indices, which for a decoding request is its only position. What mainly changes is the scheduler's bookkeeping (decoding requests take the normal _process_last_data path) and using CUDA Graph only on pure decode steps.
2. A request that has not been admitted holds nothing, since the request slot, the KV pages and the cache lock are all acquired inside _try_allocate_one, so removing it from the queue is enough; a chunked request already holds all of them and must release them.
3. Each prefill step computes 1 token, so a long prompt takes as many steps as it has tokens and rereads all the weights every step, which is extremely slow. The result is still correct.

Summary

• Causal attention makes chunked prefill equivalent to prefilling at once: a later chunk reads earlier chunks' KV through the page table, the same situation as a prefix hit.
• ChunkedReq neither samples nor joins decode; PendingReq.chunked_req remembers the progress and the next step continues with the same resources.
• Each prefill batch is bounded by max_extend_tokens, and a chunked request returns to the head of the queue to continue first.
• mini-sglang does one kind per step, so decode pauses while a long prompt is chunked; production builds solve this with a mixed batch.

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.

• 194. 分块 prefill 的块大小怎么选 (medium)

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