{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实验 04：实现并测量 KV Cache\n",
    "\n",
    "**运行条件：** CPU，约 2 分钟，内存低于 1 GB。\n",
    "\n",
    "目标不是调用现成的 `generate(use_cache=True)`，而是直接实现缓存协议，证明缓存前后结果一致，并测量计算与内存的变化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "torch.manual_seed(19)\n",
    "torch.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CausalAttention(nn.Module):\n",
    "    def __init__(self, width=128, heads=4):\n",
    "        super().__init__()\n",
    "        assert width % heads == 0\n",
    "        self.heads, self.head_dim = heads, width // heads\n",
    "        self.q = nn.Linear(width, width, bias=False)\n",
    "        self.k = nn.Linear(width, width, bias=False)\n",
    "        self.v = nn.Linear(width, width, bias=False)\n",
    "        self.out = nn.Linear(width, width, bias=False)\n",
    "    def split(self, x):\n",
    "        batch, length, width = x.shape\n",
    "        return x.view(batch, length, self.heads, self.head_dim).transpose(1, 2)\n",
    "    def merge(self, x):\n",
    "        batch, heads, length, dim = x.shape\n",
    "        return x.transpose(1, 2).contiguous().view(batch, length, heads * dim)\n",
    "    def full(self, x):\n",
    "        q, k, v = self.split(self.q(x)), self.split(self.k(x)), self.split(self.v(x))\n",
    "        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)\n",
    "        length = x.shape[1]\n",
    "        mask = torch.triu(torch.ones(length, length, dtype=torch.bool), diagonal=1)\n",
    "        weights = torch.softmax(scores.masked_fill(mask, float('-inf')), dim=-1)\n",
    "        return self.out(self.merge(weights @ v))\n",
    "    def step(self, x, cache=None):\n",
    "        q = self.split(self.q(x))\n",
    "        current_k, current_v = self.split(self.k(x)), self.split(self.v(x))\n",
    "        if cache is None:\n",
    "            k, v = current_k, current_v\n",
    "        else:\n",
    "            k = torch.cat((cache[0], current_k), dim=2)\n",
    "            v = torch.cat((cache[1], current_v), dim=2)\n",
    "        weights = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(self.head_dim), dim=-1)\n",
    "        return self.out(self.merge(weights @ v)), (k, v)\n",
    "\n",
    "attention = CausalAttention().eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 正确性：缓存前后输出必须一致"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(1, 32, 128)\n",
    "full_output = attention.full(x)\n",
    "cache, steps = None, []\n",
    "for position in range(x.shape[1]):\n",
    "    output, cache = attention.step(x[:, position:position+1], cache)\n",
    "    steps.append(output)\n",
    "cached_output = torch.cat(steps, dim=1)\n",
    "error = (full_output - cached_output).abs().max().item()\n",
    "print('cache K shape:', tuple(cache[0].shape))\n",
    "print('maximum absolute error:', error)\n",
    "assert error < 1e-5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 计算量：逐 token 生成时不再重复投影历史 K/V\n",
    "\n",
    "下面的基准包含 Python 循环开销，绝对时间不代表生产推理性能；趋势用于说明重复计算的增长方式。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def measure(length, repeats=3):\n",
    "    sample = torch.randn(1, length, 128)\n",
    "    started = time.perf_counter()\n",
    "    for _ in range(repeats):\n",
    "        for end in range(1, length + 1): attention.full(sample[:, :end])\n",
    "    no_cache = (time.perf_counter() - started) / repeats\n",
    "    started = time.perf_counter()\n",
    "    for _ in range(repeats):\n",
    "        cache = None\n",
    "        for pos in range(length): _, cache = attention.step(sample[:, pos:pos+1], cache)\n",
    "    cached = (time.perf_counter() - started) / repeats\n",
    "    return no_cache, cached\n",
    "\n",
    "lengths = [16, 32, 64, 128]\n",
    "timings = [measure(length) for length in lengths]\n",
    "for length, (plain, cached) in zip(lengths, timings):\n",
    "    print(f'{length:3d} tokens  no-cache={plain:.4f}s  cache={cached:.4f}s  ratio={plain/cached:.2f}x')\n",
    "plt.plot(lengths, [x[0] for x in timings], 'o-', label='recompute history')\n",
    "plt.plot(lengths, [x[1] for x in timings], 'o-', label='KV cache')\n",
    "plt.xlabel('generated sequence length'); plt.ylabel('seconds'); plt.legend(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 显存模型\n",
    "\n",
    "每层缓存字节数为 `2 × batch × tokens × kv_heads × head_dim × element_bytes`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cache_gib(layers, tokens, kv_heads, head_dim, element_bytes=2, batch=1):\n",
    "    return 2 * layers * tokens * kv_heads * head_dim * element_bytes * batch / 1024**3\n",
    "\n",
    "for kv_heads, name in [(32, 'MHA'), (8, 'GQA'), (1, 'MQA')]:\n",
    "    size = cache_gib(layers=32, tokens=4096, kv_heads=kv_heads, head_dim=128)\n",
    "    print(f'{name}: {size:.3f} GiB per sequence')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 练习\n",
    "\n",
    "1. 将 `batch` 和 `tokens` 分别扩大四倍，验证缓存显存的线性关系。\n",
    "2. 思考为什么 Query 不需要缓存。\n",
    "3. 当前实现每次 `torch.cat` 都会重新分配内存。设计一个预分配 Static Cache，并比较性能。\n",
    "4. 将缓存改成 FP16，检查输出误差和内存变化。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}