"""
=============================================================================
CS349E · 第 6 讲共享模块：decode_lab.py —— 优化自回归 decode：迭代级调度、选择性批处理、投机解码
=============================================================================

【这个模块要解决什么问题】
    decode 每一步只给每个序列生成 1 个 token，却要把全部权重从显存里读一遍（第 1 讲 §4.2）——
    算术强度太低，GPU 大部分时间在等数据。提高强度有两条路：
      (1) 每一步多放几个**序列**（批处理）。难点：请求随时到、长短不一。
          → Orca（OSDI 2022）：迭代级调度（iteration-level scheduling）+ 选择性批处理（selective batching）。
      (2) 每一步给同一个序列多算几个 **token**。难点：后一个 token 依赖前一个。
          → 投机解码（Leviathan et al., ICML 2023）：小模型先猜 γ 个，大模型一次验证，用一个巧妙的
            拒绝采样保证输出分布和大模型单独生成**完全一样**。
    本模块把这两件事都写成能运行、能验证的代码：

    第 1 部分（调度模拟）：
        A100 规格、GPT-3 13B 的结构（Orca 论文表 1）、Orca 的合成负载；
        run_request_level：请求级调度（Triton + FasterTransformer 的做法：一批跑完才换下一批）；
        run_iteration_level：迭代级调度（Orca 算法 1：每次迭代重新选请求，先到先服务，KV 预留）。
        每步时间用第 5 讲 serving_sim 的 StepModel（roofline + 效率折扣 + 固定开销）。
    第 2 部分（选择性批处理）：
        ragged_forward：把几个长短不一、阶段不同的请求拼成一个 [ΣL, d] 的张量做线性层，注意力按请求拆开做，
        结果和"每个请求单独算"逐位一致（在第 2 讲的 numpy Transformer 上真算）。
    第 3 部分（投机解码）：
        adjust / residual / spec_output_dist：投机采样的三个零件，以及"输出分布 = 目标分布"的精确验证；
        TransformerLM / NGramLM：给第 2 讲的模型和一个 n-gram 表套上同样的接口（KV cache + 回滚）；
        speculative_generate：带 KV cache 回滚的投机解码循环（批大小 1），统计目标模型、草稿模型各调用几次；
        spec_sample_batch：同时跑上万个独立样本的投机采样（不用 KV cache，靠因果掩码处理长短不一），
            用来做"分布不变"的蒙特卡洛检验，以及两个**错误**变体的对照；
        expected_tokens / walltime_speedup / ops_factor / best_gamma：Leviathan 论文的公式（式 1、定理 3.8、3.11）。
    第 4 部分（硬件视角）：
        verify_time：目标模型一次验证 n 个 token（批大小 B、上下文 ctx）的 roofline 时间；
        spec_round：一轮投机解码 = γ 次草稿 + 1 次验证；在不同批大小、上下文长度下算加速比；
        tree_accept_lengths：树形验证（Medusa / EAGLE / SpecInfer）的接受长度，在贪心解码下精确计算。

【模型的简化（务必知道）】
    - 调度模拟里每步时间是 roofline 估计（第 5 讲的 StepModel，60% 峰值算力、80% 峰值带宽、每步 1 ms 固定开销），
      不是测量值。只模拟单块 GPU（Orca 论文还有张量并行、流水线并行，本讲只在讲义里讨论）。
    - 投机解码的正确性检验用随机初始化的小模型（没训练过）。它们的分布形状和真实语言模型不同，
      所以测出来的接受率只说明机制，不代表真实模型的数字。

【运行】python3 decode_lab.py   # 自检（约 2 秒）
【依赖】numpy；../lecture01（hw_specs、llm_math）、../lecture02（transformer_lab）、../lecture05（serving_sim.StepModel）、
       ../prereq5（serving_sim 需要的 comm_sim）。都只读。
"""

import math
import os
import sys
from dataclasses import dataclass

import numpy as np

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

from hw_specs import Chip  # noqa: E402
from llm_math import ModelConfig  # noqa: E402
from serving_sim import StepModel  # noqa: E402
from transformer_lab import (KVCache, attention_core, apply_rope, forward, linear, lm_head_matrix,  # noqa: E402
                             rms_norm, rope_angles, silu, softmax)


# =============================================================================
# 第 1 部分：调度模拟（Orca vs 请求级调度）
# =============================================================================
# Orca 论文的实验平台是 A100 40GB（Azure ND96asr A100 v4）。第 1 讲的 hw_specs 没有收录 A100，这里单独定义，
# 避免改动第 1 讲的芯片表（那会改变第 1 讲脚本的输出）。数字来自 NVIDIA A100 数据手册（稠密、不含稀疏加速）。
A100_40G = Chip(
    "A100-SXM4-40GB", "GPU",
    peak={"bf16": 312e12, "fp16": 312e12, "int8": 624e12},
    mem_bw=1.555e12, mem_cap=40e9, mem_kind="HBM2",
    link_bw=300e9, power_w=400, year=2020,
    notes="A100 数据手册：FP16/BF16 Tensor Core 312 TFLOPS，HBM2 1555 GB/s，NVLink 3 每卡 600 GB/s 双向",
)

# GPT-3 13B（Orca 表 1：40 层、hidden 5120、单卡）。GPT-3 的 MLP 是两个矩阵 d×4d、4d×d（GELU），
# 参数量 8d²；第 1 讲的 ModelConfig 按 SwiGLU 的 3·d·d_ff 算，所以取 d_ff = 8d/3 ≈ 13653，让参数量和 FLOPs 相同。
# 多头注意力（MHA）：40 个头、每头 128 维、KV 头数 = 40。词表 50257，lm_head 和词嵌入共享（tied）。
GPT3_13B = ModelConfig("GPT-3-13B", n_layers=40, d_model=5120, n_heads=40, n_kv_heads=40, d_head=128,
                       d_ff=13653, vocab=50257, tied=True)

