"""
decode_lab.py —— 第 18 讲下半的共享模块：分页的 KV 缓存（PagedAttention）、服务模拟、推测解码
=====================================================================

【这个文件要回答的问题】
    1. 服务一个 LLM 时，每个请求的 KV 缓存随生成的 token 数增长。如果像普通数组那样给每个请求预留一整段**连续**的显存，
       会浪费很多（预留了却还没用、永远用不到、请求之间的空隙）。PagedAttention 借用操作系统的"虚拟内存 + 分页"：
       把缓存切成固定大小的块，用一张"块表"记录每个请求的逻辑块对应哪个物理块——块不必连续，用到时才分配。
    2. 解码时每生成一个 token 都要把整个大模型的权重从显存读一遍（访存受限，前置讲 P3 §6.2）。
       推测解码让一个小模型先猜 γ 个 token，大模型**一次**前向就能检查全部 γ 个；按一个巧妙的接受 / 拒绝规则，
       输出的分布和只用大模型逐个采样**完全相同**。

【本文件提供】
    分页缓存：BlockAllocator（空闲块列表 + 引用计数）、PagedKVCache（块表、追加、fork 共享、写时复制、按块做注意力）
    服务模拟：make_requests（长尾的提示 / 输出长度）、ContiguousAllocator（首次适配的连续分配）、simulate_serving
        （max / oracle / paged 三种策略：每步每个请求生成一个 token，记录显存里有多少真的存着 token、批有多大）
    推测解码：residual、first_token_distribution（一步的输出分布的闭式解）、exact_prefix_distribution（枚举所有随机性，
        精确算出前 K 个 token 的分布）、speculative_sample（批量的推测采样，带接受率统计）、ancestral_sample、
        NGramDraft（把第 1 讲的 n-gram 包装成草稿模型）、expected_tokens / walltime_speedup（Leviathan et al. 的公式）

【边界】
    - 服务模拟是"离散时间"的玩具：每一步所有在跑的请求各生成一个 token、提示的预填充瞬间完成、请求在时刻 0 全部到达。
      它只用来比较不同的显存管理策略能同时装下多少请求，不模拟真实的调度器。
    - speculative_sample 每次都把整个序列重新喂给模型（不用 KV 缓存）：我们的模型很小、句子很短，这样代码最简单；
      真实系统里草稿模型和目标模型都带 KV 缓存，被拒绝的 token 的缓存要回滚（第 3 讲 kv_cache 的截断）。
"""

import math
import random

import torch
import torch.nn.functional as F

from flash_lab import block_state, finalize, merge_states


# =====================================================================
# 1. 分页的 KV 缓存
# =====================================================================

class BlockAllocator:
    """
    物理块的分配器：一个空闲块的栈 + 每个块的引用计数。

    【参数】n_blocks：物理块的总数
    【方法】allocate() 取一个空闲块（引用计数置 1）；incref(b)；decref(b)（降到 0 时放回空闲栈）
    【边界】没有空闲块时 allocate 抛出 MemoryError——调用者（调度器）要决定是等待还是抢占别的请求。
    【为什么要引用计数】并行采样、束搜索时几个序列共享同一段提示的缓存：一个物理块可能同时属于好几个序列，
        只有最后一个序列不再需要它时才能回收。
    """

    def __init__(self, n_blocks):
        self.n_blocks = n_blocks
        self.free = list(range(n_blocks - 1, -1, -1))    # 栈顶是 0 号块（pop() 从末尾取）
        self.ref = [0] * n_blocks

    def allocate(self):
        if not self.free:
            raise MemoryError("没有空闲的物理块")
        b = self.free.pop()
        self.ref[b] = 1
        return b

    def incref(self, b):
        self.ref[b] += 1

    def decref(self, b):
        self.ref[b] -= 1
        if self.ref[b] == 0:
            self.free.append(b)

    @property
    def n_free(self):
        return len(self.free)


class PagedKVCache:
    """
    分页的 KV 缓存（一层、一个头的简化版：每个 token 存一个 d 维的 k 和一个 d 维的 v）。

    【数据结构】
        K、V        : 形状 (n_blocks, block_size, d) 的"物理显存"
        tables[sid] : 序列 sid 的块表——第 i 个逻辑块存在哪个物理块里（slides："类似页表，把逻辑地址映射到物理地址"）
        lengths[sid]: 序列 sid 已经存了多少个 token
        fill[b]     : 物理块 b 里写了几格
    【方法】
        new_seq、append（必要时分配新块；最后一块被共享时先写时复制）、fork（新序列共享父序列的全部块）、
        free_seq、gather（按块表拼成连续的 K、V，用来核对）、attention（按块表逐块做注意力，用在线 softmax 合并）
    【边界】append 在没有空闲块时抛出 MemoryError；一个序列的最后一块可能没写满——这就是分页唯一的"内部碎片"，
        每个序列最多 block_size − 1 格。
    """

    def __init__(self, n_blocks, block_size, d, dtype=torch.float64):
        self.block_size, self.d = block_size, d
        self.K = torch.zeros(n_blocks, block_size, d, dtype=dtype)
        self.V = torch.zeros(n_blocks, block_size, d, dtype=dtype)
        self.alloc = BlockAllocator(n_blocks)
        self.fill = [0] * n_blocks
        self.tables, self.lengths = {}, {}
        self.cow_copies = 0

    def new_seq(self, sid):
        self.tables[sid], self.lengths[sid] = [], 0

    def append(self, sid, k, v):
        """给序列 sid 追加一个 token 的 (k, v)。"""
        n, bs = self.lengths[sid], self.block_size
        slot = n % bs
        table = self.tables[sid]
        if slot == 0:                                       # 上一块写满了（或还没有块）：分配一个新块
            table.append(self.alloc.allocate())
        elif self.alloc.ref[table[-1]] > 1:                 # 最后一块和别的序列共享：写时复制（copy-on-write）
            old, new = table[-1], self.alloc.allocate()
            self.K[new], self.V[new] = self.K[old], self.V[old]
            self.fill[new] = self.fill[old]
            self.alloc.decref(old)
            table[-1] = new
            self.cow_copies += 1
        b = table[-1]
        self.K[b, slot], self.V[b, slot] = k, v
        self.fill[b] = max(self.fill[b], slot + 1)
        self.lengths[sid] = n + 1

    def fork(self, src, dst):
        """新序列 dst 共享 src 已有的全部块（并行采样：同一个提示生成几个不同的回答）。"""
        self.tables[dst] = list(self.tables[src])
        self.lengths[dst] = self.lengths[src]
        for b in self.tables[dst]:
            self.alloc.incref(b)

    def free_seq(self, sid):
        for b in self.tables.pop(sid):
            self.alloc.decref(b)
        del self.lengths[sid]

    def gather(self, sid):
        """按块表把序列 sid 的 K、V 拼成连续的 (n, d) 张量（只用来核对）。"""
        n = self.lengths[sid]
        idx = self.tables[sid]
        return self.K[idx].reshape(-1, self.d)[:n], self.V[idx].reshape(-1, self.d)[:n]

    def attention(self, sid, q):
        """
        单个查询 q（形状 (d,)）对序列 sid 的全部缓存做注意力，**逐块**读物理块，用 (m, l, o) 合并（01 号脚本第 5 部分）。
        【返回】(d,) 的输出
        【GenAI 联系】vLLM 的 PagedAttention 核函数就是这样：按块表找到分散的物理块，一块一块地算，
            块与块之间用在线 softmax 合并——和 FlashAttention 是同一个技巧。
        """
        n, bs = self.lengths[sid], self.block_size
        state = None
        for i, b in enumerate(self.tables[sid]):
            cnt = min(bs, n - i * bs)
            S = (self.K[b, :cnt] @ q / math.sqrt(self.d)).unsqueeze(0)
            st = block_state(S, self.V[b, :cnt])
            state = st if state is None else merge_states(state, st)
        out, _ = finalize(state)
        return out[0]

    def physical_tokens(self):
        """物理上写着 token 的格子数（共享的块只算一次）。"""
        return sum(self.fill[b] for b in range(self.alloc.n_blocks) if self.alloc.ref[b] > 0)

    def allocated_slots(self):
        return (self.alloc.n_blocks - self.alloc.n_free) * self.block_size


