目录

llama.cpp Radix Tree Prompt Cache and Resident KV Sharing

本仓库在 llama.cpp 的 Host Prompt Cache 基础上增加了 Radix Tree 后端和 Unified KV 中的 Resident KV Sharing(显存 KV 共享),用于在多个请求、会话和并发 Slot 之间保存并复用共享前缀的 KV Cache。本文面向部署和使用这些功能的用户,说明如何构建、启动、调用并验证 Host Radix Tree 与 Resident KV Sharing。

原版 llama.cpp 的项目介绍、模型支持和通用使用说明保留在 README-UPSTREAM.md。通用构建选项请参阅 docs/build.md,完整的 llama-server API 请参阅 tools/server/README.md。

缓存功能概述

本项目包含两个可以分别启用的缓存优化:Host 侧的 Radix Tree Prompt Cache,以及 Unified KV 中的 Resident KV Sharing(显存 KV 共享)。前者减少重复的 Host KV 保存和恢复,后者让多个请求复用同一份设备 KV Cell。

Host Radix Tree Prompt Cache(主机缓存)

传统 Host Prompt Cache 以完整 Prompt 状态为单位保存 KV Cache。多个 Prompt 即使拥有相同前缀,也会分别保存对应状态。Radix Tree 后端按照 Token 前缀组织缓存,将公共前缀只保存一次,并在分叉处保存各自的后缀。

一次请求的主要处理流程如下:

输入 Prompt
  -> Token 化并查找最长公共前缀
  -> 从 Host Radix Tree 直接恢复命中的 KV 片段
  -> 仅 Prefill 未命中的后缀
  -> Slot 空闲后保存新增 KV
  -> 按共享前缀分裂 Radix 边并更新节点

该实现具有以下特点:

  • 同一共享前缀可以被多个分支和并发请求重复命中,不会因一次读取而被消费。
  • 压缩边只保存实际出现的 Token 区间,公共 KV 数据由不同终止节点共享。
  • KV 通过结构化 Direct KV 接口提取、分段和恢复,不经过旧的完整状态 Blob 重组路径。
  • 容量达到 --cache-ram 上限时按 LRU 淘汰终止节点,并保留仍被其他分支引用的公共前缀。
  • Direct KV 恢复失败时清理已写入状态并回退到完整 Prefill,避免把部分恢复状态继续用于推理。
  • OpenAI 兼容接口不变。工具调用经过 Chat Template 转换为文本 Token 后,可以沿用相同的缓存路径。

该功能主要适用于 System Prompt、工具定义、RAG 文档或多轮历史较长,并且后续请求会从这些内容继续分叉的工作负载。

Resident KV Sharing(显存 KV 共享)

--cache-kv-share on 启用 Resident KV Sharing,让多个独立请求共享 Unified KV 中同一份公共前缀 K/V:只给已有 Cell 增加归属,不复制 K/V。来源请求结束或 Slot 被其它任务复用后,公共前缀仍由 Resident KV Sharing 的缓存条目保留。该功能默认关闭,关闭时服务行为与原有 Unified KV 一致。

支持条件(任一不满足会在启动时说明原因并拒绝启用,不会静默降级):

  • 普通全注意力、--kv-unified 的 KV cache:不含 SWA、循环或混合内存,位置为一维连续;
  • 仅因果自回归文本模型:编码器-解码器、仅编码器、扩散模型与 embedding/pooling 上下文不支持;
  • --context-shift 关闭且 --cache-reuse 0;
  • Resident KV Sharing 开启期间不接受请求级 LoRA 与 LoRA 适配器更新;任意 aLoRA 都不支持,固定普通 LoRA 允许;
  • Draft/MTP 和多模态上下文不支持。

参数与默认值:

参数 默认值 说明
--cache-kv-share on|off off 是否启用 Resident KV Sharing
--cache-kv-share-min-tokens N 32 Resident KV Sharing 与记录的最小前缀长度
--cache-kv-share-max-entries N 0 (auto) Resident KV 条目上限,auto = min(64, 可用缓存 ID 数)
--cache-kv-share-cells N 0 (auto) owner 覆盖并集上限 B,auto = C/2(C 为预分配 KV Cell 总数)
--cache-kv-share-headroom N 0 (auto) 请求余量目标 R,auto = min(n_ubatch, max(1, C/8))
--cache-kv-share-route longest|share longest Resident KV Sharing 选路规则,见下文

