""" ============================================================================= CS349E · 第 6 讲 · 01_continuous_batching.py —— 请求级调度 vs 迭代级调度（Orca） ============================================================================= 【这个脚本要回答的问题】 生成式模型的一个请求要跑很多次迭代（每次 1 个 token），不同请求的输入、输出长度差别很大。 传统的推理服务按"请求"调度：凑一批、整批跑完、再凑下一批。Orca（OSDI 2022）指出这有两个问题： 早结束的请求不能早返回（还在白算），新到的请求要等整批结束才能开始。 Orca 的做法：每次**迭代**都重新选请求（iteration-level scheduling，也叫连续批处理 continuous batching）。 演示 1：请求级批处理白算了多少——输入要补齐到最长的那个，早结束的请求还占着一行。 演示 2：延迟–吞吐曲线（Orca 图 10 的模拟版）：GPT-3 13B、一块 A100 40GB、Orca 的合成负载。 演示 3：所有请求一样长时（Orca 图 11）：请求级调度的"早结束"问题消失，差距缩小。 演示 4：KV cache 的预留：按模型最大长度预留（FasterTransformer）、按请求声明的上限预留（Orca 算法 1）、 不预留（会死锁）。这是第 7 讲 PagedAttention 的出发点。 【怎么读输出】每张表一行是一种调度配置或一个到达速率；"每 token 延迟"= 请求的端到端延迟 / 生成的 token 数 （Orca 图 10 的纵轴），取中位数。表后有"观察 Observations"。 【模型的假设】每步时间来自第 5 讲的 StepModel（roofline、60% 峰值算力、80% 峰值带宽、每步 1 ms 固定开销）， 是估计值，不是测量值。只模拟单块 GPU。所有输出每次运行完全相同。 【运行】python3 01_continuous_batching.py （约 3 秒） 【依赖】numpy；本目录的 decode_lab（它会导入 ../lecture01、../lecture02、../lecture05、../prereq5 的模块）。 """ import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) from decode_lab import (A100_40G, GPT3_13B, StepModel, latency_throughput, orca_workload, # noqa: E402 run_iteration_level, run_request_level) from console import fmt_bytes, print_header, print_observations, print_table # noqa: E402 STEP = StepModel(GPT3_13B, A100_40G) # 一块 A100 跑 GPT-3 13B（Orca 表 1 的第一行） N_REQ = 1000 def run(kind, bs, rate, seed=0, **wl): """ 用 kind（"ft" = 请求级，"orca" = 迭代级）和最大批大小 bs 跑一次 N_REQ 个请求的模拟。 【返回】(请求列表, 调度器返回的统计字典) 【Python 小知识】**wl 把多余的关键字参数收集成字典，再原样传给 orca_workload（例如 prompt_range=...）。 """ reqs = orca_workload(N_REQ, rate, seed=seed, **wl) if kind == "ft": st = run_request_level(reqs, STEP, bs) else: st = run_iteration_level(reqs, STEP, bs) return reqs, st # ----------------------------------------------------------------------------- def demo1_waste(): print_header("演示 1：请求级批处理白算了多少（Orca 图 3）") m = GPT3_13B kv_tok = m.kv_bytes_per_token() print(f"模型 {m.name}：参数 {m.n_params() / 1e9:.2f} B（FP16 权重 {fmt_bytes(m.weight_bytes('fp16'))}），" f"每个 token 的 KV {fmt_bytes(kv_tok)}；A100 40GB 放下权重后还能放 {STEP.kv_capacity_tokens():,} 个 token 的 KV") print(f"FasterTransformer 给每个请求按最大长度 2048 预留 KV → 一批最多 {STEP.kv_capacity_tokens() // 2048} 个请求" "（论文 §6.1：13B 模型用不了 8 以上的批）\n") rows = [] for rate in (1.0, 2.0, 4.0): _, st = run("ft", 8, rate) rows.append([f"{rate:.0f} req/s", st["batches"], f"{N_REQ / st['batches']:.2f}", f"{st['slot_tokens']:,}", f"{st['useful_tokens']:,}", f"{1 - st['useful_tokens'] / st['slot_tokens']:.1%}"]) print_table(["到达速率", "批数", "平均批大小", "算过的 token 位置", "真正需要的", "白算的比例"], rows) print("\n（算过的 token 位置 = 每批 B × 最长输入 + B × (最长输出 − 1)；真正需要的 = Σ(输入 + 输出 − 1)。）") print_observations([ "负载越高、批越满，白算的比例越大（27% → 39%）：一批里输入长度 U(32, 512)、输出长度 U(1, 128) 差别很大，" "要补齐到最长的输入，还要陪着最长的输出跑到底。", "白算的 decode 行在 decode 阶段几乎不花额外时间（受带宽限制，多算几行几乎免费）——" "真正的损失是**时间**：早结束的请求不能返回、新请求进不来、批也凑不大（预留 2048 限制到 6）。", ]) # ----------------------------------------------------------------------------- def demo2_curves(): print_header("演示 2：延迟–吞吐曲线（Orca 图 10 的模拟版：GPT-3 13B、1 块 A100、输入 U(32,512)、输出 U(1,128)）") configs = [("ft", 1), ("ft", 8), ("orca", 8), ("orca", 16), ("orca", 32)] rates = (0.5, 1, 2, 4, 6, 8, 12, 16) res = {} rows = [] for rate in rates: row = [f"{rate:g}"] for kind, bs in configs: reqs, _ = run(kind, bs, rate) thr, med, _ = latency_throughput(reqs) res[(kind, bs, rate)] = (thr, med) row.append(f"{thr:5.2f} / {med * 1e3:6.0f}") rows.append(row) heads = ["到达 req/s"] + [f"{'请求级' if k == 'ft' else '迭代级'} bs={b}" for k, b in configs] print("每格：吞吐 req/s / 每 token 延迟中位数 ms") print_table(heads, rows) base = res[("orca", 32, 0.5)][1] target = 2 * base print(f"\n空载时每 token 延迟约 {base * 1e3:.0f} ms。要求\"每 token 延迟中位数 ≤ {target * 1e3:.0f} ms（空载的 2 倍）\"，" "用二分查找找出每个配置能承受的最高到达速率，以及那时的吞吐：") rows2 = [] best = {} for kind, bs in configs: rate = max_rate_under(kind, bs, target) thr = latency_throughput(run(kind, bs, rate)[0])[0] if rate > 0 else 0.0 best[(kind, bs)] = thr peak = max(res[(kind, bs, r)][0] for r in rates) rows2.append([f"{'请求级' if kind == 'ft' else '迭代级'} bs={bs}", f"{rate:.2f}", f"{thr:.2f}", f"{peak:.2f}"]) print_table(["配置", "延迟达标的最高速率", "那时的吞吐 req/s", "表中最高吞吐（不管延迟）"], rows2) ratio = best[("orca", 32)] / max(best[("ft", 8)], 1e-9) print_observations([ "请求级调度：bs=1 时每个请求独占 GPU，吞吐不到 0.7 req/s；bs=8（实际被显存限制到 6）最高约 1.9 req/s，" "但到达速率一过 1 req/s，每 token 延迟就从几十毫秒涨到几百、几千毫秒——新请求要排队等整批结束。", f"迭代级调度：在同样的延迟要求下，bs=32 的吞吐是请求级 bs=8 的 {ratio:.0f} 倍；" "加大 max_bs 主要提高过载点，低负载时延迟几乎不变（Orca 图 10 的现象）。", "表里过载以后\"吞吐\"略有下降：有限次模拟里，最后一批请求排队排得最久，统计窗口被拉长。只看过载点之前的部分。", "Orca 论文在 175B、16 块 GPU（含流水线并行）上、按\"每 token 190 ms\"比较，报告 36.9 倍。" "我们只模拟单卡 13B、没有流水线，倍数取决于延迟要求定在哪里，方向一致。", ]) def max_rate_under(kind, bs, target, lo=0.05, hi=64.0, iters=14): """ 二分查找：每 token 延迟中位数 ≤ target 时，最高能承受多大的到达速率。 【思路】速率越高、排队越久、延迟越大（近似单调），所以可以二分。在对数尺度上二分（几何平均）， 因为速率跨好几个数量级。 【返回】速率（req/s）；连 lo 都达不到时返回 0 """ if latency_throughput(run(kind, bs, lo)[0])[1] > target: return 0.0 for _ in range(iters): mid = (lo * hi) ** 0.5 if latency_throughput(run(kind, bs, mid)[0])[1] <= target: lo = mid else: hi = mid return lo # ----------------------------------------------------------------------------- def demo3_homogeneous(): print_header("演示 3：所有请求一样长时（Orca 图 11）——请求级调度的劣势变小") rows = [] for (pin, gen) in ((32, 32), (256, 256)): wl = dict(prompt_range=(pin, pin), gen_range=(gen, gen)) for kind, bs in (("ft", 8), ("orca", 8), ("orca", 32)): peak = 0.0 for rate in (1, 2, 4, 8, 16, 32): reqs, _ = run(kind, bs, rate, **wl) peak = max(peak, latency_throughput(reqs)[0]) rows.append([f"输入 {pin}、输出 {gen}", f"{'请求级' if kind == 'ft' else '迭代级'} bs={bs}", f"{peak:.2f}"]) print_table(["负载", "配置", "最高吞吐 req/s"], rows) print_observations([ "请求一样长时没有\"早结束\"和\"补齐\"的浪费：同样 bs=8，迭代级只比请求级高 26%–37%" "（剩下的差别来自\"新请求要等整批结束\"，以及迭代级的批不受 6 的限制）。", "迭代级调度在 bs=32 时仍然高得多——但这主要来自更大的批（请求级被 2048 的预留限制在 6），不是调度本身。", "所以 Orca 的两个贡献要分开看：迭代级调度解决\"长短不一\"，按请求上限预留 KV 解决\"批凑不大\"。", ]) # ----------------------------------------------------------------------------- def demo4_reservation(): print_header("演示 4：KV cache 怎么预留（每个请求声明 max_gen = 512，实际只生成 U(1, 128) 个就遇到 EOS）") print(f"KV 容量 {STEP.kv_capacity_tokens():,} 个 token；最大批大小 64；到达速率 4、8、12 req/s\n") rows = [] for rate in (4.0, 8.0, 12.0): for reserve in ("max_seq", "max_gen", "none"): reqs = orca_workload(N_REQ, rate, seed=0, max_gen=512) st = run_iteration_level(reqs, STEP, 64, reserve=reserve, max_seq=2048) if st["deadlock"]: note = f"死锁：t = {st['deadlock_t']:.1f} s 时谁都动不了，已完成 {st['done']} 个" thr = "-" else: thr_v, med, _ = latency_throughput(reqs) thr = f"{thr_v:.2f}" note = f"每 token 延迟中位数 {med * 1e3:.0f} ms" label = {"max_seq": "按模型最大长度 2048 预留", "max_gen": "按声明的 prompt+max_gen 预留（Orca）", "none": "不预留，用多少分配多少"}[reserve] rows.append([f"{rate:g}", label, f"{st['mean_reserved']:,.0f}", f"{st['mean_used']:,.0f}", thr, note]) print_table(["速率", "预留方式", "平均占用（预留）", "平均实际写入", "吞吐 req/s", "结果"], rows) print_observations([ "按 2048 预留：显存几乎全被预留占满，真正写入的 KV 不到六分之一；同时只能跑 6 个请求，4 req/s 时每 token 延迟" "已经 144 ms，8 req/s 时超过 1 秒。", "按 prompt + max_gen 预留（Orca 算法 1）好得多，但请求实际只生成声明上限的一小部分（平均 64 / 512），" "预留量约是实际写入量的 2.5 倍——这就是第 7 讲 PagedAttention 论文说的\"预留浪费\"。12 req/s 时它已过载（384 ms）。", "不预留：显存全给真正写入的 KV，8 req/s 时延迟最低（31 ms vs 46 ms）。但 12 req/s 时，显存被正在生成的请求占满，" "每个请求都还要再写 1 个 token 的 KV，谁都动不了——第 49.7 秒死锁，只完成了 516 个请求。没有抢占（换出 / 重算）就会这样。", "第 7 讲的答案：按块（page）分配，不预留；显存不够时抢占一部分请求（换出到 CPU 或丢掉以后重算）。", ]) if __name__ == "__main__": demo1_waste() demo2_curves() demo3_homogeneous() demo4_reservation()