# Llama 3.2 1B（Meta 模型卡）：常被用作 Llama 3.1 8B / 70B 的草稿模型——同一个分词器，词表一样。
LLAMA32_1B = ModelConfig("Llama-3.2-1B", n_layers=16, d_model=2048, n_heads=32, n_kv_heads=8, d_head=64,
                         d_ff=8192, vocab=128256, tied=True)


@dataclass
class Req:
    """
    调度模拟里的一个请求。

    【字段】
        rid     ：编号；arrival：到达时间（秒）
        prompt  ：输入 token 数
        output  ：**实际**生成的 token 数（≥ 1；第一个由处理输入的那次迭代产生）
        max_gen ：请求声明的"最多生成几个"（Orca 按 prompt + max_gen 预留 KV；Orca 的实验里 output = max_gen）
        gen     ：已经生成了几个（模拟过程中更新）
        finish  ：最后一个 token 出来的时间
        admitted：是否已经开始处理（已预留 / 已分配 KV）
    """
    rid: int
    arrival: float
    prompt: int
    output: int
    max_gen: int
    gen: int = 0
    finish: float = math.nan
    admitted: bool = False

    @property
    def latency(self):
        return self.finish - self.arrival

    @property
    def norm_latency(self):
        """Orca 论文图 10 的指标：端到端延迟 / 生成的 token 数（每个 token 平均多久）。"""
        return self.latency / self.output


def orca_workload(n, rate, seed=0, prompt_range=(32, 512), gen_range=(1, 128), max_gen=None):
    """
    Orca 论文 §6.2 的合成负载：输入长度 ~ U(32, 512)，max_gen_tokens ~ U(1, 128)，泊松到达。

    【参数】n：请求数；rate：每秒到达几个；seed：随机种子
        prompt_range、gen_range：闭区间 [lo, hi] 上的均匀分布
        max_gen：None 时和论文一样 max_gen = output（每个请求恰好生成 max_gen 个）；
                 给一个整数时，所有请求都声明 max_gen 个，但实际 output 仍按 gen_range 抽（模拟"提前遇到 EOS"）
    【返回】Req 列表（按到达时间排序）
    【边界】同一个 seed 下，改变 rate 只会整体压缩 / 拉伸到达时间，长度不变——不同速率之间的比较更干净。
    【Python 小知识】rng.integers(lo, hi + 1) 抽 [lo, hi] 之间的整数（右端是开区间，所以要 +1）。
    """
    rng = np.random.default_rng(seed)
    gaps = rng.exponential(1.0, n) / rate
    t = np.cumsum(gaps)
    out = []
    for i in range(n):
        p = int(rng.integers(prompt_range[0], prompt_range[1] + 1))
        o = int(rng.integers(gen_range[0], gen_range[1] + 1))
        mg = o if max_gen is None else max(int(max_gen), o)
        out.append(Req(i, float(t[i]), p, o, mg))
    return out


def run_request_level(reqs, step, max_bs, max_seq=2048):
    """
    请求级调度（Orca 论文图 2、图 3）：引擎空闲时从队列里取最多 max_bs 个请求组成一批，
    **整批跑完**才返回、才接下一批。

    【思路】
        1. 显存：像 FasterTransformer 那样给每个请求预留 max_seq 个 token 的 KV（不管它实际多长），
           所以一批最多 ⌊KV 容量 / max_seq⌋ 个请求（GPT-3 13B 在 A100 40GB 上是 6 个——论文说 FT 用不了 8 的批）。
        2. 第一次迭代：所有输入补齐（padding）到这批里最长的那个，一起算（[B, L_max] 的矩阵）。
        3. 之后每次迭代：全批 B 行一起 decode，直到最长的那个生成完。早结束的请求那几行还在白算（图 3 的 "-"）。
        4. 整批结束时，所有请求一起返回（早结束的也要等）。
    【参数】reqs：Req 列表（会被修改：写入 finish）；step：serving_sim.StepModel；max_bs：最大批大小；
        max_seq：每个请求预留的 KV 长度
    【返回】字典：slot_tokens（算过的 token 位置总数，含补齐和白算）、useful_tokens（真正需要的）、batches（批数）
    【边界】max_bs 超过显存允许的批大小时按显存截断；显存连一个请求都放不下时抛出 ValueError。
    """
    cap = step.kv_capacity_tokens() // max_seq
    if cap < 1:
        raise ValueError("显存放不下一个 max_seq 长度的请求")
    bs = min(max_bs, cap)
    t = 0.0
    i, n = 0, len(reqs)
    slot = useful = 0
    nb = 0
    while i < n:
        t = max(t, reqs[i].arrival)                      # 引擎空闲而队列为空时，等下一个请求到达
        batch = []
        while i < n and reqs[i].arrival <= t and len(batch) < bs:
            batch.append(reqs[i])
            i += 1
        B = len(batch)
        L_max = max(r.prompt for r in batch)
        O_max = max(r.output for r in batch)
        t += step.prefill_time([L_max] * B)             # 第一次迭代：补齐到 L_max
        for j in range(1, O_max):                       # 之后 O_max − 1 次 decode 迭代，全批一起
            t += step.decode_time(B, B * (L_max + j))
        for r in batch:
            r.finish = t
            r.gen = r.output
        slot += B * L_max + B * (O_max - 1)
        useful += sum(r.prompt + r.output - 1 for r in batch)
        nb += 1
    return {"slot_tokens": slot, "useful_tokens": useful, "batches": nb, "max_bs_used": bs}


def run_iteration_level(reqs, step, max_bs, reserve="max_gen", max_seq=2048, t_limit=None):
    """
    迭代级调度（Orca 论文 §3 S1、§4.2 算法 1）：每次迭代都重新从请求池里选最多 max_bs 个请求（按到达先后），
    新请求的输入和老请求的 decode token 放在同一次迭代里算（这要靠选择性批处理，见第 2 部分）。

    【KV 预留（reserve）】
        "max_gen"：Orca 的做法，请求第一次被选中时预留 prompt + max_gen 个 token 的 KV；预留不下就停止接纳新请求
                   （算法 1 第 23–25 行）。保证已接纳的请求一定能跑完，不会死锁。
        "max_seq"：每个请求预留 max_seq（像 FasterTransformer 那样按模型最大长度）——迭代级调度 + 粗放的显存管理。
        "none"   ：不预留，用多少分配多少（接纳时只要放得下输入）。显存满了以后，正在生成的请求每一步还要再写 1 个
                   token 的 KV，可能谁都动不了——**死锁**（没有抢占时）。第 7 讲的 PagedAttention + 抢占解决这个问题。
    【每次迭代的时间】StepModel.mixed_time：decode 的请求每个 1 个 token（读各自的 KV），新请求的全部输入 token
        （注意力按因果掩码算），权重读一遍。
    【参数】reqs：Req 列表（会被修改）；step：StepModel；max_bs：最大批大小；t_limit：模拟到这个时间就停（可选）
    【返回】字典：
        deadlock     ：是否死锁（以及发生时刻 deadlock_t）
        done         ：完成的请求数
        mean_reserved / mean_used：按迭代平均的"已预留 KV token 数"和"实际写入的 KV token 数"
        iters        ：迭代次数
    【边界】一个请求的 prompt + max_gen 超过全部显存时，"max_gen" 模式下它永远进不来（会卡住队列）——本模块的负载不会出现。
    """
    slots = step.kv_capacity_tokens()
    m = step.model
    t = 0.0
    i, n = 0, len(reqs)
    pool = []                      # 已到达、未完成的请求（按到达顺序）
    reserved = used = 0            # 预留的 KV token 数、实际写入的 KV token 数
    res_sum = used_sum = 0.0
    iters = 0
    done = 0
    while done < n:
        if t_limit is not None and t > t_limit:
            break
        while i < n and reqs[i].arrival <= t:
            pool.append(reqs[i])
            i += 1
        if not pool:
            t = reqs[i].arrival
            continue
        batch = []
        for r in pool:                                   # 算法 1 的 Select：按到达顺序
            if len(batch) == max_bs:
                break
            if not r.admitted:
                if reserve == "max_gen":
                    need = r.prompt + r.max_gen
                elif reserve == "max_seq":
                    need = max_seq
                else:
                    need = r.prompt                      # 只需要放下输入
                if (reserved if reserve != "none" else used) + need > slots:
                    break                                # 放不下就停止接纳（保持先到先服务）
                if reserve != "none":
                    reserved += need
                r.admitted = True
            elif reserve == "none" and used + 1 > slots:
                continue                                 # 没有空位写这一步的 KV，这个请求这次迭代动不了
            batch.append(r)
            if reserve == "none":
                used += r.prompt if r.gen == 0 else 1    # 先占上，保证同一次迭代里不超额
        if not batch:
            if i < n and reserve != "none":
                t = max(t, reqs[i].arrival)
                continue
            return {"deadlock": True, "deadlock_t": t, "done": done, "iters": iters,
                    "mean_reserved": res_sum / max(iters, 1), "mean_used": used_sum / max(iters, 1)}
        dec = [r for r in batch if r.gen > 0]
        pre = [r for r in batch if r.gen == 0]
        ctx = sum(r.prompt + r.gen for r in dec)
        ptoks = sum(r.prompt for r in pre)
        attn = m.attn_flops_per_token(1) * sum(r.prompt * (r.prompt + 1) / 2 for r in pre)
        t += step.mixed_time(len(dec), ctx, ptoks, attn)
        iters += 1
        if reserve != "none":
            used += ptoks + len(dec)
        for r in batch:
            r.gen += 1
            if r.gen >= r.output:
                r.finish = t
                done += 1
                used -= r.prompt + r.gen - 1             # 释放它写过的 KV（最后一个 token 不写 KV）
                if reserve == "max_gen":
                    reserved -= r.prompt + r.max_gen
                elif reserve == "max_seq":
                    reserved -= max_seq
        res_sum += reserved if reserve != "none" else used
        used_sum += used
        pool = [r for r in pool if r.gen < r.output]
    return {"deadlock": False, "deadlock_t": None, "done": done, "iters": iters,
            "mean_reserved": res_sum / max(iters, 1), "mean_used": used_sum / max(iters, 1)}


def latency_throughput(reqs, warmup=0.2):
    """
    一次模拟的两个指标：吞吐（每秒完成几个请求）和"每个 token 的延迟"的中位数（Orca 图 10 的纵轴）。

    【参数】warmup：去掉最早到达的这一比例的请求（系统刚启动时是空的，偏乐观）
    【返回】(吞吐 req/s, 中位数归一化延迟 秒/token, 第 90 百分位归一化延迟)
    【边界】吞吐 = 统计窗口里完成的请求数 / 窗口长度（从第一个被统计的请求到达，到最后一个完成）。
        系统过载时队列越积越长，吞吐趋于系统能力上限，延迟随请求数无限增长——表里看到的延迟是"这 n 个请求"的结果。
    """
    k = int(len(reqs) * warmup)
    rs = reqs[k:]
    span = max(r.finish for r in rs) - rs[0].arrival
    norm = np.array([r.norm_latency for r in rs])
    return len(rs) / span, float(np.median(norm)), float(np.percentile(norm, 90))


