""" ============================================================================= CS349E · 第 7 讲 · 01_paged_attention.py —— PagedAttention：块表、按块注意力、共享与写时复制 ============================================================================= 【这个脚本要回答的问题】 vLLM 把 KV cache 像操作系统管理内存那样分页：KV 存在固定大小的物理块里，每个序列一张块表。 这样做算得对吗？能省多少？ 演示 1：vLLM 论文图 6 的例子：一个 7 个 token 的提示、块大小 4，块表怎样随着 decode 增长。 演示 2：在第 2 讲的 TINY 模型上，分页的 KV（几种块大小）和连续的 KV 跑出来的 logits 逐位比较； 按块计算的注意力（论文式 4）和"拼起来再算"比较。 演示 3：并行采样（同一个提示生成 n 个样本）：提示的块被 n 个样本共享；输出和"各自一份 KV"相同，显存省多少。 演示 4：束搜索：候选之间共享前缀、写时复制；和"每个子候选整份复制父候选的 KV"比，输出相同，复制量和显存差多少。 【运行】python3 01_paged_attention.py （约 3 秒） 【依赖】numpy；本目录的 paged_kv（它会导入 ../lecture02、../lecture06）；../lecture01 的 console。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) from paged_kv import BlockPool, PagedSequence, beam_search, blockwise_attention, paged_step, parallel_sample # noqa: E402 from console import print_header, print_observations, print_table # noqa: E402 from transformer_lab import TINY, KVCache, attention_core, forward, init_params # noqa: E402 PARAMS = init_params(TINY, seed=0) # ----------------------------------------------------------------------------- def demo1_block_table(): print_header("演示 1：块表怎样增长（vLLM 论文图 6：提示 7 个 token，块大小 4）") pool = BlockPool(TINY, n_blocks=8, block_size=4) # 先让池子里"不连续"：占用再释放几个块，模拟一个用了一阵的 GPU junk = [pool.alloc() for _ in range(5)] for b in (junk[0], junk[2], junk[4]): pool.decref(b) seq = PagedSequence(pool) words = "Four score and seven years ago our".split() toks = np.arange(len(words)) + 10 rows = [] paged_step(PARAMS, TINY, [seq], [toks]) rows.append(["prefill 7 个 token", str(seq.blocks), seq.length, f"{seq.length - (len(seq.blocks) - 1) * 4}/4"]) for name in ("fathers", "brought"): paged_step(PARAMS, TINY, [seq], [np.array([99])]) rows.append([f"decode：写入\"{name}\"的 KV", str(seq.blocks), seq.length, f"{seq.length - (len(seq.blocks) - 1) * 4}/4"]) print_table(["这一步", "块表（逻辑块 → 物理块）", "已写入的 token", "最后一块填了"], rows) print(f"\n池里一共 {pool.n_blocks} 个物理块；这个序列用了 {len(seq.blocks)} 个，另外 2 个被别的（模拟的）序列占着。") print_observations([ "物理块的编号不连续（先用了被释放出来的空块）：逻辑上连续的 KV 散在显存各处，靠块表找回来——和操作系统的页表一样。", "prefill 只分配了 2 块（7 个 token），第一次 decode 写进第 2 块剩下的那个槽位，第二次 decode 第 2 块满了才分配新块。" "任何时候浪费的最多只有最后一块里没填的槽位（< 块大小）。", ]) # ----------------------------------------------------------------------------- def demo2_equivalence(): print_header("演示 2：分页的 KV = 连续的 KV（TINY 模型，4 个长短不一的序列，prefill + 20 步 decode）") rng = np.random.default_rng(1) prompts = [rng.integers(0, TINY.vocab, L) for L in (5, 17, 33, 64)] rows = [] for B in (1, 4, 16, 64): pool = BlockPool(TINY, n_blocks=260 // B + 8, block_size=B) seqs = [PagedSequence(pool) for _ in prompts] refs = [KVCache(TINY, 1) for _ in prompts] worst = 0.0 chunks = list(prompts) for _ in range(21): outs = paged_step(PARAMS, TINY, seqs, chunks) new = [] for c, o, r in zip(chunks, outs, refs): ref = forward(PARAMS, TINY, c[None], cache=r)[0] worst = max(worst, float(np.max(np.abs(o - ref)))) new.append(np.array([int(np.argmax(o[-1]))])) chunks = new used = pool.n_used need = sum(s.length for s in seqs) rows.append([B, used, used * B, need, f"{1 - need / (used * B):.1%}", f"{worst:.1e}"]) print_table(["块大小", "用了几个物理块", "分配的槽位", "实际 token", "最后一块的浪费", "和连续 KV 的最大差"], rows) # 按块注意力 pool = BlockPool(TINY, n_blocks=16, block_size=8) s = PagedSequence(pool) paged_step(PARAMS, TINY, [s], [prompts[2]]) q = rng.standard_normal((TINY.n_heads, TINY.d_head)) K, V = s.read(1, s.length) ref = attention_core(q[None, None], K, V, s.length - 1)[0, 0].reshape(TINY.n_heads, TINY.d_head) d = float(np.max(np.abs(blockwise_attention(q, s, 1) - ref))) print(f"\n按块计算注意力（论文式 4：逐块算分数、整体 softmax、逐块乘 V）vs 拼起来再算：最大差 {d:.1e}" f"（序列长 {s.length}，块大小 8，{len(s.blocks)} 个块）") print_observations([ "不管块多大，分页的 KV 跑出来的 logits 和连续的 KV 一致到 1e-15：分页只改变 KV **放在哪里**，不改变算什么。", "块越大，最后一块平均浪费得越多（块大小 16 时 17%、64 时 38%，这里的序列都很短）；块越小，块表越长、内核每次读的数据越碎。" "vLLM 默认 16（02 脚本演示 3 详细看这个取舍）。", "真实的 PagedAttention 内核不把块拼起来：每个 warp 读一个块（论文 §5.1），直接在块上算——就是这里的按块注意力。", ]) # ----------------------------------------------------------------------------- def demo3_parallel_sampling(): print_header("演示 3：并行采样——提示的 KV 被 n 个样本共享（提示 100 个 token，各生成 40 个，温度 1，块大小 16）") prompt = np.random.default_rng(2).integers(0, TINY.vocab, 100) rows = [] for n in (2, 4, 6): ref, _ = parallel_sample(PARAMS, TINY, prompt, n, 40, 1.0, seed=7, pool=None) pool = BlockPool(TINY, n_blocks=256, block_size=16) out, st = parallel_sample(PARAMS, TINY, prompt, n, 40, 1.0, seed=7, pool=pool) used, unshared = np.array(st["blocks_used"]), np.array(st["blocks_unshared"]) saving = 1 - used.sum() / unshared.sum() rows.append([n, "相同" if out == ref else "不同", int(used[-1]), int(unshared[-1]), f"{saving:.1%}", pool.n_copies]) print_table(["样本数 n", "输出 vs 各自一份 KV", "最后一步用的块", "不共享时要的块", "平均省下的块", "写时复制的次数"], rows) print_observations([ "n 个样本共享提示的 7 个块（100 个 token 的前 96 个在 6 个满块里，第 7 块只填了 4 个）。" "往那个半满的块里写时，前 n − 1 个样本各触发一次写时复制，最后一个发现引用计数已经是 1，直接写——所以复制次数是 n − 1。", "省下的比例随 n 增加，也取决于提示和输出的长短：提示越长、输出越短，省得越多。vLLM 论文在 Alpaca 上测到 6.1%–9.8%、" "ShareGPT 上 16.2%–30.5%（图 15）。", "输出和\"各自一份 KV\"完全相同：共享是纯粹的显存优化。", ]) # ----------------------------------------------------------------------------- def demo4_beam_search(): print_header("演示 4：束搜索——候选之间共享前缀；和\"整份复制\"比（提示 64 个 token，生成 32 步，块大小 16）") prompt = np.random.default_rng(3).integers(0, TINY.vocab, 64) rows = [] for k in (2, 4, 6): ref, st_ref = beam_search(PARAMS, TINY, prompt, k, 32, pool=None) pool = BlockPool(TINY, n_blocks=512, block_size=16) out, st = beam_search(PARAMS, TINY, prompt, k, 32, pool=pool) used, unshared = np.array(st["blocks_used"]), np.array(st["blocks_unshared"]) same = [t for _, t in ref] == [t for _, t in out] rows.append([k, "相同" if same else "不同", f"{st_ref['copied_tokens']:,}", f"{pool.n_copies * 16:,}", f"{1 - used.sum() / unshared.sum():.1%}", pool.peak_used]) print_table(["束宽 k", "k 个结果", "整份复制：复制的 token", "分页：写时复制的 token", "平均省下的块", "分页的峰值块数"], rows) print_observations([ "\"整份复制\"的做法（vLLM 论文说以前的系统就是这样）每一步都要把父候选的整个 KV 复制给子候选，" "复制量随序列长度和束宽增长；分页时只有往共享的最后一块里写才复制一块（16 个 token）。", "束搜索的候选不只共享提示，还共享前面一起走过的路：这里省下 39%–68% 的块。vLLM 论文测到 37.6%–55.2%（Alpaca）、" "44.3%–66.3%（ShareGPT），比并行采样的 6.1%–30.5% 高得多（演示 3 的提示比输出长，所以那里省得也多）。", "k 个结果完全相同：分页、共享、写时复制都不改变计算。", ]) if __name__ == "__main__": demo1_block_table() demo2_equivalence() demo3_parallel_sampling() demo4_beam_search()