Profile KV-cache work across serving workloads

Compare cached and uncached attention on identical inputs. Check output equivalence first, then inspect projected rows, attention-score cells and retained state across fixed prompt and generation lengths.

Problem

When you need this. You need to compare cached and uncached generation by the computation each actually performs.

A generator can repeat substantial work while producing only one new token. Without a cache, an uncomplicated implementation runs attention over the entire growing prefix at every step. A key-value cache retains intermediate state that causal attention can reuse.

You will use two existing attention paths, verify equivalent outputs, and compare the work and retained state across four workloads. The lab uses one attention head with fixed random weights on CPU. It is small enough to inspect; it is not a pretrained language model or a GPU serving benchmark.

Concepts

Prefill processes the prompt before generation. Decode adds tokens one at a time. A key-value cache retains the keys and values of earlier positions, so a decode step can project the new token without projecting the complete prefix again.

The experiment distinguishes a work counter from elapsed time. Projected rows and query-key score cells count operations performed by this implementation. Timing also includes Python overhead. Cache payload is retained numerical state; it is not the same quantity as total process memory. The from-scratch course owns implementing the cache; this lesson owns comparing workloads.

Architecture

The two paths share weights, inputs and token choices. The full path recomputes a growing prefix; the step path appends one position to retained keys and values. The comparison checks equivalent outputs before collecting work counts and local timings. The workload sweep changes lengths while keeping that comparison rule fixed.

Working code

Fix the workload before comparing implementations

The from-scratch KV-cache lesson explains and implements caching in the trained model. This experiment asks a different question: how does the amount of saved work change when prompts or generated sequences grow?

The supplied attention layer is a small measurement fixture. Each decode step computes the new token’s query, key and value and retains earlier keys and values. The uncached path reruns the growing prefix. Keep the model, prompt and chosen tokens identical between paths before comparing their work.

Download and extract the pilot lab. In ai-system-kv-cache, run:

python3.12 -m venv .venv
source .venv/bin/activate
python cache.py
python -m unittest -v test_cache.py

Establish equivalence before measuring speed

full() projects every prefix position and calculates every causal attention row. The uncached generator uses only the final row to choose its next token. On the next iteration it repeats that work on a longer prefix.

step() projects one new position, appends its key and value, and calculates one attention output. The cached generator first processes the prompt, then processes each newly appended token. The final generated token has not yet been fed back through the layer when generation stops.

"""Single-head causal attention with fixed weights; not a pretrained language model."""
from dataclasses import dataclass, field
import hashlib
import json
import math
import platform
import random
import statistics
import time
D=8
rng=random.Random(17)
EMBED=[[rng.uniform(-1,1) for _ in range(D)] for _ in range(12)]
WEIGHTS=[[[rng.uniform(-.3,.3) for _ in range(D)] for _ in range(D)] for _ in range(3)]
MODEL_ID=hashlib.sha256(json.dumps([EMBED,WEIGHTS]).encode()).hexdigest()
def dot(a,b): return math.fsum(x*y for x,y in zip(a,b))
def project(x,w): return [dot(row,x) for row in w]
def qkv(token,position):
    if type(token) is not int or not 0<=token<len(EMBED): raise ValueError("invalid token ID")
    x=[v+math.sin(position/(10000**(i/D))) for i,v in enumerate(EMBED[token])]
    return [project(x,w) for w in WEIGHTS]
def attend(q,keys,values):
    scores=[dot(q,k)/math.sqrt(D) for k in keys]
    largest=max(scores); weights=[math.exp(x-largest) for x in scores]
    total=math.fsum(weights); weights=[x/total for x in weights]
    return [math.fsum(a*v[i] for a,v in zip(weights,values)) for i in range(D)]
def full(tokens):
    if not tokens: raise ValueError("empty prefix")
    projected=[qkv(t,p) for p,t in enumerate(tokens)]
    return [attend(row[0],[r[1] for r in projected[:p+1]],
                   [r[2] for r in projected[:p+1]]) for p,row in enumerate(projected)]
@dataclass
class Cache:
    model_id: str = MODEL_ID
    keys: list = field(default_factory=list)
    values: list = field(default_factory=list)
    tokens: list = field(default_factory=list)
def step(cache,token,position):
    if cache.model_id != MODEL_ID: raise ValueError("cache belongs to another model")
    if position != len(cache.tokens): raise ValueError("cache position mismatch")
    q,k,v=qkv(token,position)
    cache.keys.append(k);cache.values.append(v);cache.tokens.append(token)
    return attend(q,cache.keys,cache.values)
