"""
=============================================================================
CS349E · 第 7 讲 · 03_paged_serving.py —— 分页管理让服务快多少；抢占之后怎么恢复
=============================================================================

【这个脚本要回答的问题】
    02 看到分页让显存里真正存着 token 的比例从两三成提高到九成以上。换成用户和运营者关心的指标呢？
      演示 1：延迟–速率曲线（vLLM 论文图 12 的模拟版）：OPT-13B、一块 A100 40GB、ShareGPT 式和 Alpaca 式的负载。
      演示 2：在"平均每 token 延迟 ≤ 50 ms"的要求下，四种显存管理方式各能承受多大的请求速率。
      演示 3：分页的代价：过载时显存不够，要抢占（vLLM 默认丢掉 KV、之后重算），重算了多少。
      演示 4：被抢占的请求怎么恢复：从 CPU 内存换回（swap）还是重算（recompute）？和块大小的关系（论文图 19）。

【运行】python3 03_paged_serving.py      （约 50 秒：演示 2 对每种方式做二分查找）
【依赖】numpy；本目录的 kv_serving；../lecture01 的 console。
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

from kv_serving import (A100_40G, OPT_13B, StepModel, make_workload, max_rate, recompute_time,  # noqa: E402
                        simulate, swap_time)
from console import fmt_time, print_header, print_observations, print_table  # noqa: E402

STEP = StepModel(OPT_13B, A100_40G)
POLICIES = ("max", "pow2", "oracle", "paged")
NAMES = {"max": "Orca (Max)", "pow2": "Orca (Pow2)", "oracle": "Orca (Oracle)", "paged": "vLLM（块 16）"}


# -----------------------------------------------------------------------------
def demo1_curves():
    print_header("演示 1：延迟–速率曲线（vLLM 论文图 12a/12d 的模拟版）——每格：平均每 token 延迟（ms）")
    for kind, rates in (("sharegpt", (0.5, 1, 1.5, 2, 2.5, 3)), ("alpaca", (10, 20, 40, 60, 80, 100))):
        rows = []
        for p in POLICIES:
            cells = [NAMES[p]]
            for rate in rates:
                r = simulate(make_workload(1500, rate, kind), STEP, p)
                cells.append(f"{r['norm_latency'] * 1e3:.0f}")
            rows.append(cells)
        print(f"\n负载：{kind}（1500 个请求）")
        print_table(["管理方式"] + [f"{r:g} req/s" for r in rates], rows)
    print_observations([
        "每条曲线都是先平、后突然往上翘：请求速率超过系统的处理能力以后，队列无限变长（论文图 12 的形状）。",
        "翘起来的位置就是这种管理方式的容量：ShareGPT 上 Max / Pow2 在 1 req/s 已经过载，Oracle 约 1 req/s，分页到 2 req/s 才开始变慢。",
        "Alpaca 的序列短，大家都能跑到更高的速率，但顺序一样：Max ≪ Pow2 < Oracle < 分页。",
    ])


# -----------------------------------------------------------------------------
def demo2_capacity():
    print_header("演示 2：\"平均每 token 延迟 ≤ 50 ms\"时能承受的最高请求速率（二分查找，每次 1500 个请求）")
    rows = []
    for kind in ("sharegpt", "alpaca"):
        rates = {p: max_rate(kind, STEP, p, 0.05, n=1500, iters=12) for p in POLICIES}
        rows.append([kind] + [f"{rates[p]:.2f}" for p in POLICIES]
                    + [f"{rates['paged'] / rates['oracle']:.1f}×", f"{rates['paged'] / max(rates['max'], 1e-9):.1f}×"])
    print_table(["负载", "Max", "Pow2", "Oracle", "vLLM", "vLLM / Oracle", "vLLM / Max"], rows)
    print_observations([
        "分页能承受的速率约是 Oracle 的两倍、Max 的三倍以上（Alpaca 上是 Max 的 23 倍：短请求按 2048 预留浪费最大）。"
        "论文在 ShareGPT 上报告 vLLM 是 Orca (Oracle) 的 1.7–2.7 倍、Orca (Max) 的 2.7–8 倍（三种模型大小）——方向和量级一致。",
        "注意 Oracle 是\"事先知道每个请求生成多长\"的理想情况，现实中做不到；分页不需要知道输出长度就超过了它。",
        "Pow2 在 ShareGPT 上和 Max 差不多：伙伴分配器向上取整到 2 的幂，ShareGPT 的长请求很多都要 2048 个槽位的块，"
        "而 12,596 个槽位里只切得出 6 个这么大的块。",
    ])


# -----------------------------------------------------------------------------
def demo3_preemption():
    print_header("演示 3：分页的代价——过载时的抢占与重算（ShareGPT，1500 个请求）")
    rows = []
    for rate in (1.5, 2.0, 2.5, 3.0, 4.0):
        reqs = make_workload(1500, rate, "sharegpt")
        r = simulate(reqs, STEP, "paged")
        prefill_tokens = sum(q.prompt for q in reqs)
        rows.append([f"{rate:g}", f"{r['norm_latency'] * 1e3:.0f}", r["preemptions"],
                     sum(1 for q in reqs if q.preempted > 0), f"{r['recompute_tokens']:,}",
                     f"{r['recompute_tokens'] / prefill_tokens:.0%}"])
    print_table(["速率 req/s", "每 token 延迟 ms", "抢占次数", "被抢占过的请求", "重算的 token", "相对于所有输入 token"], rows)
    print_observations([
        "分页不预留，所以接纳时只看\"现在放不放得下输入\"；之后序列变长、显存不够，就要抢占最晚到的请求（先到先服务）。",
        "负载刚过容量时抢占很少；严重过载时，被抢占的请求要把\"输入 + 已生成的\"全部重算一遍，重算量可以超过原本的输入量——"
        "这是用计算换显存。vLLM 实际上还留了一点水位（watermark）、限制每步接纳的数量，减少这种来回抢占。",
        "第 6 讲的\"不预留就死锁\"在这里不会发生：显存不够时总能抢占别人，把显存腾出来。",
    ])


# -----------------------------------------------------------------------------
def demo4_swap_vs_recompute():
    print_header("演示 4：被抢占的请求怎么恢复——换出再换回（swap） vs 丢掉重算（recompute）（vLLM 论文图 19）")
    print("OPT-13B；换出 / 换回按 PCIe 有效带宽 25 GB/s、每次拷贝固定开销 10 µs（假设值）；每层的 K、V 各自成块，"
          "一个序列要拷贝 层数 × 2 × 块数 次。重算 = 把这些 token 当作新的提示做一次 prefill。")
    rows = []
    for n_tok in (256, 1024):
        for B in (1, 4, 16, 64, 256):
            t_sw = 2 * swap_time(OPT_13B, n_tok, B)
            t_rc = recompute_time(STEP, n_tok)
            rows.append([n_tok, B, fmt_time(t_sw), fmt_time(t_rc), "重算" if t_rc < t_sw else "换出"])
    print_table(["序列长度", "块大小", "换出 + 换回", "重算", "更快的"], rows)
    print_observations([
        "换出的时间强烈依赖块大小：块越小，拷贝次数越多，每次的固定开销占主导（块大小 1 时要拷几万次）。",
        "重算的时间和块大小无关（它根本不读 KV）。块小时重算快，块大时换出快——和论文图 19a 的形状一致。论文实测重算最多比换出"
        "慢约 20%，在我们的假设下块很大时换出快约一倍（PCIe 带宽、prefill 效率都是假设值）。论文的端到端实验里，"
        "块大小 16–64 时两种方式表现相当。",
        "另一个角度：重算花的是 GPU 的算力，换出花的是 PCIe 带宽和 CPU 内存。prefill 很贵（长提示、大模型）时换出更划算——"
        "这正是 Mooncake 把 KV 放到 CPU 内存和 SSD 上的出发点（05 脚本）。",
    ])


if __name__ == "__main__":
    demo1_curves()
    demo2_capacity()
    demo3_preemption()
    demo4_swap_vs_recompute()