"""第 37 章：为已有 TinyTransformer 添加显式 KV 缓存；不训练新模型。"""

from dataclasses import dataclass
from math import sqrt
import time
import torch
from tiny_transformer import TinyTransformer


@dataclass
class Cache:
    model: TinyTransformer
    layers: list
    length: int

    @property
    def nbytes(self):
        return sum(t.numel()*t.element_size() for pair in self.layers for t in pair)


@torch.inference_mode()
def cached_forward(model, ids, cache=None, use_sdpa=False):
    if model.training:
        raise ValueError("缓存接口只用于 eval 模式的推理")
    if ids.ndim != 2 or ids.shape[0] < 1 or ids.shape[1] < 1:
        raise ValueError("输入须为非空 [batch, length]")
    batch, length = ids.shape
    start = 0 if cache is None else cache.length
    config = model.config
    if cache is not None:
        if cache.model is not model or len(cache.layers) != config["layers"]:
            raise ValueError("缓存不属于当前模型")
        expected = (batch, config["heads"], start, config["width"]//config["heads"])
        for key, value in cache.layers:
            if key.shape != expected or value.shape != expected:
                raise ValueError("缓存形状或长度不一致")
            if key.device != ids.device or value.device != ids.device:
                raise ValueError("缓存与输入设备不同")
            if key.dtype != model.tokens.weight.dtype or value.dtype != key.dtype:
                raise ValueError("缓存精度不同")
    if start+length > config["max_context"]:
        raise ValueError("超过位置表长度；本例不实现滑动窗口")
    position = torch.arange(start, start+length, device=ids.device)
    x = model.tokens(ids)+model.positions(position)
    new_layers = []
    for index, block in enumerate(model.blocks):
        attention = block.attention
        normalized = block.norm1(x)
        projected = attention.qkv(normalized).chunk(3, dim=-1)
        query, key, value = [
            part.reshape(batch, length, attention.heads, -1).transpose(1, 2)
            for part in projected]
        if cache is not None:
            old_key, old_value = cache.layers[index]
            key = torch.cat([old_key, key], dim=-2)
            value = torch.cat([old_value, value], dim=-2)
        key_position = torch.arange(start+length, device=ids.device)
        allowed = key_position[None, :] <= position[:, None]
        if use_sdpa:
            # 非方形注意力显式传绝对位置掩码，不直接套 is_causal=True。
            mixed = torch.nn.functional.scaled_dot_product_attention(
                query, key, value, attn_mask=allowed, dropout_p=0.0, is_causal=False)
        else:
            scores = query@key.transpose(-2, -1)/sqrt(query.shape[-1])
            weights = scores.masked_fill(~allowed, float("-inf")).softmax(-1)
            mixed = weights@value
        mixed = mixed.transpose(1, 2).contiguous().reshape(batch, length, config["width"])
        x = x+attention.output(mixed)
        x = x+block.feedforward(block.norm2(x))
        new_layers.append((key, value))
    logits = model.output(model.final_norm(x))
    return logits, Cache(model, new_layers, start+length)


@torch.inference_mode()
def generate(model, prompt, new_tokens, cached=True):
    if new_tokens < 1 or prompt.ndim != 2:
        raise ValueError("至少生成一个词元，提示为二维 ID")
    if prompt.shape[1]+new_tokens-1 > model.config["max_context"]:
        raise ValueError("生成所需输入超过位置表长度")
    model.eval()
    sequence = prompt.clone()
    cache, positions, pairs = None, 0, 0
    stamps = []
    started = time.perf_counter()
    for index in range(new_tokens):
        inputs = sequence if not cached or cache is None else sequence[:, -1:]
        if cached:
            logits, cache = cached_forward(model, inputs, cache)
        else:
            logits = model(inputs)
        # 每序列、每层的投影位置数；分数元素按每层、每个注意力头计数。
        positions += inputs.shape[1]
        pairs += inputs.shape[1]*sequence.shape[1]
        token = logits[:, -1].argmax(-1, keepdim=True)
        sequence = torch.cat([sequence, token], dim=1)
        stamps.append(time.perf_counter()-started)
    return dict(ids=sequence, projected_positions=positions, score_elements=pairs,
                cache_bytes=cache.nbytes if cache else 0,
                first_seconds=stamps[0], total_seconds=stamps[-1], token_seconds=stamps)


def corrected_distribution(target, draft):
    """枚举一次草稿采样、接受或残差重采样，验证分布；不是多词元加速内核。"""
    p, q = torch.tensor(target, dtype=torch.float64), torch.tensor(draft, dtype=torch.float64)
    if (p.shape != q.shape or p.ndim != 1 or torch.any(p < 0) or torch.any(q < 0)
            or not torch.isfinite(p).all() or not torch.isfinite(q).all()
            or not torch.isclose(p.sum(), torch.tensor(1., dtype=p.dtype))
            or not torch.isclose(q.sum(), torch.tensor(1., dtype=q.dtype))):
        raise ValueError("需要两个等长的归一化非负概率向量")
    accepted_mass = torch.minimum(p, q)
    rejection = 1-accepted_mass.sum()
    residual = (p-q).clamp_min(0)
    correction = residual/residual.sum() if residual.sum() > 0 else torch.zeros_like(p)
    return accepted_mass+rejection*correction, float(accepted_mass.sum())


def main():
    torch.set_num_threads(1)
    torch.manual_seed(2026)
    model = TinyTransformer(35, max_context=32, width=32, heads=4, layers=2).double().eval()
    ids = torch.randint(0, 35, (2, 11))
    with torch.inference_mode():
        whole = model(ids)
        _, cache = cached_forward(model, ids[:, :8])
        chunk, _ = cached_forward(model, ids[:, 8:], cache)
        optimized, _ = cached_forward(model, ids[:, 8:], cache, use_sdpa=True)
    print("分块缓存与完整重算分数一致：", torch.allclose(chunk, whole[:, 8:], atol=1e-12))
    print("显式掩码的 SDPA 与手写注意力一致：", torch.allclose(chunk, optimized, atol=1e-12))
    prompt = ids[:1, :8]
    without, with_cache = [generate(model, prompt, 12, cached=flag) for flag in [False, True]]
    print("连续生成 12 个词元的 ID 一致：", torch.equal(without["ids"], with_cache["ids"]))
    print("每层投影位置数：", without["projected_positions"], "→", with_cache["projected_positions"])
    print("每层、每个头的密集分数元素数：", without["score_elements"], "→", with_cache["score_elements"])
    print("最终保留的 KV 张量：", with_cache["cache_bytes"], "字节")
    restored, accepted = corrected_distribution([.6, .3, .1], [.2, .5, .3])
    print("推测采样接受概率：", f"{accepted:.3f}")
    print("接受加残差重采样后的分布：", [round(x, 3) for x in restored.tolist()])
    print("本例随机权重用于核对计算；默认不报告随机波动的运行时间。")


if __name__ == "__main__":
    main()

"""第 37 章：显式 --run 测量 CPU KV 缓存；默认读取随书保存的原始测量。"""

import argparse
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import platform
import statistics
import torch
from cached_decoder import TinyTransformer, generate

DATA = Path(__file__).with_name("data")/"inference-benchmark.json"
CONFIG = dict(vocab_size=128, max_context=512, width=64, heads=4, layers=2)
LENGTHS = [16, 128]
NEW_TOKENS, WARMUPS, REPEATS = 32, 3, 7


def model_and_inputs():
    torch.set_num_threads(1)
    torch.manual_seed(2036)
    model = TinyTransformer(**CONFIG).eval()
    inputs = [torch.randint(0, CONFIG["vocab_size"], (1, length)) for length in LENGTHS]
    return model, inputs


def model_hash(model):
    digest = hashlib.sha256()
    for key, value in model.state_dict().items():
        digest.update(key.encode())
        digest.update(value.detach().numpy().tobytes())
    return digest.hexdigest()


def summarize(rows):
    result = []
    for length in LENGTHS:
        for cached in [False, True]:
            selected = [r for r in rows if r["length"] == length and r["cached"] == cached]
            totals = [1000*r["total_seconds"] for r in selected]
            first = [1000*r["first_seconds"] for r in selected]
            speed = [(NEW_TOKENS-1)/(r["total_seconds"]-r["first_seconds"]) for r in selected]
            result.append(dict(length=length, cached=cached,
                               median_total_ms=statistics.median(totals),
                               median_first_ms=statistics.median(first),
                               median_decode_tps=statistics.median(speed),
                               min_total_ms=min(totals), max_total_ms=max(totals),
                               cache_bytes=selected[0]["cache_bytes"]))
    return result


def benchmark(path):
    model, inputs = model_and_inputs()
    rows = []
    for prompt in inputs:
        for _ in range(WARMUPS):
            for cached in [False, True]:
                generate(model, prompt, NEW_TOKENS, cached)
        expected = None
        for trial in range(REPEATS):
            # 交替顺序，减少总让一个方案先运行的偏差；仍不能消除系统噪声。
            for cached in ([False, True] if trial % 2 == 0 else [True, False]):
                row = generate(model, prompt, NEW_TOKENS, cached)
                ids = row.pop("ids").tolist()
                if expected is None:
                    expected = ids
                assert ids == expected, "缓存前后生成不同，不能直接比较速度"
                rows.append(dict(length=prompt.shape[1], cached=cached, trial=trial, ids=ids, **row))
    report = dict(complete=True, measured_at=datetime.now(timezone.utc).isoformat(),
                  environment=dict(python=platform.python_version(), torch=torch.__version__,
                                   platform=platform.platform(), processor=platform.processor(),
                                   device="cpu", dtype="float32", threads=1),
                  config=CONFIG, seed=2036, new_tokens=NEW_TOKENS, warmups=WARMUPS,
                  repeats=REPEATS, prompts=[p.tolist() for p in inputs],
                  model_sha256=model_hash(model), rows=rows)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
    return report


def checked_report(path=DATA):
    report = json.loads(path.read_text())
    assert report["complete"] and report["config"] == CONFIG
    assert report["new_tokens"] == NEW_TOKENS and report["repeats"] == REPEATS
    assert report["warmups"] == WARMUPS and report["seed"] == 2036
    model, prompts = model_and_inputs()
    assert report["model_sha256"] == model_hash(model)
    assert report["prompts"] == [p.tolist() for p in prompts]
    assert len(report["rows"]) == len(LENGTHS)*2*REPEATS
    keys = {(r["length"], r["cached"], r["trial"]) for r in report["rows"]}
    assert keys == {(length, cached, i) for length in LENGTHS
                    for cached in [False, True] for i in range(REPEATS)}
    for row in report["rows"]:
        stamps = row["token_seconds"]
        assert len(stamps) == NEW_TOKENS and 0 < stamps[0] < stamps[-1]
        assert stamps == sorted(stamps)
        assert row["first_seconds"] == stamps[0] and row["total_seconds"] == stamps[-1]
    # 用当前代码复核每种输入/路径的实际输出与计数，不把它当作重新测量原时间。
    for prompt in prompts:
        for cached in [False, True]:
            current = generate(model, prompt, NEW_TOKENS, cached)
            for row in report["rows"]:
                if row["length"] == prompt.shape[1] and row["cached"] == cached:
                    assert row["ids"] == current["ids"].tolist()
                    for key in ["projected_positions", "score_elements", "cache_bytes"]:
                        assert row[key] == current[key]
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()
    if args.run and args.out is None:
        parser.error("请用 --out 保存自己的测量，不覆盖随书记录")
    report = benchmark(args.out) if args.run else checked_report()
    print("以下为本次测量：" if args.run else "以下为随书保存的 CPU 单线程测量：")
    for row in summarize(report["rows"]):
        print(f"提示 {row['length']:3d}；缓存={row['cached']}；"
              f"首词元中位数={row['median_first_ms']:.2f} ms；"
              f"总时长中位数={row['median_total_ms']:.2f} ms；"
              f"后续速度中位数={row['median_decode_tps']:.1f} token/s；"
              f"KV={row['cache_bytes']} 字节")
    print("每种条件预热 3 次、测量 7 次；同条件的生成 ID 完全一致。")
    print("随机权重只用于性能路径比较，不代表语言能力；本地结果可能不同。")


if __name__ == "__main__":
    main()

1. Prefill 的输入已经给定，只需按因果关系限制访问；自由生成时，下一个输入取决于刚选出的词元。并行处理已知位置和提前确定未知输出是两回事。
2. 编号为 6、7、8。第一行可看位置 0～6，共 7 个键；不能把新块的位置从零重启。
3. $2\times2\times1\times159\times64\times4=162816$ 字节。同样长度和布局下，批次翻倍，KV 张量字节数也翻倍；这不包括其他内存。
4. 第一个输出由提示的最后一个位置预测；最后一个输出还没有再送回模型。继续生成时，需要处理刚才的第 12 个输出，让缓存再增长一格。
5. 全部可见时为 4；只留下第一个键时为 1。形状合法仍可能语义错误。
6. $(5-1)/(1.0-0.2)=5$ token/s。$5/1$ 恰好也为 5，但它把首词元之前的等待算进去了，口径不同；把首词元时间改为 0.4 秒，两者就不相等。
7. $0.3/0.5=0.6$。已经直接接受的部分并不与 $p$ 成比例，拒绝分支需要补足缺少的概率质量。若拒绝后直接抽 $p$，本例最终会得到 $[0.44,0.42,0.14]$，而非目标分布。
8. 保存完整环境、配置和逐次结果，不只摘一个最快值。修改模型配置以后，旧记录的指纹与配置检查应拒绝把它当成同一实验；可以先使用 --run 生成自己的记录，再读取分析。
9. 用户前缀不同，会改变注意力历史，可能造成答案串扰和数据泄露。参数更新后旧 K/V 也不再对应新模型，应按模型版本与具体前缀管理缓存。
10. 不能，要检查实际分派的后端和设备。PagedAttention 改善存储组织；普通密集注意力的新查询仍然要访问允许的历史键和值。