{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Read vs Write: Measure What One AI Answer Actually Costs\n", "\n", "**Time and meter a language model on your own machine, and see where the cost of an answer really goes.**\n", "\n", "This is the companion notebook to the [DiamantAI](https://www.youtube.com/@DiamantAI) video **[Why AI Uses So Much Energy (It's Not the Thinking)](https://www.youtube.com/watch?v=w6FgmAtt7oo&list=PLBrpE2PttR2k)**. The film measures one open model on a laptop, reading a whole Sherlock Holmes story and then writing a summary of it. This notebook lets you run the same measurements on yours.\n", "\n", "## Overview\n", "\n", "A chat answer has two phases. First the model **reads** your prompt (the \"prefill\"). Then it **writes** its answer one token at a time (the \"decode\"). They look similar from the outside, but they use the hardware in very different ways, and that difference decides what an answer costs in time and in energy.\n", "\n", "This notebook measures it directly with a small open model (Qwen3-0.6B, optionally Qwen3-1.7B) through Hugging Face `transformers`. No API key, and it runs on a laptop CPU, an Apple-silicon Mac or an NVIDIA GPU:\n", "\n", "1. **Reading vs writing.** Time the read of a long prompt and the writing of an answer, per token.\n", "2. **Why writing is slow.** Every written token is one full pass over the model's weights, so tokens per second times the model's size is close to your chip's memory bandwidth. We check it on two model sizes.\n", "3. **Long chats.** Grow one conversation's memory (the KV cache) and time each written token as the chat gets longer.\n", "4. **Batching.** Answer 1, 4 and 8 different questions at once and see what each answer costs.\n", "5. **Energy (optional).** If your machine exposes a power sensor, convert time into joules.\n", "\n", "On our M4 MacBook Air (Qwen3 8B, 4-bit) we measured: a written token took about **3.5 times the energy** of a read token; each written token cost about **0.095 J per GB of model** the chip had to carry; at 32k tokens of chat each written token cost **2.6 times** what it cost at 1k; and answering **16 questions at once** cost about **half the energy per answer**. Your numbers will differ. The shape should not." ], "id": "cell-00" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> 📖 **Going deeper:** Chapter 16 of [*AI Agents Made Simple*](https://europe-west1-rag-techniques-views-tracker.cloudfunctions.net/rag-techniques-tracker?notebook=all-agents-tutorials--read-vs-write-cost&click=book-nb-agents-top&target=https%3A%2F%2Fwww.amazon.com%2Fdp%2FB0HNXRSRH2%3Fmaas%3Dmaas_adg_api_580351185840460851_static_9_573%26ref_%3Daa_maas%26tag%3Dmaas%26aa_campaignid%3Dagents-en%26aa_adgroupid%3Dgithub-notebook%26aa_creativeid%3Dgenai-read-vs-write-top&retarget=0&text=book-nb-agents-top), \"What an Agent Costs\", covers this topic. The book is by the author of this repository." ], "id": "a5f7315c" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Detailed Explanation\n", "\n", "### Motivation\n", "\n", "It is common to hear that an AI answer costs energy because the model \"thinks hard\". The measurements tell a more specific story. Most of the time and energy goes into moving the model's weights from memory to the compute units, and the amount moved depends on **how many tokens are written**, **how big the model is**, **how long the conversation already is**, and **how many answers share the trip**. Once you can see those four levers, you know which habits actually save something.\n", "\n", "### Key Components\n", "\n", "- **Read (prefill):** the whole prompt goes through the model in one pass. The weights are loaded once and used for every prompt token in parallel, so the chip is busy doing arithmetic. Cost per token is low.\n", "- **Write (decode):** each new token needs a full pass over all the weights, for just one token of work. The chip spends most of its time waiting on memory. Cost per token is high.\n", "- **KV cache:** the model keeps a key and a value vector for every token seen so far, in every layer. Each written token reads all of it, so a longer chat makes every new token more expensive.\n", "- **Batching:** one pass over the weights can serve many conversations at once. The weight trip is shared, so the cost per answer falls.\n", "- **Power sensor (optional):** NVIDIA's NVML on a GPU, or `powermetrics` on Apple silicon. We subtract the idle draw so the number is the extra energy the answer caused.\n", "\n", "### Architecture\n", "\n", "![Read vs Write Cost](../images/read_vs_write_cost.svg)\n", "\n", "### Benefits\n", "\n", "- **Honest numbers from your own hardware**, not from a vendor chart.\n", "- **A physical explanation you can check:** tokens per second times model size should land near your memory bandwidth.\n", "- **Practical habits** that follow directly from the measurements, listed at the end." ], "id": "cell-01" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Required Packages\n", "\n", "Everything runs locally with PyTorch and `transformers`. A GPU or an Apple-silicon Mac makes it faster, and a CPU also works. No API key is needed. `pynvml` is only for the optional energy step on NVIDIA GPUs." ], "id": "cell-02" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip install -q torch transformers pandas matplotlib\n", "# optional, NVIDIA only, for the energy step: !pip install -q nvidia-ml-py" ], "id": "cell-03" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Implementation\n", "\n", "### 1. Configuration\n", "\n", "The defaults finish in a few minutes on a laptop. `SECOND_MODEL_ID` adds a larger model for step 2; set it to `None` if memory is tight. If the longest chat depth runs out of memory, drop the last entry of `DEPTHS`." ], "id": "cell-04" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "MODEL_ID = \"Qwen/Qwen3-0.6B\" # about 1.2 GB of weights in 16-bit\n", "SECOND_MODEL_ID = \"Qwen/Qwen3-1.7B\" # optional second size for step 2, or None\n", "READ_TOKENS = 2048 # how long the prompt is that the model reads\n", "WRITE_TOKENS = 64 # how many tokens it writes in each timing run\n", "REPEATS = 3 # every timing runs this many times; we keep the median\n", "DEPTHS = [512, 1024, 2048, 4096, 8192] # chat lengths for step 3\n", "BATCH_SIZES = [1, 4, 8] # step 4\n", "BATCH_TOKENS = 64 # tokens written per answer in step 4" ], "id": "cell-05" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 2. Load a long text for the model to read\n", "\n", "We use *The Adventures of Sherlock Holmes* from Project Gutenberg (public domain), starting at \"A Scandal in Bohemia\", the story the film uses. It is long enough for every chat depth in step 3." ], "id": "cell-06" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import urllib.request\n", "\n", "URL = \"https://www.gutenberg.org/cache/epub/1661/pg1661.txt\"\n", "raw = urllib.request.urlopen(URL).read().decode(\"utf-8\")\n", "BOOK = raw[raw.find(\"I. A SCANDAL IN BOHEMIA\"):raw.find(\"*** END OF\")]\n", "print(f\"{len(BOOK.split()):,} words\")\n", "print(BOOK[:300])" ], "id": "cell-07" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 3. Load the model and a clock you can trust\n", "\n", "We pick CUDA, Apple's MPS or the CPU. GPUs run asynchronously: a Python call returns before the chip has finished. `sync()` waits for the chip, so every timer below measures finished work." ], "id": "cell-08" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import gc, statistics, time\n", "import torch\n", "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n", "dtype = torch.float16 if device != \"cpu\" else torch.float32\n", "\n", "def sync():\n", " # Wait until the chip has finished everything queued so far.\n", " if device == \"cuda\":\n", " torch.cuda.synchronize()\n", " elif device == \"mps\":\n", " torch.mps.synchronize()\n", "\n", "def load(model_id):\n", " tok = AutoTokenizer.from_pretrained(model_id)\n", " model = AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype).to(device).eval()\n", " return tok, model\n", "\n", "def weight_bytes(model):\n", " # Bytes of weights the chip has to read for one full pass (shared weights counted once).\n", " return sum(p.numel() * p.element_size() for p in model.parameters())\n", "\n", "tok, model = load(MODEL_ID)\n", "need = max(DEPTHS + [READ_TOKENS]) + 1024 # tokens we will use, with room to spare\n", "BOOK_IDS = tok(BOOK[:need * 6], return_tensors=\"pt\").input_ids[:, :need].to(device)\n", "print(f\"{MODEL_ID} on {device} ({dtype}) | {weight_bytes(model) / 1e9:.2f} GB of weights | \"\n", " f\"{BOOK_IDS.shape[1]:,} tokens of text ready\")" ], "id": "cell-09" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 4. The two measurements: read, then write\n", "\n", "`read_then_write` does what a chat model does for one answer, with a timer around each phase:\n", "\n", "- **Read:** one forward pass over all prompt tokens. `logits_to_keep=1` asks only for the prediction at the last position, like a real server.\n", "- **Write:** a loop that feeds back one token at a time, reusing the KV cache from the read. We pick the most likely token each step (greedy), since the choice of token does not change the cost." ], "id": "cell-10" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@torch.no_grad()\n", "def read_then_write(model, ids, n_write):\n", " # Returns (seconds to read ids, seconds to write n_write tokens, the KV cache).\n", " sync(); t0 = time.perf_counter()\n", " out = model(input_ids=ids, use_cache=True, logits_to_keep=1)\n", " sync(); read_s = time.perf_counter() - t0\n", " return (read_s, *write(model, out, n_write))\n", "\n", "@torch.no_grad()\n", "def write(model, out, n_write):\n", " # Write n_write tokens one at a time, continuing from a model output that holds a KV cache.\n", " cache, nxt = out.past_key_values, out.logits[:, -1:].argmax(-1)\n", " sync(); t0 = time.perf_counter()\n", " for _ in range(n_write):\n", " out = model(input_ids=nxt, past_key_values=cache, use_cache=True)\n", " cache, nxt = out.past_key_values, out.logits[:, -1:].argmax(-1)\n", " sync()\n", " return time.perf_counter() - t0, cache\n", "\n", "def median_runs(fn, repeats=REPEATS):\n", " # Run fn once to warm up, then `repeats` times; return the per-field medians.\n", " fn()\n", " runs = [fn() for _ in range(repeats)]\n", " return [statistics.median(r[i] for r in runs) for i in range(len(runs[0]))]\n", "\n", "results = {} # the numbers we collect for the final table" ], "id": "cell-11" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Usage Example\n", "\n", "### Step 1: Reading vs writing, one answer\n", "\n", "The model reads `READ_TOKENS` tokens of the story, then writes `WRITE_TOKENS` tokens. We compare the time **per token** in each phase." ], "id": "cell-12" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompt = BOOK_IDS[:, :READ_TOKENS]\n", "read_s, write_s = median_runs(lambda: read_then_write(model, prompt, WRITE_TOKENS)[:2])\n", "\n", "ms_read = read_s / READ_TOKENS * 1000\n", "ms_write = write_s / WRITE_TOKENS * 1000\n", "print(f\"read {READ_TOKENS:5d} tokens in {read_s:6.2f} s -> {ms_read:6.2f} ms per token read\")\n", "print(f\"write {WRITE_TOKENS:5d} tokens in {write_s:6.2f} s -> {ms_write:6.2f} ms per token written\")\n", "print(f\"one written token takes as long as {ms_write / ms_read:.0f} read tokens\")\n", "results[\"ms per token read\"] = ms_read\n", "results[\"ms per token written\"] = ms_write\n", "results[\"written / read, time per token\"] = ms_write / ms_read" ], "id": "cell-13" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Reading is fast per token because the model loads its weights once and applies them to all prompt tokens together. Writing gets no such sharing: every token pays for its own trip through the weights.\n", "\n", "The **time** ratio you see here is larger than the **energy** ratio the film reports (about 3.5 on our Mac). While reading, the chip runs flat out and draws more power; while writing it mostly waits on memory and draws less. Each written token still costs several times the energy of a read one, which is what step 5 measures if you have a power sensor." ], "id": "cell-14" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Comparison\n", "\n", "### Step 2: Why writing is slow: one pass over the weights per token\n", "\n", "If each written token needs the whole model to travel from memory, then\n", "\n", "`tokens per second x bytes of weights ≈ bytes per second the memory delivers`\n", "\n", "and that product should stay roughly constant across model sizes, near your chip's memory bandwidth (for reference: an M4 MacBook Air is rated at 120 GB/s, an NVIDIA T4 at 320 GB/s). We measure write speed on a short prompt for each model size." ], "id": "cell-15" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def write_speed(model, tok):\n", " # Median tokens per second while writing, after a short prompt.\n", " ids = tok(\"Explain how a bicycle stays upright while it moves.\", return_tensors=\"pt\").input_ids.to(device)\n", " _, write_s = median_runs(lambda: read_then_write(model, ids, WRITE_TOKENS)[:2])\n", " return WRITE_TOKENS / write_s\n", "\n", "rows = []\n", "for model_id in [MODEL_ID, SECOND_MODEL_ID]:\n", " if model_id is None:\n", " continue\n", " t, m = (tok, model) if model_id == MODEL_ID else load(model_id)\n", " tps, gb = write_speed(m, t), weight_bytes(m) / 1e9\n", " rows.append({\"model\": model_id.split(\"/\")[-1], \"weights GB\": gb, \"tokens/s written\": tps,\n", " \"GB/s moved (tokens/s x GB)\": tps * gb, \"ms per written token\": 1000 / tps})\n", " if m is not model:\n", " del m, t; gc.collect()\n", " if device == \"mps\": torch.mps.empty_cache()\n", " if device == \"cuda\": torch.cuda.empty_cache()\n", "\n", "import pandas as pd\n", "sizes = pd.DataFrame(rows)\n", "results[\"GB/s moved while writing (largest model)\"] = sizes[\"GB/s moved (tokens/s x GB)\"].iloc[-1]\n", "sizes.round(2)" ], "id": "cell-16" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The bigger model writes more slowly, by roughly the ratio of the sizes, and the GB/s column stays in the same range: the chip is limited by how fast it can move weights, not by how much arithmetic it can do. A small model in plain PyTorch often lands below the others, because a fixed per-token overhead (launching hundreds of small GPU kernels) is a bigger share of its short pass. Runtimes built for serving, such as llama.cpp, MLX or vLLM, remove most of that overhead, which is why the film's MLX run reached about 91-103 GB/s on every size.\n", "\n", "### Step 3: Long chats: the same model, a growing memory\n", "\n", "A chat model keeps a **KV cache**: for every token so far, in every layer, a key vector and a value vector. Every new token reads all of it. We feed the story into one cache, pausing at each depth in `DEPTHS` to time a few written tokens." ], "id": "cell-17" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@torch.no_grad()\n", "def write_at_depths(model, ids, depths, n_write=32, chunk=512):\n", " # Grow one KV cache through `depths` and time n_write written tokens at each depth.\n", " rows, cache, fed = [], None, 0\n", " for depth in depths:\n", " while fed < depth:\n", " n = min(chunk, depth - fed)\n", " out = model(input_ids=ids[:, fed:fed + n], past_key_values=cache, use_cache=True, logits_to_keep=1)\n", " cache, fed = out.past_key_values, fed + n\n", " length, kv_bytes = cache.get_seq_length(), cache_bytes(cache)\n", " write(model, out, 4) # warm-up at this depth\n", " times = [write(model, out, n_write)[0] for _ in range(REPEATS)]\n", " fed = cache.get_seq_length() # the probe tokens stay in the chat, like a real reply\n", " rows.append({\"chat length (tokens)\": length,\n", " \"ms per written token\": statistics.median(times) / n_write * 1000,\n", " \"KV cache MB\": kv_bytes / 1e6})\n", " return pd.DataFrame(rows)\n", "\n", "def cache_bytes(cache):\n", " # Actual bytes held in a DynamicCache (works across transformers versions).\n", " pairs = [(l.keys, l.values) for l in cache.layers] if hasattr(cache, \"layers\") \\\n", " else zip(cache.key_cache, cache.value_cache)\n", " return sum(k.numel() * k.element_size() + v.numel() * v.element_size() for k, v in pairs)\n", "\n", "depth = write_at_depths(model, BOOK_IDS, [d for d in DEPTHS if d <= BOOK_IDS.shape[1]])\n", "depth[\"vs shortest\"] = depth[\"ms per written token\"] / depth[\"ms per written token\"].iloc[0]\n", "results[f\"written token at {int(depth.iloc[-1, 0]):,} vs {int(depth.iloc[0, 0]):,} tokens of chat\"] = depth[\"vs shortest\"].iloc[-1]\n", "depth.round(2)" ], "id": "cell-18" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Plot it: the price of one written token as the chat grows" ], "id": "cell-19" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "_, ax = plt.subplots(figsize=(6, 3.5))\n", "ax.plot(depth[\"chat length (tokens)\"], depth[\"ms per written token\"], \"o-\")\n", "ax.set_xscale(\"log\", base=2)\n", "ax.set_xlabel(\"tokens already in the chat\")\n", "ax.set_ylabel(\"ms per written token\")\n", "ax.set_ylim(bottom=0)\n", "ax.set_title(f\"{MODEL_ID.split('/')[-1]} on {device}\")\n", "plt.show()" ], "id": "cell-20" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### The KV cache, by arithmetic\n", "\n", "You do not need to run a model to know how big its cache gets. From the config:\n", "\n", "`bytes per token = layers x KV heads x head size x 2 (a key and a value) x bytes per number`\n", "\n", "We compute it for the loaded model and check it against the cache we just measured. For comparison, the film's Qwen3 8B holds 147,456 bytes per token, about 4.7 GB at 32k tokens, on top of its weights." ], "id": "cell-21" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cfg = model.config\n", "head_dim = getattr(cfg, \"head_dim\", None) or cfg.hidden_size // cfg.num_attention_heads\n", "per_token = cfg.num_hidden_layers * cfg.num_key_value_heads * head_dim * 2 * torch.finfo(dtype).bits // 8\n", "print(f\"{cfg.num_hidden_layers} layers x {cfg.num_key_value_heads} KV heads x {head_dim} x 2 x \"\n", " f\"{torch.finfo(dtype).bits // 8} bytes = {per_token:,} bytes per token\")\n", "last = depth.iloc[-1]\n", "n_last = int(last[\"chat length (tokens)\"])\n", "print(f\"predicted at {n_last:,} tokens: {per_token * n_last / 1e6:.1f} MB | \"\n", " f\"measured: {last['KV cache MB']:.1f} MB\")\n", "for n in (1_000, 32_000, 128_000):\n", " print(f\"{n:>8,} tokens of chat -> {per_token * n / 1e9:5.2f} GB read for every written token, \"\n", " f\"next to {weight_bytes(model) / 1e9:.2f} GB of weights\")" ], "id": "cell-22" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Step 4: Batching: many answers share one trip through the weights\n", "\n", "We ask 1, 4 and 8 **different** questions at the same time. Prompts are padded on the left so every answer starts at the same step, and every answer writes exactly `BATCH_TOKENS` tokens so the counts are comparable." ], "id": "cell-23" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "QUESTIONS = [\"How do plants make food from sunlight?\", \"Why is the sky blue?\",\n", " \"How does a refrigerator work?\", \"What causes the seasons?\",\n", " \"How do vaccines train the immune system?\", \"Why do cats purr?\",\n", " \"How does GPS know where you are?\", \"What makes bread rise?\"]\n", "tok.padding_side = \"left\"\n", "\n", "def chat(question):\n", " return tok.apply_chat_template([{\"role\": \"user\", \"content\": question + \" Answer in about 100 words.\"}],\n", " add_generation_prompt=True, tokenize=False, enable_thinking=False)\n", "\n", "@torch.no_grad()\n", "def answer_batch(n):\n", " # Answer n different questions at once; return (seconds, tokens written).\n", " enc = tok([chat(QUESTIONS[i % len(QUESTIONS)]) for i in range(n)], return_tensors=\"pt\", padding=True).to(device)\n", " sync(); t0 = time.perf_counter()\n", " out = model.generate(**enc, max_new_tokens=BATCH_TOKENS, min_new_tokens=BATCH_TOKENS,\n", " do_sample=False, pad_token_id=tok.pad_token_id)\n", " sync()\n", " return time.perf_counter() - t0, n * (out.shape[1] - enc.input_ids.shape[1])\n", "\n", "rows = []\n", "for n in BATCH_SIZES:\n", " seconds, tokens = median_runs(lambda: answer_batch(n))\n", " rows.append({\"answers at once\": n, \"seconds\": seconds, \"tokens/s (all answers)\": tokens / seconds,\n", " \"seconds per answer\": seconds / n})\n", "batch = pd.DataFrame(rows)\n", "batch[\"cost per answer vs one at a time\"] = batch[\"seconds per answer\"] / batch[\"seconds per answer\"].iloc[0]\n", "results[f\"cost per answer, {BATCH_SIZES[-1]} at once vs 1\"] = batch[\"cost per answer vs one at a time\"].iloc[-1]\n", "batch.round(2)" ], "id": "cell-24" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The batch takes only a little longer than a single answer, because the expensive part, moving the weights, happens once per step for the whole batch. Total throughput rises and each answer's share of the chip time falls. This is why a busy server answers more cheaply per question than a lone laptop: it always has a batch to fill. With a fully busy chip, time per answer is a fair stand-in for energy per answer; step 5 measures the energy itself.\n", "\n", "### Step 5 (optional): Energy from the chip's own power sensor\n", "\n", "Time is not energy. To get joules we sample the chip's power while it works, subtract the **idle** power measured just before, and add up the difference over time. This step looks for a sensor and **skips itself if there is none**. Nothing else in the notebook needs it.\n", "\n", "- **NVIDIA GPU:** read through NVML (`pip install nvidia-ml-py`). No special permissions.\n", "- **Apple silicon:** macOS only exposes the chip's power to `powermetrics`, which needs administrator rights. The notebook never asks for them. If you want this step, start the meter yourself in a separate terminal, leave it running, and re-run this cell:\n", "\n", "```bash\n", "sudo powermetrics --samplers cpu_power,gpu_power -i 250 -o /tmp/pm.txt\n", "```\n", "\n", " The cell then reads the \"Combined Power (CPU + GPU + ANE)\" lines from that file. Stop the meter with Ctrl-C when you are done.\n", "- **Anything else** (CPU-only laptops, most cloud VMs): skipped.\n", "\n", "Close other busy apps first. If something else changes its power draw during the run, the subtraction measures that program instead of the model, so the cell refuses a baseline that is too unsteady and any result that comes out at or below zero." ], "id": "cell-25" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os, re, threading\n", "\n", "PM_FILE = \"/tmp/pm.txt\"\n", "\n", "def find_power_meter():\n", " # Return (name, read_watts) for NVML, or (\"powermetrics\", None) if the meter file is live, else (None, None).\n", " try:\n", " import pynvml\n", " pynvml.nvmlInit()\n", " handle = pynvml.nvmlDeviceGetHandleByIndex(0)\n", " return \"nvml\", lambda: pynvml.nvmlDeviceGetPowerUsage(handle) / 1000.0\n", " except Exception:\n", " pass\n", " if os.path.exists(PM_FILE) and time.time() - os.path.getmtime(PM_FILE) < 5:\n", " return \"powermetrics\", None\n", " return None, None\n", "\n", "class PowerSampler(threading.Thread):\n", " # Collect (time, watts) samples in the background from NVML or from powermetrics' output file.\n", " PM_LINE = re.compile(r\"Combined Power \\(CPU \\+ GPU \\+ ANE\\):\\s*([\\d.]+)\\s*mW\")\n", "\n", " def __init__(self, source, read_watts=None, interval=0.1):\n", " super().__init__(daemon=True)\n", " self.source, self.read_watts, self.interval = source, read_watts, interval\n", " self.samples, self.stopped = [], threading.Event()\n", "\n", " def run(self):\n", " if self.source == \"nvml\":\n", " while not self.stopped.is_set():\n", " self.samples.append((time.time(), self.read_watts()))\n", " self.stopped.wait(self.interval)\n", " return\n", " with open(PM_FILE, errors=\"ignore\") as f:\n", " f.seek(0, 2) # only samples from now on\n", " while not self.stopped.is_set():\n", " line = f.readline()\n", " if not line:\n", " time.sleep(0.05); continue\n", " m = self.PM_LINE.search(line)\n", " if m:\n", " self.samples.append((time.time(), float(m.group(1)) / 1000.0))\n", "\n", " def watts(self, t0, t1):\n", " return [w for t, w in self.samples if t0 <= t <= t1]\n", "\n", " def joules(self, t0, t1, idle_w):\n", " # Energy above idle between t0 and t1 (trapezoid rule).\n", " s = [(t, w) for t, w in self.samples if t0 <= t <= t1]\n", " if len(s) < 2:\n", " return None\n", " return sum(((wa + wb) / 2 - idle_w) * (tb - ta) for (ta, wa), (tb, wb) in zip(s, s[1:]))\n", "\n", "def idle_watts(sampler, seconds=10, max_spread=0.15):\n", " # Measure idle power; refuse if it moves by more than max_spread (coefficient of variation).\n", " time.sleep(2)\n", " t0 = time.time(); time.sleep(seconds)\n", " w = sampler.watts(t0, time.time())\n", " if len(w) < 5:\n", " raise RuntimeError(\"no power samples arrived; is the meter still running?\")\n", " mean, sd = statistics.mean(w), statistics.pstdev(w)\n", " if mean > 0 and sd / mean > max_spread and sd > 0.5:\n", " raise RuntimeError(f\"idle power moves by {sd:.2f} W around {mean:.2f} W; something else is busy. \"\n", " \"Close it and re-run this cell.\")\n", " return mean\n", "\n", "@torch.no_grad()\n", "def metered(sampler, idle_w, work, min_seconds=6):\n", " # Repeat work() until min_seconds have passed; return joules above idle and how many tokens it handled.\n", " tokens = 0\n", " sync(); t0 = time.time()\n", " while time.time() - t0 < min_seconds:\n", " tokens += work(); sync()\n", " t1 = time.time()\n", " time.sleep(1.0) # let the last samples land\n", " j = sampler.joules(t0, t1, idle_w)\n", " if j is None or j <= 0:\n", " raise RuntimeError(f\"measured {j} J above idle, which is impossible: the baseline was disturbed. Re-run.\")\n", " return j, tokens\n", "\n", "meter, read_watts = find_power_meter()\n", "if meter is None:\n", " print(\"No power sensor found, so the energy step is skipped. Everything above still stands.\\n\"\n", " \"NVIDIA: pip install nvidia-ml-py. Apple silicon: start \"\n", " \"`sudo powermetrics --samplers cpu_power,gpu_power -i 250 -o /tmp/pm.txt` \"\n", " \"in your own terminal, then re-run this cell.\")\n", "else:\n", " sampler = PowerSampler(meter, read_watts); sampler.start()\n", " prompt = BOOK_IDS[:, :READ_TOKENS]\n", " read_then_write(model, prompt, 8) # warm-up\n", "\n", " def one_read():\n", " model(input_ids=prompt, use_cache=True, logits_to_keep=1)\n", " return READ_TOKENS\n", "\n", " out = model(input_ids=prompt, use_cache=True, logits_to_keep=1)\n", " def one_write(n=16):\n", " write(model, out, n)\n", " return n\n", "\n", " idle = idle_watts(sampler)\n", " j_read, n_read = metered(sampler, idle, one_read)\n", " idle = idle_watts(sampler)\n", " j_write, n_write = metered(sampler, idle, one_write)\n", " sampler.stopped.set()\n", " jpt_read, jpt_write = j_read / n_read, j_write / n_write\n", " gb = weight_bytes(model) / 1e9\n", " print(f\"meter: {meter} | idle {idle:.2f} W\")\n", " print(f\"read : {jpt_read:.5f} J per token ({n_read:,} tokens)\")\n", " print(f\"write: {jpt_write:.5f} J per token ({n_write:,} tokens) = {jpt_write / gb:.3f} J per GB of weights\")\n", " print(f\"one written token costs {jpt_write / jpt_read:.1f}x the energy of a read token\")\n", " results[\"J per token read\"] = jpt_read\n", " results[\"J per token written\"] = jpt_write\n", " results[\"written / read, energy per token\"] = jpt_write / jpt_read" ], "id": "cell-26" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Results\n", "\n", "Everything this run measured on your machine, in one place." ], "id": "cell-27" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(f\"{MODEL_ID} on {device}\")\n", "pd.DataFrame({\"measured\": results}).round(3)" ], "id": "cell-28" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Additional Considerations\n", "\n", "### What this means when you use AI\n", "\n", "The measurements point at the same few levers the film ends on:\n", "\n", "- **Paste in what matters, not everything.** Reading is the cheap phase per token, but a long prompt is still read in full, and it stays in the chat's memory for every token written after it.\n", "- **Ask for only what you need.** Written tokens are the expensive ones, and each one is a full trip through the model. \"Answer in three sentences\" is a real saving; so is a model that does not write a long visible reasoning trace when you do not need one.\n", "- **Start a fresh chat when the old one stopped mattering.** Step 3 shows every written token getting slower as the conversation grows, because it reads the whole KV cache each time. A new chat resets that cost to zero.\n", "- **Smaller models for small jobs.** Step 2 shows the cost per written token rising with the model's size.\n", "- **Shared servers batch.** Step 4 is why a busy provider answers more cheaply per question than a single laptop. It also means one person's answer on a server is cheaper than the same answer measured alone on a laptop.\n", "\n", "### Limits of this measurement\n", "\n", "- **Runtime overhead.** Plain PyTorch adds a fixed cost per step that serving runtimes (llama.cpp, MLX, vLLM, TensorRT-LLM) mostly remove. Absolute speeds here are a floor; the ratios are what carry over.\n", "- **Precision.** We load 16-bit weights on a GPU and 32-bit on a CPU. Quantised 4-bit or 8-bit weights move fewer bytes, so they write faster and use less energy per token, which is the same bandwidth story from the other side.\n", "- **The meter's boundary.** NVML reports the GPU board; `powermetrics` reports the chip package (CPU, GPU and neural engine). Neither includes the screen, the fans, the power supply's losses or, for a data center, the cooling.\n", "- **A shared machine lies.** Anything else running changes both time and power. Close other work and run the timing cells twice; if the numbers move a lot, the machine was busy.\n", "- **Attention cost grows with the chat.** At the depths used here the KV cache read is modest next to the weights, so the slowdown in step 3 is gentle for a small model. It grows with longer chats and with models that have more layers or KV heads.\n", "\n", "## References\n", "\n", "- Pope et al., *Efficiently Scaling Transformer Inference* (MLSys 2023): prefill vs decode, memory-bandwidth-bound generation and batching.\n", "- Kwon et al., *Efficient Memory Management for Large Language Model Serving with PagedAttention* (SOSP 2023): the KV cache as the main memory cost of serving.\n", "- Williams, Waterman and Patterson, *Roofline: An Insightful Visual Performance Model for Multicore Architectures* (CACM 2009): compute-bound vs memory-bound.\n", "- Luccioni, Jernite and Strubell, *Power Hungry Processing: Watts Driving the Cost of AI Deployment?* (FAccT 2024): measured energy of model inference.\n", "- Qwen Team, *Qwen3 Technical Report* (2025), [models on Hugging Face](https://huggingface.co/Qwen).\n", "- Arthur Conan Doyle, *The Adventures of Sherlock Holmes*, [Project Gutenberg #1661](https://www.gutenberg.org/ebooks/1661) (public domain).\n", "- [NVIDIA Management Library (NVML)](https://developer.nvidia.com/management-library-nvml) and Apple's `powermetrics` manual page (`man powermetrics`)." ], "id": "cell-29" }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "📖 **Want the decisions behind agents like this one, explained without code?** [*AI Agents Made Simple*](https://europe-west1-rag-techniques-views-tracker.cloudfunctions.net/rag-techniques-tracker?notebook=all-agents-tutorials--read-vs-write-cost&click=book-nb-agents&target=https%3A%2F%2Fwww.amazon.com%2Fdp%2FB0HNXRSRH2%3Fmaas%3Dmaas_adg_api_580351185840460851_static_9_573%26ref_%3Daa_maas%26tag%3Dmaas%26aa_campaignid%3Dagents-en%26aa_adgroupid%3Dgithub-notebook%26aa_creativeid%3Dgenai-read-vs-write&retarget=0&text=book-nb-agents) by the author of this repository has a full chapter on what an agent costs: cost per finished task, input versus output tokens, prompt caching, cutting steps before tokens, and a spending cap in the loop. Available on Amazon." ], "id": "2d3ee2b2" }, { "cell_type": "markdown", "metadata": {}, "source": [ "![](https://europe-west1-genai-agents-views-tracker.cloudfunctions.net/genai-agents-tracker?notebook=all-agents-tutorials--read-vs-write-cost)" ], "id": "24a3656b" } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }