Agentrix：面向智能体的内存管理系统设计与实现

Agentrix 面向长上下文、多轮工具调用和并行分支推理，优化工具数据的保存方式、模型计算缓存的保留策略，以及共享上下文的注意力计算。系统通过 vLLM、vLLM-Ascend、SGLang 和 llama.cpp 对接不同模型、操作系统与加速硬件，已在 CUDA 和 Ascend 上开展性能与内存验证，并保留 MUSA、Apple Metal 等后端适配。

项目已实际使用 openEuler 和华为云欧拉操作系统两种国产操作系统，完成 Ascend 推理环境接入及 CUDA 模型运行适配。 国产系统与不同加速硬件通过相应推理后端衔接，具体组合见下文。

智能体执行任务时，往往需要多次处理同一份文档或工具报告，多个分支也会使用相同的背景信息。如果每个分支都复制完整数据、每轮都把全部工具结果交给模型，存储和推理开销会随任务增长。与此同时，已经结束的请求仍可能占用缓存备份空间，使稍后还要继续的会话被迫重新计算。Agentrix 结合应用掌握的任务信息与推理引擎的缓存能力，减少这些重复开销。

核心方案

工具数据共享与按需读取

工具返回的原文保存为内容固定的快照，多个分支通过引用共享数据。局部修改生成新版本，未变化的页面继续共用；当所有相关任务都不再使用某份数据时，再回收其存储空间。

模型先获得工具结果的简要信息和读取入口，需要具体证据时由应用执行读取，将选中的片段加入后续输入。只有实际送入模型的内容才会参与推理并形成计算缓存，从而减少无关工具数据占用上下文。

Agent Hints 与选择性 KV 备份

KV Cache 保存模型已经计算过的上下文状态，用于减少后续推理的重复计算。应用通过 Agent Hints 告知推理服务哪些请求还会继续、哪些上下文可能再次使用。例如，一次性汇总请求结束后不再需要恢复，就可以跳过它的新备份，把有限的 CPU 缓存空间留给后续还要继续的会话。

项目将 LangGraph 的稳定会话身份与应用明确提供的后续使用需求结合，控制新增备份的范围；无法确定是否还会使用时，保留默认行为。缓存的实际存储、淘汰和恢复由推理后端完成。

前缀复用与 ForkAttention

多个请求共享相同上下文时，系统在请求分配、输入处理和生成计算三个阶段减少重复工作。请求分配使用官方前缀亲和路由，让相关请求倾向于访问同一个模型副本；自动前缀缓存复用已经处理过的输入；ForkAttention 面向同时生成的多个分支，在满足实际缓存共享条件时，减少对公共上下文的重复读取。

路由和基础缓存能力沿用上游实现，项目主要补充面向智能体的管理策略与共享计算适配。

初赛基础与决赛演进

初赛围绕多个智能体共享上下文的特点，完成请求分配、缓存管理和 CUDA 共享计算的协同设计，为决赛的工具数据管理、Agent Hints 与异构算子适配提供基础。

前缀亲和路由与分支聚合：初赛自研 Prefix-aware DP，记录各副本已缓存的前缀，结合匹配长度与排队情况分配请求；QueryJoin 短暂聚合同前缀分支，增加共同执行的机会。决赛默认路由采用官方 consistent_hash，由应用提供稳定会话身份，初赛自研路由与合批原型不作为当前默认路径。
CUDA ForkAttention：初赛依据实际缓存页划分公共段与私有段，让多个查询复用读入的公共 KV，再分别合并各分支的计算结果。决赛保留 CUDA 实现，并将共享页识别、分段与合并方法扩展到 Ascend，新增 FIA 分段执行、查询打包、结果合并和图执行工作区复用。
缓存生命周期与工具等待原型：初赛按前缀位置、分支数量和复用情况调整缓存保留顺序，只回收引用计数为零的块；工具等待超过宽限期且缓存紧张时，释放可恢复会话的缓存引用，并保留继续执行所需的输入与位置。决赛默认使用官方缓存与卸载接口，由 Agent Hints 控制新增 CPU 备份的资格和范围；工具等待主动回收仍属于初赛原型。
上下文去重与多框架接入：初赛去除标识与内容一致的重复消息，以读取句柄保留工具原文访问能力，并对接 vLLM、SGLang 和 llama.cpp 的不同缓存布局。决赛进一步增加不可变快照、页面共享、分支引用与阶段回收，将去重扩展为工具数据从产生到释放的完整管理。

初赛核心机制的详细说明见项目说明书附录，决赛收益按各平台的具体实现与对照实验列示。

模型、操作系统与加速硬件支持

系统将应用策略与设备后端分开：工具数据管理通过通用存储和模型服务接口工作，Agent Hints 通过会话标识及后端支持的缓存接口传递需求，缓存格式、搬运方式和算子执行由各推理框架实现。支持范围分别按模型、操作系统和设备说明，具体部署选择与所固定后端版本匹配的组合。