# =============================================================================
# 第 2 部分：选择性批处理（Orca §3 S2）
# =============================================================================
def ragged_forward(params, cfg, chunks, caches, counter=None):
    """
    选择性批处理的一次迭代：几个请求（各自处于不同阶段、新 token 数不同）一起算。

    【参数】
        chunks：列表，第 i 个元素是请求 i 这次迭代的新 token（一维整数数组，长度 T_i ≥ 1）
                ——"initiation"（Orca 的叫法，即 prefill）时是全部输入，"increment"（decode）时是 1 个 token
        caches：列表，第 i 个元素是请求 i 自己的 KVCache（批大小 1），已经存着它之前的 token
        counter：transformer_lab.Counter 或 None
    【返回】列表，第 i 个元素是请求 i 每个新 token 位置的 logits，形状 [T_i, vocab]
    【思路（Orca 图 5）】
        1. 把所有新 token 拼成一个 [1, N, d] 的张量（N = Σ T_i），没有"批"这一维——
           线性层、RMSNorm、残差、SwiGLU 都是"逐 token"的运算，不需要知道哪个 token 属于哪个请求。
        2. RoPE 用每个 token 自己的绝对位置（请求 i 的第 j 个新 token 位置是 cache_i.length + j）。
        3. 注意力之前 Split：按请求切开，每个请求和自己的 KV cache 做注意力（形状各不相同）；之后 Merge 回 [1, N, ·]。
        4. 只有第 3 步需要"请求"的概念，而注意力**没有权重**——不批处理它也不会多读权重（Orca §3）。
    【边界】每个 cache 的剩余容量要放得下这次的新 token；chunks 和 caches 一一对应。
    【Python 小知识】np.concatenate(列表) 把几个数组沿第 0 维拼起来；np.cumsum 求前缀和，
        配合切片 x[:, s:e] 可以把拼起来的张量再按请求切开。
    """
    d, H, G, dh = cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.d_head
    lens = [len(c) for c in chunks]
    starts = np.concatenate([[0], np.cumsum(lens)])
    N = int(starts[-1])
    toks = np.concatenate([np.asarray(c) for c in chunks])
    pos = np.concatenate([np.arange(cache.length, cache.length + T) for cache, T in zip(caches, lens)])
    x = params["embed"][toks][None]                          # [1, N, d]
    if counter is not None:
        counter.add("embed", w_elems=N * d, act_elems=N * d)
    cos, sin = rope_angles(pos, dh, cfg.rope_theta)
    for li, lp in enumerate(params["layers"]):
        h = rms_norm(x, lp["attn_norm"], cfg.norm_eps, counter)
        q = apply_rope(linear(h, lp["wq"], "q_proj", counter).reshape(1, N, H, dh), cos, sin, counter)
        k = apply_rope(linear(h, lp["wk"], "k_proj", counter).reshape(1, N, G, dh), cos, sin, counter)
        v = linear(h, lp["wv"], "v_proj", counter).reshape(1, N, G, dh)
        outs = []
        for r, cache in enumerate(caches):                   # Split：每个请求单独做注意力
            s, e = starts[r], starts[r + 1]
            cache.write(li, k[:, s:e], v[:, s:e])
            if counter is not None:
                counter.add("kv_write", kv_elems=2 * (e - s) * G * dh)
            K, V = cache.read(li, cache.length + (e - s))
            outs.append(attention_core(q[:, s:e], K, V, cache.length, counter))
        a = np.concatenate(outs, axis=1)                     # Merge：[1, N, H·dh]
        x = x + linear(a, lp["wo"], "o_proj", counter)
        h = rms_norm(x, lp["mlp_norm"], cfg.norm_eps, counter)
        m = silu(linear(h, lp["w_gate"], "mlp_gate", counter)) * linear(h, lp["w_up"], "mlp_up", counter)
        x = x + linear(m, lp["w_down"], "mlp_down", counter)
    for cache, T in zip(caches, lens):
        cache.length += T
    x = rms_norm(x, params["final_norm"], cfg.norm_eps, counter)
    logits = linear(x, lm_head_matrix(params), "lm_head", counter)[0]   # [N, vocab]
    return [logits[starts[r]:starts[r + 1]] for r in range(len(chunks))]


# =============================================================================
# 第 3 部分：投机解码
# =============================================================================
def adjust(logits, temperature=1.0, top_k=None):
    """
    "标准化采样"（Leviathan §2.2）：把各种采样方式都写成"从一个调整过的分布里采样"。

    【参数】logits：[..., V]；temperature：温度（0 表示贪心 = argmax）；top_k：只保留最大的 k 个（None 表示不截断）
    【返回】概率 [..., V]，每行和为 1
    【思路】贪心 = 把非最大值的概率清零再归一化（一个 one-hot 分布）；温度 = logits / T 再 softmax；
        top-k = 把第 k 名之后的置为 −∞。投机采样对**调整后的** p 和 q 做就保证输出服从调整后的 p——
        所以草稿和目标必须用同一种调整（同温度、同截断）。
    【边界】并列最大值时贪心取下标最小的那个（np.argmax 的约定）。
    """
    z = np.asarray(logits, dtype=np.float64)
    if temperature == 0:
        out = np.zeros_like(z)
        idx = np.argmax(z, axis=-1)
        np.put_along_axis(out, idx[..., None], 1.0, axis=-1)
        return out
    z = z / temperature
    if top_k is not None:
        kth = np.sort(z, axis=-1)[..., -top_k][..., None]
        z = np.where(z >= kth, z, -np.inf)
    return softmax(z, axis=-1)


def residual(p, q):
    """
    被拒绝之后重新采样用的分布 p'(x) = norm(max(0, p(x) − q(x)))（Leviathan §2.3）。

    【直觉】草稿 q 在某些 x 上"给多了"（q > p）——接受规则已经按 p/q 的比例把多出来的部分拒掉；
        另一些 x 上 q "给少了"（q < p），差额 p − q 正好由重新采样补上。
    【边界】p == q 时 max(0, p − q) 全为 0（这时永远不会拒绝，用不到它）；为了数值安全，返回 p。
    """
    r = np.maximum(p - q, 0.0)
    s = r.sum(axis=-1, keepdims=True)
    return np.where(s > 0, r / np.where(s > 0, s, 1.0), p)


def spec_output_dist(p, q):
    """
    一步投机采样输出 token 的**精确**分布（不抽样，直接按概率算）：
        P(输出 x) = q(x) · min(1, p(x)/q(x))  +  (1 − β) · p'(x)，其中 β = Σ_x min(p(x), q(x)) 是接受概率。
    第一项：草稿抽到 x 并被接受；第二项：草稿被拒绝（概率 1 − β），再从 p' 里抽到 x。

    【返回】(P_out, β)。数学上 P_out ≡ p（03 脚本演示 1 验证到浮点误差）。
    【推导】q·min(1, p/q) = min(p, q)；1 − β = Σ max(0, p − q)，p' = max(0, p − q) / (1 − β)
        ⇒ P_out = min(p, q) + max(0, p − q) = p。
    """
    beta = np.minimum(p, q).sum(axis=-1)
    with np.errstate(divide="ignore", invalid="ignore"):
        ratio = np.where(q > 0, np.minimum(1.0, p / np.where(q > 0, q, 1.0)), 0.0)
    out = q * ratio + (1.0 - beta)[..., None] * residual(p, q)
    return out, beta


def sample_rows(probs, rng):
    """
    每一行按自己的概率抽一个下标（向量化）：u ~ U(0,1)，取累积概率第一次超过 u 的位置。

    【参数】probs：[N, V]；rng：np.random.Generator
    【返回】[N] 整数数组
    【边界】累积和的最后一个数可能因为舍入略小于 1：用 searchsorted 后截断到 V − 1。
    【Python 小知识】np.cumsum(axis=1) 按行求累积和；(cdf < u[:, None]).sum(1) 数出"累积概率小于 u 的有几个"，
        正好是被抽中的下标。
    """
    cdf = np.cumsum(probs, axis=1)
    u = rng.random(probs.shape[0])
    idx = (cdf < u[:, None]).sum(axis=1)
    return np.minimum(idx, probs.shape[1] - 1)


class TransformerLM:
    """
    给第 2 讲的 numpy Transformer 套一个"语言模型"接口：带 KV cache 的前向，以及回滚。

    【方法】
        new_cache()             ：新建一个批大小为 1 的 KV cache
        logits(tokens, cache)   ：把一维的 tokens 接在 cache 后面算一遍，返回每个新位置的 logits [T, V]，cache 前进 T
        rollback(cache, length) ：把 cache 退回到只保留前 length 个 token（投机解码拒绝草稿后用）
    【为什么回滚这么简单】KVCache 的 write 总是从 cache.length 往后写，所以把 length 改小，
        后面的旧数据就"作废"了，下次写会直接覆盖（真实引擎里分页的 KV 也是改一下块表里的长度，第 7 讲）。
    """

    def __init__(self, params, cfg, name="model"):
        self.params, self.cfg, self.name = params, cfg, name

    def new_cache(self):
        return KVCache(self.cfg, 1, dtype=self.params["embed"].dtype)

    def logits(self, tokens, cache, counter=None):
        toks = np.asarray(tokens, dtype=np.int64)[None, :]
        return forward(self.params, self.cfg, toks, cache=cache, counter=counter)[0]

    def full_logits(self, tokens2d):
        """不用 cache，对一批 [N, L] 的序列整段前向，返回 [N, L, V]（蒙特卡洛检验用）。"""
        return forward(self.params, self.cfg, np.asarray(tokens2d, dtype=np.int64))

    @staticmethod
    def rollback(cache, length):
        if length > cache.length:
            raise ValueError("回滚只能变短")
        cache.length = length


class _TokCache:
    """NGramLM 的"cache"：只需要记住已经喂进去的 token。"""

    def __init__(self):
        self.toks = []

    @property
    def length(self):
        return len(self.toks)


class NGramLM:
    """
    一个二元语法（bigram）草稿模型：下一个 token 的 logits 只看上一个 token，查表即可（"几乎零成本"的草稿，c ≈ 0）。

    【参数】table：[V, V] 的 log 概率表，table[a] 是"上一个是 a 时下一个 token"的 log 概率
    【用法】NGramLM.fit(样本序列列表, V) 从目标模型自己生成的文本里数出二元组频率（+ 平滑）
    """

    def __init__(self, table, name="bigram"):
        self.table, self.name = table, name

    @classmethod
    def fit(cls, seqs, vocab, smoothing=0.01, name="bigram"):
        cnt = np.full((vocab, vocab), smoothing)
        for s in seqs:
            np.add.at(cnt, (np.asarray(s[:-1]), np.asarray(s[1:])), 1.0)
        return cls(np.log(cnt / cnt.sum(axis=1, keepdims=True)), name)

    def new_cache(self):
        return _TokCache()

    def logits(self, tokens, cache, counter=None):
        cache.toks.extend(int(t) for t in tokens)
        return self.table[np.asarray(tokens, dtype=np.int64)]

    @staticmethod
    def rollback(cache, length):
        del cache.toks[length:]


