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

來源原標題

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

作者

Matthew Mayo

原始連結

https://www.kdnuggets.com/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization

摘要原文

Reusing the Prompt Prefix with a Key-Value Cache for SLM Optimization In 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. By Matthew Mayo , KDnuggets Managing Editor on September 18, 2026 in Language Models 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 . Why Reuse the Prompt Prefix with a Key-Value Cache 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 First, 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. import os import time import torch from transformers import AutoTokenizer, AutoModelForCausalLM MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct" torch.set_num_threads(os.cpu_count() or 1) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32) model.eval() # our toy data to classify (600 records) LABELS = ["billing", "technical", "account"] tickets = [ "My card was charged twice for the same invoice.", "The mobile app crashes whenever I open the settings page.", "I need to change the email address on my profile.", ] * 200 # constrained scoring, carried over from the previous article: the prompt ends at the # start of the assistant turn, so comparing the logits of each label's FIRST token # is enough to pick a winner, provided those first tokens are distinct label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS] assert len(set(label_first_ids)) == len(LABELS), ( "Labels share a first token; score full label sequences instead." ) label_first_ids = torch.tensor(label_first_ids, device=model.device) SYSTEM_PROMPT = """You classify customer support tickets into exactly one category. Categories: - billing: payments, invoices, refunds, charges, subscription costs - technical: crashes, errors, performance problems, broken features - account: profile changes, login access, permissions, account deletion Examples: Ticket: I was billed twice in March. Category: billing Ticket: The dashboard never finishes loading. Category: technical Ticket: Please remove my old phone number from my profile. Category: account Ticket: My promo code was rejected at checkout. Category: billing Ticket: Exporting to CSV throws a 500 error. Category: technical Ticket: I cannot reset my password. Category: account """ # the chat template is written out by hand so the prompt can be split at a known # token boundary; this is the exact ChatML layout Qwen2.5-Instruct expects prefix_text = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n" def suffix_text(ticket): return ( f"<|im_start|>user\nTicket: {ticket}\nCategory:<|im_end|>\n" f"<|im_start|>assistant\n" ) def encode(text): return tokenizer(text, add_special_tokens=False)["input_ids"] # the split has to be token-clean: encoding the two halves separately must give # exactly the same ids as encoding the whole prompt in one go, or the cached keys # in the next script will not line up with what the model would otherwise have seen _probe = suffix_text(tickets[0]) assert encode(prefix_text) + encode(_probe) == encode(prefix_text + _probe), ( "Prefix/suffix split is not token-clean; move the boundary." ) prefix_len = len(encode(prefix_text)) full_len = prefix_len + len(encode(_probe)) print(f"Static prefix length: {prefix_len} tokens") print(f"Full prompt length: {full_len} tokens ({100 * prefix_len / full_len:.0f}% of it static)") # time inference baseline_predictions = [] start = time.time() for n, ticket in enumerate(tickets, start=1): # this loop runs for minutes on CPU, so report progress rather than sitting silent if n % 100 == 0: print(f" {n}/{len(tickets)} tickets ({(time.time() - start) / n:.2f}s each)", flush=True) full = tokenizer( prefix_text + suffix_text(ticket), add_special_tokens=False, return_tensors="pt" ).to(model.device) with torch.no_grad(): logits = model(**full).logits[0, -1, :] baseline_predictions.append(LABELS[int(logits[label_first_ids].argmax())]) duration_full = time.time() - start # output task metrics print(f"Recomputing the full prompt every time: {duration_full:.2f} seconds") print(f" ({1000 * duration_full / len(tickets):.1f} ms per ticket)") Output: Loading weights: 100%|████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 156.40it/s] Static prefix length: 145 tokens Full prompt length: 167 tokens (87% of it static) 100/600 tickets (0.31s each) 200/600 tickets (0.30s each) 300/600 tickets (0.30s each) 400/600 tickets (0.31s each) 500/600 tickets (0.31s each) 600/600 tickets (0.31s each) Recomputing the full prompt every time: 184.85 seconds (308.1 ms per ticket) 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. import os import time import torch from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct" torch.set_num_threads(os.cpu_count() or 1) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32) model.eval() # our toy data to classify (600 records) LABELS = ["billing", "technical", "account"] tickets = [ "My card was charged twice for the same invoice.", "The mobile app crashes whenever I open the settings page.", "I need to change the email address on my profile.", ] * 200 # constrained scoring, carried over from the previous article: the prompt ends at the # start of the assistant turn, so comparing the logits of each label's FIRST token # is enough to pick a winner, provided those first tokens are distinct label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS] assert len(set(label_first_ids)) == len(LABELS), ( "Labels share a first token; score full label sequences instead." ) label_first_ids = torch.tensor(label_first_ids, device=model.