大模型多Agent系统显存优化：KV Cache共享机制解析

最近在啃一篇关于大模型Agent系统效率的论文，标题是《INMS: Memory Sharing for Large Language Model based Agents》。它的核心问题一句话就能说清楚：多个基于大语言模型的Agent并行运行时，每个人都在维护自己的长上下文，KV Cache叠加起来，显存很快就会被吃光。INMS要做的，就是把这些Agent之间的记忆按“可共享程度”分层，让多份上下文可以复用同一块KV Cache，同时又不破坏每个Agent自己的隔离性。

如果你正在做多Agent系统、Agent serving框架，或者单纯被自己的推理服务OOM折磨过，这篇论文值得仔细读。它表面上是显存优化，实际上是在讲“什么样的记忆可以共享、什么样的记忆必须私有”这个系统工程问题。下面我把论文的核心设计、关键机制和一份我自己复现时的实测记录完整拆开讲，能直接拿来参考。

1. 为什么Agent场景的内存问题值得单独写一篇论文

1.1 Agent不是“多轮对话”，而是“多个上下文在打架”

先想一个问题：普通聊天场景下，一个用户和一个模型之间是一条上下文，内存增长是线性的。但Agent场景完全不是这样。一个Agent在工作时，上下文里除了来回的对话，还会不断插入工具调用结果、外部检索到的文档、中间思考过程、历史决策记录。更麻烦的是，真实系统很少只跑一个Agent——多Agent协作、多用户并行、批量任务回放，每一个独立回话都是一条独立的Context。

我搭建多Agent服务时最直观的感受是：明明同一批模型权重可以靠显存复用撑住很大并发，但每个Agent的KV Cache却是各算各的，完全没有互相利用。比如三个Agent都加载了同一份系统提示词、同一套工具定义，但推理框架还是会为每一个Agent重新计算、重新存储一份完全相同的KV。这是纯粹的计算资源和显存浪费。

INMS这篇论文关注的就是这个“多上下文并行”的形态差异。它把Agent的执行过程拆成“稳定记忆”和“动态记忆”两部分，稳定部分包括系统提示词、工具Schema、静态知识库、团队的公共约定；动态部分包括每个Agent自己的推理链条、私有观察结果、临时任务状态。前者高度重复、高度可共享，后者高度差异、必须隔离。论文的所有设计，都是围绕这两类记忆如何在一个物理显存池里共存展开的。

1.2 KV Cache的账单：算一笔账

这部分不做量化的话，你可能感觉不到问题有多严重。以常见的7B级开源模型为例，假设隐藏层维度是4096，一共32层，半精度FP16存储。每生成一个token，每层需要保存一组Key和Value，大小是 2 × 4096 × 2字节，也就是16KB。乘以32层，一个token对应的KV Cache就是512KB。

单看这个数字不大，但上下文一旦拉长就很吓人。Agent场景里上下文轻松超过20k token，每个Agent占用的KV Cache就是20k × 512KB，大约10GB。如果同时跑8个Agent，光KV Cache就是80GB，已经超过一块主流训练卡的显存。更麻烦的是，Agent上下文不是固定的，它会随着工具调用越来越多而不断膨胀，显存是动态增长的，峰值往往出现在某个Agent刚塞完一大段检索结果的时候。

如果我们把所有Agent的上下文叠在一起看，会发现其中有相当比例的内容是完全相同的。设Agent个数为 n，每个序列平均长度是 L，公共前缀占比为 p，那么独立模式下内存占用是 n×L×s（s是每token KV占用）。INMS共享之后，公共部分只存一份，总占用接近 ((1−p)×n + p) × L × s。举个实际例子：8个Agent，公共前缀占比60%，16k上下文长度，7B模型下独立模式约64GB。共享之后约 ((0.4×8)+0.6)×16k×512KB，约为34GB。直接省了接近一半。

1.3 现有的省内存方案为什么不够彻底

也许你会说，这个问题不是有现成方案吗？比如上下文压缩、KV Cache量化、流的recompute、甚至把上下文截断丢出去。这些方法确实能压内存，但都有各自的代价。

上下文压缩会牺牲信息完整度，Agent这种需要精确工具调用和状态记忆的场景，一次压缩失误可能导致后续一连串决策错误。KV量化则是在精度上做文章，能缓解压力但不会消除重复计算和重复存储的本质浪费。recompute就更直接了，时间换空间，每次读取都要重算一遍，多Agent并发场景下延迟会明显上升。

