"""
=============================================================================
CS349E · 第 16 讲共享模块：planner.py —— 把整门课串起来：一个 token 的全栈账，和一个端到端的部署规划器
=============================================================================

【这个模块要解决什么问题】
    前 15 讲各自回答了一层的问题：roofline（第 1 讲）、Transformer 的账本（第 2 讲）、精度（第 3 讲）、切分（第 4 讲）、
    解耦（第 5 讲）、批处理与投机解码（第 6 讲）、KV 管理（第 7 讲）、硬件（第 8、9 讲）、内核（第 10、11 讲）、编译器（第 12、13 讲）。
    这一讲把它们拼起来，回答两个"系统级"的问题：

    A. 一个 token 的全栈账：Llama-3.1-70B 在 8 块 H100 上 decode 一步，时间花在哪里（读权重、读 KV、计算、张量并行的通信、
       内核启动）？从"最朴素的部署"出发，逐个加上各讲的技术，每一步省了多少——"优化阶梯"。
    B. 部署规划：给定模型、负载（提示 / 输出长度）、SLO（TTFT、TPOT 的 P90）、目标流量，在
       {GPU 型号} × {BF16 / FP8} × {张量并行度} × {共置（分块 prefill）/ 解耦} 里找每百万输出 token 最便宜的方案。
       先用解析模型筛选（快），再用第 5 讲的离散事件模拟器检验（慢、但考虑了排队和干扰）。

【关于数字】时间都是 roofline 模型加效率折扣和固定开销（第 5 讲 StepModel 的假设：60% 的峰值算力、80% 的峰值带宽、
    每步 1 ms），**不是实测**；GPU 的租用价格是假设值，只用来比较方案。解析模型里的排队近似（M/D/1 + 尾部系数）和余量都是我选的，
    03 脚本用模拟器检验它们。
【运行】python3 planner.py   # 自检（约 1 秒）
【依赖】numpy；../lecture01（hw_specs、llm_math）、../lecture05（serving_sim）、../lecture06（decode_lab 的投机解码公式）、
    ../lecture12（graph_lab：数一层 decode 有几个内核）、../prereq5（comm_sim：all-reduce 的 α–β 公式）。都是只读导入。
"""

import math
import os
import sys
from dataclasses import dataclass, field

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
for sub in ("lecture01", "lecture05", "lecture06", "lecture12", "prereq5"):
    sys.path.insert(0, os.path.join(HERE, "..", sub))

from hw_specs import B200, H100, H200  # noqa: E402
from llm_math import LLAMA31_8B, LLAMA31_70B, LLAMA31_405B  # noqa: E402
from serving_sim import (StepModel, attainment, lognormal_len, make_workload, percentile, run_colocated,  # noqa: E402
                         run_disaggregated)
from decode_lab import expected_tokens  # noqa: E402
from comm_sim import NVLINK4, t_formula  # noqa: E402

PRICES = {H100.name: 2.0, H200.name: 2.6, B200.name: 4.0}     # 每 GPU·小时的租用价格（美元，假设值，只用来比较）
CHIPS = (H100, H200, B200)


# =============================================================================
# A. 一个 token 的全栈账
# =============================================================================
_KERNELS = {}


