目录

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

高校赛题项目:面向长上下文、多分支和多 GPU 分布式智能体推理的内存管理与执行优化

本项目面向多 GPU 分布式智能体推理,重点解决共享父上下文被重复 Prefill、Prefix KV 在多卡间重复存储、兄弟分支被负载均衡策略拆散、KV Cache 持续增长,以及工具等待期间显存闲置等问题。项目实现 Prefix-aware Data Parallel Routing、ForkAttention、Tool KV Trimmer、Prompt Compaction 和分层 KV 存储,并完成国产加速卡、国产操作系统、多模型、多模态和多类 Agent 数据集适配。

核心系统位于 agentrix/ 子模块,覆盖分布式前缀感知路由、共享前缀调度、Attention 执行、应用层上下文压缩、KV Cache 分层存储、工具等待阶段缓存释放和异构平台适配。

核心创新

Agentrix 采用应用层、调度层、执行层和存储层协同设计,把智能体长生命周期推理中的重复上下文、重复 KV、重复计算和无效驻留统一纳入优化。

创新点 传统瓶颈 核心设计 代表性效果
Prefix-aware 分布式推理 普通 DP 只按负载分配请求,同一父上下文在多张 GPU 上重复 Prefill、重复保存 KV,并拆散兄弟分支 引入 Prefix Owner、最长前缀匹配、KV 容量约束、Arrival Wave 和 Cohort 聚合,使共享前缀请求优先落在同一 GPU DP 吞吐提升 4.35×~5.95×
ForkAttention 共享前缀执行 多个分支在同一 GPU 上仍会重复读取和计算相同 Prefix KV 基于物理 KV Block Table 构建 Prefix Forest,公共前缀协作计算,私有后缀独立合并 Attention Kernel 加速 1.90×~20.81×
工具等待阶段动态 KV 回收 Agent 等待工具时停止生成,但 KV 长时间占用 GPU 固定 TTL、预测 TTL 和全局压力感知的 Tool KV Trimmer Peak Live KV 降低 50.00%
可恢复 Prompt 与工具数据压缩 system prompt、工具描述、文件内容和历史结果重复进入上下文 精确去重、结构化 Backing Store、历史工具结果按需加载 完整上下文最多减少 66.40%
共享感知的分层 KV 管理 单级显存难以支撑长上下文和高并发 GPU、CPU、磁盘三级缓存,支持热度、共享度和 Fanout 感知 Peak Live KV 平均降低 39.17%
异构软硬件适配 推理优化依赖单一 CUDA 平台 CUDA、摩尔线程 MUSA、Apple Metal,以及 HCE 2.0、CentOS 7.9 构建适配 支持国产加速卡和国产操作系统环境
多模型、多模态与多任务统一适配 单模型、单数据集结果缺乏泛化性 Qwen、Llama、MiniCPM、GLM,以及 Qwen3.5/Qwen3.6 视觉语言模型;覆盖 Coding Agent、RAG、通用 Agent 覆盖 0.6B~32B 模型、文本与图文输入、6 类数据集/任务

分布式推理创新:Prefix-aware DP

普通 Data Parallel Router 通常按照请求数、队列长度或轮询策略分配请求。对 Agent 工作负载而言,相同父上下文的分支会被拆到多张 GPU,造成三类额外开销:

  • 每张 GPU 重复执行同一长前缀的 Prefill;
  • 每张 GPU 保存一份相同的 Prefix KV;
  • 兄弟分支无法在同一设备形成 ForkAttention Cohort。

Agentrix 在 vLLM Internal DP 中加入前缀感知路由,分布式调度目标由“请求数量均衡”扩展为“负载、KV 容量和前缀局部性联合优化”。

                        Prefix-aware Router
                               │
              ┌────────────────┴────────────────┐
              │                                 │
      Prefix Owner: Case A              Prefix Owner: Case B
              │                                 │
           GPU 0                              GPU 1
     A1 A2 A3 ... A32                   B1 B2 B3 ... B32
              │                                 │
      Fanout Admission                   Fanout Admission
              │                                 │
       ForkAttention                      ForkAttention

核心机制:

  1. Prefix Owner:记录长前缀已经驻留的设备,后续分支优先复用该设备上的 KV;
  2. 最长前缀匹配:根据 Prefix Hash 和物理缓存状态选择复用深度最大的 Rank;
  3. 容量感知放置:结合每张 GPU 的可用 KV Blocks,避免前缀亲和导致设备过载;
  4. Arrival Wave:短暂聚合同批兄弟请求,再按 Cohort 统一分配;
  5. Query Aggregation:让同一父上下文的分支在目标 GPU 上形成足够宽的 Decode 批次;
  6. 短请求旁路:短上下文不执行昂贵的前缀哈希和等待逻辑。

