开发者
下载
KVCache SSD分级缓存参考实践

KVCache SSD分级缓存参考实践

性能调优模型推理计算商业SMECE

发表于: 2026/06/30

1. 非商用声明

该文档提供的内容为参考实践,仅供用户参考使用,用户可参考实践文档构建自己的软件,按需进行安全、可靠性加固,但不建议直接将相关Demo或镜像文件集成到商用产品中。

2. 方案介绍

2.1 背景

在个人AI助手(OpenClaw/Hermes)多轮连续对话、AI 编程超长上下文两大典型业务场景下,KV Cache 复用率可达 80% 以上。当前仅依靠 HBM+DDR 内存承载 KV Cache,受限于内存物理容量上限,超长序列场景缓存有效命中率难以进一步提升;叠加内存硬件价格持续上行的成本约束,现有的HBM+DRAM的KVCache已无法充分发挥高复用率带来的性能收益。为此引入 NVMe SSD 作为三级存储扩容 KV Cache 存储空间,借助分层存储架构补齐超长序列场景缓存容量短板,进一步提高整体缓存命中率,降低TTFT。

2.2 方案简介

本参考实践在原有vllm-ascend+MemCache的基础上,增加UBS-IO提供SSD池化功能,支持HBM+DDR+SSD三级KVCache缓存方案。

整体架构如图所示:

相关组件代码仓链接为:

MemCache:https://gitcode.com/Ascend/memcache

MemFabric:https://gitcode.com/Ascend/memfabric_hybrid

2.3 典型配置

本特性支持的硬件为:Atlas 300I A2、Atlas 800I A2和Atlas 800I A3。

表1 软件版本

软件/镜像版本说明
Ascend HDK25.5.0NPU驱动固件,要求不低于25.5.0
CANN8.5.0昇腾异构计算架构,要求CANN 8.5.0及之后版本
vLLM Ascendv0.23.0rc1昇腾硬件插件,用于在 Ascend NPU 上运行 vLLM,建议不低于v0.23.0rc1
Docker18.09.0用于部署容器,推荐18.09.x~28.5.1版本
MemFabricv1.2.0内存池化基础软件,用于实现DRAM与显存混合池化,提供极简的内存访问接口和高性能的内存直接访问能力,支撑多种场景下的数据共享与传输
MemCachev1.2.0针对AI推理场景设计的高性能分布式KVCache存储引擎
AISBenchv3.0-20251219-master大模型推理性能测评工具,不影响服务部署
OSopenEuler 22.03 (LTS-SP4)物理机OS,以昇腾解决方案要求为准。建议系统内核不低于5.10。

表2 推荐使用的模型列表

部署模式支持模型
Atlas 800I A2 单机MinMax-M2.5-w8a8
Atlas 800I A2 双机GLM5-w8a8、GLM5.1-w8a8
Atlas 800I A3 单机MinMax-M2.5-w8a8、GLM5-w8a8、GLM5.1-w8a8
Atlas 800I A3 双机GLM5-w8a8、GLM5.1-w8a8
Atlas 300I A2 单机双卡Qwen3-32B-w8a8

3. 部署指导

3.1 前置条件

已完成OS、昇腾HDK、Docker安装,并已下载对应大模型的权重文件和测试数据集,集群(多机互联)场景下需完成参数面组网,确保网络互通。

权重下载链接:https://modelscope.cn/models/Eco-Tech/MiniMax-M2.5-w8a8-QuaRothttps://modelscope.cn/models/Eco-Tech/GLM-5-w8a8

因为有镜像拉取和安装依赖,推荐服务器处于联网环境。

KV Cache 三级缓存推荐本地 NVMe 磁盘部署,若要实现较好的性能收益,建议配置使用独占的NVMe高速盘,若没有独立的盘,也可复用有足够空间的已挂载盘,详细磁盘配置方式见部署章节。

3.2 部署流程

整体部署流程包含如下步骤:

1. 磁盘准备;

2. 创建容器;

3. 组件安装;

4. 配置文件;

5. 拉起服务;

3.3 具体步骤

3.3.1 磁盘准备

配置详情可参考此链接中的设备配置章节。

若提供给分级缓存使用的NVMe盘为独占盘且未挂载,可直接使用,可用该示例检查是否挂载(替换为真实磁盘):lsblk /dev/nvme0n1,回显结果中MOUNTPOINTS没有值则为未挂载。

如果无独占的空闲盘,可使用loop设备方案,复用已有磁盘空间,操作如下:

# 创建并预分配 1TB 镜像文件,容量可按需调整
sudo fallocate -l 1T /data/ubsio_disk.img

# 创建 loop 设备并输出系统实际分配的设备路径
LOOP_DEV=$(sudo losetup --find --show --direct-io=on /data/ubsio_disk.img)
echo "${LOOP_DEV}"
# 示例输出:/dev/loop0
 注意,需要记住此处的设备值,后续拉起服务时需要指定。