def decode_kernels_per_layer():
    """
    用第 12 讲的 graph_lab 数一层 Llama decode 有几个内核：eager（复合运算各自一个内核）和"分解 + 融合 + FlashAttention"。

    【返回】{"eager": 个数, "fused": 个数}（结构和模型大小无关，用 8B 的一层来数）
    【系统联系】第 12 讲 03 脚本演示 3：融合在 decode 时省的主要是内核个数（固定开销），而不是字节。
    """
    if not _KERNELS:
        from graph_lab import (LLAMA8B, constant_fold, cse, decode_layer, decompose, fold_norm_weight, fuse,
                               graph_cost, match_attention, simplify, trace)
        cfg = LLAMA8B
        shapes = {"attn_norm": (cfg.d,), "wq": (cfg.d, cfg.H * cfg.dh), "wk": (cfg.d, cfg.G * cfg.dh),
                  "wv": (cfg.d, cfg.G * cfg.dh), "wo": (cfg.H * cfg.dh, cfg.d), "mlp_norm": (cfg.d,),
                  "w_gate": (cfg.d, cfg.f), "w_up": (cfg.d, cfg.f), "w_down": (cfg.f, cfg.d)}
        g = trace(decode_layer(cfg, "P"), {"x": (1, cfg.d), "k_cache": (cfg.G, "S", cfg.dh),
                                           "v_cache": (cfg.G, "S", cfg.dh), "cos": (1, 1, cfg.dh), "sin": (1, 1, cfg.dh)},
                  shapes)
        _KERNELS["eager"] = graph_cost(fuse(g, enable=False, bind={"S": 4096}))["kernels"]
        g2, _ = match_attention(g)
        g2 = decompose(g2)
        g2 = cse(g2)
        g2, _ = constant_fold(g2)
        g2, _ = simplify(g2)
        g2, _ = fold_norm_weight(g2)
        _KERNELS["fused"] = graph_cost(fuse(g2, bind={"S": 4096}))["kernels"]
    return dict(_KERNELS)


@dataclass
class Deploy:
    """
    一种 decode 的部署方式（优化阶梯上的一级）。

    【字段】batch：批大小；w_bytes / kv_bytes：权重 / KV 每个数几字节；kernels：每层几个内核；launch：每个内核的固定开销（秒，
        eager 时 CPU 逐个发射约几微秒；CUDA Graphs 把整步录下来一次发射，每个内核只剩 GPU 端的一点开销——都是假设值）；
        spec：投机解码 (草稿模型, 接受率 α, 草稿长度 γ) 或 None
    """
    name: str
    batch: int = 1
    w_bytes: float = 2.0
    kv_bytes: float = 2.0
    kernels: int = 30
    launch: float = 4e-6
    spec: object = None


def tp_step(model, chip, tp, batch, ctx, n_tokens=1, w_bytes=2.0, kv_bytes=2.0, kernels=30, launch=4e-6,
            mfu=0.6, mbu=0.8, link=NVLINK4):
    """
    张量并行 tp 块 GPU 上一次前向（batch 个序列各送进 n_tokens 个 token，已有 ctx 个 token 的上下文）的时间分解。

    【公式】（第 1、2、4、6、12 讲的账本拼在一起）
        读权重 = 矩阵乘法参数 × w_bytes / (tp × 带宽 × mbu)
        读 KV  = batch × (ctx + n) × 每 token KV / (tp × 带宽 × mbu)
        计算   = FLOPs / (tp × 峰值 × mfu)，FLOPs = batch·n·(线性层 + 注意力)
        通信   = 每层 2 次 all-reduce（消息 = batch·n·d·2 字节），ring 和递归倍增里取快的（P5 的 α–β 公式）
        启动   = 层数 × 每层内核数 × 每个内核的固定开销
        一步   = max(计算, 读权重 + 读 KV) + 通信 + 启动
    【边界】tp = 1 时没有通信；不考虑通信和计算的重叠（偏保守）。
    """
    prec = "bf16" if w_bytes >= 2 else "fp8"
    bw = chip.mem_bw * mbu * tp
    t_w = model.matmul_params() * w_bytes / bw
    t_kv = batch * (ctx + n_tokens) * model.kv_bytes_per_token(kv_bytes) / bw
    flops = batch * n_tokens * model.linear_flops_per_token() + \
        batch * model.attn_flops_per_token(1) * (n_tokens * ctx + n_tokens * (n_tokens + 1) / 2)
    t_c = flops / (chip.peak_flops(prec) * mfu * tp)
    if tp > 1:
        msg = batch * n_tokens * model.d_model * 2.0
        one = min(t_formula("ring", tp, msg, link.alpha, link.bw), t_formula("rd", tp, msg, link.alpha, link.bw))
        t_comm = 2 * model.n_layers * one
    else:
        t_comm = 0.0
    t_launch = model.n_layers * kernels * launch
    t = max(t_c, t_w + t_kv) + t_comm + t_launch
    return {"weights": t_w, "kv": t_kv, "compute": t_c, "comm": t_comm, "launch": t_launch, "t": t,
            "bound": "计算" if t_c >= t_w + t_kv else "带宽"}