不同模型与参数规模

Qwen3-0.6B、Qwen3-8B：已有环境接入、工具工作流或多框架推理记录，覆盖不同参数规模的全注意力模型。
Qwen3.5-9B：已用于 CUDA 与 Ascend 的混合缓存、选择性备份和容量实验，以及 Ascend ForkAttention 模型对照；恢复时同时处理全注意力 KV 与 GDN 状态。
Llama 等其他后端已支持的模型：通过所选推理框架的模型加载和服务接口接入，复用工具数据与应用层管理策略。模型架构、权重格式、聊天模板及工具调用格式需与后端匹配。

模型能正常推理，与某项优化能否启用分别判断。当前 llama.cpp ForkAttention 面向 Qwen3 的特定 decode 形状，Ascend 在线 ForkAttention 面向已接入的 Qwen3.5 全注意力层；其他模型沿用对应后端的原生注意力路径。选择性 KV 备份还要求该服务已启用兼容的缓存连接器。

国产操作系统与跨系统支持

国产操作系统已用于实际部署和推理验证，覆盖以下组合：

openEuler 24.03 + Ascend 910B1：通过 vLLM-Ascend / CANN 接入 Qwen3-0.6B，验证图执行、前缀缓存与回收、CPU/NPU 数据搬运等能力，见Ascend 环境与接入记录。
华为云欧拉操作系统（Huawei Cloud EulerOS 2.0）+ NVIDIA Tesla T4：完成 llama.cpp CUDA 后端构建与 Qwen3 ForkAttention 运行，见华为云欧拉操作系统适配记录。

Linux 是当前 CUDA 与 Ascend 性能实验的主要环境。除上述国产操作系统外，仓库提供 Ubuntu 上的 CUDA 构建说明，以及 CentOS 7.9 上 llama.cpp + Tesla T4 的构建和运行记录。不同发行版分别匹配编译器、驱动与运行时。各组合的接入记录与性能实验分别列示。

macOS 通过 llama.cpp 的 Apple Metal 后端接入 Apple Silicon，仓库包含相应 ForkAttention 内核与构建说明。Windows 沿用所固定 llama.cpp 版本的原生构建入口接入基础推理；本项目新增优化尚无 Windows 性能对照。这里的跨系统支持按后端提供，不表示 vLLM、CANN 与 Metal 可以在所有操作系统上互换。

不同显卡与 AI 加速卡

设备	对应后端	已有支持与验证范围
NVIDIA H100、RTX 5070	vLLM / CUDA	H100 上的工具工作流、KV 管理实验；RTX 5070 上的注意力算子对照
NVIDIA RTX 4090	SGLang / CUDA	Qwen3-8B ForkAttention 接入与运行记录
NVIDIA Tesla T4	llama.cpp / CUDA	Qwen3 ForkAttention，已有 HCE 与 CentOS 两种系统的运行记录
Ascend 910B1、910B2	vLLM-Ascend / CANN	910B1 环境与推理接入；910B2 上的单卡、双卡缓存及算子实验
摩尔线程 MTT S4000 / QY2	llama.cpp / MUSA	已有 Qwen3 ForkAttention 适配与构建运行说明；不套用 CUDA 或 Ascend 的收益数据
Apple Silicon 集成 GPU	llama.cpp / Metal	已有 Qwen3 ForkAttention 内核和接入说明；按统一内存及 Metal 能力选择执行路径

海光 DCU / DTK 作为后端扩展目标保留在架构中，尚不列为已完成适配。其他设备的接入可沿用应用策略，仍需满足对应框架、驱动、算子精度和缓存布局要求。不同模型或设备之间不直接复用不兼容的 KV。

Ascend 的矩阵计算、向量计算和数据搬运由不同硬件单元承担，共享计算需要结合这些单元的分工重新组织。Ascend ForkAttention 根据分支实际共享的缓存页安排计算，将多个分支的查询打包后交给原生融合注意力算子 FIA 分段执行，再合并各段结果；辅助计算进一步优化查询打包、结果合并和图执行工作区复用，降低共享执行本身的额外开销。

各路径的具体说明见 Ascend 算子适配、SGLang 接入、llama.cpp CUDA / MUSA 接入、HCE / CentOS 适配、Apple Metal 适配和 llama.cpp 跨系统构建。

实验结果

以下为独立对照实验，分别衡量存储占用、会话恢复和模型执行，收益不能相加。

