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 对齐。