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

CacheManager：分配、释放与准入控制

KV 池的每一页，在任何时刻只有三种状态：空闲、被某个请求独占、或者在前缀缓存里（可能同时被若干请求共享）。CacheManager 负责在这三种状态之间搬运页，并回答调度器最关心的问题：再接纳一个请求，KV 还放得下吗？mini-sglang 的答案很保守——按最坏情况预留，所以它永远不需要抢占。

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

1. vLLM 在 KV 不够时会抢占（preempt）正在运行的请求，mini-sglang 为什么不需要？代价是什么？
2. available_size 为什么要把前缀缓存里"可淘汰"的部分也算进去？
3. 匹配前缀时为什么只匹配到提示词的倒数第二个 token？
4. page size 为 4 时，一个请求 cached_len=6、device_len=9，本轮要分配几页？

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

1. 准入控制按最坏情况（剩余提示词 + max_tokens）预留 KV，接纳的请求一定能跑完，永远不需要抢占；代价是预留的空间大多用不满，同时运行的请求少，并发较低。
2. 前缀缓存里没有被锁住的节点随时可以淘汰腾出空间，它们占着显存，但对新请求来说是"可用"的；不算进去会低估可用空间，白白拒绝请求。
3. 每一轮至少要算一个 token：即使整个提示词都命中，也需要最后一个位置的隐状态来算 logits，所以匹配到倒数第二个 token 为止，保证 cached_len < device_len。
4. 要为 [ceil(6/4), ceil(9/4)) = [2, 3) 这些页分配，即 1 页：位置 6、7 还在已经分配的第 1 页里（覆盖位置 4～7），只有位置 8 需要新的第 2 页。

本章要写的文件：scheduler/cache.py；补全 scheduler/prefill.py 中的 PrefillAdder._try_allocate_one。

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

minisgl/
├── __init__.py
├── core.py
├── env.py
├── attention/  （4 个文件，不变）
├── distributed/  （3 个文件，不变）
├── engine/  （5 个文件，不变）
├── kernel/  （2 个文件，不变）
├── kvcache/  （4 个文件，不变）
├── 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/ch08_cache_manager.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

空闲页列表

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 存的是每个空闲页第一个 token 的位置（page size 为 4 时是 [0, 4, 8, …]），而不是页号。这与第 4 章"page table 按 token 存位置"一致：分配出去的页只需加上页内偏移就是 token 位置，释放时从 token 位置里每隔 page size 取一个就回到页的起点。

minisgl/scheduler/cache.py · CacheManager._allocate

def _allocate(self, needed_pages: int) -> torch.Tensor:
    if needed_pages > (free_pages := len(self.free_slots)):  # 空闲页不够：从前缀缓存里淘汰
        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

分配就是从 free_slots 前面切一段。空闲页不够时，先让前缀缓存淘汰一些（第 9 章），把淘汰出来的页补进来再分配。

为本轮的 token 分配页

minisgl/scheduler/cache.py · CacheManager.allocate_paged

def allocate_paged(self, reqs: List[Req]) -> None:
    """为本轮要计算的 token 分配 KV 页，并写进 page table。已经分过页的部分不重复分配。"""
    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)

对每个请求，已经分配过的页是 [0, ceil(cached_len / page_size))，本轮结束后需要的页是 [0, ceil(device_len / page_size))，差值就是要新分配的页。page size 为 4、cached_len=6、device_len=9 时：已有 2 页（覆盖位置 0～7），需要 3 页，新分配 1 页。decode 时大多数轮次不需要新页，每 page size 轮才分配一次。

所有请求需要的页一次分配，再一次性写进 page table：

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:
    """一次性把所有新分配的位置写进 page table：先在 CPU 上算好 (行, 列) 下标，再一次索引赋值。"""
    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

先在 CPU 上算好所有 (行, 列) 下标，再用一次索引赋值写入。GPU 上这是一次异步拷贝加一个 kernel，而不是每个请求、每个 token 一次操作。

释放与"懒释放"

minisgl/scheduler/cache.py · CacheManager.lazy_free_region

@contextmanager
def lazy_free_region(self) -> Iterator[None]:
    """在这个区域内的释放先攒着，最后一次性拼进 free_slots（少做很多次 torch.cat）。"""
    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)

处理一个 batch 的结果时，可能有很多请求同时结束，每个都要释放若干页。如果每次都 torch.cat 一次 free_slots，就是 N 次拷贝。lazy_free_region 在这个区域内把 _free 换成"先记下来"，退出时一次性拼接。调度器的 _process_last_data 整个都包在这个区域里。

准入控制

第 7 章开头的动画后半段演示了下面这套准入规则：按最坏情况预留，所以永不抢占。

请求能不能被接纳，由 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
    # 准入控制：这个请求最坏情况下需要的 KV 空间（剩余输入 + 最大输出）必须放得下
    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)
    # 加锁后命中的部分不再可淘汰，可用空间可能变小，需要再检查一次
    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:  # 命中的前缀：token 和 KV 位置直接填进这一行
        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

判断标准是最坏情况：这个请求还需要的 KV = 没命中缓存的提示词 + max_tokens。加上 reserved_size（已经承诺给别人的），不能超过 available_size（空闲页 + 前缀缓存里可以淘汰的部分）。

reserved_size 有两个来源：

• 组这个 batch 之前，正在 decode 的请求未来还可能用到的空间——DecodeManager.inflight_tokens，每个请求的 remain_len 之和，再给每个请求多留一页的余量（最后一页可能只用了一部分）；
• 本 batch 里已经接纳的请求（_add_one_req 里 reserved_size += remain_len + output_len）。

检查做了两次，中间是 lock(handle)：加锁之前，命中的前缀在前缀缓存里是"可淘汰"的，被算在 available_size 里；加锁之后它变成"受保护"，available_size 相应变小。第一次检查失败可以直接返回，省掉加锁再解锁的开销；第二次检查才是准确的。

因为每个请求在接纳时就预留了最坏情况的空间，运行中的请求永远不会因为 KV 不够而停下，mini-sglang 不需要抢占（preemption）和重算。代价是保守：max_tokens 设得很大而实际很早遇到 EOS 的请求，预留的空间大部分没用上，能同时运行的请求变少。vLLM 采取相反的策略——乐观地接纳，KV 不够时抢占最晚到达的请求、释放它的 KV，之后重算。

看一个 KV 池只有 64 个 token 的例子：

examples/ch08_cache_manager.py

"""第 8 章：KV 池只有 64 个 token 时，准入控制怎样让请求排队。"""

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("提示词长度:", [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} 等待中={len(llm.prefill_manager.pending_list)} "
          f"空闲页={len(cm.free_slots):<3} 运行中预留={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("全部完成，空闲页:", len(llm.cache_manager.free_slots), "/ 64，完整性检查通过")
llm.shutdown()

运行结果

提示词长度: [5, 5, 4, 15, 5]  max_tokens = 8
  prefill uids=[0, 1, 2, 3]     等待中=1 空闲页=35  运行中预留=28
  decode  uids=[0, 1, 2, 3]     等待中=1 空闲页=31  运行中预留=24
  decode  uids=[0, 1, 2, 3]     等待中=1 空闲页=27  运行中预留=20
  decode  uids=[0, 1, 2, 3]     等待中=1 空闲页=23  运行中预留=16
  decode  uids=[0, 1, 2, 3]     等待中=1 空闲页=19  运行中预留=12
  decode  uids=[0, 1, 2, 3]     等待中=1 空闲页=15  运行中预留=8
  decode  uids=[0, 1, 2, 3]     等待中=1 空闲页=11  运行中预留=4
  decode  uids=[0, 1, 2, 3]     等待中=1 空闲页=7   运行中预留=0
  prefill uids=[4]              等待中=0 空闲页=59  运行中预留=7
  decode  uids=[4]              等待中=0 空闲页=58  运行中预留=6
  decode  uids=[4]              等待中=0 空闲页=57  运行中预留=5
  decode  uids=[4]              等待中=0 空闲页=56  运行中预留=4
  decode  uids=[4]              等待中=0 空闲页=55  运行中预留=3
  decode  uids=[4]              等待中=0 空闲页=54  运行中预留=2
  decode  uids=[4]              等待中=0 空闲页=53  运行中预留=1
  decode  uids=[4]              等待中=0 空闲页=52  运行中预留=0
全部完成，空闲页: 64 / 64，完整性检查通过

前 4 个请求最坏情况共需 13 + 13 + 12 + 23 = 61 个 token，放得下；第 5 个还需要 13 个，超过 64，只能等待。可以看到直到前 4 个请求全部结束，池里始终有空闲页（最后还剩 7 页）——这就是保守预留的代价。前 4 个结束后空间归还，第 5 个才被接纳。全部完成后 64 页全部回到空闲列表，完整性检查通过。

把 KV 交给前缀缓存

请求的 KV 什么时候、以什么方式交给前缀缓存，是 CacheManager 里最精细的一段逻辑：

minisgl/scheduler/cache.py · CacheManager.cache_req

def cache_req(self, req: Req, *, finished: bool) -> None:
    """把请求已经算好的 KV 交给前缀缓存。

    [0, old.cached_len)              prefill 之前就在缓存里（本请求锁着它）
    [old.cached_len, cached_len)     本请求自己算的，但期间别的请求也把同样的前缀插进了缓存
                                     —— 缓存里已有一份，本请求这份是重复的，释放
    [cached_len, new.cached_len)     本次新插入缓存的部分，所有权转交给缓存
    [new.cached_len, req.cached_len) 不足一页、插不进缓存的尾巴：请求结束就释放，否则留着
    """
    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)  # 所有依赖旧句柄的操作都做完了才能解锁
    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)

