PagedAttention

它解决的是 KV 物理管理，不是另一种 softmax

标准 KV Cache 在逻辑上是每层连续序列：

但在线服务事先不知道每条请求最终生成多长。若为最大长度预留连续区，会产生大量保留但未使用的容量；若每次增长都重新分配连续 tensor，又会复制历史数据并造成外部碎片。PagedAttention 把 token 维切成固定大小 $P$ 的逻辑 blocks，再为每条序列维护 block table：

attention kernel 根据序列 $i$、逻辑位置 $t$ 查表，读取物理 block 中相应 offset 的 K/V。逻辑位置仍按 $0,1,\ldots,T-1$ 排列，输出必须与 contiguous KV 相同。

shape、layout 与所有权

一种教学用物理布局为：

实际框架也可能把 head、block、token、vectorization 维度重排，例如 [num_blocks, num_kv_heads, head_size/x, block_size, x]，以匹配 kernel 的 coalesced load。block table 常是 device 上的整数 tensor：

Block table 不保存 token 内容；它只保存地址映射。PagedAttention 也不自动判断两个字符串是否相同：安全共享的 key 必须基于确定的 token ids、模型/adapter 状态、位置语义和已计算 KV 身份。

图中每个 physical block 同时代表 K 和 V 的 payload，省略了 layer 维、batch slot 重排、引用计数、free list、GPU/CPU swap、beam copy-on-write 和 kernel 内的向量化 layout。箭头表示查表关系，不表示运行时一定先 gather 成一个连续 tensor。

worked example：长度 10 与 7，block size 4

序列 A 需要 3 个逻辑 blocks：[0..3]、[4..7]、[8..9 + 2 empty]；序列 B 需要 2 个：[0..3]、[4..6 + 1 empty]。若前 4 个 token 的 KV 完全相同，可以共享一个 physical block：

A table: L0→P5, L1→P1, L2→P6
B table: L0→P5, L1→P3

唯一 live token payload 为 $4+(10-4)+(7-4)=13$ 个 positions。4 个唯一 physical blocks 提供 $4\times4=16$ 个 slots，因此最后 blocks 的内部碎片为 3 slots。若不共享 prefix，需要 5 个 blocks，即 20 slots；本例共享节省 4 slots。

一般地，一条非空序列的最后 block 最多浪费 $P-1$ 个 positions。对 $B$ 条互不共享序列：

更小的 $P$ 降低最尾部内部碎片，却增加 block table、分配操作和 kernel address translation 开销；更大的 $P$ 相反。最佳值依赖 cache layout、kernel、请求长度分布和 allocator，不是越小越好。

contiguous 与 page-table reference 对齐

实验问题：物理 blocks 非连续且 prefix 共享时，按 block table 恢复的 K/V 与 contiguous reference 是否一致？attention 输出是否相同？内部碎片与共享节省能否精确计数？

import math, platform
import torch

torch.manual_seed(53)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float32

