AgentMem — 面向智能体推理的内存管理系统

基于 vLLM PagedAttention 扩展，针对 LLM 智能体推理的 KV Cache 生命周期管理、分支推理共享、上下文压缩与分层存储优化。

AgentMem 是面向「操作系统创新赛题——智能体推理内存管理」的开源实现。它复用高性能内存工程模式（64 路分片、对象池、SeqLock、双缓冲快照、CRC32 校验、LRU/LFU 淘汰），并在 vLLM PagedAttention 基础上构建面向智能体长生命周期推理的显存优化系统。

✨ 核心特性

模块	设计章节	核心机制
KV Cache 生命周期管理器	5.2.1	4 级生命周期标签 (L0-L3) + LFU+生命周期加权淘汰 + RadixTree 前缀复用
分支推理 COW 引擎	5.2.2	PagedAttention Block 引用计数 + Copy-on-Write
上下文压缩引擎	5.2.3	重叠检测（RadixTree 前缀匹配）+ 结构化摘要（JSON/表格/文本）
工具调用数据优化	5.2.3	选择性摘要 + OffloadStore 按需加载
分层存储引擎	5.2.4	GPU(热) → CPU(温) → SSD(冷) 三级 Swap + 异步预取 + CRC32 校验
任务调度与资源分配	5.2.5	64 路分片管理 + 公平配额调度 + 优先级队列
硬件抽象层 (HAL)	5.2.6	统一 GPU 操作接口 + 运行时硬件检测 + CUDA/DTK/CANN/CPU 四后端

🚀 性能与创新点优化（v1.1+）

在原有 6 大模块基础上，针对比赛场景进一步优化（详见 docs/performance_innovations.md）：

编号	类别	名称	默认
P2	性能	EvictionPolicy 持久化版本化堆（O(n)→O(k log n) 淘汰候选选取）	开
P3	性能	原生 C 扩展热路径加速（RadixTree 前缀匹配 + 批量评分，3-10×）	可选编译
I1	创新	KV Cache int8 per-channel 量化压缩（CPU/Disk 占用降低 ~74%）	配置开
I2	创新	Markov 预测性预取（融合 RadixTree 结构预测，S8 实测命中率 40%）	开

编号说明：此处 P2/P3 为 v1.2 创新点编号（Performance 优化），与设计方案 2.2 节的性能指标编号 P1-P6 含义不同。指标 P1-P6 见上文「性能目标」表。

📊 性能目标

指标	目标	说明
显存占用降低	≥ 30%（对比原生 vLLM）	智能体多轮对话场景
推理延迟增加	≤ 10%	显存优化前提下控制延迟开销
任务成功率	≥ 95%	优化后不显著降低任务完成质量
Swap 延迟	P99 < 50ms	KV Cache 从 CPU 换入 GPU 的延迟
分支推理显存节省	≥ 40%	4 分支 Tree-of-Thought 场景
Prompt 压缩率	≥ 20% token 减少	system prompt + 工具描述场景

📦 安装

环境要求

Python >= 3.10
操作系统：openEuler 22.03 LTS（主要适配）/ Ubuntu 22.04 / openKylin 1.0
可选：NVIDIA GPU + CUDA >= 11.8（CUDA 后端）/ 海光 DCU（DTK）/ 华为昇腾 NPU（CANN）

安装方式

# 1. 核心包（纯 CPU 即可运行，含完整内存管理逻辑）
pip install -e .

# 2. 开发与测试依赖
pip install -e ".[dev]"

# 3. GPU / vLLM 集成（可选）
pip install -e ".[cuda]"

# 4. 基准测试依赖（可选）
pip install -e ".[bench]"

openEuler 一键部署

# 安装系统依赖
sudo yum install -y python3 python3-pip gcc cmake

# 安装 NVIDIA GPU 驱动（如适用）
sudo yum install -y nvidia-driver cuda

# 安装 AgentMem
git clone <repo-url> agentmem && cd agentmem
pip3 install -e ".[dev,bench]"

# 一键验证
python3 -m agentmem.verify --model mock --platform openEuler

🚀 快速开始

from agentmem import AgentMemEngine, EngineConfig

# 创建引擎（无 GPU 时自动回退到 CPUBackend）
engine = AgentMemEngine(EngineConfig(
    num_gpu_blocks=1024,
    block_size=16,
    num_cpu_blocks=4096,
    enable_compression=True,
    enable_tiered_storage=True,
    # 创新点 I1/I2：量化压缩 + Markov 预测性预取
    enable_kv_quantization=True,    # CPU/Disk 占用降低 ~74%
    enable_markov_prefetch=True,    # 预取命中率提升（S8 实测 40%）
))

# 注册 system prompt（L0 永驻，自动加入 RadixTree）
session = engine.create_session("agent-1")
session.prefill("You are a helpful assistant.", lifecycle="permanent")

# 多轮对话：system prompt 的 KV Cache 自动复用
session.prefill("[User] 你好", lifecycle="medium")
reply = session.decode(max_tokens=32)

# 工具调用：返回数据自动压缩，原始数据 offload 到磁盘
session.inject_tool_result('{"items": [...2000 行 JSON...]}', max_tokens=256)

# Tree-of-Thought 分支推理：COW 共享父分支 KV Cache
branch_a = session.fork()   # ref_count 递增，零拷贝
branch_b = session.fork()

🧪 运行测试

# 全量单元测试（248 个测试，含创新点 I1/I2/P2/P3 专项测试）
pytest -q

# 带覆盖率报告
pytest --cov=agentmem --cov-report=term-missing

# 并行测试（加速）
pytest -n auto