在 2 × RTX 5090、Qwen3-8B、32K Prefix、32 Branches 的配对实验中,Prefix-aware DP 与 ForkAttention 的组合吞吐达到普通 Flash DP 的 4.35×~5.95×。提升主要来自重复 Prefill、跨卡 KV 副本、排队和缓存抖动的减少。

详细设计:Prefix-aware DP Results

创新设计关系

应用层精确压缩
      │
      ▼
稳定的共享父上下文与 Prefix 标识
      │
      ▼
Prefix-aware DP:选择 Prefix Owner 和目标 GPU
      │
      ▼
设备内 Fanout Admission:聚合兄弟分支
      │
      ▼
ForkAttention:共享前缀协作执行
      │
      ▼
Tool KV Trimmer + GPU/CPU/Disk 分层存储

以上模块可以独立启用,也可以组合部署。组合实验会明确记录启用项,避免将系统级收益归因于单一模块。

一、赛题要求与项目实现

1.1 项目定位

本项目完成了一套面向智能体长生命周期和多 GPU 分布式推理的内存管理与执行优化系统。Prefix-aware Data Parallel Routing 是核心系统创新之一,用于在多卡环境中保持共享前缀局部性,并与 ForkAttention、工具等待阶段 KV 回收、可恢复上下文压缩和异构分层存储协同工作。

系统基于 openEuler 生态的华为云 EulerOS(HCE 2.0)完成编译和运行适配,并在 CentOS 7.9 等 Linux 发行版上验证构建流程。核心实现扩展了 vLLM、llama.cpp 和 LMCache,通过 OpenAI 兼容接口接入 Coding Agent、RAG Agent、工具调用和多阶段决策工作流。

项目重点解决以下问题:

  • 多轮推理导致上下文和 KV Cache 持续增长;
  • 多分支决策重复保存和计算相同父上下文;
  • 工具调用产生的大型中间数据进入 Prompt 和 KV Cache;
  • 工具等待期间会话停止生成,但 GPU KV 仍被占用;
  • 普通 Data Parallel 路由只关注负载均衡,导致共享前缀跨 GPU 复制、重复 Prefill 和分支拆散;
  • 单级显存无法容纳长生命周期和高并发工作负载;
  • 不同模型、推理框架和异构硬件缺少统一适配路径。

核心系统位于 agentrix/ 子模块。

1.2 赛题要求对照

赛题要求 项目实现 验证材料
基于国内主流开源操作系统开发 在基于 openEuler 的 HCE 2.0 上完成依赖安装、CUDA 环境、编译和运行适配 HCE 2.0 / CentOS 7.9 适配文档
鼓励支持更多 Linux 发行版 提供 HCE 2.0、CentOS 7.9 和常规 Linux x86_64 构建路径 构建指南、环境脚本和平台说明
扩展现有推理框架 扩展 vLLM 和 llama.cpp,集成 LMCache,并提供 SGLang 接入说明 vllm/llama.cpp/LMCache/docs/
支持典型智能体工作流 支持多轮工具调用、并行子 Agent、长代码上下文、多阶段 RAG 和分支决策 Coding-Agent、HotpotQA LangGraph、AgentBoard、AppWorld
支持分布式推理 扩展 vLLM Internal DP,实现 Prefix-aware Router、Prefix Owner、容量感知放置和 Cohort 聚合 DP=2、DP=4、DP=8 配对实验
使用开源大模型 覆盖 Qwen3、Qwen3.5、Qwen3.6、Llama 3、MiniCPM4.1 和 GLM-4 模型兼容性与 TP 测试报告
优化前后使用相同硬件 A/B 实验固定 GPU、模型、输入、输出长度、KV 容量和并发配置,只切换优化模块 单 GPU、DP、DP=8 和 KV Trimmer 配对实验
构建可复现 Benchmark 提供数据适配器、启动脚本、固定实验矩阵、指标采集和 Markdown 报告 benchmark/docs/main_experiment_matrix.md
系统评估优化效果 统计吞吐、TTFT、TPOT、端到端延迟、Live KV、KV-Time Area、缓存流量和恢复延迟 本 README 第六节及完整实验报告

1.3 推荐优化方向覆盖

推荐方向 Agentrix 实现
KV Cache 生命周期管理 Prefix Cache 复用、Tool KV Trimmer、固定/预测 TTL、缓存淘汰、动态回收和恢复
分支推理内存共享 vLLM 分页 KV 共享公共前缀,私有后缀独立分配;Prefix-aware DP 保持跨卡局部性,ForkAttention 复用共享 Prefix KV 的读取和计算
Prompt 与上下文压缩 Prompt Compaction、工具 Schema 去重、重复文件消除、精确可恢复压缩
工具调用数据优化 Tool Result Backing Store、结构化元数据、历史结果分页和按需恢复
分层内存与异构存储 GPU KV、CPU Offload、LMCache CPU/Disk、冷热分层、动态迁移和按需加载
异构 AI 加速硬件支持 NVIDIA CUDA、摩尔线程 MUSA、Apple Metal;不兼容形状自动回退到原生路径

1.4 项目能力概览

能力维度 项目实现
智能体推理流程 提供 OpenAI 兼容推理服务,支持多轮对话、工具调用、并行子 Agent、长代码上下文和多阶段 RAG
内存管理 覆盖 Prefix KV 复用、动态回收、CPU/Disk 下沉、缓存恢复和热冷分层
分布式分支推理 Prefix-aware DP 将同一父上下文的分支放置到 Prefix Owner,设备内再通过分页 KV 共享和 ForkAttention 降低重复显存访问
上下文压缩 支持 system prompt、工具描述、代码文件和历史工具结果的精确去重与按需恢复
多模型适配 覆盖 Qwen、Llama、MiniCPM、GLM,以及 Qwen3.5/Qwen3.6 多模态混合注意力模型
多任务支持 覆盖软件工程 Agent、通用 Agent、交互式 Agent、多跳 RAG 和多轮 Coding-Agent
异构平台 支持 NVIDIA CUDA、摩尔线程 MUSA、Apple Metal,并提供 HCE 2.0 和 CentOS 7.9 适配
可复现性 提供固定 A/B 矩阵、同硬件对照、算子正确性测试、API 测试、TP 测试和数据集适配器
文档与交付 提供技术方案、系统设计、部署指南、平台适配、实验报告、项目说明书、PPT 和演示视频

1.5 核心实现

  1. KV Cache 生命周期管理
    追踪请求从生成、工具等待到恢复的完整生命周期,根据 TTL 和全局 KV 压力动态释放、恢复或重新计算 KV。

  2. Prefix-aware 分布式推理
    在 vLLM Internal DP 中引入 Prefix Owner、最长前缀匹配、容量感知放置和 Arrival Wave,将同一父上下文的分支集中到可复用 Prefix KV 的设备。

  3. 共享前缀与分支隔离
    公共前缀复用相同物理 KV Blocks,分支私有后缀使用独立页,控制多路径决策带来的显存增长。

  4. ForkAttention
    根据物理 KV Block Table 构建 Prefix Forest,协作处理兄弟分支的公共前缀,减少重复显存读取和 Attention 计算。

  5. Prompt 与工具数据压缩
    精确去重 system prompt、工具描述、代码文件和历史工具结果,大型中间数据存入 Backing Store 后按需恢复。

  6. 分层内存管理
    在 GPU、CPU 和本地磁盘之间迁移 KV Cache,根据共享度、热度、容量和 I/O 成本执行淘汰与加载。

  7. 国产软硬件适配
    提供 HCE 2.0 和 CentOS 7.9 构建方案,在 llama.cpp 路径中实现 CUDA 与摩尔线程 MUSA 后端。

  8. 多模型、多模态与多数据集验证
    覆盖 Qwen、Llama、MiniCPM、GLM,以及 Qwen3.5/Qwen3.6 图文模型,并在 Coding Agent、通用 Agent、交互式 Agent和长前缀 RAG 数据集上测试。

1.6 技术实现特点

  • Prefix Owner、最长前缀匹配和容量感知的分布式路由;
  • Arrival Wave 与 Cohort 级请求聚合;
  • 物理 KV Block 级共享关系识别;
  • 分页 KV 共享与私有后缀隔离;
  • Prefix Forest 与 CUDA Graph 联合执行;
  • Prefix Owner 和 KV 容量感知的多 GPU 路由;
  • 固定 TTL、预测 TTL 和压力感知动态回收;
  • 可恢复、可审计的 Prompt 精确压缩;
  • GPU、CPU、磁盘三级 KV Cache;
  • Qwen3.5/Qwen3.6 图文输入、视觉前缀和混合注意力适配;
  • CUDA、MUSA 和 Metal 多后端支持;
  • FlashAttention、普通调度和普通 LRU 回退路径;
  • 正向结果、适用边界和负向实验完整记录。

二、项目结构

.
├── agentrix/          # Agentrix 核心系统(Git Submodule)
├── .gitmodules        # Git 子模块配置
├── README.md          # 项目总览、技术方案与实验结果说明
├── 项目说明书.pdf     # 项目背景、设计方案及实现说明
├── 详细设计.pdf       # 系统架构与核心模块的详细设计文档
└── 演示视频.mp4       # 系统功能及运行效果演示视频

2.1 根目录文件介绍

文件或目录 类型 说明
agentrix/ Git 子模块 项目核心代码仓库,包含推理框架扩展、KV Cache 管理、上下文压缩、Benchmark 和技术文档等内容。
.gitmodules 配置文件 记录 agentrix/ 子模块的仓库地址及本地路径,用于初始化和更新核心代码。
README.md 项目文档 介绍项目背景、核心创新、系统架构、平台适配、实验结果及使用方式。
项目说明书.pdf 交付文档 系统说明项目需求、总体方案、主要功能、技术实现与验证情况。
详细设计.pdf 设计文档 详细描述系统架构、模块划分、关键流程和核心算法设计。
演示视频.mp4 演示材料 展示项目主要功能、部署运行流程及实际效果。

agentrix/ 子模块的内部结构:

agentrix/
├── vllm/          # ForkAttention、调度、DP 路由和推理服务
├── LMCache/       # CPU / Disk KV 分层缓存
├── llama.cpp/     # CUDA、MUSA 和 Metal ForkAttention
├── application/   # Prompt Compaction、KV Trimmer、TTL Predictor
├── benchmark/     # 数据集、实验脚本和指标报告
├── docs/          # 架构、设计、实验和平台文档
├── Dockerfile
└── README.md

Agentrix 子项目地址:https://www.gitlink.org.cn/I1Dsk46hji/Agentrix

三、平台、模型与数据集适配

3.1 硬件平台与操作系统

平台 适配情况 验证配置
NVIDIA CUDA 支持 Turing 及更新架构;BF16 路径要求 Ampere 或更新架构 Tesla T4、RTX 5070、RTX 5090、H20
摩尔线程 MUSA 已实现 MUSA ForkAttention Kernel 和运行时分发;不支持的请求自动回退到原生路径 MTT S4000、QY2、mp_22、FP16
华为云 EulerOS(HCE)2.0 基于 openEuler;已完成包管理器、编译工具链、CUDA 环境和构建流程适配 GCC 10.3、CMake 3.22、CUDA 11.4、Tesla T4
CentOS Linux 7.9 已解决旧 GCC、旧 CMake、EOL 软件源和 Locale 等构建问题 devtoolset-10、CMake 3.26.4、CUDA 11.4、Tesla T4
Apple Metal llama.cpp 路径提供 Apple Silicon ForkAttention 支持 Metal 后端
通用 CPU / Disk 层 支持 CPU KV Offload 和本地磁盘 KV Cache vLLM OffloadingConnector、LMCache

摩尔线程后端沿用 llama.cpp 的 MUSA 兼容层,并提供独立的 MUSA Partial-Attention Kernel。当前专用路径面向 QY2 或更新架构、FP16 KV、Qwen3、2~8 个单 Token Decode 序列和 64/128 Head Dimension。超出范围的请求继续使用 MUSA 原生 Attention 路径。

操作系统适配文档:

3.2 多模型与多模态适配

模型架构 验证模型 适配内容与状态
Qwen3ForCausalLM Qwen3-0.6B、1.7B、8B、14B、32B 覆盖单 GPU、DP、TP、Coding-Agent、RAG 和 llama.cpp 后端
LlamaForCausalLM Llama-3.2-1B、Meta-Llama-3.1-8B-Instruct 完成单 GPU压力测试和 TP=2 ForkAttention 验证
MiniCPMForCausalLM MiniCPM4.1-8B 完成 TP=2、CUDA Graph、API 推理和逻辑 KV 复用验证
ChatGLMModel GLM-4-9B-chat 完成 TP=2、CUDA Graph、API 推理和逻辑 KV 复用验证
Qwen3_5ForConditionalGeneration Qwen3.5-27B、Qwen3.6-27B 适配 Head Dimension 256、GDN/Full-Attention 混合结构、图文输入和多模态因果视觉前缀

Qwen3.5 和 Qwen3.6 的 27B 模型包含 48 个 GDN/线性注意力层和 16 个 Full-Attention 层。项目对图文输入采用分层处理:视觉编码和 Prefill 继续走原有高性能路径,Full-Attention 层的共享前缀 Decode 使用 ForkAttention,GDN 层保持原有执行路径。系统同时维护多模态输入哈希、KV Block Hash 和路由命名空间,避免不同图片或截图被错误识别为可共享前缀。Head Dimension 256 专用路径要求 SM90 或更新架构。

相关文档:

3.3 多数据集与工作负载支持

类型 数据集或项目 覆盖内容
软件工程 Agent SWE-bench Verified 长代码上下文、多分支 Serving 和跨数据集性能验证
通用智能体 AgencyBench 长前缀、多分支压力实验
Agent 基准 AgentBoard 单 GPU、DP 和 TP 数据适配器验证
交互式应用 Agent AppWorld 单 GPU、DP 和 TP 数据适配器验证
多跳问答 RAG HotpotQA LangGraph Planner、十路 Fanout、工具检索、Reducer 的端到端实验
Coding-Agent 仓库任务 Django、SQLite、FFmpeg 约 30K Token 仓库父上下文、16 个异构子 Agent、多轮工具轨迹
可执行功能任务 Django、SQLite、FFmpeg 共 12 个任务 补丁范围检查、公开测试、隐藏测试和回归验证基础设施

数据集覆盖三类典型负载:

  1. 长共享前缀 Serving:SWE-bench Verified、AgencyBench、AgentBoard、AppWorld;
  2. 端到端 RAG Agent:HotpotQA + LangGraph;
  3. 多轮 Coding Agent:Django、SQLite、FFmpeg 仓库上下文和工具轨迹。

相关文档:

四、系统设计

Agent / LangGraph / RAG Workflow
                 │
                 ▼
┌──────────────────────────────────────────┐
│ Application Layer                        │
│                                          │
│  Prompt Compaction                       │
│  Tool Schema Deduplication               │
│  Recoverable Tool-result Paging          │
│  Tool KV TTL Prediction                  │
└──────────────────┬───────────────────────┘
                   │ OpenAI-compatible API
                   ▼
┌──────────────────────────────────────────┐
│ Agentrix Distributed vLLM                │
│                                          │
│  Prefix-aware DP Router                  │
│  Prefix Owner / Capacity-aware Placement │
│  Prefix Cache / Fanout Admission         │
│  Prefix Forest / CUDA Graph              │
│                                          │
│  FlashAttention      ForkAttention       │
└──────────────────┬───────────────────────┘
                   │
                   ▼
┌──────────────────────────────────────────┐
│ KV Memory Hierarchy                      │
│                                          │
│  GPU KV Pool                             │
│  Native CPU Offload                      │
│  LMCache CPU / Disk                      │
│  Tool-wait KV Trimming                   │
└──────────────────────────────────────────┘

典型执行流程:

构建共享父上下文
        │
        ▼
应用层精确压缩
        │
        ▼
建立或命中 Prefix Cache
        │
        ▼
Prefix-aware Router 选择目标 GPU
        │
        ▼
Fanout Admission 聚合兄弟分支
        │
        ├── 普通请求 ──> FlashAttention
        │
        └── 长前缀并行分支 ──> ForkAttention
                                  │
                                  ▼
                             工具调用等待
                                  │
                   Tool KV Trim / CPU / Disk
                                  │
                                  ▼
                             恢复并继续生成

五、关键模块设计

5.1 Prefix-aware Data Parallel Routing

Prefix-aware DP 是 Agentrix 的分布式推理调度核心。

普通 Data Parallel 路由主要依据请求数或队列负载分配请求,同一 Case 的兄弟分支可能落到不同 GPU。每张 GPU 随后独立执行相同长前缀的 Prefill,并保存一份 Prefix KV。设备内分支数量变少后,ForkAttention 也难以形成有效的共享执行批次。

Agentrix Router 综合以下信息选择目标 Rank:

  • 请求的 Prefix Hash 和最长可复用前缀;
  • Prefix KV 当前所在的 Owner Rank;
  • 各 Rank 的可用 KV Blocks 和活跃负载;
  • Cohort 的分支数量和预估内存需求;
  • 同批请求的到达波次;
  • 目标设备能否形成有效的 ForkAttention 批次。
普通 DP:
Case A Branch 1 ──> GPU 0
Case A Branch 2 ──> GPU 1
Case A Branch 3 ──> GPU 2
Case A Branch 4 ──> GPU 3

Prefix-aware DP:
Case A Branches ──> Prefix Owner GPU 0
Case B Branches ──> Prefix Owner GPU 1

该设计同时减少重复 Prefill、跨卡 Prefix KV 副本、设备排队和缓存抖动。短请求或低复用请求会绕过复杂路由逻辑。

