""" ============================================================================= CS349E · 第 15 讲 · 05_agent_serving_sim.py —— 很多个智能体一起跑：工具调用期间 KV 处理策略的服务模拟 ============================================================================= 【这个脚本要回答的问题】 04 脚本算的是"一次工具调用"的账。真实的服务里，几十上百个智能体会话同时在跑，它们的 KV 抢同一块显存。 哪种处理策略让同一块 GPU 每小时完成最多的会话？ 演示 1：负载：一个写内核的智能体会话长什么样；显存装得下多少。 演示 2：7 种策略 × 3 种负载（32 / 96 / 160 个同时在跑的会话）：吞吐、会话完成时间、每轮等 LLM 的时间、重算的比例。 演示 3：换到 CPU 依赖 PCIe：带宽降到 1/4、1/16 时，同步和异步的换入换出各怎样。 演示 4：如果工具调用都很快（0.3 秒），结论怎样随会话数变化。 【关于数字】Llama-3.1-8B BF16 在一块 H100 上；每次迭代的时间来自第 5 讲 StepModel 的 roofline 模型（不是实测）； 工具耗时的分布、每轮的 token 数是假设值（agent_lab.TOOL_MIX、AgentWorkload）。模拟 1200 秒，前 200 秒不计入统计； 固定随机种子，每次运行结果相同。 【运行】python3 05_agent_serving_sim.py （约 15 秒） 【依赖】numpy；本目录的 agent_lab；../lecture01 的 console。 """ import math import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from agent_lab import (KV_BYTES, PCIE_BW, POLICIES, STEP, TOOL_MIX, AgentServer, AgentWorkload) # noqa: E402 from console import print_header, print_observations, print_table # noqa: E402 LOADS = (32, 96, 160) NAMES = {"discard": "丢掉重算", "prefix_cache": "前缀缓存（可被挤掉）", "preserve": "一直保留", "ttl": "TTL = T*", "ttl_queue": "TTL，有人排队就不保留", "swap_sync": "同步换到 CPU", "swap_async": "异步换到 CPU"} def demo1_workload(): print_header("演示 1：负载与容量") wl = AgentWorkload() mean_tool = sum(m[1] * m[2] for m in TOOL_MIX) / sum(m[1] for m in TOOL_MIX) final_ctx = wl.prompt + wl.turns * wl.decode + (wl.turns - 1) * wl.tool_out avg_ctx = wl.prompt + (wl.turns / 2) * (wl.decode + wl.tool_out) cap = STEP.kv_capacity_tokens() rows = [["系统提示 + 任务", f"{wl.prompt:,} token"], ["轮数", f"{wl.turns}（每轮生成 {wl.decode} token，工具输出 {wl.tool_out} token）"], ["工具耗时的分布", "；".join(f"{m[0]} {m[1]:.0%}、平均 {m[2]:g} s" for m in TOOL_MIX) + f"（加权平均 {mean_tool:.1f} s）"], ["一个会话结束时的上下文", f"{final_ctx:,} token = {final_ctx * KV_BYTES / 1e9:.2f} GB"], ["一个会话花在等工具上的时间", f"约 {(wl.turns - 1) * mean_tool:.0f} s"], ["显存能放的 KV（80 GB × 90% − 权重）", f"{cap:,} token = {cap * KV_BYTES / 1e9:.1f} GB"]] for n in LOADS: rows.append([f"{n} 个会话的 KV 全部留在显存（按平均上下文 {avg_ctx:,.0f}）", f"{n * avg_ctx * KV_BYTES / 1e9:.1f} GB"]) print_table(["项目", "数值"], rows, aligns=["<", "<"]) print_observations([ f"一个会话的大部分时间在等工具（约 {(wl.turns - 1) * mean_tool:.0f} s），LLM 真正在为它算的时间只占一小部分——" "这段时间里，它的 KV 是\"冷\"的。", f"32 个会话的 KV 全放进显存没问题；96 个、160 个就放不下了（显存只有 {cap * KV_BYTES / 1e9:.0f} GB 给 KV）。" "显存不够时，\"谁的 KV 留在显存\"直接决定吞吐。", ]) def demo2_policies(): print_header("演示 2：7 种策略 × 3 种负载") rows, res = [], {} for n in LOADS: for p in POLICIES: m = AgentServer(p, AgentWorkload(agents=n)).run() res[(n, p)] = m jct = "—" if math.isnan(m["jct_mean"]) else f"{m['jct_mean']:.0f} s" rows.append([str(n), NAMES[p], f"{m['sessions_per_hour']:.0f}", jct, f"{m['turn_wait_mean']:.1f} s", f"{m['recompute_frac']:.0%}", f"{m['paused_gb']:.1f}", str(m["forced_evictions"])]) print_table(["会话数", "策略", "每小时完成的会话", "平均完成时间", "每轮等 LLM", "prefill 里重算的比例", "暂停中的会话占的显存（GB）", "被迫丢掉的次数"], rows) best = {n: max(POLICIES, key=lambda p: res[(n, p)]["sessions_per_hour"]) for n in LOADS} r32 = res[(32, "preserve")] print_observations([ f"32 个会话（显存够）：保留、前缀缓存、两种 TTL 完全一样（{r32['sessions_per_hour']:.0f} 个/小时），因为没有人被挤出显存；" f"\"丢掉重算\"少了 {1 - res[(32, 'discard')]['sessions_per_hour'] / r32['sessions_per_hour']:.0%}，" f"它的 prefill 里 {res[(32, 'discard')]['recompute_frac']:.0%} 都是在重算已经算过的上下文。" f"\"一直保留\"平均占着 {r32['paused_gb']:.1f} GB 显存给暂停中的会话——将近一半，和 InferCept 观察到的现象一样" "（超过 60% 的时间里，近一半的显存被暂停的请求占着）。", f"96 个会话（显存不够）：\"一直保留\"崩了（{res[(96, 'preserve')]['sessions_per_hour']:.0f} 个/小时）——显存被暂停的会话钉满，" "排在队首的会话进不来，整个引擎经常空转，直到什么都跑不动时才被迫丢掉一些。" f"按 InferCept 公式算的 TTL（{res[(96, 'ttl')]['sessions_per_hour']:.0f}）也被拖累：其他请求越多，公式算出的 T* 越长，越倾向于保留，" f"可此时正是显存最紧的时候；加上\"有人排队就不保留\"（{res[(96, 'ttl_queue')]['sessions_per_hour']:.0f}）好一些，" f"但仍不如从不保留、只留可被挤掉的缓存（{res[(96, 'prefix_cache')]['sessions_per_hour']:.0f}）。", f"换到 CPU 在这三种负载下都是最好或接近最好（最好的策略：" + "、".join(f"{n} 个会话 {NAMES[best[n]]}" for n in LOADS) + f"）：PCIe 5.0 上换 1 GB 只要二十毫秒，CPU 内存相当于一块大得多的 KV 存储；异步（和计算重叠）比同步好 " f"{res[(96, 'swap_async')]['sessions_per_hour'] / res[(96, 'swap_sync')]['sessions_per_hour'] - 1:.0%}（96 个会话）。", "160 个会话时，除了换到 CPU，所有策略都被显存卡住：大部分 prefill 都在重算（八九成），" "GPU 的时间花在了\"重新读一遍历史\"上。", ]) return res def demo3_pcie(): print_header("演示 3：换到 CPU 依赖 PCIe 的带宽（96 个会话）") rows = [] for frac in (1.0, 0.25, 1 / 16): bw = PCIE_BW * frac r = [f"{bw / 1e9:.1f} GB/s"] for p in ("swap_sync", "swap_async"): m = AgentServer(p, AgentWorkload(agents=96), pcie_bw=bw).run() r.append(f"{m['sessions_per_hour']:.0f}") rows.append(r) pc = AgentServer("prefix_cache", AgentWorkload(agents=96)).run()["sessions_per_hour"] print_table(["有效带宽", "同步换到 CPU", "异步换到 CPU"], rows) print(f"对照：前缀缓存（不用 PCIe）{pc:.0f} 个/小时") print_observations([ "带宽降到 1/4（例如 8 块 GPU 共用到 CPU 的链路），异步换入换出几乎不受影响，同步的掉了很多：同步时每一次传输都让整个引擎停下来等。", "降到 1/16，异步的也撑不住了——传输排起了长队，会话等着换入。这时\"换到 CPU\"不再比\"重算\"便宜。" "第 7 讲的 Mooncake 用 RDMA 把 KV 放到整个集群的 DRAM / SSD 里，也是在算这笔\"搬 vs 算\"的账。", "结论依赖于模型和硬件：70B 模型每 token 的 KV 是 8B 的 2.5 倍（GQA 都是 8 个 KV 头，层数 80 vs 32），" "张量并行切开之后每块 GPU 的份额又变小；上下文越长，搬得越多。具体的系统要按自己的参数算。", ]) def demo4_fast_tools(): print_header("演示 4：工具都很快时（全部是平均 0.3 秒的读文件 / 查询），随会话数变化") fast = [("快工具", 1.0, 0.3, 0.5)] pols = ("discard", "prefix_cache", "preserve", "ttl", "swap_async") rows, res = [], {} for n in (32, 48, 64): r = [str(n)] for p in pols: m = AgentServer(p, AgentWorkload(agents=n), tool_mix=fast).run() res[(n, p)] = m["sessions_per_hour"] r.append(f"{m['sessions_per_hour']:.0f}") rows.append(r) print_table(["会话数"] + [NAMES[p] for p in pols], rows) print("（每小时完成的会话数）") print_observations([ f"48 个会话时，保留（{res[(48, 'preserve')]:.0f}）和 TTL（{res[(48, 'ttl')]:.0f}）明显好于前缀缓存（{res[(48, 'prefix_cache')]:.0f}）：" "工具 0.3 秒就回来，可被挤掉的缓存刚被挤掉，它的主人就回来了，只好整个重算——这正是 InferCept 说的" "\"短拦截应该保留\"。工具慢时（演示 2），32 个会话下几种做法完全一样。", f"64 个会话时，KV 的总量超过了显存，保留（{res[(64, 'preserve')]:.0f}）和 TTL（{res[(64, 'ttl')]:.0f}）又崩了：钉住的 KV 让排在队首的会话进不来。" "保留在\"放得下\"时最好，在\"放不下\"的那一刻掉下悬崖。", f"异步换到 CPU 在三种会话数下都最好（{res[(48, 'swap_async')]:.0f}、{res[(64, 'swap_async')]:.0f}）。" "没有一种\"保留 / 丢掉\"的规则在所有负载下都对：工具耗时的分布、显存的余量、PCIe 的带宽一起决定答案——" "这正是 InferCept（按请求算浪费、动态选择）和 Continuum（预测工具耗时、把排队的代价算进 TTL）在做的事。", ]) if __name__ == "__main__": demo1_workload() demo2_policies() demo3_pcie() demo4_fast_tools()