手写 mini-sglang 阶段二 · 排得好：同时服务很多请求

分块 prefill

一个 3 万 token 的提示词如果一次 prefill，激活值（尤其是 MLP 的中间结果和 logits 之前的隐状态）可能直接撑爆显存，而且这一轮会占用 GPU 很久，其他请求都得等着。分块 prefill（chunked prefill，出自 Sarathi-Serve）把长提示词切成若干块，每轮只算一块。mini-sglang 的实现非常紧凑：一个"不能 decode 的请求"子类 ChunkedReq，加上 PrefillAdder 里的几行预算逻辑。

自测：能答上来就可以跳过本章

1. 分块 prefill 的结果与一次性 prefill 在数学上等价吗？为什么？
2. 一个请求被切成了 4 块，第 2 块要算注意力时，前面 16 个 token 的 KV 从哪来？
3. 分块中的请求这一轮产生的 logits 要不要采样？
4. mini-sglang 里，一个请求正在分块 prefill 时，其他正在 decode 的请求会怎样？

自测参考答案（先自己答，再展开对照）

1. 等价：因果注意力里每个 token 只依赖它之前的 token，后面的块通过页表读取前面块已经写好的 KV，得到的结果和一次算完相同（数值上只有浮点顺序的微小差异）。
2. 从 KV 池里：前面几块算完时它们的 KV 已经写入池中，页表记着位置，第 2 块的注意力后端按页表读出来——这和命中前缀缓存是同一种情形。
3. 不要：只有最后一块算完时，最后位置的 logits 才用来采样第一个输出 token；中间块的 logits 直接丢弃。
4. 暂停：mini-sglang 一轮只做 prefill 或 decode，分块 prefill 期间 decode 要等；正式版的 SGLang 把分块 prefill 和 decode 混在同一个 batch 里解决这个问题。

本章要写的文件：补全 scheduler/prefill.py 中的 ChunkedReq、PrefillAdder._add_one_req、PrefillAdder.try_add_one。

到这一章为止的文件：62 个（新建 0，修改 1，提前引入 0）

minisgl/
├── __init__.py
├── core.py
├── env.py
├── attention/  （4 个文件，不变）
├── distributed/  （3 个文件，不变）
├── engine/  （5 个文件，不变）
├── kernel/  （2 个文件，不变）
├── kvcache/  （5 个文件，不变）
├── layers/  （8 个文件，不变）
├── llm/  （2 个文件，不变）
├── message/  （4 个文件，不变）
├── models/  （7 个文件，不变）
├── moe/  （2 个文件，不变）
├── scheduler/
│   ├── table.py
│   ├── __init__.py
│   ├── config.py
│   ├── utils.py
│   ├── decode.py
│   ├── prefill.py  ← 本章修改
│   ├── scheduler.py
│   ├── cache.py
│   ├── io.py
├── utils/  （8 个文件，不变）

这一步的 main：examples/ch10_chunked.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

为什么等价

注意力是因果的：第 j 个位置只依赖位置 ≤ j 的 token。把提示词切成 [0,16)、[16,32)、……，第一块算完后它的 KV 已经写进了 KV 池；第二块的 query 通过 page table 读取前 32 个位置的 KV（16 个来自上一块、16 个是本块刚写入的），与一次性 prefill 时第 16～31 个位置看到的完全相同。MLP、RMSNorm 都是逐位置独立的，更不受影响。所以分块 prefill 与一次性 prefill 在数学上等价，只是换了一种计算顺序。

这正是第 5 章"前缀命中"的同一种情形：对注意力后端来说，"前 16 个 token 在缓存里、本轮算 16 个"和"前 16 个 token 命中了前缀缓存"没有任何区别。所以分块 prefill 不需要注意力后端做任何改动。

ChunkedReq

minisgl/scheduler/prefill.py · ChunkedReq

