普通 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
面向智能体的内存管理系统设计与实现
本项目面向多 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
普通 Data Parallel Router 通常按照请求数、队列长度或轮询策略分配请求。对 Agent 工作负载而言,相同父上下文的分支会被拆到多张 GPU,造成三类额外开销:
Agentrix 在 vLLM Internal DP 中加入前缀感知路由,分布式调度目标由“请求数量均衡”扩展为“负载、KV 容量和前缀局部性联合优化”。
核心机制:
在 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
创新设计关系
以上模块可以独立启用,也可以组合部署。组合实验会明确记录启用项,避免将系统级收益归因于单一模块。
一、赛题要求与项目实现
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、工具调用和多阶段决策工作流。
项目重点解决以下问题:
核心系统位于
agentrix/子模块。1.2 赛题要求对照
vllm/、llama.cpp/、LMCache/、docs/benchmark/和docs/main_experiment_matrix.md1.3 推荐优化方向覆盖
1.4 项目能力概览
1.5 核心实现
KV Cache 生命周期管理
追踪请求从生成、工具等待到恢复的完整生命周期,根据 TTL 和全局 KV 压力动态释放、恢复或重新计算 KV。
Prefix-aware 分布式推理
在 vLLM Internal DP 中引入 Prefix Owner、最长前缀匹配、容量感知放置和 Arrival Wave,将同一父上下文的分支集中到可复用 Prefix KV 的设备。
共享前缀与分支隔离
公共前缀复用相同物理 KV Blocks,分支私有后缀使用独立页,控制多路径决策带来的显存增长。
ForkAttention
根据物理 KV Block Table 构建 Prefix Forest,协作处理兄弟分支的公共前缀,减少重复显存读取和 Attention 计算。
Prompt 与工具数据压缩
精确去重 system prompt、工具描述、代码文件和历史工具结果,大型中间数据存入 Backing Store 后按需恢复。
分层内存管理
在 GPU、CPU 和本地磁盘之间迁移 KV Cache,根据共享度、热度、容量和 I/O 成本执行淘汰与加载。
国产软硬件适配
提供 HCE 2.0 和 CentOS 7.9 构建方案,在 llama.cpp 路径中实现 CUDA 与摩尔线程 MUSA 后端。
多模型、多模态与多数据集验证
覆盖 Qwen、Llama、MiniCPM、GLM,以及 Qwen3.5/Qwen3.6 图文模型,并在 Coding Agent、通用 Agent、交互式 Agent和长前缀 RAG 数据集上测试。
1.6 技术实现特点
二、项目结构
2.1 根目录文件介绍
agentrix/.gitmodulesagentrix/子模块的仓库地址及本地路径,用于初始化和更新核心代码。README.md项目说明书.pdf详细设计.pdf演示视频.mp4agentrix/子模块的内部结构:Agentrix 子项目地址:https://www.gitlink.org.cn/I1Dsk46hji/Agentrix
三、平台、模型与数据集适配
3.1 硬件平台与操作系统
mp_22、FP16摩尔线程后端沿用 llama.cpp 的 MUSA 兼容层,并提供独立的 MUSA Partial-Attention Kernel。当前专用路径面向 QY2 或更新架构、FP16 KV、Qwen3、2~8 个单 Token Decode 序列和 64/128 Head Dimension。超出范围的请求继续使用 MUSA 原生 Attention 路径。
操作系统适配文档:
3.2 多模型与多模态适配
Qwen3ForCausalLMLlamaForCausalLMMiniCPMForCausalLMChatGLMModelQwen3_5ForConditionalGenerationQwen3.5 和 Qwen3.6 的 27B 模型包含 48 个 GDN/线性注意力层和 16 个 Full-Attention 层。项目对图文输入采用分层处理:视觉编码和 Prefill 继续走原有高性能路径,Full-Attention 层的共享前缀 Decode 使用 ForkAttention,GDN 层保持原有执行路径。系统同时维护多模态输入哈希、KV Block Hash 和路由命名空间,避免不同图片或截图被错误识别为可共享前缀。Head Dimension 256 专用路径要求 SM90 或更新架构。
相关文档:
3.3 多数据集与工作负载支持
数据集覆盖三类典型负载:
相关文档:
四、系统设计
典型执行流程:
五、关键模块设计
5.1 Prefix-aware Data Parallel Routing
Prefix-aware DP 是 Agentrix 的分布式推理调度核心。
普通 Data Parallel 路由主要依据请求数或队列负载分配请求,同一 Case 的兄弟分支可能落到不同 GPU。每张 GPU 随后独立执行相同长前缀的 Prefill,并保存一份 Prefix KV。设备内分支数量变少后,ForkAttention 也难以形成有效的共享执行批次。
Agentrix Router 综合以下信息选择目标 Rank:
该设计同时减少重复 Prefill、跨卡 Prefix KV 副本、设备排队和缓存抖动。短请求或低复用请求会绕过复杂路由逻辑。
详细设计:Prefix-aware DP Results
5.2 Fanout-aware Admission
Prefix-aware Router 完成跨 GPU 放置后,Fanout-aware Admission 在目标设备内聚合兄弟分支,使它们在相近时间进入 Decode。
主要职责:
5.3 ForkAttention
ForkAttention 面向单 Token Decode 阶段。Prefix-aware DP 保证兄弟分支集中到同一 GPU,Fanout-aware Admission 让它们同时活跃,ForkAttention 随后复用这些请求共同引用的物理 Prefix KV。
主要特性:
适合长公共前缀、高 Fanout、分支集中到达的请求。短 Prompt、单分支和 Prefill-heavy 请求通常使用普通 Attention 路径。
详细设计:ForkAttention Operator Profile
5.4 Application Prompt Compaction
Prompt Compaction 执行确定性精确去重,处理应用层重复内容。
支持:
每个 Section 使用稳定的
segment_id。同一 ID 对应不同内容时直接报错,避免错误合并。自然语言内容不会被摘要、改写或模糊匹配。详细设计:Application Prompt Compaction
5.5 Tool KV Trimmer
Agent 等待文件读取、代码搜索、测试或网络请求时不生成 Token,但历史 KV 仍会占用 GPU。
Tool KV Trimmer 根据固定或预测 TTL 检查全局 KV 压力,并释放等待会话的 Live KV。
Trim 保留 Token 历史、模型输出、工具状态和恢复边界。缓存未命中时通过重新计算恢复上下文。
详细设计:
5.6 KV Offload 与 LMCache
Agentrix 支持 GPU、CPU 和本地磁盘三级 KV 存储:
主要能力:
缓存策略同时考虑共享程度、未来命中概率和 I/O 成本。
详细设计:KV Offload Experiment
六、重要实验数据
评测原则
6.1 ForkAttention 纯算子实验
测试配置:
代表性结果:
纯算子数据反映 Attention Kernel 表现,不包含模型权重、MLP、调度、采样和 API Serving。
6.2 单 GPU Serving
环境:
6.3 跨数据集单 GPU 验证
固定配置为 16K 公共前缀、16 分支和 256 输出 Token。
86 个配对 Case 的 Fork/Flash 中位加速为 2.66×。
6.4 Prefix-aware DP
环境:
吞吐提升主要来自重复 Prefill 减少、Prefix Owner 保持、请求排队缩短和缓存抖动降低。
6.5 Coding-Agent DP=8 全系统实验
环境:
对比配置:
这组结果来自 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 次真实文件读取:
全部 30 个测试单元均可从 Backing Store 精确恢复原始消息。
6.7 Tool KV Trimmer
环境为 Qwen3-0.6B、每会话 4K Prompt,两批各四个工具等待会话。
预测 TTL 相对不 Trim:
vLLM 会预留固定大小的 GPU KV Pool,因此 Live KV 下降后,
nvidia-smi中的进程显存可能保持不变。6.8 KV Offload 负向实验
实验性
cohort_lru相对普通 LRU:磁盘恢复开销抵消了部分缓存收益,因此该策略保持为可选功能,普通 LRU 仍为默认配置。
七、使用方法
7.1 初始化子模块
克隆父仓库后执行:
单独克隆 Agentrix:
7.2 环境要求
uv和ccache。7.3 构建 vLLM
运行关键测试:
7.4 安装 Benchmark
完整构建说明:AutoDL Build Guide
八、交付材料
九、文档索引
9.1 架构与设计
9.2 实验结果
9.3 构建与平台
十、结果说明
本项目已经覆盖赛题要求的智能体推理流程、KV 生命周期管理、分支共享、Prompt 压缩、工具数据按需加载、分层内存、多模型和异构硬件适配。
实验数据遵循以下口径: