AgentMem：面向智能体推理的内存管理优化系统

AgentMem 是复旦大学 CodeWisdom 系统工程组研发的、面向长生命周期多轮智能体负载的推理系统，基于 vLLM 与 vLLM Ascend 构建。系统服务于工具调用、环境交互和反复模型调用中的 Agent 任务，提供确定性轨迹回放、跨轮 KV Cache 复用、Prefill–Decode 分离调度及昇腾 NPU 全链路适配，使持续增长的任务状态能够在 GPU/NPU、CPU 和 SSD 间高效保存、迁移与恢复。

一、系统架构与能力

AgentMem 系统总体架构

系统以 Trajectory Replay、Proxy、Prefill 与 Decode 构成多轮请求主循环。Replay 通过 Token 覆盖的 Teacher-Forcing 生成确定、可复现的多轮负载；Proxy 按轮次路由请求； Prefill 节点完成新增上下文预填充，并将离散 KV Block 经 P2P/HCCL 迁移给 Decode 节点； Decode 节点完成自回归生成，工具调用结果再驱动下一轮请求，形成完整的 Agent 生命周期。

围绕这一主循环，架构中的六个模块并非彼此独立：ShadowScheduler 在工具执行间隔中预测请求未来的 Prefill 准入位置；该时间线一方面指导 Dynamic Chunk 与 Longcap-FCFS 决定长短请求如何安全进入 Prefill，另一方面只在请求进入 N+1 窗口时触发 SSD → CPU KV 预取。随后，命中的 CPU KV 按层加载至 NPU 分页 KV Cache，使数据搬运尽可能与模型 Forward 重叠。Copy-on-Write 则在存储层复用规划、反思和失败恢复分支的公共前缀；昇腾适配模块横切整个数据通路，负责 Attention、CPU–NPU KV 恢复和 HCCL 迁移。

主运行时和配套 Ascend 插件作为一个版本单元维护，使调度器、Attention 算子、HCCL 传输和分层 KV Cache 不会因独立升级而发生接口或行为漂移。

AgentMem 的核心能力包括：

确定性轨迹回放：以 Token 覆盖的 Teacher-Forcing 重放真实多轮 Agent 负载；
未来资源时间线：预测请求的 Chunk 推进、Block 分配/释放和 Prefill 准入时机；
动态 Prefill 调度：以 Dynamic Chunk 和 Longcap-FCFS 缓解 Prefill-only 节点的长短请求干扰；
预测式分层 KV 管理：在 N+1 准入窗口触发 SSD → CPU 预取，正式准入后逐层加载至 NPU；
分支状态共享：以 Copy-on-Write 复用多路径推理的公共 KV 前缀；
昇腾 NPU 优化：提供 Prefix-v2/FIA Attention、CPU–NPU staging 和 HCCL KV 迁移流水线。

二、三个核心问题

1. PD 分离之后，Prefill 节点内部再次出现共址干扰

PD 分离将计算密集的 Prefill 与访存密集、延迟敏感的 Decode 部署到不同设备，解决了两个阶段之间的共址干扰。但今天的 Agent 请求不再是形态近似的普通 Prompt：不同请求的新增 Token 长度、已命中历史 KV 长度、剩余计算量和到达时间都可能相差几个数量级。因此，原先发生在 Prefill 与 Decode 之间的干扰，又以长短 Prefill 请求相互阻塞的形式出现在 Prefill-only 节点内部。

Sarathi 为共址场景提出 Chunked Prefill，通过切分长 Prefill 为 Decode 腾出执行间隙；但将固定长度 Chunk 直接用于 PD 分离的 Prefill-only 节点会产生两个新问题：

固定 Chunk 无法适配 Agent 负载中高度变化的新增计算量与历史 KV 长度；
调度器只看当前空闲 Block，看不见请求未来各 Chunk 的资源占用。它可能在当前 Step 盲目装入一个之后无法容纳的请求，随后发生驱逐、重算或队头阻塞。

AgentMem 的解法：timeline-aware Dynamic Chunk Prefill。 系统模拟未来 Step 中运行请求的计算进度、Block 分配与释放，在准入前判断请求是否能安全走完整条时间线；同时根据队列和资源状态动态选择 Chunk Size，并用 Longcap-FCFS 控制长请求占比。这样既避免固定 Chunk 的僵化，也减少“先装入、后驱逐”造成的抖动和队头阻塞。

在 5 IPS 的调度对比中，longcap + dynamic_chunk 相对 fcfs + nochunk 将请求延迟 P50 从 59.7 s 降至 4.8 s（12.5 倍），P85 从 90.1 s 降至 8.1 s（11.1 倍）；其收益覆盖约前 92% 的请求。与此同时，最重的长请求需要经过更多 Chunk，P95 从 92.9 s 变为 116.6 s。这一结果刻画了设计目标：通过可控的极端长请求代价，消除长请求对大量短请求造成的队头阻塞，而不是把所有请求强行切成同一个固定长度。