def speculative_generate(target, draft, prompt, n_new, gamma, temperature, rng, counter_t=None, counter_d=None):
    """
    投机解码（Leviathan 算法 1），批大小 1，带 KV cache 和回滚。

    【参数】target、draft：TransformerLM 或 NGramLM；prompt：一维整数数组（长度 ≥ 1）；n_new：要生成的 token 数；
        gamma：每轮草稿几个；temperature：采样温度（0 = 贪心）；rng：np.random.Generator；
        counter_t / counter_d：可选的 Counter，分别记目标模型和草稿模型的运算
    【返回】(生成的 token 列表（长度 n_new），统计字典)
        统计：target_calls（目标模型前向次数）、draft_calls、accepted（每轮被接受的草稿数列表）
    【不变式（"待定 token"）】两个模型的 cache 里都存着"已确定序列"去掉最后一个 token 的部分；
        最后一个已确定的 token 还没喂进模型——下一轮从它开始。
    【一轮的步骤】
        1. 草稿：先把草稿 cache 缺的 token 喂进去（通常只有待定的那 1 个；上一轮全部接受时还缺 d_γ，共 2 个），
           得到 q_1、抽 d_1；再依次喂 d_1 … d_{γ−1}，得到 q_2 … q_γ、抽 d_2 … d_γ。草稿模型一共前向 γ 次。
        2. 验证：目标模型一次前向 [待定, d_1, …, d_γ]（γ+1 个 token），得到 p_1 … p_{γ+1}。
        3. 依次检验：r ~ U(0,1)，r < p_i(d_i)/q_i(d_i) 就接受 d_i，否则拒绝并从 residual(p_i, q_i) 抽一个 token、结束；
           全部接受时，再从 p_{γ+1} 白拿一个 token。
        4. 回滚：目标 cache 退回到"新的已确定序列 − 1"；草稿 cache 退回到不超过这个长度。
    【边界】贪心（temperature = 0）时 p、q 是 one-hot：接受 ⇔ d_i 正好是目标模型的 argmax，被拒绝时 residual 就是
        目标的 argmax——输出和目标模型自己贪心生成的逐 token 相同（03 脚本演示 4 验证）。
    """
    seq = [int(t) for t in prompt]
    tc, dc = target.new_cache(), draft.new_cache()
    if len(seq) > 1:                                         # 预填充：除了最后一个 token 都喂进去
        target.logits(seq[:-1], tc, counter_t)
        draft.logits(seq[:-1], dc, counter_d)
    n0 = len(seq)
    stats = {"target_calls": 0, "draft_calls": 0, "accepted": []}
    while len(seq) - n0 < n_new:
        ds, qs = [], []
        feed = seq[dc.length:]                               # 草稿 cache 缺的 token
        for i in range(gamma):
            lg = draft.logits(feed, dc, counter_d)[-1]
            stats["draft_calls"] += 1
            q = adjust(lg, temperature)
            d = int(sample_rows(q[None], rng)[0])
            ds.append(d)
            qs.append(q)
            feed = [d]
        P = adjust(target.logits([seq[-1]] + ds, tc, counter_t), temperature)   # [γ+1, V]
        stats["target_calls"] += 1
        n_acc = 0
        tok = None
        for i in range(gamma):
            if rng.random() < P[i, ds[i]] / qs[i][ds[i]]:
                n_acc += 1
            else:
                tok = int(sample_rows(residual(P[i], qs[i])[None], rng)[0])
                break
        if tok is None:
            tok = int(sample_rows(P[gamma][None], rng)[0])
        stats["accepted"].append(n_acc)
        seq.extend(ds[:n_acc] + [tok])
        target.rollback(tc, len(seq) - 1)
        draft.rollback(dc, min(dc.length, len(seq) - 1))
    return seq[n0:n0 + n_new], stats


def autoregressive_generate(model, prompt, n_new, temperature, rng, counter=None):
    """
    普通的自回归生成（基线）：每步前向 1 个 token、按调整后的分布抽样。

    【返回】(token 列表, 前向次数)
    """
    seq = [int(t) for t in prompt]
    cache = model.new_cache()
    lg = model.logits(seq, cache, counter)[-1]
    out = []
    calls = 1
    for i in range(n_new):
        tok = int(sample_rows(adjust(lg, temperature)[None], rng)[0])
        out.append(tok)
        if i == n_new - 1:
            break
        lg = model.logits([tok], cache, counter)[-1]
        calls += 1
    return out, calls


def spec_sample_batch(target, draft, prompt, n_out, gamma, temperature, n_samples, rng, variant="correct"):
    """
    同时跑 n_samples 个**独立**的投机采样，每个都生成至少 n_out 个 token（蒙特卡洛检验"输出分布不变"用）。

    【为什么不用 KV cache】每个样本被接受的草稿数不同，几轮之后长度各不相同，第 2 讲的 KVCache 只有一个公共长度。
        这里干脆每次都对整段序列 [N, L] 重新前向（序列很短，模型很小，不贵），用每个样本自己的长度去取对应位置的
        logits。因果掩码保证：某个样本长度之后的"垃圾"位置不会影响它前面的位置。
    【参数】target、draft：TransformerLM；prompt：一维数组；n_out：每个样本至少要几个 token；gamma；temperature；
        n_samples：样本数；rng
        variant："correct"      ：Leviathan 的规则；
                 "resample_p"   ：错误变体 1——被拒绝后直接从 p 重新采样（而不是 p − q 的残差）；
                 "accept_if_p_ge_q"：错误变体 2——"p(d) ≥ q(d) 就接受，否则拒绝"（去掉随机性的接受规则）。
    【返回】[n_samples, n_out] 的整数数组（每个样本的前 n_out 个输出 token）
    【边界】所有样本都从同一个 prompt 开始；n_out ≥ 1；gamma ≥ 1。
    """
    N = n_samples
    Lp = len(prompt)
    Lmax = Lp + n_out + gamma + 1
    toks = np.zeros((N, Lmax), dtype=np.int64)
    toks[:, :Lp] = prompt
    length = np.full(N, Lp)
    while (length - Lp).min() < n_out:
        idx = np.nonzero(length - Lp < n_out)[0]                           # 还没生成够的样本才参加这一轮
        work, ln = toks[idx], length[idx]
        M = len(idx)
        ar = np.arange(M)
        ds = np.zeros((M, gamma), dtype=np.int64)
        qlist = []
        for i in range(gamma):                                              # 草稿：依次抽 γ 个
            L = int((ln + i).max())
            q = adjust(draft.full_logits(work[:, :L])[ar, ln - 1 + i], temperature)
            d = sample_rows(q, rng)
            work[ar, ln + i] = d
            ds[:, i] = d
            qlist.append(q)
        qs = np.stack(qlist, axis=1)                                        # [M, γ, V]
        L = int((ln + gamma).max())
        lgt = target.full_logits(work[:, :L])                               # 目标模型一次前向验证
        P = adjust(np.stack([lgt[ar, ln - 1 + i] for i in range(gamma + 1)], axis=1), temperature)   # [M, γ+1, V]
        pd = np.take_along_axis(P[:, :gamma], ds[..., None], axis=2)[..., 0]
        qd = np.take_along_axis(qs, ds[..., None], axis=2)[..., 0]
        if variant == "accept_if_p_ge_q":
            acc = pd >= qd
        else:
            acc = rng.random((M, gamma)) < pd / qd
        n_acc = np.cumprod(acc, axis=1).sum(axis=1)                         # 前面连续接受了几个
        p_next = P[ar, n_acc]                                               # [M, V]
        q_next = np.concatenate([qs, np.zeros((M, 1, qs.shape[2]))], axis=1)[ar, n_acc]
        rejected = n_acc < gamma
        if variant == "resample_p":
            dist = p_next
        else:
            dist = np.where(rejected[:, None], residual(p_next, q_next), p_next)
        work[ar, ln + n_acc] = sample_rows(dist, rng)
        toks[idx] = work
        length[idx] = ln + n_acc + 1
    return toks[:, Lp:Lp + n_out]