def expand_kv(x,q_heads):
    return x.repeat_interleave(q_heads//x.shape[0],dim=0)

def attend_one(q,k,v):
    # q:[Hq,1,D], k/v:[Hkv,S,D]
    scores=q@expand_kv(k,q.shape[0]).transpose(-2,-1)/math.sqrt(q.shape[-1])
    return torch.softmax(scores,-1)@expand_kv(v,q.shape[0])

def gather_paged(cache,block_table,length,block_size):
    positions=torch.arange(length,device=cache.device)
    logical=positions//block_size; offsets=positions%block_size
    physical=block_table[logical]
    assert torch.all(physical>=0)
    gathered=cache[physical,:,offsets,:]  # [S,Hkv,D]
    return gathered.permute(1,0,2).contiguous()

def main():
    hq,hkv,d,p=4,2,8,4; lengths=(10,7)
    prefix=torch.randn(hkv,4,d,device=DEVICE,dtype=DTYPE)
    contiguous_k=[torch.cat((prefix,torch.randn(hkv,n-4,d,device=DEVICE)),1)
                  for n in lengths]
    contiguous_v=[torch.cat((prefix*.5,torch.randn(hkv,n-4,d,device=DEVICE)),1)
                  for n in lengths]
    tables=torch.tensor([[5,1,6],[5,3,-1]],device=DEVICE)
    physical_k=torch.zeros(7,hkv,p,d,device=DEVICE,dtype=DTYPE)
    physical_v=torch.zeros_like(physical_k)
    for seq,table in enumerate(tables):
        for logical,physical in enumerate(table.tolist()):
            if physical<0: continue
            start=logical*p; end=min(start+p,lengths[seq])
            if start<end:
                physical_k[physical,:,:end-start]=contiguous_k[seq][:,start:end]
                physical_v[physical,:,:end-start]=contiguous_v[seq][:,start:end]

    q=torch.randn(2,hq,1,d,device=DEVICE,dtype=DTYPE); max_diff=0.
    for seq,length in enumerate(lengths):
        pk=gather_paged(physical_k,tables[seq],length,p)
        pv=gather_paged(physical_v,tables[seq],length,p)
        torch.testing.assert_close(pk,contiguous_k[seq])
        torch.testing.assert_close(pv,contiguous_v[seq])
        dense=attend_one(q[seq],contiguous_k[seq],contiguous_v[seq])
        paged=attend_one(q[seq],pk,pv)
        torch.testing.assert_close(paged,dense,rtol=1e-5,atol=1e-6)
        max_diff=max(max_diff,(paged-dense).abs().max().item())

    unique_blocks=torch.unique(tables[tables>=0]).numel()
    live=4+(lengths[0]-4)+(lengths[1]-4)
    allocated=unique_blocks*p; waste=allocated-live
    assert unique_blocks==4 and waste==3
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE} Hq={hq} Hkv={hkv} D={d} block_size={p}")
    print(f"lengths={lengths} block_tables={tables.cpu().tolist()}")
    print(f"physical_cache={tuple(physical_k.shape)} shared_physical_block=5")
    print(f"max_attention_abs_diff={max_diff:.3e}")
    print(f"unique_blocks={unique_blocks} allocated_slots={allocated} unique_live_tokens={live}")
    print(f"internal_fragmentation_slots={waste} prefix_sharing_saved_slots={20-allocated}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/paged_attention_reference.py。实际运行得到最大 attention 绝对差 0.000e+00、3 个内部碎片 slots、4 个共享节省 slots。

这段 Python 先 gather 再做 dense attention，只用于证明地址映射语义，不能证明 PagedAttention kernel 性能。生产 kernel 会在 tile/load 过程中直接查表并读取 blocks，避免构造完整 contiguous copy；其速度取决于 block layout、访存合并、batch、head dimension 和实现。

它带来的成本与收益条件

收益主要来自：

• 只按实际增长分配 blocks，避免为未知最大长度一次预留整段；
• 逻辑连续不要求物理连续，降低外部碎片和搬迁需求；
• 相同 prefix、beam parent 或 copy-on-write 状态可共享完整 blocks；
• scheduler 可以用 block 数做容量预算、抢占和 swap/offload 决策。

新增成本包括 block table/metadata、分配与引用计数、每次读取的地址转换、最后 block 内部碎片，以及 kernel 为 paged layout 付出的控制流。请求很少、长度固定且 contiguous 预分配非常合适时，paged 管理未必更快；它首先是服务端容量与动态性机制，不是无条件加速 attention 数学。

Prefix sharing 也只在完整 blocks 或实现支持的边界上最自然。若两个请求只共享一个 block 的一部分，runtime 要么不共享该部分，要么采用更复杂的粒度。共享 block 一旦需要分叉写入，必须 copy-on-write，不能让一条序列覆盖另一条序列的 KV。

与上下游关系

• Tokenizer/Chat Template 决定可比较的 token prefix；Prefix Cache进一步定义 block identity、命中与 copy-on-write，相同字符串外观不足以安全共享 physical blocks。
• KV Cache 定义逻辑 K/V 内容、容量和跨 step 生命周期；PagedAttention 只改变 token 维的物理映射和分配。
• Attention Sink可把早期 sink blocks 长期 pinned、只滚动 recent block table；PagedAttention 提供物理分配与 refcount，但不定义 sink ∪ recent 可见 mask、RoPE rerotation 或模型质量。
• Prefill/Decode 在 prefill 批量写 blocks、decode 每步追加 slot；continuous batching scheduler 决定何时分配和回收。
• 缩放点积注意力 定义查表后 K/V 参与的数学；FlashAttention负责 online softmax 和 tile IO，而不是 block ownership。

参考资料

• Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention
• vLLM — Paged Attention design
• TensorRT-LLM — KV Cache Reuse