Generation and the KV cache¶

During generation a decoder reads the keys and values of earlier positions from its KV cache instead of recomputing them. This page shows how to trace a generation loop, how to edit what the cache holds, and what the cache is to TorchLens.

Trace a generation loop¶

To look inside generation rather than edit it, trace the model's own generate method. Each forward call generate makes is one call of every module, so a module's num_calls is the number of generated tokens and module:k is its output while generating token k. Call 1 is the prefill, which reads the whole prompt; each later call reads only the newest token, because the earlier ones are in the KV cache. Module addresses are the model's own, as in a plain trace of the model; the model itself is the root self, called once per token.

# docs-test: expected=[3, [1, 5, 32], [1, 1, 32]]
import torch
import torchlens as tl
from transformers import GPT2Config, GPT2LMHeadModel

torch.manual_seed(0)
lm = GPT2LMHeadModel(GPT2Config(vocab_size=128, n_embd=32, n_layer=2, n_head=2)).eval()
prompt = torch.tensor([[3, 17, 42, 8, 99]])

trace = tl.trace(lm.generate, prompt, input_kwargs={"max_new_tokens": 3})
mlp = "transformer.h.1.mlp"
RESULT = [
    trace[mlp].num_calls,  # 3: one call per generated token
    list(trace[f"{mlp}:1"].out.shape),  # the prefill: all 5 prompt positions
    list(trace[f"{mlp}:2"].out.shape),  # a decode step: the one new position
]

With a Hugging Face model whose tokenizer TorchLens can find, pass the prompt as text, as with a plain trace: tl.trace(model.generate, "The capital of France is", input_kwargs={...}). generate checks whether to stop by reading tensors as Python values, so the trace warns once with a ScalarEscapeWarning: that stopping decision is not part of the recorded graph. A trace records every operation of every step, so keep it to a few tokens; to keep chosen sites across a long generation, use tl.record(model, ids, save=..., generate_kwargs={...}), which records only the sites you select at each step.

Recipe: edit the cached past¶

"Intervene during generation" can mean three different things, and they cost very different amounts:

A bound model runs the model's real generate() with the edit installed as a hook and records no graph; bound.last_report.num_fires counts the edit's firings, so an edit on a projection that runs once per decoder forward fires once per step, and the keys and values it writes into the cache already carry the edit. Use a trace or an episode capture only when you need the recorded graph, and then for a few steps.

To change what later tokens see of the prompt without touching the new tokens, apply the edit to the key or value projection during the prefill call only. The edited values are written into the cache, and every later step reads them from there. The check below compares the result with the same edit written as a plain forward hook.

# docs-test: expected=[True, True]
import torch
import torchlens as tl
from transformers import LlamaConfig, LlamaForCausalLM

torch.manual_seed(0)
lm = LlamaForCausalLM(
    LlamaConfig(
        vocab_size=128,
        hidden_size=32,
        intermediate_size=64,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
    )
).eval()
prompt = torch.tensor([[3, 17, 42, 8, 99]])
next_token = torch.tensor([[7]])
v_proj = lm.model.layers[0].self_attn.v_proj


def blank_position_1(values):
    values = values.clone()
    values[:, 1] = 0.0  # the cached value of prompt position 1 in layer 0
    return values


def decode_after(prefill):
    """Run one decoding step that reads the cache the prefill call wrote."""
    with torch.no_grad():
        cache = prefill(prompt, use_cache=True).past_key_values
        return lm(next_token, past_key_values=cache, use_cache=True).logits


# TorchLens: the edit is bound for the prefill call only.
spec = tl.when(
    tl.module("model.layers.0.self_attn.v_proj"),
    lambda values, *, ctx: blank_position_1(values),
)
edited = decode_after(spec.bind(lm))

# The same surgery with a plain hook on the prefill call.
handle = v_proj.register_forward_hook(lambda module, args, out: blank_position_1(out))
try:
    with torch.no_grad():
        cache = lm(prompt, use_cache=True).past_key_values
finally:
    handle.remove()
with torch.no_grad():
    reference = lm(next_token, past_key_values=cache, use_cache=True).logits

clean = decode_after(lm)
RESULT = [torch.allclose(edited, reference, atol=1e-6), not torch.allclose(edited, clean)]

When you also trace these calls, remember that the whole-graph default saves every op's activation, and the attention matrix alone grows with the square of the sequence length. Pass save=tl.module(...) naming the sites you will read, and see long sequences in the performance guide.

What the cache is to TorchLens¶

When you trace a decoding step, the cache passed as past_key_values is an input container like any other: its tensors are graph inputs, addressed by their path inside the container (containers). An edit at a module inside the step changes what that step computes and what it appends to the cache; it does not rewrite earlier entries. To change earlier entries, edit them at the call that wrote them, as the recipe does.