"""
=============================================================================
CS349E · 第 15 讲 · 04_kv_during_tool_calls.py —— 智能体在等工具时，它的 KV cache 怎么办
=============================================================================

【这个脚本要回答的问题】
    今天的推理引擎把每次工具调用当成"请求结束"：KV cache 丢掉，工具回来后作为一个新请求，把整个上下文重新 prefill 一遍。
    InferCept（Abhyankar et al., ICML 2024）把三种做法的代价写成了"浪费了多少显存 × 时间"：
      演示 1：几类工具（InferCept 测的六种增强 + 写内核的智能体常用的工具）下，保留 / 丢掉重算 / 分块重算 / 同步换出的浪费。
      演示 2：保留和丢掉重算的盈亏平衡点 T*：上下文越长、同时在跑的其他请求越多，T* 越大。
      演示 3：工具要等多久事先不知道——"先保留 T*，没回来就丢"（TTL）和经典的租雪板问题：最坏也只多花一倍。
      演示 4：换到 CPU 要多久：PCIe 的账、按层流水的换入换出（InferCept 的 swap pipelining）。

【关于数字】Llama-3.1-8B 在 H100 上；T_fwd 用第 5 讲 StepModel 的 roofline 模型（不是实测）；PCIe 5.0 x16 取 63 GB/s 的 80%；
    InferCept 表 1 的工具耗时和上下文长度是论文的测量；写内核的智能体的工具耗时是假设值（agent_lab.TOOL_MIX）。
【运行】python3 04_kv_during_tool_calls.py      （不到 1 秒）
【依赖】numpy；本目录的 agent_lab；../lecture01 的 console。
"""

import math
import os
import sys

import numpy as np

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, STEP, TOOL_MIX, break_even, sample_tool, t_fwd, waste_chunked,  # noqa: E402
                       waste_discard, waste_preserve, waste_swap)
from console import print_header, print_observations, print_table  # noqa: E402

C_OTHER = 100_000        # 同时在跑的其他请求一共有多少 token 的 KV（假设：约 25 个 4k 上下文的请求）

# InferCept 表 1：(名字, 平均拦截时间 秒, 平均上下文长度)
INFERCEPT = [("算术（计算器）", 9e-5, 1422), ("问答（查维基百科）", 0.69, 1846), ("虚拟环境（ALFWorld）", 0.09, 2185),
             ("聊天（等人回复，估计）", 28.6, 753), ("生成图片（含等人，部分估计）", 20.03, 1247), ("文字转语音（含等人，部分估计）", 17.24, 1251)]


def demo1_waste():
    print_header(f"演示 1：每次拦截浪费多少显存·时间（GB·s）——其他请求共 {C_OTHER:,} 个 token 的 KV")
    rows = []
    cases = [(n, t, c, "InferCept 表 1") for n, t, c in INFERCEPT] + \
            [(f"写内核：{m[0]}", m[2], 12000, "假设") for m in TOOL_MIX]
    for name, t_int, C, src in cases:
        w = {"保留": waste_preserve(C, t_int), "丢掉重算": waste_discard(C, C_OTHER),
             "分块重算": waste_chunked(C, C_OTHER, max(1, math.ceil(C / 1024))), "同步换出换入": waste_swap(C, C + C_OTHER)}
        best = min(w, key=w.get)
        rows.append([name, f"{t_int:g} s", f"{C:,}"] + [f"{v:.3g}" for v in w.values()] + [best, src])
    print_table(["工具", "平均耗时", "上下文", "保留", "丢掉重算", "分块重算", "同步换出换入", "最省", "出处"], rows)
    print(f"公式（InferCept 式 1–4，M = 每 token 的 KV {KV_BYTES / 1024:.0f} KiB）：保留 = T_INT·C·M；丢掉重算 = T_fwd(C)·(C + C_other)·M；"
          f"分块重算 = T_fwd(C)·C·M/2 + n·T_fwd(C/n)·C_other·M；同步换出换入 = 2·T_swap(C)·(C + C_other)·M，T_swap = C·M / {PCIE_BW / 1e9:.0f} GB/s")
    print_observations([
        "毫秒级的工具（计算器、虚拟环境）：保留的浪费几乎为 0，丢掉重算反而要让所有正在跑的请求一起等一次 prefill——今天的引擎默认就是这样做的。",
        "几十秒的工具（等人回复、编译 + 测试、profiler）：保留要白占几十秒的显存，丢掉重算或换出更省。"
        "写内核的智能体的工具调用大多落在\"秒到几十秒\"这个中间地带，没有一种做法总是最好。",
        f"按 PCIe 5.0（{PCIE_BW / 1e9:.0f} GB/s）算，连**同步**换出换入在大多数行里都是最省的：8B 模型每 token 只有 128 KiB，传一遍比重算快得多。"
        "InferCept 在 A100 上的实验里，朴素的同步 Swap 浪费了 26% 的 GPU 资源、超过 25% 的时间在等传输——带宽、模型大小、"
        "同时换的请求数一变，结论就变（05 脚本扫了带宽）。",
        "分块重算把\"一次重算整个上下文\"拆进好几个迭代，利用 decode 用不满的算力：上下文长时更省（12k：4.49 vs 5.4），"
        "上下文很短时每块的固定开销让它略差。",
    ])