优化与实验条件	关键结果
工具页面共享：8 MiB 原文、八个分支，每分支修改 32 字符	唯一页面载荷从 72 MiB 降至 8.03125 MiB，减少 88.85%；更新耗时从 1.50 s 增至 2.52 s
H100 选择性备份：Qwen3.5-9B，2 GiB 设备缓存 + 2 GiB CPU 缓存，顺序压力恢复	平均恢复首 token 延迟从 460.96 ms 降至 136.90 ms，降低 70.30%
双 Ascend 选择性备份：Qwen3.5-9B，每卡 1 GiB 设备缓存 + 2 GiB CPU 缓存，顺序压力恢复	平均恢复首 token 延迟从 768.68 ms 降至 152.13 ms，降低 80.21%
双 Ascend ForkAttention：Qwen3.5-9B，每卡四并发分支、32K 共享前缀、输出 64 token	整批耗时从 3.072 s 降至 2.898 s，降低 5.67%；双卡显存峰值均值增加 760.5 MiB

恢复实验保持缓存预算相同，反映的是减少重算带来的延迟收益。ForkAttention 的模型实验体现执行时间收益，同时需要额外工作区。完整条件、统计方法与资源成本见实验记录，系统设计见项目说明书。

系统环境与依赖要求

下表列出决赛主要实验所用的 CUDA 与 Ascend 环境；SGLang、llama.cpp、MUSA 与 Metal 分别使用上节链接中的构建配置。

路径	环境与依赖
通用应用与评测	Python 3.10 或更新版本；应用包、Benchmark 及其可选依赖按各自 `pyproject.toml` 安装
CUDA	Linux x86_64、NVIDIA GPU、匹配的驱动与 CUDA Toolkit；建议 Python 3.12，GCC/G++ 11.3 或更新版本；使用子模块中的 vLLM 0.28.0 分支
Ascend	独立运行环境；固定基线为 vLLM 0.22.1 + vLLM-Ascend v0.22.1rc1，配套 CANN、Ops、torch_npu 与 Triton-Ascend；具体版本见核心仓库构建说明
通用构建工具	Git、C/C++ 编译工具、CMake、Python 包管理器；CUDA 源码构建使用 `uv`
模型与数据	模型权重、评测数据另行准备，遵循各自许可证；不随交付仓库复制

CUDA 与 Ascend 使用独立环境，不能将 CUDA 的 vLLM 子模块直接安装到上述 Ascend 基线。详细依赖见核心仓库构建说明与Ascend 架构及适配说明。

获取源码

在本仓库根目录执行：

git submodule sync --recursive
git submodule update --init --recursive
git submodule status --recursive

核心源码位于 src/agentrix/ 子模块，以上命令检出仓库记录的固定版本。

构建与安装

下文以 Linux shell 为例；其他操作系统使用对应后端文档中的构建方式。

通用应用和评测客户端

以下命令从本仓库根目录运行，不要求安装设备推理后端：

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -e './src/agentrix/application[test]'
python -m pip install -e './src/agentrix/benchmark[agent,data,demo,test]'

CUDA 推理后端

准备好 NVIDIA 驱动、CUDA Toolkit 和 uv 后，在单独终端从本仓库根目录运行：

cd src/agentrix/vllm
uv venv --python 3.12 --seed
source .venv/bin/activate
export MAX_JOBS=2
export NVCC_THREADS=1
uv pip install -v -e . --torch-backend=auto

构建目标需匹配运行 GPU；跨设备编译时按目标设备配置 TORCH_CUDA_ARCH_LIST。更完整的构建及增量编译步骤见源码构建指南。

Ascend 推理后端

使用与固定基线匹配的 Ascend 环境，先准备 vLLM 0.22.1 及配套设备依赖，再从本仓库根目录安装插件：

python -m pip install -e ./src/agentrix/vllm-ascend --no-build-isolation

驱动、CANN、Ops 和 Python 扩展需相互匹配；完整平台说明以vLLM-Ascend 文档为准。

运行方法

激活已安装所选平台推理后端的环境，将 MODEL_DIR 设置为本机模型目录，然后启动服务：

vllm serve "${MODEL_DIR}" \
  --served-model-name agentrix \
  --enable-prefix-caching \
  --host 127.0.0.1 --port 8000

这是基础服务入口。ForkAttention、CPU offload 和多卡路由需按对应技术报告单独配置，该命令不表示已启用全部优化。请求示例、工具数据演示和可视化入口见演示说明。

测试方法与验证材料

在本仓库根目录激活通用客户端环境后运行应用和评测逻辑测试：

source .venv/bin/activate
python -m pytest src/agentrix/application/tests
python -m pytest src/agentrix/benchmark/tests

设备算子测试需在对应 CUDA 或 Ascend 环境运行。实验记录直接列出框架层、CUDA、Ascend 的关键配置、对照数据、统计方法、正确性条件与资源成本。

技术文档

离线流程动画演示工具数据共享、Agent Hints 和 Ascend ForkAttention 的执行过程，支持逐步播放与截图，操作方法见演示说明。

许可证

本仓库原创内容采用 Apache License 2.0。核心仓库及其第三方子模块、模型与数据分别遵循对应许可证；工具使用及主要第三方依赖见人工智能及第三方工具使用说明。