What is KV-cache and how does prefix caching speed up LLM inference?

// KV cache from first principles: a toy 2-layer causal transformer (d=4, vocab 10).
// Generates the same tokens with and without a KV cache and counts the K/V work.
const D = 4, VOCAB = 10, LAYERS = 2;
const fmt = (n) => Math.round(n).toString().replace(/\B(?=(\d{3})+(?!\d))/g, ",");

console.log("== 1. KV cache size per token = 2 (K,V) x layers x kv_heads x head_dim x bytes ==");
for (const [name, layers, kvHeads] of [["Llama 3 8B (GQA, 8 KV heads)", 32, 8], ["same, if all 32 heads kept K/V", 32, 32], ["Llama 3 70B (GQA, 8 KV heads)", 80, 8]]) {
  const per = 2 * layers * kvHeads * 128 * 2;
  console.log(`${name.padEnd(31)} ${fmt(per / 1024)} KiB/token, 8,192-token context ${(per * 8192 / 2 ** 30).toFixed(2)} GiB`);
}

const E = [...Array(VOCAB)].map((_, t) => [...Array(D)].map((_, j) => Math.sin(t * 1.3 + j * 0.7)));
const W = [...Array(LAYERS)].map((_, l) => [0, 1, 2].map((m) =>
  [...Array(D)].map((_, i) => [...Array(D)].map((_, j) => Math.cos((l + 1) * (i * D + j) * 0.37 + m) * 0.5))));
const matvec = (x, M) => M[0].map((_, j) => x.reduce((s, xi, i) => s + xi * M[i][j], 0));
const dot = (a, b) => a.reduce((s, ai, i) => s + ai * b[i], 0);

function attend(q, keys, vals) {
  const sc = keys.map((k) => dot(q, k) / Math.sqrt(D));
  const mx = Math.max(...sc), ex = sc.map((s) => Math.exp(s - mx)), z = ex.reduce((a, b) => a + b, 0);
  return vals[0].map((_, j) => vals.reduce((s, v, i) => s + (ex[i] / z) * v[j], 0));
}
// Full forward over every position (no cache). causal=false lets a position see the future.
function forward(tokens, causal = true, counter = null) {
  let xs = tokens.map((t) => E[t].slice());
  const kByLayer = [];
  for (let l = 0; l < LAYERS; l++) {
    const q = xs.map((x) => matvec(x, W[l][0])), k = xs.map((x) => matvec(x, W[l][1])), v = xs.map((x) => matvec(x, W[l][2]));
    if (counter) counter.rows += xs.length;
    kByLayer.push(k);
    xs = xs.map((x, i) => {
      const n = causal ? i + 1 : xs.length;
      const a = attend(q[i], k.slice(0, n), v.slice(0, n));
      return x.map((xi, j) => xi + a[j]);
    });
  }
  return { last: xs[xs.length - 1], kByLayer };
}
const argmax = (h) => { let best = 0; for (let t = 1; t < VOCAB; t++) if (dot(h, E[t]) > dot(h, E[best])) best = t; return best; };

console.log("\n== 2. Generate 6 tokens after a 4-token prompt ==");
const prompt = [3, 1, 4, 1];
const noCache = { rows: 0 }, seqA = prompt.slice(), lastA = [];
for (let s = 0; s < 6; s++) { const r = forward(seqA, true, noCache); lastA.push(r.last); seqA.push(argmax(r.last)); }

const cacheK = [...Array(LAYERS)].map(() => []), cacheV = [...Array(LAYERS)].map(() => []);
let rowsCached = 0;
function step(token) { // process ONE new position, reuse cached K/V of all earlier ones
  let x = E[token].slice();
  for (let l = 0; l < LAYERS; l++) {
    cacheK[l].push(matvec(x, W[l][1])); cacheV[l].push(matvec(x, W[l][2])); rowsCached++;
    const a = attend(matvec(x, W[l][0]), cacheK[l], cacheV[l]);
    x = x.map((xi, j) => xi + a[j]);
  }
  return x;
}
const seqB = prompt.slice(), lastB = [];
let h; for (const t of prompt) h = step(t); // prefill fills the cache
for (let s = 0; s < 6; s++) { lastB.push(h); const nt = argmax(h); seqB.push(nt); if (s < 5) h = step(nt); }
let maxDiff = 0;
lastA.forEach((v, s) => v.forEach((x, j) => { maxDiff = Math.max(maxDiff, Math.abs(x - lastB[s][j])); }));
console.log(`no cache : tokens ${seqA.join(" ")}  K/V rows computed ${noCache.rows}`);
console.log(`KV cache : tokens ${seqB.join(" ")}  K/V rows computed ${rowsCached}`);
console.log(`max difference in hidden states: ${maxDiff}`);