def demo2_break_even():
    print_header("演示 2：保留 vs 丢掉重算的盈亏平衡点 T* = T_fwd(C)·(C + C_other) / C")
    Cs = (1000, 4000, 16000, 64000)
    others = (0, 50_000, 200_000)
    rows = [[f"{C:,}", f"{t_fwd(C) * 1e3:.0f} ms"] + [f"{break_even(C, o):.2f} s" for o in others] for C in Cs]
    print_table(["上下文 C", "T_fwd(C)"] + [f"C_other = {o:,}" for o in others], rows)
    print_observations([
        "没有其他请求时，T* 就是重算一遍的时间（毫秒到几秒）：工具比这还短就保留，否则丢掉。",
        "其他请求越多，丢掉重算的代价越大（重算拖慢了整批），T* 越长、越倾向于保留；"
        "可是其他请求越多，显存也越紧张——保留的代价里还应该加上\"别人因为显存不够而排队\"的时间，"
        "InferCept 的公式没有算这一项。Continuum（2025）正是把\"重算 / 重新加载的代价\"和\"被挤出显存后的排队延迟\"放在一起定 TTL。"
        "05 脚本会看到不算这一项的后果。",
    ])


def demo3_ski_rental():
    print_header("演示 3：不知道工具要等多久——TTL 与租雪板问题")
    C = 12000
    D = waste_discard(C, C_OTHER)                       # "买"的代价：丢掉重算
    T = break_even(C, C_OTHER)
    rng = np.random.default_rng(0)
    durs = np.array([sample_tool(rng)[1] for _ in range(20000)])
    pres = durs * C * KV_BYTES / 1e9
    disc = np.full_like(durs, D)
    ttl = np.where(durs <= T, pres, T * C * KV_BYTES / 1e9 + D)
    opt = np.minimum(pres, disc)
    grid = np.linspace(0.05, 3 * T, 300)
    costs = [np.where(durs <= g, pres, g * C * KV_BYTES / 1e9 + D).mean() for g in grid]
    g_best = float(grid[int(np.argmin(costs))])
    ttl_b = np.where(durs <= g_best, pres, g_best * C * KV_BYTES / 1e9 + D)
    rows = [[name, f"{v.mean():.3g}", f"{v.mean() / opt.mean():.2f}", f"{np.max(v / opt):.2f}"]
            for name, v in (("总是保留", pres), ("总是丢掉重算", disc), (f"TTL = T*（{T:.1f} s，不看分布）", ttl),
                            (f"TTL = {g_best:.1f} s（按这个分布挑的最优值）", ttl_b), ("先知（事先知道耗时）", opt))]
    print(f"上下文 C = {C:,}，其他请求 {C_OTHER:,} 个 token；工具耗时按 agent_lab.TOOL_MIX 抽 20,000 次（平均 {durs.mean():.1f} s）")
    print_table(["策略", "平均浪费（GB·s）", "相对先知", "单次最坏 / 先知"], rows)
    worst = 2.0
    print(f"对任意耗时 t：TTL = T* 的浪费 / 先知的浪费 ≤ {worst:.0f}（t ≤ T* 时两者相同；t > T* 时 TTL 付了 T*·C·M + D = 2D，先知付 D）")
    print_observations([
        "\"保留还是丢掉\"就是经典的**租雪板问题**：每天租（保留，按时间付费）还是一次买下（丢掉重算，一次性付 D），可你不知道要滑几天。"
        "确定性的最优策略是\"租到累计的租金等于买价，再买\"——也就是 TTL = T*。它的保证是**最坏情况**：任何耗时下都不超过先知的 2 倍"
        "（随机化的 TTL 可以把这个比值降到 e/(e − 1) ≈ 1.58）。",
        f"平均意义上它不一定最好：这个分布里 {np.mean(durs > T):.0%} 的调用都超过 T*（编译、测试、profiler），TTL 在它们身上付了两倍，平均 {ttl.mean() / opt.mean():.2f}× 先知，"
        f"反而不如\"总是丢掉\"（{disc.mean() / opt.mean():.2f}×）；可\"总是丢掉\"的最坏情况是先知的 {np.max(disc / opt):.0f} 倍（对 0.3 s 的读文件也去重算），"
        f"\"总是保留\"平均 {pres.mean() / opt.mean():.2f}×。",
        f"**知道分布**就能做得更好：这个分布有两个峰（0.3 s 的读文件、几秒到几十秒的编译和测试），TTL 设成 {g_best:.1f} s 时，快的保留、慢的马上丢，"
        f"平均 {ttl_b.mean() / opt.mean():.2f}× 先知。这就是为什么 InferCept 按工具的类型离线测耗时、Continuum 预测工具调用的耗时；"
        "没有预测时，InferCept 用\"已经等了多久\"（T̂_INT = 现在 − 调用开始的时刻）当估计，论文报告达到了先知 93% 的性能。",
    ])