# ---- Leviathan 论文的公式 ----
def expected_tokens(alpha, gamma):
    """
    式 (1)：一轮（γ 个草稿 + 1 次验证）平均产出几个 token = (1 − α^{γ+1}) / (1 − α)。

    【推导】产出的 token 数 = 1 + 被接受的草稿数；第 i 个草稿被接受要求前 i 个都被接受（概率 α^i），
        所以期望 = 1 + α + α² + … + α^γ（等比数列求和）。
    【边界】α = 1 时等于 γ + 1；γ = 0 时等于 1（普通 decode）。假设每个位置被接受与否**独立同分布**。
    """
    if alpha >= 1.0:
        return gamma + 1.0
    return (1.0 - alpha ** (gamma + 1)) / (1.0 - alpha)


def walltime_speedup(alpha, gamma, c):
    """
    定理 3.8：墙钟时间的期望加速比 = (1 − α^{γ+1}) / ((1 − α)(γc + 1))。

    【参数】c：代价系数 = 草稿模型跑一次的时间 / 目标模型跑一次的时间（论文假设目标模型一次验证 γ+1 个 token
        和一次只算 1 个 token 一样快——decode 受限于带宽时近似成立，第 4 部分会看到批大时不成立）
    """
    return expected_tokens(alpha, gamma) / (gamma * c + 1.0)


def ops_factor(alpha, gamma, c_hat):
    """
    定理 3.11：总运算量增加的倍数 = (1 − α)(γĉ + γ + 1) / (1 − α^{γ+1})。

    【参数】c_hat：草稿模型每个 token 的运算量 / 目标模型每个 token 的运算量
    【直觉】每轮目标模型算了 γ+1 个 token，但平均只留下 expected_tokens 个；被拒绝的那些位置的运算白做了。
    """
    return (gamma * c_hat + gamma + 1.0) / expected_tokens(alpha, gamma)


def best_gamma(alpha, c, gamma_max=64):
    """
    在 1 … gamma_max 里找让 walltime_speedup 最大的 γ（论文图 3）。返回 (γ*, 加速比)。

    【边界】如果最好的加速比 ≤ 1（草稿太贵或太不准），返回 (0, 1.0)：不用投机解码。
    """
    best = (0, 1.0)
    for g in range(1, gamma_max + 1):
        s = walltime_speedup(alpha, g, c)
        if s > best[1]:
            best = (g, s)
    return best


# =============================================================================
# 第 4 部分：硬件视角
# =============================================================================
def verify_time(model, chip, batch, ctx, n_tokens, w_bytes=2.0, kv_bytes=2.0, mfu=0.6, mbu=0.8, overhead=0.0):
    """
    一次前向：批里 batch 个序列，每个序列已有 ctx 个 token 的上下文，这次各送进 n_tokens 个 token
    （普通 decode：n_tokens = 1；投机解码验证：n_tokens = γ + 1）。roofline 时间 + 固定开销。

    【公式】
        FLOPs = batch · n · 每 token 矩阵乘法 FLOPs + batch · Σ_{j=1..n} 注意力(ctx + j)
        字节  = 权重 + batch · (ctx + n) · 每 token KV（读一遍已有的 KV；同一个序列的 n 个 query 一起读，
                 融合注意力内核让它们共享这一次读取）
        时间  = max(FLOPs / (峰值 × mfu), 字节 / (带宽 × mbu)) + overhead
    【要点】n 从 1 变成 γ+1 时：读权重不变、读 KV 几乎不变，只有运算量 ×(γ+1)。
        所以小批、长上下文时验证几乎是"免费"的；批大到运算成为瓶颈时，验证的代价就按 γ+1 增长。
    【返回】秒
    """
    prec = "bf16" if w_bytes >= 2 else "fp8"
    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)
    by = model.matmul_params() * w_bytes + batch * (ctx + n_tokens) * model.kv_bytes_per_token(kv_bytes)
    return max(flops / (chip.peak_flops(prec) * mfu), by / (chip.mem_bw * mbu)) + overhead


def spec_round(target, draft, chip, batch, ctx, alpha, gamma, overhead_t=1e-3, overhead_d=0.3e-3, **kw):
    """
    批里每个序列都做投机解码时，一轮的时间与加速比。

    【模型】一轮 = γ 次草稿（草稿模型 decode，批 batch、上下文 ctx）+ 1 次目标模型验证 γ+1 个 token；
        平均产出 expected_tokens(α, γ) 个 token / 序列。基线 = 目标模型普通 decode 一步（1 个 token / 序列）。
        （简化：批里每个序列的接受情况独立，且用同一个 γ；接受数不同的序列下一轮长度不同，这里忽略。）
    【参数】overhead_t / overhead_d：目标 / 草稿每次前向的固定开销（秒，**假设值**）；kw 传给 verify_time
    【返回】字典：t_base（普通 decode 一步）、t_draft（草稿一步）、t_verify、c（= t_draft / t_base）、
        verify_ratio（= t_verify / t_base）、speedup
    """
    t_base = verify_time(target, chip, batch, ctx, 1, overhead=overhead_t, **kw)
    t_draft = verify_time(draft, chip, batch, ctx, 1, overhead=overhead_d, **kw)
    t_ver = verify_time(target, chip, batch, ctx, gamma + 1, overhead=overhead_t, **kw)
    t_round = gamma * t_draft + t_ver
    return {"t_base": t_base, "t_draft": t_draft, "t_verify": t_ver, "c": t_draft / t_base,
            "verify_ratio": t_ver / t_base, "speedup": expected_tokens(alpha, gamma) * t_base / t_round}


