
长上下文 LLM 如何复用 KV 缓存LMCache 缓存引擎源码拆解【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache把一份两万个 token 的 RAG 文档丢给 vLLMprefill 阶段要把整段序列的 KV 缓存重新算一遍。可如果系统提示词、多轮对话历史或参考文档在多个请求间是共享的这部分计算就是纯浪费。LMCache 的定位就是给推理框架加一层 KV 缓存层让它把重复前缀的 prefill 省掉。这篇文章读 lmcache/v1/cache_engine.py 这一个文件跟着一次真实请求走通怎么知道缓存存不存在、怎么取、怎么存三步并讲清增量哈希链与前缀命中判定的取舍。读完你会拿到一张从 vLLM 适配器到存储后端的数据流地图。为什么前缀是否存在没法靠整序列哈希解决缓存系统必须回答的问题是这段前缀的 KV之前算过没有如果直接把整条序列哈希当 key行不通请求 A 是 1000 token请求 B 是 2000 token两者只有前 1000 个相同B 想复用 A 的缓存却拿不到——它的 key 和 A 完全不同。所以 key 必须前缀敏感相同的 N 个前缀 token 必须产生相同的 key而长序列要能拆成多个 key 逐段复用。这就把问题变成了两个子问题序列怎么切块、块与块之间的 key 怎么关联。这两个子问题的答案是整个引擎里信息密度最高的部分。缓存键是怎么算出来的一块一印章的增量哈希链先看一段代码注意核心逻辑不在 cache_engine.py而在 lmcache/v1/token_database.py# token_database.py L358-365 def _prefix_hash(self, token_chunks): prefix_hash self._get_init_hash() # 初始值 NONE_HASH for token_chunk in token_chunks: prefix_hash self._hash_tokens(token_chunk, prefix_hash) yield prefix_hash每块的哈希输入是上一块的哈希 本块的 token像一串盖章的链条——每一块都盖着上一块的印章。两个序列只要前 N 块 token 相同前 N 个 key 必然相同这就是前缀匹配的全部基础。cache_engine 持有一个ChunkedTokenDatabase实例分块大小来自配置缺省回退值 256见 lmcache/v1/token_database.py L327-329store/retrieve/lookup三个入口都靠它把 token 序列翻译成 key 列表。你可能会问为什么哈希函数不用自己实现看_hash_tokens的尾部# token_database.py L287-295 canon_prefix, canon_tokens, canon_extra self._canonicalize_hash_inputs( prefix_hash, tokens_tuple, extra_keys ) return _normalize_hash_to_int( self.hash_func((canon_prefix, canon_tokens, canon_extra)) )哈希函数优先从 vLLM 拿如 sha256_cbor这样 LMCache 的 key 与 vLLM 自身的 block hash 体系一致两侧不需要交换 token 来对账拿不到才回退 Python 内建hash而回退时代码会强制警告必须设置PYTHONHASHSEEDL168-175——否则每个进程哈希不同跨进程共享缓存直接失效这是部署时最容易踩的坑。拿到哈希后再包一层元数据lmcache/utils.py L388-395dataclass(slotsTrue) class CacheEngineKey: model_name: str world_size: int worker_id: int chunk_hash: int dtype: torch.dtype request_configs: Optional[dict] field(default_factorydict)为什么不止用 chunk_hash因为同一段 token在不同模型、不同张量并行度、不同 worker 上的 KV 形状完全不同key 里不带这些信息就会把 A 模型的 KV 写进 B 模型的请求。request_configs里还能塞lmcache.tag.前缀的标签L399-411用来隔离不同 LoRA 或请求类型的缓存。key 备好了。接下来看它如何被一次真实请求消费。一次请求的完整链路lookup、retrieve、storevLLM 侧的驱动方是 lmcache/integration/vllm/vllm_v1_adapter.py调度前先lookup问命中长度命中则retrieve把 KV 拉回 GPUprefill 结束再store写入。先问lookup 只认从头连续的命中# cache_engine.py lookup() L1228-1240非 layerwise 分支 hit_chunks, block_mapping self.storage_manager.batched_contains( keys, search_range, pin ) for idx, (start, end, key) in enumerate(chunk_info_list): if idx hit_chunks: res end continue return res你可能会问为什么一遇到 miss 就立刻 return而不是跳过去继续找后面的块答案还是那条哈希链——第 3 块的 key 由第 2 块的 key 推出第 2 块被逐出后第 3 块的 KV 对这条请求已经对不上位了就算还在存储里也没法用。所以命中定义是从头起最长的连续前缀batched_contains一次批量问存在性、只查元数据不动数据lookup 因此足够便宜才敢放在调度之前调用vllm_v1_adapter.py L1421 经 lookup_client 发起。再取retrieve 用 ret_mask 界定信到哪retrieve返回一个与 tokens 等长的 bool 张量L786 的 docstring 说明得很直白True 的位置表示这个 token 的 KV 已经落到 GPU。内部_process_tokens_internal按位置批量batched_get遇到取失败的块立即刹车# cache_engine.py L1762-1777 for (key, start, end), memory_obj in zip(blocks, memory_objs, strictFalse): if memory_obj is None: if last_failed_block_start is None or last_failed_block_start start: last_failed_block_start start break reordered_chunks.append((key, memory_obj, start, end)) ret_mask[start:end] True注意这里的防御姿态contains说存在不等于get时数据还在中间可能被并发逐出L1789-1790 会把失败点之后的 mask 整体翻回 False诚实收缩命中范围。取回的块随后由gpu_connector.batched_to_gpu约 L910一次性写进引擎的分页 KV内存对象走引用计数to_gpu完成后ref_count_down/unpinCPU 侧 staging 缓冲立刻可给下个请求复用。最后存store 的三步曲内存不够就提前收手# cache_engine.py store() L485-519 及 L557-568 for start, end, key in self.token_database.process_tokens( tokens, hashes, offsets, mask, request_configsrequest_configs): memory_obj self.storage_manager.allocate(kv_shapes, kv_dtypes, fmtself.fmt) if memory_obj is None: break # 内存吃紧只保留已分配的分块 ... self.gpu_connector.batched_from_gpu(memory_objs, starts, ends, **kwargs) self.storage_manager.batched_put(keys, memory_objs, locationself.store_location)注意 store 并不直接收 KV 张量而是收 token 序列加上 vLLM 分页缓冲和 slot mapping经 kwargs 传入由 GPU connector 负责把散落在分页缓冲里的 KV gather 成连续 MemoryObj——存储层因此与 vLLM 内部的页布局彻底解耦。allocate失败即break是容易被忽略的背压设计宁可少存也不让推理线程等内存漏掉的尾部会由下一次带同样前缀的请求自然补上。 还有一个 layerwise 变体store_layer/retrieve_layerL591、L972把整段处理写成生成器按层 yield让 GPU 搬运与上一层的存储写入流水重叠# cache_engine.py store_layer L751-756 for layer_id in range(self.num_layers): yield next(mem_obj_generator) self.storage_manager.batched_put( keys[layer_id], memory_objs[layer_id], locationself.store_location)代价是 key 要按层拆分、控制流复杂一截换来的是层间流水线重叠适合长序列、显存带宽紧张的场景。设计权衡哈希链为了前缀匹配放弃了什么链式哈希的代价一目了然序列中间插一个 token后面所有块的 key 全部改变旧缓存变成对不上号。这与 Merkle 式的独立块哈希恰好相反——后者局部变更只作废局部但它天然丢掉了这是同一条前缀的第几块的位置信息必须自己挂父指针而链式结构里的 prefix_hash 干的就是这件事。KV 复用的主战场就是同前缀场景且链式 key 与 vLLM 自带的前缀缓存语义一致省掉双套哈希体系的维护成本所以选链式是收益最大的取舍。另一组权衡是 lookup 与 retrieve 分离。lookup 只查元数据、可以 pin 住命中块lookup_pinsL216-219甚至走async_lookup_and_prefetchL1322-1378先把数据预取进事件管理器等请求真正被调度时 retrieve 只需等 future把磁盘或远端延迟藏进别的请求的计算时间。代价是状态管理的复杂度pin 了就必须显式释放请求被取消时要走lookup_unpinvllm_v1_adapter.py L1114-1139否则缓存会慢慢泄漏。收益边界与下一步先说缓存帮不上忙的时候存储后端有自己的逐出策略storage_manager管理 local_cpu、磁盘、远端多级层级工作集远大于缓存容量时命中率由淘汰质量决定另外收益与前缀共享度成正比请求前缀互不相干时 lookup 恒返回 0只留下哈希计算的固定开销——这种负载下不如调小 chunk 或直接关掉远端层。整个文件的风向标也很一致每个公开方法入口先is_healthy()自检unhealthy 直接跳过store L416、retrieve L808推理关键路径上宁可不用缓存也不拖慢服务。如果你要继续往下读建议按数据流顺序多级存储与逐出看 lmcache/v1/storage_backend/ 下的 storage_managergather/scatter 的具体 kernel 看 lmcache/v1/gpu_connector/按语义分段而非定长切块CacheBlend看 token_database.py L452 起的SegmentTokenDatabase端到端场景配置参考 examples/kv_cache_reuse/ 和 benchmarks/rag/。自己跑最小验证的话从 examples/basic_check/ 起步只配 local_cpu 层用两条同前缀请求观察日志里的 Retrieved X out of Y required tokens——看到第二条请求的 X 明显大于第一条就证明哈希链和前缀命中在按预期工作了。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考