The KV Cache

in progressThis lesson is still being built - content, demos, and checkpoints may change before release.

KV cache

Loading interactive demo…

A sentence decoded one token at a time. Each new token makes a query, a key, and a value; the query is scored against the key of every token so far, the row is softmaxed, and the output blends their values. With the cache the keys and values are kept and a step projects one token; without it the whole prefix is projected again every step. Take the tour, or press and explore. Legend:

The panel beside the stage is the real KvCache class from the published library executing in your browser, with the current line highlighted.

The lower-triangular weight matrix, one row per decoded token, filled as each row is softmaxed; faint until then.
The token being decoded and the cell being scored.
The tapped token; the sentence strip shades every token by its weight in that row.
An entry appended to the cache this step.
The key being scored, and an entry projected again this step without the cache.

cache entries is how many keys and values are in hand, multiply-adds the work spent so far, saved the work the cache avoided against projecting everything again, memory what the held entries take. Tap a token to read its row.

The code panel shows bare executable code by default. The bar above it names the file (hover the name for the full published path); options holds the toggles: comments, the tracer instrumentation that produces this visualization, and verbose names for the terse identifiers (display only: the default is the byte-exact published file). The view switch adds plain-English narration: both interleaves it with the code, pseudo shows only the pseudocode. hide code collapses the panel; the rail it leaves behind brings it back.

Keys: Space play/pause · ←/→ step · Home/End jump · R reset · F fullscreen · G guided tour

The causal mask. A decoder generates token $i$ before token $i + 1$ exists, so its attention row can only cover the tokens at or before it:

a_{ij} = \frac{e^{s_{ij}}}{\sum_{m \le i} e^{s_{im}}} \quad \text{for } j \le i, \qquad a_{ij} = 0 \quad \text{for } j > i .

That is the lower triangle on the stage. Training computes the whole triangle at once with a mask; decoding computes it one row per step.

Why old keys and values never change. With $k_j = W_k x_j$ and $v_j = W_v x_j$ , token $j$ 's key and value depend on $x_j$ and the fixed matrices alone, not on the step, so row $i$ needs exactly the $k_j$ and $v_j$ that row $i - 1$ already used, plus one new pair. The query $q_i$ is different: it is used to score row $i$ and nothing ever asks on token $i$ 's behalf again, so it is not kept. The cache is the keys and values; the name says what is in it.

Cost per step. Counting multiply-adds, a $d$ -vector through a $d \times d$ matrix is $d^2$ , and a dot product or a blend over $n$ keys is $n d$ . At step $i$ with $n = i + 1$ tokens so far,

\text{cache: } 3d^2 + 2nd, \qquad \text{recompute: } (2n + 1)\, d^2 + 2nd .

The recompute projects the $i$ earlier tokens into keys and values again, $2 i d^2$ extra, and at $i = 0$ the two are equal, which is why the saved readout is 0 on the first token. Over $N$ tokens the totals are $3Nd^2 + N(N+1)d$ against $(N^2 + 2N)\, d^2 + N(N+1)d$ : the cache makes the projection work linear in $N$ , while the scoring stays quadratic in both, since every row still reads every cached key.

Memory. The cache holds $2 \times n \times d$ numbers per layer; a real model has $L$ layers and a key-value width $d_{kv}$ (the number of key-value heads times the head width), so

\text{bytes} = L \times 2 \times n \times d_{kv} \times \text{bytes per number} .

Llama 3.1 8B has 32 layers and 8 key-value heads of width 128, so at two bytes a number each token costs $32 \times 2 \times 1024 \times 2 = 131{,}072$ bytes, 128 KiB. Hugging Face's Llama 3.1 post puts the cache at 15.62 GB for a 128k-token context in 16-bit, about as much as the 8B weights themselves. Multi-query and grouped-query attention shrink $d_{kv}$ by sharing key-value heads; latent attention compresses the pair; paging fights the fragmentation of many caches on one GPU.