它在两个时刻被调用：prefill 刚结束时（finished=False，提示词的 KV 进入缓存，请求继续运行），和请求结束时（finished=True）。注释里把请求的 KV 分成了四段：

cache_req 把请求的 KV 分成四段处理

中间那段"自己算的、但缓存里已经有了"是怎么来的？两个前缀相同的请求同时到达、同时 prefill，谁都没有命中缓存；先结束的那个把前缀插进了缓存，后结束的那个再插入时发现已经存在——它自己算的那份就是多余的，必须释放，否则就泄漏了。naive 缓存下 insert_prefix 什么都不保留，于是请求的页在结束时全部释放。第 9 章实现 Radix Cache 后，这四段都会真实出现。

完整性检查

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)

调度器空闲时（run_when_idle）做一次检查：此时没有运行中的请求，所以每一页要么空闲、要么在缓存里，两者之和必须等于总页数。任何一处忘了释放（或者重复释放），都会在这里暴露出来。

官方实现

• scheduler/cache.py · CacheManager（第 15～124 行）
• 准入控制：scheduler/prefill.py · PrefillAdder._try_allocate_one（第 39～63 行）
• 预留：scheduler/decode.py · DecodeManager.inflight_tokens（第 27～30 行）

官方仓库的 tests/core/test_cache_allocate.py 专门测试了 page size 大于 1 时"先淘汰、再分配"得到的页是否按页对齐、互不重叠。

测试

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 个 token 需要 2 页
    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：仍在第 2 页内，不需要新页
    cm.allocate_paged([req])
    assert len(cm.free_slots) == 6
    req.complete_one()
    req.complete_one()  # device_len 9：需要第 3 页
    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)  # 分配了却没有交给任何人：泄漏
    with pytest.raises(RuntimeError, match="integrity"):
        cm.check_integrity()

怎么讲清楚

讲 KV 管理：空闲页列表按页存起始位置，每轮只为 [ceil(cached_len / 页大小), ceil(device_len / 页大小)) 这些页分配，decode 每"页大小"轮才分配一次；不够时让前缀缓存淘汰，所以可用空间要把缓存里可淘汰的部分算进去；释放批量延迟做（lazy free）。准入控制按最坏情况（剩余提示词 + max_tokens）预留，所以永远不需要抢占，代价是并发比 vLLM 低（vLLM 乐观接纳、不够时抢占重算）。匹配前缀只匹配到提示词的倒数第二个 token，保证至少算一个 token 拿到 logits。请求结束时 KV 分段交给前缀缓存，每一页有且只有一个归属，空闲时的完整性检查能抓住泄漏。

练习

1. inflight_tokens 给每个运行中的请求多预留 page_size - 1 个 token。举一个例子说明，如果不预留，page size 为 16 时会出什么问题。
2. 实现一个"乐观"的准入策略：只预留 max_tokens 的一半，KV 不够时抢占最后接纳的请求（释放它的资源、放回等待队列队首）。需要改哪些地方？被抢占的请求重新 prefill 时，前缀缓存能帮上什么忙？
3. 如果 cache_req 忘了释放 [old.cached_len, cached_len) 这一段，会在哪里被发现？

参考答案

1. 请求的 KV 按页分配：remain_len 只按 token 计算，但下一次分配是一整页。例如某请求 device_len=17、remain_len=1，需要的却是一个完整的新页（16 个 token）。多个这样的请求加起来，按 token 估算会低估实际需要的页数，分配时可能失败。
2. 在 _prepare_batch 的 allocate_paged 之前检查空闲页，不够时从 decode 集合中挑一个请求：table_manager.free、cache_manager.cache_req(finished=True)，把它包装回 PendingReq 放到队首。它已经生成的 token 要拼进 input_ids。因为 cache_req 把它的 KV 交给了 Radix Cache，重新 prefill 时大部分前缀会直接命中（除非已被淘汰），重算的代价很小——这正是 SGLang 正式版的做法。
3. 这些页既不在空闲列表里，也不属于前缀缓存（缓存里存的是另一份相同内容的页），请求结束后就没有任何人持有它们。调度器下一次空闲时的 check_integrity 会发现"空闲页 + 缓存页 ≠ 总页数"并报错。第 9 章的端到端测试在两个相同前缀的请求同时 prefill 的场景下会覆盖到这条路径。

小结

• free_slots 按页存起始位置；分配从前面切，不够时让前缀缓存淘汰；释放在 lazy_free_region 里批量完成。
• 每轮只为 [ceil(cached_len/ps), ceil(device_len/ps)) 这些页分配，decode 每 page size 轮才分配一次。
• 准入控制按最坏情况（剩余提示词 + max_tokens）预留，因此永不抢占；代价是并发较低。
• cache_req 把请求的 KV 分成四段处理，保证每一页都有且只有一个归属；空闲时的完整性检查能抓住任何泄漏。

本章练习题（1 道）

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

• 192. cache_req：把请求的 KV 交给前缀缓存（中等）

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