推理服务 · 内存

PagedAttention

分块存放的 KV 缓存，与 vLLM 相同

块
这里 4 个词元 vLLM 16
内存
256 个词元槽位
请求
16 个，最长 64 个词元
使用率，连续分配
35%
使用率，分页
94%

全部步骤

  1. 按最坏情况预留

    由于不知道每个回答会有多长，简单的服务器会在一段连续内存中按最长可能的长度预留空间。大部分从来用不上，而新请求却在内存闲置时等待；vLLM 测得这类系统中只有 20 到 38% 的 KV 内存真正存了词元。

    每个请求按最大长度预留

  2. 按需分块

    PagedAttention 把缓存切成固定大小的块（vLLM 中是 16 个词元，这里是 4 个），只在请求的最后一个块满了时，才从内存中任意位置给它一个新块。每个请求有一张块表，按顺序列出它的块，就像虚拟内存的页表。

    逻辑块 → 物理块

  3. 跨块的注意力

    注意力内核跟着块表走：对请求的每个块，从它所在的位置取出键和值，然后做和之前一样的注意力计算。只是内存布局变了。

    对每个块：K, V ← memory[table[b]]

  4. 共享与写时复制

    同一个提示词的两个样本可以共享提示词的块；引用计数记录每个块被多少条序列使用。当一个样本要写入一个部分填充的共享块时，它先复制一份（写时复制），就像操作系统里 fork 出的进程一样。

    共享块 · 引用计数 · 写时复制

  5. 同时容纳更多请求

    同样的 16 个请求，同样的内存。分页让批次同时容纳更多请求，所以整个队列更早完成；论文中，在相同延迟下 vLLM 每秒服务的请求数是早期系统的 2 到 4 倍。

    每批更多序列 → 每秒更多词元

代码

# the KV cache as a pool of fixed-size blocks
free = list(range(num_blocks)); block_table = {seq: [] for seq in running}
# before writing token t of a sequence, make sure its block exists
if t % BLOCK == 0: block_table[seq].append(free.pop())
b = block_table[seq][t // BLOCK]; kv[b, t % BLOCK] = (k, v)
# attention gathers the sequence's blocks through its table
K = torch.cat([kv_k[b] for b in block_table[seq]])
# sharing: forked sequences copy the table, not the blocks
block_table[child] = list(block_table[parent]); ref[b] += 1 for each b
if ref[b] > 1 and writing: b = copy_block(b); ref[b_old] -= 1

延伸阅读