def next_token(state): return max(range(len(EMBED)),key=lambda i:dot(state,EMBED[i]))
def generate(prefix,steps,cached):
    if not prefix or steps<1: raise ValueError("nonempty prefix and positive steps required")
    tokens=list(prefix);cache=Cache();position=0;projected_rows=0;score_cells=0
    for _ in range(steps):
        if cached:
            while position<len(tokens):
                state=step(cache,tokens[position],position)
                position+=1;projected_rows+=1;score_cells+=position
        else:
            state=full(tokens)[-1];n=len(tokens)
            projected_rows+=n;score_cells+=n*(n+1)//2
        tokens.append(next_token(state))
    return {"tokens":tokens,"projected_token_rows":projected_rows,
            "attention_score_cells":score_cells,"cached_positions":len(cache.tokens)}
def workload_sweep():
    """Compare deterministic work for fixed, predeclared serving workloads."""
    rows=[]
    for prompt_tokens,new_tokens in ((8,8),(32,8),(8,32),(32,32)):
        prefix=[i%len(EMBED) for i in range(prompt_tokens)]
        uncached=generate(prefix,new_tokens,False)
        cached=generate(prefix,new_tokens,True)
        rows.append({"prompt_tokens":prompt_tokens,"new_tokens":new_tokens,
                     "identical_tokens":uncached["tokens"]==cached["tokens"],
                     "uncached_rows":uncached["projected_token_rows"],
                     "cached_rows":cached["projected_token_rows"],
                     "uncached_score_cells":uncached["attention_score_cells"],
                     "cached_score_cells":cached["attention_score_cells"],
                     "retained_kv_floats":2*D*cached["cached_positions"]})
    return rows

def report():
    prefix=[1,4,2,9,3,5,8,7];steps=8;records={}
    for mode in [False,True]:
        times=[]
        for _ in range(15):
            start=time.perf_counter();result=generate(prefix,steps,mode)
            times.append((time.perf_counter()-start)*1000)
        records['cached' if mode else 'uncached']={**result,"median_ms":round(statistics.median(times),3)}
    return {"implementation":"fixed-weight single attention layer; CPU Python lists",
            "python":platform.python_version(),"platform":platform.platform(),
            "model_sha256":MODEL_ID,"prompt_tokens":len(prefix),"generated_tokens":steps,
            "timing_repetitions":15,"results":records,"workload_sweep":workload_sweep(),
            "identical_tokens":records['cached']['tokens']==records['uncached']['tokens']}
if __name__ == '__main__': print(json.dumps(report(),indent=2,sort_keys=True))

The first test compares every cached position with a complete forward calculation. Another compares generated token sequences. This order matters: a faster implementation that changes the output has changed the experiment.

A cache with the wrong position is rejected before mutation. A cache from another model is rejected by its model identifier. Production cache keys may also need to bind tokenization, adapters, position settings and tenant policy. A shared cache must preserve the relevant isolation boundary.

Separate counted work from elapsed time

The recorded run starts with 8 prompt tokens and generates 8 tokens. Both paths produce the same sequence in this run.

A projected token row means three vector projections, one each for Q, K and V. A score cell means one query-key dot product. Neither quantity is a count of all floating-point operations in a language model.

The timing is the median of 15 repetitions on the Python version and machine recorded below. Small Python allocations and interpreter overhead affect it. Use the deterministic work counts to understand the mechanism; rerun timing on the system you intend to optimize.