取消 loop 设备的方式如下,注意, 仅在确认不再需要 loop 模拟盘且相关服务已经停止时执行,卸载后缓存数据将丢失。

# 1. 查看当前 loop 设备及其关联文件
losetup -a

# 2. 卸载 loop 设备,替换为实际设备路径
sudo losetup -d /dev/loop0

# 3. 删除镜像文件(可选)
sudo rm -f /data/ubsio_disk.img

3.3.2 创建容器

本文介绍使用容器部署vLLM Ascend的方式,镜像下载地址:https://quay.io/repository/ascend/vllm-ascend?tab=tags

单机部署和双机部署,均参考如下操作步骤:

# docker拉取镜像(以A2的openeuler镜像为例)
# 也可以采用下载镜像的tar包,然后使用docker load -i xx.tar的方式安装镜像
docker pull quay.io/ascend/vllm-ascend:v0.23.0rc1  

# 根据实际下载镜像版本进行替换
export IMAGE=quay.io/ascend/vllm-ascend:v0.23.0rc1

# 查看NPU设备号
ls /dev/davinci*
# 根据查询到的davinci设备号,配置容器创建时的device参数
/dev/davinci0  /dev/davinci1  /dev/davinci2  /dev/davinci3  /dev/davinci4  /dev/davinci5  /dev/davinci6  /dev/davinci7  /dev/davinci_manager

# 构建容器,注意权重挂载路径根据真实情况替换,NPU卡数据根据实际情况替换(此处以8卡为例)
docker run -itd --net=host --privileged=true \  # 若不启用特权容器,需对磁盘分区做设备映射挂载。
       --shm-size=256g \ # 如果为300I A2场景(单机),配置protocol为host_shm,需要保证shm-size值大于二级缓存总内存
       --name vllm-ascend-usbio \        # 依据实际情况配置
       -v /usr/local/dcmi:/usr/local/dcmi \
       -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
       -v /usr/local/Ascend/driver/:/usr/local/Ascend/driver \
       -v /etc/ascend_install.info:/etc/ascend_install.info \
       -v /home:/home \     # 根据真实情况替换,执行时删除注释,不要有空格
       -v /root/.cache:/root/.cache \
       --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 \
       -it $IMAGE bash

# 进入容器
docker exec -it vllm-ascend-usbio bash

3.3.3 安装memcache与memfabric

