手写 mini-sglang 阶段一 · 算得对：把骨架换成正式结构

KV 池、page table 与 token pool

引擎启动时一次性分配三块内存：存放所有请求 K、V 的 KV 池，记录"每个请求的每个 token 的 KV 存在池中哪个位置"的 page table，以及记录"每个请求的每个位置是哪个 token"的 token pool。后面的注意力后端、调度器、重叠调度都围绕这三块内存工作。这一章把它们建出来，并讲清 mini-sglang 一个独特的设计：page table 永远按 token 存位置，与 page size 无关。

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

1. Qwen3-0.6B 在 bf16 下，每个 token 的 KV 占多少字节？40 GiB 能放多少个 token？
2. 分页（page size > 1）解决了什么问题？page size 越大越好吗？
3. 为什么模型的输入 token 也要在 GPU 上保存一份（token pool），而不是每轮从 CPU 传过去？
4. KV 池为什么要多分配一页？

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

1. 2 × 28 层 × 8 个 KV 头 × 128 维 × 2 字节 = 114688 字节（112 KiB）；40 GiB 能放约 37.4 万个 token（40 × 1024 × 1024 ÷ 112）。
2. 以页为单位分配和管理，块表更短、分配释放的次数更少，也方便按页共享前缀、按页传输。page size 不是越大越好：最后一页的内部碎片更多，前缀缓存只能缓存完整的页、命中粒度更粗。
3. 采样出的 token 直接写回 GPU 上的 token pool，下一轮的输入直接从那里取，CPU 不在关键路径上——这也是重叠调度的前提。
4. 给 CUDA Graph 补齐 batch 用的 dummy 请求：它也会写 KV，多出来的最后一页专门给它写，不会踩到真实请求的数据。

本章要写的文件：kvcache/base.py、kvcache/mha_pool.py、kvcache/naive_cache.py、kvcache/__init__.py、scheduler/table.py。

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

minisgl/
├── __init__.py
├── core.py
├── distributed/  （3 个文件，不变）
├── kernel/  （2 个文件，不变）
├── kvcache/
│   ├── __init__.py  ← 本章新建
│   ├── base.py  ← 本章新建
│   ├── mha_pool.py  ← 本章新建
│   ├── naive_cache.py  ← 本章新建
├── layers/  （8 个文件，不变）
├── models/  （7 个文件，不变）
├── scheduler/
│   ├── table.py  ← 本章新建
│   ├── __init__.py  ← 提前引入（第 7 章讲）
├── utils/  （6 个文件，不变）

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

KV 池：一整块预分配的显存

推理时每个请求的 KV 缓存随生成不断增长，而且请求来来去去。如果每个请求各自申请显存，很快就会碎片化。所有现代推理引擎的做法都一样：启动时把剩余显存一次性分配成一个大池子，按固定大小的"页"管理，谁需要就分给谁几页，用完还回来。这就是 PagedAttention 的思想（原理见推理系统手册的分页 KV Cache）。

minisgl/kvcache/mha_pool.py · MHAKVCache

class MHAKVCache(BaseKVCachePool):
    """一整块预分配的 KV 显存：[2, 层数, 页数, 页大小, 本 rank 的 KV 头数, head_dim]。"""

    def __init__(self, num_kv_heads: int, num_layers: int, head_dim: int, num_pages: int,
                 page_size: int, dtype: torch.dtype, device: torch.device) -> None:
        local_kv_heads = div_even(num_kv_heads, get_tp_info().size, allow_replicate=True)
        self._kv_buffer = torch.empty(
            (2, num_layers, num_pages, page_size, local_kv_heads, head_dim),
            device=device, dtype=dtype,
        )
        self._k_buffer = self._kv_buffer[0]
        self._v_buffer = self._kv_buffer[1]
        self._num_layers = num_layers
        self._device = torch.device(device)
        # 按 token 展平后的形状：第 i 行就是 KV 池里第 i 个 token 的位置
        self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim)

    def k_cache(self, index: int) -> torch.Tensor:
        return self._k_buffer[index]

    def v_cache(self, index: int) -> torch.Tensor:
        return self._v_buffer[index]

    def store_kv(self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor,
                 layer_id: int) -> None:
        store_cache(
            k_cache=self._k_buffer[layer_id].view(self._storage_shape),
            v_cache=self._v_buffer[layer_id].view(self._storage_shape),
            indices=out_loc, k=k, v=v,
        )

    @property
    def device(self) -> torch.device:
        return self._device

    @property
    def dtype(self) -> torch.dtype:
        return self._kv_buffer.dtype

    @property
    def num_layers(self) -> int:
        return self._num_layers

布局是 [2, 层数, 页数, 页大小, 本 rank 的 KV 头数, head_dim]：

