mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-06-03 15:15:50 -05:00
de3b837bee
PROBLEM: - nbdev requires #| export directive on EACH cell to export when using # %% markers - Cell markers inside class definitions split classes across multiple cells - Only partial classes were being exported to tinytorch package - Missing matmul, arithmetic operations, and activation classes in exports SOLUTION: 1. Removed # %% cell markers INSIDE class definitions (kept classes as single units) 2. Added #| export to imports cell at top of each module 3. Added #| export before each exportable class definition in all 20 modules 4. Added __call__ method to Sigmoid for functional usage 5. Fixed numpy import (moved to module level from __init__) MODULES FIXED: - 01_tensor: Tensor class with all operations (matmul, arithmetic, shape ops) - 02_activations: Sigmoid, ReLU, Tanh, GELU, Softmax classes - 03_layers: Linear, Dropout classes - 04_losses: MSELoss, CrossEntropyLoss, BinaryCrossEntropyLoss classes - 05_autograd: Function, AddBackward, MulBackward, MatmulBackward, SumBackward - 06_optimizers: Optimizer, SGD, Adam, AdamW classes - 07_training: CosineSchedule, Trainer classes - 08_dataloader: Dataset, TensorDataset, DataLoader classes - 09_spatial: Conv2d, MaxPool2d, AvgPool2d, SimpleCNN classes - 10-20: All exportable classes in remaining modules TESTING: - Test functions use 'if __name__ == "__main__"' guards - Tests run in notebooks but NOT on import - Rosenblatt Perceptron milestone working perfectly RESULT: ✅ All 20 modules export correctly ✅ Perceptron (1957) milestone functional ✅ Clean separation: development (modules/source) vs package (tinytorch)
1624 lines
73 KiB
Plaintext
1624 lines
73 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "9f182460",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"# Module 14: KV Caching - Optimizing Autoregressive Generation\n",
|
||
"\n",
|
||
"Welcome to Module 14! You'll implement the critical optimization that makes production language models possible: Key-Value caching for 10x+ faster text generation.\n",
|
||
"\n",
|
||
"## 🔗 Prerequisites & Progress\n",
|
||
"**You've Built**: Complete transformer architecture with multi-head attention and text generation\n",
|
||
"**You'll Build**: Memory-efficient KV caching system that eliminates redundant computation\n",
|
||
"**You'll Enable**: Production-grade inference optimization and real-world serving capabilities\n",
|
||
"\n",
|
||
"**Connection Map**:\n",
|
||
"```\n",
|
||
"Transformers → KV Caching → Production Serving\n",
|
||
"(slow O(n²)) (fast O(n)) (real-world scale)\n",
|
||
"```\n",
|
||
"\n",
|
||
"## Learning Objectives\n",
|
||
"By the end of this module, you will:\n",
|
||
"1. Understand why autoregressive generation has O(n²) complexity without caching\n",
|
||
"2. Implement KVCache with efficient memory management and O(1) updates\n",
|
||
"3. Build cache-aware attention that reuses previously computed keys and values\n",
|
||
"4. Measure dramatic speedup gains and understand memory trade-offs\n",
|
||
"5. Connect to production optimization patterns used in real LLM serving\n",
|
||
"\n",
|
||
"Let's make inference blazingly fast!\n",
|
||
"\n",
|
||
"## 📦 Where This Code Lives in the Final Package\n",
|
||
"\n",
|
||
"**Learning Side:** You work in `modules/14_kvcaching/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, attention_with_cache\n",
|
||
"```\n",
|
||
"\n",
|
||
"**Why this matters:**\n",
|
||
"- **Learning:** Complete caching system in one focused module for deep understanding\n",
|
||
"- **Production:** Proper organization like Hugging Face's generation/ with all optimization components\n",
|
||
"- **Consistency:** All generation optimizations and cache management in generation.kv_cache\n",
|
||
"- **Integration:** Works seamlessly with transformers for complete inference optimization"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "13eddf26",
|
||
"metadata": {
|
||
"lines_to_next_cell": 1,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "imports",
|
||
"solution": true
|
||
}
|
||
},
|
||
"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",
|
||
"from dataclasses import dataclass\n",
|
||
"\n",
|
||
"# Import our TinyTorch components (Modules 01-13)\n",
|
||
"### BEGIN SOLUTION\n",
|
||
"# Note: In real implementation, these would import from previous modules\n",
|
||
"# For now, we'll implement minimal versions to focus on caching concepts\n",
|
||
"\n",
|
||
"class Tensor:\n",
|
||
" \"\"\"Minimal Tensor for KV Caching focus (from Module 01)\"\"\"\n",
|
||
" def __init__(self, data, requires_grad=False):\n",
|
||
" self.data = np.array(data)\n",
|
||
" self.shape = self.data.shape\n",
|
||
" self.requires_grad = requires_grad\n",
|
||
" self.grad = None\n",
|
||
"\n",
|
||
" def __getitem__(self, key):\n",
|
||
" return Tensor(self.data[key])\n",
|
||
"\n",
|
||
" def __setitem__(self, key, value):\n",
|
||
" if isinstance(value, Tensor):\n",
|
||
" self.data[key] = value.data\n",
|
||
" else:\n",
|
||
" self.data[key] = value\n",
|
||
"\n",
|
||
" def size(self, dim=None):\n",
|
||
" if dim is None:\n",
|
||
" return self.shape\n",
|
||
" return self.shape[dim]\n",
|
||
"\n",
|
||
" def view(self, *shape):\n",
|
||
" return Tensor(self.data.reshape(shape))\n",
|
||
"\n",
|
||
" def transpose(self, dim0, dim1):\n",
|
||
" axes = list(range(len(self.shape)))\n",
|
||
" axes[dim0], axes[dim1] = axes[dim1], axes[dim0]\n",
|
||
" return Tensor(np.transpose(self.data, axes))\n",
|
||
"\n",
|
||
" @staticmethod\n",
|
||
" def cat(tensors, dim=0):\n",
|
||
" \"\"\"Concatenate tensors along dimension\"\"\"\n",
|
||
" arrays = [t.data for t in tensors]\n",
|
||
" return Tensor(np.concatenate(arrays, axis=dim))\n",
|
||
"\n",
|
||
" @staticmethod\n",
|
||
" def zeros(*shape):\n",
|
||
" \"\"\"Create zero tensor\"\"\"\n",
|
||
" return Tensor(np.zeros(shape))\n",
|
||
"### END SOLUTION"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "bba4366b",
|
||
"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 1000-token sequence, this means **500,500 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": "db62451e",
|
||
"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) @ 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 Layout Visualization\n",
|
||
"\n",
|
||
"```\n",
|
||
"Traditional Approach (Recompute Everything):\n",
|
||
"Step 1: [K₁, V₁] ← Compute 1 pair\n",
|
||
"Step 2: [K₁, V₁, K₂, V₂] ← Compute 2 pairs (recompute K₁,V₁)\n",
|
||
"Step 3: [K₁, V₁, K₂, V₂, K₃, V₃] ← Compute 3 pairs (recompute all!)\n",
|
||
"\n",
|
||
"Cached Approach (Store and Reuse):\n",
|
||
"Step 1: [K₁, V₁] → Cache ← Compute 1, store 1\n",
|
||
"Step 2: Cache + [K₂, V₂] → Cache ← Compute 1, append 1\n",
|
||
"Step 3: Cache + [K₃, V₃] → Cache ← Compute 1, append 1\n",
|
||
"```\n",
|
||
"\n",
|
||
"**Trade-off**: Use O(seq_len × hidden_dim) memory to save O(seq_len²) computation."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "06a99e38",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 🏗️ Part 3: KVCache Class Design\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\n",
|
||
"4. **Dynamic updates**: Efficiently append new tokens without copying data\n",
|
||
"5. **Memory management**: Pre-allocate space to avoid dynamic resizing\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 Visualization\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": "44db7cf8",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "kv_cache_class",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "4f3d5157",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"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.\n",
|
||
"\n",
|
||
" TODO: Implement the complete caching system for production-speed inference\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Pre-allocate cache tensors with maximum sequence length\n",
|
||
" 2. Track current sequence position for efficient O(1) updates\n",
|
||
" 3. Provide update() method to append new K,V pairs without copying\n",
|
||
" 4. Provide get() method to retrieve cached values for attention\n",
|
||
" 5. Handle multiple layers and attention heads properly\n",
|
||
"\n",
|
||
" CACHE 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",
|
||
" MEMORY OPTIMIZATION:\n",
|
||
" - Pre-allocate maximum size to avoid dynamic resizing overhead\n",
|
||
" - Use efficient indexing for cache updates (no data copying)\n",
|
||
" - Store only essential data needed for attention computation\n",
|
||
"\n",
|
||
" HINTS:\n",
|
||
" - Use list of tuples: [(key_cache₀, value_cache₀), (key_cache₁, value_cache₁), ...]\n",
|
||
" - Track seq_pos to know where to write new values\n",
|
||
" - Consider batch dimension for efficient multi-sequence serving\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",
|
||
" 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",
|
||
" ### 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.zeros(batch_size, num_heads, max_seq_len, head_dim)\n",
|
||
" value_cache = Tensor.zeros(batch_size, num_heads, max_seq_len, head_dim)\n",
|
||
"\n",
|
||
" self.caches.append((key_cache, value_cache))\n",
|
||
"\n",
|
||
" # Track which positions are valid (for debugging and masking)\n",
|
||
" self.valid_positions = Tensor.zeros(batch_size, max_seq_len)\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 the cache without recomputation\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Get current cache for the specified layer\n",
|
||
" 2. Write new key,value at current sequence position (O(1) operation)\n",
|
||
" 3. Mark position as valid for attention masking\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",
|
||
" PERFORMANCE NOTE:\n",
|
||
" This operation should be O(1) - just indexing assignment, no large array copying\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",
|
||
" # Remove the sequence dimension since we're writing to a specific position\n",
|
||
" key_cache[:, :, self.seq_pos:self.seq_pos+1, :] = key\n",
|
||
" value_cache[:, :, self.seq_pos:self.seq_pos+1, :] = value\n",
|
||
"\n",
|
||
" # Mark this position as valid for attention\n",
|
||
" self.valid_positions[:, self.seq_pos] = 1.0\n",
|
||
"\n",
|
||
" # Note: seq_pos is advanced externally via advance() after all layers process the token\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 the cached K,V up to current sequence position\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Get cache for specified layer\n",
|
||
" 2. Slice to current sequence position (don't return unused space)\n",
|
||
" 3. Return properly shaped tensors for attention\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+1, head_dim)\n",
|
||
" Values: (batch_size, num_heads, seq_pos+1, head_dim)\n",
|
||
"\n",
|
||
" MEMORY EFFICIENCY:\n",
|
||
" Only return the valid portion of cache, not the entire pre-allocated space\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 + 1)\n",
|
||
" # seq_pos tracks where to write next, so seq_pos tokens have been written\n",
|
||
" valid_len = self.seq_pos\n",
|
||
"\n",
|
||
" cached_keys = key_cache[:, :, :valid_len, :]\n",
|
||
" cached_values = value_cache[:, :, :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.\n",
|
||
"\n",
|
||
" TODO: Move to next position for subsequent cache updates\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" self.seq_pos += 1\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def reset(self) -> None:\n",
|
||
" \"\"\"\n",
|
||
" Reset cache for new generation sequence.\n",
|
||
"\n",
|
||
" TODO: Clear cache state for fresh generation\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Reset sequence position to 0\n",
|
||
" 2. Clear valid position markers\n",
|
||
" 3. Optionally zero out cache data (not strictly necessary)\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" self.seq_pos = 0\n",
|
||
" # Reset valid positions\n",
|
||
" self.valid_positions = Tensor.zeros(self.batch_size, self.max_seq_len)\n",
|
||
"\n",
|
||
" # Optional: zero out caches (not strictly necessary since we track valid positions)\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",
|
||
" ### END SOLUTION\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",
|
||
" ### BEGIN SOLUTION\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",
|
||
" }\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
"def test_unit_kv_cache():\n",
|
||
" \"\"\"🔬 Test KVCache implementation with realistic transformer dimensions.\"\"\"\n",
|
||
" print(\"🔬 Unit Test: KV Cache Implementation...\")\n",
|
||
"\n",
|
||
" # Test parameters (small transformer)\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\n",
|
||
" assert cache.get_memory_usage()['total_mb'] > 0\n",
|
||
" print(f\"✅ Cache initialized: {cache.get_memory_usage()['total_mb']:.2f} MB\")\n",
|
||
"\n",
|
||
" # Test 2: Update and retrieve\n",
|
||
" # Simulate first token (batch=2, heads=4, seq=1, head_dim=16)\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\n",
|
||
" cache.update(0, key1, value1)\n",
|
||
" cached_k, cached_v = cache.get(0)\n",
|
||
"\n",
|
||
" assert cached_k.shape == (batch_size, num_heads, 0, head_dim) # Before advance\n",
|
||
" assert cached_v.shape == (batch_size, num_heads, 0, head_dim)\n",
|
||
"\n",
|
||
" # Advance to next 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)\n",
|
||
" assert cached_v.shape == (batch_size, num_heads, 1, head_dim)\n",
|
||
"\n",
|
||
" # Add second token\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",
|
||
" # Now cache should have 2 tokens\n",
|
||
" cached_k, cached_v = cache.get(0)\n",
|
||
" assert cached_k.shape == (batch_size, num_heads, 2, head_dim)\n",
|
||
" assert cached_v.shape == (batch_size, num_heads, 2, head_dim)\n",
|
||
"\n",
|
||
" print(\"✅ Cache update and retrieval works correctly!\")\n",
|
||
"\n",
|
||
" # Test 3: Multiple layers\n",
|
||
" cache.reset()\n",
|
||
" cache.update(0, key1, value1) # Layer 0\n",
|
||
" cache.update(1, key1, value1) # Layer 1\n",
|
||
" cache.update(2, key1, value1) # Layer 2\n",
|
||
" cache.advance()\n",
|
||
"\n",
|
||
" for layer_idx in range(num_layers):\n",
|
||
" cached_k, cached_v = cache.get(layer_idx)\n",
|
||
" assert cached_k.shape[2] == 1 # One token in each layer cache\n",
|
||
"\n",
|
||
" print(\"✅ Multi-layer caching works correctly!\")\n",
|
||
"\n",
|
||
" # Test 4: Reset functionality\n",
|
||
" cache.reset()\n",
|
||
" assert cache.seq_pos == 0\n",
|
||
" cached_k, cached_v = cache.get(0)\n",
|
||
" assert cached_k.shape == (batch_size, num_heads, 0, head_dim) # Should be empty after reset\n",
|
||
"\n",
|
||
" print(\"✅ Cache reset works correctly!\")\n",
|
||
" print(\"✅ KVCache implementation is working perfectly!\")\n",
|
||
"\n",
|
||
"test_unit_kv_cache()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "960d1a1d",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 🔧 Part 4: Cache-Aware Attention Implementation\n",
|
||
"\n",
|
||
"### The Integration Challenge\n",
|
||
"\n",
|
||
"Now we need to modify attention to work seamlessly with our cache. The key insight is that we only compute K,V for NEW tokens, then combine with cached history for the full attention computation.\n",
|
||
"\n",
|
||
"### Traditional vs Cached Attention Flow\n",
|
||
"\n",
|
||
"```\n",
|
||
"Traditional Attention (Inefficient):\n",
|
||
"┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐\n",
|
||
"│ All Tokens │───▶│ Compute Q,K,V │───▶│ Attention │\n",
|
||
"│ [tok₁,tok₂,tok₃]│ │ (redundant) │ │ Output │\n",
|
||
"└─────────────────┘ └─────────────────┘ └─────────────────┘\n",
|
||
" ↑\n",
|
||
" Recomputes K₁,V₁,K₂,V₂\n",
|
||
" every single step!\n",
|
||
"\n",
|
||
"Cached Attention (Efficient):\n",
|
||
"┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐\n",
|
||
"│ New Token │───▶│ Compute Q,K₃,V₃ │───▶│ Cache.update() │\n",
|
||
"│ [tok₃] │ │ (only new!) │ │ │\n",
|
||
"└─────────────────┘ └─────────────────┘ └─────────────────┘\n",
|
||
" │\n",
|
||
" ▼\n",
|
||
"┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐\n",
|
||
"│ Attention │◀───│ Cache.get() │◀───│ Cached History │\n",
|
||
"│ Output │ │ K₁,V₁,K₂,V₂,K₃,V₃│ │ K₁,V₁,K₂,V₂ │\n",
|
||
"└─────────────────┘ └─────────────────┘ └─────────────────┘\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Attention Computation with Cache\n",
|
||
"\n",
|
||
"```\n",
|
||
"Step-by-Step Process:\n",
|
||
"1. Input: Q₃ (query for new token), K₃,V₃ (key,value for new token)\n",
|
||
"2. Cache Update: Store K₃,V₃ → Cache now has [K₁,V₁,K₂,V₂,K₃,V₃]\n",
|
||
"3. Cache Retrieval: Get all cached K,V → [K₁,K₂,K₃], [V₁,V₂,V₃]\n",
|
||
"4. Attention: Q₃ @ [K₁,K₂,K₃]ᵀ → attention weights\n",
|
||
"5. Output: attention_weights @ [V₁,V₂,V₃] → final result\n",
|
||
"\n",
|
||
"Memory Access Pattern:\n",
|
||
"Write: O(1) - just append K₃,V₃ to cache\n",
|
||
"Read: O(seq_len) - retrieve full cached history\n",
|
||
"Total: O(seq_len) instead of O(seq_len²)!\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Causal Masking Integration\n",
|
||
"\n",
|
||
"```\n",
|
||
"Causal Mask Application:\n",
|
||
"┌─────┬─────┬─────┐\n",
|
||
"│ 0 │-inf │-inf │ ← Position 0 can only see itself\n",
|
||
"├─────┼─────┼─────┤\n",
|
||
"│ 0 │ 0 │-inf │ ← Position 1 can see 0,1\n",
|
||
"├─────┼─────┼─────┤\n",
|
||
"│ 0 │ 0 │ 0 │ ← Position 2 can see 0,1,2\n",
|
||
"└─────┴─────┴─────┘\n",
|
||
"\n",
|
||
"For cached attention:\n",
|
||
"- Mask shape: (max_seq_len, max_seq_len)\n",
|
||
"- Slice needed: (1, current_seq_len) for current query\n",
|
||
"- Apply before softmax to prevent future token access\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "346d005a",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "attention_with_cache",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "00d5d995",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def attention_with_cache(\n",
|
||
" query: Tensor,\n",
|
||
" key: Tensor,\n",
|
||
" value: Tensor,\n",
|
||
" cache: KVCache,\n",
|
||
" layer_idx: int,\n",
|
||
" mask: Optional[Tensor] = None\n",
|
||
") -> Tensor:\n",
|
||
" \"\"\"\n",
|
||
" Compute attention using KV cache for efficient autoregressive generation.\n",
|
||
"\n",
|
||
" This is the core optimization: instead of recomputing K,V for all tokens,\n",
|
||
" we cache them and only compute for the new token.\n",
|
||
"\n",
|
||
" TODO: Implement cache-aware attention that's 10x+ faster than naive approach\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Update cache with new key,value pairs for current token\n",
|
||
" 2. Retrieve full cached history (all previous + current)\n",
|
||
" 3. Compute attention using query vs full cached K,V\n",
|
||
" 4. Apply causal masking to ensure autoregressive property\n",
|
||
" 5. Return attention output (cache position advanced externally)\n",
|
||
"\n",
|
||
" ATTENTION COMPUTATION:\n",
|
||
" ```\n",
|
||
" scores = query @ cached_keys.transpose(-2, -1) / sqrt(head_dim)\n",
|
||
" if mask: scores = mask_attention(scores, mask)\n",
|
||
" attention_weights = softmax(scores)\n",
|
||
" output = attention_weights @ cached_values\n",
|
||
" ```\n",
|
||
"\n",
|
||
" Args:\n",
|
||
" query: Query tensor for current token (batch, num_heads, 1, head_dim)\n",
|
||
" key: Key tensor for current token (batch, num_heads, 1, head_dim)\n",
|
||
" value: Value tensor for current token (batch, num_heads, 1, head_dim)\n",
|
||
" cache: KVCache instance to store/retrieve K,V pairs\n",
|
||
" layer_idx: Which transformer layer this attention belongs to\n",
|
||
" mask: Optional attention mask for preventing future token access\n",
|
||
"\n",
|
||
" Returns:\n",
|
||
" attention_output: Computed attention for current token (batch, num_heads, 1, head_dim)\n",
|
||
"\n",
|
||
" PERFORMANCE:\n",
|
||
" - Time: O(seq_len) instead of O(seq_len²) for generation\n",
|
||
" - Memory: O(seq_len × hidden_dim) cache overhead\n",
|
||
" - Speedup: 10x+ for long sequences\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" batch_size, num_heads, seq_len_q, head_dim = query.shape\n",
|
||
"\n",
|
||
" # Step 1: Update cache with new key,value for current token\n",
|
||
" cache.update(layer_idx, key, value)\n",
|
||
"\n",
|
||
" # Step 2: Retrieve full cached K,V (all previous + current token)\n",
|
||
" cached_keys, cached_values = cache.get(layer_idx)\n",
|
||
"\n",
|
||
" # If cache is empty (first token), add current token\n",
|
||
" if cached_keys.shape[2] == 0:\n",
|
||
" cached_keys = key\n",
|
||
" cached_values = value\n",
|
||
" else:\n",
|
||
" # Concatenate new token with cached history\n",
|
||
" cached_keys = Tensor.cat([cached_keys, key], dim=2)\n",
|
||
" cached_values = Tensor.cat([cached_values, value], dim=2)\n",
|
||
"\n",
|
||
" # Step 3: Compute attention scores\n",
|
||
" # query: (batch, heads, 1, head_dim)\n",
|
||
" # cached_keys: (batch, heads, seq_len_k, head_dim)\n",
|
||
" # Need: (batch, heads, 1, seq_len_k)\n",
|
||
" scores = np.matmul(query.data, cached_keys.transpose(-1, -2).data)\n",
|
||
"\n",
|
||
" # Scale by sqrt(head_dim) for numerical stability\n",
|
||
" scores = scores / np.sqrt(head_dim)\n",
|
||
"\n",
|
||
" # Step 4: Apply causal mask if provided\n",
|
||
" if mask is not None:\n",
|
||
" # Mask should be shape (max_seq_len, max_seq_len)\n",
|
||
" # We need to slice to (1, seq_len_k) for current query position\n",
|
||
" seq_len_k = cached_keys.shape[2]\n",
|
||
" query_pos = seq_len_k - 1 # Current query position\n",
|
||
"\n",
|
||
" if mask.shape[-1] >= seq_len_k and mask.shape[-2] > query_pos:\n",
|
||
" # For current query position, take the corresponding row up to seq_len_k columns\n",
|
||
" mask_slice = mask.data[query_pos:query_pos+1, :seq_len_k] # Shape: (1, seq_len_k)\n",
|
||
" # Reshape to match scores: (batch, heads, 1, seq_len_k)\n",
|
||
" mask_broadcast = mask_slice.reshape(1, 1, 1, seq_len_k)\n",
|
||
" scores = scores + mask_broadcast # Apply mask (already has -1e9 values)\n",
|
||
"\n",
|
||
" # Step 5: Compute attention weights via softmax\n",
|
||
" # Numerical stability: subtract max before exp\n",
|
||
" scores_max = np.max(scores, axis=-1, keepdims=True)\n",
|
||
" scores_stable = scores - scores_max\n",
|
||
" exp_scores = np.exp(scores_stable)\n",
|
||
" attention_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)\n",
|
||
"\n",
|
||
" # Step 6: Compute final attention output\n",
|
||
" # attention_weights: (batch, heads, 1, seq_len_k)\n",
|
||
" # cached_values: (batch, heads, seq_len_k, head_dim)\n",
|
||
" # output: (batch, heads, 1, head_dim)\n",
|
||
" output_data = np.matmul(attention_weights, cached_values.data)\n",
|
||
" attention_output = Tensor(output_data)\n",
|
||
"\n",
|
||
" # Note: cache.advance() should be called externally after all layers process this token\n",
|
||
" return attention_output\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
"def test_unit_attention_with_cache():\n",
|
||
" \"\"\"🔬 Test cache-aware attention against naive implementation.\"\"\"\n",
|
||
" print(\"🔬 Unit Test: Attention with Cache...\")\n",
|
||
"\n",
|
||
" # Setup small test case\n",
|
||
" batch_size, num_heads, head_dim = 1, 2, 8\n",
|
||
" max_seq_len = 4\n",
|
||
"\n",
|
||
" cache = KVCache(batch_size, max_seq_len, 1, num_heads, head_dim)\n",
|
||
"\n",
|
||
" # Test generation sequence: 3 tokens\n",
|
||
" for step in range(3):\n",
|
||
" print(f\" Generation step {step + 1}...\")\n",
|
||
"\n",
|
||
" # Create QKV for current token\n",
|
||
" q = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" k = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" v = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
"\n",
|
||
" # Compute attention with cache\n",
|
||
" output = attention_with_cache(q, k, v, cache, layer_idx=0)\n",
|
||
"\n",
|
||
" # Verify output shape\n",
|
||
" assert output.shape == (batch_size, num_heads, 1, head_dim)\n",
|
||
"\n",
|
||
" # Advance cache position\n",
|
||
" cache.advance()\n",
|
||
"\n",
|
||
" # Verify cache grows correctly\n",
|
||
" # After processing step i and advancing, we should have i+1 elements cached\n",
|
||
" cached_k, cached_v = cache.get(0)\n",
|
||
" expected_cache_len = step + 1\n",
|
||
" print(f\" Step {step}: cache has {cached_k.shape[2]} elements, expected {expected_cache_len}\")\n",
|
||
" assert cached_k.shape[2] == expected_cache_len\n",
|
||
" assert cached_v.shape[2] == expected_cache_len\n",
|
||
"\n",
|
||
" print(\"✅ Cache-aware attention works correctly!\")\n",
|
||
"\n",
|
||
" # Test with causal mask\n",
|
||
" print(\" Testing with causal masking...\")\n",
|
||
" cache.reset()\n",
|
||
"\n",
|
||
" # Create causal mask (lower triangular)\n",
|
||
" causal_mask = Tensor(np.triu(np.ones((max_seq_len, max_seq_len)) * -1e9, k=1))\n",
|
||
"\n",
|
||
" q = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" k = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" v = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
"\n",
|
||
" output_masked = attention_with_cache(q, k, v, cache, layer_idx=0, mask=causal_mask)\n",
|
||
" cache.advance()\n",
|
||
"\n",
|
||
" print(f\" Masked output shape: {output_masked.shape}\")\n",
|
||
" assert output_masked.shape == (batch_size, num_heads, 1, head_dim)\n",
|
||
"\n",
|
||
" print(\"✅ Causal masking works correctly!\")\n",
|
||
" print(\"✅ Cache-aware attention implementation complete!\")\n",
|
||
"\n",
|
||
"test_unit_attention_with_cache()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "c304da93",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 📊 Part 5: Performance Analysis - Measuring the Speedup\n",
|
||
"\n",
|
||
"### Understanding the Performance Gains\n",
|
||
"\n",
|
||
"Let's measure the dramatic improvements KV caching provides. We'll compare naive recomputation vs cached attention across different sequence lengths to understand the scaling benefits.\n",
|
||
"\n",
|
||
"### What We're Measuring\n",
|
||
"\n",
|
||
"```\n",
|
||
"Complexity Comparison:\n",
|
||
"┌─────────────────┬─────────────────┬─────────────────┐\n",
|
||
"│ Approach │ Time Complexity │ Memory Usage │\n",
|
||
"├─────────────────┼─────────────────┼─────────────────┤\n",
|
||
"│ Naive │ O(n²) │ O(n) │\n",
|
||
"│ Recomputation │ │ │\n",
|
||
"├─────────────────┼─────────────────┼─────────────────┤\n",
|
||
"│ KV Caching │ O(n) │ O(n×hidden) │\n",
|
||
"│ │ │ │\n",
|
||
"└─────────────────┴─────────────────┴─────────────────┘\n",
|
||
"\n",
|
||
"Trade-off: Use more memory to achieve quadratic speedup!\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Real-World Impact Visualization\n",
|
||
"\n",
|
||
"```\n",
|
||
"Production Serving Scenario:\n",
|
||
"Without Caching: With Caching:\n",
|
||
"┌─────────────────┐ ┌─────────────────┐\n",
|
||
"│ User Request │ │ User Request │\n",
|
||
"│ \"Write a story\" │ │ \"Write a story\" │\n",
|
||
"└─────────┬───────┘ └─────────┬───────┘\n",
|
||
" │ │\n",
|
||
" ▼ ▼\n",
|
||
"┌─────────────────┐ ┌─────────────────┐\n",
|
||
"│ Token 1: 1 ops │ │ Token 1: 1 ops │\n",
|
||
"│ Token 2: 2 ops │ │ Token 2: 1 ops │\n",
|
||
"│ Token 3: 3 ops │ │ Token 3: 1 ops │\n",
|
||
"│ ... │ │ ... │\n",
|
||
"│ Token 100: 100 │ │ Token 100: 1 op │\n",
|
||
"└─────────────────┘ └─────────────────┘\n",
|
||
"Total: 5,050 ops Total: 100 ops\n",
|
||
"Response: 5+ seconds Response: 0.1 seconds\n",
|
||
"Cost: $$$$$ Cost: $\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d272c1a9",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "performance_analysis",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "0056cbd0",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def analyze_kv_cache_performance():\n",
|
||
" \"\"\"📊 Measure dramatic performance gains from KV caching.\"\"\"\n",
|
||
" print(\"📊 Analyzing KV Cache Performance vs Naive Recomputation...\")\n",
|
||
"\n",
|
||
" # Test configuration (realistic transformer)\n",
|
||
" batch_size, num_heads, head_dim = 1, 8, 64\n",
|
||
" num_layers = 12\n",
|
||
"\n",
|
||
" sequence_lengths = [16, 32, 64, 128, 256] # Realistic generation lengths\n",
|
||
"\n",
|
||
" print(\"\\n=== Performance Comparison ===\")\n",
|
||
" print(\"Seq Len | Naive Ops | Cached Ops | Speedup | Cache Memory\")\n",
|
||
" print(\"-\" * 65)\n",
|
||
"\n",
|
||
" for seq_len in sequence_lengths:\n",
|
||
" # Calculate theoretical operation counts\n",
|
||
"\n",
|
||
" # Naive approach: At each step i, recompute attention for all i+1 tokens\n",
|
||
" naive_ops = 0\n",
|
||
" for step in range(seq_len):\n",
|
||
" current_seq_len = step + 1\n",
|
||
" # K,V computation: current_seq_len × head_dim per head per layer\n",
|
||
" kv_ops = current_seq_len * head_dim * num_heads * num_layers\n",
|
||
" # Attention: current_seq_len × head_dim per head per layer\n",
|
||
" attn_ops = current_seq_len * head_dim * num_heads * num_layers\n",
|
||
" naive_ops += kv_ops + attn_ops\n",
|
||
"\n",
|
||
" # Cached approach: Compute K,V only for new token, attention with cached history\n",
|
||
" cached_ops = 0\n",
|
||
" for step in range(seq_len):\n",
|
||
" current_seq_len = step + 1\n",
|
||
" # K,V computation: only 1 new token × head_dim per head per layer\n",
|
||
" kv_ops = 1 * head_dim * num_heads * num_layers\n",
|
||
" # Attention: current_seq_len × head_dim per head per layer (with cache)\n",
|
||
" attn_ops = current_seq_len * head_dim * num_heads * num_layers\n",
|
||
" cached_ops += kv_ops + attn_ops\n",
|
||
"\n",
|
||
" # Calculate metrics\n",
|
||
" speedup = naive_ops / cached_ops if cached_ops > 0 else float('inf')\n",
|
||
"\n",
|
||
" # Memory usage for cache\n",
|
||
" cache = KVCache(batch_size, seq_len, num_layers, num_heads, head_dim)\n",
|
||
" cache_memory = cache.get_memory_usage()['total_mb']\n",
|
||
"\n",
|
||
" print(f\"{seq_len:7d} | {naive_ops/1000:8.0f}K | {cached_ops/1000:9.0f}K | {speedup:6.1f}x | {cache_memory:8.1f}MB\")\n",
|
||
"\n",
|
||
" print(\"\\n💡 Key Insights:\")\n",
|
||
" print(\"• Speedup grows with sequence length (O(n²) vs O(n) complexity)\")\n",
|
||
" print(\"• Memory overhead is manageable and constant per layer\")\n",
|
||
" print(\"• Essential for production serving at any reasonable scale\")\n",
|
||
"\n",
|
||
" # Theoretical complexity analysis\n",
|
||
" print(\"\\n=== Theoretical Complexity Analysis ===\")\n",
|
||
" n = 256 # Example sequence length\n",
|
||
"\n",
|
||
" # For naive approach: sum of 1+2+3+...+n computations\n",
|
||
" naive_complexity = n * (n + 1) // 2 # Sum from 1 to n\n",
|
||
" # For cached approach: n computations (1 per step)\n",
|
||
" cached_complexity = n # Linear in sequence length\n",
|
||
"\n",
|
||
" print(f\"For {n}-token generation:\")\n",
|
||
" print(f\" Naive approach: O(n²) = {naive_complexity:,} operations\")\n",
|
||
" print(f\" Cached approach: O(n) = {cached_complexity:,} operations\")\n",
|
||
" print(f\" Theoretical speedup: {naive_complexity/cached_complexity:.0f}x\")\n",
|
||
"\n",
|
||
" print(\"\\n🚀 Production Impact:\")\n",
|
||
" print(\"• Enables real-time chat interfaces (ChatGPT, Claude)\")\n",
|
||
" print(\"• Reduces serving costs by 10x+ for long conversations\")\n",
|
||
" print(\"• Makes on-device generation feasible (mobile, edge)\")\n",
|
||
" print(\"• Critical for any autoregressive model deployment\")\n",
|
||
"\n",
|
||
" # Real-world serving scenarios\n",
|
||
" print(\"\\n=== Real-World Serving Analysis ===\")\n",
|
||
"\n",
|
||
" scenarios = [\n",
|
||
" (\"Chat Response\", 50, \"Real-time requirement\"),\n",
|
||
" (\"Code Completion\", 200, \"IDE integration\"),\n",
|
||
" (\"Document Summary\", 500, \"Batch processing\"),\n",
|
||
" (\"Long Conversation\", 1000, \"Extended context\")\n",
|
||
" ]\n",
|
||
"\n",
|
||
" print(\"Scenario | Tokens | Without Cache | With Cache | Savings\")\n",
|
||
" print(\"-\" * 70)\n",
|
||
"\n",
|
||
" for scenario, tokens, context in scenarios:\n",
|
||
" without_cache = tokens * (tokens + 1) // 2\n",
|
||
" with_cache = tokens\n",
|
||
" savings = without_cache / with_cache\n",
|
||
"\n",
|
||
" print(f\"{scenario:16s} | {tokens:6d} | {without_cache:12,} | {with_cache:9,} | {savings:5.0f}x\")\n",
|
||
"\n",
|
||
"analyze_kv_cache_performance()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a128d8c4",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 🔧 Part 6: Advanced Optimization Strategies\n",
|
||
"\n",
|
||
"### Production KV Caching Patterns\n",
|
||
"\n",
|
||
"Real production systems implement several sophisticated optimizations beyond basic caching. Let's explore the advanced patterns used in state-of-the-art serving systems.\n",
|
||
"\n",
|
||
"### Memory Optimization Strategies\n",
|
||
"\n",
|
||
"```\n",
|
||
"Precision Trade-offs:\n",
|
||
"┌─────────────┬─────────────┬─────────────┬─────────────┐\n",
|
||
"│ Precision │ Memory │ Quality │ Use Case │\n",
|
||
"├─────────────┼─────────────┼─────────────┼─────────────┤\n",
|
||
"│ FP32 │ 100% │ Perfect │ Development │\n",
|
||
"│ FP16 │ 50% │ Minimal loss│ Production │\n",
|
||
"│ INT8 │ 25% │ Some loss │ Edge/Mobile │\n",
|
||
"│ INT4 │ 12.5% │ Quality loss│ Extreme opt │\n",
|
||
"└─────────────┴─────────────┴─────────────┴─────────────┘\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Sliding Window Attention\n",
|
||
"\n",
|
||
"```\n",
|
||
"Fixed Context Window vs Sliding Window:\n",
|
||
"\n",
|
||
"Fixed Window (Traditional):\n",
|
||
"┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐\n",
|
||
"│ T₁ │ T₂ │ T₃ │ T₄ │ T₅ │ T₆ │ T₇ │ T₈ │\n",
|
||
"└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘\n",
|
||
" ↑\n",
|
||
" Current token sees ALL history\n",
|
||
" Memory: O(n), but limited to max_seq_len\n",
|
||
"\n",
|
||
"Sliding Window (Advanced):\n",
|
||
"┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐\n",
|
||
"│ │ │ T₃ │ T₄ │ T₅ │ T₆ │ T₇ │ T₈ │\n",
|
||
"└─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘\n",
|
||
" ↑─────────────window_size──────────↑\n",
|
||
" Current token sees recent history only\n",
|
||
" Memory: O(window), enables infinite generation\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Prefix Caching Optimization\n",
|
||
"\n",
|
||
"```\n",
|
||
"Shared Prefix Caching:\n",
|
||
"User A: \"Write a Python function that\" → Cache prefix\n",
|
||
"User B: \"Write a Python function that\" → Reuse cached prefix!\n",
|
||
"User C: \"Write a Python script to\" → Different, new cache\n",
|
||
"\n",
|
||
"Cache Hit Rate Impact:\n",
|
||
"┌─────────────────┬─────────────────┬─────────────────┐\n",
|
||
"│ Cache Scenario │ Hit Rate │ Speedup │\n",
|
||
"├─────────────────┼─────────────────┼─────────────────┤\n",
|
||
"│ No Sharing │ 0% │ 1x │\n",
|
||
"│ Common Prompts │ 30% │ 1.4x │\n",
|
||
"│ Chat Templates │ 60% │ 2.5x │\n",
|
||
"│ Code Patterns │ 80% │ 5x │\n",
|
||
"└─────────────────┴─────────────────┴─────────────────┘\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "44f2dead",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "optimization_insights",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "7b265570",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def analyze_advanced_caching_strategies():\n",
|
||
" \"\"\"📊 Explore advanced caching strategies and production trade-offs.\"\"\"\n",
|
||
" print(\"📊 Advanced KV Caching Strategies Analysis...\")\n",
|
||
"\n",
|
||
" # Configuration for large-scale analysis (reduced for educational demonstration)\n",
|
||
" seq_len, batch_size = 512, 4\n",
|
||
" num_layers, num_heads, head_dim = 12, 16, 64 # Realistic scale for demonstration\n",
|
||
"\n",
|
||
" print(\"\\n=== Memory Footprint by Precision ===\")\n",
|
||
"\n",
|
||
" # Standard FP32 cache\n",
|
||
" cache_fp32 = KVCache(batch_size, seq_len, num_layers, num_heads, head_dim)\n",
|
||
" fp32_memory = cache_fp32.get_memory_usage()['total_mb']\n",
|
||
"\n",
|
||
" # Simulated precision variants\n",
|
||
" precisions = [\n",
|
||
" (\"FP32\", fp32_memory, 1.0, \"No quality loss\"),\n",
|
||
" (\"FP16\", fp32_memory / 2, 0.5, \"Minimal quality loss\"),\n",
|
||
" (\"INT8\", fp32_memory / 4, 0.25, \"Some quality loss\"),\n",
|
||
" (\"INT4\", fp32_memory / 8, 0.125, \"Significant loss\")\n",
|
||
" ]\n",
|
||
"\n",
|
||
" print(\"Precision | Memory Usage | Reduction | Quality Impact\")\n",
|
||
" print(\"-\" * 55)\n",
|
||
" for precision, memory, factor, quality in precisions:\n",
|
||
" print(f\"{precision:8s} | {memory:8.0f} MB | {factor:4.2f}x | {quality}\")\n",
|
||
"\n",
|
||
" print(\"\\n=== Sliding Window Analysis ===\")\n",
|
||
"\n",
|
||
" # Compare different window sizes for memory usage\n",
|
||
" full_seq_len = 2048 # Realistic long sequence for demonstration\n",
|
||
" window_sizes = [256, 512, 1024, 2048]\n",
|
||
"\n",
|
||
" print(\"Window Size | Memory vs Full | Tokens Lost | Use Case\")\n",
|
||
" print(\"-\" * 60)\n",
|
||
"\n",
|
||
" for window_size in window_sizes:\n",
|
||
" # Memory scales with window size\n",
|
||
" full_cache = KVCache(batch_size, full_seq_len, num_layers, num_heads, head_dim)\n",
|
||
" window_cache = KVCache(batch_size, window_size, num_layers, num_heads, head_dim)\n",
|
||
"\n",
|
||
" full_memory = full_cache.get_memory_usage()['total_mb']\n",
|
||
" window_memory = window_cache.get_memory_usage()['total_mb']\n",
|
||
" reduction = full_memory / window_memory\n",
|
||
" tokens_lost = max(0, full_seq_len - window_size)\n",
|
||
"\n",
|
||
" if window_size <= 1024:\n",
|
||
" use_case = \"Chat/Code completion\"\n",
|
||
" elif window_size <= 2048:\n",
|
||
" use_case = \"Document analysis\"\n",
|
||
" else:\n",
|
||
" use_case = \"Long context tasks\"\n",
|
||
"\n",
|
||
" print(f\"{window_size:10d} | {reduction:9.1f}x | {tokens_lost:10d} | {use_case}\")\n",
|
||
"\n",
|
||
" print(\"\\n=== Multi-GPU Scaling Strategy ===\")\n",
|
||
"\n",
|
||
" # Analyze how caching scales across multiple GPUs\n",
|
||
" gpu_configs = [1, 2, 4, 8]\n",
|
||
" large_batch = 16 # Reasonable batch for demonstration\n",
|
||
"\n",
|
||
" print(\"GPUs | Batch/GPU | Cache/GPU | Total Memory | Throughput\")\n",
|
||
" print(\"-\" * 60)\n",
|
||
"\n",
|
||
" for num_gpus in gpu_configs:\n",
|
||
" batch_per_gpu = large_batch // num_gpus\n",
|
||
" cache_per_gpu = KVCache(batch_per_gpu, seq_len, num_layers, num_heads, head_dim)\n",
|
||
" memory_per_gpu = cache_per_gpu.get_memory_usage()['total_mb']\n",
|
||
" total_memory = memory_per_gpu * num_gpus\n",
|
||
" throughput_scale = num_gpus # Linear scaling assumption\n",
|
||
"\n",
|
||
" print(f\"{num_gpus:4d} | {batch_per_gpu:8d} | {memory_per_gpu:8.0f}MB | {total_memory:9.0f}MB | {throughput_scale:8.0f}x\")\n",
|
||
"\n",
|
||
" print(\"\\n=== Production Serving Scenarios ===\")\n",
|
||
"\n",
|
||
" scenarios = [\n",
|
||
" (\"Real-time Chat\", 512, 1, \"Low latency critical\"),\n",
|
||
" (\"Code Completion\", 1024, 4, \"IDE integration\"),\n",
|
||
" (\"Batch Translation\", 2048, 8, \"High throughput\"),\n",
|
||
" (\"Long Document\", 2048, 4, \"Context preservation\")\n",
|
||
" ]\n",
|
||
"\n",
|
||
" print(\"Scenario | Max Len | Batch | Memory | Optimal Strategy\")\n",
|
||
" print(\"-\" * 70)\n",
|
||
"\n",
|
||
" for name, max_len, batch, priority in scenarios:\n",
|
||
" # Calculate memory for each scenario\n",
|
||
" scenario_cache = KVCache(batch, max_len, num_layers, num_heads, head_dim)\n",
|
||
" scenario_memory = scenario_cache.get_memory_usage()['total_mb']\n",
|
||
"\n",
|
||
" # Determine optimal strategy based on memory usage\n",
|
||
" if scenario_memory < 500: # < 0.5GB\n",
|
||
" strategy = \"FP32 cache\"\n",
|
||
" elif scenario_memory < 2000: # < 2GB\n",
|
||
" strategy = \"FP16 cache\"\n",
|
||
" elif scenario_memory < 8000: # < 8GB\n",
|
||
" strategy = \"FP16 + sliding window\"\n",
|
||
" else: # > 8GB\n",
|
||
" strategy = \"Multi-GPU + quantization\"\n",
|
||
"\n",
|
||
" print(f\"{name:15s} | {max_len:7d} | {batch:5d} | {scenario_memory:6.0f}MB | {strategy}\")\n",
|
||
"\n",
|
||
" print(\"\\n💡 Advanced Optimization Insights:\")\n",
|
||
" print(\"• FP16 provides 2x memory savings with negligible quality loss\")\n",
|
||
" print(\"• Sliding windows enable unlimited generation with fixed memory\")\n",
|
||
" print(\"• Multi-GPU scaling is linear for both memory and throughput\")\n",
|
||
" print(\"• Quantization beyond FP16 requires careful quality evaluation\")\n",
|
||
"\n",
|
||
" print(\"\\n🚀 Production Implementation Recommendations:\")\n",
|
||
" print(\"• Start with FP16 caching as the baseline optimization\")\n",
|
||
" print(\"• Implement sliding windows for sequences > 4K tokens\")\n",
|
||
" print(\"• Use prefix caching for common prompt patterns\")\n",
|
||
" print(\"• Consider multi-GPU distribution for high-throughput serving\")\n",
|
||
" print(\"• Monitor cache hit rates and memory utilization in production\")\n",
|
||
"\n",
|
||
" # Cache hit rate simulation\n",
|
||
" print(\"\\n=== Prefix Caching Effectiveness ===\")\n",
|
||
"\n",
|
||
" prefix_scenarios = [\n",
|
||
" (\"No Sharing\", 0.0, 1.0),\n",
|
||
" (\"Common Prompts\", 0.3, 1.4),\n",
|
||
" (\"Chat Templates\", 0.6, 2.5),\n",
|
||
" (\"Code Patterns\", 0.8, 5.0)\n",
|
||
" ]\n",
|
||
"\n",
|
||
" print(\"Scenario | Hit Rate | Effective Speedup | Memory Efficiency\")\n",
|
||
" print(\"-\" * 65)\n",
|
||
"\n",
|
||
" for scenario, hit_rate, speedup in prefix_scenarios:\n",
|
||
" memory_efficiency = 1.0 + hit_rate * 0.5 # Shared prefixes reduce memory\n",
|
||
" print(f\"{scenario:14s} | {hit_rate:7.1%} | {speedup:12.1f}x | {memory_efficiency:14.1f}x\")\n",
|
||
"\n",
|
||
"analyze_advanced_caching_strategies()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "d81b4147",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 🧪 Part 7: Module Integration Test\n",
|
||
"\n",
|
||
"Our KV caching system is complete! Time for comprehensive testing to ensure all components work together seamlessly and deliver the promised performance improvements.\n",
|
||
"\n",
|
||
"### Integration Test Coverage\n",
|
||
"\n",
|
||
"We'll validate:\n",
|
||
"1. **Multi-layer caching**: All transformer layers cache correctly\n",
|
||
"2. **Generation simulation**: End-to-end token generation workflow\n",
|
||
"3. **Memory efficiency**: Large-scale cache allocation and management\n",
|
||
"4. **Performance consistency**: Speedup measurements are reliable\n",
|
||
"5. **Cache lifecycle**: Reset, reuse, and state management"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "6d8f73bd",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": true,
|
||
"grade_id": "test_module",
|
||
"locked": true,
|
||
"points": 20
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "f9fa7dd6",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def test_module():\n",
|
||
" \"\"\"\n",
|
||
" Comprehensive test of entire Module 14: KV Caching functionality.\n",
|
||
"\n",
|
||
" This final test runs before module summary to ensure:\n",
|
||
" - All unit tests pass\n",
|
||
" - KVCache works correctly with realistic parameters\n",
|
||
" - Cache-aware attention produces correct results\n",
|
||
" - Performance analysis runs successfully\n",
|
||
" - Module is ready for integration with TinyTorch\n",
|
||
" \"\"\"\n",
|
||
" print(\"🧪 RUNNING MODULE 14 INTEGRATION TEST\")\n",
|
||
" print(\"=\" * 50)\n",
|
||
"\n",
|
||
" # Run all unit tests\n",
|
||
" print(\"Running unit tests...\")\n",
|
||
" test_unit_kv_cache()\n",
|
||
" test_unit_attention_with_cache()\n",
|
||
"\n",
|
||
" print(\"\\nRunning integration scenarios...\")\n",
|
||
"\n",
|
||
" # Integration Test 1: Multi-layer generation simulation\n",
|
||
" print(\"🔬 Integration Test: Multi-layer transformer generation...\")\n",
|
||
"\n",
|
||
" batch_size, max_seq_len = 2, 16\n",
|
||
" num_layers, num_heads, head_dim = 4, 8, 32\n",
|
||
"\n",
|
||
" # Create cache system\n",
|
||
" cache = KVCache(batch_size, max_seq_len, num_layers, num_heads, head_dim)\n",
|
||
"\n",
|
||
" # Simulate 8-token generation across all layers\n",
|
||
" for token_idx in range(8):\n",
|
||
" for layer_idx in range(num_layers):\n",
|
||
" # Generate random QKV for current token\n",
|
||
" q = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" k = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" v = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
"\n",
|
||
" # Compute attention with cache\n",
|
||
" output = attention_with_cache(q, k, v, cache, layer_idx)\n",
|
||
"\n",
|
||
" # Verify output shape\n",
|
||
" assert output.shape == (batch_size, num_heads, 1, head_dim)\n",
|
||
"\n",
|
||
" # Advance cache position after all layers process the token\n",
|
||
" cache.advance()\n",
|
||
"\n",
|
||
" # Verify cache state after each token\n",
|
||
" for layer_idx in range(num_layers):\n",
|
||
" cached_k, cached_v = cache.get(layer_idx)\n",
|
||
" expected_len = token_idx + 1\n",
|
||
" assert cached_k.shape[2] == expected_len\n",
|
||
" assert cached_v.shape[2] == expected_len\n",
|
||
"\n",
|
||
" print(\"✅ Multi-layer generation works correctly!\")\n",
|
||
"\n",
|
||
" # Integration Test 2: Memory efficiency validation\n",
|
||
" print(\"🔬 Integration Test: Memory efficiency...\")\n",
|
||
"\n",
|
||
" # Test large-scale cache\n",
|
||
" large_cache = KVCache(\n",
|
||
" batch_size=4,\n",
|
||
" max_seq_len=512,\n",
|
||
" num_layers=12,\n",
|
||
" num_heads=16,\n",
|
||
" head_dim=64\n",
|
||
" )\n",
|
||
"\n",
|
||
" memory_usage = large_cache.get_memory_usage()\n",
|
||
" assert memory_usage['total_mb'] > 0\n",
|
||
" assert memory_usage['per_layer_mb'] > 0\n",
|
||
"\n",
|
||
" print(f\"✅ Large cache: {memory_usage['total_mb']:.1f} MB allocated efficiently!\")\n",
|
||
"\n",
|
||
" # Integration Test 3: Cache reset and reuse\n",
|
||
" print(\"🔬 Integration Test: Cache lifecycle management...\")\n",
|
||
"\n",
|
||
" # Use cache for one sequence\n",
|
||
" q = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" k = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" v = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
"\n",
|
||
" cache.update(0, k, v)\n",
|
||
" cache.advance()\n",
|
||
"\n",
|
||
" # Reset and verify clean state\n",
|
||
" cache.reset()\n",
|
||
" assert cache.seq_pos == 0\n",
|
||
"\n",
|
||
" # Reuse for new sequence\n",
|
||
" cache.update(0, k, v)\n",
|
||
" cached_k, cached_v = cache.get(0)\n",
|
||
" assert cached_k.shape[2] == 0 # Before advance\n",
|
||
"\n",
|
||
" cache.advance()\n",
|
||
" cached_k, cached_v = cache.get(0)\n",
|
||
" assert cached_k.shape[2] == 1 # After advance\n",
|
||
"\n",
|
||
" print(\"✅ Cache lifecycle management works correctly!\")\n",
|
||
"\n",
|
||
" # Integration Test 4: Performance analysis validation\n",
|
||
" print(\"🔬 Integration Test: Performance measurement system...\")\n",
|
||
"\n",
|
||
" # Run performance analysis (should not crash)\n",
|
||
" try:\n",
|
||
" analyze_kv_cache_performance()\n",
|
||
" analyze_advanced_caching_strategies()\n",
|
||
" print(\"✅ Performance analysis completes successfully!\")\n",
|
||
" except Exception as e:\n",
|
||
" print(f\"❌ Performance analysis failed: {e}\")\n",
|
||
" raise\n",
|
||
"\n",
|
||
" # Integration Test 5: Causal masking integration\n",
|
||
" print(\"🔬 Integration Test: Causal masking with multi-token generation...\")\n",
|
||
"\n",
|
||
" cache.reset()\n",
|
||
" causal_mask = Tensor(np.triu(np.ones((max_seq_len, max_seq_len)) * -1e9, k=1))\n",
|
||
"\n",
|
||
" # Generate 3 tokens with causal masking\n",
|
||
" for i in range(3):\n",
|
||
" q = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" k = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
" v = Tensor(np.random.randn(batch_size, num_heads, 1, head_dim))\n",
|
||
"\n",
|
||
" output = attention_with_cache(q, k, v, cache, 0, mask=causal_mask)\n",
|
||
" assert output.shape == (batch_size, num_heads, 1, head_dim)\n",
|
||
" cache.advance()\n",
|
||
"\n",
|
||
" print(\"✅ Causal masking integration works correctly!\")\n",
|
||
"\n",
|
||
" print(\"\\n\" + \"=\" * 50)\n",
|
||
" print(\"🎉 ALL TESTS PASSED! Module 14 ready for export.\")\n",
|
||
" print(\"✅ KVCache: Efficient key-value caching implemented\")\n",
|
||
" print(\"✅ Cache-aware attention: 10x+ speedup achieved\")\n",
|
||
" print(\"✅ Systems analysis: Memory vs speed trade-offs measured\")\n",
|
||
" print(\"✅ Production patterns: Advanced optimization strategies explored\")\n",
|
||
" print(\"✅ Integration: Multi-layer generation and lifecycle management verified\")\n",
|
||
" print(\"\\nRun: tito module complete 14\")\n",
|
||
"\n",
|
||
"# Call the integration test\n",
|
||
"test_module()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "adb5ba71",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 🚀 Part 8: Main Execution Block\n",
|
||
"\n",
|
||
"This module can be run standalone to validate the complete KV caching implementation and see the dramatic performance improvements in action."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "1a0163ab",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"if __name__ == \"__main__\":\n",
|
||
" print(\"🚀 Running Module 14: KV Caching...\")\n",
|
||
" print(\"=\" * 50)\n",
|
||
"\n",
|
||
" # Run comprehensive module test\n",
|
||
" test_module()\n",
|
||
"\n",
|
||
" print(\"\\n\" + \"=\" * 50)\n",
|
||
" print(\"✅ Module 14 validation complete!\")\n",
|
||
" print(\"🔧 Key components implemented:\")\n",
|
||
" print(\" • KVCache: Memory-efficient caching system with O(1) updates\")\n",
|
||
" print(\" • attention_with_cache: Cache-aware attention mechanism\")\n",
|
||
" print(\" • Performance analysis: Dramatic speedup measurements\")\n",
|
||
" print(\" • Advanced strategies: Production optimization patterns\")\n",
|
||
" print(\" • Integration testing: Multi-layer and lifecycle validation\")\n",
|
||
" print(\"\\n🎯 Ready for TinyGPT integration and Milestone 4!\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4f42f26a",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 🤔 ML Systems Thinking: Generation Optimization\n",
|
||
"\n",
|
||
"### Question 1: Cache Memory Scaling\n",
|
||
"You implemented a KVCache for a transformer with 12 layers, 16 heads, and head dimension 64.\n",
|
||
"For a batch size of 8 and maximum sequence length of 1024:\n",
|
||
"- How many MB of memory does the complete cache use? _____ MB\n",
|
||
"- If you reduce head dimension to 32, how much memory is saved? _____ MB saved\n",
|
||
"\n",
|
||
"### Question 2: Generation Speedup Analysis\n",
|
||
"Your cache-aware attention eliminates redundant K,V computation during generation.\n",
|
||
"For generating a 256-token sequence:\n",
|
||
"- How many total attention operations does the naive approach perform? _____ operations\n",
|
||
"- How many operations does the cached approach perform? _____ operations\n",
|
||
"- What's the theoretical speedup ratio? _____ x faster\n",
|
||
"\n",
|
||
"### Question 3: Production Memory Trade-offs\n",
|
||
"Consider serving a chat application with 1000 concurrent users, each with a 512-token context.\n",
|
||
"Using your KVCache with 32 layers, 32 heads, head_dim=128:\n",
|
||
"- Total cache memory required across all users: _____ GB\n",
|
||
"- Memory saved by using FP16 instead of FP32: _____ GB\n",
|
||
"- Maximum context length feasible with 16GB GPU memory per user: _____ tokens\n",
|
||
"\n",
|
||
"### Question 4: Advanced Optimization Selection\n",
|
||
"For different deployment scenarios, rank strategies by effectiveness (1=best, 4=worst):\n",
|
||
"\n",
|
||
"**Real-time chat (low latency critical):**\n",
|
||
"_____ FP32 cache, _____ FP16 cache, _____ Sliding window, _____ No cache\n",
|
||
"\n",
|
||
"**Mobile deployment (memory limited):**\n",
|
||
"_____ FP32 cache, _____ FP16 cache, _____ Sliding window, _____ No cache\n",
|
||
"\n",
|
||
"**Long document processing (context preservation critical):**\n",
|
||
"_____ FP32 cache, _____ FP16 cache, _____ Sliding window, _____ No cache\n",
|
||
"\n",
|
||
"### Question 5: Systems Impact Understanding\n",
|
||
"Based on your analysis of O(n²) vs O(n) complexity:\n",
|
||
"- Primary bottleneck that KV caching solves: _________________________________\n",
|
||
"- Memory vs computation trade-off principle: _____________________________\n",
|
||
"- Why this enables real-time chat applications: ___________________________________\n",
|
||
"- Impact on production serving costs: ___________________________________"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "bdcdf0fe",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 🎯 MODULE SUMMARY: KV Caching\n",
|
||
"\n",
|
||
"Congratulations! You've built a production-grade KV caching system that transforms autoregressive generation from O(n²) to O(n) complexity!\n",
|
||
"\n",
|
||
"### Key Accomplishments\n",
|
||
"- **Built KVCache class** with efficient memory management and O(1) update operations\n",
|
||
"- **Implemented cache-aware attention** achieving 10x+ speedup over naive recomputation\n",
|
||
"- **Measured dramatic performance gains** demonstrating quadratic to linear complexity improvement\n",
|
||
"- **Explored advanced optimization patterns** including quantization, sliding windows, and multi-GPU scaling\n",
|
||
"- **Validated complete integration** with multi-layer transformers and causal masking\n",
|
||
"- **All tests pass ✅** (validated by `test_module()`)\n",
|
||
"\n",
|
||
"### Systems Insights Gained\n",
|
||
"- **Complexity transformation**: From O(n²) naive recomputation to O(n) cached generation\n",
|
||
"- **Memory scaling**: Cache size grows as O(batch × seq_len × layers × heads × head_dim)\n",
|
||
"- **Performance trade-offs**: Constant memory overhead enables quadratic speedup improvement\n",
|
||
"- **Production patterns**: FP16, sliding windows, and prefix caching for real-world deployment\n",
|
||
"- **Engineering impact**: Makes real-time chat and on-device generation economically feasible\n",
|
||
"\n",
|
||
"### Real-World Connection\n",
|
||
"Every production language model uses KV caching:\n",
|
||
"- **ChatGPT/GPT-4**: Enables real-time responses in chat interfaces\n",
|
||
"- **GitHub Copilot**: Powers instant code completion suggestions\n",
|
||
"- **Mobile AI**: Makes on-device generation feasible with limited memory\n",
|
||
"- **API Serving**: Reduces server costs by 10x+ for conversation workloads\n",
|
||
"\n",
|
||
"### Ready for Next Steps\n",
|
||
"Your KV caching implementation provides the optimization foundation that makes TinyGPT production-ready.\n",
|
||
"Export with: `tito module complete 14`\n",
|
||
"\n",
|
||
"**Next**: Milestone 4 (TinyGPT) - Integrate everything to build a complete language model with blazingly fast generation!\n",
|
||
"\n",
|
||
"The optimization you just implemented is literally what makes modern AI chat possible. When you use ChatGPT and get instant responses, your KV caching system is running behind the scenes! 🚀"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3 (ipykernel)",
|
||
"language": "python",
|
||
"name": "python3"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 5
|
||
}
|