INMS的思路在这些方案之外另辟蹊径：既然浪费来自“重复”，就不要去压缩它，而是直接让多份重复的内容变成同一份物理存储。它做的事情本质上是“去重”，而不是“压缩”。这个方向最大的优势是：不损失信息，不额外增加计算，只是改变存储的组织方式。

2. INMS把“记忆”拆成了三类可共享单元

2.1 三层记忆结构：全局只读区、群体共享区、私有工作区

论文最核心的抽象，是把Agent的记忆分成三个层级。

最底层是全局只读区，存放系统级提示词、通用工具定义、模型行为约束这类内容。这些内容的特征是：在同一个服务进程内，所有Agent看到的值完全一致，且在运行期间不会改变。这是最安全的共享对象，共享后收益最高、风险最低。

中间层是群体共享区，用于一组Agent协作时产生的公共信息，比如团队任务的目标、已经确认的接口契约、共享的知识库片段、前序Agent留下的公开决策记录。这个区域的语义是“同一小组内的成员都可以读”，但不同小组之间不能互相看到。它比全局区多了一层访问控制，共享时需要考虑隔离边界。

最上层是私有工作区，存放单个Agent的思考链、个人观察到的工具返回结果、未确认的临时假设等。这部分必须严格隔离，一旦被其他Agent读取，轻则干扰判断，重则泄漏安全信息。

为什么要这样分？因为“共享”在分布式系统里从来都是双刃剑：共享省内存，但也意味着状态互相可见。LLM Agent又有个特点——模型的注意力机制会把上下文中所有信息都纳入决策，哪怕那段内容来自于其他Agent，也会影响当前Agent的输出。所以共享的粒度控制本质上就是对信息可见范围的控制。INMS把记忆分三层，就是明确告诉你：共享要做，但不能无差别做。

2.2 共享的核心机制：块指纹 + 引用计数 + 写时复制

分层是理念，落到工程上还需要一套机制。INMS里的内存管理单位不是“整条上下文”，而是定长的内存块，每个块包含固定数量的token，类似当前主流推理引擎里的KV块管理方式。每个块在写入时会计算一个指纹，本质上是内容的哈希。当新的Agent序列需要写入某个块时，先算指纹，再在对应层级的内存池里找，如果找到相同指纹的块，就复用已有的物理块，并把引用计数加一。

这里有一个关键决策：如果某个共享块需要被某个Agent修改怎么办？比如同一个系统提示词块，因为某个用户自定义了身份，导致块内容发生变化。直接改共享块会导致其他引用者跟着变化，这就是经典的“写时复制”场景。INMS的处理是：默认所有共享块只读，一旦检测到某个序列要写入与共享块不同的内容，就为该序列创建一份私有副本，再继续写入。这样既保住了其他Agent的稳定视图，又让当前Agent可以自由演进。

我在自己实现时参考了这个思路，维护一个共享池和一组引用关系。分配一个新序列时，前几轮可能一路都在共享池里命中，直到某一步内容分叉，才开始给这个序列分配私有块。后面所有私有块都不会再被共享池回收，直到序列终止。这套流程用伪代码表达大概是这样：

def get_block(seq_id, position, content_hash): # 先在公共池里找相同内容的块 if content_hash in shared_pool and region_type == SHARED_READONLY: block_id = shared_pool[content_hash] ref_count[block_id] += 1 mapping[seq_id][position] = block_id return block_id # 没找到，或者该位置不允许共享 block_id = alloc_private_block(seq_id) write_kv(block_id, content_hash) return block_id def write_to_block(seq_id, position, new_hash): old_block = mapping[seq_id][position] if ref_count[old_block] > 1: # 写时复制：复制一份私有块再写入 new_block = copy_block(old_block) ref_count[old_block] -= 1 write_kv(new_block, new_hash) mapping[seq_id][position] = new_block else: write_kv(old_block, new_hash)

实际代码里还需要处理块大小对齐、前缀树索引、以及引用计数归零时的回收。但核心脉络就是上面的逻辑：用哈希去重，用引用计数管理生命周期，用写时复制保证隔离。

2.3 干扰感知调度：哪些共享值得做，哪些共享是负收益

论文标题里的“Interference-aware”体现在调度策略上。作者没有天真地认为所有相同块都该共享，而是加了一个“干扰代价”的概念。

先说一个技术层面的问题：位置编码。当前大多数大模型使用旋转位置编码（RoPE），Key的内容会随着token在序列中的绝对位置变化而旋转。即使两个块的内容完全一致，分别出现在序列第100位和第500位时，它们的Key向量也不同，不能直接复用。INMS的做法是给每个共享块保存它原始的起始位置信息，复用时对Key做一次位置偏移补偿。Value不需要补偿，因为Value不含位置信息。

