LMCache缓存引擎源码解析:KV存取全链路

发布时间:2026/9/14 1:51:15
LMCache缓存引擎源码解析:KV存取全链路 LMCache缓存引擎源码解析KV存取全链路【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCacheLMCache 是一个面向 LLM 推理的 KV 缓存管理层负责把推理引擎算出的 KV 缓存注意力中间结果长文本推理里最贵的部分落到可复用的存储里让后续请求跳过重复的 prefill 计算。本文沿一条请求的数据流拆解核心文件 lmcache/v1/cache_engine.py约 2200 行讲透分块、键设计、查找、存储与检索的完整链路回答一个问题LMCache 的 KV 缓存复用到底是怎么做到的。缓存引擎实际管理的是哪三块LMCacheEngine类本身不碰磁盘也不碰 GPU 显存它更像调度台真正的活由三个协作者完成初始化见 lmcache/v1/cache_engine.py L100-L246token_database把 token 序列转成起点, 终点, 缓存键三元组列表负责定址storage_manager统一管理 CPU 内存、本地磁盘、Redis、P2P 等多级后端见 lmcache/v1/storage_backend/管分配、淘汰和引用计数gpu_connector推理引擎 GPU 上的 KV 与 LMCache 侧 CPU 内存对象MemoryObj之间的搬运工。类比token_database 算收件地址storage_manager 管仓库gpu_connector 负责取送。这张图展示的是 LMCache 数据面在分离式推理中的位置P 侧引擎读/写分布式 KV 存储D 侧引擎从同一存储里取回缓存——而读写两侧走的正是本文要讲的同一条引擎链路。token 序列是怎么切成可缓存的块整段序列当一个键的问题是一个 token 不同就整体失效。LMCache 默认每 256 个 token 切一块ChunkedTokenDatabase的chunk_size256lmcache/v1/token_database.py L298。更关键的是每块的哈希不是独立算的而是链式的前缀哈希# lmcache/v1/token_database.pyChunkedTokenDatabase._prefix_hash 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 # 每块产出一个链式哈希这段代码要解决同一文本出现在不同位置不能互相命中的问题。为什么这样写KV 里的注意力结果与位置强相关若每块哈希只依赖本块 token不同请求里相同文本段会撞上同一个键直接复用就是错的。链式哈希把序列开头到本块的完整前缀编码进一个数天然保证只有前缀完全相同的块才可能命中。块键之外还拼上model_nameworld_sizeworker_idchunk_hashdtypeCacheEngineKeylmcache/utils.py L389把不同模型、不同并行拓扑的缓存隔离开。请求还没调度时lookup 如何算出可复用的前缀长度vLLM 调度器在入队前会调engine.lookup(tokens)lmcache/v1/cache_engine.py L1130问一句我的前缀有多少 token 在缓存里引擎先用process_tokens分块拿到键序列再调storage_manager.batched_contains批量查存在性。核心在返回值逻辑——只认连续命中# lmcache/v1/cache_engine.py L1236-1240lookup 非 layerwise 分支 for idx, (start, end, key) in enumerate(chunk_info_list): if idx hit_chunks: # 前 hit_chunks 块连续命中 res end continue return res # 遇到断点只返回断点前的前缀长度它要解决的是8 块命中 5 块能用几块。为什么只返回一个整数缓存只能按前缀来跳过 prefill中间的命中没有意义给上层一个命中的前缀 token 数就足够它决定跳过多少计算调度接口因此保持极简。prefill 算完KV 从 GPU 进存储的三步引擎算完 prefill 后回调engine.store(tokens, **kwargs)同文件 L388核心就三步# lmcache/v1/cache_engine.py L485-569store 核心 for start, end, key in self.token_database.process_tokens(...): memory_obj self.storage_manager.allocate(...) # 1. 分配 CPU 内存对象 if memory_obj is None: break # 内存不足存得下多少存多少 self.gpu_connector.batched_from_gpu(...) # 2. GPU KV 批量拷入 CPU self.storage_manager.batched_put(keys, memory_objs) # 3. 批量异步落到各级后端它要解决存缓存不能拖慢在线服务的问题。为什么这样写分配失败只丢弃尾部、绝不阻塞请求batched_put是异步提交store返回时数据可能还没落盘所以关键路径上只剩一次 GPU→CPU 拷贝。整个 store 被stats_monitor.on_store_request/on_store_finished包住分段计时算键、搬 GPU、put三个阶段L469-L574这就是文档里提到的请求级可观测性埋点。第二次请求来了缓存里的 KV 怎么回到 GPUretrieve(tokens)L780是 store 的逆过程同样的 token 序列经链式哈希产生同样的键序列——这是能命中的前提——然后get_block_mapping定位键所在的层batched_get逐批取出 MemoryObj最后gpu_connector.batched_to_gpu写回 GPU KV 缓存。返回值是一个 bool 掩码ret_mask哪些位置已补齐推理引擎对这些位置跳过 prefill。两个保证正确性的细节值得注意。一是连续性_process_tokens_internal里若某块索引在但取不出来立即 break并把该位置之后的掩码统一清回 FalseL1762-L1790绝不向上传递带洞的前缀。二是引用计数每次batched_get给 MemoryObj 加引用retrieve 结束后统一ref_count_down/unpin异步预取async_lookup_and_prefetchL1322取到却没被使用的对象走cleanup_memory_objs归还防止 CPU 钉住内存池被占满。这个文件里藏着哪些工程取舍三个取舍解释了 LMCache 为什么能塞进推理关键路径逐层流水store_layer/retrieve_layerL593、L974用生成器按层 yield第 i 层的 GPU→CPU 拷贝与第 i-1 层的落盘并行把 PCIe 传输时间藏进计算间隙只认前缀主链路 lookup 只返回连续前缀长度任意位置命中交给非前缀复用路线CacheBlend见 README 的 Non-prefix KV reuse 条目主路径语义简单、正确性好证明MLA 只写一份save_only_first_rank模式下仅 rank 0 存储再广播给其他 rankL856-L871避免 N 份冗余写入广播时甚至用 GPU 侧副本替换 CPU 源缓冲把 leader 的 to_gpu 从 PCIe 受限压回 HBM 速度。收尾LMCache 的做法 vs 朴素做法维度朴素做法LMCache 的做法分块粒度整段序列一个键一 token 差异全失效256 token 一块前缀可部分复用哈希方式每块独立哈希分不清位置链式前缀哈希位置编码进键内存压力OOM 或拒绝服务分配失败丢弃尾部块存得下多少存多少落盘时机同步写完才返回batched_put 异步关键路径只剩 GPU→CPU命中判定返回命中块集合上层自己拼返回单一前缀长度调度器直接可用多层 KV整层串行拷贝生成器逐层流水传输与落盘重叠把这条链路走完就能回答开头的问题LMCache 的复用不是靠某个精巧的算法而是链式哈希定址 批量异步搬运 严格前缀语义 引用计数兜底四件事的组合。想继续看各后端的实现入口在 lmcache/v1/storage_backend/端到端用法可看 examples/kv_cache_reuse/。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考