"""
=============================================================================
CS349E · 第 7 讲 · 02_fragmentation.py —— KV 显存去哪了：预留、内部碎片、外部碎片
=============================================================================

【这个脚本要回答的问题】
    vLLM 论文图 2 说：以前的系统里，真正存着 token 的 KV 只占显存的 20.4%–38.2%，vLLM 是 96.3%。
    剩下的去哪了？
      演示 1：伙伴分配器（buddy allocator）的一个小例子：内部碎片（要 5 给 8）和外部碎片（空闲的加起来够，却没有一块连续的够大）。
      演示 2：OPT-13B、一块 A100 40GB、过载时，四种管理方式下显存的去向（论文图 2）：
              按最大长度预留（Orca Max）、按 2 的幂预留输出（Orca Pow2）、恰好按真实长度预留（Orca Oracle）、分页（vLLM）。
      演示 3：同样的负载下，平均同时在跑几个请求（论文图 13）。
      演示 4：块大小的取舍：块越大，最后一块浪费越多；块越小，块表越长、每次读的数据越碎。

【运行】python3 02_fragmentation.py      （约 10 秒）
【依赖】numpy；本目录的 kv_serving（它会导入 ../lecture06 的 decode_lab 和更早的模块）；../lecture01 的 console。
"""

import os
import sys

import numpy as np

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

from kv_serving import A100_40G, OPT_13B, Buddy, StepModel, make_workload, simulate  # noqa: E402
from console import fmt_bytes, print_header, print_observations, print_table  # noqa: E402

STEP = StepModel(OPT_13B, A100_40G)
LABELS = {"max": "Orca (Max)：按 2048 预留", "pow2": "Orca (Pow2)：输出按 2 的幂预留",
          "oracle": "Orca (Oracle)：恰好按真实长度", "paged": "vLLM：分页，块 16"}


# -----------------------------------------------------------------------------
def demo1_buddy():
    print_header("演示 1：伙伴分配器的两种碎片（总共 16 个槽位）")
    b = Buddy(16)
    rows = []
    a = b.alloc(5)
    rows.append(["请求 A 要 5 个", f"给了 {1 << a[1]} 个（起点 {a[0]}）", f"{(1 << a[1]) - 5} 个内部碎片", b.free_slots(), b.largest_free()])
    c = b.alloc(3)
    rows.append(["请求 B 要 3 个", f"给了 {1 << c[1]} 个（起点 {c[0]}）", f"{(1 << c[1]) - 3} 个内部碎片", b.free_slots(), b.largest_free()])
    d = b.alloc(2)
    rows.append(["请求 C 要 2 个", f"给了 {1 << d[1]} 个（起点 {d[0]}）", "0", b.free_slots(), b.largest_free()])
    b.release(*c)
    rows.append(["B 结束，释放", "-", "-", b.free_slots(), b.largest_free()])
    e = b.alloc(6)
    rows.append(["请求 D 要 6 个", "失败" if e is None else f"给了 {1 << e[1]} 个", "-", b.free_slots(), b.largest_free()])
    print_table(["事件", "结果", "内部碎片", "空闲槽位合计", "最大的连续空闲块"], rows)
    print_observations([
        "内部碎片：伙伴分配器只给 2 的幂，要 5 给 8、要 3 给 4。",
        "外部碎片：B 释放后一共有 6 个空闲槽位，正好够 D，可它们是 4 + 2 两块、不连续，D 进不来。",
        "以前的系统要求每个请求的 KV 在显存里连续，所以这两种碎片都躲不掉；分页之后，一个请求的块可以散在任何地方，"
        "外部碎片消失，内部碎片只剩最后一块里没填满的部分。",
    ])


# -----------------------------------------------------------------------------
def run_all(kind, rate, n=2000):
    return {p: simulate(make_workload(n, rate, kind), STEP, p) for p in ("max", "pow2", "oracle", "paged")}


