xiaoO AgentOS记忆系统 x LMCache:语义感知 KV Cache 协同说明文档
功能介绍
xiaoO 已经内嵌了 KVCache 感知调度能力,通过配置文件中的 kvcache_enabled 标志位开启,在请求前会根据请求中的 chunk_hashes 来判断是否需要从 KVCache 中读取缓存,在上下文变化时会根据 diff_deleted 方法来判断是否需要删除 KVCache 中的缓存
示例配置
[llm]
provider = "local"
model = "glm4.7"
api_base = "http://localhost:8080/v1"
max_tokens = 20000
context_window = 128000
reasoning_effort = "off"
kvcache_enabled = true
kvcache_debug_enabled = false # 打开 debug 文件生成
LMCache-Ascend 部分说明
环境配置
| 组件 | 版本 / 分支 | 说明 |
|---|---|---|
| 容器镜像 | quay.io/ascend/vllm-ascend:v0.18.0rc1-a3 |
基于 Ascend NPU 的 vLLM 推理容器 |
| LMCache | v0.4.3 |
上游 LMCache 核心,不做任何修改 |
| LMCache-Ascend | main |
本次新增功能所在仓库 |
启动推理服务
# 启动容器
export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3
docker run --privileged \
--name agentos-lmcache-ascend \
--net=host \
--ipc=host \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci6 \
--device /dev/davinci7 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /dev/shm:/dev/shm \
-v /etc/hccn.conf:/etc/hccn.conf \
-v /usr/bin/hccn_tool:/usr/bin/hccn_tool \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /root/.cache:/root/.cache \
-v /root/.pip:/root/.pip \
-v /workspace:/workspace \
-v /mnt/sdb/models:/mnt/sdb/models \
-it $IMAGE bash
# 安装 LMCache
git clone https://github.com/LMCache/LMCache.git
cd LMCache && git checkout v0.4.3 && export NO_CUDA_EXT=1
python3 -m pip install -v --no-build-isolation -e .
# 安装 LMCache-Ascend(语义感知 KV Cache 调度特性已合入master)
git clone https://github.com/LMCache/LMCache-Ascend.git
cd LMCache-Ascend && unset SOC_VERSION
git submodule update --init --recursive
python3 -m pip install -v --no-build-isolation -e .
# 启动glm4.7服务
#!/bin/bash
export PYTHONHASHSEED=0
export HCCL_BUFFSIZE=512
export OMP_PROC_BIND=false
export OMP_NUM_THREADS=1
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export HCCL_OP_EXPANSION_MODE=AIV
export VLLM_ASCEND_BALANCE_SCHEDULING=1
export VLLM_ASCEND_ENABLE_TOPK_OPTIMIZE=1
export VLLM_ASCEND_ENABLE_FLASHCOMM1=1
export VLLM_ASCEND_ENABLE_FUSED_MC2=1
export LMCACHE_CONFIG_FILE="/home/lmcache_config.yaml"
timestamp=$(date "+%Y%m%d%H%M%S")
vllm serve /mnt/sdb/models/GLM-4.7 \
--host 0.0.0.0 \
--port 8080 \
--data-parallel-size 1 \
--tensor-parallel-size 16 \
--seed 1024 \
--served-model-name glm4.7 \
--enable-expert-parallel \
--max-num-seqs 16 \
--max-model-len 133000 \
--max-num-batched-tokens 8192 \
--trust-remote-code \
--gpu-memory-utilization 0.92 \
--async-scheduling \
--tool-call-parser glm47 \
--reasoning-parser glm45 \
--enable-auto-tool-choice \
--load-format=dummy \
--no-enable-prefix-caching \
--compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY"}' \
--kv-transfer-config '{"kv_connector":"LMCacheAscendConnectorV1Dynamic","kv_role":"kv_both", "kv_connector_module_path":"lmcache_ascend.integration.vllm.lmcache_ascend_connector_v1"}' \
2>&1 | tee logs/vllm_base_glm47_$timestamp.log
# 驱逐接口调用示例,从ddr驱逐三个 chunk,如果要从ssd驱逐,需要添加 location,或者不指定 location 将chunk完全删除
curl -s -X POST http://localhost:6999/memory/evict \
-H "Content-Type: application/json" \
-d "{\"chunk_hashes\": [\"-6ed9a448787a198f\",\"-4ca98d8c9e52152e\",\"3f5c45ecbab5e9fd\"], \"locations\": [\"LocalCPUBackend\"]}"
# 预取接口调用示例,从ssd预取三个 chunk到ddr
curl -s -X POST http://localhost:6999/memory/prefetch \
-H "Content-Type: application/json" \
-d "{\"chunk_hashes\": [\"-6ed9a448787a198f\",\"-4ca98d8c9e52152e\",\"3f5c45ecbab5e9fd\"], \"lookup_id\": \"my_prefetch_001\"}"
一、问题背景
大模型推理中,KV Cache 是决定延迟的核心瓶颈。当前 LMCache-Ascend 已经实现了多层存储(GPU HBM → CPU DDR → Local SSD)的自动搬运,但存在一个根本架构短板:
LMCache 知道"缓存存了哪些 KV chunk",但上层记忆系统才知道"哪些 chunk 即将被复用"。
缺失的能力:
- 请求结束后,上层记忆系统无法获知这一批 prompt 对应哪些 KV chunk
- 下一个请求到来前,上层记忆系统无法主动预取即将复用的 chunk
- 无法主动驱逐不再需要的 chunk 以释放 DDR 空间
本方案以 chunk_hash 为全局标识符,在 LMCache-Ascend 与上层记忆系统之间建立两条管控通道:
- 回写通道:请求结束时通过
kv_transfer_params返回chunk_hashes - 控制通道:通过 REST API
/memory/prefetch和/memory/evict接收上层调度指令
二、总体架构
上层记忆系统
│
│ "继续讨论刚才的话题" → intent → 查映射表 → [hash_a, hash_b]
│
├── POST /memory/evict(hash_old) ──→ 驱逐无用 chunk
├── POST /memory/prefetch(hash_a, hash_b) ──→ 预取复用 chunk
│
▼
vLLM 推理引擎
│ 推理结束后返回 kv_transfer_params:
│ {"chunk_hashes": ["hash_x", "hash_y", ...]}
│
▼
LMCache-Ascend 存储引擎 (GPU HBM ←→ CPU DDR ←→ SSD)
chunk_hash 是贯穿全链路的全局标识符:
- 确定性:相同 token 序列 + 相同
PYTHONHASHSEED→ 相同 hash - 前缀依赖:hash[i] = blake3(hash[i-1] || tokens[cs:ce]),链式结构防篡改
- 纯 CPU:blake3 微秒级完成,请求结束重算亚毫秒开销
三、chunk_hash 的返回机制
3.1 数据流
Prefill 阶段:
lookup_client.lookup(token_ids=prompt_tokens)
├── token_database.process_tokens() → 计算 [hash_0, hash_1, ...]
├── _lookup_hashes_cache[req_id] = [hex(h) for h in hashes] ← 缓存到内存
└── ZMQ → Worker → 后端查询/预取
... 推理执行 ...
请求结束:
request_finished()
├── lookup_client.get_cached_hashes(req_id) ← 直接 pop 缓存,零计算
└── 返回 {"chunk_hashes": ["abc01", "def02", ...]} → kv_transfer_params
3.2 实现细节
改动 1:LMCacheAscendAsyncLookupClient(新建文件)
继承 LMCacheAsyncLookupClient,在 lookup() 方法中缓存 chunk hashes:
class LMCacheAscendAsyncLookupClient(LMCacheAsyncLookupClient):
def lookup(self, token_ids, lookup_id, ...):
# ... 原有逻辑 ...
# 新增一行:缓存 hex hashes
self._lookup_hashes_cache[lookup_id] = [f"{h:x}" for h in hashes]
# ... 其余不变 ...
def get_cached_hashes(self, lookup_id):
# pop = 取走即清理,不泄漏内存
return self._lookup_hashes_cache.pop(lookup_id, None)
关键设计决策:
get_cached_hashes使用.pop()— 取走即清理clear_lookup_status不清 hashes 缓存 — 因为它在update_state_after_alloc()中被调用(早于request_finished),提前清理会导致缓存丢失__init__禁止直接实例化 — 通过from_existing()原地升级已创建的实例,避免 ZMQ socket/线程状态丢失
改动 2:vllm_v1_adapter.py — _upgrade_lookup_client()
在 connector 初始化时将已创建的 LMCacheAsyncLookupClient 原地升级:
LMCacheManager.__init__()
└─ factory.create_lookup_client()
→ LMCacheAsyncLookupClient() ← 普通类
→ HitLimitLookupClient(client) ← 装饰器
→ ChunkStatisticsLookupClient(client) ← 装饰器
LMCacheAscendConnectorV1Impl.__init__()
└─ super().__init__() # 上述 lookup_client 创建完成
└─ _upgrade_lookup_client()
└─ 穿透装饰链 → 找到 LMCacheAsyncLookupClient
└─ from_existing() → __class__ = LMCacheAscendAsyncLookupClient
改动 3:vllm_v1_adapter.py — request_finished() 返回 chunk_hashes
def request_finished(self, request, block_ids):
_, return_params = super().request_finished(request, block_ids)
if hasattr(request, "all_token_ids") and request.all_token_ids:
# 优先:从 lookup_client 缓存取(零成本)
inner = self.lookup_client
while hasattr(inner, "actual_lookup_client"):
inner = inner.actual_lookup_client
if hasattr(inner, "get_cached_hashes"):
new_hashes = inner.get_cached_hashes(request.request_id)
# 兜底:缓存未命中时才全量重算
if new_hashes is None:
new_hashes = [
f"{hash_val:x}"
for _, _, hash_val in
self._token_db.process_tokens(
tokens=request.prompt_token_ids or request.all_token_ids
)
]
if new_hashes:
return_params = return_params or {}
return_params["chunk_hashes"] = new_hashes
return delay_free, return_params
两阶段覆概率分析:
- 首页 | 命中
get_cached_hashes()| ~100% - 兜底 |
ChunkedTokenDatabase.process_tokens()| 极低(首次预填逻辑异常时触发)
四、REST API:prefetch 与 evict
4.1 接口定义
POST /memory/prefetch
上层记忆系统将已命中的 chunk_hashes 传给 LMCache-Ascend,LMCache-Ascend 提前从 SSD 加载到 DDR。
// Request
POST /memory/prefetch
{
"chunk_hashes": ["a1b2c3d4...", "e5f6g7h8..."],
"lookup_id": "req_001"
}
// Response
{
"status": "prefetch_started",
"lookup_id": "req_001",
"num_chunks": 2
}
内部调用链:
prefetch()
→ _hashes_to_keys() # hex hash → CacheEngineKey
→ storage_manager.async_lookup_and_prefetch()
→ batched_async_contains() # 检查 key 是否存在
→ batched_get_non_blocking() # 异步 SSD→DDR 传输
→ prefetch_all_done_callback()
→ hot_cache 注册 # 后续检索直读 DDR
→ selective unpin # 按 backend 释放 pin
→ send_response_to_scheduler()
POST /memory/evict
上层记忆系统驱逐不需要的 KV cache。
// Request
POST /memory/evict
{
"chunk_hashes": ["a1b2c3d4..."],
"locations": ["LocalCPUBackend"] // 可选,默认所有 backend
}
// Response
{
"status": "success",
"num_evicted": 1
}
4.2 多进程转发
在 vLLM v1 多进程模式下,Scheduler 进程不持有 LMCacheEngine,但持有 API Server。转发判断逻辑:
_get_engine(request):
if lmcache_engine exists: → 直接处理(Worker 进程)
if port_offset == 0: → 转发到所有 Worker(Scheduler 进程)
else: → 返回 503
转发实现在 _forward_to_all_workers() 中完成——对所有 Worker 端口并行发送 urllib 请求,合并结果:
AgentOS → POST scheduler:6999/memory/prefetch
├── Worker 0:7000/memory/prefetch
└── Worker 1:7001/memory/prefetch
← 合并响应
4.3 端口编排
port_start = 6999 (默认)
Scheduler (port_offset=0) → :6999
Worker 0 (port_offset=1) → :7000
Worker 1 (port_offset=2) → :7001
...
五、storage_manager 增强
5.1 hot_cache 注册
prefetch 完成后,将 MemoryObj 注册到 LocalCPUBackend.hot_cache:后续同步检索和 vLLM lookup 通过前缀匹配直接命中,绕过磁盘 I/O。
5.2 selective unpin
问题:batched_async_contains(pin=True) 只在特定 backend 上 pin 了 key,如果全局 batched_unpin(all_keys) 会对所有 backend 尝试 unpin,引发 double-unpin 错误。
解决:在 async_lookup_and_prefetch 中追踪 loading_task_backends,prefetch 完成后只在这些 backend 上调 unpin()。
5.3 DiskBackend pin 释放
LocalDiskBackend.batched_get_non_blocking() 会在新分配的 DDR buffer 上调 memory_obj.pin()。对象注册到 hot_cache 后,后续检索不再经过 DiskBackend,此 pin 可以释放。
六、交互时序全景
上层记忆系统 vLLM LMCache
│ │ │
│ ① 语义分析 │ │
│ prompt → intent │ │
│ 查映射表 → [h1, h2] │ │
│ │ │
│ ② 驱逐旧话题 chunk │ │
│ ─── POST /memory/evict ────────────────────────────▶ │
│ │ batched_remove() │
│ │ │
│ ③ 预取复用 chunk │ │
│ ─── POST /memory/prefetch ──────────────────────────▶ │
│ │ async_lookup_and_prefetch │
│ │ SSD→DDR (异步, 非阻塞) │
│ │ hot_cache 注册 │
│ │ │
│ ④ 发起推理请求 │ │
│ ─── POST /chat ────────▶ │ │
│ │ lookup + retrieve ──────▶ │
│ │ cache hit │
│ ◄── SSE stream ──────── │ │
│ │ │
│ ◄── [DONE] + │ │
│ {"chunk_hashes": │ │
│ ["hx","hy",...]} │ │
│ │ │
│ ⑤ 回写映射表 │ │
│ intent → [hx, hy, ...] │ │
步骤 ②③ 在推理请求之前执行,prefetch 是异步非阻塞的,不影响任何推理主流程。
七、修改清单
LMCache 核心:零改动
LMCache v0.4.3 代码不做任何修改。所有新增能力通过 lmcache-ascend 的 monkey-patch 机制注入。
lmcache-ascend 新增文件
| 文件 | 行数 | 功能 |
|---|---|---|
v1/lookup_client/lmcache_ascend_async_lookup_client.py |
~70 | 子类:lookup 时缓存 chunk hashes,提供 get_cached_hashes() |
v1/lookup_client/__init__.py |
1 | 包初始化 |
v1/internal_api_server/__init__.py |
1 | 包初始化 |
v1/internal_api_server/memory/__init__.py |
1 | 包初始化 |
v1/internal_api_server/memory/memory_api.py |
~299 | /memory/prefetch + /memory/evict REST API |
lmcache-ascend 修改文件
| 文件 | 改动 | 功能 |
|---|---|---|
__init__.py |
+2 个 _patch_* 函数 + 调用 |
_patch_storage_manager_agentos:hot_cache + selective unpin;_patch_api_server_agentos:注册 memory router + app.state 端口变量 |
integration/vllm/vllm_v1_adapter.py |
_upgrade_lookup_client() + request_finished() 增强 |
注入 Ascend 子类 + 返回 chunk_hashes |
八、测试验证
代码质量
pre-commit run --all-files 所有 hook 通过:
- check-spdx-header ✅
- ruff (E/F/B) ✅
- isort ✅
- ruff-format ✅
- codespell ✅
- clang-format ✅
单元测试
| 测试文件 | 覆盖内容 |
|---|---|
tests/v1/lookup_client/test_lmcache_ascend_async_lookup_client.py |
from_existing 升级、get_cached_hashes 取/弹/空/选择性清理、lookup 缓存 hex hash/覆盖/空 tokens、clear_lookup_status 不清 hashes |
tests/v1/internal_api_server/test_memory_api.py |
_hashes_to_keys 键构造、prefetch 缺字段/无 storage/成功路径、evict 缺字段/无 storage/成功路径 |