# Benchmark（CPU mock 模式，S1-S8 全场景基线对照）
python -m benchmarks.run_benchmark --scenarios all --backend cpu --report markdown

# 指定场景 Benchmark（S7 量化 / S8 Markov 预取）
python -m benchmarks.run_benchmark --scenarios s7,s8 --report markdown

# 可选：编译原生 C 扩展（Linux/gcc，加速 RadixTree 与淘汰评分 3-10×）
python setup_native.py build_ext --inplace

# 运行示例
python -m examples.multi_turn_chat      # 多轮对话（前缀复用）
python -m examples.tot_branching        # Tree-of-Thought（COW 共享）
python -m examples.tool_calling         # 工具调用（数据压缩）

测试覆盖

模块	测试文件	测试数	覆盖率
common	test_common.py	26	90%+
core	test_core.py	36	90%+
lifecycle	test_lifecycle.py	16	93%
cow	test_cow.py	16	94%
compression	test_compression.py	21	80%+
storage	test_storage.py	25	70%+
scheduler	test_scheduler.py	17	94%
hal	test_hal.py	13	70%+
engine	test_engine.py	28	87%
extended	test_extended.py	21	—
innovations	test_innovations.py	29	84%
合计	11 文件	248	80%+

Benchmark 结果摘要

场景	显存降低	复用率	关键指标
S1 多轮对话	86.1%	89.6%	P1 目标 ≥30% ✅
S2 工具调用	80.9%	59.3%	工具数据压缩 91.5% ✅
S3 ToT 分支	73.5%	—	COW 节省 67.5% ✅
S4 多任务并发	83.2%	88.8%	公平性 Jain 指数 ✅
S5 长上下文	—	74.0%	Swap 延迟 0.090ms ✅
S6 分层存储	—	80.8%	CRC 校验 0 失败 ✅
S7 KV 量化	—	66.2%	I1 量化节省 74% / CRC 0 失败 ✅
S8 Markov 预取	—	83.3%	I2 Markov 命中率 40% / 预取命中率 80% ✅

�️ 可视化管理界面（Dashboard）

AgentMem 提供基于 FastAPI + Chart.js 的可视化管理界面，适合比赛现场演示：

实时监控：WebSocket 推送 engine.stats()，仪表盘卡片 + 显存/Swap 延迟折线图
场景交互：网页上直接触发 prefill / decode / 工具调用 / fork 分支，实时观察指标变化
基准对照：勾选 S1-S8 场景一键运行，柱状图对照 baseline vs optimized + P1-P6 达标表
系统架构：6 大基础模块 + 4 个创新点（I1/I2/P2/P3）卡片 + 数据流图

# 安装依赖
pip install fastapi uvicorn[standard]

# 启动 Dashboard（默认 127.0.0.1:8000）
python -m agentmem.dashboard

# 比赛演示（监听所有网卡）
python -m agentmem.dashboard --host 0.0.0.0 --port 8080

# 强制 CPU 后端
python -m agentmem.dashboard --backend cpu

浏览器打开 http://127.0.0.1:8000 即可访问。环境变量 AGENTMEM_BACKEND / AGENTMEM_GPU_BLOCKS / AGENTMEM_CPU_BLOCKS 可调整引擎参数。

� 项目结构

agentmem/
├── agentmem/                 # 核心包
│   ├── common/               # 公共层：类型、SeqLock、CRC32、哈希
│   ├── core/                 # KVBlockPool、RadixTree、淘汰策略、PayloadStore
│   ├── lifecycle/            # 5.2.1 生命周期管理器
│   ├── cow/                  # 5.2.2 COW 引擎
│   ├── compression/          # 5.2.3 上下文压缩 + 工具数据优化
│   ├── storage/              # 5.2.4 分层存储引擎（TieredStorage/SwapCache/OffloadStore/PrefetchEngine）
│   ├── scheduler/            # 5.2.5 任务调度（64 路分片 + 公平配额）
│   ├── hal/                  # 5.2.6 硬件抽象层（CUDA/DTK/CANN/CPU）
│   ├── engine.py             # AgentMemEngine 主编排器
│   ├── tokenizer.py          # Mock 推理（真实部署替换为 vLLM）
│   ├── verify.py             # 一键验证工具
│   └── dashboard/            # 可视化管理界面（FastAPI + Chart.js）
├── benchmarks/               # S1-S8 Benchmark 场景 + 基线对照 + 报告
│   ├── framework.py          # 基线/优化配置、指标采集
│   ├── run_benchmark.py      # CLI 入口
│   ├── report.py             # Markdown/JSON 报告生成
│   └── scenarios/            # S1-S8 场景实现（S7 量化 / S8 Markov 预取为 v1.2 新增）
├── tests/                    # pytest 单元测试（248 个，覆盖率 80%+）
├── docs/                     # 架构/部署/API/Benchmark 文档
├── examples/                 # 使用示例（多轮对话/ToT/工具调用）
├── pyproject.toml            # 项目配置（依赖/pytest/coverage）
└── README.md

📚 文档

🏗️ 设计哲学

AgentMem 采用插件式架构，作为独立 Python 包通过 vLLM 的 Scheduler/BlockManager 接口注入，核心改动全部位于 agentmem/ 目录下，最小化对 vLLM 核心的侵入（对应风险 R1 应对策略）。

📄 许可证

MIT License — 见 LICENSE。

👥 团队

角色	负责模块
dushun（项目负责人 / 系统架构）	总体架构、分层存储引擎、Benchmark
TBD	GPU 显存管理 & vLLM 集成
TBD	算法 & 前缀共享
TBD	上下文压缩 & 工具优化
TBD	国产化适配 & HAL