def demo2_breakdown():
    print_header("演示 2：显存的去向（vLLM 论文图 2）——OPT-13B、A100 40GB、过载（请求比系统能处理的多）")
    cap = STEP.kv_capacity_tokens()
    print(f"放下 {fmt_bytes(OPT_13B.weight_bytes('fp16'))} 的权重后，KV 能用 {cap:,} 个 token 的槽位"
          f"（每个 token {fmt_bytes(OPT_13B.kv_bytes_per_token())}；论文表 1：KV 用 12 GB、15.7K 个槽位，其余约 2 GB 留给激活等；"
          "我们按 10% 即 4 GB 留）")
    for kind, rate in (("sharegpt", 3.0), ("alpaca", 60.0)):
        res = run_all(kind, rate)
        rows = []
        for p, r in res.items():
            mem = r["mem"]
            rows.append([LABELS[p], f"{mem['used']:.1%}", f"{mem['reserved']:.1%}", f"{mem['internal']:.1%}",
                         f"{mem['external']:.1%}", f"{mem['free']:.1%}"])
        print(f"\n负载：{kind}（{rate:g} req/s）")
        print_table(["管理方式", "存着 token", "预留（将来会用）", "内部碎片", "外部碎片", "其他空闲"], rows)
    print_observations([
        "ShareGPT（长、方差大）：按 2048 预留时真正存着 token 的只有约四分之一，大头是\"内部碎片\"——预留了却永远用不到；"
        "恰好按真实长度预留也只有约四成（伙伴分配器向上取整 + 还没生成到的部分被提前占着）。分页时九成以上存着 token。",
        "Alpaca（短）：按 2048 预留几乎全是浪费；Pow2 已经接近装满（空闲 + 外部碎片不到两成），Oracle 和分页还有大量空闲，"
        "瓶颈不在显存。论文在 OPT-175B + Alpaca（显存充裕、序列短）上也看到 vLLM 相对 Oracle / Pow2 的优势变小（图 12f）。",
        "\"预留（将来会用）\"不是纯浪费，但它在请求的整个生命周期里占着显存，本来可以给别的请求用。",
        "论文图 2 的数字（以前的系统 20.4%–38.2%，vLLM 96.3%）是在它自己的负载上测的；这里的负载是合成的，只比较趋势。",
    ])


# -----------------------------------------------------------------------------
def demo3_batch():
    print_header("演示 3：平均同时在跑几个请求（vLLM 论文图 13：ShareGPT 2 req/s、Alpaca 30 req/s）")
    rows = []
    for kind, rate in (("sharegpt", 2.0), ("alpaca", 30.0)):
        res = run_all(kind, rate)
        b = {p: r["batch"] for p, r in res.items()}
        rows.append([f"{kind} {rate:g} req/s"] + [f"{b[p]:.1f}" for p in ("max", "pow2", "oracle", "paged")]
                    + [f"{b['paged'] / b['oracle']:.1f}× / {b['paged'] / b['max']:.1f}×"])
    print_table(["负载", "Max", "Pow2", "Oracle", "vLLM", "vLLM 是 Oracle / Max 的"], rows)
    print_observations([
        "ShareGPT：分页能同时跑的请求数约是 Oracle 的 1.8 倍、Max 的 3 倍（论文：OPT-13B 上 2.2 倍和 4.3 倍）。"
        "decode 受带宽限制时，一步的时间几乎不随批大小变化，所以批大小就是吞吐。",
        "Alpaca：请求很短，除了 Max 以外显存都不是瓶颈，批大小由到达速率决定，几种方式差不多。",
    ])


# -----------------------------------------------------------------------------
def demo4_block_size():
    print_header("演示 4：块大小的取舍（ShareGPT 式的长度分布，每个请求的最终长度 = 输入 + 输出）")
    reqs = make_workload(5000, 1.0, "sharegpt")
    lens = np.array([r.prompt + r.output for r in reqs])
    rows = []
    for B in (1, 4, 8, 16, 32, 64, 128, 256):
        alloc = np.ceil(lens / B) * B
        waste = 1 - lens.sum() / alloc.sum()
        entries = np.ceil(lens / B).mean()
        per_block = B * OPT_13B.n_kv_heads * OPT_13B.d_head * 2
        rows.append([B, f"{waste:.2%}", f"{entries:.0f}", fmt_bytes(per_block)])
    print(f"5000 个请求，最终长度平均 {lens.mean():.0f}、中位数 {np.median(lens):.0f} 个 token")
    print_table(["块大小（token）", "最后一块的浪费", "每个请求平均几个块（块表长度）", "一个块（一层的 K 或 V）的字节"], rows)
    print_observations([
        "块大小 16 时最后一块平均只浪费约 1.5%（平均半块 = 8 个 token，而请求平均几百个 token）；块大小 256 时浪费 20% 以上。",
        "块太小也不好：块表变长；注意力内核每次只读一小块（16 个 token、一层的 K 是 160 KB），GPU 的并行度用不满；"
        "换出到 CPU 时要做很多次小拷贝（03 脚本演示 4）。vLLM 论文图 18b：块大小 16–128 在 ShareGPT 上最好，"
        "Alpaca（短序列）上 16、32 最好，默认取 16。",
    ])


if __name__ == "__main__":
    demo1_buddy()
    demo2_breakdown()
    demo3_batch()
    demo4_block_size()