← 全部 AgentX 优化

AgentX 行业影响 · 推理引擎

vLLM

我们与来自 Inferact、Red Hat、NVIDIA 和 AMD 的 vLLM maintainer 一起,以 AgentX 的真实回放器为北极星开展工作,相关修复均已合入上游,其中大部分可以直接迁移到生产环境。

1M 上下文下的 prefix cache 命中率
>95%
hybrid CPU offload 带来的输出吞吐量提升
+81.7%
平均端到端延迟下降
−46.6%

混合注意力 prefix caching

vLLM 改进了 hybrid attention 的 prefix caching,使短生命周期的 sliding-window 分配不再挤掉有价值的长上下文 checkpoint。选择性保留(selective retention)会保住稀疏的回放边界,在 14 个并发请求、上下文最长 100 万 token 的条件下,报告 prefix cache 命中率超过 95%。

同样的可达性策略也应用到了 Mooncake,并移除了不可达的 sliding-window 查找。更早的后续工作还停止 offload 那些永远不会被复用的 sliding-window block,并把 speculative lookahead block 保留在被保住的 prefix 内。

前后对比示意图:未做保留时每个 sliding-window 尾部都会被释放,整个 prefix 必须重算;采用选择性保留后少量 checkpoint 得以存活,prefix 仍可复用。
改动前没有任何 window 尾部存活,因此没有位置可恢复,full-attention KV 虽然驻留但无法使用;采用选择性保留后少量尾部得以保留,prefix cache 命中率超过 95%。查看原始分辨率图片

面向 hybrid 模型的 CPU KV offload

高并发智能体负载必须依赖 offload,而 AgentX 推动的工作让 CPU KV offload 从只支持统一 full-attention 模型,扩展到可用于 hybrid 模型。这一区别很关键:统一模型每个 token 只有一种 KV 布局,connector 用单一 block 几何形状就能描述要保存什么;hybrid 模型同时携带多个 cache group,形状与生命周期各不相同,而假设单一布局的 connector 无法表达某个 block 属于哪一组。结果就是,最需要 offload 的长会话模型反而用不了 offload。

通用的 SimpleCPU connector 最先落地,随后在 ROCm 上启用,并进一步扩展到 DeepSeek-V4 的 hybrid attention:与 prefix 超出 HBM 后直接重算相比,报告输出吞吐量提高 81.7%、平均端到端延迟降低 46.6%。Mooncake 也获得了等价的 hybrid 内存分配支持。

收窄 store 路径

在真实负载上做 profiling 后发现,offload 一旦能跑通,开销就转移到了 store 路径上——写得太多、也太频繁。三条规则收窄了它。

现在,若已有相同内容的传输在途,就跳过这次 store,于是共享同一 prefix 的并发会话只需付一次代价,而不是各付一次;store 只覆盖新生成的 KV 区间,因此延长历史的会话只写增量,而不是每轮重写整个 prefix;store 也不再取决于相同 block 是否仍在 HBM 中,因此已经排定的工作不会因为底下发生淘汰而被丢弃。

vLLM connector 层位于 GPU HBM 与 CPU DRAM 池之间的示意图,标注了相关 connector PR 以及收窄写入的三条 store 侧规则。
位于 GPU HBM 与 CPU DRAM 池之间的 connector 层。store 侧规则会跳过在途的重复传输、只写新生成的 KV,并与 block 是否仍在 HBM 解耦;调度器侧的查找则改为异步。查看原始分辨率图片

把查找移出关键路径

load 路径需要单独调优,因为查找发生在每一次调度决策上,而不只是在真正搬运数据时。把调度器路径中的查找改为异步,可以让 connector 离开单步执行的关键路径,一个 step 不必等待 CPU 侧 cache 查询就能开始接纳新工作。

此后,紧凑的零拷贝查找 key、接收侧并行加载,以及预先构建的 Mooncake key 字符串,进一步消除了残留的 CPU 与传输开销。

长生命周期 hybrid 状态下的正确性与记账

长生命周期的 hybrid 状态还带来了一批固定形状请求很少触及的正确性与记账修复。vLLM 现在按 hybrid cache group 分别发出 cache 事件,正确处理分布式 context 下 store 的 stride,并在分布式 context 与 prefill 下正确计算查找 prefix;相关的 context-parallel 记账改动则让 cache 归属与分片后的 token 区间对齐。

Speculative 状态现在会在合并后的 Mooncake group 之间以及 SimpleCPU coordinator 中正确传递,避免多轮复用的 cache 悄悄丢失 EAGLE 状态。

ROCm:cache 层之下

在 ROCm 侧,近期工作继续深入到 cache 层以下,那里的开销是按层而非按请求计的。当 prefix 能够存活并及时到达之后,剩下的就是 decode 步骤本身;而一次会话中要执行数千次的 decode,会为每一次可以避免的拷贝、每一个不匹配的内核付出代价。三个仍为 open 的改动针对的正是这一层。

一个 Kimi-K3 改动把 KDA decode 结果直接写入该层的输出 buffer,为每个 KDA 层省去一次设备内拷贝;单看收益很小,但它会在每个解码 token 的每一层上重复发生。第二个改动用 AITER 的 sparse-MLA decode 内核替代通用路径,报告 AgentX 输出吞吐量提升 5.22%,token 间延迟大幅下降。

第三个改动很好地说明了"测量形状"有多重要:配套改动把 full-graph 注意力投影切换到调优过的 AITER GEMM,在低并发固定序列上取得 2.3% 的提升。内核级改动可以在统一形状上给出干净的收益,却在智能体 trace 上被 cache 与调度波动淹没。

其他项目