深入解析MemOS KV Cache:源码级Agent短期记忆存储设计
1. 先把 KV Cache 在 MemOS 里的坐标讲清楚从一次 Agent“断片”事故说起事情的起因是挺普通的。当时我在调一个基于 MemOS 的多轮对话 Agent功能本身不复杂就是让 Agent 记住用户前面几轮说过的话在后续回答里引用。最开始的做法很粗暴把所有对话历史塞进一个全局 dictkey 是会话 IDvalue 是整段历史列表。本地单线程跑测试一切正常一旦接上真实流量——多个会话并发读写同一个会话里历史越积越多——问题就来了内存只涨不降重启之后所有记忆清零最要命的是偶尔会出现“这个会话明明存在却读不到数据”的诡异现象。后来我去翻 MemOS 源码里跟记忆存储相关的模块才发现真正落地的方案根本不是我想的那种“一个大字典走天下”而是拆成了好几层其中专门有一层叫 KV Cache负责管理短期内的键值数据。这篇笔记系列第 4 篇就是把这块的源码实现从头到尾梳理一遍重点记录 KV Cache 的定位、数据结构、读写路径、淘汰策略和并发处理。先说结论KV Cache 在 MemOS 里的角色不是替你做长期持久化而是给“高频读、低频写、生命周期短”的数据提供一个高速暂存区。它解决的是一类很具体的问题——同一个 Agent 任务流里工具调用结果、中间状态、临时会话上下文会被反复读取每次都去查全量存储或者重新计算开销大到没法接受。与之相对的用户的长期偏好、历史对话归档这些“必须永久保留”的数据根本不应该放进 KV Cache 里管那是另一层存储系统的事情。这个边界如果不提前想清楚KV Cache 很容易被用成“垃圾堆”什么数据都往里塞最后缓存层膨胀到比主存储还大命中率反而低得可怜。MemOS 源码里面对这个问题的做法非常务实它只承诺三件事数据能按 key 查到、过期数据能清掉、容量有上限。超出这个范围的需求都交给上层业务自己处理。1.1 Agent 记忆的分层逻辑为什么旧代码把所有东西塞进一个大字典在继续往下读源码之前得先把 Agent 的记忆分层搞清楚否则看 KV Cache 的实现会觉得很突兀。Agent 的记忆需求粗略分三种。第一种是长期记忆比如用户的名字、偏好、历史项目信息这类数据以“天”甚至“年”为生命周期必须持久化通常落到数据库或者向量存储里。第二种是短期记忆比如当前任务上下文、最近几轮的对话内容、某个工具刚返回的结构化结果生命周期只有几分钟到几小时等任务结束就可以丢。第三种是瞬时状态比如某个函数执行到一半的中间变量、一次请求的 trace 数据只在一个请求周期内有效甚至不需要跨请求保存。我最初的错误就在于把第二种和第三种都当成第一种来存用一个全局 dict 全部带走。短期数据和长期数据处理方式完全不同混在一个容器里会导致一个经典问题长期数据占着 key 不放短期数据不断插入dict 无限膨胀最终内存爆炸。而 MemOS 源码给我的第一个启发就是——不同生命周期的数据应该交给不同策略的存储层KV Cache 只服务于“短期但会跨几次请求复用”的那一层。MemOS 的源码里对 KV Cache 的定位写得很直白它就是短期记忆的物理载体。对话过程中产生的临时实体、工具输出摘要、最近 N 轮的对话快照会被序列化后以 KV 对的形式放进去一旦过期或者容量不够淘汰掉也不心疼因为上游随时可以从原始数据源重新算出来。1.2 KV Cache 在源码里的模块边界它到底该管什么、不该管什么读 MemOS 源码有个很明显的感觉它的模块边界削得很干净。KV Cache 这个模块对外暴露的接口非常克制核心就几个方法get、set、delete、clear再加一个周期性清理的 sweep。它不感知业务不关心 value 是什么结构也完全不管数据从哪来、到哪去。你给它一个 key、一个 value、一个过期时间它保证在这个时间点之前你 get 同一个 key 能拿到当时 set 进去的东西或者等价的东西时间一过它就默认你可以再来 set 一次新的。它不管什么不管持久化。Cache 之所以叫 Cache就是因为数据丢了可以重建。MemOS 的源码里所有写操作都是内存操作异步的落盘和同步到外部存储是另外的模块在干KV Cache 本身连文件句柄都没持有。这一点很多二次开发的同学习惯了把缓存当数据库用遇到缓存重启丢数据就骂框架其实是用错了模块。同时它也不管跨进程的一致性。单机多线程模型下KV Cache 负责把并发访问理清楚多实例部署时它会让位给 Redis 或者其他外部缓存实现。源码里这块通过一个抽象接口切开了边界后续我会在并发章节详细展示这个设计。2. KV Cache 的数据结构设计从基础 dict 到带 TTL 的分桶读完 MemOS 的 KV Cache 实现我最直观的感受是它没有发明任何新东西每个设计决策都是在解决一个具体问题。整体数据结构可以概括成两句话用哈希表保证 O(1) 的 key 查找用双向链表维护访问顺序。2.1 最初用 dict 实现的问题无过期、无上限、无淘汰先说说很多人包括我的第一步Python 里直接拿 dict 做缓存。dict 很方便查询 O(1)但是有四个天然缺陷。第一没有过期机制。你只能自己额外维护一个“key - 到期时间”的字典然后每次 get 的时候先查时间再做一次比较。这等于每次读取都做了两次哈希查找而且时间字典会和数据字典慢慢不一致。第二没有容量上限。dict 会一直增长直到内存耗尽。第三没有淘汰策略。即使你知道缓存已经满了dict 也不知道该丢谁。第四并发访问需要自己做锁。如果只靠一个 dict你基本上是把一套缓存系统应该做的事情全部推给业务代码每个使用方自己实现一遍最后必然百花齐放地写错。MemOS 源码里第一版 KV Cache 其实也是这么写的注释里还留着当时的一行 TODO换成支持淘汰策略的结构。后来版本的代码就进化成了哈希索引加访问序链表的结构。2.2 源码中 KVEntry 的字段拆解真正在源码里担负存储职责的是一个叫 KVEntry 的类。我按照自己的理解把核心字段重写成了可读性更强的版本dataclass class KVEntry: key: str value: Any expire_at: float # 绝对过期时间time.time() ttl access_count: int 0 # 累计访问次数用于 LFU 加权 last_access: float 0.0 # 最近访问时间用于 LRU 排序 size_bytes: int 0 # value 序列化后的字节数 prev: KVEntry | None None next: KVEntry | None None有几个字段值得展开讲。expire_at存的是绝对时间不是 TTL。这是源码里一个很小但很重要的细节。如果存 TTL那每次判断还要拿当前时间做一次减法存在精度损失和比较不一致的问题存绝对时间的话get 的时候只需一次if now expire_at就能判定过期。代价是 set 的时候必须取一次当前时间但这个开销完全可以接受。access_count和last_access是两个供淘汰策略使用的统计字段。一开始我以为只会保留一个后来看到源码里同时维护了它们才意识到它的淘汰策略不是纯 LRU也不是纯 LFU而是两者结合。size_bytes字段是我补充进去的但 MemOS 的容量控制思路确实考虑到了 value 的大小而不只是条目数量。如果一个缓存只限制条目数那某个 key 塞进去一个 100MB 的列表其他 key 就全被挤掉了这显然不合理。prev和next是双向链表的指针。哈希表负责按 key 找到对应的 entry双向链表负责记录所有 entry 的访问先后顺序。这两个结构通过同一个 entry 对象关联起来也就是所谓“哈希表 链表”的经典组合。2.3 哈希索引与双向链表为什么这种组合是标配很多人第一次见到“哈希表 双向链表”可能会想为什么不能只用 dict 存 key 和 value然后另外搞一个存储访问顺序的列表可以但问题在于“更新访问顺序”这个操作。每访问一个 key你都要把这个 key 挪到列表头部表示最近被使用。如果列表是普通的数组或者单链表你要先找到这个节点在列表中的位置然后做删除和插入。如果只存 value 不存节点指针查找位置需要 O(n) 遍历。但如果链表节点本身就是你哈希表里存的 entry那你在 get 到某个 key 的时候已经拿到了这个 entry通过 entry 的prev和next指针做摘除和插入都是 O(1) 操作。这就是 MemOS 这个结构设计的核心动机哈希表保证“找得到”双向链表保证“挪得动”。顺便说一句这个结构跟大家在操作系统课程里学过的 LRU 缓存实现思路一脉相承。MemOS 没有另辟蹊径而是选了一条最稳的路。源码笔记如果只看“新奇”的地方很容易忽略这种“看似平淡但非常重要”的设计但我自己写代码的经验是缓存这类基础模块越保守的设计越可靠。3. 读写主路径的源码实现细节get / set / delete 背后各有各的坑数据结构定下来之后读写路径就是围绕它做文章。源码里这三个方法的实现都不长但每个都埋了一些需要仔细品的细节。我先用简化伪代码把核心逻辑列出来再逐个解释为什么要这么写。3.1 set 路径TTL 写入、覆盖策略与“写穿”问题def set(self, key: str, value: Any, ttl: float) - None: now time.time() entry self._index.get(key) if entry is not None: # 覆盖旧值先移除旧节点 self._unlink(entry) self._total_size - entry.size_bytes else: entry self._new_entry(key, value, ttl, now) entry.value value entry.expire_at now ttl entry.size_bytes self._estimate_size(value) entry.access_count 0 entry.last_access now self._link_to_head(entry) self._index[key] entry self._total_size entry.size_bytes self._evict_if_needed()第一眼看上去这就是标准的 LRU 缓存写入流程查旧值、更新节点、移到头部、超容量就淘汰。这里藏在细节里的决策有三个。第一个决策覆盖已经存在的 key 时access_count被重置为 0。有这个细节我挺意外的因为一般的 LRU 实现覆盖时不会重置访问计数。MemOS 的考虑是一旦数据内容变了旧的访问频率统计已经没有参考价值了——你访问的是旧数据跟新数据没关系。如果保留旧的高频计数那么这个 key 会一直占据淘汰保护的有利位置排挤掉其他真正热门的数据。第二个决策覆盖操作同样会刷新过期时间。这是一个很合理的业务语义既然我重新 set 了这个 key那它的生命周期应当从重新写入的时刻开始算而不是继承旧的剩余时间。否则就会出现很尴尬的情况一个 key 已经剩 1 秒过期我重新写入了新值结果这新值的 TTL 只有 1 秒完全不符合预期。第三个决策所有 set 路径都会走到_evict_if_needed()。这意味着缓存容量不够时写入操作本身就是触发淘汰的时机不需要等后台任务慢慢跑。这种主动淘汰的设计可以保证写入完成后缓存一定处于合法容量范围内不会出现写了再说的中间状态。3.2 get 路径惰性删除、访问计数与热点提升def get(self, key: str) - Any: entry self._index.get(key) if entry is None: self._miss_count 1 return None now time.time() if now entry.expire_at: self._delete_entry(entry) self._miss_count 1 return None entry.access_count 1 entry.last_access now self._move_to_head(entry) self._hit_count 1 return self._deserialize(entry.value)get 这里最值得说的是“惰性删除”和“命中态处理”。惰性删除的逻辑是不专门起一个线程周期扫描所有 key 是否过期而是在每次 get 访问到某个 key 的时候发现它过期了才顺手删掉。这样做的好处很明显避免无意义的遍历。如果整个缓存的 key 都还没过期定时扫描就是白做功。坏处也很明显如果某个过期 key 永远不会再被访问那它就占着内存不释放。所以在实际运行时MemOS 还配了一个低频 sweep 定时器来处理那些“躲过所有访问”的过期垃圾这个我在下一节讲淘汰策略时会展开。命中态处理里包含了两个动作递增访问计数和移动到链表头部。递增access_count是为了给后续的加权淘汰提供数据移动到链表头部则是让这个 entry 成为“最近使用”的节点。链表头的语义就是“最热”链表尾就是“最冷”淘汰的时候优先从尾部摘除。还有一个细节源码里 get 返回的是反序列化后的值而不是直接返回存储的原始对象。这是为了避免外部调用方直接持有内部引用破坏了缓存内部的一致性。我知道有些会图省事直接返回内部对象结果外部改了对象属性缓存里的数据也被悄悄改了排查的时候非常痛苦。3.3 delete 与 clear手动失效的几种真实触发场景delete 和 clear 实现上没什么悬念就是摘链表、删哈希、更新计数重点在于它们会被什么场景触发。常见的主动删除场景有三个一是业务明确知道某个 key 的数据已经失效。比如 Agent 调用工具获取了最新天气此时缓存里还存着昨天查的旧天气业务应该主动 delete 旧 key避免下次命中过期数据。二是数据更新流程里存在“先删后写”的模式。有些场景下新值还没算出来但旧值已经确认不能用了那就先 delete让访问穿透到下层存储。三是一致性兜底。比如缓存背后的数据库发生了批量变更源码里会明确调用 clear 让整个缓存的局部或全部失效。4. 淘汰策略的工程取舍为什么只用 LRU 会在 Agent 场景翻车淘汰策略是 KV Cache 源码笔记里真正值得逐行展开的部分。MemOS 最终采用的是一个“以 LRU 为骨架、用访问次数做加权、以容量字节数为准入条件”的混合策略。4.1 定时清理与惰性清理的配合机制前面提过 get 路径里的惰性删除它负责处理单个 key 的过期判断。但这不够源码里还有一个周期性任务_sweep_expired()它起到一个很朴素但是很关键的作用兜底。def _sweep_expired(self): now time.time() for entry in self._iterate_all(): if now entry.expire_at: self._delete_entry(entry)需要注意的是这个循环不是每次全量扫描而是基于一个“过期时间最小堆”的优化。源码中实际上维护了一个按expire_at排序的小根堆每次只检查堆顶元素是否过期过期就弹出并删除对应 entry直到堆顶元素还未过期时停止。这个方案能把扫描的时间复杂度降下来避免定期全表遍历的性能损耗。我之前做类似模块时就是单纯地定时遍历整个 dict几万个 key 还好到几十万个 key 的时候那个定时任务就很吃 CPU。学了这个小根堆做法之后我把自己的项目也改成了这个方案效果立竿见影。4.2 LFU 与 LRU 的加权计分方案纯 LRU 的问题在 Agent 场景下很典型一个 key 在短时间内被密集访问然后好长一段时间没人碰它。按照纯 LRU 的逻辑这个 key 会因为“长时间没访问”而被排到链表尾部随时可能被淘汰。但如果过一会儿用户又重新问起同样的问题这个 key 又要重新计算一次浪费资源。纯 LFU 也有问题一个 key 曾经被访问过一万次但最近再也用不到了它的访问计数依然很高会一直占着缓存位置不出去。这个在 Agent 场景里也很常见——某次任务中反复读取了某个中间结果任务结束后这个结果就彻底没用了但它的访问计数让它在缓存里赖着不走。MemOS 源码里淘汰时用的排序指标是同时考虑了“最近访问时间”和“累计访问次数”的加权分。核心逻辑可以简化为def _score(self, entry: KVEntry) - float: age time.time() - entry.last_access # 让“最近访问时间”的权重大于“累计访问次数” return (1.0 entry.access_count * 0.1) / (1.0 age / 1000.0)这个公式的效果是如果两个 key 访问时间接近访问次数更高的获胜如果两个 key 访问次数相同最近访问过的获胜。它不像 LFU 那样只看计数也不像 LRU 那样只看时间而是在两者之间找平衡。具体系数每个业务场景都可以调关键是淘汰动作本身必须基于性价比——保留那些“未来最可能被再次访问”的条目。当然实际源码中淘汰操作是发生在_evict_if_needed()里它会从链表尾部开始往前逐个计算 score淘汰掉分数最低的条目每淘汰一个就检查一次当前总大小是否已降到阈值以下。4.3 容量控制条目数、字节数与单条上限这块是我认为 MemOS 做得比较完善的地方。它把容量控制拆成了三个维度第一个维度是最大条目数防止 key 数量过多导致哈希表膨胀和链表遍历开销。第二个维度是总字节数上限防止单个大 value 挤爆内存。第三个维度是单条 value 大小上限任何超过这个阈值的 value 在 set 的时候会直接拒绝缓存。def _evict_if_needed(self): while self._total_size self._max_bytes or len(self._index) self._max_entries: victim self._tail if victim is None: return self._delete_entry(victim)判断是否超限时是“或”的关系也就是任何一个维度触线就淘汰。这样设计的好处是大量小 key 的场景靠最大条目数兜底少量大 value 的场景靠总字节数兜底极端数据被单条上限直接拦住。这个设计对我最大的启发是容量控制不能只拍脑袋定一个“最多存 10000 条”。要结合 value 的典型大小和 Agent 实际运行的内存余量倒推最大字节数再结合业务上最多允许缓存多少个独立的原子状态确定最大条目数。5. 并发与多实例一致性缓存从单机走向共享时遇到的坑KV Cache 如果只在一个进程里跑并发问题还不算严重。一旦跨线程、跨进程乃至多实例部署情况就复杂了。源码里有一层适配抽象我把它单独拉出来说。5.1 锁粒度演进全局锁到分片锁最开始的实现所有 get/set 操作包在一把全局锁里。对于小缓存这没问题但一旦缓存命中率高、读请求密集全局锁就变成了性能瓶颈。所有读操作排队等一把锁缓存本身的快速优势被抵消了大半。MemOS 后续引入了分片锁的方案把哈希桶分成了 N 个分片每个分片一把独立的读写锁。访问不同的 key 时如果两个 key 落在不同的分片就可以并行执行只有落在同一分片的访问才需要竞争锁。哈希函数负责把 key 分散到各个分片尽量让访问热度均匀分布。这个改动带来的提升在单机压测场景下非常明显。不过这要求对 key 的分布有大概的判断如果业务 key 有明显的热点前缀哈希分片可能把热点全打到一个分片上导致某些分片锁竞争激烈而其他分片空转。5.2 进程间共享如何把本地缓存接口抽象成可替换后端单机多线程场景下进程内共享没问题。但 Agent 服务通常不会只跑一个实例。多个 worker 进程各持一份本地缓存就会遇到一致性问题用户请求打到实例 A写入了缓存下一次请求被负载均衡转到实例 BB 的本地缓存里压根没有这个 key于是又去查了一次底层数据。缓存存在的意义就减弱了。MemOS 源码里抽象了一个CacheBackend接口本地实现叫MemoryBackend同时还有一个RedisBackend。业务层代码只依赖CacheBackend不需要感知底层实现。这样单机本地调试用 MemoryBackend多实例部署时切换到 RedisBackend主体业务代码一行都不用改。class CacheBackend(ABC): abstractmethod def get(self, key: str): ... abstractmethod def set(self, key: str, value: Any, ttl: float): ... abstractmethod def delete(self, key: str): ... abstractmethod def clear(self): ...接口的抽象力度也很有讲究没有暴露淘汰策略、没有暴露内部容量配置。因为这些东西是底层实现细节一旦暴露出来上层代码就会开始依赖它以后的演进空间就被堵死了。5.3 缓存穿透、击穿、雪崩在 Agent 对话场景里的具体表现这三个缓存经典问题放到 Agent 场景里也都能找到对应的例子。缓存穿透Agent 的每个请求都会携带一个会话 ID如果某个异常请求伪造了一个不存在的会话 ID且 KV Cache 对“查不到”的结果没有做负缓存那么每次请求都会穿透到下层存储反复做无效查询。这个问题用“缓存空结果 短 TTL”就能解决。缓存击穿某个热点 key比如 Agent 系统消息里配置的一段全局工具说明在过期的一瞬间大量并发请求同时发现缓存里没数据于是一起冲进底层存储重建缓存造成瞬时压力。解决办法是加互斥锁只允许一个请求去重建其他请求等待结果。缓存雪崩在批量导入工具配置或批量更新系统提示词之后如果这些 key 恰好设置了相同的 TTL它们会在同一时刻集体过期下一波所有请求都会穿透到底层存储。解决办法是在 TTL 上增加随机抖动让过期时间分散开同时配合持久化文件把部分热点数据常驻内存。6. MemOS 的 KV Cache 和 LLM 推理层 KV Cache 有什么不同这是一个非常容易混淆的概念。跟我聊过的不少同学看到“KV Cache”就直接联想到大模型推理框架里的那套 KV Cache然后拿那套思路来套 MemOS 的缓存结果越看越糊涂。6.1 管理对象不同Token 状态 vs 应用语义记忆LLM 推理层的 KV Cache管理对象是 Transformer 模型每一层的 Key 和 Value 矩阵。它服务于注意力机制目的是避免每次生成新 token 时重新计算所有历史 token 的 K、V 向量。这是一套极其依赖显存的底层缓存系统生命周期以毫秒级计算一个对话生成结束就可能被清掉。MemOS 的 KV Cache 则完全不同。它管理的是应用层的数据对象会话上下文、工具调用结果、用户偏好片段。它不关心 token 之间的注意力关系只关心“这个 key 对应的数据还在不在有效期内、能不能查得到”。两者唯一的共同点只是名字里都有“KV Cache”五个字母。6.2 为什么不能直接用 transformers 框架的缓存对象有人可能会想既然 Agent 底层接的是 LLM那直接复用 transformers 推理时的 KV Cache 不就行了技术上不行逻辑上更没必要。transformers 的 KV Cache 对象的生命周期由推理引擎管理数据形态是 GPU 上的张量序列接口完全围绕注意力机制设计。如果拿它来缓存 Agent 的会话状态首先你没法持久化——显存一释放就全没了其次你没法跨模型共享——每个模型的 Layer 数量和 Head 维度都不一样最重要的是你根本没有必要把应用层的字典、列表形态的数据塞进张量里这等于用开拖拉机的方式去送外卖工具和场景不匹配。7. 实测中踩过的坑和排查过程照抄源码但照样翻车源码笔记写到这如果不记录一些实际翻车的经历总觉得少了点什么。下面三个坑是我在基于 KV Cache 落地 Agent 功能时真实遇到的每个都花费了不小的排查代价。7.1 缓存返回旧对象引用共享引起的幽灵更新现象是我缓存了一个对象 A然后从缓存里 get 出来修改了某些字段紧跟着下一次 get 发现拿到的对象已经被改过了。第一反应以为是并发问题后来加锁也没解决。排查过程最后落到 get 方法的返回值上。源码里 get 返回的是反序列化后的新对象但我当时复刻实现时为了省事直接把内部对象引用返回了。外部代码改的是同一个对象缓存里的数据自然跟着变。解决办法就是老老实实做一层深拷贝——如果 value 是复杂嵌套对象深拷贝开销会变高这种情况下可以退而求其次明确要求 value 里的数据不可变或者让外部只读取不修改。这个教训让我意识到缓存模块的“边界感”很重要不能图快把内部实现细节暴露出去。7.2 TTL 被每次 get 刷新热点数据永不淘汰我一开始实现的 get 路径在每次读到时都会顺手把expire_at往后推一段。本意是让那些“一直在被使用的 key”多活一会结果造成了一个问题某个 key 被一个高频轮询任务反复 get它的过期时间被不断推后永远不淘汰缓存里堆满了这种高频轮询留下的“僵尸数据”。解决方案很简单跟 MemOS 源码保持一致get 可以更新last_access和access_count但绝不更新expire_at。过期时间只在 set/覆盖的时候重新计算。要延长某个 key 的生命周期业务上应该主动重新 set 它的数据而不是依赖读取操作隐式续期。7.3 大 value 序列化耗时导致主链路超时还有一次我把 Agent 全量会话历史直接塞进 KV Cachekey 是会话 IDvalue 是一整个几十万字符的 JSON 字符串。单次 set 的序列化耗时超过 300msget 反序列化也接近 200ms。这直接导致 Agent 主链路的请求超时。排查后发现两个问题一是大 value 本身就不该整个塞缓存应该拆分成更小的键值对或者只缓存经过摘要处理的精简版本二是 set 之前没有检查单条 value 的大小上限如果源码里有这个限制set 一开始就应该拒绝掉。后来我把大 value 拆成按轮次缓存的多个小块只保留最近几轮的完整内容更早的对话只存摘要请求耗时立刻降了下来。8. 源码之外的一些后续想法这次读 KV Cache 源码读下来我自己最受用的一个习惯是先把数据边界看清楚再动手。一个模块该做什么、不该做什么代码里其实写得很清楚只是很多时候我没有认真去看凭感觉就把所有东西往里面塞。KV Cache 管的是短期数据如果你往里塞长期数据那问题不在缓存在自己。后面我准备接着往下看 MemOS 里跟持久化相关的存储层重点是想搞清楚短期缓存和长期归档之间的数据搬运是怎么做的。如果大家在实际使用 KV Cache 的过程中遇到过其他有意思的坑欢迎一起交流下次笔记还可以把这些真实案例一起整理进去。