class ChunkedReq(Req):
    """只 prefill 了一部分的请求：这一轮不采样，也不进入 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

分块中的请求是 Req 的一个子类，只改了两处：不能被采样（本轮 LM head 算出的 logits 属于提示词的中间位置，没有意义），也不能进入 decode 集合。

预算

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
    # 只把本块的 token 写进 token pool；KV 页由调度器在 _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,
    )

每个 prefill batch 有一个 token 预算 token_budget（即 max_extend_tokens，命令行 --max-prefill-length，默认 8192）。接纳一个请求时，本块的大小是"预算剩余"和"请求剩余"中较小的那个；如果本块没能覆盖请求剩下的全部提示词，就创建 ChunkedReq，否则创建普通的 Req。

reserved_size 按请求剩余的全部提示词加 max_tokens 预留，而不只是本块——请求第一次被接纳时，它未来需要的全部空间就已经承诺了，后续的块不会因为空间不够而卡住。

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:  # 分块进行中：资源早已分好，接着做下一块
        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 记住分块进行到哪里。下一轮再遇到这个请求时，直接沿用它已经分配好的请求槽和缓存句柄，从上次的 cached_len 继续，不再做前缀匹配和准入检查。

回到上一章的 PrefillManager.schedule_next_batch：没做完的分块请求被放回等待队列的队首，下一轮优先继续：

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

调度器处理结果时跳过 ChunkedReq（第 7 章 _process_last_data 开头的 isinstance 判断），也不把它交给前缀缓存；写回 token pool 时它的列号是 −1（不需要把采样结果当作下一轮输入）。

看一个例子

max_extend_tokens=16 时三个请求的 prefill 轮次

examples/ch10_chunked.py

"""第 10 章：max_extend_tokens=16 时，一个长提示词怎样被切块，以及它与其他请求的交错。"""

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("提示词长度:", [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 = "分块" if isinstance(r, ChunkedReq) else "完整"
        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} 个 token: {', '.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("与 HF（整段 prefill）的输出一致:", [o["token_ids"] for o in out] == refs)

运行结果

提示词长度: [61, 5, 5]
  prefill 16 个 token: uid0[0:16]分块
  prefill 16 个 token: uid0[16:32]分块
  prefill 16 个 token: uid0[32:48]分块
  prefill 16 个 token: uid0[48:61]完整, uid1[0:3]分块
  prefill  7 个 token: uid1[3:5]完整, uid2[0:5]完整
  decode   3 个 token: uid0, uid1, uid2
  decode   3 个 token: uid0, uid1, uid2
与 HF（整段 prefill）的输出一致: True

61 个 token 的长提示词被切成 16 + 16 + 16 + 13；第四轮剩下 3 个 token 的预算，分给了排在后面的 uid1，它也成了一个分块请求；第五轮 uid1 做完剩下的 2 个 token，uid2 一次做完。三个请求都完成 prefill 后才开始 decode。最终输出与 Hugging Face 整段 prefill 的结果相同。

代价

注意上面的轨迹：在长提示词分块 prefill 的整个过程中，没有任何 decode。这是"prefill 优先"策略的直接结果——只要等待队列不空，调度器就做 prefill，而分块中的请求始终在队首。如果此时有别的请求正在 decode，它们的下一个 token 要等所有块都做完才会出来，TPOT 出现一个尖峰。

SGLang 正式版和 vLLM 的做法是把 prefill 块和 decode 请求放进同一个 batch：每轮 decode 请求照常各算一个 token，剩余的预算给 prefill 块。这样长提示词不会阻塞 decode，代价是 batch 里同时有两种形状，CUDA Graph 只能覆盖纯 decode 的轮次，调度逻辑也更复杂。mini-sglang 为了简单选择了"一轮只做一种"。

官方实现

• scheduler/prefill.py · ChunkedReq（第 23～29 行）
• scheduler/prefill.py · PrefillAdder._add_one_req（第 65～90 行）
• 官方文档 docs/features.md 提醒：块太小（例如 128）会明显降低性能——每块都要完整地过一遍模型，权重被重复读取，块越小 GPU 越吃不满。

测试

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)
    # 每个 prefill batch 恰好用满 6 个 token 的预算（最后一个除外），所有提示词 token 各算一次
    assert sum(sizes) == total and all(s == 6 for s in sizes[:-1]) and sizes[-1] <= 6
    llm.cache_manager.check_integrity()
    llm.shutdown()

每个 prefill batch 恰好用满 6 个 token 的预算（最后一个除外），所有提示词 token 恰好各算一次，输出与 HF 一致。第 11 章的测试还会在"分块 + 重叠调度 + page size 4"的组合下再验证一遍。

怎么讲清楚

讲分块 prefill：因果注意力让分块和一次性 prefill 在数学上等价——后面的块通过 page table 读前面块已经写好的 KV，和命中前缀缓存是同一种情形；中间的块不采样（logits 丢掉），最后一块才采样并进入 decode。每个 prefill batch 受 max_extend_tokens 限制，未完成的分块请求回到队首优先继续。作用是限制单轮的时长和显存峰值；但 mini-sglang 一轮只做一种，分块期间 decode 仍会暂停，正式的 SGLang / vLLM 把分块 prefill 和 decode 混在同一个 batch 里，decode 每一步都能前进，只是变慢一些。

练习

1. 实现"混合 batch"：在 _schedule_next_batch 里，如果 decode 集合不空，先把所有 decode 请求放进 batch（每个 1 个 token），剩余的预算再给 prefill。Batch.phase 该怎么设？注意力后端和 LM head 需要改什么？
2. 分块请求被 abort（客户端断开）时，PrefillManager.abort_req 返回了它的 chunked_req，调度器随后调用 _free_req_resources。为什么非分块、还在等待的请求 abort 时不需要释放资源？
3. max_extend_tokens 设为 1 会发生什么？

参考答案

1. 可以沿用 phase="prefill"（因为有 query 数大于 1 的请求），把 decode 请求也当作 extend_len = 1 的"prefill"：注意力后端的元数据本来就支持任意的 cu_seqlens_q；LM head 用 get_last_indices 取每个请求的最后位置，对 decode 请求就是它唯一的位置。要改的主要是调度器的记账（decode 请求走 _process_last_data 的正常路径）和 CUDA Graph 只在纯 decode 时使用。
2. 还没被接纳的请求没有分配任何东西（请求槽、KV 页、缓存锁都是在 _try_allocate_one 里才获得的），从队列里删掉即可；分块中的请求已经拿到了这些资源，必须释放。
3. 每轮 prefill 只算 1 个 token，长提示词要跑和它长度一样多的轮次，每轮都要读一遍全部权重，极慢；但结果依然正确。

小结

• 因果注意力让分块 prefill 与一次性 prefill 等价：后面的块通过 page table 读取前面块的 KV，与前缀命中是同一种情形。
• ChunkedReq 不采样、不进 decode；PendingReq.chunked_req 记住进度，下一轮沿用资源继续。
• 每个 prefill batch 受 max_extend_tokens 限制；分块请求回到队首优先继续。
• mini-sglang 一轮只做一种，长提示词分块期间 decode 会暂停；正式版用混合 batch 解决。

本章练习题（1 道）

在浏览器里直接写代码、跑测试；也可以在本地用 python practice/judge.py 判题。

• 194. 分块 prefill 的块大小怎么选（中等）

发现错误或有更好的写法？在 GitHub 上编辑此页 · 提交问题