可直接使用pip命令进行安装,参考命令如下(如果有网络问题,可加-i https://pypi.tuna.tsinghua.edu.cn/simple/参数使用镜像源):

# 1. 先安装 MemFabric
pip install memfabric_hybrid==1.2.0

# 2. 再安装包含 UBSIO 的 MemCache
pip install memcache_hybrid==1.2.0

# 3. 验证 Python 包安装成功,正常时回显为“import-ok”
python3 -c "import memfabric_hybrid, memcache_hybrid; print('import-ok')"

3.3.4 修改配置文件

详细的参数含义可参考此链接

配置文件会随着pip install自动下载,下载路径可通过如下命令获取:

# 相关配置文件会安装到对应目录下,获取命令如下
MMC_CONFIG_DIR="$(python3 -c 'from pathlib import Path; import memcache_hybrid; print(Path(memcache_hybrid.__file__).resolve().parent / "config")')"
echo "$MMC_CONFIG_DIR"

目录下会有mmc-meta.conf和mmc-local.conf配置文件,最小修改如下:

1. mmc-meta.conf配置(修改命令:vim $MMC_CONFIG_DIR/mmc-meta.conf):

# 将IP配置为主节点IP,端口可根据实际情况修改
ock.mmc.meta_service_url = tcp://xx.xx.xx.xx:5000
ock.mmc.local_service.config_store_url = tcp://xx.xx.xx.xx:6000
ock.mmc.meta_service.metrics_url = http://xx.xx.xx.xx:8000

# 空间使用率达到高水位时,Put 操作触发数据淘汰
ock.mmc.evict_threshold_high = 70

# 数据淘汰完成后的目标空间使用率,即低水位
ock.mmc.evict_threshold_low = 60

# DRAM 空间使用率达到该水位时,淘汰时会快速删除数据释放内存,降低内存占用压力。
ock.mmc.rewarm.dram_watermark = 95

2. mmc-local.conf配置(修改命令:vim $MMC_CONFIG_DIR/mmc-local.conf):

# 将IP配置为主节点IP,端口可根据实际情况修改
ock.mmc.meta_service_url = tcp://xx.xx.xx.xx:5000
ock.mmc.local_service.config_store_url = tcp://xx.xx.xx.xx:6000

# 根据不同硬件设备,配置为对应值:
# 800I A3场景(单双机),配置为device_sdma
# 800I A2场景(单双机),配置为device_rdma
# 300I A2场景(单机),配置为host_shm
ock.mmc.local_service.protocol = device_rdma

# 每device分配的内存空间,即DRAM二级缓存大小,可根据实际场景调整,调整方式后文单独说明
ock.mmc.local_service.dram.size = 20GB  

# 启用 SSD 池化特性开关
ock.mmc.local_service.storage.enabled = true

# 配置 UBS IO 用作本地 SSD 缓存的块设备路径,支持整盘、分区或 loop 设备;多个路径使用英文冒号分隔;配置的设备必须由 UBS IO 独占,且不能存在挂载点
ubsio.disk.path = /dev/loop0:/dev/loop1:/dev/loop2

# UBS IO 单进程内存池大小,单位 GB。典型配置:推荐配置为10GB,建议不小于5G。
ubsio.mem.size_in_gb = 10

# 启动vLLM服务的device数量
ubsio.standalone.device_count = 8

# 非灵渠场景,建议配置此功能为false
ubsio.bdm.batch_read.standalone.use_scratch_pool = false

此外,当前版本建议new_disk参数配置为true,后续版本会进行更新。

当前方案数据传输路径为从HBM(NPU)到DRAM(二级缓存),再到SSD(三级缓存)。所以使能三级缓存需要同时使能二级缓存,也就是服务器除了需要预留SSD磁盘空间,还需要预留一定的内存空间。

对内存的主要空间占用有二级缓存大小(ock.mmc.local_service.dram.size)和ubsio进程占用(ubsio.mem.size_in_gb),2个参数配置的都是每device占用,总空间占用需要乘以device数量,计算公式如下:

总内存占用 = device_count *(dram.size + mem.size_in_gb)

在SSD缓存场景,DRAM主要用于KVCache数据传输,当前给的推荐值是20GB,可以继续增加。如果想要减少此配置节约内存,可根据实际服务和请求所需的KVCache大小进行调整,从vLLM服务日志中抓取相关信息进行估算。

# 每device的KVCache存储空间大小
Available KV cache memory: 13.88 GiB

# 服务总KVCache大小
GPU KV cache size: 1,217,457 tokens

# 使用KVCache信息和请求负载,预设淘汰水位系数为1.5,推荐使用如下公式估算所需内存空间
每device所需内存 = 并发 * 请求平均长度 * 1.5 * KVC_memory / KVC_size

3.3.5 拉起服务

场景一:单机服务

拉起单机vLLM Ascend服务,增加开启三级缓存特性的参数,原有业务的配置参数(如卡号,权重等)保持不变。

1. 启动meta_service

MMC_CONFIG_DIR="$(python3 -c 'from pathlib import Path; import memcache_hybrid; print(Path(memcache_hybrid.__file__).resolve().parent / "config")')"
export MMC_META_CONFIG_PATH="$MMC_CONFIG_DIR/mmc-meta.conf"
python3 -c "from memcache_hybrid import MetaService; MetaService.main()"

2. 启动vllm服务

在原有vLLM启动脚本中,增加如下内容:

# 增加环境变量
MMC_CONFIG_DIR="$(python3 -c 'from pathlib import Path; import memcache_hybrid; print(Path(memcache_hybrid.__file__).resolve().parent / "config")')"
export MMC_LOCAL_CONFIG_PATH="$MMC_CONFIG_DIR/mmc-local.conf"

# vLLM服务启动命令增加如下参数
--kv-transfer-config \
'{
    "kv_connector": "AscendStoreConnector",
    "kv_role": "kv_both",
    "kv_load_failure_policy": "recompute", # 需要模型支持,当前Qwen3.5和DeepSeekV4不支持,需要等待vLLM版本更新
    "kv_connector_extra_config":{
    "backend": "memcache",
    "lookup_rpc_port": "0"
    }
}'

当遇到KVCache卸载错误时,可在kv-transfer-config中配置"kv_load_failure_policy": "recompute"参数,自动使能重计算替代报错,参考链接。注意,该功能当前并非所有模型都支持。

以下给出完整示例:

nic_name="eth0"        # 根据实际网卡名配置
local_ip="xx.xx.xx.xx" # 配置为当前节点IP
export HCCL_IF_IP=$local_ip
export GLOO_SOCKET_IFNAME=$nic_name
export TP_SOCKET_IFNAME=$nic_name
export HCCL_SOCKET_IFNAME=$nic_name

export HCCL_OP_EXPANSION_MODE="AIV"
export HCCL_BUFFSIZE=512

sysctl -w vm.swappiness=0
sysctl -w kernel.numa_balancing=0
sysctl kernel.sched_migration_cost_ns=50000

export TASK_QUEUE_ENABLE=1
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export VLLM_ASCEND_ENABLE_FLASHCOMM1=1
export HCCL_INTRA_PCIE_ENABLE=1
export HCCL_INTRA_ROCE_ENABLE=0
export OMP_PROC_BIND=false
export OMP_NUM_THREADS=8
export VLLM_TORCH_PROFILER_WITH_STACK=0
export VLLM_TORCH_PROFILER_DIR="./profile"

