昇腾社区首页
中文
注册
开发者
下载
让大模型推理“起跑”更快:基于HBM+DRAM的KVCache分级缓存实践,Prefill性能飙升!

让大模型推理“起跑”更快:基于HBM+DRAM的KVCache分级缓存实践,Prefill性能飙升!

vLLM昇腾部署模型服务服务化技术文章

发表于 2026/03/11

1. 非商用声明

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

2. 方案介绍

2.1 背景

在大模型驱动的智能对话和内容生成应用中,如客服机器人或长文档摘要系统,用户往往需要处理大量文本或与模型进行多轮交互。这类场景有着大量重复内容且对响应速度有较高要求。为提升用户体验且控制部署成本,业界常用Prefix-cache技术,将历史数据缓存到HBM显存,以查代算提升性能。当前存在HBM空间有限的挑战,使用大显存推理卡的成本较高。通过HBM+DRAM的分级缓存可以大幅增加KVCache的缓存空间,提升命中率,有效缓解计算资源压力,显著降低TTFT(首token响应时延)。

本文将基于MemCache技术方案,介绍分级缓存的参考实践,帮助读者提高大模型推理服务性能。

2.2 方案简介

使用Prefix-cache特性,可以复用HBM显存中缓存的历史对话KVCache,从而在遇到相同内容时以查代算加速首token的输出。基于MemCache池化软件将KVCache存储空间从有限的HBM显存空间动态扩展到更大的系统内存DRAM,实现多卡多节点间的KV缓存共享,等效实现更大规模的Prefix-cache。本参考实践主要内容涵盖:

(1)KVCache分级缓存方案的安装部署,包括MemCache及其依赖组件。

(2)典型使用场景及提升效果,以多轮对话及长文本处理为例进行展开,基于开源数据集进行性能对比测试。

MemCache依赖MemFabric组件,相关代码仓链接为:

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

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

2.3 部署方案配置

2.3.1 软件版本

本参考实践使用的软件配套版本如下:

表1 软件版本表

软件/镜像

版本

说明

Ascend HDK

25.5.0

NPU驱动固件,要求不低于25.5.0

CANN

8.5.0

昇腾异构计算架构,要求CANN 8.1.RC1及之后版本

vLLM Ascend

0.13.0

推理框架,用于推理部署模型,建议不低于0.13.0

Docker

18.09.0

用于部署容器,推荐18.09.x~28.5.1版本;

MemFabric

1.0.6

内存池化基础软件, 用于实现DRAM与显存混合池化,提供极简的内存访问接口和高性能的内存直接访问能力,支撑多种场景下的数据共享与传输。

MemCache

1.0.0

针对AI推理场景设计的高性能分布式KVCache存储引擎。

AISBench

v3.0-20251219-master

大模型推理性能测评工具,不影响服务部署

OS

openEuler release 22.03 (LTS-SP4)

物理机操作系统(请使用昇腾官方兼容的OS)

2.3.2 典型配置

本特性支持的硬件为:300I A2(单机),800I A2(单双机)。

本参考实践使用的硬件配置如下:

表1 服务器参考配置

节点类型

服务器配置

数量

部署内容

单机

CPU:鲲鹏处理器

NPU: Atlas 300I A2 (64G) 2卡或以上

内存:512GB(基础占用+KVcache缓存,可按需配置)

1台

使用vLLm Ascend部署Qwen3-32B-w8a8推理服务

双机

服务器:Atlas 800I A2(64G)

内存:512GB(基础占用+KVcache缓存,可按需配置)

2台

使用vLLm Ascend部署DeepSeek-V3.1-w8a8推理服务

3. 部署指导

3.1 前置条件

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

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

1. 本实践使用Qwen3-32B的int8量化为例进行测试。模型权重可从Hugging Face等社区下载后按照昇腾msmodelslim工具要求,进行W8A8量化。

Hugging Face权重下载链接:https://huggingface.co/Qwen/Qwen3-32Bhttps://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B

2. HDK安装参考链接:https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/850/softwareinst/instg/instg_0005.html?Mode=PmIns&InstallType=local&OS=openEuler

3. 组网参考链接:https://support.huawei.com/enterprise/zh/doc/EDOC1100543351/ebeddbba?idPath=23710424|251366513|254884019|261408772|258915651

3.2 部署流程汇总

安装步骤:

表1 安装步骤

安装步骤

步骤说明

步骤1:配置大页内存(物理机)

按需配置大页内存空间,供后续MemCache服务使用

步骤2:创建vLLM Ascend容器

拉取vLLM Ascend镜像,创建容器

步骤3:安装MemFabric

容器内安装MemFabric

步骤4:安装MemCache

容器内安装MemCache

步骤5:拉起MetaService

配置mmc-meta.conf+mmc-local.conf文件,拉起MetaService

步骤6:拉起开启MemCache特性的vLLM Ascend服务

拉起命令中,配置相关环境变量并配置开启MemCache

3.3 详细部署操作

3.3.1 配置大页内存(物理机)

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

在物理机的os下,根据计划分配给MemCache内存空间的大小,来配置大页内存。注意该大页内存空间为公共空间,如果其它程序也需要占用内存大页,则按需增加大页内存空间,保证给MemCache的剩余大页内存空间充足。

# 查询大页内存
grep -i huge /proc/meminfo

# 回显示例
AnonHugePages:    665600 kB
ShmemHugePages:        0 kB
FileHugePages:         0 kB
HugePages_Total:       0
HugePages_Free:        0
HugePages_Rsvd:        0
HugePages_Surp:        0
Hugepagesize:       2048 kB
Hugetlb:               0 kB