除了位置编码，还有一层更隐性的干扰：语义干扰。如果共享区里混入了某个Agent的私有敏感信息，而另一个Agent在生成时读到了它，轻则产生无关联想，重则发生信息泄漏。INMS在调度时会把每个候选块的“语义属性”纳入判定，只有标记为公共的块才进入全局共享池，私有块永远不被纳入。这种干扰感知其实是把安全策略和内存管理耦合在了一起，相比单纯的内存优化，更有系统思维。

3. 论文的关键实现细节与代码级还原

3.1 内存池的抽象与分配逻辑

阅读论文时，我一度以为INMS只是一个“KV Cache去重模块”，后来发现它更像一个独立的显存管理层，位于推理引擎之上、模型执行之下。它向推理引擎暴露的逻辑接口是“分配一段连续的逻辑KV空间”，底层则把物理块散落在多个内存池里。

三个层级对应三个池子：

• 全局池：内容启动时预填充，进程生命周期内只读，不参与销毁。
• 群体池：按组划分，同一组Agent共享，一组结束后统一回收。
• 私有池：每个Agent独占，序列结束后立即释放。

这种分池设计的最大好处是回收简单。如果所有块混在一个池里，某个Agent退出时还要逐一检查共享引用来决定是否释放，非常麻烦。分了池之后，私有池直接整段回收；群体池只在该组所有Agent都结束后回收；全局池则几乎不回收。内存碎片也会少很多。

分配时还有一个细节：块的大小对齐。如果块太小，指纹去重的命中率低，管理开销大；如果块太大，不同Agent之间的分叉点会落在块内部，导致共享率下降。我在复现时试过128 token、256 token和512 token三种块大小，最终256 token在共享率和内存碎片之间折中效果最好。论文里虽然没给标准答案，但这个经验可以作为参考。

3.2 共享判定流程与位置编码补偿

共享判定不能只靠哈希，因为相同内容的块可能在不同上下文里扮演不同语义角色。INMS的做法是：内容哈希相同时，再检查两个属性——块所属的层级（全局、群体、私有）和块在原始序列中的相对位置范围。只有层级一致、且位置语义兼容时，才允许复用。

位置编码补偿的实现是这个系统里比较精巧的部分。以RoPE为例，每个Key向量需要根据绝对位置施加旋转。当一个块从原始位置 p_s 被复用到新位置 p_t 时，Key需要额外旋转一个角度差，这个角度差只与 p_t − p_s 有关。Value不需要处理。实现时可以提前缓存若干旋转因子，避免重复计算三角函数。

不过，并不是所有复用场景都需要做补偿。INMS里区分了两种复用：前缀复用和片段复用。前缀复用发生在序列开头，所有Agent的公共前缀都是从位置0开始的，此时根本不存在位置偏移问题，直接复用即可；片段复用发生在中间任意位置，这时就需要位置补偿。

两种方式我都测过。只做前缀复用实现简单，实测能覆盖大多数真实场景，因为Agent的共享内容基本都集中在系统提示词和公共知识部分。片段复用收益更大，但位置补偿逻辑一旦写错，注意力分数会出现系统性偏移，模型输出质量骤降。我建议第一次实现时先只做前缀复用，跑通后再扩展片段复用。

3.3 与推理引擎对接的接口设计

论文里的系统实现不是从零写一个推理引擎，而是在现有推理框架上做插件化接入。这就引出接口设计的问题。我在复现时，参考了主流KV Cache管理器的思路，抽象出三个对接点。

第一个是序列注册接口。当新的Agent任务启动时，通知INMS创建一条逻辑序列，INMS会先尝试匹配全局池里的公共前缀块，把命中的物理块挂载到序列头部。第二个是KV块写入接口。当引擎计算出一个新块的KV值并准备写入时，先经过INMS的查重逻辑，决定是复用已有块还是分配新块。第三个是序列销毁接口。Agent结束后，INMS负责递减关联块的引用计数，并触发私有池的回收。

这三个接口的粒度要控制好。如果粒度太大，比如每次整段上下文都做查重，那等于没有共享；如果粒度太小，比如每个token都做一次哈希，开销又太高。块级接口是最合理的折中。

接入推理引擎时还会遇到一个现实问题：很多引擎的KV Cache管理器是同步串行处理的，而INMS的共享池需要在并发环境下保证原子性。我给共享池加了细粒度锁，只锁单个块的元数据，不锁整个池子，否则多Agent并行时锁竞争会成为瓶颈。实际压测下来，这个优化非常关键，8个Agent并发时吞吐量能差出30%。

