{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM10-L3: Implement a KV cache and rotary embeddings in PyTorch\n",
    "\n",
    "In this notebook you add two ingredients of modern decoder-only models to attention: **rotary position embeddings (RoPE)** and a **key-value (KV) cache**,\n",
    "including grouped-query attention (several query heads sharing one key-value head). You will prove that decoding with a cache gives exactly the same\n",
    "result as recomputing the whole sequence, and count how much work the cache saves.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, a few minutes. No downloads; the model is tiny and has random weights, so the *text* it produces is meaningless.\n",
    "What is tested is the *equality of two ways of computing the same thing*, which is exactly what a real inference engine must guarantee.\n",
    "\n",
    "**How it works:** cells marked `TODO` raise `NotImplementedError` until you write the code. Each is followed by a self-check (`assert`).\n",
    "Theory: Lessons LM10-T2 (rotary embeddings), LM10-T3 (grouped-query attention) and LM10-T5 (KV caching) on this site."
   ]
  },
  {
   "cell_type": "code",
   "id": "c002",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "torch.manual_seed(0)\n",
    "print(\"torch\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. TODO: rotary position embeddings\n",
    "\n",
    "RoPE rotates each (query or key) vector in 2-D planes by an angle that grows with the token position. A head vector of width `hd` is split into two halves\n",
    "`x1 = x[..., :hd/2]` and `x2 = x[..., hd/2:]`; the pair `(x1[i], x2[i])` is rotated by the angle `pos * theta_i` with `theta_i = base ** (-2 i / hd)`, `i = 0 ... hd/2 - 1`.\n",
    "\n",
    "1. `rope_tables(hd, n_pos, base=10000.0)` returns `cos, sin`, each of shape `(n_pos, hd/2)`, holding `cos(pos * theta_i)` and `sin(pos * theta_i)`.\n",
    "2. `apply_rope(x, cos, sin)` takes `x` of shape `(B, H, T, hd)` and tables of shape `(T, hd/2)` and returns\n",
    "   `cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def rope_tables(hd, n_pos, base=10000.0):\n",
    "    raise NotImplementedError(\"TODO: build the cos and sin tables, each of shape (n_pos, hd // 2)\")\n",
    "\n",
    "def apply_rope(x, cos, sin):\n",
    "    raise NotImplementedError(\"TODO: rotate the two halves of x by the angles in cos and sin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "### Self-check 1: shapes, norm, position 0, and relative position"
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "hd, n_pos = 8, 16\n",
    "cos, sin = rope_tables(hd, n_pos)\n",
    "assert cos.shape == (n_pos, hd // 2) and sin.shape == (n_pos, hd // 2), \"tables must have shape (n_pos, hd // 2)\"\n",
    "assert torch.allclose(cos[0], torch.ones(hd // 2)) and torch.allclose(sin[0], torch.zeros(hd // 2)), \"position 0 is a rotation by angle 0\"\n",
    "assert abs(cos[1, 0].item() - torch.tensor(1.0).cos().item()) < 1e-6, \"theta_0 = 1, so position 1 has angle 1 radian in the first plane\"\n",
    "\n",
    "q = torch.randn(1, 1, 1, hd)\n",
    "rotated = apply_rope(q.expand(1, 1, n_pos, hd), cos, sin)\n",
    "assert rotated.shape == (1, 1, n_pos, hd), \"apply_rope must keep the shape\"\n",
    "assert torch.allclose(rotated[:, :, 0], q[:, :, 0], atol=1e-6), \"position 0 must leave the vector unchanged\"\n",
    "assert torch.allclose(rotated.norm(dim=-1), q.norm(dim=-1).expand(1, 1, n_pos), atol=1e-5), \"a rotation must preserve the length of the vector\"\n",
    "\n",
    "# relative position: the score between a query at m and a key at n depends only on m - n\n",
    "k = torch.randn(1, 1, 1, hd)\n",
    "qs = apply_rope(q.expand(1, 1, n_pos, hd), cos, sin)\n",
    "ks = apply_rope(k.expand(1, 1, n_pos, hd), cos, sin)\n",
    "scores = (qs @ ks.transpose(2, 3))[0, 0]\n",
    "assert torch.allclose(scores[5, 3], scores[9, 7], atol=1e-5), \"score(m=5, n=3) must equal score(m=9, n=7): the offset is 2 in both\"\n",
    "assert torch.allclose(scores[4, 4], scores[11, 11], atol=1e-5), \"score at offset 0 must not depend on the absolute position\"\n",
    "assert not torch.allclose(scores[5, 3], scores[5, 2], atol=1e-3), \"different offsets should give different scores\"\n",
    "print(\"self-check 1 passed: RoPE is a rotation and scores depend only on the offset\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "## 2. TODO: the cache and the mask\n",
    "\n",
    "During decoding the model feeds only the *new* tokens. Their keys and values must be appended to the stored ones, and the causal mask must account for the offset.\n",
    "\n",
    "1. `update_cache(cache, k, v)`: `cache` is `None` or a tuple `(k_past, v_past)`, with tensors of shape `(B, H_kv, T_past, hd)`. `k`, `v` have shape `(B, H_kv, n_new, hd)`.\n",
    "   Return the new tuple with the new keys and values concatenated along the *time* dimension (`dim=2`).\n",
    "2. `offset_causal_mask(n_query, n_total)`: a boolean tensor of shape `(n_query, n_total)` that is `True` where attention is **forbidden**.\n",
    "   The queries are the last `n_query` of `n_total` positions: query `i` has absolute position `n_total - n_query + i` and may see key `j` only if `j` is at most that position."
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def update_cache(cache, k, v):\n",
    "    raise NotImplementedError(\"TODO: append k and v to the cache along the time dimension\")\n",
    "\n",
    "def offset_causal_mask(n_query, n_total):\n",
    "    raise NotImplementedError(\"TODO: True where a query must NOT attend to a key\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "### Self-check 2: cache growth and mask"
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "k1, v1 = torch.randn(1, 2, 3, hd), torch.randn(1, 2, 3, hd)\n",
    "k2, v2 = torch.randn(1, 2, 1, hd), torch.randn(1, 2, 1, hd)\n",
    "c = update_cache(None, k1, v1)\n",
    "assert c[0].shape == (1, 2, 3, hd), \"an empty cache becomes the new keys and values\"\n",
    "c = update_cache(c, k2, v2)\n",
    "assert c[0].shape == (1, 2, 4, hd) and c[1].shape == (1, 2, 4, hd), \"the time dimension (dim=2) must grow from 3 to 4\"\n",
    "assert torch.equal(c[0][:, :, :3], k1) and torch.equal(c[0][:, :, 3:], k2), \"old keys first, new keys last\"\n",
    "\n",
    "full = offset_causal_mask(4, 4)\n",
    "assert torch.equal(full, torch.triu(torch.ones(4, 4, dtype=torch.bool), diagonal=1)), \"with no past the mask is the usual causal mask\"\n",
    "step = offset_causal_mask(1, 7)\n",
    "assert step.shape == (1, 7) and not step.any(), \"a single new token may attend to every cached position and to itself\"\n",
    "chunk = offset_causal_mask(2, 5)\n",
    "assert chunk.tolist() == [[False, False, False, False, True], [False, False, False, False, False]], \"two new tokens at positions 3 and 4\"\n",
    "print(\"self-check 2 passed: the cache grows along time and the mask respects the offset\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "## 3. A small decoder with grouped-query attention (given)\n",
    "\n",
    "The layer below uses your functions. It has `n_heads` query heads and `n_kv_heads` key-value heads (`n_kv_heads` divides `n_heads`); each key-value head\n",
    "serves `n_heads // n_kv_heads` query heads. RoPE is applied to queries and keys, and the **rotated keys are cached** (rotating at the right absolute position\n",
    "once is cheaper than rotating the whole cache at every step). `start_pos` is the absolute position of the first input token."
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class GQAttention(nn.Module):\n",
    "    def __init__(self, d, n_heads, n_kv_heads, max_pos=256):\n",
    "        super().__init__()\n",
    "        assert n_heads % n_kv_heads == 0\n",
    "        self.H, self.Hkv, self.hd = n_heads, n_kv_heads, d // n_heads\n",
    "        self.Wq = nn.Linear(d, n_heads * self.hd, bias=False)\n",
    "        self.Wk = nn.Linear(d, n_kv_heads * self.hd, bias=False)\n",
    "        self.Wv = nn.Linear(d, n_kv_heads * self.hd, bias=False)\n",
    "        self.Wo = nn.Linear(n_heads * self.hd, d, bias=False)\n",
    "        cos, sin = rope_tables(self.hd, max_pos)\n",
    "        self.register_buffer(\"cos\", cos)\n",
    "        self.register_buffer(\"sin\", sin)\n",
    "\n",
    "    def forward(self, x, cache=None, start_pos=0):\n",
    "        B, n, _ = x.shape\n",
    "        q = self.Wq(x).view(B, n, self.H, self.hd).transpose(1, 2)\n",
    "        k = self.Wk(x).view(B, n, self.Hkv, self.hd).transpose(1, 2)\n",
    "        v = self.Wv(x).view(B, n, self.Hkv, self.hd).transpose(1, 2)\n",
    "        c, s = self.cos[start_pos:start_pos + n], self.sin[start_pos:start_pos + n]\n",
    "        q, k = apply_rope(q, c, s), apply_rope(k, c, s)\n",
    "        k, v = update_cache(cache, k, v)\n",
    "        new_cache = (k, v)\n",
    "        rep = self.H // self.Hkv\n",
    "        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)\n",
    "        scores = q @ k.transpose(2, 3) / self.hd ** 0.5\n",
    "        scores = scores.masked_fill(offset_causal_mask(n, k.shape[2]), float(\"-inf\"))\n",
    "        ctx = (torch.softmax(scores, dim=-1) @ v).transpose(1, 2).reshape(B, n, self.H * self.hd)\n",
    "        return self.Wo(ctx), new_cache\n",
    "\n",
    "\n",
    "class TinyLM(nn.Module):\n",
    "    def __init__(self, vocab=50, d=32, n_heads=4, n_kv_heads=2):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(vocab, d)\n",
    "        self.attn = GQAttention(d, n_heads, n_kv_heads)\n",
    "        self.head = nn.Linear(d, vocab, bias=False)\n",
    "\n",
    "    def forward(self, ids, cache=None, start_pos=0):\n",
    "        h = self.emb(ids)\n",
    "        a, cache = self.attn(h, cache, start_pos)\n",
    "        return self.head(h + a), cache"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "### Self-check 3: cached decoding equals full recomputation\n",
    "\n",
    "Run a sequence of 12 tokens in one pass, then run the same sequence as a 8-token *prefill* followed by four one-token *decode* steps that reuse the cache.\n",
    "The logits must agree, and the cache of a model with fewer key-value heads must be proportionally smaller."
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(3)\n",
    "model = TinyLM(vocab=50, d=32, n_heads=4, n_kv_heads=2).eval()\n",
    "ids = torch.randint(0, 50, (1, 12))\n",
    "with torch.no_grad():\n",
    "    full_logits, _ = model(ids)\n",
    "    part, cache = model(ids[:, :8])\n",
    "    pieces = [part]\n",
    "    for t in range(8, 12):\n",
    "        step_logits, cache = model(ids[:, t:t + 1], cache, start_pos=t)\n",
    "        pieces.append(step_logits)\n",
    "inc_logits = torch.cat(pieces, dim=1)\n",
    "err = (full_logits - inc_logits).abs().max().item()\n",
    "print(f\"largest difference between full and cached logits: {err:.2e}\")\n",
    "assert err < 1e-5, \"cached decoding must reproduce the full forward pass: check the offset of the mask and of the RoPE positions\"\n",
    "assert cache[0].shape == (1, 2, 12, 8), f\"the cache should hold 12 positions of 2 key-value heads of width 8, got {tuple(cache[0].shape)}\"\n",
    "\n",
    "def cache_bytes(cache):\n",
    "    return sum(t.numel() * t.element_size() for t in cache)\n",
    "mha_cache = TinyLM(vocab=50, d=32, n_heads=4, n_kv_heads=4).eval()(ids)[1]\n",
    "mqa_cache = TinyLM(vocab=50, d=32, n_heads=4, n_kv_heads=1).eval()(ids)[1]\n",
    "print(\"cache bytes: MHA\", cache_bytes(mha_cache), \" GQA(2)\", cache_bytes(cache), \" MQA\", cache_bytes(mqa_cache))\n",
    "assert cache_bytes(mha_cache) == 2 * cache_bytes(cache) == 4 * cache_bytes(mqa_cache), \"the cache must scale with the number of key-value heads\"\n",
    "print(\"self-check 3 passed: the cache is exact and its size scales with the key-value heads\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "## 4. TODO: greedy generation with and without a cache\n",
    "\n",
    "Implement `generate(model, prompt, n_new, use_cache)` for a `prompt` of shape `(1, T)` of token ids. At every step take the arg-max of the logits at the last position\n",
    "and append it. Also count `processed`: the total number of tokens that the model is given as input over the whole generation.\n",
    "\n",
    "- Without a cache every step feeds the **whole** sequence so far (`start_pos=0`, `cache=None`).\n",
    "- With a cache the first step feeds the whole prompt (the prefill); every later step feeds **only the last generated token**, with `start_pos` equal to the number of tokens already in the cache.\n",
    "\n",
    "Return `(ids, processed)`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def generate(model, prompt, n_new, use_cache):\n",
    "    raise NotImplementedError(\"TODO: greedy decoding, with the prefill/decode split when use_cache is True\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "### Self-check 4: identical text, much less work"
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "prompt = torch.randint(0, 50, (1, 6))\n",
    "a, work_a = generate(model, prompt, 10, use_cache=False)\n",
    "b, work_b = generate(model, prompt, 10, use_cache=True)\n",
    "assert torch.equal(a, b), \"the cache must not change the generated tokens\"\n",
    "assert a.shape == (1, 16), \"6 prompt tokens plus 10 new ones\"\n",
    "assert work_a == sum(6 + i for i in range(10)), f\"without a cache the model reads 6 + 7 + ... + 15 = 105 tokens, got {work_a}\"\n",
    "assert work_b == 6 + 9, f\"with a cache: 6 prompt tokens and then 9 single tokens = 15, got {work_b}\"\n",
    "print(f\"tokens processed: {work_a} without a cache, {work_b} with a cache ({work_a / work_b:.0f}x fewer)\")\n",
    "print(\"self-check 4 passed: same output, less work\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "## 5. Experiment: how does the saving grow?\n",
    "\n",
    "Time both ways for a longer prompt. The numbers depend on your machine and on this tiny model, so there are no assertions; read them and compare with Lesson LM10-T5.\n",
    "Note how the count of processed tokens grows with the *square* of the length without a cache and linearly with one."
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "big = TinyLM(vocab=50, d=128, n_heads=8, n_kv_heads=2).eval()\n",
    "prompt = torch.randint(0, 50, (1, 64))\n",
    "for use_cache in (False, True):\n",
    "    t0 = time.perf_counter()\n",
    "    out, work = generate(big, prompt, 96, use_cache)\n",
    "    print(f\"use_cache={use_cache!s:5}: {time.perf_counter() - t0:6.2f} s, {work:5d} tokens processed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## 6. Reflect\n",
    "\n",
    "Write short answers in your own notes, then compare with Lessons LM10-T2, LM10-T3 and LM10-T5.\n",
    "\n",
    "1. Why are the *rotated* keys cached, and what would go wrong if you cached unrotated keys and forgot the position when reusing them?\n",
    "2. A model has 32 layers, 8 key-value heads of width 128 and stores the cache in 16 bits. How many bytes does the cache need per token? (Compare Lab LM10-L1.)\n",
    "3. Why does the cache not change the output, but change the memory and the time?\n",
    "4. Which of the two phases, the prefill of the prompt or one decode step, does more arithmetic per byte of weights read?\n",
    "\n",
    "**Done when** all four self-checks print \"passed\"."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}