# 其中,Hugepagesize表示1个大页内存大小,Hugetlb表示内存大页总空间,Hugepage_Total表示大页内存总数,HugePages_Free表示剩余内存大页数量
# 假设系统中,预计分配给MemCache的内存空间为256GB,以及留有一些buffer空间(建议不小于1GB空间,例如预留4GB空间),可按照如下方式配置大页内存
# 从上述命令查询到Hugepagesize为2048kB,则需将系统的大页数量设置为 (256GB+4GB)*1024*1024 / 2048kB ≈ 133120个
# 在系统中配置133120个大页内存(如需持久化配置,可将该命令配置到启动命令)
echo 133120 > /proc/sys/vm/nr_hugepages

# 查询配置后效果
grep -i huge /proc/meminfo

# 回显示例,配置了133120个大页内存,总空间约260GB,剩余大页数量133120个
# 如果其它程序占用了大页内存,可增加大页内存总数,保证使能MemCache前,HugePages_Free的数量充足
AnonHugePages:    622592 kB
ShmemHugePages:        0 kB
FileHugePages:         0 kB
HugePages_Total:   133120
HugePages_Free:    133120
HugePages_Rsvd:        0
HugePages_Surp:        0
Hugepagesize:       2048 kB
Hugetlb:        272629760 kB

3.3.2 创建vLLM Ascend容器

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

若需采用其它安装方式(使用Python安装),请参考官方安装指导:https://docs.vllm.ai/projects/ascend/zh-cn/v0.13.0/installation.html#set-up-using-docker

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

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

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

# 查看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 -dit \
    --name vllm-ascend-memcache \# 依据实际情况配置
    --shm-size=500g \
    --net=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 /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 /data/weights:/root/.cache \# 根据真实情况替换,执行时删除注释,不要有空格
    -it $IMAGE bash

# 进入容器
docker exec -ti vllm-ascend-memcache bash

3.3.3 Memfabric安装

单机部署和双机部署,均按如下操作步骤。

支持pip install安装和源码编译安装,选择一个方式完成安装即可。

方式1:pip install安装

pip install安装方式如下:

# 进入容器,容器名根据实际情况配置
docker exec -ti vllm-ascend-memcache bash

# 查看MemFabric当前可用版本;若没有可用版本,则使用源码编译的安装方式
pip index versions memfabric-hybrid

# 安装指定版本的MemFabric
pip install memfabric-hybrid==1.0.6

# 查询MemFabric的安装信息,Location为安装位置
pip show memfabric_hybrid
# 会有类似回显,记录实际安装位置,后续会用到
Name: memfabric_hybrid
Version: 1.0.6
Summary: python api for memfabric hybrid
Home-page: https://gitcode.com/Ascend/memfabric_hybrid
Author:
Author-email:
License: Mulan PSL v2
Location: /usr/local/python3.11.14/lib/python3.11/site-packages
Requires:
Required-by:

# 后续使用时,根据查询到的MemFabric实际安装位置,配置环境变量
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memfabric_hybrid/lib/:$LD_LIBRARY_PATH

方式2:源码编译安装

源码编译安装方式如下,注意在vLLM Ascend容器内进行操作,且关闭vLLM Ascend服务:

# 进入容器,容器名根据实际情况配置
docker exec -ti vllm-ascend-memcache bash

安装MemFabric_hybrid,官方安装指南https://gitcode.com/Ascend/memfabric_hybrid/blob/v1.0.0/doc/build.md

官方建议编译环境:

cmake: 3.20.x
gcc: 11.4+
python 3.11.10
pybind11 2.10.3
cmake: 3.20.x
make 4.3 or ninja 1.10.1

前置环境准备好后,进行操作:

# 拉取memfabric_hybrid源码
git clone --branch e1809aa7a227817a7fa3d3609cb685a0cbcef4f7  https://gitcode.com/Ascend/memfabric_hybrid

cd memfabric_hybrid 
git clean -xdf 
git reset --hard

# 拉取三方库
git submodule update --recursive --init

# 编译
bash script/build_and_pack_run.sh --build_mode RELEASE --build_python ON --xpu_type NPU --build_test OFF
# 编译完后会生成run包到output目录,格式为 memfabric-hybrid-${version}_${os}_${arch}.run,根据实际run包进行执行
# 其中,versin表示memfabric_hybrid的版本;os表示操作系统,如linux;arch表示架构,如x86或aarch64

# 安装
cd output
bash memfabric_hybrid-1.0.5_linux_aarch64.run
source /usr/local/memfabric_hybrid/set_env.sh

# 使用pip show检查是否安装成功
pip show memfabric_hybrid

# 会有类似回显
Name: memfabric_hybrid
Version: 1.0.5
Summary: python api for memfabric hybrid
Home-page: https://gitcode.com/Ascend/memfabric_hybrid
Author:
Author-email:
License: Mulan PSL v2
Location: /usr/local/python3.11.14/lib/python3.11/site-packages
Requires:
Required-by:

3.3.4 MemCache安装

单机部署和双机部署,均按如下操作步骤。

支持pip install安装和源码编译安装,选择一个方式完成安装即可。

方式1:pip install安装

pip install安装方式如下:

# 进入容器,容器名根据实际情况配置
docker exec -ti vllm-ascend-memcache bash

# 查看MemFabric当前可用版本;若没有可用版本,则使用源码编译的安装方式
pip index versions memcache-hybrid

# 安装指定版本的MemFabric
pip install memcache-hybrid==1.0.0

# 查询MemFabric的安装信息,Location为安装位置
pip show memcache-hybrid
# 会有类似回显,记录实际安装位置,后续会用到
Name: memcache_hybrid
Version: 1.0.0
Summary: python api for memcache_hybrid
Home-page: https://gitcode.com/Ascend/memcache
Author:
Author-email:
License: Mulan PSL v2
Location: /usr/local/python3.11.14/lib/python3.11/site-packages
Requires:
Required-by:

# 后续使用时,根据查询到的MemFabric实际安装位置,配置环境变量
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memcache_hybrid/lib/:$LD_LIBRARY_PATH

方式2:源码编译安装