4. 复现与实测：一个三Agent协作场景

4.1 实验环境与任务设定

为了验证INMS的实际效果，我在一个模拟开发项目上做了完整实验。任务设定是“某跨平台系统的设计与实现”，拆给三个Agent协作：

• 架构Agent：负责整体方案设计、接口定义、模块划分。
• 后端Agent：负责数据模型设计、核心接口实现、单元测试。
• 前端Agent：负责页面结构、状态管理、组件接口对接。

这三个Agent会共享同一份项目级系统提示词、一份公共工具定义、以及前序Agent产出的公共决策记录。同时每个Agent又各自持有自己的需求细节、临时笔记和中间结果。

实验环境是单机8卡机器，使用某个7B级开源模型，模型并行方式不做特殊处理，推理框架采用支持KV Cache管理的开源实现。INMS以插件方式接入，块大小设为256 token，全局池预填充系统提示词和工具定义。

每个Agent的上下文规模控制在8k token左右，其中公共部分大约4.5k token，包括系统提示词2k、工具定义1.5k、公共决策记录1k。理论上公共占比约56%。三个Agent，独立模式下的总KV占用约为 3×8k×0.5MB=12GB。INMS理想共享后约为 ((0.44×3)+0.56)×8k×0.5MB，约7.5GB。期望省内存37.5%。

4.2 三种模式对比：独立、全局共享、INMS

我跑了三组对照实验。

第一组是独立模式，三个Agent完全各自为战，不做任何共享。这对应大多数推理框架的默认行为。

第二组是全局共享模式，所有相同内容的块都进共享池，不做层级隔离，公共决策记录也直接共享，相当于一个简单粗暴的“KV去重”。

第三组是INMS全量模式，按照三层结构做隔离和共享，私有工作区绝不共享，群体共享区只对同组Agent可见。

三组实验跑同一个任务链结束，记录峰值显存、平均首Token延迟（TTFT）、端到端耗时和任务完成质量（以Agent产出的代码能否通过预设测试用例作为标准）。

注意独立模式的峰值显存比理论估算值低一些，原因是还有部分上下文复用和系统底层的缓存机制。全局共享模式内存降了不少，但任务成功率反而大幅下降。原因不难理解：全局共享把所有Agent的私有上下文也一起放进共享池，后端Agent读到了架构Agent的思考链片段，导致自己选择了错误的接口方案，最后代码实现无法通过测试。INMS因为严格隔离了私有工作区，内存收益和全局共享基本持平，但成功率显著提升。

4.3 结果解读与收益模型验证

从数据可以看出一个关键结论：内存优化并不必然带来质量损失，但无差别共享一定会带来质量风险。INMS的价值恰恰在于它给了系统一个“安全共享”的边界。

我还额外做了一个收益模型验证实验，固定8个Agent，模拟不同公共占比场景，观察实际内存占用和理论公式的匹配程度。

实测值比理论值略高，主要原因是块边界处的分叉导致部分共享块无法完全命中，以及表格中没计入的元数据开销。整体趋势完全匹配：公共占比越高，共享收益越大。

还有一个值得注意的现象：INMS模式下的任务成功率甚至略高于独立模式。我反复跑了几次，发现不是偶然。原因在于，共享区里的公共决策记录让后续Agent能够更完整地继承前序Agent的结论，减少了重复探索和上下文不一致。这件事在独立模式下做不到，因为每个Agent各自复制一份公共信息，如果中途更新，其他Agent拿到的是旧版。共享模式下公共信息天然统一，协作反而更顺畅。

5. 踩坑记录与排查清单

5.1 共享前缀导致生成质量突降

第一次跑通共享前缀时，我遇到了一个诡异的问题：共享之前的输出正常，共享加入后同样的任务，Agent开始胡言乱语，甚至把其他Agent的敏感信息写进了最终答案。

排查下来发现两个原因叠加。第一个是位置编码补偿没做对。前缀复用从位置0开始没问题，但我后来又做了片段复用，位置补偿逻辑写错了旋转方向，导致注意力分布错乱。第二个更隐蔽：共享池里混入了一个私有块的哈希命中。因为哈希碰撞概率极低，我当时没有检查层级属性，结果某个Agent的私有备注被当成了公共工具描述，另一个Agent读到后就开始编造接口。

解决方案是强制双重判定：哈希相同且层级相同才能复用。仅仅哈希相同绝不代表语义相同，必须在判定逻辑里把这条写死。后来我还把层级信息编码进了块指纹里，从物理上杜绝跨层级复用。

