LMCache

LMCache 是大模型推理的 KV cache 管理层,把 KV cache 卸到 CPU 等存储里复用。昇腾实现在 LMCache-Ascend。

前置条件

硬件

Atlas 800T / 900 A2 训练系列,Ascend 910B。本文示例为单卡。

软件

类别

要求

CANN

toolkit + 驱动固件已安装,并可 source set_env.sh

ATB

Ascend Transformer Boost。vLLM EngineCore 子进程要加载 libatb.so

Python

3.10 到 3.13,见 vLLM-Ascend 安装文档

编译依赖

cmake、ninja、libnuma-dev,见下文安装

vLLM-Ascend

装带轮子的最新 vllm-ascend,vllm 装与之主版本号相同的版本,见下文安装

模型

Qwen/Qwen2.5-0.5B-Instruct

阅读本文前,请先按 快速安装昇腾环境 准备好 CANN 与驱动。

本文验证环境:镜像 swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-910b-ubuntu22.04-py3.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

完整输出较长,其中应包含:

Python 3...

3. 安装编译依赖

安装 cmake、ninja 和 NUMA 头文件。

apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y cmake ninja-build libnuma-dev
cmake --version

完整输出较长,其中应包含:

...
cmake version 3...

4. 安装 vLLM-Ascend

按 vLLM-Ascend 安装文档 安装配套的 vllm、vllm-ascend 和 triton-ascend。vllm-ascend 会把 torch 换成昇腾索引的 2.10.0 构建,最后再装回与 torch_npu 同版本的 CPU 构建。

python -m pip install --retries 3 vllm==0.23.0
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.23.0
python -m pip install --retries 3 \
  --index-url https://download.pytorch.org/whl/cpu \
  --extra-index-url https://pypi.org/simple \
  torch==2.10.0+cpu torchvision==0.25.0+cpu torchaudio==2.10.0+cpu
python -m pip install --force-reinstall --no-deps \
  --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi \
  triton-ascend==3.2.2
python -c "import importlib.metadata as m
for n in ['torch', 'torch-npu', 'vllm', 'vllm-ascend']:
    print(n, m.version(n))"

完整输出较长,其中应包含:

...
torch ...+cpu
torch-npu ...
vllm ...
vllm-ascend ...

5. 安装 LMCache

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 cachetools cryptography fastapi \
  httpx numba uvicorn opentelemetry-api opentelemetry-sdk \
  opentelemetry-exporter-otlp opentelemetry-exporter-prometheus \
  grpcio protobuf
python -c "import importlib.metadata as m; print('lmcache', m.version('lmcache'))"

完整输出较长,其中应包含:

...
lmcache <ver>

Note

<ver> 是 LMCache 最新的 Release tag。

6. 克隆并编译 LMCache-Ascend

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>

删除已有的 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

Note

<ref> 是 LMCache-Ascend 的最新 Release 的 tag;<ver> 是 LMCache 最新的 Release tag。

soc 应和本机 npu-smi info -t board 的 Chip Name 一致。如果编译缺少 numaif.h,回到第 3 节执行 apt-get install,装上 cmake、ninja-build 和 libnuma-dev,然后从本节重新编译,再继续第 7 节。

7. 启动 LMCache 服务并做一次 KV 卸载

先启动独立的 lmcache server,再用 vllm.LLM 加 LMCacheMPConnector 做一次离线 KV 卸载。模型是 Hugging Face 上的 Qwen/Qwen2.5-0.5B-Instruct。服务端的 --chunk-size 128 与 vLLM 的 block size 对齐。

参数

含义

temperature

0,固定本次生成

top_p

0.95

max_tokens

8

max_model_len

512

gpu_memory_utilization

0.4

用 python 运行下面的脚本。

import os
import socket
import subprocess
import sys
import time

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


def wait_for_port(port, timeout=180):
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=2):
                return
        except OSError:
            time.sleep(2)
    raise RuntimeError(f"lmcache server port {port} not ready in {timeout}s")


def main():
    server = subprocess.Popen(
        [
            "lmcache",
            "server",
            "--chunk-size",
            "128",
            "--l1-size-gb",
            "5",
            "--eviction-policy",
            "LRU",
        ]
    )
    try:
        wait_for_port(5555)
        model = snapshot_download("Qwen/Qwen2.5-0.5B-Instruct")
        ktc = KVTransferConfig(
            kv_connector="LMCacheMPConnector",
            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)
    finally:
        server.terminate()
        try:
            server.wait(timeout=30)
        except subprocess.TimeoutExpired:
            server.kill()


if __name__ == "__main__":
    main()

完整输出较长,其中应包含:

...Platform plugin ascend is activated...Creating transfer context (device_type=npu, mode=lmcache_driven)...
reply  Alex and I am a 17
reply  about a person who discovers they have the
current_device npu:0
npu_available True
workload_ok True
...

脚本会打印两条 reply。结束时若出现 Assertion failed: pfd.revents,退出码为 0 且上面几行都在即可。

8. 更多文档

存储后端、框架对接和多机部署与社区相同。见 LMCache 文档。昇腾侧仓库见 LMCache-Ascend。