• 第 0 维区分 K 和 V，k_cache(i)、v_cache(i) 取出第 i 层的 [页数, 页大小, 头数, head_dim]；
• 把"页数 × 页大小"展平后，第 j 行就是池中第 j 个 token 的位置，这就是 _storage_shape；
• store_kv 把本轮新算出的 K、V 按 out_loc（每个新 token 的池内位置）写进去。它由注意力后端在计算注意力之前调用。

minisgl/kernel/torch_ops.py · store_cache

def store_cache(
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    indices: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
) -> None:
    """把本轮新算出的 K、V 按 indices 写进 KV 池。k_cache: [num_tokens, H, D]；k: [n, H*D] 或 [n, H, D]。"""
    idx = indices.long()
    k_cache[idx] = k.view(k.shape[0], *k_cache.shape[1:])
    v_cache[idx] = v.view(v.shape[0], *v_cache.shape[1:])

这是个纯粹的"按下标搬运"，GPU 上由第 19 章的自定义 kernel 完成：一个 warp 负责一个 token，每次搬 16 字节。

examples/ch04_kvcache.py

"""第 4 章：KV 池的布局、容量估算，以及 page table / token pool 的样子。"""

import torch
from minisgl.distributed import set_tp_info
from minisgl.kvcache.mha_pool import MHAKVCache
from minisgl.models import ModelConfig
from minisgl.utils import cached_load_hf_config

set_tp_info(0, 1)
c = ModelConfig.from_hf(cached_load_hf_config("models/Qwen3-0.6B"))
per_token = 2 * c.num_layers * c.num_kv_heads * c.head_dim * 2  # K+V，bf16
print(f"Qwen3-0.6B 每个 token 的 KV：{per_token} 字节 = {per_token / 1024:.0f} KiB")
for gib in (4, 40):
    print(f"  {gib:>2} GiB 显存能放 {gib * 2**30 // per_token:>8} 个 token")

pool = MHAKVCache(num_kv_heads=c.num_kv_heads, num_layers=c.num_layers, head_dim=c.head_dim,
                  num_pages=64 + 1, page_size=16, dtype=torch.bfloat16, device=torch.device("cpu"))
print("KV 池形状 [K/V, 层, 页, 页内位置, 头, 维度]:", tuple(pool._kv_buffer.shape))
print("第 3 层的 K:", tuple(pool.k_cache(3).shape), " 按 token 展平:", pool._storage_shape)

# 两个请求的 page table（page_size = 16，按 token 存位置）
page_table = torch.zeros((3, 64), dtype=torch.int32)
page_table[0, :20] = torch.cat([torch.arange(32, 48), torch.arange(80, 84)])  # 第 2 页 + 第 5 页的前 4 个
page_table[1, :5] = torch.arange(160, 165)  # 第 10 页的前 5 个
print("请求 0 的前 20 个位置:", page_table[0, :20].tolist())
print("请求 1 的前 5 个位置: ", page_table[1, :5].tolist())

运行结果

Qwen3-0.6B 每个 token 的 KV：114688 字节 = 112 KiB
   4 GiB 显存能放    37449 个 token
  40 GiB 显存能放   374491 个 token
KV 池形状 [K/V, 层, 页, 页内位置, 头, 维度]: (2, 28, 65, 16, 8, 128)
第 3 层的 K: (65, 16, 8, 128)  按 token 展平: (1040, 8, 128)
请求 0 的前 20 个位置: [32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 80, 81, 82, 83]
请求 1 的前 5 个位置:  [160, 161, 162, 163, 164]

Qwen3-0.6B 用了 GQA（16 个 q 头、8 个 KV 头），每个 token 的 KV 是 112 KiB。模型权重只有约 1.2 GB（bf16），一张 80 GB 的卡除去权重和激活，几乎全部显存都给了 KV 池——能同时服务多少请求、每个请求能多长，都由这个池子的大小决定。

为什么多一页？ 引擎为 CUDA Graph 准备了一个 dummy 请求（第 18 章），它补齐 batch 时也会写 KV。多出来的最后一页专门给它写，不会踩到真实请求的数据。

page table：按 token 存位置

KV 池、page table 与 token pool：两个请求的 KV 散落在不同的页里

page table 是一个二维 int32 张量，形状 [max_running_req + 1, max_seq_len 向上对齐到 32]，每行对应一个运行中的请求（最后一行给 dummy 请求），page_table[r, j] 是这个请求第 j 个 token 的 KV 在池中的位置。

mini-sglang 的一个特别之处：无论 page size 是多少，page table 都按 token 存位置，而不是按页存页号。上面的例子里 page size 是 16，请求 0 占了第 2 页和第 5 页，它的 page table 行是 [32, 33, …, 47, 80, 81, 82, 83]，而不是 [2, 5]。

