S06显存与 KV 缓存·283 行

PagedAttention

为每个请求分配连续 KV 缓冲区，会因碎片浪费掉 60–80% 的 KV 缓存显存。把缓存分页成固定大小的 block，几乎能全部收回。

block table
块管理器
内部碎片
写时复制
gather kernel

为什么重要

问题

S05 里那个连续缓存有一个致命缺陷，缺陷出在分配器上，而不在 kernel 上。要给一个请求一块连续缓冲区，你必须在准入时就决定它的大小，而那时你还不知道这次生成会有多长。于是只能按最大值预留。

在真实流量上测量，这会产生三种截然不同的浪费：

内部碎片：每个请求「已预留但从未写入」的那段尾巴。一个按 2048 token 预留、却在 40 个 token 处停下的请求，浪费掉其中 98%。
外部碎片：空闲 block 有的是，但没有相邻的空闲 block；总空闲量绰绰有余，新请求依然放不进去。
无法共享的副本：两个使用同一段 2000 token system prompt 的请求会存下两份一模一样的拷贝，因为各自拥有私有缓冲区。

vLLM 那篇论文测得当时的系统只有 20–40% 的 KV 显存装着有用的 token，其余 60–80% 都是碎片。

核心思路

解法

操作系统在 1960 年代就解决了这个问题。别再给进程连续的物理内存；给它连续的虚拟内存，背后由分散的物理页支撑，再用一张页表做翻译。

PagedAttention 就是把这个想法用到 KV 缓存上。把缓存切成固定大小的 block，vLLM 默认每块 16 个 token；再给每个请求一张 block table，把它的逻辑 block i 映射到任何一个恰好空闲的物理 block，随着请求增长一次分配一块。

block 管理器python

class BlockManager:
    def __init__(self, num_blocks, block_size):
        self.block_size = block_size
        self.free = list(range(num_blocks))     # a free list. That is all.
        self.tables: dict[int, list[int]] = {}  # request id -> physical blocks

    def append_token(self, req_id, n_tokens_now) -> bool:
        table = self.tables.setdefault(req_id, [])
        needed = ceil(n_tokens_now / self.block_size)
        while len(table) < needed:
            if not self.free:
                return False                    # out of memory -> preempt (S10)
            table.append(self.free.pop())       # ANY free block will do
        return True

    def free_request(self, req_id):
        self.free.extend(self.tables.pop(req_id, []))

十五行就是整个分配器。难的是另一半：注意力再也不能假设它的 key 和 value 住在一块平坦的数组里。

为什么这需要一个定制 CUDA kernel

普通的注意力 kernel 用 stride 遍历内存。分页 kernel 每读一块之前，都要先去 block table 里查它的物理地址。这层间接寻址必须发生在 kernel 内部：事先把这些 block 收集成一块连续缓冲区，付出的恰好是你想避免的那次拷贝。所以 PagedAttention 既是一个内存管理想法，也是一篇 kernel 论文。

工作原理

示意图连续分配 vs 分页分配

同样的十个 token。上面它们占了 36 个预留槽位，还挡住了第四个请求；下面它们占 10 块，剩下 14 块空闲。

如何选择 block 大小

block 大小是唯一真正的调优旋钮，它在两种代价之间做交换：

更大的 block → block table 更短、每块的间接寻址更少、kernel 里的访存合并更好。代价是每个请求最后那个不满的 block 浪费更多，前缀共享的粒度也更粗。
更小的 block → 浪费接近于零、共享粒度更细，代价是表更长，每个 token 的 kernel 开销更大。

期望的内部浪费是每请求 block_size / 2 个 token，与请求长度无关。于是浪费不再随上下文增长。

vLLM 默认 block 大小: 16; 平均每请求浪费 8 个 token
典型利用率: >96%; 而连续预留只有 20–40%
吞吐提升: 2–4×; 完全来自同一张 GPU 上能塞下更多请求

写时复制，白送的

因为 block 是通过表来引用的，两个请求可以指向同一个物理 block。加上引用计数，你就得到了写时复制：同一 prompt 的并行采样（n=4）或 beam search 的分支，会共享公共前缀的每一个 block，只有在分叉处才复制那一块。

把一段 2000 token 的 prompt 在四个采样间共享，会把 8000 token 的缓存变成大约 2000。下一章我们把这个思路再推一步：在不同请求之间共享 block。

动手观察

动手试试

一个真实的 block 管理器，跑在贴近现实的负载上：大多数是短生成，夹杂少数长生成。先用分页模式跑一遍，再把分配器切成连续模式跑一遍。看看在完全相同的显存下，两种模式各能同时维持多少个请求。

模拟器block 管理器先跑分页，再切到连续

分配器

block 大小（token）4负载随机种子5

第 0 / 90 步

运行中: 14
已完成: 0
block: 58/64
浪费 / 抢占次数: 17% · 0

