当前位置：首页 > news >正文

Prompt 压缩与 KV Cache 复用：降低推理延迟的上下文窗口优化实践

news 2026/7/14 11:06:59

Prompt 压缩与 KV Cache 复用：降低推理延迟的上下文窗口优化实践

一、上下文即成本：为什么 4K Token 的上下文比 512 Token 慢 5 倍？

推理延迟与 Prompt Token 数呈超线性关系。Self-Attention 的计算复杂度为 O(N²)，其中 N 是序列总长度。当上下文从 512 Token 增长到 4096 Token 时，Attention 计算量增加到 64 倍——但实际延迟只增加到约 5 倍，因为 GPU 的并行计算能力部分掩盖了复杂度增长。

更隐蔽的成本是Prefill 延迟。每个新请求都需要对 Prompt 中的所有 Token 做一次完整的 Prefill（计算 KV Cache）。如果多轮对话中每次请求都携带完整历史（"历史 3000 Token + 当前问题 100 Token"），那 95% 的 Prefill 计算都在重复处理历史对话。

二、Prompt 压缩与 KV Cache 复用的协同策略

flowchart TD subgraph 请求到达 A[新用户消息] --> B{Prompt 压缩} end subgraph 压缩策略 B --> C[语义冗余检测<br/>移除重复/套话] B --> D[对话摘要<br/>长历史 → 短摘要] B --> E[重要性评分<br/>保留关键信息] end C & D & E --> F[压缩后的 Prompt<br/>Token 减少 40%~70%] subgraph KV Cache 管理 F --> G{KV Cache 命中?} G -->|命中| H[跳过 Prefill<br/>直接 Decode] G -->|未命中| I[Prefill 新 Token<br/>仅计算增量] end H --> J[推理延迟降低 50%~80%] I --> J

两套独立但互补的优化：

Prompt 压缩：减少每个请求携带的 Token 数量 → 降低 Prefill 和 Attention 计算量
KV Cache 复用：利用历史请求的 KV Cache → 跳过重复 Token 的 Prefill

三、Prompt 压缩的实现策略

# prompt_compressor.py — Prompt 压缩引擎 from typing import List import re class PromptCompressor: """ Prompt 压缩器：在不丢失关键信息的前提下减少 Token 数 """ def __init__(self, max_context_tokens: int = 4096): self.max_tokens = max_context_tokens def compress(self, messages: List[dict]) -> List[dict]: """ 对对话历史做多层压缩 策略优先级：去重 → 摘要 → 截断 """ original_est = self._estimate_tokens(messages) # 策略一：去除语义重复的消息 messages = self._deduplicate(messages) # 策略二：将长历史转换为摘要 if self._estimate_tokens(messages) > self.max_tokens * 0.8: messages = self._summarize_old_messages(messages) # 策略三：硬截断（兜底） messages = self._truncate(messages) compressed_est = self._estimate_tokens(messages) compression_ratio = 1 - compressed_est / original_est if original_est > 0 else 0 print(f"Prompt 压缩: {original_est} → {compressed_est} tokens " f"(压缩率 {compression_ratio:.0%})") return messages def _deduplicate(self, messages: List[dict]) -> List[dict]: """ 去除连续的语义重复消息 例如：用户连续发送两次"你好" → 只保留一次 """ if len(messages) < 2: return messages deduped = [messages[0]] for i in range(1, len(messages)): prev = messages[i-1].get("content", "") curr = messages[i].get("content", "") # 简化的重复检测：基于 Jaccard 相似度 if self._text_similarity(prev, curr) < 0.85: deduped.append(messages[i]) return deduped def _summarize_old_messages(self, messages: List[dict], keep_recent: int = 3) -> List[dict]: """ 将超过 keep_recent 轮的历史消息替换为摘要 摘要本身是一个 system 消息 """ if len(messages) <= keep_recent: return messages recent = messages[-keep_recent:] # 保留最近 N 条 old = messages[:-keep_recent] # 需要摘要的部分 # 实际生产中调用一个小模型生成摘要 # 这里用规则模拟：提取用户问题作为摘要 user_questions = [ m.get("content", "")[:80] + "..." for m in old if m.get("role") == "user" ] summary = "对话历史摘要：用户曾问了以下问题：" + "; ".join(user_questions[:5]) return [{"role": "system", "content": summary}] + recent def _truncate(self, messages: List[dict]) -> List[dict]: """ 硬截断：按 Token 预算裁剪 策略：保留 system prompt + 最近的消息 """ result = [] token_count = 0 # 从最新到最旧遍历（保留最近的） for msg in reversed(messages): msg_tokens = self._estimate_tokens([msg]) if token_count + msg_tokens > self.max_tokens: break result.insert(0, msg) token_count += msg_tokens return result def _estimate_tokens(self, messages: List[dict]) -> int: """ 粗略 Token 估算（比实际少 5%~10%） 中文: ~1.5 字符/Token 英文: ~4 字符/Token """ total_chars = sum( len(m.get("content", "")) for m in messages ) return total_chars // 2 # 综合中英文的粗略估算 def _text_similarity(self, a: str, b: str) -> float: """基于词集的 Jaccard 相似度""" set_a = set(a.split()) set_b = set(b.split()) if not set_a or not set_b: return 0.0 return len(set_a & set_b) / len(set_a | set_b)

四、KV Cache 复用的工程约束

Hash-based 复用 vs Prefix-based 复用：

Hash-based：对整个 Prompt 做 Hash，如果与历史请求相同则复用。实现简单但命中率极低（< 2%）
Prefix-based：如果当前 Prompt 的前 N 个 Token 与某历史 Prompt 相同，复用这 N 个 Token 的 KV Cache。命中率高（多轮对话中 system prompt 几乎不变），但需要额外的前缀匹配数据结构（Trie 或 Radix Tree）

显存开销：保留历史 KV Cache 需要占用显存。对于 4096 Token 的上下文，单层 KV Cache 约 32MB（Llama-2-7B）。推荐保留最近 5 个用户的 KV Cache 作为"热缓存"，超出的写入 SSD 或直接丢弃。

跨 GPU 的 KV Cache 迁移：如果用户的请求被路由到不同 GPU（非一致性哈希），KV Cache 复用会失效。解决方案是路由层使用"用户 ID 一致性哈希"，确保同一用户的请求始终落到同一 GPU。