""" ============================================================================= CS349E · 第 7 讲 · 04_prefix_caching.py —— 前缀缓存：算过的 KV 存下来给后面的请求用 ============================================================================= 【这个脚本要回答的问题】 很多请求的开头是一样的：同一个系统提示、同一组 few-shot 例子、同一篇长文档、多轮对话的历史。 这些前缀的 KV 算一次就够了——只要能**认出**"这一段前缀以前算过"，并且**存得下**。 演示 1：怎样认出来：链式哈希（每块的哈希包含它前面的全部内容，Mooncake 图 3 / vLLM 的自动前缀缓存）。 演示 2：在第 2 讲的模型上真的复用：8 个请求共享一段 few-shot 前缀（vLLM 论文图 16：1-shot 80 个 token、 5-shot 341 个 token），后来的请求直接映射到缓存里的块，只算自己的那部分；logits 和从头算的一致。 演示 3：一个仿 Mooncake 公开 trace 的合成负载：块有多热门、缓存容量和命中率的关系、LRU vs LFU（论文表 1）。 演示 4：命中率换算成 prefill 的计算量：LLaMA2-70B、8 块 A800，前缀缓存省了多少 GPU 时间。 【运行】python3 04_prefix_caching.py （约 10 秒） 【依赖】numpy；本目录的 paged_kv、kv_store；../lecture02 的 transformer_lab；../lecture01 的 console。 """ import os import sys from collections import Counter import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) from kv_store import (BLOCK, LLAMA2_70B, TIERS, block_hashes, hit_ratio, make_trace, # noqa: E402 prefill_compute_time, prefill_time) from paged_kv import BlockPool, PagedSequence, paged_step # noqa: E402 from console import print_header, print_observations, print_table # noqa: E402 from transformer_lab import TINY, KVCache, config_like, forward, init_params # noqa: E402 # ----------------------------------------------------------------------------- def demo1_hashing(): print_header("演示 1：链式哈希——每块的哈希 = hash(前一块的哈希 ‖ 这一块的 token)（块大小 4）") sys_prompt = [101, 7, 7, 3, 55, 2, 9, 14] # 两块的"系统提示" a = sys_prompt + [21, 22, 23, 24, 25, 26, 27, 28, 29] b = sys_prompt + [21, 22, 23, 24, 31, 32, 33, 34] c = [101, 7, 7, 3, 55, 2, 9, 99] + [21, 22, 23, 24, 25, 26, 27, 28] rows = [] for name, toks in (("请求 A", a), ("请求 B", b), ("请求 C（系统提示改了最后一个 token）", c)): rows.append([name, len(toks), " ".join(block_hashes(toks, block=4))]) print_table(["请求", "token 数", "每个完整块的哈希（前 12 位）"], rows) print_observations([ "A、B 的前 3 块一样（系统提示 + 共同的开头），哈希也一样——B 可以直接用 A 算过的前 3 块 KV；第 4 块不同，从这里开始要自己算。", "C 只在第 2 块改了一个 token，第 2 块以及**之后所有块**的哈希都变了：KV 依赖前面的全部内容，链式哈希正好表达这一点。" "所以前缀缓存只能复用\"从开头开始连续相同\"的部分。", "A 的最后一个 token 不满一块，不算哈希、不进缓存（vLLM 和 Mooncake 都只缓存完整的块）。", ]) # ----------------------------------------------------------------------------- def demo2_shared_prefix(): print_header("演示 2：8 个请求共享 few-shot 前缀（vLLM 论文图 16 的设置），块大小 16") cfg = config_like(TINY, max_seq=512) params = init_params(cfg, seed=0) rng = np.random.default_rng(5) rows = [] for shots, plen in (("1-shot", 80), ("5-shot", 341)): prefix = rng.integers(0, cfg.vocab, plen) suffixes = [rng.integers(0, cfg.vocab, int(rng.integers(10, 30))) for _ in range(8)] pool = BlockPool(cfg, n_blocks=256, block_size=16) cache = {} # 哈希 → 物理块（缓存自己持有一份引用） computed = 0 worst = 0.0 for suf in suffixes: toks = np.concatenate([prefix, suf]) hs = block_hashes(toks, block=16) seq = PagedSequence(pool) n_hit = 0 for h in hs: # 从头开始连续命中几块 if h not in cache: break seq.blocks.append(cache[h]) pool.incref(cache[h]) n_hit += 1 seq.length = n_hit * 16 new = toks[seq.length:] out = paged_step(params, cfg, [seq], [new])[0][-1] computed += len(new) for i, h in enumerate(hs): # 把这个请求的完整块登记进缓存 if h not in cache: cache[h] = seq.blocks[i] pool.incref(seq.blocks[i]) ref = forward(params, cfg, toks[None], cache=KVCache(cfg, 1))[0, -1] worst = max(worst, float(np.max(np.abs(out - ref)))) seq.free() total = sum(plen + len(s) for s in suffixes) rows.append([shots, plen, total, computed, f"{1 - computed / total:.0%}", f"{total / computed:.1f}×", len(cache), f"{worst:.1e}"]) print_table(["前缀", "前缀长度", "8 个请求的输入 token", "实际算的", "省下的", "prefill 计算量少", "缓存里的块", "和从头算的最大差"], rows) print_observations([ "第一个请求从头算并把完整的块登记进缓存；后面 7 个请求的前缀直接映射到缓存里的物理块（引用计数 +1，不复制），" "只算自己的后缀——logits 和从头算的完全一致（这里的最大差是 0）。", "5-shot 的前缀 341 个 token = 21 个满块 + 5 个零头；零头不在缓存里，每个请求要和后缀一起重算。前缀越长、后缀越短，省得越多。", "vLLM 论文图 16：同样的翻译负载上，共享前缀让吞吐比 Orca (Oracle) 高 1.67 倍（1-shot）和 3.58 倍（5-shot）——" "除了省 prefill 的计算，前缀的 KV 只存一份，也省显存、让批更大。", ]) # ----------------------------------------------------------------------------- def demo3_trace(): print_header("演示 3：仿 Mooncake 公开 trace 的合成负载——块的热度、缓存容量与命中率（论文 §4、表 1）") tr = make_trace(seed=0) ins = np.array([r.input_len for r in tr]) outs = np.array([r.output_len for r in tr]) cnt = Counter(b for r in tr for b in r.blocks) v = np.array(sorted(cnt.values(), reverse=True)) print(f"{len(tr)} 个请求，{tr[-1].arrival / 60:.0f} 分钟；平均输入 {ins.mean():,.0f}、平均输出 {outs.mean():.0f} 个 token" f"（论文的真实 trace：8,590 和 182）；块大小 {BLOCK}；一共 {len(cnt):,} 个不同的块") print(f"块的热度：{np.mean(v == 1):.0%} 的块只出现过一次（之后再没被用到）；最热的块被访问 {v[0]:,} 次；" f"最热的 1% 的块占了全部访问的 {v[:max(1, len(v) // 100)].sum() / v.sum():.0%}") rows = [] caps = (None, 100000, 50000, 30000, 10000, 3000, 1000) for pol in ("lru", "lfu"): rows.append([pol.upper()] + [f"{hit_ratio(tr, c, pol):.2f}" for c in caps]) rows.append(["论文 LRU"] + ["0.55", "0.55", "0.55", "0.54", "0.46", "-", "0.34"]) rows.append(["论文 LFU"] + ["0.55", "0.55", "0.55", "0.50", "0.40", "-", "0.34"]) print_table(["策略 \\ 容量（块）"] + ["无限" if c is None else f"{c:,}" for c in caps], rows) print_observations([ "命中率随容量先快速上升、再趋于平缓：无限容量时约 0.58（论文 0.55）——剩下的块本来就是第一次出现，再大的缓存也没用。", "LRU 比 LFU 好（论文也是 LRU 最好）：复用多半发生在时间上相近的请求之间（同一个会话的下一轮），LFU 会让" "\"历史上很热、现在已经没人用\"的块赖着不走。", "超过一半的块只出现一次，而少数块被访问几千次：热门的块（公共的系统提示、热门文档）值得复制到多台机器上，" "否则所有请求都去同一台机器拉 KV，网络会堵（Mooncake §6.2，05 脚本演示 2）。", "这是合成负载，只模仿了论文给出的几个统计量；论文开源了真实 trace，可以替换进来。", ]) # ----------------------------------------------------------------------------- def demo4_compute_saved(): print_header("演示 4：命中率换算成 prefill 的 GPU 时间（LLaMA2-70B、8 块 A800 张量并行、60% 峰值算力）") tr = make_trace(seed=0) total_nocache = sum(prefill_compute_time(LLAMA2_70B, r.input_len, 0) for r in tr) rows = [["不缓存", "0%", f"{total_nocache / 3600:.2f}", "1.00×"]] from kv_store import make_cache for cap in (None, 10000, 1000): cache = make_cache(cap, "lru") tot = 0.0 hit_tok = 0 all_tok = 0 for r in tr: n_hit = min(cache.match(r.blocks) * BLOCK, r.input_len - 1) tot += prefill_time(LLAMA2_70B, r.input_len, n_hit, TIERS[1]) hit_tok += n_hit all_tok += r.input_len cache.insert(r.blocks) rows.append([f"LRU，容量 {'无限' if cap is None else f'{cap:,} 块'}", f"{hit_tok / all_tok:.0%}", f"{tot / 3600:.2f}", f"{total_nocache / tot:.2f}×"]) print("（命中的 KV 从本机 CPU 内存读，和计算按层重叠；每块的 KV = 512 × 320 KB = 168 MB。）") print_table(["缓存", "命中的输入 token", "prefill 的总 GPU 时间（小时，8 卡）", "省了多少"], rows) kv_blk = BLOCK * LLAMA2_70B.kv_bytes_per_token() print(f"\n容量 10,000 块需要 {10000 * kv_blk / 1e12:.2f} TB 的存储——一台 8 卡机器的 HBM 只有 0.64 TB，" "所以这么大的前缀缓存只能放在 CPU 内存、SSD 或者整个集群的内存池里。") print_observations([ "命中了多少输入 token，prefill 的计算就省了差不多多少（注意力那部分稍有出入）：无限缓存时 prefill 的 GPU 时间少一半以上。", "但缓存很大：一块 512 个 token 的 KV 就有 168 MB，一万块要 1.7 TB。GPU 显存装不下，这就是 Mooncake 把 KV 放到" "整个集群的 CPU 内存和 SSD 上、按 KV 的位置来调度请求的原因（05 脚本）。", ]) if __name__ == "__main__": demo1_hashing() demo2_shared_prefix() demo3_trace() demo4_compute_saved()