""" ============================================================================= CS349E · 第 2 讲 · 04_kv_cache_decode.py —— 带 KV cache 的 decode 循环：正确性、省了多少、代价是什么 ============================================================================= 【这个脚本要回答的问题】 原课作业 1 的主题之一：写一个带 KV cache 的自回归生成循环，并确认它和"每步重算整个序列"的结果一样。 演示 1：正确性：带 cache 和不带 cache 生成的 token 完全相同；以及两个常见 bug（提示没写进 cache、 位置编号写错）会让输出在哪一步开始跑偏。 演示 2：省了多少：不带 cache 时每一步都要把整个序列重算一遍，总运算量随长度平方增长； 带 cache 时每一步只算 1 个新 token。小模型上数 FLOPs、测时间，再用公式推到 Llama-3.1-8B。 演示 3：代价是什么：KV cache 占的显存随"批大小 × 上下文"线性增长，多长的上下文就和权重一样大。 演示 4：批量 decode 的账：权重字节不随批大小变化，KV 字节和 FLOPs 随批大小线性增长—— 强度为什么停在一个上限（第 1 讲 §4.3 的"每个样本一个矩阵"，这里一行行数出来）。 【怎么读输出】演示 1、3、4 是确定性的计算；演示 2 的墙钟时间和机器有关，只看趋势（FLOPs 是确定的）。 【运行】python3 04_kv_cache_decode.py （约 6 秒） 【依赖】numpy；同目录的 transformer_lab；从 ../lecture01 导入 console、hw_specs、llm_math。 """ import os import sys import time import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_num, fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from llm_math import LLAMA31_8B, LLAMA31_70B, decode_step, prefill_time # noqa: E402 from transformer_lab import (ATTN_CORE_OPS, MATMUL_OPS, TINY, Config, Counter, KVCache, analytic_ops, # noqa: E402 forward, generate, init_params, sum_rows) def greedy_decode(params, cfg, prompt, n_new, bug=None): """ 一个手写的 decode 循环（和 transformer_lab.generate 一样），可以故意注入两种常见 bug。 【参数】bug： None ：正确做法 "prompt_not_cached" ：prefill 时没有把提示写进 cache（算出了第一个 token，但后面的步骤看不到提示） "position_reset" ：每一步 decode 都把新 token 当成位置 0（RoPE 的位置没有加上已有长度） 【返回】(tokens [B, T0 + n_new], 每步的 logits 列表) """ B = prompt.shape[0] cache = KVCache(cfg, B, dtype=params["embed"].dtype) if bug == "prompt_not_cached": logits = forward(params, cfg, prompt, last_only=True)[:, -1] else: logits = forward(params, cfg, prompt, cache=cache, last_only=True)[:, -1] toks, step_logits = prompt.copy(), [] for i in range(n_new): step_logits.append(logits) nxt = np.argmax(logits, axis=-1)[:, None] toks = np.concatenate([toks, nxt], axis=1) if i == n_new - 1: break rp0 = 0 if bug == "position_reset" else None logits = forward(params, cfg, nxt, cache=cache, last_only=True, rope_pos0=rp0)[:, -1] return toks, step_logits # --------------------------------------------------------------------------- # 演示 1：正确性 # --------------------------------------------------------------------------- def demo1_correctness(): print_header("演示 1：带 KV cache 的生成 = 每步重算整个序列（以及两个常见 bug）") cfg = TINY params = init_params(cfg, seed=0) rng = np.random.default_rng(1) prompt = rng.integers(0, cfg.vocab, size=(4, 16)) n_new = 24 ref_toks, ref_logits = generate(params, cfg, prompt, n_new, use_cache=False) rows = [] for bug, label in ((None, "正确的 KV cache"), ("prompt_not_cached", "bug：提示没写进 cache"), ("position_reset", "bug：decode 的位置都当成 0")): toks, logits = greedy_decode(params, cfg, prompt, n_new, bug=bug) same = np.array_equal(toks, ref_toks) diffs = [float(np.max(np.abs(a - b))) for a, b in zip(logits, ref_logits)] first_bad = next((i for i, dd in enumerate(diffs) if dd > 1e-9), None) tok_diff = np.nonzero(np.any(toks[:, 16:] != ref_toks[:, 16:], axis=0))[0] rows.append([label, "是" if same else "否", f"{max(diffs):.1e}", "-" if first_bad is None else f"第 {first_bad} 步", "-" if len(tok_diff) == 0 else f"第 {int(tok_diff[0])} 个"]) print(f"4 个序列，提示 16 个 token，贪心生成 {n_new} 个；参照 = 不用 cache、每步重算整个序列") print_table(["做法", "token 完全相同", "logits 最大差", "logits 从哪一步开始不同", "token 从第几个开始不同"], rows) print_observations([ "正确的 KV cache 和重算整个序列的 logits 只差 1e-14 量级（float64 的舍入误差）：两者在数学上是同一个计算，" "只是 cache 把旧 token 的 K、V 存了起来（因果掩码保证它们不会变，02 脚本演示 3）。", "提示没写进 cache：第 0 步（由 prefill 直接算出）是对的，从第 1 步起模型就'忘了'提示，输出立刻跑偏。", "位置编号写错：第 0 步也是对的，之后每个新 token 都以为自己在位置 0，RoPE 算出的相对位置全错。" "这类 bug 不会报错，只会让输出'有点怪'——所以作业要求和不用 cache 的参照逐步比较 logits，而不只是看生成的文字。", ]) # --------------------------------------------------------------------------- # 演示 2：省了多少 # --------------------------------------------------------------------------- def demo2_savings(): print_header("演示 2：KV cache 省了多少——FLOPs 从平方变成线性") cfg = Config(vocab=1024, d_model=256, n_layers=4, n_heads=8, n_kv_heads=2, d_head=32, d_ff=688, max_seq=1024, name="small") params = init_params(cfg, seed=0, dtype=np.float32) rng = np.random.default_rng(2) T0 = 64 prompt = rng.integers(0, cfg.vocab, size=(1, T0)) rows = [] for n_new in (64, 128, 256): res = [] for use_cache in (False, True): cp, cd = Counter(), Counter() t0 = time.perf_counter() toks, _ = generate(params, cfg, prompt, n_new, use_cache=use_cache, counter_prefill=cp, counter_decode=cd) dt = time.perf_counter() - t0 fl = cp.total("useful", MATMUL_OPS + ATTN_CORE_OPS) + cd.total("useful", MATMUL_OPS + ATTN_CORE_OPS) res.append((fl, dt, toks)) assert np.array_equal(res[0][2], res[1][2]) rows.append([n_new, fmt_num(res[0][0]), fmt_num(res[1][0]), f"{res[0][0] / res[1][0]:.0f}×", fmt_time(res[0][1]), fmt_time(res[1][1]), f"{res[0][1] / res[1][1]:.0f}×"]) print(f"小模型（d = {cfg.d_model}，{cfg.n_layers} 层，float32），提示 {T0} 个 token，生成 n 个（两种做法的 token 完全相同）：") print_table(["生成 n", "FLOPs：不用 cache", "用 cache", "比值", "时间：不用 cache", "用 cache", "比值"], rows) m = LLAMA31_8B rows = [] for T0, n in ((1000, 1000), (4000, 1000), (1000, 8000)): with_cache = prefill_time(m, H100, T0)["flops"] + sum(m.flops_per_token(T0 + t) for t in range(n - 1)) # 不用 cache：第 t 步要对 T0 + t 个 token 做一次完整的前向（相当于一次 prefill） without = sum(prefill_time(m, H100, T0 + t)["flops"] for t in range(n)) rows.append([f"提示 {T0}，生成 {n}", fmt_num(with_cache), fmt_num(without), f"{without / with_cache:.0f}×"]) print("\nLlama-3.1-8B（公式；'不用 cache'按每步一次完整 prefill 估算）：") print_table(["场景", "FLOPs：用 cache", "不用 cache", "比值"], rows) print_observations([ "不用 cache 时，第 t 步要重算 T0 + t 个 token，总运算量 ≈ n·T0 + n²/2 个 token 的前向；用 cache 时只有 T0 + n 个。" "生成越长，差距越大（小模型上 n = 256 时约 140 倍的 FLOPs）。", "墙钟时间的比值比 FLOPs 的比值小得多：用 cache 时每一步只处理 1 个 token，numpy 的每次调用都有固定开销，" "矩阵乘法也退化成了向量乘矩阵（02 脚本演示 6）——FLOPs 省了，但每个 FLOP 变贵了。这就是 decode 的本质。", "Llama-3.1-8B 生成 8000 个 token 时，不用 cache 要多花约 4000 倍的运算。没有 KV cache，长输出的推理模型根本跑不起来。", ]) # --------------------------------------------------------------------------- # 演示 3：KV cache 占多少显存 # --------------------------------------------------------------------------- def demo3_memory(): print_header("演示 3：KV cache 的代价——显存随'批大小 × 上下文'线性增长") cfg = TINY params = init_params(cfg, seed=0) rng = np.random.default_rng(3) cache = KVCache(cfg, batch=3) forward(params, cfg, rng.integers(0, cfg.vocab, size=(3, 40)), cache=cache) per_tok = cfg.to_model_config().kv_bytes_per_token() print(f"小模型：3 个序列 × 40 个 token，cache.nbytes = {fmt_bytes(cache.nbytes())}；" f"公式 3 × 40 × {per_tok:.0f} B = {fmt_bytes(3 * 40 * per_tok)}") rows = [] for m in (LLAMA31_8B, LLAMA31_70B): w = m.weight_bytes("bf16") kv = m.kv_bytes_per_token() cells = [m.name, fmt_bytes(w), fmt_bytes(kv, binary=True)] for B, ctx in ((1, 8192), (32, 8192), (8, 131072)): cells.append(fmt_bytes(B * ctx * kv)) cells.append(f"{w / kv / 1000:.0f} K") rows.append(cells) print() print_table(["模型", "BF16 权重", "KV / token", "B=1, 8K", "B=32, 8K", "B=8, 128K", "KV = 权重时的 token 数"], rows) print_observations([ "KV cache 是用显存换运算：每个 token 一份 K、V，每层每组 KV 头各 d_head 个数。", "8B 模型只要约 12 万个 token（比如 32 个序列 × 4K）的 KV 就和全部权重一样大；128K 上下文时，8 个序列的 KV 是权重的 8 倍多。", "所以推理系统里最紧张的资源常常不是算力，而是**放 KV cache 的显存**：它决定了能同时服务几个请求（第 1 讲 §4.4）。" "缩小 KV（GQA/MQA/MLA，05 脚本；KV 量化，第 3 讲）、管理 KV（分页、前缀共享、换出，第 7 讲）都是在攻这一点。", ]) # --------------------------------------------------------------------------- # 演示 4：批量 decode 的账 # --------------------------------------------------------------------------- def demo4_batched_decode(): print_header("演示 4：批量 decode——权重读一次给大家用，KV 各读各的") cfg = TINY params = init_params(cfg, seed=0) rng = np.random.default_rng(4) ctx = 128 rows = [] for B in (1, 4, 16, 64): cache = KVCache(cfg, B) forward(params, cfg, rng.integers(0, cfg.vocab, size=(B, ctx - 1)), cache=cache) c = Counter() forward(params, cfg, rng.integers(0, cfg.vocab, size=(B, 1)), cache=cache, counter=c, last_only=True) w = c.total("w", MATMUL_OPS) kv = c.total("kv", ATTN_CORE_OPS) fl = c.total("useful", MATMUL_OPS + ATTN_CORE_OPS) rows.append([B, fmt_bytes(w), fmt_bytes(kv), fmt_num(fl), f"{fl / (w + kv):.2f}"]) print(f"小模型（实际运行并记账），每个序列上下文 {ctx}：") print_table(["批大小 B", "读权重", "读 KV", "FLOPs", "强度 FLOPs/字节"], rows) m = LLAMA31_8B rows = [] for B in (1, 16, 64, 256, 1024): a = analytic_ops(m, B, 1, 4095, last_only=True) w = sum_rows(a, "w", MATMUL_OPS) kv = sum_rows(a, "kv", ATTN_CORE_OPS) fl = sum_rows(a, "useful", MATMUL_OPS + ATTN_CORE_OPS) st = decode_step(m, H100, batch=B, ctx=4096) rows.append([B, fmt_bytes(w), fmt_bytes(kv), f"{fl / (w + kv):.1f}", fmt_time(st["t_step"]), f"{B / st['t_step']:,.0f}"]) lim = m.flops_per_token(4096) / (4096 * m.kv_bytes_per_token()) print(f"\nLlama-3.1-8B（公式），上下文 4096，H100：") print_table(["B", "读权重", "读 KV", "强度", "一步下界", "tok/s/芯片"], rows) print_observations([ "读权重的字节不随 B 变化（所有序列共享同一份权重），读 KV 的字节和 FLOPs 都和 B 成正比。", f"所以强度 = (B·f) / (W + B·k) 随 B 增大，但有上限 f / k（f：每 token 运算量，k：每序列 KV 字节）。" f"Llama-3.1-8B 在 4096 上下文时上限约 {lim:.0f} FLOP/byte，远低于 H100 的屋脊点 295：批再大也到不了 compute-bound。", "这就是第 1 讲 §4.3 的结论，这里由实际运行的模型逐字节数出来。要突破这个上限，只能让每个序列的 KV 变小或少读：" "GQA/MQA/MLA（05 脚本）、KV 量化（第 3 讲）、稀疏注意力。", ]) if __name__ == "__main__": demo1_correctness() demo2_savings() demo3_memory() demo4_batched_decode()