MMC_CONFIG_DIR="$(python3 -c 'from pathlib import Path; import memcache_hybrid; print(Path(memcache_hybrid.__file__).resolve().parent / "config")')"
export MMC_LOCAL_CONFIG_PATH="$MMC_CONFIG_DIR/mmc-local.conf"

LOG_FILE="logs/mix_$(date +%Y-%m-%d-%H-%M).log"
# vllm启动参数参考出口基线配置
vllm serve /data/weights/MiniMax-M2.5-W8A8 \    # 权重路径根据实际调整
    --served-model-name MiniMax-M2.5 \
    --host 0.0.0.0 \
    --port 7000 \
    --trust-remote-code \
    --tensor-parallel-size 8 \
    --quantization ascend \
    --enable-expert-parallel \
    --enable-chunked-prefill \
    --enable-prefix-caching \
    --max-num-seqs 200 \
    --seed 1024 \
    --max-num-batched-tokens 6144 \
    --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY","cudagraph_capture_sizes": [4,16,40,80,160,256,400,600,800]}' \
    --gpu-memory-utilization 0.9 \
    --enable-auto-tool-choice \
    --tool-call-parser minimax_m2 \
    --reasoning-parser minimax_m2_append_think \
    --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./profile", "torch_profiler_with_stack": false}' \
    --additional-config '{"enable_cpu_binding":true}' \
    --model-loader-extra-config '{"enable_multithread_load":true,"num_threads":16}' \
    --speculative_config '{"method": "eagle3", "model": "/data/weights/MiniMax-M2.5-eagel-0318/",  "num_speculative_tokens":3}' \
    --kv-transfer-config \
    '{
        "kv_connector": "AscendStoreConnector",
        "kv_role": "kv_both",
        "kv_load_failure_policy": "recompute",
        "kv_connector_extra_config":{
        "backend": "memcache",
        "lookup_rpc_port": "0"
        }
    }' 2>&1 | tee ${LOG_FILE} &

场景二:双机服务

双机服务开启三级缓存特性,只需要在主节点拉起meta_service,然后在原有双机拉起命令上,增加开启三级缓存特性的参数。拉起命令参考如下,meta_service只有主节点拉起,vLLM服务主从节点都需要拉起。

主从节点vLLM服务参数配置基本相同,不同的有:网卡(nic_name),本机IP(local_ip),data-parallel-start-rank参数(主节点为0,从节点为1),headless参数(只有从节点配置)。相关内容已标记到注释,执行时需从模板删除注释。

主节点:

1. 启动meta_service(只在主节点启动)

MMC_CONFIG_DIR="$(python3 -c 'from pathlib import Path; import memcache_hybrid; print(Path(memcache_hybrid.__file__).resolve().parent / "config")')"
export MMC_META_CONFIG_PATH="$MMC_CONFIG_DIR/mmc-meta.conf"
python3 -c "from memcache_hybrid import MetaService; MetaService.main()"

2. 启动vllm服务(主从节点都启动)

nic_name="eth0"  # 根据实际网卡名配置,主从节点不同
local_ip="xx.xx.xx.xx" # 配置为当前节点IP,主从节点不同
MASTER_IP="xx.xx.xx.xx" # 配置为主节点IP,主从节点相同
export HCCL_IF_IP=$local_ip
export GLOO_SOCKET_IFNAME=$nic_name
export TP_SOCKET_IFNAME=$nic_name
export HCCL_SOCKET_IFNAME=$nic_name
export HCCL_CONNECT_TIMEOUT=300
export HCCL_OP_EXPANSION_MODE="AIV"
export OMP_PROC_BIND=false
export OMP_NUM_THREADS=10
export HCCL_BUFFSIZE=200
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export VLLM_ASCEND_ENABLE_MLAPO=1
export HCCL_INTRA_ROCE_ENABLE=1
export TASK_QUEUE_ENABLE=1
export CPU_AFFINITY_CONF=1

LOG_FILE="logs/mix_$(date +%Y-%m-%d-%H-%M).log"

MMC_CONFIG_DIR="$(python3 -c 'from pathlib import Path; import memcache_hybrid; print(Path(memcache_hybrid.__file__).resolve().parent / "config")')"
export MMC_LOCAL_CONFIG_PATH="$MMC_CONFIG_DIR/mmc-local.conf"