def demo4_swap():
    print_header("演示 4：换到 CPU 要多久")
    rows = []
    for C in (2000, 12000, 64000):
        byts = C * KV_BYTES
        rows.append([f"{C:,}", f"{byts / 1e9:.2f} GB", f"{byts / PCIE_BW * 1e3:.0f} ms", f"{t_fwd(C) * 1e3:.0f} ms"])
    print_table(["上下文", "KV 大小", f"单向传输（{PCIE_BW / 1e9:.0f} GB/s）", "对比：重算一遍（T_fwd）"], rows)
    it = STEP.decode_time(64, 64 * 4000)
    hidden = it * PCIE_BW / KV_BYTES
    print(f"一次 decode 迭代（批 64、每个上下文 4k）约 {it * 1e3:.1f} ms；这段时间里 PCIe 能搬 {hidden:,.0f} 个 token 的 KV——"
          f"InferCept 的 \"swap budget\"：每次迭代能\"免费\"（藏在计算后面）换入换出的 token 数。")
    print_observations([
        f"对 Llama-3.1-8B（每 token 128 KiB），换出 12k 个 token 只要几十毫秒，比重算还快：只要 PCIe 不拥挤、CPU 内存够大，"
        "换出通常是最划算的。代价在于：同步地换（换的时候整个引擎停下来等），或者很多请求同时换、挤在一条 PCIe 上。",
        "InferCept 的做法：按层流水（换第 i + 2 层的同时传第 i + 1 层、第 i 层的显存已经给别人用）；把一个序列的换入换出拆到好几个迭代里，"
        "每个迭代只换\"藏得住\"的那么多（swap budget）；把预算优先给\"浪费最大\"的请求，其余的按式 (5) 在保留和分块重算之间选。",
        "这和第 7 讲的 Mooncake 是同一个方向：把 KV 当成可以在 HBM、DRAM、SSD 之间搬动的数据，按\"搬的代价 vs 重算的代价\"决定放在哪里。",
    ])


if __name__ == "__main__":
    demo1_waste()
    demo2_break_even()
    demo3_ski_rental()
    demo4_swap()