详细设计:Prefix-aware DP Results

5.2 Fanout-aware Admission

Prefix-aware Router 完成跨 GPU 放置后,Fanout-aware Admission 在目标设备内聚合兄弟分支,使它们在相近时间进入 Decode。

主要职责:

  • 识别同一父上下文的兄弟分支;
  • 在短时间窗口内聚合请求;
  • 保持高 Fanout Prefix 的活跃度;
  • 控制无关请求对 Cohort 的干扰;
  • 在聚合收益较低时使用普通调度。

5.3 ForkAttention

ForkAttention 面向单 Token Decode 阶段。Prefix-aware DP 保证兄弟分支集中到同一 GPU,Fanout-aware Admission 让它们同时活跃,ForkAttention 随后复用这些请求共同引用的物理 Prefix KV。

                 Shared Prefix KV
                        │
        ┌───────────────┼───────────────┐
        │               │               │
     Query A         Query B         Query C
        │               │               │
  Private KV A    Private KV B    Private KV C

主要特性:

  • 根据物理 KV Block 判断共享关系;
  • 支持多共享根和嵌套前缀;
  • 支持 Prefix Forest CUDA Graph;
  • 支持自适应 Prefix Split;
  • 形状不兼容或共享收益较低时回退到 FlashAttention。

适合长公共前缀、高 Fanout、分支集中到达的请求。短 Prompt、单分支和 Prefill-heavy 请求通常使用普通 Attention 路径。

详细设计:ForkAttention Operator Profile

5.4 Application Prompt Compaction

Prompt Compaction 执行确定性精确去重,处理应用层重复内容。

支持:

  • 删除空 Section;
  • 删除身份和内容完全一致的重复 Section;
  • 规范化 JSON 表示;
  • 合并重复工具 Schema;
  • 对大型历史工具结果执行可恢复分页;
  • 统计压缩前后字符数和 Token 数。

每个 Section 使用稳定的 segment_id。同一 ID 对应不同内容时直接报错,避免错误合并。自然语言内容不会被摘要、改写或模糊匹配。

详细设计:Application Prompt Compaction

5.5 Tool KV Trimmer

Agent 等待文件读取、代码搜索、测试或网络请求时不生成 Token,但历史 KV 仍会占用 GPU。

Tool KV Trimmer 根据固定或预测 TTL 检查全局 KV 压力,并释放等待会话的 Live KV。

进入工具等待
      │
      ▼
固定或预测 TTL
      │
      ▼
检查全局 KV 压力
      │
      ├── 压力低:保留热 KV
      │
      └── 压力高:释放 Live KV
                         │
                         ▼
                    工具结果返回
                         │
       Prefix Cache / Connector / Recompute

Trim 保留 Token 历史、模型输出、工具状态和恢复边界。缓存未命中时通过重新计算恢复上下文。

详细设计:

5.6 KV Offload 与 LMCache

Agentrix 支持 GPU、CPU 和本地磁盘三级 KV 存储:

GPU KV ──> CPU KV ──> Local Disk / NVMe

主要能力:

  • vLLM Native CPU Offload;
  • LMCache CPU / Disk Cache;
  • 普通 LRU;
  • Fork-aware 热前缀保护;
  • HOT、COOLING、COLD 生命周期;
  • CPU 淘汰后的异步磁盘写入;
  • 缓存恢复失败时重新计算。

缓存策略同时考虑共享程度、未来命中概率和 I/O 成本。

详细设计:KV Offload Experiment

六、重要实验数据

实验结果对应指定模型、硬件和受控工作负载。不同场景需要重新测量。

评测原则

  • 同硬件配对对照: 优化前后使用相同 GPU、模型权重、精度、上下文长度、输出长度、并发数和 KV 容量;
  • 控制单一或明确组合变量: 算子实验只比较 Attention Kernel,系统实验明确列出调度、压缩和缓存等组合项;
  • 工作量一致: 对比相同请求集合、输入 Token 和输出 Token,记录完成请求数;
  • 正确性检查: 覆盖算子数值一致性、API 请求完成、Prefix 恢复和 Prompt Backing Store 精确恢复;
  • 任务成功率评测: Coding-Agent 基础设施支持补丁范围检查、公开测试、隐藏测试和回归测试;仅衡量 Serving 性能的实验不会替代任务成功率;
  • 完整指标: 同时记录吞吐、TTFT、TPOT、端到端延迟、Live KV、KV-Time Area、Host RSS、外部缓存流量和磁盘 I/O。

