利用键值缓存复用提示词前缀优化 SLM

📌 One-Sentence Summary 本文演示了如何通过 Hugging Face 的 `DynamicCache` 复用静态提示词前缀，将 0.5B SLM 分类任务的单工单推理时间缩短约 57%，并详细说明了保持结果一致所需的缓存、掩码和位置细节。 📝 Summary 本系列文章的第二篇聚焦于面向窄域自动化的小语言模型（SLM）优化，主题是前缀缓存。作者在 M2 MacBook Air 上使用 float16 精度的 Qwen2.5-0.5B-Instruct，对一项支持工单分类任务进行了基准测试，其中每个提示词有 87%（167 个 token 中的 145 个）是静态指令块。基线方法对每张工单都重新编码完整提示词，在 600 条记录上平均每张工单耗时 308 ms。优化版本将静态前缀一次性预填充到 `DynamicCache` 中，之后每张工单只输入变化的后缀 token，同时传入覆盖前缀加后缀的注意力掩码，以及 `cache_position` 偏移量，使旋转嵌入与完整提示词计算保持一致。每次前向传播后调用 `prefix_cache.crop(prefix_len)` 可防止缓存增长，并阻止工单之间相互关注。经验证，缓存路径产生的预测与完整重新编码完全相同，并将运行时间降至每张工单 133.5 ms，减少约 57%。文章指出，收益随静态与动态 token 比例的增加而扩大，并详述了常见陷阱：token 干净的提示词切分、掩码宽度、位置偏移，以及使用 `torch.no_grad()` 而非 `torch.inference_mode()`。 💡 Main Points 通过键值缓存复用静态提示词前缀，可将单工单延迟降低约 57%，且不改变模型预测结果。 在基准测试中，完整重新编码平均每张工单耗时 308.1 ms（600 张工单总计 184.85 s），而缓存前缀路径平均每张工单耗时 133.5 ms（总计 80.07 s）。经验证，缓存路径在每个不同的工单上都与完整重新编码结果一致，确认这是一种纯计算优化。 加速效果随静态与动态 token 比例的增加而扩大，因此冗长、详细的指令块会带来回报而非负担。 在示例中，167 个 token 中有 145 个（87%）是静态的，这正是收益显著的原因。提示词中在多次调用间保持固定的部分越多，可跳过的预填充工作就越多。 正确的缓存复用要求注意力掩码覆盖已缓存前缀加上新 token，并且 `cache_position` 将新 token 偏移 `prefix_len`。 尽管只传入了 `suffix_len` 个 id，掩码宽度仍为 `prefix_len + suffix_len`。`cache_position` 确保旋转嵌入与完整提示词所产生的结果一致；较新的 Transformers 版本可以推断这一点，但显式传入可防止版本漂移。 每次前向传播后必须将缓存裁剪回前缀长度，以避免无限增长和跨工单注意力泄漏。 每次前向传播都会将后缀的键和值追加到缓存中。如果不调用 `prefix_cache.crop(prefix_len)`，第二张工单就会关注第一张工单的 token，缓存也会无限增长。 提示词必须在 token 干净的边界处切分，并且流水线应在 `torch.no_grad()` 而非 `torch.inference_mode()` 下运行。 如果切分点落在单词中间，分别对两半进行分词可能产生与拼接后分词不同的 token 序列，从而破坏缓存对齐。在推理模式下创建的张量带有标志，会使 `crop()` 执行的切片和重新赋值操作变得棘手。 💬 Key Quotes 窄域自动化提示词相对静态。任务指令、分类体系定义和少量示例构成了 token 的大部分，只有很短的提示词尾部会随条目变化。 收益随静态内容与动态内容之间的比例而扩大，这就是为什么这项技术奖励冗长、详细的指令块，而不是惩罚它们。 调用 `prefix_cache.crop(prefix_len)` 不是可选项。前向传播会将后缀的键和值追加到缓存中，因此如果不裁剪，第二张工单就会关注第一张工单的 token，缓存也会无限增长。 一旦围绕它的代码不再将每次调用视为孤立事件，小语言模型——比如我们今天使用的 0.5B 参数模型——就会成为窄域自动化的实用生产选择。 📊 Article Meta AI Screening： 88 Source： KDnuggets Author： Matthew Mayo Category： 人工智能 Language： 英文 Read Time： 10 min Word Count： 2301 Tags： AI 与智能应用 , 性能优化 , LLM 推理优化 , 提示工程 , AI 工作流 Read Full Article