console.log("\n== 3. Why caching is allowed: old K/V never change in a causal model ==");
for (const causal of [true, false]) {
  const short = forward(seqA.slice(0, 3), causal).kByLayer[1][1], long = forward(seqA, causal).kByLayer[1][1];
  const d = Math.max(...short.map((x, j) => Math.abs(x - long[j])));
  console.log(`${causal ? "causal mask  " : "no mask      "} layer-2 K of position 1, context 3 vs 10 tokens: max change ${d.toFixed(4)}`);
}

console.log("\n== 4. Same counting at real scale: 2,000-token prompt, 300 new tokens, 80 layers ==");
let without = 0; for (let k = 1; k <= 300; k++) without += 2000 + k - 1;
const withc = 2000 + 299;
console.log(`token-positions processed: no cache ${fmt(without)}, with cache ${fmt(withc)} (${(without / withc).toFixed(0)}x), x 80 layers each`);

// Block-hash prefix cache (the idea behind automatic prefix caching), LRU eviction.
// Block size 4 tokens for readability; capacity 8 blocks. Words stand for token ids.
const BLOCK = 4, CAP = 8;

function fnv1a(s) {
  let h = 0x811c9dc5;
  for (let i = 0; i < s.length; i++) { h ^= s.charCodeAt(i); h = Math.imul(h, 0x01000193) >>> 0; }
  return h.toString(16).padStart(8, "0");
}

const cache = new Set(); // block hashes whose K/V tensors are resident
let order = [];          // LRU order: front = evict first
let totalTokens = 0, totalComputed = 0;

function blockHashes(tokens) {
  const out = []; let parent = "root";
  for (let i = 0; i + BLOCK <= tokens.length; i += BLOCK) {
    parent = fnv1a(parent + "|" + tokens.slice(i, i + BLOCK).join(" "));
    out.push(parent); // a block hash covers ALL tokens before it, not just its own 4
  }
  return out;
}

function serve(name, text) {
  const tokens = text.split(" ");
  const hashes = blockHashes(tokens);
  let hit = 0;
  while (hit < hashes.length && cache.has(hashes[hit])) hit++; // stop at the first miss
  const evicted = [];
  for (let b = hit; b < hashes.length; b++) {
    if (cache.has(hashes[b])) continue;
    while (cache.size >= CAP) {
      const victim = order.find((h) => !hashes.includes(h));
      order = order.filter((h) => h !== victim); cache.delete(victim); evicted.push(victim.slice(0, 4));
    }
    cache.add(hashes[b]);
  }
  // touch tail first so the first block ends up most recent: parents outlive children
  for (let b = hashes.length - 1; b >= 0; b--) { order = order.filter((h) => h !== hashes[b]); order.push(hashes[b]); }
  const computed = tokens.length - hit * BLOCK;
  totalTokens += tokens.length; totalComputed += computed;
  console.log(`${name}: ${tokens.length} tokens, ${hashes.length} full blocks, hit ${hit} blocks (${hit * BLOCK} tokens), prefill ${computed}` +
    (evicted.length ? `, evicted ${evicted.join(",")}` : "") + `, cached ${cache.size}/${CAP}`);
  return hashes;
}

const SYS = "you are the support bot for acme corp . be brief .";
console.log(`system prompt = ${SYS.split(" ").length} tokens = 3 blocks of ${BLOCK}`);
const h1 = serve("r1 first user  ", SYS + " how do i reset my password ?");
console.log(`r1 block hashes ${h1.map((h) => h.slice(0, 4)).join(" -> ")}`);
serve("r2 other user  ", SYS + " where is my order ?");
serve("r3 r1 retried  ", SYS + " how do i reset my password ?");
serve("r4 one word off", SYS.replace("support", "sales") + " where is my order ?");
serve("r5 r1 turn two ", SYS + " how do i reset my password ? click forgot password on the login page . thanks , and my email ?");
serve("r6 new topic   ", "translate to french : the cat sleeps on the warm mat today");
serve("r7 back to sys ", SYS + " can i change my plan ?");
console.log(`total: ${totalTokens} prompt tokens, ${totalComputed} computed, ${totalTokens - totalComputed} served from cache (${(100 * (totalTokens - totalComputed) / totalTokens).toFixed(1)}%)`);