"""
=============================================================================
CS349E · 第 7 讲共享模块：paged_kv.py —— 分页的 KV cache：块池、块表、共享与写时复制
=============================================================================

【这个模块要解决什么问题】
    第 2 讲的 KVCache 给每个序列预先分配一整块连续的 [max_len, G, dh] 数组。第 6 讲看到：
    预留得多了浪费显存，预留得少了会死锁。vLLM 的 PagedAttention（SOSP 2023）借用了操作系统的虚拟内存：
        - 显存被切成固定大小的**物理块**（block，每块存 B 个 token 的 K 和 V）；
        - 每个序列有一张**块表**（block table）：第 i 个逻辑块 → 哪个物理块；
        - 需要的时候才分配新块（最后一块没填满的部分是唯一的浪费）；
        - 多个序列可以指向同一个物理块（并行采样共享提示、束搜索共享前缀、不同请求共享系统提示），
          用**引用计数**记录有几个序列在用；要往一个共享的块里写时，先复制一份（**写时复制**，copy-on-write）。
    本模块用 numpy 把这些写出来，并且**真的接到第 2 讲的 Transformer 上**：
    第 6 讲的 ragged_forward（选择性批处理）只要求"cache"有 length / write / read 三样东西，
    PagedSequence 实现这三样——于是同一个前向函数就能跑在分页的 KV 上，结果和连续的 KVCache 逐位比较。

【模块里有什么】
    BlockPool      ：物理块池（K、V 数组 + 空闲链表 + 引用计数 + 复制计数）
    PagedSequence  ：一个序列的块表（fork 共享、append_slots 按需分配 + 写时复制、write / read、free）
    paged_step     ：一次迭代：给每个序列分配好这次的槽位，再调用第 6 讲的 ragged_forward
    blockwise_attention：按块计算一个 query 的注意力（PagedAttention 论文式 4 的写法）
    beam_search / parallel_sample：束搜索、并行采样，可以用分页共享的 KV，也可以用"每个候选一份完整拷贝"的 KV

【运行】python3 paged_kv.py   # 自检：分页前向 = 连续前向；束搜索两种实现输出相同
【依赖】numpy；../lecture02（transformer_lab）、../lecture06（decode_lab.ragged_forward；它又会导入 ../lecture01、
       ../lecture05、../prereq5）。都只读。
"""

import math
import os
import sys

import numpy as np

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

from decode_lab import adjust, ragged_forward, sample_rows  # noqa: E402
from transformer_lab import KVCache, forward, softmax  # noqa: E402


class BlockPool:
    """
    物理块池：所有序列的 KV 都存在这里。

    【字段】
        k, v      ：[层数, 块数, 块大小, KV 头数, 每头维度] 的数组（vLLM 也是每层一套块，但所有层共用同一张块表）
        ref       ：每个物理块的引用计数（有几个序列的块表指向它）；0 表示空闲
        free      ：空闲块的编号（当作栈用：最后放回去的最先被拿走）
        n_copies  ：写时复制一共复制了几个块（统计用）
        peak_used ：同时在用的块数的最大值（统计用）
    【边界】块用完时 alloc 抛出 MemoryError——真实引擎会在这之前就抢占某些序列（kv_serving 里模拟）。
    【Python 小知识】list.pop() 取出并删除最后一个元素（O(1)），所以把列表当栈用。
    """

    def __init__(self, cfg, n_blocks, block_size, dtype=np.float64):
        shape = (cfg.n_layers, n_blocks, block_size, cfg.n_kv_heads, cfg.d_head)
        self.k = np.zeros(shape, dtype=dtype)
        self.v = np.zeros(shape, dtype=dtype)
        self.block_size = block_size
        self.n_blocks = n_blocks
        self.ref = np.zeros(n_blocks, dtype=np.int64)
        self.free = list(range(n_blocks - 1, -1, -1))      # 先分配编号小的块，方便看输出
        self.n_copies = 0
        self.peak_used = 0

    @property
    def n_used(self):
        return self.n_blocks - len(self.free)

    def alloc(self):
        """拿一个空闲块，引用计数置 1。【边界】没有空闲块时抛出 MemoryError。"""
        if not self.free:
            raise MemoryError("块池用完了")
        b = self.free.pop()
        self.ref[b] = 1
        self.peak_used = max(self.peak_used, self.n_used)
        return b

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

    def decref(self, b):
        """引用计数减 1；减到 0 时放回空闲链表。"""
        self.ref[b] -= 1
        if self.ref[b] == 0:
            self.free.append(b)
        elif self.ref[b] < 0:
            raise RuntimeError(f"块 {b} 的引用计数变成负数：重复释放")

    def copy_block(self, src):
        """
        写时复制：分配一个新块，把 src 的内容（所有层的 K、V）复制过去，返回新块的编号。

        【系统联系】vLLM 把一次迭代里所有要复制的块攒起来，用一个内核批量复制（论文 §5.1 "fused block copy"），
            避免很多次小的 cudaMemcpy。
        """
        dst = self.alloc()
        self.k[:, dst] = self.k[:, src]
        self.v[:, dst] = self.v[:, src]
        self.n_copies += 1
        return dst


class PagedSequence:
    """
    一个序列的块表：逻辑块 i → 物理块 self.blocks[i]。逻辑位置 p 在第 p // B 块的第 p % B 个槽位。

    【和第 2 讲 KVCache 的接口一样】length、write(layer, k_new, v_new)、read(layer, upto)；
        所以第 6 讲的 ragged_forward 可以直接把它当 cache 用。区别：
        - 写之前要先调用 append_slots(T)（确保这 T 个新位置所在的块都已分配、而且不和别人共享）；
        - read 返回的是把各个块"拼起来"的结果（真实的 PagedAttention 内核不拼，按块读，见 blockwise_attention）。
    【字段】pool：所属的块池；blocks：物理块编号列表；length：已写入的 token 数
    """

    def __init__(self, pool):
        self.pool = pool
        self.blocks = []
        self.length = 0

    # ---- 块表的维护 ----
    def fork(self):
        """
        复制出一个新序列，和自己共享全部物理块（引用计数 +1），不复制任何数据。

        【用途】并行采样：提示算完以后 fork 出 n 个样本；束搜索：一个候选扩展出几个子候选。
        【系统联系】和操作系统的 fork() 一样：子进程和父进程共享内存页，谁先写谁复制。
        """
        child = PagedSequence(self.pool)
        child.blocks = list(self.blocks)
        child.length = self.length
        for b in self.blocks:
            self.pool.incref(b)
        return child

    def append_slots(self, n):
        """
        为接下来要写的 n 个 token 准备好槽位：
            1. 如果最后一块没满、但被别的序列共享（引用计数 > 1），先写时复制，换成自己独占的一块；
            2. 不够的话分配新块。
        【返回】这次新分配（含复制）的块数
        【边界】n = 0 时什么也不做。块池用完时抛出 MemoryError（此时块表可能已经部分扩展）。
        """
        if n <= 0:
            return 0
        B = self.pool.block_size
        added = 0
        if self.length % B != 0:                      # 最后一块没满，接下来要往里写
            last = self.blocks[-1]
            if self.pool.ref[last] > 1:
                new = self.pool.copy_block(last)
                self.pool.decref(last)
                self.blocks[-1] = new
                added += 1
        need = math.ceil((self.length + n) / B) - len(self.blocks)
        for _ in range(need):
            self.blocks.append(self.pool.alloc())
            added += 1
        return added

    def free(self):
        """释放这个序列：每个块的引用计数 −1（共享的块要等最后一个使用者释放才真正回到空闲链表）。"""
        for b in self.blocks:
            self.pool.decref(b)
        self.blocks = []
        self.length = 0

    # ---- 和 KVCache 相同的接口 ----
    def write(self, layer, k_new, v_new):
        """
        把这一层 T 个新 token 的 K、V（形状 [1, T, G, dh]）写到位置 length … length+T−1。

        【边界】调用前必须先 append_slots(T)；写入的块必须是独占的（引用计数 = 1），否则会改坏别人的数据——
            这里检查并报错，帮助发现忘记写时复制的 bug。
        """
        B = self.pool.block_size
        T = k_new.shape[1]
        for j in range(T):
            p = self.length + j
            blk = self.blocks[p // B]
            if self.pool.ref[blk] != 1:
                raise RuntimeError(f"往共享块 {blk} 里写（引用计数 {self.pool.ref[blk]}）：忘了写时复制")
            self.pool.k[layer, blk, p % B] = k_new[0, j]
            self.pool.v[layer, blk, p % B] = v_new[0, j]

    def read(self, layer, upto):
        """
        读这一层位置 0 … upto−1 的 K、V，拼成 [1, upto, G, dh]。

        【Python 小知识】pool.k[layer, self.blocks] 用一个整数列表做"花式索引"，一次取出好几个块，
            得到 [块数, B, G, dh]；reshape 成 [块数·B, G, dh] 再截取前 upto 个。花式索引会复制数据（不是视图）。
        """
        k = self.pool.k[layer, self.blocks].reshape(-1, *self.pool.k.shape[3:])[:upto]
        v = self.pool.v[layer, self.blocks].reshape(-1, *self.pool.v.shape[3:])[:upto]
        return k[None], v[None]

    def n_blocks_unshared_equiv(self):
        """如果不共享，这个序列需要几个块 = ⌈length / B⌉（算"共享省了多少"用）。"""
        return math.ceil(self.length / self.pool.block_size)


def paged_step(params, cfg, seqs, chunks, counter=None):
    """
    一次迭代：先给每个序列分配好这次的槽位（append_slots），再用第 6 讲的 ragged_forward（选择性批处理）算。

    【参数】seqs：PagedSequence 列表；chunks：每个序列这次的新 token（一维数组）
    【返回】每个序列每个新位置的 logits 列表（同 ragged_forward）
    【系统联系】vLLM 的调度器每次迭代先决定"哪些序列这次跑、各自需要几个新块"，块不够就抢占（§4.5），
        然后才把块表发给 GPU 上的 worker。这里的顺序一样：先分配，后计算。
    """
    for s, c in zip(seqs, chunks):
        s.append_slots(len(c))
    return ragged_forward(params, cfg, chunks, seqs, counter)


def blockwise_attention(q, seq, layer):
    """
    PagedAttention 论文式 (4)：按块计算一个 query 对整个序列的注意力，不把 KV 拼成连续数组。

    【参数】q：[H, dh]（一个 decode token 的 query，已加 RoPE）；seq：PagedSequence；layer：层号
    【返回】[H, dh]
    【思路】第一遍：对每个块 j 算分数 A_j = q·K_jᵀ/√dh（只算这个块里真正写了的槽位）；
        把所有块的分数拼起来做 softmax（真实内核用在线 softmax 只扫一遍，第 11 讲）；
        第二遍：按块把 V_j 和对应的权重相乘累加。GQA：H 个 query 头分成 G 组，每组用同一组 K/V。
    【边界】序列长度至少为 1。
    """
    H, dh = q.shape
    G = seq.pool.k.shape[3]
    R = H // G
    B = seq.pool.block_size
    qg = q.reshape(G, R, dh)
    scores = []
    for j, blk in enumerate(seq.blocks):
        n = min(B, seq.length - j * B)                       # 这个块里写了几个
        if n <= 0:
            break
        Kj = seq.pool.k[layer, blk, :n]                      # [n, G, dh]
        scores.append(np.einsum("grd,ngd->grn", qg, Kj) / math.sqrt(dh))
    w = softmax(np.concatenate(scores, axis=-1), axis=-1)    # [G, R, length]
    out = np.zeros((G, R, dh))
    pos = 0
    for j, blk in enumerate(seq.blocks[:len(scores)]):
        n = scores[j].shape[-1]
        out += np.einsum("grn,ngd->grd", w[..., pos:pos + n], seq.pool.v[layer, blk, :n])
        pos += n
    return out.reshape(H, dh)


# -----------------------------------------------------------------------------
# 束搜索与并行采样：两种 KV 管理方式
# -----------------------------------------------------------------------------
def _contig_clone(cache):
    """"以前的系统"：每个子候选复制父候选的整个 KV cache（numpy 数组整份复制）。"""
    c = KVCache(cache.cfg, 1, max_len=cache.max_len, dtype=cache.k.dtype)
    c.k[:] = cache.k
    c.v[:] = cache.v
    c.length = cache.length
    return c


def beam_search(params, cfg, prompt, width, steps, pool=None):
    """
    束搜索：每一步把每个候选扩展成"它 + 词表里每个 token"，按累计 log 概率保留最好的 width 个。

    【参数】prompt：一维 token 数组；width：束宽 k；steps：生成几步
        pool：给一个 BlockPool 时用分页共享的 KV（fork + 写时复制）；None 时每个候选一份完整的连续 KVCache，
              候选被选中多次时整份复制（vLLM 论文说的"以前的系统要频繁复制 KV"）。
    【返回】(最好的 width 个序列（按分数从高到低，每个是 token 列表）, 统计字典)
        统计：copied_tokens（复制了多少个 token 的 KV）、blocks_used / blocks_unshared（每一步结束时，
        分页方式实际用的块数 vs 不共享时需要的块数，列表）
    【思路】
        1. prefill：算一次提示，得到第一步的 logits；初始候选只有 1 个。
        2. 每一步：所有候选的 log 概率 + 各自的累计分数 → 从 (候选数 × V) 个组合里取前 width 个
           （分数相同时按下标顺序，保证两种实现选得一样）。
        3. 新候选 = 父候选 fork（分页：共享块）或复制（连续：整份拷贝）+ 追加选中的 token；
           没被选中的父候选释放。
        4. 下一步前向：每个新候选喂进它最后的那个 token（用第 6 讲的 ragged_forward，一次迭代算所有候选）。
    【边界】steps ≥ 1；width ≤ 词表大小。
    """
    stats = {"copied_tokens": 0, "blocks_used": [], "blocks_unshared": []}
    prompt = [int(t) for t in prompt]
    if pool is not None:
        root = PagedSequence(pool)
        logits = paged_step(params, cfg, [root], [np.array(prompt)])[0][-1:]
    else:
        root = KVCache(cfg, 1)
        logits = forward(params, cfg, np.array([prompt]), cache=root)[0][-1:]
    beams = [(0.0, [], root)]                                 # (累计 log 概率, 生成的 token, cache)
    for step in range(steps):
        lp = np.log(softmax(logits, axis=-1))                 # [候选数, V]
        cand = np.array([b[0] for b in beams])[:, None] + lp
        flat = np.argsort(-cand.ravel(), kind="stable")[:width]
        V = lp.shape[1]
        chosen = [(int(i) // V, int(i) % V) for i in flat]
        new_beams = []
        for parent, tok in chosen:
            _, toks, cache = beams[parent]
            if pool is not None:
                child = cache.fork()
            else:
                child = _contig_clone(cache)
                stats["copied_tokens"] += cache.length
            new_beams.append((float(cand[parent, tok]), toks + [tok], child))
        for _, _, cache in beams:                              # 父候选都不要了（被选中的已经 fork / 复制过）
            if pool is not None:
                cache.free()
        beams = new_beams
        if step == steps - 1:
            break
        chunks = [np.array([b[1][-1]]) for b in beams]
        caches = [b[2] for b in beams]
        if pool is not None:
            outs = paged_step(params, cfg, caches, chunks)
            stats["blocks_used"].append(pool.n_used)
            stats["blocks_unshared"].append(sum(c.n_blocks_unshared_equiv() for c in caches))
        else:
            outs = ragged_forward(params, cfg, chunks, caches)
        logits = np.concatenate([o[-1:] for o in outs], axis=0)
    result = [(s, t) for s, t, _ in beams]
    if pool is not None:
        for _, _, c in beams:
            c.free()
    return result, stats


def parallel_sample(params, cfg, prompt, n, steps, temperature, seed, pool=None):
    """
    并行采样：同一个提示生成 n 个独立的样本（代码补全里"给我 n 个候选"）。

    【参数】pool：给一个 BlockPool 时，提示只算一次、n 个样本 fork 共享提示的块；None 时每个样本各自一份连续 KV。
    【返回】(n 个样本的 token 列表, 统计字典：blocks_used / blocks_unshared 每一步的列表)
    【边界】两种方式用同一个随机数序列（每步一次为 n 个样本抽样），所以输出应该相同。
    """
    rng = np.random.default_rng(seed)
    stats = {"blocks_used": [], "blocks_unshared": []}
    prompt = np.array([int(t) for t in prompt])
    if pool is not None:
        root = PagedSequence(pool)
        last = paged_step(params, cfg, [root], [prompt])[0][-1:]
        seqs = [root.fork() for _ in range(n)]
        root.free()
    else:
        seqs = []
        for _ in range(n):
            c = KVCache(cfg, 1)
            last = forward(params, cfg, prompt[None], cache=c)[0][-1:]
            seqs.append(c)
    logits = np.repeat(last, n, axis=0)
    outs = [[] for _ in range(n)]
    for step in range(steps):
        toks = sample_rows(adjust(logits, temperature), rng)
        for i in range(n):
            outs[i].append(int(toks[i]))
        if step == steps - 1:
            break
        chunks = [np.array([t]) for t in toks]
        if pool is not None:
            res = paged_step(params, cfg, seqs, chunks)
            stats["blocks_used"].append(pool.n_used)
            stats["blocks_unshared"].append(sum(s.n_blocks_unshared_equiv() for s in seqs))
        else:
            res = ragged_forward(params, cfg, chunks, seqs)
        logits = np.concatenate([r[-1:] for r in res], axis=0)
    if pool is not None:
        for s in seqs:
            s.free()
    return outs, stats


def _selftest():
    from transformer_lab import Config, init_params
    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=64)
    params = init_params(cfg, seed=0)
    # 1. 分页前向 = 连续前向（含跨块、不同长度、prefill + decode 混合）
    pool = BlockPool(cfg, n_blocks=32, block_size=4)
    seqs = [PagedSequence(pool) for _ in range(3)]
    refs = [KVCache(cfg, 1) for _ in range(3)]
    rng = np.random.default_rng(0)
    worst = 0.0
    for step in range(5):
        chunks = [rng.integers(0, cfg.vocab, size=(7 if step == 0 else 1) + i) for i in range(3)]
        outs = paged_step(params, cfg, seqs, chunks)
        for c, o, r in zip(chunks, outs, refs):
            worst = max(worst, float(np.max(np.abs(o - forward(params, cfg, c[None], cache=r)[0]))))
    assert worst < 1e-12, worst
    # 2. 写时复制：fork 之后往共享块里写，不影响父序列
    a = seqs[0]
    b = a.fork()
    paged_step(params, cfg, [b], [np.array([5])])
    assert pool.n_copies == (1 if a.length % 4 else 0)
    # 3. 按块注意力 = 拼起来的注意力
    q = rng.standard_normal((cfg.n_heads, cfg.d_head))
    K, V = a.read(0, a.length)
    from transformer_lab import attention_core
    ref = attention_core(q[None, None], K, V, a.length - 1)[0, 0].reshape(cfg.n_heads, cfg.d_head)
    assert np.max(np.abs(blockwise_attention(q, a, 0) - ref)) < 1e-12
    # 4. 束搜索：分页共享 = 整份复制
    r1, _ = beam_search(params, cfg, [1, 2, 3], 4, 6, pool=None)
    pool2 = BlockPool(cfg, n_blocks=64, block_size=4)
    r2, _ = beam_search(params, cfg, [1, 2, 3], 4, 6, pool=pool2)
    assert [t for _, t in r1] == [t for _, t in r2]
    assert pool2.n_used == 0, "束搜索结束后所有块都应该被释放"
    print(f"paged_kv 自检通过：分页前向 = 连续前向（最大差 {worst:.1e}）、写时复制、按块注意力、束搜索两种实现一致")


if __name__ == "__main__":
    _selftest()