← 全部 AgentX 优化

AgentX 行业影响 · KV cache 层

LMCache

LMCache 是一个开源 KV cache 层,位于 vLLM 等推理引擎之下,按 prefix hash 为 key,把可复用的 KV chunk 存放在 CPU DRAM、本地 NVMe 以及 Mooncake、Redis、S3 等远端后端上。它可以作为 vLLM 原生 offload connector 的替代方案。

旧路径死锁前后完成的请求数对比
120 vs 28
hybrid group 每 token 存储量的下降幅度
~20×
MI350X 上通过的 KV 传输内核测试
56/56

分块的外部 cache 加载

LMCache 的多进程路径针对智能体场景下 cache 搬运的规模与形态做了修改,起点是一种不是"变慢"而是"卡死"的失败:当大量上下文超过 100,000 token 的请求各自在开始前就为整次加载预留 block 时,池会被一群全都在等待、没有一个在推进的请求耗尽。

分块的外部 cache 加载改为按分块预留,于是各次加载可以交错进行并逐步排空。在并发 32 的验证中,旧路径在第 28 个请求处死锁,新路径完成了 120 个请求;并发 48 时即使 KV 池已占用 98.5%,运行依然继续。

示意图:多个 vLLM 实例之上共享一个 LMCache 层,包含 CPU DRAM、经 GDS 访问的本地 NVMe 以及远端存储层,后到的请求可以在另一个实例上读到先前写入的内容。
LMCache 作为按 hash 寻址的 KV 层位于所有实例之下,横跨 CPU DRAM、本地 NVMe 与远端存储,因此后到的请求即使落在另一个实例上,也能读到先前写入的内容。查看原始分辨率图片

搬得更少,也更少挡路

其余改动的目标是减少搬运量以及运行时的干扰。只存储 DeepSeek-V4 hybrid group 中真正有用的部分,把每 token 的存储量降低了近二十倍;sliding-window 预取现在只加载仍然有效的 window,而不是那些永远不会被读取的 window 状态——这与 vLLM 在 offload 上采用的可达性论证相同,只是从存储一侧切入。

随后,按对象组只做一次原生传输调用,消除了 staging 拷贝与内核启动之间反复的 Python 锁交接;这类开销正比于片段数量,而不是其中的字节数。

Open:hybrid 锁记账

当前有两个 LMCache 改动与 AgentX 特别相关,但都仍为 open。hybrid 锁记账修复防止一个请求释放掉另一个请求在共享 sliding-window 或 recurrent 状态 chunk 上的读锁。要复现它需要三个条件同时成立:多个请求共享同一批 chunk、记账按持有者而非按 chunk 进行、并且淘汰确实开始发生。

开启 DRAM offload 的长时间 Kimi-K3 运行同时满足了这三点,产生了数万条告警、损坏的生成结果,并在淘汰开始后最终导致 GPU 崩溃。只要不是长时间、共享且内存吃紧的运行,这个缺陷就一直潜伏。

AMD Instinct 支持

另一条并行的工作线让上述能力在 AMD Instinct 硬件上真正可用。CacheBlend 的非 prefix 复用依赖仅支持 CUDA 的 FlashInfer,因此一个 Triton block-sparse 注意力后端重新实现了它需要的三个内核:带 CSR 索引并输出 log-sum-exp 的 block-sparse attention、causal prefill,以及基于 log-sum-exp 的输出融合;在检测到 ROCm 或缺少 FlashInfer 时会自动切换到这些实现。ROCm Dockerfile 则对齐了 CUDA 的构建与轻量镜像。一个 AMD hipFile 后端扩展了此前只能通过 NVIDIA cuFile 访问存储的 GDS L1 slab 文件层:它通过 ctypes 绑定 ROCm 的 hipFile,并按 torch.version.hip 分发,cuFile 路径保持不变。

最后的缺口是分发方式。CUDA 用户安装预编译 wheel,AMD 用户却要从源码构建。我们与 AMD 一起发布了预编译的 gfx942 与 gfx950 wheel 来补上这一环:它可以装入上游镜像,并在 MI350X 上通过全部 56 项 KV 传输内核测试;它发布到 GitHub release 而非 PyPI,因此直接执行 pip install lmcache 仍会安装 CUDA 版本。一个单行的后续修复把 bind mount 进来的仓库标记为 git safe directory——这个问题只在 CI 中出现,因为容器以 root 身份运行在 runner 所有的检出目录上,setup.py 中的版本探测会拒绝读取它。

DCP 感知的 CPU offload

DCP 感知的 CPU offload 解决了两个在长上下文下必须同时启用的特性之间的直接冲突。启用 decode context parallelism 后,每个 rank 只持有 KV 的一个 stride,因此单个 rank 能保存的内容并不构成可用的 prefix。修复方案在保存前汇集这些带 stride 的分片,并在加载后重新分发。如果不这样做,启用 context parallelism 就会悄悄让 CPU cache 命中失效,而失效的恰恰是催生这两个特性的长 prefix。其验证记录了超过 30,000 次 CPU 命中事件,单个请求的加载量可达数十万 token。

其他项目