AI 推理性能调优：Prefix Caching 前缀缓存的推理加速实践-平芜编程栈

AI 推理性能调优：Prefix Caching 前缀缓存的推理加速实践

一、重复前缀的浪费：当系统提示词吃掉一半算力

大语言模型的推理分为两个阶段：Prefill（预填充）和 Decode（解码）。Prefill 阶段处理输入 Prompt 的所有 Token，计算 KV Cache；Decode 阶段逐个生成输出 Token。对于长 Prompt（如包含系统指令、知识库文档的 RAG 请求），Prefill 阶段的计算量远大于 Decode 阶段。

在 RAG 场景中，系统提示词和知识库文档通常占 Prompt 的 80% 以上，且在不同请求间高度重复。如果每次请求都重新计算这些重复前缀的 KV Cache，大量的 GPU 算力被浪费在重复计算上。Prefix Caching（前缀缓存）通过缓存共享前缀的 KV Cache，使得后续请求只需计算新增部分的 KV Cache，Prefill 阶段的计算量减少 60-80%。

flowchart TB subgraph 无前缀缓存 R1[请求1: 系统提示+文档A+问题1] --> P1[Prefill: 计算全部KV Cache<br/>耗时: 800ms] R2[请求2: 系统提示+文档A+问题2] --> P2[Prefill: 重新计算全部KV Cache<br/>耗时: 800ms] Note1[重复计算系统提示+文档A<br/>浪费60%算力] -.-> P2 end subgraph 有前缀缓存 R3[请求1: 系统提示+文档A+问题1] --> P3[Prefill: 计算全部KV Cache<br/>耗时: 800ms<br/>同时缓存前缀KV] R4[请求2: 系统提示+文档A+问题2] --> P4[Prefill: 仅计算问题2的KV<br/>耗时: 200ms<br/>复用缓存的前缀KV] Note2[前缀缓存命中<br/>Prefill加速4x] -.-> P4 end

二、前缀缓存的核心机制

2.1 KV Cache 的结构

KV Cache 是 Transformer 推理的核心数据结构，存储了每一层、每个注意力头的 Key 和 Value 矩阵。对于 L 层、H 个注意力头的模型，每个 Token 的 KV Cache 大小为2 × L × H × d_head × sizeof(dtype)字节。一个 7B 模型（32 层、32 头、d_head=128、FP16），每个 Token 的 KV Cache 约为 512KB。

2.2 前缀匹配与缓存复用

前缀缓存的关键是识别请求间的共享前缀。在 RAG 场景中，系统提示词 + 检索到的文档构成共享前缀，用户问题是变化部分。缓存系统以 Token 序列的哈希值作为缓存键，当新请求的前缀与缓存中的某个条目匹配时，直接复用其 KV Cache，跳过 Prefill 计算。

sequenceDiagram participant Client as 客户端 participant Server as 推理服务 participant Cache as 前缀缓存 participant GPU as GPU Client->>Server: 请求1: [系统提示+文档A+问题1] Server->>Cache: 查询前缀缓存 Cache-->>Server: 未命中 Server->>GPU: 完整 Prefill（800ms） GPU->>Cache: 缓存 [系统提示+文档A] 的 KV Cache Server->>Client: 返回回答1 Client->>Server: 请求2: [系统提示+文档A+问题2] Server->>Cache: 查询前缀缓存 Cache-->>Server: 命中！返回 [系统提示+文档A] 的 KV Cache Server->>GPU: 仅 Prefill [问题2]（200ms） Note over GPU: 复用缓存的 KV Cache<br/>跳过重复前缀的计算 Server->>Client: 返回回答2

三、生产级代码实现

3.1 前缀缓存管理器