Dynamic Chunk 与无切分 FCFS 的请求延迟对比（5 IPS）

2. Agent 推理正在从计算瓶颈转向 KV Cache 搬运瓶颈

Agent 每轮通常只追加少量新 Token，却要复用很长的历史上下文。将 KV Cache 放到 SSD 可以避免重复 Prefill，但也把问题从“重新计算”变成“何时搬运”。2026 年 2 月由 DeepSeek-AI 等机构参与的 DualPath 论文指出，多轮 Agent 推理正逐渐由 KV Cache 存储 I/O 而非计算主导；在分离式架构中，Prefill 侧的存储网络尤其容易成为带宽瓶颈。

如果等到请求真正被准入后才执行 SSD → CPU → NPU 恢复，存储访问会直接进入前台关键路径；如果过早预取，又会浪费 SSD 带宽和 CPU Arena 容量。

AgentMem 的解法：复用调度时间线做预测式 KV 预取。 ShadowScheduler 前瞻 N+1 至 N+5 个调度步；只有请求进入 N+1 准入窗口时，才异步将对应 KV 从 SSD 提升到 CPU Arena。请求真正被准入后，CPU → NPU 传输再按层与模型 Forward 流水重叠。由此， timeline 不仅回答“下一个 Chunk 应该多大”，还回答“哪份 KV 应该现在开始搬”。

在 Llama-3.1-8B 全量实验中，N+1 预测准确率为 **95.88%**、精确率为 **91.24%**、召回率为 **93.07%**；预取 Block 的消费精确率为 **99.19%**，所需 Block 在准入前的存活率为 **99.50%**。这说明系统既能较准地预测未来准入，也很少为了错误目标浪费搬运。

下图展示的是 AgentMem 启用预测式 SSD 预取与 CPU → NPU 逐层流水后的 Prefill Step 关键路径。此时模型 NPU Forward 的 P50/P95 分别约为 4.16/5.30 秒，而仍暴露在前台的 KV 阶段仅约 15/472 毫秒；时间线中绿色 Forward 占据绝大部分关键路径。换言之，原本会在请求准入后串行阻塞的 SSD 恢复与 CPU → NPU 加载，已在 AgentMem 中被提前发起或与 Forward 重叠隐藏，仅剩少量未及时完成预取的恢复操作进入前台。

Prefill 各阶段延迟分位数

Prefill Step 关键路径时间线

3. 通用 CUDA 路径无法直接变成高效的昇腾实现

昇腾适配并不是将接口名称从 CUDA 替换为 NPU。原版 vLLM Ascend 在 Chunked Prefill 场景会先物化 [C, max_model_len] 的巨大临时 Mask，再切出 [C, D+C]；长上下文下该张量可接近 GB 级，并在每层重复读取。同时，原路径使用通用 SplitFuse 算子、 64-Token 逻辑 Block，而昇腾 Prefix-v2 要求 128-Token 物理页和对齐 Prefix，三者在 Mask、Block 布局和动态边界上并不兼容。

CPU 与 NPU 之间也存在平台差异：CUDA Kernel 可以访问注册后的 mapped host memory，但 Ascend 910B3 的 AIV 直接解引用 host pointer 会触发 vector core exception；逐 Block、逐层提交小 DMA 又会产生大量 runtime 开销。

AgentMem 的解法：算子、布局、调度和传输协同适配。

直接构造有效形状 [C, D+C]，消除 [C, max_model_len] 临时 Mask；
保留 64-Token Kernel 视图，在算子入口恢复 128-Token 物理 KV 视图和 Block Table；
对齐场景走 Ascend 原生 Prefix-v2，非对齐动态边界回退 FIA Paged Prompt；
调度器主动保证非末尾 Chunk 的 (D+C) % 128 == 0，提高 Prefix-v2 命中率；
使用有界 pinned CPU/NPU staging、AscendC fused scatter/gather 和双缓冲流水完成 CPU ↔ NPU KV 搬运，不再尝试不受硬件支持的 AIV 直读 Host 内存。

最终，长上下文 Chunked Prefill 从原始 Mask + SplitFuse 路径到 Prefix-v2 路径的模型 Step 获得约 5.4–7.2 倍加速；单独的 Attention Kernel 获得约 1.9–2.0 倍 加速。16 Blocks、32 Layers 的传输微基准中，tiled staged load 为 0.91 ms，而 1024 次小 ACL DMA load 为 12.81 ms。

三、从三个问题到一条统一时间线

复杂 Agent 请求到达
        │
        ▼
