Reusing the Prompt Prefix with a Key-Value Cache for SLM Optimization

In a previous article we discussed constraining output space for small language model (SLM) narrow automation optimization. We mentioned at the time that this was the first in a short series of SLM optimization strategy articles. We are now at the second of those. This time we will focus on the reuse of the prompt prefix with a key-value cache. Let's not waste any further time on the niceties and get right to it instead.

As in our previous article, all benchmarks below 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.

First, make sure you set up a Python environment and install your requirements:

pip install torch transformers accelerate

We will continue to use the support ticket framing from our first article.

Narrow automation prompts are relatively static. A task instruction, a taxonomy definition, and a few examples make up the bulk of the tokens, and only a short prompt tail changes from item to item. If your instruction block runs to a couple of hundred tokens and each ticket adds twenty or thirty, then the overwhelming majority of every prompt is byte-for-byte identical to the last one. It's clearly wasteful to be recomputing all of it, for every layer, on every call.

Transformers compute a key and a value vector for each token at each layer, and these depend only on the tokens to the left. This means that for a fixed prefix, they are identical on every call. Computing them once and holding onto them shrinks the per-item pre-fill down to just the tokens that actually changed.

Re-encoding Every Ticket

Output:

Reusing the Prompt Prefix

We now 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, the tokenizer, the prompt halves, and the baseline timing are all still in scope.

Output:

Some further explanation of the code above: