LMCache
LMCache 是大模型推理的 KV cache 管理层,把已经算过的 KV cache 卸到 CPU 等存储里再复用。昇腾实现在 LMCache-Ascend。主仓 LMCache 没有昇腾实现。这篇在单卡上安装 vLLM-Ascend 和 LMCache-Ascend,再用 vllm.LLM 做一次离线 KV 卸载。
前置条件
硬件
Atlas 800T / 900 A2 训练系列,Ascend 910B。本文示例为单卡。编译会读 npu-smi 得到芯片型号,A2 上常见 Ascend910B3 或 Ascend910B4。没有 npu-smi 时再设 SOC_VERSION。
软件
类别 |
要求 |
|---|---|
CANN |
toolkit + 驱动固件已安装,并可 |
ATB |
Ascend Transformer Boost。vLLM EngineCore 子进程要加载 |
Python |
3.10 到 3.13,见 vLLM-Ascend 安装文档 |
编译依赖 |
|
vLLM-Ascend |
装带轮子的最新 |
LMCache |
PyPI |
模型 |
阅读本文前,请先按 快速安装昇腾环境 准备好 CANN 与驱动。
本文验证环境:镜像 swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-910b-ubuntu22.04-py3.12,Python 3.12。这一行记录的是本文跑通时的环境。
1. 加载 CANN 环境
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/latest/atb/set_env.sh
export PATH=/usr/local/sbin:/usr/sbin:$PATH
2. 检查环境是否就绪
2.1 确认 NPU 在线
npu-smi info
输出类似:
+------------------------------------------------------------------------------------------------+
| npu-smi 25.5.2 Version: 25.5.2 |
+---------------------------+---------------+----------------------------------------------------+
| NPU Name | Health | Power(W) Temp(C) Hugepages-Usage(page)|
| Chip | Bus-Id | AICore(%) Memory-Usage(MB) HBM-Usage(MB) |
+===========================+===============+====================================================+
| 0 910B4 | OK | 89.9 39 0 / 0 |
| 0 | 0000:c1:00.0 | 0 0 / 0 2922 / 32768 |
+===========================+===============+====================================================+
如果 npu-smi 找不到,回到 快速安装昇腾环境 检查驱动与设备挂载。
2.2 确认工具可用
下面确认 CANN 已加载,并且 python 在 PATH 里。
test -n "$ASCEND_HOME_PATH"
python --version
3. 安装编译依赖
安装 cmake、ninja 和 NUMA 头文件。
apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y cmake ninja-build libnuma-dev
cmake --version
4. 安装 vLLM-Ascend
按 vLLM-Ascend 安装文档 安装配套的 vllm、vllm-ascend 和 triton-ascend。
python -m pip install --retries 3 vllm==0.22.1
python -m pip install \
--extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi/variant \
--extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi \
vllm-ascend==0.22.1rc1
python -m pip install --force-reinstall --no-deps \
--extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi \
triton-ascend==3.2.1
python -c "import importlib.metadata as m
for n in ['torch', 'torch-npu', 'vllm', 'vllm-ascend']:
print(n, m.version(n))"
5. 安装 lmcache
将 <ver> 换成最新 Release 的 tag 去掉开头的 v。
NO_CUDA_EXT=1 python -m pip install --no-build-isolation lmcache==<ver> --no-deps
python -m pip install \
aiofile aiofiles blake3 aiohttp msgspec numpy psutil pyyaml pyzmq \
redis safetensors sortedcontainers transformers huggingface_hub \
prometheus_client py-cpuinfo
python -c "import importlib.metadata as m; print('lmcache', m.version('lmcache'))"
完整输出较长,其中应包含:
...
lmcache <ver>
<ver>是最新 Release 去掉开头的v。
6. 克隆并编译 LMCache-Ascend
将 <ref> 换成最新 Release 的 tag。主仓在 GitHub。子模块和主仓分开拉。
if [ ! -f LMCache-Ascend/csrc/hixl/CMakeLists.txt ]; then
rm -rf LMCache-Ascend
for _ in 1 2 3; do
GIT_TERMINAL_PROMPT=0 GIT_HTTP_VERSION=HTTP/1.1 \
git clone --depth 1 --branch <ref> https://github.com/LMCache/LMCache-Ascend.git LMCache-Ascend && break
rm -rf LMCache-Ascend
sleep 5
done
fi
GIT_TERMINAL_PROMPT=0 GIT_HTTP_VERSION=HTTP/1.1 \
git -C LMCache-Ascend submodule update --init --recursive
git -C LMCache-Ascend describe --tags --exact-match
完整输出较长,其中应包含:
...
<ref>
<ref>是最新 Release 的 tag。
删除已有的 LMCache-Ascend/build,再在当前目录安装 ./LMCache-Ascend。
rm -rf LMCache-Ascend/build
python -m pip install -v --no-build-isolation -e ./LMCache-Ascend
python -c "import lmcache, lmcache_ascend, torch, torch_npu; from lmcache_ascend import _build_info as b; import lmcache_ascend.c_ops; print('lmcache', lmcache.__version__); print('soc', b.__soc_version__); print('c_ops_ok', True); print('npu_available', torch.npu.is_available())"
完整输出较长,其中应包含:
...
lmcache <ver>
soc Ascend910B...
c_ops_ok True
npu_available True
<ver>是最新 Release 去掉开头的v。
soc 应和本机 npu-smi info -t board 的 Chip Name 一致。如果编译缺少 numaif.h,回到第 3 节执行 apt-get install,装上 cmake、ninja-build 和 libnuma-dev,然后从本节重新编译,再继续第 7 节。
7. 用离线 LLM 做一次 KV 卸载
用 vllm.LLM 和 vLLM-Ascend 内置的 LMCacheAscendConnector 做一次离线 KV 卸载。模型是 Hugging Face 上的 Qwen/Qwen2.5-0.5B-Instruct。
参数 |
含义 |
|---|---|
|
|
|
|
|
|
|
|
|
|
用 python 运行下面的脚本。
import os
import sys
os.environ["PYTHONHASHSEED"] = "0"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_LOGGING_LEVEL"] = "INFO"
os.dup2(sys.stdout.fileno(), sys.stderr.fileno())
from huggingface_hub import snapshot_download
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
import lmcache_ascend
import torch
import torch_npu
from lmcache.integration.vllm.utils import ENGINE_NAME
from lmcache.v1.cache_engine import LMCacheEngineBuilder
def main():
os.environ.setdefault("LMCACHE_CHUNK_SIZE", "256")
os.environ.setdefault("LMCACHE_LOCAL_CPU", "True")
os.environ.setdefault("LMCACHE_MAX_LOCAL_CPU_SIZE", "2")
model = snapshot_download("Qwen/Qwen2.5-0.5B-Instruct")
ktc = KVTransferConfig(
kv_connector="LMCacheAscendConnector",
kv_role="kv_both",
)
llm = LLM(
model=model,
enforce_eager=True,
kv_transfer_config=ktc,
max_model_len=512,
gpu_memory_utilization=0.4,
trust_remote_code=True,
)
params = SamplingParams(temperature=0, top_p=0.95, max_tokens=8)
outputs = llm.generate(
["Hello, my name is", "Tell me a short story"],
params,
)
for item in outputs:
print("reply", item.outputs[0].text)
print("current_device", f"npu:{torch.npu.current_device()}")
print("npu_available", torch.npu.is_available())
print("workload_ok", True)
LMCacheEngineBuilder.destroy(ENGINE_NAME)
if __name__ == "__main__":
main()
完整输出较长,其中应包含:
...Platform plugin ascend is activated...Using NPU for LMCache engine...
current_device npu:0
npu_available True
workload_ok True
...
脚本会打印两条 reply。结束时若出现 Assertion failed: pfd.revents,退出码为 0 且上面几行都在即可。
8. 更多文档
存储后端、框架对接和多机部署与社区相同。见 LMCache 文档。昇腾侧仓库见 LMCache-Ascend。