未来资源时间线：预测每个请求的 Chunk 推进、Block 分配/释放和准入 Step
        │
        ├── 计算侧：Dynamic Chunk + Longcap-FCFS
        │           └── 避免盲目装填、驱逐和队头阻塞
        │
        └── 存储侧：N+1 至 N+5 Shadow Admission
                    └── N+1 触发 SSD → CPU，准入后逐层 CPU → NPU
                                      │
                                      ▼
                       Ascend Prefix-v2/FIA + KV 传输流水线

这条时间线是 AgentMem 区别于普通分层缓存或固定 Chunk 系统的核心：调度器不再只对 “当前能否放下”作局部判断，而是将未来计算进度变成存储预取和异构执行的共同依据。

四、支撑模块

1. 可复现的轨迹回放评测

AgentMem 通过基于 Token 覆盖的 Teacher-Forcing 机制精确重放真实智能体轨迹，固定多轮请求的 Token 序列和轮间时序。这样可以在相同负载下比较不同调度和内存管理方案，并保证回放产生的 KV Cache 与对应 Prefill 过程保持一致。

2. 预测驱动的 KV Cache 生命周期管理

系统使用 ShadowScheduler 模拟未来调度状态，前瞻请求在 N+1 至 N+5 个调度步内的准入机会。预测结果用于指导缓存的保留、迁移和释放；只有当 SSD 中的目标状态进入 N+1 窗口时，系统才触发 SSD 到 CPU Arena 的异步提升，减少过早预取和无效搬运。

3. GPU/NPU—CPU—SSD 三级分层存储

系统按照访问热度管理 KV Cache：

GPU/NPU 保存当前计算所需的热数据；
CPU Arena 保存近期可能复用的温数据；
SSD 保存容量较大、访问频率较低的冷数据。

缓存支持跨层迁移、按需恢复、内容寻址去重和跨轮增量复用。当优化路径失败时，请求可退化为常规重算，不影响推理正确性。

4. Copy-on-Write 树状状态共享

针对智能体规划、反思和失败恢复产生的多条执行分支，AgentMem 在存储层维护状态树。不同分支共享公共历史前缀，仅在状态发生修改时执行延迟复制，从而降低重复 KV Cache 的存储和搬运开销。

5. Timeline-aware 动态预填充调度

在 Prefill—Decode 分离架构中，AgentMem 结合两项调度机制：

Longcap-FCFS：按剩余 Prefill 长度对请求分带准入，控制长请求占用比例；
Dynamic Chunk：根据运行时队列状态动态选择 Chunk Size，使短请求能够在长请求的 Chunk 间隙中获得调度机会。

该设计重点改善长短请求共址时的资源竞争和尾延迟。

6. 昇腾 NPU 全链路适配

项目面向 openEuler 和昇腾 910B3 完成了关键数据通路适配，包括：

Chunk Prefill Attention 的渐进式优化与 Prefix Attention 快速路径；
CPU—NPU KV Cache 逐层 staging 传输；
ATB Attention 动态形状稳定性修复；
基于 HCCL 的 Prefill—Decode KV Cache 迁移流水线。

五、实验效果

报告中的主要实验使用 4 张昇腾 910B3 NPU、1P3D 部署、1999 条混合轨迹和 Poisson 5 JPS 负载，并在 Qwen3-4B 与 Meta-Llama-3.1-8B-Instruct 上分别与 4 个独立 vLLM 实例的基线对比：

指标	Qwen3-4B：基线 → AgentMem	变化	Llama-3.1-8B：基线 → AgentMem	变化
任务成功率	97.75% → 100.00%	+2.25 个百分点	98.20% → 100.00%	+1.80 个百分点
完成轮数	11,942 → 12,740	完成全部请求	12,055 → 12,740	完成全部请求
请求吞吐	1.213 → 2.075 req/s	+71.06%	1.073 → 2.208 req/s	+105.78%
Token 吞吐	213.97 → 375.64 tok/s	+75.56%	188.44 → 395.06 tok/s	+109.65%
Scheduler Prefix Hit	0.25% → 92.85%	超过 371 倍	0.24% → 92.91%	超过 394 倍
请求 E2E P95	275.90 → 94.18 s	-65.86%	278.76 → 93.69 s	-66.39%
请求 E2E P99	468.48 → 165.02 s	-64.78%	492.56 → 138.29 s	-71.92%
ITL P99	1.923 → 0.127 s	-93.40%	1.356 → 0.103 s	-92.40%

两种模型均在相同实验负载下完成全部 12,740 轮请求，并显著改善吞吐和尾延迟。需要注意的是，1P3D 将 Prefill 集中到单张 NPU，因此两种模型的 TTFT、排队和 Prefill P50 相比 4P 对等基线有所增加；该架构以可控的中位延迟代价换取更高吞吐、完整任务成功率和更稳定的 P95/P99 表现。

两种模型均完成全部 1,999 条轨迹；图中展示任务成功率和轮次完成率的双模型对比。

两模型任务成功率与轮次完成率

跨轮 KV 复用、PD 分离和动态调度共同提升了请求与 Token 吞吐：

两模型请求与 Token 吞吐

两种模型的 P95 延迟也呈现一致趋势：尽管 Prefill-only 拓扑会增加部分中位延迟，但排队、Decode、端到端延迟和 ITL 的尾部指标均显著下降。

两模型 P95 延迟对比

下图再以 Llama-3.1-8B 展示全分位延迟，便于观察 P50 与 P95/P99 之间的具体权衡：部分 Prefill 侧 P50 因 1P3D 单 Prefill 节点集中化而上升，但排队、Decode、端到端延迟和 ITL 的 P95/P99 均显著下降。

Llama-3.1-8B 全分位延迟对比

六、系统工作流程

一次多轮请求主要经过以下阶段：

轨迹回放客户端产生确定、可复现的多轮负载；
Proxy 将请求路由至 Prefill 节点；
Longcap-FCFS 和 Dynamic Chunk 完成请求准入与动态切分；
Prefill 生成的 KV Cache 通过 HCCL 迁移至 Decode 节点；
工具执行期间，ShadowScheduler 预测后续请求的准入时机；
KV Cache 在 NPU、CPU Arena 和 SSD 之间按预测结果迁移，并在下一轮增量复用；
任务结束后统一回收各层缓存资源。

当任务产生规划分支或失败恢复路径时，Copy-on-Write 状态树在上述主流程中按需启用。

七、仓库结构

agentmem_exp/       轨迹回放、结果分析与可视化工具
dynamic_chunk/      BRB 校准与昇腾 Prefill 性能分析
examples/           在线服务和 P/D 分离部署入口
tool_call_time/     工具调用耗时预测模型及运行包
trace_data/mixed/   随仓库交付的默认混合轨迹评测集（1999 条去重轨迹）
vllm/               集成 AgentMem 调度与 KV 管理的 vLLM 运行时
third_party/
  vllm-ascend/      与主运行时版本匹配的 Ascend 插件

八、安装与部署

已验证环境

组件	版本或配置
操作系统	openEuler 22.03 LTS-SP4，Linux aarch64
加速设备	Ascend 910B3，单卡 64 GiB HBM
Python	3.10.16
CANN	8.2 RC1
PyTorch	2.7.1
torch-npu	2.7.1.dev20250724
vLLM Ascend	0.10.2rc1，本仓库配套修改版本

在仓库根目录执行：

SOC_VERSION=Ascend910B3 bash tools/install_agentmem_ascend.sh

安装脚本会安装当前仓库中的 vLLM 运行时和配套 Ascend 插件，并构建、检查以下本地扩展：

vllm_ascend_C
licht_arena_atomic
licht_fused_scatter_npu

Python 发行包名称仍保留为 vllm 和 vllm-ascend，用于兼容上游插件接口。不要让 pip 的 build isolation 自动替换已经与 CANN 配套的 PyTorch/torch-npu 版本。完整的资源、端口、目录和验收要求见昇腾单机部署指南。

九、启动服务

生产环境入口如下：

bash examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/\
run_disagg_p2p_nccl_xpyd_prod.sh --help

当前优化后的 1P3D 配置使用启动器的 --licht-v3 和 --prefill-opt 开关。这两个名称是为兼容既有内部运行时接口而保留的标识，系统和仓库统一称为 AgentMem。

十、轨迹回放与结果分析

实验工具位于 agentmem_exp/。仓库已随附默认数据集 trace_data/mixed/，评委无需额外下载即可直接使用客户端默认路径。该数据集由 LMSYS Chat（600）、WildChat（600）、 MTRAG（300）、SWE-bench（300）和 LongBench（200）组成；加载时会去除 1 条重复轨迹，最终包含 1,999 条轨迹。使用方式如下：

# 查看轨迹回放客户端参数
python agentmem_exp/multiturn_trace_client.py --help

# 分析实验结果
python agentmem_exp/analyze.py --help

# 生成可视化结果
python agentmem_exp/visualize_results.py --help

运行结果写入每次实验指定的 RUN_OUTPUT_DIR。时间戳目录、日志、图表、模型缓存以及 CPU/SSD Arena 等运行时数据不应提交到仓库。

十一、设计文档

十二、上游项目与许可证

AgentMem 的主运行时基于 vLLM，仓库内的设备插件基于 vLLM Ascend。项目保留了对应的 Apache-2.0 许可证和第三方声明，详见：