这样做的好处是简单：

• 写 KV 时，out_loc = page_table[行, cached_len:device_len]，直接得到每个新 token 的位置，不用做"页号 × 页大小 + 页内偏移"的换算；
• FlashInfer 后端把 KV 池当成 page size = 1 来用，page table 的一行直接就是它需要的 indices；
• FlashAttention 需要页号，就每隔 page size 取一个位置再除以 page size：page_table[r, ::page_size] // page_size（第 17 章）。

代价是 page table 大了 page size 倍，但它只是 int32，最多几十 MB。对齐到 32 个 int32（128 字节）是为了让每一行的起始地址对齐，GPU 读取时效率更高。

page size 大于 1 的意义在于前缀缓存和分配的粒度：Radix Cache 只缓存完整的页（第 9 章），分配器每次分配一整页（第 8 章），页越大管理的开销越小，但最后一页平均浪费半页，前缀匹配也只能按页对齐。一些 kernel（例如 TensorRT-LLM 的注意力）只支持特定的 page size。

token pool：输入也放在设备上

minisgl/scheduler/table.py · TableManager

class TableManager:
    """管理 page table / token pool 的行：每个运行中的请求占一行。"""

    def __init__(self, max_running_reqs: int, page_table: torch.Tensor) -> None:
        self._free_slots = list(range(max_running_reqs))
        self.page_table = page_table
        # token_pool[i, j]：第 i 行请求第 j 个位置的 token id。它和 page_table 形状相同、放在设备上，
        # 这样模型输入可以直接在设备上用下标取出，采样结果也可以直接写回（重叠调度的关键，见第 11 章）。
        # dummy 请求也从这里取输入，所以要初始化成合法的 token id（0）。
        self.token_pool = torch.zeros_like(page_table, dtype=torch.int32)

    @property
    def available_size(self) -> int:
        return len(self._free_slots)

    def allocate(self) -> int:
        return self._free_slots.pop()

    def free(self, slot: int) -> None:
        self._free_slots.append(slot)

TableManager 管理 page table 的行：请求被接纳时分配一行，结束时归还。它还持有 token pool：与 page table 形状相同的 int32 张量，token_pool[r, j] 是这个请求第 j 个位置的 token id。

为什么要在设备上再存一份 token？因为本轮的输入是上一轮采样的输出。引擎采样得到 next_tokens 后，调度器直接在设备上执行 token_pool[行, device_len] = next_tokens；下一轮组 batch 时，input_ids = token_pool[行, 位置] 也在设备上完成。整个过程 CPU 不需要知道新 token 是什么。这正是第 11 章重叠调度的前提：GPU 在算第 N+1 轮时，CPU 可能还没收到第 N 轮的结果。

dummy 请求也从 token pool 取输入，所以 token pool 要初始化为合法的 token id（0），而不是随机值。

前缀缓存的接口

KV 池只管"存"，不管"哪个位置属于谁"。决定哪些位置可以复用的是前缀缓存，接口在 kvcache/base.py：

minisgl/kvcache/base.py · BasePrefixCache

class BasePrefixCache(ABC):
    @abstractmethod
    def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
        """加锁（或解锁）一个句柄：加锁后这段前缀不会被淘汰。只改变计数，不改变缓存内容。"""

    @abstractmethod
    def match_prefix(self, input_ids: torch.Tensor) -> MatchResult:
        """查找最长的已缓存前缀。返回的句柄在加锁之前随时可能被淘汰。"""

    @abstractmethod
    def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> InsertResult:
        """把一段 token 及其 KV 位置插入缓存。"""

    @abstractmethod
    def evict(self, size: int) -> torch.Tensor:
        """淘汰至少 size 个 token，返回被释放的 KV 位置。evict(0) 总是安全的空操作。"""

    @abstractmethod
    def reset(self) -> None: ...

    @property
    @abstractmethod
    def size_info(self) -> SizeInfo: ...

    @abstractmethod
    def check_integrity(self) -> None: ...

本章先实现最简单的 NaivePrefixCache：永远不命中，插入什么都不保留（请求结束后，它的页全部释放）。

minisgl/kvcache/naive_cache.py · NaivePrefixCache