6.1 ForkAttention 纯算子实验

测试配置:

  • 公共前缀:4K、8K、16K、32K、64K;
  • 并行 Query:2、4、8、16、32;
  • 共 25 个测试单元;
  • Qwen3-0.6B Attention 几何;
  • 每个 Query 包含 128 Token 私有后缀。
指标 结果
ForkAttention 获胜 25 / 25
Attention Kernel 加速 1.90×~20.81×
L2 Read Traffic 降低 47.8%~96.6%
Tensor Pipe 指令降低 86.2%~96.8%
FMA Pipe 指令降低 71.2%~94.7%

代表性结果:

Prefix Query 数 Flash / Fork 加速
4K 2 72.3 / 37.7 μs 1.92×
8K 16 881.1 / 103.8 μs 8.49×
16K 32 3432.0 / 203.2 μs 16.89×
64K 32 13955.9 / 670.8 μs 20.81×

纯算子数据反映 Attention Kernel 表现,不包含模型权重、MLP、调度、采样和 API Serving。

6.2 单 GPU Serving

环境:

  • NVIDIA RTX 5070 12 GiB;
  • Qwen3-1.7B;
  • 8K 公共前缀;
  • 16 个并行分支;
  • 每个分支生成 256 Token。
指标 FlashAttention ForkAttention 变化
输出吞吐 448.19 tok/s 923.95 tok/s +106.15%
请求延迟 P50 7823.08 ms 3056.12 ms -60.9%
TPOT P50 29.59 ms 11.30 ms -61.8%
Fork 物理激活率 77.51%

6.3 跨数据集单 GPU 验证

固定配置为 16K 公共前缀、16 分支和 256 输出 Token。

数据集 Flash Fork 吞吐提升
SWE-bench Verified 257.10 tok/s 667.12 tok/s +159.47%
AgencyBench 251.91 tok/s 684.74 tok/s +171.82%
AgentBoard 247.41 tok/s 649.72 tok/s +162.61%
AppWorld 250.23 tok/s 676.96 tok/s +170.54%

86 个配对 Case 的 Fork/Flash 中位加速为 2.66×

6.4 Prefix-aware DP

环境:

  • 2 × NVIDIA RTX 5090;
  • Qwen3-8B;
  • 32K Prefix;
  • 32 Branches。
数据集 Flash 普通 DP Fork Prefix-aware DP 加速
AgentBoard 231.4 tok/s 1007.1 tok/s 4.35×
AppWorld 173.8 tok/s 1033.4 tok/s 5.95×
SWE-bench Verified 185.8 tok/s 1032.7 tok/s 5.56×

吞吐提升主要来自重复 Prefill 减少、Prefix Owner 保持、请求排队缩短和缓存抖动降低。

6.5 Coding-Agent DP=8 全系统实验

环境:

  • 8 × NVIDIA H20;
  • Qwen3-32B FP16;
  • Internal DP=8;
  • Django、SQLite、FFmpeg;
  • 每个 Case 约 30K Token 代码仓库上下文;
  • 每个 Case 派生 16 个 Coding Subagents;
  • 每个 Variant 处理 2304 个分支请求。

对比配置:

Baseline:
FlashAttention
+ 普通 Internal DP
+ 未启用 Prompt Compaction

Agentrix:
ForkAttention
+ Prefix-aware Internal DP
+ Exact Prompt Compaction
范围 Baseline Wall Time Agentrix Wall Time 配对加速 Peak Live KV 降低 TTFT 降低
Django 1562.16 s 76.09 s 20.53× 41.84% 96.52%
SQLite 1542.86 s 74.56 s 20.70× 36.81% 96.77%
FFmpeg 1611.36 s 72.09 s 22.35× 38.85% 96.79%
总计 4716.38 s 222.75 s 21.19× 39.17% 96.69%

这组结果来自 Attention Backend、DP Placement、Fanout Scheduling、CUDA Graph 和 Prompt Compaction 的共同作用。

详细结果:Coding-Agent DP=8 Results

6.6 Prompt Compaction

在 Django、SQLite 和 FFmpeg Coding-Agent 上下文中追加 16 次真实文件读取:

工具输出上限 完整上下文减少 平均节省 Token 工具后缀减少
4 KiB 34.19% 18177 80%~86%
32 KiB 66.40% 69926 93%~95%

全部 30 个测试单元均可从 Backing Store 精确恢复原始消息。

6.7 Tool KV Trimmer

环境为 Qwen3-0.6B、每会话 4K Prompt,两批各四个工具等待会话。

