{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f167b85e",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Module 15: Memoization - Computational Reuse for Inference\n",
    "\n",
    "Welcome to Module 15! You'll implement memoization - a fundamental optimization pattern. We'll apply it to transformers through KV caching for 10-15x faster text generation.\n",
    "\n",
    "## 🔗 Prerequisites & Progress\n",
    "**You've Built**: Complete transformer architecture (Module 13) and profiling tools (Module 14)\n",
    "**You'll Build**: Memoization system that eliminates redundant computation through caching\n",
    "**You'll Enable**: Production-grade inference optimization using computational reuse\n",
    "\n",
    "**Connection Map**:\n",
    "```\n",
    "Profiling (14) → Memoization (15) → Quantization (16)\n",
    "(measure O(n²))  (cache K,V → O(n))  (reduce precision)\n",
    "```\n",
    "\n",
    "## Learning Objectives\n",
    "By the end of this module, you will:\n",
    "1. Understand memoization as a general optimization pattern (cache results, avoid recomputation)\n",
    "2. Apply memoization to transformers through KV caching\n",
    "3. Implement KVCache with efficient memory management and O(1) updates\n",
    "4. Build cache-aware attention that reuses previously computed keys and values\n",
    "5. Measure dramatic speedup gains (10-15x) and understand memory trade-offs\n",
    "\n",
    "Let's make inference blazingly fast through computational reuse!\n",
    "\n",
    "## 📦 Where This Code Lives in the Final Package\n",
    "\n",
    "**Learning Side:** You work in `modules/15_memoization/kvcaching_dev.py`  \n",
    "**Building Side:** Code exports to `tinytorch.generation.kv_cache`\n",
    "\n",
    "```python\n",
    "# How to use this module:\n",
    "from tinytorch.generation.kv_cache import KVCache, enable_kv_cache\n",
    "```\n",
    "\n",
    "**Why this matters:**\n",
    "- **Learning:** Complete caching system demonstrating production optimization techniques\n",
    "- **Production:** Proper organization matching Hugging Face's generation/ module structure\n",
    "- **Consistency:** All generation optimizations in generation.kv_cache\n",
    "- **Integration:** Works seamlessly with transformers for complete inference optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b34fcf1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp generation.kv_cache\n",
    "#| export\n",
    "\n",
    "import numpy as np\n",
    "import time\n",
    "from typing import Tuple, Optional, Dict, List\n",
    "\n",
    "# Import TinyTorch components from previous modules\n",
    "from tinytorch.core.tensor import Tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "560eefc2",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🔬 Motivation: Why Memoization Matters for Transformers\n",
    "\n",
    "Before we learn KV caching, let's profile transformer generation to understand \n",
    "the problem we're solving. We'll see O(n²) growth in latency as we generate text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d66ae97",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Profile transformer generation to discover the bottleneck\n",
    "from tinytorch.profiling.profiler import Profiler\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "profiler = Profiler()\n",
    "\n",
    "def naive_attention_step(seq_len, hidden_dim=64):\n",
    "    \"\"\"\n",
    "    Simulates one step of attention computation.\n",
    "    Without caching, this processes ALL previous tokens every time.\n",
    "    \"\"\"\n",
    "    # Q, K, V for entire sequence\n",
    "    q = Tensor(np.random.randn(1, seq_len, hidden_dim))\n",
    "    k = Tensor(np.random.randn(1, seq_len, hidden_dim))\n",
    "    v = Tensor(np.random.randn(1, seq_len, hidden_dim))\n",
    "    \n",
    "    # Attention: Q @ K.T then @ V\n",
    "    # This is O(seq_len²) in complexity\n",
    "    scores = q @ k.T  # (1, seq_len, seq_len)\n",
    "    output = scores @ v\n",
    "    \n",
    "    return output\n",
    "\n",
    "# Profile at increasing sequence lengths\n",
    "print(\"🔬 Profiling Transformer Generation (Without Caching):\\n\")\n",
    "print(\"   Seq Len  |  Latency (ms)  |  Growth\")\n",
    "print(\"   ---------|----------------|----------\")\n",
    "\n",
    "sequence_lengths = [10, 20, 40, 80, 160]\n",
    "latencies = []\n",
    "\n",
    "for seq_len in sequence_lengths:\n",
    "    # Measure latency for this sequence length\n",
    "    latency = profiler.measure_latency(\n",
    "        lambda: naive_attention_step(seq_len),\n",
    "        None,\n",
    "        warmup=5,\n",
    "        iterations=20\n",
    "    )\n",
    "    latencies.append(latency)\n",
    "    \n",
    "    # Calculate growth rate\n",
    "    if len(latencies) > 1:\n",
    "        growth = latencies[-1] / latencies[-2]\n",
    "        print(f\"   {seq_len:3d}      |  {latency:6.2f}        |  {growth:.2f}×\")\n",
    "    else:\n",
    "        print(f\"   {seq_len:3d}      |  {latency:6.2f}        |  baseline\")\n",
    "\n",
    "print(\"\\n💡 Key Observations:\")\n",
    "print(\"   • Latency grows QUADRATICALLY with sequence length\")\n",
    "print(\"   • Each new token forces recomputation of ALL previous K,V pairs\")\n",
    "print(\"   • For 160 tokens: ~4× time vs 80 tokens (2² growth)\")\n",
    "\n",
    "print(\"\\n🎯 The Problem:\")\n",
    "print(\"   K and V values for previous tokens NEVER change,\")\n",
    "print(\"   yet we recompute them every single step!\")\n",
    "\n",
    "print(\"\\n✨ The Solution:\")\n",
    "print(\"   CACHE the K,V values! (That's memoization)\")\n",
    "print(\"   • First compute: Calculate and store K,V\")\n",
    "print(\"   • Later steps: Reuse stored K,V\")\n",
    "print(\"   • Complexity: O(n²) → O(n)\")\n",
    "print(\"   • Speedup: 10-15× for typical generation\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cad5a0e9",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 Part 1: Understanding the Autoregressive Generation Problem\n",
    "\n",
    "### The Core Inefficiency\n",
    "\n",
    "When generating text token by token, transformers face a fundamental computational bottleneck. Let's visualize what happens during naive generation:\n",
    "\n",
    "```\n",
    "Token Generation Process (Without Caching):\n",
    "\n",
    "Step 1: Generate \"Hello\"\n",
    "Input: [START]\n",
    "Attention: Q₁ × [K₁] × [V₁]               ← 1 computation\n",
    "\n",
    "Step 2: Generate \"world\"\n",
    "Input: [START, Hello]\n",
    "Attention: Q₂ × [K₁, K₂] × [V₁, V₂]       ← 2 computations (K₁,V₁ RECOMPUTED!)\n",
    "\n",
    "Step 3: Generate \"!\"\n",
    "Input: [START, Hello, world]\n",
    "Attention: Q₃ × [K₁, K₂, K₃] × [V₁, V₂, V₃] ← 3 computations (K₁,V₁,K₂,V₂ RECOMPUTED!)\n",
    "```\n",
    "\n",
    "**The Problem**: For each new token, we recompute ALL previous key-value pairs even though they never change!\n",
    "\n",
    "### Computational Complexity Analysis\n",
    "\n",
    "```\n",
    "Naive Generation Complexity:\n",
    "Step 1: 1 K,V computation\n",
    "Step 2: 2 K,V computations\n",
    "Step 3: 3 K,V computations\n",
    "...\n",
    "Step n: n K,V computations\n",
    "\n",
    "Total: 1 + 2 + 3 + ... + n = n(n+1)/2 = O(n²) complexity!\n",
    "```\n",
    "\n",
    "For a 100-token sequence, this means **5,050 redundant computations**!\n",
    "\n",
    "### Real-World Impact\n",
    "\n",
    "This inefficiency makes production LLM serving economically impossible without optimization:\n",
    "- **ChatGPT/GPT-4**: Would be too slow for real-time chat without caching\n",
    "- **Code completion**: IDEs couldn't provide instant suggestions\n",
    "- **Mobile deployment**: On-device generation would drain batteries instantly\n",
    "- **API serving**: Server costs would be 10x+ higher\n",
    "\n",
    "**The Solution**: Cache key-value pairs after computing them once, transforming O(n²) into O(n)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "045c13d9",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🧮 Part 2: The Key-Value Caching Insight\n",
    "\n",
    "### Mathematical Foundation\n",
    "\n",
    "The core insight comes from understanding what changes during autoregressive generation:\n",
    "\n",
    "```\n",
    "Attention Computation Breakdown:\n",
    "\n",
    "Q = new_token @ W_q        ← Only new token (changes each step)\n",
    "K = all_tokens @ W_k       ← Includes old tokens (mostly redundant!)\n",
    "V = all_tokens @ W_v       ← Includes old tokens (mostly redundant!)\n",
    "\n",
    "attention_output = softmax(Q @ K.T / √d_k) @ V\n",
    "```\n",
    "\n",
    "**Key Insight**: K and V matrices for previous tokens NEVER change!\n",
    "\n",
    "```\n",
    "Token Dependencies:\n",
    "K₁ = token₁ @ W_k  ← Computed once, never changes\n",
    "K₂ = token₂ @ W_k  ← Computed once, never changes\n",
    "K₃ = token₃ @ W_k  ← Computed once, never changes\n",
    "\n",
    "Same for V₁, V₂, V₃...\n",
    "```\n",
    "\n",
    "### Cache-Optimized Generation\n",
    "\n",
    "```\n",
    "Optimized Generation Process (With Caching):\n",
    "\n",
    "Step 1: Generate \"Hello\"\n",
    "Compute: K₁, V₁ → Store in cache\n",
    "Attention: Q₁ × cached[K₁] × cached[V₁]\n",
    "\n",
    "Step 2: Generate \"world\"\n",
    "Compute: K₂, V₂ → Append to cache\n",
    "Attention: Q₂ × cached[K₁, K₂] × cached[V₁, V₂]\n",
    "\n",
    "Step 3: Generate \"!\"\n",
    "Compute: K₃, V₃ → Append to cache\n",
    "Attention: Q₃ × cached[K₁, K₂, K₃] × cached[V₁, V₂, V₃]\n",
    "```\n",
    "\n",
    "**Result**: Each step computes only ONE new K,V pair instead of recomputing ALL!\n",
    "\n",
    "### Memory vs Compute Trade-off\n",
    "\n",
    "```\n",
    "Traditional Approach:\n",
    "Memory: O(1)          (no storage needed)\n",
    "Compute: O(n²)        (recompute everything)\n",
    "\n",
    "Cached Approach:\n",
    "Memory: O(n × d_k)    (store all K,V pairs)\n",
    "Compute: O(n)         (only compute new pairs)\n",
    "\n",
    "For n=100, d_k=64:\n",
    "Memory cost: 6.4 KB per layer\n",
    "Compute savings: 50x reduction in K,V computations\n",
    "```\n",
    "\n",
    "**Trade-off Winner**: Memory is cheap, compute is expensive! Use O(n) memory to save O(n²) compute."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c85596c",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🏗️ Part 3: KVCache Class Implementation\n",
    "\n",
    "### Core Requirements\n",
    "\n",
    "Our KVCache needs to efficiently handle:\n",
    "\n",
    "1. **Multi-layer storage**: Each transformer layer needs its own K,V cache\n",
    "2. **Multi-head attention**: Each attention head has separate K,V pairs\n",
    "3. **Batch processing**: Support multiple sequences simultaneously (batch inference)\n",
    "4. **Dynamic updates**: Efficiently append new tokens without copying data\n",
    "5. **Memory management**: Pre-allocate space to avoid dynamic resizing overhead\n",
    "\n",
    "### Cache Architecture Visualization\n",
    "\n",
    "```\n",
    "KVCache Memory Layout:\n",
    "┌─────────────────────────────────────────────────────────┐\n",
    "│                    KVCache Object                       │\n",
    "├─────────────────────────────────────────────────────────┤\n",
    "│ Layer 0: ┌─────────────┬─────────────┐                 │\n",
    "│          │ Key Cache   │ Value Cache │                 │\n",
    "│          │ (B,H,S,D)   │ (B,H,S,D)   │                 │\n",
    "│          └─────────────┴─────────────┘                 │\n",
    "├─────────────────────────────────────────────────────────┤\n",
    "│ Layer 1: ┌─────────────┬─────────────┐                 │\n",
    "│          │ Key Cache   │ Value Cache │                 │\n",
    "│          │ (B,H,S,D)   │ (B,H,S,D)   │                 │\n",
    "│          └─────────────┴─────────────┘                 │\n",
    "├─────────────────────────────────────────────────────────┤\n",
    "│   ...    ┌─────────────┬─────────────┐                 │\n",
    "│ Layer N: │ Key Cache   │ Value Cache │                 │\n",
    "│          │ (B,H,S,D)   │ (B,H,S,D)   │                 │\n",
    "│          └─────────────┴─────────────┘                 │\n",
    "└─────────────────────────────────────────────────────────┘\n",
    "\n",
    "Where:\n",
    "B = batch_size    (number of sequences)\n",
    "H = num_heads     (attention heads per layer)\n",
    "S = max_seq_len   (maximum sequence length)\n",
    "D = head_dim      (dimension per attention head)\n",
    "```\n",
    "\n",
    "### Update Operation Flow\n",
    "\n",
    "```\n",
    "Cache Update Process:\n",
    "                      seq_pos = 2\n",
    "                         ↓\n",
    "┌─────┬─────┬─────┬─────┬─────┬─────┐\n",
    "│ K₁  │ K₂  │ ??? │ ??? │ ??? │ ??? │ ← Key Cache\n",
    "├─────┼─────┼─────┼─────┼─────┼─────┤\n",
    "│ V₁  │ V₂  │ ??? │ ??? │ ??? │ ??? │ ← Value Cache\n",
    "└─────┴─────┴─────┴─────┴─────┴─────┘\n",
    "\n",
    "New token arrives: K₃, V₃\n",
    "\n",
    "                      seq_pos = 2\n",
    "                         ↓\n",
    "┌─────┬─────┬─────┬─────┬─────┬─────┐\n",
    "│ K₁  │ K₂  │ K₃  │ ??? │ ??? │ ??? │ ← Write K₃ here\n",
    "├─────┼─────┼─────┼─────┼─────┼─────┤\n",
    "│ V₁  │ V₂  │ V₃  │ ??? │ ??? │ ??? │ ← Write V₃ here\n",
    "└─────┴─────┴─────┴─────┴─────┴─────┘\n",
    "\n",
    "Then: seq_pos += 1 (advance to position 3)\n",
    "```\n",
    "\n",
    "This design enables **O(1) updates** - just write to the next position!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3f7baa6",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "kvcache-class",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class KVCache:\n",
    "    \"\"\"\n",
    "    Efficient key-value cache for autoregressive generation.\n",
    "\n",
    "    Stores K,V matrices for each transformer layer to avoid recomputation\n",
    "    during sequential token generation. This is THE critical optimization\n",
    "    that makes production language model serving economically viable.\n",
    "    \n",
    "    ⚠️  IMPORTANT: INFERENCE-ONLY (No Gradient Tracking)\n",
    "    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n",
    "    KV caching is designed ONLY for inference (generation), NOT training.\n",
    "    - During generation: No gradients computed (model.eval() mode)\n",
    "    - Cache operations use .data (no gradient tracking)\n",
    "    - This is correct and intentional for maximum speed\n",
    "    - DO NOT use caching during training (use standard forward pass)\n",
    "    \n",
    "    Architecture:\n",
    "    - Pre-allocates cache tensors with maximum sequence length\n",
    "    - Tracks current sequence position for efficient O(1) updates\n",
    "    - Provides update() method to append new K,V pairs without copying\n",
    "    - Provides get() method to retrieve cached values for attention\n",
    "    - Handles multiple layers and attention heads properly\n",
    "    \n",
    "    Memory Layout:\n",
    "    ```\n",
    "    Layer 0: [Key_cache, Value_cache]  # Shape: (batch, num_heads, max_seq, head_dim)\n",
    "    Layer 1: [Key_cache, Value_cache]\n",
    "    ...\n",
    "    Layer N: [Key_cache, Value_cache]\n",
    "    ```\n",
    "\n",
    "    Performance:\n",
    "    - Update: O(1) - just index assignment\n",
    "    - Get: O(1) - just slicing (no data copy)\n",
    "    - Memory: O(num_layers × batch × heads × max_seq × head_dim)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, batch_size: int, max_seq_len: int, num_layers: int,\n",
    "                 num_heads: int, head_dim: int):\n",
    "        \"\"\"\n",
    "        Initialize KV cache for efficient generation.\n",
    "\n",
    "        TODO: Set up pre-allocated cache storage for all transformer layers\n",
    "\n",
    "        APPROACH:\n",
    "        1. Store configuration parameters (batch_size, max_seq_len, etc.)\n",
    "        2. Initialize sequence position counter to 0\n",
    "        3. Create empty list for cache storage\n",
    "        4. For each layer, pre-allocate zero-filled key and value caches\n",
    "        5. Store each layer's (key_cache, value_cache) tuple in the list\n",
    "\n",
    "        Args:\n",
    "            batch_size: Number of sequences to generate simultaneously\n",
    "            max_seq_len: Maximum sequence length to support\n",
    "            num_layers: Number of transformer layers\n",
    "            num_heads: Number of attention heads per layer\n",
    "            head_dim: Dimension of each attention head\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> cache = KVCache(batch_size=2, max_seq_len=128, num_layers=4,\n",
    "        ...                 num_heads=8, head_dim=64)\n",
    "        >>> cache.seq_pos  # 0 (no tokens cached yet)\n",
    "        >>> len(cache.caches)  # 4 (one per layer)\n",
    "        >>> cache.caches[0][0].shape  # (2, 8, 128, 64) - key cache for layer 0\n",
    "\n",
    "        HINTS:\n",
    "        - Cache shape: (batch_size, num_heads, max_seq_len, head_dim)\n",
    "        - Use Tensor(np.zeros(...)) to create cache tensors\n",
    "        - Store caches as list of tuples: [(key_0, val_0), (key_1, val_1), ...]\n",
    "        - Pre-allocation avoids dynamic resizing overhead during generation\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.batch_size = batch_size\n",
    "        self.max_seq_len = max_seq_len\n",
    "        self.num_layers = num_layers\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = head_dim\n",
    "\n",
    "        # Current sequence position (how many tokens are cached)\n",
    "        self.seq_pos = 0\n",
    "\n",
    "        # Cache storage: list of (key_cache, value_cache) tuples per layer\n",
    "        self.caches = []\n",
    "\n",
    "        for layer_idx in range(num_layers):\n",
    "            # Pre-allocate cache tensors with maximum size\n",
    "            # Shape: (batch_size, num_heads, max_seq_len, head_dim)\n",
    "            key_cache = Tensor(np.zeros((batch_size, num_heads, max_seq_len, head_dim)))\n",
    "            value_cache = Tensor(np.zeros((batch_size, num_heads, max_seq_len, head_dim)))\n",
    "\n",
    "            self.caches.append((key_cache, value_cache))\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def update(self, layer_idx: int, key: Tensor, value: Tensor) -> None:\n",
    "        \"\"\"\n",
    "        Update cache with new key-value pairs for given layer.\n",
    "\n",
    "        TODO: Efficiently append new K,V to cache without data copying\n",
    "\n",
    "        APPROACH:\n",
    "        1. Validate layer_idx is in range [0, num_layers-1]\n",
    "        2. Validate seq_pos hasn't exceeded max_seq_len\n",
    "        3. Retrieve the (key_cache, value_cache) tuple for this layer\n",
    "        4. Write new key to position seq_pos in key_cache using indexed assignment\n",
    "        5. Write new value to position seq_pos in value_cache using indexed assignment\n",
    "        6. Note: seq_pos is advanced externally via advance() after all layers\n",
    "\n",
    "        This is the core caching operation - efficiently append new K,V\n",
    "        to the cache without recomputation. This operation is O(1) because\n",
    "        it's just an indexed assignment.\n",
    "\n",
    "        IMPORTANT: KV caching is designed for INFERENCE (generation) only,\n",
    "        not training. During generation, gradients are not computed. If you\n",
    "        need gradients, don't use caching (use standard forward pass instead).\n",
    "\n",
    "        Args:\n",
    "            layer_idx: Which transformer layer (0 to num_layers-1)\n",
    "            key: New key tensor, shape (batch_size, num_heads, 1, head_dim)\n",
    "            value: New value tensor, shape (batch_size, num_heads, 1, head_dim)\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> cache = KVCache(batch_size=1, max_seq_len=10, num_layers=2,\n",
    "        ...                 num_heads=4, head_dim=64)\n",
    "        >>> new_k = Tensor(np.random.randn(1, 4, 1, 64))\n",
    "        >>> new_v = Tensor(np.random.randn(1, 4, 1, 64))\n",
    "        >>> cache.update(layer_idx=0, key=new_k, value=new_v)\n",
    "        >>> cache.seq_pos  # Still 0 (update doesn't advance position)\n",
    "        >>> cache.advance()\n",
    "        >>> cache.seq_pos  # Now 1\n",
    "\n",
    "        HINTS:\n",
    "        - Use slicing: cache[:, :, seq_pos:seq_pos+1, :] to write to position\n",
    "        - Use .data for direct NumPy access (no gradient tracking needed)\n",
    "        - Raise ValueError with helpful messages for invalid inputs\n",
    "        - This is an in-place operation (modifies cache, returns None)\n",
    "\n",
    "        Raises:\n",
    "            ValueError: If layer_idx is out of range or sequence is full\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        if layer_idx >= self.num_layers:\n",
    "            raise ValueError(f\"Layer index {layer_idx} >= num_layers {self.num_layers}\")\n",
    "\n",
    "        if self.seq_pos >= self.max_seq_len:\n",
    "            raise ValueError(f\"Sequence position {self.seq_pos} >= max_seq_len {self.max_seq_len}\")\n",
    "\n",
    "        # Get cache for this layer\n",
    "        key_cache, value_cache = self.caches[layer_idx]\n",
    "\n",
    "        # Update cache at current position (efficient O(1) write)\n",
    "        # Note: We use .data here because caching is inference-only (no gradients needed)\n",
    "        # This avoids gradient tracking overhead during generation\n",
    "        key_cache.data[:, :, self.seq_pos:self.seq_pos+1, :] = key.data\n",
    "        value_cache.data[:, :, self.seq_pos:self.seq_pos+1, :] = value.data\n",
    "\n",
    "        # Note: seq_pos is advanced externally via advance() after all layers process\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def get(self, layer_idx: int) -> Tuple[Tensor, Tensor]:\n",
    "        \"\"\"\n",
    "        Retrieve cached key-value pairs for attention computation.\n",
    "\n",
    "        TODO: Return only the valid cached portion for this layer\n",
    "\n",
    "        APPROACH:\n",
    "        1. Validate layer_idx is in range\n",
    "        2. Retrieve the (key_cache, value_cache) tuple for this layer\n",
    "        3. Calculate valid_len = seq_pos (number of tokens currently cached)\n",
    "        4. Slice key_cache to get [:, :, :valid_len, :] (only filled portion)\n",
    "        5. Slice value_cache to get [:, :, :valid_len, :] (only filled portion)\n",
    "        6. Wrap sliced data in new Tensor objects and return\n",
    "\n",
    "        Returns only the valid portion of the cache (up to current seq_pos).\n",
    "        This is O(1) because we're just slicing NumPy arrays (view, not copy).\n",
    "\n",
    "        IMPORTANT: Returns Tensors without gradient tracking since caching\n",
    "        is inference-only. The returned tensors can be used in attention\n",
    "        computation but won't propagate gradients backward.\n",
    "\n",
    "        Args:\n",
    "            layer_idx: Which transformer layer to get cache for\n",
    "\n",
    "        Returns:\n",
    "            (cached_keys, cached_values): Tensors shaped for attention\n",
    "            Keys: (batch_size, num_heads, seq_pos, head_dim)\n",
    "            Values: (batch_size, num_heads, seq_pos, head_dim)\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> cache = KVCache(batch_size=1, max_seq_len=100, num_layers=2,\n",
    "        ...                 num_heads=4, head_dim=64)\n",
    "        >>> # After processing 3 tokens\n",
    "        >>> cache.seq_pos = 3\n",
    "        >>> cached_k, cached_v = cache.get(layer_idx=0)\n",
    "        >>> cached_k.shape  # (1, 4, 3, 64) - only first 3 positions\n",
    "        >>> cached_v.shape  # (1, 4, 3, 64)\n",
    "\n",
    "        HINTS:\n",
    "        - valid_len = self.seq_pos (how many tokens have been cached so far)\n",
    "        - Use slicing: cache.data[:, :, :valid_len, :] to get valid portion\n",
    "        - Wrap result in Tensor() for consistency with TinyTorch API\n",
    "        - If seq_pos=0, returns empty cache (shape with 0 in sequence dimension)\n",
    "\n",
    "        Raises:\n",
    "            ValueError: If layer_idx is out of range\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        if layer_idx >= self.num_layers:\n",
    "            raise ValueError(f\"Layer index {layer_idx} >= num_layers {self.num_layers}\")\n",
    "\n",
    "        # Get cache for this layer\n",
    "        key_cache, value_cache = self.caches[layer_idx]\n",
    "\n",
    "        # Return only the valid portion (up to current sequence position)\n",
    "        # seq_pos tracks where to write next, so we have seq_pos valid tokens\n",
    "        valid_len = self.seq_pos\n",
    "\n",
    "        # Note: Creating new Tensors from .data (no gradient tracking)\n",
    "        # This is correct for inference-only caching\n",
    "        cached_keys = Tensor(key_cache.data[:, :, :valid_len, :])\n",
    "        cached_values = Tensor(value_cache.data[:, :, :valid_len, :])\n",
    "\n",
    "        return cached_keys, cached_values\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def advance(self) -> None:\n",
    "        \"\"\"\n",
    "        Advance sequence position after processing current token.\n",
    "\n",
    "        Call this after all layers have processed the current token and\n",
    "        updated their caches. This moves the write pointer forward.\n",
    "        \"\"\"\n",
    "        self.seq_pos += 1\n",
    "\n",
    "    def reset(self) -> None:\n",
    "        \"\"\"\n",
    "        Reset cache for new generation sequence.\n",
    "\n",
    "        Call this when starting a new generation (new prompt).\n",
    "        Resets the sequence position counter and optionally zeros cache data.\n",
    "        \"\"\"\n",
    "        self.seq_pos = 0\n",
    "\n",
    "        # Zero out caches for clean state (helps with debugging)\n",
    "        for layer_idx in range(self.num_layers):\n",
    "            key_cache, value_cache = self.caches[layer_idx]\n",
    "            key_cache.data.fill(0.0)\n",
    "            value_cache.data.fill(0.0)\n",
    "\n",
    "    def get_memory_usage(self) -> Dict[str, float]:\n",
    "        \"\"\"\n",
    "        Calculate memory usage of the cache system.\n",
    "\n",
    "        Returns:\n",
    "            Dictionary with memory statistics in MB\n",
    "        \"\"\"\n",
    "        # Calculate size of one cache tensor\n",
    "        cache_size = self.batch_size * self.num_heads * self.max_seq_len * self.head_dim\n",
    "        bytes_per_float = 4  # float32\n",
    "\n",
    "        # Each layer has key_cache + value_cache\n",
    "        total_cache_tensors = self.num_layers * 2\n",
    "        total_elements = cache_size * total_cache_tensors\n",
    "        total_bytes = total_elements * bytes_per_float\n",
    "        total_mb = total_bytes / (1024 * 1024)\n",
    "\n",
    "        return {\n",
    "            'total_mb': total_mb,\n",
    "            'per_layer_mb': total_mb / self.num_layers,\n",
    "            'cache_tensors': total_cache_tensors,\n",
    "            'total_elements': total_elements\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63c67a40",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Unit Test: KVCache Implementation\n",
    "\n",
    "Let's test that our cache correctly stores and retrieves key-value pairs across multiple layers and sequence positions.\n",
    "\n",
    "**This is a unit test** - it tests the KVCache class in isolation with simulated attention keys and values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "553ced7f",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-kvcache",
     "locked": true,
     "points": 10
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_kvcache():\n",
    "    \"\"\"🔬 Unit Test: KVCache Implementation\"\"\"\n",
    "    print(\"🔬 Unit Test: KVCache Implementation...\")\n",
    "\n",
    "    # Test parameters (small transformer for testing)\n",
    "    batch_size, max_seq_len = 2, 8\n",
    "    num_layers, num_heads, head_dim = 3, 4, 16\n",
    "\n",
    "    # Create cache\n",
    "    cache = KVCache(batch_size, max_seq_len, num_layers, num_heads, head_dim)\n",
    "\n",
    "    # Test 1: Initial state\n",
    "    assert cache.seq_pos == 0, \"Cache should start at position 0\"\n",
    "    mem_usage = cache.get_memory_usage()\n",
    "    assert mem_usage['total_mb'] > 0, \"Cache should have non-zero memory usage\"\n",
    "    print(f\"   Cache initialized: {mem_usage['total_mb']:.2f} MB\")\n",
    "\n",
    "    # Test 2: Single token update and retrieval\n",
    "    key1 = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "    value1 = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "\n",
    "    # Update layer 0 with first token\n",
    "    cache.update(0, key1, value1)\n",
    "\n",
    "    # Before advance, get() should return empty (seq_pos=0)\n",
    "    cached_k, cached_v = cache.get(0)\n",
    "    assert cached_k.shape == (batch_size, num_heads, 0, head_dim), \"Before advance, cache should be empty\"\n",
    "\n",
    "    # Advance position\n",
    "    cache.advance()\n",
    "\n",
    "    # Now cache should have 1 token\n",
    "    cached_k, cached_v = cache.get(0)\n",
    "    assert cached_k.shape == (batch_size, num_heads, 1, head_dim), f\"Expected shape (2,4,1,16), got {cached_k.shape}\"\n",
    "    assert cached_v.shape == (batch_size, num_heads, 1, head_dim), f\"Expected shape (2,4,1,16), got {cached_v.shape}\"\n",
    "\n",
    "    # Test 3: Multi-token sequence\n",
    "    key2 = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "    value2 = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "    cache.update(0, key2, value2)\n",
    "    cache.advance()\n",
    "\n",
    "    cached_k, cached_v = cache.get(0)\n",
    "    assert cached_k.shape == (batch_size, num_heads, 2, head_dim), \"Should have 2 tokens cached\"\n",
    "    assert cached_v.shape == (batch_size, num_heads, 2, head_dim), \"Should have 2 tokens cached\"\n",
    "\n",
    "    # Test 4: Multiple layers\n",
    "    cache.reset()\n",
    "    key_test = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "    value_test = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "\n",
    "    # Update all layers with same token\n",
    "    cache.update(0, key_test, value_test)  # Layer 0\n",
    "    cache.update(1, key_test, value_test)  # Layer 1\n",
    "    cache.update(2, key_test, value_test)  # Layer 2\n",
    "    cache.advance()\n",
    "\n",
    "    # Each layer should have the cached token\n",
    "    for layer_idx in range(num_layers):\n",
    "        cached_k, cached_v = cache.get(layer_idx)\n",
    "        assert cached_k.shape[2] == 1, f\"Layer {layer_idx} should have 1 token\"\n",
    "\n",
    "    # Test 5: Reset functionality\n",
    "    cache.reset()\n",
    "    assert cache.seq_pos == 0, \"Reset should clear sequence position\"\n",
    "    cached_k, cached_v = cache.get(0)\n",
    "    assert cached_k.shape == (batch_size, num_heads, 0, head_dim), \"Reset should clear cache\"\n",
    "\n",
    "    print(\"✅ KVCache implementation works correctly!\")\n",
    "\n",
    "# Run test immediately when developing this module\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_kvcache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f84f91ca",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🎯 Part 4: Enabling KV Caching for Model Generation\n",
    "\n",
    "### Integration Strategy\n",
    "\n",
    "Now we need a clean way to enable KV caching in our existing transformer models without breaking the existing code. We'll create an `enable_kv_cache()` function that:\n",
    "\n",
    "1. Creates a KVCache instance sized for the model\n",
    "2. Returns a flag to indicate caching is enabled\n",
    "3. Can be called before generation starts\n",
    "\n",
    "The actual integration with attention will happen in the milestone code where we:\n",
    "1. Check if cache is enabled\n",
    "2. Only compute K,V for new token (not all tokens)\n",
    "3. Update cache with new K,V\n",
    "4. Use cached K,V for attention computation\n",
    "\n",
    "### Generation Flow Comparison\n",
    "\n",
    "```\n",
    "Without Cache (Current):\n",
    "for each new token:\n",
    "    input_seq = [all tokens so far]        # Length grows: 1, 2, 3, ...\n",
    "    logits = model.forward(input_seq)       # Recomputes everything!\n",
    "    next_token = sample(logits[-1])\n",
    "    append next_token\n",
    "\n",
    "With Cache (New):\n",
    "cache = enable_kv_cache(model)\n",
    "for each new token:\n",
    "    input_token = [just new token]          # Length always 1\n",
    "    logits = model.forward_cached(input_token, cache)  # Only new computation\n",
    "    next_token = sample(logits[-1])\n",
    "    append next_token\n",
    "```\n",
    "\n",
    "**Key Difference**: Input changes from growing sequence to single token, with cache providing history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc4b9e1",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "def enable_kv_cache(batch_size: int, max_seq_len: int, num_layers: int,\n",
    "                    num_heads: int, head_dim: int) -> KVCache:\n",
    "    \"\"\"\n",
    "    Create and return a KVCache instance for model generation.\n",
    "    \n",
    "    This function creates a properly sized cache for the model architecture.\n",
    "    Call this before starting generation, then pass the cache to your\n",
    "    generation loop.\n",
    "\n",
    "    Args:\n",
    "        batch_size: Number of sequences to generate simultaneously\n",
    "        max_seq_len: Maximum sequence length to support\n",
    "        num_layers: Number of transformer layers in model\n",
    "        num_heads: Number of attention heads per layer\n",
    "        head_dim: Dimension per attention head (usually embed_dim // num_heads)\n",
    "\n",
    "    Returns:\n",
    "        KVCache instance ready for use\n",
    "    \n",
    "    Example:\n",
    "        ```python\n",
    "        # Enable caching for generation\n",
    "        cache = enable_kv_cache(\n",
    "            batch_size=1,\n",
    "            max_seq_len=100,\n",
    "            num_layers=4,\n",
    "            num_heads=4,\n",
    "            head_dim=32\n",
    "        )\n",
    "        \n",
    "        # Use in generation loop (pseudocode)\n",
    "        for step in range(max_new_tokens):\n",
    "            # Only process new token with cache\n",
    "            logits = model.forward_cached(new_token, cache)\n",
    "            next_token = sample(logits)\n",
    "        ```\n",
    "    \"\"\"\n",
    "    cache = KVCache(batch_size, max_seq_len, num_layers, num_heads, head_dim)\n",
    "    \n",
    "    print(f\"⚡ KV Cache enabled:\")\n",
    "    print(f\"   Batch size: {batch_size}\")\n",
    "    print(f\"   Max sequence: {max_seq_len}\")\n",
    "    print(f\"   Layers: {num_layers}\")\n",
    "    print(f\"   Heads: {num_heads}\")\n",
    "    print(f\"   Head dim: {head_dim}\")\n",
    "    \n",
    "    mem_info = cache.get_memory_usage()\n",
    "    print(f\"   Memory: {mem_info['total_mb']:.2f} MB\")\n",
    "    print()\n",
    "    \n",
    "    return cache"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd144e88",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Unit Test: Cache Enablement\n",
    "\n",
    "Let's verify that we can create caches for realistic model configurations.\n",
    "\n",
    "**This is a unit test** - it tests the cache creation and memory calculation for different model sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ea3206",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-cache-enablement",
     "locked": true,
     "points": 10
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_cache_enablement():\n",
    "    \"\"\"🔬 Unit Test: Cache Enablement for Different Models\"\"\"\n",
    "    print(\"🔬 Unit Test: Cache Enablement for Different Models...\")\n",
    "\n",
    "    # Test 1: Small model (fast generation)\n",
    "    print(\"   Test 1: Small Model (Tiny Transformer)\")\n",
    "    cache_small = KVCache(\n",
    "        batch_size=1,\n",
    "        max_seq_len=64,\n",
    "        num_layers=2,\n",
    "        num_heads=4,\n",
    "        head_dim=32\n",
    "    )\n",
    "    mem_small = cache_small.get_memory_usage()\n",
    "    assert mem_small['total_mb'] < 1.0, \"Small model should use < 1 MB\"\n",
    "    print(f\"   Small model cache: {mem_small['total_mb']:.3f} MB\")\n",
    "\n",
    "    # Test 2: Medium model (balanced performance)\n",
    "    print(\"   Test 2: Medium Model (Standard Transformer)\")\n",
    "    cache_medium = KVCache(\n",
    "        batch_size=1,\n",
    "        max_seq_len=128,\n",
    "        num_layers=4,\n",
    "        num_heads=8,\n",
    "        head_dim=64\n",
    "    )\n",
    "    mem_medium = cache_medium.get_memory_usage()\n",
    "    assert 1.0 < mem_medium['total_mb'] < 10.0, \"Medium model should use 1-10 MB\"\n",
    "    print(f\"   Medium model cache: {mem_medium['total_mb']:.3f} MB\")\n",
    "\n",
    "    # Test 3: Batch inference (multiple sequences)\n",
    "    print(\"   Test 3: Batch Inference (4 sequences)\")\n",
    "    cache_batch = KVCache(\n",
    "        batch_size=4,  # Generate 4 sequences in parallel\n",
    "        max_seq_len=64,\n",
    "        num_layers=2,\n",
    "        num_heads=4,\n",
    "        head_dim=32\n",
    "    )\n",
    "    mem_batch = cache_batch.get_memory_usage()\n",
    "    assert mem_batch['total_mb'] > mem_small['total_mb'], \"Batch cache should be larger\"\n",
    "    print(f\"   Batch cache: {mem_batch['total_mb']:.3f} MB (4x batch size)\")\n",
    "\n",
    "    print(\"✅ Cache enablement works correctly!\")\n",
    "\n",
    "# Run test immediately when developing this module\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_cache_enablement()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f454d7a9",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 Part 5: Using KV Cache in Practice\n",
    "\n",
    "### Practical Integration Checklist\n",
    "\n",
    "To use KV caching in your transformer generation:\n",
    "\n",
    "**✅ Before Generation:**\n",
    "1. Create cache with `enable_kv_cache()`\n",
    "2. Set cache dimensions to match your model architecture\n",
    "3. Verify memory usage is acceptable\n",
    "\n",
    "**✅ During Generation (Modified Forward Pass):**\n",
    "1. For the first token (prompt), process normally and populate cache\n",
    "2. For subsequent tokens:\n",
    "   - Only process the NEW token (not entire sequence)\n",
    "   - Update cache with new K,V pairs\n",
    "   - Retrieve full cached K,V for attention\n",
    "   - Use cached values in attention computation\n",
    "   - Advance cache position after all layers\n",
    "\n",
    "**✅ After Generation:**\n",
    "1. Reset cache if generating another sequence\n",
    "2. Monitor memory usage for production deployment\n",
    "\n",
    "### Performance Expectations\n",
    "\n",
    "```\n",
    "Expected Speedup by Sequence Length:\n",
    "┌───────────┬──────────┬───────────┬──────────┐\n",
    "│ Seq Len   │ No Cache │ With Cache│ Speedup  │\n",
    "├───────────┼──────────┼───────────┼──────────┤\n",
    "│  10 tokens│ ~80 tok/s│ ~600 tok/s│   7.5x   │\n",
    "│  25 tokens│ ~40 tok/s│ ~500 tok/s│  12.5x   │\n",
    "│  50 tokens│ ~25 tok/s│ ~400 tok/s│  16.0x   │\n",
    "│ 100 tokens│ ~12 tok/s│ ~200 tok/s│  16.7x   │\n",
    "└───────────┴──────────┴───────────┴──────────┘\n",
    "\n",
    "Key Insight: Speedup increases with sequence length!\n",
    "Why? Longer sequences = more redundant computation without cache.\n",
    "```\n",
    "\n",
    "### Production Considerations\n",
    "\n",
    "**Memory Management:**\n",
    "- Cache memory = `batch_size × num_layers × num_heads × max_seq_len × head_dim × 4 bytes`\n",
    "- For GPT-2 (12 layers, 12 heads, seq_len=1024, head_dim=64): ~37 MB per sequence\n",
    "- For GPT-3 (96 layers, 96 heads, seq_len=2048, head_dim=128): ~4.7 GB per sequence\n",
    "\n",
    "**Trade-off Analysis:**\n",
    "- **10x+ speedup** for typical generation lengths (50-200 tokens)\n",
    "- **Modest memory cost** compared to model parameters (often <1% of model size)\n",
    "- **Enables real-time interaction** that's impossible without caching\n",
    "\n",
    "**Best Practices:**\n",
    "1. Always use caching for production serving\n",
    "2. Tune `max_seq_len` to expected generation length (don't over-allocate)\n",
    "3. Consider batch inference to amortize model loading costs\n",
    "4. Monitor cache memory usage in production"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d10b23",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🎯 Part 5: Non-Invasive Integration with Existing Models\n",
    "\n",
    "### The Challenge\n",
    "\n",
    "We built KV caching in Module 14, but our transformer (Modules 12-13) doesn't know about it!\n",
    "\n",
    "**❌ BAD Solution**: Go back and modify Module 12 (MultiHeadAttention)\n",
    "- Breaks \"forward-only\" learning (students shouldn't revisit old modules)\n",
    "- Makes Module 12 depend on Module 14 (wrong dependency direction!)\n",
    "- Violates clean module boundaries\n",
    "\n",
    "**✅ GOOD Solution**: Module 14 ADDS caching to existing models without modification!\n",
    "- Use composition + monkey-patching (like `enable_autograd()`)\n",
    "- Module 14 wraps/enhances Module 12, not modifies it\n",
    "- Students learn systems engineering: \"Add capabilities, don't break old code\"\n",
    "\n",
    "### Implementation Strategy\n",
    "\n",
    "We'll create `enable_kv_cache(model)` that:\n",
    "1. Creates cache for the model's architecture\n",
    "2. Wraps each attention layer with caching logic\n",
    "3. Intercepts attention calls and manages cache automatically\n",
    "4. Returns the cache for manual control if needed\n",
    "\n",
    "This is **non-invasive enhancement** - a critical ML systems pattern!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44c5bdff",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "enable-kv-cache",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "def enable_kv_cache(model):\n",
    "    \"\"\"\n",
    "    Enable KV caching for a transformer model WITHOUT modifying Module 12/13 code.\n",
    "\n",
    "    TODO: Create cache and non-invasively patch attention layers\n",
    "\n",
    "    APPROACH:\n",
    "    1. Validate model has required attributes (embed_dim, num_layers, num_heads, max_seq_len, blocks)\n",
    "    2. Calculate head_dim from embed_dim and num_heads\n",
    "    3. Create KVCache instance sized for this model's architecture\n",
    "    4. Store cache on model as model._kv_cache and set model._cache_enabled flag\n",
    "    5. For each transformer block, wrap its attention forward method with caching logic\n",
    "    6. Print confirmation message with cache statistics\n",
    "    7. Return the cache object\n",
    "\n",
    "    This function demonstrates **non-invasive optimization** - adding capabilities\n",
    "    to existing systems without breaking them. Similar to how Module 05 (Autograd)\n",
    "    uses enable_autograd() to add gradient tracking to Tensors.\n",
    "\n",
    "    Args:\n",
    "        model: A GPT-style transformer model with:\n",
    "               - model.embed_dim (int)\n",
    "               - model.num_layers (int)\n",
    "               - model.num_heads (int)\n",
    "               - model.max_seq_len (int)\n",
    "               - model.blocks (list of TransformerBlock objects)\n",
    "\n",
    "    Returns:\n",
    "        cache: KVCache object for this model\n",
    "\n",
    "    EXAMPLE:\n",
    "    >>> from tinytorch.models.transformer import GPT\n",
    "    >>> model = GPT(vocab_size=100, embed_dim=128, num_layers=4, num_heads=4)\n",
    "    >>> cache = enable_kv_cache(model)\n",
    "    >>> hasattr(model, '_kv_cache')  # True\n",
    "    >>> model._cache_enabled  # True\n",
    "    >>> cache.num_layers  # 4 (matches model)\n",
    "\n",
    "    HINTS:\n",
    "    - Use hasattr() to validate model attributes exist\n",
    "    - head_dim = model.embed_dim // model.num_heads\n",
    "    - Store cache on model with model._kv_cache = cache\n",
    "    - Set flag with model._cache_enabled = True\n",
    "    - Save original forward with block._original_attention_forward\n",
    "    - Use a factory function to create patched forwards (closure captures layer_idx)\n",
    "\n",
    "    Pedagogical Note:\n",
    "        This teaches students that optimizations can be LAYERED on top of\n",
    "        working systems. Module 14 doesn't break Modules 12-13; it enhances them!\n",
    "    \"\"\"\n",
    "    ### BEGIN SOLUTION\n",
    "    import types\n",
    "\n",
    "    # Validate model has required attributes\n",
    "    required_attrs = ['embed_dim', 'num_layers', 'num_heads', 'max_seq_len', 'blocks']\n",
    "    for attr in required_attrs:\n",
    "        if not hasattr(model, attr):\n",
    "            raise AttributeError(\n",
    "                f\"Model missing '{attr}' - enable_kv_cache() requires a GPT-style model \"\n",
    "                f\"with {', '.join(required_attrs)}\"\n",
    "            )\n",
    "\n",
    "    # Calculate head dimension\n",
    "    head_dim = model.embed_dim // model.num_heads\n",
    "    if model.embed_dim % model.num_heads != 0:\n",
    "        raise ValueError(\n",
    "            f\"embed_dim ({model.embed_dim}) must be divisible by num_heads ({model.num_heads})\"\n",
    "        )\n",
    "\n",
    "    # Create cache for this model\n",
    "    cache = KVCache(\n",
    "        batch_size=1,  # Default to single sequence; can be reset for batch inference\n",
    "        max_seq_len=model.max_seq_len,\n",
    "        num_layers=model.num_layers,\n",
    "        num_heads=model.num_heads,\n",
    "        head_dim=head_dim\n",
    "    )\n",
    "\n",
    "    # Store cache on model for easy access\n",
    "    model._kv_cache = cache\n",
    "    model._cache_enabled = True\n",
    "\n",
    "    # Patch each transformer block's attention\n",
    "    for layer_idx, block in enumerate(model.blocks):\n",
    "        # Store original attention forward method\n",
    "        if not hasattr(block, '_original_attention_forward'):\n",
    "            block._original_attention_forward = block.attention.forward\n",
    "\n",
    "        # Create cached version\n",
    "        def make_cached_forward(layer_idx, original_forward, cache_obj):\n",
    "            \"\"\"Factory to create cached forward with correct layer_idx closure\"\"\"\n",
    "            def cached_forward(x, mask=None):\n",
    "                \"\"\"\n",
    "                Cached attention forward pass with REAL speedup!\n",
    "                \n",
    "                PATH SELECTION STRATEGY (Key to Understanding KV Caching):\n",
    "                ──────────────────────────────────────────────────────────\n",
    "                \n",
    "                We have THREE possible paths through attention:\n",
    "                \n",
    "                1️⃣ TRAINING PATH (seq_len > 1):\n",
    "                   - Input: Full sequence of tokens (e.g., 64 tokens)\n",
    "                   - Action: Use ORIGINAL attention (no caching)\n",
    "                   - Why: Need full gradient flow for backpropagation\n",
    "                   - Complexity: O(n²) but that's fine for training\n",
    "                   - Example: x.shape = (batch=1, seq=64, embed=128)\n",
    "                \n",
    "                2️⃣ FIRST TOKEN PATH (seq_len == 1 AND cache empty):\n",
    "                   - Input: Single token (the first one in generation)\n",
    "                   - Action: Use ORIGINAL attention (initialize cache)\n",
    "                   - Why: Cache is empty, nothing to retrieve yet\n",
    "                   - Complexity: O(1) - only one token\n",
    "                   - Example: x.shape = (batch=1, seq=1, embed=128), cache.seq_pos=0\n",
    "                \n",
    "                3️⃣ CACHED GENERATION PATH (seq_len == 1 AND cache populated):\n",
    "                   - Input: Single NEW token (during generation)\n",
    "                   - Action: Compute K,V for new token ONLY, retrieve history from cache\n",
    "                   - Why: This is where the speedup happens! O(n²) → O(n)\n",
    "                   - Complexity: O(n) - only compute for new token, reuse cache\n",
    "                   - Example: x.shape = (batch=1, seq=1, embed=128), cache.seq_pos=5\n",
    "                \n",
    "                \n",
    "                WHY .data INSTEAD OF TENSOR OPERATIONS?\n",
    "                ────────────────────────────────────────\n",
    "                \n",
    "                In the cached path, we use numpy via .data for three reasons:\n",
    "                \n",
    "                1. **Explicit Intent**: Makes it crystal clear this is inference-only\n",
    "                   - Training: Uses Tensor operations → gradients tracked\n",
    "                   - Inference: Uses .data → no gradient overhead\n",
    "                \n",
    "                2. **Performance**: Avoids any autograd bookkeeping\n",
    "                   - Even if small, every bit counts in generation\n",
    "                   - Production LLMs (vLLM, llama.cpp) use similar patterns\n",
    "                \n",
    "                3. **Educational Clarity**: Shows students the distinction\n",
    "                   - \"When do I need gradients?\" (training)\n",
    "                   - \"When can I skip them?\" (inference)\n",
    "                \n",
    "                We COULD use Tensor operations with requires_grad=False, but .data\n",
    "                is more explicit and is the industry-standard pattern.\n",
    "                \n",
    "                \n",
    "                THE O(n²) → O(n) TRANSFORMATION:\n",
    "                ─────────────────────────────────\n",
    "                \n",
    "                WITHOUT Cache (Standard Attention):\n",
    "                  Step 1: Process token 1  → Compute attention for 1 token  (1²  = 1 op)\n",
    "                  Step 2: Process tokens 1-2 → Compute attention for 2 tokens (2²  = 4 ops)\n",
    "                  Step 3: Process tokens 1-3 → Compute attention for 3 tokens (3²  = 9 ops)\n",
    "                  ...\n",
    "                  Step N: Process tokens 1-N → Compute attention for N tokens (N² ops)\n",
    "                  \n",
    "                  Total: 1 + 4 + 9 + ... + N² = O(N³) across all steps!\n",
    "                \n",
    "                WITH Cache (Our Implementation):\n",
    "                  Step 1: Process token 1 → Compute K,V for token 1, cache it      (1 op)\n",
    "                  Step 2: Process token 2 → Compute K,V for token 2, retrieve 1    (2 ops)\n",
    "                  Step 3: Process token 3 → Compute K,V for token 3, retrieve 1-2  (3 ops)\n",
    "                  ...\n",
    "                  Step N: Process token N → Compute K,V for token N, retrieve 1-(N-1) (N ops)\n",
    "                  \n",
    "                  Total: 1 + 2 + 3 + ... + N = O(N²) across all steps!\n",
    "                \n",
    "                That's why we see 5-7x speedup on short sequences, and 10-15x on longer ones!\n",
    "                \"\"\"\n",
    "                from tinytorch.core.tensor import Tensor\n",
    "                import numpy as np\n",
    "                \n",
    "                seq_len = x.shape[1]\n",
    "                \n",
    "                # ═══════════════════════════════════════════════════════════════\n",
    "                # PATH SELECTION: Choose between training, first token, or cached\n",
    "                # ═══════════════════════════════════════════════════════════════\n",
    "                \n",
    "                # PATH 1: TRAINING (seq_len > 1)\n",
    "                # ───────────────────────────────────\n",
    "                # Input is a full sequence (e.g., 64 tokens during training)\n",
    "                # We MUST use original attention to preserve gradient flow\n",
    "                # No caching during training - we need backprop through everything\n",
    "                if seq_len > 1:\n",
    "                    return original_forward(x, mask)  # O(n²) but preserves gradients\n",
    "                \n",
    "                # PATH 2: FIRST TOKEN (seq_len == 1, cache empty)\n",
    "                # ────────────────────────────────────────────────\n",
    "                # This is the very first token in generation (cache.seq_pos == 0)\n",
    "                # Cache is empty, so there's nothing to retrieve yet\n",
    "                # Use original attention to process this token, which will populate cache\n",
    "                if cache_obj.seq_pos == 0:\n",
    "                    return original_forward(x, mask)  # O(1) - just one token\n",
    "                \n",
    "                # PATH 3: CACHED GENERATION (seq_len == 1, cache populated)\n",
    "                # ──────────────────────────────────────────────────────────\n",
    "                # This is a NEW token during generation (cache has history)\n",
    "                # We can now use the cache for massive speedup!\n",
    "                # Compute K,V for ONLY this new token, retrieve cached history\n",
    "                \n",
    "                # Get attention layer (assumes block.attention has the attention object)\n",
    "                attention = block.attention\n",
    "                \n",
    "                # Step 1: Compute Q, K, V for NEW token only\n",
    "                # Access the linear projection layers\n",
    "                Q_new = attention.q_proj.forward(x)  # (batch, 1, embed_dim)\n",
    "                K_new = attention.k_proj.forward(x)  # (batch, 1, embed_dim)\n",
    "                V_new = attention.v_proj.forward(x)  # (batch, 1, embed_dim)\n",
    "                \n",
    "                # Step 2: Reshape to multi-head format\n",
    "                batch_size = x.shape[0]\n",
    "                num_heads = attention.num_heads\n",
    "                head_dim = attention.head_dim\n",
    "                \n",
    "                # Reshape: (batch, 1, embed_dim) → (batch, num_heads, 1, head_dim)\n",
    "                Q_heads = Q_new.reshape(batch_size, 1, num_heads, head_dim)\n",
    "                Q_heads = Tensor(np.transpose(Q_heads.data, (0, 2, 1, 3)))  # (batch, num_heads, 1, head_dim)\n",
    "                \n",
    "                K_heads = K_new.reshape(batch_size, 1, num_heads, head_dim)\n",
    "                K_heads = Tensor(np.transpose(K_heads.data, (0, 2, 1, 3)))\n",
    "                \n",
    "                V_heads = V_new.reshape(batch_size, 1, num_heads, head_dim)\n",
    "                V_heads = Tensor(np.transpose(V_heads.data, (0, 2, 1, 3)))\n",
    "                \n",
    "                # Step 3: Update cache with new K, V (using .data for performance)\n",
    "                cache_obj.update(layer_idx, K_heads, V_heads)\n",
    "                \n",
    "                # Step 4: Retrieve ALL cached K, V (includes history + new token)\n",
    "                K_all, V_all = cache_obj.get(layer_idx)\n",
    "                \n",
    "                # Step 5: Compute attention using new Q with ALL cached K, V\n",
    "                # ─────────────────────────────────────────────────────────\n",
    "                # Scaled dot-product attention: softmax(Q @ K^T / sqrt(d_k)) @ V\n",
    "                #\n",
    "                # NOTE: We use .data (numpy arrays) here instead of Tensor operations\n",
    "                # Why? This is INFERENCE-ONLY code (no gradients needed):\n",
    "                # - Explicit: Makes it clear this is inference, not training\n",
    "                # - Fast: Avoids autograd overhead (even if small)\n",
    "                # - Standard: Production LLMs (vLLM, llama.cpp) do the same\n",
    "                #\n",
    "                # If this were training, we'd use Tensor operations for gradient flow.\n",
    "                # But in generation (inference), .data is the right choice.\n",
    "                \n",
    "                # Q @ K^T: (batch, num_heads, 1, head_dim) @ (batch, num_heads, head_dim, seq_len)\n",
    "                #        → (batch, num_heads, 1, seq_len)\n",
    "                K_transposed = np.transpose(K_all.data, (0, 1, 3, 2))  # .data = numpy array\n",
    "                scores = np.matmul(Q_heads.data, K_transposed)  # Pure numpy matmul\n",
    "                \n",
    "                # Scale by sqrt(head_dim)\n",
    "                scores = scores / np.sqrt(head_dim)\n",
    "                \n",
    "                # Apply mask if provided (causal mask for generation)\n",
    "                if mask is not None:\n",
    "                    # Mask should be (1, 1, 1, seq_len) for this token\n",
    "                    # In generation, we can attend to all previous tokens\n",
    "                    pass  # No masking needed in generation (we see all history)\n",
    "                \n",
    "                # Softmax over key dimension\n",
    "                scores_max = np.max(scores, axis=-1, keepdims=True)\n",
    "                exp_scores = np.exp(scores - scores_max)\n",
    "                attention_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)\n",
    "                \n",
    "                # Apply attention weights to values\n",
    "                # (batch, num_heads, 1, seq_len) @ (batch, num_heads, seq_len, head_dim)\n",
    "                # → (batch, num_heads, 1, head_dim)\n",
    "                attention_output = np.matmul(attention_weights, V_all.data)\n",
    "                \n",
    "                # Step 6: Reshape back and apply output projection\n",
    "                # (batch, num_heads, 1, head_dim) → (batch, 1, num_heads, head_dim)\n",
    "                attention_output_transposed = np.transpose(attention_output, (0, 2, 1, 3))\n",
    "                \n",
    "                # Concatenate heads: (batch, 1, num_heads * head_dim)\n",
    "                concat_data = attention_output_transposed.reshape(batch_size, 1, num_heads * head_dim)\n",
    "                concat_output = Tensor(concat_data)\n",
    "                \n",
    "                # Output projection\n",
    "                output = attention.out_proj.forward(concat_output)\n",
    "                \n",
    "                return output\n",
    "            \n",
    "            return cached_forward\n",
    "\n",
    "        # Patch this block's attention\n",
    "        block.attention.forward = make_cached_forward(layer_idx, block._original_attention_forward, cache)\n",
    "\n",
    "    print(f\"⚡ KV Cache enabled for model!\")\n",
    "    print(f\"   Architecture: {model.num_layers} layers × {model.num_heads} heads × {head_dim}D\")\n",
    "    print(f\"   Memory: {cache.get_memory_usage()['total_mb']:.2f} MB\")\n",
    "    print(f\"   Cache stored in: model._kv_cache\")\n",
    "    print()\n",
    "    print(f\"💡 To disable: call disable_kv_cache(model)\")\n",
    "    print()\n",
    "\n",
    "    return cache\n",
    "    ### END SOLUTION\n",
    "\n",
    "\n",
    "#| export  \n",
    "def disable_kv_cache(model):\n",
    "    \"\"\"\n",
    "    Disable KV caching and restore original attention behavior.\n",
    "    \n",
    "    Args:\n",
    "        model: Model with caching enabled\n",
    "    \n",
    "    Example:\n",
    "        ```python\n",
    "        cache = enable_kv_cache(model)\n",
    "        # ... do cached generation ...\n",
    "        disable_kv_cache(model)  # Back to normal\n",
    "        ```\n",
    "    \"\"\"\n",
    "    if not hasattr(model, '_cache_enabled') or not model._cache_enabled:\n",
    "        print(\"⚠️  KV cache not enabled on this model\")\n",
    "        return\n",
    "    \n",
    "    # Restore original attention forwards\n",
    "    for block in model.blocks:\n",
    "        if hasattr(block, '_original_attention_forward'):\n",
    "            block.attention.forward = block._original_attention_forward\n",
    "    \n",
    "    # Clean up\n",
    "    model._cache_enabled = False\n",
    "    if hasattr(model, '_kv_cache'):\n",
    "        delattr(model, '_kv_cache')\n",
    "    \n",
    "    print(\"✓ KV cache disabled, original attention restored\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ea98b51",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Unit Test: Non-Invasive Cache Integration\n",
    "\n",
    "Let's verify that `enable_kv_cache()` works without breaking the model!\n",
    "\n",
    "**This is an integration test** - it tests Module 14 enhancing Modules 12-13 without modification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87a4e516",
   "metadata": {
    "lines_to_next_cell": 2,
    "nbgrader": {
     "grade": true,
     "grade_id": "test-noninvasive",
     "locked": true,
     "points": 10
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_noninvasive_integration():\n",
    "    \"\"\"🔬 Unit Test: Non-Invasive Cache Integration\"\"\"\n",
    "    print(\"🔬 Unit Test: Non-Invasive Cache Integration...\")\n",
    "\n",
    "    # Create a mock transformer-like object for testing\n",
    "    class MockTransformerBlock:\n",
    "        def __init__(self):\n",
    "            self.attention = self\n",
    "\n",
    "        def forward(self, x):\n",
    "            # Simple pass-through for testing\n",
    "            return x\n",
    "\n",
    "    class MockGPT:\n",
    "        def __init__(self):\n",
    "            self.vocab_size = 100\n",
    "            self.embed_dim = 128\n",
    "            self.num_layers = 4\n",
    "            self.num_heads = 4\n",
    "            self.max_seq_len = 64\n",
    "            self.blocks = [MockTransformerBlock() for _ in range(self.num_layers)]\n",
    "\n",
    "    # Test 1: Enable caching\n",
    "    model = MockGPT()\n",
    "    print(\"   Test 1: Enable caching on model\")\n",
    "    cache = enable_kv_cache(model)\n",
    "    assert hasattr(model, '_kv_cache'), \"Model should have _kv_cache attribute\"\n",
    "    assert hasattr(model, '_cache_enabled'), \"Model should have _cache_enabled flag\"\n",
    "    assert model._cache_enabled == True, \"Cache should be enabled\"\n",
    "    assert cache is model._kv_cache, \"Returned cache should match model._kv_cache\"\n",
    "\n",
    "    # Test 2: Attention forward still works\n",
    "    print(\"   Test 2: Attention forward pass still works\")\n",
    "    test_input = Tensor(np.random.randn(1, 10, 128))\n",
    "    for block in model.blocks:\n",
    "        output = block.attention.forward(test_input)\n",
    "        assert output.shape == test_input.shape, \"Forward pass should preserve shape\"\n",
    "\n",
    "    # Test 3: Disable caching\n",
    "    print(\"   Test 3: Disable caching\")\n",
    "    disable_kv_cache(model)\n",
    "    assert model._cache_enabled == False, \"Cache should be disabled\"\n",
    "    assert not hasattr(model, '_kv_cache'), \"Cache object should be removed\"\n",
    "\n",
    "    # Test 4: Can re-enable\n",
    "    print(\"   Test 4: Re-enable caching\")\n",
    "    _ = enable_kv_cache(model)\n",
    "    assert model._cache_enabled == True, \"Cache should be re-enabled\"\n",
    "\n",
    "    print(\"✅ Non-invasive cache integration works correctly!\")\n",
    "\n",
    "# Run test immediately when developing this module\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_noninvasive_integration()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0326e8e",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🧪 Module Integration Test\n",
    "\n",
    "Final validation that everything works together correctly before module completion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef08eafe",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": true,
     "grade_id": "module-integration",
     "locked": true,
     "points": 20
    }
   },
   "outputs": [],
   "source": [
    "def test_module():\n",
    "    \"\"\"\n",
    "    Comprehensive test of entire KV Caching module functionality.\n",
    "\n",
    "    This final test runs before module summary to ensure:\n",
    "    - All unit tests pass\n",
    "    - Functions work together correctly\n",
    "    - Module is ready for integration with TinyTorch\n",
    "    \"\"\"\n",
    "    print(\"🧪 RUNNING MODULE INTEGRATION TEST\")\n",
    "    print(\"=\" * 50)\n",
    "    print()\n",
    "\n",
    "    # Run all unit tests\n",
    "    print(\"Running unit tests...\")\n",
    "    test_unit_kvcache()\n",
    "    print()\n",
    "    test_unit_cache_enablement()\n",
    "    print()\n",
    "    test_unit_noninvasive_integration()\n",
    "    print()\n",
    "\n",
    "    print(\"Running integration scenarios...\")\n",
    "    print()\n",
    "\n",
    "    # Integration Test: Complete KV Cache Workflow\n",
    "    print(\"🔬 Integration Test: Complete KV Cache Workflow...\")\n",
    "    batch_size, max_seq_len = 1, 128\n",
    "    num_layers, num_heads, head_dim = 4, 8, 64\n",
    "\n",
    "    cache = KVCache(batch_size, max_seq_len, num_layers, num_heads, head_dim)\n",
    "\n",
    "    # Simulate generation loop (processing multiple tokens)\n",
    "    for _ in range(5):\n",
    "        for layer_idx in range(num_layers):\n",
    "            # Simulate new key-value pairs\n",
    "            new_key = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "            new_value = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
    "\n",
    "            # Update cache\n",
    "            cache.update(layer_idx, new_key, new_value)\n",
    "\n",
    "        # Advance position after all layers processed\n",
    "        cache.advance()\n",
    "\n",
    "    # Verify cache state\n",
    "    assert cache.seq_pos == 5, f\"Expected seq_pos=5, got {cache.seq_pos}\"\n",
    "\n",
    "    # Verify retrieval\n",
    "    for layer_idx in range(num_layers):\n",
    "        cached_k, cached_v = cache.get(layer_idx)\n",
    "        assert cached_k.shape == (batch_size, num_heads, 5, head_dim)\n",
    "        assert cached_v.shape == (batch_size, num_heads, 5, head_dim)\n",
    "\n",
    "    print(\"✅ Complete KV cache workflow validated!\")\n",
    "    print()\n",
    "\n",
    "    # Integration Test: Memory Tracking\n",
    "    print(\"🔬 Integration Test: Memory Tracking...\")\n",
    "    mem_info = cache.get_memory_usage()\n",
    "    assert mem_info['total_mb'] > 0\n",
    "    assert mem_info['cache_tensors'] == num_layers * 2\n",
    "    print(f\"✅ Memory tracking: {mem_info['total_mb']:.2f} MB for {mem_info['cache_tensors']} tensors\")\n",
    "    print()\n",
    "\n",
    "    print(\"=\" * 50)\n",
    "    print(\"🎉 ALL TESTS PASSED! Module ready for export.\")\n",
    "    print(\"Run: tito module complete 14\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "736d019f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    test_module()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff0d2a86",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎓 Module 14 Complete!\n",
    "\n",
    "You've implemented KV caching - the critical optimization that makes production language models economically viable!\n",
    "\n",
    "### What You Built\n",
    "\n",
    "✅ **KVCache Class**: Efficient memory management for key-value pairs across layers\n",
    "✅ **O(1) Updates**: Fast cache updates without data copying\n",
    "✅ **Memory Tracking**: Understanding cache size and memory trade-offs\n",
    "✅ **Non-Invasive Integration**: `enable_kv_cache()` adds optimization WITHOUT breaking modules\n",
    "✅ **Production Patterns**: Integration strategy for real transformer models\n",
    "\n",
    "### Key Systems Engineering Lesson\n",
    "\n",
    "**Module 14 doesn't modify Modules 12-13 - it ENHANCES them!**\n",
    "\n",
    "This teaches the critical principle: **Add capabilities forward, never break backward.**\n",
    "- Old code keeps working (Module 12 unchanged)\n",
    "- New code adds optimization (Module 14 layers on top)\n",
    "- Clean separation of concerns (caching is separate from attention logic)\n",
    "\n",
    "### Performance Impact\n",
    "\n",
    "```\n",
    "Without Cache: O(n²) complexity → slow, expensive, impractical\n",
    "With Cache:    O(n) complexity  → fast, cheap, production-ready\n",
    "\n",
    "Real Impact: 10-15x speedup for typical generation!\n",
    "```\n",
    "\n",
    "### What's Next\n",
    "\n",
    "**Module 15 (Profiling)**: Now that you've seen a concrete optimization, learn how to systematically measure and find more optimizations using professional profiling tools.\n",
    "\n",
    "### Try It Yourself\n",
    "\n",
    "Run the chatbot milestone with and without caching:\n",
    "\n",
    "```bash\n",
    "# Without cache (slow - baseline)\n",
    "python milestones/05_2017_transformer/vaswani_chatgpt.py\n",
    "\n",
    "# With cache (fast - 10-15x speedup!)\n",
    "python milestones/05_2017_transformer/vaswani_chatgpt.py --use-cache\n",
    "```\n",
    "\n",
    "Watch the tokens/sec metric jump from ~40 to ~500! 🚀\n",
    "\n",
    "---\n",
    "\n",
    "**Congratulations! You've completed Module 14: KV Caching!**\n",
    "\n",
    "You now understand the optimization that makes ChatGPT, Claude, and all production LLMs possible. This is THE technique that transformed language models from research toys into products used by millions of people every day.\n",
    "\n",
    "**From Theory to Practice**: You've gone from O(n²) naive generation to O(n) optimized generation. This is real ML engineering!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
