LMCache 源码解析LMCacheEngine 如何用分块 KV 缓存复用压低 LLM 首字延迟【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCacheLMCache 是架在推理引擎之上的 KV 缓存复用层。本文解析 lmcache/v1/cache_engine.py讲清 LMCacheEngine 如何把 token 序列变成可复用的 KV 缓存块以及缓存引擎的分块前缀哈希、L1 内存分配和 layerwise 流水线是怎么配合的。30 秒速览TL;DRLMCacheEngine是 KV 缓存的搬运调度中枢store 把 GPU 上的 KV 转成 CPU 侧MemoryObjretrieve 再把它写回 GPU缓存键 chunk_size默认 256切块的链式前缀哈希块间像链子改一个 token后续块的键全部变化分配失败不抛错而是少存/少命中几个块直接降级layerwise 模式用生成器流水线让第 L 层的 GPU 拷贝与第 L-1 层的落盘重叠执行它要解决什么问题长上下文推理的大头开销在 prefill同样的系统提示词、文档、多轮对话前缀每个请求都重算一遍 attention 并生成 KV。LMCache 把推理引擎里的 GPU KV 缓存搬到 CPU/磁盘/P2P 后端复用下次命中前缀直接回载跳过重复 prefill。cache_engine.py在整体架构里居中上层LMCacheManager和 vLLM/SGLang 适配器通过它进出它左手按 token 分块算键右手调StorageManager落盘同时用GPUConnector做 GPU 与 L1CPU 固定内存之间的批量搬运。vLLM 的请求路径是调度前lookup问命中多长前缀prefill 结束store命中时retrieve。跟着一条数据走完全程以一次 1024-token 的请求为例走retrieve主链路chunk_size256共 4 块1) 算键。retrievecache_engine.py L780-L854 先健康检查、开统计计时再调token_database.process_tokens。分块在前 token_database.py L334-L356哈希在 L358-L365每块产出(start, end, CacheEngineKey)yield 出去不落地成列表。2) 定位块。_process_tokens_internalL1708-L1759 拿全部键调storage_manager.get_block_mapping按后端位置分组再对每个位置batched_get批量取回MemoryObjKV 张量在 L1 内存里的句柄。3) 写回 GPU。L877-L912gpu_connector.batched_to_gpu把整批数据按starts/ends一次性写进引擎的 paged KV buffer全程不经过 Python 张量拼接。4) 收尾。L925-L943 逐块ref_count_down释放 L1 占用返回布尔掩码ret_mask哪些位置命中调度器按它裁剪 prefill。store 侧反向走同一条路 L388-L569process_tokens生成键 → 逐块storage_manager.allocate→gpu_connector.batched_from_gpu一次拷贝 →storage_manager.batched_put异步分发到各后端。这条链的关键KV 张量从头到尾不以整体出现全部以 256-token 的块为寻址、分配、转移的最小单位。3 个硬设计决策① 链式前缀哈希键一致性靠链换命中检查的 O(1)它解决什么跨请求、跨进程、跨节点找到同一份 KV。怎么实现每块哈希输入是(上一块哈希, 本块 token 序列)def _prefix_hash(self, token_chunks): prefix_hash self._get_init_hash() for token_chunk in token_chunks: prefix_hash self._hash_tokens(token_chunk, prefix_hash) yield prefix_hash改一个 token后面每块键全变等价于链子动一环后面全变。代价是跨进程共享缓存时哈希必须逐位一致所以项目专门做了 vLLM 多版本哈希函数适配token_database.py L123-L148并在未设PYTHONHASHSEED时打 warning/errorL310-L326。换掉这种设计的替代方案是每块独立哈希 显式前缀元数据键稳定了但命中多长前缀就得额外比对链结构batched_contains的顺序前缀语义也没法直接保留。② 逐块分配 MemoryObj内存压力就截断降级它解决什么GB 级 KV 不可能整体torch.emptyCPU 内存又随时可能被别的请求占满。怎么实现store 循环里每块单独allocate返回None就 break只保留已分配的部分for start, end, key in self.token_database.process_tokens(...): memory_obj self.storage_manager.allocate( kv_shapes, kv_dtypes, fmtself.fmt) if memory_obj is None: logger.warning(Local cpu memory under pressure so choosing to store only %d total chunks of KV cache., len(memory_objs)) break付出的是半截前缀——尾部块丢了换来的是永不 OOM、永不阻塞请求。retrieve 侧对称处理某个块batched_get返回None就break并回滚 L1762-L1790失败点之后的ret_mask清零避免引擎按掩码读到脏位置。如果换全量分配或整体失败的做法长序列请求会把整段缓存一起丢掉。③ layerwise 生成器拷贝与落盘重叠它解决什么非 layerwise 路径要等整个序列的 KV 从 GPU 拷完才开始batched_putPCIe 传输和存储 IO 串行。怎么实现store_layerL593-L776 是生成器第一次next启动第 0 层拷贝之后每层拷 L 层 → put L-1 层交错执行for layer_id in range(self.num_layers): yield next(mem_obj_generator) # 拷贝第 layer_id 层 self.storage_manager.batched_put( keys[layer_id], memory_objs[layer_id], ...) # 落盘第 layer_id-1 层代价是键要拆成每层一个split_layers且层式检索要求所有块来自同一位置L1050-L1054 直接 assert。如果换每层独立任务 线程池的实现能去掉生成器这种手工状态机但每层的同步点和 ref 计数都要自己兜复杂度并不更低。容易忽略的工程细节健康降级三件套store 直接 returnL416、retrieve 返回全 False 掩码L808-L810、mark_init_failed让is_healthy()恒为 False整体退回全重算而不会崩PD 分离的 sync backend 的remove不自动减引用要手动ref_count_down否则 PD buffer pool 泄漏L929-L933异步加载async_lookup_and_prefetchL1322-L1378 在调度阶段就发起预取retrieve 时从EventManager的 future 直接拿 L1 数据把 IO 延迟藏在排队里MLA 的save_only_first_rankpassive rank 跳过 store/retrieve由 leader 广播 KV广播时顺手保留 GPU 副本供batched_to_gpu从 HBM 读省掉二次走 PCIeL1816-L1915埋点_lmcache_nvtx_annotatestore_stats/retrieve_stats按 process_tokens / from_gpu / put 三段计时日志直接给 GB 与 GB/sL577-L589想上手用/扩展入口是LMCacheEngineBuilder.get_or_createL2092-L2143按instance_id建单例并校验 config/metadata 一致性销毁走destroy。配置面在 lmcache/v1/config.pychunk_size、max_local_cpu_size、enable_blending、save_only_first_rank是主旋钮。扩展点按策略接口挂钩推理引擎适配实现GPUConnectorInterfacebatched_from_gpu/batched_to_gpu新增存储介质实现 StorageBackend 并注册给StorageManager改分块方式替换TokenDatabaseL2083-L2090 是工厂blend 场景换成SegmentTokenDatabase。跑通示例可看 examples/kv_cache_reuse/压缩链路参考 docs/source/kv_cache_optimizations/。数据佐证注释里写了实测MLA 场景 leader rank 的batched_to_gpu走 PCIe 约 9mspassive rank 从 HBM 读约 0.5ms这个 18 倍差就是 GPU 副本替换优化的动机L886-L888TTFT 估算工具可复现命中前缀省多少延迟benchmarks/ttft-estimator/ttft-estimator.pyLLaMA/H100 示例输出总结LMCacheEngine的三条主线前缀哈希链解决键一致性分块分配解决内存压力下的降级layerwise 生成器解决传输与落盘的重叠。落地可做的优化batched_contains的 layerwise 分支仍是逐块循环L1191 有 TODO批量化后 lookup 延迟能再降多位置检索L1049 的 TODO打通后热块在 L1 与远端并存时不必强选一个位置store 侧内存压力丢尾块策略可以演进成按访问频次保留热点前缀块提高降级后的有效命中率。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考