import hashlib import time import threading from typing import Dict, List, Optional, Tuple from dataclasses import dataclass, field from collections import OrderedDict import logging logger = logging.getLogger(__name__) @dataclass class CacheEntry: """前缀缓存条目""" prefix_hash: str # 前缀 Token 序列的哈希值 token_ids: List[int] # 前缀 Token ID 序列 kv_cache: any # KV Cache 数据（实际为 GPU Tensor） token_count: int # 前缀 Token 数量 size_bytes: int # KV Cache 占用字节数 created_at: float # 创建时间 last_accessed: float # 最后访问时间 hit_count: int = 0 # 命中次数 class PrefixCacheManager: """前缀缓存管理器 设计考量： - LRU 淘汰策略：显存不足时淘汰最久未使用的缓存 - 哈希匹配：使用 Token 序列的滚动哈希，支持增量计算 - 显存预算：限制缓存总大小，防止 OOM - 统计指标：命中率、平均前缀长度、显存利用率 """ def __init__( self, max_cache_size_bytes: int = 4 * 1024 * 1024 * 1024, # 默认 4GB kv_cache_per_token_bytes: int = 512 * 1024, # 每个 Token 约 512KB ): self.max_cache_size = max_cache_size_bytes self.kv_per_token = kv_cache_per_token_bytes self._cache: OrderedDict[str, CacheEntry] = OrderedDict() self._current_size = 0 self._lock = threading.Lock() # 统计 self._total_requests = 0 self._cache_hits = 0 def lookup( self, token_ids: List[int], ) -> Optional[Tuple[CacheEntry, int]]: """查找前缀缓存 Returns: (cache_entry, matched_length): 缓存条目和匹配的 Token 数量 None: 无缓存命中 """ self._total_requests += 1 # 逐步缩短前缀，寻找最长匹配 for length in range(len(token_ids), 0, -1): prefix = token_ids[:length] prefix_hash = self._compute_hash(prefix) with self._lock: if prefix_hash in self._cache: entry = self._cache[prefix_hash] # 更新访问时间和 LRU 顺序 entry.last_accessed = time.time() entry.hit_count += 1 self._cache.move_to_end(prefix_hash) self._cache_hits += 1 logger.info( f"前缀缓存命中: hash={prefix_hash[:12]}..., " f"匹配 {length}/{len(token_ids)} tokens" ) return entry, length return None def store( self, token_ids: List[int], kv_cache: any, prefix_length: Optional[int] = None, ) -> None: """存储前缀缓存 Args: token_ids: 完整的 Token 序列 kv_cache: 对应的 KV Cache 数据 prefix_length: 要缓存的前缀长度，默认缓存全部 """ if prefix_length is None: prefix_length = len(token_ids) prefix = token_ids[:prefix_length] prefix_hash = self._compute_hash(prefix) size_bytes = prefix_length * self.kv_per_token # 检查是否超出显存预算 with self._lock: while self._current_size + size_bytes > self.max_cache_size and self._cache: self._evict_oldest() entry = CacheEntry( prefix_hash=prefix_hash, token_ids=prefix, kv_cache=kv_cache, token_count=prefix_length, size_bytes=size_bytes, created_at=time.time(), last_accessed=time.time(), ) self._cache[prefix_hash] = entry self._current_size += size_bytes logger.info( f"缓存前缀: hash={prefix_hash[:12]}..., " f"tokens={prefix_length}, size={size_bytes / 1024 / 1024:.1f}MB" ) def _evict_oldest(self) -> None: """淘汰最久未使用的缓存条目""" if not self._cache: return # OrderedDict 的第一个元素就是最久未访问的 oldest_hash, oldest_entry = next(iter(self._cache.items())) self._current_size -= oldest_entry.size_bytes del self._cache[oldest_hash] logger.debug(f"淘汰缓存: hash={oldest_hash[:12]}..., 释放 {oldest_entry.size_bytes / 1024 / 1024:.1f}MB") def _compute_hash(self, token_ids: List[int]) -> str: """计算 Token 序列的哈希值""" data = ",".join(str(t) for t in token_ids) return hashlib.sha256(data.encode()).hexdigest() def get_stats(self) -> Dict: """获取缓存统计信息""" hit_rate = self._cache_hits / max(self._total_requests, 1) return { "total_requests": self._total_requests, "cache_hits": self._cache_hits, "hit_rate": round(hit_rate, 4), "cached_entries": len(self._cache), "used_size_mb": round(self._current_size / 1024 / 1024, 1), "max_size_mb": round(self.max_cache_size / 1024 / 1024, 1), "utilization": round(self._current_size / self.max_cache_size, 4), }

3.2 与推理引擎集成

class PrefixCacheInferenceEngine: """集成前缀缓存的推理引擎 设计考量： - 请求到达时先查询前缀缓存 - 缓存命中时，仅 Prefill 新增部分，复用缓存的 KV Cache - 缓存未命中时，完整 Prefill 并存储前缀 - 支持自定义前缀切分策略（如按系统提示/文档/问题切分） """ def __init__(self, model, tokenizer, cache_manager: PrefixCacheManager): self.model = model self.tokenizer = tokenizer self.cache = cache_manager async def generate( self, messages: List[Dict[str, str]], max_tokens: int = 512, ) -> Dict: """带前缀缓存的推理生成""" # 1. Tokenize full_text = self._format_messages(messages) token_ids = self.tokenizer.encode(full_text) # 2. 查询前缀缓存 cache_result = self.cache.lookup(token_ids) if cache_result: # 缓存命中：仅 Prefill 新增部分 cached_entry, matched_length = cache_result new_token_ids = token_ids[matched_length:] logger.info( f"前缀缓存命中: 匹配 {matched_length} tokens, " f"新增 {len(new_token_ids)} tokens" ) # 使用缓存的 KV Cache + Prefill 新增部分 output = await self._generate_with_cached_prefix( cached_entry.kv_cache, matched_length, new_token_ids, max_tokens, ) else: # 缓存未命中：完整 Prefill output = await self._generate_full(token_ids, max_tokens) # 存储前缀缓存（缓存系统提示 + 文档部分） prefix_length = self._compute_prefix_length(messages) if prefix_length < len(token_ids): self.cache.store(token_ids, output.kv_cache, prefix_length) return output def _compute_prefix_length(self, messages: List[Dict[str, str]]) -> int: """计算应缓存的前缀长度 设计考量： - 系统提示和文档内容作为前缀缓存 - 用户问题不缓存（每次不同） - 按消息边界切分，避免在 Token 中间截断 """ prefix_text = "" for msg in messages: if msg["role"] in ("system", "document"): prefix_text += self._format_message(msg) else: break # 遇到用户消息，停止 return len(self.tokenizer.encode(prefix_text))