vllm serve /data/weights/GLM-5-w8a8-new \
--host 0.0.0.0 \
--port 7000 \
--tool-call-parser glm47 \
--reasoning-parser glm45 \
--enable-auto-tool-choice \
--data-parallel-size 2 \
--data-parallel-size-local 1 \
--data-parallel-start-rank 0 \ # 主节点配置为0,从节点配置为1
--headless \ # 只有从节点需要headless参数
--data-parallel-address ${master_ip} \
--data-parallel-rpc-port 14891 \
--tensor-parallel-size 8 \
--quantization ascend \
--enable-expert-parallel \
--seed 1024 \
--max-num-seqs 64 \
--served-model-name glm-5 \
--max-model-len 35000 \
--max-num-batched-tokens 4096 \
--trust-remote-code \
--gpu-memory-utilization 0.92 \
--enable-chunked-prefill \
--enable-prefix-caching \
--async-scheduling \
--additional-config '{
  "multistream_overlap_shared_expert":true,
  "fuse_qknorm_rope": false,
  "fuse_muls_add": true,
  "enable_npugraph_ex": true
}' \
--compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [1,4,8,12,16,20,24,28,32,36,48,60,72]}' \
--speculative-config '{"num_speculative_tokens": 3, "method": "deepseek_mtp"}' \
--kv-transfer-config \
    '{
        "kv_connector": "AscendStoreConnector",
        "kv_role": "kv_both",
        "kv_load_failure_policy": "recompute",
        "kv_connector_extra_config":{
        "backend": "memcache",
        "lookup_rpc_port": "0"
        }
    }' 2>&1 | tee ${LOG_FILE} &

3.3.6 停止服务

# 容器内,停止服务进程
pkill -9 -f python
pkill -9 -f VLLM
pkill -9 -f mmc

4. 典型场景参考实践

按照部署章节指导,拉起服务后,可进行场景化测试。在推理过程中输入上下文序列越长,生成的KV Cache 越多,vLLM 自带前缀缓存仅能使用显存存储,无法满足长序列问答场景下的存储需求。为此采用HBM+DDR+SSD三级KVCache缓存方案,以存代算,借助DRAM与SSD扩充缓存容量,提升缓存命中率,实现推理性能优化。

为充分验证三级缓存架构的性能优势,本次选取 DRAM 最小配置场景开展测试验证。

4.1 测试准备

1. 安装aisbench工具

# 拉取aisbench源码
git clone https://gitee.com/aisbench/benchmark.git
cd benchmark/
pip3 install -e ./ --use-pep517

# 安装核心依赖
pip3 install -r requirements/api.txt
pip3 install -r requirements/extra.txt

# 验证安装,如果打印出AISBench评测工具的所有命令行的帮助信息,说明安装成功。
ais_bench -h

2. 修改aisbenc测试文件

使用aisench时,需要配置好模型服务信息和数据集信息,进行如下操作:

# 查询配置文件路径
ais_bench --models vllm_api_stream_chat --custom-dataset-path /datasets/qa.jsonl --mode perf --debug --search

# 找到模型服务配置文件路径
│ --models    │ vllm_api_stream_chat │ /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py │

# 根据服务信息,修改配置文件
vim /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py

根据服务信息,修改以下参数:

models = [
    dict(
        attr="service",
        type=VLLMCustomAPIChatStream,
        abbr='vllm-api-stream-chat',
        path="/data/weights/GLM-5-w8a8-new",     # 模型权重路径,配置为实际权重位置
        model="qwen3",              # 指定服务端已加载模型名称,与拉起服务时的名称保持一致 
        request_rate = 0,           # 请求发送频率
        retry = 2,                  # 每个请求最大重试次数
        host_ip = "xx.xx.xx.xx",    # 指定推理服务的IP
        host_port = 8080,           # 指定推理服务的端口
        max_out_len = 1024,          # 推理服务输出的token的最大数量
        batch_size=4,               # 请求发送的最大并发数
        trust_remote_code=False,    # tokenizer是否信任远程代码,默认False;
        generation_kwargs = dict(   # 模型推理参数
            temperature = 0.01,
            ignore_eos=False,
        )
    )
]

3. 数据集构造

aisbench支持支持的任务类型包括选择 (mcq) 和问答 (qa) 两种,数据集格式支持 .jsonl 和 .csv 两种格式的数据集,这里简单构造问答qa类型的.jsonl格式的数据集(data_test.jsonl)用于测试环境能否正常运行,读者可换成自己的数据集或者开源数据集

{"question": "Q1", "answer": ""} 
{"question": "Q2", "answer": ""} 
{"question": "Q3", "answer": ""} 
{"question": "Q4", "answer": ""}

4. 配置完成后,测试查看是否能正常运行

ais_bench --models vllm_api_stream_chat --custom-dataset-path /path/to/data_test.jsonl --mode perf

4.2 长序列问答场景性能测试

针对长序列问答场景,通过请求内容多轮循环复用来模拟重复查询,首次请求全量推理写入 KV Cache,二次请求缓存命中,体现缓存命中带来的性能提升。

使用附录中的数据集生成脚本(data_gen.py),构建高前缀复用率的超长上下文问答数据集;如单条输入序列 token 长度设置为 64k,将每条序列重复两次,形成 50% 请求重复率数据集,同时将公共前缀占比参数设定为 90%,提升 KV Cache 前缀命中概率。