def ladder_point(model, chip, tp, ctx, d, draft=LLAMA31_8B):
    """
    优化阶梯上的一级：每个用户每秒几个 token、每块 GPU 每秒几个 token、每百万输出 token 的成本。

    【投机解码】d.spec = (α, γ)：一轮 = γ 次草稿 decode（草稿也在同样的 tp 块 GPU 上）+ 1 次目标模型验证 γ+1 个 token，
        平均产出 expected_tokens(α, γ) 个 token（第 6 讲的式 1）。
    【返回】dict：t_token（每个用户的平均 token 间隔）、tok_s_user、tok_s_gpu、usd_per_mtok、step（tp_step 的分解）
    """
    kw = dict(w_bytes=d.w_bytes, kv_bytes=d.kv_bytes, kernels=d.kernels, launch=d.launch)
    base = tp_step(model, chip, tp, d.batch, ctx, 1, **kw)
    if d.spec is None:
        t_tok = base["t"]
    else:
        alpha, gamma = d.spec
        t_draft = tp_step(draft, chip, tp, d.batch, ctx, 1, **kw)["t"]
        t_ver = tp_step(model, chip, tp, d.batch, ctx, gamma + 1, **kw)["t"]
        t_tok = (gamma * t_draft + t_ver) / expected_tokens(alpha, gamma)
    tok_gpu = d.batch / t_tok / tp
    return {"t_token": t_tok, "tok_s_user": 1 / t_tok, "tok_s_gpu": tok_gpu,
            "usd_per_mtok": PRICES[chip.name] / 3600 / tok_gpu * 1e6, "step": base}


# =============================================================================
# B. 部署规划：解析模型筛选
# =============================================================================
@dataclass
class Scenario:
    """
    一个部署需求。

    【字段】model：模型；prompt / output：提示 / 输出长度的平均值（对数正态，σ = sigma）；ttft / tpot：P90 的 SLO（秒）；
        rate：目标流量（请求 / 秒）；precisions：允许的精度（质量要求决定能不能用 FP8，第 3 讲）
    """
    name: str
    model: object
    prompt: int
    output: int
    ttft: float
    tpot: float
    rate: float
    sigma: float = 0.8
    precisions: tuple = ("bf16", "fp8")
    note: str = ""

    def prompt_quantile(self, q=0.9):
        """
        提示长度的 q 分位数（对数正态：均值 = prompt，σ = sigma，所以 μ = ln(prompt) − σ²/2）。

        【为什么需要】SLO 是按 P90 定的：TTFT 的 P90 主要由**长**提示决定，用平均长度检查会漏掉长尾（03 脚本）。
        【Python 小知识】正态分布的 q 分位数：q = 0.9 时 z ≈ 1.2816（用误差函数的反函数太麻烦，这里只支持几个常用的 q）。
        """
        z = {0.5: 0.0, 0.9: 1.2816, 0.95: 1.6449, 0.99: 2.3263}[q]
        mu = math.log(self.prompt) - self.sigma ** 2 / 2
        return int(round(math.exp(mu + z * self.sigma)))


def make_step(model, chip, tp, prec):
    """一个实例（tp 块 GPU）的每步时间模型；权重放不下时返回 None。"""
    wb = 2.0 if prec == "bf16" else 1.0
    st = StepModel(model, chip, tp=tp, w_bytes=wb, kv_bytes=wb)
    if st.kv_capacity_tokens() <= 0:
        return None
    return st