# =====================================================================
# 2. 服务模拟：连续预留 vs 分页
# =====================================================================

def make_requests(n, seed, max_len=2048):
    """
    n 个请求的 (提示长度, 输出长度)，都是长尾分布（对数正态）：大多数几十到几百个 token，偶尔很长。
    【边界】提示长度截到 [4, 1024]，输出长度截到 [1, max_len − 提示长度]，保证每个请求都不超过 max_len。
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        p = int(min(1024, max(4, rng.lognormvariate(math.log(150), 0.8))))
        o = int(min(max_len - p, max(1, rng.lognormvariate(math.log(250), 0.9))))
        out.append((p, o))
    return out


class ContiguousAllocator:
    """
    在一段连续的"显存"（capacity 个 token 的格子）里，用首次适配（first fit）分配连续的区间。
    【方法】allocate(n) 返回起点或 None（找不到够大的空隙）；free(start)
    【为什么会有外部碎片】区间有长有短、释放的时间不同，空闲的格子会被切成许多小段：
        总的空闲可能够，却没有哪一段连续地够大。
    """

    def __init__(self, capacity):
        self.capacity = capacity
        self.used = []                                   # 按起点排序的 (start, size)

    def allocate(self, n):
        prev = 0
        for i, (s, size) in enumerate(self.used):
            if s - prev >= n:
                self.used.insert(i, (prev, n))
                return prev
            prev = s + size
        if self.capacity - prev >= n:
            self.used.append((prev, n))
            return prev
        return None

    def free(self, start):
        for i, (s, _) in enumerate(self.used):
            if s == start:
                del self.used[i]
                return
        raise KeyError(start)

    def largest_gap(self):
        prev, best = 0, 0
        for s, size in self.used:
            best = max(best, s - prev)
            prev = s + size
        return max(best, self.capacity - prev)


def simulate_serving(requests, capacity, policy, block_size=16, max_len=2048, watermark=0.01):
    """
    离散时间的服务模拟：请求按顺序排队（先来先服务），每一步：
        1. 从队头开始，能放下就接纳（预填充瞬间完成：提示的 KV 立刻写进缓存）；放不下就停（不跳过队头）；
        2. 每个在跑的请求生成 1 个 token（KV 缓存长 1）；
        3. 生成完的请求释放显存。

    【参数】
        requests : [(提示长度, 输出长度), ...]
        capacity : KV 缓存一共能放多少个 token
        policy   : "max"    —— 每个请求预留 max_len 个连续的格子（不知道会生成多长，只能按上限留）
                   "oracle" —— 预留恰好"提示 + 输出"那么多（假设事先知道输出长度；现实中做不到）
                   "paged"  —— 按 block_size 分块，用到时才分配；没有空闲块时抢占最晚到达的请求：
                               释放它的块、放回队头，之后重新预填充（vLLM 的"重算"；已生成的 token 不丢，只是要重新算一遍 KV）
        watermark: 分页时，接纳新请求后至少还要留下这么多比例的空闲块（vLLM 默认 1%），给在跑的请求长大用，减少抢占
    【返回】字典：steps（总步数）、busy_steps（还有请求在排队的步数：这时显存是瓶颈，下面的平均值都只在这些步上算）、
        mean_batch（平均同时在跑的请求数）、token_share（平均有多少比例的显存真的存着 token）、
        reserved_share（预留了、以后会用到的比例；分页时是最后一块没写满的格子）、never_share（预留了、永远用不到的比例）、
        free_share（没有分配出去的比例）、preemptions（抢占次数）、
        kv_token_steps（每步所有请求的上下文长度之和，累加起来——解码时要从显存读的 KV 总量，以 token 计）
    """
    queue = list(range(len(requests)))
    length = {rid: p for rid, (p, _) in enumerate(requests)}     # 每个请求当前的总长度（提示 + 已生成）
    running = {}                    # 请求号 -> {"start": 连续分配的起点, "reserved": 预留格子} 或 {"blocks": 块数}
    done, steps, preempt = 0, 0, 0
    sum_batch = sum_tok = sum_res = sum_never = sum_free = kv_token_steps = 0
    cont = ContiguousAllocator(capacity)
    free_blocks = capacity // block_size
    keep = int(watermark * (capacity // block_size))
    busy = 0
    while done < len(requests):
        # 1. 接纳
        while queue:
            rid = queue[0]
            p, o = requests[rid]
            n = length[rid]
            if policy in ("max", "oracle"):
                need = max_len if policy == "max" else p + o
                start = cont.allocate(need)
                if start is None:
                    break
                running[rid] = {"start": start, "reserved": need}
            else:
                if math.ceil((n + 1) / block_size) > free_blocks - keep:   # 放得下已有的 token 和下一个新 token，还留出水位线
                    break
                free_blocks -= math.ceil(n / block_size)
                running[rid] = {"blocks": math.ceil(n / block_size)}
            queue.pop(0)
        if not running:
            raise RuntimeError("一个请求都放不下：capacity 太小")
        # 2. 每个请求生成一个 token
        for rid in sorted(running):
            if rid not in running:                                   # 这一步里已经被抢占了
                continue
            r = running[rid]
            if policy == "paged" and length[rid] % block_size == 0:  # 当前块写满，需要新块
                while free_blocks == 0 and max(running) != rid:      # 抢占最晚到达的请求
                    victim = max(running)
                    free_blocks += running.pop(victim)["blocks"]
                    queue.insert(0, victim)
                    preempt += 1
                if free_blocks == 0:                                 # 自己就是最晚的：这一步先不生成
                    continue
                free_blocks -= 1
                r["blocks"] += 1
            length[rid] += 1
        steps += 1
        tok = sum(length[rid] for rid in running)
        kv_token_steps += tok
        if queue:                                                    # 统计只在"还有人排队"的步上做（在释放之前）
            busy += 1
        else:
            tok = None
        if tok is not None:
            sum_batch += len(running)
            sum_tok += tok
            if policy in ("max", "oracle"):
                final = sum(requests[rid][0] + requests[rid][1] for rid in running)
                reserved = sum(r["reserved"] for r in running.values())
                sum_res += final - tok
                sum_never += reserved - final
                sum_free += capacity - reserved
            else:
                allocated = sum(r["blocks"] for r in running.values()) * block_size
                sum_res += allocated - tok
                sum_free += capacity - allocated
        # 3. 释放
        for rid in list(running):
            p, o = requests[rid]
            if length[rid] >= p + o:
                r = running.pop(rid)
                if policy in ("max", "oracle"):
                    cont.free(r["start"])
                else:
                    free_blocks += r["blocks"]
                done += 1
    total = busy * capacity
    return {"steps": steps, "busy_steps": busy, "mean_batch": sum_batch / busy, "token_share": sum_tok / total,
            "reserved_share": sum_res / total, "never_share": sum_never / total, "free_share": sum_free / total,
            "preemptions": preempt, "kv_token_steps": kv_token_steps}


# =====================================================================
# 3. 推测解码
# =====================================================================

def residual(p, q):
    """
    拒绝之后改从 norm(max(0, p − q)) 采样（Leviathan et al. 2023 的"调整后的分布"）。
    【边界】p ≤ q 处处成立时（只可能在 p = q 时发生，此时根本不会被拒绝）分母为 0，返回 p 本身以免除以 0。
    """
    r = torch.clamp(p - q, min=0)
    s = r.sum()
    return r / s if s > 0 else p


def first_token_distribution(p, q):
    """
    推测解码一步里**第一个**输出 token 的分布（闭式解）：
        P(输出 x) = q(x)·min(1, p(x)/q(x))  +  (1 − Σ_y min(p(y), q(y)))·norm(max(0, p − q))(x)
                  = min(p(x), q(x)) + (p(x) − min(p(x), q(x)))       = p(x)
    【返回】(V,) 张量；理论上恰好等于 p（02 号练习的证明）
    """
    acc = torch.minimum(p, q)
    return acc + (1 - acc.sum()) * residual(p, q)


def exact_prefix_distribution(p_fn, q_fn, K, gamma, vocab_size, rule="leviathan"):
    """
    把推测解码的**全部**随机性（草稿的采样、每次的接受 / 拒绝、拒绝后的重采样、全部接受后的额外 token）枚举出来，
    精确地算出输出序列前 K 个 token 的分布。

    【参数】p_fn(prefix) / q_fn(prefix)：给一个 token 元组，返回下一个 token 的概率 (V,)（float64）；
        K：看前几个 token；gamma：每轮草稿的长度
        rule：检验规则。"leviathan" —— 正确的规则；两个**错误**的对照：
            "resample_p"  —— 拒绝后从 p（而不是 norm(max(0, p − q))）重采样；
            "accept_all"  —— 不检验，草稿全部接受（等于直接从草稿模型采样，只是每轮多一个奖励 token）
    【返回】{长度为 K 的元组: 概率}
    【思路】一轮推测可以按位置依次看：第 i 个草稿 token x ~ q；以 min(1, p(x)/q(x)) 接受，接受就看下一个位置；
        被拒绝就从 residual 采一个、这一轮结束；γ 个都接受就再从 p 采一个。后面的草稿在前面被拒绝时会被丢掉，
        所以"先全部草拟、再逐个检验"和"边草拟边检验"的输出分布相同——枚举时用后者，分支少得多。
    【边界】只适合很小的词表和 K（分支数约 V^K）；不处理"遇到结束符就停"（把它当普通 token）。
    """
    result = {}

    def one_round(prefix, i, prob):
        if prob == 0.0:
            return
        if len(prefix) >= K:                       # 已经够 K 个：后面的随机性不影响前 K 个
            result[prefix[:K]] = result.get(prefix[:K], 0.0) + prob
            return
        p = p_fn(prefix)
        if i == gamma:                              # γ 个都接受了：从 p 再采一个，这一轮结束
            for x in range(vocab_size):
                one_round_next(prefix + (x,), prob * p[x].item())
            return
        q = q_fn(prefix)
        acc = q if rule == "accept_all" else torch.minimum(p, q)
        for x in range(vocab_size):                 # 草稿 x 被接受的概率是 q(x)·min(1, p/q) = min(p, q)
            if acc[x] > 0:
                one_round(prefix + (x,), i + 1, prob * acc[x].item())
        rej = 1.0 - acc.sum().item()
        if rej > 1e-15:
            res = p if rule == "resample_p" else residual(p, q)
            for y in range(vocab_size):
                one_round_next(prefix + (y,), prob * rej * res[y].item())

    def one_round_next(prefix, prob):               # 一轮结束、开始新的一轮
        one_round(prefix, 0, prob)

    one_round((), 0, 1.0)
    return result


@torch.no_grad()
def ancestral_sample(target_fn, prompts, gen, max_len, eos_id):
    """
    普通的逐个采样（基准）：每一步把所有序列喂给目标模型，从最后一个位置的分布采一个 token。
    【参数】target_fn(seqs) -> (B, T, V) 概率（每个位置的下一个 token 分布，右侧填充）；prompts：整数列表的列表
    【返回】(序列列表, 目标模型被调用的次数)
    """
    seqs = [list(p) for p in prompts]
    alive = [True] * len(seqs)
    calls = 0
    while any(alive):
        idx = [b for b in range(len(seqs)) if alive[b]]
        probs = target_fn([seqs[b] for b in idx])
        calls += 1
        for row, b in enumerate(idx):
            x = int(torch.multinomial(probs[row, len(seqs[b]) - 1], 1, generator=gen))
            seqs[b].append(x)
            if x == eos_id or len(seqs[b]) >= max_len:
                alive[b] = False
    return seqs, calls


@torch.no_grad()
def speculative_sample(target_fn, draft_fn, prompts, gamma, gen, max_len, eos_id, greedy=False):
    """
    批量的推测采样（Leviathan et al. 2023；Chen et al. 2023）。每一轮：
        1. 草稿模型连续草拟 γ 个 token：x_i ~ q_i（greedy=True 时取 argmax）；
        2. 目标模型对"序列 + γ 个草稿"做**一次**前向，得到 γ + 1 个位置的分布 p_1 … p_{γ+1}；
        3. 依次检验：u ~ U(0, 1)，u < p_i(x_i)/q_i(x_i) 就接受 x_i；第一次拒绝时从 norm(max(0, p_i − q_i)) 采一个代替它、本轮结束；
           γ 个都接受就再从 p_{γ+1} 采一个"奖励 token"。
       greedy=True 时检验变成"x_i 是否等于 p_i 的 argmax"，拒绝时取 argmax p_i——输出和目标模型的贪心解码逐字相同。
    【参数】target_fn(seqs) -> (B, T, V)；draft_fn(seqs) -> (B, V)（每个序列下一个 token 的分布）；gamma；gen：torch.Generator
    【返回】(序列列表, 统计字典)：target_calls（批量调用的次数）、verifies（目标模型检验了多少次"一个序列的一轮草稿"——
        批里每个还活着的序列算一次，相当于不批量时调用目标模型的次数）、draft_calls、drafted（草拟的 token 数）、accepted（被接受的草稿数）、
        generated（生成的 token 总数）、beta_sum / beta_n（Σ_x min(p, q) 在每个被检验位置上的和与个数：
        它的平均值就是论文的接受率 α）
    【边界】草稿遇到结束符就停止这一轮的草拟；接受的 token 是结束符时序列结束，后面的草稿作废。
    """
    seqs = [list(p) for p in prompts]
    alive = [True] * len(seqs)
    st = {"target_calls": 0, "verifies": 0, "draft_calls": 0, "drafted": 0, "accepted": 0, "generated": 0, "beta_sum": 0.0, "beta_n": 0}
    while any(alive):
        idx = [b for b in range(len(seqs)) if alive[b]]
        drafts = {b: [] for b in idx}
        qs = {b: [] for b in idx}
        active = list(idx)
        for _ in range(gamma):                                    # 1. 草拟
            if not active:
                break
            q = draft_fn([seqs[b] + drafts[b] for b in active])
            st["draft_calls"] += 1
            nxt = []
            for row, b in enumerate(active):
                x = int(q[row].argmax()) if greedy else int(torch.multinomial(q[row], 1, generator=gen))
                drafts[b].append(x)
                qs[b].append(q[row])
                if x != eos_id and len(seqs[b]) + len(drafts[b]) < max_len:
                    nxt.append(b)
            active = nxt
        probs = target_fn([seqs[b] + drafts[b] for b in idx])      # 2. 目标模型一次前向
        st["target_calls"] += 1
        for row, b in enumerate(idx):                             # 3. 逐个检验
            st["verifies"] += 1
            base = len(seqs[b]) - 1
            finished = False
            for i, x in enumerate(drafts[b]):
                p, q = probs[row, base + i], qs[b][i]
                st["drafted"] += 1
                st["beta_sum"] += torch.minimum(p, q).sum().item()
                st["beta_n"] += 1
                ok = (x == int(p.argmax())) if greedy else (torch.rand(1, generator=gen).item() < (p[x] / q[x]).item())
                if ok:
                    seqs[b].append(x)
                    st["accepted"] += 1
                    st["generated"] += 1
                    if x == eos_id or len(seqs[b]) >= max_len:
                        finished = True
                        break
                else:
                    y = int(p.argmax()) if greedy else int(torch.multinomial(residual(p, q), 1, generator=gen))
                    seqs[b].append(y)
                    st["generated"] += 1
                    finished = y == eos_id or len(seqs[b]) >= max_len
                    break
            else:                                                 # 全部接受：奖励 token
                p = probs[row, base + len(drafts[b])]
                y = int(p.argmax()) if greedy else int(torch.multinomial(p, 1, generator=gen))
                seqs[b].append(y)
                st["generated"] += 1
                finished = y == eos_id or len(seqs[b]) >= max_len
            if finished:
                alive[b] = False
    return seqs, st


def model_probs_fn(model, pad_id, temperature=1.0):
    """
    把一个 `logits, _ = model(ids)` 接口的语言模型包装成 target_fn：右侧填充成一个批，返回 (B, T, V) 的概率（float64）。
    【边界】因果掩码保证右侧的填充不会影响前面位置的输出，所以填充无害。
    """
    def fn(seqs):
        T = max(len(s) for s in seqs)
        ids = torch.full((len(seqs), T), pad_id, dtype=torch.long)
        for b, s in enumerate(seqs):
            ids[b, :len(s)] = torch.tensor(s)
        logits, _ = model(ids)
        return F.softmax(logits.double() / temperature, dim=-1)
    return fn


def last_position(fn):
    """把 target_fn 变成 draft_fn：只取每个序列最后一个位置的分布。"""
    def draft(seqs):
        probs = fn(seqs)
        return torch.stack([probs[b, len(s) - 1] for b, s in enumerate(seqs)])
    return draft


class NGramDraft:
    """
    把第 1 讲的 NGramLM 包装成草稿模型：draft(seqs) -> (B, V) 的概率（按第 2 讲 GPT 的词表编号）。
    【参数】ngram：已经 fit 好的 NGramLM；vocab：toy_corpus.Vocab
    【边界】n-gram 的预测词表不含 <pad> 和 <s>，这两个编号的概率为 0；同一个上下文的结果会被缓存。
    """

    def __init__(self, ngram, vocab):
        self.ngram, self.vocab = ngram, vocab
        self.cache = {}

    def dist(self, seq):
        words = self.vocab.decode(seq[1:])                      # 去掉开头的 <s>
        ctx = self.ngram.context_of(["<s>"] * (self.ngram.n - 1) + words)
        if ctx not in self.cache:
            v = torch.zeros(len(self.vocab), dtype=torch.float64)
            for w, pr in self.ngram.distribution(ctx).items():
                v[self.vocab.stoi[w]] = pr
            self.cache[ctx] = v / v.sum()
        return self.cache[ctx]

    def __call__(self, seqs):
        return torch.stack([self.dist(s) for s in seqs])


def expected_tokens(alpha, gamma):
    """每次调用目标模型平均产生的 token 数（Leviathan et al. 公式 1）：(1 − α^{γ+1}) / (1 − α)。α = 1 时为 γ + 1。"""
    return gamma + 1 if alpha >= 1 else (1 - alpha ** (gamma + 1)) / (1 - alpha)


def walltime_speedup(alpha, gamma, c):
    """
    总耗时的期望加速比（Leviathan et al. 定理 3.8）：(1 − α^{γ+1}) / ((1 − α)(γc + 1))。
    c = 草稿模型跑一次的时间 / 目标模型跑一次的时间；假设目标模型检验 γ + 1 个 token 和生成 1 个 token 一样快（访存受限）。
    """
    return expected_tokens(alpha, gamma) / (gamma * c + 1)