源码编译安装方式如下,注意在vLLM Ascend容器内进行操作:

安装MemCache,官方安装指南:https://gitcode.com/Ascend/memcache/blob/v1.0.0/doc/build.md

官方推荐编译环境

cmake: 3.20.x
gcc: 11.4+
python 3.11.10
pybind11 2.10.3
cmake: 3.20.x
make 4.3 or ninja 1.10.1

编译安装方式:

# 拉取MemCache源码
git clone --branch ead5e114a30fe52ce7b0fb88f3c8fcebf82ba0e5 https://gitcode.com/Ascend/memcache
cd memcache 
git clean -xdf 
git reset --hard

# 拉取三方库
git submodule update --recursive --init 
git submodule update --remote 3rdparty/memfabric_hybrid

# 编译
bash script/build_and_pack_run.sh --build_mode RELEASE --build_python ON
# 编译完后会生成一个run包到output目录,格式为 memcache_hybrid-${version}_${os}_${arch}.run,根据实际run包进行执行

# 安装
cd output
bash memcache_hybrid-1.0.0_linux_aarch64.run
source /usr/local/memcache_hybrid/set_env.sh

# 检查是否安装成功
pip show memcache_hybrid

# 回显内容
Name: memcache_hybrid
Version: 1.0.0
Summary: python api for memcache_hybrid
Home-page: https://gitcode.com/Ascend/memcache
Author:
Author-email:
License: Mulan PSL v2
Location: /usr/local/python3.11.14/lib/python3.11/site-packages
Requires:
Required-by:

3.3.5 修改配置文件+拉起MetaService服务

MemCache特性有2个配置文件(mmc-meta.conf和mmc-local.conf)需要配置。其中mmc-meta.conf需要在MetaService拉起前完成配置,mmc-local.conf需要在拉起vLLM Ascend前完成配置。

双机场景,双机使用相同的配置文件内容,并只需要在主机拉起MetaService。

1. 进入容器:

# 进入容器,容器名根据实际情况配置
docker exec -ti vllm-ascend-memcache bash

2. 找到配置文件路径,并进行修改:

配置文件参数含义:https://gitcode.com/Ascend/memcache/blob/v1.0.0/doc/memcache_config.md

# 默认配置文件目录
cd /usr/local/memcache_hybrid/latest/config

# 使用源码编译方式,该路径下会有默认配置文件,根据实际情况,进行相应修改
mmc-local.conf  mmc-meta.conf

# 如果使用pip install安装方式,需要自行下载conf文件
https://gitcode.com/Ascend/memcache/blob/ead5e114a30fe52ce7b0fb88f3c8fcebf82ba0e5/config/mmc-meta.conf
https://gitcode.com/Ascend/memcache/blob/ead5e114a30fe52ce7b0fb88f3c8fcebf82ba0e5/config/mmc-local.conf

3. 修改mmc-meta.conf配置文件,单机使用时,可直接使用mmc-meta.conf的默认配置。

双机使用时,需要修改mmc-meta.conf文件,具体的修改内容如下:

# 修改配置文件
vim mmc-meta.conf

# 双机场景,在主节点和从节点上,以下IP地址都需配置为主节点的IP地址
# 端口号可按需修改,需要同步修改mmc-local.conf,保持一致
ock.mmc.meta_service_url = tcp://xx.xx.xx.xx:5000
ock.mmc.meta_service.config_store_url = tcp://xx.xx.xx.xx:6000
ock.mmc.meta_service.metrics_url = xx.xx.xx.xx:8000

4. 修改mmc-local.conf配置文件,按需配置NPU卡数和MemCache占用的DRAM内存大小。

单机可使用默认IP配置。双机修改IP配置,和mmc-meta.conf配置相同。

# 修改配置文件
vim mmc-local.conf

# 配置总NPU卡数(多机场景配所有设备的总NPU卡数)
ock.mmc.local_service.world_size = 8
# 表示使用CPU的内存进行KVCache缓存
ock.mmc.local_service.protocol = device_rdma
# 表示给每卡分配的内存空间,例如配置128G,则使用2卡的服务会占用共256G系统内存
# 注意预留足够的空间给系统应用,剩余内存空间可分配给MemCache使用
ock.mmc.local_service.dram.size = 128GB

# 双机场景,在主节点和从节点上,以下IP地址都需配置为主节点的IP地址
# 端口号可按需修改,需同mmc-meta.conf保持一致
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.local_service.hcom_url = tcp://xx.xx.xx.xx:7000

5. 拉起metaservice服务(双机部署时,只有主机需要拉起)。

使用pypi安装方式,拉起metaservice方式为脚本拉起方式:

# 通过环境变量,激活memcache和指定conf文件
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memfabric_hybrid/lib/:$LD_LIBRARY_PATH
export MMC_META_CONFIG_PATH=/usr/local/memcache_hybrid/latest/config/mmc-meta.conf

# 使用python脚本方式,拉起metaservice服务
python3 metaservice.py

脚本内容如下(metaservice.py):

from memcache_hybrid import MetaService
MetaService.main()

使用源码编译安装方式,拉起metaservice方式为二进制拉起:

# 源码安装方式,通过source激活,通过环境变量指定配置文件(可将配置文件方式放至任意路径然后指定)
# 使用二进制拉起metaserfice服务
source /usr/local/memcache_hybrid/set_env.sh
source /usr/local/memfabric_hybrid/set_env.sh
export MMC_META_CONFIG_PATH=/usr/local/memcache_hybrid/latest/config/mmc-meta.conf
/usr/local/memcache_hybrid/latest/aarch64-linux/bin/mmc_meta_service

3.3.6 拉起开启MemCache特性的vLLM Ascend服务

场景1:单机服务

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

source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

# 使用pip install的方式,使能方式为配置环境变量,根据实际安装情况选择一种使能方式
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memfabric_hybrid/lib/:$LD_LIBRARY_PATH
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memcache_hybrid/lib:$LD_LIBRARY_PATH