Choices worth reading closely in the implementation (the code panel's comments and tracer checkboxes reveal the full published file):

Composition over public weights: the class owns a SelfAttention for its seeded wq, wk, and wv and never calls attend, which would score every token against every other. The decode is a different loop over the same matrices, and project is its own one-liner because the attention's is private.
One loop for both modes: the projection loop runs from k.length to the step. With the cache the list already holds the prefix, so it projects one token; the recompute empties the list first, so the same loop projects everything again, with recomputed marking the repeats.
No query for a repeat: an earlier token's query is null in the recompute, since nothing asks on its behalf now; only its key and value are needed, which is why the naive cost is $2 d^2$ per earlier token, not $3 d^2$ .
The ledger counts what runs, $d^2$ in project and $d$ per dot and per value in the blend, and compare puts it beside the two closed forms; the tests check the ledger equals the closed form for the run's mode.
The temporary list in recompute mode is the same k and v arrays truncated to zero each step, so the code panel shows what not caching means: the same numbers computed again and thrown away.

A transformer decoder writes one token at a time, and each new token attends to every token before it. Run naively, every step projects the whole prefix into keys and values again. But a token's key and value depend only on its own vector and the fixed matrices, so they are the same at every step; keeping them turns a step into one projection plus a scan of the cache. The cache was in the earliest released decoder code, and as contexts grew to hundreds of thousands of tokens its memory, not the weights, became what limits how many conversations a machine can serve at once.

2017Ashish Vaswani and colleagues at Google - 'Attention Is All You Need' publishes the transformer, on June 12; its decoder produces one token at a time, each attending to the tokens before it.source
2019OpenAI - The released GPT-2 code carries a 'past' tensor of keys and values, and each step concatenates its own key and value onto it: a KV cache in a few lines.source
2019Noam Shazeer - 'Fast Transformer Decoding: One Write-Head is All You Need' shares one key and one value across all the heads, shrinking the cache and the memory bandwidth incremental decoding needs.source
2022Reiner Pope, Sholto Douglas, and colleagues at Google - 'Efficiently Scaling Transformer Inference' shows that the lower memory of multi-query attention lets serving scale to 32 times longer contexts.source
2023Joshua Ainslie and colleagues at Google - Grouped-query attention: an intermediate number of key-value heads, between one per head and one shared by all, and a recipe for converting trained checkpoints.source
2023Meta - Llama 2's 34B and 70B models use grouped-query attention for inference scalability: a smaller cache per token.source
2023Woosuk Kwon, Zhuohan Li, and colleagues at UC Berkeley - PagedAttention and vLLM manage the cache in blocks like virtual memory pages, after measuring that existing systems used only 20 to 38 percent of their cache memory for actual token states.source
2024DeepSeek-AI - DeepSeek-V2's multi-head latent attention compresses the keys and values into a latent vector, cutting the cache by 93.3 percent against DeepSeek 67B.source

Used today

Chat serving - Every reply a chat model writes is decoded token by token over a cache; Hugging Face's transformers keeps it in past_key_values and grows it as generation proceeds.
Prefix caching - A shared prompt's keys and values are computed once and copied for every request that starts with it, so a long system prompt is paid for one time.
Living within memory - Offloading a cache to CPU memory and quantizing its numbers to fewer bits trade speed for room, which is what long contexts on small GPUs need.

length8

dim

mode

cache entries: 8 of 8
multiply-adds: 672
saved: 896 (x2.33)
memory: 256 B

kv-cache.ts

import type { Tracer } from "../trace";import { Kernels } from "./kernels";import { SelfAttention } from "./self-attention";/** Keep every token's key and value, or project the whole prefix again at every step. */export type DecodeMode = "cache" | "recompute";export interface DecodeResult {  /** weights[query][key] over the keys at or before the query: a lower triangle, each row summing to one. */  weights: number[][];  /** One mixed value vector per decoded token. */  outputs: number[][];  /** Multiply-adds spent at each step, by the run's own ledger. */  work: number[];  /** The ledger's sum over every step. */  total: number;  /** What the same run costs without a cache, by the closed form. */  naiveTotal: number;}/** * Causal decoding one token at a time over a self-attention's public * weights. Every step projects the new token into a query, a key, and a * value, scores the query against the key of every token so far, and * blends their values. The keys and values of earlier tokens never * change, because each one depends only on its own vector and the fixed * matrices, so a cache keeps them and a step projects one token; the * query is used once and is never kept. Without the cache the same * step projects the whole prefix again. A ledger counts multiply-adds: * a vector through a d by d matrix is d², a dot product or a blend over * n keys is n·d. * * Time per step: O(d² + n·d) with the cache, O(n·d²) without. * Space: O(n·d) for the cache, whatever the mode spends per step. */export class KvCache {  readonly attention: SelfAttention;  private ledger = 0;  constructor(    readonly dim: number,    readonly mode: DecodeMode = "cache",    seed = 1,    private readonly tracer?: Tracer,  ) {    this.attention = new SelfAttention(dim, "identity", seed);  }  /** One d-vector through a d by d matrix: d² multiply-adds on the step's ledger. */  private project(m: readonly (readonly number[])[], x: readonly number[]): number[] {    this.ledger += this.dim * this.dim;    return m.map((row) => row.reduce((s, weight, i) => s + weight * (x[i] ?? 0), 0));  }  decode(inputs: readonly (readonly number[])[]): DecodeResult {    if (inputs.some((x) => x.length !== this.dim)) throw new RangeError(`every input must have ${this.dim} numbers`);    const scale = 1 / Math.sqrt(this.dim);    const k: number[][] = [];    const keyCache: number[][] = [];    const v: number[][] = [];    const valueCache: number[][] = [];    const weights: number[][] = [];    const outputs: number[][] = [];    const work: number[] = [];    let naiveTotal = 0;    for (let step = 0; step < inputs.length; step++) {      this.ledger = 0;      if (this.mode === "recompute") {        // Nothing survives a step without the cache: the list is rebuilt from the first token.        k.length = 0;        keyCache.length = 0;        v.length = 0;        valueCache.length = 0;      }      let q: number[] = [];      let queryVector: number[] = [];      for (let index = k.length; index <= step; index++) {      for (let index = keyCache.length; index <= step; index++) {        const x = inputs[index]!;        const recomputed = index < step;        const key = this.project(this.attention.wk, x);        const value = this.project(this.attention.wv, x);        // Only the newest token needs a query: nothing asks on behalf of an earlier one again.        const query = recomputed ? null : this.project(this.attention.wq, x);        if (query) q = query;        if (query) queryVector = query;        k.push(key);        keyCache.push(key);        v.push(value);        valueCache.push(value);        // ▸ a token's vector through the fixed matrices; an old one needs no query        this.tracer?.emit("project", { step, index, recomputed, query: query ? [...query] : null, key: [...key], value: [...value] });      }      // ▸ the new key and value join the cache for every later step      if (this.mode === "cache") this.tracer?.emit("append", { step, entries: k.length });      if (this.mode === "cache") this.tracer?.emit("append", { step, entries: keyCache.length });      const scores: number[] = [];      for (let index = 0; index < k.length; index++) {      for (let index = 0; index < keyCache.length; index++) {        const dot = q.reduce((s, value, i) => s + value * k[index]![i]!, 0);        const dot = queryVector.reduce((s, value, i) => s + value * keyCache[index]![i]!, 0);        this.ledger += this.dim;        const scaled = dot * scale;        scores.push(scaled);        // ▸ the new query against one cached key, scaled down by root d        this.tracer?.emit("score", { query: step, key: index, dot, scaled });      }      const row = Kernels.softmax(scores);      const top = row.indexOf(Math.max(...row));      weights.push(row);      // ▸ the scores over the tokens so far become weights that sum to one      this.tracer?.emit("softmax", { query: step, top, weights: [...row] });      const output = Array.from({ length: this.dim }, (_, i) => row.reduce((s, w, index) => s + w * v[index]![i]!, 0));      const output = Array.from({ length: this.dim }, (_, i) => row.reduce((s, w, index) => s + w * valueCache[index]![i]!, 0));      this.ledger += k.length * this.dim;      this.ledger += keyCache.length * this.dim;      outputs.push(output);      // ▸ the output is the cached values blended by those weights      this.tracer?.emit("mix", { query: step, output: [...output] });      const cachedWork = 3 * this.dim * this.dim + 2 * (step + 1) * this.dim;      const naiveWork = (2 * (step + 1) + 1) * this.dim * this.dim + 2 * (step + 1) * this.dim;      work.push(this.ledger);      naiveTotal += naiveWork;      // ▸ the step's multiply-adds beside what projecting everything again costs      this.tracer?.emit("compare", { step, work: this.ledger, cachedWork, naiveWork, saved: naiveWork - this.ledger });    }    return { weights, outputs, work, total: work.reduce((s, w) => s + w, 0), naiveTotal };  }}

k =: 8 entries
saved =: 896

Every token decoded: 672 multiply-adds against 1,568 without a cache, 2.33 times less work, for 256 bytes of cached keys and values.

import type { Tracer } from "../trace";import { Kernels } from "./kernels";import { SelfAttention } from "./self-attention";/** Keep every token's key and value, or project the whole prefix again at every step. */export type DecodeMode = "cache" | "recompute";export interface DecodeResult { /** weights[query][key] over the keys at or before the query: a lower triangle, each row summing to one. */ weights: number[][]; /** One mixed value vector per decoded token. */ outputs: number[][]; /** Multiply-adds spent at each step, by the run's own ledger. */ work: number[]; /** The ledger's sum over every step. */ total: number; /** What the same run costs without a cache, by the closed form. */ naiveTotal: number;}/** * Causal decoding one token at a time over a self-attention's public * weights. Every step projects the new token into a query, a key, and a * value, scores the query against the key of every token so far, and * blends their values. The keys and values of earlier tokens never * change, because each one depends only on its own vector and the fixed * matrices, so a cache keeps them and a step projects one token; the * query is used once and is never kept. Without the cache the same * step projects the whole prefix again. A ledger counts multiply-adds: * a vector through a d by d matrix is d², a dot product or a blend over * n keys is n·d. * * Time per step: O(d² + n·d) with the cache, O(n·d²) without. * Space: O(n·d) for the cache, whatever the mode spends per step. */export class KvCache { readonly attention: SelfAttention; private ledger = 0; constructor( readonly dim: number, readonly mode: DecodeMode = "cache", seed = 1, private readonly tracer?: Tracer, ) { this.attention = new SelfAttention(dim, "identity", seed); } /** One d-vector through a d by d matrix: d² multiply-adds on the step's ledger. */ private project(m: readonly (readonly number[])[], x: readonly number[]): number[] { this.ledger += this.dim * this.dim; return m.map((row) => row.reduce((s, weight, i) => s + weight * (x[i] ?? 0), 0)); } decode(inputs: readonly (readonly number[])[]): DecodeResult { if (inputs.some((x) => x.length !== this.dim)) throw new RangeError(`every input must have ${this.dim} numbers`); const scale = 1 / Math.sqrt(this.dim); const k: number[][] = []; const keyCache: number[][] = []; const v: number[][] = []; const valueCache: number[][] = []; const weights: number[][] = []; const outputs: number[][] = []; const work: number[] = []; let naiveTotal = 0; for (let step = 0; step < inputs.length; step++) { this.ledger = 0; if (this.mode === "recompute") { // Nothing survives a step without the cache: the list is rebuilt from the first token. k.length = 0; keyCache.length = 0; v.length = 0; valueCache.length = 0; } let q: number[] = []; let queryVector: number[] = []; for (let index = k.length; index <= step; index++) { for (let index = keyCache.length; index <= step; index++) { const x = inputs[index]!; const recomputed = index < step; const key = this.project(this.attention.wk, x); const value = this.project(this.attention.wv, x); // Only the newest token needs a query: nothing asks on behalf of an earlier one again. const query = recomputed ? null : this.project(this.attention.wq, x); if (query) q = query; if (query) queryVector = query; k.push(key); keyCache.push(key); v.push(value); valueCache.push(value); // ▸ a token's vector through the fixed matrices; an old one needs no query this.tracer?.emit("project", { step, index, recomputed, query: query ? [...query] : null, key: [...key], value: [...value] }); } // ▸ the new key and value join the cache for every later step if (this.mode === "cache") this.tracer?.emit("append", { step, entries: k.length }); if (this.mode === "cache") this.tracer?.emit("append", { step, entries: keyCache.length }); const scores: number[] = []; for (let index = 0; index < k.length; index++) { for (let index = 0; index < keyCache.length; index++) { const dot = q.reduce((s, value, i) => s + value * k[index]![i]!, 0); const dot = queryVector.reduce((s, value, i) => s + value * keyCache[index]![i]!, 0); this.ledger += this.dim; const scaled = dot * scale; scores.push(scaled); // ▸ the new query against one cached key, scaled down by root d this.tracer?.emit("score", { query: step, key: index, dot, scaled }); } const row = Kernels.softmax(scores); const top = row.indexOf(Math.max(...row)); weights.push(row); // ▸ the scores over the tokens so far become weights that sum to one this.tracer?.emit("softmax", { query: step, top, weights: [...row] }); const output = Array.from({ length: this.dim }, (_, i) => row.reduce((s, w, index) => s + w * v[index]![i]!, 0)); const output = Array.from({ length: this.dim }, (_, i) => row.reduce((s, w, index) => s + w * valueCache[index]![i]!, 0)); this.ledger += k.length * this.dim; this.ledger += keyCache.length * this.dim; outputs.push(output); // ▸ the output is the cached values blended by those weights this.tracer?.emit("mix", { query: step, output: [...output] }); const cachedWork = 3 * this.dim * this.dim + 2 * (step + 1) * this.dim; const naiveWork = (2 * (step + 1) + 1) * this.dim * this.dim + 2 * (step + 1) * this.dim; work.push(this.ledger); naiveTotal += naiveWork; // ▸ the step's multiply-adds beside what projecting everything again costs this.tracer?.emit("compare", { step, work: this.ledger, cachedWork, naiveWork, saved: naiveWork - this.ledger }); } return { weights, outputs, work, total: work.reduce((s, w) => s + w, 0), naiveTotal }; }}