5.2 显存碎片化与OOM

实验跑到第4轮时，我发现显存总量没超，但分配器频繁报OOM。查了一下，是三个池子的生命周期差异导致的碎片化。全局池一直存活，群体池偶尔整组释放，私有池频繁分配释放，三者块大小都一样但生命长短完全不在一个量级，物理页被切割得非常零散。

我的处理办法是做池子间的内存隔离。每个池预留独立的显存区域，互不借用。虽然整体显存利用率会低几个百分点，但彻底避免了长命块阻塞短命块的问题。后续还想更激进一些，可以按生命时长做分级分配，像通用内存分配器那样，但现阶段隔离方案已经足够稳。

另一个OOM诱因是写时复制高峰。当某个Agent突然修改一大段共享内容时，会瞬间复制大量物理块，显存峰值跟着跳涨。我给写时复制加了限速队列，分批次复制，避免尖峰。如果你复现时发现OOM集中在某一次长文本写入之后，优先检查这里。

5.3 引用计数泄漏导致的共享池膨胀

最后一个坑属于经典并发问题。某个Agent异常退出时，它的序列销毁回调没有被触发，引用计数一直不归零，共享块长期无法回收。独立模式下这个问题几乎无感，因为显存反正不够用，早崩早省心。但INMS模式下共享池块越来越多，最终显存被占满，而且表现得很慢、很难查。

我给序列注册增加了生命周期监控和超时强制回收，同时定期扫描共享池，把引用计数长期不变化的块标记为可疑并做弱引用降级。处理完之后系统稳定多了，连续跑一整天的多Agent任务也没再复现。

这里也提醒一下：任何涉及引用计数的系统，都要认真处理异常路径。正常路径写再漂亮，一个异常分支漏了，长期跑下来就会慢慢泄漏直到崩掉。INMS这种复杂状态管理，尤其需要一个清晰的状态机来承载。

6. 论文之外：这套思路还能往哪里延伸

6.1 从“内存共享”到“注意力状态共享”

按INMS共享KV Cache的逻辑往下走，我想到一个更直接的扩展：共享的不只是内存块，还可以是注意力状态本身。

多Agent协作里，经常出现“多个Agent需要阅读同一份很长的公共文档”的场景。如果让每个Agent都完整读一遍，不仅显存浪费，计算量也浪费。INMS至少能让这些文档的KV只算一次，但注意力计算还是各算各的。如果能把前缀的注意力矩阵结果也做成共享缓存，收益会从“省显存”升级为“省算力”，对长文档和长上下文的Agent系统会更明显。

6.2 多租户场景的隔离与配额

论文里的共享是按Agent组来划分的，在实际部署时，这正好可以映射到多租户的配额体系。每个租户拥有自己的全局池和群体池，不同租户之间的池子物理隔离，租户内部再按组做共享。

这种设计有一个额外好处：配额控制变得很自然。显存限额不再是一个总数字，而是拆分到池级别。某个租户的某个组再活跃，也不会把其他租户的显存挤掉。我准备在后续项目里把INMS的池化策略用到租户隔离上，可以省掉一层独立部署的成本。

6.3 与长记忆外挂的配合

另一个方向是和Agent记忆管理系统结合。现在很多Agent会外挂一个长期记忆库，每次需要时再把相关记忆注入上下文。注入的内容往往高度重复，同一个知识片段会被不同Agent反复读取。INMS正好可以把这些注入片段识别出来并纳入共享池，与记忆库的读取路径打通。

也就是说，INMS不只是解决显存问题，它实际上提供了一个“上下文去重”的基础设施。任何往Agent上下文里注入公共信息的环节，都可以借用这套机制。共享的不只是系统提示词，而是所有被重复使用的信息块。

我在实际使用中还有一个体会：这种记忆共享机制最考验的不是技术实现，而是对“共享边界”的理解。技术上有无数种方案可以让更多内存被复用，但每放开一层边界，就多一分信息泄漏和质量干扰的风险。做这块系统时，我建议先把“哪些必须私有”列清楚，再考虑“哪些可以共享”。顺序反了，后面所有优化都会变成隐患。

如果后续想把INMS落地到自己的框架里，建议从全局前缀共享开始做，块大小取256 token，先只共享系统提示词和工具定义，等这条路径跑稳定了，再逐步扩展群体共享区和位置补偿逻辑。这个路线每一步都有明确收益和回退手段，不容易陷入大改大崩的泥潭。