"""
=============================================================================
第 11 讲 · 02：KV cache 的复用与压缩 —— CacheBlend 和 Cartridges 的账
=============================================================================

对应讲义：lectures/lecture11_memory.md 第 3、4 节（Yao et al. 2024 CacheBlend；Eyuboglu et al. 2025 Cartridges）

-----------------------------------------------------------------------------
【第一部分：CacheBlend —— 预先算好的 KV cache 怎样拼起来用】
-----------------------------------------------------------------------------
RAG 的输入 = 6 篇检索到的文档（每篇 512 个 token）+ 问题。每篇文档的 KV cache 可以**预先算好、存起来**，
下次检索到同一篇文档时直接加载，省掉 prefill 的计算。
问题在于：第 2 篇以后的文档，预先计算时"前面没有别的文档"；真正使用时前面有别的文档。
它们之间的注意力（cross-attention）在预先算好的 KV 里是缺失的。直接拼接复用（full reuse），
回答需要综合多篇文档的问题时，质量会下降。

卡通的质量模型：
    · 每个 token 有一个"跨文档依赖的强度" d：大多数接近 0，少数很大（注意力是稀疏的）。第 1 篇文档的 d = 0，
      因为它本来就在开头，预先算好的 KV 和真实的完全一样。
    · 重算一个 token，就消除它的偏差。剩下的偏差占比 R = 没重算的 token 的 d 之和 / 所有 d 之和。
    · 回答质量（类似 F1）= 0.80 × (1 − R) + 0.62 × R：全部重算 0.80，完全复用 0.62。
      （CacheBlend 论文报告，相对完全复用，F1 高 0.1~0.2；这里取 0.18 的差距。）
挑哪些 token 重算：
    · 随机：随机挑；
    · HKVD（high-KV-deviation）：按"估计的偏差"挑最大的那些。估计来自前几层的实际偏差（有噪声），
      依据是"在一层偏差大的 token，在下一层也往往偏差大"；
    · oracle：按真实的 d 挑（上帝视角，是 HKVD 的上限）。
时间模型（每层流水线化：加载第 l+1 层的 KV 和重算第 l 层同时进行）：
    · 重算：全部 3072 个 token 重算一层约 20 ms（于是重算 15% 约 3 ms/层，和论文里 Llama-7B、4K 上下文的数字一致）；
    · 加载：从 NVMe SSD 加载一层约 16 ms（论文的数字），从 CPU 内存约 4 ms、从 GPU 显存约 0.5 ms（假设）。
    · 首 token 延迟（TTFT）≈ 32 层 × max(这一层的加载时间, 这一层的重算时间)。

-----------------------------------------------------------------------------
【第二部分：Cartridges 的账 —— 把整个语料压成一小块训练出来的 KV cache】
-----------------------------------------------------------------------------
Llama 3 8B 的 KV cache：32 层 × 8 个 KV 头 × 128 维 × 2（K 和 V）× 2 字节（bf16）= 128 KiB / token。
一个 10 万 token 的语料（例如一个代码库）放进上下文，每个用户就要占约 12 GiB 的 KV cache。
Cartridge = 为这个语料离线训练一个只有 p 个"虚拟 token"的 KV cache，推理时加载它代替整个语料。
这一部分只是算术（没有随机性），帮你建立"显存才是长上下文的瓶颈"的直觉。

运行时间：不到 1 秒（只用标准库）。
"""

import math
import random

N_DOCS, DOC_LEN, N_LAYERS = 6, 512, 32
Q_FULL, Q_REUSE = 0.80, 0.62                # 全部重算 / 完全复用时的回答质量
TAIL = 10                                   # d = u^TAIL：TAIL 越大，偏差越集中在少数 token 上
HKVD_NOISE = 0.7                            # HKVD 估计的噪声（对数尺度上的标准差）
RECOMPUTE_MS_PER_LAYER_ALL = 20.0           # 全部 3072 个 token 重算一层的时间
LOAD_MS_PER_LAYER = {"GPU 显存": 0.5, "CPU 内存": 4.0, "NVMe SSD": 16.0}   # 加载一层全部 KV 的时间


def display_width(text):
    """
    文字在终端里占的宽度：中文字符占 2 格，英文和数字占 1 格。用来对齐表格。

    【Python 小知识】ord(ch) 返回字符的编号；大于 127 的不是 ASCII（这里粗略地当作中文、占 2 格）。
    """
    return sum(2 if ord(ch) > 127 else 1 for ch in text)


