I·Infra Notes AI Infra 工程实践指南 ← 返回主博客
模型与缓存 / READING NOTES

6.1 分页 KV 的最小状态机

KV 块分配、分页寻址与前缀复用

每层 K/V 的逻辑 shape 可写成 [request, token, kv_head, head_dim];分页后物理池按 Block 分配,block_table[request][logical_block] = physical_block。地址 physical_block * block_size + token_pos % block_size,实际内存布局还含 layer、K/V、head 和 dtype stride。

from collections import deque

class BlockPool:
    def __init__(self, nblocks, block_size):
        self.free = deque(range(nblocks))
        self.refcount = [0] * nblocks
        self.tables = {}
        self.block_size = block_size

    def reserve(self, req_id, new_total_tokens):
        table = self.tables.setdefault(req_id, [])
        need = (new_total_tokens + self.block_size - 1) // self.block_size
        missing = need - len(table)
        if missing > len(self.free):
            raise MemoryError('KV blocks exhausted')
        for _ in range(missing):
            block = self.free.popleft()
            self.refcount[block] = 1
            table.append(block)
        return table

    def release(self, req_id):
        for block in self.tables.pop(req_id, []):
            self.refcount[block] -= 1
            if self.refcount[block] == 0:
                self.free.append(block)

这段只展示块分配,不含 Prefix Cache、共享块写时复制、跨层混合缓存或异步释放。真正管理器应以引用计数维护物理块总量守恒;活跃引用块与零引用可分配块互斥,而 cached 是可与两者重叠的索引属性(详见 6.10),释放不能重复入 free list,复用时要校验模型、LoRA、token 序列、位置、KV dtype/布局等影响结果的上下文。Prefix Cache 命中后要区分完整前缀和待计算尾部;是否需要重新计算一个 token 取决于引擎如何获得采样 logits,不能一概规定“必须留一个块重算”。

源码对照:vLLM V1 KVCacheManager.get_computed_blocks / allocate_slots 先查已计算块,再判断新增槽位与可用块;当前实现对完整 Prefix 命中会限制最大命中长度为 prompt_length - 1,为取得 logits 重算末尾 token。这是该实现的明确选择;阅读其他引擎时仍要查看其 logits 取得方式。SGLang RadixCache 则从压缩前缀树的匹配、引用与淘汰入口读起。

Paged Attention Kernel:Q 由本轮计算得到,Kernel 依次从 Block Table 查 K/V 物理块,按有效长度 mask 读取,在线累加 softmax。它解决 KV 的非连续物理布局;与 FlashAttention 的 tiling 可以组合。常见 Bug:最后不满块的 mask 读入旧数据、Prefix 共享块被覆盖、抢占后表与长度不一致。调试时打印逻辑 token → (block_id, slot)、引用计数和前 3 个 K/V 值,与连续缓存 reference 逐 token 对齐。