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

两套独立但互补的优化:

  1. Prompt 压缩:减少每个请求携带的 Token 数量 → 降低 Prefill 和 Attention 计算量
  2. 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。

五、总结

Prompt 压缩解决的是"携带太多历史"的问题——通过去重、摘要和截断减少 Token 量 40%~70%。KV Cache 复用解决的是"重复计算相同前缀"的问题——通过 Prefix 匹配跳过已计算的 Prefill。

两者结合可将多轮对话的推理延迟降低 50%~80%。建议在编排层集成 PromptCompressor 模块,在 vLLM 侧开启 enable_prefix_caching=True。监控两个核心指标:压缩率(偏高说明历史过长)和 KV Cache 命中率(偏低说明路由一致性需要改善)。

Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