class NaivePrefixCache(BasePrefixCache):
    """不做任何前缀复用：永远不命中，插入什么都不保留（调用方随后会把页全部释放）。"""

    def __init__(self, device: torch.device):
        self.empty_tensor = torch.empty(0, dtype=torch.int32, device=device)

    def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
        pass

    def match_prefix(self, input_ids: torch.Tensor) -> MatchResult:
        return MatchResult(NaiveCacheHandle(self.empty_tensor))

    def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> InsertResult:
        return InsertResult(0, NaiveCacheHandle(self.empty_tensor))

    def evict(self, size: int) -> torch.Tensor:
        if size == 0:
            return self.empty_tensor
        raise RuntimeError("NaivePrefixCache has nothing to evict")

    def reset(self) -> None:
        pass

    @property
    def size_info(self) -> SizeInfo:
        return SizeInfo(evictable_size=0, protected_size=0)

    def check_integrity(self) -> None:
        pass

真正的 Radix Cache 在第 9 章实现。两者通过注册表 SUPPORTED_CACHE_MANAGER 按名字（naive、radix）创建。

官方实现

• KV 池：kvcache/mha_pool.py · MHAKVCache（第 10～68 行），写 KV 由自定义 CUDA kernel 完成（kernel/csrc/jit/store.cu）
• page table 的创建与对齐：engine/engine.py · Engine.__init__（第 30～110 行） 中 "Page table initialization" 一段，注释写着 "aligned to 128 bytes; store raw locations instead of pages"；core.py 的 Context 里也注明 "this table always treat page_size = 1"
• token pool：scheduler/table.py · TableManager（第 4～21 行）

测试

tests/test_ch04_kvcache.py · test_kv_pool_layout_and_store

def test_kv_pool_layout_and_store():
    set_tp_info(0, 1)
    pool = MHAKVCache(num_kv_heads=2, num_layers=3, head_dim=4, num_pages=5, page_size=2,
                      dtype=torch.float32, device=torch.device("cpu"))
    assert pool._kv_buffer.shape == (2, 3, 5, 2, 2, 4)  # K/V, 层, 页, 页内位置, 头, 维度
    k = torch.arange(3 * 8, dtype=torch.float32).view(3, 8)  # 3 个 token，每个 2 头 x 4 维
    pool.store_kv(k, -k, out_loc=torch.tensor([9, 0, 4], dtype=torch.int32), layer_id=1)
    flat_k = pool.k_cache(1).reshape(10, 2, 4)  # 按 token 展平：第 9 个 token 在第 4 页的第 2 个位置
    assert torch.equal(flat_k[9].flatten(), k[0]) and torch.equal(flat_k[4].flatten(), k[2])
    assert torch.equal(pool.v_cache(1)[4, 1].flatten(), -k[0])

第二个测试检查张量并行时的 KV 头复制：2 个 KV 头分给 4 个 rank，每个 rank 保存 1 个头（第 16 章）。

怎么讲清楚

讲 KV 池：启动时一次性分配一整块显存 [2, 层数, 页数, 页大小, KV 头数, head_dim]，按页分给请求，避免运行时分配和碎片；每个 token 的 KV = 2 × 层数 × KV 头数 × head_dim × 2 字节（Qwen3-0.6B 是 112 KB）。page table 每行对应一个请求、按 token 存 KV 的位置，注意力后端需要页号时自己换算；分页让显存按需增长、支持前缀共享，页越大元数据越少、kernel 越高效，但尾部浪费和前缀复用的粒度变差。token pool 把输入 token 也放在设备上，采样结果直接写回、下一轮直接读，CPU 不在关键路径上。多分配一页给补齐用的 dummy 请求。

练习

1. Llama-3.1-8B（32 层、8 个 KV 头、head_dim 128）在 bf16 下每个 token 的 KV 多大？80 GB 的卡放完 16 GB 权重后，最多能容纳多少个 token？
2. 如果 page table 按页存页号，写 KV 时 out_loc 要怎么算？写出 PyTorch 表达式。
3. page size 为 16 时，一个 17 个 token 的请求占几页？平均每个请求浪费多少个 token 的空间？

参考答案

1. 2 × 32 × 8 × 128 × 2 = 131072 字节（128 KiB）；约 64 GiB 可用时能放约 52 万个 token（实际还要扣掉激活和 CUDA Graph 的显存）。
2. page_ids[pos // page_size] * page_size + pos % page_size，其中 pos 是 [cached_len, device_len)。
3. 2 页；最后一页平均用一半，平均浪费约 page_size / 2 = 8 个 token。

小结

• KV 池是启动时一次性分配的大张量 [2, 层, 页, 页大小, 头, 维]，按页分给请求；多一页给 dummy 请求。
• page table 每行对应一个请求，按 token 存 KV 位置，与 page size 无关；需要页号的后端自己换算。
• token pool 在设备上保存每个请求每个位置的 token，采样结果直接写回，下一轮输入直接取出，CPU 不在关键路径上。
• 前缀缓存决定哪些位置可以复用；先用永远不命中的 naive 版本。

本章练习题（1 道）

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

• 188. 按页分配 KV 并写进 page table（中等）

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