""" ============================================================================= CS349E · 第 7 讲 · 05_mooncake.py —— Mooncake：以 KV cache 为中心的集群调度 ============================================================================= 【这个脚本要回答的问题】 Mooncake（Kimi 的服务平台）的三个关键设计，各用一个小模拟来理解： 演示 1："用存储换计算"——把算过的 KV 存在 CPU 内存、SSD、别的机器上，读回来比重算快多少？和模型结构（MHA vs GQA）的关系。 演示 2：按 KV 调度 prefill（论文算法 1、图 8）：8 个 prefill 实例各有自己的前缀缓存，请求分给谁？ 随机 / 按负载 / 按缓存命中 / 按缓存命中 + 热点 KV 复制。 演示 3：过载时提前拒绝（论文 §7、表 3）：prefill 做完才发现 decode 放不下，prefill 就白做了；提前拒绝能省多少， 代价是什么（prefill 和 decode 的负载此起彼伏）。 【模型的假设】LLaMA2-70B（Mooncake 的实验用同结构的哑模型）、每个实例 8 块 A800、prefill 按 60% 峰值算力； 各层存储的带宽是假设值（见 kv_store.TIERS）。负载是 kv_store.make_trace 生成的合成 trace（仿论文公开 trace 的统计量）。 【运行】python3 05_mooncake.py （约 5 秒） 【依赖】numpy；本目录的 kv_store；../lecture01 的 console、llm_math。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) from kv_store import (LLAMA2_70B, TIERS, iid_requests, make_trace, overload_cluster, # noqa: E402 prefill_cluster, prefill_compute_time, prefill_time) from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 from llm_math import ModelConfig # noqa: E402 OPT_13B = ModelConfig("OPT-13B", n_layers=40, d_model=5120, n_heads=40, n_kv_heads=40, d_head=128, d_ff=13653, vocab=50272, tied=True) # ----------------------------------------------------------------------------- def demo1_tiers(): print_header("演示 1：用存储换计算——读回 KV vs 重算（每个 token）") rows = [] for model, n_gpu, label in ((LLAMA2_70B, 8, "LLaMA2-70B（GQA，8 组 KV），8 块 A800"), (OPT_13B, 1, "OPT-13B（MHA，40 组 KV），1 块 A100")): rc = (prefill_compute_time(model, 8192, 0, n_gpu=n_gpu, overhead=0.0)) / 8192 cells = [label, fmt_bytes(model.kv_bytes_per_token()), fmt_time(rc)] for tier in TIERS[1:]: bw = tier.bw * n_gpu / 8 # 带宽是按 8 卡机器给的；1 块 GPU 只分到 1/8 ld = model.kv_bytes_per_token() / bw cells.append(f"{fmt_time(ld)}（{rc / ld:.1f}×）") rows.append(cells) print_table(["模型与硬件", "每 token 的 KV", "重算（8K 提示的平均）"] + [t.name for t in TIERS[1:]], rows) print("（括号里是\"重算 / 读回\"：大于 1 表示读回更快。带宽是一台 8 卡机器的合计，1 块 GPU 按 1/8 算；都是假设值。）") print("\n一个 8.7K token 的请求，前一半命中缓存，prefill 要多久（LLaMA2-70B，8 块 A800）：") rows2 = [] n_in, n_hit = 8704, 4352 rows2.append(["不用缓存（全部重算）", fmt_time(prefill_compute_time(LLAMA2_70B, n_in, 0))]) for tier in TIERS[1:]: rows2.append([f"从{tier.name}读，按层重叠", fmt_time(prefill_time(LLAMA2_70B, n_in, n_hit, tier))]) rows2.append([f"从{tier.name}读，不重叠", fmt_time(prefill_time(LLAMA2_70B, n_in, n_hit, tier, overlap=False))]) print_table(["做法", "prefill 时间"], rows2) print_observations([ "LLaMA2-70B 每个 token 的 KV 只有 320 KB（GQA 把 KV 头从 64 减到 8），重算一个 token 却要约 140 GFLOP：" "从哪一层存储读回来都比重算快一个数量级以上——这就是 Mooncake 说的\"用存储换计算\"。", "OPT-13B 是 MHA，每个 token 的 KV 有 819 KB，计算量却只有 70B 的五分之一：读回来只比重算快几倍，从 SSD 读反而比重算慢。" "所以 KV 缓存值不值得放到慢的存储上，取决于模型结构——GQA、MLA 这类\"KV 小\"的设计让分层缓存更划算（第 2 讲 §6）。", "按层重叠（第 l 层算的时候读第 l+1 层的 KV，Mooncake §5.2）让读 KV 的时间几乎完全藏在计算后面。", ]) # ----------------------------------------------------------------------------- def demo2_scheduling(): print_header("演示 2：prefill 请求分给谁（Mooncake 算法 1、图 8）——8 个 prefill 实例，各有 3,000 块的本地缓存") names = {"random": "随机", "load": "按负载（队列最短）", "cache": "按缓存（估计 TTFT 最小）", "kvcentric": "按缓存 + 热点 KV 复制（KV-centric）"} for rate in (4, 8): tr = make_trace(seed=0, rate=rate) print(f"\n会话到达速率 {rate}/s（平均 {len(tr) / tr[-1].arrival:.1f} 个请求/s），TTFT 的 SLO 定为 5 秒") rows = [] for pol in ("random", "load", "cache", "kvcentric"): r = prefill_cluster(tr, policy=pol, ttft_slo=5.0) rows.append([names[pol], f"{r['mean']:.2f}", f"{r['p90']:.2f}", f"{r['attain']:.1%}", f"{r['hit']:.0%}", f"{r['computed'] / 1e6:.1f} M", f"{r['fetched'] / 1e6:.1f} M"]) print_table(["调度策略", "平均 TTFT（秒）", "P90", "TTFT ≤ 5 s 的比例", "命中的输入 token", "实际算的 token", "从别的实例拉的 token"], rows) print_observations([ "随机和按负载分配：同一篇文档的请求散到 8 个实例上，每个实例的缓存都只命中一小部分（约 27%），要算的 token 多，" "负载一高就排长队。", "按缓存分配（估计每个实例的 TTFT = 排队 + 用它本地命中的前缀做 prefill，选最小的）：同一篇文档的后续请求自然落到" "算过它的那个实例上，命中率翻倍，要算的 token 少了约 40%——负载高时 TTFT 差一个数量级。", "再加上 KV 复制：如果别的实例命中的前缀长得多、而那个实例很忙，就把 KV 通过 RDMA 拉过来（比重算快几十倍），" "顺便存进本地——热门的前缀自动复制到多台机器上。TTFT 再降一些，代价是网络流量。和论文图 8 的排序一致。", ]) # ----------------------------------------------------------------------------- def demo3_overload(): print_header("演示 3：过载时提前拒绝（Mooncake §7、表 3）——8 个 prefill 实例 + 8 个 decode 实例，请求到达速率 = decode 容量的 2 倍") base = make_trace(seed=0) occ = float(np.mean([(r.input_len + r.output_len) * r.output_len * 0.05 for r in base])) dec_rate = 8 * 60_000 / occ tr = iid_requests(base, 6000, 2 * dec_rate) end = tr[-1].arrival print(f"decode 池大约每秒能完成 {dec_rate:.1f} 个请求；平稳的泊松到达 {2 * dec_rate:.1f} 个/秒，共 {len(tr)} 个" "（用平稳到达，是为了把\"负载本身的波动\"和\"调度策略引起的波动\"分开）") rows = [] series = {} names = {"baseline": "基线：prefill 后才检查 decode", "early": "提前拒绝：到达时看现在的 decode 负载", "predict": "按预测提前拒绝：估计 prefill 完成时的 decode 负载"} for pol in ("baseline", "early", "predict"): r = overload_cluster(tr, policy=pol) k = int(end / 5.0) d, p = r["decode_load"][20:k], r["prefill_load"][20:k] series[pol] = (d, p) rows.append([names[pol], r["completed"], r["rejected"], r["rejected_after_prefill"], f"{r['wasted_prefill_s']:.0f}", f"{d.mean():.0%}", f"{np.corrcoef(d, p)[0, 1]:+.2f}"]) print_table(["策略", "完成", "拒绝", "其中 prefill 后才拒绝", "白做的 prefill（秒）", "decode 平均占用", "prefill 与 decode 负载的相关系数"], rows) print("\n前 60 秒的负载（每 5 秒一个采样；P = prefill 实例忙的比例，D = decode 池的占用）：") rows2 = [] for pol in ("baseline", "early"): d, p = series[pol] rows2.append([pol, "P"] + [f"{x:.2f}" for x in p[:12]]) rows2.append(["", "D"] + [f"{x:.2f}" for x in d[:12]]) print_table(["策略", ""] + [f"{5 * i}s" for i in range(12)], rows2) print_observations([ "基线：decode 池一直接近满，prefill 照常接请求，做完了才发现 decode 放不下——约三成多的请求 prefill 白做了。", "提前拒绝：到达时就看 decode 池现在满不满，白做的 prefill 减少约 70%。但**完成的请求几乎一样多**：decode 池才是瓶颈，" "完成数由它决定；省下来的是 prefill 的算力（可以少开几台 prefill 机器，或者去做别的事）。", "提前拒绝的副作用：prefill 和 decode 的负载变成**反相**（相关系数为负）。decode 满时拒绝一切、prefill 闲下来；" "decode 空出来以后一下子接很多，等它们做完 prefill 涌进 decode，decode 又满了——论文图 9、10a 描述的振荡。", "论文的解法是预测 prefill 完成时 decode 的负载。我们按论文的办法假设每个请求的 decode 时间相同，" "在这个输出长度长尾分布的负载上预测不准：白做的 prefill 只少了约三分之一，完成数还少了约 5%。" "论文报告的拒绝数从 4,183（基线）降到 3,771（提前）和 3,589（按预测），我们的玩具复现不出\"完成更多\"这一点——" "它的 SLO 统计、decode 的接纳条件、真实负载的形状都和这里不同。", ]) if __name__ == "__main__": demo1_tiers() demo2_scheduling() demo3_overload()