python data_gen.py \
--tokenizer-dir /data/weights/GLM-5-w8a8-new \        # 权重路径
--group_size 1 \        # 同一前缀下的重复样本数
--num_groups 200 \      # 数据集条数
--tokens 64000 \        # 序列长度
--num_diff_prefix 100 \ # 重复间隔
--fraction 90           # 前缀占比

生成后的数据集jsonl文件,内容格式如下示例:

{"question": "Q1", "answer": ""}
{"question": "Q2", "answer": ""}
{"question": "Q3", "answer": ""}
{"question": "Q4", "answer": ""}
...
{"question": "Q4", "answer": ""}
{"question": "Q3", "answer": ""}
{"question": "Q2", "answer": ""}
{"question": "Q1", "answer": ""}

性能测试

# 查询配置文件路径
ais_bench --models vllm_api_stream_chat --custom-dataset-path /path/to/32000_202508167_bs222_90_prefix.jsonl --mode perf --debug --search
# 找到测试配置文件路径
│ --models    │ vllm_api_general_chat │ /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py │
# 修改测试配置文件
vim /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py
# 使用数据集进行测试
ais_bench --models vllm_api_stream_chat --custom-dataset-path /path/to/32000_202508167_bs222_90_prefix.jsonl --mode perf

4.3 测试结果

性能结果的指标解析参考文档:https://ais-bench-benchmark-rf.readthedocs.io/zh-cn/latest/base_tutorials/results_intro/performance_metric.html

为充分验证三级缓存架构的性能优势,本次选取 DRAM 最小配置场景开展测试验证,以下为 A3 单机部署 GLM5.1-w8a8 模型的实测性能数据,其中external prefix-cache hit hate在开启多级缓存后,才会在vLLM服务日志中打印。总命中率为公式计算,不为vLLM打印值,计算公式可参考附录。

表1 长序列问答测试结果参考

模型prefix-cachememcache + ubsio输入长度数据集重复率prefix-cache

   hit rate
external prefix-

   cache hit hate
总命中率TTFT平均(ms)TTFT减少比值
GLM5.1-w8a8开启关闭32k72%21.4%/21.4%11716/
开启开启32k72%21.4%64.3%71.5%480459%

随着多级缓存特性的开启,KVCache缓存空间增大,缓存命中率逐步提高,核心收益体现在首Tokens时延TTFT的优化,随着命中率的增高而不断降低。

5. 附录

5.1 命中率计算公式

请求首先使用prefix-cache特性进行命中(prefix-cache hit rate,记该命中率为h1)。若prefix-cache未命中(比例为1-h1),则使用分级缓存进行命中(external prefix-cache hit rate,记该命中率为h2)。

由于2种命中率的分母(查询基数)不同,总命中率不能直接相加,应采用的计算公式为:

h_all = h1 + (1-h1)*h2 =  h1 + h2 - h1*h2

其中prefix-cache hit rate和external prefix-cache hit rate可在vLLM服务端日志中查询得到(开启多级缓存后,才会打印external prefix-cache hit rate)。

5.2 数据集生成脚本内容

数据集生成脚本data_gen.py内容:

import argparse
import json
import os
import random
import re
from typing import List, Optional
from transformers import AutoTokenizer
class DatasetTokenGuard:
    LEGAL_PATTERN = re.compile(r'^[A-Za-z0-9 ]+$')
    def __init__(self, tokenizer):
        self.tkz = tokenizer
        self.special_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
        self.vocab_total = getattr(tokenizer, "vocab_size", 0) or 0
        self.no_space_toks: List[int] = []
        self.space_toks: List[int] = []
        self.pad_token: Optional[int] = None
    def text_filter(self, raw_txt: str) -> str:
        clean_chars = (c for c in raw_txt if c.isalnum() or c == " ")
        s = "".join(clean_chars)
        for sym in ("\n", "\r", "\t"):
            s = s.replace(sym, " ")
        return s
    def check_valid(self, txt: str) -> bool:
        if not txt:
            return False
        if any(bad in txt for bad in ("\n", "\r", "\t")):
            return False
        return self.LEGAL_PATTERN.fullmatch(txt) is not None
    def ids_to_text(self, tid_list: List[int]) -> str:
        try:
            return self.tkz.decode(tid_list, clean_up_tokenization_spaces=False)
        except Exception:
            return ""
    def text_to_ids(self, txt: str) -> List[int]:
        try:
            return self.tkz.encode(txt, add_special_tokens=False)
        except Exception:
            return []
    def get_token_count(self, txt: str) -> int:
        return len(self.text_to_ids(txt))
    def build_token_pool(self):
        # 遍历词表筛选合法token
        for tid in range(self.vocab_total):
            if tid in self.special_ids:
                continue
            piece = self.ids_to_text([tid])
            if not piece or not self.check_valid(piece):
                continue
            if piece.startswith(" "):
                self.space_toks.append(tid)
            else:
                self.no_space_toks.append(tid)
        # 确定填充token
        if self.space_toks:
            self.pad_token = self.space_toks[0]
        else:
            test_strs = (" 0", " 1", " a", " A", " B", " 2")
            for ts in test_strs:
                tids = self.text_to_ids(ts)
                if tids and tids[-1] not in self.special_ids:
                    self.pad_token = tids[-1]
                    break
            if self.pad_token is None and self.no_space_toks:
                self.pad_token = self.no_space_toks[0]
        # 补充token池容量下限
        def fill_space_pool():
            supply_list = [" 0", " 1", " a", " b", " A", " B", " 2", " 3", " X", " Z"]
            for s in supply_list:
                tids = self.text_to_ids(s)
                if not tids:
                    continue
                t = tids[-1]
                if self.check_valid(self.ids_to_text([t])):
                    self.space_toks.append(t)
                    if len(self.space_toks) >= 3:
                        return
        def fill_nospace_pool():
            supply_list = ["a", "b", "A", "B", "Z", "X", "0", "1", "2"]
            for s in supply_list:
                tids = self.text_to_ids(s)
                if not tids:
                    continue
                t = tids[-1]
                dec = self.ids_to_text([t])
                if self.check_valid(dec) and not dec.startswith(" "):
                    self.no_space_toks.append(t)
                    if len(self.no_space_toks) >= 2:
                        return
        if len(self.space_toks) < 3:
            fill_space_pool()
        if len(self.no_space_toks) < 2:
            fill_nospace_pool()
        # 兜底补齐空池
        if not self.space_toks and self.pad_token is not None:
            self.space_toks.append(self.pad_token)
        if not self.no_space_toks and self.pad_token is not None:
            self.no_space_toks.append(self.pad_token)
        if self.pad_token is None:
            for tid in range(self.vocab_total):
                if tid not in self.special_ids:
                    self.pad_token = tid
                    break
        if self.pad_token is None:
            self.pad_token = 0
    def calibrate_length(self, tid_seq: List[int], target_len: int, append_tok: int) -> List[int]:
        tmp_ids = tid_seq.copy()
        max_loop = 4096
        for _ in range(max_loop):
            curr_txt = self.ids_to_text(tmp_ids)
            if not self.check_valid(curr_txt):
                curr_txt = self.text_filter(curr_txt)
            real_len = self.get_token_count(curr_txt)
            if real_len == target_len:
                return self.text_to_ids(curr_txt)
            if real_len < target_len:
                tmp_ids.append(append_tok)
            else:
                if tmp_ids:
                    tmp_ids.pop()
                else:
                    tmp_ids.append(append_tok)
        # 迭代不收敛,强制截断补齐
        if not tmp_ids:
            tmp_ids = [append_tok]
        tmp_ids = tmp_ids[:target_len]
        while len(tmp_ids) < target_len:
            tmp_ids.append(append_tok)
        final_txt = self.text_filter(self.ids_to_text(tmp_ids))
        fixed = self.text_to_ids(final_txt)
        for _ in range(256):
            if len(fixed) == target_len:
                break
            if len(fixed) < target_len:
                fixed.append(append_tok)
            else:
                fixed.pop()
        return fixed
def make_prefix_sequence(guard: DatasetTokenGuard, fix_len: int, rng: random.Random) -> List[int]:
    seq = []
    first_tok = guard.no_space_toks[0] if guard.no_space_toks else guard.pad_token
    seq.append(first_tok)
    remain = fix_len - 1
    for _ in range(remain):
        pick = rng.choice(guard.space_toks) if guard.space_toks else guard.pad_token
        seq.append(pick)
    add_t = guard.space_toks[0] if guard.space_toks else guard.pad_token
    return guard.calibrate_length(seq, fix_len, add_t)
def num_to_bit_tokens(val: int, bit_width: int, zero_tok: int, one_tok: int) -> List[int]:
    bin_str = format(val, f"0{bit_width}b")
    res = []
    for ch in bin_str:
        res.append(zero_tok if ch == "0" else one_tok)
    return res
def safe_mkdir(path: str):
    try:
        os.makedirs(path, exist_ok=True)
    except Exception:
        pass