选路规则:

  • longest:在 A(目标 Slot 自己已验证的前缀)、S(Resident KV 索引)、H(Host Radix 可恢复前缀)中取最长可执行前缀,等长优先 A、S、H;
  • share:A/S 有效候选时取两者最长(等长取 A),没有 A/S 才使用 H;
  • 未选中 H 时不会为查询做 Host payload 分裂或搬运;完整命中时最后一个 Token 回退到 M-1 重算,产生一份末尾私有 KV,不覆盖共享 Cell。

记录、保留与回收:

  • 成功完成的完整输入 Prompt [0,M) 与成功恢复的 Host 前缀 [0,H) 各自最多自动记录一次为可共享前缀;
  • Host Radix 的保存键与导出区间截到已确认完成的 Token 快照,未计算的末尾采样 Token 不进入 Host 树;
  • 记录准入按 owner 覆盖并集 O 的真实增量与 B 判断,同键先去重;容量不足时优先回收能释放物理 Cell 的未 pin 条目;
  • Resident KV 公共前缀只读,请求的后缀、采样器与生成计数各自独立。

当前不支持“Resident KV 前缀 + Host Radix 后缀”的混合恢复(append/mixed)。

Resident KV Sharing 的收益口径是池内物理 Cell 占用与搬运减少,公共前缀只保留一份 K/V,不等同于整卡 HBM 占用下降。/metrics 中的 kv_cache_used_cells、kv_cache_owner_cells、kv_cache_owner_only_cells、kv_cache_business_logical_cells 和 kv_share_path_{a,s,h,prefill}_* 可用于区分物理占用与请求路径。

启用 Resident KV Sharing 的最小启动示例:

build/bin/llama-server -m "$MODEL" --host 127.0.0.1 --port 8080 \
  --kv-unified --no-context-shift --cache-reuse 0 \
  --cache-kv-share on --cache-kv-share-route longest \
  --ctx-size 4096 --parallel 4

分支与版本

本仓库保留了上游基线、早期开发过程、当前定版和后续开发四个分支:

分支 定位
master 保留原始上游 llama.cpp 基线,不包含本项目的 Radix Tree Prompt Cache 实现。
feature/prompt-cache-radix-tree 早期开发主分支,Radix Tree 的核心设计、主要实现过程和阶段性实验均在该分支完成,适合追溯功能演进。
radixTree 当前定版分支,包含 Host Radix Tree 与显存 KV 共享,面向交付和部署。
dev 后续开发分支,resident kv share就是在此分支开发。

本文中的构建、启动和测试命令均以 radixTree 分支为准。

获取代码

项目发布代码位于 radixTree 分支:

git clone -b radixTree https://gitee.com/frankPointer/llama.cpp-radix-tree.git
cd llama.cpp-radix-tree

如果已经克隆仓库,请先确认当前分支:

git switch radixTree
git status --short --branch

构建 llama-server

Ascend CANN

先安装 CMake、C/C++ 编译器和与设备匹配的 CANN Toolkit,并加载 CANN 运行环境。set_env.sh 的位置取决于本机安装目录:

source /path/to/ascend-toolkit/set_env.sh

使用 Release 模式构建 llama-server:

cmake -S . -B build \
  -DGGML_CANN=on \
  -DCMAKE_BUILD_TYPE=Release

cmake --build build --target llama-server -j

构建产物位于 build/bin/llama-server。启动日志中出现 CANN0 model buffer size 和模型层卸载信息,表示模型已使用 CANN 后端。

其他计算后端

Host Radix Tree 位于 llama-server 层,不限定模型计算后端;Resident KV Sharing 还要求启用受支持的 Unified KV。CPU、CUDA、Metal 等后端仍按上游 构建文档 编译。例如 CPU Release 构建:

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build --target llama-server -j

启动缓存服务

下面的示例同时启用 Host Radix Tree Prompt Cache 和 Resident KV Sharing,使用 4 个 Slot、Unified KV 和 4096 MiB Host Prompt Cache。请根据模型上下文长度、并发量和主机内存调整参数:

MODEL=/path/to/model.gguf

./build/bin/llama-server \
  --model "$MODEL" \
  --alias local \
  --host 0.0.0.0 \
  --port 8080 \
  --ctx-size 8192 \
  --parallel 4 \
  --kv-unified \
  --n-gpu-layers 99 \
  --batch-size 2048 \
  --ubatch-size 512 \
  --cache-ram 4096 \
  --cache-radix-prompt on \
  --cache-kv-share on \
  --cache-kv-share-route longest \
  --cache-idle-slots \
  --no-context-shift \
  --cache-reuse 0 \
  --log-file /tmp/llama-server-radix.log