def make_dependence(rng):
    """
    生成一次 RAG 输入里每个 token 的跨文档依赖强度 d。

    【思路】注意力是稀疏的：一个 token 只和少数几个前面的 token 有强的注意力。所以只有少数 token 的 KV
        会因为"前面多了别的文档"而明显改变。d = u^10（u 在 0~1 上均匀）是一个简单的"重尾"分布：
        最大的 10% 的 token 贡献了约七成的 d 之和。
    【返回】长度 6 × 512 = 3072 的列表；前 512 个（第 1 篇文档）全是 0。
    【Python 小知识】[0.0] * 512 得到 512 个 0.0 组成的列表；两个列表用 + 拼接。
    """
    return [0.0] * DOC_LEN + [rng.random() ** TAIL for _ in range((N_DOCS - 1) * DOC_LEN)]


def quality(d, recomputed):
    """
    重算了 recomputed 这些 token 之后的回答质量。

    【参数】d：make_dependence 的结果；recomputed：被重算的 token 下标的**集合**（set）
    【返回】介于 Q_REUSE 和 Q_FULL 之间的数：剩下的偏差占比 R 越小，越接近 Q_FULL
    【边界】如果所有 d 都是 0（实际上不会发生），total = 0 会除以零；这里不做处理。
    【Python 小知识】`i not in recomputed` 对集合是 O(1) 的；如果 recomputed 是列表，就要一个个比较，慢很多。
    """
    total = math.fsum(d)
    remaining = math.fsum(x for i, x in enumerate(d) if i not in recomputed)
    r = remaining / total
    return Q_FULL * (1 - r) + Q_REUSE * r


def select_tokens(d, fraction, method, rng):
    """
    挑出要重算的 token（只在第 2 篇及以后的文档里挑：第 1 篇不需要重算）。

    【参数】fraction：重算的比例（相对**全部** token，所以 15% = 460 个 token）；method："random" / "hkvd" / "oracle"
    【返回】被选中的 token 下标的集合
    【思路】HKVD 的估计 = 真实的 d × exp(高斯噪声)：前几层的偏差和最终的偏差高度相关，但不完全一样。
        真实的 CacheBlend 是逐层"逐步筛选"的：第 1 层重算全部 token，挑出偏差大的；之后每层只在上一层挑出的
        token 里再挑。这里用一次"有噪声的估计"概括这个过程。
    【边界】fraction 超过 5/6 时，要挑的 token 比候选还多，rng.sample 会报错（ValueError）。
    """
    candidates = list(range(DOC_LEN, len(d)))
    k = int(fraction * len(d))
    if method == "random":
        return set(rng.sample(candidates, k))
    if method == "oracle":
        score = {i: d[i] for i in candidates}
    else:
        score = {i: d[i] * math.exp(rng.gauss(0, HKVD_NOISE)) for i in candidates}
    return set(sorted(candidates, key=lambda i: -score[i])[:k])


def ttft_ms(recompute_fraction, load_fraction, load_ms):
    """
    首 token 延迟（毫秒）。

    【参数】recompute_fraction：要重算的 token 比例；load_fraction：要从存储里加载 KV 的 token 比例；
        load_ms：加载一层**全部** token 的 KV 需要的时间（取决于存在哪个设备上）
    【思路】每一层"加载"和"重算"同时进行（流水线），所以这一层花的时间是两者里较大的那个。
        全部重算：load_fraction = 0；完全复用：recompute_fraction = 0；
        CacheBlend：加载全部 KV，再重算其中的一小部分；前缀缓存：只加载第 1 篇文档，其余重算。
    """
    per_layer = max(load_ms * load_fraction, RECOMPUTE_MS_PER_LAYER_ALL * recompute_fraction)
    return N_LAYERS * per_layer