def main_workflow():
    parser = argparse.ArgumentParser(description="测试数据集生成工具")
    parser.add_argument("--tokenizer-dir", type=str, default="/mnt/nfs/levis/DeepSeek-R1_w8a8_vllm")
    parser.add_argument("--seed", type=int, default=202508167)
    parser.add_argument("--tokens", type=int, default=3500)
    parser.add_argument("--num_diff_prefix", type=int, default=10)
    parser.add_argument("--group_size", type=int, default=5)
    parser.add_argument("--num_groups", type=int, default=40)
    parser.add_argument("--fraction", type=int, default=50)
    args = parser.parse_args()
    total_tok = args.tokens
    rng_engine = random.Random(args.seed)
    prefix_cnt = args.num_diff_prefix
    per_group_sample = args.group_size
    group_total = args.num_groups
    ratio_pct = args.fraction
    # 加载分词器
    try:
        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir, use_fast=True)
    except Exception:
        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir, use_fast=False)
    guard = DatasetTokenGuard(tokenizer)
    guard.build_token_pool()
    # 分配bit 0/1 token,保证不重复
    tok0 = guard.space_toks[0] if len(guard.space_toks)>=1 else guard.pad_token
    tok1 = guard.space_toks[1] if len(guard.space_toks)>=2 else (guard.no_space_toks[0] if guard.no_space_toks else guard.pad_token)
    if tok1 == tok0:
        combined = guard.space_toks + guard.no_space_toks
        for t in combined:
            if t != tok0:
                tok1 = t
                break
    prefix_len = max(1, total_tok * ratio_pct // 100)
    max_bit_len = 16
    bit_len = max(1, min(max_bit_len, total_tok - prefix_len))
    # 路径构造
    out_root = f"{total_tok}_{args.seed}_ds_p{prefix_cnt}_gs{per_group_sample}_ng{group_total}"
    safe_mkdir(out_root)
    path_prefix_ref = os.path.join(out_root, f"{total_tok}_{args.seed}_pre_{ratio_pct}.jsonl")
    path_full_data = os.path.join(out_root, f"{total_tok}_{args.seed}_full_{per_group_sample*group_total}_{ratio_pct}.jsonl")
    stored_prefix_list = []
    buf_prefix_lines = []
    # 批量生成前缀池
    for _ in range(prefix_cnt):
        raw_pre = make_prefix_sequence(guard, prefix_len, rng_engine)
        try:
            txt = guard.ids_to_text(raw_pre)
            if not guard.check_valid(txt):
                txt = guard.text_filter(txt)
            add_t = guard.space_toks[0] if guard.space_toks else guard.pad_token
            fixed_pre_ids = guard.calibrate_length(guard.text_to_ids(txt), prefix_len, add_t)
            pre_txt = guard.ids_to_text(fixed_pre_ids)
            buf_prefix_lines.append(json.dumps({"question": pre_txt, "answer": ""}, ensure_ascii=True))
            stored_prefix_list.append(fixed_pre_ids)
        except Exception:
            fallback_text = " ".join([f"A{x%10}" for x in range(max(1, prefix_len))])
            add_t = guard.space_toks[0] if guard.space_toks else guard.pad_token
            fall_ids = guard.calibrate_length(guard.text_to_ids(fallback_text), prefix_len, add_t)
            safe_text = guard.ids_to_text(fall_ids)
            buf_prefix_lines.append(json.dumps({"question": safe_text, "answer": ""}, ensure_ascii=True))
            stored_prefix_list.append(fall_ids)
    # 一次性写入前缀文件
    with open(path_prefix_ref, "w", encoding="utf-8") as f:
        f.write("\n".join(buf_prefix_lines)+"\n")
    # 生成完整数据集,行缓存
    data_lines = []
    for g_idx in range(group_total):
        for s_idx in range(per_group_sample):
            bit_id_seq = num_to_bit_tokens(s_idx, bit_len, tok0, tok1)
            flip_flag = (g_idx // prefix_cnt) % 2
            if flip_flag == 0:
                sel_idx = g_idx % prefix_cnt
            else:
                sel_idx = prefix_cnt - 1 - (g_idx % prefix_cnt)
            full_ids = stored_prefix_list[sel_idx].copy() + bit_id_seq
            curr_text = guard.ids_to_text(full_ids)
            used_tok = guard.get_token_count(curr_text)
            fill_num = max(0, total_tok - used_tok)
            for _ in range(fill_num):
                fill_t = rng_engine.choice(guard.space_toks) if guard.space_toks else guard.pad_token
                full_ids.append(fill_t)
            add_tok = guard.space_toks[0] if guard.space_toks else guard.pad_token
            final_ids = guard.calibrate_length(full_ids, total_tok, add_tok)
            q_text = guard.ids_to_text(final_ids)
            if not guard.check_valid(q_text):
                q_text = guard.text_filter(q_text)
            data_lines.append(json.dumps({"question": q_text, "answer": ""}, ensure_ascii=True))
    # 批量写入数据文件
    with open(path_full_data, "w", encoding="utf-8") as fw:
        fw.write("\n".join(data_lines)+"\n")
    print("数据集生成完成")
    print(f"前缀参考文件:{path_prefix_ref}")
    print(f"完整数据集:{path_full_data}")
    print(f"总token数:{total_tok}, 前缀token长:{prefix_len}, 前缀占比:{ratio_pct}%, bit编码长度:{bit_len}")
if __name__ == "__main__":
    main_workflow()


本页内容