物理 KV block 池 · 64 块 × 4 token

空闲
被真实 token 占用
已预留但为空（浪费）

block 按需一块一块地从池中任意位置发放。一个请求的 block 不必相邻；下面的 block table 会让注意力 kernel 看起来它们是连续的。浪费被控制在每请求不到一个 block。

block table（逻辑 → 物理）

r017 token01234
r116 token5678
r218 token910111213
r38 token1415
r417 token1617181920
r510 token212223

这张表就是全部的间接寻址。注意力 kernel 沿着它找到逻辑 block i 实际所在的位置，和 CPU 页表把虚拟页映射到物理帧是一回事。

内部碎片

有效 token: 213
已用 block 中的浪费: 8%

把 block 大小调到 16，看着浪费上升：每个请求的最后一个 block 都只填了一部分，所以期望浪费是每请求 blockSize/2 个 token。调到 1 浪费就消失了，但 block table 会长 16 倍，kernel 的间接寻址也多 16 倍。vLLM 默认取 16，正是这两条曲线的交点。

然后把 block 大小推到 16，再拉回 1。16 时浪费条明显变长；1 时浪费消失，每张 block table 却长了十六倍。默认值买到的，是这条曲线上一个不错的位置。

亲手实现

实现

下面是让分页注意力得以成立的那一步 gather。在 NumPy 里它是一次索引操作；在 CUDA 里它是 kernel 的内层循环。

code/s06_paged_attention.py（节选）python

def paged_attention(q, block_table, k_cache, v_cache, seq_len, block_size):
    """
    k_cache: [num_blocks, block_size, n_kv_heads, head_dim]  — the pool
    block_table: [num_logical_blocks]                        — this request's map
    """
    out_scores = []
    for logical, physical in enumerate(block_table):
        start = logical * block_size
        n = min(block_size, seq_len - start)
        if n <= 0:
            break
        k_blk = k_cache[physical, :n]        # the indirection, in one line
        out_scores.append(q @ k_blk.transpose(0, 2, 1) / sqrt(head_dim))

    scores = np.concatenate(out_scores, axis=-1)   # looks contiguous again
    weights = softmax(scores)
    return gather_values(weights, block_table, v_cache, seq_len, block_size)

引用计数必须精确

释放了一个仍被别的请求指向的 block，你会得到无声的数据损坏：注意力读到别人的 key，模型产出流畅的胡话，而任何地方都不会报错。释放得过于保守，则会一直泄漏直到池子耗尽。要直接测试分配器：每次操作之后都断言：持有每个物理 block 的 block table 数量，恰好等于它的引用计数所声称的数目。

在本地运行

实现一个 block 管理器和一个分页注意力函数，断言分页注意力与连续注意力数值完全一致，然后在两种分配器下回放同一份合成负载，报告利用率、浪费，以及各自能同时驻留多少个请求。

$ python code/s06_paged_attention.py

预期输出：一项数值等价断言、一次利用率对比（连续约 26% vs 分页约 52%），以及一次 block 大小扫描，显示浪费确实跟随 block_size/2。

只需要 NumPy — 查看环境准备.

生产实践

在生产环境中

vLLM —— 最初的实现。v1 引擎还支持一条「cascade」路径，把共享前缀与分叉后缀分开批处理。
SGLang —— 用 RadixAttention 树取代了扁平的 block table，把分页与前缀共享统一进同一个结构（S07）。
baseRT —— 在宣传里把「分页 KV 缓存」和连续批处理并列；两者本就配套，因为正是分页让可变大小的 batch 变得负担得起。
FlashInfer —— SGLang 等引擎做分页注意力所依赖的注意力 kernel 库，让你不必自己写 CUDA 也能拿到这层间接寻址。

练习

1
加上引用计数，为 n>1 采样实现写时复制。验证从一段 1000 token 的 prompt 采样四份，用掉的缓存约为 1000 个 block 而不是 4000。
2
把 block 大小从 1 扫到 32，画出利用率与 block table 长度的关系曲线。找到你自己负载下两条曲线的交点，再和 vLLM 的默认值 16 比较。
3
让分配器失败：把它逼到耗尽，再实现两种恢复方案，从头重算该请求，或把它的 block 换出到主机内存。在不同上下文长度下测量哪个更便宜；这正是 S10 的抢占决策。

继续学习

接下来

block 是可共享的，但到目前为止我们只把它用在同一个请求的分支上。一个对话服务每小时会把同一段 system prompt 发送成千上万次。S07 会对 block 内容做哈希，让任意两个拥有公共前缀的请求共享同一批物理 block。prefill 成本常常因此下降一个数量级。

自测

习题

先作答，再看解析。答错比答对更有价值，因为解析会指出你该回头重读哪一部分。

自测第 1 题 / 共 6

block 大小为 16 时，每个请求的期望内部碎片是多少？它如何随上下文长度增长？

得分 0/6