def best_round(target, draft, chip, batch, ctx, alpha, gamma_max=16, **kw):
    """在 γ = 0 … gamma_max 里找加速比最大的（γ = 0 表示不投机，加速比 1）。返回 (γ*, spec_round 的字典或 None)。"""
    best = (0, None)
    best_s = 1.0
    for g in range(1, gamma_max + 1):
        r = spec_round(target, draft, chip, batch, ctx, alpha, g, **kw)
        if r["speedup"] > best_s:
            best, best_s = (g, r), r["speedup"]
    return best


def draft_ranks(target_seq_logits_draft, seq):
    """
    贪心解码下，树形验证的接受长度只取决于：在每个位置，目标模型的下一个 token 在草稿模型的分布里排第几。

    【参数】target_seq_logits_draft：草稿模型在一条序列上每个位置的 logits [L, V]（输入是目标模型贪心生成的真实序列）；
        seq：这条序列 [L]（位置 t 的"正确答案"是 seq[t+1]）
    【返回】ranks：长度 L−1 的整数数组，ranks[t] = 正确答案在草稿 logits[t] 里的名次（0 = 第一名）
    """
    lg = np.asarray(target_seq_logits_draft)[:-1]
    ans = np.asarray(seq)[1:]
    ans_score = lg[np.arange(len(ans)), ans]
    return (lg > ans_score[:, None]).sum(axis=1)


def tree_accept_lengths(ranks, branching):
    """
    静态树（每层的分支数 branching = (k_1, …, k_D)）在贪心解码下，从每个位置出发能接受几个草稿 token。

    【思路】贪心时目标模型的继续是确定的。树的第 d 层有草稿在"正确前缀"上的前 k_d 名候选；
        只要正确答案在前 k_d 名里，这一层就有一个节点被接受，继续往下一层。
        所以从位置 t 出发的接受长度 = 最大的 m ≤ D，使得 ranks[t+d−1] < k_d 对 d = 1 … m 都成立。
        链式草稿（普通投机解码）就是 k_d 全为 1 的特例。
    【返回】(每个起点的接受长度数组, 树里的节点数 = 每轮目标模型要验证的 token 数 − 1)
    【边界】序列末尾不够 D 个位置的起点被丢掉。
    【系统联系】树形验证要用"树形注意力掩码"：每个节点只看得见自己的祖先（Medusa、EAGLE、SpecInfer 都这么做）。
        节点越多，接受长度越长，但每次验证的 token 也越多——批小时几乎免费，批大时就不是了（05 脚本演示 3）。
    """
    D = len(branching)
    L = len(ranks) - D + 1
    ok = np.stack([ranks[d:d + L] < branching[d] for d in range(D)], axis=1)   # [L, D]
    acc = np.cumprod(ok, axis=1).sum(axis=1)
    nodes = int(sum(np.prod(branching[:d + 1]) for d in range(D)))
    return acc, nodes


# =============================================================================
# 自检
# =============================================================================
def _selftest():
    from transformer_lab import Config, init_params
    # 1. 投机采样的精确分布 = p
    rng = np.random.default_rng(0)
    p = rng.dirichlet(np.ones(50))
    q = rng.dirichlet(np.ones(50))
    out, beta = spec_output_dist(p, q)
    assert np.max(np.abs(out - p)) < 1e-15, "投机采样的输出分布应该等于 p"
    assert abs(beta - (1 - 0.5 * np.abs(p - q).sum())) < 1e-15, "β = 1 − TV(p, q)"
    # 2. 公式
    assert abs(expected_tokens(0.8, 5) - (1 - 0.8 ** 6) / 0.2) < 1e-12
    assert abs(walltime_speedup(0.9, 10, 0.0) - 6.861894039) < 1e-6    # Leviathan 表 1 的 6.86X
    # 3. 选择性批处理 = 逐个请求单独算
    cfg = Config(vocab=64, d_model=32, n_layers=2, n_heads=4, n_kv_heads=2, d_head=8, d_ff=96, max_seq=32)
    params = init_params(cfg, seed=1)
    chunks = [np.array([3, 5, 7]), np.array([9]), np.array([1, 2])]
    caches = [KVCache(cfg, 1) for _ in chunks]
    forward(params, cfg, np.array([[4, 4]]), cache=caches[1])        # 请求 1 已经有 2 个 token 的上下文
    ref_caches = [KVCache(cfg, 1) for _ in chunks]
    forward(params, cfg, np.array([[4, 4]]), cache=ref_caches[1])
    outs = ragged_forward(params, cfg, chunks, caches)
    for c, o, rc in zip(chunks, outs, ref_caches):
        ref = forward(params, cfg, c[None], cache=rc)[0]
        assert np.max(np.abs(o - ref)) < 1e-12
    # 4. 贪心投机解码 = 贪心解码
    tgt = TransformerLM(params, cfg, "target")
    dp = dict(params)
    dp["layers"] = params["layers"][:1]
    from transformer_lab import config_like
    drf = TransformerLM(dp, config_like(cfg, n_layers=1), "draft")
    prompt = np.array([1, 2, 3])
    a, _ = autoregressive_generate(tgt, prompt, 20, 0, np.random.default_rng(0))
    b, st = speculative_generate(tgt, drf, prompt, 20, 3, 0, np.random.default_rng(0))
    assert a == b, "贪心投机解码必须和贪心解码逐 token 相同"
    # 5. 调度：请求级的 FT 批大小上限 = 6（论文：13B 用不了 8）
    sm = StepModel(GPT3_13B, A100_40G)
    assert sm.kv_capacity_tokens() // 2048 == 6
    print("decode_lab 自检通过：投机采样分布 = p、β = 1 − TV、公式 = 论文表 1、选择性批处理 = 逐个计算、"
          f"贪心投机 = 贪心（目标模型前向 {st['target_calls']} 次 vs 20 次）、FT 批上限 6")


if __name__ == "__main__":
    _selftest()