{
  "generated_tokens": 8,
  "identical_tokens": true,
  "implementation": "fixed-weight single attention layer; CPU Python lists",
  "model_sha256": "e305ca8f0ebe35367f04b0976405ec1357ff940690e0b62fa12d34713d216554",
  "platform": "macOS-27.0.1-arm64-arm-64bit-Mach-O",
  "prompt_tokens": 8,
  "python": "3.15.0",
  "results": {
    "cached": {
      "attention_score_cells": 120,
      "cached_positions": 15,
      "median_ms": 0.352,
      "projected_token_rows": 15,
      "tokens": [
        1,
        4,
        2,
        9,
        3,
        5,
        8,
        7,
        8,
        8,
        8,
        8,
        3,
        3,
        3,
        3
      ]
    },
    "uncached": {
      "attention_score_cells": 596,
      "cached_positions": 0,
      "median_ms": 1.797,
      "projected_token_rows": 92,
      "tokens": [
        1,
        4,
        2,
        9,
        3,
        5,
        8,
        7,
        8,
        8,
        8,
        8,
        3,
        3,
        3,
        3
      ]
    }
  },
  "timing_repetitions": 15,
  "workload_sweep": [
    {
      "cached_rows": 15,
      "cached_score_cells": 120,
      "identical_tokens": true,
      "new_tokens": 8,
      "prompt_tokens": 8,
      "retained_kv_floats": 240,
      "uncached_rows": 92,
      "uncached_score_cells": 596
    },
    {
      "cached_rows": 39,
      "cached_score_cells": 780,
      "identical_tokens": true,
      "new_tokens": 8,
      "prompt_tokens": 32,
      "retained_kv_floats": 624,
      "uncached_rows": 284,
      "uncached_score_cells": 5204
    },
    {
      "cached_rows": 39,
      "cached_score_cells": 780,
      "identical_tokens": true,
      "new_tokens": 32,
      "prompt_tokens": 8,
      "retained_kv_floats": 624,
      "uncached_rows": 752,
      "uncached_score_cells": 10576
    },
    {
      "cached_rows": 63,
      "cached_score_cells": 2016,
      "identical_tokens": true,
      "new_tokens": 32,
      "prompt_tokens": 32,
      "retained_kv_floats": 1008,
      "uncached_rows": 1520,
      "uncached_score_cells": 38224
    }
  ]
}

Account for the state you retain

For a conventional uncompressed KV cache, tensor payload grows with batch size, cached positions, layer count, KV-head count, head dimension and bytes per stored element. Both K and V consume space. Python lists in this lab also have object overhead, so multiplying element count by a numeric dtype size would not report the process’s actual memory use.

Within one decode step, the new query still reads the retained keys and values. KV caching removes repeated prefix computation; it does not make attention over an ever-growing context free.

Grouped-query attention changes the number of KV heads. KV quantization changes how their values are represented. Prefix caching can reuse eligible prompt computation across requests. These are different mechanisms and need separate correctness, quality and isolation checks.

Compare the four workload shapes

workload_sweep() fixes prompt lengths of 8 and 32 and generation lengths of 8 and 32 before measuring anything. Its counters come from the executed generation loops. They count rows and score cells, rather than inferring a speedup from a formula.

Read across rows with the same prompt length to isolate generation length. Then compare rows with the same generation length to isolate prompt length. Longer histories save more repeated computation but retain more state. These counters help explain a serving profile; they do not account for scheduling, GPU utilization or memory bandwidth.

Design the next measurement

Change the prompt length while keeping the generation length fixed. Before running the experiment, predict which counted quantities grow. Then increase generation length while keeping the prompt fixed.

For a serving system, measure queueing, time to first token, latency between output tokens, throughput and memory under a specified concurrency. Time to first token includes more than the prefill computation. This local program supplies neither a request queue nor GPU kernels, so its timing cannot establish production capacity.

Design exercise: two users share the same system instruction but have different retrieved documents. Which prefix could a serving engine reuse, and what must be included in its cache identity? Explain why a matching English sentence alone is insufficient.

Failure cases

A cache is valid only for the model and position sequence that produced it. The tests check that a wrong-position update refuses before changing state and that another model cannot consume the cache. A speed measurement is uninterpretable if either path computes a different result.

Deployment

Run this as a local measurement program. Preserve the prompt lengths, generation lengths, weights and counter definitions with any comparison. Before transferring the approach into a serving runtime, test request isolation, cache invalidation and numerical equivalence there; the lab does not implement a shared serving cache.

Scaling

Longer prompts increase prefill work and retained state. Longer generations add decode steps, each of which still reads earlier keys and values. Use the supplied sweep to inspect those two effects separately. Batch size, layers, KV-head count and stored dtype also affect payload in a real model; this single-head Python fixture does not measure those serving configurations.

Key takeaways

• Equivalence is a prerequisite for comparing saved work.
• Cached decode avoids repeated prefix projection while retaining a growing state.
• Operation counts, local timing and process memory answer different measurement questions.

Best practices

• Keep the token sequence and weights identical between the two paths.
• Report prompt and generation lengths beside results.
• Measure the target runtime before turning a small Python timing ratio into a serving claim.

References

• Hugging Face: How caching works: causal attention, cached keys and values, and the new position’s query.
• GQA paper: grouped-query attention and its KV-head tradeoff.
• PagedAttention paper: managing growing KV storage and sharing it in a serving system.

Sources reviewed October 23, 2026. Operation counts and local timings are generated by cache.py.