十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

LMCache 源码解析:一个文件讲透KV缓存复用机制

LMCache 源码解析:一个文件讲透KV缓存复用机制 LMCache 源码解析一个文件讲透KV缓存复用机制【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache同一段两万个 token 的系统提示词用户每发一次请求推理引擎都要把 prefill首个 token 输出前的预填充计算重新算一遍——首 token 延迟TTFT被这种重复计算吃得所剩无几。这篇 LMCache 源码解析从 lmcache/v1/cache_engine.py 讲起跟踪一块 KV 从写入到被复用的完整旅程帮你看懂它的 KV 缓存机制理解 LLM 推理加速是怎么落到代码里的。全景LMCacheEngine 站在哪一环LMCacheEnginecache_engine.py L83是承上启下的翻译官上游的 serving 引擎vLLM/SGLang通过它把 token 序列和 GPU 上的 KV 张量交进来下游是StorageManager管理的一堆存储后端本地 CPU、本地磁盘、P2P 等。它自己不存数据只做三件事——把 token 切成块并生成键委托给TokenDatabase、搬运张量委托给GPUConnector的batched_from_gpu/batched_to_gpu、落盘/取数委托给StorageManager的batched_put/batched_get。引擎实例由LMCacheEngineBuilder.get_or_create按instance_id缓存复用L2093并按配置选择分块策略默认ChunkedTokenDatabase等长分块开启 CacheBlend 混合推理时用SegmentTokenDatabase按分隔符切段L2083-2090。store()一条KV序列如何被分块、键控并批量写入写入入口store()L388的主循环值得逐行看with store_stats.profile_process_tokens(): for start, end, key in self.token_database.process_tokens( tokens, hashes, offsets, mask): memory_obj self.storage_manager.allocate(kv_shapes, kv_dtypes, fmtself.fmt) if memory_obj is None: break # L1 内存紧张只存到当前块为止 keys.append(key); memory_objs.append(memory_obj) self.gpu_connector.batched_from_gpu(memory_objs, starts, ends, **kwargs) # GPU→CPU 批量拷贝 self.storage_manager.batched_put(keys, memory_objs, locationself.store_location) # 批量提交各后端三个关键点分块粒度由chunk_size控制默认 256token_database.py L329尾部不满一块的 token 是否保存由save_unfull_chunk决定。分配失败即截断allocate返回None时直接break宁可少存也不阻塞请求——这是典型的内存压力下降级写法。全程批量先攒齐所有块的键和MemoryObj一次batched_from_gpu完成 GPU→CPU 搬运再一次batched_put提交给所有后端每个后端各自batched_submit_put_task见 storage_manager.py L386-431。逐块 put 会放大 I/O 开销批量是这条路径上最直接的优化。MLA 模型下只有首个 rank 参与写入save_only_first_rankL116-120其余 rank 在_is_passive()判断后直接跳过storeL424-426。retrieve()前缀匹配如何在首个未命中处停下读路径分两步调度侧先调lookup()L1130拿到连续前缀命中了多少 token再在真正加载时调retrieve()L780。两者共享同一套token_database.process_tokens生成的键序列语义是严格的从第 0 块开始连续命中for start, end, key in chunk_info_list: if idx hit_chunks: # 前 hit_chunks 块视为命中 res end continue return res # 第一个未命中处即停res 停在上一块的 endretrieve内部的_process_tokens_internalL1708对未命中的处理更讲究批量batched_get后逐块检查一旦某块取不到元数据在、数据没了就记录失败位置并break然后把失败点之后的命中标记全部清零if memory_obj is None: last_failed_block_start start # 记录首个失败位置 break ... if last_failed_block_start is not None: ret_mask[last_failed_block_start:] False # 失败点后的假命中作废最后被丢弃的块会走ref_count_down/unpin释放L925-937避免 L1 内存池泄漏。一个容易踩坑的细节写在注释里L945-956retrieve返回的 token 数可能大于请求所需——例如 page_size16 而 chunk_size256 时整块回载会覆盖一小段引擎已有的 token调用方不能假设取回长度 需要长度。CacheEngineKey 与前缀哈希链不同模型、不同 rank 为何不串键键由两部分组成。分块哈希在 token_database.py 中以链式方式计算L358-365prefix_hash self._get_init_hash() # NONE_HASH默认 0 for token_chunk in token_chunks: prefix_hash self._hash_tokens(token_chunk, prefix_hash) # hash(前缀, 当前块) yield prefix_hash每个块的哈希 hash_func((前缀哈希, token 元组, 附加键))L289-295所以序列中任意一个 token 变化其后所有块的键都会变——这既保证了前缀匹配的准确性也让缓存命中必须从第 0 块连续成立成为结构性约束。键本身是 5 元组utils.py L389-396dataclass(slotsTrue) class CacheEngineKey: model_name: str # 模型名隔离不同模型 world_size: int # 并行规模隔离不同张量切分 worker_id: int # rank 编号隔离 TP 下各卡的分片 chunk_hash: int # 上文的链式分块哈希 dtype: torch.dtype # KV 精度fp8/bf16 互不复用序列化格式为modelworld_sizeworker_idchunk_hashdtypeto_stringL438-446可选追加tag%value形式的请求级标签如 LoRA ID。注意 MLA 场景下world_size会被折叠为 1token_database.py L234-248使键与部署规模无关。另外若 vLLM 的哈希函数不可用而回退到 Python 内建hash跨进程一致性依赖PYTHONHASHSEED——源码在 P/D 分离模式下会直接报错incorrect KV cache transferL311-326。性能细节批量、跳过已存在与延迟入 GPU这条路径上的几个不显眼但省时间的设计跳过已存在块层式写入store_layer中先用storage_manager.contains检查首层键已存在则continue不重复占内存、不重复传输L674-679。延迟入 GPUretrieve只在reordered_chunks非空时才调用batched_to_gpuL877-912全未命中时一次 PCIe 传输都不发生。HBM 替代读save_only_first_rank下 leader rank 广播时在 GPU 上保留一份副本_leader_gpu_substitute_objs让随后的batched_to_gpu从显存读而不是重新走 PCIe——注释里给出实测差异约 9 msPCIe 受限对 0.5 msHBM 受限L884-888。命中即回写 L1从远端后端取到数据时batched_get会自动把结果提交给LocalCPUBackend做热缓存storage_manager.py L494-513下次命中走本地内存。运行时开关freeze()冻结 L1 热缓存禁止写入、只读本地set_hot_cache()动态开关热缓存L335-384供在线调优而不重启。监控则通过LMCStatsMonitor.GetOrCreate()单例以on_store_request/on_store_finished等钩子埋点L222按 store/process_tokens/from_gpu/put 分段计时并输出 GB/s 吞吐日志里能直接看到每段耗时占比。工程实践读这个文件先记住四件事跨进程共享缓存前先固定哈希用内建 hash 回退时设置PYTHONHASHSEED0否则各进程算出的chunk_hash不同缓存形同虚设token_database.py L168-175。mask 有硬约束mask中 False 的数量必须是chunk_size的整数倍否则process_tokens直接ValueErrorL404-412——已算前缀按 256 对齐再传进去。同 instance_id 重复构建要同配置get_or_create发现 config/metadata 与既有实例不一致会抛ValueErrorL2135-2142热重载配置时先destroy再建。异步预取是 lookup 与 retrieve 之间的第三路async_lookup_and_prefetchL1322把查键和取数都丢进 event loopretrieve里enable_async_loading时直接取 future 结果而不是同步阻塞L841-854——调试取数问题时注意区分这两条路。场景与数据这套机制收益来自哪里最典型的场景是 benchmarks/rag 里的多文档 RAG多份长文档被预计算precompute.py成 KV 块落盘新请求的公共前缀直接命中多轮对话与 PD 分离 部署prefill 与 decode 分开KV 经 P2P 后端搬运也是同款机制可验证的量化数据主要来自代码内的实测注释leader rank 的batched_to_gpu走 PCIe 约 9 ms改用 HBM 副本后降至约 0.5 msL886-888——单是这一个优化就在每次 retrieve 的关键路径上省了近 8 ms。评测方法仓库也自带benchmarks/rag统计吞吐、平均 TTFT 与平均质量三项指标READMEbenchmarks/ttft-estimator则用于估算不同上下文长度下的 TTFT 收益接下来你可以做的两件事部署后打开enable_kv_events并用 lmcache lookup 类接口观察前缀命中长度确认PYTHONHASHSEED、chunk_size与 serving 端对齐2. 跑一遍 benchmarks/rag 的 precompute launch 脚本建立 TTFT 基线再用freeze/set_hot_cache开关对比命中率数据比直觉可靠。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表