在我们关于SLM优化技术的短系列文章的第二篇中，我们重点关注使用键值缓存复用提示前缀。
在上一篇文章中，我们讨论了为小型语言模型（SLM）窄域自动化优化约束输出空间。我们当时提到，这是SLM优化策略短系列文章的第一篇。现在我们来到了第二篇。这次我们将重点关注使用键值缓存复用提示前缀。让我们不再浪费时间寒暄，直接进入正题。
与上一篇文章一样，以下所有基准测试均通过Hugging Face Transformers使用Qwen2.5-0.5B-Instruct的float16精度，在配备24GB RAM和16核神经引擎的M2 Macbook Air上运行。
首先，确保你设置好Python环境并安装所需依赖：
pip install torch transformers accelerate
我们将继续使用第一篇文章中的支持工单框架。
为什么要使用键值缓存复用提示前缀
窄域自动化提示相对静态。任务指令、分类定义和几个示例构成了大部分token，只有一小段提示尾部会因条目而变化。如果你的指令块达到几百个token，而每个工单增加二三十个，那么每个提示的绝大多数内容都与上一个逐字节完全相同。在每次调用时、对每一层重新计算所有这些内容显然是浪费的。
Transformer（深度学习中的注意力模型架构）在每一层为每个token计算一个键向量和一个值向量，而这些仅依赖于左侧的token。这意味着对于固定的前缀，它们在每次调用时都是相同的。计算一次并保留它们，可以将每个条目的预填充缩减到仅实际发生变化的token。
重新编码每个工单
首先是基线：一个真实的少样本提示，为每个工单完整重新编码。我们之前文章中的约束评分技巧在此延续使用，因此决策只需一次前向传播，唯一需要优化的就是预填充。注意，聊天模板是手写出来的，而不是通过apply_chat_template()，因为下一个脚本需要在已知边界处拆分提示。
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"system\n{SYSTEM_PROMPT}\n"
def suffix_text(ticket):
return (
f"user\nTicket: {ticket}\nCategory:\n"
f"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)")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)
复用提示前缀
现在我们让前缀只通过模型运行一次，保留生成的缓存，然后每个条目只输入其自身的 token。这继续在同一个脚本中进行，因此模型、分词器、提示的两部分以及基线计时都仍在作用域内。
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.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"system\n{SYSTEM_PROMPT}\n"
def suffix_text(ticket):
return (
f"user\nTicket: {ticket}\nCategory:\n"
f"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
# 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 = tokenizer(prefix_text, add_special_tokens=False, return_tensors="pt").to(model.device)
prefix_ids = prefix["input_ids"]
prefix_len = prefix_ids.shape[1]
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)")
# Populate the cache once with the static instruction block
prefix_cache = DynamicCache()
with torch.no_grad():
model(
input_ids=prefix_ids,
attention_mask=torch.ones_like(prefix_ids),
past_key_values=prefix_cache,
use_cache=True,
)
def classify_cached(ticket):
suffix = tokenizer(
suffix_text(ticket), add_special_tokens=False, return_tensors="pt"
).to(model.device)
suffix_ids = suffix["input_ids"]
suffix_len = suffix_ids.shape[1]
# the mask must cover the cached prefix as well as the new tokens, and the new
# tokens must be told they start at position prefix_len, not at position 0
attention_mask = torch.ones((1, prefix_len + suffix_len), device=model.device, dtype=torch.long)
cache_position = torch.arange(prefix_len, prefix_len + suffix_len, device=model.device)
with torch.no_grad():
out = model(
input_ids=suffix_ids,
attention_mask=attention_mask,
past_key_values=prefix_cache,
cache_position=cache_position,
use_cache=True,
)
logits = out.logits[0, -1, :]
label = LABELS[int(logits[label_first_ids].argmax())]
# roll the cache back so the next ticket starts from the prefix alone
prefix_cache.crop(prefix_len)
return label
def classify_full(ticket):
"""Reference path: re-encode the whole prompt, no cache. Used only to verify."""
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, :]
return LABELS[int(logits[label_first_ids].argmax())]
# correctness first: the cached path must agree with the uncached one on every
# distinct ticket, otherwise the speedup below is measuring the wrong computation
mismatches = [t for t in dict.fromkeys(tickets) if classify_cached(t) != classify_full(t)]
assert not mismatches, f"Cached path disagrees with full re-encoding on: {mismatches}"
print(f"Cache verified against full re-encoding on {len(set(tickets))} distinct tickets")
# time inference
predictions = []
start = time.time()
for n, ticket in enumerate(tickets, start=1):
# this loop is fast, but report progress anyway so the two scripts look alike
if n % 100 == 0:
print(f" {n}/{len(tickets)} tickets ({(time.time() - start) / n:.2f}s each)", flush=True)
predictions.append(classify_cached(ticket))
duration_cached = time.time() - start
# output task metrics
print(f"Reusing the cached prefix: {duration_cached:.2f} seconds")
print(f" ({1000 * duration_cached / len(tickets):.1f} ms per ticket)")Static prefix length: 145 tokens Full prompt length: 167 tokens (87% of it static) Cache verified against full re-encoding on 3 distinct tickets 100/600 tickets (0.13s each) 200/600 tickets (0.13s each) 300/600 tickets (0.13s each) 400/600 tickets (0.13s each) 500/600 tickets (0.13s each) 600/600 tickets (0.13s each) Reusing the cached prefix: 80.07 seconds (133.5 ms per ticket)
我们的基线脚本版本总共运行了 184.85 秒，每张工单平均约 0.3 秒。缓存提示前缀版本的脚本总共运行 80.07 秒，每张工单平均 0.13 秒。整体运行时间减少了约 57%。收益随静态内容与动态内容的比例而扩展，这就是为什么这项技术奖励长而详细的指令块，而不是惩罚它们。另外请注意：每张工单的预测结果应该完全相同，这正是你想要的结果：这是纯粹的计算优化，而不是对模型行为的改变。
对上述代码的进一步解释：
DynamicCache 保存前缀的逐层键和值张量。将其作为 past_key_values 传入会告诉模型这些位置已经计算过，因此前向传播只处理新 token，同时仍然在完整上下文上向后进行注意力计算。
有两个参数必须与缓存一致，否则结果会出现具有欺骗性的错误。注意力掩码覆盖缓存前缀和新 token，因此即使只传入 suffix_len 个 id，其宽度也是 prefix_len + suffix_len。cache_position 告诉模型新 token 从偏移量 prefix_len 开始，因此旋转嵌入与完整提示会产生的结果相匹配。较新版本的 Transformers 会根据缓存长度推断位置，但显式传入它们可以记录意图并防止版本漂移。
对 prefix_cache.crop(prefix_len) 的调用不是可选的。前向传播会将后缀的键和值追加到缓存中，因此如果不进行裁剪，第二张工单就会关注第一张工单的 token，并且缓存会无限制增长。
在干净的边界处拆分提示——行尾或聊天模板分隔符。如果拆分发生在单词中间，分别对两部分进行分词可能会产生与对拼接结果进行分词不同的 token 序列，然后缓存的键就不再对应于模型原本会看到的内容。第一个脚本中的断言直接检查了这一点。
将整个流水线保持在 torch.no_grad() 下，而不是 torch.inference_mode() 下。在推理模式下创建的张量带有一个标志，使得之后对它们进行切片和重新赋值变得很别扭，而这正是 crop() 对缓存所做的事情。
总结
这是我们尝试针对窄域自动化优化 SLM 的第二篇文章，这次的目标技术是前缀缓存。这项技术用一次预填充替代对静态指令块的完整重新编码，其中的键和值张量只计算一次并复用，让我们只需考虑不同的 token。通过实现它，我们以远低得多的计算成本获得了与朴素循环相同的预测结果。我们的指令越长、越详细，这种权衡就越好。
一个小语言模型，例如我们今天使用的 0.5B 参数模型，一旦其周围的代码不再把每次调用都当作孤立事件，它就会成为窄域自动化的实用生产选择。提示每次几乎都是相同的。一旦你的循环知道这一点，小模型就不再是一种妥协，而开始成为显而易见的答案。
Matthew Mayo（@mattmayo13）拥有计算机科学硕士学位和数据挖掘研究生文凭。作为 KDnuggets 与 Statology 的管理编辑，以及 Machine Learning Mastery 的特约编辑，Matthew 致力于让复杂的数据科学概念变得易于理解。他的专业兴趣包括自然语言处理、语言模型、机器学习算法以及探索新兴 AI。他致力于推动数据科学社区的知识民主化。Matthew 从 6 岁起就开始编程。
转载声明：本文为转载整理，原文版权归原作者及原站点所有。
原文来源：KDnuggets
原创地址：https://www.kdnuggets.com/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization
原文作者：https://www.facebook.com/kdnuggets；如有侵权，请联系管理员删除。