def part_one():
    print("=" * 78)
    print("第一部分：CacheBlend —— 重算多少、重算哪些 token？（6 篇文档 × 512 token，100 次随机实例取平均）")
    print("=" * 78)
    rng = random.Random(250)
    instances = [make_dependence(rng) for _ in range(100)]
    print(f"全部重算：质量 {Q_FULL:.3f}；完全复用（不重算）：质量 {Q_REUSE:.3f}")
    print(f"{'重算比例':10}" + "".join(f"{m:>12}" for m in ("随机", "HKVD", "oracle")))
    for frac in (0.05, 0.10, 0.15, 0.30):
        row = f"{frac:<12.0%}"
        for method in ("random", "hkvd", "oracle"):
            mean = sum(quality(d, select_tokens(d, frac, method, rng)) for d in instances) / len(instances)
            row += f"{mean:12.3f}"
        print(row)

    print("\n首 token 延迟（TTFT，毫秒），KV cache 存在不同的设备上：")
    print(" " * 26 + "".join(dev.rjust(14 - display_width(dev) + len(dev)) for dev in LOAD_MS_PER_LAYER))
    rows = (("全部重算（不复用）", 1.0, 0.0),
            ("前缀缓存（只复用第 1 篇）", 5 / 6, 1 / 6),
            ("完全复用（不重算）", 0.0, 1.0),
            ("CacheBlend（重算 15%）", 0.15, 1.0))
    for name, recompute, load in rows:
        cells = "".join(f"{ttft_ms(recompute, load, ms):14.0f}" for ms in LOAD_MS_PER_LAYER.values())
        print(name + " " * max(1, 26 - display_width(name)) + cells)

    print("\n观察：")
    print("  · 完全复用最快，但丢掉了跨文档的注意力，质量明显下降；全部重算质量最好，但最慢。")
    print("  · 只重算一小部分 token 就能恢复大部分质量 —— 前提是挑对：按偏差挑（HKVD）的 15% 恢复了八成多的质量差距，")
    print("    随机挑的 15% 只恢复不到两成。原因是注意力是稀疏的：只有少数 token 的 KV 会因为\"前面多了别的文档\"而明显改变。")
    print("    HKVD 的估计有噪声，比 oracle 略差，但差得不多。CacheBlend 论文：重算 10%~20% 的 token 就够了。")
    print("  · 前缀缓存（prefix caching）在 RAG 里帮助有限：只有第 1 篇文档在开头，后面的都得重算，只快了约 17%。")
    print("  · 重算 15% 的时间（每层约 3 ms）可以藏在加载 KV 的时间后面：KV 存在 CPU 内存或 SSD 上时，")
    print("    CacheBlend 和完全复用一样快 —— 延迟只由加载决定，重算是\"免费\"的。")
    print("    KV 存在 GPU 显存里时，重算反而成了瓶颈，但仍比全部重算快约 6.7 倍。")
    print("    这让 KV cache 可以放在更便宜、容量更大的设备上（LMCache 等系统的设计思想）。")
    print("    CacheBlend 论文报告：相对全部重算，TTFT 降低 2.2~3.3 倍，吞吐提高 2.8~5 倍，质量几乎不变。\n")


def part_two():
    print("=" * 78)
    print("第二部分：Cartridges 的账（Llama 3 8B，bf16，一张 80 GB 的 H100）")
    print("=" * 78)
    bytes_per_token = 32 * 8 * 128 * 2 * 2       # 层 × KV 头 × 每个头的维度 × (K, V) × 每个数 2 字节
    mib, gib = 1024 ** 2, 1024 ** 3
    corpus_tokens = 100_000
    free_bytes = (80 - 16) * 10 ** 9            # 80 GB 显存，减去约 16 GB 的模型权重
    print(f"每个 token 的 KV cache：{bytes_per_token // 1024} KiB；留给 KV cache 的显存：约 {free_bytes / gib:.1f} GiB")
    print(" " * 38 + "每个用户的 KV cache" + "    能同时服务的用户")
    for name, tokens in (("ICL：把 10 万 token 的语料放进上下文", corpus_tokens),
                         ("Cartridge：p = 2048 个虚拟 token", 2048),
                         ("Cartridge：p = 512", 512)):
        size = tokens * bytes_per_token
        users = free_bytes // size
        print(name + " " * max(1, 38 - display_width(name)) + f"{size / mib:>13,.0f} MiB" + f"{users:>18,d}")

    print("\n训练一个 cartridge 的成本什么时候能被分摊回来？（只算 prefill 的计算，假设 prefill 速度 2 万 token/秒/GPU）")
    train_gpu_seconds = 0.5 * 3600 * 8          # 论文：8 × H100 上约 30 分钟
    prefill_seconds = corpus_tokens / 20_000
    print(f"  训练：约 {train_gpu_seconds / 3600:.0f} GPU 小时；每次查询省下的 prefill：约 {prefill_seconds:.0f} GPU 秒；"
          f"打平需要约 {train_gpu_seconds / prefill_seconds:,.0f} 次查询")
    print("\n观察：")
    print("  · 长上下文最大的成本往往不是计算，而是**显存**：每个用户约 12 GiB 的 KV cache，一张卡只能同时服务 4 个用户；")
    print("    解码阶段主要受显存带宽限制，同时服务的用户越少（batch 越小），吞吐越低。")
    print("  · 2048 个虚拟 token 的 cartridge 只占 256 MiB，同一张卡能同时服务的用户从 4 个变成 238 个。Cartridges 论文报告：")
    print("    和 ICL 质量相当的 cartridge，显存少 38.6 倍、吞吐高 26.4 倍。")
    print("  · 代价是离线训练（论文：8 × H100 上约 30 分钟）。只看 prefill 的计算，要几千次查询才能打平；")
    print("    但如果算上显存和吞吐的收益，以及\"很多用户反复查询同一个语料\"（代码库、病历、财报），")
    print("    而训练可以用空闲的算力离线完成，这笔账通常是划算的。语料只被查几次的话，直接放进上下文更好。\n")


if __name__ == "__main__":
    part_one()
    part_two()