# 使用源码编译方式,使能方式为source脚本,根据实际安装情况选择一种使能方式
source /usr/local/memcache_hybrid/set_env.sh
source /usr/local/memfabric_hybrid/set_env.sh

source /usr/local/memfabric_hybrid/set_env.sh
export MMC_LOCAL_CONFIG_PATH=/usr/local/memcache_hybrid/latest/config/mmc-local.conf
export VLLM_USE_V1=1
export ASCEND_RT_VISIBLE_DEVICES=2,3 # 使用的NPU卡号
export ACL_OP_INIT_MODE=1
export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True"
export PYTHONHASHSEED=0
export HCCL_BUFFSIZE=1024

python -m vllm.entrypoints.openai.api_server \
  --model=/root/.cache/DeepSeek-R1-Distill-Qwen-32B-w8a8 \# 修改实际的权重路径
  --served-model-name qwen3 \
  --trust-remote-code \
  --async-scheduling \
  --quantization ascend \
  --distributed-executor-backend mp \
  --tensor-parallel-size 2 \
  --port 8113 \
  --max-model-len 32K \
  --compilation_config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' \
  --additional_config='{"ascend_scheduler_config":{"enabled":false}, "enable_shared_expert_dp":false, "chunked_prefill_for_mla":true}' \
  --gpu-memory-utilization 0.90 \
  --kv-transfer-config \# 使能MemCache特性
  '{
      "kv_connector": "AscendStoreConnector",
      "kv_role": "kv_both",
      "kv_connector_extra_config": {
        "backend": "memcache",
        "mooncake_rpc_port":"0"
      }
  }' > log_vllm_memcache.log 2>&1

开启MemCache特性的vLLM Ascend服务,在服务的运行日志中会新增External prefix cache hit rate关键字,用于指示MemCache命中率(优先显存命中,显存未命中而内存命中时,统计此命中率)。

单机多个服务开启MemCache场景,保持mmc-meta.conf和mmc-local.conf文件内容相同,可以只拉起1个MetaService,然后拉起多个推理服务。

场景2:双机服务

双机DeepSeek-V3服务开启MemCache场景,只需要在主机拉起MetaService,然后在原有双机拉起命令上,增加开启MemCache特性的参数。

双机拉起命令可参考官方文档:https://docs.vllm.ai/projects/ascend/zh-cn/v0.13.0/tutorials/DeepSeek-V3.1.html

权重可参考前置准备,使用modelslim进行转换。

双机拉起DeepSeek的参考命令如下,假设主节点的IP为xx.xx.xx.0,对应网卡名为xxx0;从节点IP为xx.xx.xx.1,对应网卡名为xxx1。

主节点拉起命令(需保证按照上一章节拉起MetaService和配置大页内存):

source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

# 使用pip install的方式,使能方式为配置环境变量,根据实际安装情况选择一种使能方式
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memfabric_hybrid/lib/:$LD_LIBRARY_PATH
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memcache_hybrid/lib:$LD_LIBRARY_PATH

# 使用源码编译方式,使能方式为source脚本,根据实际安装情况选择一种使能方式
source /usr/local/memcache_hybrid/set_env.sh
source /usr/local/memfabric_hybrid/set_env.sh

export MMC_LOCAL_CONFIG_PATH=/usr/local/memcache_hybrid/latest/config/mmc-local.conf

# 可通过ifconfig获得以下参数
# nic_name为local_ip对应的网卡名,例如enp***,local_ip为当前节点的通信IP
nic_name="xxx0"
local_ip="xx.xx.xx.0"

# node0_ip配置为主节点IP,
node0_ip="xx.xx.xx.0"

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 OMP_PROC_BIND=false
export OMP_NUM_THREADS=10
export PYTHONHASHSEED=0
export HCCL_BUFFSIZE=1024
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export VLLM_USE_V1=1

vllm serve /root/.cache/DeepSeek-V3.1-w8a8-mtp-QuaRot \
  --host 0.0.0.0 \
  --port 8004 \
  --enforce-eager \
  --data-parallel-size 2 \
  --data-parallel-size-local 1 \
  --api-server-count 2 \
  --data-parallel-address $node0_ip \
  --data-parallel-rpc-port 13389  \
  --tensor-parallel-size 8 \
  --seed 1024 \
  --served-model-name deepseek \
  --max-model-len 65536 \
  --max-num-batched-tokens 16384 \
  --trust-remote-code \
  --gpu-memory-utilization 0.9 \
  --quantization ascend \
  --max-num_seqs 20 \
  --enable-expert-parallel \
  --no-enable-prefix-caching \
  --additional_config='{"ascend_scheduler_config":{"enabled":false}, "enable_shared_expert_dp":false, "chunked_prefill_for_mla":true}' \
  --kv-transfer-config \
  '{
        "kv_connector": "AscendStoreConnector",
        "kv_role": "kv_both",
        "kv_connector_extra_config": {
                "backend": "memcache",
                "mooncake_rpc_port":"0"
           }
  }' > log_hunbu_0.log 2>&1

从节点拉起命令:

source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

# 使用pip install的方式,使能方式为配置环境变量,根据实际安装情况选择一种使能方式
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memfabric_hybrid/lib/:$LD_LIBRARY_PATH
export LD_LIBRARY_PATH=/usr/local/python3.11.14/lib/python3.11/site-packages/memcache_hybrid/lib:$LD_LIBRARY_PATH

# 使用源码编译方式,使能方式为source脚本,根据实际安装情况选择一种使能方式
source /usr/local/memcache_hybrid/set_env.sh
source /usr/local/memfabric_hybrid/set_env.sh

export MMC_LOCAL_CONFIG_PATH=/usr/local/memcache_hybrid/latest/config/mmc-local.conf

