LLM serving 中,请求不是整齐的矩阵。 每个用户 prompt 长度不同,decode 输出长度也不同,请求会随时进来、结束、被抢占或扩展。 如果给每个请求预留一大段连续 KV Cache,就会出现内部碎片、外部碎片和重复复制。 当显存被碎片吃掉,系统能放进 batch 的请求数下降,吞吐也随之下降。

显存估算脚本

def kv_bytes_per_token(layers, kv_heads, head_dim, bytes_per_elem=2):
    return 2 * layers * kv_heads * head_dim * bytes_per_elem

def gib(x):
    return x / 2**30

layers = 32
kv_heads = 32
head_dim = 128
tokens = 128 * 4096  # 128 个请求,每个 4K context
total = kv_bytes_per_token(layers, kv_heads, head_dim) * tokens
print(f"KV cache ~= {gib(total):.1f} GiB")

PagedAttention 借鉴操作系统分页: 请求看到的是连续的逻辑 token 序列,但底层 KV Cache 可以存放在不连续的物理 blocks 中。 每个请求维护自己的 block table,把第 0 个逻辑块、第 1 个逻辑块映射到具体物理块。 attention kernel 读取 KV 时不再假设 KV 是一段连续内存,而是按 block table 找到每个 token 所在的物理位置。

设 block size 是 16 tokens。请求当前长度 50,它需要 4 个逻辑 block: 前三个满块各 16 tokens,最后一个只用了 2 tokens。 当 decode 继续生成 token 时,系统先写入最后一个未满块; 如果最后一个块满了,再从 free block pool 分配一个新的物理 block, 把它追加到 block table 里。

toy block allocator

class KVBlockPool:
    def __init__(self, num_blocks):
        self.free = list(range(num_blocks))
        self.ref = [0] * num_blocks

    def alloc(self):
        if not self.free:
            raise RuntimeError("KV cache exhausted")
        block = self.free.pop()
        self.ref[block] = 1
        return block

    def incref(self, block):
        self.ref[block] += 1

    def decref(self, block):
        self.ref[block] -= 1
        if self.ref[block] == 0:
            self.free.append(block)

class SequenceKV:
    def __init__(self, pool, block_size):
        self.pool = pool
        self.block_size = block_size
        self.block_table = []
        self.length = 0

    def append_token(self):
        if self.length % self.block_size == 0:
            self.block_table.append(self.pool.alloc())
        self.length += 1

并行采样和 beam search 会从同一个 prompt 分叉出多个候选。 分叉时,这些候选的 prompt KV 完全相同。 PagedAttention 不需要为每个候选复制一份物理 KV blocks, 而是让多个 block table 指向同一批 physical blocks,并增加 ref count。 当某个候选要写入一个仍被共享的 block,才触发 copy-on-write:分配新块、复制旧块内容、更新该候选的 block table。

toy COW append

def fork_sequence(src, pool):
    dst = SequenceKV(pool, src.block_size)
    dst.block_table = list(src.block_table)
    dst.length = src.length
    for block in dst.block_table:
        pool.incref(block)
    return dst

def ensure_private_last_block(seq):
    block = seq.block_table[-1]
    if seq.pool.ref[block] == 1:
        return
    new_block = seq.pool.alloc()
    seq.pool.decref(block)
    seq.block_table[-1] = new_block
    # 真实系统还要把旧 block 的有效 KV bytes 复制到 new_block

def append_with_cow(seq):
    if seq.length % seq.block_size == 0:
        seq.block_table.append(seq.pool.alloc())
    else:
        ensure_private_last_block(seq)
    seq.length += 1

PagedAttention 是 vLLM 的底层能力之一。 它让 serving engine 能更灵活地管理 KV Cache,从而容纳更多并发请求。 但真正的高吞吐还来自调度:请求进入等待队列,调度器根据剩余 KV blocks、 prefill/decode 预算和正在运行的序列,决定这一轮 batch 放哪些请求。

今天不用真的部署 7B 模型,先用一个小脚本把“请求长度 → logical blocks → physical blocks” 的关系跑清楚。你要能看到:短请求只浪费最后一个 block 的尾部, 多个候选共享前缀时,物理 block 数不会线性翻倍。

import math

def blocks_needed(tokens, block_size):
    return math.ceil(tokens / block_size)

def allocate_requests(lengths, block_size):
    next_block = 0
    tables = {}
    for req_id, length in enumerate(lengths):
        n = blocks_needed(length, block_size)
        tables[req_id] = list(range(next_block, next_block + n))
        next_block += n
    return tables, next_block

lengths = [7, 16, 31, 48, 63]
block_size = 16
tables, physical_blocks = allocate_requests(lengths, block_size)

for req, table in tables.items():
    print(f"request={req} len={lengths[req]} block_table={table}")
print("physical_blocks", physical_blocks)

waste_tokens = sum(blocks_needed(x, block_size) * block_size - x for x in lengths)
print("tail_waste_tokens", waste_tokens)

1. KV Cache 的显存大小和 layers、kv_heads、head_dim、context length 分别是什么关系?
2. 为什么传统连续 KV 分配会在 LLM serving 里产生严重碎片?
3. block table 中 logical block 和 physical block 的区别是什么?
4. copy-on-write 为什么能降低 parallel sampling 和 beam search 的 KV 内存占用?
5. PagedAttention 解决的是 attention 算法复杂度、kernel IO,还是 serving 内存管理?