{"slug": "reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization", "title": "Reusing the Prompt Prefix with a Key-Value Cache for SLM Optimization", "summary": "A second article in a KDnuggets series on small language model (SLM) optimization details reusing the prompt prefix with a key-value cache to cut per-item pre-fill costs, benchmarked on Qwen2.5-0.5B-Instruct in float16 via Hugging Face Transformers on an M2 MacBook Air with 24GB RAM and a 16-core Neural Engine. The technique exploits the fact that key and value vectors for a fixed prefix depend only on tokens to the left and are identical on every call, so computing them once shrinks pre-fill to only the changed tokens; the baseline re-encodes a few-shot support-ticket prompt in full for each of 600 records across the labels billing, technical, and account.", "body_md": "# Reusing the Prompt Prefix with a Key-Value Cache for SLM Optimization\n\nIn this second article in our short series on SLM optimization techniques we focus on the reuse of the prompt prefix with a key-value cache.\n\nIn [a previous article](https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization) 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.\n\nAs in our previous article, all benchmarks below use [**Qwen2.5-0.5B-Instruct**](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct) in float16 through **[Hugging Face Transformers](https://huggingface.co/docs/transformers/index)**, running on an M2 Macbook Air with 24GB RAM and a 16-core Neural Engine.\n\nFirst, make sure you set up a Python environment and install your requirements:\n\n```\npip install torch transformers accelerate\n```\n\nWe will continue to use the support ticket framing [from our first article](https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization).\n\n## Why Reuse the Prompt Prefix with a Key-Value Cache\n\nNarrow 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.\n\nTransformers 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.\n\n## Re-encoding Every Ticket\n\nFirst, the baseline: a realistic few-shot prompt, re-encoded in full for every ticket. The constrained scoring trick from our earlier article is carried forward here, so the decision is a single forward pass and the only thing left to optimize is the pre-fill. Note that the chat template is written out by hand rather than going through `apply_chat_template()`, because the next script needs to split the prompt at a known boundary.\n\n``` python\nimport os\nimport time\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\nMODEL_ID = \"Qwen/Qwen2.5-0.5B-Instruct\"\n\ntorch.set_num_threads(os.cpu_count() or 1)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)\nmodel.eval()\n\n# our toy data to classify (600 records)\nLABELS = [\"billing\", \"technical\", \"account\"]\ntickets = [\n    \"My card was charged twice for the same invoice.\",\n    \"The mobile app crashes whenever I open the settings page.\",\n    \"I need to change the email address on my profile.\",\n] * 200\n\n# constrained scoring, carried over from the previous article: the prompt ends at the\n# start of the assistant turn, so comparing the logits of each label's FIRST token\n# is enough to pick a winner, provided those first tokens are distinct\nlabel_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]\nassert len(set(label_first_ids)) == len(LABELS), (\n    \"Labels share a first token; score full label sequences instead.\"\n)\nlabel_first_ids = torch.tensor(label_first_ids, device=model.device)\n\nSYSTEM_PROMPT = \"\"\"You classify customer support tickets into exactly one category.\n\nCategories:\n- billing: payments, invoices, refunds, charges, subscription costs\n- technical: crashes, errors, performance problems, broken features\n- account: profile changes, login access, permissions, account deletion\n\nExamples:\nTicket: I was billed twice in March.\nCategory: billing\nTicket: The dashboard never finishes loading.\nCategory: technical\nTicket: Please remove my old phone number from my profile.\nCategory: account\nTicket: My promo code was rejected at checkout.\nCategory: billing\nTicket: Exporting to CSV throws a 500 error.\nCategory: technical\nTicket: I cannot reset my password.\nCategory: account\n\"\"\"\n\n# the chat template is written out by hand so the prompt can be split at a known\n# token boundary; this is the exact ChatML layout Qwen2.5-Instruct expects\nprefix_text = f\"<|im_start|>system\\n{SYSTEM_PROMPT}<|im_end|>\\n\"\n\ndef suffix_text(ticket):\n    return (\n        f\"<|im_start|>user\\nTicket: {ticket}\\nCategory:<|im_end|>\\n\"\n        f\"<|im_start|>assistant\\n\"\n    )\n\ndef encode(text):\n    return tokenizer(text, add_special_tokens=False)[\"input_ids\"]\n\n# the split has to be token-clean: encoding the two halves separately must give\n# exactly the same ids as encoding the whole prompt in one go, or the cached keys\n# in the next script will not line up with what the model would otherwise have seen\n_probe = suffix_text(tickets[0])\nassert encode(prefix_text) + encode(_probe) == encode(prefix_text + _probe), (\n    \"Prefix/suffix split is not token-clean; move the boundary.\"\n)\n\nprefix_len = len(encode(prefix_text))\nfull_len = prefix_len + len(encode(_probe))\nprint(f\"Static prefix length: {prefix_len} tokens\")\nprint(f\"Full prompt length:   {full_len} tokens ({100 * prefix_len / full_len:.0f}% of it static)\")\n\n# time inference\nbaseline_predictions = []\nstart = time.time()\n\nfor n, ticket in enumerate(tickets, start=1):\n\n    # this loop runs for minutes on CPU, so report progress rather than sitting silent\n    if n % 100 == 0:\n        print(f\"  {n}/{len(tickets)} tickets ({(time.time() - start) / n:.2f}s each)\", flush=True)\n\n    full = tokenizer(\n        prefix_text + suffix_text(ticket), add_special_tokens=False, return_tensors=\"pt\"\n    ).to(model.device)\n    with torch.no_grad():\n        logits = model(**full).logits[0, -1, :]\n    baseline_predictions.append(LABELS[int(logits[label_first_ids].argmax())])\n\nduration_full = time.time() - start\n\n# output task metrics\nprint(f\"Recomputing the full prompt every time: {duration_full:.2f} seconds\")\nprint(f\"  ({1000 * duration_full / len(tickets):.1f} ms per ticket)\")\n```\n\nOutput:\n\n```\nLoading weights: 100%|████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 156.40it/s]\nStatic prefix length: 145 tokens\nFull prompt length:   167 tokens (87% of it static)\n  100/600 tickets (0.31s each)\n  200/600 tickets (0.30s each)\n  300/600 tickets (0.30s each)\n  400/600 tickets (0.31s each)\n  500/600 tickets (0.31s each)\n  600/600 tickets (0.31s each)\nRecomputing the full prompt every time: 184.85 seconds\n  (308.1 ms per ticket)\n```\n\n## Reusing the Prompt Prefix\n\nWe 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.\n\n``` python\nimport os\nimport time\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache\n\nMODEL_ID = \"Qwen/Qwen2.5-0.5B-Instruct\"\n\ntorch.set_num_threads(os.cpu_count() or 1)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)\nmodel.eval()\n\n# our toy data to classify (600 records)\nLABELS = [\"billing\", \"technical\", \"account\"]\ntickets = [\n    \"My card was charged twice for the same invoice.\",\n    \"The mobile app crashes whenever I open the settings page.\",\n    \"I need to change the email address on my profile.\",\n] * 200\n\n# constrained scoring, carried over from the previous article: the prompt ends at the\n# start of the assistant turn, so comparing the logits of each label's FIRST token\n# is enough to pick a winner, provided those first tokens are distinct\nlabel_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]\nassert len(set(label_first_ids)) == len(LABELS), (\n    \"Labels share a first token; score full label sequences instead.\"\n)\nlabel_first_ids = torch.tensor(label_first_ids, device=model.device)\n\nSYSTEM_PROMPT = \"\"\"You classify customer support tickets into exactly one category.\n\nCategories:\n- billing: payments, invoices, refunds, charges, subscription costs\n- technical: crashes, errors, performance problems, broken features\n- account: profile changes, login access, permissions, account deletion\n\nExamples:\nTicket: I was billed twice in March.\nCategory: billing\nTicket: The dashboard never finishes loading.\nCategory: technical\nTicket: Please remove my old phone number from my profile.\nCategory: account\nTicket: My promo code was rejected at checkout.\nCategory: billing\nTicket: Exporting to CSV throws a 500 error.\nCategory: technical\nTicket: I cannot reset my password.\nCategory: account\n\"\"\"\n\n# the chat template is written out by hand so the prompt can be split at a known\n# token boundary; this is the exact ChatML layout Qwen2.5-Instruct expects\nprefix_text = f\"<|im_start|>system\\n{SYSTEM_PROMPT}<|im_end|>\\n\"\n\ndef suffix_text(ticket):\n    return (\n        f\"<|im_start|>user\\nTicket: {ticket}\\nCategory:<|im_end|>\\n\"\n        f\"<|im_start|>assistant\\n\"\n    )\n\ndef encode(text):\n    return tokenizer(text, add_special_tokens=False)[\"input_ids\"]\n\n# the split has to be token-clean: encoding the two halves separately must give\n# exactly the same ids as encoding the whole prompt in one go, or the cached keys\n# will not line up with what the model would otherwise have seen\n_probe = suffix_text(tickets[0])\nassert encode(prefix_text) + encode(_probe) == encode(prefix_text + _probe), (\n    \"Prefix/suffix split is not token-clean; move the boundary.\"\n)\n\nprefix = tokenizer(prefix_text, add_special_tokens=False, return_tensors=\"pt\").to(model.device)\nprefix_ids = prefix[\"input_ids\"]\nprefix_len = prefix_ids.shape[1]\nfull_len = prefix_len + len(encode(_probe))\nprint(f\"Static prefix length: {prefix_len} tokens\")\nprint(f\"Full prompt length:   {full_len} tokens ({100 * prefix_len / full_len:.0f}% of it static)\")\n\n# Populate the cache once with the static instruction block\nprefix_cache = DynamicCache()\nwith torch.no_grad():\n    model(\n        input_ids=prefix_ids,\n        attention_mask=torch.ones_like(prefix_ids),\n        past_key_values=prefix_cache,\n        use_cache=True,\n    )\n\ndef classify_cached(ticket):\n    suffix = tokenizer(\n        suffix_text(ticket), add_special_tokens=False, return_tensors=\"pt\"\n    ).to(model.device)\n    suffix_ids = suffix[\"input_ids\"]\n    suffix_len = suffix_ids.shape[1]\n\n    # the mask must cover the cached prefix as well as the new tokens, and the new\n    # tokens must be told they start at position prefix_len, not at position 0\n    attention_mask = torch.ones((1, prefix_len + suffix_len), device=model.device, dtype=torch.long)\n    cache_position = torch.arange(prefix_len, prefix_len + suffix_len, device=model.device)\n\n    with torch.no_grad():\n        out = model(\n            input_ids=suffix_ids,\n            attention_mask=attention_mask,\n            past_key_values=prefix_cache,\n            cache_position=cache_position,\n            use_cache=True,\n        )\n        logits = out.logits[0, -1, :]\n        label = LABELS[int(logits[label_first_ids].argmax())]\n\n        # roll the cache back so the next ticket starts from the prefix alone\n        prefix_cache.crop(prefix_len)\n\n    return label\n\ndef classify_full(ticket):\n    \"\"\"Reference path: re-encode the whole prompt, no cache. Used only to verify.\"\"\"\n    full = tokenizer(\n        prefix_text + suffix_text(ticket), add_special_tokens=False, return_tensors=\"pt\"\n    ).to(model.device)\n    with torch.no_grad():\n        logits = model(**full).logits[0, -1, :]\n    return LABELS[int(logits[label_first_ids].argmax())]\n\n# correctness first: the cached path must agree with the uncached one on every\n# distinct ticket, otherwise the speedup below is measuring the wrong computation\nmismatches = [t for t in dict.fromkeys(tickets) if classify_cached(t) != classify_full(t)]\nassert not mismatches, f\"Cached path disagrees with full re-encoding on: {mismatches}\"\nprint(f\"Cache verified against full re-encoding on {len(set(tickets))} distinct tickets\")\n\n# time inference\npredictions = []\nstart = time.time()\n\nfor n, ticket in enumerate(tickets, start=1):\n\n    # this loop is fast, but report progress anyway so the two scripts look alike\n    if n % 100 == 0:\n        print(f\"  {n}/{len(tickets)} tickets ({(time.time() - start) / n:.2f}s each)\", flush=True)\n\n    predictions.append(classify_cached(ticket))\n\nduration_cached = time.time() - start\n\n# output task metrics\nprint(f\"Reusing the cached prefix: {duration_cached:.2f} seconds\")\nprint(f\"  ({1000 * duration_cached / len(tickets):.1f} ms per ticket)\")\n```\n\nOutput:\n\n```\nStatic prefix length: 145 tokens\nFull prompt length:   167 tokens (87% of it static)\nCache verified against full re-encoding on 3 distinct tickets\n  100/600 tickets (0.13s each)\n  200/600 tickets (0.13s each)\n  300/600 tickets (0.13s each)\n  400/600 tickets (0.13s each)\n  500/600 tickets (0.13s each)\n  600/600 tickets (0.13s each)\nReusing the cached prefix: 80.07 seconds\n  (133.5 ms per ticket)\n```\n\nOur baseline script version ran for 184.85 seconds total, with an average of ~0.3s for each ticket. The cached prompt prefix version of the script is 80.07 seconds total, with an average of 0.13s per ticket. This is an overall runtime reduction of ~57%. The gain scales with the ratio between your static and dynamic content, which is why this technique rewards long, detailed instruction blocks rather than punishing them. And just to note: the predictions should be identical on every ticket, which is the result you want: this is a pure compute optimization, not a change to the model's behavior.\n\nSome further explanation of the code above:\n\n- `DynamicCache` 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 only processes the new tokens while still attending backward across the full context.\n- Two arguments have to agree with the cache or the results will be deceivingly wrong. The attention mask covers the cached prefix *and* the new tokens, so its width is`prefix_len + suffix_len` even though only`suffix_len` ids are passed in.`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 infer the positions from the cache length, but passing them explicitly documents the intent and guards against version drift.\n- The call to `prefix_cache.crop(prefix_len)` is not optional. The forward pass appends the suffix keys and values to the cache, so without the crop the second ticket would attend to the first ticket's tokens and the cache would grow without bound.\n- Split the prompt at a clean boundary — 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 comes mid-word, and then the cached keys no longer correspond to what the model would have seen. The assertion in the first script checks this directly.\n- Keep the whole pipeline under `torch.no_grad()` rather than`torch.inference_mode()` . Tensors created inside inference mode carry a flag that makes them awkward to slice and reassign afterwards, which is exactly what`crop()` does to the cache.\n\n## Wrapping Up\n\nThis was the second entry in our attempts at optimizing SLMs for narrow automation, and our target technique this time was **prefix caching**. This technique replaces the full and complete re-encoding of a static instruction block with a single pre-fill whose key and value tensors are computed once and reused, leaving us to only account for the tokens that differ. By implementing it, we get the same predictions the naive loop produced at a far lower compute cost. The longer and more detailed our instructions grow, the better the trade becomes.\n\nA small language model, such as our use of a 0.5B parameter model today, becomes a practical production choice for narrow automation once the code around it 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.\n\n \n\n \n\n[**\\[Matthew Mayo\\](https://www.kdnuggets.com/wp-content/uploads/./profile-pic.jpg)**](https://www.linkedin.com/in/mattmayo13/) ([**@mattmayo13**](https://twitter.com/mattmayo13)) holds a master's degree in computer science and a graduate diploma in data mining. As managing editor of [KDnuggets](https://www.kdnuggets.com/) & [Statology](https://www.statology.org/), and contributing editor at [Machine Learning Mastery](https://machinelearningmastery.com/), Matthew aims to make complex data science concepts accessible. His professional interests include natural language processing, language models, machine learning algorithms, and exploring emerging AI. He is driven by a mission to democratize knowledge in the data science community. Matthew has been coding since he was 6 years old.", "url": "https://wpnews.pro/news/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization", "canonical_source": "https://www.kdnuggets.com/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization", "published_at": "2026-09-18 14:00:34+00:00", "updated_at": "2026-09-18 16:56:01.268979+00:00", "lang": "en", "topics": ["large-language-models", "ai-tools", "natural-language-processing", "ai-research"], "entities": ["Qwen2.5-0.5B-Instruct", "Hugging Face Transformers", "KDnuggets", "M2 MacBook Air"], "alternates": {"html": "https://wpnews.pro/news/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization", "markdown": "https://wpnews.pro/news/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization.md", "text": "https://wpnews.pro/news/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization.txt", "jsonld": "https://wpnews.pro/news/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization.jsonld"}}