{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0", "mimetype": "text/x-python", "file_extension": ".py", "codemirror_mode": { "name": "ipython", "version": 3 }, "pygments_lexer": "ipython3", "nbconvert_exporter": "python" }, "colab": { "provenance": [], "gpuType": "T4", "name": "The Good and Evil of the Key Value Cache โ€” Case Study Notebook" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ๐Ÿ”ง Setup: Run this cell first!\n", "# Check GPU availability and install dependencies\n", "\n", "import torch\n", "import sys\n", "\n", "# Check GPU\n", "if torch.cuda.is_available():\n", " device = torch.device('cuda')\n", " print(f\"โœ… GPU available: {torch.cuda.get_device_name(0)}\")\n", " print(f\" Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n", "else:\n", " device = torch.device('cpu')\n", " print(\"โš ๏ธ No GPU detected. Some cells may run slowly.\")\n", " print(\" Go to Runtime โ†’ Change runtime type โ†’ GPU\")\n", "\n", "print(f\"\\n๐Ÿ“ฆ Python {sys.version.split()[0]}\")\n", "print(f\"๐Ÿ”ฅ PyTorch {torch.__version__}\")\n", "\n", "# Set random seeds for reproducibility\n", "import random\n", "import numpy as np\n", "\n", "SEED = 42\n", "random.seed(SEED)\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)\n", "if torch.cuda.is_available():\n", " torch.cuda.manual_seed_all(SEED)\n", "\n", "print(f\"๐ŸŽฒ Random seed set to {SEED}\")\n", "\n", "%matplotlib inline" ], "id": "setup_cell" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Case Study: The Hidden Tax of Fast Inference โ€” Profiling and Optimizing KV Cache Behavior in a Legal Document AI System\n", "## Implementation Notebook\n", "\n", "*Vizuara Case Study Series*\n", "*Estimated time: 60โ€“90 minutes*\n", "\n", "---" ], "id": "cell_1" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Context\n", "\n", "Imagine you are a senior ML engineer at **Lex Machina Legal Systems**, a LegalTech startup whose flagship product, *LexDraft*, helps attorneys draft contracts in real time. The system is powered by a fine-tuned 7-billion parameter transformer and must return a first token in under 300 milliseconds โ€” even when the document already contains 8,000 tokens of prior context.\n", "\n", "Recently the system started cracking. First-token latency is spiking to 1.4 seconds on long documents. During peak hours the GPU cluster is crashing with out-of-memory (OOM) errors. And here is the paradox that makes this mystery worth solving: **GPU compute utilization sits at only 22โ€“35% at the exact moments when things are worst.** The GPU is idle, yet also overwhelmed.\n", "\n", "The CFO has blocked a quick hardware fix (the cluster already costs \\$38k/month). We need to *understand* the system before we can fix it.\n", "\n", "This notebook is our diagnostic journey. We will:\n", "1. Build a first-principles model of KV cache memory cost and arithmetic intensity\n", "2. Implement a minimal (but real) transformer with and without KV caching and measure the difference\n", "3. Profile how VRAM and throughput scale with sequence length and batch size\n", "4. Visualize the roofline model to classify our inference regime\n", "5. Explore concrete mitigation strategies\n", "\n", "Let us begin." ], "id": "cell_2" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.1 Environment Setup and Imports\n", "\n", "Before anything else, let us make sure we are on a GPU and install the one lightweight profiling helper we will need." ], "id": "cell_3" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip install -q matplotlib numpy torch torchinfo psutil" ], "id": "cell_4" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# โ”€โ”€ Core imports โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "import math\n", "import time\n", "import gc\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib.ticker as mticker\n", "from dataclasses import dataclass\n", "from typing import Optional, Tuple, List, Dict\n", "\n", "# โ”€โ”€ Reproducibility โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "SEED = 42\n", "torch.manual_seed(SEED)\n", "np.random.seed(SEED)\n", "\n", "# โ”€โ”€ Device detection โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "print(f\"๐Ÿ–ฅ๏ธ Running on: {DEVICE}\")\n", "\n", "if DEVICE.type == \"cuda\":\n", " props = torch.cuda.get_device_properties(0)\n", " VRAM_GB = props.total_memory / 1e9\n", " GPU_NAME = props.name\n", " print(f\" GPU : {GPU_NAME}\")\n", " print(f\" VRAM : {VRAM_GB:.1f} GB\")\n", " # Peak memory bandwidth for a T4 is ~320 GB/s; A100-80GB is ~2,000 GB/s\n", " # We will parameterise the roofline later\n", "else:\n", " VRAM_GB = 16.0 # CPU fallback โ€” profiling still runs, just slower\n", " GPU_NAME = \"CPU (no CUDA)\"\n", " print(\" โš ๏ธ No GPU found โ€” results will be slow but correct\")\n", "\n", "print(\"\\nโœ… Environment ready.\")" ], "id": "cell_5" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.2 First-Principles: What Does the KV Cache Actually Cost?\n", "\n", "Before touching any code, let us derive the memory cost of the KV cache analytically. Understanding the formula is more important than running the benchmark โ€” it tells us *why* the system behaves the way it does.\n", "\n", "### The attention mechanism and the caching opportunity\n", "\n", "In multi-head self-attention, for each token $i$ in the sequence, the model computes:\n", "\n", "$$\\mathbf{q}_i, \\mathbf{k}_i, \\mathbf{v}_i = W_Q x_i,\\; W_K x_i,\\; W_V x_i$$\n", "\n", "When generating token $n+1$, we need keys and values for **all** previous tokens. In naive inference we recompute them from scratch every step โ€” $O(t^2)$ redundant work. The KV cache stores them once.\n", "\n", "### The memory cost formula\n", "\n", "For a model with:\n", "- $L$ transformer layers\n", "- $H$ attention heads \n", "- $d$ head dimension (typically $d_{model}/H$)\n", "- Sequence length $n$ (prefix) + $m$ (generated tokens)\n", "- Batch size $B$\n", "- Data type bytes $b$ (2 for float16)\n", "\n", "The KV cache stores **both** keys **and** values at **every** layer:\n", "\n", "$$\\text{KV Cache Bytes} = 2 \\times L \\times H \\times d \\times (n + m) \\times B \\times b$$\n", "\n", "The factor of 2 accounts for keys and values. Let us implement this formula and see what it predicts for our 7B model." ], "id": "cell_6" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@dataclass\n", "class ModelConfig:\n", " \"\"\"Architecture hyperparameters for a 7B-class transformer.\n", " \n", " These match the LexDraft production model described in the case study.\n", " We keep the model frozen โ€” only serving-layer behaviour is in scope.\n", " \"\"\"\n", " d_model: int = 4096 # Residual stream dimension\n", " n_heads: int = 32 # Number of attention heads\n", " n_layers: int = 32 # Number of transformer blocks\n", " vocab_size: int = 32000 # BPE tokenizer vocabulary\n", " max_seq_len: int = 12000 # Maximum supported context\n", "\n", " @property\n", " def head_dim(self) -> int:\n", " return self.d_model // self.n_heads # = 128\n", "\n", "\n", "def kv_cache_bytes(\n", " config: ModelConfig,\n", " seq_len: int,\n", " batch_size: int = 1,\n", " dtype_bytes: int = 2, # float16 โ†’ 2 bytes\n", ") -> int:\n", " \"\"\"Return the exact KV cache memory in bytes for given settings.\n", " \n", " Formula: 2 * L * H * d_head * seq_len * B * dtype_bytes\n", " ^ ^ ^ ^ ^ ^\n", " K&V lay head dim tokens batch\n", " \"\"\"\n", " return (\n", " 2 # keys + values\n", " * config.n_layers\n", " * config.n_heads\n", " * config.head_dim\n", " * seq_len\n", " * batch_size\n", " * dtype_bytes\n", " )\n", "\n", "\n", "# โ”€โ”€ Quick sanity check โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "cfg = ModelConfig()\n", "example_bytes = kv_cache_bytes(cfg, seq_len=8000, batch_size=1)\n", "print(f\"KV cache for n=8000, B=1 : {example_bytes / 1e9:.2f} GB\")\n", "print(f\"KV cache for n=8000, B=80 : {kv_cache_bytes(cfg, 8000, 80) / 1e9:.2f} GB\")\n", "print(f\"Model weights (FP16) : {cfg.d_model * cfg.n_layers * 12 * 2 / 1e9:.1f} GB approx\")" ], "id": "cell_7" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now let us **visualise** how KV cache memory grows across the ranges we care about. This single chart explains the OOM crashes better than any log file." ], "id": "cell_8" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "seq_lengths = np.arange(256, 12_001, 256)\n", "batch_sizes = [1, 8, 16, 32, 64, 80]\n", "\n", "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", "\n", "# โ”€โ”€ Left: KV cache per batch size โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "ax = axes[0]\n", "for B in batch_sizes:\n", " mem_gb = [kv_cache_bytes(cfg, n, B) / 1e9 for n in seq_lengths]\n", " ax.plot(seq_lengths, mem_gb, label=f\"B={B}\")\n", "\n", "# Mark the A100-40GB VRAM line (our p4d.24xlarge node)\n", "a100_vram = 40.0\n", "weight_footprint_gb = 14.0 # ~14 GB for a 7B model in FP16\n", "remaining_for_kv = a100_vram - weight_footprint_gb\n", "ax.axhline(remaining_for_kv, color=\"red\", linestyle=\"--\", linewidth=1.5,\n", " label=f\"VRAM budget for KV\\n({remaining_for_kv:.0f} GB after weights)\")\n", "ax.axvline(8000, color=\"gray\", linestyle=\":\", linewidth=1, label=\"Target max context\")\n", "\n", "ax.set_xlabel(\"Sequence Length (tokens)\", fontsize=11)\n", "ax.set_ylabel(\"KV Cache Memory (GB)\", fontsize=11)\n", "ax.set_title(\"KV Cache Growth vs. Sequence Length\\n(7B model, FP16)\", fontsize=12)\n", "ax.legend(fontsize=8, loc=\"upper left\")\n", "ax.grid(True, alpha=0.3)\n", "\n", "# โ”€โ”€ Right: Fraction of VRAM consumed by KV cache โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "ax2 = axes[1]\n", "for B in [1, 8, 32, 80]:\n", " frac = [kv_cache_bytes(cfg, n, B) / (a100_vram * 1e9) * 100\n", " for n in seq_lengths]\n", " ax2.plot(seq_lengths, frac, label=f\"B={B}\")\n", "\n", "ax2.axhline(100 * remaining_for_kv / a100_vram, color=\"red\", linestyle=\"--\",\n", " linewidth=1.5, label=\"Available after weights\")\n", "ax2.set_xlabel(\"Sequence Length (tokens)\", fontsize=11)\n", "ax2.set_ylabel(\"% of Total VRAM\", fontsize=11)\n", "ax2.set_title(\"KV Cache as % of A100-40GB VRAM\", fontsize=12)\n", "ax2.legend(fontsize=9)\n", "ax2.grid(True, alpha=0.3)\n", "\n", "plt.tight_layout()\n", "plt.savefig(\"kv_cache_memory_scaling.png\", dpi=150, bbox_inches=\"tight\")\n", "plt.show()\n", "\n", "print(\"\\n๐Ÿ“Š Key insight: At B=80 and n=8000, the KV cache alone exceeds \"\n", " f\"{kv_cache_bytes(cfg, 8000, 80)/1e9:.1f} GB โ€” \"\n", " \"completely overwhelming a 40 GB card after weights are loaded.\")" ], "id": "cell_9" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.3 Building a Minimal Transformer โ€” No Caching First\n", "\n", "We will now implement a **small but architecturally faithful** transformer. We deliberately scale the dimensions down so everything fits comfortably on a T4 (16 GB VRAM), but the shapes and operations mirror the production 7B model exactly โ€” just with fewer layers and a smaller hidden dimension.\n", "\n", "We start with the multi-head attention block because that is the heart of everything we care about." ], "id": "cell_10" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class MultiHeadAttention(nn.Module):\n", " \"\"\"Standard scaled dot-product multi-head attention.\n", " \n", " Intentionally written to be readable rather than maximally efficient โ€”\n", " we want to see exactly what happens at each step.\n", " \"\"\"\n", "\n", " def __init__(self, d_model: int, n_heads: int):\n", " super().__init__()\n", " assert d_model % n_heads == 0, \"d_model must be divisible by n_heads\"\n", " self.d_model = d_model\n", " self.n_heads = n_heads\n", " self.head_dim = d_model // n_heads\n", "\n", " # Projection matrices โ€” one linear for Q, K, V together for efficiency\n", " self.W_qkv = nn.Linear(d_model, 3 * d_model, bias=False)\n", " self.W_out = nn.Linear(d_model, d_model, bias=False)\n", "\n", " def forward(\n", " self,\n", " x: torch.Tensor, # (B, T, d_model)\n", " kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n", " use_causal_mask: bool = True,\n", " ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:\n", " \"\"\"\n", " Args:\n", " x : Input tensor of shape (B, T, d_model).\n", " During prefill, T = full context length.\n", " During decode, T = 1 (single new token).\n", " kv_cache : Optional (key_cache, value_cache) tensors from\n", " previous steps, each of shape (B, n_heads, T_prev, head_dim).\n", " use_causal_mask: Mask future positions (always True for autoregressive).\n", "\n", " Returns:\n", " out : (B, T, d_model)\n", " new_cache : Updated (key_cache, value_cache) for next step.\n", " \"\"\"\n", " B, T, _ = x.shape\n", "\n", " # โ”€โ”€ 1. Project to Q, K, V โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " qkv = self.W_qkv(x) # (B, T, 3*d_model)\n", " q, k, v = qkv.split(self.d_model, dim=-1) # each: (B, T, d_model)\n", "\n", " # โ”€โ”€ 2. Reshape into heads โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " def split_heads(t):\n", " return t.view(B, -1, self.n_heads, self.head_dim).transpose(1, 2)\n", " # โ†’ (B, n_heads, T, head_dim)\n", "\n", " q, k, v = split_heads(q), split_heads(k), split_heads(v)\n", "\n", " # โ”€โ”€ 3. Concatenate with KV cache if provided โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " if kv_cache is not None:\n", " k_cache, v_cache = kv_cache\n", " k = torch.cat([k_cache, k], dim=2) # append new key\n", " v = torch.cat([v_cache, v], dim=2) # append new value\n", " \n", " new_cache = (k, v) # store full history for next step\n", "\n", " # โ”€โ”€ 4. Scaled dot-product attention โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " T_kv = k.shape[2]\n", " scale = math.sqrt(self.head_dim)\n", " attn_scores = torch.matmul(q, k.transpose(-2, -1)) / scale\n", " # (B, n_heads, T_q, T_kv)\n", "\n", " if use_causal_mask:\n", " # Lower-triangular mask: position i can attend to positions โ‰ค i\n", " # When decoding with cache, T=1 so this mask is a 1ร—T_kv slice\n", " mask = torch.ones(T, T_kv, device=x.device, dtype=torch.bool)\n", " if T == T_kv:\n", " mask = torch.tril(mask)\n", " # else: decoding step โ€” the single new token can attend to all cache\n", " attn_scores = attn_scores.masked_fill(~mask, float(\"-inf\"))\n", "\n", " attn_weights = F.softmax(attn_scores, dim=-1)\n", " out = torch.matmul(attn_weights, v) # (B, n_heads, T, head_dim)\n", "\n", " # โ”€โ”€ 5. Merge heads and project โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)\n", " out = self.W_out(out)\n", "\n", " return out, new_cache" ], "id": "cell_11" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Good. Now let us build a single transformer block around this attention layer, and then stack blocks into a full (small) model." ], "id": "cell_12" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class TransformerBlock(nn.Module):\n", " \"\"\"One transformer layer: LayerNorm โ†’ Attention โ†’ Residual โ†’ FFN โ†’ Residual.\"\"\"\n", "\n", " def __init__(self, d_model: int, n_heads: int, ffn_mult: int = 4):\n", " super().__init__()\n", " self.ln1 = nn.LayerNorm(d_model)\n", " self.attn = MultiHeadAttention(d_model, n_heads)\n", " self.ln2 = nn.LayerNorm(d_model)\n", " self.ffn = nn.Sequential(\n", " nn.Linear(d_model, ffn_mult * d_model, bias=False),\n", " nn.GELU(),\n", " nn.Linear(ffn_mult * d_model, d_model, bias=False),\n", " )\n", "\n", " def forward(\n", " self,\n", " x: torch.Tensor,\n", " kv_cache: Optional[Tuple] = None,\n", " ) -> Tuple[torch.Tensor, Tuple]:\n", " # Pre-norm architecture (used by most modern LLMs)\n", " attn_out, new_cache = self.attn(self.ln1(x), kv_cache=kv_cache)\n", " x = x + attn_out\n", " x = x + self.ffn(self.ln2(x))\n", " return x, new_cache\n", "\n", "\n", "class SmallTransformer(nn.Module):\n", " \"\"\"A small but architecturally faithful transformer for profiling.\n", " \n", " Scaled-down 7B: d_model=512, n_heads=8, n_layers=6.\n", " Same attention patterns and caching logic; fits on a T4.\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " vocab_size: int = 32000,\n", " d_model: int = 512,\n", " n_heads: int = 8,\n", " n_layers: int = 6,\n", " max_seq_len: int = 4096,\n", " ):\n", " super().__init__()\n", " self.d_model = d_model\n", " self.n_heads = n_heads\n", " self.n_layers = n_layers\n", "\n", " self.embed = nn.Embedding(vocab_size, d_model)\n", " self.pos_embed = nn.Embedding(max_seq_len, d_model)\n", " self.blocks = nn.ModuleList([\n", " TransformerBlock(d_model, n_heads) for _ in range(n_layers)\n", " ])\n", " self.ln_final = nn.LayerNorm(d_model)\n", " self.lm_head = nn.Linear(d_model, vocab_size, bias=False)\n", "\n", " def forward(\n", " self,\n", " token_ids: torch.Tensor, # (B, T)\n", " kv_caches: Optional[List] = None, # one cache per layer\n", " position_offset: int = 0, # for decode steps: offset into pos table\n", " ) -> Tuple[torch.Tensor, List]:\n", " B, T = token_ids.shape\n", " positions = torch.arange(\n", " position_offset, position_offset + T,\n", " device=token_ids.device\n", " )\n", "\n", " x = self.embed(token_ids) + self.pos_embed(positions)\n", "\n", " new_caches = []\n", " for i, block in enumerate(self.blocks):\n", " cache_i = kv_caches[i] if kv_caches is not None else None\n", " x, new_cache = block(x, kv_cache=cache_i)\n", " new_caches.append(new_cache)\n", "\n", " logits = self.lm_head(self.ln_final(x)) # (B, T, vocab_size)\n", " return logits, new_caches\n", "\n", "\n", "# โ”€โ”€ Instantiate and inspect โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "model = SmallTransformer().to(DEVICE)\n", "param_count = sum(p.numel() for p in model.parameters())\n", "print(f\"๐Ÿ“ฆ Model parameter count : {param_count:,}\")\n", "print(f\" In memory (FP32) : {param_count * 4 / 1e6:.1f} MB\")\n", "print(f\" In memory (FP16) : {param_count * 2 / 1e6:.1f} MB\")" ], "id": "cell_13" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.4 Implementing Autoregressive Generation โ€” With and Without KV Cache\n", "\n", "This is the central experiment. We will implement two generation functions that are **mathematically identical** but computationally very different. The difference will make the business problem tangible.\n", "\n", "### TODO 3.4.1 โ€” Implement naive generation (no cache)\n", "\n", "Your task is to implement `generate_naive`, which regenerates **all** keys and values from scratch at every decoding step. This is the expensive baseline." ], "id": "cell_14" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def generate_naive(\n", " model: SmallTransformer,\n", " prompt_ids: torch.Tensor, # (B, n_prefix) โ€” the context tokens\n", " n_new_tokens: int = 32,\n", ") -> Tuple[torch.Tensor, List[float]]:\n", " \"\"\"Generate tokens WITHOUT any KV caching.\n", " \n", " At each step t, we feed the ENTIRE sequence (prefix + generated so far)\n", " through the model. This means all keys and values are recomputed from\n", " scratch at every step โ€” O(t) matrix multiplications per step, O(tยฒ) total.\n", " \n", " Args:\n", " model : The SmallTransformer instance.\n", " prompt_ids : Integer token ids of shape (B, n_prefix).\n", " n_new_tokens : How many new tokens to generate.\n", " \n", " Returns:\n", " generated_ids : (B, n_prefix + n_new_tokens)\n", " step_times_ms : Wall-clock time per generation step in milliseconds.\n", " \n", " TODO: Implement this function.\n", " \n", " HINTS:\n", " Step 1 โ€” Convert prompt_ids to the right device.\n", " Step 2 โ€” Start a loop over range(n_new_tokens).\n", " Step 3 โ€” Inside the loop, call model.forward() on the FULL current sequence\n", " (no kv_caches argument โ€” pass None).\n", " Step 4 โ€” Take the logits at the LAST position: logits[:, -1, :].\n", " Step 5 โ€” Greedy decode: argmax over vocabulary dimension.\n", " Step 6 โ€” Append the new token id to the sequence (torch.cat along dim=1).\n", " Step 7 โ€” Time each step using time.perf_counter() and append to a list.\n", " Step 8 โ€” Return (full_sequence, step_times).\n", " \"\"\"\n", " # โ”€โ”€ YOUR CODE HERE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " raise NotImplementedError(\"Implement generate_naive โ€” see hints above.\")" ], "id": "cell_15" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# โœ… Verification 3.4.1\n", "torch.manual_seed(SEED)\n", "_test_prompt = torch.randint(0, 32000, (1, 16)).to(DEVICE)\n", "model.eval()\n", "with torch.no_grad():\n", " try:\n", " _out_ids, _times = generate_naive(model, _test_prompt, n_new_tokens=5)\n", " assert _out_ids.shape == (1, 21), \\\n", " f\"โŒ Expected shape (1, 21), got {_out_ids.shape}\"\n", " assert len(_times) == 5, \\\n", " f\"โŒ Expected 5 timing entries, got {len(_times)}\"\n", " print(f\"โœ… generate_naive works! Output shape: {_out_ids.shape}\")\n", " print(f\" Step times (ms): {[f'{t:.1f}' for t in _times]}\")\n", " except NotImplementedError:\n", " print(\"โณ Not yet implemented โ€” complete the TODO above first.\")" ], "id": "cell_16" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### TODO 3.4.2 โ€” Implement cached generation\n", "\n", "Now implement the efficient version. The key insight: during the **prefill** phase we process the entire prompt in a single parallel forward pass and store the resulting KV cache. During the **decode** phase we only feed the single new token at each step." ], "id": "cell_17" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def generate_with_cache(\n", " model: SmallTransformer,\n", " prompt_ids: torch.Tensor, # (B, n_prefix)\n", " n_new_tokens: int = 32,\n", ") -> Tuple[torch.Tensor, List[float], float]:\n", " \"\"\"Generate tokens WITH KV caching.\n", " \n", " Phase 1 โ€” PREFILL: Run the full prompt through the model once.\n", " This populates the KV cache and gives us the first logits.\n", " Phase 2 โ€” DECODE: At each step, pass only the single new token.\n", " The model appends to the cache; we never recompute old K/V.\n", " \n", " Args:\n", " model : The SmallTransformer instance.\n", " prompt_ids : Integer token ids of shape (B, n_prefix).\n", " n_new_tokens : How many new tokens to generate.\n", " \n", " Returns:\n", " generated_ids : (B, n_prefix + n_new_tokens)\n", " step_times_ms : Wall-clock time per DECODE step (not prefill).\n", " prefill_time_ms : Time for the prefill phase (this is your TTFT proxy).\n", " \n", " TODO: Implement this function.\n", " \n", " HINTS:\n", " Step 1 โ€” PREFILL: Call model.forward(prompt_ids, kv_caches=None).\n", " Time this separately โ€” it is your TTFT proxy.\n", " Save the returned kv_caches as `cache`.\n", " Step 2 โ€” Greedy-decode the LAST logit position to get the first new token.\n", " Append it to prompt_ids to form `current_ids`.\n", " Step 3 โ€” DECODE LOOP over range(n_new_tokens - 1):\n", " a. Extract just the LAST token: current_ids[:, -1:] โ€” shape (B,1)\n", " b. Call model.forward() with kv_caches=cache AND\n", " position_offset = current_ids.shape[1] - 1\n", " (so positional embeddings continue from the right index)\n", " c. Update cache with the new kv_caches returned by the model.\n", " d. Greedy-decode logits[:, 0, :] (only one position now).\n", " e. Append to current_ids.\n", " f. Record step time.\n", " Step 4 โ€” Return (current_ids, step_times, prefill_time).\n", " \"\"\"\n", " # โ”€โ”€ YOUR CODE HERE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " raise NotImplementedError(\"Implement generate_with_cache โ€” see hints above.\")" ], "id": "cell_18" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# โœ… Verification 3.4.2\n", "torch.manual_seed(SEED)\n", "_test_prompt2 = torch.randint(0, 32000, (1, 16)).to(DEVICE)\n", "model.eval()\n", "with torch.no_grad():\n", " try:\n", " _out_ids2, _times2, _prefill = generate_with_cache(\n", " model, _test_prompt2, n_new_tokens=5\n", " )\n", " assert _out_ids2.shape == (1, 21), \\\n", " f\"โŒ Expected shape (1, 21), got {_out_ids2.shape}\"\n", " assert isinstance(_prefill, float), \"โŒ prefill_time_ms should be a float\"\n", " print(f\"โœ… generate_with_cache works! Output shape: {_out_ids2.shape}\")\n", " print(f\" Prefill time : {_prefill:.2f} ms\")\n", " print(f\" Decode times : {[f'{t:.1f}' for t in _times2]} ms\")\n", " except NotImplementedError:\n", " print(\"โณ Not yet implemented โ€” complete the TODO above first.\")" ], "id": "cell_19" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.5 Profiling: How Does Latency Scale With Sequence Length?\n", "\n", "With both generators implemented, let us run the core profiling experiment. We sweep over a range of prompt lengths and measure how step times scale โ€” this directly mirrors the production problem where TTFT spikes at long context." ], "id": "cell_20" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def profile_generation_methods(\n", " model: SmallTransformer,\n", " context_lengths: List[int],\n", " n_new_tokens: int = 20,\n", " n_warmup: int = 2,\n", " batch_size: int = 1,\n", ") -> Dict:\n", " \"\"\"Sweep over context lengths and record timing for both methods.\n", " \n", " We run a few warmup passes first so that CUDA kernels are compiled\n", " and cached โ€” otherwise our first measurements would be artificially slow.\n", " \"\"\"\n", " model.eval()\n", " results = {\n", " \"context_lengths\" : context_lengths,\n", " \"naive_ttft_ms\" : [],\n", " \"naive_mean_step_ms\" : [],\n", " \"cached_prefill_ms\" : [],\n", " \"cached_mean_step_ms\": [],\n", " }\n", "\n", " # Warmup\n", " print(\"๐Ÿ”ฅ Warming up CUDA kernels...\")\n", " _w = torch.randint(0, 32000, (batch_size, 64)).to(DEVICE)\n", " for _ in range(n_warmup):\n", " with torch.no_grad():\n", " model(_w)\n", " if DEVICE.type == \"cuda\":\n", " torch.cuda.synchronize()\n", "\n", " print(f\"\\n๐Ÿ“ Profiling {len(context_lengths)} context lengths \"\n", " f\"({n_new_tokens} tokens generated each)...\\n\")\n", "\n", " for ctx_len in context_lengths:\n", " prompt = torch.randint(0, 32000, (batch_size, ctx_len)).to(DEVICE)\n", "\n", " with torch.no_grad():\n", " # โ”€โ”€ Naive (no cache) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " if DEVICE.type == \"cuda\":\n", " torch.cuda.synchronize()\n", " t0 = time.perf_counter()\n", " _, naive_times = generate_naive(model, prompt, n_new_tokens)\n", " if DEVICE.type == \"cuda\":\n", " torch.cuda.synchronize()\n", "\n", " # Naive has no explicit prefill phase โ€” the first step IS the prefill\n", " results[\"naive_ttft_ms\"].append(naive_times[0])\n", " results[\"naive_mean_step_ms\"].append(np.mean(naive_times))\n", "\n", " # โ”€โ”€ Cached โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " _, cached_times, prefill_ms = generate_with_cache(\n", " model, prompt, n_new_tokens\n", " )\n", " results[\"cached_prefill_ms\"].append(prefill_ms)\n", " results[\"cached_mean_step_ms\"].append(\n", " np.mean(cached_times) if cached_times else 0.0\n", " )\n", "\n", " print(f\" n={ctx_len:5d} | naive TTFT: {naive_times[0]:6.1f} ms |\"\n", " f\" cached prefill: {prefill_ms:6.1f} ms |\"\n", " f\" cached step: {np.mean(cached_times) if cached_times else 0:.2f} ms\")\n", "\n", " return results\n", "\n", "\n", "# โ”€โ”€ Run the profiling sweep โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "CONTEXT_LENGTHS = [64, 128, 256, 512, 768, 1024]\n", "# Note: On a T4 we keep lengths modest. The analytical curves extrapolate to 8K.\n", "\n", "profile_data = profile_generation_methods(model, CONTEXT_LENGTHS, n_new_tokens=15)\n", "print(\"\\nโœ… Profiling complete.\")" ], "id": "cell_21" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let us now visualise the profiling results. This is the chart that, in the real case study, convinced the CFO that the problem was architectural, not hardware-capacity." ], "id": "cell_22" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ctx = profile_data[\"context_lengths\"]\n", "\n", "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", "\n", "# โ”€โ”€ Left: First-token latency (TTFT proxy) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "ax = axes[0]\n", "ax.plot(ctx, profile_data[\"naive_ttft_ms\"], \"o-\",\n", " color=\"crimson\", label=\"Naive (no cache) โ€” first step\")\n", "ax.plot(ctx, profile_data[\"cached_prefill_ms\"], \"s-\",\n", " color=\"steelblue\", label=\"Cached โ€” prefill phase\")\n", "\n", "ax.set_xlabel(\"Context Length (tokens)\", fontsize=11)\n", "ax.set_ylabel(\"Time (ms)\", fontsize=11)\n", "ax.set_title(\"First-Token Latency vs. Context Length\\n\"\n", " \"(lower is better; production target โ‰ค 300 ms)\", fontsize=12)\n", "ax.legend(fontsize=10)\n", "ax.grid(True, alpha=0.3)\n", "\n", "# Annotate the production threshold\n", "ax.axhline(300, color=\"orange\", linestyle=\"--\", linewidth=1.5,\n", " label=\"300 ms SLA\")\n", "ax.legend(fontsize=9)\n", "\n", "# โ”€โ”€ Right: Per-step decode latency โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "ax2 = axes[1]\n", "ax2.plot(ctx, profile_data[\"naive_mean_step_ms\"], \"o-\",\n", " color=\"crimson\", label=\"Naive โ€” avg step (grows with ctx!)\")\n", "ax2.plot(ctx, profile_data[\"cached_mean_step_ms\"], \"s-\",\n", " color=\"steelblue\", label=\"Cached โ€” avg decode step (flat!)\")\n", "\n", "ax2.set_xlabel(\"Context Length (tokens)\", fontsize=11)\n", "ax2.set_ylabel(\"Avg Step Time (ms)\", fontsize=11)\n", "ax2.set_title(\"Per-Step Decode Latency vs. Context Length\\n\"\n", " \"(cached should be approximately flat)\", fontsize=12)\n", "ax2.legend(fontsize=9)\n", "ax2.grid(True, alpha=0.3)\n", "\n", "plt.tight_layout()\n", "plt.savefig(\"latency_profiling.png\", dpi=150, bbox_inches=\"tight\")\n", "plt.show()\n", "\n", "# โ”€โ”€ Print speedup table โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "print(\"\\n๐Ÿ“Š Speedup Table (Cached vs. Naive):\")\n", "print(f\"{'Context':>10} {'Naive TTFT':>12} {'Cached Prefill':>16} {'Decode Speedup':>16}\")\n", "print(\"-\" * 60)\n", "for i, n in enumerate(ctx):\n", " naive_step = profile_data[\"naive_mean_step_ms\"][i]\n", " cached_step = profile_data[\"cached_mean_step_ms\"][i]\n", " speedup = naive_step / cached_step if cached_step > 0 else float(\"inf\")\n", " print(f\"{n:>10} {profile_data['naive_ttft_ms'][i]:>11.1f}ms\"\n", " f\" {profile_data['cached_prefill_ms'][i]:>15.1f}ms\"\n", " f\" {speedup:>14.1f}x\")" ], "id": "cell_23" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.6 The Roofline Model โ€” Classifying Our Inference Regime\n", "\n", "The most important insight in this case study is that the GPU showing **low utilisation during high latency** is not a paradox โ€” it is a diagnostic. The roofline model explains exactly why.\n", "\n", "### Two ceilings, one roof\n", "\n", "A GPU has two fundamental performance limits:\n", "1. **Compute ceiling**: Peak FLOP/s (how fast it can do arithmetic)\n", "2. **Memory bandwidth ceiling**: Peak GB/s (how fast it can read/write memory)\n", "\n", "An operation's **arithmetic intensity** (AI) โ€” measured in FLOPs per byte โ€” determines which ceiling it hits:\n", "\n", "$$\\text{Arithmetic Intensity} = \\frac{\\text{Total FLOPs}}{\\text{Total bytes moved}}$$\n", "\n", "- **High AI (compute-bound)**: matrix multiplies with large dimensions โ€” the GPU's arithmetic units are the bottleneck\n", "- **Low AI (memory-bound)**: vector operations, loading large tensors โ€” the memory bus is the bottleneck\n", "\n", "Let us compute the arithmetic intensity of prefill vs. decode and see which regime each falls into." ], "id": "cell_24" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def compute_arithmetic_intensity(\n", " config: ModelConfig,\n", " seq_len: int,\n", " is_prefill: bool,\n", " batch_size: int = 1,\n", ") -> Tuple[float, float]:\n", " \"\"\"Estimate arithmetic intensity for one transformer layer.\n", " \n", " Returns (flops, bytes_moved) so we can plot both dimensions.\n", " \n", " For simplicity we focus on the attention QKV projections and the\n", " attention score computation โ€” these dominate at long sequence lengths.\n", " \n", " PREFILL: Process seq_len tokens in parallel.\n", " QKV projection: 2 * seq_len * d_model * (3*d_model) FLOPs\n", " Attention scores: 2 * n_heads * seq_len^2 * head_dim FLOPs\n", " Bytes: weights + activations loaded from HBM\n", " \n", " DECODE: Process 1 token, read entire KV cache.\n", " QKV projection: 2 * 1 * d_model * (3*d_model) FLOPs (trivial)\n", " Attention scores: 2 * n_heads * seq_len * head_dim FLOPs\n", " Bytes: KV cache of ALL previous tokens must be read from HBM\n", " \"\"\"\n", " d = config.d_model\n", " H = config.n_heads\n", " dh = config.head_dim\n", " L = config.n_layers\n", " B = batch_size\n", " n = seq_len\n", "\n", " if is_prefill:\n", " # QKV projection across all tokens\n", " flops_qkv = 2 * n * d * (3 * d) * L * B\n", " # Attention score matrix: QยทKแต€ for all heads\n", " flops_attn = 2 * H * n * n * dh * L * B\n", " total_flops = flops_qkv + flops_attn\n", "\n", " # Bytes: load model weights (QKV projection matrix) + input activations\n", " bytes_weights = 3 * d * d * 2 * L # FP16 weights per layer\n", " bytes_activations = n * d * 2 * B * L # input/output activations\n", " total_bytes = bytes_weights + bytes_activations\n", "\n", " else: # decode\n", " # QKV for single token โ€” tiny compute\n", " flops_qkv = 2 * 1 * d * (3 * d) * L * B\n", " # Attention: single query vs. all n cached keys\n", " flops_attn = 2 * H * 1 * n * dh * L * B\n", " total_flops = flops_qkv + flops_attn\n", "\n", " # Bytes: MUST load entire KV cache from HBM for every step\n", " bytes_kv_cache = kv_cache_bytes(config, n, B, dtype_bytes=2)\n", " bytes_weights = 3 * d * d * 2 * L\n", " total_bytes = bytes_kv_cache + bytes_weights\n", "\n", " ai = total_flops / total_bytes if total_bytes > 0 else 0\n", " return total_flops, total_bytes, ai\n", "\n", "\n", "# โ”€โ”€ Sweep AI over sequence lengths โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "seq_range = np.arange(256, 8_001, 256)\n", "\n", "prefill_ai = [compute_arithmetic_intensity(cfg, n, is_prefill=True)[2]\n", " for n in seq_range]\n", "decode_ai = [compute_arithmetic_intensity(cfg, n, is_prefill=False)[2]\n", " for n in seq_range]\n", "\n", "print(\"Sample Arithmetic Intensities (FLOPs/Byte):\")\n", "for n in [512, 2048, 4096, 8000]:\n", " _, _, ai_p = compute_arithmetic_intensity(cfg, n, is_prefill=True)\n", " _, _, ai_d = compute_arithmetic_intensity(cfg, n, is_prefill=False)\n", " print(f\" n={n:5d} prefill AI: {ai_p:6.1f} decode AI: {ai_d:5.2f}\")" ], "id": "cell_25" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now let us draw the roofline chart. This single visualisation is arguably the most important output of the entire case study." ], "id": "cell_26" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# โ”€โ”€ Hardware specs โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "# A100-40GB (production). T4 shown for comparison.\n", "HW = {\n", " \"A100-40GB\": {\"peak_tflops\": 312.0, \"bandwidth_tbps\": 1.555}, # FP16\n", " \"T4\" : {\"peak_tflops\": 65.0, \"bandwidth_tbps\": 0.320},\n", "}\n", "HW_NAME = \"A100-40GB\" # We model the production system\n", "peak_tflops = HW[HW_NAME][\"peak_tflops\"] # tera-FLOP/s\n", "bandwidth_tbps = HW[HW_NAME][\"bandwidth_tbps\"] # tera-Byte/s\n", "\n", "# Ridge point: AI where compute = bandwidth * AI โ†’ AI* = peak_FLOPS / bandwidth\n", "ridge_point = (peak_tflops * 1e12) / (bandwidth_tbps * 1e12)\n", "print(f\"Hardware: {HW_NAME}\")\n", "print(f\" Peak compute : {peak_tflops} TFLOP/s (FP16)\")\n", "print(f\" Memory BW : {bandwidth_tbps * 1e3:.0f} GB/s\")\n", "print(f\" Ridge point : {ridge_point:.1f} FLOP/Byte\")\n", "print(f\"\\n โ†’ Operations with AI < {ridge_point:.0f} are MEMORY-BOUND\")\n", "print(f\" โ†’ Operations with AI > {ridge_point:.0f} are COMPUTE-BOUND\")\n", "\n", "# โ”€โ”€ Build roofline plot โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "ai_range = np.logspace(-2, 4, 500)\n", "\n", "# Memory-bound ceiling: perf = BW ร— AI\n", "mem_roof = bandwidth_tbps * 1e3 * ai_range # GFLOP/s\n", "# Compute ceiling: flat at peak\n", "comp_roof = np.full_like(ai_range, peak_tflops * 1e3) # GFLOP/s\n", "roofline = np.minimum(mem_roof, comp_roof)\n", "\n", "fig, ax = plt.subplots(figsize=(11, 7))\n", "ax.loglog(ai_range, roofline, \"k-\", linewidth=2.5, label=\"Roofline\")\n", "ax.fill_between(ai_range, roofline, alpha=0.05, color=\"black\")\n", "\n", "# Shade the memory-bound region\n", "ax.axvline(ridge_point, color=\"gray\", linestyle=\"--\", alpha=0.7,\n", " label=f\"Ridge point ({ridge_point:.0f} FLOP/Byte)\")\n", "ax.text(ridge_point * 0.3, peak_tflops * 30, \"MEMORY\\nBOUND\",\n", " fontsize=11, color=\"steelblue\", ha=\"center\", alpha=0.8)\n", "ax.text(ridge_point * 3.0, peak_tflops * 30, \"COMPUTE\\nBOUND\",\n", " fontsize=11, color=\"darkgreen\", ha=\"center\", alpha=0.8)\n", "\n", "# โ”€โ”€ Plot our operations โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "# Use the 7B config values at representative sequence lengths\n", "sample_ns = [512, 2048, 4096, 8000]\n", "colors_p = plt.cm.Blues(np.linspace(0.4, 0.9, len(sample_ns)))\n", "colors_d = plt.cm.Reds( np.linspace(0.4, 0.9, len(sample_ns)))\n", "\n", "for i, n in enumerate(sample_ns):\n", " _, _, ai_p = compute_arithmetic_intensity(cfg, n, is_prefill=True)\n", " _, _, ai_d = compute_arithmetic_intensity(cfg, n, is_prefill=False)\n", "\n", " perf_p = min(bandwidth_tbps * 1e3 * ai_p, peak_tflops * 1e3)\n", " perf_d = min(bandwidth_tbps * 1e3 * ai_d, peak_tflops * 1e3)\n", "\n", " ax.scatter(ai_p, perf_p, s=100, color=colors_p[i], zorder=5,\n", " marker=\"^\", label=f\"Prefill n={n}\" if i == 0 else \"\")\n", " ax.scatter(ai_d, perf_d, s=100, color=colors_d[i], zorder=5,\n", " marker=\"v\", label=f\"Decode n={n}\" if i == 0 else \"\")\n", " ax.annotate(f\"P{n}\", (ai_p, perf_p), textcoords=\"offset points\",\n", " xytext=(4, 4), fontsize=7, color=colors_p[i])\n", " ax.annotate(f\"D{n}\", (ai_d, perf_d), textcoords=\"offset points\",\n", " xytext=(4, -12), fontsize=7, color=colors_d[i])\n", "\n", "ax.set_xlabel(\"Arithmetic Intensity (FLOP/Byte)\", fontsize=12)\n", "ax.set_ylabel(\"Attainable Performance (GFLOP/s)\", fontsize=12)\n", "ax.set_title(f\"Roofline Model โ€” {HW_NAME}\\n\"\n", " \"โ–ฒ = Prefill โ–ผ = Decode (7B model config, varying context length)\",\n", " fontsize=12)\n", "ax.legend(fontsize=9, loc=\"lower right\")\n", "ax.grid(True, which=\"both\", alpha=0.2)\n", "ax.set_xlim([1e-2, 1e4])\n", "\n", "plt.tight_layout()\n", "plt.savefig(\"roofline_model.png\", dpi=150, bbox_inches=\"tight\")\n", "plt.show()" ], "id": "cell_27" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The roofline chart tells the whole story. Notice that all **decode** points (โ–ผ) cluster in the deep memory-bound region โ€” their arithmetic intensity is far below the ridge point. The GPU is doing very little arithmetic per byte of memory it moves. This is why compute utilisation appears low: the GPU is spending most of its time *waiting for the KV cache to be streamed from HBM*, not performing FLOPs.\n", "\n", "The **prefill** points (โ–ฒ) are higher but still memory-bound at long contexts โ€” and their absolute GFLOP cost grows quadratically with sequence length, which explains the TTFT spike." ], "id": "cell_28" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.7 Memory Occupancy Profiling\n", "\n", "Let us now measure actual VRAM consumption as a function of context length and batch size on our small model. On a T4 this gives us real numbers; we will then analytically extrapolate to the 7B production model.\n", "\n", "### TODO 3.7.1 โ€” Implement VRAM measurement during inference" ], "id": "cell_29" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def measure_vram_mb(\n", " model: SmallTransformer,\n", " batch_size: int,\n", " context_len: int,\n", " n_new_tokens: int = 10,\n", " use_cache: bool = True,\n", ") -> Dict[str, float]:\n", " \"\"\"Profile peak VRAM usage during cached generation.\n", " \n", " Returns a dict with keys:\n", " 'peak_mb' : peak VRAM allocated during generation (MB)\n", " 'model_mb' : VRAM for model parameters alone (MB)\n", " 'cache_mb' : Inferred KV cache footprint (MB) = peak - model\n", " \n", " TODO: Implement this function.\n", " \n", " HINTS:\n", " Step 1 โ€” Compute model_mb by summing (p.numel() * 4 / 1e6)\n", " for all parameters (we run in FP32 here for simplicity).\n", " Step 2 โ€” If DEVICE is CPU, return estimated values using kv_cache_bytes().\n", " Step 3 โ€” For CUDA:\n", " a. Flush cached memory: torch.cuda.empty_cache()\n", " b. Reset peak stats: torch.cuda.reset_peak_memory_stats()\n", " c. Create a random prompt of shape (batch_size, context_len).\n", " d. Run generate_with_cache (or generate_naive) inside torch.no_grad().\n", " e. Read peak memory: torch.cuda.max_memory_allocated() / 1e6 (MB)\n", " f. Compute cache_mb = peak_mb - model_mb (clip to โ‰ฅ 0).\n", " Step 4 โ€” Return the dict.\n", " \n", " Note: This is an approximation โ€” activations are also included in peak_mb.\n", " For our purposes the decomposition is illustrative.\n", " \"\"\"\n", " # โ”€โ”€ YOUR CODE HERE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " raise NotImplementedError(\"Implement measure_vram_mb โ€” see hints above.\")" ], "id": "cell_30" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# โœ… Verification 3.7.1\n", "model.eval()\n", "try:\n", " _vram = measure_vram_mb(model, batch_size=1, context_len=128, n_new_tokens=5)\n", " assert \"peak_mb\" in _vram, \"โŒ Missing 'peak_mb' key\"\n", " assert \"model_mb\" in _vram, \"โŒ Missing 'model_mb' key\"\n", " assert \"cache_mb\" in _vram, \"โŒ Missing 'cache_mb' key\"\n", " assert _vram[\"peak_mb\"] > 0, \"โŒ peak_mb should be > 0\"\n", " print(f\"โœ… measure_vram_mb works!\")\n", " print(f\" Model : {_vram['model_mb']:.1f} MB\")\n", " print(f\" Peak : {_vram['peak_mb']:.1f} MB\")\n", " print(f\" Cache : {_vram['cache_mb']:.1f} MB\")\n", "except NotImplementedError:\n", " print(\"โณ Not yet implemented โ€” complete the TODO above first.\")" ], "id": "cell_31" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# โ”€โ”€ Run VRAM sweep (small model, representative batch/context combinations) โ”€โ”€\n", "print(\"๐Ÿ“Š Running VRAM sweep...\\n\")\n", "\n", "vram_context_lengths = [64, 128, 256, 512, 1024]\n", "vram_batch_sizes = [1, 4, 8]\n", "\n", "vram_results = {}\n", "model.eval()\n", "\n", "for B in vram_batch_sizes:\n", " vram_results[B] = []\n", " for ctx in vram_context_lengths:\n", " try:\n", " result = measure_vram_mb(model, B, ctx, n_new_tokens=8)\n", " vram_results[B].append(result[\"peak_mb\"])\n", " print(f\" B={B:2d}, n={ctx:5d} โ†’ peak {result['peak_mb']:6.1f} MB \"\n", " f\"(model {result['model_mb']:.0f} MB + \"\n", " f\"cache ~{result['cache_mb']:.0f} MB)\")\n", " except NotImplementedError:\n", " vram_results[B].append(0.0)\n", "\n", "# โ”€โ”€ Visualise: measured small model + analytical 7B projection โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", "\n", "# Left: measured small model\n", "ax = axes[0]\n", "for B in vram_batch_sizes:\n", " if any(v > 0 for v in vram_results[B]):\n", " ax.plot(vram_context_lengths, vram_results[B], \"o-\", label=f\"B={B}\")\n", "ax.set_xlabel(\"Context Length (tokens)\", fontsize=11)\n", "ax.set_ylabel(\"Peak VRAM (MB)\", fontsize=11)\n", "ax.set_title(\"Measured VRAM: SmallTransformer\\n(d=512, 6 layers)\", fontsize=12)\n", "ax.legend(); ax.grid(True, alpha=0.3)\n", "\n", "# Right: analytical 7B projection\n", "ax2 = axes[1]\n", "a7b_ctx = np.arange(256, 12_001, 256)\n", "for B in [1, 8, 32, 80]:\n", " weight_gb = 14.0 # ~14 GB for 7B FP16\n", " kv_gb = [kv_cache_bytes(cfg, n, B) / 1e9 for n in a7b_ctx]\n", " total_gb = [weight_gb + k for k in kv_gb]\n", " ax2.plot(a7b_ctx, total_gb, label=f\"B={B}\")\n", "\n", "ax2.axhline(40.0, color=\"red\", linestyle=\"--\", linewidth=1.5,\n", " label=\"A100-40GB limit\")\n", "ax2.axvline(8000, color=\"gray\", linestyle=\":\", linewidth=1, label=\"Target ctx\")\n", "ax2.set_xlabel(\"Context Length (tokens)\", fontsize=11)\n", "ax2.set_ylabel(\"Estimated VRAM (GB)\", fontsize=11)\n", "ax2.set_title(\"Analytical VRAM: 7B Production Model\\n(weights + KV cache, FP16)\",\n", " fontsize=12)\n", "ax2.legend(fontsize=8); ax2.grid(True, alpha=0.3)\n", "\n", "plt.tight_layout()\n", "plt.savefig(\"vram_occupancy.png\", dpi=150, bbox_inches=\"tight\")\n", "plt.show()" ], "id": "cell_32" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.8 Mitigation Strategies โ€” Three Levers We Can Actually Pull\n", "\n", "Now that we have diagnosed the problem rigorously, let us explore the mitigation strategies available at the serving layer. Remember: the model weights are frozen. We can only touch how we *serve* the model.\n", "\n", "### Strategy 1: KV Cache Quantisation\n", "\n", "Instead of storing keys and values in FP16 (2 bytes each), we can quantise them to INT8 (1 byte) or even INT4 (0.5 bytes). The memory savings are proportional, at the cost of a small quality degradation." ], "id": "cell_33" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def quantise_kv_cache_int8(\n", " cache: Tuple[torch.Tensor, torch.Tensor]\n", ") -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n", " \"\"\"Quantise a KV cache entry from FP16/FP32 to INT8.\n", " \n", " We use symmetric per-tensor quantisation: \n", " q = round(x / scale), scale = max(|x|) / 127\n", " \n", " Args:\n", " cache: (keys, values) each of shape (B, n_heads, T, head_dim)\n", " \n", " Returns:\n", " (keys_int8, values_int8, key_scale, val_scale)\n", " โ€” the scales are needed for dequantisation during attention.\n", " \"\"\"\n", " k, v = cache\n", "\n", " def quant(t: torch.Tensor):\n", " scale = t.abs().max() / 127.0 + 1e-8\n", " quantised = (t / scale).round().clamp(-128, 127).to(torch.int8)\n", " return quantised, scale\n", "\n", " k_q, k_s = quant(k)\n", " v_q, v_s = quant(v)\n", " return k_q, v_q, k_s, v_s\n", "\n", "\n", "def dequantise_kv_cache_int8(\n", " k_q: torch.Tensor, v_q: torch.Tensor,\n", " k_s: torch.Tensor, v_s: torch.Tensor,\n", " target_dtype: torch.dtype = torch.float32,\n", ") -> Tuple[torch.Tensor, torch.Tensor]:\n", " \"\"\"Dequantise INT8 KV cache back to float for attention computation.\"\"\"\n", " k = k_q.to(target_dtype) * k_s\n", " v = v_q.to(target_dtype) * v_s\n", " return k, v\n", "\n", "\n", "# โ”€โ”€ Demonstrate memory reduction โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "k_demo = torch.randn(1, 8, 512, 64, device=DEVICE) # a mini KV cache\n", "v_demo = torch.randn(1, 8, 512, 64, device=DEVICE)\n", "\n", "k_q, v_q, k_s, v_s = quantise_kv_cache_int8((k_demo, v_demo))\n", "\n", "bytes_fp32 = k_demo.nbytes + v_demo.nbytes\n", "bytes_int8 = k_q.nbytes + v_q.nbytes\n", "\n", "print(\"KV Cache Quantisation Demo\")\n", "print(f\" FP32 cache size : {bytes_fp32 / 1e6:.2f} MB\")\n", "print(f\" INT8 cache size : {bytes_int8 / 1e6:.2f} MB ({bytes_int8/bytes_fp32*100:.0f}% of original)\")\n", "print(f\" Memory saved : {(bytes_fp32 - bytes_int8) / 1e6:.2f} MB ({(1-bytes_int8/bytes_fp32)*100:.0f}% reduction)\")\n", "\n", "# Verify reconstruction quality\n", "k_rec, v_rec = dequantise_kv_cache_int8(k_q, v_q, k_s, v_s,\n", " target_dtype=k_demo.dtype)\n", "k_err = (k_rec - k_demo).abs().mean().item()\n", "v_err = (v_rec - v_demo).abs().mean().item()\n", "print(f\"\\n Reconstruction error (mean absolute):\")\n", "print(f\" Keys : {k_err:.5f}\")\n", "print(f\" Values : {v_err:.5f}\")\n", "print(f\" โ†’ Errors are small relative to typical activation magnitudes โœ…\")" ], "id": "cell_34" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Strategy 2: Sliding Window / Cache Eviction\n", "\n", "For very long documents, we may not need the full history. A sliding window keeps only the most recent $W$ tokens in the cache, bounding memory at $O(W)$ rather than $O(n)$." ], "id": "cell_35" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def apply_sliding_window(\n", " cache: Tuple[torch.Tensor, torch.Tensor],\n", " window_size: int,\n", ") -> Tuple[torch.Tensor, torch.Tensor]:\n", " \"\"\"Truncate the KV cache to the most recent `window_size` tokens.\n", " \n", " In legal drafting, the immediately preceding clauses matter more than\n", " clauses from page 1 โ€” this heuristic often holds in practice.\n", " \"\"\"\n", " k, v = cache\n", " T = k.shape[2]\n", " if T <= window_size:\n", " return cache # Nothing to evict\n", " return k[:, :, -window_size:, :], v[:, :, -window_size:, :]\n", "\n", "\n", "# โ”€โ”€ Visualise the memory trade-off โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "window_sizes = [256, 512, 1024, 2048, 4096]\n", "ctx_full = np.arange(256, 8001, 256)\n", "\n", "fig, ax = plt.subplots(figsize=(10, 5))\n", "\n", "# Full cache (no eviction)\n", "full_gb = [kv_cache_bytes(cfg, n, batch_size=80) / 1e9 for n in ctx_full]\n", "ax.plot(ctx_full, full_gb, \"k-\", linewidth=2, label=\"Full cache (B=80)\")\n", "\n", "# Sliding windows\n", "colors = plt.cm.viridis(np.linspace(0.2, 0.85, len(window_sizes)))\n", "for w, c in zip(window_sizes, colors):\n", " windowed_gb = [kv_cache_bytes(cfg, min(n, w), batch_size=80) / 1e9\n", " for n in ctx_full]\n", " ax.plot(ctx_full, windowed_gb, \"--\", color=c, label=f\"Window W={w}\")\n", "\n", "ax.axhline(40 - 14, color=\"red\", linestyle=\":\", linewidth=1.5,\n", " label=\"Available VRAM after weights\")\n", "ax.set_xlabel(\"Document Length (tokens)\", fontsize=11)\n", "ax.set_ylabel(\"KV Cache Memory (GB) [B=80]\", fontsize=11)\n", "ax.set_title(\"Memory Bounded by Sliding Window Strategy\\n(Production 7B model, 80 concurrent sessions)\",\n", " fontsize=12)\n", "ax.legend(fontsize=8, loc=\"upper left\")\n", "ax.grid(True, alpha=0.3)\n", "plt.tight_layout()\n", "plt.savefig(\"sliding_window_memory.png\", dpi=150, bbox_inches=\"tight\")\n", "plt.show()" ], "id": "cell_36" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Strategy 3: Analytical Concurrency Limits\n", "\n", "Let us put everything together and compute the **maximum number of concurrent sessions** we can support on a single A100-40GB node as a function of context length and chosen mitigation strategy." ], "id": "cell_37" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def max_concurrent_sessions(\n", " config: ModelConfig,\n", " context_len: int,\n", " vram_gb: float = 40.0,\n", " weight_gb: float = 14.0,\n", " quantisation_factor: float = 1.0, # 0.5 = INT8 halves KV memory\n", " window_size: Optional[int] = None, # None = no sliding window\n", ") -> int:\n", " \"\"\"Compute the maximum batch size B that fits in VRAM.\n", " \n", " vram_budget = vram_gb - weight_gb\n", " kv_per_session = kv_cache_bytes(config, effective_ctx, B=1) * quant_factor\n", " max_B = floor(vram_budget / kv_per_session)\n", " \"\"\"\n", " available_gb = (vram_gb - weight_gb) * 1e9\n", " effective_ctx = min(context_len, window_size) if window_size else context_len\n", " kv_per_session = (\n", " kv_cache_bytes(config, effective_ctx, batch_size=1)\n", " * quantisation_factor\n", " )\n", " if kv_per_session == 0:\n", " return 0\n", " return max(1, int(available_gb // kv_per_session))\n", "\n", "\n", "# โ”€โ”€ Comparison table โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "strategies = {\n", " \"Baseline (FP16, no window)\" : {\"qf\": 1.0, \"ws\": None},\n", " \"INT8 KV cache\" : {\"qf\": 0.5, \"ws\": None},\n", " \"FP16 + Window W=2048\" : {\"qf\": 1.0, \"ws\": 2048},\n", " \"INT8 + Window W=2048\" : {\"qf\": 0.5, \"ws\": 2048},\n", "}\n", "target_ctxs = [2000, 4000, 6000, 8000, 10000, 12000]\n", "\n", "print(\"Maximum Concurrent Sessions per A100-40GB Node\\n\")\n", "print(f\"{'Strategy':<35}\", end=\"\")\n", "for c in target_ctxs:\n", " print(f\" {c:>6}\", end=\"\")\n", "print(\"\\n\" + \"-\" * (35 + len(target_ctxs) * 8))\n", "\n", "for name, params in strategies.items():\n", " print(f\"{name:<35}\", end=\"\")\n", " for c in target_ctxs:\n", " B = max_concurrent_sessions(cfg, c, quantisation_factor=params[\"qf\"],\n", " window_size=params[\"ws\"])\n", " flag = \" โœ…\" if B >= 80 else (\" โš ๏ธ\" if B >= 40 else \" โŒ\")\n", " print(f\" {B:>5}{flag[1]}\", end=\"\")\n", " print()\n", "\n", "print(f\"\\nTarget: โ‰ฅ 80 concurrent sessions (from business constraints)\")\n", "print(f\"โœ… = โ‰ฅ80 โš  = 40-79 โœ— = <40\")" ], "id": "cell_38" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# โ”€โ”€ Visualise concurrency limits โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "ctx_sweep = np.arange(256, 12_001, 128)\n", "\n", "fig, ax = plt.subplots(figsize=(11, 6))\n", "\n", "style_map = {\n", " \"Baseline (FP16, no window)\" : (\"crimson\", \"-\", 2.5),\n", " \"INT8 KV cache\" : (\"darkorange\", \"--\", 2.0),\n", " \"FP16 + Window W=2048\" : (\"steelblue\", \"-.\", 2.0),\n", " \"INT8 + Window W=2048\" : (\"seagreen\", \":\", 2.5),\n", "}\n", "\n", "for name, params in strategies.items():\n", " concurrencies = [\n", " max_concurrent_sessions(cfg, n, quantisation_factor=params[\"qf\"],\n", " window_size=params[\"ws\"])\n", " for n in ctx_sweep\n", " ]\n", " color, ls, lw = style_map[name]\n", " ax.plot(ctx_sweep, concurrencies, linestyle=ls, linewidth=lw,\n", " color=color, label=name)\n", "\n", "ax.axhline(80, color=\"black\", linestyle=\"--\", linewidth=1.5,\n", " label=\"Target: 80 concurrent sessions\")\n", "ax.set_xlabel(\"Context Length (tokens)\", fontsize=12)\n", "ax.set_ylabel(\"Max Concurrent Sessions\", fontsize=12)\n", "ax.set_title(\"Concurrency Limits by Mitigation Strategy\\n(A100-40GB, 7B model)\",\n", " fontsize=13)\n", "ax.legend(fontsize=9, loc=\"upper right\")\n", "ax.set_ylim(0, 250)\n", "ax.grid(True, alpha=0.3)\n", "ax.set_xlim(256, 12000)\n", "\n", "plt.tight_layout()\n", "plt.savefig(\"concurrency_strategies.png\", dpi=150, bbox_inches=\"tight\")\n", "plt.show()" ], "id": "cell_39" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3.9 Summary Experiment โ€” Putting It All Together\n", "\n", "Let us run one final end-to-end experiment that ties the profiling, the roofline insight, and the mitigation strategies into a single dashboard. This is the kind of output we would present to the engineering team and (in simplified form) to the CFO." ], "id": "cell_40" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def run_full_benchmark(\n", " model: SmallTransformer,\n", " context_lengths: List[int] = [64, 128, 256, 512],\n", " n_tokens: int = 20,\n", ") -> None:\n", " \"\"\"Run a comprehensive benchmark and print a structured report.\"\"\"\n", " model.eval()\n", " \n", " print(\"=\" * 70)\n", " print(\" INFERENCE BENCHMARK REPORT โ€” LexDraft Serving Analysis\")\n", " print(\"=\" * 70)\n", " print(f\"\\n Model : SmallTransformer (proxy for 7B production model)\")\n", " print(f\" Device : {GPU_NAME}\")\n", " print(f\" Tokens : {n_tokens} generated per request\\n\")\n", " print(f\" {'ctx':>6} {'TTFT-naive':>12} {'TTFT-cache':>12} \"\n", " f\"{'Decode-naive':>14} {'Decode-cache':>13} {'Speedup':>8}\")\n", " print(\" \" + \"-\" * 70)\n", "\n", " for ctx in context_lengths:\n", " prompt = torch.randint(0, 32000, (1, ctx)).to(DEVICE)\n", " with torch.no_grad():\n", " _, naive_times = generate_naive(model, prompt, n_tokens)\n", " _, cached_times, prefill_ms = generate_with_cache(model, prompt, n_tokens)\n", "\n", " naive_ttft = naive_times[0]\n", " cached_ttft = prefill_ms\n", " naive_step = np.mean(naive_times)\n", " cached_step = np.mean(cached_times) if cached_times else 1e-9\n", " speedup = naive_step / cached_step\n", "\n", " print(f\" {ctx:>6} {naive_ttft:>10.1f}ms {cached_ttft:>10.1f}ms \"\n", " f\"{naive_step:>12.2f}ms {cached_step:>11.2f}ms {speedup:>7.1f}x\")\n", "\n", " print(\"\\n\" + \"=\" * 70)\n", " print(\" KEY FINDINGS (from analytical model, 7B + A100-40GB)\")\n", " print(\"=\" * 70)\n", "\n", " # Analytical findings\n", " for n in [2000, 8000]:\n", " _, _, ai_d = compute_arithmetic_intensity(cfg, n, is_prefill=False)\n", " B_baseline = max_concurrent_sessions(cfg, n)\n", " B_int8 = max_concurrent_sessions(cfg, n, quantisation_factor=0.5)\n", " B_window = max_concurrent_sessions(cfg, n, window_size=2048)\n", " B_combined = max_concurrent_sessions(cfg, n, quantisation_factor=0.5,\n", " window_size=2048)\n", " kv_gb = kv_cache_bytes(cfg, n, batch_size=80) / 1e9\n", "\n", " print(f\"\\n Context n={n} tokens:\")\n", " print(f\" Decode arithmetic intensity : {ai_d:.2f} FLOP/Byte \"\n", " f\"({'MEMORY-BOUND โš ๏ธ' if ai_d < ridge_point else 'compute-bound โœ…'})\")\n", " print(f\" KV cache @ B=80 : {kv_gb:.1f} GB\")\n", " print(f\" Max concurrent (baseline) : {B_baseline}\")\n", " print(f\" Max concurrent (INT8 KV) : {B_int8}\")\n", " print(f\" Max concurrent (W=2048) : {B_window}\")\n", " print(f\" Max concurrent (INT8+W2048) : {B_combined}\")\n", "\n", " print(\"\\n RECOMMENDATION: Apply INT8 KV cache quantisation + W=2048 window.\")\n", " print(\" This achieves โ‰ฅ80 concurrent sessions up to n=8,000 tokens\")\n", " print(\" without additional GPU hardware, staying within the $45k/mo budget.\")\n", " print(\"=\" * 70)\n", "\n", "\n", "run_full_benchmark(model, CONTEXT_LENGTHS, n_tokens=15)" ], "id": "cell_41" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Notebook Summary\n", "\n", "Let us step back and review what we built โ€” and connect it back to the crisis at Lex Machina.\n", "\n", "### What we built\n", "\n", "| Section | What we implemented |\n", "|---|---|\n", "| 3.2 | `kv_cache_bytes()` โ€” first-principles formula for KV cache memory |\n", "| 3.3 | `MultiHeadAttention`, `TransformerBlock`, `SmallTransformer` โ€” a complete, cacheable transformer |\n", "| 3.4 | `generate_naive()` and `generate_with_cache()` โ€” the two worlds |\n", "| 3.5 | Profiling sweep: TTFT and per-step latency vs. context length |\n", "| 3.6 | Roofline model: arithmetic intensity for prefill vs. decode |\n", "| 3.7 | VRAM occupancy measurement and 7B analytical projection |\n", "| 3.8 | Three mitigation strategies: INT8 quantisation, sliding window, concurrency analysis |\n", "| 3.9 | End-to-end benchmark report |\n", "\n", "### The central insight\n", "\n", "The paradox โ€” **low GPU utilisation during high latency and OOM crashes** โ€” is fully explained by the roofline model. Decoding with a KV cache is deeply **memory-bandwidth-bound**. The GPU's compute units are starved because the bottleneck is reading tens of gigabytes of cached keys and values from HBM on every single decode step. At B=80 concurrent sessions with n=8,000 token contexts, the 7B model's KV cache alone requires more than 26 GB โ€” on top of 14 GB of weights โ€” completely overwhelming a 40 GB card.\n", "\n", "### The fix โ€” without new hardware\n", "\n", "By combining:\n", "1. **INT8 KV cache quantisation** (halves cache memory with negligible quality loss)\n", "2. **Sliding window eviction at W=2,048** (bounds per-session memory regardless of document length)\n", "\n", "we bring the peak KV cache footprint to under 3 GB per node at B=80, comfortably within the 26 GB VRAM budget remaining after weights. This keeps TTFT below 300 ms, sustains โ‰ฅ15 tokens/sec decode throughput, and protects the \\$3.9M ARR from the at-risk enterprise accounts โ€” all without touching the frozen model or exceeding the \\$45k/month compute budget.\n", "\n", "### Further directions\n", "\n", "- **Paged attention** (as in vLLM): instead of one contiguous KV cache block per session, allocate cache in small pages so fragmentation does not waste capacity\n", "- **Multi-Query Attention (MQA) / Grouped-Query Attention (GQA)**: reduce the number of KV heads at training time โ€” an architectural change, but worth flagging for the next fine-tuning cycle\n", "- **Speculative decoding**: use a small draft model to propose multiple tokens; the large model verifies in one parallel pass, improving throughput without reducing quality\n", "\n", "We hope this notebook gave you not just code, but a way of *seeing* the system โ€” the ability to hold a formula, a chart, and a hardware constraint in the same thought.\n", "\n", "*โ€” Vizuara Case Study Series*" ], "id": "cell_42" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(\"๐ŸŽ“ Notebook complete!\")\n", "print(\"\\nFiles saved:\")\n", "for f in [\"kv_cache_memory_scaling.png\", \"latency_profiling.png\",\n", " \"roofline_model.png\", \"vram_occupancy.png\",\n", " \"sliding_window_memory.png\", \"concurrency_strategies.png\"]:\n", " print(f\" ๐Ÿ“Š {f}\")\n", "\n", "print(\"\\n๐Ÿ’ก Business impact recap:\")\n", "print(\" โ€ข $3.9M ARR protected from churn\")\n", "print(\" โ€ข TTFT target (โ‰ค300ms) achievable at n=8,000 without new hardware\")\n", "print(\" โ€ข 80+ concurrent sessions per A100 node โ†’ stays within $45k/mo budget\")\n", "print(\" โ€ข Root cause: memory-bandwidth-bound decode, not compute starvation\")" ], "id": "cell_43" } ] }