# 可通过ifconfig获得以下参数
# nic_name为local_ip对应的网卡名,例如enp***,local_ip为当前节点的通信IP
nic_name="xxx1"
local_ip="xx.xx.xx.1"

# node0_ip配置为主节点IP,
node0_ip="xx.xx.xx.0"

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 OMP_PROC_BIND=false
export OMP_NUM_THREADS=10
export PYTHONHASHSEED=0
export HCCL_BUFFSIZE=1024
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export VLLM_USE_V1=1

vllm serve /weights/DeepSeek-V3.1-w8a8-mtp-QuaRot \
  --host 0.0.0.0 \
  --port 8004 \
  --headless  \
  --enforce-eager \
  --data-parallel-size 2 \
  --data-parallel-size-local 1 \
  --data-parallel-start-rank 1 \
  --data-parallel-address $node0_ip \
  --data-parallel-rpc-port 13389  \
  --tensor-parallel-size 8 \
  --seed 1024 \
  --served-model-name deepseek \
  --max-model-len 65536 \
  --max-num-batched-tokens 16384 \
  --trust-remote-code \
  --gpu-memory-utilization 0.9 \
  --quantization ascend \
  --max-num_seqs 20 \
  --enable-expert-parallel \
  --no-enable-prefix-caching \
  --additional_config='{"ascend_scheduler_config":{"enabled":false}, "enable_shared_expert_dp":false, "chunked_prefill_for_mla":true}' \
  --kv-transfer-config \
   '{
        "kv_connector": "AscendStoreConnector",
        "kv_role": "kv_both",
        "kv_connector_extra_config": {
                "backend": "memcache",
                "mooncake_rpc_port":"0"
           }
  }' > log_hunbu_1.log 2>&1

4. 典型场景参考实践

4.1 典型场景

按照部署章节指导,拉起服务后,可进行场景化测试。本实践采用2卡300I A2 64G,配置256G内存空间,vLLM Ascend服务进行测试。