服务启动后可以检查健康状态:

curl http://127.0.0.1:8080/health

日志中应包含以下信息:

prompt cache radix backend is enabled by --cache-radix-prompt
resident KV sharing enabled: N = ...
idle slots will be saved to prompt cache and cleared upon starting a new task

缓存组合

以下启动组合可用于部署或对照测试:

模式 启动参数 行为
No Cache --cache-ram 0 --cache-radix-prompt off 关闭 Host Prompt Cache
Legacy --cache-ram 4096 --cache-radix-prompt off 使用原有完整状态缓存
Radix --cache-ram 4096 --cache-radix-prompt on 使用 Radix Tree Prompt Cache
Radix + Resident --cache-ram 4096 --cache-radix-prompt on --kv-unified --cache-kv-share on 同时使用 Host Radix Tree 和 Resident KV Sharing

--cache-radix-prompt on 只有在 --cache-ram 非零时才会生效。修改缓存后端需要重启服务器;缓存保存在当前服务器进程的主机内存中,服务器退出后不会保留。

Radix 后端默认关闭,因此不传 --cache-radix-prompt on 时仍保持 Legacy 行为。严格的 No Cache 对照还应在请求中设置 "cache_prompt": false,仓库内的 Benchmark 已自动完成这项设置。

Resident KV Sharing 的参数和行为见缓存功能概述;“缓存组合”表格同时列出 Host Prompt Cache 的模式和两项功能的联合模式。两个功能可以单独开启,也可以同时开启。

调用接口与缓存观测

这些缓存功能不改变 llama-server 的 OpenAI 兼容 API。下面的请求显式启用请求级 Prompt Cache:

curl http://127.0.0.1:8080/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "local",
    "messages": [
      {
        "role": "system",
        "content": "You are a document assistant. Always answer from the supplied context."
      },
      {
        "role": "user",
        "content": "Summarize document A."
      }
    ],
    "cache_prompt": true,
    "temperature": 0,
    "max_tokens": 32
  }'

后续请求保持较长的 System Prompt、工具定义或文档前缀不变,只修改末尾问题,即可复用已保存前缀。cache_prompt 默认为 true,示例中显式填写是为了说明缓存行为。需要强制当前请求执行完整 Prefill 时可将其设置为 false;需要完全关闭 Host Prompt Cache 时,应在服务器启动时使用 --cache-ram 0。

响应 timings 中与缓存相关的主要字段为:

字段 含义
cache_n 本次请求直接复用的 Prompt Token 数
prompt_n 本次仍需执行 Prefill 的 Prompt Token 数
prompt_ms 服务端 Prompt Eval 时间
predicted_n 生成 Token 数
predicted_ms 服务端 Decode 时间

当相同或共享前缀请求中的 cache_n > 0 时,说明请求复用了已有 KV。该字段也可能包含活动 Slot 中的前缀复用;确认 Host Radix Restore 时,还应检查服务器日志中的 Radix loads 是否增加。日志会周期性输出 Radix 终止节点数、Host Cache 占用、unique tokens、logical tokens、保存/加载次数、淘汰次数和错误回退次数。

参数与容量建议

  • --cache-ram 的单位为 MiB,限制的是 Host Prompt Cache。容量应结合主机内存和预期工作集设置。
  • 多请求部署建议显式设置 --parallel 和 --kv-unified。默认空闲 Slot 清理依赖 Unified KV 与非零 Host Cache。
  • Resident KV Sharing 使用 Unified KV 的同一 Cell 池;--ctx-size 决定总池容量,--cache-kv-share-cells 只限制 resident owner 覆盖的物理 Cell,并不会单独分配第二个池。
  • 需要比较 Resident KV Sharing 的容量收益时,应同时记录 kv_cache_used_cells、kv_cache_owner_cells 和 kv_cache_owner_only_cells,不要用整卡 HBM 占用替代 Cell 统计。
  • --ctx-size 应覆盖并发 Slot 的实际上下文需求。长共享前缀场景需要为多个同时恢复的请求预留足够 KV 容量。
  • 对照不同缓存后端时,应固定模型、量化、构建类型、计算后端、Context、Slot、Batch、UBatch、Host Cache 容量和请求集合。
  • 工具调用使用最终 Chat Template 生成的文本 Token 进行匹配。只要公共 System Prompt 和工具定义保持稳定,就可以复用相应前缀。