def decode_capacity(st, sc, margin=0.85):
    """
    解耦时一个 decode 实例能承受多少请求 / 秒。

    【做法】平均上下文 ≈ 提示 + 输出 / 2；KV 按"提示 + 全部输出"预留（第 5 讲 DecodeInstance 的做法），所以批大小 ≤ 容量 / (P + O)；
        在此之下找最大的批 b，使一步的时间 ≤ margin × TPOT（留 15% 给长短不一、P90 的尾巴）。
        吞吐 = b / (一步的时间 × 平均输出长度)。
    【返回】dict(batch, rps, t_step) 或 None（批 1 也满足不了 TPOT，或者放不下一个请求）
    """
    P, O = sc.prompt, sc.output
    b_mem = int(st.kv_capacity_tokens() // (P + O))
    if b_mem < 1:
        return None
    t = lambda b: st.decode_time(b, b * (P + O / 2))       # noqa: E731
    if t(1) > margin * sc.tpot:
        return None
    lo, hi = 1, b_mem
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if t(mid) <= margin * sc.tpot:
            lo = mid
        else:
            hi = mid - 1
    return {"batch": lo, "rps": lo / (t(lo) * O), "t_step": t(lo)}


def prefill_capacity(st, sc, tail=2.3, max_tokens=4096, p_ttft=None):
    """
    解耦时一个 prefill 实例能承受多少请求 / 秒（排队的近似）。

    【做法】单个请求的 prefill 时间 S；短提示可以凑成一批（最多 max_tokens 个 token），平均每个请求占 S_b；
        M/D/1 的平均等待 W = ρ·S_b / (2(1 − ρ))；用"P90 ≈ S + tail × W"近似尾部（tail = 2.3 ≈ ln 10，指数尾巴的 90% 分位）。
        解出满足 P90 ≤ TTFT 的最大利用率 ρ（再封顶 0.9），吞吐 = ρ / S_b。
    【参数】p_ttft：检查 TTFT 时用的提示长度（默认用平均长度；plan(tail_aware=True) 传 P90 的长度）
    【返回】dict(rps, rho, S) 或 None（单个请求的 prefill 就超过 TTFT）
    【边界】这是一个粗略的近似：提示长度不一、批的组合、P90 的形状都没精确建模——03 脚本用模拟器检验。
    """
    P = sc.prompt
    S = st.prefill_time([p_ttft or P])
    if S >= sc.ttft:
        return None
    k = max(1, max_tokens // P)
    Sb = st.prefill_time([P] * k) / k
    a = sc.ttft - S
    rho = min(0.9, a / (a + tail * Sb / 2))
    return {"rps": rho / Sb, "rho": rho, "S": S}


def colocated_capacity(st, sc, margin=0.85, chunks=(256, 512, 1024, 2048), p_ttft=None):
    """
    共置（分块 prefill）时一个实例能承受多少请求 / 秒。

    【做法】稳态下每一步 = b 个 decode + 一块 c 个 prefill token；一步的时间 T = StepModel.mixed_time(...) ≤ margin × TPOT；
        decode 跟得上要求 R·O ≤ b / T，prefill 跟得上要求 R·P ≤ c / T，所以 R = min(b/O, c/P) / T；
        TTFT 粗估为 ⌈P / c⌉ 步，要求 ≤ TTFT 的一半（另一半留给排队）；显存：b × (P + O) ≤ 容量。
        在 c 的几个取值和 b 上搜索 R 最大的组合。
    【参数】p_ttft：检查 TTFT 时用的提示长度（默认平均长度）
    【返回】dict(batch, chunk, rps, t_step) 或 None
    """
    P, O = sc.prompt, sc.output
    cap = st.kv_capacity_tokens()
    best = None
    for c in chunks:
        b_mem = int(cap // (P + O))
        for b in sorted({int(x) for x in np.unique(np.geomspace(1, max(1, b_mem), 40).round())}):
            ctx = b * (P + O / 2)
            attn = st.model.attn_flops_per_token(1) * c * (P / 2)
            T = st.mixed_time(b, ctx, c, attn)
            if T > margin * sc.tpot:
                break
            if math.ceil((p_ttft or P) / c) * T > sc.ttft / 2:
                continue
            R = min(b / O, c / P) / T
            if best is None or R > best["rps"]:
                best = {"batch": b, "chunk": c, "rps": R, "t_step": T}
    return best


@dataclass
class Plan:
    """一个部署方案。mode："colo"（共置）或 "disagg"（解耦）；instances：[(tp, 个数, 角色)]。"""
    chip: object
    prec: str
    mode: str
    instances: list
    gpus: int
    usd_per_mtok: float
    detail: dict = field(default_factory=dict)
    price: float = 0.0          # 每 GPU·小时的价格（美元）

    def capacity(self):
        """方案按解析模型能承受的流量（请求 / 秒）：共置 = 实例数 × 每实例吞吐；解耦 = prefill、decode 两边较小的那个。"""
        if self.mode == "colo":
            return self.instances[0][1] * self.detail["colo"]["rps"]
        return min(self.instances[0][1] * self.detail["prefill"]["rps"], self.instances[1][1] * self.detail["decode"]["rps"])

    def label(self):
        inst = " + ".join(f"{n}×TP{tp}{'' if role == 'colo' else ('P' if role == 'prefill' else 'D')}"
                          for tp, n, role in self.instances)
        return f"{self.chip.name} · {self.prec.upper()} · {'共置' if self.mode == 'colo' else '解耦'} · {inst}"


def plan(sc, chips=CHIPS, tps=(1, 2, 4, 8), modes=("colo", "disagg"), tail_aware=False, prices=None):
    """
    枚举所有方案，按每百万输出 token 的成本排序。

    【解耦】和 DistServe 的算法 1 一样：prefill、decode 各自挑"每块 GPU 吞吐最高"的张量并行度，再各自复制到满足流量
        （同一种 GPU、同一种精度——KV 的格式要一致）。
    【成本】GPU 数 × 每 GPU·小时价格 / (流量 × 平均输出长度 × 3600) × 10^6。
    【tail_aware】True 时用 P90 的提示长度检查 TTFT（SLO 是 P90）；False 时用平均长度（第一版的做法，03 脚本对比两者）。
    【prices】{芯片名: 每 GPU·小时的价格}；默认用 PRICES（04 脚本用它做价格的假设分析）。
    【返回】Plan 列表（便宜的在前；成本一样时，解析模型给的余量大的在前）；某种组合满足不了 SLO 时跳过
    """
    out = []
    tok_per_h = sc.rate * sc.output * 3600
    p_ttft = sc.prompt_quantile(0.9) if tail_aware else None
    prices = prices or PRICES
    for chip in chips:
        price = prices[chip.name]
        for prec in sc.precisions:
            steps = {tp: make_step(sc.model, chip, tp, prec) for tp in tps}
            steps = {tp: s for tp, s in steps.items() if s is not None}
            if "colo" in modes:
                for tp, st in steps.items():
                    c = colocated_capacity(st, sc, p_ttft=p_ttft)
                    if c:
                        n = math.ceil(sc.rate / c["rps"])
                        g = n * tp
                        out.append(Plan(chip, prec, "colo", [(tp, n, "colo")], g, g * price / tok_per_h * 1e6,
                                        {"colo": c}, price))
            if "disagg" in modes:
                pc = {tp: prefill_capacity(st, sc, p_ttft=p_ttft) for tp, st in steps.items()}
                dc = {tp: decode_capacity(st, sc) for tp, st in steps.items()}
                pc = {tp: v for tp, v in pc.items() if v}
                dc = {tp: v for tp, v in dc.items() if v}
                if pc and dc:
                    tpp = max(pc, key=lambda tp: pc[tp]["rps"] / tp)
                    tpd = max(dc, key=lambda tp: dc[tp]["rps"] / tp)
                    n_p = math.ceil(sc.rate / pc[tpp]["rps"])
                    n_d = math.ceil(sc.rate / dc[tpd]["rps"])
                    g = n_p * tpp + n_d * tpd
                    out.append(Plan(chip, prec, "disagg", [(tpp, n_p, "prefill"), (tpd, n_d, "decode")], g,
                                    g * price / tok_per_h * 1e6, {"prefill": pc[tpp], "decode": dc[tpd]}, price))
    return sorted(out, key=lambda p: (round(p.usd_per_mtok, 9), -p.capacity()))


# =============================================================================
# C. 用第 5 讲的模拟器检验
# =============================================================================
def workload(sc, n=1500, rate=None, seed=0):
    """按场景生成 n 个请求（泊松到达，提示 / 输出长度对数正态）。"""
    return make_workload(n, rate or sc.rate, lognormal_len(sc.prompt, sc.sigma, 16, 32 * sc.prompt),
                         lognormal_len(sc.output, sc.sigma, 2, 16 * sc.output), seed=seed)


def simulate(p, sc, rate=None, n=1500, seed=0, duration=None):
    """
    用离散事件模拟器跑一个方案：返回 dict(attain, ttft_p90, tpot_p90)。

    【参数】duration：给定时，请求数 = max(n, 流量 × duration)，保证模拟覆盖足够长的到达时间。
        过载的系统要过一段时间队列才会长到违反 SLO；请求太少（模拟太短）会把一个不稳定的系统误判成"满足 SLO"。

    【共置】run_colocated(..., policy="chunked", chunk_budget = 方案里的块大小)；【解耦】run_disaggregated（KV 走 InfiniBand）。
    """
    if duration is not None:
        n = max(n, int((rate or sc.rate) * duration))
    reqs = workload(sc, n, rate, seed)
    if p.mode == "colo":
        tp, k, _ = p.instances[0]
        st = make_step(sc.model, p.chip, tp, p.prec)
        run_colocated(reqs, st, n_instances=k, policy="chunked", chunk_budget=p.detail["colo"]["chunk"])
    else:
        (tpp, n_p, _), (tpd, n_d, _) = p.instances
        run_disaggregated(reqs, make_step(sc.model, p.chip, tpp, p.prec), make_step(sc.model, p.chip, tpd, p.prec),
                          n_prefill=n_p, n_decode=n_d)
    return {"attain": attainment(reqs, sc.ttft, sc.tpot), "ttft_p90": percentile(reqs, "ttft", 90),
            "tpot_p90": percentile(reqs, "tpot", 90)}


def with_counts(p, sc, counts):
    """换一组实例个数，重新算 GPU 数和成本（counts 和 p.instances 一一对应）。"""
    inst = [(tp, n, role) for (tp, _, role), n in zip(p.instances, counts)]
    g = sum(tp * n for tp, n, _ in inst)
    return Plan(p.chip, p.prec, p.mode, inst, g, g * p.price / (sc.rate * sc.output * 3600) * 1e6, p.detail, p.price)


def repair(p, sc, target=0.9, max_steps=30, n=1500, duration=300.0):
    """
    "检验 → 修复"：用模拟器跑方案；SLO 达成率不到 target 就加实例，直到达标。

    【加在哪里】共置：再加一个实例。解耦：TTFT 的 P90 超了就加一个 prefill 实例，否则加一个 decode 实例。
    【返回】(修复后的方案 或 None, 最后一次模拟的结果, 历史 [(GPU 数, 达成率)])；连加三次实例达成率都没涨 1 个百分点时
        判定"加机器也没用"，返回 None（例如长提示在这种配置上单独 prefill 就超过 TTFT）
    【为什么需要】解析模型的排队近似对长短不一的请求不够准（03 脚本）；DistServe 的放置算法也是"解析筛选 + 模拟验证"。
    【duration】每次模拟至少覆盖这么多秒的到达（见 simulate 的说明）。
    """
    cur, hist = p, []
    for _ in range(max_steps):
        r = simulate(cur, sc, n=n, duration=duration)
        hist.append((cur.gpus, r["attain"]))
        if r["attain"] >= target:
            return cur, r, hist
        if len(hist) >= 4 and hist[-1][1] - hist[-4][1] < 0.01:
            return None, r, hist          # 连加三次实例都没有进步：加机器解决不了（例如单个长请求本身就超时）
        counts = [k for _, k, _ in cur.instances]
        if cur.mode == "colo":
            counts[0] += 1
        elif r["ttft_p90"] > sc.ttft:
            counts[0] += 1
        else:
            counts[1] += 1
        cur = with_counts(cur, sc, counts)
    return cur, r, hist


def verified_best(sc, top=4, **kw):
    """解析模型挑前 top 个方案（按 P90 提示长度检查 TTFT），逐个模拟检验 + 修复，返回最便宜的 (方案, 模拟结果)；都不可行时返回 (None, None)。"""
    best = (None, None)
    for p in plan(sc, tail_aware=True, **kw)[:top]:
        q, r, _ = repair(p, sc)
        if q is not None and (best[0] is None or q.usd_per_mtok < best[0].usd_per_mtok):
            best = (q, r)
    return best


SCENARIOS = {
    "chat": Scenario("聊天", LLAMA31_70B, 1000, 250, 1.0, 0.05, 20.0, note="提示约 1000、输出约 250；TTFT ≤ 1 s、TPOT ≤ 50 ms（P90）"),
    "agent": Scenario("写代码的智能体", LLAMA31_70B, 8000, 400, 3.0, 0.05, 4.0,
                      note="提示约 8000（长上下文、不算前缀缓存）、输出约 400；TTFT ≤ 3 s、TPOT ≤ 50 ms"),
    "batch": Scenario("离线批处理", LLAMA31_70B, 2000, 300, 60.0, 0.5, 20.0, note="提示约 2000、输出约 300；只要求 TTFT ≤ 60 s、TPOT ≤ 0.5 s"),
    "fast": Scenario("低延迟（小模型）", LLAMA31_8B, 500, 100, 0.2, 0.015, 40.0,
                     note="提示约 500、输出约 100；TTFT ≤ 0.2 s、TPOT ≤ 15 ms"),
}


# =============================================================================
# 自检
# =============================================================================
def _selftest():
    k = decode_kernels_per_layer()
    assert k["fused"] < k["eager"], k
    d0 = tp_step(LLAMA31_70B, H100, 8, 1, 4096)
    assert d0["bound"] == "带宽" and d0["comm"] > 0
    assert tp_step(LLAMA31_70B, H100, 1, 1, 4096)["comm"] == 0.0
    base = ladder_point(LLAMA31_70B, H100, 8, 4096, Deploy("基线", kernels=k["eager"]))
    big = ladder_point(LLAMA31_70B, H100, 8, 4096, Deploy("批 64", batch=64, kernels=k["fused"], launch=1e-6))
    assert big["tok_s_gpu"] > 10 * base["tok_s_gpu"]
    assert make_step(LLAMA31_405B, H100, 1, "bf16") is None              # 405B 放不进一块 H100
    sc = SCENARIOS["chat"]
    ps = plan(sc, chips=(H100,))
    assert ps and all(p.gpus > 0 for p in ps)
    assert ps[0].usd_per_mtok <= ps[-1].usd_per_mtok
    small = Scenario("小", LLAMA31_8B, 300, 50, 0.5, 0.05, 2.0)
    p0 = plan(small, chips=(H100,))[0]
    r = simulate(p0, small, n=300)
    assert 0.0 <= r["attain"] <= 1.0
    p1, r1, _ = repair(p0, small, n=300, duration=60)
    assert r1["attain"] >= 0.9 and p1.gpus >= p0.gpus
    print("planner 自检通过。")


if __name__ == "__main__":
    print("== planner 自检 ==")
    _selftest()