模式 Peak Live KV KV-Time Area 恢复延迟
不 Trim 3.500 GiB 3.316 GiB·s 25.59 ms
固定 TTL 2.844 GiB 2.384 GiB·s 26.77 ms
预测 TTL 1.750 GiB 1.395 GiB·s 26.80 ms

预测 TTL 相对不 Trim:

  • Peak Live KV 降低 **50.00%**;
  • KV-Time Area 降低 **57.94%**;
  • 恢复延迟增加约 1.21 ms

vLLM 会预留固定大小的 GPU KV Pool,因此 Live KV 下降后,nvidia-smi 中的进程显存可能保持不变。

6.8 KV Offload 负向实验

实验性 cohort_lru 相对普通 LRU:

指标 变化
GPU → CPU Store 流量 -26.79%
Store Admission Failure 28 → 16
外部 Prefix Cache Hit Rate 16.3% → 19.2%
文件系统读取 +19.59%
端到端吞吐 -8.91%

磁盘恢复开销抵消了部分缓存收益,因此该策略保持为可选功能,普通 LRU 仍为默认配置。

七、使用方法

7.1 初始化子模块

克隆父仓库后执行:

git submodule sync --recursive
git submodule update --init --recursive
cd agentrix

单独克隆 Agentrix:

git clone --recurse-submodules https://www.gitlink.org.cn/I1Dsk46hji/Agentrix.git
cd Agentrix

7.2 环境要求

  • Linux x86_64;
  • NVIDIA GPU + CUDA,或摩尔线程 GPU + MUSA;
  • NVIDIA Driver 和 CUDA Toolkit,或兼容的 MUSA SDK 与驱动;
  • GCC/G++ 11.3 或更新版本;
  • Python 3.12;
  • Git、curl、CMake 和标准 C/C++ 构建工具;
  • 推荐使用 uvccache
sudo apt update
sudo apt install -y build-essential git curl ccache

curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"

7.3 构建 vLLM

cd agentrix/vllm

uv venv --python 3.12 --seed
source .venv/bin/activate

export CUDA_HOME="$(dirname "$(dirname "$(readlink -f "$(command -v nvcc)")")")"
export PATH="$CUDA_HOME/bin:$PWD/.venv/bin:$PATH"
export MAX_JOBS=2
export NVCC_THREADS=1

# SM120 / Blackwell;其他 GPU 架构需要调整
export TORCH_CUDA_ARCH_LIST=12.0

uv pip install -v -e . --torch-backend=auto

运行关键测试:

uv pip install -r requirements/test/cuda.txt

python -m pytest -q \
  tests/v1/worker/test_gpu_block_table.py \
  tests/v1/worker/test_fork_cudagraph_dispatch.py \
  tests/kernels/test_fork_attention.py \
  tests/kernels/test_fork_attention_backend.py

7.4 安装 Benchmark

cd ../benchmark

uv venv --python 3.12 --seed
source .venv/bin/activate

uv pip install -e ../application
uv pip install -e ".[data,test]"

python -m pytest
agentrix-bench inspect-data

完整构建说明:AutoDL Build Guide

八、交付材料

九、文档索引

9.1 架构与设计

9.2 实验结果

9.3 构建与平台

十、结果说明

本项目已经覆盖赛题要求的智能体推理流程、KV 生命周期管理、分支共享、Prompt 压缩、工具数据按需加载、分层内存、多模型和异构硬件适配。

实验数据遵循以下口径:

  • 纯算子数据用于验证 Kernel 机制和显存流量变化;
  • 单 GPU、DP 和 DP=8 数据用于验证完整推理服务效果;
  • DP=8 的 21.19× 来自 ForkAttention、Prefix-aware DP、Fanout Scheduling、CUDA Graph 和 Prompt Compaction 的共同作用;
  • Live KV 下降表示 KV Pool 内部 Block 被回收,CUDA 总分配显存可能保持不变;
  • Serving Benchmark 统计性能和资源利用率,任务成功率通过可执行任务测试单独评估;
  • Prompt Compaction 保证字节级可恢复,模型任务质量仍需使用原始 Prompt 和压缩 Prompt 进行配对测试;
  • FlashAttention、普通调度和普通 LRU 均保留为回退路径。
关于
208.5 MB
邀请码
    Gitlink(确实开源)
  • 加入我们
  • 官网邮箱:gitlink@ccf.org.cn
  • QQ群
  • QQ群
  • 公众号
  • 公众号

版权所有:中国计算机学会技术支持:开源发展技术委员会
京ICP备13000930号-9 京公网安备 11010802047560号