功能边界

  • Host Radix Direct KV 路径当前面向纯文本 Prompt。
  • Resident KV Sharing 对 Draft/MTP、多模态和不支持的 KV 类型会在启动时拒绝启用。
  • 缓存只在单个 llama-server 进程内有效,不在不同服务器实例之间共享,也不持久化到磁盘。
  • 实际收益取决于共享前缀长度、分支数量、命中频率、Host KV 恢复成本和并发调度。无共享前缀的负载不会因 Radix Tree 自动获得加速。

正确性与验证

项目新增的测试按实现层、真实模型层和 HTTP 服务层组织:

测试 入口 验证内容
Radix Tree 单元测试 tests/test-server-prompt-cache-radix.cpp 插入、压缩边分裂、最长前缀、LRU、容量限制和失败回滚
Direct KV Roundtrip tests/test-server-prompt-cache-direct-roundtrip.cpp KV 提取、分段、恢复、逐层字节对照和 logits 对照
Resident KV Sharing 单元测试 tests/test-server-kv-resident-cache.cpp Resident KV 索引、owner ID、公共 Cell 去重、pin/LRU 和回收生命周期
生命周期测试 tools/server/tests/radix/test_radix_lifecycle.py 功能开关、路由、空闲 Slot 保存和后续恢复
CANN Smoke tools/server/tests/radix/test_qwen_cann_smoke.py CANN 设备、模型卸载、Unified KV、4 Slot 和真实请求
HTTP Token 等价性 tools/server/tests/radix/test_radix_http_token_equivalence.py No Cache 确定性、Full Prefill 与 Radix Restore 的逐 Token 对照
异常与回退 tools/server/tests/radix/test_radix_http_fallback.py Direct KV 写入后故障、完整 Prefill 回退及后续缓存可用性

详细说明位于 tools/server/tests/radix/README.md。

构建与运行测试目标

cmake --build build --target \
  llama-server \
  test-server-prompt-cache-radix \
  test-server-kv-resident-cache \
  test-server-prompt-cache-direct-roundtrip \
  -j

python3 -m pip install -r tools/server/tests/requirements.txt

测试使用本地 GGUF 模型,不会自动下载模型,也不在代码中保存机器相关路径:

export LLAMA_TEST_MODEL_FILE=/path/to/model.gguf

其中 test-server-kv-resident-cache 是不依赖模型的 Resident KV 索引单元测试;索引元数据位于 Host 侧,可以单独运行,或与其它目标一起通过 CTest 执行。

ctest --test-dir build \
  -R '^test-server-(prompt-cache-radix|kv-resident-cache) \
  --output-on-failure

运行 CPU 结构、Direct KV 和生命周期测试:

tools/server/tests/radix/run_correctness.sh cpu

在上述测试基础上增加 CANN 部署和 HTTP Token 等价性测试:

tools/server/tests/radix/run_correctness.sh cann

异常回退测试使用独立的故障注入构建。故障注入默认不进入正常 Release 二进制:

cmake -S . -B build-radix-fault \
  -DGGML_CANN=on \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLAMA_SERVER_RADIX_TEST_FAULT_INJECTION=on

cmake --build build-radix-fault --target \
  llama-server \
  test-server-prompt-cache-radix \
  test-server-prompt-cache-direct-roundtrip \
  -j

LLAMA_TEST_BUILD_DIR="$PWD/build-radix-fault" \
  tools/server/tests/radix/run_correctness.sh fault

测试的完整日志、HTTP JSON 和 JUnit XML 默认写入 /tmp/llama-radix-correctness-<UTC timestamp>/。可以通过 LLAMA_TEST_EVIDENCE_DIR 指定其他输出目录。

源码与测试索引

路径 内容
tools/server/server-prompt-cache-radix.h/.cpp Radix Tree 数据结构与 LRU/容量管理
tools/server/server-kv-resident-cache.h/.cpp Resident KV 索引、owner 生命周期与回收管理
tools/server/server-task.h/.cpp Host Prompt Cache、Direct KV 保存/恢复和统计
tests/test-server-prompt-cache-radix.cpp Radix Tree 主机侧单元测试
tests/test-server-prompt-cache-direct-roundtrip.cpp Direct KV 结构和 logits 测试
tests/test-server-kv-resident-cache.cpp Resident KV 索引和生命周期单元测试
tools/server/tests/radix/ 服务器级正确性测试
邀请码
    Gitlink(确实开源)
  • 加入我们
  • 官网邮箱:gitlink@ccf.org.cn
  • QQ群
  • QQ群
  • 公众号
  • 公众号

版权所有:中国计算机学会技术支持:开源发展技术委员会
京ICP备13000930号-9 京公网安备 11010802047560号