""" ============================================================================= CS349E · 第 5 讲 · 04_placement_and_kv.py —— DistServe 的放置搜索（70B 模型）与 KV cache 的搬运账 ============================================================================= 【这个脚本要回答的问题】 解耦的第二层好处：prefill 和 decode 可以各自选**不同的并行方式**。DistServe 的算法 1（高带宽集群）： 对 prefill 和 decode 分别枚举 (张量并行度, 流水线段数)，用模拟器 + 二分查找求每种配置的 goodput， 各自选"每块 GPU 的 goodput 最高"的那个；再按目标流量 R 复制足够多的实例： prefill 实例数 = ⌈R / prefill goodput⌉，decode 实例数 = ⌈R / decode goodput⌉。 代价是把 KV cache 从 prefill 实例搬到 decode 实例。 演示 1：Llama-3.1-70B、H100（每台机器 8 块），聊天类负载：枚举配置、求 goodput、选最优，和共置系统比较。 演示 2：KV 的搬运账：每个请求要搬多少字节（复现论文 OPT-66B 的 1.13 GB / 90 Gbps）；在不同链路上要多久，和 prefill 比； 带宽不够时，DistServe 的算法 2 怎么办（同一流水线段的 prefill 和 decode 放在同一台机器上，走 NVLink）。 【怎么读输出】每步时间是 roofline 模型（假设值），负载是合成的（提示约 1000、输出约 250，对数正态）。goodput 有约 ±10% 的 模拟噪声。结果是确定性的，每次运行相同。 【运行】python3 04_placement_and_kv.py （约 30 秒——它要对十几种配置各做一次二分查找） 【依赖】numpy；同目录的 serving_sim；../prereq5/comm_sim；../lecture01 的 console、hw_specs、llm_math。 """ import math import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) sys.path.insert(0, os.path.join(HERE, "..", "prereq5")) from comm_sim import ETH100, IB_NDR, NVLINK4, LinkSpec # noqa: E402 from console import fmt_bytes, 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, ModelConfig # noqa: E402 from serving_sim import (StepModel, attainment, lognormal_len, make_workload, max_rate, run_colocated, # noqa: E402 run_decode_only, run_disaggregated, run_prefill_only) M = LLAMA31_70B PROMPT = lognormal_len(1000, 1.0, 16, 8000) OUTPUT = lognormal_len(250, 1.0, 2, 2000) TTFT, TPOT = 1.0, 0.05 CONFIGS = [(tp, pp) for tp in (2, 4, 8) for pp in (1, 2) if tp * pp <= 8] def wl(n, rate): return make_workload(n, rate, PROMPT, OUTPUT, seed=0) # --------------------------------------------------------------------------- # 演示 1：放置搜索 # --------------------------------------------------------------------------- def demo1_search(): print_header(f"演示 1：DistServe 算法 1——Llama-3.1-70B（BF16）、H100，SLO：TTFT ≤ {TTFT:.1f} s、TPOT ≤ {TPOT * 1e3:.0f} ms") rows = [] best_p = best_d = (0.0, None) for tp, pp in CONFIGS: st = StepModel(M, H100, tp=tp, pp=pp) cap = st.kv_capacity_tokens() gp = max_rate(lambda r: attainment(run_prefill_only(wl(2000, r), st), TTFT, None), hi=500, iters=10) gd = max_rate(lambda r: attainment(run_decode_only(wl(5000, r), st), None, TPOT), hi=2000, iters=10) \ if cap >= 50000 else 0.0 n = tp * pp rows.append([f"TP{tp} × PP{pp}", n, f"{cap / 1e3:,.0f} K", f"{gp:.2f}", f"{gp / n:.2f}", f"{gd:.1f}" if gd else "KV 太少", f"{gd / n:.2f}" if gd else "-"]) if gp / n > best_p[0]: best_p = (gp / n, (tp, pp), gp) if gd / n > best_d[0]: best_d = (gd / n, (tp, pp), gd) print("每种配置单独作为 prefill 实例 / decode 实例时的 goodput（请求/秒）：") print_table(["配置", "GPU 数", "KV 容量（token）", "prefill goodput", "每 GPU", "decode goodput", "每 GPU"], rows) R = 100.0 (gp_gpu, (ptp, ppp), gp), (gd_gpu, (dtp, dpp), gd) = best_p, best_d n_p, n_d = math.ceil(R / gp), math.ceil(R / gd) gpus = n_p * ptp * ppp + n_d * dtp * dpp print(f"\n最好的 prefill 配置：TP{ptp} × PP{ppp}（每 GPU {gp_gpu:.2f}）；最好的 decode 配置：TP{dtp} × PP{dpp}（每 GPU {gd_gpu:.2f}）") print(f"目标流量 R = {R:.0f} 请求/秒 → ⌈{R:.0f} / {gp:.1f}⌉ = {n_p} 个 prefill 实例 + ⌈{R:.0f} / {gd:.1f}⌉ = {n_d} 个 decode 实例 " f"= {gpus} 块 GPU（按目标流量算每 GPU {R / gpus:.2f} 请求/秒）") p_st = StepModel(M, H100, tp=ptp, pp=ppp) d_st = StepModel(M, H100, tp=dtp, pp=dpp) full = max_rate(lambda r: attainment(run_disaggregated(wl(5000, r), p_st, d_st, n_p, n_d), TTFT, TPOT), hi=500, iters=10) print(f"用完整模拟检查这个部署（含 KV 传输）：goodput {full:.1f} 请求/秒（目标 {R:.0f}）") rows = [] for tp in (4, 8): st = StepModel(M, H100, tp=tp) for policy, kw in (("prefill 优先", {}), ("分块（预算 1024）", {"policy": "chunked", "chunk_budget": 1024})): g = max_rate(lambda r: attainment(run_colocated(wl(4000, r), st, **kw), TTFT, TPOT), hi=500, iters=10) rows.append([f"共置 TP{tp}，{policy}", tp, f"{g:.2f}", f"{g / tp:.2f}"]) rows.append([f"解耦：{n_p}×(TP{ptp}×PP{ppp}) + {n_d}×(TP{dtp}×PP{dpp})", gpus, f"{full:.2f}", f"{full / gpus:.2f}"]) print() print_table(["系统", "GPU 数", "goodput（请求/秒）", "每 GPU"], rows) print_observations([ "两个阶段对配置的**排序**不一样：prefill 只要放得下权重就行，TP2×PP2、TP4、TP2 这些 KV 很少的配置也排得靠前；" "decode 要 KV 容量（批开得大）和读权重的带宽，TP2 根本做不了，TP4 也明显不如 8 卡的配置。" "在这个模型里两者的第一名恰好都是 TP4×PP2；论文里 OPT-175B 的 prefill 选了 3×3、decode 选了 3×4，并不相同。", "每块 GPU 的 decode 能力约是 prefill 的 5 倍：按流量配比时 prefill 实例要多得多。共置系统里这个比例是写死的（每块 GPU 两种活都干）。", "在这个设置下，解耦部署的每 GPU goodput 比共置（含分块 prefill）高约 10%–40%。注意两点：实例数要取整（⌈R / goodput⌉），" "目标流量越大，取整的浪费越小；完整模拟的 goodput 高于按单实例估算的目标 R，部分来自多实例的排队合并效应（02 脚本）。" "而且解耦部署一共 48 块 GPU，共置只模拟了 4–8 块的一个实例，比较对共置不完全公平。" "另外，我们的共置模拟只支持张量并行（不支持流水线），和论文里 vLLM 的情况一样。", ]) # --------------------------------------------------------------------------- # 演示 2：KV 的搬运账 # --------------------------------------------------------------------------- def demo2_kv_transfer(): print_header("演示 2：KV cache 的搬运账——每个请求多少字节、要多快的网络、和 prefill 比要多久") opt66 = ModelConfig("OPT-66B", n_layers=64, d_model=9216, n_heads=72, n_kv_heads=72, d_head=128, d_ff=36864, vocab=50272) kv512 = 512 * opt66.kv_bytes_per_token(2) print(f"复现论文 §3.3：OPT-66B（MHA）一个 512 token 的请求，KV = {fmt_bytes(kv512)} = {kv512 / 2 ** 30:.2f} GiB；" f"每秒 10 个请求要 {kv512 * 10 / 1e9:.1f} GB/s = {kv512 * 10 * 8 / 1e9:.0f} Gb/s（论文写作 1.13 GB、90 Gbps，" f"大约是把 GiB 当作 GB 来算的）") links = [("NVLink 4（机器内）", NVLINK4), ("InfiniBand NDR（每块 GPU 一张网卡）", IB_NDR), ("8 张网卡一起用", LinkSpec("8×IB", 8 * IB_NDR.bw, IB_NDR.alpha)), ("100 Gb 以太网", ETH100)] rows = [] for m, tp in ((LLAMA31_8B, 1), (M, 8)): st = StepModel(m, H100, tp=tp) for L in (1000, 8000): kv = L * m.kv_bytes_per_token(2) pre = st.prefill_time([L]) cells = [m.name, L, fmt_bytes(kv), fmt_time(pre)] for _, link in links: t = link.alpha + kv / link.bw cells.append(f"{fmt_time(t)}（{t / pre * 100:.0f}%）") rows.append(cells) print("\n一个请求的 KV（BF16）在不同链路上的搬运时间（括号里是相对 prefill 时间的比例；prefill 用单卡（8B）或 TP8（70B））：") print_table(["模型", "提示长度", "KV 大小", "prefill 时间"] + [name for name, _ in links], rows) R = 30 need = R * 1000 * M.kv_bytes_per_token(2) print(f"\nLlama-3.1-70B、平均提示 1000 个 token、每秒 {R} 个请求：KV 搬运要 {need / 1e9:.1f} GB/s = {need * 8 / 1e9:.0f} Gb/s" f"（一张 400 Gb/s 的 InfiniBand 网卡就够）") print_observations([ "GQA 让 KV 小了很多：Llama-3.1-70B 每 token 320 KiB，OPT-66B（MHA）每 token 2.25 MiB——同样 512 个 token，70B 只要 0.17 GB。" "第 2 讲的'为推理而改的注意力'在解耦里又省了一次：搬得少。", "KV 搬运时间和 prefill 时间同比例增长（都 ∝ 提示长度），所以比例几乎是常数：在 InfiniBand 上约是 prefill 时间的几到十几个百分点，" "8 张网卡一起用或者走 NVLink 时只有 1%–2%。论文的测量：OPT-175B 的 KV 传输不到总延迟的 0.1%（它们是同一台机器内走 NVLink）。", "100 Gb 以太网就紧张了：搬运要 prefill 时间的 36%–62%（TTFT 之后 decode 还要再等这么久才能开始），" "而且一张网卡每秒只能搬约 38 个 70B 请求（1000 token）的 KV。这时 DistServe 的算法 2（低带宽集群）把模型按层切成流水线段，" "把**同一段**的 prefill 实例和 decode 实例放在同一台机器上——KV 只在对应的层之间传，于是全部走机器内的 NVLink。", "【系统联系】KV 的搬运量正是第 7 讲 Mooncake 这类'以 KV cache 为中心'的架构要管理的东西：KV 不只是在 prefill 和 decode 之间搬，" "还可以缓存起来给后续请求复用。", ]) if __name__ == "__main__": demo1_search() demo2_kv_transfer()