参考vLLM官方Prefix-caching文档(https://docs.vllm.ai/en/v0.13.0/features/automatic_prefix_caching/),有2个主要适用场景:长文档问答和多轮对话。

在这2个场景下,Prefix-caching可以实现较大的性能提升。vLLM自带的Prefix-caching只能使用显存对KVCache进行存储,在此基础上,使用分级缓存方案将内存也用于存储KVCache,提升KVCache空间,可以获得更高的缓存命中率,从而提升推理的性能。

这2个场景在实际生产中有着许多应用,例如长文档(财报分析的投研报告生成,合同与合规审查等),用户多次处理同一文档时,相同内容无需重复计算;在智能助手(如医疗的预问诊等)的多轮问答中,历史对话前缀可被复用,避免重复计算,从而大幅降低首token时延(TTFT),提升用户体验,让系统更高效响应用户。

4.2 测试准备

测试准备,需要准备好数据集文件和测试工具。本文将基于开源的LongBench数据集(该数据集平均长度为几千字,贴合场景),构造长文档问答和多轮对话的场景化测试,有助于读者快速评估KVCache分级缓存带来的收益,从而评估场景的匹配性和适用性。

步骤1:数据集下载

LongBench数据集下载链接(https://huggingface.co/datasets/zai-org/LongBench),下载后解压出jsonl文件待使用,读者也可基于自己的jsonl格式数据集,使用后文的预处理脚本进行格式转换。

步骤2:准备测试工具

本文采用开源的AISBench工具进行测试,可使用docker容器的方式部署。

AISBench代码仓:https://github.com/AISBench/benchmark

AISBench镜像使用指导和下载链接:https://github.com/AISBench/benchmark/wiki/AISBench-Docker-Images-Guidance

使用单独的AISBench镜像,先从下载链接中,拉取镜像压缩文件到服务器。

# 导入镜像
docker load -i aisbench_benchmark_v3.0-20251219-master_aarch64_py_310.tar.gz

# 创建容器,给网络权限,注意根据实际情况挂载权重和数据集目录(-v为挂载目录)
docker run --name ais_bench_container -itd --net=host \
-w /benchmark \
--ipc=host \
-v /model_path:/model_path \
-v /datasets:/datasets \
--entrypoint /bin/bash \
ghcr.io/aisbench/aisbench_benchmark:v3.0-20251219-master_aarch64_py_310

# 进入容器,容器内为安装好的AISBench环境
docker exec -ti ais_bench_container bash

步骤3:进行测试配置

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

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

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

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

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

models = [
    dict(
        attr="service",
        type=VLLMCustomAPIChat,
        abbr="vllm-api-general-chat",
        path="/data/weights/DeepSeek-R1-Distill-Qwen-32B-w8a8", # 模型权重路径,测试工具需要读取词表等文件
        model="qwen3", # 模型名,与拉起服务时的名称保持一致
        stream=True, # 是否开启流式接口
        request_rate=0, # 发送速率,配0则为并发测试
        retry=2,
        api_key="",
        host_ip="xx.xx.xx.xx", # 根据实际IP填写
        host_port=8113, # 根据实际配置的端口填写
        url="",
        max_out_len=1024, # 最大输出长度配置
        batch_size=32, # 请求并发
        trust_remote_code=False,
        generation_kwargs=dict(
            temperature=0.01,
            ignore_eos=True, # 为True则输出会打满最大输出长度,建议在性能测试时开启,防止输出长度不同带来的性能波动
        ),
        pred_postprocessor=dict(type=extract_non_reasoning_content),
    )
]

步骤4:检查测试是否成功

简易创建一个QA对数据集(qa.jsonl),用于测试环境能否正常运行,读者可换成自己的数据集或者开源数据集

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

完成配置后,进行测试,查看能否正常运行:

# 修改配置,进行请求测试
ais_bench --models vllm_api_general_chat --custom-dataset-path /datasets/qa.jsonl --mode perf --debug --num-warmups 0

4.3 长文档问答参考测试

测试思路为,将不同的问题进行多轮循环,模拟重复出现的请求内容,体现缓存命中带来的性能提升。

基于原始数据集,使用预处理脚本(完整脚本见附录),将问题重复2次(问题第一次请求时正常推理,第二次请求时实现命中),构造50%重复率的测试集。

当KVCache缓存空间足够大时,总缓存命中率(Prefix-cache+MemCache)会接近50%。因为存储的最小单位为block,所以问题的尾部会丢失一些,真实命中率会比理论值稍低一些。

具体操作为,将原始数据集(例如multi_news.jsonl)转换为AISBench支持的QA对格式数据集,参考命令如下,预处理脚本内容参考附录,执行时删除注释。

1. 处理前:原始数据格式

假设,预处理前的原始数据集jsonl文件的内容如下所示,其中Q1~Q4代表不同的原始问题:

{"context": "Q1", "xx": "xx"} 
{"context": "Q2", "xx": "xx"} 
{"context": "Q3", "xx": "xx"} 
{"context": "Q4", "xx": "xx"}

2. 进行数据处理

注意,预处理前的原始数据集为jsonl后缀,长文本处理后的输出文件需指定jsonl后缀。

# 举例
python dataset_convert.py \
--input multi_news.jsonl \ # 输入文件,支持多个文件
--output multi_news_convert.jsonl \ # 输出文件
--top-n 128 \ # 选取个数
--min-length 0 \ # 最小长度过滤,小于此长度则不选取
--max-length 8192 \ # 最大长度过滤,大于此长度则不选取
--repeat 2 \ # 重复2轮
--source-fields context \ # 原jsonl文件中,输入内容字段,支持多个字段拼接
--count-mode valid-only \ # 长度过滤掉的用例,不进入计数
--alternate-order True \ # 每轮重复采用逆序排序
--output-format qa \ # 输出格式为QA对
--rounds 4 # 多轮对话轮数,QA对格式下不生效

3. 处理后:QA对数据格式

则处理后的数据集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": ""}

4. 进行测试

参考测试命令:

# 查询配置文件路径
ais_bench --models vllm_api_general_chat --custom-dataset-path /datasets/multi_news_convert.json --mode perf --debug --num-warmups 0 --search

# 找到测试配置文件路径
│ --models    │ vllm_api_general_chat │ /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_general_chat.py │

# 修改测试配置文件
vim /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_general_chat.py

# 使用multi_news_convert.json数据集进行测试
ais_bench --models vllm_api_general_chat --custom-dataset-path /datasets/multi_news_convert.json --mode perf --debug --num-warmups 0
# 因为预热时会缓存住部分数据集,所以推荐配置 --num-warmups 0关闭预热,避免影响缓存命中;
# 如果担心未预热会影响性能,可使用其它数据集预热后进行测试
# 测试完成后,如果进行下一轮测试,建议重启服务,将已有缓存清理掉,避免影响

4.4 多轮对话参考测试

测试思路为,将多个不同的问题,组成多轮对话,模拟真实场景的多轮对话请求,从第二轮开始,可以命中之前轮次出现过的问题和回答,体现缓存命中带来的性能提升。

将原始数据集(例如multi_news.jsonl)转换为AISBench已支持的ShareGPT格式数据集(参考手册:https://github.com/AISBench/benchmark/blob/master/ais_bench/benchmark/configs/datasets/sharegpt/README.md),参考命令如下,执行时删除注释。

1. 处理前:原始数据格式

假设,预处理前的原始数据集jsonl文件的内容如下所示,其中Q1~Q4代表不同的原始问题:

{"context": "Q1", "xx": "xx"} 
{"context": "Q2", "xx": "xx"} 
{"context": "Q3", "xx": "xx"} 
{"context": "Q4", "xx": "xx"}

2. 数据处理

注意,预处理前的原始数据集为jsonl后缀,多轮对话处理后的输出文件需指定json后缀。

# 举例
python dataset_convert.py \
--input multi_news.jsonl \ # 输入文件,支持多个文件
--output multi_news_convert.json \ # 输出文件
--top-n 32\ # 选取个数
--min-length 1024 \ # 最小长度过滤,小于此长度则不选取
--max-length 4096\ # 最大长度过滤,大于此长度则不选取
--repeat 1 \ # 只运行1轮
--source-fields context \ # 原jsonl文件中,输入内容字段,支持多个字段拼接
--count-mode valid-only \ # 长度过滤掉的用例,不进入计数
--alternate-order True \ # 每轮重复采用逆序排序,只运行1轮时不生效
--output-format multi-turn\ # 输出格式为多轮对话
--rounds 4 # 多轮对话轮数

3. 处理后:对话数据格式

则处理后的数据集json文件,内容格式如下示例(为简化示例,按照2轮对话举例):

[
  {
    "id": "id1",
    "conversations": [
      {
        "from": "human",
        "value": "Q1"
      },
      {
        "from": "gpt",
        "value": " "
      },
      {
        "from": "human",
        "value": "Q2"
      },
      {
        "from": "gpt",
        "value": " "
      }
    ]
  },
  {
    "id": "id3",
    "conversations": [
      {
        "from": "human",
        "value": "Q3"
      },
      {
        "from": "gpt",
        "value": " "
      },
      {
        "from": "human",
        "value": "Q4"
      },
      {
        "from": "gpt",
        "value": " "
      }
    ]
  }
]

4. 进行测试

参考测试命令:

# 查询配置文件路径
ais_bench --models vllm_api_general_chat --datasets sharegpt_gen --mode perf --debug --num-warmups 0 --search

# 找到测试配置文件路径
│ --models    │ vllm_api_general_chat │ /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_general_chat.py │

# 修改配置文件
vim /benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_general_chat.py

# 找到数据集配置文件路径,进行修改
│ --datasets  │ sharegpt_gen          │ /benchmark/ais_bench/benchmark/configs/datasets/sharegpt/sharegpt_gen.py        │

# 将path修改为真实测试机路径,或者将转换好的文件按照放到指定位置;
sharegpt_datasets = [
dict(
abbr='sharegpt',
type=ShareGPTDataset,
disable_shuffle=True,
path='ais_bench/datasets/sharegpt/ShareGPT_V3_unfiltered_cleaned_split.json', # 数据集路径,使用相对路径时相对于源码根路径,支持绝对路径
reader_cfg=sharegpt_reader_cfg,
infer_cfg=sharegpt_infer_cfg,
eval_cfg=sharegpt_eval_cfg
)

# 进行多轮对话测试
ais_bench --models vllm_api_general_chat --datasets sharegpt_gen --mode perf --debug --num-warmups 0
# 因为预热时会缓存住部分数据集,所以推荐配置 --num-warmups 0关闭预热,避免影响缓存命中;
# 如果担心未预热印象性能,可使用其它数据集预热后进行测试
# 测试完成后,如果进行下一轮测试,建议重启服务,将已有缓存清理掉,避免影响

4.5 性能结果和日志解析

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

给出一组长文档测试结果参考:

表1 长文本测试结果参考


Prefix-Cache

MemCache

重复率

总命中率

TTFT平均(ms)

测试1

关闭

关闭

50%

0%

3057.8

测试2

开启

关闭

50%

17%

2807

测试3

开启

开启

50%

48.8%

2070.8

随着不同特性的开启,KVCache缓存空间增大,Cache的命中率逐步提高,主要收益在TTFT(Time To First Token)的优化,从不命中提升到约50%的命中率下,TTFT可以降低30%以上,并且还会随着命中率的增高而继续降低。

在vllm的打屏日志中,可以实时观察到当前的缓存命中率,关于vllm自带的命中率的统计方式可参考官方链接,默认统计最近1000个请求中的命中率:https://docs.vllm.ai/en/v0.13.0/design/metrics/#prefix-cache-metrics

如果希望提前预估对应数据集可能占用的Memcache内存空间大小,可先使用vllm启动一次服务,然后找到关于空间占用的日志打屏,然后进行近似换算,例如:

# 不同版本下,可能有不同格式的打屏,假设某4卡运行有如下打屏
Available KV cache memory: 9.21 GiB
GPU KV cache size: 150,800 tokens
# 换算方式为:4*9.21/150≈0.25GB/k-token,也就是每1K个token,大概配置0.25GB空间

# 另外某2卡服务有如下日志,也可进行换算
Available memory: 36408063180,
GPU KV cache size: 277,760 tokens
# 换算方式:(36*1000/1024)*2/277≈0.25GB/k-token,每1K个token,大概需要配置0.25GB空间
# 假设目标数据集有100条,预期输入4K token,输出1K token,共计5K token(注意token和文本长度不为1:1换算),全部缓存下则需要500K token,建议配置的空间不小于:500K*0.25GB/k-token=125GB

注意,需要修改mmc-meta.conf中的驱逐阈值配置如下,才可利用全部的内存空间;如果不使用全部内存空间,则需要配置更大空间,保证阈值下的空间满足上升计算值:

# 开启驱逐的空间使用阈值
ock.mmc.evict_threshold_high = 100
# 驱逐后,剩余空间的阈值
ock.mmc.evict_threshold_low = 90

在开启MemCache后,如果想判断运行完某数据集后,缓存空间是否足够,可通过在MemCache的日志中,查询关键字Evicted,找到是否有驱逐行为。

# 找到MemCache日志
/usr/local/memcache_hybrid/latest/aarch64-linux/logs/mmc-meta.log

# 查询关键词Evicted,如果有日志,则代表发生驱逐
grep 'Evicted' mmc-meta.log | tail -n 50

5. 附录

5.1 数据集预处理脚本内容

参数含义解析如下:

表1 参数含义表

参数

类型

默认值

说明

--input

字符串(可多个)

["multi_news.jsonl"]

输入的 .jsonl 文件路径列表。支持多个文件,按顺序读取合并。

--output

字符串

"multi_news_convert.json"

输出文件路径。

--top-n

整数

128

选取的有效样本数量,受--count-mode影响。若总有效样本不足,则实际输出数量少于设定值。

--min-length

整数

1024

样本字段 "length" 的最小允许值。低于此值的样本将被过滤,注意是字符数而非token数。

--max-length

整数

8192

样本字段 "length" 的最大允许值。高于此值的样本将被过滤,注意是字符数而非token数。

--repeat

整数

1

将选中的样本序列重复多少次。例如 repeat=2 表示 [A,B] → [A,B,A,B](若启用交替则为 [A,B,B,A])。

--source-fields

字符串(可多个)

["context"]

指定输入 JSON 中哪些字段用于拼接生成 "question"。按命令行顺序拼接。

--count-mode

枚举

"valid-only"

控制 --top-n 的计数逻辑:

"valid-only":仅满足长度条件的样本计入 top-n(推荐)

"all-lines":所有读取行(包括无效行)都计入。

--alternate-order

字符串

"True"

是否在重复时交替顺序:

"True":奇数轮正序,偶数轮逆序(如 [A,B] → [A,B,B,A])

"False":始终正序(如 [A,B] → [A,B,A,B])

--output-format

枚举

"multi-turn"

输出格式:

"qa":每行一个 {"question": "...", "answer": ""}(JSONL)

"multi-turn":输出多轮对话 JSON 数组,按照sharegpt数据集格式(见multi-turn 示例)

--rounds

整数

4

仅在 multi-turn 模式下生效。每段对话包含多少轮(每轮 = 1 个 human + 1 个 gpt)。总样本数需为 rounds 的整数倍,否则末尾不足部分会被丢弃。

执行脚本时指定参数:

# 举例
python dataset_convert.py \
--input multi_news.jsonl \
--output multi_news_convert.jsonl \
--top-n 128 \
--min-length 1024 \
--max-length 8192 \
--repeat 1 \
--source-fields context \
--count-mode valid-only \
--alternate-order True \
--output-format multi-turn \
--rounds 4

脚本dataset_convert.py内容:

import json
import argparse
from pathlib import Path
from typing import List
def convert_jsonl(
    input_paths: List[str],
    output_path: str,
    select_top_n: int = 200,
    min_length: int = 1024,
    max_length: int = 32768,
    repeat_times: int = 2,
    alternate_order: bool = True,
    source_fields: List[str] = None,
    count_mode: str = "valid-only",
    output_format: str = "qa",
    rounds: int = 2
):
    if source_fields is None:
        source_fields = ["context"]
    
    if output_format not in ("qa", "multi-turn"):
        raise ValueError("--output-format must be 'qa' or 'multi-turn'")
    
    if output_format == "multi-turn" and rounds < 1:
        raise ValueError("--rounds must be >= 1 in multi-turn mode")
    valid_records = []
    total_read_lines = 0
    done = False  # Unified termination flag
    for input_path in input_paths:
        if done:
            break
        input_file = Path(input_path)
        if not input_file.exists():
            raise FileNotFoundError(f"Input file {input_file} does not exist.")
        
        with open(input_file, 'r', encoding='utf-8') as f_in:
            for line in f_in:
                line = line.strip()
                if not line:
                    continue
                # In "all-lines" mode, count every non-empty line
                if count_mode == "all-lines":
                    total_read_lines += 1
                    if total_read_lines > select_top_n:
                        done = True
                        break  # Exit inner loop
                # Parse JSON
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue  # Skip malformed lines
                # Check length constraint
                length = record.get("length", 0)
                if length < min_length or length > max_length:
                    continue
                # Build question from source fields
                question_parts = []
                for field in source_fields:
                    value = record.get(field, "")
                    if not isinstance(value, str):
                        value = str(value)
                    question_parts.append(value)
                question = "".join(question_parts)
                valid_records.append({"question": question, "answer": ""})
                # In "valid-only" mode, stop once we have enough valid records
                if count_mode == "valid-only":
                    if len(valid_records) >= select_top_n:
                        done = True
                        break  # Exit inner loop
            # If done inside inner loop, exit outer loop too
            if done:
                break
    # Truncate to exact top-N only in "valid-only" mode
    if count_mode == "valid-only":
        valid_records = valid_records[:select_top_n]
    # Apply repetition with optional order alternation
    final_records = []
    for i in range(repeat_times):
        if alternate_order and i % 2 == 1:
            chunk = reversed(valid_records)
        else:
            chunk = valid_records
        final_records.extend(chunk)
    # Write output
    with open(output_path, 'w', encoding='utf-8') as f_out:
        if output_format == "qa":
            for rec in final_records:
                f_out.write(json.dumps({"question": rec["question"], "answer": rec["answer"]}, ensure_ascii=False) + '\n')
        elif output_format == "multi-turn":
            conversations = []
            idx = 0
            while idx + rounds <= len(final_records):
                group = final_records[idx:idx + rounds]
                conv = []
                for item in group:
                    conv.append({"from": "human", "value": item["question"]})
                    conv.append({"from": "gpt", "value": " "})
                conversations.append({
                    "id": f"id{len(conversations) + 1}",
                    "conversations": conv
                })
                idx += rounds
            json.dump(conversations, f_out, ensure_ascii=False, indent=2)
    # Print summary
    print(f"✅ Processed {len(valid_records)} valid records from {len(input_paths)} file(s).")
    print(f"📏 Length filter: {min_length} ≤ length ≤ {max_length}")
    print(f"🔗 Source fields: {source_fields}")
    print(f"🔢 Count mode: {count_mode}")
    if output_format == "qa":
        print(f"📤 Output format: QA pairs ({len(final_records)} lines)")
    else:
        print(f"📤 Output format: Multi-turn (rounds={rounds}, {len(conversations)} conversations)")
def main():
    parser = argparse.ArgumentParser(
        description="Convert .jsonl files to QA or multi-turn format with filtering and merging."
    )
    parser.add_argument("--input", type=str, nargs='+', default=["multi_news.jsonl"], help="Input .jsonl file(s)")
    parser.add_argument("--output", type=str, default="test_py_multi_news.jsonl", help="Output file path")
    parser.add_argument("--top-n", type=int, default=64, help="Number of records to select")
    parser.add_argument("--min-length", type=int, default=1024, help="Minimum allowed 'length' ")
    parser.add_argument("--max-length", type=int, default=4096, help="Maximum allowed 'length' ")
    parser.add_argument("--repeat", type=int, default=2, help="Repeat times ")
    parser.add_argument("--source-fields", type=str, nargs='+', default=["context"], help="Fields to concatenate into 'question' ")
    parser.add_argument("--count-mode", choices=["valid-only", "all-lines"], default="valid-only", help="Counting mode ")
    parser.add_argument("--alternate-order", type=str, choices=["True", "False"], default="True", help="Alternate order on repeat ")
    parser.add_argument("--output-format", choices=["qa", "multi-turn"], default="multi-turn", help="Output format ")
    parser.add_argument("--rounds", type=int, default=4, help="Rounds per conversation in multi-turn mode ")
    args = parser.parse_args()
    alternate_order = (args.alternate_order == "True")
    convert_jsonl(
        input_paths=args.input,
        output_path=args.output,
        select_top_n=args.top_n,
        min_length=args.min_length,
        max_length=args.max_length,
        repeat_times=args.repeat,
        alternate_order=alternate_order,
        source_fields=args.source_fields,
        count_mode=args.count_mode,
        output_format=args.output_format,
        rounds=args.rounds
    )
if __name__ == "__main__":
    main()


本页内容