RAG 服务的内存池设计复用 embedding 模型推理结果减少重复计算一、深度引言与场景痛点大家好我是赵咕咕。做 RAG 系统两年我发现一个尴尬的事实embedding 推理是 RAG 服务中最贵、最频繁、同时也是最浪费的操作。说它贵——一次 embedding 推理尤其在 GPU 上跑大模型如text-embedding-3-large的延迟在 20-50ms如果调用 OpenAI API还要加网络延迟和每 Token 的费用。说它频繁——一个典型的 RAG 问答流程至少要做 3 次 embedding用户问题 1 次、检索文档 N 条、可能还有重排序。说它浪费——高频问题如怎么重置密码的 embedding 结果每小时被重复计算了上百次。每次结果都是一样的——文本相同 → embedding 向量相同这是确定性计算。这就是embedding 缓存要解决的问题。把计算过的向量存进内存池相同文本下次直接命中缓存绕过推理。二、底层机制与原理深度剖析2.1 Embedding 缓存的命门哈希 vs 语义最简单的做法是用文本的 MD5 作为 key计算过的向量作为 value。这就行了不因为用户可能会问怎么重置密码如何重置密码密码重置方法这三个问题语义几乎相同但文本不同MD5 完全不同——哈希缓存命中不了。这就引出了 embedding 缓存的核心权衡精确匹配哈希 vs 语义匹配向量相似度。2.2 两级缓存的架构最优方案是两级缓存L1 做精确哈希匹配命中快L2 做语义相似度匹配命中率更高。当请求文本到达时首先计算其 MD5 或 SHA256 哈希值并检查 L1 缓存是否命中。如果命中则直接返回缓存向量如果未命中则进入 L2 语义匹配阶段。在 L2 阶段系统不会立即调用完整的 embedding 模型而是先计算请求文本的粗粒度 embedding例如使用fast-embed的 128 维版本。随后在 L2 缓存中进行语义匹配判断最相似缓存向量的 cosine 相似度是否超过阈值如 0.95。如果超过阈值则返回最相似的缓存向量如果未达到阈值则调用完整 embedding 模型进行计算。计算完成后新得到的向量会被同时写入 L1 和 L2 缓存并返回给请求方。L1 精确缓存用哈希表dict实现key 是文本的 SHA256value 是向量。它支持 O(1) 查找命中延迟几乎为 0。为了限制内存上限采用 LRULeast Recently Used淘汰策略。L2 语义缓存的关键设计是不对请求文本做完整 embedding那会失去缓存的意义而是用一个极轻量级的粗粒度模型如fast-embed的 128 维版本做个快速 embedding。然后在缓存中使用 FAISS 索引做最近邻搜索。如果最相似缓存的 cosine 相似度超过阈值如 0.95就返回缓存中的向量。如果不到阈值就调用完整 embedding 模型然后把新结果写入两级缓存。这里有个关键点L2 缓存里存的向量是完整 embedding 结果比如 1536 维不是粗粒度 embedding128 维。粗粒度 embedding 只是用来做相似度搜索的指针真正的 value 还是完整向量。2.3 为什么不用向量数据库做缓存你可能会想直接用 Qdrant/Redis 做 embedding 缓存不就行了可以但不推荐。向量数据库的查/写延迟通常在 2-10ms加上网络 RTT 就是 5-20ms。而内存中的 LRU dictL1延迟是纳秒级FAISS 内存索引L2延迟是 0.1-1ms。对缓存来说延迟就是一切。如果查缓存的时间和做一次推理差不多缓存就没意义了。三、生产级代码实现import asyncio import hashlib import logging import time from collections import OrderedDict from dataclasses import dataclass from typing import Any import numpy as np import faiss logger logging.getLogger(__name__) dataclass class CacheStats: 缓存统计信息。 l1_hits: int 0 l2_hits: int 0 misses: int 0 l1_size: int 0 l2_size: int 0 property def total(self) - int: return self.l1_hits self.l2_hits self.misses property def hit_rate(self) - float: total self.total if total 0: return 0.0 return (self.l1_hits self.l2_hits) / total class EmbeddingCache: 两级 embedding 缓存L1 精确 L2 语义。 def __init__( self, l1_capacity: int 10000, l2_capacity: int 50000, dim: int 1536, similarity_threshold: float 0.95, coarse_dim: int 128, ): # L1: LRU 精确缓存 self._l1_capacity l1_capacity self._l1_cache: OrderedDict[str, np.ndarray] OrderedDict() # L2: 语义缓存 self._l2_capacity l2_capacity self._l2_dim dim self._l2_index faiss.IndexFlatIP(dim) # Inner Product Cosine (归一化后) self._l2_texts: list[str] [] self._l2_vectors: list[np.ndarray] [] self._coarse_dim coarse_dim self._similarity_threshold similarity_threshold self._stats CacheStats() self._lock asyncio.Lock() async def get(self, text: str) - np.ndarray | None: 查找缓存。命中返回向量未命中返回 None。 async with self._lock: # L1: 精确匹配 key self._hash(text) if key in self._l1_cache: self._l1_cache.move_to_end(key) self._stats.l1_hits 1 return self._l1_cache[key].copy() # L2: 语义匹配 if len(self._l2_vectors) 0: coarse_vec await self._coarse_embed(text) coarse_vec coarse_vec.reshape(1, -1).astype(np.float32) # L2 的索引是完整向量先用粗粒度搜 # 直接搜索粗粒度不太合适这里用 L2 向量的前 coarse_dim 维近似 distances, indices self._l2_index.search( coarse_vec[:, :self._coarse_dim], k1 ) if len(indices[0]) 0 and distances[0][0] self._similarity_threshold: idx indices[0][0] vector self._l2_vectors[idx].copy() # 升入 L1 self._set_l1(key, vector) self._stats.l2_hits 1 return vector self._stats.misses 1 return None async def set(self, text: str, vector: np.ndarray) - None: 写入缓存向量。 async with self._lock: key self._hash(text) self._set_l1(key, vector.astype(np.float32)) self._set_l2(text, vector.astype(np.float32)) def _set_l1(self, key: str, vector: np.ndarray) - None: 写入 L1LRU 淘汰。 if key in self._l1_cache: self._l1_cache.move_to_end(key) self._l1_cache[key] vector if len(self._l1_cache) self._l1_capacity: self._l1_cache.popitem(lastFalse) # 淘汰最旧的 self._stats.l1_size len(self._l1_cache) def _set_l2(self, text: str, vector: np.ndarray) - None: 写入 L2容量满时淘汰最旧的。 if len(self._l2_vectors) self._l2_capacity: remove_idx 0 self._l2_vectors.pop(remove_idx) self._l2_texts.pop(remove_idx) self._rebuild_index() self._l2_vectors.append(vector) self._l2_texts.append(text) self._l2_index.add(vector.reshape(1, -1)) self._stats.l2_size len(self._l2_vectors) def _rebuild_index(self) - None: 重建 FAISS 索引容量满时重建。 if not self._l2_vectors: self._l2_index faiss.IndexFlatIP(self._l2_dim) return matrix np.stack(self._l2_vectors).astype(np.float32) # 归一化以使用 Inner Product 作为 Cosine Similarity faiss.normalize_L2(matrix) self._l2_index faiss.IndexFlatIP(self._l2_dim) self._l2_index.add(matrix) async def _coarse_embed(self, text: str) - np.ndarray: 生成粗粒度 embedding用于 L2 相似度搜索。 # 生产环境替换为 fast-embed 或 sentence-transformers 轻量模型 # 这里给出一个基于哈希的快速近似仅供演示 h hashlib.sha256(text.encode()).digest() vec np.frombuffer(h, dtypenp.uint8).astype(np.float32) / 255.0 # 填充或截断到 coarse_dim if len(vec) self._coarse_dim: vec np.pad(vec, (0, self._coarse_dim - len(vec))) else: vec vec[:self._coarse_dim] return vec staticmethod def _hash(text: str) - str: return hashlib.sha256(text.encode(utf-8)).hexdigest() def get_stats(self) - CacheStats: return self._stats class CachedEmbedder: 带缓存的 embedding 推理器。 def __init__( self, model, # 替换为你的 embedding 模型 cache: EmbeddingCache, warm_up_texts: list[str] | None None, ): self._model model self._cache cache async def embed(self, text: str) - np.ndarray: 嵌入一条文本带缓存。 # 检查缓存 cached await self._cache.get(text) if cached is not None: return cached # 缓存未命中推理 try: vector await asyncio.wait_for( self._model.aembed(text), timeout10.0 ) # 确保是 numpy 数组 if not isinstance(vector, np.ndarray): vector np.array(vector, dtypenp.float32) # 写入缓存 await self._cache.set(text, vector) return vector except asyncio.TimeoutError: logger.error(Embedding 推理超时: %.50s, text) raise async def embed_batch( self, texts: list[str], batch_size: int 32 ) - list[np.ndarray]: 批量嵌入先去缓存查未命中才推理。 results: list[np.ndarray | None] [None] * len(texts) uncached_indices: list[int] [] uncached_texts: list[str] [] # 先查缓存 for i, text in enumerate(texts): cached await self._cache.get(text) if cached is not None: results[i] cached else: uncached_indices.append(i) uncached_texts.append(text) if not uncached_texts: return [r for r in results if r is not None] # type: ignore # 批量推理未缓存的文本 new_vectors [] for start in range(0, len(uncached_texts), batch_size): batch uncached_texts[start:start batch_size] try: batch_vectors await asyncio.wait_for( self._model.aembed_batch(batch), timeout30.0 ) except Exception as e: logger.error(批量 Embedding 推理失败 (batch %d): %s, start, e) raise for j, vec in enumerate(batch_vectors): if not isinstance(vec, np.ndarray): vec np.array(vec, dtypenp.float32) orig_idx uncached_indices[start j] results[orig_idx] vec # 写入缓存 await self._cache.set(batch[j], vec) return [r for r in results if r is not None] # type: ignore property def stats(self) - CacheStats: return self._cache.get_stats() async def main(): # 模拟一个 embedding 模型 class MockEmbedder: async def aembed(self, text: str) - np.ndarray: await asyncio.sleep(0.05) # 模拟 50ms 推理 return np.random.randn(1536).astype(np.float32) async def aembed_batch(self, texts: list[str]) - list[np.ndarray]: await asyncio.sleep(0.1) return [np.random.randn(1536).astype(np.float32) for _ in texts] cache EmbeddingCache(l1_capacity5000, l2_capacity20000, dim1536) embedder CachedEmbedder(MockEmbedder(), cache) # 第一次调用cache miss t0 time.monotonic() v1 await embedder.embed(怎么重置密码) t1 time.monotonic() print(f第一次 (miss): {1000*(t1-t0):.1f}ms) # 第二次调用相同文本L1 hit v2 await embedder.embed(怎么重置密码) t2 time.monotonic() print(f第二次 (L1 hit): {1000*(t2-t1):.1f}ms) print(f缓存命中率: {embedder.stats.hit_rate:.1%}) if __name__ __main__: asyncio.run(main())实现要点L1 用 OrderedDict 实现 LRUmove_to_end在命中时将元素移到末尾popitem(lastFalse)淘汰最旧的。Python 内置实现零依赖。L2 用 FAISS IndexFlatIPInner Product 索引 faiss.normalize_L2归一化后等价于 Cosine Similarity。FAISS 的搜索是 C 实现单次查询微秒级。粗粒度 embedding 的性能取舍我给出了一个基于 SHA256 的近似方案纯演示用。生产环境应该用一个真正的轻量模型如all-MiniLM-L6-v2384 维单次推理 5ms而不是完整的大模型。缓存写回L2 命中后写入 L1这叫缓存晋升。一个热门文本会在 L1 中常驻享受纳秒级命中。统计可观测CacheStats提供 L1/L2 命中率和总量建议接入 Prometheus观察缓存效率随时间的退化情况。四、边界分析与架构权衡4.1 内存容量规划假设你的向量维度是 1536text-embedding-3-small每个向量 1536 × 4 bytes 6KB。另外 FAISS 索引本身也有内存开销IndexFlatIP 存储所有向量与向量数据大小相当。L1 10000 条 ≈ 60MBL2 50000 条 ≈ 300MB 300MB索引 600MB。总计约 660MB。如果内存紧张优先砍 L2。L1 的命中率通常能覆盖 60%-80% 的场景高频问题精确重复L2 贡献额外的 10%-20%。4.2 缓存失效策略文本的 embedding 向量不应该变化同一模型 同一文本 同一向量。但如果模型升级了比如从text-embedding-3-small换成text-embedding-3-large所有缓存必须失效。建议在缓存中存储模型版本号或者在模型升级时清空整个缓存async def invalidate_all(self): async with self._lock: self._l1_cache.clear() self._l2_vectors.clear() self._l2_texts.clear() self._l2_index faiss.IndexFlatIP(self._l2_dim) self._stats CacheStats()4.3 什么时候不应该缓存文本极其长且唯一如用户上传了一整篇论文5000 字这种文本基本不会重复出现缓存只是在浪费内存。流式 embedding 场景如果 embedding 模型是在线的如调用 OpenAI API并且文本是流式生成的缓存的管理复杂度超过收益。要求严格一致性的场景如果你无法接受缓存返回的是旧模型的向量这种可能就不要用缓存或者在模型更新时做全量清空。4.4 分布式缓存的考量如果你的 RAG 服务是水平扩展的多进程每个进程有自己的内存缓存。同一文本在两个进程中各做了一次推理。解决方案有两个方向一致性哈希路由在负载均衡层按文本哈希路由到固定实例。相同文本总是命中同一个实例的缓存。Redis 共享缓存用 Redis 做 L1HSET md5 - vector_binaryFAISS 索引仍然在本机内存。但这引入了网络延迟L1 命中从纳秒级变成毫秒级。对于大多数场景一致性哈希路由比共享 Redis 更简单也更高效。五、总结Embedding 缓存是 RAG 服务优化中最投资回报比最高的手段之一几十行代码几十 MB 内存能把 embedding 推理量降低 50%-80%。两层缓存的取舍很清晰L1 精确匹配搞定高频重复问题。命中率高、延迟极低、实现最简单。L2 语义匹配搞定换个说法问同一个问题。命中率补充代价是额外的 FAISS 索引和粗粒度 embedding 推理。生产建议先从 L1 开始上线观察命中率。如果命中率已经超过 70%L2 的价值会打折扣。如果命中率不到 50%说明你的问题多样性较高L2 会是一个好的补充。内存和 GPU 时间之间的权衡永远是缓存设计的核心命题。在 embedding 这个场景下内存太便宜了GPU 太贵了——答案呼之欲出。下一篇预告向量检索的分布式索引构建MapReduce 模式在大规模向量建库中的应用。
