Leveraging a Key-Value Cache to Reuse Prompt Prefixes for Optimizing Small Language Models

If you run small language models for narrow automation tasks, you may be paying a hidden compute tax on every single call. This tutorial, the second in a series on SLM optimization strategies, shows you how to eliminate that waste using prompt prefix caching with a key-value cache. By the end, you will understand why most of your prompt tokens never need to be recomputed, and how to cut per-item processing time by more than half. If you have not yet read the first article on constraining output space, it is worth a look, though this piece stands on its own.

As before, all benchmarks use Qwen2.5-0.5B-Instruct in float16 through Hugging Face Transformers, running on an M2 Macbook Air with 24GB RAM and a 16-core Neural Engine. Set up a Python environment and install the requirements, and note that we continue the support-ticket framing from the first article.

Why Reuse the Prompt Prefix with a Key-Value Cache

Narrow automation prompts tend to be quite static. A task instruction, a taxonomy definition, and a handful of examples account for most of the tokens, while only a short tail actually changes from one item to the next. If your instruction block spans a couple of hundred tokens and each ticket contributes twenty or thirty, then nearly all of every prompt is identical, byte for byte, to the previous one. Recomputing all of that, for every layer, on every call, is clearly wasteful.

Here is the key insight: Transformers compute a key and a value vector for every token at every layer, and these depend only on the tokens to the left. For a fixed prefix, they are therefore identical across every call. Computing them once and storing them shrinks each item's pre-fill down to just the tokens that actually changed.

Re-encoding Every Ticket (The Baseline)

We begin with a baseline: a realistic few-shot prompt, re-encoded in full for each ticket. The constrained scoring technique from the earlier article carries forward here, so each decision is a single forward pass and the only remaining thing to optimize is the pre-fill. One detail matters: the chat template is written out by hand rather than generated with `apply_chat_template()`, because the next script needs to split the prompt at a known boundary.

Reusing the Prompt Prefix

Now we run the prefix through the model exactly once, keep the resulting cache, and feed each item only its own tokens. This continues in the same script, so the model, tokenizer, prompt halves, and baseline timing all remain in scope.

The results speak for themselves. The baseline ran for 184.85 seconds in total, averaging roughly 0.3 seconds per ticket. The cached prefix version finished in 80.07 seconds, averaging 0.13 seconds per ticket — an overall runtime reduction of about 57%. The gain scales with the ratio of static to dynamic content, which means this technique rewards long, detailed instruction blocks rather than punishing them. And importantly, predictions should be identical for every ticket. That is exactly what you want: this is a pure compute optimization, not a change to model behavior.

How the Code Works

A few details deserve closer explanation:

DynamicCache. This object holds the per-layer key and value tensors for the prefix. Passing it as `past_key_values` tells the model those positions are already computed, so the forward pass processes only the new tokens while still attending backward across the full context.

Two arguments must agree with the cache. The attention mask spans the cached prefix and the new tokens, so its width is `prefix_len + suffix_len` even though only `suffix_len` ids are passed in. Meanwhile, `cache_position` tells the model the new tokens begin at offset `prefix_len`, so the rotary embeddings match what the full prompt would have produced. Recent versions of Transformers can infer positions from cache length, but passing them explicitly documents intent and guards against version drift.

Cropping is mandatory. The call to `prefix_cache.crop(prefix_len)` is not optional. Each forward pass appends the suffix keys and values to the cache, so without cropping, the second ticket would attend to the first ticket's tokens, and the cache would grow without bound.

Split at a clean boundary. Cut the prompt at the end of a line or a chat template delimiter. Tokenizing two halves separately can produce a different token sequence than tokenizing the concatenation if the split lands mid-word — and then the cached keys no longer correspond to what the model would actually have seen. The assertion in the first script checks this directly.

Use `torch.no_grad()`, not `torch.inference_mode()`. Tensors created inside inference mode carry a flag that makes them awkward to slice and reassign afterward — which is precisely what `crop()` does to the cache.

Wrapping Up

This was the second entry in our series on optimizing SLMs for narrow automation, and the technique in focus was prefix caching. It replaces the complete re-encoding of a static instruction block with a single pre-fill whose key and value tensors are computed once and reused, leaving only the differing tokens to account for. Implementing it delivers the same predictions the naive loop produced, at a far lower compute cost — and the longer and more detailed your instructions grow, the better the trade becomes.

A small language model, like the 0.5B parameter model used here, becomes a practical production choice for narrow automation once the surrounding code stops treating every call as an isolated event. The prompt is mostly the same every time. Once your loop knows that, the small model stops being a compromise and starts being the obvious answer.

Common Mistakes

• Forgetting to crop the cache, causing tickets to attend to earlier tickets' tokens and the cache to grow endlessly.
• Splitting the prompt mid-word, which produces token sequences that do not match the full concatenation.
• Mismatched attention mask or cache positions, yielding silently wrong results.
• Using `torch.inference_mode()` instead of `torch.no_grad()`, which complicates slicing and reassignment of cache tensors.